diff --git a/.github/codecov.yml b/.github/codecov.yml index f1ec60595..c5c2373a8 100644 --- a/.github/codecov.yml +++ b/.github/codecov.yml @@ -37,7 +37,11 @@ flags: carryforward: true integration: paths: - - integration + - pgdog + - pgdog-plugin + - pgdog-macros + - pgdog-plugin-build + - plugins carryforward: false ignore: @@ -45,4 +49,7 @@ ignore: - "**/examples/**" - integration/python/venv +fixes: + - "/home/runner/_work/pgdog/pgdog/::" + slack_app: true diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 382ba0745..5b41d46a2 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -1,6 +1,10 @@ name: ci on: push: + branches: + - main + pull_request: + types: [opened, synchronize, reopened] jobs: fmt: @@ -76,7 +80,7 @@ jobs: RUSTFLAGS: "-C link-dead-code" run: | cargo llvm-cov clean --workspace - cargo llvm-cov nextest --lcov --output-path lcov.info --no-fail-fast --test-threads=1 --filter-expr "package(pgdog)" + cargo llvm-cov nextest --lcov --output-path lcov.info --no-fail-fast --test-threads=1 --filter-expr "package(pgdog) | package(pgdog-config) | package(pgdog-vector)" - name: Run documentation tests run: cargo test --doc # Requires CODECOV_TOKEN secret for upload @@ -89,6 +93,7 @@ jobs: fail_ci_if_error: true integration: runs-on: blacksmith-4vcpu-ubuntu-2404 + timeout-minutes: 30 env: LLVM_PROFILE_FILE: ${{ github.workspace }}/target/llvm-cov-target/profiles/pgdog-%p-%m.profraw steps: @@ -108,6 +113,7 @@ jobs: sudo -u postgres createdb $USER sudo -u postgres psql -c 'ALTER SYSTEM SET max_connections TO 1000;' sudo -u postgres psql -c 'ALTER SYSTEM SET max_prepared_transactions TO 1000;' + sudo -u postgres psql -c 'ALTER SYSTEM SET wal_level TO logical;' sudo service postgresql restart bash integration/setup.sh sudo apt update && sudo apt install -y python3-virtualenv mold @@ -140,6 +146,8 @@ jobs: echo "PGDOG_BIN=$(realpath "$BIN_PATH")" >> "$GITHUB_ENV" - name: pgbench run: bash integration/pgbench/run.sh + - name: schema-sync + run: bash integration/schema_sync/run.sh - name: Go run: bash integration/go/run.sh - name: JavaScript @@ -156,6 +164,8 @@ jobs: run: bash integration/rust/run.sh - name: Stop shared PgDog run: bash -lc 'source integration/common.sh; stop_pgdog' + - name: Data sync + run: bash integration/copy_data/dev.sh - name: Python run: bash integration/python/run.sh - name: Load balancer @@ -164,8 +174,6 @@ jobs: run: bash integration/complex/run.sh - name: Dry run run: bash integration/dry_run/run.sh - # - name: Plugins - # run: bash integration/plugins/run.sh - name: Ensure PgDog stopped run: | if pgrep -x pgdog > /dev/null; then @@ -182,3 +190,61 @@ jobs: files: integration.lcov flags: integration fail_ci_if_error: true + plugin-unit-tests: + runs-on: blacksmith-4vcpu-ubuntu-2404 + continue-on-error: true + timeout-minutes: 30 + steps: + - uses: actions/checkout@v4 + - uses: actions-rs/toolchain@v1 + with: + toolchain: stable + override: true + - uses: useblacksmith/rust-cache@v3 + with: + prefix-key: "plugin-unit-v1" + - name: Install CMake 3.31 + run: | + sudo apt-get install -y mold + sudo apt remove cmake + sudo pip3 install cmake==3.31.6 + cmake --version + - name: Install test dependencies + run: cargo install cargo-nextest --version "0.9.78" --locked + - name: Run plugin unit tests + run: cargo nextest run -E 'package(pgdog-example-plugin)' --no-fail-fast + plugin-integration-tests: + runs-on: blacksmith-4vcpu-ubuntu-2404 + continue-on-error: true + timeout-minutes: 30 + steps: + - uses: actions/checkout@v4 + - uses: actions-rs/toolchain@v1 + with: + toolchain: stable + override: true + - uses: useblacksmith/rust-cache@v3 + with: + prefix-key: "plugin-integration-v1" + - name: Setup PostgreSQL + run: | + sudo service postgresql start + sudo -u postgres createuser --superuser --login $USER + sudo -u postgres createdb $USER + bash integration/setup.sh + - name: Install dependencies + run: | + sudo apt update && sudo apt install -y python3-virtualenv mold + sudo gem install bundler + sudo apt remove -y cmake + sudo pip3 install cmake==3.31.6 + cmake --version + cargo install cargo-nextest --version "0.9.78" --locked + - name: Build plugin + run: | + cargo build --release + pushd plugins/pgdog-example-plugin + cargo build --release + popd + - name: Run plugin integration tests + run: bash integration/plugins/run.sh diff --git a/AGENTS.md b/AGENTS.md index 801a1693e..619abcee1 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -8,6 +8,7 @@ Run `cargo check` for a quick type pass, and `cargo build` to compile local bina ## Coding Style & Naming Conventions Follow Rust conventions: modules and functions in `snake_case`, types in `UpperCamelCase`, constants in `SCREAMING_SNAKE_CASE`. Keep modules under ~200 lines unless justified. Format with `cargo fmt` and lint using `cargo clippy --all-targets --all-features` before posting a PR. +Prefer keeping `#[cfg(test)]` blocks at the end of a file; only place `#[cfg(test)]` imports directly beneath normal imports when that keeps the module tidy. ## Testing Guidelines Adhere to TDDβ€”write the failing test first, implement minimally, then refactor. Co-locate unit tests with their crates, reserving heavier scenarios for `integration/` against the prepared-transaction Postgres stack. Invoke `cargo nextest run --test-threads=1 ` for focused iterations; gate Kerberos coverage behind `--features gssapi`. Do **not** run `cargo test`; Nextest with a single-thread budget is the required harness. diff --git a/CLAUDE.md b/CLAUDE.md index 5ab59889a..5e70dfc80 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -2,7 +2,7 @@ - `cargo check` to test that the code compiles. It shouldn't contain warnings. This is quicker than `cargo build`. - `cargo fmt` to reformat code according to Rust standards. -- `cargo nextest run --test-threads=1 ` to run a specific test +- `cargo nextest run --test-threads=1 ` to run a specific test. Run `pgdog` tests from the `pgdog` directory (`cd pgdog` first). - `cargo nextest run --test-threads=1 --no-fail-fast` to run all tests. Make sure to use `--test-threads=1` because some tests conflict with each other. # Code style diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 118ff0f9e..3cbb39047 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -4,6 +4,7 @@ Contributions are welcome. If you see a bug, feel free to submit a PR with a fix ## Necessary crates - cargo install +(if you use mise, these can be installed with `mise install`) - cargo-nextest - cargo-watch diff --git a/Cargo.lock b/Cargo.lock index e58a6fd3f..245826f36 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -114,6 +114,15 @@ version = "1.0.98" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e16d2d3311acee920a9eb8d33b8cbc1787ce4a264e85f964c2404b969bdcd487" +[[package]] +name = "ar_archive_writer" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0c269894b6fe5e9d7ada0cf69b5bf847ff35bc25fc271f08e1d080fce80339a" +dependencies = [ + "object 0.32.2", +] + [[package]] name = "arc-swap" version = "1.7.1" @@ -200,7 +209,7 @@ dependencies = [ "cfg-if", "libc", "miniz_oxide", - "object", + "object 0.36.7", "rustc-demangle", "windows-targets 0.52.6", ] @@ -433,10 +442,11 @@ dependencies = [ [[package]] name = "cc" -version = "1.2.22" +version = "1.2.51" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "32db95edf998450acc7881c932f94cd9b05c87b4b2599e8bab064753da4acfd1" +checksum = "7a0aeaff4ff1a90589618835a598e545176939b97874f7abc7851caa0618f203" dependencies = [ + "find-msvc-tools", "jobserver", "libc", "shlex", @@ -825,6 +835,37 @@ dependencies = [ "powerfmt", ] +[[package]] +name = "derive_builder" +version = "0.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "507dfb09ea8b7fa618fcf76e953f4f5e192547945816d5358edffe39f6f94947" +dependencies = [ + "derive_builder_macro", +] + +[[package]] +name = "derive_builder_core" +version = "0.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d5bcf7b024d6835cfb3d473887cd966994907effbe9227e8c8219824d06c4e8" +dependencies = [ + "darling", + "proc-macro2", + "quote", + "syn 2.0.101", +] + +[[package]] +name = "derive_builder_macro" +version = "0.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ab63b0e2bf4d5928aff72e83a7dace85d7bba5fe12dcc3c5a572d78caffd3f3c" +dependencies = [ + "derive_builder_core", + "syn 2.0.101", +] + [[package]] name = "derive_more" version = "2.0.1" @@ -1006,6 +1047,12 @@ dependencies = [ "winapi", ] +[[package]] +name = "find-msvc-tools" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "645cbb3a84e60b7531617d5ae4e57f7e27308f6445f5abf653209ea76dec8dff" + [[package]] name = "finl_unicode" version = "1.3.0" @@ -1324,7 +1371,7 @@ dependencies = [ "idna", "ipnet", "once_cell", - "rand 0.9.1", + "rand 0.9.2", "ring 0.17.14", "thiserror 2.0.12", "tinyvec", @@ -1346,7 +1393,7 @@ dependencies = [ "moka", "once_cell", "parking_lot", - "rand 0.9.1", + "rand 0.9.2", "resolv-conf", "smallvec", "thiserror 2.0.12", @@ -1717,15 +1764,6 @@ dependencies = [ "either", ] -[[package]] -name = "itertools" -version = "0.14.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2b192c782037fadd9cfa75548310488aabdbf3d2da73885b31bd0abd03351285" -dependencies = [ - "either", -] - [[package]] name = "itoa" version = "1.0.15" @@ -2142,6 +2180,15 @@ dependencies = [ "libc", ] +[[package]] +name = "object" +version = "0.32.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6a622008b6e321afc04970976f62ee297fdbaa6f95318ca343e3eebb9648441" +dependencies = [ + "memchr", +] + [[package]] name = "object" version = "0.36.7" @@ -2334,8 +2381,7 @@ dependencies = [ [[package]] name = "pg_query" version = "6.1.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6ca6fdb8f9d32182abf17328789f87f305dd8c8ce5bf48c5aa2b5cffc94e1c04" +source = "git+https://github.com/pgdogdev/pg_query.rs.git?rev=4f79b92fe4d630b1f253f27f13c9096c77530fd6#4f79b92fe4d630b1f253f27f13c9096c77530fd6" dependencies = [ "bindgen 0.66.1", "cc", @@ -2346,12 +2392,13 @@ dependencies = [ "prost-build", "serde", "serde_json", + "stacker", "thiserror 1.0.69", ] [[package]] name = "pgdog" -version = "0.1.8" +version = "0.1.22" dependencies = [ "arc-swap", "async-trait", @@ -2362,6 +2409,7 @@ dependencies = [ "clap", "csv-core", "dashmap", + "derive_builder", "fnv", "futures", "hickory-resolver", @@ -2375,9 +2423,11 @@ dependencies = [ "once_cell", "parking_lot", "pg_query", + "pgdog-config", "pgdog-plugin", + "pgdog-vector", "pin-project", - "rand 0.8.5", + "rand 0.9.2", "ratatui", "regex", "rmp-serde", @@ -2389,6 +2439,8 @@ dependencies = [ "serde_json", "sha1", "socket2", + "stats_alloc", + "tempfile", "thiserror 2.0.12", "tikv-jemallocator", "tokio", @@ -2401,6 +2453,22 @@ dependencies = [ "uuid", ] +[[package]] +name = "pgdog-config" +version = "0.1.0" +dependencies = [ + "pgdog-vector", + "rand 0.9.2", + "serde", + "serde_json", + "tempfile", + "thiserror 2.0.12", + "toml 0.8.22", + "tracing", + "url", + "uuid", +] + [[package]] name = "pgdog-example-plugin" version = "0.1.0" @@ -2440,6 +2508,14 @@ dependencies = [ "toml 0.9.5", ] +[[package]] +name = "pgdog-vector" +version = "0.1.0" +dependencies = [ + "serde", + "serde_json", +] + [[package]] name = "phf" version = "0.11.3" @@ -2570,7 +2646,7 @@ dependencies = [ "hmac", "md-5", "memchr", - "rand 0.9.1", + "rand 0.9.2", "sha2", "stringprep", ] @@ -2656,7 +2732,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "be769465445e8c1474e9c5dac2018218498557af32d9ed057325ec9a41ae81bf" dependencies = [ "heck", - "itertools 0.14.0", + "itertools 0.13.0", "log", "multimap", "once_cell", @@ -2676,7 +2752,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8a56d757972c98b346a9b766e3f02746cde6dd1cd1d1d563472929fdd74bec4d" dependencies = [ "anyhow", - "itertools 0.14.0", + "itertools 0.13.0", "proc-macro2", "quote", "syn 2.0.101", @@ -2691,6 +2767,16 @@ dependencies = [ "prost", ] +[[package]] +name = "psm" +version = "0.1.28" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d11f2fedc3b7dafdc2851bc52f277377c5473d378859be234bc7ebb593144d01" +dependencies = [ + "ar_archive_writer", + "cc", +] + [[package]] name = "ptr_meta" version = "0.1.4" @@ -2745,9 +2831,9 @@ dependencies = [ [[package]] name = "rand" -version = "0.9.1" +version = "0.9.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9fbfd9d094a40bf3ae768db9361049ace4c0e04a4fd6b359518bd7b73a73dd97" +checksum = "6db2770f06117d490610c7488547d543617b21bfa07796d7a12f6f1bd53850d1" dependencies = [ "rand_chacha 0.9.0", "rand_core 0.9.3", @@ -3093,6 +3179,8 @@ version = "0.1.0" dependencies = [ "chrono", "futures-util", + "libc", + "ordered-float", "parking_lot", "reqwest", "serde_json", @@ -3100,6 +3188,7 @@ dependencies = [ "sqlx", "tokio", "tokio-postgres", + "tokio-rustls", "uuid", ] @@ -3275,8 +3364,7 @@ checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" [[package]] name = "scram" version = "0.6.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7679a5e6b97bac99b2c208894ba0d34b17d9657f0b728c1cd3bf1c5f7f6ebe88" +source = "git+https://github.com/pgdogdev/scram.git?rev=848003d#848003da80ef7d421ce6bf540f9e1ea01089caf8" dependencies = [ "base64 0.13.1", "rand 0.8.5", @@ -3770,12 +3858,32 @@ version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a8f112729512f8e442d81f95a8a7ddf2b7c6b8a1a6f509a95864142b30cab2d3" +[[package]] +name = "stacker" +version = "0.1.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e1f8b29fb42aafcea4edeeb6b2f2d7ecd0d969c48b4cf0d2e64aafc471dd6e59" +dependencies = [ + "cc", + "cfg-if", + "libc", + "psm", + "windows-sys 0.52.0", + "windows-sys 0.59.0", +] + [[package]] name = "static_assertions" version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f" +[[package]] +name = "stats_alloc" +version = "0.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c0e04424e733e69714ca1bbb9204c1a57f09f5493439520f9f68c132ad25eec" + [[package]] name = "stringprep" version = "0.1.5" @@ -3898,9 +4006,9 @@ checksum = "55937e1799185b12863d447f42597ed69d9928686b8d88a1df17376a097d8369" [[package]] name = "tempfile" -version = "3.20.0" +version = "3.23.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e8a64e3985349f2441a1a9ef0b853f869006c3855f2cda6862a94d26ebb9d6a1" +checksum = "2d31c77bdf42a745371d260a26ca7163f1e0924b64afa0b688e61b5a9fa02f16" dependencies = [ "fastrand", "getrandom 0.3.3", @@ -4146,7 +4254,7 @@ dependencies = [ "pin-project-lite", "postgres-protocol", "postgres-types", - "rand 0.9.1", + "rand 0.9.2", "socket2", "tokio", "tokio-util", diff --git a/Cargo.toml b/Cargo.toml index 95cdffe8b..c4d3b6477 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -3,14 +3,17 @@ members = [ "examples/demo", "integration/rust", - "pgdog", "pgdog-macros", - "pgdog-plugin", "pgdog-plugin-build", "plugins/pgdog-example-plugin", + "pgdog", "pgdog-config", "pgdog-macros", + "pgdog-plugin", "pgdog-plugin-build", "pgdog-vector", "plugins/pgdog-example-plugin", ] resolver = "2" # [patch.crates-io] # tokio = { path = "../tokio/tokio" } +# [patch."https://github.com/pgdogdev/pg_query.rs.git"] +# pg_query = { path = "../pg_query.rs" } + [profile.release] codegen-units = 1 lto = true diff --git a/Dockerfile b/Dockerfile index 3d369f4a9..b4ff5a871 100644 --- a/Dockerfile +++ b/Dockerfile @@ -12,14 +12,27 @@ WORKDIR /build RUN rm /bin/sh && ln -s /bin/bash /bin/sh RUN source ~/.cargo/env && \ + if [ "$(uname -m)" = "aarch64" ] || [ "$(uname -m)" = "arm64" ]; then \ + export RUSTFLAGS="-Ctarget-feature=+lse"; \ + fi && \ + cd pgdog && \ cargo build --release FROM ubuntu:latest ENV RUST_LOG=info +ARG PSQL_VERSION=18 +ENV PSQL_VERSION=${PSQL_VERSION} RUN apt update && \ - apt install -y ca-certificates postgresql-client && \ + apt install -y curl ca-certificates ssl-cert && \ update-ca-certificates +RUN install -d /usr/share/postgresql-common/pgdg && \ + curl -o /usr/share/postgresql-common/pgdg/apt.postgresql.org.asc --fail https://www.postgresql.org/media/keys/ACCC4CF8.asc && \ + . /etc/os-release && \ + sh -c "echo 'deb [signed-by=/usr/share/postgresql-common/pgdg/apt.postgresql.org.asc] https://apt.postgresql.org/pub/repos/apt $VERSION_CODENAME-pgdg main' > /etc/apt/sources.list.d/pgdg.list" + +RUN apt update && apt install -y postgresql-client-${PSQL_VERSION} + COPY --from=builder /build/target/release/pgdog /usr/local/bin/pgdog WORKDIR /pgdog diff --git a/README.md b/README.md index 6159040e0..50c43467b 100644 --- a/README.md +++ b/README.md @@ -220,7 +220,7 @@ psql "postgres://pgdog:pgdog@127.0.0.1:6432/pgdog?gssencmode=disable" ## 🚦 Status 🚦 -This project is just getting started and early adopters are welcome to try PgDog internally. Status on features stability will be [updated regularly](https://docs.pgdog.dev/features/). Most features have tests and are benchmarked regularly for performance regressions. +PgDog is used in production and at scale. Most features are stable, while some are experimental. Check [documentation](https://docs.pgdog.dev/features/) for more details. ## Performance diff --git a/SECURITY.md b/SECURITY.md new file mode 100644 index 000000000..fe8f7c05a --- /dev/null +++ b/SECURITY.md @@ -0,0 +1,13 @@ +# Security Policy + +## Supported Versions + +Since the project is currently pre-v1.0, only the latest released version +is supported with security updates. Users are expected to upgrade to the +newest version to receive security fixes. + +## Reporting a Vulnerability + +Please report any vulnerabilities to security@pgdog.dev. You'll get a acknowledgement +within 24 hours. Any high priority vulnerabilities will be fixed within 72 hours. + diff --git a/cli.sh b/cli.sh new file mode 100755 index 000000000..c568f2870 --- /dev/null +++ b/cli.sh @@ -0,0 +1,47 @@ +#!/bin/bash +# +# Dev CLI. +# + +# Connect to the admin database. +function admin() { + PGPASSWORD=pgdog psql -h 127.0.0.1 -p 6432 -U admin admin +} + +# Run pgbench. +# +# Arguments: +# +# - protocol: simple|extended|prepared +# +function bench() { + PGPASSWORD=pgdog pgbench -h 127.0.0.1 -p 6432 -U pgdog pgdog --protocol ${2:-simple} -t 100000 -c 10 -P 1 -S +} + +function bench_init() { + PGPASSWORD=pgdog pgbench -h 127.0.0.1 -p 6432 -U pgdog pgdog -i +} + +function psql_cmd() { + PGPASSWORD=pgdog psql -h 127.0.0.1 -p 6432 -U pgdog $1 +} + +# Parse command +case "$1" in + admin) + admin + ;; + psql) + psql_cmd $2 + ;; + binit) + bench_init + ;; + bench) + bench $2 + ;; + *) + echo "Usage: $0 {admin} {bench}" + exit 1 + ;; +esac diff --git a/docker-compose.yml b/docker-compose.yml index a006cc72f..2db1181dd 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -11,7 +11,7 @@ services: environment: RUST_LOG: debug shard_0: - image: postgres:17 + image: postgres:18 environment: POSTGRES_PASSWORD: postgres volumes: @@ -19,7 +19,7 @@ services: networks: - postgres shard_1: - image: postgres:17 + image: postgres:18 environment: POSTGRES_PASSWORD: postgres networks: @@ -27,7 +27,7 @@ services: volumes: - ./docker/setup.sql:/docker-entrypoint-initdb.d/setup.sql shard_2: - image: postgres:17 + image: postgres:18 environment: POSTGRES_PASSWORD: postgres networks: diff --git a/example.pgdog.toml b/example.pgdog.toml index eb38c9720..9e44070cf 100644 --- a/example.pgdog.toml +++ b/example.pgdog.toml @@ -106,6 +106,11 @@ tls_certificate = "relative/or/absolute/path/to/certificate.pem" # Path to PEM-encoded TLS certificate private key tls_private_key = "relative/or/absolute/path/to/private_key.pem" +# Require clients to use TLS when connecting +# +# Default: false +tls_client_required = false + # TLS mode for Postgres server connections. # # Default: disabled @@ -127,6 +132,12 @@ tls_server_ca_certificate = "relative/or/absolute/path/to/certificate.pem" # Default: 60 seconds shutdown_timeout = 60_000 +# How long to wait for active connections to be forcibly terminated +# after shutdown_timeout expires. +# +# Default: disabled +shutdown_termination_timeout = 60_000 + # OpenMetrics server port. # # If set, enables Prometheus-style metrics exporter. @@ -153,6 +164,13 @@ openmetrics_namespace = "pgdog_" # Default: extended prepared_statements = "extended" +# Will override the default and enable query parsing. This can be useful if +# you don't have a primary/replica, have one shard, and don't want +# prepared_statements = 'full' but still want to enable query parsing +# in order to synchronise your clients `set` commands for instance. +# NOTE: true enables, and false will just fallback to default behaviour. +query_parser_enabled = true + # Limit on the number of prepared statements active on # any Postgres server connection. # @@ -317,6 +335,11 @@ port = 5432 # Role set to replica. role = "replica" +[rewrite] +enabled = false +shard_key = "ignore" +split_inserts = "error" + # # TCP tweaks. # diff --git a/examples/demo/docker-compose.yml b/examples/demo/docker-compose.yml index d79f54560..a54d758bc 100644 --- a/examples/demo/docker-compose.yml +++ b/examples/demo/docker-compose.yml @@ -1,6 +1,6 @@ services: db_0: - image: postgres:16 + image: postgres:18 environment: POSTGRES_PASSWORD: postgres ports: @@ -8,7 +8,7 @@ services: volumes: - shard_0:/var/lib/postgresql/data db_1: - image: postgres:16 + image: postgres:18 environment: POSTGRES_PASSWORD: postgres ports: @@ -16,7 +16,7 @@ services: volumes: - shard_1:/var/lib/postgresql/data db_2: - image: postgres:16 + image: postgres:18 environment: POSTGRES_PASSWORD: postgres ports: diff --git a/examples/multi_tenant/docker-compose.yml b/examples/multi_tenant/docker-compose.yml index fa85f64be..5e599160b 100644 --- a/examples/multi_tenant/docker-compose.yml +++ b/examples/multi_tenant/docker-compose.yml @@ -1,6 +1,6 @@ services: shard_0: - image: postgres:17 + image: postgres:18 environment: POSTGRES_PASSWORD: postgres ports: @@ -8,7 +8,7 @@ services: volumes: - shard_0:/var/lib/postgresql/data shard_1: - image: postgres:17 + image: postgres:18 environment: POSTGRES_PASSWORD: postgres ports: @@ -16,7 +16,7 @@ services: volumes: - shard_1:/var/lib/postgresql/data shard_2: - image: postgres:17 + image: postgres:18 environment: POSTGRES_PASSWORD: postgres ports: diff --git a/examples/pgbouncer_benchmark/docker-compose.yml b/examples/pgbouncer_benchmark/docker-compose.yml index 43f580573..323cfde0f 100644 --- a/examples/pgbouncer_benchmark/docker-compose.yml +++ b/examples/pgbouncer_benchmark/docker-compose.yml @@ -1,6 +1,6 @@ services: postgres: - image: postgres:17 + image: postgres:18 environment: POSTGRES_PASSWORD: postgres pgbouncer: diff --git a/integration/common.sh b/integration/common.sh index 8001131ed..812acc69f 100644 --- a/integration/common.sh +++ b/integration/common.sh @@ -3,6 +3,7 @@ # N.B.: Scripts using this are expected to define $SCRIPT_DIR # correctly. # +export NODE_ID=pgdog-dev-1 COMMON_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd ) function wait_for_pgdog() { echo "Waiting for PgDog" diff --git a/integration/complex/cancel_query/pgdog.toml b/integration/complex/cancel_query/pgdog.toml new file mode 100644 index 000000000..51e3decaf --- /dev/null +++ b/integration/complex/cancel_query/pgdog.toml @@ -0,0 +1,16 @@ +[general] +query_timeout = 60000 +shutdown_timeout = 0 +shutdown_termination_timeout = 1000 + +[rewrite] +enabled = false +shard_key = "ignore" +split_inserts = "error" + +[[databases]] +name = "pgdog" +host = "127.0.0.1" + +[admin] +password = "pgdog" diff --git a/integration/complex/cancel_query/run.py b/integration/complex/cancel_query/run.py new file mode 100644 index 000000000..e63f93b73 --- /dev/null +++ b/integration/complex/cancel_query/run.py @@ -0,0 +1,60 @@ +import asyncio +import sys + +import asyncpg +import psycopg + + +SHUTDOWN_TIMEOUT = 1 +SLEEP_SECONDS = 10000 +APPLICATION_NAME = "pgdog_cancel_query" + + +async def trigger_shutdown() -> None: + conn = await asyncpg.connect( + host="127.0.0.1", + port=6432, + database="pgdog", + user="pgdog", + password="pgdog", + ) + + try: + await conn.execute(f"SET application_name = '{APPLICATION_NAME}'") + sleep_task = asyncio.create_task(conn.execute(f"SELECT pg_sleep({SLEEP_SECONDS})")) + + # Give the backend time to register the long running query. + await asyncio.sleep(1) + + admin = psycopg.connect( + "dbname=admin user=admin host=127.0.0.1 port=6432 password=pgdog" + ) + admin.autocommit = True + try: + admin.execute("SHUTDOWN") + finally: + admin.close() + + try: + await asyncio.wait_for(sleep_task, timeout=SHUTDOWN_TIMEOUT) + except asyncio.TimeoutError: + print("pg_sleep query did not terminate after PgDog shutdown", file=sys.stderr) + raise SystemExit(1) + except (asyncpg.exceptions.PostgresError, asyncpg.exceptions.InterfaceError): + # Expected: connection terminates as PgDog shuts down. + return + except (ConnectionError, psycopg.Error): + # psycopg errors propagate through asyncpg when connection drops. + return + else: + print("pg_sleep query completed without interruption", file=sys.stderr) + raise SystemExit(1) + finally: + try: + await conn.close() + except Exception: + pass + + +if __name__ == "__main__": + asyncio.run(trigger_shutdown()) diff --git a/integration/complex/cancel_query/run.sh b/integration/complex/cancel_query/run.sh new file mode 100644 index 000000000..a50f0db37 --- /dev/null +++ b/integration/complex/cancel_query/run.sh @@ -0,0 +1,30 @@ +#!/bin/bash +set -e +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd ) +source ${SCRIPT_DIR}/../../common.sh + +APPLICATION_NAME="pgdog_cancel_query" +QUERY="SELECT COUNT(*) FROM pg_stat_activity WHERE application_name = '${APPLICATION_NAME}'" + +export PGPASSWORD=pgdog + +active_venv + +run_pgdog "${SCRIPT_DIR}" +wait_for_pgdog + +pushd ${SCRIPT_DIR} +python run.py +popd + +attempts=0 +until [[ "$(psql -h localhost -U pgdog -tAq -c "${QUERY}")" == "0" ]]; do + if [[ ${attempts} -ge 5 ]]; then + echo "Found lingering sessions with application_name='${APPLICATION_NAME}'" >&2 + exit 1 + fi + attempts=$((attempts + 1)) + sleep 5 +done + +stop_pgdog diff --git a/integration/complex/cancel_query/users.toml b/integration/complex/cancel_query/users.toml new file mode 100644 index 000000000..1347d9833 --- /dev/null +++ b/integration/complex/cancel_query/users.toml @@ -0,0 +1,57 @@ +[[users]] +name = "pgdog" +database = "pgdog" +password = "pgdog" + +[[users]] +name = "pgdog" +database = "pgdog_sharded" +password = "pgdog" + +[[users]] +name = "pgdog_session" +database = "pgdog" +password = "pgdog" +server_user = "pgdog" +pooler_mode = "session" + +[[users]] +name = "pgdog_2pc" +database = "pgdog" +password = "pgdog" +server_user = "pgdog" +two_phase_commit = true +min_pool_size = 0 + +[[users]] +name = "pgdog_2pc" +database = "pgdog_sharded" +password = "pgdog" +server_user = "pgdog" +two_phase_commit = true +min_pool_size = 0 + +[[users]] +name = "pgdog_migrator" +database = "pgdog_sharded" +password = "pgdog" +server_user = "pgdog" +schema_admin = true + +[[users]] +name = "pgdog" +database = "failover" +password = "pgdog" + +[[users]] +name = "pgdog" +database = "single_sharded_list" +password = "pgdog" + +[[users]] +name = "pgdog_no_cross_shard" +database = "single_sharded_list" +password = "pgdog" +server_user = "pgdog" +cross_shard_disabled = true +min_pool_size = 0 diff --git a/integration/complex/passthrough_auth/run.sh b/integration/complex/passthrough_auth/run.sh index 37ae5d350..71c5ad79b 100644 --- a/integration/complex/passthrough_auth/run.sh +++ b/integration/complex/passthrough_auth/run.sh @@ -12,6 +12,7 @@ PGDOG_BIN_PATH="${PGDOG_BIN:-${SCRIPT_DIR}/../../../target/release/pgdog}" "${PGDOG_BIN_PATH}" \ --config ${SCRIPT_DIR}/pgdog-enabled.toml \ --users ${SCRIPT_DIR}/users.toml & +PGDOG_PID=$! until pg_isready -h 127.0.0.1 -p 6432 -U pgdog -d pgdog; do sleep 1 @@ -32,10 +33,12 @@ if [[ "$statement_timeout" != *"100ms"* ]]; then fi killall -TERM pgdog +wait "${PGDOG_PID}" 2> /dev/null || true "${PGDOG_BIN_PATH}" \ --config ${SCRIPT_DIR}/pgdog-disabled.toml \ --users ${SCRIPT_DIR}/users.toml & +PGDOG_PID=$! until pg_isready -h 127.0.0.1 -p 6432 -U pgdog -d pgdog; do sleep 1 @@ -49,3 +52,4 @@ fi psql -U pgdog pgdog -c 'SELECT 1' > /dev/null killall -TERM pgdog +wait "${PGDOG_PID}" 2> /dev/null || true diff --git a/integration/complex/run.sh b/integration/complex/run.sh old mode 100644 new mode 100755 index a761375d7..6726daa0a --- a/integration/complex/run.sh +++ b/integration/complex/run.sh @@ -5,4 +5,5 @@ SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd ) pushd ${SCRIPT_DIR} bash shutdown.sh bash passthrough_auth/run.sh +bash cancel_query/run.sh popd diff --git a/integration/complex/shutdown.sh b/integration/complex/shutdown.sh index 1f1b7b333..41d8e1d20 100644 --- a/integration/complex/shutdown.sh +++ b/integration/complex/shutdown.sh @@ -12,6 +12,8 @@ pushd ${SCRIPT_DIR} python shutdown.py pgdog popd +sleep 1 + if pgrep pgdog; then echo "Shutdown failed" exit 1 @@ -24,6 +26,8 @@ pushd ${SCRIPT_DIR} python shutdown.py pgdog_sharded popd +sleep 1 + if pgrep pgdog; then echo "Shutdown failed" exit 1 diff --git a/integration/copy_data/dev.sh b/integration/copy_data/dev.sh new file mode 100644 index 000000000..6486d4711 --- /dev/null +++ b/integration/copy_data/dev.sh @@ -0,0 +1,20 @@ +#!/bin/bash +set -e +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd ) +DEFAULT_BIN="${SCRIPT_DIR}/../../target/release/pgdog" +PGDOG_BIN=${PGDOG_BIN:-$DEFAULT_BIN} + +export PGUSER=pgdog +export PGDATABASE=pgdog +export PGHOST=127.0.0.1 +export PGPORT=5432 +export PGPASSWORD=pgdog + +pushd ${SCRIPT_DIR} + +psql -f init.sql + +${PGDOG_BIN} schema-sync --from-database source --to-database destination --publication pgdog +${PGDOG_BIN} data-sync --sync-only --from-database source --to-database destination --publication pgdog --replication-slot copy_data +${PGDOG_BIN} schema-sync --from-database source --to-database destination --publication pgdog --cutover +popd diff --git a/integration/copy_data/init.sql b/integration/copy_data/init.sql new file mode 100644 index 000000000..5e8dfb34f --- /dev/null +++ b/integration/copy_data/init.sql @@ -0,0 +1,8 @@ +\c pgdog1 +DROP SCHEMA IF EXISTS copy_data CASCADE; +\c pgdog2 +DROP SCHEMA IF EXISTS copy_data CASCADE; +\c pgdog +DROP SCHEMA IF EXISTS copy_data CASCADE; +SELECT pg_drop_replication_slot(slot_name) FROM pg_replication_slots; +\i setup.sql diff --git a/integration/copy_data/pgdog.toml b/integration/copy_data/pgdog.toml new file mode 100644 index 000000000..91fa253d0 --- /dev/null +++ b/integration/copy_data/pgdog.toml @@ -0,0 +1,21 @@ +[[databases]] +name = "source" +host = "127.0.0.1" +database_name = "pgdog" + +[[databases]] +name = "destination" +host = "127.0.0.1" +database_name = "pgdog1" +shard = 0 + +[[databases]] +name = "destination" +host = "127.0.0.1" +database_name = "pgdog2" +shard = 1 + +[[sharded_tables]] +database = "destination" +column = "tenant_id" +data_type = "bigint" diff --git a/integration/copy_data/psql.sh b/integration/copy_data/psql.sh new file mode 100644 index 000000000..68bf3e41e --- /dev/null +++ b/integration/copy_data/psql.sh @@ -0,0 +1,2 @@ +#!/bin/bash +PGPASSWORD=pgdog psql -h 127.0.0.1 -p 5432 -U pgdog $1 diff --git a/integration/copy_data/setup.sql b/integration/copy_data/setup.sql new file mode 100644 index 000000000..b2de7dcac --- /dev/null +++ b/integration/copy_data/setup.sql @@ -0,0 +1,161 @@ +CREATE SCHEMA IF NOT EXISTS copy_data; + +CREATE TABLE IF NOT EXISTS copy_data.users ( + id BIGINT NOT NULL, + tenant_id BIGINT NOT NULL, + email VARCHAR NOT NULL, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + settings JSONB NOT NULL DEFAULT '{}'::jsonb, + PRIMARY KEY(id, tenant_id) +) PARTITION BY HASH(tenant_id); + +CREATE TABLE IF NOT EXISTS copy_data.users_0 PARTITION OF copy_data.users + FOR VALUES WITH (MODULUS 2, REMAINDER 0); + +CREATE TABLE IF NOT EXISTS copy_data.users_1 PARTITION OF copy_data.users + FOR VALUES WITH (MODULUS 2, REMAINDER 1); + +TRUNCATE TABLE copy_data.users; + +INSERT INTO copy_data.users (id, tenant_id, email, created_at, settings) +SELECT + gs.id, + ((gs.id - 1) % 20) + 1 AS tenant_id, -- distribute across 20 tenants + format('user_%s_tenant_%s@example.com', gs.id, ((gs.id - 1) % 20) + 1) AS email, + NOW() - (random() * interval '365 days') AS created_at, -- random past date + jsonb_build_object( + 'theme', CASE (random() * 3)::int + WHEN 0 THEN 'light' + WHEN 1 THEN 'dark' + ELSE 'auto' + END, + 'notifications', (random() > 0.5) + ) AS settings +FROM generate_series(1, 10000) AS gs(id); + +DROP TABLE copy_data.orders; +CREATE TABLE IF NOT EXISTS copy_data.orders ( + id BIGSERIAL PRIMARY KEY, + user_id BIGINT NOT NULL, + tenant_id BIGINT NOT NULL, + amount DOUBLE PRECISION NOT NULL DEFAULT 0.0, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + refunded_at TIMESTAMPTZ +); + +CREATE TABLE IF NOT EXISTS copy_data.order_items ( + user_id BIGINT NOT NULL, + tenant_id BIGINT NOT NULL, + order_id BIGINT NOT NULL REFERENCES copy_data.orders(id), + amount DOUBLE PRECISION NOT NULL DEFAULT 0.0, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + refunded_at TIMESTAMPTZ +); + +-- --- Fix/define schema (safe to run if you're starting fresh) --- +-- Adjust/drop statements as needed if the tables already exist. +TRUNCATE TABLE copy_data.order_items CASCADE; +TRUNCATE TABLE copy_data.orders CASCADE; + +WITH u AS ( + -- Pull the 10k users we inserted earlier + SELECT id AS user_id, tenant_id + FROM copy_data.users + WHERE id BETWEEN 1 AND 10000 + ORDER BY id +), +orders_base AS ( + -- One order per user (10k orders), deterministic order_id = user_id + SELECT + u.user_id AS order_id, + u.user_id, + u.tenant_id, + -- random created_at in last 365 days + NOW() - (random() * INTERVAL '365 days') AS created_at, + -- ~10% refunded + CASE WHEN random() < 0.10 + THEN NOW() - (random() * INTERVAL '180 days') + ELSE NULL + END AS refunded_at + FROM u +), +items_raw AS ( + -- 1–5 items per order, random amounts $5–$200 + SELECT + ob.order_id, + ob.user_id, + ob.tenant_id, + -- skew item counts 1..5 (uniform) + gs.i AS item_index, + -- random item amount with cents + ROUND((5 + random() * 195)::numeric, 2)::float8 AS item_amount, + -- item created_at: on/after order created_at by up to 3 hours + ob.created_at + (random() * INTERVAL '3 hours') AS item_created_at, + -- if order refunded, item refunded too (optionally jitter within 2 hours) + CASE WHEN ob.refunded_at IS NOT NULL + THEN ob.refunded_at + (random() * INTERVAL '2 hours') + ELSE NULL + END AS item_refunded_at + FROM orders_base ob + CROSS JOIN LATERAL generate_series(1, 1 + (floor(random()*5))::int) AS gs(i) +), +order_totals AS ( + SELECT + order_id, + user_id, + tenant_id, + MIN(item_created_at) AS created_at, + -- sum of item amounts per order + ROUND(SUM(item_amount)::numeric, 2)::float8 AS order_amount, + -- carry refund state from items_raw (same per order) + MAX(item_refunded_at) AS refunded_at + FROM items_raw + GROUP BY order_id, user_id, tenant_id +), +ins_orders AS ( + INSERT INTO copy_data.orders (id, user_id, tenant_id, amount, created_at, refunded_at) + SELECT + ot.order_id, -- id = user_id = 1..10000 + ot.user_id, + ot.tenant_id, + ot.order_amount, + ot.created_at, + ot.refunded_at + FROM order_totals ot + RETURNING id +) +INSERT INTO copy_data.order_items (user_id, tenant_id, order_id, amount, created_at, refunded_at) +SELECT + ir.user_id, + ir.tenant_id, + ir.order_id, + ir.item_amount, + ir.item_created_at, + ir.item_refunded_at +FROM items_raw ir; + +CREATE TABLE copy_data.log_actions( + id BIGSERIAL PRIMARY KEY, + tenant_id BIGINT, + action VARCHAR, + created_at TIMESTAMPTZ NOT NULL DEFAULT now() +); + +INSERT INTO copy_data.log_actions (tenant_id, action) +SELECT + CASE WHEN random() < 0.2 THEN NULL ELSE (floor(random() * 10000) + 1)::bigint END AS tenant_id, + (ARRAY['login', 'logout', 'click', 'purchase', 'view', 'error'])[ + floor(random() * 6 + 1)::int + ] AS action +FROM generate_series(1, 10000); + +CREATE TABLE copy_data.with_identity( + id BIGINT GENERATED ALWAYS AS identity, + tenant_id BIGINT NOT NULL +); + +INSERT INTO copy_data.with_identity (tenant_id) +SELECT floor(random() * 10000)::bigint FROM generate_series(1, 10000); + +DROP PUBLICATION IF EXISTS pgdog; +CREATE PUBLICATION pgdog FOR TABLES IN SCHEMA copy_data; diff --git a/integration/copy_data/users.toml b/integration/copy_data/users.toml new file mode 100644 index 000000000..67142d309 --- /dev/null +++ b/integration/copy_data/users.toml @@ -0,0 +1,11 @@ +[[users]] +database = "source" +name = "pgdog" +password = "pgdog" +schema_admin = true + +[[users]] +database = "destination" +name = "pgdog" +password = "pgdog" +schema_admin = true diff --git a/integration/d_plus/pgdog.toml b/integration/d_plus/pgdog.toml new file mode 100644 index 000000000..a7f5ca716 --- /dev/null +++ b/integration/d_plus/pgdog.toml @@ -0,0 +1,26 @@ +[general] +expanded_explain = true + +[[databases]] +name = "pgdog" +host = "127.0.0.1" +shard = 0 +database_name = "shard_0" + +[[databases]] +name = "pgdog" +host = "127.0.0.1" +shard = 1 +database_name = "shard_1" + +[admin] +password = "pgdog" + +[[omnisharded_tables]] +database = "pgdog" +tables = [ + "pg_class", + "pg_namespace", + "pg_am" +] +sticky = true diff --git a/integration/d_plus/users.toml b/integration/d_plus/users.toml new file mode 100644 index 000000000..581cdb75b --- /dev/null +++ b/integration/d_plus/users.toml @@ -0,0 +1,4 @@ +[[users]] +name = "pgdog" +database = "pgdog" +password = "pgdog" diff --git a/integration/dev-server.sh b/integration/dev-server.sh index d425f0ca9..2f17f6de2 100755 --- a/integration/dev-server.sh +++ b/integration/dev-server.sh @@ -4,6 +4,7 @@ THIS_SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && p source ${THIS_SCRIPT_DIR}/setup.sh source ${THIS_SCRIPT_DIR}/toxi/setup.sh pushd ${THIS_SCRIPT_DIR}/../ +export NODE_ID=pgdog-dev-1 CMD="cargo run -- --config ${THIS_SCRIPT_DIR}/pgdog.toml --users ${THIS_SCRIPT_DIR}/users.toml" if [[ -z "$1" ]]; then diff --git a/integration/dry_run/pgdog.toml b/integration/dry_run/pgdog.toml index a774568f4..2fc702c19 100644 --- a/integration/dry_run/pgdog.toml +++ b/integration/dry_run/pgdog.toml @@ -1,6 +1,11 @@ [general] dry_run = true +[rewrite] +enabled = false +shard_key = "ignore" +split_inserts = "error" + [[databases]] name = "pgdog" host = "127.0.0.1" diff --git a/integration/failover/dev-server.sh b/integration/failover/dev-server.sh new file mode 100644 index 000000000..513f06312 --- /dev/null +++ b/integration/failover/dev-server.sh @@ -0,0 +1,7 @@ +#!/bin/bash +set -e +THIS_SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd ) +source ${THIS_SCRIPT_DIR}/../toxi/setup.sh +pushd ${THIS_SCRIPT_DIR} +cargo run +popd diff --git a/integration/failover/pgdog.toml b/integration/failover/pgdog.toml new file mode 100644 index 000000000..7d213af81 --- /dev/null +++ b/integration/failover/pgdog.toml @@ -0,0 +1,47 @@ +[general] +checkout_timeout = 1_000 +connect_timeout = 1_000 +ban_timeout = 30_000 +query_timeout = 1_000 +idle_healthcheck_interval = 1_000 +client_login_timeout = 1_000 +load_balancing_algorithm = "round_robin" +[rewrite] +enabled = false +shard_key = "ignore" +split_inserts = "error" + +[[databases]] +name = "failover" +host = "127.0.0.1" +port = 5435 +role = "primary" +database_name = "pgdog" + +[[databases]] +name = "failover" +host = "127.0.0.1" +port = 5436 +role = "replica" +database_name = "pgdog" +read_only = true + +[[databases]] +name = "failover" +host = "127.0.0.1" +port = 5437 +role = "replica" +database_name = "pgdog" +read_only = true + +[[databases]] +name = "failover" +host = "127.0.0.1" +port = 5438 +role = "replica" +database_name = "pgdog" +read_only = true + +[admin] +password = "pgdog" +user = "admin" diff --git a/integration/failover/psql.sh b/integration/failover/psql.sh new file mode 100644 index 000000000..e124e3392 --- /dev/null +++ b/integration/failover/psql.sh @@ -0,0 +1,3 @@ +#!/bin/bash +export PGPASSWORD=pgdog +psql -h 127.0.0.1 -p 6432 -U pgdog ${1} diff --git a/integration/failover/users.toml b/integration/failover/users.toml new file mode 100644 index 000000000..167934ef9 --- /dev/null +++ b/integration/failover/users.toml @@ -0,0 +1,4 @@ +[[users]] +database = "failover" +name = "pgdog" +password = "pgdog" diff --git a/integration/go/go_pgx/connectivity_test.go b/integration/go/go_pgx/connectivity_test.go new file mode 100644 index 000000000..de18e37c7 --- /dev/null +++ b/integration/go/go_pgx/connectivity_test.go @@ -0,0 +1,50 @@ +package main + +import ( + "context" + "testing" + + "github.com/jackc/pgx/v5" + "github.com/stretchr/testify/assert" +) + +func TestConnectivityWithoutTLS(t *testing.T) { + ctx := context.Background() + conn, err := pgx.Connect(ctx, "postgres://pgdog:pgdog@127.0.0.1:6432/pgdog?sslmode=disable") + assert.NoError(t, err) + defer conn.Close(ctx) + + err = conn.Ping(ctx) + assert.NoError(t, err) +} + +func TestConnectivityWithTLS(t *testing.T) { + ctx := context.Background() + conn, err := pgx.Connect(ctx, "postgres://pgdog:pgdog@127.0.0.1:6432/pgdog?sslmode=require") + assert.NoError(t, err) + defer conn.Close(ctx) + + err = conn.Ping(ctx) + assert.NoError(t, err) +} + +func TestConnectivityWithPassthrough(t *testing.T) { + ctx := context.Background() + + adminConn, err := pgx.Connect(ctx, "postgres://admin:pgdog@127.0.0.1:6432/admin") + assert.NoError(t, err) + defer adminConn.Close(ctx) + + _, err = adminConn.Exec(ctx, "SET passthrough_auth TO 'enabled'", pgx.QueryExecModeSimpleProtocol) + assert.NoError(t, err) + + conn, err := pgx.Connect(ctx, "postgres://pgdog:pgdog@127.0.0.1:6432/pgdog?sslmode=disable") + assert.NoError(t, err) + defer conn.Close(ctx) + + err = conn.Ping(ctx) + assert.NoError(t, err) + + _, err = pgx.Connect(ctx, "postgres://pgdog:wrong_password@127.0.0.1:6432/pgdog?sslmode=disable") + assert.Error(t, err) +} diff --git a/integration/go/go_pq/go_pq_test.go b/integration/go/go_pq/go_pq_test.go index 045316c15..9744dc153 100644 --- a/integration/go/go_pq/go_pq_test.go +++ b/integration/go/go_pq/go_pq_test.go @@ -27,7 +27,67 @@ func PqConnections() []*sql.DB { return []*sql.DB{normal, sharded} } +func TestAuthenticationWithoutTLS(t *testing.T) { + conn, err := sql.Open("postgres", "postgres://pgdog:pgdog@127.0.0.1:6432/pgdog?sslmode=disable") + assert.Nil(t, err) + defer conn.Close() + + err = conn.Ping() + assert.Nil(t, err) + + // Reset config + adminConn, err := sql.Open("postgres", "postgres://admin:pgdog@127.0.0.1:6432/admin?sslmode=disable") + assert.Nil(t, err) + defer adminConn.Close() + + _, err = adminConn.Exec("RELOAD") + assert.Nil(t, err) +} + +func TestAuthenticationWithTLS(t *testing.T) { + conn, err := sql.Open("postgres", "postgres://pgdog:pgdog@127.0.0.1:6432/pgdog?sslmode=require") + assert.Nil(t, err) + defer conn.Close() + + err = conn.Ping() + assert.Nil(t, err) +} + +func TestAuthenticationWithPassthrough(t *testing.T) { + adminConn, err := sql.Open("postgres", "postgres://admin:pgdog@127.0.0.1:6432/admin?sslmode=disable") + assert.Nil(t, err) + defer adminConn.Close() + + _, err = adminConn.Exec("SET passthrough_auth TO 'enabled'") + assert.Nil(t, err) + + conn, err := sql.Open("postgres", "postgres://pgdog:pgdog@127.0.0.1:6432/pgdog?sslmode=disable") + assert.Nil(t, err) + defer conn.Close() + + err = conn.Ping() + assert.Nil(t, err) + + badConn, err := sql.Open("postgres", "postgres://pgdog:wrong_password@127.0.0.1:6432/pgdog?sslmode=disable") + assert.Nil(t, err) + defer badConn.Close() + + err = badConn.Ping() + assert.NotNil(t, err) + + // Reset config + _, err = adminConn.Exec("RELOAD") + assert.Nil(t, err) +} + func TestPqCrud(t *testing.T) { + adminConn, err := sql.Open("postgres", "postgres://admin:pgdog@127.0.0.1:6432/admin?sslmode=disable") + assert.Nil(t, err) + defer adminConn.Close() + + _, err = adminConn.Exec("SET prepared_statements TO 'extended_anonymous'") + assert.Nil(t, err) + conns := PqConnections() for _, conn := range conns { diff --git a/integration/load_balancer/pgdog.toml b/integration/load_balancer/pgdog.toml index 57088ec7e..8c8a8aa1c 100644 --- a/integration/load_balancer/pgdog.toml +++ b/integration/load_balancer/pgdog.toml @@ -16,6 +16,12 @@ pooler_mode = "transaction" load_balancing_strategy = "round_robin" auth_type = "trust" read_write_split = "exclude_primary" +lsn_check_delay = 5_000 + +[rewrite] +enabled = false +shard_key = "ignore" +split_inserts = "error" # [replica_lag] # check_interval = 2000 diff --git a/integration/logical/pgdog.toml b/integration/logical/pgdog.toml index 60586b9d5..2910809b9 100644 --- a/integration/logical/pgdog.toml +++ b/integration/logical/pgdog.toml @@ -1,3 +1,10 @@ +[general] + +[rewrite] +enabled = false +shard_key = "ignore" +split_inserts = "error" + [[databases]] name = "pgdog" host = "127.0.0.1" diff --git a/integration/mirror/pgdog.toml b/integration/mirror/pgdog.toml index 466217d0b..1e60ea827 100644 --- a/integration/mirror/pgdog.toml +++ b/integration/mirror/pgdog.toml @@ -2,6 +2,11 @@ mirror_exposure = 1.0 openmetrics_port = 9090 +[rewrite] +enabled = false +shard_key = "ignore" +split_inserts = "error" + [[databases]] name = "pgdog" host = "127.0.0.1" diff --git a/integration/pgbench/run.sh b/integration/pgbench/run.sh index 83c730a20..870729f19 100644 --- a/integration/pgbench/run.sh +++ b/integration/pgbench/run.sh @@ -7,5 +7,6 @@ run_pgdog wait_for_pgdog bash ${SCRIPT_DIR}/dev.sh +bash ${SCRIPT_DIR}/stress.sh stop_pgdog diff --git a/integration/pgbench/stress.sh b/integration/pgbench/stress.sh new file mode 100644 index 000000000..ac08aa258 --- /dev/null +++ b/integration/pgbench/stress.sh @@ -0,0 +1,12 @@ +#!/bin/bash +# +# This will create a decent amount of churn on the network & message +# buffers, by generating different size responses. +# +set -e +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd ) +export PGPASSWORD=pgdog + +pushd ${SCRIPT_DIR} +pgbench -h 127.0.0.1 -U pgdog -p 6432 pgdog -t 100 -c 10 --protocol simple -f stress.sql -P 1 +popd diff --git a/integration/pgbench/stress.sql b/integration/pgbench/stress.sql new file mode 100644 index 000000000..f45cb2515 --- /dev/null +++ b/integration/pgbench/stress.sql @@ -0,0 +1,7 @@ +\set response_size (random(2, 1000000)) + +-- Generate large response from server. +-- Range: 2 bytes to 1M +SELECT + string_agg(substr('abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789', (random()*61+1)::int, 1), '') +FROM generate_series(1, :response_size); diff --git a/integration/pgdog.toml b/integration/pgdog.toml index 25a15093d..01d789e2c 100644 --- a/integration/pgdog.toml +++ b/integration/pgdog.toml @@ -11,11 +11,25 @@ read_write_strategy = "aggressive" openmetrics_port = 9090 openmetrics_namespace = "pgdog_" prepared_statements_limit = 500 +prepared_statements = "extended" +expanded_explain = true # dns_ttl = 15_000 query_cache_limit = 500 pub_sub_channel_size = 4098 two_phase_commit = false healthcheck_port = 8080 +tls_certificate = "integration/tls/cert.pem" +tls_private_key = "integration/tls/key.pem" +query_parser_engine = "pg_query_raw" + +[memory] +net_buffer = 8096 +message_buffer = 8096 + +[rewrite] +enabled = false +shard_key = "ignore" +split_inserts = "error" # ------------------------------------------------------------------------------ # ----- Database :: pgdog ------------------------------------------------------ @@ -356,6 +370,37 @@ shard = 1 database = "pgdog_sharded" tables = ["sharded_omni"] +# ------------------------------------------------------------------------------ +# ----- Schema-based sharding -------------------------------------------------- + +[[databases]] +name = "pgdog_schema" +host = "127.0.0.1" +database_name = "shard_0" +shard = 0 + +[[databases]] +name = "pgdog_schema" +host = "127.0.0.1" +database_name = "shard_1" +shard = 1 + +[[sharded_schemas]] +database = "pgdog_schema" +name = "shard_0" +shard = 0 + +[[sharded_schemas]] +database = "pgdog_schema" +name = "shard_1" +shard = 1 + +# By default, all queries go to shard 0 +[[sharded_schemas]] +database = "pgdog_schema" +shard = 0 + + # ------------------------------------------------------------------------------ # ----- Admin ------------------------------------------------------------------ diff --git a/integration/plugins/dev.sh b/integration/plugins/dev.sh index 7392d18d8..9ac000e4f 100644 --- a/integration/plugins/dev.sh +++ b/integration/plugins/dev.sh @@ -2,5 +2,7 @@ set -e SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd ) pushd ${SCRIPT_DIR} +bundle config set --local path vendor/bundle +bundle install bundle exec rspec *_spec.rb popd diff --git a/integration/plugins/pgdog.toml b/integration/plugins/pgdog.toml index 10a323e29..1359e9547 100644 --- a/integration/plugins/pgdog.toml +++ b/integration/plugins/pgdog.toml @@ -1,3 +1,10 @@ +[general] + +[rewrite] +enabled = false +shard_key = "ignore" +split_inserts = "error" + [[plugins]] name = "pgdog_example_plugin" diff --git a/integration/prepared_statements_full/pgdog.toml b/integration/prepared_statements_full/pgdog.toml new file mode 100644 index 000000000..37ae8403d --- /dev/null +++ b/integration/prepared_statements_full/pgdog.toml @@ -0,0 +1,9 @@ +[general] +prepared_statements = "full" + +[[databases]] +name = "pgdog" +host = "127.0.0.1" + +[admin] +password = "pgdog" diff --git a/integration/prepared_statements_full/users.toml b/integration/prepared_statements_full/users.toml new file mode 100644 index 000000000..9a8205f04 --- /dev/null +++ b/integration/prepared_statements_full/users.toml @@ -0,0 +1,4 @@ +[[users]] +database = "pgdog" +name = "pgdog" +password = "pgdog" diff --git a/integration/pub_sub/pgdog.toml b/integration/pub_sub/pgdog.toml index 26ac29f6e..2509a0645 100644 --- a/integration/pub_sub/pgdog.toml +++ b/integration/pub_sub/pgdog.toml @@ -1,6 +1,11 @@ [general] pub_sub_channel_size = 4098 +[rewrite] +enabled = false +shard_key = "ignore" +split_inserts = "error" + [[databases]] name = "pgdog" host = "127.0.0.1" diff --git a/integration/python/globals.py b/integration/python/globals.py index 564bdde9b..6c42c1924 100644 --- a/integration/python/globals.py +++ b/integration/python/globals.py @@ -37,6 +37,12 @@ def normal_sync(): ) +def direct_sync(): + return psycopg.connect( + user="pgdog", password="pgdog", dbname="pgdog", host="127.0.0.1", port=5432 + ) + + async def sharded_async(): return await asyncpg.connect( user="pgdog", @@ -57,3 +63,13 @@ async def normal_async(): port=6432, statement_cache_size=250, ) + +async def schema_sharded_async(): + return await asyncpg.connect( + user="pgdog", + password="pgdog", + database="pgdog_schema", + host="127.0.0.1", + port=6432, + statement_cache_size=250 + ) diff --git a/integration/python/test_asyncpg.py b/integration/python/test_asyncpg.py index 1ab30e828..c48f173f2 100644 --- a/integration/python/test_asyncpg.py +++ b/integration/python/test_asyncpg.py @@ -1,10 +1,11 @@ import asyncpg import pytest from datetime import datetime -from globals import normal_async, sharded_async, no_out_of_sync, admin +from globals import normal_async, sharded_async, no_out_of_sync, admin, schema_sharded_async import random import string import pytest_asyncio +from io import BytesIO @pytest_asyncio.fixture @@ -231,6 +232,7 @@ async def test_copy(conns): records = 250 for i in range(50): for conn in conns: + # Test COPY FROM (TO table) rows = [[x, f"value_{x}", datetime.now()] for x in range(records)] await conn.copy_records_to_table( "sharded", records=rows, columns=["id", "value", "created_at"] @@ -238,6 +240,16 @@ async def test_copy(conns): count = await conn.fetch("SELECT COUNT(*) FROM sharded") assert len(count) == 1 assert count[0][0] == records + + # Test COPY TO STDOUT + buffer = BytesIO() + copied_data = await conn.copy_from_table( + "sharded", columns=["id", "value", "created_at"], output=buffer + ) + buffer.seek(0) + lines = buffer.read().decode('utf-8').strip().split('\n') + assert len(lines) == records, f"expected {records} lines in COPY output, got {len(lines)}" + await conn.execute("DELETE FROM sharded") @@ -400,3 +412,135 @@ async def test_copy_jsonb(): finally: await conn.execute("DROP TABLE IF EXISTS jsonb_copy_test") await conn.close() + +@pytest.mark.asyncio +async def test_schema_sharding(): + admin().cursor().execute("SET cross_shard_disabled TO true") + conn = await schema_sharded_async() + + for _ in range(25): + for schema in ["shard_0", "shard_1"]: + await conn.execute(f"DROP SCHEMA IF EXISTS {schema} CASCADE") + await conn.execute(f"CREATE SCHEMA {schema}") + for shard in [0, 1]: + schema = f"shard_{shard}" + await conn.execute(f"CREATE TABLE {schema}.test(id BIGINT, created_at TIMESTAMPTZ DEFAULT NOW())") + await conn.fetch(f"SELECT * FROM {schema}.test WHERE id = $1", 1) + + insert = await conn.fetch(f"INSERT INTO {schema}.test VALUES ($1, NOW()), ($2, NOW()) RETURNING *", 1, 2) + assert len(insert) == 2 + assert insert[0][0] == 1 + assert insert[1][0] == 2 + + update = await conn.fetch(f"UPDATE {schema}.test SET id = $1 WHERE id = $2 RETURNING *", 3, 2) + assert len(update) == 1 + assert update[0][0] == 3 + + delete = await conn.execute(f"DELETE FROM {schema}.test WHERE id = $1", 3) + assert delete == "DELETE 1" + + await conn.execute(f"TRUNCATE {schema}.test") + admin().cursor().execute("SET cross_shard_disabled TO false") + +@pytest.mark.asyncio +async def test_schema_sharding_transactions(): + admin().cursor().execute("SET cross_shard_disabled TO true") + conn = await schema_sharded_async() + + for _ in range(25): + for schema in ["shard_0", "shard_1"]: + await conn.execute(f"DROP SCHEMA IF EXISTS {schema} CASCADE") + await conn.execute(f"CREATE SCHEMA {schema}") + + for shard in [0, 1]: + async with conn.transaction(): + await conn.execute("SET LOCAL statement_timeout TO '10s'") + schema = f"shard_{shard}" + await conn.execute(f"CREATE TABLE {schema}.test(id BIGINT, created_at TIMESTAMPTZ DEFAULT NOW())") + await conn.fetch(f"SELECT * FROM {schema}.test WHERE id = $1", 1) + + insert = await conn.fetch(f"INSERT INTO {schema}.test VALUES ($1, NOW()), ($2, NOW()) RETURNING *", 1, 2) + assert len(insert) == 2 + assert insert[0][0] == 1 + assert insert[1][0] == 2 + + update = await conn.fetch(f"UPDATE {schema}.test SET id = $1 WHERE id = $2 RETURNING *", 3, 2) + assert len(update) == 1 + assert update[0][0] == 3 + + delete = await conn.execute(f"DELETE FROM {schema}.test WHERE id = $1", 3) + assert delete == "DELETE 1" + admin().cursor().execute("SET cross_shard_disabled TO false") + +@pytest.mark.asyncio +async def test_schema_sharding_default(): + admin().cursor().execute("SET cross_shard_disabled TO true") + conn = await schema_sharded_async() + + for _ in range(25): + # Note no schema specified. + await conn.execute("DROP TABLE IF EXISTS test_schema_sharding_default") + await conn.execute("CREATE TABLE test_schema_sharding_default(id BIGINT)") + + await conn.execute("SELECT * FROM test_schema_sharding_default") + try: + await conn.execute("/* pgdog_shard: 1 */ SELECT * FROM test_schema_sharding_default") + raise Exception("table shouldn't exist on shard 1") + except Exception as e: + assert "relation \"test_schema_sharding_default\" does not exist" == str(e) + admin().cursor().execute("SET cross_shard_disabled TO false") + + +@pytest.mark.asyncio +async def test_schema_sharding_search_path(): + admin().cursor().execute("SET cross_shard_disabled TO true") + + conn = await schema_sharded_async() + + for schema in ["shard_0", "shard_1"]: + await conn.execute(f"CREATE SCHEMA IF NOT EXISTS {schema}") + await conn.execute(f"CREATE TABLE IF NOT EXISTS {schema}.test(id BIGINT, created_at TIMESTAMPTZ DEFAULT NOW())") + + + import asyncio + + async def run_test(): + conn = await schema_sharded_async() + + for _ in range(10): + for schema in ["shard_0", "shard_1"]: + await conn.execute(f"SET search_path TO {schema}") + + async with conn.transaction(): + await conn.execute("SET LOCAL statement_timeout TO '10s'") + await conn.fetch(f"SELECT * FROM {schema}.test WHERE id = $1", 1) + + await conn.close() + + await asyncio.gather(*[run_test() for _ in range(10)]) + admin().cursor().execute("SET cross_shard_disabled TO false") + + +@pytest.mark.asyncio +async def test_pgdog_role_selection(): + conn = await asyncpg.connect( + user="pgdog", + password="pgdog", + database="pgdog", + host="127.0.0.1", + port=6432, + statement_cache_size=250, + server_settings={ + "pgdog.role": "replica", + } + ) + + got_err = False + for _ in range(10): + try: + await conn.execute("CREATE TABLE IF NOT EXISTS test_pgdog_role_selection(id BIGINT)") + except asyncpg.exceptions.ReadOnlySQLTransactionError: + got_err = True + pass + + assert got_err diff --git a/integration/python/test_prepared.py b/integration/python/test_prepared.py index d7d8df249..4a36accf8 100644 --- a/integration/python/test_prepared.py +++ b/integration/python/test_prepared.py @@ -1,25 +1,33 @@ -from globals import normal_sync -import pytest +from globals import normal_sync, no_out_of_sync, admin +from multiprocessing import Pool +import time -@pytest.mark.skip(reason="These are not working") -def test_prepared_full(): - for _ in range(5): - conn = normal_sync() - conn.autocommit = True - - cur = conn.cursor() - cur.execute("PREPARE test_stmt AS SELECT 1") - cur.execute("PREPARE test_stmt AS SELECT 2") - +def run_prepare_execute(worker_id): conn = normal_sync() conn.autocommit = True - for _ in range(5): - cur = conn.cursor() - - for i in range(5): - cur.execute(f"PREPARE test_stmt_{i} AS SELECT $1::bigint") - cur.execute(f"EXECUTE test_stmt_{i}({i})") - result = cur.fetchone() - assert result[0] == i - conn.commit() + cur = conn.cursor() + + stmt_name = f"stmt_{worker_id}" + cur.execute(f"PREPARE {stmt_name} AS SELECT $1::bigint * 2") + time.sleep(0.01) + + for i in range(100): + cur.execute(f"EXECUTE {stmt_name}({i})") + result = cur.fetchone() + assert result[0] == i * 2 + time.sleep(0.01) + + cur.execute(f"DEALLOCATE {stmt_name}") + conn.close() + return True + + +def test_prepare_execute_parallel(): + admin().execute("SET prepared_statements TO 'full'") + + with Pool(5) as pool: + results = pool.map(run_prepare_execute, range(5)) + assert all(results) + no_out_of_sync() + admin().execute("RELOAD") diff --git a/integration/python/test_psycopg.py b/integration/python/test_psycopg.py index 23ae11d97..a9df46f22 100644 --- a/integration/python/test_psycopg.py +++ b/integration/python/test_psycopg.py @@ -1,5 +1,6 @@ import psycopg -from globals import no_out_of_sync, sharded_sync, normal_sync +import pytest +from globals import direct_sync, no_out_of_sync, normal_sync, sharded_sync def setup(conn): @@ -69,3 +70,50 @@ def _run_insert_test(conn): assert len(results) == 1 assert results[0][0] == id no_out_of_sync() + + +def _execute_parameter_count(conn, count: int) -> int: + placeholders = ", ".join("%s" for _ in range(count)) + query = f"SELECT array_length(ARRAY[{placeholders}], 1)" + params = list(range(count)) + cur = conn.cursor() + cur.execute(query, params) + value = cur.fetchone()[0] + conn.commit() + return value + + +@pytest.mark.parametrize( + "count, expect_error_keywords", + [ + (65_535, ()), + (65_536, ("between 0 and 65535", "too many")), + ], +) +def test_postgres_variants_parameter_limits(count, expect_error_keywords): + successes = [] + errors = [] + for connector in (direct_sync, normal_sync): + conn = connector() + try: + if expect_error_keywords: + with pytest.raises(psycopg.Error) as excinfo: + _execute_parameter_count(conn, count) + + message = str(excinfo.value).lower() + errors.append(message) + try: + conn.rollback() + except psycopg.Error: + pass + else: + successes.append(_execute_parameter_count(conn, count)) + finally: + conn.close() + + if expect_error_keywords: + assert len(errors) == 2 + for message in errors: + assert any(keyword in message for keyword in expect_error_keywords) + else: + assert successes == [count, count] diff --git a/integration/python/test_psycopg2.py b/integration/python/test_psycopg2.py new file mode 100644 index 000000000..d3c7fd6a2 --- /dev/null +++ b/integration/python/test_psycopg2.py @@ -0,0 +1,131 @@ +import psycopg2 +import pytest +from globals import admin + +queries = [ + "SELECT 1", + "CREATE TABLE IF NOT EXISTS test_conn_reads(id BIGINT)", + "INSERT INTO test_conn_reads (id) VALUES (1)", + "INSERT INTO test_conn_reads (id) VALUES (2)", + "INSERT INTO test_conn_reads (id) VALUES (3)", + "SELECT * FROM test_conn_reads WHERE id = 1", + "SELECT * FROM test_conn_reads WHERE id = 2", + "SELECT * FROM test_conn_reads WHERE id = 3", + "SET work_mem TO '4MB'", + "SET work_mem TO '6MB'", + "SET work_mem TO '8MB'", +] + + +@pytest.fixture +def conn_reads(): + return psycopg2.connect( + "host=127.0.0.1 port=6432 user=pgdog password=pgdog " + "options='-c pgdog.role=replica'" + ) + + +@pytest.fixture +def conn_writes(): + return psycopg2.connect( + "host=127.0.0.1 port=6432 user=pgdog password=pgdog " + "options='-c pgdog.role=primary'" + ) + + +@pytest.fixture +def conn_default(): + return psycopg2.connect("host=127.0.0.1 port=6432 user=pgdog password=pgdog") + + +def test_conn_writes(conn_writes): + admin().execute("SET query_parser TO 'off'") + for query in queries: + conn_writes.autocommit = True + cursor = conn_writes.cursor() + cursor.execute(query) + admin().execute("SET query_parser TO 'auto'") + + +def test_conn_reads(conn_reads, conn_writes): + admin().execute("SET query_parser TO 'off'") + + conn_writes.autocommit = True + conn_reads.autocommit = True + + conn_writes.cursor().execute( + "CREATE TABLE IF NOT EXISTS test_conn_reads(id BIGINT)" + ) + + read = False + for query in queries: + cursor = conn_reads.cursor() + try: + cursor.execute(query) + except psycopg2.errors.ReadOnlySqlTransaction: + # Some will succeed because we allow reads + # on the primary. + read = True + admin().execute("SET query_parser TO 'auto'") + + conn_writes.cursor().execute("DROP TABLE IF EXISTS test_conn_reads") + assert read, "expected some queries to hit replicas and fail" + + +def test_transactions_writes(conn_writes): + admin().execute("SET query_parser TO 'off'") + + for query in queries: + conn_writes.cursor().execute(query) + conn_writes.commit() + + admin().execute("SET query_parser TO 'auto'") + + +def test_transactions_reads(conn_reads): + admin().execute("SET query_parser TO 'off'") + read = False + + for query in queries: + try: + conn_reads.cursor().execute(query) + except psycopg2.errors.ReadOnlySqlTransaction: + # Some will succeed because we allow reads + # on the primary. + read = True + conn_reads.commit() + + assert read, "expected some queries to hit replicas and fail" + admin().execute("SET query_parser TO 'auto'") + + +def test_transaction_reads_explicit(conn_reads, conn_writes): + conn_reads.autocommit = True + admin().execute("SET query_parser TO 'off'") + + conn_writes.cursor().execute( + "CREATE TABLE IF NOT EXISTS test_conn_reads(id BIGINT)" + ) + conn_writes.commit() + + cursor = conn_reads.cursor() + + read = False + + for _ in range(15): + cursor.execute("BEGIN") + try: + cursor.execute("INSERT INTO test_conn_reads (id) VALUES (1)") + cursor.execute("COMMIT") + except psycopg2.errors.ReadOnlySqlTransaction: + read = True + cursor.execute("ROLLBACK") + + assert read, "expected some queries to hit replicas and fail" + + for _ in range(15): + cursor.execute("BEGIN READ ONLY") # Won't be parsed, doesn't matter to PgDog + cursor.execute("SELECT 1") + cursor.execute("ROLLBACK") + + admin().execute("SET query_parser TO 'on'") diff --git a/integration/python/test_sqlalchemy.py b/integration/python/test_sqlalchemy.py index 60ea7695c..c5f9e2ee4 100644 --- a/integration/python/test_sqlalchemy.py +++ b/integration/python/test_sqlalchemy.py @@ -10,7 +10,7 @@ import pytest_asyncio import pytest from sqlalchemy.sql.expression import delete -from globals import admin +from globals import admin, schema_sharded_async class Base(AsyncAttrs, DeclarativeBase): @@ -31,16 +31,23 @@ class User(Base): email: Mapped[str] +# class Products(Base): +# __tablename__ == "products" + +# id: Mapped[int] = mapped_column(primary_key=True) +# name: Mapped[str] + + @pytest_asyncio.fixture async def engines(): # Configure connection pool for stress testing normal = create_async_engine( "postgresql+asyncpg://pgdog:pgdog@127.0.0.1:6432/pgdog", - pool_size=20, # Number of connections to maintain in pool - max_overflow=30, # Additional connections beyond pool_size - pool_timeout=30, # Timeout when getting connection from pool - pool_recycle=3600, # Recycle connections after 1 hour - pool_pre_ping=True, # Verify connections before use + pool_size=20, # Number of connections to maintain in pool + max_overflow=30, # Additional connections beyond pool_size + pool_timeout=30, # Timeout when getting connection from pool + pool_recycle=3600, # Recycle connections after 1 hour + pool_pre_ping=True, # Verify connections before use ) normal_sessions = async_sessionmaker(normal, expire_on_commit=True) @@ -57,6 +64,50 @@ async def engines(): return [(normal, normal_sessions), (sharded, sharded_sessions)] +@pytest_asyncio.fixture +async def schema_sharding_engine(): + from sqlalchemy import event + + pool = create_async_engine( + "postgresql+asyncpg://pgdog:pgdog@127.0.0.1:6432/pgdog_schema", + pool_size=20, # Number of connections to maintain in pool + max_overflow=30, # Additional connections beyond pool_size + pool_timeout=30, # Timeout when getting connection from pool + pool_recycle=3600, # Recycle connections after 1 hour + pool_pre_ping=True, # Verify connections before use + ) + + @event.listens_for(pool.sync_engine, "connect") + def set_search_path(dbapi_connection, connection_record): + cursor = dbapi_connection.cursor() + cursor.execute("SET search_path TO shard_0, public") + cursor.close() + dbapi_connection.commit() + + session = async_sessionmaker(pool, expire_on_commit=True) + return pool, session + + +@pytest_asyncio.fixture +async def schema_sharding_startup_param(): + pool = create_async_engine( + "postgresql+asyncpg://pgdog:pgdog@127.0.0.1:6432/pgdog_schema", + connect_args={ + "server_settings": { + "search_path": "shard_0,public", + } + }, + pool_size=20, # Number of connections to maintain in pool + max_overflow=30, # Additional connections beyond pool_size + pool_timeout=30, # Timeout when getting connection from pool + pool_recycle=3600, # Recycle connections after 1 hour + pool_pre_ping=True, # Verify connections before use + ) + + session = async_sessionmaker(pool, expire_on_commit=True) + return pool, session + + @pytest.mark.asyncio async def test_session_manager(engines): for engine, session_factory in engines: @@ -192,7 +243,9 @@ async def test_connection_pool_stress(engines): # Setup test table async with normal() as session: await session.execute(text("DROP TABLE IF EXISTS stress_test")) - await session.execute(text(""" + await session.execute( + text( + """ CREATE TABLE stress_test ( id BIGSERIAL PRIMARY KEY, name VARCHAR(50), @@ -201,20 +254,27 @@ async def test_connection_pool_stress(engines): active BOOLEAN, created_at TIMESTAMP DEFAULT NOW() ) - """)) + """ + ) + ) await session.commit() # Insert initial data async with normal() as session: for i in range(100): - name = ''.join(random.choices(string.ascii_letters, k=10)) + name = "".join(random.choices(string.ascii_letters, k=10)) age = random.randint(18, 80) score = round(random.uniform(0, 100), 2) active = random.choice([True, False]) - await session.execute(text(""" + await session.execute( + text( + """ INSERT INTO stress_test (name, age, score, active) VALUES (:name, :age, :score, :active) - """), {"name": name, "age": age, "score": score, "active": active}) + """ + ), + {"name": name, "age": age, "score": score, "active": active}, + ) await session.commit() async def run_varied_queries(engine, task_id): @@ -230,134 +290,196 @@ async def run_varied_queries(engine, task_id): # Vary the query patterns to create different prepared statements # Use task_id and iteration to create more unique queries query_type = random.randint(1, 12) - variation = (task_id * 20 + i) % 10 # Creates 200 unique variations + variation = ( + task_id * 20 + i + ) % 10 # Creates 200 unique variations if query_type == 1: # Simple select with WHERE - vary the column to create unique statements age_filter = 20 + variation if variation % 3 == 0: - result = await conn.execute(text( - "SELECT COUNT(*) FROM stress_test WHERE age > :age" - ), {"age": age_filter}) + result = await conn.execute( + text( + "SELECT COUNT(*) FROM stress_test WHERE age > :age" + ), + {"age": age_filter}, + ) elif variation % 3 == 1: - result = await conn.execute(text( - "SELECT COUNT(*) FROM stress_test WHERE age >= :age" - ), {"age": age_filter}) + result = await conn.execute( + text( + "SELECT COUNT(*) FROM stress_test WHERE age >= :age" + ), + {"age": age_filter}, + ) else: - result = await conn.execute(text( - "SELECT COUNT(*) FROM stress_test WHERE age = :age" - ), {"age": age_filter}) + result = await conn.execute( + text( + "SELECT COUNT(*) FROM stress_test WHERE age = :age" + ), + {"age": age_filter}, + ) elif query_type == 2: # Complex WHERE with multiple conditions min_age = random.randint(18, 40) max_score = random.uniform(50, 100) - result = await conn.execute(text(""" + result = await conn.execute( + text( + """ SELECT name, age, score FROM stress_test WHERE age >= :min_age AND score <= :max_score AND active = true LIMIT 10 - """), {"min_age": min_age, "max_score": max_score}) + """ + ), + {"min_age": min_age, "max_score": max_score}, + ) elif query_type == 3: # Aggregation with GROUP BY - result = await conn.execute(text(""" + result = await conn.execute( + text( + """ SELECT active, AVG(score) as avg_score, COUNT(*) as count FROM stress_test GROUP BY active HAVING COUNT(*) > :min_count - """), {"min_count": random.randint(1, 10)}) + """ + ), + {"min_count": random.randint(1, 10)}, + ) elif query_type == 4: # ORDER BY with different columns - use variation for uniqueness - order_col = ['age', 'score', 'name', 'created_at'][variation % 4] - order_dir = ['ASC', 'DESC'][variation % 2] - result = await conn.execute(text(f""" + order_col = ["age", "score", "name", "created_at"][ + variation % 4 + ] + order_dir = ["ASC", "DESC"][variation % 2] + result = await conn.execute( + text( + f""" SELECT * FROM stress_test WHERE score > :score ORDER BY {order_col} {order_dir} LIMIT :limit - """), {"score": variation * 5, "limit": 5 + variation}) + """ + ), + {"score": variation * 5, "limit": 5 + variation}, + ) elif query_type == 5: # JOIN with subquery - result = await conn.execute(text(""" + result = await conn.execute( + text( + """ SELECT s.name, s.age FROM stress_test s WHERE s.score > ( SELECT AVG(score) FROM stress_test WHERE active = :active ) LIMIT :limit - """), {"active": random.choice([True, False]), "limit": random.randint(3, 8)}) + """ + ), + { + "active": random.choice([True, False]), + "limit": random.randint(3, 8), + }, + ) elif query_type == 6: # UPDATE with different conditions - use task_id to avoid conflicts - min_age_base = 20 + (task_id * 5) # Each task gets different age range - await conn.execute(text(""" + min_age_base = 20 + ( + task_id * 5 + ) # Each task gets different age range + await conn.execute( + text( + """ UPDATE stress_test SET score = score + :bonus WHERE age BETWEEN :min_age AND :max_age - """), { - "bonus": random.uniform(-5, 5), - "min_age": min_age_base, - "max_age": min_age_base + 4 - }) + """ + ), + { + "bonus": random.uniform(-5, 5), + "min_age": min_age_base, + "max_age": min_age_base + 4, + }, + ) # Transaction auto-commits with engine.begin() elif query_type == 7: # INSERT with varying values - name = ''.join(random.choices(string.ascii_letters, k=8)) - await conn.execute(text(""" + name = "".join(random.choices(string.ascii_letters, k=8)) + await conn.execute( + text( + """ INSERT INTO stress_test (name, age, score, active) VALUES (:name, :age, :score, :active) - """), { - "name": f"stress_{name}_{task_id}", - "age": random.randint(18, 80), - "score": round(random.uniform(0, 100), 2), - "active": random.choice([True, False]) - }) + """ + ), + { + "name": f"stress_{name}_{task_id}", + "age": random.randint(18, 80), + "score": round(random.uniform(0, 100), 2), + "active": random.choice([True, False]), + }, + ) # Transaction auto-commits with engine.begin() elif query_type == 8: # DELETE with different conditions - await conn.execute(text(""" + await conn.execute( + text( + """ DELETE FROM stress_test WHERE name LIKE :pattern AND score < :max_score - """), { - "pattern": f"stress_%_{task_id}", - "max_score": random.uniform(10, 30) - }) + """ + ), + { + "pattern": f"stress_%_{task_id}", + "max_score": random.uniform(10, 30), + }, + ) # Transaction auto-commits with engine.begin() elif query_type == 9: # Different SELECT with JOIN-like pattern - result = await conn.execute(text(f""" + result = await conn.execute( + text( + f""" SELECT name, score FROM stress_test WHERE active = :active AND score BETWEEN :min_score AND :max_score ORDER BY score {['ASC', 'DESC'][variation % 2]} LIMIT :limit - """), { - "active": variation % 2 == 0, - "min_score": variation * 10, - "max_score": variation * 10 + 20, - "limit": 5 + variation - }) + """ + ), + { + "active": variation % 2 == 0, + "min_score": variation * 10, + "max_score": variation * 10 + 20, + "limit": 5 + variation, + }, + ) elif query_type == 10: # Window function queries - result = await conn.execute(text(f""" + result = await conn.execute( + text( + f""" SELECT name, age, score, ROW_NUMBER() OVER (ORDER BY score {['ASC', 'DESC'][variation % 2]}) as rank FROM stress_test WHERE age > :min_age LIMIT :limit - """), { - "min_age": 20 + variation, - "limit": 10 + variation - }) + """ + ), + {"min_age": 20 + variation, "limit": 10 + variation}, + ) elif query_type == 11: # CASE statement variations - result = await conn.execute(text(f""" + result = await conn.execute( + text( + f""" SELECT name, CASE WHEN score > :high_threshold THEN 'High' @@ -369,17 +491,22 @@ async def run_varied_queries(engine, task_id): WHERE active = :active ORDER BY {['age', 'score', 'name'][variation % 3]} LIMIT :limit - """), { - "high_threshold": 70 + variation, - "med_threshold": 40 + variation, - "active": variation % 2 == 0, - "limit": 8 + variation - }) + """ + ), + { + "high_threshold": 70 + variation, + "med_threshold": 40 + variation, + "active": variation % 2 == 0, + "limit": 8 + variation, + }, + ) elif query_type == 12: # Advanced aggregation with different GROUP BY if variation % 2 == 0: - result = await conn.execute(text(""" + result = await conn.execute( + text( + """ SELECT CASE WHEN age < :age_threshold THEN 'Young' ELSE 'Old' END as age_group, AVG(score) as avg_score, @@ -389,21 +516,27 @@ async def run_varied_queries(engine, task_id): FROM stress_test GROUP BY CASE WHEN age < :age_threshold THEN 'Young' ELSE 'Old' END HAVING COUNT(*) > :min_count - """), { - "age_threshold": 30 + variation, - "min_count": variation + 1 - }) + """ + ), + { + "age_threshold": 30 + variation, + "min_count": variation + 1, + }, + ) else: - result = await conn.execute(text(""" + result = await conn.execute( + text( + """ SELECT active, PERCENTILE_CONT(0.5) WITHIN GROUP (ORDER BY score) as median_score, COUNT(*) as total FROM stress_test WHERE score > :min_score GROUP BY active - """), { - "min_score": variation * 5 - }) + """ + ), + {"min_score": variation * 5}, + ) queries_run += 1 break # Success, break out of retry loop @@ -436,3 +569,85 @@ async def run_varied_queries(engine, task_id): result = await session.execute(text("SELECT COUNT(*) FROM stress_test")) count = result.scalar() assert count > 0 # Should have some data left + + +@pytest.mark.asyncio +async def test_schema_sharding_with_startup_param_and_set_after_query( + schema_sharding_startup_param, +): + import asyncio + + admin().cursor().execute("SET cross_shard_disabled TO true") + + (pool, session_factory) = schema_sharding_startup_param + + async def run_test(task_id): + async with session_factory() as session: + async with session.begin(): + await session.execute(text("SELECT 1")) + await session.execute(text("SET LOCAL work_mem TO '4MB'")) + for row in await session.execute(text("SHOW search_path")): + print(row) + + tasks = [asyncio.create_task(run_test(i)) for i in range(10)] + await asyncio.gather(*tasks) + + admin().cursor().execute("SET cross_shard_disabled TO false") + + +@pytest.mark.asyncio +async def test_schema_sharding(schema_sharding_engine): + import asyncio + + admin().cursor().execute("SET cross_shard_disabled TO true") + # All queries should touch shard_0 only. + # Set it up separately + conn = await schema_sharded_async() + await conn.execute("SET search_path TO shard_0, public") + await conn.execute("CREATE SCHEMA IF NOT EXISTS shard_0") + await conn.execute("DROP TABLE IF EXISTS shard_0.test_schema_sharding") + await conn.execute("CREATE TABLE shard_0.test_schema_sharding(id BIGINT)") + await conn.close() + + (pool, session_factory) = schema_sharding_engine + + async def run_schema_sharding_test(task_id): + for _ in range(10): + async with session_factory() as session: + async with session.begin(): + await session.execute(text("SET LOCAL work_mem TO '4MB'")) + for row in await session.execute(text("SHOW search_path")): + print(row) + await session.execute(text("SELECT 1")) + await session.execute(text("SELECT * FROM test_schema_sharding")) + + # Run 10 concurrent executions in parallel + tasks = [asyncio.create_task(run_schema_sharding_test(i)) for i in range(1)] + await asyncio.gather(*tasks) + + admin().cursor().execute("SET cross_shard_disabled TO false") + + +@pytest.mark.asyncio +async def test_role_selection(): + engine = create_async_engine( + "postgresql+asyncpg://pgdog:pgdog@127.0.0.1:6432/pgdog", + pool_size=20, + max_overflow=30, + pool_timeout=30, + pool_recycle=3600, + pool_pre_ping=True, + connect_args={"server_settings": {"pgdog.role": "primary"}}, + ) + session_factory = async_sessionmaker(engine, expire_on_commit=True) + + for _ in range(1): + async with session_factory() as session: + async with session.begin(): + await session.execute( + text("CREATE TABLE IF NOT EXISTS test_role_selection(id BIGINT)") + ) + await session.execute( + text("CREATE TABLE IF NOT EXISTS test_role_selection(id BIGINT)") + ) + await session.rollback() diff --git a/integration/rewrite/pgdog.toml b/integration/rewrite/pgdog.toml new file mode 100644 index 000000000..7012ff002 --- /dev/null +++ b/integration/rewrite/pgdog.toml @@ -0,0 +1,22 @@ +[[databases]] +name = "pgdog" +shard = 0 +host = "127.0.0.1" +database_name = "shard_0" + +[[databases]] +name = "pgdog" +shard = 1 +host = "127.0.0.1" +database_name = "shard_1" + +[[sharded_tables]] +column = "org_id" +database = "pgdog" + +[rewrite] +enabled = true +split_inserts = "rewrite" + +[admin] +password = "pgdog" diff --git a/integration/rewrite/users.toml b/integration/rewrite/users.toml new file mode 100644 index 000000000..581cdb75b --- /dev/null +++ b/integration/rewrite/users.toml @@ -0,0 +1,4 @@ +[[users]] +name = "pgdog" +database = "pgdog" +password = "pgdog" diff --git a/integration/role_detector/failover.sh b/integration/role_detector/failover.sh new file mode 100644 index 000000000..1118e2e94 --- /dev/null +++ b/integration/role_detector/failover.sh @@ -0,0 +1,3 @@ +#!/bin/bash +docker stop $(docker ps | grep 45000 | awk '{print $1}') +PGPASSWORD=postgres psql -h 127.0.0.1 -p 45001 -U postgres postgres -c 'SELECT pg_promote()' diff --git a/integration/role_detector/pgdog.toml b/integration/role_detector/pgdog.toml new file mode 100644 index 000000000..61301e3fe --- /dev/null +++ b/integration/role_detector/pgdog.toml @@ -0,0 +1,25 @@ +[general] +lsn_check_delay = 0 +lsn_check_interval = 5_000 +idle_timeout = 1_000 + +[[databases]] +name = "postgres" +role = "auto" +host = "127.0.0.1" +port = 45000 + +[[databases]] +name = "postgres" +role = "auto" +host = "127.0.0.1" +port = 45001 + +[[databases]] +name = "postgres" +role = "auto" +host = "127.0.0.1" +port = 45002 + +[admin] +password = "pgdog" diff --git a/integration/role_detector/users.toml b/integration/role_detector/users.toml new file mode 100644 index 000000000..6a7546826 --- /dev/null +++ b/integration/role_detector/users.toml @@ -0,0 +1,4 @@ +[[users]] +name = "postgres" +database = "postgres" +password = "postgres" diff --git a/integration/ruby/ar_spec.rb b/integration/ruby/ar_spec.rb index ecae2ba75..800a211f5 100644 --- a/integration/ruby/ar_spec.rb +++ b/integration/ruby/ar_spec.rb @@ -2,6 +2,7 @@ require_relative 'rspec_helper' require 'pp' +require 'timeout' class Sharded < ActiveRecord::Base self.table_name = 'sharded' @@ -187,4 +188,129 @@ def conn(db, prepared) end end end + + describe 'chaos testing with interrupted queries' do + before do + conn('failover', false) + ActiveRecord::Base.connection.execute 'DROP TABLE IF EXISTS sharded' + ActiveRecord::Base.connection.execute 'CREATE TABLE sharded (id BIGSERIAL PRIMARY KEY, value TEXT)' + end + + it 'handles interrupted queries and continues operating normally' do + interrupted_count = 0 + successful_count = 0 + mutex = Mutex.new + + # Apply latency toxic to slow down query transmission, + # making it easier to interrupt queries mid-flight + Toxiproxy[:primary].toxic(:latency, latency: 100, jitter: 50).apply do + # Phase 1: Chaos - interrupt queries randomly with thread kills + chaos_threads = [] + killer_threads = [] + + # Start 10 query threads + 10.times do |thread_id| + t = Thread.new do + 100.times do |i| + begin + case rand(3) + when 0 + # SELECT query + Sharded.where('id > ?', 0).limit(10).to_a + when 1 + # INSERT query + Sharded.create value: "thread_#{thread_id}_iter_#{i}" + when 2 + # Transaction with multiple operations + Sharded.transaction do + rec = Sharded.create value: "tx_#{thread_id}_#{i}" + Sharded.where(id: rec.id).first if rec.id + end + end + mutex.synchronize { successful_count += 1 } + rescue StandardError => e + # Killed mid-query or other error + mutex.synchronize { interrupted_count += 1 } + end + end + end + chaos_threads << t + end + + # Start killer thread that randomly kills query threads + killer = Thread.new do + 50.times do + sleep(rand(0.01..0.05)) + alive_threads = chaos_threads.select(&:alive?) + if alive_threads.any? + victim = alive_threads.sample + victim.kill + mutex.synchronize { interrupted_count += 1 } + end + end + end + killer_threads << killer + + # Wait for killer to finish + killer_threads.each(&:join) + + # Wait for remaining threads (with timeout) + chaos_threads.each { |t| t.join(0.1) } + + puts "Chaos phase complete: #{successful_count} successful, #{interrupted_count} interrupted" + expect(interrupted_count).to be > 0 + end # End toxiproxy latency + + # Give PgDog time to clean up broken connections + sleep(0.5) + + # Disconnect all connections to clear bad state + ActiveRecord::Base.connection_pool.disconnect! + + # Wait a bit more for cleanup + sleep(0.5) + + # Phase 2: Verify database continues to operate normally + verification_errors = [] + errors_mutex = Mutex.new + + verification_threads = 10.times.map do |thread_id| + Thread.new do + 20.times do |i| + begin + # Simple queries that don't depend on finding specific records + # INSERT + rec = Sharded.create value: "verify_#{thread_id}_#{i}" + expect(rec.id).to be > 0 + + # SELECT with basic query + results = Sharded.where('value LIKE ?', 'verify_%').limit(5).to_a + expect(results).to be_a(Array) + + # COUNT query + count = Sharded.where('id > ?', 0).count + expect(count).to be >= 0 + rescue PG::Error => e + # PG errors should fail the test + raise + rescue StandardError => e + errors_mutex.synchronize { verification_errors << e } + end + end + end + end + + verification_threads.each(&:join) + + # Verify no errors occurred during verification + expect(verification_errors).to be_empty, "Verification errors: #{verification_errors.map(&:message).join(', ')}" + + # Verify we can still execute basic queries + ActiveRecord::Base.connection.execute('SELECT 1') + + # Verify count works + count = Sharded.count + expect(count).to be >= 0 + end + end end diff --git a/integration/ruby/lb_spec.rb b/integration/ruby/lb_spec.rb index f0844af20..a3a6c3c39 100644 --- a/integration/ruby/lb_spec.rb +++ b/integration/ruby/lb_spec.rb @@ -12,6 +12,8 @@ describe 'random' do it 'distributes traffic evenly' do conn = failover + # Reset stats and bans + admin.exec "RECONNECT" before = admin_stats('failover') 250.times do diff --git a/integration/ruby/rspec_helper.rb b/integration/ruby/rspec_helper.rb index 2069433a3..5eb8317ae 100644 --- a/integration/ruby/rspec_helper.rb +++ b/integration/ruby/rspec_helper.rb @@ -31,8 +31,10 @@ def ensure_done expect(pool['cl_waiting']).to eq('0') expect(pool['out_of_sync']).to eq('0') end + current_client_id = conn.backend_pid clients = conn.exec 'SHOW CLIENTS' clients.each do |client| + next if client['id'].to_i == current_client_id expect(client['state']).to eq('idle') end servers = conn.exec 'SHOW SERVERS' diff --git a/integration/rust/Cargo.toml b/integration/rust/Cargo.toml index db57f55a3..a485aa50f 100644 --- a/integration/rust/Cargo.toml +++ b/integration/rust/Cargo.toml @@ -17,3 +17,6 @@ chrono = "0.4" parking_lot = "*" reqwest = "*" serde_json = "*" +ordered-float = "4.2" +tokio-rustls = "0.26" +libc = "0.2" diff --git a/integration/rust/src/lib.rs b/integration/rust/src/lib.rs index 138906d09..23f5525c1 100644 --- a/integration/rust/src/lib.rs +++ b/integration/rust/src/lib.rs @@ -1 +1,2 @@ pub mod setup; +pub mod utils; diff --git a/integration/rust/src/setup.rs b/integration/rust/src/setup.rs index 2fa02b996..7268f8f0b 100644 --- a/integration/rust/src/setup.rs +++ b/integration/rust/src/setup.rs @@ -44,6 +44,14 @@ pub async fn connections_sqlx() -> Vec> { pools } +pub async fn connection_sqlx_direct() -> Pool { + PgPoolOptions::new() + .max_connections(1) + .connect("postgres://pgdog:pgdog@127.0.0.1:5432/pgdog?application_name=sqlx_direct") + .await + .unwrap() +} + #[derive(Debug, PartialEq, Clone)] pub struct Backend { pub pid: i32, diff --git a/integration/rust/src/utils.rs b/integration/rust/src/utils.rs new file mode 100644 index 000000000..c05185344 --- /dev/null +++ b/integration/rust/src/utils.rs @@ -0,0 +1,19 @@ +use super::setup::admin_sqlx; +use sqlx::{Executor, Row}; + +pub async fn assert_setting_str(name: &str, expected: &str) { + let admin = admin_sqlx().await; + let rows = admin.fetch_all("SHOW CONFIG").await.unwrap(); + let mut found = false; + for row in rows { + let db_name: String = row.get(0); + let value: String = row.get(1); + + if name == db_name { + found = true; + assert_eq!(value, expected); + } + } + + assert!(found); +} diff --git a/integration/rust/tests/data/tls/initial_cert.pem b/integration/rust/tests/data/tls/initial_cert.pem new file mode 100644 index 000000000..463db3c09 --- /dev/null +++ b/integration/rust/tests/data/tls/initial_cert.pem @@ -0,0 +1,18 @@ +-----BEGIN CERTIFICATE----- +MIIC0TCCAbmgAwIBAgIJAOrzAMPPjSh2MA0GCSqGSIb3DQEBCwUAMBgxFjAUBgNV +BAMMDXBnZG9nLWluaXRpYWwwHhcNMjUwOTMwMTYxNDE4WhcNMjYwOTMwMTYxNDE4 +WjAYMRYwFAYDVQQDDA1wZ2RvZy1pbml0aWFsMIIBIjANBgkqhkiG9w0BAQEFAAOC +AQ8AMIIBCgKCAQEAnCZjNgngYmFtQeYoWI3gKFkR2bvtJUpsI2feBDew6cZR2tDY +vjeGm60qK0H4rYE0WELmr2qAuuN/Rb6+vZjaQdT8v77MEZPbXhAqPcAuPRTmJm4G +OadJiErf9/W+wsOPbDau6Ko5wKqR9qjiySQ4B0D3UgCEIC/hYjv3XHMsYCjGvRZ/ +sfYA5BSlHxlQ6odcLb2IX/W0230vrcybB+OjfE+9hPzIek95gwqzaK6K4dzQfkTo +w/ap4+cAwoXtGJbvQnHvVSPmL+MlfKhGiKOyxk1gehTXz5MgjrW6G31DyXy3Zk3v +dsOuOmCcd041/XqCIWZcsMBmWy/xBh/+/xfcIQIDAQABox4wHDAaBgNVHREEEzAR +gglsb2NhbGhvc3SHBH8AAAEwDQYJKoZIhvcNAQELBQADggEBABRdN0sNRlFcypPi +Rw8MQbDbhBf83cwNoDKdhEraA25ERL7bXhUcCoD0mcrGkPSYMn/N39hCOrXG9vxm +6KpFWGxMELca8v4LyV/p0v316azLPfZ+9UJJlfElPJPqdfMnUFgmxylt8icB3Cfi +sKFB4ExsNZL0rMJWXc+xrLOAvQQqNWVwCAJFWn0SwY/UkEI/oJai8Wm1aBRvvZ2i +js9yBKk8LJ9dzIq1NC+wCEK9h6JkeckB4ROZzWyHdviQjQrmTXxsiCetCLfvkHaA +ErtQJev2Lv7HMXaoWoP+U+mTG4oXpD1h53wGsdr8eaK3WKp9HrfdlVqCs4Cuq7z0 +3DwFSjU= +-----END CERTIFICATE----- diff --git a/integration/rust/tests/data/tls/initial_key.pem b/integration/rust/tests/data/tls/initial_key.pem new file mode 100644 index 000000000..4e483862f --- /dev/null +++ b/integration/rust/tests/data/tls/initial_key.pem @@ -0,0 +1,28 @@ +-----BEGIN PRIVATE KEY----- +MIIEvQIBADANBgkqhkiG9w0BAQEFAASCBKcwggSjAgEAAoIBAQCcJmM2CeBiYW1B +5ihYjeAoWRHZu+0lSmwjZ94EN7DpxlHa0Ni+N4abrSorQfitgTRYQuavaoC6439F +vr69mNpB1Py/vswRk9teECo9wC49FOYmbgY5p0mISt/39b7Cw49sNq7oqjnAqpH2 +qOLJJDgHQPdSAIQgL+FiO/dccyxgKMa9Fn+x9gDkFKUfGVDqh1wtvYhf9bTbfS+t +zJsH46N8T72E/Mh6T3mDCrNororh3NB+ROjD9qnj5wDChe0Ylu9Cce9VI+Yv4yV8 +qEaIo7LGTWB6FNfPkyCOtbobfUPJfLdmTe92w646YJx3TjX9eoIhZlywwGZbL/EG +H/7/F9whAgMBAAECggEAV+W39SRMHbUQBodjcK20X6H7zV/e1x30j12ZeTBMMtwD +GbR0PWcOK7WnRiBltm1DpOdL6bR+8DS9YOpFfn57ZZFaESl6v+5GDsX0sTvsC1An +WbyXXn7Pgpv7RR4dGo9wvY5umOOxjMW3Umyw9F6h91tXnN5TgbbSHTT6Qh1G/n00 +WIqmwbQrYfp9WeamU8MlQ5bE+SCPHHXzACxVO2uH0GREcMSoCzEHeI/sCTY/IFkc +wKKWS85w/9auNzpTdxV1nwMbtQT2m8aDCFPVl4a4xVSddvGA+v+ohIX+lrrO8Rzb +xYQV6Nh7isnfXnnNNIuZ4lm3bNM6Rx4RhgN7Hgc5EQKBgQDPs8IsC13napZ8EHP4 +nP8Jgu62zXsAmbiHEFfZrPyAnj3ZEG6bl9c3sY7srb8aXigYoTK2HbLbSa3JsWmI +NbNI/kjNpyhHQO3tNi15QFzKDnWTwVV1ovCmhdvq4U3NRBgF7Q1askxrCUocRG6u +xhUcNcOqfmO19+eAy0lXg52cfQKBgQDAdcqpYI5AgCJ5qYEyOwy5y3+9n0rpPSHd +bgCdhZ3Fil+puiud68SbNsgVlM0yxxz9BeLUisQYq9Z7QrWvzQDGqbE9/cfpSl6I +bRIGZ5eXSMvvT33XfNswQs4IWNrIwo7pw12bA2rDEYSSZC0XdmZfpAKdDkCFfH3D +UHKYO6djdQKBgQCpfz6UBuqo8XjA4gRh/Gy8bFc2YtVgFhJaVmH6x4p/w6MhQqGg +4/bEAmhqiReNAw2hm9rwd6gAAE6Ma/V9LKWUib8L5L+f9kKz9CSD8JxIYChfXcTJ +7SCKJG7lbNu7CTi5jUv6mcp3BuutycKxagDMNqvotJ/WXepUVpERk9zJWQKBgFK9 +kT4WK7HhJHEnhUqiBkuOCEHuTJdPV9LJauxNuFFnts7SIdRHuwN7nrNggINXBMhm +kmkLq1hr786YFGIbAT1nULK0+w/5kACY24nzWUGJ41rj0tckb1slLUx7Xru2oRgw +jHqLEogAbP0+ogAXP9XYPeNlcCmzJqIkYM+/vavNAoGAPPrTBzlndW+2eoK5KVsi +FT8OSDaYbkujsNlwZLtiyY69RLu80Av2M2roDdepSaKp8kSxwxyQSAjZuqiGbljL +xqXXGFeUDkTZ95Fs/aU2gqW3InqCU8vFH0FDlfZdfW0BKJLEMwFL6+hiEYYRJdTU +fdbHERo2fiqZCfJ858jiwX8= +-----END PRIVATE KEY----- diff --git a/integration/rust/tests/data/tls/rotated_cert.pem b/integration/rust/tests/data/tls/rotated_cert.pem new file mode 100644 index 000000000..6431e92d2 --- /dev/null +++ b/integration/rust/tests/data/tls/rotated_cert.pem @@ -0,0 +1,18 @@ +-----BEGIN CERTIFICATE----- +MIIC0TCCAbmgAwIBAgIJANODytKPdkDrMA0GCSqGSIb3DQEBCwUAMBgxFjAUBgNV +BAMMDXBnZG9nLXJvdGF0ZWQwHhcNMjUwOTMwMTYxNDI1WhcNMjYwOTMwMTYxNDI1 +WjAYMRYwFAYDVQQDDA1wZ2RvZy1yb3RhdGVkMIIBIjANBgkqhkiG9w0BAQEFAAOC +AQ8AMIIBCgKCAQEA0tgdo+T7uVq7guG7oY8pjqRuS4jfSrfGFtjA+MKnaIujtmAw +IPQdI2Hr9pnyT78AYHg/2enkoBdUvjqHhW+SZv/VuatiiI2A7Xz7N7XItWZ80Si5 +eYaFSwBY1HOfENVW3XEgb37GwcuwYZQi7NweyrrSbkJZ3MZwGsn+Ktip7L8HAqJ2 +Dssj+F++dcJhpgcK20qsZ8yO8tpqv1q2384fSviWvj1ff0PB5jnL0XD9hy44xIl6 +9voo3VDPAx2tgJ9EkemJbWikEdxEChOrK5cuGSu9y3j4zJk+SVJ6hzqxwsXBh0I5 +vRvzSs5z3m9qVMHtnZ99lcNV+JZ1nci+B5KKLQIDAQABox4wHDAaBgNVHREEEzAR +gglsb2NhbGhvc3SHBH8AAAEwDQYJKoZIhvcNAQELBQADggEBAImlxlhZKL6FhsyV +sEAB5tMfpSG5ZpPOdi4i1lyefEn1B94ynEsFb0wP26DqqnTQ6aMOR7do8omIrVpq +TpQICxrJQGDAA9wArSXjRpie1TS0BAcSC8gQZ7kykaoWGoaiRqy96EBYSugFAUae +mUutN/+DSttupzwfWGo+iIPrEMcqo//Sqt4CPyBOjI1FMGBVaX3uRUNlprW1EfwP +Z8O9Y2RDeZPFUK8RcxYKE90AaEjrMm7pKWGFIS2oanvg0org2ckRf2HwuZXhV32m +CNWz7uHmK50FO6o2ukVglOSCwSwWsuo0YHbiABHn5I7J3wEMqUnkuSAjdAIDg1+R +PZ47fxg= +-----END CERTIFICATE----- diff --git a/integration/rust/tests/data/tls/rotated_key.pem b/integration/rust/tests/data/tls/rotated_key.pem new file mode 100644 index 000000000..018c2191c --- /dev/null +++ b/integration/rust/tests/data/tls/rotated_key.pem @@ -0,0 +1,28 @@ +-----BEGIN PRIVATE KEY----- +MIIEvQIBADANBgkqhkiG9w0BAQEFAASCBKcwggSjAgEAAoIBAQDS2B2j5Pu5WruC +4buhjymOpG5LiN9Kt8YW2MD4wqdoi6O2YDAg9B0jYev2mfJPvwBgeD/Z6eSgF1S+ +OoeFb5Jm/9W5q2KIjYDtfPs3tci1ZnzRKLl5hoVLAFjUc58Q1VbdcSBvfsbBy7Bh +lCLs3B7KutJuQlncxnAayf4q2KnsvwcConYOyyP4X751wmGmBwrbSqxnzI7y2mq/ +Wrbfzh9K+Ja+PV9/Q8HmOcvRcP2HLjjEiXr2+ijdUM8DHa2An0SR6YltaKQR3EQK +E6srly4ZK73LePjMmT5JUnqHOrHCxcGHQjm9G/NKznPeb2pUwe2dn32Vw1X4lnWd +yL4HkootAgMBAAECggEAXlMO146OSrrbnk7sSPeqCMVpDmO6OUwD056+ncs/Z5bo +86MOhP+QtY6OKLFwZNq3CXFiZ1Oq0y/82mmGzVw/q9KSQ9D3cM2VOympnZ+2neiu +uEe2yjYzFX2fP9RF+hrnFIQSla6qrnI4gz7pbPuAzwNLNsZ6OzmPV3y8N2DcjCuN +xhQq8VvKThaA5F5iCOX1imOMX7EnE4vgVOIATLummcWmGpo4dQd6aSUDDWdFuExU +b9EZ9/4NPbLEzVJSSlgPUAXwRN3YQ1bfSlToIVFw724/lNSq4sTkEDPQqIeFb3eb +dBGdDcwWqMIOOwGccKBpyOXZJg1bprD6rrPtUtN/BQKBgQDs/cMPHLRf1VurmmdR +yE9b/S1cHtKCz7luQD8If7leVFjPAKQib8zIZfKjAj9TZFwxvWQcosD7XOehWC9F +SNrsireysN45yBl3k0msO0d0ljjpcmmz/PTl+uYLVZmRQJp4hwE3sc0GE0Nl3Swu +9/PsuCZ6Pr2wUtiaoT2jyUuuMwKBgQDjwXc3m33qLmCX3ZI09Ih3woWmGvK9uwpj +4UqYGEF//z1wIEwNGsgIzVJF334yCZX5Z8pp9WYOSd1mSD8QT8aQjixhCXtlU7a1 +b+plAimVRJzMKm5ZWAfTKIVhruP3nCvWAC757dIWe8qYHrB1pLpIumCFStQSbKzu +MJ8mX5/GHwKBgBMi+plBzB7g76IPucAU2LOo4fzKUF1XwLVyYqShC6reTL2KY7aU +KIkWEl2vVMW7GOa7UFYvnj2t5tZUdJy3oVXwbZz0Qz2PNt88+Xn6324+oyHWp0pt +ZqkbdW/83YWpHdAVtrd0mAWhkJOtJGA2jW/T/udoIZEXX348/uk22/GZAoGAMz0b +LQ92THESmhfnBLLe4NKKbswxQC4MMFxHA+CxG7K4h7k8YtZbml9W2xFkuq0daHbJ +Ov1ScHR9sr0eMvU/ntXddhdEA4/J0xfSi9botAQzolsJaGA9omvDVi6aauJfmk2A +RAoU8an38jE1UcI1hpcnj9U90MdSQGP/6gopT9ECgYEAq/Ndb69Xszfsfhkk1Cc1 +SLdZcLVpIYuSxYpDZR1WDeaOV9Ingw/9OpykLlI3utwX7wQooujxeGAk1Uy3pT+s +9/MbeVYTmQ2flERHBfzy04lS6zt5DJz+ryzWzT6BFWsSmJTxX8pqdIJJNFEKSH14 +WP7OKXfkY2ZO+IYP4thbIuY= +-----END PRIVATE KEY----- diff --git a/integration/rust/tests/integration/auth.rs b/integration/rust/tests/integration/auth.rs index 4c2d2b0dd..8556137e8 100644 --- a/integration/rust/tests/integration/auth.rs +++ b/integration/rust/tests/integration/auth.rs @@ -1,6 +1,7 @@ use rust::setup::admin_sqlx; +use rust::utils::assert_setting_str; use serial_test::serial; -use sqlx::{Connection, Executor, PgConnection, Row}; +use sqlx::{Connection, Executor, PgConnection}; #[tokio::test] #[serial] @@ -25,23 +26,6 @@ async fn test_auth() { assert!(PgConnection::connect(bad_password).await.is_err()); } -async fn assert_setting_str(name: &str, expected: &str) { - let admin = admin_sqlx().await; - let rows = admin.fetch_all("SHOW CONFIG").await.unwrap(); - let mut found = false; - for row in rows { - let db_name: String = row.get(0); - let value: String = row.get(1); - - if name == db_name { - found = true; - assert_eq!(value, expected); - } - } - - assert!(found); -} - #[tokio::test] async fn test_passthrough_auth() { let admin = admin_sqlx().await; diff --git a/integration/rust/tests/integration/ban.rs b/integration/rust/tests/integration/ban.rs index b4d5e01e8..eed84ee58 100644 --- a/integration/rust/tests/integration/ban.rs +++ b/integration/rust/tests/integration/ban.rs @@ -70,30 +70,4 @@ async fn test_ban_unban() { } ensure_client_state("idle").await; - - for (pool, database) in conns - .into_iter() - .zip(["pgdog", "pgdog_sharded"].into_iter()) - { - for _ in 0..25 { - pool.execute("SELECT 1").await.unwrap(); - } - - ban_unban(database, true, false).await; - - for _ in 0..25 { - let err = pool.execute("CREATE TABLE test (id BIGINT)").await; - assert!(err.err().unwrap().to_string().contains("pool is banned")); - } - - ban_unban(database, false, false).await; - - let mut t = pool.begin().await.unwrap(); - t.execute("CREATE TABLE test_ban_unban (id BIGINT)") - .await - .unwrap(); - t.rollback().await.unwrap(); - - pool.close().await; - } } diff --git a/integration/rust/tests/integration/client_ids.rs b/integration/rust/tests/integration/client_ids.rs new file mode 100644 index 000000000..7402ac4a9 --- /dev/null +++ b/integration/rust/tests/integration/client_ids.rs @@ -0,0 +1,104 @@ +use futures_util::future::join_all; +use rust::setup::admin_sqlx; +use serial_test::serial; +use sqlx::{Executor, Row}; +use std::collections::HashSet; +use tokio::task::JoinHandle; +use tokio_postgres::{Client, NoTls}; + +/// Number of client connections to create for testing unique IDs. +const NUM_CLIENTS: usize = 500; + +#[tokio::test] +#[serial] +async fn test_client_ids_unique() { + // Set auth type to md5 for faster connection setup. + let admin = admin_sqlx().await; + admin.execute("SET auth_type TO 'md5'").await.unwrap(); + admin.close().await; + + // Spawn all connection attempts in parallel. + let connect_futures: Vec<_> = (0..NUM_CLIENTS) + .map(|_| async { + let (client, connection) = tokio_postgres::connect( + "host=127.0.0.1 user=pgdog dbname=pgdog password=pgdog port=6432 application_name=test_client_ids", + NoTls, + ) + .await + .unwrap(); + + let handle = tokio::spawn(async move { + if let Err(e) = connection.await { + eprintln!("connection error: {}", e); + } + }); + + (client, handle) + }) + .collect(); + + let results: Vec<(Client, JoinHandle<()>)> = join_all(connect_futures).await; + let (clients, connection_handles): (Vec<_>, Vec<_>) = results.into_iter().unzip(); + + // Connect to admin DB and run SHOW CLIENTS. + let admin = admin_sqlx().await; + let rows = admin.fetch_all("SHOW CLIENTS").await.unwrap(); + + // Collect IDs for our test clients (filter by application_name). + let mut client_ids: Vec = Vec::new(); + for row in &rows { + let application_name: String = row.get("application_name"); + if application_name == "test_client_ids" { + let id: i64 = row.get("id"); + client_ids.push(id); + } + } + + // Verify we have exactly NUM_CLIENTS test clients. + assert_eq!( + client_ids.len(), + NUM_CLIENTS, + "expected {} clients, found {}", + NUM_CLIENTS, + client_ids.len() + ); + + // Verify all client IDs are unique. + let unique_ids: HashSet = client_ids.iter().copied().collect(); + assert_eq!( + unique_ids.len(), + NUM_CLIENTS, + "expected {} unique client IDs, found {} (duplicates exist)", + NUM_CLIENTS, + unique_ids.len() + ); + + // Drop clients to close connections gracefully. + drop(clients); + + // Wait for all connection handlers to complete. + for handle in connection_handles { + let _ = handle.await; + } + + // Verify disconnected clients no longer appear in SHOW CLIENTS. + let rows = admin.fetch_all("SHOW CLIENTS").await.unwrap(); + let remaining_test_clients: Vec = rows + .iter() + .filter(|row| { + let app_name: String = row.get("application_name"); + app_name == "test_client_ids" + }) + .map(|row| row.get("id")) + .collect(); + + assert!( + remaining_test_clients.is_empty(), + "expected 0 test clients after disconnect, found {}", + remaining_test_clients.len() + ); + + // Restore auth type to scram. + admin.execute("SET auth_type TO 'scram'").await.unwrap(); + admin.close().await; +} diff --git a/integration/rust/tests/integration/explain.rs b/integration/rust/tests/integration/explain.rs new file mode 100644 index 000000000..badd412f6 --- /dev/null +++ b/integration/rust/tests/integration/explain.rs @@ -0,0 +1,92 @@ +use rust::setup::connections_tokio; +use tokio_postgres::SimpleQueryMessage; + +#[tokio::test] +async fn explain_routing_annotations_surface() -> Result<(), Box> { + let mut clients = connections_tokio().await; + let sharded = clients.swap_remove(1); + + for shard in [0, 1] { + let drop = format!( + "/* pgdog_shard: {} */ DROP TABLE IF EXISTS explain_avg_test", + shard + ); + sharded.simple_query(drop.as_str()).await.ok(); + } + + for shard in [0, 1] { + let create = format!( + "/* pgdog_shard: {} */ CREATE TABLE explain_avg_test(price DOUBLE PRECISION)", + shard + ); + sharded.simple_query(create.as_str()).await?; + } + + sharded + .simple_query( + "/* pgdog_shard: 0 */ INSERT INTO explain_avg_test(price) VALUES (10.0), (14.0)", + ) + .await?; + sharded + .simple_query( + "/* pgdog_shard: 1 */ INSERT INTO explain_avg_test(price) VALUES (18.0), (22.0)", + ) + .await?; + + let rows = sharded + .simple_query("EXPLAIN SELECT AVG(price) FROM explain_avg_test") + .await?; + + let mut plan_lines = vec![]; + for message in rows { + if let SimpleQueryMessage::Row(row) = message { + plan_lines.push(row.get(0).unwrap_or_default().to_string()); + } + } + + assert!( + plan_lines.iter().any(|line| line.contains("Aggregate")), + "missing Aggregate node in EXPLAIN output: {:?}", + plan_lines + ); + let routing_header = plan_lines + .iter() + .find(|line| line.contains("PgDog Routing:")); + assert!( + routing_header.is_some(), + "missing PgDog Routing header in EXPLAIN output: {:?}", + plan_lines + ); + let summary_line = plan_lines + .iter() + .find(|line| line.contains("Summary:")) + .cloned(); + assert!( + summary_line + .as_ref() + .map(|line| line.contains("Summary: shard=all")) + .unwrap_or(false), + "unexpected summary line: {:?} (all lines: {:?})", + summary_line, + plan_lines + ); + let broadcast_line = plan_lines + .iter() + .find(|line| line.contains("no sharding key matched")) + .cloned(); + assert!( + broadcast_line.is_some(), + "missing broadcast note in EXPLAIN output: {:?}", + plan_lines + ); + + for shard in [0, 1] { + let drop = format!( + "/* pgdog_shard: {} */ DROP TABLE IF EXISTS explain_avg_test", + shard + ); + sharded.simple_query(drop.as_str()).await.ok(); + } + + Ok(()) +} diff --git a/integration/rust/tests/integration/idle_in_transaction.rs b/integration/rust/tests/integration/idle_in_transaction.rs new file mode 100644 index 000000000..9905fee31 --- /dev/null +++ b/integration/rust/tests/integration/idle_in_transaction.rs @@ -0,0 +1,48 @@ +use std::time::Duration; + +use rust::setup::{admin_sqlx, connection_sqlx_direct, connections_sqlx}; +use sqlx::{Executor, Row}; +use tokio::time::sleep; + +#[tokio::test] +async fn test_idle_in_transaction_timeout() { + let admin = admin_sqlx().await; + admin + .execute("SET client_idle_in_transaction_timeout TO 500") + .await + .unwrap(); + + let conn_direct = connection_sqlx_direct().await; + + for conn in connections_sqlx().await { + let mut conn = conn.acquire().await.unwrap(); + + conn.execute("BEGIN").await.unwrap(); + let pid_before = conn + .fetch_one("SELECT pg_backend_pid()") + .await + .unwrap() + .get::(0); + sleep(Duration::from_millis(750)).await; + let err = conn.execute("SELECT 1").await.unwrap_err(); + assert!(err.to_string().contains("idle in transaction")); + + sleep(Duration::from_millis(500)).await; + + let (pid_after, query): (i32, String) = + sqlx::query_as("SELECT pid, query FROM pg_stat_activity WHERE pid = $1") + .bind(pid_before) + .fetch_one(&conn_direct) + .await + .unwrap(); + + assert_eq!( + pid_before, pid_after, + "expexted pooler not to cycle connection" + ); + assert_eq!(query, "ROLLBACK", "expected a rollback on the connection",); + } + + // Reset settings. + admin.execute("RELOAD").await.unwrap(); +} diff --git a/integration/rust/tests/integration/mod.rs b/integration/rust/tests/integration/mod.rs index 6ba4da7bd..a3e936056 100644 --- a/integration/rust/tests/integration/mod.rs +++ b/integration/rust/tests/integration/mod.rs @@ -1,15 +1,25 @@ pub mod auth; pub mod avg; pub mod ban; +pub mod client_ids; pub mod cross_shard_disabled; pub mod distinct; +pub mod explain; pub mod fake_transactions; +pub mod idle_in_transaction; pub mod maintenance_mode; pub mod notify; pub mod per_stmt_routing; pub mod prepared; pub mod reload; +pub mod rewrite; +pub mod savepoint; +pub mod set_in_transaction; pub mod set_sharding_key; pub mod shard_consistency; +pub mod stddev; pub mod syntax_error; pub mod timestamp_sorting; +pub mod tls_enforced; +pub mod tls_reload; +pub mod transaction_state; diff --git a/integration/rust/tests/integration/prepared.rs b/integration/rust/tests/integration/prepared.rs index 4c03eb964..d27908464 100644 --- a/integration/rust/tests/integration/prepared.rs +++ b/integration/rust/tests/integration/prepared.rs @@ -1,6 +1,8 @@ +use std::time::Duration; + use rust::setup::*; use sqlx::{Executor, Row, postgres::PgPoolOptions, types::BigDecimal}; -use tokio::spawn; +use tokio::{spawn, time::sleep}; #[tokio::test] async fn test_prepared_cache() { @@ -79,6 +81,136 @@ async fn test_prepared_cache() { }); } +#[tokio::test] +async fn test_prepared_cache_respects_limit() { + let admin = admin_sqlx().await; + + // Start from a clean state so results aren't influenced by previous tests. + admin.execute("RECONNECT").await.unwrap(); + admin + .execute("SET prepared_statements_limit TO 100") + .await + .unwrap(); + + // Run the average helper query a few times to populate the cache. + for _ in 0..3 { + let pools = connections_sqlx().await; + for pool in &pools { + sqlx::query("/* test_prepared_cache_rust */ SELECT $1") + .bind(5) + .fetch_one(pool) + .await + .unwrap(); + } + + for pool in pools { + pool.close().await; + } + } + + let mut prepared = admin.fetch_all("SHOW PREPARED").await.unwrap(); + prepared.retain(|row| { + row.get::("statement") + .contains("/* test_prepared_cache_rust") + }); + assert_eq!( + prepared.len(), + 2, + "expected the original statement plus its helper" + ); + + // Tighten the cache limit and ensure unused statements are evicted. + admin + .execute("SET prepared_statements_limit TO 1") + .await + .unwrap(); + + let mut prepared = admin.fetch_all("SHOW PREPARED").await.unwrap(); + prepared.retain(|row| { + row.get::("statement") + .contains("/* test_prepared_cache_rust") + }); + assert!( + prepared.len() <= 1, + "expected helper statements to be evicted when limit drops" + ); + + admin.execute("RELOAD").await.unwrap(); +} + +#[tokio::test] +async fn test_prepared_cache_helper_not_evicted_on_close() { + let admin = admin_sqlx().await; + admin.execute("RECONNECT").await.unwrap(); + admin + .execute("SET prepared_statements_limit TO 100") + .await + .unwrap(); + + let mut pools = connections_sqlx().await; + let sharded = pools.remove(1); + let primary = pools.remove(0); + + // Build a deterministic dataset per shard to trigger AVG rewrite. + for shard in [0, 1] { + let drop = format!( + "/* pgdog_shard: {} */ DROP TABLE IF EXISTS avg_helper_cleanup", + shard + ); + sharded.execute(drop.as_str()).await.ok(); + + let create = format!( + "/* pgdog_shard: {} */ CREATE TABLE avg_helper_cleanup(price DOUBLE PRECISION)", + shard + ); + sharded.execute(create.as_str()).await.unwrap(); + } + + sharded + .execute("/* pgdog_shard: 0 */ INSERT INTO avg_helper_cleanup(price) VALUES (10.0), (14.0)") + .await + .unwrap(); + sharded + .execute("/* pgdog_shard: 1 */ INSERT INTO avg_helper_cleanup(price) VALUES (18.0), (22.0)") + .await + .unwrap(); + + sqlx::query("/* test_avg_helper_cleanup */ SELECT AVG(price) FROM avg_helper_cleanup") + .fetch_one(&sharded) + .await + .unwrap(); + + // Clean up tables so subsequent tests start fresh. + for shard in [0, 1] { + let drop = format!( + "/* pgdog_shard: {} */ DROP TABLE IF EXISTS avg_helper_cleanup", + shard + ); + sharded.execute(drop.as_str()).await.ok(); + } + + sharded.close().await; + primary.close().await; + + let prepared = admin + .fetch_all("SHOW PREPARED") + .await + .unwrap() + .into_iter() + .filter(|row| { + row.get::("statement") + .contains("test_avg_helper_cleanup") + }) + .collect::>(); + + assert!( + !prepared.is_empty(), + "helper rewrite statements should be not evicted once the connection closes" + ); + + admin.execute("RELOAD").await.unwrap(); +} + #[tokio::test] async fn test_prepard_cache_eviction() { let admin = admin_sqlx().await; @@ -121,6 +253,9 @@ async fn test_prepard_cache_eviction() { conn.close().await; } + // Let maintenance clean up. + sleep(Duration::from_secs(2)).await; + // Evicted only when clients disconnect or manually close the statement. let prepared = admin.fetch_all("SHOW PREPARED").await.unwrap(); assert_eq!(prepared.len(), 2); diff --git a/integration/rust/tests/integration/rewrite.rs b/integration/rust/tests/integration/rewrite.rs new file mode 100644 index 000000000..1141a9dd9 --- /dev/null +++ b/integration/rust/tests/integration/rewrite.rs @@ -0,0 +1,353 @@ +use rust::setup::{admin_sqlx, connections_sqlx}; +use sqlx::{Executor, Pool, Postgres}; + +const TEST_TABLE: &str = "sharded_list"; +const SHARDED_INSERT_TABLE: &str = "sharded"; +const NON_SHARDED_TABLE: &str = "shard_key_rewrite_non_sharded"; + +struct RewriteConfigGuard { + admin: Pool, +} + +impl RewriteConfigGuard { + async fn enable(admin: Pool) -> Self { + admin + .execute("SET two_phase_commit TO false") + .await + .expect("disable two_phase_commit"); + admin + .execute("SET rewrite_enabled TO true") + .await + .expect("enable rewrite features"); + admin + .execute("SET rewrite_shard_key_updates TO rewrite") + .await + .expect("enable shard key rewrite"); + Self { admin } + } +} + +impl Drop for RewriteConfigGuard { + fn drop(&mut self) { + let admin = self.admin.clone(); + tokio::spawn(async move { + let _ = admin + .execute("SET rewrite_shard_key_updates TO ignore") + .await; + let _ = admin.execute("SET rewrite_split_inserts TO error").await; + let _ = admin.execute("SET rewrite_enabled TO false").await; + let _ = admin.execute("SET two_phase_commit TO false").await; + }); + } +} + +#[tokio::test] +#[ignore] +async fn sharded_multi_row_insert_rejected() { + let admin = admin_sqlx().await; + let _guard = RewriteConfigGuard::enable(admin.clone()).await; + + let mut pools = connections_sqlx().await; + let sharded = pools.swap_remove(1); + + let err = sqlx::query(&format!( + "INSERT INTO {SHARDED_INSERT_TABLE} (id, value) VALUES (1, 'one'), (2, 'two')" + )) + .execute(&sharded) + .await + .expect_err("multi-row insert into sharded table should fail for now"); + + let db_err = err + .as_database_error() + .expect("expected database error from proxy"); + assert_eq!( + db_err.message(), + format!( + "multi-row INSERT into sharded table \"{SHARDED_INSERT_TABLE}\" is not supported when rewrite.split_inserts=error" + ) + ); +} + +#[tokio::test] +async fn non_sharded_multi_row_insert_allowed() { + let admin = admin_sqlx().await; + let _guard = RewriteConfigGuard::enable(admin.clone()).await; + + let mut pools = connections_sqlx().await; + let primary = pools.swap_remove(0); + + prepare_non_sharded_table(&primary).await; + + sqlx::query(&format!( + "INSERT INTO {NON_SHARDED_TABLE} (id, name) VALUES (1, 'one'), (2, 'two')" + )) + .execute(&primary) + .await + .expect("multi-row insert should succeed for non-sharded table"); + + let rows: Vec<(i64, String)> = sqlx::query_as(&format!( + "SELECT id, name FROM {NON_SHARDED_TABLE} ORDER BY id" + )) + .fetch_all(&primary) + .await + .expect("read back non-sharded rows"); + assert_eq!(rows.len(), 2, "expected two rows in non-sharded table"); + assert_eq!(rows[0], (1, "one".to_string())); + assert_eq!(rows[1], (2, "two".to_string())); + + cleanup_non_sharded_table(&primary).await; +} + +#[tokio::test] +async fn split_inserts_rewrite_moves_rows_across_shards() { + let admin = admin_sqlx().await; + let _guard = RewriteConfigGuard::enable(admin.clone()).await; + + admin + .execute("SET rewrite_split_inserts TO rewrite") + .await + .expect("enable split insert rewrite"); + + let mut pools = connections_sqlx().await; + let pool = pools.swap_remove(1); + + prepare_split_table(&pool).await; + + sqlx::query(&format!( + "INSERT INTO {SHARDED_INSERT_TABLE} (id, value) VALUES (1, 'one'), (11, 'eleven')" + )) + .execute(&pool) + .await + .expect("split insert should succeed"); + + let shard0: Option = sqlx::query_scalar(&format!( + "SELECT value FROM {SHARDED_INSERT_TABLE} WHERE id = 1" + )) + .fetch_optional(&pool) + .await + .expect("fetch shard 0 row"); + let shard1: Option = sqlx::query_scalar(&format!( + "SELECT value FROM {SHARDED_INSERT_TABLE} WHERE id = 11" + )) + .fetch_optional(&pool) + .await + .expect("fetch shard 1 row"); + + assert_eq!(shard0.as_deref(), Some("one"), "expected row on shard 0"); + assert_eq!(shard1.as_deref(), Some("eleven"), "expected row on shard 1"); + + cleanup_split_table(&pool).await; +} + +#[tokio::test] +async fn update_moves_row_between_shards() { + let admin = admin_sqlx().await; + let _guard = RewriteConfigGuard::enable(admin.clone()).await; + + let mut pools = connections_sqlx().await; + let pool = pools.swap_remove(1); + + prepare_table(&pool).await; + + let insert = format!("INSERT INTO {TEST_TABLE} (id, value) VALUES (1, 'old')"); + pool.execute(insert.as_str()) + .await + .expect("insert initial row"); + + assert_eq!(count_on_shard(&pool, 0, 1).await, 1, "row on shard 0"); + assert_eq!(count_on_shard(&pool, 1, 1).await, 0, "no row on shard 1"); + + let mut txn = pool.begin().await.unwrap(); + let update = format!("UPDATE {TEST_TABLE} SET id = 11 WHERE id = 1"); + let result = txn.execute(update.as_str()).await.expect("rewrite update"); + assert_eq!(result.rows_affected(), 1, "exactly one row updated"); + txn.commit().await.unwrap(); + + assert_eq!( + count_on_shard(&pool, 0, 1).await, + 0, + "row removed from shard 0" + ); + assert_eq!( + count_on_shard(&pool, 1, 11).await, + 1, + "row inserted on shard 1" + ); + + cleanup_table(&pool).await; +} + +#[tokio::test] +async fn update_rejects_multiple_rows() { + let admin = admin_sqlx().await; + let _guard = RewriteConfigGuard::enable(admin.clone()).await; + + let mut pools = connections_sqlx().await; + let pool = pools.swap_remove(1); + + prepare_table(&pool).await; + + let insert_first = format!("INSERT INTO {TEST_TABLE} (id, value) VALUES (1, 'old')"); + pool.execute(insert_first.as_str()) + .await + .expect("insert first row"); + + let insert_second = format!("INSERT INTO {TEST_TABLE} (id, value) VALUES (2, 'older')"); + pool.execute(insert_second.as_str()) + .await + .expect("insert second row"); + + let mut txn = pool.begin().await.unwrap(); + + let update = format!("UPDATE {TEST_TABLE} SET id = 11 WHERE id IN (1, 2)"); + let err = txn + .execute(update.as_str()) + .await + .expect_err("expected multi-row rewrite to fail"); + let db_err = err + .as_database_error() + .expect("expected database error from proxy"); + assert!( + db_err + .message() + .contains("sharding key update changes more than one row (2)"), + "unexpected error message: {}", + db_err.message() + ); + txn.rollback().await.unwrap(); + + assert_eq!( + count_on_shard(&pool, 0, 1).await, + 1, + "row 1 still on shard 0" + ); + assert_eq!( + count_on_shard(&pool, 0, 2).await, + 1, + "row 2 still on shard 0" + ); + assert_eq!( + count_on_shard(&pool, 1, 11).await, + 0, + "no row inserted on shard 1" + ); + + cleanup_table(&pool).await; +} + +#[tokio::test] +async fn update_expects_transactions() { + let admin = admin_sqlx().await; + let _guard = RewriteConfigGuard::enable(admin.clone()).await; + + let mut pools = connections_sqlx().await; + let pool = pools.swap_remove(1); + + prepare_table(&pool).await; + + let insert = format!("INSERT INTO {TEST_TABLE} (id, value) VALUES (1, 'old')"); + pool.execute(insert.as_str()) + .await + .expect("insert initial row"); + + let mut conn = pool.acquire().await.expect("acquire connection"); + + let update = format!("UPDATE {TEST_TABLE} SET id = 11 WHERE id = 1"); + let err = conn + .execute(update.as_str()) + .await + .expect_err("sharding key update must be executed inside a transaction"); + let db_err = err + .as_database_error() + .expect("expected database error from proxy"); + assert!( + db_err + .message() + .contains("sharding key update must be executed inside a transaction"), + "unexpected error message: {}", + db_err.message() + ); + + drop(conn); + + assert_eq!(count_on_shard(&pool, 0, 1).await, 1, "row still on shard 0"); + assert_eq!( + count_on_shard(&pool, 1, 11).await, + 0, + "no row inserted on shard 1" + ); + + cleanup_table(&pool).await; +} + +async fn prepare_table(pool: &Pool) { + for shard in [0, 1] { + let drop = format!("/* pgdog_shard: {shard} */ DROP TABLE IF EXISTS {TEST_TABLE}"); + pool.execute(drop.as_str()).await.unwrap(); + let create = format!( + "/* pgdog_shard: {shard} */ CREATE TABLE {TEST_TABLE} (id BIGINT PRIMARY KEY, value TEXT)" + ); + pool.execute(create.as_str()).await.unwrap(); + } +} + +async fn cleanup_table(pool: &Pool) { + for shard in [0, 1] { + let drop = format!("/* pgdog_shard: {shard} */ DROP TABLE IF EXISTS {TEST_TABLE}"); + pool.execute(drop.as_str()).await.ok(); + } +} + +async fn prepare_split_table(pool: &Pool) { + for shard in [0, 1] { + let drop = + format!("/* pgdog_shard: {shard} */ DROP TABLE IF EXISTS {SHARDED_INSERT_TABLE}"); + pool.execute(drop.as_str()).await.unwrap(); + let create = format!( + "/* pgdog_shard: {shard} */ CREATE TABLE {SHARDED_INSERT_TABLE} (id BIGINT PRIMARY KEY, value TEXT)" + ); + pool.execute(create.as_str()).await.unwrap(); + } +} + +async fn cleanup_split_table(pool: &Pool) { + for shard in [0, 1] { + let drop = + format!("/* pgdog_shard: {shard} */ DROP TABLE IF EXISTS {SHARDED_INSERT_TABLE}"); + pool.execute(drop.as_str()).await.ok(); + } +} + +async fn count_on_shard(pool: &Pool, shard: i32, id: i64) -> i64 { + let sql = format!( + "/* pgdog_shard: {shard} */ SELECT COUNT(*)::bigint FROM {TEST_TABLE} WHERE id = {id}" + ); + sqlx::query_scalar(sql.as_str()) + .fetch_one(pool) + .await + .unwrap() +} + +async fn prepare_non_sharded_table(pool: &Pool) { + let _ = sqlx::query(&format!("DROP TABLE IF EXISTS {NON_SHARDED_TABLE}")) + .execute(pool) + .await; + + sqlx::query(&format!( + "CREATE TABLE IF NOT EXISTS {NON_SHARDED_TABLE} (id BIGINT PRIMARY KEY, name TEXT)" + )) + .execute(pool) + .await + .expect("create non-sharded table"); + + sqlx::query(&format!("TRUNCATE TABLE {NON_SHARDED_TABLE}")) + .execute(pool) + .await + .expect("truncate non-sharded table"); +} + +async fn cleanup_non_sharded_table(pool: &Pool) { + let _ = sqlx::query(&format!("DROP TABLE IF EXISTS {NON_SHARDED_TABLE}")) + .execute(pool) + .await; +} diff --git a/integration/rust/tests/integration/savepoint.rs b/integration/rust/tests/integration/savepoint.rs new file mode 100644 index 000000000..13973f297 --- /dev/null +++ b/integration/rust/tests/integration/savepoint.rs @@ -0,0 +1,26 @@ +use rust::setup::connections_sqlx; +use sqlx::Executor; + +#[tokio::test] +async fn test_savepoint() { + let conns = connections_sqlx().await; + + for conn in conns { + let mut transaction = conn.begin().await.unwrap(); + transaction + .execute("CREATE TABLE test_savepoint (id BIGINT)") + .await + .unwrap(); + transaction.execute("SAVEPOINT test").await.unwrap(); + assert!(transaction.execute("SELECT sdfsf").await.is_err()); + transaction + .execute("ROLLBACK TO SAVEPOINT test") + .await + .unwrap(); + transaction + .execute("SELECT * FROM test_savepoint") + .await + .unwrap(); + transaction.rollback().await.unwrap(); + } +} diff --git a/integration/rust/tests/integration/set_in_transaction.rs b/integration/rust/tests/integration/set_in_transaction.rs new file mode 100644 index 000000000..7a327cd36 --- /dev/null +++ b/integration/rust/tests/integration/set_in_transaction.rs @@ -0,0 +1,234 @@ +use rust::setup::{admin_sqlx, connections_sqlx}; +use sqlx::Executor; + +async fn run_set_in_transaction_reset_after_commit() { + let pools = connections_sqlx().await; + let sharded = &pools[1]; + + let mut conn = sharded.acquire().await.unwrap(); + + // Get the original lock_timeout before any transaction + let original_timeout: String = sqlx::query_scalar("SHOW lock_timeout") + .fetch_one(&mut *conn) + .await + .unwrap(); + + // Make sure we set it to something different + let new_timeout = if original_timeout == "45s" { + "30s" + } else { + "45s" + }; + + // Start a transaction and change lock_timeout + conn.execute("BEGIN").await.unwrap(); + conn.execute(format!("SET lock_timeout TO '{}'", new_timeout).as_str()) + .await + .unwrap(); + + // Verify lock_timeout is set inside transaction + let timeout_in_tx: String = sqlx::query_scalar("SHOW lock_timeout") + .fetch_one(&mut *conn) + .await + .unwrap(); + assert_eq!( + timeout_in_tx, new_timeout, + "lock_timeout should be {} inside transaction", + new_timeout + ); + + conn.execute("SELECT 1").await.unwrap(); + + conn.execute("SET statement_timeout TO '1234s'") + .await + .unwrap(); + + let statement_timeout: String = sqlx::query_scalar("SHOW statement_timeout") + .fetch_one(&mut *conn) + .await + .unwrap(); + assert_eq!( + statement_timeout, "1234s", + "statement_timeout should be 1234s inside transaction", + ); + + conn.execute("COMMIT").await.unwrap(); + + // Verify lock_timeout is preserved after commit + let timeout_after_commit: String = sqlx::query_scalar("SHOW lock_timeout") + .fetch_one(&mut *conn) + .await + .unwrap(); + assert_ne!( + timeout_after_commit, original_timeout, + "lock_timeout should be preserved after commit" + ); + assert_eq!( + timeout_after_commit, new_timeout, + "lock_timeout should be preserved after commit" + ); + + let statement_timeout: String = sqlx::query_scalar("SHOW statement_timeout") + .fetch_one(&mut *conn) + .await + .unwrap(); + assert_eq!( + statement_timeout, "1234s", + "statement_timeout should be captured after query inside transaction", + ); +} + +#[tokio::test] +async fn test_set_in_transaction_reset_after_commit() { + let mut handles = Vec::new(); + for _ in 0..10 { + handles.push(tokio::spawn(run_set_in_transaction_reset_after_commit())); + } + for handle in handles { + handle.await.unwrap(); + } +} + +async fn run_set_in_transaction_reset_after_rollback() { + let pools = connections_sqlx().await; + let sharded = &pools[1]; + + let mut conn = sharded.acquire().await.unwrap(); + + // Get the original statement_timeout before any transaction + let original_timeout: String = sqlx::query_scalar("SHOW statement_timeout") + .fetch_one(&mut *conn) + .await + .unwrap(); + + // Make sure we set it to something different + let new_timeout = if original_timeout == "30s" { + "45s" + } else { + "30s" + }; + + // Start a transaction and change statement_timeout + conn.execute("BEGIN").await.unwrap(); + conn.execute(format!("SET statement_timeout TO '{}'", new_timeout).as_str()) + .await + .unwrap(); + + // Verify statement_timeout is set inside transaction + let timeout_in_tx: String = sqlx::query_scalar("SHOW statement_timeout") + .fetch_one(&mut *conn) + .await + .unwrap(); + assert_eq!( + timeout_in_tx, new_timeout, + "statement_timeout should be {} inside transaction", + new_timeout + ); + + conn.execute("ROLLBACK").await.unwrap(); + + // Verify statement_timeout is back to original after rollback + let timeout_after_rollback: String = sqlx::query_scalar("SHOW statement_timeout") + .fetch_one(&mut *conn) + .await + .unwrap(); + assert_eq!( + timeout_after_rollback, original_timeout, + "statement_timeout should be reset to original after rollback" + ); +} + +#[tokio::test] +async fn test_set_in_transaction_reset_after_rollback() { + let admin = admin_sqlx().await; + admin + .execute("SET cross_shard_disabled TO true") + .await + .unwrap(); + + let mut handles = Vec::new(); + for _ in 0..10 { + handles.push(tokio::spawn(run_set_in_transaction_reset_after_rollback())); + } + for handle in handles { + handle.await.unwrap(); + } + + admin + .execute("SET cross_shard_disabled TO false") + .await + .unwrap(); +} + +async fn run_set_local_in_transaction_reset_after_commit() { + let pools = connections_sqlx().await; + let sharded = &pools[1]; + + let mut conn = sharded.acquire().await.unwrap(); + + // Get the original work_mem before any transaction + let original_work_mem: String = sqlx::query_scalar("SHOW work_mem") + .fetch_one(&mut *conn) + .await + .unwrap(); + + // Make sure we set it to something different + let new_work_mem = if original_work_mem == "8MB" { + "16MB" + } else { + "8MB" + }; + + // Start a transaction and change work_mem using SET LOCAL + conn.execute("BEGIN").await.unwrap(); + conn.execute(format!("SET LOCAL work_mem TO '{}'", new_work_mem).as_str()) + .await + .unwrap(); + + // Verify work_mem is set inside transaction + let work_mem_in_tx: String = sqlx::query_scalar("SHOW work_mem") + .fetch_one(&mut *conn) + .await + .unwrap(); + assert_eq!( + work_mem_in_tx, new_work_mem, + "work_mem should be {} inside transaction", + new_work_mem + ); + + conn.execute("COMMIT").await.unwrap(); + + // Verify work_mem is reset to original after commit (SET LOCAL is transaction-scoped) + let work_mem_after_commit: String = sqlx::query_scalar("SHOW work_mem") + .fetch_one(&mut *conn) + .await + .unwrap(); + assert_eq!( + work_mem_after_commit, original_work_mem, + "work_mem should be reset to original after commit (SET LOCAL is transaction-scoped)" + ); +} + +#[tokio::test] +async fn test_set_local_in_transaction_reset_after_commit() { + let admin = admin_sqlx().await; + admin + .execute("SET cross_shard_disabled TO true") + .await + .unwrap(); + + let mut handles = Vec::new(); + for _ in 0..10 { + handles.push(tokio::spawn( + run_set_local_in_transaction_reset_after_commit(), + )); + } + for handle in handles { + handle.await.unwrap(); + } + + admin + .execute("SET cross_shard_disabled TO false") + .await + .unwrap(); +} diff --git a/integration/rust/tests/integration/set_sharding_key.rs b/integration/rust/tests/integration/set_sharding_key.rs index d645ee347..915c21a6c 100644 --- a/integration/rust/tests/integration/set_sharding_key.rs +++ b/integration/rust/tests/integration/set_sharding_key.rs @@ -53,12 +53,10 @@ async fn test_single_sharding_function() { #[tokio::test] async fn test_single_sharding_function_rejected() { - let mut conns = connections_sqlx().await; - { - let sharded = &mut conns[1]; - - sharded.execute("BEGIN").await.unwrap(); - let result = sharded.execute("SET pgdog.sharding_key TO '1'").await; + for conn in connections_sqlx().await { + conn.execute("BEGIN").await.unwrap(); + conn.execute("SET pgdog.sharding_key TO '1'").await.unwrap(); + let result = conn.execute("SELECT 1").await; assert!( result .err() @@ -67,29 +65,21 @@ async fn test_single_sharding_function_rejected() { .contains("config has more than one sharding function") ); } - - { - let normal = &mut conns[0]; - normal.execute("BEGIN").await.unwrap(); - let _ = normal - .execute("SET pgdog.sharding_key TO '1'") - .await - .unwrap(); - } } #[tokio::test] -async fn test_set_require_begin() { +async fn test_set_no_require_begin() { let conns = connections_sqlx().await; for conn in conns { for query in ["SET pgdog.sharding_key TO '1'", "SET pgdog.shard TO 0"] { - let result = conn.execute(query).await.err().unwrap(); - assert!( - result - .to_string() - .contains("this command requires a transaction") - ); + match conn.execute(query).await { + Ok(_) => (), + Err(err) => assert!( + err.to_string() + .contains("config has more than one sharding function") + ), + } } } } diff --git a/integration/rust/tests/integration/stddev.rs b/integration/rust/tests/integration/stddev.rs new file mode 100644 index 000000000..9263b93d6 --- /dev/null +++ b/integration/rust/tests/integration/stddev.rs @@ -0,0 +1,566 @@ +use std::collections::BTreeSet; + +use ordered_float::OrderedFloat; +use rust::setup::{connections_sqlx, connections_tokio}; +use sqlx::{Connection, Executor, PgConnection, Row, postgres::PgPool}; + +const SHARD_URLS: [&str; 2] = [ + "postgres://pgdog:pgdog@127.0.0.1:5432/shard_0", + "postgres://pgdog:pgdog@127.0.0.1:5432/shard_1", +]; + +struct Moments { + sum: f64, + sum_sq: f64, + count: i64, +} + +impl Moments { + fn variance_population(&self) -> Option { + if self.count == 0 { + return None; + } + + let n = self.count as f64; + let numerator = self.sum_sq - (self.sum * self.sum) / n; + let variance = numerator / n; + Some(if variance < 0.0 { 0.0 } else { variance }) + } + + fn variance_sample(&self) -> Option { + if self.count <= 1 { + return None; + } + + let n = self.count as f64; + let numerator = self.sum_sq - (self.sum * self.sum) / n; + let variance = numerator / (n - 1.0); + Some(if variance < 0.0 { 0.0 } else { variance }) + } + + fn stddev_population(&self) -> Option { + self.variance_population().map(|variance| variance.sqrt()) + } + + fn stddev_sample(&self) -> Option { + self.variance_sample().map(|variance| variance.sqrt()) + } +} + +#[tokio::test] +async fn stddev_pop_merges_with_helpers() -> Result<(), Box> { + let conns = connections_sqlx().await; + let sharded = conns.get(1).cloned().unwrap(); + + reset_table( + &sharded, + "stddev_pop_reduce_test", + "(price DOUBLE PRECISION)", + ) + .await?; + + seed_stat_data( + &sharded, + &[ + ( + 0, + "INSERT INTO stddev_pop_reduce_test(price) VALUES (10.0), (14.0)", + ), + ( + 1, + "INSERT INTO stddev_pop_reduce_test(price) VALUES (18.0), (22.0)", + ), + ], + ) + .await?; + + let rows = sharded + .fetch_all("SELECT STDDEV_POP(price) AS stddev_pop FROM stddev_pop_reduce_test") + .await?; + + assert_eq!(rows.len(), 1); + let pgdog_stddev: f64 = rows[0].get("stddev_pop"); + + let stats = combined_moments("stddev_pop_reduce_test", "price").await?; + let expected = stats + .stddev_population() + .expect("population stddev should exist"); + + assert!( + (pgdog_stddev - expected).abs() < 1e-9, + "STDDEV_POP mismatch: pgdog={}, expected={}", + pgdog_stddev, + expected + ); + + cleanup_table(&sharded, "stddev_pop_reduce_test").await; + + Ok(()) +} + +#[tokio::test] +async fn stddev_samp_aliases_should_merge() -> Result<(), Box> { + let conns = connections_sqlx().await; + let sharded = conns.get(1).cloned().unwrap(); + + reset_table( + &sharded, + "stddev_sample_reduce_test", + "(price DOUBLE PRECISION)", + ) + .await?; + + seed_stat_data( + &sharded, + &[ + ( + 0, + "INSERT INTO stddev_sample_reduce_test(price) VALUES (10.0), (14.0)", + ), + ( + 1, + "INSERT INTO stddev_sample_reduce_test(price) VALUES (18.0), (22.0)", + ), + ], + ) + .await?; + + let rows = sharded + .fetch_all("SELECT STDDEV(price) AS stddev_price, STDDEV_SAMP(price) AS stddev_samp_price FROM stddev_sample_reduce_test") + .await?; + + assert_eq!(rows.len(), 1); + let stddev_alias: f64 = rows[0].get("stddev_price"); + let stddev_samp: f64 = rows[0].get("stddev_samp_price"); + + let stats = combined_moments("stddev_sample_reduce_test", "price").await?; + let expected = stats + .stddev_sample() + .expect("sample stddev should exist for n > 1"); + + for value in [stddev_alias, stddev_samp] { + assert!( + (value - expected).abs() < 1e-9, + "STDDEV_SAMP mismatch: value={}, expected={}", + value, + expected + ); + } + + cleanup_table(&sharded, "stddev_sample_reduce_test").await; + + Ok(()) +} + +#[tokio::test] +async fn variance_variants_should_merge() -> Result<(), Box> { + let conns = connections_sqlx().await; + let sharded = conns.get(1).cloned().unwrap(); + + reset_table(&sharded, "variance_reduce_test", "(price DOUBLE PRECISION)").await?; + + seed_stat_data( + &sharded, + &[ + ( + 0, + "INSERT INTO variance_reduce_test(price) VALUES (10.0), (14.0)", + ), + ( + 1, + "INSERT INTO variance_reduce_test(price) VALUES (18.0), (22.0)", + ), + ], + ) + .await?; + + let rows = sharded + .fetch_all( + "SELECT VAR_POP(price) AS variance_pop, VAR_SAMP(price) AS variance_samp, VARIANCE(price) AS variance_alias FROM variance_reduce_test", + ) + .await?; + + assert_eq!(rows.len(), 1); + let variance_pop: f64 = rows[0].get("variance_pop"); + let variance_samp: f64 = rows[0].get("variance_samp"); + let variance_alias: f64 = rows[0].get("variance_alias"); + + let stats = combined_moments("variance_reduce_test", "price").await?; + let expected_pop = stats + .variance_population() + .expect("population variance should exist for n > 0"); + let expected_samp = stats + .variance_sample() + .expect("sample variance should exist for n > 1"); + + assert!( + (variance_pop - expected_pop).abs() < 1e-9, + "VAR_POP mismatch: value={}, expected={}", + variance_pop, + expected_pop + ); + + for value in [variance_samp, variance_alias] { + assert!( + (value - expected_samp).abs() < 1e-9, + "VAR_SAMP mismatch: value={}, expected={}", + value, + expected_samp + ); + } + + cleanup_table(&sharded, "variance_reduce_test").await; + + Ok(()) +} + +#[tokio::test] +async fn stddev_multiple_columns_should_merge() -> Result<(), Box> { + let conns = connections_sqlx().await; + let sharded = conns.get(1).cloned().unwrap(); + + reset_table( + &sharded, + "stddev_multi_column", + "(price DOUBLE PRECISION, discount DOUBLE PRECISION)", + ) + .await?; + + seed_stat_data( + &sharded, + &[ + ( + 0, + "INSERT INTO stddev_multi_column(price, discount) VALUES (10.0, 1.0), (14.0, 3.0)", + ), + ( + 1, + "INSERT INTO stddev_multi_column(price, discount) VALUES (18.0, 2.0), (22.0, 6.0)", + ), + ], + ) + .await?; + + let rows = sharded + .fetch_all( + "SELECT STDDEV_POP(price) AS stddev_price_pop, STDDEV_SAMP(discount) AS stddev_discount_samp FROM stddev_multi_column", + ) + .await?; + + assert_eq!(rows.len(), 1); + let price_pop: f64 = rows[0].get("stddev_price_pop"); + let discount_samp: f64 = rows[0].get("stddev_discount_samp"); + + let price_stats = combined_moments("stddev_multi_column", "price").await?; + let discount_stats = combined_moments("stddev_multi_column", "discount").await?; + + let expected_price_pop = price_stats + .stddev_population() + .expect("population stddev should exist"); + let expected_discount_samp = discount_stats + .stddev_sample() + .expect("sample stddev should exist"); + + assert!( + (price_pop - expected_price_pop).abs() < 1e-9, + "STDDEV_POP(price) mismatch: value={}, expected={}", + price_pop, + expected_price_pop + ); + + assert!( + (discount_samp - expected_discount_samp).abs() < 1e-9, + "STDDEV_SAMP(discount) mismatch: value={}, expected={}", + discount_samp, + expected_discount_samp + ); + + cleanup_table(&sharded, "stddev_multi_column").await; + + Ok(()) +} + +#[tokio::test] +async fn stddev_prepared_statement_should_merge() -> Result<(), Box> { + let conns = connections_sqlx().await; + let sharded = conns.get(1).cloned().unwrap(); + let pg_clients = connections_tokio().await; + let tokio_client = pg_clients.into_iter().nth(1).unwrap(); + + reset_table( + &sharded, + "stddev_prepared_params", + "(price DOUBLE PRECISION)", + ) + .await?; + + seed_stat_data( + &sharded, + &[ + ( + 0, + "INSERT INTO stddev_prepared_params(price) VALUES (10.0), (14.0)", + ), + ( + 1, + "INSERT INTO stddev_prepared_params(price) VALUES (18.0), (22.0)", + ), + ], + ) + .await?; + + let statement = tokio_client + .prepare("SELECT STDDEV_SAMP(price) AS stddev_price FROM stddev_prepared_params WHERE price >= $1") + .await?; + let column_names: Vec<_> = statement + .columns() + .iter() + .map(|c| c.name().to_string()) + .collect(); + println!("tokio column names: {:?}", column_names); + let tokio_rows = tokio_client.query(&statement, &[&12.0_f64]).await?; + println!("tokio row count: {}", tokio_rows.len()); + + let stddev_rows = sqlx::query( + "SELECT STDDEV_SAMP(price) AS stddev_price FROM stddev_prepared_params WHERE price >= $1", + ) + .bind(12.0_f64) + .fetch_all(&sharded) + .await?; + + assert_eq!(stddev_rows.len(), 1); + let pgdog_stddev: f64 = stddev_rows[0].get("stddev_price"); + + let stats = + combined_moments_with_filter("stddev_prepared_params", "price", "price >= $1", 12.0_f64) + .await?; + let expected = stats + .stddev_sample() + .expect("sample stddev should exist for filtered rows"); + + assert!( + (pgdog_stddev - expected).abs() < 1e-9, + "Prepared STDDEV_SAMP mismatch: pgdog={}, expected={}", + pgdog_stddev, + expected + ); + + cleanup_table(&sharded, "stddev_prepared_params").await; + + Ok(()) +} + +#[tokio::test] +async fn stddev_distinct_should_error_until_supported() -> Result<(), Box> { + let conns = connections_sqlx().await; + let sharded = conns.get(1).cloned().unwrap(); + + reset_table( + &sharded, + "stddev_distinct_error", + "(price DOUBLE PRECISION)", + ) + .await?; + + seed_stat_data( + &sharded, + &[ + ( + 0, + "INSERT INTO stddev_distinct_error(price) VALUES (10.0), (14.0)", + ), + ( + 1, + "INSERT INTO stddev_distinct_error(price) VALUES (18.0), (22.0)", + ), + ], + ) + .await?; + + let result = sharded + .fetch_all("SELECT STDDEV_SAMP(DISTINCT price) FROM stddev_distinct_error") + .await; + assert!( + result.is_err(), + "DISTINCT STDDEV should error until supported" + ); + + cleanup_table(&sharded, "stddev_distinct_error").await; + + Ok(()) +} + +#[tokio::test] +#[ignore] +async fn stddev_distinct_future_expectation() -> Result<(), Box> { + let conns = connections_sqlx().await; + let sharded = conns.get(1).cloned().unwrap(); + + reset_table(&sharded, "stddev_distinct_test", "(price DOUBLE PRECISION)").await?; + + seed_stat_data( + &sharded, + &[ + ( + 0, + "INSERT INTO stddev_distinct_test(price) VALUES (10.0), (14.0), (14.0)", + ), + ( + 1, + "INSERT INTO stddev_distinct_test(price) VALUES (14.0), (18.0), (22.0), (22.0)", + ), + ], + ) + .await?; + + let rows = sharded + .fetch_all( + "SELECT STDDEV_SAMP(DISTINCT price) AS stddev_distinct FROM stddev_distinct_test", + ) + .await?; + + assert_eq!(rows.len(), 1); + let pgdog_stddev: f64 = rows[0].get("stddev_distinct"); + + let stats = distinct_moments("stddev_distinct_test", "price").await?; + let expected = stats + .stddev_sample() + .expect("sample stddev should exist for distinct data"); + + assert!( + (pgdog_stddev - expected).abs() < 1e-9, + "DISTINCT STDDEV future expectation mismatch: pgdog={}, expected={}", + pgdog_stddev, + expected + ); + + cleanup_table(&sharded, "stddev_distinct_test").await; + + Ok(()) +} + +async fn reset_table(pool: &PgPool, table: &str, schema: &str) -> Result<(), sqlx::Error> { + for shard in [0, 1] { + let drop_stmt = format!( + "/* pgdog_shard: {} */ DROP TABLE IF EXISTS {}", + shard, table + ); + pool.execute(drop_stmt.as_str()).await.ok(); + } + + for shard in [0, 1] { + let create_stmt = format!( + "/* pgdog_shard: {} */ CREATE TABLE {} {}", + shard, table, schema + ); + pool.execute(create_stmt.as_str()).await?; + } + + Ok(()) +} + +async fn cleanup_table(pool: &PgPool, table: &str) { + for shard in [0, 1] { + let drop_stmt = format!("/* pgdog_shard: {} */ DROP TABLE {}", shard, table); + let _ = pool.execute(drop_stmt.as_str()).await; + } +} + +async fn seed_stat_data(pool: &PgPool, inserts: &[(usize, &str)]) -> Result<(), sqlx::Error> { + for (shard, statement) in inserts { + let command = format!("/* pgdog_shard: {} */ {}", shard, statement); + pool.execute(command.as_str()).await?; + } + + Ok(()) +} + +async fn combined_moments(table: &str, column: &str) -> Result { + let mut totals = Moments { + sum: 0.0, + sum_sq: 0.0, + count: 0, + }; + + for url in SHARD_URLS { + let mut conn = PgConnection::connect(url).await?; + let query = format!( + "SELECT COALESCE(SUM({col}), 0)::float8, COALESCE(SUM(POWER({col}, 2)), 0)::float8, COUNT(*) FROM {table}", + col = column, + table = table + ); + let (sum, sum_sq, count): (f64, f64, i64) = sqlx::query_as::<_, (f64, f64, i64)>(&query) + .fetch_one(&mut conn) + .await?; + totals.sum += sum; + totals.sum_sq += sum_sq; + totals.count += count; + } + + Ok(totals) +} + +async fn combined_moments_with_filter( + table: &str, + column: &str, + predicate: &str, + value: f64, +) -> Result { + let mut totals = Moments { + sum: 0.0, + sum_sq: 0.0, + count: 0, + }; + + for url in SHARD_URLS { + let mut conn = PgConnection::connect(url).await?; + let query = format!( + "SELECT COALESCE(SUM({col}), 0)::float8, COALESCE(SUM(POWER({col}, 2)), 0)::float8, COUNT(*) FROM {table} WHERE {predicate}", + col = column, + table = table, + predicate = predicate + ); + let (sum, sum_sq, count): (f64, f64, i64) = sqlx::query_as::<_, (f64, f64, i64)>(&query) + .bind(value) + .fetch_one(&mut conn) + .await?; + totals.sum += sum; + totals.sum_sq += sum_sq; + totals.count += count; + } + + Ok(totals) +} + +async fn distinct_moments(table: &str, column: &str) -> Result { + let mut distinct_values: BTreeSet> = BTreeSet::new(); + + for url in SHARD_URLS { + let mut conn = PgConnection::connect(url).await?; + let query = format!( + "SELECT DISTINCT {col} FROM {table}", + col = column, + table = table + ); + let rows: Vec> = sqlx::query_scalar::<_, Option>(&query) + .fetch_all(&mut conn) + .await?; + for value in rows.into_iter().flatten() { + distinct_values.insert(OrderedFloat(value)); + } + } + + let mut sum = 0.0; + let mut sum_sq = 0.0; + for value in &distinct_values { + let v = value.0; + sum += v; + sum_sq += v * v; + } + + Ok(Moments { + sum, + sum_sq, + count: distinct_values.len() as i64, + }) +} diff --git a/integration/rust/tests/integration/tls_enforced.rs b/integration/rust/tests/integration/tls_enforced.rs new file mode 100644 index 000000000..058f2a266 --- /dev/null +++ b/integration/rust/tests/integration/tls_enforced.rs @@ -0,0 +1,33 @@ +use rust::setup::admin_sqlx; +use rust::utils::assert_setting_str; +use sqlx::{ConnectOptions, Connection, Executor, postgres::PgSslMode}; + +#[tokio::test] +async fn test_tls_enforced() { + let admin = admin_sqlx().await; + let conn_base: sqlx::postgres::PgConnectOptions = + "postgres://pgdog:pgdog@127.0.0.1:6432/pgdog?application_name=sqlx" + .parse() + .unwrap(); + + let opts_notls = conn_base.clone().ssl_mode(PgSslMode::Disable); + let opts_tls = conn_base.clone().ssl_mode(PgSslMode::Require); + + { + let tls = opts_tls.connect().await.unwrap(); + let notls = opts_notls.connect().await.unwrap(); + tls.close().await.unwrap(); + notls.close().await.unwrap(); + } + + admin + .execute("SET tls_client_required TO 'true'") + .await + .unwrap(); + assert_setting_str("tls_client_required", "true").await; + + opts_tls.connect().await.unwrap(); + assert!(opts_notls.connect().await.is_err()); + // Reset settings. + admin.execute("RELOAD").await.unwrap(); +} diff --git a/integration/rust/tests/integration/tls_reload.rs b/integration/rust/tests/integration/tls_reload.rs new file mode 100644 index 000000000..fa9eb44eb --- /dev/null +++ b/integration/rust/tests/integration/tls_reload.rs @@ -0,0 +1,238 @@ +use std::{ + fs, + path::{Path, PathBuf}, + sync::Arc, + time::Duration, +}; + +use rust::setup::admin_tokio; +use serial_test::serial; +use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + net::TcpStream, + time::sleep, +}; +use tokio_rustls::rustls::pki_types::pem::PemObject; +use tokio_rustls::{TlsConnector, rustls}; + +const HOST: &str = "127.0.0.1"; +const PORT: u16 = 6432; +const INITIAL_CERT_PLACEHOLDER: &str = "integration/tls/cert.pem"; +const INITIAL_KEY_PLACEHOLDER: &str = "integration/tls/key.pem"; + +fn asset_path(name: &str) -> PathBuf { + PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("tests/data/tls") + .join(name) +} + +fn rotated_cert_path() -> PathBuf { + asset_path("rotated_cert.pem") +} + +fn rotated_key_path() -> PathBuf { + asset_path("rotated_key.pem") +} + +#[tokio::test] +#[serial] +async fn tls_acceptor_swaps_after_sighup() { + let mut guard = ConfigGuard::new().expect("config guard"); + + let initial_cert = fetch_server_cert_der().await.expect("initial cert"); + let rotated_cert = load_cert_der(&rotated_cert_path()).expect("rotated cert load"); + assert_ne!( + initial_cert.as_ref(), + rotated_cert.as_ref(), + "test requires distinct certificates" + ); + + guard + .rewrite_config(&rotated_cert_path(), &rotated_key_path()) + .expect("rewrite config for rotated cert"); + + guard.signal_sighup().expect("send sighup"); + sleep(Duration::from_millis(500)).await; + + let reloaded_cert = fetch_server_cert_der().await.expect("cert after sighup"); + + assert_eq!( + reloaded_cert.as_ref(), + rotated_cert.as_ref(), + "TLS acceptor should switch to the rotated certificate after SIGHUP" + ); +} + +#[tokio::test] +#[serial] +async fn tls_acceptor_swaps_after_admin_reload() { + let mut guard = ConfigGuard::new().expect("config guard"); + + let rotated_cert = load_cert_der(&rotated_cert_path()).expect("rotated cert load"); + + guard + .rewrite_config(&rotated_cert_path(), &rotated_key_path()) + .expect("rewrite config for rotated cert"); + + let admin = admin_tokio().await; + admin.simple_query("RELOAD").await.expect("admin reload"); + sleep(Duration::from_millis(500)).await; + + let reloaded_cert = fetch_server_cert_der().await.expect("cert after RELOAD"); + + assert_eq!( + reloaded_cert.as_ref(), + rotated_cert.as_ref(), + "TLS acceptor should switch to the rotated certificate after RELOAD" + ); +} + +fn load_cert_der( + path: &Path, +) -> Result, std::io::Error> { + rustls::pki_types::CertificateDer::from_pem_file(path) + .map_err(|err| std::io::Error::new(std::io::ErrorKind::InvalidData, err)) +} + +async fn fetch_server_cert_der() +-> Result, Box> { + let mut stream = TcpStream::connect((HOST, PORT)).await?; + + let mut request = Vec::with_capacity(8); + request.extend_from_slice(&8u32.to_be_bytes()); + request.extend_from_slice(&80877103u32.to_be_bytes()); + stream.write_all(&request).await?; + stream.flush().await?; + + let mut reply = [0u8; 1]; + stream.read_exact(&mut reply).await?; + if reply[0] != b'S' { + return Err("server rejected TLS negotiation".into()); + } + + let verifier = Arc::new(AllowAllVerifier); + let config = rustls::ClientConfig::builder() + .dangerous() + .with_custom_certificate_verifier(verifier) + .with_no_client_auth(); + + let connector = TlsConnector::from(Arc::new(config)); + let tls_stream = connector + .connect( + rustls::pki_types::ServerName::try_from("localhost")?, + stream, + ) + .await?; + + let (_, connection) = tls_stream.get_ref(); + let presented = connection + .peer_certificates() + .and_then(|certs| certs.first()) + .ok_or("server did not present a certificate")?; + + Ok(presented.clone().into_owned()) +} + +struct ConfigGuard { + path: PathBuf, + original: String, + pid: libc::pid_t, +} + +impl ConfigGuard { + fn new() -> Result { + let manifest_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR")); + let config_path = manifest_dir.join("../pgdog.toml").canonicalize()?; + let original = fs::read_to_string(&config_path)?; + let pid_path = manifest_dir.join("../pgdog.pid").canonicalize()?; + let pid_contents = fs::read_to_string(pid_path)?; + let pid: libc::pid_t = pid_contents.trim().parse().map_err(|err| { + std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!("invalid pid: {err}"), + ) + })?; + + Ok(Self { + path: config_path, + original, + pid, + }) + } + + fn rewrite_config(&mut self, cert: &Path, key: &Path) -> Result<(), std::io::Error> { + let mut updated = self.original.clone(); + let cert_str = cert.to_string_lossy(); + let key_str = key.to_string_lossy(); + + updated = updated.replace(INITIAL_CERT_PLACEHOLDER, cert_str.as_ref()); + updated = updated.replace(INITIAL_KEY_PLACEHOLDER, key_str.as_ref()); + fs::write(&self.path, updated) + } + + fn signal_sighup(&self) -> Result<(), std::io::Error> { + let res = unsafe { libc::kill(self.pid, libc::SIGHUP) }; + if res != 0 { + return Err(std::io::Error::last_os_error()); + } + Ok(()) + } +} + +impl Drop for ConfigGuard { + fn drop(&mut self) { + let _ = fs::write(&self.path, &self.original); + let _ = unsafe { libc::kill(self.pid, libc::SIGHUP) }; + std::thread::sleep(Duration::from_millis(500)); + } +} + +#[derive(Debug)] +struct AllowAllVerifier; + +impl rustls::client::danger::ServerCertVerifier for AllowAllVerifier { + fn verify_server_cert( + &self, + _end_entity: &rustls::pki_types::CertificateDer<'_>, + _intermediates: &[rustls::pki_types::CertificateDer<'_>], + _server_name: &rustls::pki_types::ServerName<'_>, + _ocsp_response: &[u8], + _now: rustls::pki_types::UnixTime, + ) -> Result { + Ok(rustls::client::danger::ServerCertVerified::assertion()) + } + + fn verify_tls12_signature( + &self, + _message: &[u8], + _cert: &rustls::pki_types::CertificateDer<'_>, + _dss: &rustls::DigitallySignedStruct, + ) -> Result { + Ok(rustls::client::danger::HandshakeSignatureValid::assertion()) + } + + fn verify_tls13_signature( + &self, + _message: &[u8], + _cert: &rustls::pki_types::CertificateDer<'_>, + _dss: &rustls::DigitallySignedStruct, + ) -> Result { + Ok(rustls::client::danger::HandshakeSignatureValid::assertion()) + } + + fn supported_verify_schemes(&self) -> Vec { + vec![ + rustls::SignatureScheme::RSA_PKCS1_SHA256, + rustls::SignatureScheme::RSA_PKCS1_SHA384, + rustls::SignatureScheme::RSA_PKCS1_SHA512, + rustls::SignatureScheme::RSA_PSS_SHA256, + rustls::SignatureScheme::RSA_PSS_SHA384, + rustls::SignatureScheme::RSA_PSS_SHA512, + rustls::SignatureScheme::ECDSA_NISTP256_SHA256, + rustls::SignatureScheme::ECDSA_NISTP384_SHA384, + rustls::SignatureScheme::ECDSA_NISTP521_SHA512, + rustls::SignatureScheme::ED25519, + rustls::SignatureScheme::ED448, + ] + } +} diff --git a/integration/rust/tests/integration/transaction_state.rs b/integration/rust/tests/integration/transaction_state.rs new file mode 100644 index 000000000..9fabb05d8 --- /dev/null +++ b/integration/rust/tests/integration/transaction_state.rs @@ -0,0 +1,135 @@ +use rust::setup::admin_sqlx; +use serial_test::serial; +use sqlx::{Executor, Pool, Postgres, Row}; +use tokio::time::{Duration, Instant, sleep}; + +const APP_NAME: &str = "test_transaction_state_flow"; + +#[tokio::test] +#[serial] +async fn test_transaction_state_transitions() { + let admin = admin_sqlx().await; + assert!(fetch_client_state(&admin, APP_NAME).await.is_none()); + + let (client, connection) = tokio_postgres::connect( + "host=127.0.0.1 user=pgdog dbname=pgdog password=pgdog port=6432", + tokio_postgres::NoTls, + ) + .await + .unwrap(); + + let connection_handle = tokio::spawn(async move { + if let Err(e) = connection.await { + panic!("connection error: {}", e); + } + }); + + client + .batch_execute(&format!("SET application_name TO '{}';", APP_NAME)) + .await + .unwrap(); + client + .batch_execute("SET statement_timeout TO '10s';") + .await + .unwrap(); + client.batch_execute("SELECT 1;").await.unwrap(); + + wait_for_client_state(&admin, APP_NAME, "idle").await; + + client.batch_execute("BEGIN;").await.unwrap(); + wait_for_client_state(&admin, APP_NAME, "idle in transaction").await; + + { + let query = client.simple_query("SELECT pg_sleep(0.25);"); + tokio::pin!(query); + + let deadline = Instant::now() + Duration::from_secs(5); + let mut saw_active = false; + loop { + tokio::select! { + result = &mut query => { + result.unwrap(); + break; + } + _ = sleep(Duration::from_millis(10)) => { + if let Some(state) = fetch_client_state(&admin, APP_NAME).await { + if state == "active" { + saw_active = true; + } + } + + if Instant::now() >= deadline { + panic!("timed out waiting for client to become active"); + } + } + } + } + + assert!( + saw_active, + "client never reported active state during query" + ); + } + + wait_for_client_state(&admin, APP_NAME, "idle in transaction").await; + + client.batch_execute("COMMIT;").await.unwrap(); + wait_for_client_state(&admin, APP_NAME, "idle").await; + + drop(client); + + wait_for_no_client(&admin, APP_NAME).await; + + admin.close().await; + connection_handle.await.unwrap(); +} + +async fn fetch_client_state(admin: &Pool, application_name: &str) -> Option { + let rows = admin.fetch_all("SHOW CLIENTS").await.unwrap(); + for row in rows { + let db: String = row.get::("database"); + let app: String = row.get::("application_name"); + if db == "pgdog" && app == application_name { + return Some(row.get::("state")); + } + } + None +} + +async fn wait_for_client_state(admin: &Pool, application_name: &str, expected: &str) { + let deadline = Instant::now() + Duration::from_secs(5); + loop { + if let Some(state) = fetch_client_state(admin, application_name).await { + if state == expected { + return; + } + } + + if Instant::now() >= deadline { + panic!( + "timed out waiting for client state '{}' (expected '{}')", + fetch_client_state(admin, application_name) + .await + .unwrap_or_else(|| "".to_string()), + expected + ); + } + + sleep(Duration::from_millis(25)).await; + } +} + +async fn wait_for_no_client(admin: &Pool, application_name: &str) { + let deadline = Instant::now() + Duration::from_secs(5); + loop { + if fetch_client_state(admin, application_name).await.is_none() { + return; + } + + if Instant::now() >= deadline { + panic!("client '{}' still present", application_name); + } + + sleep(Duration::from_millis(25)).await; + } +} diff --git a/integration/rust/tests/sqlx/bad_auth.rs b/integration/rust/tests/sqlx/bad_auth.rs index 35328f421..3b78deb26 100644 --- a/integration/rust/tests/sqlx/bad_auth.rs +++ b/integration/rust/tests/sqlx/bad_auth.rs @@ -1,9 +1,12 @@ +use rust::setup::admin_sqlx; use serial_test::serial; -use sqlx::{Connection, PgConnection}; +use sqlx::{Connection, Executor, PgConnection}; #[tokio::test] #[serial] async fn test_bad_auth() { + admin_sqlx().await.execute("RELOAD").await.unwrap(); + for user in ["pgdog", "pgdog_bad_user"] { for password in ["bad_password", "another_password", ""] { for db in ["random_db", "pgdog"] { @@ -15,11 +18,14 @@ async fn test_bad_auth() { .err() .unwrap(); println!("{}", err); - assert!(err.to_string().contains(&format!( - "user \"{}\" and database \"{}\" is wrong, or the database does not exist", - user, db - ))); + let err = err.to_string(); + assert!( + err.contains("is down") + || err.contains("is wrong, or the database does not exist") + ); } } } + + admin_sqlx().await.execute("RELOAD").await.unwrap(); } diff --git a/integration/rust/tests/sqlx/mod.rs b/integration/rust/tests/sqlx/mod.rs index a67973237..9dc97955a 100644 --- a/integration/rust/tests/sqlx/mod.rs +++ b/integration/rust/tests/sqlx/mod.rs @@ -1,3 +1,4 @@ pub mod bad_auth; pub mod params; pub mod select; +pub mod unique_id; diff --git a/integration/rust/tests/sqlx/unique_id.rs b/integration/rust/tests/sqlx/unique_id.rs new file mode 100644 index 000000000..af553c673 --- /dev/null +++ b/integration/rust/tests/sqlx/unique_id.rs @@ -0,0 +1,60 @@ +use sqlx::{Postgres, pool::Pool, postgres::PgPoolOptions}; + +async fn sharded_pool() -> Pool { + PgPoolOptions::new() + .max_connections(1) + .connect("postgres://pgdog:pgdog@127.0.0.1:6432/pgdog_sharded?application_name=sqlx") + .await + .unwrap() +} + +#[tokio::test] +async fn test_unique_id() { + let conn = sharded_pool().await; + + let row: (i64,) = sqlx::query_as("SELECT pgdog.unique_id()") + .fetch_one(&conn) + .await + .unwrap(); + + assert!(row.0 > 0, "unique_id should be positive"); + + conn.close().await; +} + +#[tokio::test] +async fn test_unique_id_multiple() { + let conn = sharded_pool().await; + + let row: (i64, i64) = sqlx::query_as("SELECT pgdog.unique_id(), pgdog.unique_id()") + .fetch_one(&conn) + .await + .unwrap(); + + assert!(row.0 > 0, "first unique_id should be positive"); + assert!(row.1 > 0, "second unique_id should be positive"); + assert_ne!(row.0, row.1, "unique_ids should be different"); + + conn.close().await; +} + +#[tokio::test] +async fn test_unique_id_uniqueness() { + let conn = sharded_pool().await; + + let mut ids = Vec::new(); + + for _ in 0..100 { + let row: (i64,) = sqlx::query_as("SELECT pgdog.unique_id()") + .fetch_one(&conn) + .await + .unwrap(); + ids.push(row.0); + } + + ids.sort(); + ids.dedup(); + assert_eq!(ids.len(), 100, "all unique_ids should be unique"); + + conn.close().await; +} diff --git a/integration/schema_sharding/pgdog.toml b/integration/schema_sharding/pgdog.toml new file mode 100644 index 000000000..be97af746 --- /dev/null +++ b/integration/schema_sharding/pgdog.toml @@ -0,0 +1,38 @@ +[general] +expanded_explain = true +cross_shard_disabled = true + +[[databases]] +name = "pgdog" +host = "127.0.0.1" +shard = 0 +database_name = "shard_0" + +[[databases]] +name = "pgdog" +host = "127.0.0.1" +shard = 1 +database_name = "shard_1" + +[[sharded_schemas]] +database = "pgdog" +name = "shard_0" +shard = 0 + +[[sharded_schemas]] +database = "pgdog" +name = "shard_1" +shard = 1 + +[[sharded_schemas]] +database = "pgdog" +name = "acustomer" +shard = 0 + +# All other schemas go to shard 0. +[[sharded_schemas]] +database = "pgdog" +shard = 0 + +[admin] +password = "pgdog" diff --git a/integration/schema_sharding/users.toml b/integration/schema_sharding/users.toml new file mode 100644 index 000000000..539bb1832 --- /dev/null +++ b/integration/schema_sharding/users.toml @@ -0,0 +1,4 @@ +[[users]] +name = "pgdog" +password = "pgdog" +database = "pgdog" diff --git a/integration/schema_sync/.gitignore b/integration/schema_sync/.gitignore new file mode 100644 index 000000000..5eec98606 --- /dev/null +++ b/integration/schema_sync/.gitignore @@ -0,0 +1 @@ +.claude diff --git a/integration/schema_sync/dev.sh b/integration/schema_sync/dev.sh new file mode 100644 index 000000000..9370eb84d --- /dev/null +++ b/integration/schema_sync/dev.sh @@ -0,0 +1,52 @@ +#!/bin/bash +set -e +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd ) +PGDOG_BIN_PATH="${PGDOG_BIN:-${SCRIPT_DIR}/../../target/release/pgdog}" +pushd ${SCRIPT_DIR} + +export PGPASSWORD=pgdog +export PGUSER=pgdog +export PGHOST=127.0.0.1 +export PGPORT=5432 +psql -f ${SCRIPT_DIR}/ecommerce_schema.sql pgdog1 +psql -c 'CREATE PUBLICATION pgdog FOR ALL TABLES' pgdog1 || true + +${PGDOG_BIN_PATH} \ + schema-sync \ + --from-database source \ + --to-database destination \ + --publication pgdog + +${PGDOG_BIN_PATH} \ + schema-sync \ + --from-database source \ + --to-database destination \ + --publication pgdog \ + --data-sync-complete + +${PGDOG_BIN_PATH} \ + schema-sync \ + --from-database source \ + --to-database destination \ + --publication pgdog \ + --cutover + +pg_dump \ + --schema-only \ + --exclude-schema pgdog \ + --no-publications pgdog1 > source.sql + +pg_dump \ + --schema-only \ + --exclude-schema pgdog \ + --no-publications pgdog2 > destination.sql + +for f in source.sql destination.sql; do + sed -i '/^\\restrict.*$/d' $f + sed -i '/^\\unrestrict.*$/d' $f +done + +diff source.sql destination.sql +rm source.sql +rm destination.sql +popd diff --git a/integration/schema_sync/ecommerce_schema.sql b/integration/schema_sync/ecommerce_schema.sql new file mode 100644 index 000000000..3e4086377 --- /dev/null +++ b/integration/schema_sync/ecommerce_schema.sql @@ -0,0 +1,994 @@ +-- E-Commerce Database Schema +-- Demonstrates: Multiple schemas, partitioning, foreign keys, indexes, views, materialized views, +-- composite types, domains, enums, arrays, JSONB, full-text search, and more + +-- Enable required extensions +CREATE EXTENSION IF NOT EXISTS ltree; +CREATE EXTENSION IF NOT EXISTS btree_gist; +CREATE EXTENSION IF NOT EXISTS pg_trgm; + +-- Drop existing schemas if they exist +DROP SCHEMA IF EXISTS core CASCADE; +DROP SCHEMA IF EXISTS inventory CASCADE; +DROP SCHEMA IF EXISTS sales CASCADE; +DROP SCHEMA IF EXISTS analytics CASCADE; +DROP SCHEMA IF EXISTS audit CASCADE; + +-- Create schemas +CREATE SCHEMA core; -- Core entities (users, addresses) +CREATE SCHEMA inventory; -- Products and inventory +CREATE SCHEMA sales; -- Orders and transactions +CREATE SCHEMA analytics; -- Analytics views and materialized views +CREATE SCHEMA audit; -- Audit trails + +-- ============================================================================ +-- DOMAINS AND TYPES +-- ============================================================================ + +-- Custom domains +CREATE DOMAIN core.email AS VARCHAR(255) + CHECK (VALUE ~ '^[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}$'); + +CREATE DOMAIN core.phone AS VARCHAR(20) + CHECK (VALUE ~ '^\+?[0-9\s\-\(\)]+$'); + +CREATE DOMAIN inventory.money_positive AS NUMERIC(12,2) + CHECK (VALUE >= 0); + +CREATE DOMAIN inventory.weight_kg AS NUMERIC(8,3) + CHECK (VALUE > 0); + +-- Enums +CREATE TYPE core.user_status AS ENUM ('active', 'inactive', 'suspended', 'deleted'); +CREATE TYPE core.user_role AS ENUM ('customer', 'vendor', 'admin', 'support'); +CREATE TYPE core.address_type AS ENUM ('billing', 'shipping', 'both'); + +CREATE TYPE inventory.product_status AS ENUM ('draft', 'active', 'discontinued', 'out_of_stock'); +CREATE TYPE inventory.stock_movement_type AS ENUM ('purchase', 'sale', 'return', 'adjustment', 'damage', 'transfer'); + +CREATE TYPE sales.order_status AS ENUM ('pending', 'confirmed', 'processing', 'shipped', 'delivered', 'cancelled', 'refunded'); +CREATE TYPE sales.payment_status AS ENUM ('pending', 'authorized', 'captured', 'failed', 'refunded', 'partially_refunded'); +CREATE TYPE sales.payment_method AS ENUM ('credit_card', 'debit_card', 'paypal', 'bank_transfer', 'cryptocurrency', 'cash_on_delivery'); +CREATE TYPE sales.shipment_status AS ENUM ('pending', 'picked', 'packed', 'in_transit', 'delivered', 'returned'); + +-- Composite types +CREATE TYPE core.geo_point AS ( + latitude NUMERIC(9,6), + longitude NUMERIC(9,6) +); + +CREATE TYPE inventory.dimensions AS ( + length_cm NUMERIC(6,2), + width_cm NUMERIC(6,2), + height_cm NUMERIC(6,2) +); + +CREATE TYPE sales.money_with_currency AS ( + amount NUMERIC(12,2), + currency_code CHAR(3) +); + +-- ============================================================================ +-- CORE SCHEMA: Users, Addresses, Authentication +-- ============================================================================ + +CREATE TABLE core.users ( + user_id BIGSERIAL PRIMARY KEY, + email core.email NOT NULL UNIQUE, + username VARCHAR(50) NOT NULL UNIQUE, + password_hash VARCHAR(255) NOT NULL, + first_name VARCHAR(100), + last_name VARCHAR(100), + full_name VARCHAR(201) GENERATED ALWAYS AS ( + TRIM(COALESCE(first_name, '') || ' ' || COALESCE(last_name, '')) + ) STORED, + phone core.phone, + date_of_birth DATE, + status core.user_status NOT NULL DEFAULT 'active', + roles core.user_role[] NOT NULL DEFAULT ARRAY['customer']::core.user_role[], + preferences JSONB DEFAULT '{}', + metadata JSONB DEFAULT '{}', + last_login_at TIMESTAMPTZ, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + CONSTRAINT users_age_check CHECK (date_of_birth IS NULL OR date_of_birth < CURRENT_DATE - INTERVAL '13 years') +); + +CREATE INDEX idx_users_email ON core.users(email); +CREATE INDEX idx_users_username ON core.users(username); +CREATE INDEX idx_users_status ON core.users(status) WHERE status = 'active'; +CREATE INDEX idx_users_roles ON core.users USING GIN(roles); +CREATE INDEX idx_users_preferences ON core.users USING GIN(preferences); +CREATE INDEX idx_users_created_at ON core.users(created_at); + +CREATE TABLE core.countries ( + country_code CHAR(2) PRIMARY KEY, + country_name VARCHAR(100) NOT NULL, + iso3 CHAR(3) NOT NULL UNIQUE, + phone_prefix VARCHAR(10), + currency_code CHAR(3), + region VARCHAR(50) +); + +CREATE TABLE core.addresses ( + address_id BIGSERIAL PRIMARY KEY, + user_id BIGINT NOT NULL REFERENCES core.users(user_id) ON DELETE CASCADE, + address_type core.address_type NOT NULL, + recipient_name VARCHAR(200), + address_line1 VARCHAR(255) NOT NULL, + address_line2 VARCHAR(255), + city VARCHAR(100) NOT NULL, + state_province VARCHAR(100), + postal_code VARCHAR(20), + country_code CHAR(2) NOT NULL REFERENCES core.countries(country_code), + full_address TEXT GENERATED ALWAYS AS ( + address_line1 || + COALESCE(', ' || address_line2, '') || + ', ' || city || + COALESCE(', ' || state_province, '') || + COALESCE(' ' || postal_code, '') + ) STORED, + search_vector TSVECTOR, + coordinates core.geo_point, + is_default BOOLEAN NOT NULL DEFAULT FALSE, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW() +); + +CREATE INDEX idx_addresses_user_id ON core.addresses(user_id); +CREATE INDEX idx_addresses_country ON core.addresses(country_code); +CREATE INDEX idx_addresses_default ON core.addresses(user_id, is_default) WHERE is_default = TRUE; +CREATE INDEX idx_addresses_search ON core.addresses USING GIN(search_vector); + +-- Trigger to update search vector for addresses +CREATE OR REPLACE FUNCTION core.addresses_search_vector_update() +RETURNS TRIGGER AS $$ +BEGIN + NEW.search_vector := to_tsvector('english', + COALESCE(NEW.recipient_name, '') || ' ' || + COALESCE(NEW.address_line1, '') || ' ' || + COALESCE(NEW.address_line2, '') || ' ' || + COALESCE(NEW.city, '') || ' ' || + COALESCE(NEW.state_province, '') || ' ' || + COALESCE(NEW.postal_code, '') + ); + RETURN NEW; +END; +$$ LANGUAGE plpgsql; + +CREATE TRIGGER trigger_addresses_search_vector + BEFORE INSERT OR UPDATE OF recipient_name, address_line1, address_line2, city, state_province, postal_code + ON core.addresses + FOR EACH ROW + EXECUTE FUNCTION core.addresses_search_vector_update(); + +-- ============================================================================ +-- INVENTORY SCHEMA: Products, Categories, Stock +-- ============================================================================ + +CREATE TABLE inventory.categories ( + category_id BIGSERIAL PRIMARY KEY, + parent_category_id BIGINT REFERENCES inventory.categories(category_id) ON DELETE SET NULL, + category_name VARCHAR(100) NOT NULL, + category_slug VARCHAR(100) NOT NULL UNIQUE, + description TEXT, + level INTEGER NOT NULL DEFAULT 0, + path LTREE, + metadata JSONB DEFAULT '{}', + is_active BOOLEAN NOT NULL DEFAULT TRUE, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW() +); + +CREATE INDEX idx_categories_parent ON inventory.categories(parent_category_id); +CREATE INDEX idx_categories_slug ON inventory.categories(category_slug); +CREATE INDEX idx_categories_path ON inventory.categories USING GIST(path); + +CREATE TABLE inventory.brands ( + brand_id BIGSERIAL PRIMARY KEY, + brand_name VARCHAR(100) NOT NULL UNIQUE, + brand_slug VARCHAR(100) NOT NULL UNIQUE, + description TEXT, + logo_url VARCHAR(500), + website_url VARCHAR(500), + is_active BOOLEAN NOT NULL DEFAULT TRUE, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW() +); + +CREATE INDEX idx_brands_slug ON inventory.brands(brand_slug); + +CREATE TABLE inventory.products ( + product_id BIGSERIAL PRIMARY KEY, + sku VARCHAR(100) NOT NULL UNIQUE, + product_name VARCHAR(255) NOT NULL, + product_slug VARCHAR(255) NOT NULL, + brand_id BIGINT REFERENCES inventory.brands(brand_id) ON DELETE SET NULL, + category_id BIGINT REFERENCES inventory.categories(category_id) ON DELETE SET NULL, + description TEXT, + long_description TEXT, + status inventory.product_status NOT NULL DEFAULT 'draft', + base_price inventory.money_positive NOT NULL, + currency_code CHAR(3) NOT NULL DEFAULT 'USD', + cost_price inventory.money_positive, + weight inventory.weight_kg, + dimensions inventory.dimensions, + attributes JSONB DEFAULT '{}', + tags TEXT[], + search_vector TSVECTOR, + image_urls TEXT[], + is_digital BOOLEAN NOT NULL DEFAULT FALSE, + requires_shipping BOOLEAN NOT NULL DEFAULT TRUE, + is_taxable BOOLEAN NOT NULL DEFAULT TRUE, + tax_category VARCHAR(50), + min_order_quantity INTEGER NOT NULL DEFAULT 1, + max_order_quantity INTEGER, + profit_margin NUMERIC(5,2) GENERATED ALWAYS AS ( + CASE + WHEN cost_price IS NULL OR cost_price = 0 THEN NULL + ELSE ((base_price - cost_price) / cost_price * 100) + END + ) STORED, + total_volume_cm3 NUMERIC(12,2) GENERATED ALWAYS AS ( + CASE + WHEN dimensions IS NULL THEN NULL + ELSE (dimensions).length_cm * (dimensions).width_cm * (dimensions).height_cm + END + ) STORED, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + CONSTRAINT products_price_check CHECK (base_price >= 0 AND (cost_price IS NULL OR cost_price <= base_price)), + CONSTRAINT products_quantity_check CHECK (min_order_quantity > 0 AND (max_order_quantity IS NULL OR max_order_quantity >= min_order_quantity)) +); + +CREATE INDEX idx_products_sku ON inventory.products(sku); +CREATE INDEX idx_products_slug ON inventory.products(product_slug); +CREATE INDEX idx_products_brand ON inventory.products(brand_id); +CREATE INDEX idx_products_category ON inventory.products(category_id); +CREATE INDEX idx_products_status ON inventory.products(status) WHERE status = 'active'; +CREATE INDEX idx_products_search ON inventory.products USING GIN(search_vector); +CREATE INDEX idx_products_tags ON inventory.products USING GIN(tags); +CREATE INDEX idx_products_attributes ON inventory.products USING GIN(attributes); +CREATE INDEX idx_products_price ON inventory.products(base_price); +CREATE INDEX idx_products_profit_margin ON inventory.products(profit_margin) WHERE profit_margin IS NOT NULL; + +-- Trigger to update search vector for products +CREATE OR REPLACE FUNCTION inventory.products_search_vector_update() +RETURNS TRIGGER AS $$ +BEGIN + NEW.search_vector := + setweight(to_tsvector('english', COALESCE(NEW.product_name, '')), 'A') || + setweight(to_tsvector('english', COALESCE(NEW.description, '')), 'B') || + setweight(to_tsvector('english', COALESCE(array_to_string(NEW.tags, ' '), '')), 'C'); + RETURN NEW; +END; +$$ LANGUAGE plpgsql; + +CREATE TRIGGER trigger_products_search_vector + BEFORE INSERT OR UPDATE OF product_name, description, tags + ON inventory.products + FOR EACH ROW + EXECUTE FUNCTION inventory.products_search_vector_update(); + +-- Product variants (e.g., different sizes, colors) +CREATE TABLE inventory.product_variants ( + variant_id BIGSERIAL PRIMARY KEY, + product_id BIGINT NOT NULL REFERENCES inventory.products(product_id) ON DELETE CASCADE, + variant_sku VARCHAR(100) NOT NULL UNIQUE, + variant_name VARCHAR(255) NOT NULL, + variant_attributes JSONB NOT NULL, -- e.g., {"size": "L", "color": "red"} + price_adjustment NUMERIC(12,2) NOT NULL DEFAULT 0, + weight_adjustment NUMERIC(8,3) NOT NULL DEFAULT 0, + image_urls TEXT[], + is_active BOOLEAN NOT NULL DEFAULT TRUE, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW() +); + +CREATE INDEX idx_variants_product ON inventory.product_variants(product_id); +CREATE INDEX idx_variants_sku ON inventory.product_variants(variant_sku); +CREATE INDEX idx_variants_attributes ON inventory.product_variants USING GIN(variant_attributes); + +-- Warehouses +CREATE TABLE inventory.warehouses ( + warehouse_id BIGSERIAL PRIMARY KEY, + warehouse_code VARCHAR(20) NOT NULL UNIQUE, + warehouse_name VARCHAR(100) NOT NULL, + address_line1 VARCHAR(255) NOT NULL, + address_line2 VARCHAR(255), + city VARCHAR(100) NOT NULL, + state_province VARCHAR(100), + postal_code VARCHAR(20), + country_code CHAR(2) NOT NULL REFERENCES core.countries(country_code), + coordinates core.geo_point, + capacity_cubic_meters NUMERIC(10,2), + is_active BOOLEAN NOT NULL DEFAULT TRUE, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW() +); + +CREATE INDEX idx_warehouses_country ON inventory.warehouses(country_code); + +-- Stock levels (partitioned by warehouse) +CREATE TABLE inventory.stock_levels ( + stock_id BIGSERIAL, + warehouse_id BIGINT NOT NULL REFERENCES inventory.warehouses(warehouse_id) ON DELETE CASCADE, + product_id BIGINT REFERENCES inventory.products(product_id) ON DELETE CASCADE, + variant_id BIGINT REFERENCES inventory.product_variants(variant_id) ON DELETE CASCADE, + quantity_available INTEGER NOT NULL DEFAULT 0, + quantity_reserved INTEGER NOT NULL DEFAULT 0, + quantity_damaged INTEGER NOT NULL DEFAULT 0, + reorder_point INTEGER NOT NULL DEFAULT 10, + reorder_quantity INTEGER NOT NULL DEFAULT 100, + last_counted_at TIMESTAMPTZ, + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + PRIMARY KEY (stock_id, warehouse_id), + CONSTRAINT stock_product_or_variant CHECK ( + (product_id IS NOT NULL AND variant_id IS NULL) OR + (product_id IS NULL AND variant_id IS NOT NULL) + ), + CONSTRAINT stock_quantities_positive CHECK ( + quantity_available >= 0 AND + quantity_reserved >= 0 AND + quantity_damaged >= 0 + ) +) PARTITION BY HASH (warehouse_id); + +CREATE TABLE inventory.stock_levels_p0 PARTITION OF inventory.stock_levels + FOR VALUES WITH (MODULUS 4, REMAINDER 0); +CREATE TABLE inventory.stock_levels_p1 PARTITION OF inventory.stock_levels + FOR VALUES WITH (MODULUS 4, REMAINDER 1); +CREATE TABLE inventory.stock_levels_p2 PARTITION OF inventory.stock_levels + FOR VALUES WITH (MODULUS 4, REMAINDER 2); +CREATE TABLE inventory.stock_levels_p3 PARTITION OF inventory.stock_levels + FOR VALUES WITH (MODULUS 4, REMAINDER 3); + +CREATE INDEX idx_stock_warehouse ON inventory.stock_levels(warehouse_id); +CREATE INDEX idx_stock_product ON inventory.stock_levels(product_id) WHERE product_id IS NOT NULL; +CREATE INDEX idx_stock_variant ON inventory.stock_levels(variant_id) WHERE variant_id IS NOT NULL; +CREATE INDEX idx_stock_low ON inventory.stock_levels(warehouse_id, quantity_available) + WHERE quantity_available <= reorder_point; + +-- Stock movements (partitioned by date) +CREATE TABLE inventory.stock_movements ( + movement_id BIGSERIAL, + warehouse_id BIGINT NOT NULL REFERENCES inventory.warehouses(warehouse_id), + product_id BIGINT REFERENCES inventory.products(product_id), + variant_id BIGINT REFERENCES inventory.product_variants(variant_id), + movement_type inventory.stock_movement_type NOT NULL, + quantity INTEGER NOT NULL, + reference_type VARCHAR(50), -- e.g., 'order', 'purchase_order', 'adjustment' + reference_id BIGINT, + notes TEXT, + performed_by BIGINT REFERENCES core.users(user_id), + movement_date TIMESTAMPTZ NOT NULL DEFAULT NOW(), + PRIMARY KEY (movement_id, movement_date), + CONSTRAINT movement_product_or_variant CHECK ( + (product_id IS NOT NULL AND variant_id IS NULL) OR + (product_id IS NULL AND variant_id IS NOT NULL) + ) +) PARTITION BY RANGE (movement_date); + +-- Create partitions for last 12 months and next 12 months +CREATE TABLE inventory.stock_movements_2024_q4 PARTITION OF inventory.stock_movements + FOR VALUES FROM ('2024-10-01') TO ('2025-01-01'); +CREATE TABLE inventory.stock_movements_2025_q1 PARTITION OF inventory.stock_movements + FOR VALUES FROM ('2025-01-01') TO ('2025-04-01'); +CREATE TABLE inventory.stock_movements_2025_q2 PARTITION OF inventory.stock_movements + FOR VALUES FROM ('2025-04-01') TO ('2025-07-01'); +CREATE TABLE inventory.stock_movements_2025_q3 PARTITION OF inventory.stock_movements + FOR VALUES FROM ('2025-07-01') TO ('2025-10-01'); +CREATE TABLE inventory.stock_movements_2025_q4 PARTITION OF inventory.stock_movements + FOR VALUES FROM ('2025-10-01') TO ('2026-01-01'); +CREATE TABLE inventory.stock_movements_2026_q1 PARTITION OF inventory.stock_movements + FOR VALUES FROM ('2026-01-01') TO ('2026-04-01'); + +CREATE INDEX idx_movements_warehouse ON inventory.stock_movements(warehouse_id, movement_date); +CREATE INDEX idx_movements_product ON inventory.stock_movements(product_id, movement_date) WHERE product_id IS NOT NULL; +CREATE INDEX idx_movements_variant ON inventory.stock_movements(variant_id, movement_date) WHERE variant_id IS NOT NULL; +CREATE INDEX idx_movements_reference ON inventory.stock_movements(reference_type, reference_id); + +-- ============================================================================ +-- SALES SCHEMA: Orders, Payments, Shipments +-- ============================================================================ + +-- Shopping carts +CREATE TABLE sales.carts ( + cart_id BIGSERIAL PRIMARY KEY, + user_id BIGINT REFERENCES core.users(user_id) ON DELETE CASCADE, + session_id VARCHAR(255), -- For guest carts + currency_code CHAR(3) NOT NULL DEFAULT 'USD', + coupon_code VARCHAR(50), + discount_amount NUMERIC(12,2) NOT NULL DEFAULT 0, + tax_amount NUMERIC(12,2) NOT NULL DEFAULT 0, + shipping_amount NUMERIC(12,2) NOT NULL DEFAULT 0, + total_amount NUMERIC(12,2) NOT NULL DEFAULT 0, + metadata JSONB DEFAULT '{}', + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + expires_at TIMESTAMPTZ, + CONSTRAINT cart_user_or_session CHECK (user_id IS NOT NULL OR session_id IS NOT NULL) +); + +CREATE INDEX idx_carts_user ON sales.carts(user_id); +CREATE INDEX idx_carts_session ON sales.carts(session_id); +CREATE INDEX idx_carts_expires ON sales.carts(expires_at) WHERE expires_at IS NOT NULL; + +CREATE TABLE sales.cart_items ( + cart_item_id BIGSERIAL PRIMARY KEY, + cart_id BIGINT NOT NULL REFERENCES sales.carts(cart_id) ON DELETE CASCADE, + product_id BIGINT REFERENCES inventory.products(product_id) ON DELETE CASCADE, + variant_id BIGINT REFERENCES inventory.product_variants(variant_id) ON DELETE CASCADE, + quantity INTEGER NOT NULL, + unit_price NUMERIC(12,2) NOT NULL, + total_price NUMERIC(12,2) NOT NULL, + added_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + CONSTRAINT cart_items_product_or_variant CHECK ( + (product_id IS NOT NULL AND variant_id IS NULL) OR + (product_id IS NULL AND variant_id IS NOT NULL) + ), + CONSTRAINT cart_items_quantity_positive CHECK (quantity > 0) +); + +CREATE INDEX idx_cart_items_cart ON sales.cart_items(cart_id); + +-- Orders (partitioned by created_at) +CREATE TABLE sales.orders ( + order_id BIGSERIAL, + order_number VARCHAR(50) NOT NULL, + user_id BIGINT NOT NULL REFERENCES core.users(user_id), + status sales.order_status NOT NULL DEFAULT 'pending', + billing_address_id BIGINT REFERENCES core.addresses(address_id), + shipping_address_id BIGINT REFERENCES core.addresses(address_id), + currency_code CHAR(3) NOT NULL DEFAULT 'USD', + subtotal_amount NUMERIC(12,2) NOT NULL, + discount_amount NUMERIC(12,2) NOT NULL DEFAULT 0, + tax_amount NUMERIC(12,2) NOT NULL DEFAULT 0, + shipping_amount NUMERIC(12,2) NOT NULL DEFAULT 0, + total_amount NUMERIC(12,2) NOT NULL, + discount_percentage NUMERIC(5,2) GENERATED ALWAYS AS ( + CASE + WHEN subtotal_amount > 0 THEN (discount_amount / subtotal_amount * 100) + ELSE 0 + END + ) STORED, + effective_tax_rate NUMERIC(5,2) GENERATED ALWAYS AS ( + CASE + WHEN subtotal_amount > 0 THEN (tax_amount / subtotal_amount * 100) + ELSE 0 + END + ) STORED, + coupon_code VARCHAR(50), + notes TEXT, + ip_address INET, + user_agent TEXT, + metadata JSONB DEFAULT '{}', + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + confirmed_at TIMESTAMPTZ, + shipped_at TIMESTAMPTZ, + delivered_at TIMESTAMPTZ, + cancelled_at TIMESTAMPTZ, + PRIMARY KEY (order_id, created_at), + CONSTRAINT orders_amounts_positive CHECK ( + subtotal_amount >= 0 AND + discount_amount >= 0 AND + tax_amount >= 0 AND + shipping_amount >= 0 AND + total_amount >= 0 + ) +) PARTITION BY RANGE (created_at); + +CREATE TABLE sales.orders_2024_q4 PARTITION OF sales.orders + FOR VALUES FROM ('2024-10-01') TO ('2025-01-01'); +CREATE TABLE sales.orders_2025_q1 PARTITION OF sales.orders + FOR VALUES FROM ('2025-01-01') TO ('2025-04-01'); +CREATE TABLE sales.orders_2025_q2 PARTITION OF sales.orders + FOR VALUES FROM ('2025-04-01') TO ('2025-07-01'); +CREATE TABLE sales.orders_2025_q3 PARTITION OF sales.orders + FOR VALUES FROM ('2025-07-01') TO ('2025-10-01'); +CREATE TABLE sales.orders_2025_q4 PARTITION OF sales.orders + FOR VALUES FROM ('2025-10-01') TO ('2026-01-01'); +CREATE TABLE sales.orders_2026_q1 PARTITION OF sales.orders + FOR VALUES FROM ('2026-01-01') TO ('2026-04-01'); + +CREATE UNIQUE INDEX idx_orders_order_number ON sales.orders(order_number, created_at); +CREATE INDEX idx_orders_user ON sales.orders(user_id, created_at); +CREATE INDEX idx_orders_status ON sales.orders(status, created_at); +CREATE INDEX idx_orders_created ON sales.orders(created_at); + +CREATE TABLE sales.order_items ( + order_item_id BIGSERIAL PRIMARY KEY, + order_id BIGINT NOT NULL, + product_id BIGINT REFERENCES inventory.products(product_id), + variant_id BIGINT REFERENCES inventory.product_variants(variant_id), + product_name VARCHAR(255) NOT NULL, -- Snapshot at order time + product_sku VARCHAR(100) NOT NULL, + quantity INTEGER NOT NULL, + unit_price NUMERIC(12,2) NOT NULL, + discount_amount NUMERIC(12,2) NOT NULL DEFAULT 0, + tax_amount NUMERIC(12,2) NOT NULL DEFAULT 0, + total_price NUMERIC(12,2) NOT NULL, + attributes JSONB DEFAULT '{}', -- Snapshot of product attributes + CONSTRAINT order_items_product_or_variant CHECK ( + (product_id IS NOT NULL AND variant_id IS NULL) OR + (product_id IS NULL AND variant_id IS NOT NULL) + ), + CONSTRAINT order_items_quantity_positive CHECK (quantity > 0), + CONSTRAINT order_items_amounts_positive CHECK ( + unit_price >= 0 AND + discount_amount >= 0 AND + tax_amount >= 0 AND + total_price >= 0 + ) +); + +CREATE INDEX idx_order_items_order ON sales.order_items(order_id); +CREATE INDEX idx_order_items_product ON sales.order_items(product_id) WHERE product_id IS NOT NULL; + +-- Payments +CREATE TABLE sales.payments ( + payment_id BIGSERIAL PRIMARY KEY, + order_id BIGINT NOT NULL, + payment_method sales.payment_method NOT NULL, + payment_status sales.payment_status NOT NULL DEFAULT 'pending', + amount NUMERIC(12,2) NOT NULL, + currency_code CHAR(3) NOT NULL DEFAULT 'USD', + transaction_id VARCHAR(255), + gateway_name VARCHAR(100), + gateway_response JSONB, + card_last4 VARCHAR(4), + card_brand VARCHAR(50), + payer_email core.email, + payer_name VARCHAR(200), + metadata JSONB DEFAULT '{}', + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + authorized_at TIMESTAMPTZ, + captured_at TIMESTAMPTZ, + failed_at TIMESTAMPTZ, + refunded_at TIMESTAMPTZ, + CONSTRAINT payments_amount_positive CHECK (amount > 0) +); + +CREATE INDEX idx_payments_order ON sales.payments(order_id); +CREATE INDEX idx_payments_status ON sales.payments(payment_status); +CREATE INDEX idx_payments_transaction ON sales.payments(transaction_id); + +-- Shipments +CREATE TABLE sales.shipments ( + shipment_id BIGSERIAL PRIMARY KEY, + order_id BIGINT NOT NULL, + warehouse_id BIGINT NOT NULL REFERENCES inventory.warehouses(warehouse_id), + tracking_number VARCHAR(100), + carrier_name VARCHAR(100), + carrier_service VARCHAR(100), + status sales.shipment_status NOT NULL DEFAULT 'pending', + estimated_delivery_date DATE, + actual_delivery_date DATE, + weight_kg NUMERIC(8,3), + dimensions inventory.dimensions, + shipping_cost NUMERIC(12,2), + notes TEXT, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + picked_at TIMESTAMPTZ, + packed_at TIMESTAMPTZ, + shipped_at TIMESTAMPTZ, + delivered_at TIMESTAMPTZ +); + +CREATE INDEX idx_shipments_order ON sales.shipments(order_id); +CREATE INDEX idx_shipments_tracking ON sales.shipments(tracking_number); +CREATE INDEX idx_shipments_status ON sales.shipments(status); +CREATE INDEX idx_shipments_warehouse ON sales.shipments(warehouse_id); + +CREATE TABLE sales.shipment_items ( + shipment_item_id BIGSERIAL PRIMARY KEY, + shipment_id BIGINT NOT NULL REFERENCES sales.shipments(shipment_id) ON DELETE CASCADE, + order_item_id BIGINT NOT NULL REFERENCES sales.order_items(order_item_id), + quantity INTEGER NOT NULL, + CONSTRAINT shipment_items_quantity_positive CHECK (quantity > 0) +); + +CREATE INDEX idx_shipment_items_shipment ON sales.shipment_items(shipment_id); +CREATE INDEX idx_shipment_items_order_item ON sales.shipment_items(order_item_id); + +-- Reviews and ratings +CREATE TABLE sales.product_reviews ( + review_id BIGSERIAL PRIMARY KEY, + product_id BIGINT NOT NULL REFERENCES inventory.products(product_id) ON DELETE CASCADE, + user_id BIGINT NOT NULL REFERENCES core.users(user_id) ON DELETE CASCADE, + order_id BIGINT, + rating INTEGER NOT NULL CHECK (rating BETWEEN 1 AND 5), + title VARCHAR(200), + review_text TEXT, + search_vector TSVECTOR, + is_verified_purchase BOOLEAN NOT NULL DEFAULT FALSE, + is_approved BOOLEAN NOT NULL DEFAULT FALSE, + helpful_count INTEGER NOT NULL DEFAULT 0, + unhelpful_count INTEGER NOT NULL DEFAULT 0, + helpfulness_score NUMERIC(5,2) GENERATED ALWAYS AS ( + CASE + WHEN (helpful_count + unhelpful_count) > 0 + THEN (helpful_count::NUMERIC / (helpful_count + unhelpful_count) * 100) + ELSE NULL + END + ) STORED, + images TEXT[], + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + UNIQUE(product_id, user_id, order_id) +); + +CREATE INDEX idx_reviews_product ON sales.product_reviews(product_id, is_approved) WHERE is_approved = TRUE; +CREATE INDEX idx_reviews_user ON sales.product_reviews(user_id); +CREATE INDEX idx_reviews_rating ON sales.product_reviews(product_id, rating); +CREATE INDEX idx_reviews_search ON sales.product_reviews USING GIN(search_vector) WHERE is_approved = TRUE; + +-- Trigger to update search vector for reviews +CREATE OR REPLACE FUNCTION sales.reviews_search_vector_update() +RETURNS TRIGGER AS $$ +BEGIN + NEW.search_vector := + setweight(to_tsvector('english', COALESCE(NEW.title, '')), 'A') || + setweight(to_tsvector('english', COALESCE(NEW.review_text, '')), 'B'); + RETURN NEW; +END; +$$ LANGUAGE plpgsql; + +CREATE TRIGGER trigger_reviews_search_vector + BEFORE INSERT OR UPDATE OF title, review_text + ON sales.product_reviews + FOR EACH ROW + EXECUTE FUNCTION sales.reviews_search_vector_update(); + +-- ============================================================================ +-- AUDIT SCHEMA: Change tracking +-- ============================================================================ + +CREATE TABLE audit.user_activity_log ( + log_id BIGSERIAL, + user_id BIGINT REFERENCES core.users(user_id) ON DELETE SET NULL, + activity_type VARCHAR(50) NOT NULL, + entity_type VARCHAR(50), + entity_id BIGINT, + ip_address INET, + user_agent TEXT, + request_data JSONB, + response_data JSONB, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + PRIMARY KEY (log_id, created_at) +) PARTITION BY RANGE (created_at); + +CREATE TABLE audit.user_activity_log_2025_q1 PARTITION OF audit.user_activity_log + FOR VALUES FROM ('2025-01-01') TO ('2025-04-01'); +CREATE TABLE audit.user_activity_log_2025_q2 PARTITION OF audit.user_activity_log + FOR VALUES FROM ('2025-04-01') TO ('2025-07-01'); +CREATE TABLE audit.user_activity_log_2025_q3 PARTITION OF audit.user_activity_log + FOR VALUES FROM ('2025-07-01') TO ('2025-10-01'); +CREATE TABLE audit.user_activity_log_2025_q4 PARTITION OF audit.user_activity_log + FOR VALUES FROM ('2025-10-01') TO ('2026-01-01'); + +CREATE INDEX idx_activity_user ON audit.user_activity_log(user_id, created_at); +CREATE INDEX idx_activity_type ON audit.user_activity_log(activity_type, created_at); +CREATE INDEX idx_activity_entity ON audit.user_activity_log(entity_type, entity_id); + +-- ============================================================================ +-- ANALYTICS SCHEMA: Views and Materialized Views +-- ============================================================================ + +-- View: Current stock availability across all warehouses +CREATE VIEW analytics.product_stock_summary AS +SELECT + COALESCE(sl.product_id, pv.product_id) as product_id, + sl.variant_id, + p.product_name, + p.sku, + pv.variant_sku, + SUM(sl.quantity_available) as total_available, + SUM(sl.quantity_reserved) as total_reserved, + SUM(sl.quantity_damaged) as total_damaged, + COUNT(DISTINCT sl.warehouse_id) as warehouse_count, + MAX(sl.updated_at) as last_updated +FROM inventory.stock_levels sl +LEFT JOIN inventory.products p ON sl.product_id = p.product_id +LEFT JOIN inventory.product_variants pv ON sl.variant_id = pv.variant_id +GROUP BY COALESCE(sl.product_id, pv.product_id), sl.variant_id, p.product_name, p.sku, pv.variant_sku; + +-- View: Order summary with customer info +CREATE VIEW analytics.order_summary AS +SELECT + o.order_id, + o.order_number, + o.status, + o.created_at, + u.user_id, + u.email, + u.first_name, + u.last_name, + o.subtotal_amount, + o.discount_amount, + o.tax_amount, + o.shipping_amount, + o.total_amount, + o.currency_code, + COUNT(oi.order_item_id) as item_count, + SUM(oi.quantity) as total_quantity +FROM sales.orders o +JOIN core.users u ON o.user_id = u.user_id +LEFT JOIN sales.order_items oi ON o.order_id = oi.order_id +GROUP BY o.order_id, o.order_number, o.status, o.created_at, + u.user_id, u.email, u.first_name, u.last_name, + o.subtotal_amount, o.discount_amount, o.tax_amount, + o.shipping_amount, o.total_amount, o.currency_code; + +-- Materialized View: Daily sales metrics +CREATE MATERIALIZED VIEW analytics.daily_sales_metrics AS +SELECT + DATE(created_at) as sale_date, + status, + currency_code, + COUNT(*) as order_count, + SUM(total_amount) as total_revenue, + SUM(subtotal_amount) as subtotal_revenue, + SUM(discount_amount) as total_discounts, + SUM(tax_amount) as total_tax, + SUM(shipping_amount) as total_shipping, + AVG(total_amount) as avg_order_value, + PERCENTILE_CONT(0.5) WITHIN GROUP (ORDER BY total_amount) as median_order_value +FROM sales.orders +WHERE created_at >= CURRENT_DATE - INTERVAL '365 days' +GROUP BY DATE(created_at), status, currency_code +WITH DATA; + +CREATE UNIQUE INDEX idx_daily_sales_metrics ON analytics.daily_sales_metrics(sale_date, status, currency_code); +CREATE INDEX idx_daily_sales_date ON analytics.daily_sales_metrics(sale_date); + +-- Materialized View: Product performance metrics +CREATE MATERIALIZED VIEW analytics.product_performance AS +SELECT + p.product_id, + p.product_name, + p.sku, + p.status, + c.category_name, + b.brand_name, + COUNT(DISTINCT oi.order_id) as order_count, + SUM(oi.quantity) as units_sold, + SUM(oi.total_price) as total_revenue, + AVG(pr.rating) as avg_rating, + COUNT(pr.review_id) as review_count +FROM inventory.products p +LEFT JOIN inventory.categories c ON p.category_id = c.category_id +LEFT JOIN inventory.brands b ON p.brand_id = b.brand_id +LEFT JOIN sales.order_items oi ON p.product_id = oi.product_id +LEFT JOIN sales.product_reviews pr ON p.product_id = pr.product_id AND pr.is_approved = TRUE +GROUP BY p.product_id, p.product_name, p.sku, p.status, c.category_name, b.brand_name +WITH DATA; + +CREATE UNIQUE INDEX idx_product_performance ON analytics.product_performance(product_id); +CREATE INDEX idx_product_performance_revenue ON analytics.product_performance(total_revenue DESC NULLS LAST); + +-- Materialized View: Customer lifetime value +CREATE MATERIALIZED VIEW analytics.customer_lifetime_value AS +SELECT + u.user_id, + u.email, + u.first_name, + u.last_name, + u.created_at as customer_since, + COUNT(DISTINCT o.order_id) as total_orders, + SUM(o.total_amount) as lifetime_value, + AVG(o.total_amount) as avg_order_value, + MAX(o.created_at) as last_order_date, + MIN(o.created_at) as first_order_date +FROM core.users u +LEFT JOIN sales.orders o ON u.user_id = o.user_id AND o.status NOT IN ('cancelled', 'refunded') +WHERE 'customer' = ANY(u.roles) +GROUP BY u.user_id, u.email, u.first_name, u.last_name, u.created_at +WITH DATA; + +CREATE UNIQUE INDEX idx_customer_ltv ON analytics.customer_lifetime_value(user_id); +CREATE INDEX idx_customer_ltv_value ON analytics.customer_lifetime_value(lifetime_value DESC NULLS LAST); + +-- View: Low stock alerts +CREATE VIEW analytics.low_stock_alerts AS +SELECT + w.warehouse_id, + w.warehouse_name, + p.product_id, + p.product_name, + p.sku, + sl.quantity_available, + sl.quantity_reserved, + sl.reorder_point, + sl.reorder_quantity, + (sl.quantity_available - sl.quantity_reserved) as available_for_sale +FROM inventory.stock_levels sl +JOIN inventory.warehouses w ON sl.warehouse_id = w.warehouse_id +LEFT JOIN inventory.products p ON sl.product_id = p.product_id +WHERE sl.quantity_available <= sl.reorder_point + AND w.is_active = TRUE; + +-- View: Monthly revenue by category +CREATE VIEW analytics.category_revenue_monthly AS +SELECT + DATE_TRUNC('month', o.created_at) as month, + c.category_id, + c.category_name, + COUNT(DISTINCT o.order_id) as order_count, + SUM(oi.quantity) as units_sold, + SUM(oi.total_price) as total_revenue +FROM sales.orders o +JOIN sales.order_items oi ON o.order_id = oi.order_id +LEFT JOIN inventory.products p ON oi.product_id = p.product_id +LEFT JOIN inventory.categories c ON p.category_id = c.category_id +WHERE o.status NOT IN ('cancelled', 'refunded') + AND o.created_at >= CURRENT_DATE - INTERVAL '24 months' +GROUP BY DATE_TRUNC('month', o.created_at), c.category_id, c.category_name; + +-- ============================================================================ +-- FUNCTIONS AND PROCEDURES +-- ============================================================================ + +-- Function to calculate cart total +CREATE OR REPLACE FUNCTION sales.calculate_cart_total(p_cart_id BIGINT) +RETURNS TABLE( + subtotal NUMERIC(12,2), + tax NUMERIC(12,2), + shipping NUMERIC(12,2), + total NUMERIC(12,2) +) AS $$ +DECLARE + v_subtotal NUMERIC(12,2); + v_tax NUMERIC(12,2); + v_shipping NUMERIC(12,2); + v_total NUMERIC(12,2); +BEGIN + SELECT COALESCE(SUM(total_price), 0) INTO v_subtotal + FROM sales.cart_items + WHERE cart_id = p_cart_id; + + v_tax := v_subtotal * 0.08; -- 8% tax + v_shipping := CASE + WHEN v_subtotal > 100 THEN 0 + ELSE 9.99 + END; + + v_total := v_subtotal + v_tax + v_shipping; + + RETURN QUERY SELECT v_subtotal, v_tax, v_shipping, v_total; +END; +$$ LANGUAGE plpgsql; + +-- Function to reserve stock +CREATE OR REPLACE FUNCTION inventory.reserve_stock( + p_warehouse_id BIGINT, + p_product_id BIGINT, + p_variant_id BIGINT, + p_quantity INTEGER +) +RETURNS BOOLEAN AS $$ +DECLARE + v_available INTEGER; +BEGIN + -- Lock the row for update + SELECT quantity_available INTO v_available + FROM inventory.stock_levels + WHERE warehouse_id = p_warehouse_id + AND (product_id = p_product_id OR variant_id = p_variant_id) + FOR UPDATE; + + IF v_available >= p_quantity THEN + UPDATE inventory.stock_levels + SET + quantity_available = quantity_available - p_quantity, + quantity_reserved = quantity_reserved + p_quantity, + updated_at = NOW() + WHERE warehouse_id = p_warehouse_id + AND (product_id = p_product_id OR variant_id = p_variant_id); + + RETURN TRUE; + ELSE + RETURN FALSE; + END IF; +END; +$$ LANGUAGE plpgsql; + +-- Procedure to refresh analytics materialized views +CREATE OR REPLACE PROCEDURE analytics.refresh_all_metrics() +LANGUAGE plpgsql +AS $$ +BEGIN + REFRESH MATERIALIZED VIEW CONCURRENTLY analytics.daily_sales_metrics; + REFRESH MATERIALIZED VIEW CONCURRENTLY analytics.product_performance; + REFRESH MATERIALIZED VIEW CONCURRENTLY analytics.customer_lifetime_value; + RAISE NOTICE 'All analytics metrics refreshed at %', NOW(); +END; +$$; + +-- ============================================================================ +-- TRIGGERS +-- ============================================================================ + +-- Trigger to update updated_at timestamp +CREATE OR REPLACE FUNCTION update_updated_at_column() +RETURNS TRIGGER AS $$ +BEGIN + NEW.updated_at = NOW(); + RETURN NEW; +END; +$$ LANGUAGE plpgsql; + +CREATE TRIGGER trigger_users_updated_at BEFORE UPDATE ON core.users + FOR EACH ROW EXECUTE FUNCTION update_updated_at_column(); + +CREATE TRIGGER trigger_addresses_updated_at BEFORE UPDATE ON core.addresses + FOR EACH ROW EXECUTE FUNCTION update_updated_at_column(); + +CREATE TRIGGER trigger_products_updated_at BEFORE UPDATE ON inventory.products + FOR EACH ROW EXECUTE FUNCTION update_updated_at_column(); + +CREATE TRIGGER trigger_carts_updated_at BEFORE UPDATE ON sales.carts + FOR EACH ROW EXECUTE FUNCTION update_updated_at_column(); + +-- ============================================================================ +-- SAMPLE DATA +-- ============================================================================ + +-- Insert sample countries +INSERT INTO core.countries (country_code, country_name, iso3, phone_prefix, currency_code, region) VALUES +('US', 'United States', 'USA', '+1', 'USD', 'North America'), +('CA', 'Canada', 'CAN', '+1', 'CAD', 'North America'), +('GB', 'United Kingdom', 'GBR', '+44', 'GBP', 'Europe'), +('DE', 'Germany', 'DEU', '+49', 'EUR', 'Europe'), +('FR', 'France', 'FRA', '+33', 'EUR', 'Europe'), +('JP', 'Japan', 'JPN', '+81', 'JPY', 'Asia'), +('CN', 'China', 'CHN', '+86', 'CNY', 'Asia'), +('AU', 'Australia', 'AUS', '+61', 'AUD', 'Oceania'); + +-- Insert sample users +INSERT INTO core.users (email, username, password_hash, first_name, last_name, phone, date_of_birth, status, roles) +VALUES +('john.doe@example.com', 'johndoe', '$2a$10$abcdefghijklmnopqrstuv', 'John', 'Doe', '+1234567890', '1990-05-15', 'active', ARRAY['customer']::core.user_role[]), +('jane.smith@example.com', 'janesmith', '$2a$10$abcdefghijklmnopqrstuv', 'Jane', 'Smith', '+1234567891', '1985-08-22', 'active', ARRAY['customer', 'vendor']::core.user_role[]), +('admin@example.com', 'admin', '$2a$10$abcdefghijklmnopqrstuv', 'Admin', 'User', '+1234567892', '1980-01-01', 'active', ARRAY['admin']::core.user_role[]); + +-- Insert sample brands +INSERT INTO inventory.brands (brand_name, brand_slug, description, is_active) +VALUES +('TechPro', 'techpro', 'Premium technology products', true), +('StyleWear', 'stylewear', 'Fashion and apparel', true), +('HomeComfort', 'homecomfort', 'Home and living essentials', true); + +-- Insert sample categories +INSERT INTO inventory.categories (category_name, category_slug, level, path) +VALUES +('Electronics', 'electronics', 0, 'electronics'), +('Computers', 'computers', 1, 'electronics.computers'), +('Laptops', 'laptops', 2, 'electronics.computers.laptops'), +('Clothing', 'clothing', 0, 'clothing'), +('Men', 'men', 1, 'clothing.men'), +('Women', 'women', 1, 'clothing.women'); + +-- Insert sample warehouses +INSERT INTO inventory.warehouses (warehouse_code, warehouse_name, address_line1, city, state_province, postal_code, country_code, is_active) +VALUES +('WH-NYC-01', 'New York Main Warehouse', '123 Industrial Blvd', 'New York', 'NY', '10001', 'US', true), +('WH-LAX-01', 'Los Angeles Distribution Center', '456 Logistics Way', 'Los Angeles', 'CA', '90001', 'US', true), +('WH-LON-01', 'London Fulfillment Center', '789 Commerce Street', 'London', 'England', 'SW1A 1AA', 'GB', true); + +-- Insert sample products +INSERT INTO inventory.products (sku, product_name, product_slug, brand_id, category_id, description, status, base_price, weight, tags) +VALUES +('LAPTOP-001', 'TechPro UltraBook Pro 15', 'techpro-ultrabook-pro-15', 1, 3, 'High-performance laptop with 15-inch display', 'active', 1299.99, 1.8, ARRAY['laptop', 'computer', 'tech']), +('LAPTOP-002', 'TechPro PowerStation 17', 'techpro-powerstation-17', 1, 3, 'Workstation laptop for professionals', 'active', 2499.99, 2.5, ARRAY['laptop', 'workstation', 'professional']), +('SHIRT-001', 'StyleWear Classic Cotton Tee', 'stylewear-classic-tee', 2, 5, 'Comfortable cotton t-shirt', 'active', 29.99, 0.2, ARRAY['clothing', 'shirt', 'casual']); + +COMMENT ON SCHEMA core IS 'Core business entities including users, addresses, and authentication'; +COMMENT ON SCHEMA inventory IS 'Product catalog and inventory management'; +COMMENT ON SCHEMA sales IS 'Orders, payments, and transaction processing'; +COMMENT ON SCHEMA analytics IS 'Business intelligence views and metrics'; +COMMENT ON SCHEMA audit IS 'Activity logging and audit trails'; + +COMMENT ON TABLE core.users IS 'User accounts with role-based access control'; +COMMENT ON TABLE inventory.products IS 'Product catalog with full-text search capabilities'; +COMMENT ON TABLE sales.orders IS 'Customer orders partitioned by creation date'; +COMMENT ON TABLE inventory.stock_levels IS 'Inventory levels partitioned by warehouse'; diff --git a/integration/schema_sync/pgdog.toml b/integration/schema_sync/pgdog.toml new file mode 100644 index 000000000..3b9cfaede --- /dev/null +++ b/integration/schema_sync/pgdog.toml @@ -0,0 +1,16 @@ +[general] + +[rewrite] +enabled = false +shard_key = "ignore" +split_inserts = "error" + +[[databases]] +name = "source" +host = "127.0.0.1" +database_name = "pgdog1" + +[[databases]] +name = "destination" +host = "127.0.0.1" +database_name = "pgdog2" diff --git a/integration/schema_sync/run.sh b/integration/schema_sync/run.sh new file mode 100644 index 000000000..d67c3a5c4 --- /dev/null +++ b/integration/schema_sync/run.sh @@ -0,0 +1,4 @@ +#!/bin/bash +set -e +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd ) +bash ${SCRIPT_DIR}/dev.sh diff --git a/integration/schema_sync/users.toml b/integration/schema_sync/users.toml new file mode 100644 index 000000000..098c31715 --- /dev/null +++ b/integration/schema_sync/users.toml @@ -0,0 +1,11 @@ +[[users]] +name = "pgdog" +database = "source" +password = "pgdog" +schema_admin = true + +[[users]] +name = "pgdog" +database = "destination" +password = "pgdog" +schema_admin = true diff --git a/integration/setup.sh b/integration/setup.sh index 913732dd1..2a168e23f 100644 --- a/integration/setup.sh +++ b/integration/setup.sh @@ -24,7 +24,7 @@ fi export PGPASSWORD='pgdog' export PGHOST=127.0.0.1 export PGPORT=5432 -#export PGUSER='pgdog' +export PGUSER='pgdog' for db in pgdog shard_0 shard_1; do psql -c "DROP DATABASE $db" || true @@ -38,7 +38,16 @@ done for db in pgdog shard_0 shard_1; do for table in sharded sharded_omni; do psql -c "DROP TABLE IF EXISTS ${table}" ${db} -U pgdog - psql -c "CREATE TABLE IF NOT EXISTS ${table} (id BIGINT PRIMARY KEY, value TEXT)" ${db} -U pgdog + psql -c "CREATE TABLE IF NOT EXISTS ${table} ( + id BIGINT PRIMARY KEY, + value TEXT, + created_at TIMESTAMPTZ DEFAULT NOW(), + enabled BOOLEAN DEFAULT false, + user_id BIGINT, + region_id INTEGER DEFAULT 10, + country_id SMALLINT DEFAULT 5, + options JSONB DEFAULT '{}'::jsonb + )" ${db} -U pgdog done psql -c "CREATE TABLE IF NOT EXISTS sharded_varchar (id_varchar VARCHAR)" ${db} -U pgdog diff --git a/integration/tls/cert.pem b/integration/tls/cert.pem new file mode 100644 index 000000000..dc7f5453f --- /dev/null +++ b/integration/tls/cert.pem @@ -0,0 +1,34 @@ +-----BEGIN CERTIFICATE----- +MIIF7zCCA9egAwIBAgIUbzkoynbPz2lq2enZqe00RMSeQccwDQYJKoZIhvcNAQEL +BQAwgYYxCzAJBgNVBAYTAlhYMRIwEAYDVQQIDAlTdGF0ZU5hbWUxETAPBgNVBAcM +CENpdHlOYW1lMRQwEgYDVQQKDAtDb21wYW55TmFtZTEbMBkGA1UECwwSQ29tcGFu +eVNlY3Rpb25OYW1lMR0wGwYDVQQDDBRDb21tb25OYW1lT3JIb3N0bmFtZTAeFw0y +NDEyMjkwMDEzMjVaFw0zNDEyMjcwMDEzMjVaMIGGMQswCQYDVQQGEwJYWDESMBAG +A1UECAwJU3RhdGVOYW1lMREwDwYDVQQHDAhDaXR5TmFtZTEUMBIGA1UECgwLQ29t +cGFueU5hbWUxGzAZBgNVBAsMEkNvbXBhbnlTZWN0aW9uTmFtZTEdMBsGA1UEAwwU +Q29tbW9uTmFtZU9ySG9zdG5hbWUwggIiMA0GCSqGSIb3DQEBAQUAA4ICDwAwggIK +AoICAQCXeyKPV5zhBok+yeGgTYQjO+SakHkJ80NSxkKqXi1bIwlHZukROEt4T7h5 +qiV3TlomC5d6iUnLOG8EDFuZfIfRZOsvF99aCzk92MQ+EcezLSm3EXZm2+LjVH5s +hWcDxc7X6T8ZTcOTRnlO0RamYgUtuI/0kPK0mj+cJNF8OKXrz74BdwKvn+VmLY4H +fW1Y49FylSKIZGyb56ki24yzcwSBaSuwW8XkknuTW/w3ScmUTlwKvg55HmAqi2XC +OV0+I09XtoXGxde80DU1e/5NYjXxJ3IokiEByKMLBpVi3/B+p3VxsIHskTKjDSh0 +AyNZyacnqHtoWoHaU1iD9k/aJjK3StTJZl3JZr/SrV1u1evRZF9IlayZagaYauTq +Ms5MHie9M822x108viYMfWTtzXai5VNTWBz3Agpr459JD5cL1XFxeuGj1csRQ5dV +SbfeZF8s/lE6aMxUn1uLQyd2W+44RchgGfb1ek6KyDFsS1YB0qQNktaGQVnoKtKC +Y0dhUon6r+DN8biO3zm9v95NFql45BiJFLESImF1Qs/FR7LzU/YeFTVffjEpvuVy +kqOEBPqtnNLKXzxpky6qmuSnI5jD3w5Nw68peSyFcTB5Vc0zBfofEraUN8a81mg9 +SL491tkI7/6ZkEL+zEKdoJKS3QfTHa03AOIRLuALJH9m1x0P1QIDAQABo1MwUTAd +BgNVHQ4EFgQU6TmTbgrh4q68CxQbs6D1Ql1UN8UwHwYDVR0jBBgwFoAU6TmTbgrh +4q68CxQbs6D1Ql1UN8UwDwYDVR0TAQH/BAUwAwEB/zANBgkqhkiG9w0BAQsFAAOC +AgEAHWca690XikBloy6WEqm6xz8JAe/+Q9SAImvgkdCsw5k07GWUWAW4vEjNtJK5 +6D3Ky+zY13gfDy2n7xyum+Kwt01/tXmRLSTv+OE7tA75reBG2ZY9YzV5sLiW9Fza ++ukJCbLdYOuNG7eOr7rQWFrm7ARmRkqPAKA5QCNKszFYUj8C/nmKvrT6N9CUzryv +xGXEXxcQLIf+S7v1yhI18Vbe0B1aDvwruwuULcMmW+OpKHF37NxbcL5dGLMZv1CZ +ezW+zqOurPNvDXHvz/TSDfmUyP8Kl4wnsHYiX+on4ewcPY/yxwMJKOsdPemP6Neg +anJAL15JDC0k2c0YDS8T2JdSt6YmsczD8EUQgWzOBgvM4K4j3qFbe5z1qCGCoMIG +veEWKschmCUu3WBlYZjye4BnAdKAuM54+PJm3Y385zAPhP4aaoQJuJWE9eTEAOAB +s4yBQyHSMq2W2D/Ku1Rt48kU2/2MdSXr+G3vvs7XzpBcrYF341S+MGDF3ErKxgrT +IDRHZBPPiggkJbCjJSoK2MA7grXnxdNb8RU3ZVPM0tMlB3EojXtZesdGadLP6ej9 +fE7VMix81WayoY7JrRKN9Q5jGKYhcmVWie4M0W2Ypl6pZ4Y7L/7jNijYvDeh6j14 +m5Z+/IpLsssLk6K1Kdloj3FHZYgHkt1Lh1XphOJ/yL8xmOw= +-----END CERTIFICATE----- diff --git a/integration/tls/key.pem b/integration/tls/key.pem new file mode 100644 index 000000000..3731e2d28 --- /dev/null +++ b/integration/tls/key.pem @@ -0,0 +1,52 @@ +-----BEGIN PRIVATE KEY----- +MIIJQQIBADANBgkqhkiG9w0BAQEFAASCCSswggknAgEAAoICAQCXeyKPV5zhBok+ +yeGgTYQjO+SakHkJ80NSxkKqXi1bIwlHZukROEt4T7h5qiV3TlomC5d6iUnLOG8E +DFuZfIfRZOsvF99aCzk92MQ+EcezLSm3EXZm2+LjVH5shWcDxc7X6T8ZTcOTRnlO +0RamYgUtuI/0kPK0mj+cJNF8OKXrz74BdwKvn+VmLY4HfW1Y49FylSKIZGyb56ki +24yzcwSBaSuwW8XkknuTW/w3ScmUTlwKvg55HmAqi2XCOV0+I09XtoXGxde80DU1 +e/5NYjXxJ3IokiEByKMLBpVi3/B+p3VxsIHskTKjDSh0AyNZyacnqHtoWoHaU1iD +9k/aJjK3StTJZl3JZr/SrV1u1evRZF9IlayZagaYauTqMs5MHie9M822x108viYM +fWTtzXai5VNTWBz3Agpr459JD5cL1XFxeuGj1csRQ5dVSbfeZF8s/lE6aMxUn1uL +Qyd2W+44RchgGfb1ek6KyDFsS1YB0qQNktaGQVnoKtKCY0dhUon6r+DN8biO3zm9 +v95NFql45BiJFLESImF1Qs/FR7LzU/YeFTVffjEpvuVykqOEBPqtnNLKXzxpky6q +muSnI5jD3w5Nw68peSyFcTB5Vc0zBfofEraUN8a81mg9SL491tkI7/6ZkEL+zEKd +oJKS3QfTHa03AOIRLuALJH9m1x0P1QIDAQABAoICABtzVrqzpYv4v3/HnVHLok2x +QZbJ5glJ0lIqd/PAL8dzdK/CBCvY8AI8LiGsFfCGHBuHX7q2rM79Kc8Jv0Kz+LfX +KjBlStYaMRQWV0+pMK91WHkimrp+j+Hi0qsvTJD4NGjXjZX0C+RBMeP4y3o4ypfz +uXCYIMdeKXdOC8FPUbAHPDcvPicd2ngW+sU8M0fXtwGk6XZefnkNNM8KireNOQyL +hr2Fj/nBGtBEK9NIFZXA0nim4uALg2FKVBUriIxlYTAztQ/lk9gVQgMwdk/HI5/R +JmSYQI9+cJ9joMgjbUVCauvAkPbSBCNck89cLziK7LXo1/474oqyLljxlpxhbjCV +8i9BL6la06DXZULrBryKPc2Ds0eRDW45Sp7oYMR0AzE1iC1d11SvtyVVL3QFxyMj +NMTZdt13ypL3x4JCtWSq+cOpmwaQkRhUJMGCeGjBPD37q3wDJW6FePLoIBu4ZeY2 +mRj/qTsfuUjqWB29VyqhnVoBuU4ifsmxNJO2Or2NI/sd1mCA9pYVKK2oPNzSPg78 +e9USRTOc8XO9XFUgnGUFdGQYNRNEQ4zLf0RnKZGGQodUQXMA8Px01tW0/eIf/aCt +f1xiXpoqM1cFjjPxvpWc6gGgYlhEfcqMpWZ3XjJ4a9XrR5KVW+dWc44JgXKObgmI +0lk9TDAljQr5yfnkEvSRAoIBAQDU2OiMRabtmtFHndreOCiE61B3YJQ1PpxQn/u+ +BJ9KZ3TiE+Z4M4yRXy6HqGVPY2mXeSf6eOtXmmRbQX86YBatsj7ibHDX1f54Ijr3 +auw+/ywwt7SwXfHeJpr4+HxluF9+A/NrBQoZeyyU0TxhbxHBYQH2RHySwnSidYW9 +l2PgoaVfEYc+/cuadwB333UuKNdPXFY6mhQt9NjqofflkEP720icSIfCzFUvRIgd +3+H3SXb9ry5lXNn8b/TUTPQyA1Ni+lp+6p8bT6rD7VanuEUeKtCb/Ie2xwoTbGd7 +sTQRURdG3is2y94kLRwdP7cjGO+M5vZITkexGV4m7km0CsuRAoIBAQC2MTljmXOM +sMNcHx3WKGzQYNPjf2XWmW3wP/rrt5lVt7WYlkart+whgcDWJVqMDOuviJdHbFCh +LxfIjHoe8T/ZLDT0CynUMfsoxRkLedYEp1TVkLD+9P89ZnGYUSJ28uBflQ3RHzOg +1kkne1LqYjyOkKuBFI9oGHlN6MsHxD8KkbY6cMIXZdPFABGwgewATW5/pvs4UztS +Dhte0ma7NU/A68K+aVUXm35+akIhNxd535afz+XWuiBZc13kThhExFBZEh19upCc +e1DLCVhHefKLnoO0AS89KtoNs7aHXQucr/MEI93imNz/IMC32YPUmzHQw8tN6o4Y +U2lu81KgTrYFAoIBAGUneM1BROXjD9bDVIMLmWYiFynEwmrTiKJghdl2hOVtaYUQ +BBXYGdP0sj5Sb2NdUY9lSvSkhuQpQcyEwhxSEjUWYwBknPRWhQs+6Vswe3os9ylo +BP1UiGAVZM0x+py1FNzkr8iKqpQVj8hh8Bo2GPAYVEBfp/xvYdLbm2XRDuxwphEa +WXY8U4jjSVuu3RfE3R6gOXK8Sx7UIErSEugMueJ2AnoTlkGjrlA6d54LCm7lgSFr +IdeWWxq3cll7AQrLvdNqO5vZkSf/op5eqzImRuLhYibfyve4fDdi64NDYgVgznkl +mM//72Ct95CG+Vg6v43tLdqLKVMnRTGnSWvBPaECggEAKtTRqA+YMZgQpWSPUBx6 +0FYjGhWGLHgvd06jP60O+C7TG0cg4BfCBHKLkgyAB/K1qbOT1O+q2OnITpZv0zxm +BTk2TbUeJUuGvyPu6lq/LKLl97snURjptFaUF/ni/1HD29Sfxezu5z3ZPtXoPT/Q ++rcaCqN5v0AZrG4w5OeG5oYw7/Y4OuXubh7BCdzRTZTmiE4KO0id5oF4f8c47YPv +9uu2AaujnIQqra9vUn2wIC+nKnTmlJ93IXBUv2p4nBoGxZnTow4sFw2KheDxhwQt +OBOQ5M1ufJPJZXU9UP9XzoMyv2NrM206byQVCmOxcVb21BxjfDLLKv7ZB4Nehl9a +vQKCAQAvj/s5kaxv5BKGOce0GYkMUgEHXqHLp5hg07EDuuJbEev5Vq1zZdFffBEc +xWxXU8gChIsfiTzYDZdoAxJN4M6OoVPFz2LKSPlxRkesD8+bZA5xhGphLk0jR6Ly +lZJD7lqgG6F1NfzcqYnjZoPlYeaSLvkCoNkJmwYFNTXoK6b+wZj/ghsFj98sZZv/ +daNN/0BACowwrvJX6cJfN1MOkbe4rvuMBdgUMG2nb4kOr9uB2yHt4cPDfoLBDiSw +0hsPmpvOhof9VnUyQFFSzgr0Au5943DjLVMAtbAdnoqhjePnyNteq2grIcmVZ/oL +WwV6cBrpOmlqZKPrar6DXY2NLWf/ +-----END PRIVATE KEY----- diff --git a/integration/tls/pgdog.toml b/integration/tls/pgdog.toml new file mode 100644 index 000000000..544afd300 --- /dev/null +++ b/integration/tls/pgdog.toml @@ -0,0 +1,12 @@ +[general] +tls_certificate = "cert.pem" +tls_private_key = "key.pem" + +[rewrite] +enabled = false +shard_key = "ignore" +split_inserts = "error" + +[[databases]] +name = "pgdog" +host = "127.0.0.1" diff --git a/integration/tls/users.toml b/integration/tls/users.toml new file mode 100644 index 000000000..539bb1832 --- /dev/null +++ b/integration/tls/users.toml @@ -0,0 +1,4 @@ +[[users]] +name = "pgdog" +password = "pgdog" +database = "pgdog" diff --git a/integration/toxi/toxi_spec.rb b/integration/toxi/toxi_spec.rb index 0588a77d2..ce5f2a248 100644 --- a/integration/toxi/toxi_spec.rb +++ b/integration/toxi/toxi_spec.rb @@ -43,9 +43,8 @@ def warm_up it 'some connections survive' do threads = [] - errors = 0 + errors = Concurrent::AtomicFixnum.new(0) sem = Concurrent::Semaphore.new(0) - (5.0 / 25 * 25.0).ceil 25.times do t = Thread.new do c = 1 @@ -54,13 +53,13 @@ def warm_up c = conn break rescue StandardError - errors += 1 + errors.increment end 25.times do c.exec 'SELECT 1' rescue PG::SystemError c = conn # reconnect - errors += 1 + errors.increment end end threads << t @@ -69,7 +68,7 @@ def warm_up sem.release(25) threads.each(&:join) end - expect(errors).to be < 25 # 5% error rate (instead of 100%) + expect(errors.value).to be < 25 # 5% error rate (instead of 100%) end it 'active record works' do @@ -78,22 +77,93 @@ def warm_up # Connect (the pool is lazy) Sharded.where(id: 1).first errors = 0 + ok = 0 # Can't ban primary because it issues SET queries # that we currently route to primary. Toxiproxy[role].toxic(toxic).apply do 25.times do Sharded.where(id: 1).first + ok += 1 rescue StandardError errors += 1 end end - expect(errors).to eq(1) + expect(errors).to be <= 1 + expect(25 - ok).to eq(errors) + end +end + +describe 'healthcheck' do + before :each do + admin_conn = admin + admin_conn.exec 'RECONNECT' + admin_conn.exec "SET read_write_split TO 'exclude_primary'" + admin_conn.exec 'SET ban_timeout TO 1' + end + + describe 'will heal itself' do + def health(role, field = 'healthy') + admin.exec('SHOW POOLS').select do |pool| + pool['database'] == 'failover' && pool['role'] == role + end.map { |pool| pool[field] } + end + + 10.times do + it 'replica' do + # Cache connect params. + conn.exec 'SELECT 1' + + Toxiproxy[:replica].toxic(:reset_peer).apply do + errors = 0 + 4.times do + conn.exec 'SELECT 1' + rescue PG::Error + errors += 1 + end + expect(errors).to be >= 1 + expect(health('replica')).to include('f') + sleep(0.4) # ban maintenance runs every 333ms + expect(health('replica', 'banned')).to include('t') + end + + 4.times do + conn.exec 'SELECT 1' + end + + admin.exec 'HEALTHCHECK' + sleep(0.4) + + expect(health('replica')).to eq(%w[t t t]) + expect(health('replica', 'banned')).to eq(%w[f f f]) + end + end + + it 'primary' do + # Cache connect params. + conn.exec 'DELETE FROM sharded' + + Toxiproxy[:primary].toxic(:reset_peer).apply do + begin + conn.exec 'DELETE FROM sharded' + rescue PG::Error + end + expect(health('primary')).to eq(['f']) + end + + conn.exec 'DELETE FROM sharded' + + expect(health('primary')).to eq(%w[t]) + end + end + + after do + admin.exec 'RELOAD' end end describe 'tcp' do around :each do |example| - Timeout.timeout(10) do + Timeout.timeout(30) do example.run end end @@ -138,8 +208,23 @@ def warm_up end describe 'both down' do - it 'unbans all pools' do - 25.times do + 10.times do + it 'unbans all pools' do + rw_config = admin.exec('SHOW CONFIG').select do |config| + config['name'] == 'read_write_split' + end[0]['value'] + expect(rw_config).to eq('include_primary') + + def pool_stat(field, value) + failover = admin.exec('SHOW POOLS').select do |pool| + pool['database'] == 'failover' + end + entries = failover.select { |item| item[field] == value } + entries.size + end + + admin.exec 'SET checkout_timeout TO 100' + Toxiproxy[:primary].toxic(:reset_peer).apply do Toxiproxy[:replica].toxic(:reset_peer).apply do Toxiproxy[:replica2].toxic(:reset_peer).apply do @@ -148,19 +233,16 @@ def warm_up conn.exec_params 'SELECT $1::bigint', [1] rescue StandardError end - banned = admin.exec('SHOW POOLS').select do |pool| - pool['database'] == 'failover' - end.select { |item| item['banned'] == 't' } - expect(banned.size).to eq(4) + + expect(pool_stat('healthy', 'f')).to eq(4) end end end end - conn.exec 'SELECT $1::bigint', [25] - banned = admin.exec('SHOW POOLS').select do |pool| - pool['database'] == 'failover' - end.select { |item| item['banned'] == 't' } - expect(banned.size).to eq(0) + + 4.times do + conn.exec 'SELECT $1::bigint', [25] + end end end end @@ -172,25 +254,21 @@ def warm_up Toxiproxy[:primary].toxic(:reset_peer).apply do c = conn c.exec 'BEGIN' - c.exec 'CREATE TABLE test(id BIGINT)' + c.exec 'CREATE TABLE IF NOT EXISTS test(id BIGINT)' c.exec 'ROLLBACK' rescue StandardError end + banned = admin.exec('SHOW POOLS').select do |pool| pool['database'] == 'failover' && pool['role'] == 'primary' end - expect(banned[0]['banned']).to eq('t') + expect(banned[0]['healthy']).to eq('f') c = conn c.exec 'BEGIN' - c.exec 'CREATE TABLE test(id BIGINT)' + c.exec 'CREATE TABLE IF NOT EXISTS test(id BIGINT)' c.exec 'SELECT * FROM test' c.exec 'ROLLBACK' - - banned = admin.exec('SHOW POOLS').select do |pool| - pool['database'] == 'failover' && pool['role'] == 'primary' - end - expect(banned[0]['banned']).to eq('f') end it 'active record works' do @@ -199,16 +277,58 @@ def warm_up # Connect (the pool is lazy) Sharded.where(id: 1).first errors = 0 + ok = 0 # Can't ban primary because it issues SET queries # that we currently route to primary. Toxiproxy[:primary].toxic(:reset_peer).apply do 25.times do Sharded.where(id: 1).first + ok += 1 rescue StandardError errors += 1 end end - expect(errors).to eq(1) + expect(errors).to be <= 1 + expect(25 - ok).to eq(errors) + end + + it 'clients can connect when all servers are down after caching connection params' do + # First, establish a connection to cache connection parameters + c = conn + c.exec 'SELECT 1' + c.close + + # Verify initial state - all pools should be healthy before toxics + pools = admin.exec('SHOW POOLS').select do |pool| + pool['database'] == 'failover' + end + expect(pools.all? { |p| p['healthy'] == 't' }).to be true + + # Now bring down all servers + Toxiproxy[:primary].toxic(:reset_peer).apply do + Toxiproxy[:replica].toxic(:reset_peer).apply do + Toxiproxy[:replica2].toxic(:reset_peer).apply do + Toxiproxy[:replica3].toxic(:reset_peer).apply do + # Try to establish many connections + connections = [] + 50.times do + c = conn + expect(c).not_to be_nil + connections << c + end + + # Check internal state - verify we have active client connections + clients = admin.exec('SHOW CLIENTS').select do |client| + client['database'] == 'failover' + end + expect(clients.size).to be >= 50 + + # Clean up connections without executing queries to avoid timeouts + connections.each { |c| c.close rescue nil } + end + end + end + end end end end diff --git a/integration/users.toml b/integration/users.toml index 6d7178088..87273f7c6 100644 --- a/integration/users.toml +++ b/integration/users.toml @@ -62,3 +62,8 @@ password = "pgdog" server_user = "pgdog" cross_shard_disabled = true min_pool_size = 0 + +[[users]] +name = "pgdog" +password = "pgdog" +database = "pgdog_schema" diff --git a/mise.toml b/mise.toml new file mode 100644 index 000000000..0d2f0ed0b --- /dev/null +++ b/mise.toml @@ -0,0 +1,3 @@ +[tools] +"cargo:cargo-nextest" = "latest" +"cargo:cargo-watch" = "latest" \ No newline at end of file diff --git a/pgdog-config/Cargo.toml b/pgdog-config/Cargo.toml new file mode 100644 index 000000000..0b34ec2c5 --- /dev/null +++ b/pgdog-config/Cargo.toml @@ -0,0 +1,18 @@ +[package] +name = "pgdog-config" +version = "0.1.0" +edition = "2021" + +[dependencies] +serde = { version = "1", features = ["derive"] } +serde_json = "1" +tracing = "0.1" +thiserror = "2" +toml = "0.8" +url = "2" +uuid = { version = "1", features = ["v4", "serde"] } +rand = "*" +pgdog-vector = { path = "../pgdog-vector" } + +[dev-dependencies] +tempfile = "3.23.0" diff --git a/pgdog-config/src/auth.rs b/pgdog-config/src/auth.rs new file mode 100644 index 000000000..d397fb01b --- /dev/null +++ b/pgdog-config/src/auth.rs @@ -0,0 +1,60 @@ +use serde::{Deserialize, Serialize}; +use std::{fmt::Display, str::FromStr}; + +#[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq)] +#[serde(rename_all = "snake_case")] +pub enum PassthoughAuth { + #[default] + Disabled, + Enabled, + EnabledPlain, +} + +#[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq)] +#[serde(rename_all = "snake_case")] +pub enum AuthType { + Md5, + #[default] + Scram, + Trust, + Plain, +} + +impl Display for AuthType { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Md5 => write!(f, "md5"), + Self::Scram => write!(f, "scram"), + Self::Trust => write!(f, "trust"), + Self::Plain => write!(f, "plain"), + } + } +} + +impl AuthType { + pub fn md5(&self) -> bool { + matches!(self, Self::Md5) + } + + pub fn scram(&self) -> bool { + matches!(self, Self::Scram) + } + + pub fn trust(&self) -> bool { + matches!(self, Self::Trust) + } +} + +impl FromStr for AuthType { + type Err = String; + + fn from_str(s: &str) -> Result { + match s.to_lowercase().as_str() { + "md5" => Ok(Self::Md5), + "scram" => Ok(Self::Scram), + "trust" => Ok(Self::Trust), + "plain" => Ok(Self::Plain), + _ => Err(format!("Invalid auth type: {}", s)), + } + } +} diff --git a/pgdog-config/src/core.rs b/pgdog-config/src/core.rs new file mode 100644 index 000000000..44d369347 --- /dev/null +++ b/pgdog-config/src/core.rs @@ -0,0 +1,779 @@ +use serde::{Deserialize, Serialize}; +use std::collections::{HashMap, HashSet}; +use std::fs::read_to_string; +use std::path::PathBuf; +use tracing::{info, warn}; + +use crate::sharding::ShardedSchema; +use crate::{ + EnumeratedDatabase, Memory, OmnishardedTable, PassthoughAuth, PreparedStatements, + QueryParserEngine, QueryParserLevel, ReadWriteSplit, RewriteMode, Role, +}; + +use super::database::Database; +use super::error::Error; +use super::general::General; +use super::networking::{MultiTenant, Tcp}; +use super::pooling::{PoolerMode, Stats}; +use super::replication::{MirrorConfig, Mirroring, ReplicaLag, Replication}; +use super::rewrite::Rewrite; +use super::sharding::{ManualQuery, OmnishardedTables, ShardedMapping, ShardedTable}; +use super::users::{Admin, Plugin, Users}; + +#[derive(Debug, Clone)] +pub struct ConfigAndUsers { + /// pgdog.toml + pub config: Config, + /// users.toml + pub users: Users, + /// Path to pgdog.toml. + pub config_path: PathBuf, + /// Path to users.toml. + pub users_path: PathBuf, +} + +impl ConfigAndUsers { + /// Load configuration from disk or use defaults. + pub fn load(config_path: &PathBuf, users_path: &PathBuf) -> Result { + let mut config: Config = if let Ok(config) = read_to_string(config_path) { + let config = match toml::from_str(&config) { + Ok(config) => config, + Err(err) => return Err(Error::config(&config, err)), + }; + info!("loaded \"{}\"", config_path.display()); + config + } else { + warn!( + "\"{}\" doesn't exist, loading defaults instead", + config_path.display() + ); + Config::default() + }; + + if config.multi_tenant.is_some() { + info!("multi-tenant protection enabled"); + } + + let mut users: Users = if let Ok(users) = read_to_string(users_path) { + let mut users: Users = toml::from_str(&users)?; + users.check(&config); + info!("loaded \"{}\"", users_path.display()); + users + } else { + warn!( + "\"{}\" doesn't exist, loading defaults instead", + users_path.display() + ); + Users::default() + }; + + // Override admin set in pgdog.toml + // with what's in users.toml. + if let Some(admin) = users.admin.take() { + config.admin = admin; + } + + if config.admin.random() { + #[cfg(debug_assertions)] + info!("[debug only] admin password: {}", config.admin.password); + #[cfg(not(debug_assertions))] + warn!("admin password has been randomly generated"); + } + + Ok(ConfigAndUsers { + config, + users, + config_path: config_path.to_owned(), + users_path: users_path.to_owned(), + }) + } + + /// Prepared statements are enabled. + pub fn prepared_statements(&self) -> PreparedStatements { + // Disable prepared statements automatically in session mode + if self.config.general.pooler_mode == PoolerMode::Session { + PreparedStatements::Disabled + } else { + self.config.general.prepared_statements + } + } + + /// Prepared statements are in "full" mode (used for query parser decision). + pub fn prepared_statements_full(&self) -> bool { + self.config.general.prepared_statements.full() + } + + pub fn query_parser_enabled(&self) -> bool { + self.config.general.query_parser_enabled + } + + pub fn pub_sub_enabled(&self) -> bool { + self.config.general.pub_sub_channel_size > 0 + } +} + +impl Default for ConfigAndUsers { + fn default() -> Self { + Self { + config: Config::default(), + users: Users::default(), + config_path: PathBuf::from("pgdog.toml"), + users_path: PathBuf::from("users.toml"), + } + } +} + +/// Configuration. +#[derive(Serialize, Deserialize, Debug, Clone, Default)] +#[serde(deny_unknown_fields)] +pub struct Config { + /// General configuration. + #[serde(default)] + pub general: General, + + /// Rewrite configuration. + #[serde(default)] + pub rewrite: Rewrite, + + /// Statistics. + #[serde(default)] + pub stats: Stats, + + /// TCP settings + #[serde(default)] + pub tcp: Tcp, + + /// Multi-tenant + pub multi_tenant: Option, + + /// Servers. + #[serde(default)] + pub databases: Vec, + + #[serde(default)] + pub plugins: Vec, + + #[serde(default)] + pub admin: Admin, + + /// List of sharded tables. + #[serde(default)] + pub sharded_tables: Vec, + + /// Queries routed manually to a single shard. + #[serde(default)] + pub manual_queries: Vec, + + /// List of omnisharded tables. + #[serde(default)] + pub omnisharded_tables: Vec, + + /// Explicit sharding key mappings. + #[serde(default)] + pub sharded_mappings: Vec, + + #[serde(default)] + pub sharded_schemas: Vec, + + /// Replica lag configuration. + #[serde(default, deserialize_with = "ReplicaLag::deserialize_optional")] + pub replica_lag: Option, + + /// Replication config. + #[serde(default)] + pub replication: Replication, + + /// Mirroring configurations. + #[serde(default)] + pub mirroring: Vec, + + /// Memory tweaks + #[serde(default)] + pub memory: Memory, +} + +impl Config { + /// Organize all databases by name for quicker retrieval. + pub fn databases(&self) -> HashMap>> { + let mut databases = HashMap::new(); + let mut number = 0; + for database in &self.databases { + let entry = databases + .entry(database.name.clone()) + .or_insert_with(Vec::new); + while entry.len() <= database.shard { + entry.push(vec![]); + } + entry + .get_mut(database.shard) + .unwrap() + .push(EnumeratedDatabase { + number, + database: database.clone(), + }); + number += 1; + } + databases + } + + /// Organize sharded tables by database name. + pub fn sharded_tables(&self) -> HashMap> { + let mut tables = HashMap::new(); + + for table in &self.sharded_tables { + let entry = tables + .entry(table.database.clone()) + .or_insert_with(Vec::new); + entry.push(table.clone()); + } + + tables + } + + pub fn omnisharded_tables(&self) -> HashMap> { + let mut tables = HashMap::new(); + + for table in &self.omnisharded_tables { + let entry = tables + .entry(table.database.clone()) + .or_insert_with(Vec::new); + for t in &table.tables { + entry.push(OmnishardedTable { + name: t.clone(), + sticky_routing: table.sticky, + }); + } + } + + tables + } + + pub fn sharded_schemas(&self) -> HashMap> { + let mut schemas = HashMap::new(); + + for schema in &self.sharded_schemas { + let entry = schemas + .entry(schema.database.clone()) + .or_insert_with(Vec::new); + entry.push(schema.clone()); + } + + schemas + } + + /// Manual queries. + pub fn manual_queries(&self) -> HashMap { + let mut queries = HashMap::new(); + + for query in &self.manual_queries { + queries.insert(query.fingerprint.clone(), query.clone()); + } + + queries + } + + /// Sharded mappings. + pub fn sharded_mappings( + &self, + ) -> HashMap<(String, String, Option), Vec> { + let mut mappings = HashMap::new(); + + for mapping in &self.sharded_mappings { + let mapping = mapping.clone(); + let entry = mappings + .entry(( + mapping.database.clone(), + mapping.column.clone(), + mapping.table.clone(), + )) + .or_insert_with(Vec::new); + entry.push(mapping); + } + + mappings + } + + pub fn check(&mut self) { + // Check databases. + let mut duplicate_dbs = HashSet::new(); + for database in self.databases.clone() { + let id = ( + database.name.clone(), + database.role, + database.shard, + database.port, + database.host.clone(), + ); + let new = duplicate_dbs.insert(id); + if !new { + warn!( + "database \"{}\" (shard={}) has a duplicate {}", + database.name, database.shard, database.role, + ); + } + } + + struct Check { + pooler_mode: Option, + role: Role, + role_warned: bool, + parser_warned: bool, + have_replicas: bool, + sharded: bool, + } + + // Check identical configs. + let mut checks = HashMap::::new(); + for database in &self.databases { + if let Some(existing) = checks.get_mut(&database.name) { + if existing.pooler_mode != database.pooler_mode { + warn!( + "database \"{}\" (shard={}, role={}) has a different \"pooler_mode\" setting, ignoring", + database.name, database.shard, database.role, + ); + } + let auto = existing.role == Role::Auto || database.role == Role::Auto; + if auto && existing.role != database.role && !existing.role_warned { + warn!( + r#"database "{}" has a mix of auto and specific roles, automatic role detection will be disabled"#, + database.name + ); + existing.role_warned = true; + } + if !existing.have_replicas { + existing.have_replicas = database.role == Role::Replica; + } + if !existing.sharded { + existing.sharded = database.shard > 0; + } + + if (existing.sharded || existing.have_replicas) + && self.general.query_parser == QueryParserLevel::Off + && !existing.parser_warned + { + existing.parser_warned = true; + warn!( + r#"database "{}" may need the query parser for load balancing/sharding, but it's disabled"#, + database.name + ); + } + } else { + checks.insert( + database.name.clone(), + Check { + pooler_mode: database.pooler_mode, + role: database.role, + role_warned: false, + parser_warned: false, + have_replicas: database.role == Role::Replica, + sharded: database.shard > 0, + }, + ); + } + } + + // Check that idle_healthcheck_interval is shorter than ban_timeout. + if self.general.ban_timeout > 0 + && self.general.idle_healthcheck_interval >= self.general.ban_timeout + { + warn!( + "idle_healthcheck_interval ({}ms) should be shorter than ban_timeout ({}ms) to ensure health checks are triggered before a ban expires", + self.general.idle_healthcheck_interval, self.general.ban_timeout + ); + } + + // Warn about plain auth and TLS + match self.general.passthrough_auth { + PassthoughAuth::Enabled if !self.general.tls_client_required => { + warn!( + "consider setting \"tls_client_required\" while \"passthrough_auth\" is enabled to prevent clients from exposing plaintext passwords" + ); + } + PassthoughAuth::EnabledPlain => { + warn!( + "\"passthrough_auth\" is set to \"plain\", network traffic may expose plaintext passwords" + ) + } + _ => (), + } + + if !self.general.two_phase_commit && self.rewrite.enabled { + if self.rewrite.shard_key == RewriteMode::Rewrite { + warn!( + r#"rewrite.shard_key = "rewrite" may apply non-atomic sharding key rewrites; enabling "two_phase_commit" is strongly recommended"# + ); + } + + if self.rewrite.split_inserts == RewriteMode::Rewrite { + warn!( + r#"rewrite.split_inserts = "rewrite" may commit partial multi-row inserts; enabling "two_phase_commit" is strongly recommended"# + ); + } + } + + for (database, check) in checks { + if !check.have_replicas + && self.general.read_write_split == ReadWriteSplit::ExcludePrimary + { + warn!( + r#"database "{}" has no replicas and "read_write_split" is set to "{}": read queries will be rejected"#, + database, self.general.read_write_split + ); + } + } + + if self.general.query_parser_enabled { + warn!(r#""query_parser_enabled" is deprecated, use "query_parser" = "on" instead"#); + self.general.query_parser = QueryParserLevel::On; + } + + if self.general.query_parser_engine == QueryParserEngine::PgQueryRaw { + if self.memory.stack_size < 32 * 1024 * 1024 { + self.memory.stack_size = 32 * 1024 * 1024; + warn!( + r#""pg_query_raw" parser engine requires a large thread stack, setting it to 32MiB for each Tokio worker"# + ); + } + } + } + + /// Multi-tenancy is enabled. + pub fn multi_tenant(&self) -> &Option { + &self.multi_tenant + } + + /// Get mirroring configuration for a specific source/destination pair. + pub fn get_mirroring_config( + &self, + source_db: &str, + destination_db: &str, + ) -> Option { + self.mirroring + .iter() + .find(|m| m.source_db == source_db && m.destination_db == destination_db) + .map(|m| MirrorConfig { + queue_length: m.queue_length.unwrap_or(self.general.mirror_queue), + exposure: m.exposure.unwrap_or(self.general.mirror_exposure), + }) + } + + /// Get all mirroring configurations mapped by source database. + pub fn mirroring_by_source(&self) -> HashMap> { + let mut result = HashMap::new(); + + for mirror in &self.mirroring { + let config = MirrorConfig { + queue_length: mirror.queue_length.unwrap_or(self.general.mirror_queue), + exposure: mirror.exposure.unwrap_or(self.general.mirror_exposure), + }; + + result + .entry(mirror.source_db.clone()) + .or_insert_with(Vec::new) + .push((mirror.destination_db.clone(), config)); + } + + result + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{PoolerMode, PreparedStatements}; + use std::time::Duration; + + #[test] + fn test_basic() { + let source = r#" +[general] +host = "0.0.0.0" +port = 6432 +default_pool_size = 15 +pooler_mode = "transaction" + +[[databases]] +name = "production" +role = "primary" +host = "127.0.0.1" +port = 5432 +database_name = "postgres" + +[tcp] +keepalive = true +interval = 5000 +time = 1000 +user_timeout = 1000 +retries = 5 + +[[plugins]] +name = "pgdog_routing" + +[multi_tenant] +column = "tenant_id" +"#; + + let config: Config = toml::from_str(source).unwrap(); + assert_eq!(config.databases[0].name, "production"); + assert_eq!(config.plugins[0].name, "pgdog_routing"); + assert!(config.tcp.keepalive()); + assert_eq!(config.tcp.interval().unwrap(), Duration::from_millis(5000)); + assert_eq!( + config.tcp.user_timeout().unwrap(), + Duration::from_millis(1000) + ); + assert_eq!(config.tcp.time().unwrap(), Duration::from_millis(1000)); + assert_eq!(config.tcp.retries().unwrap(), 5); + assert_eq!(config.multi_tenant.unwrap().column, "tenant_id"); + } + + #[test] + fn test_prepared_statements_disabled_in_session_mode() { + let mut config = ConfigAndUsers::default(); + + // Test transaction mode (default) - prepared statements should be enabled + config.config.general.pooler_mode = PoolerMode::Transaction; + config.config.general.prepared_statements = PreparedStatements::Extended; + assert_eq!( + config.prepared_statements(), + PreparedStatements::Extended, + "Prepared statements should be enabled in transaction mode" + ); + + // Test session mode - prepared statements should be disabled + config.config.general.pooler_mode = PoolerMode::Session; + config.config.general.prepared_statements = PreparedStatements::Extended; + assert_eq!( + config.prepared_statements(), + PreparedStatements::Disabled, + "Prepared statements should be disabled in session mode" + ); + + // Test session mode with full prepared statements - should still be disabled + config.config.general.pooler_mode = PoolerMode::Session; + config.config.general.prepared_statements = PreparedStatements::Full; + assert_eq!( + config.prepared_statements(), + PreparedStatements::Disabled, + "Prepared statements should be disabled in session mode even when set to Full" + ); + + // Test transaction mode with disabled prepared statements - should remain disabled + config.config.general.pooler_mode = PoolerMode::Transaction; + config.config.general.prepared_statements = PreparedStatements::Disabled; + assert_eq!( + config.prepared_statements(), + PreparedStatements::Disabled, + "Prepared statements should remain disabled when explicitly set to Disabled in transaction mode" + ); + } + + #[test] + fn test_mirroring_config() { + let source = r#" +[general] +host = "0.0.0.0" +port = 6432 +mirror_queue = 128 +mirror_exposure = 1.0 + +[[databases]] +name = "source_db" +host = "127.0.0.1" +port = 5432 + +[[databases]] +name = "destination_db1" +host = "127.0.0.1" +port = 5433 + +[[databases]] +name = "destination_db2" +host = "127.0.0.1" +port = 5434 + +[[mirroring]] +source_db = "source_db" +destination_db = "destination_db1" +queue_length = 256 +exposure = 0.5 + +[[mirroring]] +source_db = "source_db" +destination_db = "destination_db2" +exposure = 0.75 +"#; + + let config: Config = toml::from_str(source).unwrap(); + + // Verify we have 2 mirroring configurations + assert_eq!(config.mirroring.len(), 2); + + // Check first mirroring config + assert_eq!(config.mirroring[0].source_db, "source_db"); + assert_eq!(config.mirroring[0].destination_db, "destination_db1"); + assert_eq!(config.mirroring[0].queue_length, Some(256)); + assert_eq!(config.mirroring[0].exposure, Some(0.5)); + + // Check second mirroring config + assert_eq!(config.mirroring[1].source_db, "source_db"); + assert_eq!(config.mirroring[1].destination_db, "destination_db2"); + assert_eq!(config.mirroring[1].queue_length, None); // Should use global default + assert_eq!(config.mirroring[1].exposure, Some(0.75)); + + // Verify global defaults are still set + assert_eq!(config.general.mirror_queue, 128); + assert_eq!(config.general.mirror_exposure, 1.0); + + // Test get_mirroring_config method + let mirror_config = config + .get_mirroring_config("source_db", "destination_db1") + .unwrap(); + assert_eq!(mirror_config.queue_length, 256); + assert_eq!(mirror_config.exposure, 0.5); + + let mirror_config2 = config + .get_mirroring_config("source_db", "destination_db2") + .unwrap(); + assert_eq!(mirror_config2.queue_length, 128); // Uses global default + assert_eq!(mirror_config2.exposure, 0.75); + + // Non-existent mirror config should return None + assert!(config + .get_mirroring_config("source_db", "non_existent") + .is_none()); + } + + #[test] + fn test_admin_override_from_users_toml() { + use std::io::Write; + use tempfile::NamedTempFile; + + let pgdog_config = r#" +[admin] +name = "pgdog_admin" +user = "pgdog_admin_user" +password = "pgdog_admin_password" +"#; + + let users_config = r#" +[admin] +name = "users_admin" +user = "users_admin_user" +password = "users_admin_password" +"#; + + let mut pgdog_file = NamedTempFile::new().unwrap(); + let mut users_file = NamedTempFile::new().unwrap(); + + pgdog_file.write_all(pgdog_config.as_bytes()).unwrap(); + users_file.write_all(users_config.as_bytes()).unwrap(); + + pgdog_file.flush().unwrap(); + users_file.flush().unwrap(); + + let config_and_users = + ConfigAndUsers::load(&pgdog_file.path().into(), &users_file.path().into()).unwrap(); + + assert_eq!(config_and_users.config.admin.name, "users_admin"); + assert_eq!(config_and_users.config.admin.user, "users_admin_user"); + assert_eq!( + config_and_users.config.admin.password, + "users_admin_password" + ); + assert!(config_and_users.users.admin.is_none()); + } + + #[test] + fn test_admin_override_with_default_config() { + use std::io::Write; + use tempfile::NamedTempFile; + + let pgdog_config = r#" +[general] +host = "0.0.0.0" +port = 6432 +"#; + + let users_config = r#" +[admin] +name = "users_admin" +user = "users_admin_user" +password = "users_admin_password" +"#; + + let mut pgdog_file = NamedTempFile::new().unwrap(); + let mut users_file = NamedTempFile::new().unwrap(); + + pgdog_file.write_all(pgdog_config.as_bytes()).unwrap(); + users_file.write_all(users_config.as_bytes()).unwrap(); + + pgdog_file.flush().unwrap(); + users_file.flush().unwrap(); + + let config_and_users = + ConfigAndUsers::load(&pgdog_file.path().into(), &users_file.path().into()).unwrap(); + + assert_eq!(config_and_users.config.admin.name, "users_admin"); + assert_eq!(config_and_users.config.admin.user, "users_admin_user"); + assert_eq!( + config_and_users.config.admin.password, + "users_admin_password" + ); + assert!(config_and_users.users.admin.is_none()); + } + + #[test] + fn test_omnisharded_tables() { + let source = r#" +[general] +host = "0.0.0.0" +port = 6432 + +[[databases]] +name = "db1" +host = "127.0.0.1" +port = 5432 + +[[databases]] +name = "db2" +host = "127.0.0.1" +port = 5433 + +[[omnisharded_tables]] +database = "db1" +tables = ["table_a", "table_b"] + +[[omnisharded_tables]] +database = "db1" +tables = ["table_c"] +sticky = true + +[[omnisharded_tables]] +database = "db2" +tables = ["table_x"] +"#; + + let config: Config = toml::from_str(source).unwrap(); + + assert_eq!(config.omnisharded_tables.len(), 3); + + let tables = config.omnisharded_tables(); + + assert_eq!(tables.len(), 2); + + let db1_tables = tables.get("db1").unwrap(); + assert_eq!(db1_tables.len(), 3); + assert_eq!(db1_tables[0].name, "table_a"); + assert!(!db1_tables[0].sticky_routing); + assert_eq!(db1_tables[1].name, "table_b"); + assert!(!db1_tables[1].sticky_routing); + assert_eq!(db1_tables[2].name, "table_c"); + assert!(db1_tables[2].sticky_routing); + + let db2_tables = tables.get("db2").unwrap(); + assert_eq!(db2_tables.len(), 1); + assert_eq!(db2_tables[0].name, "table_x"); + assert!(!db2_tables[0].sticky_routing); + } +} diff --git a/pgdog-config/src/data_types.rs b/pgdog-config/src/data_types.rs new file mode 100644 index 000000000..3ee2f7106 --- /dev/null +++ b/pgdog-config/src/data_types.rs @@ -0,0 +1,186 @@ +use std::{ + cmp::Ordering, + fmt::Display, + hash::{Hash, Hasher}, +}; + +use serde::{ + de::{self, Visitor}, + ser::SerializeSeq, + Deserialize, Serialize, +}; + +/// Wrapper type for f32 that implements Ord for PostgreSQL compatibility +#[derive(Debug, Clone, Copy, Serialize, Deserialize)] +pub struct Float(pub f32); + +impl PartialOrd for Float { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl Ord for Float { + fn cmp(&self, other: &Self) -> Ordering { + // PostgreSQL ordering: NaN is greater than all other values + match (self.0.is_nan(), other.0.is_nan()) { + (true, true) => Ordering::Equal, + (true, false) => Ordering::Greater, + (false, true) => Ordering::Less, + (false, false) => self.0.partial_cmp(&other.0).unwrap_or(Ordering::Equal), + } + } +} + +impl PartialEq for Float { + fn eq(&self, other: &Self) -> bool { + // PostgreSQL treats NaN as equal to NaN for indexing purposes + if self.0.is_nan() && other.0.is_nan() { + true + } else { + self.0 == other.0 + } + } +} + +impl Eq for Float {} + +impl Hash for Float { + fn hash(&self, state: &mut H) { + if self.0.is_nan() { + // All NaN values hash to the same value + 0u8.hash(state); + } else { + // Use bit representation for consistent hashing + self.0.to_bits().hash(state); + } + } +} + +impl Display for Float { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + if self.0.is_nan() { + write!(f, "NaN") + } else if self.0.is_infinite() { + if self.0.is_sign_positive() { + write!(f, "Infinity") + } else { + write!(f, "-Infinity") + } + } else { + write!(f, "{}", self.0) + } + } +} + +impl From for Float { + fn from(value: f32) -> Self { + Float(value) + } +} + +impl From for f32 { + fn from(value: Float) -> Self { + value.0 + } +} + +#[derive(Clone, PartialEq, PartialOrd, Ord, Eq, Hash, Debug)] +#[repr(C)] +pub struct Vector { + pub values: Vec, +} + +impl Vector { + /// Length of the vector. + pub fn len(&self) -> usize { + self.values.len() + } + + /// Is the vector empty? + pub fn is_empty(&self) -> bool { + self.len() == 0 + } +} + +struct VectorVisitor; + +impl<'de> Visitor<'de> for VectorVisitor { + type Value = Vector; + + fn visit_seq(self, mut seq: A) -> Result + where + A: de::SeqAccess<'de>, + { + let mut results = vec![]; + while let Some(n) = seq.next_element::()? { + results.push(n); + } + + Ok(Vector::from(results.as_slice())) + } + + fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + formatter.write_str("expected a list of floating points") + } +} + +impl<'de> Deserialize<'de> for Vector { + fn deserialize(deserializer: D) -> Result + where + D: de::Deserializer<'de>, + { + deserializer.deserialize_seq(VectorVisitor) + } +} + +impl Serialize for Vector { + fn serialize(&self, serializer: S) -> Result + where + S: serde::Serializer, + { + let mut seq = serializer.serialize_seq(Some(self.len()))?; + for v in &self.values { + seq.serialize_element(v)?; + } + seq.end() + } +} + +impl From<&[f64]> for Vector { + fn from(value: &[f64]) -> Self { + Self { + values: value.iter().map(|v| Float(*v as f32)).collect(), + } + } +} + +impl From<&[f32]> for Vector { + fn from(value: &[f32]) -> Self { + Self { + values: value.iter().map(|v| Float(*v)).collect(), + } + } +} + +impl From> for Vector { + fn from(value: Vec) -> Self { + Self { + values: value.into_iter().map(Float::from).collect(), + } + } +} + +impl From> for Vector { + fn from(value: Vec) -> Self { + Self { + values: value.into_iter().map(|v| Float(v as f32)).collect(), + } + } +} + +impl From> for Vector { + fn from(value: Vec) -> Self { + Self { values: value } + } +} diff --git a/pgdog-config/src/database.rs b/pgdog-config/src/database.rs new file mode 100644 index 000000000..2bb483c0f --- /dev/null +++ b/pgdog-config/src/database.rs @@ -0,0 +1,193 @@ +use serde::{Deserialize, Serialize}; +use std::{ + fmt::Display, + ops::{Deref, DerefMut}, + str::FromStr, +}; + +use super::pooling::PoolerMode; + +#[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq, Copy)] +#[serde(rename_all = "snake_case")] +pub enum ReadWriteStrategy { + #[default] + Conservative, + Aggressive, +} + +impl FromStr for ReadWriteStrategy { + type Err = String; + + fn from_str(s: &str) -> Result { + match s.to_lowercase().as_str() { + "conservative" => Ok(Self::Conservative), + "aggressive" => Ok(Self::Aggressive), + _ => Err(format!("Invalid read-write strategy: {}", s)), + } + } +} + +#[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq, Copy)] +#[serde(rename_all = "snake_case")] +pub enum LoadBalancingStrategy { + #[default] + Random, + RoundRobin, + LeastActiveConnections, +} + +impl FromStr for LoadBalancingStrategy { + type Err = String; + + fn from_str(s: &str) -> Result { + match s.to_lowercase().replace(['_', '-'], "").as_str() { + "random" => Ok(Self::Random), + "roundrobin" => Ok(Self::RoundRobin), + "leastactiveconnections" => Ok(Self::LeastActiveConnections), + _ => Err(format!("Invalid load balancing strategy: {}", s)), + } + } +} + +#[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq, Copy)] +#[serde(rename_all = "snake_case")] +pub enum ReadWriteSplit { + #[default] + IncludePrimary, + ExcludePrimary, + IncludePrimaryIfReplicaBanned, +} + +impl FromStr for ReadWriteSplit { + type Err = String; + + fn from_str(s: &str) -> Result { + match s.to_lowercase().replace(['_', '-'], "").as_str() { + "includeprimary" => Ok(Self::IncludePrimary), + "excludeprimary" => Ok(Self::ExcludePrimary), + "includeprimaryifreplicabanned" => Ok(Self::IncludePrimaryIfReplicaBanned), + _ => Err(format!("Invalid read-write split: {}", s)), + } + } +} + +impl Display for ReadWriteSplit { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let display = match self { + Self::ExcludePrimary => "exclude_primary", + Self::IncludePrimary => "include_primary", + Self::IncludePrimaryIfReplicaBanned => "include_primary_if_replica_banned", + }; + + write!(f, "{}", display) + } +} + +/// Database server proxied by pgDog. +#[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq, Ord, PartialOrd, Eq)] +#[serde(deny_unknown_fields)] +pub struct Database { + /// Database name visible to the clients. + pub name: String, + /// Database role, e.g. primary. + #[serde(default)] + pub role: Role, + /// Database host or IP address, e.g. 127.0.0.1. + pub host: String, + /// Database port, e.g. 5432. + #[serde(default = "Database::port")] + pub port: u16, + /// Shard. + #[serde(default)] + pub shard: usize, + /// PostgreSQL database name, e.g. "postgres". + pub database_name: Option, + /// Use this user to connect to the database, overriding the userlist. + pub user: Option, + /// Use this password to login, overriding the userlist. + pub password: Option, + // Maximum number of connections to this database from this pooler. + // #[serde(default = "Database::max_connections")] + // pub max_connections: usize, + /// Pool size for this database pools, overriding `default_pool_size`. + pub pool_size: Option, + /// Minimum pool size for this database pools, overriding `min_pool_size`. + pub min_pool_size: Option, + /// Pooler mode. + pub pooler_mode: Option, + /// Statement timeout. + pub statement_timeout: Option, + /// Idle timeout. + pub idle_timeout: Option, + /// Read-only mode. + pub read_only: Option, + /// Server lifetime. + pub server_lifetime: Option, +} + +impl Database { + #[allow(dead_code)] + fn max_connections() -> usize { + usize::MAX + } + + fn port() -> u16 { + 5432 + } +} + +#[derive( + Serialize, Deserialize, Debug, Clone, Default, PartialEq, Ord, PartialOrd, Eq, Hash, Copy, +)] +#[serde(rename_all = "snake_case")] +pub enum Role { + #[default] + Primary, + Replica, + Auto, +} + +impl std::fmt::Display for Role { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Primary => write!(f, "primary"), + Self::Replica => write!(f, "replica"), + Self::Auto => write!(f, "auto"), + } + } +} + +impl FromStr for Role { + type Err = String; + + fn from_str(s: &str) -> Result { + match s.to_lowercase().as_str() { + "primary" => Ok(Self::Primary), + "replica" => Ok(Self::Replica), + "auto" => Ok(Self::Auto), + _ => Err(format!("Invalid role: {}", s)), + } + } +} + +/// Database with a unique number, identifying it +/// in the config. +#[derive(Debug, Clone)] +pub struct EnumeratedDatabase { + pub number: usize, + pub database: Database, +} + +impl Deref for EnumeratedDatabase { + type Target = Database; + + fn deref(&self) -> &Self::Target { + &self.database + } +} + +impl DerefMut for EnumeratedDatabase { + fn deref_mut(&mut self) -> &mut Self::Target { + &mut self.database + } +} diff --git a/pgdog-config/src/error.rs b/pgdog-config/src/error.rs new file mode 100644 index 000000000..b90a6af0a --- /dev/null +++ b/pgdog-config/src/error.rs @@ -0,0 +1,67 @@ +//! Configuration errors. + +use thiserror::Error; + +/// Configuration error. +#[derive(Debug, Error)] +pub enum Error { + #[error("{0}")] + Io(#[from] std::io::Error), + + #[error("{0}")] + Deser(#[from] toml::de::Error), + + #[error("{0}, line {1}")] + MissingField(String, usize), + + #[error("{0}")] + Url(#[from] url::ParseError), + + #[error("{0}")] + Json(#[from] serde_json::Error), + + #[error("incomplete startup")] + IncompleteStartup, + + #[error("no database urls in environment")] + NoDbsInEnv, + + #[error("parse error: {0}")] + ParseError(String), +} + +impl Error { + pub fn config(source: &str, err: toml::de::Error) -> Self { + let span = err.span(); + let message = err.message(); + + let span = if let Some(span) = span { + span + } else { + return Self::MissingField(message.into(), 0); + }; + + let mut lines = vec![]; + let mut line = 1; + for (i, c) in source.chars().enumerate() { + if c == '\n' { + lines.push((line, i)); + line += 1; + } + } + + let mut lines = lines.into_iter().peekable(); + + while let Some(line) = lines.next() { + if span.start < line.1 { + if let Some(next) = lines.peek() { + if next.1 > span.start { + return Self::MissingField(message.into(), line.0); + } + } + } + } + + Self::MissingField(message.into(), 0) + } +} diff --git a/pgdog-config/src/general.rs b/pgdog-config/src/general.rs new file mode 100644 index 000000000..dff1e092b --- /dev/null +++ b/pgdog-config/src/general.rs @@ -0,0 +1,965 @@ +use serde::{Deserialize, Serialize}; +use std::env; +use std::net::Ipv4Addr; +use std::path::PathBuf; +use std::time::Duration; + +use crate::pooling::ConnectionRecovery; +use crate::{QueryParserEngine, QueryParserLevel}; + +use super::auth::{AuthType, PassthoughAuth}; +use super::database::{LoadBalancingStrategy, ReadWriteSplit, ReadWriteStrategy}; +use super::networking::TlsVerifyMode; +use super::pooling::{PoolerMode, PreparedStatements}; + +#[derive(Serialize, Deserialize, Debug, Clone)] +#[serde(deny_unknown_fields)] +pub struct General { + /// Run on this address. + #[serde(default = "General::host")] + pub host: String, + /// Run on this port. + #[serde(default = "General::port")] + pub port: u16, + /// Spawn this many Tokio threads. + #[serde(default = "General::workers")] + pub workers: usize, + /// Default pool size, e.g. 10. + #[serde(default = "General::default_pool_size")] + pub default_pool_size: usize, + /// Minimum number of connections to maintain in the pool. + #[serde(default = "General::min_pool_size")] + pub min_pool_size: usize, + /// Pooler mode, e.g. transaction. + #[serde(default)] + pub pooler_mode: PoolerMode, + /// How often to check a connection. + #[serde(default = "General::healthcheck_interval")] + pub healthcheck_interval: u64, + /// How often to issue a healthcheck via an idle connection. + #[serde(default = "General::idle_healthcheck_interval")] + pub idle_healthcheck_interval: u64, + /// Delay idle healthchecks by this time at startup. + #[serde(default = "General::idle_healthcheck_delay")] + pub idle_healthcheck_delay: u64, + /// Healthcheck timeout. + #[serde(default = "General::healthcheck_timeout")] + pub healthcheck_timeout: u64, + /// HTTP health check port. + pub healthcheck_port: Option, + /// Maximum duration of a ban. + #[serde(default = "General::ban_timeout")] + pub ban_timeout: u64, + /// Rollback timeout. + #[serde(default = "General::rollback_timeout")] + pub rollback_timeout: u64, + /// Load balancing strategy. + #[serde(default = "General::load_balancing_strategy")] + pub load_balancing_strategy: LoadBalancingStrategy, + /// How aggressive should the query parser be in determining reads. + #[serde(default)] + pub read_write_strategy: ReadWriteStrategy, + /// Read write split. + #[serde(default)] + pub read_write_split: ReadWriteSplit, + /// TLS certificate. + pub tls_certificate: Option, + /// TLS private key. + pub tls_private_key: Option, + #[serde(default)] + pub tls_client_required: bool, + /// TLS verification mode (for connecting to servers) + #[serde(default = "General::default_tls_verify")] + pub tls_verify: TlsVerifyMode, + /// TLS CA certificate (for connecting to servers). + pub tls_server_ca_certificate: Option, + /// Shutdown timeout. + #[serde(default = "General::default_shutdown_timeout")] + pub shutdown_timeout: u64, + /// Shutdown termination timeout (after shutdown_timeout expires, forcibly terminate). + #[serde(default = "General::default_shutdown_termination_timeout")] + pub shutdown_termination_timeout: Option, + /// Broadcast IP. + pub broadcast_address: Option, + /// Broadcast port. + #[serde(default = "General::broadcast_port")] + pub broadcast_port: u16, + /// Load queries to file (warning: slow, don't use in production). + #[serde(default)] + pub query_log: Option, + /// Enable OpenMetrics server on this port. + pub openmetrics_port: Option, + /// OpenMetrics prefix. + pub openmetrics_namespace: Option, + /// Prepared statatements support. + #[serde(default)] + pub prepared_statements: PreparedStatements, + /// Parse Queries override. + #[serde(default = "General::query_parser_enabled")] + pub query_parser_enabled: bool, + /// Query parser. + #[serde(default)] + pub query_parser: QueryParserLevel, + /// Query parser engine. + #[serde(default)] + pub query_parser_engine: QueryParserEngine, + /// Limit on the number of prepared statements in the server cache. + #[serde(default = "General::prepared_statements_limit")] + pub prepared_statements_limit: usize, + #[serde(default = "General::query_cache_limit")] + pub query_cache_limit: usize, + /// Automatically add connection pools for user/database pairs we don't have. + #[serde(default = "General::default_passthrough_auth")] + pub passthrough_auth: PassthoughAuth, + /// Server connect timeout. + #[serde(default = "General::default_connect_timeout")] + pub connect_timeout: u64, + /// Attempt connections multiple times on bad networks. + #[serde(default = "General::connect_attempts")] + pub connect_attempts: u64, + /// How long to wait between connection attempts. + #[serde(default = "General::default_connect_attempt_delay")] + pub connect_attempt_delay: u64, + /// How long to wait for a query to return the result before aborting. Dangerous: don't use unless your network is bad. + #[serde(default = "General::default_query_timeout")] + pub query_timeout: u64, + /// Checkout timeout. + #[serde(default = "General::checkout_timeout")] + pub checkout_timeout: u64, + /// Login timeout. + #[serde(default = "General::client_login_timeout")] + pub client_login_timeout: u64, + /// Dry run for sharding. Parse the query, route to shard 0. + #[serde(default)] + pub dry_run: bool, + /// Idle timeout. + #[serde(default = "General::idle_timeout")] + pub idle_timeout: u64, + /// Client idle timeout. + #[serde(default = "General::default_client_idle_timeout")] + pub client_idle_timeout: u64, + /// Client idle in transaction timeout. + #[serde(default = "General::default_client_idle_in_transaction_timeout")] + pub client_idle_in_transaction_timeout: u64, + /// Server lifetime. + #[serde(default = "General::server_lifetime")] + pub server_lifetime: u64, + /// Mirror queue size. + #[serde(default = "General::mirror_queue")] + pub mirror_queue: usize, + /// Mirror exposure + #[serde(default = "General::mirror_exposure")] + pub mirror_exposure: f32, + #[serde(default)] + pub auth_type: AuthType, + /// Disable cross-shard queries. + #[serde(default)] + pub cross_shard_disabled: bool, + /// How often to refresh DNS entries, in ms. + #[serde(default)] + pub dns_ttl: Option, + /// LISTEN/NOTIFY channel size. + #[serde(default)] + pub pub_sub_channel_size: usize, + /// Log client connections. + #[serde(default = "General::log_connections")] + pub log_connections: bool, + /// Log client disconnections. + #[serde(default = "General::log_disconnections")] + pub log_disconnections: bool, + /// Two-phase commit. + #[serde(default)] + pub two_phase_commit: bool, + /// Two-phase commit automatic transactions. + #[serde(default)] + pub two_phase_commit_auto: Option, + /// Enable expanded EXPLAIN output. + #[serde(default = "General::expanded_explain")] + pub expanded_explain: bool, + /// Stats averaging period (in milliseconds). + #[serde(default = "General::stats_period")] + pub stats_period: u64, + /// Connection cleanup algorithm. + #[serde(default = "General::connection_recovery")] + pub connection_recovery: ConnectionRecovery, + /// LSN check interval. + #[serde(default = "General::lsn_check_interval")] + pub lsn_check_interval: u64, + /// LSN check timeout. + #[serde(default = "General::lsn_check_timeout")] + pub lsn_check_timeout: u64, + /// LSN check delay. + #[serde(default = "General::lsn_check_delay")] + pub lsn_check_delay: u64, + /// Minimum ID for unique ID generator. + #[serde(default)] + pub unique_id_min: u64, +} + +impl Default for General { + fn default() -> Self { + Self { + host: Self::host(), + port: Self::port(), + workers: Self::workers(), + default_pool_size: Self::default_pool_size(), + min_pool_size: Self::min_pool_size(), + pooler_mode: Self::pooler_mode(), + healthcheck_interval: Self::healthcheck_interval(), + idle_healthcheck_interval: Self::idle_healthcheck_interval(), + idle_healthcheck_delay: Self::idle_healthcheck_delay(), + healthcheck_timeout: Self::healthcheck_timeout(), + healthcheck_port: Self::healthcheck_port(), + ban_timeout: Self::ban_timeout(), + rollback_timeout: Self::rollback_timeout(), + load_balancing_strategy: Self::load_balancing_strategy(), + read_write_strategy: Self::read_write_strategy(), + read_write_split: Self::read_write_split(), + tls_certificate: Self::tls_certificate(), + tls_private_key: Self::tls_private_key(), + tls_client_required: bool::default(), + tls_verify: Self::default_tls_verify(), + tls_server_ca_certificate: Self::tls_server_ca_certificate(), + shutdown_timeout: Self::default_shutdown_timeout(), + shutdown_termination_timeout: Self::default_shutdown_termination_timeout(), + broadcast_address: Self::broadcast_address(), + broadcast_port: Self::broadcast_port(), + query_log: Self::query_log(), + openmetrics_port: Self::openmetrics_port(), + openmetrics_namespace: Self::openmetrics_namespace(), + prepared_statements: Self::prepared_statements(), + query_parser_enabled: Self::query_parser_enabled(), + query_parser: QueryParserLevel::default(), + query_parser_engine: QueryParserEngine::default(), + prepared_statements_limit: Self::prepared_statements_limit(), + query_cache_limit: Self::query_cache_limit(), + passthrough_auth: Self::default_passthrough_auth(), + connect_timeout: Self::default_connect_timeout(), + connect_attempt_delay: Self::default_connect_attempt_delay(), + connect_attempts: Self::connect_attempts(), + query_timeout: Self::default_query_timeout(), + checkout_timeout: Self::checkout_timeout(), + client_login_timeout: Self::client_login_timeout(), + dry_run: Self::dry_run(), + idle_timeout: Self::idle_timeout(), + client_idle_timeout: Self::default_client_idle_timeout(), + client_idle_in_transaction_timeout: Self::default_client_idle_in_transaction_timeout(), + mirror_queue: Self::mirror_queue(), + mirror_exposure: Self::mirror_exposure(), + auth_type: Self::auth_type(), + cross_shard_disabled: Self::cross_shard_disabled(), + dns_ttl: Self::default_dns_ttl(), + pub_sub_channel_size: Self::pub_sub_channel_size(), + log_connections: Self::log_connections(), + log_disconnections: Self::log_disconnections(), + two_phase_commit: bool::default(), + two_phase_commit_auto: None, + expanded_explain: Self::expanded_explain(), + server_lifetime: Self::server_lifetime(), + stats_period: Self::stats_period(), + connection_recovery: Self::connection_recovery(), + lsn_check_interval: Self::lsn_check_interval(), + lsn_check_timeout: Self::lsn_check_timeout(), + lsn_check_delay: Self::lsn_check_delay(), + unique_id_min: u64::default(), + } + } +} + +impl General { + fn env_or_default(env_var: &str, default: T) -> T { + env::var(env_var) + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(default) + } + + fn env_string_or_default(env_var: &str, default: &str) -> String { + env::var(env_var).unwrap_or_else(|_| default.to_string()) + } + + fn env_bool_or_default(env_var: &str, default: bool) -> bool { + env::var(env_var) + .ok() + .and_then(|v| match v.to_lowercase().as_str() { + "true" | "1" | "yes" | "on" => Some(true), + "false" | "0" | "no" | "off" => Some(false), + _ => None, + }) + .unwrap_or(default) + } + + fn env_option(env_var: &str) -> Option { + env::var(env_var).ok().and_then(|v| v.parse().ok()) + } + + fn env_option_string(env_var: &str) -> Option { + env::var(env_var).ok().filter(|s| !s.is_empty()) + } + + fn env_enum_or_default(env_var: &str) -> T { + env::var(env_var) + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or_default() + } + + fn host() -> String { + Self::env_string_or_default("PGDOG_HOST", "0.0.0.0") + } + + pub fn port() -> u16 { + Self::env_or_default("PGDOG_PORT", 6432) + } + + fn workers() -> usize { + Self::env_or_default("PGDOG_WORKERS", 2) + } + + fn default_pool_size() -> usize { + Self::env_or_default("PGDOG_DEFAULT_POOL_SIZE", 10) + } + + fn min_pool_size() -> usize { + Self::env_or_default("PGDOG_MIN_POOL_SIZE", 1) + } + + fn healthcheck_interval() -> u64 { + Self::env_or_default("PGDOG_HEALTHCHECK_INTERVAL", 30_000) + } + + fn idle_healthcheck_interval() -> u64 { + Self::env_or_default("PGDOG_IDLE_HEALTHCHECK_INTERVAL", 30_000) + } + + fn idle_healthcheck_delay() -> u64 { + Self::env_or_default("PGDOG_IDLE_HEALTHCHECK_DELAY", 5_000) + } + + fn healthcheck_port() -> Option { + Self::env_option("PGDOG_HEALTHCHECK_PORT") + } + + fn ban_timeout() -> u64 { + Self::env_or_default( + "PGDOG_BAN_TIMEOUT", + Duration::from_secs(300).as_millis() as u64, + ) + } + + fn rollback_timeout() -> u64 { + Self::env_or_default("PGDOG_ROLLBACK_TIMEOUT", 5_000) + } + + fn idle_timeout() -> u64 { + Self::env_or_default( + "PGDOG_IDLE_TIMEOUT", + Duration::from_secs(60).as_millis() as u64, + ) + } + + fn client_login_timeout() -> u64 { + Self::env_or_default( + "PGDOG_CLIENT_LOG_TIMEOUT", + Duration::from_secs(60).as_millis() as u64, + ) + } + + fn default_client_idle_timeout() -> u64 { + Self::env_or_default( + "PGDOG_CLIENT_IDLE_TIMEOUT", + Duration::MAX.as_millis() as u64, + ) + } + + fn default_client_idle_in_transaction_timeout() -> u64 { + Self::env_or_default( + "PGDOG_CLIENT_IDLE_IN_TRANSACTION_TIMEOUT", + Duration::MAX.as_millis() as u64, + ) + } + + fn default_query_timeout() -> u64 { + Self::env_or_default("PGDOG_QUERY_TIMEOUT", Duration::MAX.as_millis() as u64) + } + + pub fn query_timeout(&self) -> Duration { + Duration::from_millis(self.query_timeout) + } + + pub fn dns_ttl(&self) -> Option { + self.dns_ttl.map(Duration::from_millis) + } + + pub fn client_idle_timeout(&self) -> Duration { + Duration::from_millis(self.client_idle_timeout) + } + + pub fn connect_attempt_delay(&self) -> Duration { + Duration::from_millis(self.connect_attempt_delay) + } + + pub fn client_idle_in_transaction_timeout(&self) -> Duration { + Duration::from_millis(self.client_idle_in_transaction_timeout) + } + + fn load_balancing_strategy() -> LoadBalancingStrategy { + Self::env_enum_or_default("PGDOG_LOAD_BALANCING_STRATEGY") + } + + fn default_tls_verify() -> TlsVerifyMode { + env::var("PGDOG_TLS_VERIFY") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(TlsVerifyMode::Prefer) + } + + fn default_shutdown_timeout() -> u64 { + Self::env_or_default("PGDOG_SHUTDOWN_TIMEOUT", 60_000) + } + + fn default_shutdown_termination_timeout() -> Option { + Self::env_option("PGDOG_SHUTDOWN_TERMINATION_TIMEOUT") + } + + fn default_connect_timeout() -> u64 { + Self::env_or_default("PGDOG_CONNECT_TIMEOUT", 5_000) + } + + fn default_connect_attempt_delay() -> u64 { + Self::env_or_default("PGDOG_CONNECT_ATTEMPT_DELAY", 0) + } + + fn connect_attempts() -> u64 { + Self::env_or_default("PGDOG_CONNECT_ATTEMPTS", 1) + } + + fn pooler_mode() -> PoolerMode { + Self::env_enum_or_default("PGDOG_POOLER_MODE") + } + + fn lsn_check_timeout() -> u64 { + Self::env_or_default("PGDOG_LSN_CHECK_TIMEOUT", 5_000) + } + + fn lsn_check_interval() -> u64 { + Self::env_or_default("PGDOG_LSN_CHECK_INTERVAL", 5_000) + } + + fn lsn_check_delay() -> u64 { + Self::env_or_default("PGDOG_LSN_CHECK_DELAY", Duration::MAX.as_millis() as u64) + } + + fn read_write_strategy() -> ReadWriteStrategy { + Self::env_enum_or_default("PGDOG_READ_WRITE_STRATEGY") + } + + fn read_write_split() -> ReadWriteSplit { + Self::env_enum_or_default("PGDOG_READ_WRITE_SPLIT") + } + + fn prepared_statements() -> PreparedStatements { + Self::env_enum_or_default("PGDOG_PREPARED_STATEMENTS") + } + + fn query_parser_enabled() -> bool { + Self::env_bool_or_default("PGDOG_QUERY_PARSER_ENABLED", false) + } + + fn auth_type() -> AuthType { + Self::env_enum_or_default("PGDOG_AUTH_TYPE") + } + + fn tls_certificate() -> Option { + Self::env_option_string("PGDOG_TLS_CERTIFICATE").map(PathBuf::from) + } + + fn tls_private_key() -> Option { + Self::env_option_string("PGDOG_TLS_PRIVATE_KEY").map(PathBuf::from) + } + + fn tls_server_ca_certificate() -> Option { + Self::env_option_string("PGDOG_TLS_SERVER_CA_CERTIFICATE").map(PathBuf::from) + } + + fn query_log() -> Option { + Self::env_option_string("PGDOG_QUERY_LOG").map(PathBuf::from) + } + + pub fn openmetrics_port() -> Option { + Self::env_option("PGDOG_OPENMETRICS_PORT") + } + + pub fn openmetrics_namespace() -> Option { + Self::env_option_string("PGDOG_OPENMETRICS_NAMESPACE") + } + + fn default_dns_ttl() -> Option { + Self::env_option("PGDOG_DNS_TTL") + } + + pub fn pub_sub_channel_size() -> usize { + Self::env_or_default("PGDOG_PUB_SUB_CHANNEL_SIZE", 0) + } + + pub fn dry_run() -> bool { + Self::env_bool_or_default("PGDOG_DRY_RUN", false) + } + + pub fn cross_shard_disabled() -> bool { + Self::env_bool_or_default("PGDOG_CROSS_SHARD_DISABLED", false) + } + + pub fn broadcast_address() -> Option { + Self::env_option("PGDOG_BROADCAST_ADDRESS") + } + + pub fn broadcast_port() -> u16 { + Self::env_or_default("PGDOG_BROADCAST_PORT", Self::port() + 1) + } + + fn healthcheck_timeout() -> u64 { + Self::env_or_default( + "PGDOG_HEALTHCHECK_TIMEOUT", + Duration::from_secs(5).as_millis() as u64, + ) + } + + fn checkout_timeout() -> u64 { + Self::env_or_default( + "PGDOG_CHECKOUT_TIMEOUT", + Duration::from_secs(5).as_millis() as u64, + ) + } + + pub fn mirror_queue() -> usize { + Self::env_or_default("PGDOG_MIRROR_QUEUE", 128) + } + + pub fn mirror_exposure() -> f32 { + Self::env_or_default("PGDOG_MIRROR_EXPOSURE", 1.0) + } + + pub fn prepared_statements_limit() -> usize { + Self::env_or_default("PGDOG_PREPARED_STATEMENTS_LIMIT", usize::MAX) + } + + pub fn query_cache_limit() -> usize { + Self::env_or_default("PGDOG_QUERY_CACHE_LIMIT", 50_000) + } + + pub fn log_connections() -> bool { + Self::env_bool_or_default("PGDOG_LOG_CONNECTIONS", true) + } + + pub fn log_disconnections() -> bool { + Self::env_bool_or_default("PGDOG_LOG_DISCONNECTIONS", true) + } + + pub fn expanded_explain() -> bool { + Self::env_bool_or_default("PGDOG_EXPANDED_EXPLAIN", false) + } + + pub fn server_lifetime() -> u64 { + Self::env_or_default( + "PGDOG_SERVER_LIFETIME", + Duration::from_secs(3600 * 24).as_millis() as u64, + ) + } + + pub fn connection_recovery() -> ConnectionRecovery { + Self::env_enum_or_default("PGDOG_CONNECTION_RECOVERY") + } + + fn stats_period() -> u64 { + Self::env_or_default("PGDOG_STATS_PERIOD", 15_000) + } + + fn default_passthrough_auth() -> PassthoughAuth { + if let Ok(auth) = env::var("PGDOG_PASSTHROUGH_AUTH") { + // TODO: figure out why toml::from_str doesn't work. + match auth.as_str() { + "enabled" => PassthoughAuth::Enabled, + "disabled" => PassthoughAuth::Disabled, + "enabled_plain" => PassthoughAuth::EnabledPlain, + _ => PassthoughAuth::default(), + } + } else { + PassthoughAuth::default() + } + } + + /// Get shutdown timeout as a duration. + pub fn shutdown_timeout(&self) -> Duration { + Duration::from_millis(self.shutdown_timeout) + } + + pub fn shutdown_termination_timeout(&self) -> Option { + self.shutdown_termination_timeout.map(Duration::from_millis) + } + + /// Get TLS config, if any. + pub fn tls(&self) -> Option<(&PathBuf, &PathBuf)> { + if let Some(cert) = &self.tls_certificate { + if let Some(key) = &self.tls_private_key { + return Some((cert, key)); + } + } + + None + } + + pub fn passthrough_auth(&self) -> bool { + self.tls().is_some() && self.passthrough_auth == PassthoughAuth::Enabled + || self.passthrough_auth == PassthoughAuth::EnabledPlain + } + + /// Support for LISTEN/NOTIFY. + pub fn pub_sub_enabled(&self) -> bool { + self.pub_sub_channel_size > 0 + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::env; + + #[test] + fn test_env_workers() { + env::set_var("PGDOG_WORKERS", "8"); + assert_eq!(General::workers(), 8); + env::remove_var("PGDOG_WORKERS"); + assert_eq!(General::workers(), 2); + } + + #[test] + fn test_env_pool_sizes() { + env::set_var("PGDOG_DEFAULT_POOL_SIZE", "50"); + env::set_var("PGDOG_MIN_POOL_SIZE", "5"); + + assert_eq!(General::default_pool_size(), 50); + assert_eq!(General::min_pool_size(), 5); + + env::remove_var("PGDOG_DEFAULT_POOL_SIZE"); + env::remove_var("PGDOG_MIN_POOL_SIZE"); + + assert_eq!(General::default_pool_size(), 10); + assert_eq!(General::min_pool_size(), 1); + } + + #[test] + fn test_env_timeouts() { + env::set_var("PGDOG_HEALTHCHECK_INTERVAL", "60000"); + env::set_var("PGDOG_HEALTHCHECK_TIMEOUT", "10000"); + env::set_var("PGDOG_CONNECT_TIMEOUT", "10000"); + env::set_var("PGDOG_CHECKOUT_TIMEOUT", "15000"); + env::set_var("PGDOG_IDLE_TIMEOUT", "120000"); + + assert_eq!(General::healthcheck_interval(), 60000); + assert_eq!(General::healthcheck_timeout(), 10000); + assert_eq!(General::default_connect_timeout(), 10000); + assert_eq!(General::checkout_timeout(), 15000); + assert_eq!(General::idle_timeout(), 120000); + + env::remove_var("PGDOG_HEALTHCHECK_INTERVAL"); + env::remove_var("PGDOG_HEALTHCHECK_TIMEOUT"); + env::remove_var("PGDOG_CONNECT_TIMEOUT"); + env::remove_var("PGDOG_CHECKOUT_TIMEOUT"); + env::remove_var("PGDOG_IDLE_TIMEOUT"); + + assert_eq!(General::healthcheck_interval(), 30000); + assert_eq!(General::healthcheck_timeout(), 5000); + assert_eq!(General::default_connect_timeout(), 5000); + assert_eq!(General::checkout_timeout(), 5000); + assert_eq!(General::idle_timeout(), 60000); + } + + #[test] + fn test_env_invalid_values() { + env::set_var("PGDOG_WORKERS", "invalid"); + env::set_var("PGDOG_DEFAULT_POOL_SIZE", "not_a_number"); + + assert_eq!(General::workers(), 2); + assert_eq!(General::default_pool_size(), 10); + + env::remove_var("PGDOG_WORKERS"); + env::remove_var("PGDOG_DEFAULT_POOL_SIZE"); + } + + #[test] + fn test_env_host_port() { + // Test existing env var functionality + env::set_var("PGDOG_HOST", "192.168.1.1"); + env::set_var("PGDOG_PORT", "8432"); + + assert_eq!(General::host(), "192.168.1.1"); + assert_eq!(General::port(), 8432); + + env::remove_var("PGDOG_HOST"); + env::remove_var("PGDOG_PORT"); + + assert_eq!(General::host(), "0.0.0.0"); + assert_eq!(General::port(), 6432); + } + + #[test] + fn test_env_enum_fields() { + // Test pooler mode + env::set_var("PGDOG_POOLER_MODE", "session"); + assert_eq!(General::pooler_mode(), PoolerMode::Session); + env::remove_var("PGDOG_POOLER_MODE"); + assert_eq!(General::pooler_mode(), PoolerMode::Transaction); + + // Test load balancing strategy + env::set_var("PGDOG_LOAD_BALANCING_STRATEGY", "round_robin"); + assert_eq!( + General::load_balancing_strategy(), + LoadBalancingStrategy::RoundRobin + ); + env::remove_var("PGDOG_LOAD_BALANCING_STRATEGY"); + assert_eq!( + General::load_balancing_strategy(), + LoadBalancingStrategy::Random + ); + + // Test read-write strategy + env::set_var("PGDOG_READ_WRITE_STRATEGY", "aggressive"); + assert_eq!( + General::read_write_strategy(), + ReadWriteStrategy::Aggressive + ); + env::remove_var("PGDOG_READ_WRITE_STRATEGY"); + assert_eq!( + General::read_write_strategy(), + ReadWriteStrategy::Conservative + ); + + // Test read-write split + env::set_var("PGDOG_READ_WRITE_SPLIT", "exclude_primary"); + assert_eq!(General::read_write_split(), ReadWriteSplit::ExcludePrimary); + env::remove_var("PGDOG_READ_WRITE_SPLIT"); + assert_eq!(General::read_write_split(), ReadWriteSplit::IncludePrimary); + + // Test TLS verify mode + env::set_var("PGDOG_TLS_VERIFY", "verify_full"); + assert_eq!(General::default_tls_verify(), TlsVerifyMode::VerifyFull); + env::remove_var("PGDOG_TLS_VERIFY"); + assert_eq!(General::default_tls_verify(), TlsVerifyMode::Prefer); + + // Test prepared statements + env::set_var("PGDOG_PREPARED_STATEMENTS", "full"); + assert_eq!(General::prepared_statements(), PreparedStatements::Full); + env::remove_var("PGDOG_PREPARED_STATEMENTS"); + assert_eq!(General::prepared_statements(), PreparedStatements::Extended); + + // Test auth type + env::set_var("PGDOG_AUTH_TYPE", "md5"); + assert_eq!(General::auth_type(), AuthType::Md5); + env::remove_var("PGDOG_AUTH_TYPE"); + assert_eq!(General::auth_type(), AuthType::Scram); + } + + #[test] + fn test_env_additional_timeouts() { + env::set_var("PGDOG_IDLE_HEALTHCHECK_INTERVAL", "45000"); + env::set_var("PGDOG_IDLE_HEALTHCHECK_DELAY", "10000"); + env::set_var("PGDOG_BAN_TIMEOUT", "600000"); + env::set_var("PGDOG_ROLLBACK_TIMEOUT", "10000"); + env::set_var("PGDOG_SHUTDOWN_TIMEOUT", "120000"); + env::set_var("PGDOG_SHUTDOWN_TERMINATION_TIMEOUT", "15000"); + env::set_var("PGDOG_CONNECT_ATTEMPT_DELAY", "1000"); + env::set_var("PGDOG_QUERY_TIMEOUT", "30000"); + env::set_var("PGDOG_CLIENT_IDLE_TIMEOUT", "3600000"); + + assert_eq!(General::idle_healthcheck_interval(), 45000); + assert_eq!(General::idle_healthcheck_delay(), 10000); + assert_eq!(General::ban_timeout(), 600000); + assert_eq!(General::rollback_timeout(), 10000); + assert_eq!(General::default_shutdown_timeout(), 120000); + assert_eq!( + General::default_shutdown_termination_timeout(), + Some(15_000) + ); + assert_eq!(General::default_connect_attempt_delay(), 1000); + assert_eq!(General::default_query_timeout(), 30000); + assert_eq!(General::default_client_idle_timeout(), 3600000); + + env::remove_var("PGDOG_IDLE_HEALTHCHECK_INTERVAL"); + env::remove_var("PGDOG_IDLE_HEALTHCHECK_DELAY"); + env::remove_var("PGDOG_BAN_TIMEOUT"); + env::remove_var("PGDOG_ROLLBACK_TIMEOUT"); + env::remove_var("PGDOG_SHUTDOWN_TIMEOUT"); + env::remove_var("PGDOG_SHUTDOWN_TERMINATION_TIMEOUT"); + env::remove_var("PGDOG_CONNECT_ATTEMPT_DELAY"); + env::remove_var("PGDOG_QUERY_TIMEOUT"); + env::remove_var("PGDOG_CLIENT_IDLE_TIMEOUT"); + + assert_eq!(General::idle_healthcheck_interval(), 30000); + assert_eq!(General::idle_healthcheck_delay(), 5000); + assert_eq!(General::ban_timeout(), 300000); + assert_eq!(General::rollback_timeout(), 5000); + assert_eq!(General::default_shutdown_timeout(), 60000); + assert_eq!(General::default_shutdown_termination_timeout(), None); + assert_eq!(General::default_connect_attempt_delay(), 0); + } + + #[test] + fn test_env_path_fields() { + env::set_var("PGDOG_TLS_CERTIFICATE", "/path/to/cert.pem"); + env::set_var("PGDOG_TLS_PRIVATE_KEY", "/path/to/key.pem"); + env::set_var("PGDOG_TLS_SERVER_CA_CERTIFICATE", "/path/to/ca.pem"); + env::set_var("PGDOG_QUERY_LOG", "/var/log/pgdog/queries.log"); + + assert_eq!( + General::tls_certificate(), + Some(PathBuf::from("/path/to/cert.pem")) + ); + assert_eq!( + General::tls_private_key(), + Some(PathBuf::from("/path/to/key.pem")) + ); + assert_eq!( + General::tls_server_ca_certificate(), + Some(PathBuf::from("/path/to/ca.pem")) + ); + assert_eq!( + General::query_log(), + Some(PathBuf::from("/var/log/pgdog/queries.log")) + ); + + env::remove_var("PGDOG_TLS_CERTIFICATE"); + env::remove_var("PGDOG_TLS_PRIVATE_KEY"); + env::remove_var("PGDOG_TLS_SERVER_CA_CERTIFICATE"); + env::remove_var("PGDOG_QUERY_LOG"); + + assert_eq!(General::tls_certificate(), None); + assert_eq!(General::tls_private_key(), None); + assert_eq!(General::tls_server_ca_certificate(), None); + assert_eq!(General::query_log(), None); + } + + #[test] + fn test_env_numeric_fields() { + env::set_var("PGDOG_BROADCAST_PORT", "7432"); + env::set_var("PGDOG_OPENMETRICS_PORT", "9090"); + env::set_var("PGDOG_PREPARED_STATEMENTS_LIMIT", "1000"); + env::set_var("PGDOG_QUERY_CACHE_LIMIT", "500"); + env::set_var("PGDOG_CONNECT_ATTEMPTS", "3"); + env::set_var("PGDOG_MIRROR_QUEUE", "256"); + env::set_var("PGDOG_MIRROR_EXPOSURE", "0.5"); + env::set_var("PGDOG_DNS_TTL", "60000"); + env::set_var("PGDOG_PUB_SUB_CHANNEL_SIZE", "100"); + + assert_eq!(General::broadcast_port(), 7432); + assert_eq!(General::openmetrics_port(), Some(9090)); + assert_eq!(General::prepared_statements_limit(), 1000); + assert_eq!(General::query_cache_limit(), 500); + assert_eq!(General::connect_attempts(), 3); + assert_eq!(General::mirror_queue(), 256); + assert_eq!(General::mirror_exposure(), 0.5); + assert_eq!(General::default_dns_ttl(), Some(60000)); + assert_eq!(General::pub_sub_channel_size(), 100); + + env::remove_var("PGDOG_BROADCAST_PORT"); + env::remove_var("PGDOG_OPENMETRICS_PORT"); + env::remove_var("PGDOG_PREPARED_STATEMENTS_LIMIT"); + env::remove_var("PGDOG_QUERY_CACHE_LIMIT"); + env::remove_var("PGDOG_CONNECT_ATTEMPTS"); + env::remove_var("PGDOG_MIRROR_QUEUE"); + env::remove_var("PGDOG_MIRROR_EXPOSURE"); + env::remove_var("PGDOG_DNS_TTL"); + env::remove_var("PGDOG_PUB_SUB_CHANNEL_SIZE"); + + assert_eq!(General::broadcast_port(), General::port() + 1); + assert_eq!(General::openmetrics_port(), None); + assert_eq!(General::prepared_statements_limit(), usize::MAX); + assert_eq!(General::query_cache_limit(), 50_000); + assert_eq!(General::connect_attempts(), 1); + assert_eq!(General::mirror_queue(), 128); + assert_eq!(General::mirror_exposure(), 1.0); + assert_eq!(General::default_dns_ttl(), None); + assert_eq!(General::pub_sub_channel_size(), 0); + } + + #[test] + fn test_env_boolean_fields() { + env::set_var("PGDOG_DRY_RUN", "true"); + env::set_var("PGDOG_CROSS_SHARD_DISABLED", "yes"); + env::set_var("PGDOG_LOG_CONNECTIONS", "false"); + env::set_var("PGDOG_LOG_DISCONNECTIONS", "0"); + + assert_eq!(General::dry_run(), true); + assert_eq!(General::cross_shard_disabled(), true); + assert_eq!(General::log_connections(), false); + assert_eq!(General::log_disconnections(), false); + + env::remove_var("PGDOG_DRY_RUN"); + env::remove_var("PGDOG_CROSS_SHARD_DISABLED"); + env::remove_var("PGDOG_LOG_CONNECTIONS"); + env::remove_var("PGDOG_LOG_DISCONNECTIONS"); + + assert_eq!(General::dry_run(), false); + assert_eq!(General::cross_shard_disabled(), false); + assert_eq!(General::log_connections(), true); + assert_eq!(General::log_disconnections(), true); + } + + #[test] + fn test_env_other_fields() { + env::set_var("PGDOG_BROADCAST_ADDRESS", "192.168.1.100"); + env::set_var("PGDOG_OPENMETRICS_NAMESPACE", "pgdog_metrics"); + + assert_eq!( + General::broadcast_address(), + Some("192.168.1.100".parse().unwrap()) + ); + assert_eq!( + General::openmetrics_namespace(), + Some("pgdog_metrics".to_string()) + ); + + env::remove_var("PGDOG_BROADCAST_ADDRESS"); + env::remove_var("PGDOG_OPENMETRICS_NAMESPACE"); + + assert_eq!(General::broadcast_address(), None); + assert_eq!(General::openmetrics_namespace(), None); + } + + #[test] + fn test_env_invalid_enum_values() { + env::set_var("PGDOG_POOLER_MODE", "invalid_mode"); + env::set_var("PGDOG_AUTH_TYPE", "not_an_auth"); + env::set_var("PGDOG_TLS_VERIFY", "bad_verify"); + + // Should fall back to defaults for invalid values + assert_eq!(General::pooler_mode(), PoolerMode::Transaction); + assert_eq!(General::auth_type(), AuthType::Scram); + assert_eq!(General::default_tls_verify(), TlsVerifyMode::Prefer); + + env::remove_var("PGDOG_POOLER_MODE"); + env::remove_var("PGDOG_AUTH_TYPE"); + env::remove_var("PGDOG_TLS_VERIFY"); + } + + #[test] + fn test_general_default_uses_env_vars() { + // Set some environment variables + env::set_var("PGDOG_WORKERS", "8"); + env::set_var("PGDOG_POOLER_MODE", "session"); + env::set_var("PGDOG_AUTH_TYPE", "trust"); + env::set_var("PGDOG_DRY_RUN", "true"); + + let general = General::default(); + + assert_eq!(general.workers, 8); + assert_eq!(general.pooler_mode, PoolerMode::Session); + assert_eq!(general.auth_type, AuthType::Trust); + assert_eq!(general.dry_run, true); + + env::remove_var("PGDOG_WORKERS"); + env::remove_var("PGDOG_POOLER_MODE"); + env::remove_var("PGDOG_AUTH_TYPE"); + env::remove_var("PGDOG_DRY_RUN"); + } +} diff --git a/pgdog-config/src/lib.rs b/pgdog-config/src/lib.rs new file mode 100644 index 000000000..6e47ff147 --- /dev/null +++ b/pgdog-config/src/lib.rs @@ -0,0 +1,34 @@ +// Submodules +pub mod auth; +pub mod core; +pub mod data_types; +pub mod database; +pub mod error; +pub mod general; +pub mod memory; +pub mod networking; +pub mod overrides; +pub mod pooling; +pub mod replication; +pub mod rewrite; +pub mod sharding; +pub mod url; +pub mod users; +pub mod util; + +pub use auth::{AuthType, PassthoughAuth}; +pub use core::{Config, ConfigAndUsers}; +pub use data_types::*; +pub use database::{ + Database, EnumeratedDatabase, LoadBalancingStrategy, ReadWriteSplit, ReadWriteStrategy, Role, +}; +pub use error::Error; +pub use general::General; +pub use memory::*; +pub use networking::{MultiTenant, Tcp, TlsVerifyMode}; +pub use overrides::Overrides; +pub use pooling::{PoolerMode, PreparedStatements, Stats}; +pub use replication::*; +pub use rewrite::{Rewrite, RewriteMode}; +pub use sharding::*; +pub use users::{Admin, Plugin, User, Users}; diff --git a/pgdog-config/src/memory.rs b/pgdog-config/src/memory.rs new file mode 100644 index 000000000..f2d50c64a --- /dev/null +++ b/pgdog-config/src/memory.rs @@ -0,0 +1,35 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Serialize, Deserialize, Debug, Clone)] +#[serde(deny_unknown_fields)] +pub struct Memory { + #[serde(default = "default_net_buffer")] + pub net_buffer: usize, + #[serde(default = "default_message_buffer")] + pub message_buffer: usize, + #[serde(default = "default_stack_size")] + pub stack_size: usize, +} + +impl Default for Memory { + fn default() -> Self { + Self { + net_buffer: default_net_buffer(), + message_buffer: default_message_buffer(), + stack_size: default_stack_size(), + } + } +} + +fn default_net_buffer() -> usize { + 4096 +} + +fn default_message_buffer() -> usize { + default_net_buffer() +} + +// Default: 2MiB. +fn default_stack_size() -> usize { + 2 * 1024 * 1024 +} diff --git a/pgdog-config/src/networking.rs b/pgdog-config/src/networking.rs new file mode 100644 index 000000000..830f49603 --- /dev/null +++ b/pgdog-config/src/networking.rs @@ -0,0 +1,113 @@ +use serde::{Deserialize, Serialize}; +use std::str::FromStr; +use std::time::Duration; + +use crate::util::human_duration_optional; + +#[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq, Copy)] +#[serde(rename_all = "snake_case")] +pub enum TlsVerifyMode { + #[default] + Disabled, + Prefer, + VerifyCa, + VerifyFull, +} + +impl FromStr for TlsVerifyMode { + type Err = String; + + fn from_str(s: &str) -> Result { + match s.to_lowercase().replace(['_', '-'], "").as_str() { + "disabled" => Ok(Self::Disabled), + "prefer" => Ok(Self::Prefer), + "verifyca" => Ok(Self::VerifyCa), + "verifyfull" => Ok(Self::VerifyFull), + _ => Err(format!("Invalid TLS verify mode: {}", s)), + } + } +} + +#[derive(Serialize, Deserialize, PartialEq, Debug, Clone)] +#[serde(deny_unknown_fields)] +pub struct Tcp { + #[serde(default = "Tcp::default_keepalive")] + keepalive: bool, + user_timeout: Option, + time: Option, + interval: Option, + retries: Option, + congestion_control: Option, +} + +impl std::fmt::Display for Tcp { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "keepalive={} user_timeout={} time={} interval={}, retries={}, congestion_control={}", + self.keepalive(), + human_duration_optional(self.user_timeout()), + human_duration_optional(self.time()), + human_duration_optional(self.interval()), + if let Some(retries) = self.retries() { + retries.to_string() + } else { + "default".into() + }, + if let Some(ref c) = self.congestion_control { + c.as_str() + } else { + "" + }, + ) + } +} + +impl Default for Tcp { + fn default() -> Self { + Self { + keepalive: Self::default_keepalive(), + user_timeout: None, + time: None, + interval: None, + retries: None, + congestion_control: None, + } + } +} + +impl Tcp { + fn default_keepalive() -> bool { + true + } + + pub fn keepalive(&self) -> bool { + self.keepalive + } + + pub fn time(&self) -> Option { + self.time.map(Duration::from_millis) + } + + pub fn interval(&self) -> Option { + self.interval.map(Duration::from_millis) + } + + pub fn user_timeout(&self) -> Option { + self.user_timeout.map(Duration::from_millis) + } + + pub fn retries(&self) -> Option { + self.retries + } + + pub fn congestion_control(&self) -> &Option { + &self.congestion_control + } +} + +#[derive(Serialize, Deserialize, PartialEq, Debug, Clone)] +#[serde(rename_all = "snake_case")] +pub struct MultiTenant { + pub column: String, +} diff --git a/pgdog-config/src/overrides.rs b/pgdog-config/src/overrides.rs new file mode 100644 index 000000000..1769713dc --- /dev/null +++ b/pgdog-config/src/overrides.rs @@ -0,0 +1,6 @@ +#[derive(Debug, Clone, Default)] +pub struct Overrides { + pub default_pool_size: Option, + pub min_pool_size: Option, + pub session_mode: Option, +} diff --git a/pgdog-config/src/pooling.rs b/pgdog-config/src/pooling.rs new file mode 100644 index 000000000..eebaefbd3 --- /dev/null +++ b/pgdog-config/src/pooling.rs @@ -0,0 +1,110 @@ +use serde::{Deserialize, Serialize}; +use std::str::FromStr; + +#[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq, Copy)] +#[serde(rename_all = "snake_case")] +pub enum PreparedStatements { + Disabled, + #[default] + Extended, + ExtendedAnonymous, + Full, +} + +impl PreparedStatements { + pub fn full(&self) -> bool { + matches!(self, PreparedStatements::Full) + } + + pub fn enabled(&self) -> bool { + !matches!(self, PreparedStatements::Disabled) + } + + pub fn rewrite_anonymous(&self) -> bool { + matches!(self, PreparedStatements::ExtendedAnonymous) + } +} + +impl FromStr for PreparedStatements { + type Err = String; + + fn from_str(s: &str) -> Result { + match s.to_lowercase().as_str() { + "disabled" => Ok(Self::Disabled), + "extended" => Ok(Self::Extended), + "extended_anonymous" => Ok(Self::ExtendedAnonymous), + "full" => Ok(Self::Full), + _ => Err(format!("Invalid prepared statements mode: {}", s)), + } + } +} + +/// Empty struct for stats +#[derive(Serialize, Deserialize, Debug, Clone, Default)] +pub struct Stats {} + +#[derive(Serialize, Deserialize, Debug, Clone, Copy, Default, PartialEq, Eq, Ord, PartialOrd)] +#[serde(rename_all = "snake_case")] +pub enum PoolerMode { + #[default] + Transaction, + Session, + Statement, +} + +impl std::fmt::Display for PoolerMode { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Transaction => write!(f, "transaction"), + Self::Session => write!(f, "session"), + Self::Statement => write!(f, "statement"), + } + } +} + +impl FromStr for PoolerMode { + type Err = String; + + fn from_str(s: &str) -> Result { + match s.to_lowercase().as_str() { + "transaction" => Ok(Self::Transaction), + "session" => Ok(Self::Session), + _ => Err(format!("Invalid pooler mode: {}", s)), + } + } +} + +#[derive(Serialize, Deserialize, Debug, Clone, Copy, Default, PartialEq, Eq, Ord, PartialOrd)] +#[serde(rename_all = "snake_case")] +pub enum ConnectionRecovery { + #[default] + Recover, + RollbackOnly, + Drop, +} + +impl ConnectionRecovery { + pub fn can_recover(&self) -> bool { + matches!(self, ConnectionRecovery::Recover) + } + + pub fn can_rollback(&self) -> bool { + matches!( + self, + ConnectionRecovery::Recover | ConnectionRecovery::RollbackOnly + ) + } +} + +impl FromStr for ConnectionRecovery { + type Err = String; + + fn from_str(s: &str) -> Result { + match s.to_lowercase().as_str() { + "recover" => Ok(Self::Recover), + "rollbackonly" => Ok(Self::RollbackOnly), + "drop" => Ok(Self::Drop), + _ => Err(format!("Invalid pooler mode: {}", s)), + } + } +} diff --git a/pgdog-config/src/replication.rs b/pgdog-config/src/replication.rs new file mode 100644 index 000000000..e28e68b1f --- /dev/null +++ b/pgdog-config/src/replication.rs @@ -0,0 +1,190 @@ +use serde::{Deserialize, Serialize}; +use std::path::PathBuf; +use std::str::FromStr; +use std::time::Duration; + +#[derive(Deserialize)] +struct RawReplicaLag { + #[serde(default)] + check_interval: Option, + #[serde(default)] + max_age: Option, +} + +#[derive(Debug, Clone)] +pub struct ReplicaLag { + pub check_interval: Duration, + pub max_age: Duration, +} + +impl ReplicaLag { + fn default_max_age() -> Duration { + Duration::from_millis(25) + } + + fn default_check_interval() -> Duration { + Duration::from_millis(1000) + } + + /// Custom "all-or-none" deserializer that returns Option. + pub fn deserialize_optional<'de, D>(de: D) -> Result, D::Error> + where + D: serde::de::Deserializer<'de>, + { + let maybe: Option = Option::deserialize(de)?; + + Ok(match maybe { + None => None, + + Some(RawReplicaLag { + check_interval: None, + max_age: None, + }) => None, + + Some(RawReplicaLag { + check_interval: Some(ci_u64), + max_age: Some(ma_u64), + }) => Some(ReplicaLag { + check_interval: Duration::from_millis(ci_u64), + max_age: Duration::from_millis(ma_u64), + }), + + Some(RawReplicaLag { + check_interval: None, + max_age: Some(ma_u64), + }) => Some(ReplicaLag { + check_interval: Self::default_check_interval(), + max_age: Duration::from_millis(ma_u64), + }), + + _ => { + return Err(serde::de::Error::custom( + "replica_lag: cannot set check_interval without max_age", + )) + } + }) + } +} + +// NOTE: serialize and deserialize are not inverses. +// - Normally you'd expect ser <-> deser to round-trip, but here deser applies defaults... +// for missing fields +// - Serializes takes those applied defaults into account so that ReplicaLag always reflects... +// the actual effective values. +// - This ensures pgdog.admin sees the true config that is applied, not just what was configured. + +impl Serialize for ReplicaLag { + fn serialize(&self, serializer: S) -> Result + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut state = serializer.serialize_struct("ReplicaLag", 2)?; + state.serialize_field("check_interval", &(self.check_interval.as_millis() as u64))?; + state.serialize_field("max_age", &(self.max_age.as_millis() as u64))?; + state.end() + } +} + +impl Default for ReplicaLag { + fn default() -> Self { + Self { + check_interval: Self::default_check_interval(), + max_age: Self::default_max_age(), + } + } +} + +/// Replication configuration. +#[derive(Serialize, Deserialize, PartialEq, Debug, Clone)] +#[serde(rename_all = "snake_case")] +pub struct Replication { + /// Path to the pg_dump executable. + #[serde(default = "Replication::pg_dump_path")] + pub pg_dump_path: PathBuf, +} + +impl Replication { + fn pg_dump_path() -> PathBuf { + PathBuf::from("pg_dump") + } +} + +impl Default for Replication { + fn default() -> Self { + Self { + pg_dump_path: Self::pg_dump_path(), + } + } +} + +/// Mirroring configuration. +#[derive(Serialize, Deserialize, Debug, Clone, Default)] +#[serde(deny_unknown_fields)] +pub struct Mirroring { + /// Source database name to mirror from. + pub source_db: String, + /// Destination database name to mirror to. + pub destination_db: String, + /// Queue length for this mirror (overrides global mirror_queue). + pub queue_length: Option, + /// Exposure for this mirror (overrides global mirror_exposure). + pub exposure: Option, +} + +impl FromStr for Mirroring { + type Err = String; + + fn from_str(s: &str) -> Result { + let mut source_db = None; + let mut destination_db = None; + let mut queue_length = None; + let mut exposure = None; + + for pair in s.split('&') { + let parts: Vec<&str> = pair.split('=').collect(); + if parts.len() != 2 { + return Err(format!("Invalid key=value pair: {}", pair)); + } + + match parts[0] { + "source_db" => source_db = Some(parts[1].to_string()), + "destination_db" => destination_db = Some(parts[1].to_string()), + "queue_length" => { + queue_length = Some( + parts[1] + .parse::() + .map_err(|_| format!("Invalid queue_length: {}", parts[1]))?, + ); + } + "exposure" => { + exposure = Some( + parts[1] + .parse::() + .map_err(|_| format!("Invalid exposure: {}", parts[1]))?, + ); + } + _ => return Err(format!("Unknown parameter: {}", parts[0])), + } + } + + let source_db = source_db.ok_or("Missing required parameter: source_db")?; + let destination_db = destination_db.ok_or("Missing required parameter: destination_db")?; + + Ok(Mirroring { + source_db, + destination_db, + queue_length, + exposure, + }) + } +} + +/// Runtime mirror configuration with resolved values. +#[derive(Debug, Clone)] +pub struct MirrorConfig { + /// Queue length for this mirror. + pub queue_length: usize, + /// Exposure for this mirror. + pub exposure: f32, +} diff --git a/pgdog-config/src/rewrite.rs b/pgdog-config/src/rewrite.rs new file mode 100644 index 000000000..957633194 --- /dev/null +++ b/pgdog-config/src/rewrite.rs @@ -0,0 +1,76 @@ +use serde::{Deserialize, Serialize}; +use std::fmt; +use std::str::FromStr; + +#[derive(Serialize, Deserialize, Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +#[serde(rename_all = "lowercase")] +pub enum RewriteMode { + Ignore, + Error, + Rewrite, +} + +impl Default for RewriteMode { + fn default() -> Self { + Self::Error + } +} + +impl fmt::Display for RewriteMode { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let value = match self { + RewriteMode::Error => "error", + RewriteMode::Rewrite => "rewrite", + RewriteMode::Ignore => "ignore", + }; + f.write_str(value) + } +} + +impl FromStr for RewriteMode { + type Err = (); + + fn from_str(s: &str) -> Result { + match s.to_ascii_lowercase().as_str() { + "error" => Ok(RewriteMode::Error), + "rewrite" => Ok(RewriteMode::Rewrite), + "ignore" => Ok(RewriteMode::Ignore), + _ => Err(()), + } + } +} + +#[derive(Serialize, Deserialize, Debug, Clone)] +#[serde(deny_unknown_fields)] +pub struct Rewrite { + /// Global rewrite toggle. When disabled, rewrite-specific features remain + /// inactive, even if individual policies request rewriting. + #[serde(default)] + pub enabled: bool, + /// Policy for handling shard-key updates. + #[serde(default = "Rewrite::default_shard_key")] + pub shard_key: RewriteMode, + /// Policy for handling multi-row INSERT statements that target sharded tables. + #[serde(default = "Rewrite::default_split_inserts")] + pub split_inserts: RewriteMode, +} + +impl Default for Rewrite { + fn default() -> Self { + Self { + enabled: false, + shard_key: Self::default_shard_key(), + split_inserts: Self::default_split_inserts(), + } + } +} + +impl Rewrite { + const fn default_shard_key() -> RewriteMode { + RewriteMode::Error + } + + const fn default_split_inserts() -> RewriteMode { + RewriteMode::Error + } +} diff --git a/pgdog-config/src/sharding.rs b/pgdog-config/src/sharding.rs new file mode 100644 index 000000000..4dd177e40 --- /dev/null +++ b/pgdog-config/src/sharding.rs @@ -0,0 +1,316 @@ +use serde::{Deserialize, Serialize}; +use std::collections::HashMap; +use std::collections::{hash_map::DefaultHasher, HashSet}; +use std::hash::{Hash, Hasher as StdHasher}; +use std::path::PathBuf; +use tracing::{info, warn}; + +use super::error::Error; +use pgdog_vector::Vector; + +/// Sharded table. +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Default)] +#[serde(rename_all = "snake_case", deny_unknown_fields)] +pub struct ShardedTable { + /// Database this table belongs to. + pub database: String, + /// Table name. If none specified, all tables with the specified + /// column are considered sharded. + #[serde(default)] + pub name: Option, + /// Schema name. If not specified, will match all schemas. + #[serde(default)] + pub schema: Option, + /// Table sharded on this column. + #[serde(default)] + pub column: String, + /// This table is the primary sharding anchor (e.g. "users"). + #[serde(default)] + pub primary: bool, + /// Centroids for vector sharding. + #[serde(default)] + pub centroids: Vec, + #[serde(default)] + pub centroids_path: Option, + /// Data type of the column. + #[serde(default)] + pub data_type: DataType, + /// How many centroids to probe. + #[serde(default)] + pub centroid_probes: usize, + /// Hasher function. + #[serde(default)] + pub hasher: Hasher, + /// Explicit routing rules. + #[serde(skip, default)] + pub mapping: Option, +} + +impl ShardedTable { + /// Load centroids from file, if provided. + /// + /// Centroids can be very large vectors (1000+ columns). + /// Hardcoding them in pgdog.toml is then impractical. + pub fn load_centroids(&mut self) -> Result<(), Error> { + if let Some(centroids_path) = &self.centroids_path { + if let Ok(f) = std::fs::read_to_string(centroids_path) { + let centroids: Vec = serde_json::from_str(&f)?; + self.centroids = centroids; + info!("loaded {} centroids", self.centroids.len()); + } else { + warn!( + "centroids at path \"{}\" not found", + centroids_path.display() + ); + } + } + + if self.centroid_probes < 1 { + self.centroid_probes = (self.centroids.len() as f32).sqrt().ceil() as usize; + if self.centroid_probes > 0 { + info!("setting centroid probes to {}", self.centroid_probes); + } + } + + Ok(()) + } +} + +#[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq)] +#[serde(rename_all = "snake_case")] +pub enum Hasher { + #[default] + Postgres, + Sha1, +} + +#[derive(Serialize, Deserialize, PartialEq, Debug, Clone, Default, Copy, Eq, Hash)] +#[serde(rename_all = "snake_case")] +pub enum DataType { + #[default] + Bigint, + Uuid, + Vector, + Varchar, +} + +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Default, Eq)] +#[serde(rename_all = "snake_case", deny_unknown_fields)] +pub struct ShardedMapping { + pub database: String, + pub column: String, + pub table: Option, + pub schema: Option, + pub kind: ShardedMappingKind, + pub start: Option, + pub end: Option, + #[serde(default)] + pub values: HashSet, + pub shard: usize, +} + +impl Hash for ShardedMapping { + fn hash(&self, state: &mut H) { + self.database.hash(state); + self.column.hash(state); + self.table.hash(state); + self.kind.hash(state); + self.start.hash(state); + self.end.hash(state); + + // Hash the values in a deterministic way by XORing their individual hashes + let mut values_hash = 0u64; + for value in &self.values { + let mut hasher = DefaultHasher::new(); + value.hash(&mut hasher); + values_hash ^= hasher.finish(); + } + values_hash.hash(state); + + self.shard.hash(state); + } +} + +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Default, Eq, Hash)] +pub struct ShardedMappingKey { + database: String, + column: String, + table: Option, +} + +#[derive(Serialize, Deserialize, PartialEq, Debug, Clone, Default, Hash, Eq)] +#[serde(rename_all = "snake_case", deny_unknown_fields)] +pub enum ShardedMappingKind { + #[default] + List, + Range, +} + +#[derive(Serialize, Deserialize, PartialEq, Debug, Clone, Eq, Hash)] +#[serde(untagged)] +pub enum FlexibleType { + Integer(i64), + Uuid(uuid::Uuid), + String(String), +} + +impl From for FlexibleType { + fn from(value: i64) -> Self { + Self::Integer(value) + } +} + +impl From for FlexibleType { + fn from(value: uuid::Uuid) -> Self { + Self::Uuid(value) + } +} + +impl From for FlexibleType { + fn from(value: String) -> Self { + Self::String(value) + } +} + +#[derive(Serialize, Deserialize, PartialEq, Debug, Clone, Default)] +#[serde(rename_all = "snake_case", deny_unknown_fields)] +pub struct OmnishardedTables { + pub database: String, + pub tables: Vec, + #[serde(default)] + pub sticky: bool, +} + +#[derive(PartialEq, Debug, Clone, Default)] +pub struct OmnishardedTable { + pub name: String, + pub sticky_routing: bool, +} + +/// Queries with manual routing rules. +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, Hash)] +pub struct ManualQuery { + pub fingerprint: String, +} + +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, Hash, Default)] +pub struct ShardedSchema { + /// Database name. + pub database: String, + /// Schema name. + pub name: Option, + #[serde(default)] + pub shard: usize, + /// All shards. + #[serde(default)] + pub all: bool, +} + +impl ShardedSchema { + /// This schema mapping is used to route all other queries. + pub fn is_default(&self) -> bool { + self.name.is_none() + } + + pub fn name(&self) -> &str { + self.name.as_deref().unwrap_or("*") + } + + pub fn shard(&self) -> Option { + if self.all { + None + } else { + Some(self.shard) + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ListShards { + mapping: HashMap, +} + +#[derive(Debug, Clone, Eq, PartialEq, Hash)] +pub enum Mapping { + Range(Vec), // TODO: optimize with a BTreeMap. + List(ListShards), // Optimized. +} + +impl Hash for ListShards { + fn hash(&self, state: &mut H) { + // Hash the mapping in a deterministic way by XORing individual key-value hashes + let mut mapping_hash = 0u64; + for (key, value) in &self.mapping { + let mut hasher = DefaultHasher::new(); + key.hash(&mut hasher); + value.hash(&mut hasher); + mapping_hash ^= hasher.finish(); + } + mapping_hash.hash(state); + } +} + +impl Mapping { + pub fn new(mappings: &[ShardedMapping]) -> Option { + let range = mappings + .iter() + .filter(|m| m.kind == ShardedMappingKind::Range) + .cloned() + .collect::>(); + let list = mappings.iter().any(|m| m.kind == ShardedMappingKind::List); + + if !range.is_empty() { + Some(Self::Range(range)) + } else if list { + Some(Self::List(ListShards::new(mappings))) + } else { + None + } + } +} + +impl ListShards { + pub fn is_empty(&self) -> bool { + self.mapping.is_empty() + } + + pub fn new(mappings: &[ShardedMapping]) -> Self { + let mut mapping = HashMap::new(); + + for map in mappings + .iter() + .filter(|m| m.kind == ShardedMappingKind::List) + { + for value in &map.values { + mapping.insert(value.clone(), map.shard); + } + } + + Self { mapping } + } + + pub fn shard(&self, value: &FlexibleType) -> Result, Error> { + if let Some(shard) = self.mapping.get(value) { + Ok(Some(*shard)) + } else { + Ok(None) + } + } +} + +#[derive(Serialize, Deserialize, Debug, Copy, Clone, PartialEq, Eq, Hash, Default)] +#[serde(rename_all = "snake_case", deny_unknown_fields)] +pub enum QueryParserLevel { + On, + #[default] + Auto, + Off, +} + +#[derive(Serialize, Deserialize, Debug, Copy, Clone, PartialEq, Eq, Hash, Default)] +#[serde(rename_all = "snake_case", deny_unknown_fields)] +pub enum QueryParserEngine { + #[default] + PgQueryProtobuf, + PgQueryRaw, +} diff --git a/pgdog-config/src/url.rs b/pgdog-config/src/url.rs new file mode 100644 index 000000000..6d746fccb --- /dev/null +++ b/pgdog-config/src/url.rs @@ -0,0 +1,304 @@ +//! Parse URL and convert to config struct. +use std::{collections::BTreeSet, env::var, str::FromStr}; +use url::Url; + +use super::{ConfigAndUsers, Database, Error, PoolerMode, Role, User, Users}; + +pub fn database_name(url: &Url) -> String { + let database = url.path().chars().skip(1).collect::(); + if database.is_empty() { + "postgres".into() + } else { + database + } +} + +impl From<&Url> for Database { + fn from(value: &Url) -> Self { + let host = value + .host() + .map(|host| host.to_string()) + .unwrap_or("127.0.0.1".into()); + let port = value.port().unwrap_or(5432); + + let mut database = Database { + name: database_name(value), + host, + port, + ..Default::default() + }; + + for (key, val) in value.query_pairs() { + match key.as_ref() { + "database_name" => database.database_name = Some(val.to_string()), + "role" => { + if let Ok(role) = Role::from_str(&val) { + database.role = role; + } + } + "shard" => { + if let Ok(shard) = val.parse::() { + database.shard = shard; + } + } + "user" => database.user = Some(val.to_string()), + "password" => database.password = Some(val.to_string()), + "pool_size" => { + if let Ok(size) = val.parse::() { + database.pool_size = Some(size); + } + } + "min_pool_size" => { + if let Ok(size) = val.parse::() { + database.min_pool_size = Some(size); + } + } + "pooler_mode" => { + if let Ok(mode) = PoolerMode::from_str(&val) { + database.pooler_mode = Some(mode); + } + } + "statement_timeout" => { + if let Ok(timeout) = val.parse::() { + database.statement_timeout = Some(timeout); + } + } + "idle_timeout" => { + if let Ok(timeout) = val.parse::() { + database.idle_timeout = Some(timeout); + } + } + "read_only" => { + if let Ok(read_only) = val.parse::() { + database.read_only = Some(read_only); + } + } + "server_lifetime" => { + if let Ok(lifetime) = val.parse::() { + database.server_lifetime = Some(lifetime); + } + } + _ => {} + } + } + + database + } +} + +impl From<&Url> for User { + fn from(value: &Url) -> Self { + let user = value.username(); + let user = if user.is_empty() { + var("USER").unwrap_or("postgres".into()) + } else { + user.to_string() + }; + let password = value.password().unwrap_or("postgres").to_owned(); + User { + name: user, + password: Some(password), + database: database_name(value), + ..Default::default() + } + } +} + +impl ConfigAndUsers { + /// Load from database URLs. + pub fn databases_from_urls(mut self, urls: &[String]) -> Result { + let urls = urls + .iter() + .map(|url| Url::parse(url)) + .collect::, url::ParseError>>()?; + let databases = urls + .iter() + .map(Database::from) + .collect::>() // Make sure we only have unique entries. + .into_iter() + .collect::>(); + let users = urls + .iter() + .map(User::from) + .collect::>() // Make sure we only have unique entries. + .into_iter() + .collect::>(); + + self.users = Users { users, admin: None }; + self.config.databases = databases; + + Ok(self) + } + + /// Load from mirroring strings. + pub fn mirroring_from_strings(mut self, mirror_strs: &[String]) -> Result { + use super::Mirroring; + + let mirroring = mirror_strs + .iter() + .map(|s| Mirroring::from_str(s).map_err(Error::ParseError)) + .collect::, _>>()?; + + self.config.mirroring = mirroring; + + Ok(self) + } +} + +#[cfg(test)] +mod test { + use super::*; + + #[test] + fn test_url() { + let url = Url::parse("postgres://user:password@host:5432/name").unwrap(); + println!("{:#?}", url); + } + + #[test] + fn test_database_name_from_query_param() { + let url = + Url::parse("postgres://user:password@host:5432/name?database_name=dbname").unwrap(); + let database = Database::from(&url); + + assert_eq!(database.name, "name"); + assert_eq!(database.database_name, Some("dbname".to_string())); + } + + #[test] + fn test_role_from_query_param() { + let url = Url::parse("postgres://user:password@host:5432/name?role=replica").unwrap(); + let database = Database::from(&url); + + assert_eq!(database.role, super::super::Role::Replica); + } + + #[test] + fn test_shard_from_query_param() { + let url = Url::parse("postgres://user:password@host:5432/name?shard=5").unwrap(); + let database = Database::from(&url); + + assert_eq!(database.shard, 5); + } + + #[test] + fn test_numeric_fields_from_query_params() { + let url = Url::parse("postgres://user:password@host:5432/name?pool_size=10&min_pool_size=2&statement_timeout=5000&idle_timeout=300&server_lifetime=3600").unwrap(); + let database = Database::from(&url); + + assert_eq!(database.pool_size, Some(10)); + assert_eq!(database.min_pool_size, Some(2)); + assert_eq!(database.statement_timeout, Some(5000)); + assert_eq!(database.idle_timeout, Some(300)); + assert_eq!(database.server_lifetime, Some(3600)); + } + + #[test] + fn test_bool_field_from_query_param() { + let url = Url::parse("postgres://user:password@host:5432/name?read_only=true").unwrap(); + let database = Database::from(&url); + + assert_eq!(database.read_only, Some(true)); + } + + #[test] + fn test_pooler_mode_from_query_param() { + let url = + Url::parse("postgres://user:password@host:5432/name?pooler_mode=session").unwrap(); + let database = Database::from(&url); + + assert_eq!( + database.pooler_mode, + Some(super::super::PoolerMode::Session) + ); + } + + #[test] + fn test_string_fields_from_query_params() { + let url = Url::parse("postgres://user:password@host:5432/name?user=admin&password=secret") + .unwrap(); + let database = Database::from(&url); + + assert_eq!(database.user, Some("admin".to_string())); + assert_eq!(database.password, Some("secret".to_string())); + } + + #[test] + fn test_multiple_query_params() { + let url = Url::parse("postgres://user:password@host:5432/name?database_name=realdb&role=replica&shard=3&pool_size=20&read_only=true").unwrap(); + let database = Database::from(&url); + + assert_eq!(database.name, "name"); + assert_eq!(database.database_name, Some("realdb".to_string())); + assert_eq!(database.role, super::super::Role::Replica); + assert_eq!(database.shard, 3); + assert_eq!(database.pool_size, Some(20)); + assert_eq!(database.read_only, Some(true)); + } + + #[test] + fn test_basic_mirroring_string() { + let mirror_str = "source_db=primary&destination_db=backup"; + let mirroring = super::super::Mirroring::from_str(mirror_str).unwrap(); + + assert_eq!(mirroring.source_db, "primary"); + assert_eq!(mirroring.destination_db, "backup"); + assert_eq!(mirroring.queue_length, None); + assert_eq!(mirroring.exposure, None); + } + + #[test] + fn test_mirroring_with_queue_length() { + let mirror_str = "source_db=db1&destination_db=db2&queue_length=256"; + let mirroring = super::super::Mirroring::from_str(mirror_str).unwrap(); + + assert_eq!(mirroring.source_db, "db1"); + assert_eq!(mirroring.destination_db, "db2"); + assert_eq!(mirroring.queue_length, Some(256)); + assert_eq!(mirroring.exposure, None); + } + + #[test] + fn test_mirroring_with_exposure() { + let mirror_str = "source_db=prod&destination_db=staging&exposure=0.5"; + let mirroring = super::super::Mirroring::from_str(mirror_str).unwrap(); + + assert_eq!(mirroring.source_db, "prod"); + assert_eq!(mirroring.destination_db, "staging"); + assert_eq!(mirroring.queue_length, None); + assert_eq!(mirroring.exposure, Some(0.5)); + } + + #[test] + fn test_mirroring_with_both_overrides() { + let mirror_str = "source_db=main&destination_db=backup&queue_length=512&exposure=0.75"; + let mirroring = super::super::Mirroring::from_str(mirror_str).unwrap(); + + assert_eq!(mirroring.source_db, "main"); + assert_eq!(mirroring.destination_db, "backup"); + assert_eq!(mirroring.queue_length, Some(512)); + assert_eq!(mirroring.exposure, Some(0.75)); + } + + #[test] + fn test_config_mirroring_from_strings() { + let config = ConfigAndUsers::default(); + let mirror_strs = vec![ + "source_db=db1&destination_db=db1_mirror".to_string(), + "source_db=db2&destination_db=db2_mirror&queue_length=256&exposure=0.5".to_string(), + ]; + + let config = config.mirroring_from_strings(&mirror_strs).unwrap(); + + assert_eq!(config.config.mirroring.len(), 2); + assert_eq!(config.config.mirroring[0].source_db, "db1"); + assert_eq!(config.config.mirroring[0].destination_db, "db1_mirror"); + assert_eq!(config.config.mirroring[0].queue_length, None); + assert_eq!(config.config.mirroring[0].exposure, None); + + assert_eq!(config.config.mirroring[1].source_db, "db2"); + assert_eq!(config.config.mirroring[1].destination_db, "db2_mirror"); + assert_eq!(config.config.mirroring[1].queue_length, Some(256)); + assert_eq!(config.config.mirroring[1].exposure, Some(0.5)); + } +} diff --git a/pgdog-config/src/users.rs b/pgdog-config/src/users.rs new file mode 100644 index 000000000..3052c0599 --- /dev/null +++ b/pgdog-config/src/users.rs @@ -0,0 +1,179 @@ +use serde::{Deserialize, Serialize}; +use std::collections::HashMap; +use std::env; +use tracing::warn; + +use super::core::Config; +use super::pooling::PoolerMode; +use crate::util::random_string; + +/// pgDog plugin. +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)] +#[serde(deny_unknown_fields)] +pub struct Plugin { + /// Plugin name. + pub name: String, +} + +/// Users and passwords. +#[derive(Serialize, Deserialize, Debug, Clone, Default)] +#[serde(deny_unknown_fields)] +pub struct Users { + pub admin: Option, + /// Users and passwords. + #[serde(default)] + pub users: Vec, +} + +impl Users { + /// Organize users by database name. + pub fn users(&self) -> HashMap> { + let mut users = HashMap::new(); + + for user in &self.users { + let entry = users.entry(user.database.clone()).or_insert_with(Vec::new); + entry.push(user.clone()); + } + + users + } + + pub fn check(&mut self, config: &Config) { + for user in &mut self.users { + if user.password().is_empty() { + if !config.general.passthrough_auth() { + warn!( + "user \"{}\" doesn't have a password and passthrough auth is disabled", + user.name + ); + } + + if let Some(min_pool_size) = user.min_pool_size { + if min_pool_size > 0 { + warn!("user \"{}\" (database \"{}\") doesn't have a password configured, \ + so we can't connect to the server to maintain min_pool_size of {}; setting it to 0", user.name, user.database, min_pool_size); + user.min_pool_size = Some(0); + } + } + } + } + } +} + +/// User allowed to connect to pgDog. +#[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq, Eq, Ord, PartialOrd)] +#[serde(deny_unknown_fields)] +pub struct User { + /// User name. + pub name: String, + /// Database name, from pgdog.toml. + pub database: String, + /// User's password. + pub password: Option, + /// Pool size for this user pool, overriding `default_pool_size`. + pub pool_size: Option, + /// Minimum pool size for this user pool, overriding `min_pool_size`. + pub min_pool_size: Option, + /// Pooler mode. + pub pooler_mode: Option, + /// Server username. + pub server_user: Option, + /// Server password. + pub server_password: Option, + /// Statement timeout. + pub statement_timeout: Option, + /// Relication mode. + #[serde(default)] + pub replication_mode: bool, + /// Sharding into this database. + pub replication_sharding: Option, + /// Idle timeout. + pub idle_timeout: Option, + /// Read-only mode. + pub read_only: Option, + /// Schema owner. + #[serde(default)] + pub schema_admin: bool, + /// Disable cross-shard queries for this user. + pub cross_shard_disabled: Option, + /// Two-pc. + pub two_phase_commit: Option, + /// Automatic transactions. + pub two_phase_commit_auto: Option, + /// Server lifetime. + pub server_lifetime: Option, +} + +impl User { + pub fn password(&self) -> &str { + if let Some(ref s) = self.password { + s.as_str() + } else { + "" + } + } + + /// New user from user, password and database. + pub fn new(user: &str, password: &str, database: &str) -> Self { + Self { + name: user.to_owned(), + database: database.to_owned(), + password: Some(password.to_owned()), + ..Default::default() + } + } +} + +/// Admin database settings. +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)] +#[serde(deny_unknown_fields)] +pub struct Admin { + /// Admin database name. + #[serde(default = "Admin::name")] + pub name: String, + /// Admin user name. + #[serde(default = "Admin::user")] + pub user: String, + /// Admin user's password. + #[serde(default = "Admin::password")] + pub password: String, +} + +impl Default for Admin { + fn default() -> Self { + Self { + name: Self::name(), + user: Self::user(), + password: admin_password(), + } + } +} + +impl Admin { + fn name() -> String { + "admin".into() + } + + fn user() -> String { + "admin".into() + } + + fn password() -> String { + admin_password() + } + + /// The password has been randomly generated. + pub fn random(&self) -> bool { + let prefix = "_pgdog_"; + self.password.starts_with(prefix) && self.password.len() == prefix.len() + 12 + } +} + +fn admin_password() -> String { + if let Ok(password) = env::var("PGDOG_ADMIN_PASSWORD") { + password + } else { + let pw = random_string(12); + format!("_pgdog_{}", pw) + } +} diff --git a/pgdog-config/src/util.rs b/pgdog-config/src/util.rs new file mode 100644 index 000000000..a7f3c61ec --- /dev/null +++ b/pgdog-config/src/util.rs @@ -0,0 +1,54 @@ +use std::time::Duration; + +use rand::{distr::Alphanumeric, Rng}; + +pub fn human_duration_optional(duration: Option) -> String { + if let Some(duration) = duration { + human_duration(duration) + } else { + "default".into() + } +} + +/// Get a human-readable duration for amounts that +/// a human would use. +pub fn human_duration(duration: Duration) -> String { + let second = 1000; + let minute = second * 60; + let hour = minute * 60; + let day = hour * 24; + let week = day * 7; + // Ok that's enough. + + let ms = duration.as_millis(); + let ms_fmt = |ms: u128, unit: u128, name: &str| -> String { + if !ms.is_multiple_of(unit) { + format!("{}ms", ms) + } else { + format!("{}{}", ms / unit, name) + } + }; + + if ms < second { + format!("{}ms", ms) + } else if ms < minute { + ms_fmt(ms, second, "s") + } else if ms < hour { + ms_fmt(ms, minute, "m") + } else if ms < day { + ms_fmt(ms, hour, "h") + } else if ms < week { + ms_fmt(ms, day, "d") + } else { + ms_fmt(ms, 1, "ms") + } +} + +/// Generate a random string of length n. +pub fn random_string(n: usize) -> String { + rand::rng() + .sample_iter(&Alphanumeric) + .take(n) + .map(char::from) + .collect() +} diff --git a/pgdog-plugin/Cargo.toml b/pgdog-plugin/Cargo.toml index 454c3026a..65193addb 100644 --- a/pgdog-plugin/Cargo.toml +++ b/pgdog-plugin/Cargo.toml @@ -17,7 +17,7 @@ crate-type = ["rlib", "cdylib"] libloading = "0.8" libc = "0.2" tracing = "0.1" -pg_query = "6.1.1" +pg_query = { git = "https://github.com/pgdogdev/pg_query.rs.git", rev = "4f79b92fe4d630b1f253f27f13c9096c77530fd6" } pgdog-macros = { path = "../pgdog-macros", version = "0.1.1" } toml = "0.9" diff --git a/pgdog-vector/Cargo.toml b/pgdog-vector/Cargo.toml new file mode 100644 index 000000000..5e0385cae --- /dev/null +++ b/pgdog-vector/Cargo.toml @@ -0,0 +1,8 @@ +[package] +name = "pgdog-vector" +version = "0.1.0" +edition = "2021" + +[dependencies] +serde = { version = "1", features = ["derive"] } +serde_json = "1" diff --git a/pgdog-vector/src/distance_simd_rust.rs b/pgdog-vector/src/distance_simd_rust.rs new file mode 100644 index 000000000..3432b526a --- /dev/null +++ b/pgdog-vector/src/distance_simd_rust.rs @@ -0,0 +1,276 @@ +use super::Float; + +#[cfg(target_arch = "x86_64")] +use std::arch::x86_64::*; + +#[cfg(target_arch = "aarch64")] +use std::arch::aarch64::*; + +/// Helper function to convert Float slice to f32 slice +/// SAFETY: Float is a newtype wrapper around f32, so the memory layout is identical +#[inline(always)] +unsafe fn float_slice_to_f32(floats: &[Float]) -> &[f32] { + // This is safe because Float is a transparent wrapper around f32 + std::slice::from_raw_parts(floats.as_ptr() as *const f32, floats.len()) +} + +/// Scalar reference implementation - no allocations +#[inline] +pub fn euclidean_distance_scalar(p: &[Float], q: &[Float]) -> f32 { + debug_assert_eq!(p.len(), q.len()); + + let mut sum = 0.0f32; + for i in 0..p.len() { + let diff = q[i].0 - p[i].0; + sum += diff * diff; + } + sum.sqrt() +} + +/// SIMD implementation for x86_64 with SSE - optimized version +#[cfg(all(target_arch = "x86_64", target_feature = "sse"))] +#[inline] +pub fn euclidean_distance_sse(p: &[Float], q: &[Float]) -> f32 { + debug_assert_eq!(p.len(), q.len()); + + unsafe { + // Convert Float slices to f32 slices to avoid temporary arrays + let p_f32 = float_slice_to_f32(p); + let q_f32 = float_slice_to_f32(q); + + let mut sum1 = _mm_setzero_ps(); + let mut sum2 = _mm_setzero_ps(); + let chunks = p.len() / 8; // Process 8 at a time (2 SSE vectors) + + // Unroll loop to process 8 floats per iteration + for i in 0..chunks { + let idx = i * 8; + + // First 4 floats - direct load from slice + let p_vec1 = _mm_loadu_ps(p_f32.as_ptr().add(idx)); + let q_vec1 = _mm_loadu_ps(q_f32.as_ptr().add(idx)); + let diff1 = _mm_sub_ps(q_vec1, p_vec1); + sum1 = _mm_add_ps(sum1, _mm_mul_ps(diff1, diff1)); + + // Second 4 floats + let p_vec2 = _mm_loadu_ps(p_f32.as_ptr().add(idx + 4)); + let q_vec2 = _mm_loadu_ps(q_f32.as_ptr().add(idx + 4)); + let diff2 = _mm_sub_ps(q_vec2, p_vec2); + sum2 = _mm_add_ps(sum2, _mm_mul_ps(diff2, diff2)); + } + + // Combine accumulators + let sum = _mm_add_ps(sum1, sum2); + + // More efficient horizontal sum using shuffles + let shuf = _mm_shuffle_ps(sum, sum, 0b01_00_11_10); // [2,3,0,1] + let sums = _mm_add_ps(sum, shuf); + let shuf = _mm_shuffle_ps(sums, sums, 0b10_11_00_01); // [1,0,3,2] + let result = _mm_add_ps(sums, shuf); + let mut total = _mm_cvtss_f32(result); + + // Handle remaining elements + for i in (chunks * 8)..p.len() { + let diff = q[i].0 - p[i].0; + total += diff * diff; + } + + total.sqrt() + } +} + +/// SIMD implementation for x86_64 with AVX2 - optimized version +#[cfg(all(target_arch = "x86_64", target_feature = "avx2"))] +#[inline] +pub fn euclidean_distance_avx2(p: &[Float], q: &[Float]) -> f32 { + debug_assert_eq!(p.len(), q.len()); + + unsafe { + // Convert Float slices to f32 slices to avoid temporary arrays + let p_f32 = float_slice_to_f32(p); + let q_f32 = float_slice_to_f32(q); + + let mut sum1 = _mm256_setzero_ps(); + let mut sum2 = _mm256_setzero_ps(); + let chunks = p.len() / 16; // Process 16 at a time (2 AVX2 vectors) + + // Unroll loop to process 16 floats per iteration + for i in 0..chunks { + let idx = i * 16; + + // First 8 floats - direct load from slice + let p_vec1 = _mm256_loadu_ps(p_f32.as_ptr().add(idx)); + let q_vec1 = _mm256_loadu_ps(q_f32.as_ptr().add(idx)); + let diff1 = _mm256_sub_ps(q_vec1, p_vec1); + + #[cfg(target_feature = "fma")] + { + sum1 = _mm256_fmadd_ps(diff1, diff1, sum1); + } + #[cfg(not(target_feature = "fma"))] + { + sum1 = _mm256_add_ps(sum1, _mm256_mul_ps(diff1, diff1)); + } + + // Second 8 floats + let p_vec2 = _mm256_loadu_ps(p_f32.as_ptr().add(idx + 8)); + let q_vec2 = _mm256_loadu_ps(q_f32.as_ptr().add(idx + 8)); + let diff2 = _mm256_sub_ps(q_vec2, p_vec2); + + #[cfg(target_feature = "fma")] + { + sum2 = _mm256_fmadd_ps(diff2, diff2, sum2); + } + #[cfg(not(target_feature = "fma"))] + { + sum2 = _mm256_add_ps(sum2, _mm256_mul_ps(diff2, diff2)); + } + } + + // Combine accumulators + let sum = _mm256_add_ps(sum1, sum2); + + // More efficient horizontal sum using extract and SSE + let high = _mm256_extractf128_ps(sum, 1); + let low = _mm256_castps256_ps128(sum); + let sum128 = _mm_add_ps(low, high); + + // Use SSE horizontal sum (more efficient than storing to array) + let shuf = _mm_shuffle_ps(sum128, sum128, 0b01_00_11_10); + let sums = _mm_add_ps(sum128, shuf); + let shuf = _mm_shuffle_ps(sums, sums, 0b10_11_00_01); + let result = _mm_add_ps(sums, shuf); + let mut total = _mm_cvtss_f32(result); + + // Handle remaining elements + for i in (chunks * 16)..p.len() { + let diff = q[i].0 - p[i].0; + total += diff * diff; + } + + total.sqrt() + } +} + +/// SIMD implementation for ARM NEON - optimized version +#[cfg(target_arch = "aarch64")] +#[inline] +pub fn euclidean_distance_neon(p: &[Float], q: &[Float]) -> f32 { + debug_assert_eq!(p.len(), q.len()); + + unsafe { + // Convert Float slices to f32 slices to avoid temporary arrays + let p_f32 = float_slice_to_f32(p); + let q_f32 = float_slice_to_f32(q); + + let mut sum1 = vdupq_n_f32(0.0); + let mut sum2 = vdupq_n_f32(0.0); + let chunks = p.len() / 8; // Process 8 at a time (2 vectors) + + // Unroll loop to process 8 floats per iteration (2x4) + for i in 0..chunks { + let idx = i * 8; + + // First 4 floats - direct load from slice + let p_vec1 = vld1q_f32(p_f32.as_ptr().add(idx)); + let q_vec1 = vld1q_f32(q_f32.as_ptr().add(idx)); + let diff1 = vsubq_f32(q_vec1, p_vec1); + sum1 = vfmaq_f32(sum1, diff1, diff1); + + // Second 4 floats + let p_vec2 = vld1q_f32(p_f32.as_ptr().add(idx + 4)); + let q_vec2 = vld1q_f32(q_f32.as_ptr().add(idx + 4)); + let diff2 = vsubq_f32(q_vec2, p_vec2); + sum2 = vfmaq_f32(sum2, diff2, diff2); + } + + // Combine the two accumulators + let sum = vaddq_f32(sum1, sum2); + + // Horizontal sum - more efficient version + let sum_pairs = vpaddq_f32(sum, sum); + let sum_final = vpaddq_f32(sum_pairs, sum_pairs); + let mut total = vgetq_lane_f32(sum_final, 0); + + // Handle remaining elements + for i in (chunks * 8)..p.len() { + let diff = q[i].0 - p[i].0; + total += diff * diff; + } + + total.sqrt() + } +} + +/// Auto-select best implementation based on CPU features +#[inline] +pub fn euclidean_distance(p: &[Float], q: &[Float]) -> f32 { + #[cfg(all(target_arch = "x86_64", target_feature = "avx2"))] + { + euclidean_distance_avx2(p, q) + } + + #[cfg(all( + target_arch = "x86_64", + target_feature = "sse", + not(target_feature = "avx2") + ))] + { + euclidean_distance_sse(p, q) + } + + #[cfg(target_arch = "aarch64")] + { + euclidean_distance_neon(p, q) + } + + #[cfg(not(any( + all(target_arch = "x86_64", target_feature = "sse"), + target_arch = "aarch64" + )))] + { + euclidean_distance_scalar(p, q) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::Vector; + + #[test] + fn test_no_allocations() { + // This test verifies we're not allocating + let v1 = Vector::from(&[1.0, 2.0, 3.0][..]); + let v2 = Vector::from(&[4.0, 5.0, 6.0][..]); + + // These calls should not allocate + let dist = euclidean_distance(&v1, &v2); + assert!(dist > 0.0); + } + + #[test] + fn test_correctness() { + let v1: Vec = vec![Float(3.0), Float(4.0)]; + let v2: Vec = vec![Float(0.0), Float(0.0)]; + + let dist_scalar = euclidean_distance_scalar(&v1, &v2); + let dist_simd = euclidean_distance(&v1, &v2); + + assert!((dist_scalar - 5.0).abs() < 1e-6); + assert!((dist_simd - 5.0).abs() < 1e-6); + } + + #[test] + fn test_large_vectors() { + // Test with 1536-dimensional vectors (OpenAI embeddings) + let v1: Vec = (0..1536).map(|i| Float(i as f32 * 0.001)).collect(); + let v2: Vec = (0..1536).map(|i| Float(i as f32 * 0.001 + 0.5)).collect(); + + let dist_scalar = euclidean_distance_scalar(&v1, &v2); + let dist_simd = euclidean_distance(&v1, &v2); + + let relative_error = ((dist_simd - dist_scalar).abs() / dist_scalar).abs(); + assert!(relative_error < 1e-5); + } +} diff --git a/pgdog-vector/src/float.rs b/pgdog-vector/src/float.rs new file mode 100644 index 000000000..79d0b743b --- /dev/null +++ b/pgdog-vector/src/float.rs @@ -0,0 +1,82 @@ +use std::{ + cmp::Ordering, + fmt::Display, + hash::{Hash, Hasher}, +}; + +use serde::{Deserialize, Serialize}; + +/// Wrapper type for f32 that implements Ord for PostgreSQL compatibility +#[derive(Debug, Clone, Copy, Serialize, Deserialize)] +pub struct Float(pub f32); + +impl PartialOrd for Float { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl Ord for Float { + fn cmp(&self, other: &Self) -> Ordering { + // PostgreSQL ordering: NaN is greater than all other values + match (self.0.is_nan(), other.0.is_nan()) { + (true, true) => Ordering::Equal, + (true, false) => Ordering::Greater, + (false, true) => Ordering::Less, + (false, false) => self.0.partial_cmp(&other.0).unwrap_or(Ordering::Equal), + } + } +} + +impl PartialEq for Float { + fn eq(&self, other: &Self) -> bool { + // PostgreSQL treats NaN as equal to NaN for indexing purposes + if self.0.is_nan() && other.0.is_nan() { + true + } else { + self.0 == other.0 + } + } +} + +impl Eq for Float {} + +impl Hash for Float { + fn hash(&self, state: &mut H) { + if self.0.is_nan() { + // All NaN values hash to the same value + 0u8.hash(state); + } else { + // Use bit representation for consistent hashing + self.0.to_bits().hash(state); + } + } +} + +impl Display for Float { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + if self.0.is_nan() { + write!(f, "NaN") + } else if self.0.is_infinite() { + if self.0.is_sign_positive() { + write!(f, "Infinity") + } else { + write!(f, "-Infinity") + } + } else { + write!(f, "{}", self.0) + } + } +} + +impl From for Float { + fn from(value: f32) -> Self { + Float(value) + } +} + +impl From for f32 { + fn from(value: Float) -> Self { + value.0 + } +} diff --git a/pgdog-vector/src/lib.rs b/pgdog-vector/src/lib.rs new file mode 100644 index 000000000..f36c40120 --- /dev/null +++ b/pgdog-vector/src/lib.rs @@ -0,0 +1,358 @@ +use std::{fmt::Debug, ops::Deref}; + +use serde::{ + de::{self, Visitor}, + ser::SerializeSeq, + Deserialize, Serialize, +}; + +pub mod distance_simd_rust; +pub mod float; + +pub use float::*; + +#[derive(Clone, PartialEq, PartialOrd, Ord, Eq, Hash)] +#[repr(C)] +pub struct Vector { + pub values: Vec, +} + +impl Vector { + /// Length of the vector. + pub fn len(&self) -> usize { + self.values.len() + } + + /// Is the vector empty? + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + /// Compute L2 distance between the vectors. + pub fn distance_l2(&self, other: &Self) -> f32 { + Distance::Euclidean(self, other).distance() + } +} + +pub enum Distance<'a> { + Euclidean(&'a Vector, &'a Vector), +} + +impl Distance<'_> { + pub fn distance(&self) -> f32 { + match self { + Self::Euclidean(p, q) => { + assert_eq!(p.len(), q.len()); + // Avoids temporary array allocations by working directly with the Float slices + distance_simd_rust::euclidean_distance(p, q) + } + } + } + + // Fallback implementation for testing + pub fn distance_scalar(&self) -> f32 { + match self { + Self::Euclidean(p, q) => { + assert_eq!(p.len(), q.len()); + // Use scalar implementation + distance_simd_rust::euclidean_distance_scalar(p, q) + } + } + } +} + +impl Debug for Vector { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + if self.values.len() > 3 { + f.debug_struct("Vector") + .field( + "values", + &format!( + "[{}..{}]", + self.values[0], + self.values[self.values.len() - 1] + ), + ) + .finish() + } else { + f.debug_struct("Vector") + .field("values", &self.values) + .finish() + } + } +} + +impl Deref for Vector { + type Target = Vec; + + fn deref(&self) -> &Self::Target { + &self.values + } +} + +impl From<&[f64]> for Vector { + fn from(value: &[f64]) -> Self { + Self { + values: value.iter().map(|v| Float(*v as f32)).collect(), + } + } +} + +impl From<&[f32]> for Vector { + fn from(value: &[f32]) -> Self { + Self { + values: value.iter().map(|v| Float(*v)).collect(), + } + } +} + +impl From> for Vector { + fn from(value: Vec) -> Self { + Self { + values: value.into_iter().map(Float::from).collect(), + } + } +} + +impl From> for Vector { + fn from(value: Vec) -> Self { + Self { + values: value.into_iter().map(|v| Float(v as f32)).collect(), + } + } +} + +impl From> for Vector { + fn from(value: Vec) -> Self { + Self { values: value } + } +} + +#[derive(Debug)] +pub struct Centroids<'a> { + centroids: &'a [Vector], +} + +impl Centroids<'_> { + /// Find the shards with the closest centroids, + /// according to the number of probes. + pub fn shard(&self, vector: &Vector, shards: usize, probes: usize) -> Vec { + let mut selected = vec![]; + let mut centroids = self.centroids.iter().enumerate().collect::>(); + centroids.sort_by(|(_, a), (_, b)| { + a.distance_l2(vector) + .partial_cmp(&b.distance_l2(vector)) + .unwrap_or(std::cmp::Ordering::Equal) + }); + let centroids = centroids.into_iter().take(probes); + for (i, _) in centroids { + selected.push(i % shards); + } + + selected + } +} + +impl<'a> From<&'a Vec> for Centroids<'a> { + fn from(centroids: &'a Vec) -> Self { + Centroids { centroids } + } +} + +struct VectorVisitor; + +impl<'de> Visitor<'de> for VectorVisitor { + type Value = Vector; + + fn visit_seq(self, mut seq: A) -> Result + where + A: de::SeqAccess<'de>, + { + let mut results = vec![]; + while let Some(n) = seq.next_element::()? { + results.push(n); + } + + Ok(Vector::from(results.as_slice())) + } + + fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + formatter.write_str("expected a list of floating points") + } +} + +impl<'de> Deserialize<'de> for Vector { + fn deserialize(deserializer: D) -> Result + where + D: de::Deserializer<'de>, + { + deserializer.deserialize_seq(VectorVisitor) + } +} + +impl Serialize for Vector { + fn serialize(&self, serializer: S) -> Result + where + S: serde::Serializer, + { + let mut seq = serializer.serialize_seq(Some(self.len()))?; + for v in &self.values { + seq.serialize_element(v)?; + } + seq.end() + } +} + +#[cfg(test)] +mod test { + + use super::*; + + #[test] + fn test_euclidean() { + let v1 = Vector::from(&[1.0, 2.0, 3.0][..]); + let v2 = Vector::from(&[1.5, 2.0, 3.0][..]); + let distance = Distance::Euclidean(&v1, &v2).distance(); + assert_eq!(distance, 0.5); + } + + #[test] + fn test_simd_features() { + println!("SIMD features available:"); + #[cfg(target_arch = "x86_64")] + { + println!(" x86_64 architecture detected"); + #[cfg(target_feature = "sse")] + println!(" SSE: enabled"); + #[cfg(target_feature = "avx2")] + println!(" AVX2: enabled"); + #[cfg(target_feature = "fma")] + println!(" FMA: enabled"); + } + #[cfg(target_arch = "aarch64")] + { + println!(" ARM64 architecture detected"); + println!(" NEON: enabled (always available on ARM64)"); + } + } + + #[test] + fn test_simd_vs_scalar() { + // Test small vectors + let v1 = Vector::from(&[3.0, 4.0][..]); + let v2 = Vector::from(&[0.0, 0.0][..]); + let simd_dist = Distance::Euclidean(&v1, &v2).distance(); + let scalar_dist = Distance::Euclidean(&v1, &v2).distance_scalar(); + assert!((simd_dist - scalar_dist).abs() < 1e-6); + assert!((simd_dist - 5.0).abs() < 1e-6); + + // Test medium vectors + let v3: Vec = (0..128).map(|i| i as f32).collect(); + let v4: Vec = (0..128).map(|i| (i + 1) as f32).collect(); + let v3 = Vector::from(v3); + let v4 = Vector::from(v4); + let simd_dist = Distance::Euclidean(&v3, &v4).distance(); + let scalar_dist = Distance::Euclidean(&v3, &v4).distance_scalar(); + assert!((simd_dist - scalar_dist).abs() < 1e-4); + } + + #[test] + fn test_openai_embedding_size() { + // Test with OpenAI text-embedding-3-small dimension (1536) + let v1: Vec = (0..1536).map(|i| (i as f32) * 0.001).collect(); + let v2: Vec = (0..1536).map(|i| (i as f32) * 0.001 + 0.5).collect(); + let v1 = Vector::from(v1); + let v2 = Vector::from(v2); + + let simd_dist = Distance::Euclidean(&v1, &v2).distance(); + let scalar_dist = Distance::Euclidean(&v1, &v2).distance_scalar(); + + // Check that both implementations produce very similar results + let relative_error = ((simd_dist - scalar_dist).abs() / scalar_dist).abs(); + assert!( + relative_error < 1e-5, + "Relative error: {}, SIMD: {}, Scalar: {}", + relative_error, + simd_dist, + scalar_dist + ); + } + + #[test] + fn test_special_values() { + // Test with NaN + let v_nan = Vector::from(vec![Float(f32::NAN), Float(1.0)]); + let v_normal = Vector::from(&[1.0, 1.0][..]); + let simd_dist = Distance::Euclidean(&v_nan, &v_normal).distance(); + let scalar_dist = Distance::Euclidean(&v_nan, &v_normal).distance_scalar(); + assert!(simd_dist.is_nan()); + assert!(scalar_dist.is_nan()); + + // Test with Infinity + let v_inf = Vector::from(vec![Float(f32::INFINITY), Float(1.0)]); + let simd_dist = Distance::Euclidean(&v_inf, &v_normal).distance(); + let scalar_dist = Distance::Euclidean(&v_inf, &v_normal).distance_scalar(); + assert!(simd_dist.is_infinite()); + assert!(scalar_dist.is_infinite()); + } + + #[test] + fn test_zero_distance() { + let v1 = Vector::from(&[1.0, 2.0, 3.0, 4.0, 5.0][..]); + let simd_dist = Distance::Euclidean(&v1, &v1).distance(); + let scalar_dist = Distance::Euclidean(&v1, &v1).distance_scalar(); + assert_eq!(simd_dist, 0.0); + assert_eq!(scalar_dist, 0.0); + } + + #[test] + fn test_various_sizes() { + // Test various vector sizes to ensure correct handling of tail elements + for size in [ + 1, 3, 4, 7, 8, 15, 16, 31, 32, 63, 64, 127, 128, 255, 256, 512, 1024, + ] { + let v1: Vec = (0..size).map(|i| i as f32).collect(); + let v2: Vec = (0..size).map(|i| (i * 2) as f32).collect(); + let v1 = Vector::from(v1); + let v2 = Vector::from(v2); + + let simd_dist = Distance::Euclidean(&v1, &v2).distance(); + let scalar_dist = Distance::Euclidean(&v1, &v2).distance_scalar(); + + let relative_error = if scalar_dist != 0.0 { + ((simd_dist - scalar_dist).abs() / scalar_dist).abs() + } else { + (simd_dist - scalar_dist).abs() + }; + + assert!( + relative_error < 1e-5, + "Size {}: relative error {}, SIMD: {}, Scalar: {}", + size, + relative_error, + simd_dist, + scalar_dist + ); + } + } + + #[test] + fn test_vector_distance_with_special_values() { + // Test distance calculation with normal values + let v1 = Vector::from(&[3.0, 4.0][..]); + let v2 = Vector::from(&[0.0, 0.0][..]); + let distance = v1.distance_l2(&v2); + assert_eq!(distance, 5.0); // 3-4-5 triangle + + // Test distance with NaN + let v_nan = Vector::from(vec![Float(f32::NAN), Float(1.0)]); + let v_normal = Vector::from(&[1.0, 1.0][..]); + let distance_nan = v_nan.distance_l2(&v_normal); + assert!(distance_nan.is_nan()); + + // Test distance with Infinity + let v_inf = Vector::from(vec![Float(f32::INFINITY), Float(1.0)]); + let distance_inf = v_inf.distance_l2(&v_normal); + assert!(distance_inf.is_infinite()); + } +} diff --git a/pgdog/Cargo.lock b/pgdog/Cargo.lock index 35650a353..4bdb3d17d 100644 --- a/pgdog/Cargo.lock +++ b/pgdog/Cargo.lock @@ -546,7 +546,7 @@ dependencies = [ [[package]] name = "rand" -version = "0.8.5" +version = "0.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404" dependencies = [ diff --git a/pgdog/Cargo.toml b/pgdog/Cargo.toml index e16090d8b..84d59e74e 100644 --- a/pgdog/Cargo.toml +++ b/pgdog/Cargo.toml @@ -1,12 +1,12 @@ [package] name = "pgdog" -version = "0.1.8" +version = "0.1.22" edition = "2021" description = "Modern PostgreSQL proxy, pooler and load balancer." -authors = ["Lev Kokotov "] +authors = ["PgDog "] license = "AGPL-3.0" homepage = "https://pgdog.dev" -repository = "https://github.com/levkk/pgdog" +repository = "https://github.com/pgdogdev/pgdog" readme = "README.md" default-run = "pgdog" @@ -28,7 +28,7 @@ clap = { version = "4", features = ["derive"] } serde = { version = "1", features = ["derive"] } serde_json = "1" async-trait = "0.1" -rand = "0.8" +rand = "0.9.2" once_cell = "1" tokio-rustls = "0.26" rustls-native-certs = "0.8" @@ -38,12 +38,12 @@ toml = "0.8" pgdog-plugin = { path = "../pgdog-plugin", version = "0.1.8" } tokio-util = { version = "0.7", features = ["rt"] } fnv = "1" -scram = "0.6" +scram = { git = "https://github.com/pgdogdev/scram.git", rev = "848003d" } base64 = "0.22" md5 = "0.7" futures = "0.3" csv-core = "0.1" -pg_query = "6.1.1" +pg_query = { git = "https://github.com/pgdogdev/pg_query.rs.git", rev = "4f79b92fe4d630b1f253f27f13c9096c77530fd6" } regex = "1" uuid = { version = "1", features = ["v4", "serde"] } url = "2" @@ -61,6 +61,9 @@ lru = "0.16" hickory-resolver = "0.25.2" lazy_static = "1" dashmap = "6" +derive_builder = "0.20.2" +pgdog-config = { path = "../pgdog-config" } +pgdog-vector = { path = "../pgdog-vector" } [target.'cfg(not(target_env = "msvc"))'.dependencies] tikv-jemallocator = "0.6" @@ -68,3 +71,7 @@ tikv-jemallocator = "0.6" [build-dependencies] cc = "1" + +[dev-dependencies] +tempfile = "3.23.0" +stats_alloc = "0.1.10" diff --git a/pgdog/src/admin/ban.rs b/pgdog/src/admin/ban.rs index 96c57db69..c55effcf2 100644 --- a/pgdog/src/admin/ban.rs +++ b/pgdog/src/admin/ban.rs @@ -1,3 +1,5 @@ +use std::time::Duration; + use super::prelude::*; use crate::backend::{databases::databases, pool}; @@ -43,7 +45,7 @@ impl Command for Ban { async fn execute(&self) -> Result, Error> { for database in databases().all().values() { for shard in database.shards() { - for pool in shard.pools() { + for (_role, ban, pool) in shard.pools_with_roles_and_bans() { if let Some(id) = self.id { if id != pool.id() { continue; @@ -51,9 +53,9 @@ impl Command for Ban { } if self.unban { - pool.unban(); + ban.unban(false); } else { - pool.ban(pool::Error::ManualBan); + ban.ban(pool::Error::ManualBan, Duration::MAX); } } } diff --git a/pgdog/src/admin/healthcheck.rs b/pgdog/src/admin/healthcheck.rs new file mode 100644 index 000000000..ff019924a --- /dev/null +++ b/pgdog/src/admin/healthcheck.rs @@ -0,0 +1,54 @@ +use super::prelude::*; +use crate::backend::{databases::databases, pool::monitor::Monitor}; + +#[derive(Default)] +pub struct Healthcheck { + id: Option, +} + +#[async_trait] +impl Command for Healthcheck { + fn name(&self) -> String { + "HEALTHCHECK".into() + } + + fn parse(sql: &str) -> Result { + let parts = sql.split(" ").collect::>(); + + match parts[..] { + ["healthcheck"] => Ok(Self::default()), + ["healthcheck", id] => Ok(Self { + id: Some(id.parse()?), + }), + _ => Err(Error::Syntax), + } + } + + async fn execute(&self) -> Result, Error> { + for database in databases().all().values() { + for shard in database.shards() { + for pool in shard.pools() { + if let Some(id) = self.id { + if id != pool.id() { + continue; + } + } + + let _ = Monitor::healthcheck(&pool).await; + } + } + } + Ok(vec![]) + } +} + +#[cfg(test)] +mod test { + use super::*; + + #[test] + fn test_parse_healthcheck() { + let cmd = Healthcheck::parse("healthcheck").unwrap(); + assert!(cmd.id.is_none()); + } +} diff --git a/pgdog/src/admin/mod.rs b/pgdog/src/admin/mod.rs index 21c5f0712..521119c8c 100644 --- a/pgdog/src/admin/mod.rs +++ b/pgdog/src/admin/mod.rs @@ -6,6 +6,7 @@ use crate::net::messages::Message; pub mod ban; pub mod error; +pub mod healthcheck; pub mod maintenance_mode; pub mod named_row; pub mod parser; @@ -18,6 +19,7 @@ pub mod reset_query_cache; pub mod server; pub mod set; pub mod setup_schema; +pub mod show_client_memory; pub mod show_clients; pub mod show_config; pub mod show_instance_id; @@ -27,6 +29,8 @@ pub mod show_peers; pub mod show_pools; pub mod show_prepared_statements; pub mod show_query_cache; +pub mod show_replication; +pub mod show_server_memory; pub mod show_servers; pub mod show_stats; pub mod show_transactions; @@ -35,6 +39,9 @@ pub mod shutdown; pub use error::Error; +#[cfg(test)] +mod tests; + /// All pooler commands implement this trait. #[async_trait] pub trait Command: Sized { diff --git a/pgdog/src/admin/parser.rs b/pgdog/src/admin/parser.rs index db8e83c0b..5d94ab445 100644 --- a/pgdog/src/admin/parser.rs +++ b/pgdog/src/admin/parser.rs @@ -1,12 +1,14 @@ //! Admin command parser. use super::{ - ban::Ban, maintenance_mode::MaintenanceMode, pause::Pause, prelude::Message, probe::Probe, - reconnect::Reconnect, reload::Reload, reset_query_cache::ResetQueryCache, set::Set, - setup_schema::SetupSchema, show_clients::ShowClients, show_config::ShowConfig, + ban::Ban, healthcheck::Healthcheck, maintenance_mode::MaintenanceMode, pause::Pause, + prelude::Message, probe::Probe, reconnect::Reconnect, reload::Reload, + reset_query_cache::ResetQueryCache, set::Set, setup_schema::SetupSchema, + show_client_memory::ShowClientMemory, show_clients::ShowClients, show_config::ShowConfig, show_instance_id::ShowInstanceId, show_lists::ShowLists, show_mirrors::ShowMirrors, show_peers::ShowPeers, show_pools::ShowPools, show_prepared_statements::ShowPreparedStatements, - show_query_cache::ShowQueryCache, show_servers::ShowServers, show_stats::ShowStats, + show_query_cache::ShowQueryCache, show_replication::ShowReplication, + show_server_memory::ShowServerMemory, show_servers::ShowServers, show_stats::ShowStats, show_transactions::ShowTransactions, show_version::ShowVersion, shutdown::Shutdown, Command, Error, }; @@ -34,10 +36,14 @@ pub enum ParseResult { Shutdown(Shutdown), ShowLists(ShowLists), ShowPrepared(ShowPreparedStatements), + ShowReplication(ShowReplication), + ShowServerMemory(ShowServerMemory), + ShowClientMemory(ShowClientMemory), Set(Set), Ban(Ban), Probe(Probe), MaintenanceMode(MaintenanceMode), + Healthcheck(Healthcheck), } impl ParseResult { @@ -65,10 +71,14 @@ impl ParseResult { Shutdown(shutdown) => shutdown.execute().await, ShowLists(show_lists) => show_lists.execute().await, ShowPrepared(cmd) => cmd.execute().await, + ShowReplication(show_replication) => show_replication.execute().await, + ShowServerMemory(show_server_memory) => show_server_memory.execute().await, + ShowClientMemory(show_client_memory) => show_client_memory.execute().await, Set(set) => set.execute().await, Ban(ban) => ban.execute().await, Probe(probe) => probe.execute().await, MaintenanceMode(maintenance_mode) => maintenance_mode.execute().await, + Healthcheck(healthcheck) => healthcheck.execute().await, } } @@ -96,10 +106,14 @@ impl ParseResult { Shutdown(shutdown) => shutdown.name(), ShowLists(show_lists) => show_lists.name(), ShowPrepared(show) => show.name(), + ShowReplication(show_replication) => show_replication.name(), + ShowServerMemory(show_server_memory) => show_server_memory.name(), + ShowClientMemory(show_client_memory) => show_client_memory.name(), Set(set) => set.name(), Ban(ban) => ban.name(), Probe(probe) => probe.name(), MaintenanceMode(maintenance_mode) => maintenance_mode.name(), + Healthcheck(healthcheck) => healthcheck.name(), } } } @@ -119,11 +133,26 @@ impl Parser { "reconnect" => ParseResult::Reconnect(Reconnect::parse(&sql)?), "reload" => ParseResult::Reload(Reload::parse(&sql)?), "ban" | "unban" => ParseResult::Ban(Ban::parse(&sql)?), + "healthcheck" => ParseResult::Healthcheck(Healthcheck::parse(&sql)?), "show" => match iter.next().ok_or(Error::Syntax)?.trim() { "clients" => ParseResult::ShowClients(ShowClients::parse(&sql)?), "pools" => ParseResult::ShowPools(ShowPools::parse(&sql)?), "config" => ParseResult::ShowConfig(ShowConfig::parse(&sql)?), "servers" => ParseResult::ShowServers(ShowServers::parse(&sql)?), + "server" => match iter.next().ok_or(Error::Syntax)?.trim() { + "memory" => ParseResult::ShowServerMemory(ShowServerMemory::parse(&sql)?), + command => { + debug!("unknown admin show server command: '{}'", command); + return Err(Error::Syntax); + } + }, + "client" => match iter.next().ok_or(Error::Syntax)?.trim() { + "memory" => ParseResult::ShowClientMemory(ShowClientMemory::parse(&sql)?), + command => { + debug!("unknown admin show client command: '{}'", command); + return Err(Error::Syntax); + } + }, "peers" => ParseResult::ShowPeers(ShowPeers::parse(&sql)?), "query_cache" => ParseResult::ShowQueryCache(ShowQueryCache::parse(&sql)?), "stats" => ParseResult::ShowStats(ShowStats::parse(&sql)?), @@ -133,6 +162,7 @@ impl Parser { "instance_id" => ParseResult::ShowInstanceId(ShowInstanceId::parse(&sql)?), "lists" => ParseResult::ShowLists(ShowLists::parse(&sql)?), "prepared" => ParseResult::ShowPrepared(ShowPreparedStatements::parse(&sql)?), + "replication" => ParseResult::ShowReplication(ShowReplication::parse(&sql)?), command => { debug!("unknown admin show command: '{}'", command); return Err(Error::Syntax); @@ -187,4 +217,16 @@ mod tests { let result = Parser::parse("FOO BAR"); assert!(matches!(result, Err(Error::Syntax))); } + + #[test] + fn parses_show_server_memory_command() { + let result = Parser::parse("SHOW SERVER MEMORY;"); + assert!(matches!(result, Ok(ParseResult::ShowServerMemory(_)))); + } + + #[test] + fn parses_show_client_memory_command() { + let result = Parser::parse("SHOW CLIENT MEMORY;"); + assert!(matches!(result, Ok(ParseResult::ShowClientMemory(_)))); + } } diff --git a/pgdog/src/admin/probe.rs b/pgdog/src/admin/probe.rs index 2218f2904..571056c5f 100644 --- a/pgdog/src/admin/probe.rs +++ b/pgdog/src/admin/probe.rs @@ -2,7 +2,7 @@ use tokio::time::{timeout, Duration, Instant}; use url::Url; use crate::{ - backend::{pool::Address, Server, ServerOptions}, + backend::{pool::Address, ConnectReason, Server, ServerOptions}, config::config, }; @@ -30,6 +30,7 @@ impl Command for Probe { Server::connect( &Address::try_from(self.url.clone()).map_err(|_| Error::InvalidAddress)?, ServerOptions::default(), + ConnectReason::Probe, ), ) .await? diff --git a/pgdog/src/admin/reconnect.rs b/pgdog/src/admin/reconnect.rs index b7f18e4c7..20af53546 100644 --- a/pgdog/src/admin/reconnect.rs +++ b/pgdog/src/admin/reconnect.rs @@ -21,7 +21,7 @@ impl Command for Reconnect { } async fn execute(&self) -> Result, Error> { - reconnect(); + reconnect()?; Ok(vec![]) } } diff --git a/pgdog/src/admin/set.rs b/pgdog/src/admin/set.rs index 99cbee90a..c89c944b2 100644 --- a/pgdog/src/admin/set.rs +++ b/pgdog/src/admin/set.rs @@ -1,6 +1,6 @@ use crate::{ backend::databases, - config::{self, config}, + config::{self, config, RewriteMode}, frontend::PreparedStatements, }; @@ -91,6 +91,10 @@ impl Command for Set { .close_unused(config.config.general.prepared_statements_limit); } + "prepared_statements" => { + config.config.general.prepared_statements = Self::from_json(&self.value)?; + } + "cross_shard_disabled" => { config.config.general.cross_shard_disabled = Self::from_json(&self.value)?; } @@ -103,11 +107,57 @@ impl Command for Set { config.config.general.two_phase_commit_auto = Self::from_json(&self.value)?; } + "rewrite_shard_key_updates" => { + config.config.rewrite.shard_key = self + .value + .parse::() + .map_err(|_| Error::Syntax)?; + } + + "rewrite_split_inserts" => { + config.config.rewrite.split_inserts = self + .value + .parse::() + .map_err(|_| Error::Syntax)?; + } + + "rewrite_enabled" => { + config.config.rewrite.enabled = Self::from_json(&self.value)?; + } + + "healthcheck_interval" => { + config.config.general.healthcheck_interval = self.value.parse()?; + } + + "idle_healthcheck_interval" => { + config.config.general.idle_healthcheck_interval = self.value.parse()?; + } + + "idle_healthcheck_delay" => { + config.config.general.idle_healthcheck_delay = self.value.parse()?; + } + + "ban_timeout" => { + config.config.general.ban_timeout = self.value.parse()?; + } + + "tls_client_required" => { + config.config.general.tls_client_required = Self::from_json(&self.value)?; + } + + "query_parser" => { + config.config.general.query_parser = Self::from_json(&self.value)?; + } + + "client_idle_in_transaction_timeout" => { + config.config.general.client_idle_in_transaction_timeout = self.value.parse()?; + } + _ => return Err(Error::Syntax), } config::set(config)?; - databases::init(); + databases::init()?; Ok(vec![]) } diff --git a/pgdog/src/admin/show_client_memory.rs b/pgdog/src/admin/show_client_memory.rs new file mode 100644 index 000000000..30a14b799 --- /dev/null +++ b/pgdog/src/admin/show_client_memory.rs @@ -0,0 +1,63 @@ +use crate::{ + frontend::comms::comms, + net::messages::{DataRow, Field, Protocol, RowDescription}, +}; + +use super::prelude::*; + +pub struct ShowClientMemory; + +#[async_trait] +impl Command for ShowClientMemory { + fn name(&self) -> String { + "SHOW CLIENT MEMORY".into() + } + + fn parse(_sql: &str) -> Result { + Ok(ShowClientMemory {}) + } + + async fn execute(&self) -> Result, Error> { + let rd = RowDescription::new(&[ + Field::bigint("client_id"), + Field::text("database"), + Field::text("user"), + Field::text("addr"), + Field::numeric("port"), + Field::numeric("buffer_reallocs"), + Field::numeric("buffer_reclaims"), + Field::numeric("buffer_bytes_used"), + Field::numeric("buffer_bytes_alloc"), + Field::numeric("prepared_statements_bytes"), + Field::numeric("net_buffer_bytes"), + Field::numeric("total_bytes"), + ]); + let mut messages = vec![rd.message()?]; + + let clients = comms().clients(); + for (_, client) in clients { + let mut row = DataRow::new(); + let memory = &client.stats.memory_stats; + + let user = client.paramters.get_default("user", "postgres"); + let database = client.paramters.get_default("database", user); + + row.add(client.id.pid as i64) + .add(database) + .add(user) + .add(client.addr.ip().to_string().as_str()) + .add(client.addr.port() as i64) + .add(memory.buffer.reallocs as i64) + .add(memory.buffer.reclaims as i64) + .add(memory.buffer.bytes_used as i64) + .add(memory.buffer.bytes_alloc as i64) + .add(memory.prepared_statements as i64) + .add(memory.stream as i64) + .add((memory.total()) as i64); + + messages.push(row.message()?); + } + + Ok(messages) + } +} diff --git a/pgdog/src/admin/show_clients.rs b/pgdog/src/admin/show_clients.rs index ae4be4bb3..022724143 100644 --- a/pgdog/src/admin/show_clients.rs +++ b/pgdog/src/admin/show_clients.rs @@ -26,6 +26,7 @@ impl Command for ShowClients { .collect::>(); let fields = vec![ + Field::bigint("id"), Field::text("user"), Field::text("database"), Field::text("addr"), @@ -44,7 +45,6 @@ impl Command for ShowClients { Field::numeric("bytes_sent"), Field::numeric("errors"), Field::text("application_name"), - Field::numeric("memory_used"), Field::bool("locked"), Field::numeric("prepared_statements"), ]; @@ -77,6 +77,7 @@ impl Command for ShowClients { let row = self .filter .clone() + .add("id", client.id.pid as i64) .add("user", user) .add("database", client.paramters.get_default("database", user)) .add("addr", client.addr.ip().to_string()) @@ -117,7 +118,6 @@ impl Command for ShowClients { "application_name", client.paramters.get_default("application_name", ""), ) - .add("memory_used", client.stats.memory_used) .add("locked", client.stats.locked) .add("prepared_statements", client.stats.prepared_statements) .data_row(); diff --git a/pgdog/src/admin/show_config.rs b/pgdog/src/admin/show_config.rs index 4d57cc73d..afa0d2fa1 100644 --- a/pgdog/src/admin/show_config.rs +++ b/pgdog/src/admin/show_config.rs @@ -33,7 +33,12 @@ impl Command for ShowConfig { // Reflection using JSON. let general = serde_json::to_value(&config.config.general)?; let tcp = serde_json::to_value(config.config.tcp.clone())?; - let objects = [("", general.as_object()), ("tcp_", tcp.as_object())]; + let memory = serde_json::to_value(config.config.memory.clone())?; + let objects = [ + ("", general.as_object()), + ("tcp_", tcp.as_object()), + ("memory_", memory.as_object()), + ]; for (prefix, object) in objects.iter() { if let Some(object) = object { diff --git a/pgdog/src/admin/show_pools.rs b/pgdog/src/admin/show_pools.rs index fb966f955..5366dc5c4 100644 --- a/pgdog/src/admin/show_pools.rs +++ b/pgdog/src/admin/show_pools.rs @@ -1,5 +1,5 @@ use crate::{ - backend::databases::databases, + backend::{self, databases::databases}, net::messages::{DataRow, Field, Protocol, RowDescription}, }; @@ -30,27 +30,30 @@ impl Command for ShowPools { Field::numeric("cl_waiting"), Field::numeric("sv_idle"), Field::numeric("sv_active"), + Field::numeric("sv_idle_xact"), Field::numeric("sv_total"), Field::numeric("maxwait"), Field::numeric("maxwait_us"), Field::text("pool_mode"), Field::bool("paused"), Field::bool("banned"), + Field::bool("healthy"), Field::numeric("errors"), Field::numeric("re_synced"), Field::numeric("out_of_sync"), + Field::numeric("force_closed"), Field::bool("online"), - Field::text("replica_lag"), Field::bool("schema_admin"), ]); let mut messages = vec![rd.message()?]; for (user, cluster) in databases().all() { for (shard_num, shard) in cluster.shards().iter().enumerate() { - for (role, pool) in shard.pools_with_roles() { + for (role, ban, pool) in shard.pools_with_roles_and_bans() { let mut row = DataRow::new(); let state = pool.state(); let maxwait = state.maxwait.as_secs() as i64; let maxwait_us = state.maxwait.subsec_micros() as i64; + let idle_in_transaction = backend::stats::idle_in_transaction(&pool); row.add(pool.id() as i64) .add(user.database.as_str()) @@ -62,17 +65,19 @@ impl Command for ShowPools { .add(state.waiting) .add(state.idle) .add(state.checked_out) + .add(idle_in_transaction) .add(state.total) .add(maxwait) .add(maxwait_us) .add(state.pooler_mode.to_string()) .add(state.paused) - .add(state.banned) + .add(ban.banned()) + .add(pool.healthy()) .add(state.errors) .add(state.re_synced) .add(state.out_of_sync) + .add(state.force_close) .add(state.online) - .add(state.replica_lag.simple_display()) .add(cluster.schema_admin()); messages.push(row.message()?); diff --git a/pgdog/src/admin/show_prepared_statements.rs b/pgdog/src/admin/show_prepared_statements.rs index 4329d7dc1..a448e2d96 100644 --- a/pgdog/src/admin/show_prepared_statements.rs +++ b/pgdog/src/admin/show_prepared_statements.rs @@ -1,4 +1,8 @@ -use crate::{frontend::PreparedStatements, stats::memory::MemoryUsage}; +use crate::{ + frontend::PreparedStatements, + net::{data_row::Data, ToDataRowColumn}, + stats::memory::MemoryUsage, +}; use super::prelude::*; @@ -20,11 +24,15 @@ impl Command for ShowPreparedStatements { let mut messages = vec![RowDescription::new(&[ Field::text("name"), Field::text("statement"), + Field::text("rewrite"), Field::numeric("used_by"), Field::numeric("memory_used"), ]) .message()?]; for (key, stmt) in statements.statements() { + let name = stmt.name(); + let rewrite = statements.rewritten_parse(&name).ok_or(Error::Empty)?; + let rewritten = statements.is_rewritten(&name); let name_memory = statements .names() .get(&stmt.name()) @@ -33,6 +41,11 @@ impl Command for ShowPreparedStatements { let mut dr = DataRow::new(); dr.add(stmt.name()) .add(key.query()?) + .add(if rewritten { + rewrite.query().to_data_row_column() + } else { + Data::null() + }) .add(stmt.used) .add(name_memory); messages.push(dr.message()?); diff --git a/pgdog/src/admin/show_query_cache.rs b/pgdog/src/admin/show_query_cache.rs index db69f3651..96e6aed49 100644 --- a/pgdog/src/admin/show_query_cache.rs +++ b/pgdog/src/admin/show_query_cache.rs @@ -43,7 +43,7 @@ impl Command for ShowQueryCache { continue; } let mut data_row = DataRow::new(); - let stats = { query.1.stats.lock().clone() }; + let stats = { *query.1.stats.lock() }; data_row .add(query.0) .add(stats.hits) @@ -60,7 +60,8 @@ impl Command for ShowQueryCache { mod test { use crate::{ backend::ShardingSchema, - net::{FromBytes, ToBytes}, + frontend::{BufferedQuery, PreparedStatements}, + net::{FromBytes, Parse, ToBytes}, }; use super::*; @@ -68,12 +69,16 @@ mod test { #[tokio::test] async fn test_show_query_cache() { let cache = Cache::get(); + let mut prepared_statements = PreparedStatements::default(); for q in 0..5 { cache - .parse( - format!("SELECT $1::bigint, {}", q).as_str(), + .query( + &BufferedQuery::Prepared(Parse::new_anonymous( + format!("SELECT $1::bigint, {}", q).as_str(), + )), &ShardingSchema::default(), + &mut prepared_statements, ) .unwrap(); } diff --git a/pgdog/src/admin/show_replication.rs b/pgdog/src/admin/show_replication.rs new file mode 100644 index 000000000..31dce8b99 --- /dev/null +++ b/pgdog/src/admin/show_replication.rs @@ -0,0 +1,89 @@ +use tokio::time::Instant; + +use crate::{ + backend::databases::databases, + net::{ + data_row::Data, + messages::{DataRow, Field, Protocol, RowDescription}, + ToDataRowColumn, + }, +}; + +use super::prelude::*; + +pub struct ShowReplication; + +#[async_trait] +impl Command for ShowReplication { + fn name(&self) -> String { + "SHOW REPLICATION".into() + } + + fn parse(_sql: &str) -> Result { + Ok(ShowReplication {}) + } + + async fn execute(&self) -> Result, Error> { + let rd = RowDescription::new(&[ + Field::bigint("id"), + Field::text("database"), + Field::text("user"), + Field::text("addr"), + Field::numeric("port"), + Field::numeric("shard"), + Field::text("role"), + Field::text("replica_lag"), + Field::text("pg_lsn"), + Field::text("lsn_age"), + Field::text("pg_is_in_recovery"), + ]); + let mut messages = vec![rd.message()?]; + let now = Instant::now(); + for (user, cluster) in databases().all() { + for (shard_num, shard) in cluster.shards().iter().enumerate() { + for (role, _ban, pool) in shard.pools_with_roles_and_bans() { + let mut row = DataRow::new(); + let state = pool.state(); + + let valid = state.lsn_stats.valid(); + let lsn_age = now.duration_since(state.lsn_stats.fetched); + + row.add(pool.id() as i64) + .add(user.database.as_str()) + .add(user.user.as_str()) + .add(pool.addr().host.as_str()) + .add(pool.addr().port as i64) + .add(shard_num as i64) + .add(role.to_string()) + .add(if valid { + state + .replica_lag + .as_millis() + .to_string() + .to_data_row_column() + } else { + Data::null() + }) + .add(if valid { + state.lsn_stats.lsn.to_string().to_data_row_column() + } else { + Data::null() + }) + .add(if valid { + lsn_age.as_millis().to_string().to_data_row_column() + } else { + Data::null() + }) + .add(if valid { + state.lsn_stats.replica.to_data_row_column() + } else { + Data::null() + }); + + messages.push(row.message()?); + } + } + } + Ok(messages) + } +} diff --git a/pgdog/src/admin/show_server_memory.rs b/pgdog/src/admin/show_server_memory.rs new file mode 100644 index 000000000..d6f6168d3 --- /dev/null +++ b/pgdog/src/admin/show_server_memory.rs @@ -0,0 +1,63 @@ +use crate::{ + backend::stats::stats, + net::messages::{DataRow, Field, Protocol, RowDescription}, +}; + +use super::prelude::*; + +pub struct ShowServerMemory; + +#[async_trait] +impl Command for ShowServerMemory { + fn name(&self) -> String { + "SHOW SERVER MEMORY".into() + } + + fn parse(_sql: &str) -> Result { + Ok(ShowServerMemory {}) + } + + async fn execute(&self) -> Result, Error> { + let rd = RowDescription::new(&[ + Field::bigint("pool_id"), + Field::text("database"), + Field::text("user"), + Field::text("addr"), + Field::numeric("port"), + Field::numeric("remote_pid"), + Field::numeric("buffer_reallocs"), + Field::numeric("buffer_reclaims"), + Field::numeric("buffer_bytes_used"), + Field::numeric("buffer_bytes_alloc"), + Field::numeric("prepared_statements_bytes"), + Field::numeric("net_buffer_bytes"), + Field::numeric("total_bytes"), + ]); + let mut messages = vec![rd.message()?]; + + let stats = stats(); + for (_, server) in stats { + let mut row = DataRow::new(); + let stats = server.stats; + let memory = &stats.memory; + + row.add(stats.pool_id as i64) + .add(server.addr.database_name.as_str()) + .add(server.addr.user.as_str()) + .add(server.addr.host.as_str()) + .add(server.addr.port as i64) + .add(stats.id.pid as i64) + .add(memory.buffer.reallocs as i64) + .add(memory.buffer.reclaims as i64) + .add(memory.buffer.bytes_used as i64) + .add(memory.buffer.bytes_alloc as i64) + .add(memory.prepared_statements as i64) + .add(memory.stream as i64) + .add((memory.total()) as i64); + + messages.push(row.message()?); + } + + Ok(messages) + } +} diff --git a/pgdog/src/admin/show_servers.rs b/pgdog/src/admin/show_servers.rs index 5f8548ccb..76dd026dc 100644 --- a/pgdog/src/admin/show_servers.rs +++ b/pgdog/src/admin/show_servers.rs @@ -44,6 +44,7 @@ impl Command for ShowServers { Ok(Self { row: NamedRow::new( &[ + Field::bigint("pool_id"), Field::text("database"), Field::text("user"), Field::text("addr"), @@ -52,6 +53,7 @@ impl Command for ShowServers { Field::text("connect_time"), Field::text("request_time"), Field::numeric("remote_pid"), + // Field::bigint("client_id"), Field::numeric("transactions"), Field::numeric("queries"), Field::numeric("rollbacks"), @@ -62,7 +64,6 @@ impl Command for ShowServers { Field::numeric("bytes_sent"), Field::numeric("age"), Field::text("application_name"), - Field::text("memory_used"), ], &mandatory, ), @@ -77,38 +78,34 @@ impl Command for ShowServers { let now_time = SystemTime::now(); for (_, server) in stats { - let age = now.duration_since(server.stats.created_at); - let request_age = now.duration_since(server.stats.last_used); + let stats = server.stats; + let age = now.duration_since(stats.created_at); + let request_age = now.duration_since(stats.last_used); let request_time = now_time - request_age; let dr = self .row .clone() + .add("pool_id", stats.pool_id) .add("database", server.addr.database_name) .add("user", server.addr.user) .add("addr", server.addr.host.as_str()) .add("port", server.addr.port.to_string()) - .add("state", server.stats.state.to_string()) - .add( - "connect_time", - format_time(server.stats.created_at_time.into()), - ) + .add("state", stats.state.to_string()) + .add("connect_time", format_time(stats.created_at_time.into())) .add("request_time", format_time(request_time.into())) - .add("remote_pid", server.stats.id.pid as i64) - .add("transactions", server.stats.total.transactions) - .add("queries", server.stats.total.queries) - .add("rollbacks", server.stats.total.rollbacks) - .add( - "prepared_statements", - server.stats.total.prepared_statements, - ) - .add("healthchecks", server.stats.total.healthchecks) - .add("errors", server.stats.total.errors) - .add("bytes_received", server.stats.total.bytes_received) - .add("bytes_sent", server.stats.total.bytes_sent) + .add("remote_pid", stats.id.pid as i64) + // .add("client_id", stats.client_id.map(|client| client.pid as i64)) + .add("transactions", stats.total.transactions) + .add("queries", stats.total.queries) + .add("rollbacks", stats.total.rollbacks) + .add("prepared_statements", stats.total.prepared_statements) + .add("healthchecks", stats.total.healthchecks) + .add("errors", stats.total.errors) + .add("bytes_received", stats.total.bytes_received) + .add("bytes_sent", stats.total.bytes_sent) .add("age", age.as_secs() as i64) .add("application_name", server.application_name.as_str()) - .add("memory_used", server.stats.total.memory_used) .data_row(); messages.push(dr.message()?); } diff --git a/pgdog/src/admin/show_stats.rs b/pgdog/src/admin/show_stats.rs index 93cf96cd5..febd5df52 100644 --- a/pgdog/src/admin/show_stats.rs +++ b/pgdog/src/admin/show_stats.rs @@ -36,12 +36,18 @@ impl Command for ShowStats { Field::numeric(&format!("{}_received", prefix)), Field::numeric(&format!("{}_sent", prefix)), Field::numeric(&format!("{}_xact_time", prefix)), + Field::numeric(&format!("{}_idle_xact_time", prefix)), Field::numeric(&format!("{}_query_time", prefix)), Field::numeric(&format!("{}_wait_time", prefix)), // Field::numeric(&format!("{}_client_parse_count", prefix)), Field::numeric(&format!("{}_server_parse_count", prefix)), Field::numeric(&format!("{}_bind_count", prefix)), Field::numeric(&format!("{}_close_count", prefix)), + Field::numeric(&format!("{}_errors", prefix)), + Field::numeric(&format!("{}_cleaned", prefix)), + Field::numeric(&format!("{}_rollbacks", prefix)), + Field::numeric(&format!("{}_connect_time", prefix)), + Field::numeric(&format!("{}_connect_count", prefix)), ] }) .collect::>(), @@ -78,11 +84,17 @@ impl Command for ShowStats { .add(stat.received) .add(stat.sent) .add(stat.xact_time.as_millis() as u64) + .add(stat.idle_xact_time.as_millis() as u64) .add(stat.query_time.as_millis() as u64) .add(stat.wait_time.as_millis() as u64) .add(stat.parse_count) .add(stat.bind_count) - .add(stat.close); + .add(stat.close) + .add(stat.errors) + .add(stat.cleaned) + .add(stat.rollbacks) + .add(stat.connect_time.as_millis() as u64) + .add(stat.connect_count); } messages.push(dr.message()?); diff --git a/pgdog/src/admin/tests/mod.rs b/pgdog/src/admin/tests/mod.rs new file mode 100644 index 000000000..4fe94b82a --- /dev/null +++ b/pgdog/src/admin/tests/mod.rs @@ -0,0 +1,467 @@ +use crate::admin::Command; +use crate::backend::databases::{databases, from_config, replace_databases, Databases}; +use crate::backend::pool::mirror_stats::Counts; +use crate::config::{self, ConfigAndUsers, Database, Role, User as ConfigUser}; +use crate::net::messages::{DataRow, DataType, FromBytes, Protocol, RowDescription}; + +use super::show_client_memory::ShowClientMemory; +use super::show_config::ShowConfig; +use super::show_lists::ShowLists; +use super::show_mirrors::ShowMirrors; +use super::show_pools::ShowPools; +use super::show_server_memory::ShowServerMemory; + +#[derive(Clone)] +struct SavedState { + config: ConfigAndUsers, + databases: Databases, +} + +/// Minimal harness to snapshot and restore admin singletons touched by SHOW commands. +pub(crate) struct TestAdminContext { + original: SavedState, +} + +impl TestAdminContext { + pub(crate) fn new() -> Self { + let config = (*config::config()).clone(); + let databases = (*databases()).clone(); + + Self { + original: SavedState { config, databases }, + } + } + + pub(crate) fn set_config(&self, config: ConfigAndUsers) { + // Update the live config first so downstream readers pick up test values. + config::set(config.clone()).expect("failed to install test config"); + + // Rebuild the in-process database registry from the supplied config. + replace_databases(from_config(&config), false).unwrap(); + } +} + +impl Default for TestAdminContext { + fn default() -> Self { + Self::new() + } +} + +impl Drop for TestAdminContext { + fn drop(&mut self) { + // Restore the original configuration and database registry so other tests remain isolated. + let _ = config::set(self.original.config.clone()); + replace_databases(self.original.databases.clone(), false).unwrap(); + } +} + +#[tokio::test(flavor = "current_thread")] +async fn show_pools_reports_schema_admin_flag() { + let context = TestAdminContext::new(); + + let mut config = ConfigAndUsers::default(); + config.config.databases.push(Database { + name: "app".into(), + host: "127.0.0.1".into(), + role: Role::Primary, + shard: 0, + ..Default::default() + }); + config.users.users.push(ConfigUser { + name: "alice".into(), + database: "app".into(), + password: Some("secret".into()), + schema_admin: true, + ..Default::default() + }); + + context.set_config(config); + + let command = ShowPools; + let messages = command + .execute() + .await + .expect("show pools execution failed"); + + assert!( + messages.len() >= 2, + "expected row description plus data row" + ); + + let row_description = RowDescription::from_bytes(messages[0].payload()) + .expect("row description message should parse"); + let actual_names: Vec<&str> = row_description + .fields + .iter() + .map(|field| field.name.as_str()) + .collect(); + let expected_names = vec![ + "id", + "database", + "user", + "addr", + "port", + "shard", + "role", + "cl_waiting", + "sv_idle", + "sv_active", + "sv_idle_xact", + "sv_total", + "maxwait", + "maxwait_us", + "pool_mode", + "paused", + "banned", + "healthy", + "errors", + "re_synced", + "out_of_sync", + "force_closed", + "online", + "schema_admin", + ]; + assert_eq!(actual_names, expected_names); + + let schema_admin_field = row_description + .field_index("schema_admin") + .and_then(|idx| row_description.field(idx)) + .expect("schema_admin field present"); + assert_eq!(schema_admin_field.data_type(), DataType::Bool); + + let data_row = DataRow::from_bytes(messages[1].payload()).expect("data row should parse"); + let schema_admin_index = row_description + .field_index("schema_admin") + .expect("schema_admin column index"); + let schema_admin_value = data_row + .get_text(schema_admin_index) + .expect("schema_admin value should be textual"); + assert_eq!(schema_admin_value.as_str(), "t"); +} + +#[tokio::test(flavor = "current_thread")] +async fn show_config_pretty_prints_general_settings() { + let context = TestAdminContext::new(); + + let mut config = ConfigAndUsers::default(); + config.config.general.default_pool_size = 42; + config.config.general.connect_timeout = 2_000; + + context.set_config(config); + + let command = ShowConfig; + let messages = command + .execute() + .await + .expect("show config execution failed"); + + assert!(messages.len() > 1, "expected row description and rows"); + + let row_description = RowDescription::from_bytes(messages[0].payload()) + .expect("row description message should parse"); + let column_names: Vec<&str> = row_description + .fields + .iter() + .map(|field| field.name.as_str()) + .collect(); + assert_eq!(column_names, vec!["name", "value"]); + + for field in row_description.fields.iter() { + assert_eq!(field.data_type(), DataType::Text); + } + + let rows: Vec<(String, String)> = messages + .iter() + .skip(1) + .map(|message| { + let row = + DataRow::from_bytes(message.payload()).expect("data row should parse into columns"); + let name = row.get_text(0).expect("name column should be text"); + let value = row.get_text(1).expect("value column should be text"); + (name, value) + }) + .collect(); + + let pool_size = rows + .iter() + .find(|(name, _)| name == "default_pool_size") + .map(|(_, value)| value) + .expect("default_pool_size row should be present"); + assert_eq!(pool_size, "42"); + + let connect_timeout = rows + .iter() + .find(|(name, _)| name == "connect_timeout") + .map(|(_, value)| value) + .expect("connect_timeout row should be present"); + assert_eq!(connect_timeout, "2s"); +} + +#[tokio::test(flavor = "current_thread")] +async fn show_mirrors_reports_counts() { + let context = TestAdminContext::new(); + + let mut config = ConfigAndUsers::default(); + config.config.databases.push(Database { + name: "app".into(), + host: "127.0.0.1".into(), + role: Role::Primary, + shard: 0, + ..Default::default() + }); + config.users.users.push(ConfigUser { + name: "alice".into(), + database: "app".into(), + password: Some("secret".into()), + ..Default::default() + }); + + context.set_config(config); + + // Seed synthetic mirror stats for the configured cluster. + let databases = databases(); + let (user, cluster) = databases.all().iter().next().expect("cluster should exist"); + assert_eq!(user.database, "app"); + assert_eq!(user.user, "alice"); + { + let cluster_stats = cluster.stats(); + let mut stats = cluster_stats.lock(); + stats.counts = Counts { + total_count: 5, + mirrored_count: 4, + dropped_count: 1, + error_count: 2, + queue_length: 3, + }; + } + + let command = ShowMirrors; + let messages = command + .execute() + .await + .expect("show mirrors execution failed"); + + assert_eq!(messages[0].code(), 'T'); + let row_description = + RowDescription::from_bytes(messages[0].payload()).expect("row description should parse"); + let expected_columns = [ + "database", + "user", + "total_count", + "mirrored_count", + "dropped_count", + "error_count", + "queue_length", + ]; + let actual_columns: Vec<&str> = row_description + .fields + .iter() + .map(|field| field.name.as_str()) + .collect(); + assert_eq!(actual_columns, expected_columns); + + assert!(messages.len() > 1, "expected at least one data row"); + let data_row = DataRow::from_bytes(messages[1].payload()).expect("data row should parse"); + assert_eq!(data_row.get_text(0).as_deref(), Some("app")); + assert_eq!(data_row.get_text(1).as_deref(), Some("alice")); + assert_eq!(data_row.get_int(2, true), Some(5)); + assert_eq!(data_row.get_int(3, true), Some(4)); + assert_eq!(data_row.get_int(4, true), Some(1)); + assert_eq!(data_row.get_int(5, true), Some(2)); + assert_eq!(data_row.get_int(6, true), Some(3)); +} + +#[tokio::test(flavor = "current_thread")] +async fn show_lists_reports_basic_counts() { + let context = TestAdminContext::new(); + + let mut config = ConfigAndUsers::default(); + config.config.databases.push(Database { + name: "app".into(), + host: "127.0.0.1".into(), + role: Role::Primary, + shard: 0, + ..Default::default() + }); + config.users.users.push(ConfigUser { + name: "alice".into(), + database: "app".into(), + password: Some("secret".into()), + ..Default::default() + }); + + context.set_config(config.clone()); + + let command = ShowLists; + let messages = command + .execute() + .await + .expect("show lists execution failed"); + + assert_eq!(messages.len(), 2, "expected row description plus one row"); + + let row_description = + RowDescription::from_bytes(messages[0].payload()).expect("row description should parse"); + let column_names: Vec<&str> = row_description + .fields + .iter() + .map(|field| field.name.as_str()) + .collect(); + let expected_columns = vec![ + "databases", + "users", + "pools", + "used_clients", + "used_clients", + "free_servers", + "used_servers", + ]; + assert_eq!(column_names, expected_columns); + + for field in row_description.fields.iter() { + assert_eq!(field.data_type(), DataType::Numeric); + } + + let data_row = DataRow::from_bytes(messages[1].payload()).expect("data row should parse"); + + let database_count = data_row.get_int(0, true).expect("databases value"); + let user_count = data_row.get_int(1, true).expect("users value"); + let pool_count = data_row.get_int(2, true).expect("pools value"); + let used_clients_placeholder = data_row.get_int(3, true).expect("used clients placeholder"); + let used_clients = data_row.get_int(4, true).expect("used clients value"); + let free_servers = data_row.get_int(5, true).expect("free servers value"); + let used_servers = data_row.get_int(6, true).expect("used servers value"); + + assert_eq!(database_count, config.config.databases.len() as i64); + assert_eq!(user_count, config.users.users.len() as i64); + + let expected_pools: usize = databases() + .all() + .values() + .map(|cluster| { + cluster + .shards() + .iter() + .map(|shard| shard.pools().len()) + .sum::() + }) + .sum(); + assert_eq!(pool_count, expected_pools as i64); + + assert_eq!(used_clients_placeholder, 0); + assert_eq!(used_clients, 0); + assert_eq!(free_servers, 0); + assert_eq!(used_servers, 0); +} + +#[tokio::test(flavor = "current_thread")] +async fn show_server_memory_reports_memory_stats() { + let command = ShowServerMemory; + let messages = command + .execute() + .await + .expect("show server memory execution failed"); + + assert!(messages.len() >= 1, "expected at least row description"); + + let row_description = RowDescription::from_bytes(messages[0].payload()) + .expect("row description message should parse"); + let actual_names: Vec<&str> = row_description + .fields + .iter() + .map(|field| field.name.as_str()) + .collect(); + let expected_names = vec![ + "pool_id", + "database", + "user", + "addr", + "port", + "remote_pid", + "buffer_reallocs", + "buffer_reclaims", + "buffer_bytes_used", + "buffer_bytes_alloc", + "prepared_statements_bytes", + "net_buffer_bytes", + "total_bytes", + ]; + assert_eq!(actual_names, expected_names); + + for (idx, expected_type) in [ + DataType::Bigint, + DataType::Text, + DataType::Text, + DataType::Text, + DataType::Numeric, + DataType::Numeric, + DataType::Numeric, + DataType::Numeric, + DataType::Numeric, + DataType::Numeric, + DataType::Numeric, + DataType::Numeric, + DataType::Numeric, + ] + .iter() + .enumerate() + { + let field = row_description.field(idx).expect("field should exist"); + assert_eq!(field.data_type(), *expected_type); + } +} + +#[tokio::test(flavor = "current_thread")] +async fn show_client_memory_reports_memory_stats() { + let command = ShowClientMemory; + let messages = command + .execute() + .await + .expect("show client memory execution failed"); + + assert!(messages.len() >= 1, "expected at least row description"); + + let row_description = RowDescription::from_bytes(messages[0].payload()) + .expect("row description message should parse"); + let actual_names: Vec<&str> = row_description + .fields + .iter() + .map(|field| field.name.as_str()) + .collect(); + let expected_names = vec![ + "client_id", + "database", + "user", + "addr", + "port", + "buffer_reallocs", + "buffer_reclaims", + "buffer_bytes_used", + "buffer_bytes_alloc", + "prepared_statements_bytes", + "net_buffer_bytes", + "total_bytes", + ]; + assert_eq!(actual_names, expected_names); + + for (idx, expected_type) in [ + DataType::Bigint, + DataType::Text, + DataType::Text, + DataType::Text, + DataType::Numeric, + DataType::Numeric, + DataType::Numeric, + DataType::Numeric, + DataType::Numeric, + DataType::Numeric, + DataType::Numeric, + DataType::Numeric, + ] + .iter() + .enumerate() + { + let field = row_description.field(idx).expect("field should exist"); + assert_eq!(field.data_type(), *expected_type); + } +} diff --git a/pgdog/src/auth/md5.rs b/pgdog/src/auth/md5.rs index 51a917a0f..7c95ddf66 100644 --- a/pgdog/src/auth/md5.rs +++ b/pgdog/src/auth/md5.rs @@ -22,7 +22,7 @@ impl<'a> Client<'a> { Self { password, user, - salt: rand::thread_rng().gen(), + salt: rand::rng().random(), } } diff --git a/pgdog/src/auth/scram/client.rs b/pgdog/src/auth/scram/client.rs index 915a3d5dc..8a669579e 100644 --- a/pgdog/src/auth/scram/client.rs +++ b/pgdog/src/auth/scram/client.rs @@ -66,3 +66,70 @@ impl<'a> Client<'a> { Ok(()) } } + +#[cfg(test)] +mod tests { + use super::*; + use scram::{ + hash_password, AuthenticationProvider, AuthenticationStatus, PasswordInfo, ScramServer, + }; + use std::num::NonZeroU32; + + struct TestProvider { + password: String, + } + + impl AuthenticationProvider for TestProvider { + fn get_password_for(&self, _user: &str) -> Option { + let iterations = NonZeroU32::new(4096).unwrap(); + let salt = b"testsalt".to_vec(); + let hash = hash_password(&self.password, iterations, &salt); + Some(PasswordInfo::new( + hash.to_vec(), + iterations.get() as u16, + salt, + )) + } + } + + #[test] + fn scram_client_full_handshake_succeeds() { + let mut client = Client::new("user", "secret"); + let provider = TestProvider { + password: "secret".into(), + }; + let server = ScramServer::new(provider); + + let client_first = client.first().expect("client first message"); + let server = server + .handle_client_first(&client_first) + .expect("server handle client first"); + let (server_client_final, server_first) = server.server_first(); + + client + .server_first(&server_first) + .expect("client handles server first"); + + let client_final = client.last().expect("client final message"); + let server_final = server_client_final + .handle_client_final(&client_final) + .expect("server handles client final"); + let (status, server_final_msg) = server_final.server_final(); + assert_eq!(status, AuthenticationStatus::Authenticated); + + client + .server_last(&server_final_msg) + .expect("client validates server final"); + } + + #[test] + fn scram_client_enforces_call_order() { + let mut client = Client::new("user", "secret"); + let err = client + .last() + .expect_err("last without handshake should fail"); + matches!(err, Error::OutOfOrder) + .then_some(()) + .expect("expected out-of-order error"); + } +} diff --git a/pgdog/src/auth/scram/server.rs b/pgdog/src/auth/scram/server.rs index 3cb6a3a36..34fdeabe8 100644 --- a/pgdog/src/auth/scram/server.rs +++ b/pgdog/src/auth/scram/server.rs @@ -52,7 +52,7 @@ impl AuthenticationProvider for UserPassword { fn get_password_for(&self, _user: &str) -> Option { // TODO: This is slow. We should move it to its own thread pool. let iterations = 4096; - let salt = rand::thread_rng().gen::<[u8; 16]>().to_vec(); + let salt = rand::rng().random::<[u8; 16]>().to_vec(); let hash = hash_password(&self.password, NonZeroU32::new(iterations).unwrap(), &salt); Some(PasswordInfo::new(hash.to_vec(), iterations as u16, salt)) } @@ -204,6 +204,43 @@ impl Server { } } +#[cfg(test)] +mod tests { + use super::*; + use base64::engine::general_purpose::STANDARD; + + #[test] + fn user_password_provider_generates_info() { + let server = Server::new("secret"); + let provider = match server.provider { + Provider::Plain(ref inner) => inner.clone(), + _ => unreachable!(), + }; + + assert!( + provider.get_password_for("user").is_some(), + "plain provider should produce password info" + ); + } + + #[test] + fn hashed_password_provider_parses_scram_hash() { + let iterations = std::num::NonZeroU32::new(4096).unwrap(); + let salt = b"testsalt"; + let hash = hash_password("secret", iterations, salt); + let salt_b64 = STANDARD.encode(salt); + let hash_b64 = STANDARD.encode(hash.as_ref()); + let scram_hash = format!("SCRAM-SHA-256${}:{salt_b64}:${hash_b64}", iterations.get()); + + let provider = HashedPassword { hash: scram_hash }; + + assert!( + provider.get_password_for("user").is_some(), + "hashed provider should produce password info" + ); + } +} + #[cfg(test)] mod test { use super::*; diff --git a/pgdog/src/backend/connect_reason.rs b/pgdog/src/backend/connect_reason.rs new file mode 100644 index 000000000..e9500d6ea --- /dev/null +++ b/pgdog/src/backend/connect_reason.rs @@ -0,0 +1,29 @@ +use std::fmt::Display; + +#[derive(Debug, Clone, Copy, Default, PartialEq)] +pub enum ConnectReason { + BelowMin, + ClientWaiting, + Replication, + PubSub, + Probe, + Healthcheck, + #[default] + Other, +} + +impl Display for ConnectReason { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let reason = match self { + Self::BelowMin => "min", + Self::ClientWaiting => "client", + Self::Replication => "replication", + Self::PubSub => "pub/sub", + Self::Probe => "probe", + Self::Healthcheck => "healthcheck", + Self::Other => "other", + }; + + write!(f, "{}", reason) + } +} diff --git a/pgdog/src/backend/databases.rs b/pgdog/src/backend/databases.rs index 8d7f2480b..56fdde93b 100644 --- a/pgdog/src/backend/databases.rs +++ b/pgdog/src/backend/databases.rs @@ -9,15 +9,17 @@ use parking_lot::lock_api::MutexGuard; use parking_lot::{Mutex, RawMutex}; use tracing::{debug, error, info, warn}; +use crate::backend::replication::ShardedSchemas; use crate::config::PoolerMode; use crate::frontend::client::query_engine::two_pc::Manager; use crate::frontend::router::parser::Cache; +use crate::frontend::router::sharding::mapping::mapping_valid; use crate::frontend::router::sharding::Mapping; use crate::frontend::PreparedStatements; use crate::{ backend::pool::PoolConfig, config::{config, load, ConfigAndUsers, ManualQuery, Role}, - net::messages::BackendKeyData, + net::{messages::BackendKeyData, tls}, }; use super::{ @@ -44,7 +46,7 @@ pub fn databases() -> Arc { } /// Replace databases pooler-wide. -pub fn replace_databases(new_databases: Databases, reload: bool) { +pub fn replace_databases(new_databases: Databases, reload: bool) -> Result<(), Error> { // Order of operations is important // to ensure zero downtime for clients. let old_databases = databases(); @@ -52,29 +54,49 @@ pub fn replace_databases(new_databases: Databases, reload: bool) { reload_notify::started(); if reload { // Move whatever connections we can over to new pools. - old_databases.move_conns_to(&new_databases); + old_databases.move_conns_to(&new_databases)?; } new_databases.launch(); DATABASES.store(new_databases); old_databases.shutdown(); reload_notify::done(); + + Ok(()) } /// Re-create all connections. -pub fn reconnect() { - replace_databases(databases().duplicate(), false); +pub fn reconnect() -> Result<(), Error> { + let config = config(); + let databases = from_config(&config); + + replace_databases(databases, false)?; + Ok(()) +} + +/// Re-create databases from existing config, +/// preserving connections. +pub fn reload_from_existing() -> Result<(), Error> { + let _lock = lock(); + + let config = config(); + let databases = from_config(&config); + + replace_databases(databases, true)?; + Ok(()) } /// Initialize the databases for the first time. -pub fn init() { +pub fn init() -> Result<(), Error> { let config = config(); - replace_databases(from_config(&config), false); + replace_databases(from_config(&config), false)?; // Resize query cache Cache::resize(config.config.general.query_cache_limit); // Start two-pc manager. let _monitor = Manager::get(); + + Ok(()) } /// Shutdown all databases. @@ -88,7 +110,9 @@ pub fn reload() -> Result<(), Error> { let new_config = load(&old_config.config_path, &old_config.users_path)?; let databases = from_config(&new_config); - replace_databases(databases, true); + replace_databases(databases, true)?; + + tls::reload()?; // Remove any unused prepared statements. PreparedStatements::global() @@ -172,6 +196,15 @@ impl ToUser for (&str, Option<&str>) { } } +impl ToUser for &pgdog_config::User { + fn to_user(&self) -> User { + User { + user: self.name.clone(), + database: self.database.clone(), + } + } +} + /// Databases. #[derive(Default, Clone)] pub struct Databases { @@ -289,34 +322,20 @@ impl Databases { /// Move all connections we can from old databases config to new /// databases config. - pub(crate) fn move_conns_to(&self, destination: &Databases) -> usize { + pub(crate) fn move_conns_to(&self, destination: &Databases) -> Result { let mut moved = 0; for (user, cluster) in &self.databases { let dest = destination.databases.get(user); if let Some(dest) = dest { if cluster.can_move_conns_to(dest) { - cluster.move_conns_to(dest); + cluster.move_conns_to(dest)?; moved += 1; } } } - moved - } - - /// Create new identical databases. - fn duplicate(&self) -> Databases { - Self { - databases: self - .databases - .iter() - .map(|(k, v)| (k.clone(), v.duplicate())) - .collect(), - manual_queries: self.manual_queries.clone(), - mirrors: self.mirrors.clone(), - mirror_configs: self.mirror_configs.clone(), - } + Ok(moved) } /// Shutdown all pools. @@ -363,84 +382,89 @@ pub(crate) fn new_pool( let sharded_tables = config.sharded_tables(); let omnisharded_tables = config.omnisharded_tables(); let sharded_mappings = config.sharded_mappings(); + let sharded_schemas = config.sharded_schemas(); let general = &config.general; let databases = config.databases(); - let shards = databases.get(&user.database); - - if let Some(shards) = shards { - let mut shard_configs = vec![]; - for user_databases in shards { - let has_single_replica = user_databases.len() == 1; - let primary = user_databases - .iter() - .find(|d| d.role == Role::Primary) - .map(|primary| PoolConfig { - address: Address::new(primary, user), - config: Config::new(general, primary, user, has_single_replica), - }); - let replicas = user_databases - .iter() - .filter(|d| d.role == Role::Replica) - .map(|replica| PoolConfig { - address: Address::new(replica, user), - config: Config::new(general, replica, user, has_single_replica), - }) - .collect::>(); - - shard_configs.push(ClusterShardConfig { primary, replicas }); - } - let mut sharded_tables = sharded_tables - .get(&user.database) - .cloned() - .unwrap_or(vec![]); + let shards = databases.get(&user.database).cloned()?; - for sharded_table in &mut sharded_tables { - let mappings = sharded_mappings.get(&( - sharded_table.database.clone(), - sharded_table.column.clone(), - sharded_table.name.clone(), - )); + let mut shard_configs = vec![]; + for user_databases in shards { + let has_single_replica = user_databases.len() == 1; + let primary = user_databases + .iter() + .find(|d| d.role == Role::Primary) + .map(|primary| PoolConfig { + address: Address::new(primary, user, primary.number), + config: Config::new(general, primary, user, has_single_replica), + }); + let replicas = user_databases + .iter() + .filter(|d| matches!(d.role, Role::Replica | Role::Auto)) // Auto role is assumed read-only until proven otherwise. + .map(|replica| PoolConfig { + address: Address::new(replica, user, replica.number), + config: Config::new(general, replica, user, has_single_replica), + }) + .collect::>(); + + shard_configs.push(ClusterShardConfig { primary, replicas }); + } - if let Some(mappings) = mappings { - sharded_table.mapping = Mapping::new(mappings); - - if let Some(ref mapping) = sharded_table.mapping { - if !mapping.valid() { - warn!( - "sharded table name=\"{}\", column=\"{}\" has overlapping ranges", - sharded_table.name.as_ref().unwrap_or(&String::from("")), - sharded_table.column - ); - } + let mut sharded_tables = sharded_tables + .get(&user.database) + .cloned() + .unwrap_or_default(); + let sharded_schemas = sharded_schemas + .get(&user.database) + .cloned() + .unwrap_or_default(); + + for sharded_table in &mut sharded_tables { + let mappings = sharded_mappings.get(&( + sharded_table.database.clone(), + sharded_table.column.clone(), + sharded_table.name.clone(), + )); + + if let Some(mappings) = mappings { + sharded_table.mapping = Mapping::new(mappings); + + if let Some(ref mapping) = sharded_table.mapping { + if !mapping_valid(mapping) { + warn!( + "sharded table name=\"{}\", column=\"{}\" has overlapping ranges", + sharded_table.name.as_ref().unwrap_or(&String::from("")), + sharded_table.column + ); } } } + } - let omnisharded_tables = omnisharded_tables - .get(&user.database) - .cloned() - .unwrap_or(vec![]); - let sharded_tables = ShardedTables::new(sharded_tables, omnisharded_tables); - - let cluster_config = ClusterConfig::new( - general, - user, - &shard_configs, - sharded_tables, - config.multi_tenant(), - ); + let omnisharded_tables = omnisharded_tables + .get(&user.database) + .cloned() + .unwrap_or(vec![]); + let sharded_tables = ShardedTables::new(sharded_tables, omnisharded_tables); + let sharded_schemas = ShardedSchemas::new(sharded_schemas); + + let cluster_config = ClusterConfig::new( + general, + user, + &shard_configs, + sharded_tables, + config.multi_tenant(), + sharded_schemas, + &config.rewrite, + ); - Some(( - User { - user: user.name.clone(), - database: user.database.clone(), - }, - Cluster::new(cluster_config), - )) - } else { - None - } + Some(( + User { + user: user.name.clone(), + database: user.database.clone(), + }, + Cluster::new(cluster_config), + )) } /// Load databases from config. @@ -617,6 +641,7 @@ mod tests { ..Default::default() }, ], + ..Default::default() }; let databases = from_config(&ConfigAndUsers { @@ -691,6 +716,7 @@ mod tests { }, // Note: user2 missing for dest_db - this should disable mirroring ], + ..Default::default() }; let databases = from_config(&ConfigAndUsers { @@ -760,6 +786,7 @@ mod tests { ..Default::default() }, ], + ..Default::default() }; let databases = from_config(&ConfigAndUsers { @@ -837,6 +864,7 @@ mod tests { ..Default::default() }, ], + ..Default::default() }; let databases = from_config(&ConfigAndUsers { @@ -929,6 +957,7 @@ mod tests { ..Default::default() }, ], + ..Default::default() }; let databases = from_config(&ConfigAndUsers { @@ -1006,6 +1035,7 @@ mod tests { ..Default::default() }, ], + ..Default::default() }; let databases = from_config(&ConfigAndUsers { @@ -1056,7 +1086,10 @@ mod tests { }]; // No users at all - let users = crate::config::Users { users: vec![] }; + let users = crate::config::Users { + users: vec![], + ..Default::default() + }; let databases = from_config(&ConfigAndUsers { config: config.clone(), @@ -1083,6 +1116,7 @@ mod tests { }, // No user for dest_db! ], + ..Default::default() }; let databases_partial = from_config(&ConfigAndUsers { @@ -1110,6 +1144,7 @@ mod tests { }, // No user for source_db! ], + ..Default::default() }; let databases_dest_only = from_config(&ConfigAndUsers { diff --git a/pgdog/src/backend/disconnect_reason.rs b/pgdog/src/backend/disconnect_reason.rs new file mode 100644 index 000000000..f49edaa33 --- /dev/null +++ b/pgdog/src/backend/disconnect_reason.rs @@ -0,0 +1,39 @@ +use std::fmt::Display; + +#[derive(Debug, Clone, Copy, Default)] +pub enum DisconnectReason { + Idle, + Old, + Error, + Offline, + ForceClose, + Paused, + ReplicationMode, + OutOfSync, + Unhealthy, + Healthcheck, + PubSub, + #[default] + Other, +} + +impl Display for DisconnectReason { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let reason = match self { + Self::Idle => "idle", + Self::Old => "max age", + Self::Error => "error", + Self::Other => "other", + Self::ForceClose => "force close", + Self::Paused => "pool paused", + Self::Offline => "pool offline", + Self::OutOfSync => "out of sync", + Self::ReplicationMode => "in replication mode", + Self::Unhealthy => "unhealthy", + Self::Healthcheck => "standalone healthcheck", + Self::PubSub => "pub/sub", + }; + + write!(f, "{}", reason) + } +} diff --git a/pgdog/src/backend/error.rs b/pgdog/src/backend/error.rs index f5e673368..7cdf6536e 100644 --- a/pgdog/src/backend/error.rs +++ b/pgdog/src/backend/error.rs @@ -33,9 +33,18 @@ pub enum Error { #[error("server not connected")] NotConnected, + #[error("direct-to-shard not connected")] + DirectToShardNotConnected, + + #[error("multi-shard not connected")] + MultiShardNotConnected, + #[error("multi shard copy not connected")] CopyNotConnected, + #[error("cluster not connected")] + ClusterNotConnected, + #[error("{0}")] Pool(#[from] crate::backend::pool::Error), @@ -84,6 +93,9 @@ pub enum Error { #[error("protocol is out of sync")] ProtocolOutOfSync, + #[error("rollback left server in inconsistent state")] + RollbackFailed, + #[error("decoder is missing required data to decode row")] DecoderRowError, @@ -116,6 +128,9 @@ pub enum Error { #[error("2pc commit supported with multi-shard binding only")] TwoPcMultiShardOnly, + + #[error("unsupported aggregation {function}: {reason}")] + UnsupportedAggregation { function: String, reason: String }, } impl From for Error { diff --git a/pgdog/src/backend/mod.rs b/pgdog/src/backend/mod.rs index 7bb4db713..2bc400685 100644 --- a/pgdog/src/backend/mod.rs +++ b/pgdog/src/backend/mod.rs @@ -1,6 +1,8 @@ //! pgDog backend managers connections to PostgreSQL. +pub mod connect_reason; pub mod databases; +pub mod disconnect_reason; pub mod error; pub mod maintenance_mode; pub mod pool; @@ -14,8 +16,10 @@ pub mod server; pub mod server_options; pub mod stats; +pub use connect_reason::ConnectReason; +pub use disconnect_reason::DisconnectReason; pub use error::Error; -pub use pool::{Cluster, ClusterShardConfig, Pool, Replicas, Shard, ShardingSchema}; +pub use pool::{Cluster, ClusterShardConfig, LoadBalancer, Pool, Shard, ShardingSchema}; pub use prepared_statements::PreparedStatements; pub use protocol::*; pub use pub_sub::{PubSubClient, PubSubListener}; diff --git a/pgdog/src/backend/pool/address.rs b/pgdog/src/backend/pool/address.rs index f438f76d3..7b86fc0c9 100644 --- a/pgdog/src/backend/pool/address.rs +++ b/pgdog/src/backend/pool/address.rs @@ -20,11 +20,13 @@ pub struct Address { pub user: String, /// Password. pub password: String, + /// Database number (in the config). + pub database_number: usize, } impl Address { /// Create new address from config values. - pub fn new(database: &Database, user: &User) -> Self { + pub fn new(database: &Database, user: &User, database_number: usize) -> Self { Address { host: database.host.clone(), port: database.port, @@ -47,6 +49,7 @@ impl Address { } else { user.password().to_string() }, + database_number, } } @@ -74,13 +77,18 @@ impl Address { user: "pgdog".into(), password: "pgdog".into(), database_name: "pgdog".into(), + database_number: 0, } } } impl std::fmt::Display for Address { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "{}:{}, {}", self.host, self.port, self.database_name) + write!( + f, + "{}@{}:{}/{}", + self.user, self.host, self.port, self.database_name + ) } } @@ -100,6 +108,7 @@ impl TryFrom for Address { password, user, database_name, + database_number: 0, }) } } @@ -124,7 +133,7 @@ mod test { ..Default::default() }; - let address = Address::new(&database, &user); + let address = Address::new(&database, &user, 0); assert_eq!(address.host, "127.0.0.1"); assert_eq!(address.port, 6432); @@ -136,7 +145,7 @@ mod test { database.password = Some("hunter3".into()); database.user = Some("alice".into()); - let address = Address::new(&database, &user); + let address = Address::new(&database, &user, 0); assert_eq!(address.database_name, "not_pgdog"); assert_eq!(address.user, "alice"); diff --git a/pgdog/src/backend/pool/ban.rs b/pgdog/src/backend/pool/ban.rs deleted file mode 100644 index df6e7597b..000000000 --- a/pgdog/src/backend/pool/ban.rs +++ /dev/null @@ -1,64 +0,0 @@ -//! Pool ban. -use std::time::Duration; -use tokio::time::Instant; - -use super::Error; - -/// Pool ban. -#[derive(Debug, Copy, Clone)] -pub struct Ban { - /// When the banw as created. - pub(super) created_at: Instant, - /// Why it was created. - pub(super) reason: Error, - /// Ban timeout - pub(super) ban_timeout: Duration, -} - -impl std::fmt::Display for Ban { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "{} ({:.3}ms)", self.reason, self.ban_timeout.as_millis()) - } -} - -impl Ban { - /// Check if the ban has expired. - pub(super) fn expired(&self, now: Instant) -> bool { - if self.reason == Error::ManualBan { - false - } else { - let duration = now.duration_since(self.created_at); - - duration > self.ban_timeout - } - } - - pub(super) fn manual(&self) -> bool { - self.reason == Error::ManualBan - } -} - -#[cfg(test)] -mod test { - use super::*; - - #[test] - fn test_expired() { - let ban_timeout = Duration::from_secs(300); - let created_at = Instant::now(); - - let mut ban = Ban { - created_at, - reason: Error::CheckoutTimeout, - ban_timeout, - }; - - let later = created_at + ban_timeout + Duration::from_secs(1); - - assert!(!ban.expired(Instant::now())); - assert!(ban.expired(later)); - - ban.reason = Error::ManualBan; - assert!(!ban.expired(later)); - } -} diff --git a/pgdog/src/backend/pool/cleanup.rs b/pgdog/src/backend/pool/cleanup.rs index d96ef9e45..f5146ee6a 100644 --- a/pgdog/src/backend/pool/cleanup.rs +++ b/pgdog/src/backend/pool/cleanup.rs @@ -71,6 +71,11 @@ impl Cleanup { clean } + /// Number of queries to run for cleanup. + pub fn len(&self) -> usize { + self.queries.len() + } + /// Cleanup prepared statements. pub fn prepared_statements() -> Self { Self { diff --git a/pgdog/src/backend/pool/cluster.rs b/pgdog/src/backend/pool/cluster.rs index aebda7c71..8fb2e2b1f 100644 --- a/pgdog/src/backend/pool/cluster.rs +++ b/pgdog/src/backend/pool/cluster.rs @@ -1,26 +1,34 @@ //! A collection of replicas and a primary. use parking_lot::{Mutex, RwLock}; -use std::sync::Arc; +use pgdog_config::{PreparedStatements, QueryParserEngine, QueryParserLevel, Rewrite, RewriteMode}; +use std::{ + sync::{ + atomic::{AtomicBool, Ordering}, + Arc, + }, + time::Duration, +}; use tokio::spawn; use tracing::{error, info}; use crate::{ backend::{ databases::{databases, User as DatabaseUser}, - replication::{ReplicationConfig, ShardedColumn}, + replication::{ReplicationConfig, ShardedSchemas}, Schema, ShardedTables, }, config::{ - General, MultiTenant, PoolerMode, ReadWriteSplit, ReadWriteStrategy, ShardedTable, User, + ConnectionRecovery, General, MultiTenant, PoolerMode, ReadWriteSplit, ReadWriteStrategy, + ShardedTable, User, }, net::{messages::BackendKeyData, Query}, }; -use super::{Address, Config, Error, Guard, MirrorStats, Request, Shard}; +use super::{Address, Config, Error, Guard, MirrorStats, Request, Shard, ShardConfig}; use crate::config::LoadBalancingStrategy; -#[derive(Clone, Debug)] +#[derive(Clone, Debug, Default)] /// Database configuration. pub struct PoolConfig { /// Database address. @@ -38,16 +46,25 @@ pub struct Cluster { password: String, pooler_mode: PoolerMode, sharded_tables: ShardedTables, + sharded_schemas: ShardedSchemas, replication_sharding: Option, schema: Arc>, multi_tenant: Option, rw_strategy: ReadWriteStrategy, - rw_split: ReadWriteSplit, schema_admin: bool, stats: Arc>, cross_shard_disabled: bool, two_phase_commit: bool, two_phase_commit_auto: bool, + online: Arc, + rewrite: Rewrite, + prepared_statements: PreparedStatements, + dry_run: bool, + expanded_explain: bool, + pub_sub_channel_size: usize, + query_parser: QueryParserLevel, + connection_recovery: ConnectionRecovery, + query_parser_engine: QueryParserEngine, } /// Sharding configuration from the cluster. @@ -57,6 +74,12 @@ pub struct ShardingSchema { pub shards: usize, /// Sharded tables. pub tables: ShardedTables, + /// Scemas. + pub schemas: ShardedSchemas, + /// Rewrite config. + pub rewrite: Rewrite, + /// Query parser engine. + pub query_parser_engine: QueryParserEngine, } impl ShardingSchema { @@ -65,12 +88,29 @@ impl ShardingSchema { } } +#[derive(Debug)] pub struct ClusterShardConfig { pub primary: Option, pub replicas: Vec, } +impl ClusterShardConfig { + pub fn pooler_mode(&self) -> PoolerMode { + // One of these will exist. + + if let Some(ref primary) = self.primary { + return primary.config.pooler_mode; + } + + self.replicas + .first() + .map(|replica| replica.config.pooler_mode) + .unwrap_or_default() + } +} + /// Cluster creation config. +#[derive(Debug)] pub struct ClusterConfig<'a> { pub name: &'a str, pub shards: &'a [ClusterShardConfig], @@ -87,6 +127,16 @@ pub struct ClusterConfig<'a> { pub cross_shard_disabled: bool, pub two_pc: bool, pub two_pc_auto: bool, + pub sharded_schemas: ShardedSchemas, + pub rewrite: &'a Rewrite, + pub prepared_statements: &'a PreparedStatements, + pub dry_run: bool, + pub expanded_explain: bool, + pub pub_sub_channel_size: usize, + pub query_parser: QueryParserLevel, + pub query_parser_engine: QueryParserEngine, + pub connection_recovery: ConnectionRecovery, + pub lsn_check_interval: Duration, } impl<'a> ClusterConfig<'a> { @@ -96,13 +146,20 @@ impl<'a> ClusterConfig<'a> { shards: &'a [ClusterShardConfig], sharded_tables: ShardedTables, multi_tenant: &'a Option, + sharded_schemas: ShardedSchemas, + rewrite: &'a Rewrite, ) -> Self { + let pooler_mode = shards + .first() + .map(|shard| shard.pooler_mode()) + .unwrap_or(user.pooler_mode.unwrap_or(general.pooler_mode)); + Self { name: &user.database, password: user.password(), user: &user.name, replication_sharding: user.replication_sharding.clone(), - pooler_mode: user.pooler_mode.unwrap_or(general.pooler_mode), + pooler_mode, lb_strategy: general.load_balancing_strategy, shards, sharded_tables, @@ -117,6 +174,16 @@ impl<'a> ClusterConfig<'a> { two_pc_auto: user .two_phase_commit_auto .unwrap_or(general.two_phase_commit_auto.unwrap_or(false)), // Disable by default. + sharded_schemas, + rewrite, + prepared_statements: &general.prepared_statements, + dry_run: general.dry_run, + expanded_explain: general.expanded_explain, + pub_sub_channel_size: general.pub_sub_channel_size, + query_parser: general.query_parser, + query_parser_engine: general.query_parser_engine, + connection_recovery: general.connection_recovery, + lsn_check_interval: Duration::from_millis(general.lsn_check_interval), } } } @@ -140,33 +207,75 @@ impl Cluster { cross_shard_disabled, two_pc, two_pc_auto, + sharded_schemas, + rewrite, + prepared_statements, + dry_run, + expanded_explain, + pub_sub_channel_size, + query_parser, + connection_recovery, + lsn_check_interval, + query_parser_engine, } = config; + let identifier = Arc::new(DatabaseUser { + user: user.to_owned(), + database: name.to_owned(), + }); + Self { - identifier: Arc::new(DatabaseUser { - user: user.to_owned(), - database: name.to_owned(), - }), + identifier: identifier.clone(), shards: shards .iter() - .map(|config| Shard::new(&config.primary, &config.replicas, lb_strategy, rw_split)) + .enumerate() + .map(|(number, config)| { + Shard::new(ShardConfig { + number, + primary: &config.primary, + replicas: &config.replicas, + lb_strategy, + rw_split, + identifier: identifier.clone(), + lsn_check_interval, + }) + }) .collect(), password: password.to_owned(), pooler_mode, sharded_tables, + sharded_schemas, replication_sharding, schema: Arc::new(RwLock::new(Schema::default())), multi_tenant: multi_tenant.clone(), rw_strategy, - rw_split, schema_admin, stats: Arc::new(Mutex::new(MirrorStats::default())), cross_shard_disabled, two_phase_commit: two_pc && shards.len() > 1, two_phase_commit_auto: two_pc_auto && shards.len() > 1, + online: Arc::new(AtomicBool::new(false)), + rewrite: rewrite.clone(), + prepared_statements: *prepared_statements, + dry_run, + expanded_explain, + pub_sub_channel_size, + query_parser, + connection_recovery, + query_parser_engine, } } + /// Change config to work with logical replication streaming. + pub fn logical_stream(&self) -> Self { + let mut cluster = self.clone(); + // Disable rewrites, we are only sending valid statements. + cluster.rewrite.enabled = false; + cluster.rewrite.shard_key = RewriteMode::Ignore; + cluster.rewrite.split_inserts = RewriteMode::Ignore; + cluster + } + /// Get a connection to a primary of the given shard. pub async fn primary(&self, shard: usize, request: &Request) -> Result { let shard = self.shards.get(shard).ok_or(Error::NoShard(shard))?; @@ -190,34 +299,12 @@ impl Cluster { } /// Move connections from cluster to another, saving them. - pub(crate) fn move_conns_to(&self, other: &Cluster) { + pub(crate) fn move_conns_to(&self, other: &Cluster) -> Result<(), Error> { for (from, to) in self.shards.iter().zip(other.shards.iter()) { - from.move_conns_to(to); + from.move_conns_to(to)?; } - } - /// Create new identical cluster connection pool. - /// - /// This will allocate new server connections. Use when reloading configuration - /// and you expect to drop the current Cluster entirely. - pub fn duplicate(&self) -> Self { - Self { - identifier: self.identifier.clone(), - shards: self.shards.iter().map(|s| s.duplicate()).collect(), - password: self.password.clone(), - pooler_mode: self.pooler_mode, - sharded_tables: self.sharded_tables.clone(), - replication_sharding: self.replication_sharding.clone(), - schema: self.schema.clone(), - multi_tenant: self.multi_tenant.clone(), - rw_strategy: self.rw_strategy, - rw_split: self.rw_split, - schema_admin: self.schema_admin, - stats: Arc::new(Mutex::new(MirrorStats::default())), - cross_shard_disabled: self.cross_shard_disabled, - two_phase_commit: self.two_phase_commit, - two_phase_commit_auto: self.two_phase_commit_auto, - } + Ok(()) } /// Cancel a query executed by one of the shards. @@ -264,9 +351,33 @@ impl Cluster { self.sharded_tables.tables() } - /// Find sharded column position, if the table and columns match the configuration. - pub fn sharded_column(&self, table: &str, columns: &[&str]) -> Option { - self.sharded_tables.sharded_column(table, columns) + /// Get query rewrite config. + pub fn rewrite(&self) -> &Rewrite { + &self.rewrite + } + + pub fn query_parser(&self) -> QueryParserLevel { + self.query_parser + } + + pub fn prepared_statements(&self) -> &PreparedStatements { + &self.prepared_statements + } + + pub fn connection_recovery(&self) -> &ConnectionRecovery { + &self.connection_recovery + } + + pub fn dry_run(&self) -> bool { + self.dry_run + } + + pub fn expanded_explain(&self) -> bool { + self.expanded_explain + } + + pub fn pub_sub_enabled(&self) -> bool { + self.pub_sub_channel_size > 0 } /// A cluster is read_only if zero shards have a primary. @@ -311,6 +422,21 @@ impl Cluster { !(self.shards().len() == 1 && (self.read_only() || self.write_only())) } + /// Use the query parser. + pub fn use_query_parser(&self) -> bool { + match self.query_parser() { + QueryParserLevel::Off => return false, + QueryParserLevel::On => return true, + QueryParserLevel::Auto => { + self.multi_tenant().is_some() + || self.router_needed() + || self.dry_run() + || self.prepared_statements() == &PreparedStatements::Full + || self.pub_sub_enabled() + } + } + } + /// Multi-tenant config. pub fn multi_tenant(&self) -> &Option { &self.multi_tenant @@ -328,6 +454,9 @@ impl Cluster { ShardingSchema { shards: self.shards.len(), tables: self.sharded_tables.clone(), + schemas: self.sharded_schemas.clone(), + rewrite: self.rewrite.clone(), + query_parser_engine: self.query_parser_engine, } } @@ -388,6 +517,8 @@ impl Cluster { } }); } + + self.online.store(true, Ordering::Relaxed); } /// Shutdown the connection pools. @@ -395,6 +526,12 @@ impl Cluster { for shard in self.shards() { shard.shutdown(); } + + self.online.store(false, Ordering::Relaxed); + } + + pub(crate) fn online(&self) -> bool { + self.online.load(Ordering::Relaxed) } /// Execute a query on every primary in the cluster. @@ -413,11 +550,16 @@ impl Cluster { #[cfg(test)] mod test { - use std::sync::Arc; + use std::{sync::Arc, time::Duration}; + + use pgdog_config::{ConfigAndUsers, OmnishardedTable, ShardedSchema}; use crate::{ - backend::pool::{Address, Config, PoolConfig}, - backend::{Shard, ShardedTables}, + backend::{ + pool::{Address, Config, PoolConfig, ShardConfig}, + replication::ShardedSchemas, + Shard, ShardedTables, + }, config::{ DataType, Hasher, LoadBalancingStrategy, ReadWriteSplit, ReadWriteStrategy, ShardedTable, @@ -427,7 +569,35 @@ mod test { use super::{Cluster, DatabaseUser}; impl Cluster { - pub fn new_test() -> Self { + pub fn new_test(config: &ConfigAndUsers) -> Self { + let identifier = Arc::new(DatabaseUser { + user: "pgdog".into(), + database: "pgdog".into(), + }); + let primary = &Some(PoolConfig { + address: Address::new_test(), + config: Config::default(), + }); + let replicas = &[PoolConfig { + address: Address::new_test(), + config: Config::default(), + }]; + + let shards = (0..2) + .into_iter() + .map(|number| { + Shard::new(ShardConfig { + number, + primary, + replicas, + lb_strategy: LoadBalancingStrategy::Random, + rw_split: ReadWriteSplit::IncludePrimary, + identifier: identifier.clone(), + lsn_check_interval: Duration::MAX, + }) + }) + .collect::>(); + Cluster { sharded_tables: ShardedTables::new( vec![ShardedTable { @@ -440,46 +610,48 @@ mod test { centroids_path: None, centroid_probes: 1, hasher: Hasher::Postgres, - mapping: None, + ..Default::default() }], - vec!["sharded_omni".into()], + vec![ + OmnishardedTable { + name: "sharded_omni".into(), + sticky_routing: false, + }, + OmnishardedTable { + name: "sharded_omni_sticky".into(), + sticky_routing: true, + }, + ], ), - shards: vec![ - Shard::new( - &Some(PoolConfig { - address: Address::new_test(), - config: Config::default(), - }), - &[PoolConfig { - address: Address::new_test(), - config: Config::default(), - }], - LoadBalancingStrategy::Random, - ReadWriteSplit::default(), - ), - Shard::new( - &Some(PoolConfig { - address: Address::new_test(), - config: Config::default(), - }), - &[PoolConfig { - address: Address::new_test(), - config: Config::default(), - }], - LoadBalancingStrategy::Random, - ReadWriteSplit::default(), - ), - ], - identifier: Arc::new(DatabaseUser { - user: "pgdog".into(), - database: "pgdog".into(), - }), + sharded_schemas: ShardedSchemas::new(vec![ + ShardedSchema { + database: "pgdog".into(), + name: Some("shard_0".into()), + shard: 0, + ..Default::default() + }, + ShardedSchema { + database: "pgdog".into(), + name: Some("shard_1".into()), + shard: 1, + ..Default::default() + }, + ]), + shards, + identifier, + prepared_statements: config.config.general.prepared_statements, + dry_run: config.config.general.dry_run, + expanded_explain: config.config.general.expanded_explain, + query_parser: config.config.general.query_parser, + rewrite: config.config.rewrite.clone(), + two_phase_commit: config.config.general.two_phase_commit, + two_phase_commit_auto: config.config.general.two_phase_commit_auto.unwrap_or(false), ..Default::default() } } - pub fn new_test_single_shard() -> Cluster { - let mut cluster = Self::new_test(); + pub fn new_test_single_shard(config: &ConfigAndUsers) -> Cluster { + let mut cluster = Self::new_test(config); cluster.shards.pop(); cluster diff --git a/pgdog/src/backend/pool/config.rs b/pgdog/src/backend/pool/config.rs index e38462585..7287b4ad9 100644 --- a/pgdog/src/backend/pool/config.rs +++ b/pgdog/src/backend/pool/config.rs @@ -2,6 +2,7 @@ use std::time::Duration; +use pgdog_config::{pooling::ConnectionRecovery, Role}; use serde::{Deserialize, Serialize}; use crate::config::{Database, General, PoolerMode, User}; @@ -57,6 +58,18 @@ pub struct Config { pub read_only: bool, /// Maximum prepared statements per connection. pub prepared_statements_limit: usize, + /// Stats averaging period. + pub stats_period: Duration, + /// Recovery algo. + pub connection_recovery: ConnectionRecovery, + /// LSN check interval. + pub lsn_check_interval: Duration, + /// LSN check timeout. + pub lsn_check_timeout: Duration, + /// LSN check delay. + pub lsn_check_delay: Duration, + /// Automatic role detection enabled. + pub role_detection: bool, } impl Config { @@ -182,7 +195,13 @@ impl Config { .read_only .unwrap_or(user.read_only.unwrap_or_default()), prepared_statements_limit: general.prepared_statements_limit, + stats_period: Duration::from_millis(general.stats_period), bannable: !is_only_replica, + connection_recovery: general.connection_recovery, + lsn_check_interval: Duration::from_millis(general.lsn_check_interval), + lsn_check_timeout: Duration::from_millis(general.lsn_check_timeout), + lsn_check_delay: Duration::from_millis(general.lsn_check_delay), + role_detection: database.role == Role::Auto, ..Default::default() } } @@ -214,7 +233,61 @@ impl Default for Config { pooler_mode: PoolerMode::default(), read_only: false, prepared_statements_limit: usize::MAX, + stats_period: Duration::from_millis(15_000), dns_ttl: Duration::from_millis(60_000), + connection_recovery: ConnectionRecovery::Recover, + lsn_check_interval: Duration::from_millis(5_000), + lsn_check_timeout: Duration::from_millis(5_000), + lsn_check_delay: Duration::from_millis(5_000), + role_detection: false, } } } + +#[cfg(test)] +mod test { + use super::*; + + fn create_database(role: Role) -> Database { + Database { + name: "test".into(), + role, + host: "localhost".into(), + port: 5432, + ..Default::default() + } + } + + #[test] + fn test_role_auto_enables_role_detection() { + let general = General::default(); + let user = User::default(); + let database = create_database(Role::Auto); + + let config = Config::new(&general, &database, &user, false); + + assert!(config.role_detection); + } + + #[test] + fn test_role_primary_disables_role_detection() { + let general = General::default(); + let user = User::default(); + let database = create_database(Role::Primary); + + let config = Config::new(&general, &database, &user, false); + + assert!(!config.role_detection); + } + + #[test] + fn test_role_replica_disables_role_detection() { + let general = General::default(); + let user = User::default(); + let database = create_database(Role::Replica); + + let config = Config::new(&general, &database, &user, false); + + assert!(!config.role_detection); + } +} diff --git a/pgdog/src/backend/pool/connection/aggregate.rs b/pgdog/src/backend/pool/connection/aggregate.rs index 5fa0f8251..ecd6b1450 100644 --- a/pgdog/src/backend/pool/connection/aggregate.rs +++ b/pgdog/src/backend/pool/connection/aggregate.rs @@ -3,7 +3,10 @@ use std::collections::{HashMap, VecDeque}; use crate::{ - frontend::router::parser::{Aggregate, AggregateFunction, AggregateTarget, RewritePlan}, + frontend::router::parser::{ + rewrite::statement::aggregate::{AggregateRewritePlan, HelperKind}, + Aggregate, AggregateFunction, AggregateTarget, + }, net::{ messages::{ data_types::{Double, Float, Numeric}, @@ -12,6 +15,7 @@ use crate::{ Decoder, }, }; +use rust_decimal::prelude::{FromPrimitive, ToPrimitive}; use rust_decimal::Decimal; use super::Error; @@ -45,35 +49,54 @@ struct Accumulator<'a> { target: &'a AggregateTarget, datum: Datum, avg: Option, + variance: Option, } impl<'a> Accumulator<'a> { pub fn from_aggregate( aggregate: &'a Aggregate, - counts: &HashMap<(usize, bool), usize>, + helpers: &HashMap<(usize, bool), HelperColumns>, ) -> Vec { aggregate .targets() .iter() .map(|target| { + let helper = helpers + .get(&(target.expr_id(), target.is_distinct())) + .copied() + .unwrap_or_default(); + let mut accumulator = match target.function() { AggregateFunction::Count => Accumulator { target, datum: Datum::Bigint(0), avg: None, + variance: None, }, _ => Accumulator { target, datum: Datum::Null, avg: None, + variance: None, }, }; if matches!(target.function(), AggregateFunction::Avg) { - let count_column = counts - .get(&(target.expr_id(), target.is_distinct())) - .copied(); - accumulator.avg = Some(AvgState::new(count_column)); + accumulator.avg = Some(AvgState::new(helper.count)); + } + + if matches!( + target.function(), + AggregateFunction::StddevPop + | AggregateFunction::StddevSamp + | AggregateFunction::VarPop + | AggregateFunction::VarSamp + ) { + accumulator.variance = Some(VarianceState::new( + target.function().clone(), + helper, + target.is_distinct(), + )); } accumulator @@ -143,6 +166,20 @@ impl<'a> Accumulator<'a> { } } } + AggregateFunction::StddevPop + | AggregateFunction::StddevSamp + | AggregateFunction::VarPop + | AggregateFunction::VarSamp => { + if let Some(state) = self.variance.as_mut() { + if !state.supported { + return Ok(false); + } + + if !state.accumulate(row, decoder)? { + return Ok(false); + } + } + } } Ok(true) @@ -171,6 +208,16 @@ impl<'a> Accumulator<'a> { return Ok(false); } + if let Some(state) = self.variance.as_mut() { + match state.finalize()? { + Some(result) => { + self.datum = result; + return Ok(true); + } + None => return Ok(false), + } + } + Ok(true) } } @@ -194,14 +241,204 @@ impl AvgState { } } +#[derive(Debug, Default, Clone, Copy)] +struct HelperColumns { + count: Option, + sum: Option, + sumsq: Option, +} + +#[derive(Debug, Clone)] +struct UnsupportedAggregate { + function: String, + reason: String, +} + +#[derive(Debug)] +struct VarianceState { + function: AggregateFunction, + count_column: Option, + sum_column: Option, + sumsq_column: Option, + total_count: Datum, + total_sum: Datum, + total_sumsq: Datum, + supported: bool, +} + +impl VarianceState { + fn new(function: AggregateFunction, helper: HelperColumns, distinct: bool) -> Self { + let supported = + !distinct && helper.count.is_some() && helper.sum.is_some() && helper.sumsq.is_some(); + Self { + function, + count_column: helper.count, + sum_column: helper.sum, + sumsq_column: helper.sumsq, + total_count: Datum::Null, + total_sum: Datum::Null, + total_sumsq: Datum::Null, + supported, + } + } + + fn accumulate(&mut self, row: &DataRow, decoder: &Decoder) -> Result { + if !self.supported { + return Ok(false); + } + + let Some(count_column) = self.count_column else { + self.supported = false; + return Ok(false); + }; + + let count = row + .get_column(count_column, decoder)? + .ok_or(Error::DecoderRowError)?; + + if count.value.is_null() { + return Ok(true); + } + + let Some(count_i128) = datum_as_i128(&count.value) else { + self.supported = false; + return Ok(false); + }; + + if count_i128 == 0 { + return Ok(true); + } + + self.total_count = self.total_count.clone() + count.value.clone(); + + let Some(sum_column) = self.sum_column else { + self.supported = false; + return Ok(false); + }; + let sum = row + .get_column(sum_column, decoder)? + .ok_or(Error::DecoderRowError)?; + if sum.value.is_null() { + self.supported = false; + return Ok(false); + } + self.total_sum = self.total_sum.clone() + sum.value.clone(); + + let Some(sumsq_column) = self.sumsq_column else { + self.supported = false; + return Ok(false); + }; + let sumsq = row + .get_column(sumsq_column, decoder)? + .ok_or(Error::DecoderRowError)?; + if sumsq.value.is_null() { + self.supported = false; + return Ok(false); + } + self.total_sumsq = self.total_sumsq.clone() + sumsq.value.clone(); + + Ok(true) + } + + fn finalize(&mut self) -> Result, Error> { + if !self.supported { + return Ok(None); + } + + if self.total_sum.is_null() || self.total_sumsq.is_null() { + return Ok(Some(Datum::Null)); + } + + let Some(count) = datum_as_i128(&self.total_count) else { + return Ok(None); + }; + + let sample = matches!( + self.function, + AggregateFunction::StddevSamp | AggregateFunction::VarSamp + ); + + if count == 0 { + return Ok(Some(Datum::Null)); + } + + if sample && count <= 1 { + return Ok(Some(Datum::Null)); + } + + let result = match (&self.total_sum, &self.total_sumsq) { + (Datum::Double(sum), Datum::Double(sumsq)) => { + let Some(variance) = compute_variance_double(sum.0, sumsq.0, count, sample) else { + return Ok(None); + }; + match self.function { + AggregateFunction::StddevPop | AggregateFunction::StddevSamp => { + Some(Datum::Double(Double(variance.max(0.0).sqrt()))) + } + AggregateFunction::VarPop | AggregateFunction::VarSamp => { + Some(Datum::Double(Double(variance))) + } + _ => None, + } + } + (Datum::Float(sum), Datum::Float(sumsq)) => { + let Some(variance) = + compute_variance_double(sum.0 as f64, sumsq.0 as f64, count, sample) + else { + return Ok(None); + }; + match self.function { + AggregateFunction::StddevPop | AggregateFunction::StddevSamp => { + Some(Datum::Float(Float(variance.max(0.0).sqrt() as f32))) + } + AggregateFunction::VarPop | AggregateFunction::VarSamp => { + Some(Datum::Float(Float(variance as f32))) + } + _ => None, + } + } + (Datum::Numeric(sum), Datum::Numeric(sumsq)) => { + let Some(sum_dec) = sum.as_decimal() else { + return Ok(None); + }; + let Some(sumsq_dec) = sumsq.as_decimal() else { + return Ok(None); + }; + let sum_dec = sum_dec.to_owned(); + let sumsq_dec = sumsq_dec.to_owned(); + let Some(variance) = compute_variance_decimal(sum_dec, sumsq_dec, count, sample) + else { + return Ok(None); + }; + match self.function { + AggregateFunction::StddevPop | AggregateFunction::StddevSamp => { + let Some(stddev) = sqrt_decimal(variance) else { + return Ok(None); + }; + Some(Datum::Numeric(Numeric::from(stddev))) + } + AggregateFunction::VarPop | AggregateFunction::VarSamp => { + Some(Datum::Numeric(Numeric::from(variance))) + } + _ => None, + } + } + _ => None, + }; + + Ok(result) + } +} + #[derive(Debug)] pub(super) struct Aggregates<'a> { rows: &'a VecDeque, mappings: HashMap>>, decoder: &'a Decoder, aggregate: &'a Aggregate, - count_columns: HashMap<(usize, bool), usize>, - avg_supported: bool, + helper_columns: HashMap<(usize, bool), HelperColumns>, + merge_supported: bool, + unsupported: Option, } impl<'a> Aggregates<'a> { @@ -209,55 +446,116 @@ impl<'a> Aggregates<'a> { rows: &'a VecDeque, decoder: &'a Decoder, aggregate: &'a Aggregate, - plan: &RewritePlan, + plan: &AggregateRewritePlan, ) -> Self { - let mut count_columns = aggregate - .targets() - .iter() - .filter_map(|target| { - if matches!(target.function(), AggregateFunction::Count) { - Some(((target.expr_id(), target.is_distinct()), target.column())) - } else { - None + let mut helper_columns: HashMap<(usize, bool), HelperColumns> = HashMap::new(); + let mut unsupported: Option = None; + + for target in aggregate.targets() { + let key = (target.expr_id(), target.is_distinct()); + match target.function() { + AggregateFunction::Count => { + helper_columns.entry(key).or_default().count = Some(target.column()); } - }) - .collect::>(); + AggregateFunction::Sum => { + helper_columns.entry(key).or_default().sum = Some(target.column()); + } + _ => {} + } + } for helper in plan.helpers() { - count_columns + let Some(index) = decoder.rd().field_index(&helper.alias) else { + unsupported.get_or_insert(UnsupportedAggregate { + function: "aggregate".to_string(), + reason: format!("missing helper column alias '{}'", helper.alias), + }); + continue; + }; + + let entry = helper_columns .entry((helper.expr_id, helper.distinct)) - .or_insert(helper.helper_column); + .or_default(); + match helper.kind { + HelperKind::Count => entry.count = Some(index), + HelperKind::Sum => entry.sum = Some(index), + HelperKind::SumSquares => entry.sumsq = Some(index), + } } - let avg_supported = aggregate - .targets() - .iter() - .filter(|target| matches!(target.function(), AggregateFunction::Avg)) - .all(|target| count_columns.contains_key(&(target.expr_id(), target.is_distinct()))); + let merge_supported = aggregate.targets().iter().all(|target| { + let key = (target.expr_id(), target.is_distinct()); + match target.function() { + AggregateFunction::Avg => helper_columns + .get(&key) + .and_then(|columns| columns.count) + .is_some(), + AggregateFunction::StddevPop + | AggregateFunction::StddevSamp + | AggregateFunction::VarPop + | AggregateFunction::VarSamp => { + if target.is_distinct() { + unsupported.get_or_insert(UnsupportedAggregate { + function: target.function().as_str().to_string(), + reason: "DISTINCT is not supported".to_string(), + }); + return false; + } + helper_columns + .get(&key) + .map(|columns| { + columns.count.is_some() + && columns.sum.is_some() + && columns.sumsq.is_some() + }) + .unwrap_or_else(|| { + unsupported.get_or_insert(UnsupportedAggregate { + function: target.function().as_str().to_string(), + reason: "missing helper columns".to_string(), + }); + false + }) + } + _ => true, + } + }); Self { rows, decoder, mappings: HashMap::new(), aggregate, - count_columns, - avg_supported, + helper_columns, + merge_supported, + unsupported, } } pub(super) fn aggregate(mut self) -> Result, Error> { - if !self.avg_supported { + if !self.merge_supported { + if let Some(info) = self.unsupported { + return Err(Error::UnsupportedAggregation { + function: info.function, + reason: info.reason, + }); + } return Ok(self.rows.clone()); } for row in self.rows { let grouping = Grouping::new(row, self.aggregate.group_by(), self.decoder)?; let entry = self.mappings.entry(grouping).or_insert_with(|| { - Accumulator::from_aggregate(self.aggregate, &self.count_columns) + Accumulator::from_aggregate(self.aggregate, &self.helper_columns) }); for aggregate in entry { if !aggregate.accumulate(row, self.decoder)? { + if let Some(info) = self.unsupported.clone() { + return Err(Error::UnsupportedAggregation { + function: info.function, + reason: info.reason, + }); + } return Ok(self.rows.clone()); } } @@ -284,6 +582,12 @@ impl<'a> Aggregates<'a> { } for mut acc in accumulator { if !acc.finalize()? { + if let Some(info) = self.unsupported.clone() { + return Err(Error::UnsupportedAggregation { + function: info.function, + reason: info.reason, + }); + } return Ok(self.rows.clone()); } row.insert( @@ -312,7 +616,7 @@ fn multiply_for_average(value: &Datum, count: &Datum) -> Option { Some(Datum::Float(Float((float.0 as f64 * multiplier) as f32))) } Datum::Numeric(numeric) => { - let decimal = numeric.as_decimal()?; + let decimal = numeric.as_decimal()?.to_owned(); let product = decimal * Decimal::from_i128_with_scale(multiplier_i128, 0); Some(Datum::Numeric(Numeric::from(product))) } @@ -357,10 +661,58 @@ fn datum_as_i128(datum: &Datum) -> Option { } } +fn compute_variance_double(sum: f64, sumsq: f64, count: i128, sample: bool) -> Option { + let n = count as f64; + let numerator = sumsq - (sum * sum) / n; + let denominator = if sample { (count - 1) as f64 } else { n }; + if denominator == 0.0 { + return None; + } + let variance = numerator / denominator; + Some(if variance < 0.0 { 0.0 } else { variance }) +} + +fn compute_variance_decimal( + sum: Decimal, + sumsq: Decimal, + count: i128, + sample: bool, +) -> Option { + let n = Decimal::from_i128_with_scale(count, 0); + let denominator = if sample { + Decimal::from_i128_with_scale(count - 1, 0) + } else { + n + }; + + if denominator == Decimal::ZERO { + return None; + } + + let numerator = sumsq - (sum * sum) / n; + + let variance = numerator / denominator; + Some(if variance < Decimal::ZERO { + Decimal::ZERO + } else { + variance + }) +} + +fn sqrt_decimal(value: Decimal) -> Option { + if value < Decimal::ZERO { + return None; + } + let value_f64 = value.to_f64()?; + Decimal::from_f64(value_f64.sqrt()) +} + #[cfg(test)] mod test { use super::*; - use crate::frontend::router::parser::HelperMapping; + use crate::frontend::router::parser::rewrite::statement::aggregate::{ + HelperKind, HelperMapping, + }; use crate::net::{ messages::{Field, Format, RowDescription}, Decoder, @@ -407,7 +759,7 @@ mod test { .cloned() .unwrap(); let aggregate = match stmt.stmt.unwrap().node.unwrap() { - pg_query::NodeEnum::SelectStmt(stmt) => Aggregate::parse(&stmt).unwrap(), + pg_query::NodeEnum::SelectStmt(stmt) => Aggregate::parse(&stmt), _ => panic!("expected select stmt"), }; @@ -423,10 +775,9 @@ mod test { shard1.add(3_i64).add(18.0_f64); rows.push_back(shard1); - let plan = RewritePlan::new(); + let plan = AggregateRewritePlan::default(); let aggregates = Aggregates::new(&rows, &decoder, &aggregate, &plan); - assert!(aggregates.avg_supported); - assert_eq!(aggregates.count_columns.len(), 1); + assert!(aggregates.merge_supported); let mut result = aggregates.aggregate().unwrap(); assert_eq!(result.len(), 1); @@ -448,7 +799,7 @@ mod test { .cloned() .unwrap(); let aggregate = match stmt.stmt.unwrap().node.unwrap() { - pg_query::NodeEnum::SelectStmt(stmt) => Aggregate::parse(&stmt).unwrap(), + pg_query::NodeEnum::SelectStmt(stmt) => Aggregate::parse(&stmt), _ => panic!("expected select stmt"), }; @@ -463,9 +814,9 @@ mod test { shard1.add(18.0_f64); rows.push_back(shard1); - let plan = RewritePlan::new(); + let plan = AggregateRewritePlan::default(); let aggregates = Aggregates::new(&rows, &decoder, &aggregate, &plan); - assert!(!aggregates.avg_supported); + assert!(!aggregates.merge_supported); let result = aggregates.aggregate().unwrap(); assert_eq!(result.len(), 2); @@ -485,11 +836,14 @@ mod test { .cloned() .unwrap(); let aggregate = match stmt.stmt.unwrap().node.unwrap() { - pg_query::NodeEnum::SelectStmt(stmt) => Aggregate::parse(&stmt).unwrap(), + pg_query::NodeEnum::SelectStmt(stmt) => Aggregate::parse(&stmt), _ => panic!("expected select stmt"), }; - let rd = RowDescription::new(&[Field::double("avg"), Field::bigint("__pgdog_count")]); + let rd = RowDescription::new(&[ + Field::double("avg"), + Field::bigint("__pgdog_count_expr0_col0"), + ]); let decoder = Decoder::from(&rd); let mut rows = VecDeque::new(); @@ -500,13 +854,15 @@ mod test { shard1.add(20.0_f64).add(2_i64); rows.push_back(shard1); - let mut plan = RewritePlan::new(); + let mut plan = AggregateRewritePlan::default(); plan.add_drop_column(1); plan.add_helper(HelperMapping { - avg_column: 0, + target_column: 0, helper_column: 1, expr_id: 0, distinct: false, + kind: HelperKind::Count, + alias: "__pgdog_count_expr0_col0".into(), }); let mut result = Aggregates::new(&rows, &decoder, &aggregate, &plan) @@ -529,15 +885,15 @@ mod test { .cloned() .unwrap(); let aggregate = match stmt.stmt.unwrap().node.unwrap() { - pg_query::NodeEnum::SelectStmt(stmt) => Aggregate::parse(&stmt).unwrap(), + pg_query::NodeEnum::SelectStmt(stmt) => Aggregate::parse(&stmt), _ => panic!("expected select stmt"), }; let rd = RowDescription::new(&[ Field::double("avg_price"), Field::double("avg_discount"), - Field::bigint("__pgdog_count_2"), - Field::bigint("__pgdog_count_3"), + Field::bigint("__pgdog_count_expr0_col0"), + Field::bigint("__pgdog_count_expr1_col1"), ]); let decoder = Decoder::from(&rd); @@ -549,20 +905,24 @@ mod test { shard1.add(20.0_f64).add(4.0_f64).add(2_i64).add(2_i64); rows.push_back(shard1); - let mut plan = RewritePlan::new(); + let mut plan = AggregateRewritePlan::default(); plan.add_drop_column(2); plan.add_drop_column(3); plan.add_helper(HelperMapping { - avg_column: 0, + target_column: 0, helper_column: 2, expr_id: 0, distinct: false, + kind: HelperKind::Count, + alias: "__pgdog_count_expr0_col0".into(), }); plan.add_helper(HelperMapping { - avg_column: 1, + target_column: 1, helper_column: 3, expr_id: 1, distinct: false, + kind: HelperKind::Count, + alias: "__pgdog_count_expr1_col1".into(), }); let mut result = Aggregates::new(&rows, &decoder, &aggregate, &plan) @@ -578,6 +938,152 @@ mod test { assert!((avg_discount - 3.0).abs() < f64::EPSILON); } + #[test] + fn aggregate_stddev_samp_with_helpers() { + let stmt = pg_query::parse("SELECT STDDEV(price) FROM menu") + .unwrap() + .protobuf + .stmts + .first() + .cloned() + .unwrap(); + let aggregate = match stmt.stmt.unwrap().node.unwrap() { + pg_query::NodeEnum::SelectStmt(stmt) => Aggregate::parse(&stmt), + _ => panic!("expected select stmt"), + }; + + let rd = RowDescription::new(&[ + Field::double("stddev_price"), + Field::bigint("__pgdog_count_expr0_col0"), + Field::double("__pgdog_sum_expr0_col0"), + Field::double("__pgdog_sumsq_expr0_col0"), + ]); + let decoder = Decoder::from(&rd); + + let mut rows = VecDeque::new(); + let mut shard0 = DataRow::new(); + shard0 + .add(2.8284271247461903_f64) + .add(2_i64) + .add(24.0_f64) + .add(296.0_f64); + rows.push_back(shard0); + let mut shard1 = DataRow::new(); + shard1 + .add(2.8284271247461903_f64) + .add(2_i64) + .add(40.0_f64) + .add(808.0_f64); + rows.push_back(shard1); + + let mut plan = AggregateRewritePlan::default(); + plan.add_drop_column(1); + plan.add_drop_column(2); + plan.add_drop_column(3); + plan.add_helper(HelperMapping { + target_column: 0, + helper_column: 1, + expr_id: 0, + distinct: false, + kind: HelperKind::Count, + alias: "__pgdog_count_expr0_col0".into(), + }); + plan.add_helper(HelperMapping { + target_column: 0, + helper_column: 2, + expr_id: 0, + distinct: false, + kind: HelperKind::Sum, + alias: "__pgdog_sum_expr0_col0".into(), + }); + plan.add_helper(HelperMapping { + target_column: 0, + helper_column: 3, + expr_id: 0, + distinct: false, + kind: HelperKind::SumSquares, + alias: "__pgdog_sumsq_expr0_col0".into(), + }); + + let mut result = Aggregates::new(&rows, &decoder, &aggregate, &plan) + .aggregate() + .unwrap(); + + assert_eq!(result.len(), 1); + let row = result.pop_front().unwrap(); + let stddev = row.get::(0, Format::Text).unwrap().0; + assert!((stddev - 5.163977794943222).abs() < 1e-9); + } + + #[test] + fn aggregate_var_pop_with_helpers() { + let stmt = pg_query::parse("SELECT VAR_POP(price) FROM menu") + .unwrap() + .protobuf + .stmts + .first() + .cloned() + .unwrap(); + let aggregate = match stmt.stmt.unwrap().node.unwrap() { + pg_query::NodeEnum::SelectStmt(stmt) => Aggregate::parse(&stmt), + _ => panic!("expected select stmt"), + }; + + let rd = RowDescription::new(&[ + Field::double("var_price"), + Field::bigint("__pgdog_count_expr0_col0"), + Field::double("__pgdog_sum_expr0_col0"), + Field::double("__pgdog_sumsq_expr0_col0"), + ]); + let decoder = Decoder::from(&rd); + + let mut rows = VecDeque::new(); + let mut shard0 = DataRow::new(); + shard0.add(4.0_f64).add(2_i64).add(24.0_f64).add(296.0_f64); + rows.push_back(shard0); + let mut shard1 = DataRow::new(); + shard1.add(4.0_f64).add(2_i64).add(40.0_f64).add(808.0_f64); + rows.push_back(shard1); + + let mut plan = AggregateRewritePlan::default(); + plan.add_drop_column(1); + plan.add_drop_column(2); + plan.add_drop_column(3); + plan.add_helper(HelperMapping { + target_column: 0, + helper_column: 1, + expr_id: 0, + distinct: false, + kind: HelperKind::Count, + alias: "__pgdog_count_expr0_col0".into(), + }); + plan.add_helper(HelperMapping { + target_column: 0, + helper_column: 2, + expr_id: 0, + distinct: false, + kind: HelperKind::Sum, + alias: "__pgdog_sum_expr0_col0".into(), + }); + plan.add_helper(HelperMapping { + target_column: 0, + helper_column: 3, + expr_id: 0, + distinct: false, + kind: HelperKind::SumSquares, + alias: "__pgdog_sumsq_expr0_col0".into(), + }); + + let mut result = Aggregates::new(&rows, &decoder, &aggregate, &plan) + .aggregate() + .unwrap(); + + assert_eq!(result.len(), 1); + let row = result.pop_front().unwrap(); + let variance = row.get::(0, Format::Text).unwrap().0; + assert!((variance - 20.0).abs() < 1e-9); + } + #[test] fn aggregate_distinct_count_not_paired() { let stmt = pg_query::parse("SELECT COUNT(DISTINCT price), AVG(price) FROM menu") @@ -588,7 +1094,7 @@ mod test { .cloned() .unwrap(); let aggregate = match stmt.stmt.unwrap().node.unwrap() { - pg_query::NodeEnum::SelectStmt(stmt) => Aggregate::parse(&stmt).unwrap(), + pg_query::NodeEnum::SelectStmt(stmt) => Aggregate::parse(&stmt), _ => panic!("expected select stmt"), }; @@ -603,9 +1109,9 @@ mod test { shard1.add(3_i64).add(18.0_f64); rows.push_back(shard1); - let plan = RewritePlan::new(); + let plan = AggregateRewritePlan::default(); let aggregates = Aggregates::new(&rows, &decoder, &aggregate, &plan); - assert!(!aggregates.avg_supported); // no matching COUNT without DISTINCT + assert!(!aggregates.merge_supported); // no matching COUNT without DISTINCT let result = aggregates.aggregate().unwrap(); // No merge should happen; rows should remain per shard @@ -615,4 +1121,98 @@ mod test { assert_eq!(avg0, 12.0); assert_eq!(avg1, 18.0); } + + #[test] + fn aggregate_errors_when_helper_alias_missing() { + let stmt = pg_query::parse("SELECT AVG(price) FROM menu") + .unwrap() + .protobuf + .stmts + .first() + .cloned() + .unwrap(); + let aggregate = match stmt.stmt.unwrap().node.unwrap() { + pg_query::NodeEnum::SelectStmt(stmt) => Aggregate::parse(&stmt), + _ => panic!("expected select stmt"), + }; + + let rd = RowDescription::new(&[Field::double("avg")]); + let decoder = Decoder::from(&rd); + + let mut rows = VecDeque::new(); + let mut shard0 = DataRow::new(); + shard0.add(12.0_f64); + rows.push_back(shard0); + + let mut plan = AggregateRewritePlan::default(); + plan.add_helper(HelperMapping { + target_column: 0, + helper_column: 1, + expr_id: 0, + distinct: false, + kind: HelperKind::Count, + alias: "__pgdog_count_expr0_col0".into(), + }); + + let result = Aggregates::new(&rows, &decoder, &aggregate, &plan).aggregate(); + + match result { + Err(Error::UnsupportedAggregation { function, reason }) => { + assert_eq!(function, "aggregate"); + assert!(reason.contains("missing helper column alias")); + } + other => panic!("expected unsupported aggregation error, got {other:?}"), + } + } + + #[test] + fn aggregate_group_by_merges_rows() { + let stmt = pg_query::parse("SELECT price, SUM(quantity) FROM menu GROUP BY 1") + .unwrap() + .protobuf + .stmts + .first() + .cloned() + .unwrap(); + let aggregate = match stmt.stmt.unwrap().node.unwrap() { + pg_query::NodeEnum::SelectStmt(stmt) => Aggregate::parse(&stmt), + _ => panic!("expected select stmt"), + }; + + let rd = RowDescription::new(&[Field::double("price"), Field::bigint("sum")]); + let decoder = Decoder::from(&rd); + + let mut rows = VecDeque::new(); + let mut shard0 = DataRow::new(); + shard0.add(10.0_f64).add(5_i64); + rows.push_back(shard0); + let mut shard1 = DataRow::new(); + shard1.add(10.0_f64).add(7_i64); + rows.push_back(shard1); + let mut shard2 = DataRow::new(); + shard2.add(20.0_f64).add(4_i64); + rows.push_back(shard2); + + let mut result = Aggregates::new( + &rows, + &decoder, + &aggregate, + &AggregateRewritePlan::default(), + ) + .aggregate() + .unwrap(); + + assert_eq!(result.len(), 2); + let mut groups: Vec<(f64, i64)> = result + .drain(..) + .map(|row| { + let price = row.get::(0, Format::Text).unwrap().0; + let sum = row.get::(1, Format::Text).unwrap(); + (price, sum) + }) + .collect(); + groups.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap()); + assert_eq!(groups[0], (10.0, 12)); + assert_eq!(groups[1], (20.0, 4)); + } } diff --git a/pgdog/src/backend/pool/connection/binding.rs b/pgdog/src/backend/pool/connection/binding.rs index 8f351ea51..5e1168895 100644 --- a/pgdog/src/backend/pool/connection/binding.rs +++ b/pgdog/src/backend/pool/connection/binding.rs @@ -2,7 +2,7 @@ use crate::{ frontend::{client::query_engine::TwoPcPhase, ClientRequest}, - net::{parameter::Parameters, ProtocolMessage}, + net::{parameter::Parameters, BackendKeyData, ProtocolMessage, Query}, state::State, }; @@ -97,6 +97,7 @@ impl Binding { } let message = server.read().await?; + read = true; if let Some(message) = state.forward(message)? { return Ok(message); @@ -127,7 +128,7 @@ impl Binding { if let Some(server) = server { server.send(client_request).await } else { - Err(Error::NotConnected) + Err(Error::DirectToShardNotConnected) } } @@ -251,7 +252,10 @@ impl Binding { } /// Execute a query on all servers. - pub async fn execute(&mut self, query: &str) -> Result, Error> { + pub async fn execute( + &mut self, + query: impl Into + Clone, + ) -> Result, Error> { let mut result = vec![]; match self { Binding::Direct(Some(ref mut server)) => { @@ -259,7 +263,9 @@ impl Binding { } Binding::MultiShard(ref mut servers, _) => { - let futures = servers.iter_mut().map(|server| server.execute(query)); + let futures = servers + .iter_mut() + .map(|server| server.execute(query.clone())); let results = join_all(futures).await; for server_result in results { @@ -273,65 +279,73 @@ impl Binding { Ok(result) } - /// Execute two-phase commit transaction control statements. - pub async fn two_pc(&mut self, name: &str, phase: TwoPcPhase) -> Result<(), Error> { - match self { - Binding::MultiShard(ref mut servers, _) => { - let skip_missing = match phase { - TwoPcPhase::Phase1 => false, - TwoPcPhase::Phase2 | TwoPcPhase::Rollback => true, - }; + pub(crate) async fn two_pc_on_guards( + servers: &mut [Guard], + name: &str, + phase: TwoPcPhase, + ) -> Result<(), Error> { + let skip_missing = matches!(phase, TwoPcPhase::Phase2 | TwoPcPhase::Rollback); - // Build futures for all servers - let mut futures = Vec::new(); - for (shard, server) in servers.iter_mut().enumerate() { - // Each shard has its own transaction name. - // This is to make this work on sharded databases that use the same - // database underneath. - let shard_name = format!("{}_{}", name, shard); - - let query = match phase { - TwoPcPhase::Phase1 => format!("PREPARE TRANSACTION '{}'", shard_name), - TwoPcPhase::Phase2 => format!("COMMIT PREPARED '{}'", shard_name), - TwoPcPhase::Rollback => format!("ROLLBACK PREPARED '{}'", shard_name), - }; + let mut futures = Vec::new(); + for (shard, server) in servers.iter_mut().enumerate() { + let shard_name = format!("{}_{}", name, shard); - futures.push(server.execute(query)); - } + let query = match phase { + TwoPcPhase::Phase1 => format!("PREPARE TRANSACTION '{}'", shard_name), + TwoPcPhase::Phase2 => format!("COMMIT PREPARED '{}'", shard_name), + TwoPcPhase::Rollback => format!("ROLLBACK PREPARED '{}'", shard_name), + }; - // Execute all operations in parallel - let results = join_all(futures).await; + futures.push(server.execute(query)); + } - // Process results and handle errors - for (shard, result) in results.into_iter().enumerate() { - match result { - Err(Error::ExecutionError(err)) => { - // Undefined object, transaction doesn't exist. - if !(skip_missing && err.code == "42704") { - return Err(Error::ExecutionError(err)); - } - } - Err(err) => return Err(err), - Ok(_) => { - if phase == TwoPcPhase::Phase2 { - servers[shard].stats_mut().transaction_2pc(); - } - } + let results = join_all(futures).await; + + for (shard, result) in results.into_iter().enumerate() { + match result { + Err(Error::ExecutionError(err)) => { + if !(skip_missing && err.code == "42704") { + return Err(Error::ExecutionError(err)); } } + Err(err) => return Err(err), + Ok(_) => { + if phase == TwoPcPhase::Phase2 { + servers[shard].stats_mut().transaction_2pc(); + } + } + } + } - Ok(()) + Ok(()) + } + + /// Execute two-phase commit transaction control statements. + pub async fn two_pc(&mut self, name: &str, phase: TwoPcPhase) -> Result<(), Error> { + match self { + Binding::MultiShard(ref mut servers, _) => { + Self::two_pc_on_guards(servers, name, phase).await } _ => Err(Error::TwoPcMultiShardOnly), } } - pub async fn link_client(&mut self, params: &Parameters) -> Result { + /// Link client to server. + pub async fn link_client( + &mut self, + id: &BackendKeyData, + params: &Parameters, + transaction_start_stmt: Option<&str>, + ) -> Result { match self { - Binding::Direct(Some(ref mut server)) => server.link_client(params).await, + Binding::Direct(Some(ref mut server)) => { + server.link_client(id, params, transaction_start_stmt).await + } Binding::MultiShard(ref mut servers, _) => { - let futures = servers.iter_mut().map(|server| server.link_client(params)); + let futures = servers + .iter_mut() + .map(|server| server.link_client(id, params, transaction_start_stmt)); let results = join_all(futures).await; let mut max = 0; @@ -348,6 +362,17 @@ impl Binding { } } + /// Handle transaction end. + pub fn transaction_params_hook(&mut self, rollback: bool) { + match self { + Binding::Direct(Some(ref mut server)) => server.transaction_params_hook(rollback), + Binding::MultiShard(ref mut servers, _) => servers + .iter_mut() + .for_each(|server| server.transaction_params_hook(rollback)), + _ => (), + } + } + pub fn changed_params(&mut self) -> Parameters { match self { Binding::Direct(Some(ref mut server)) => server.changed_params().clone(), @@ -362,6 +387,18 @@ impl Binding { } } + /// Reset changed params on all servers, disabling parameter tracking + /// for this request. + pub fn reset_changed_params(&mut self) { + match self { + Binding::Direct(Some(ref mut server)) => server.reset_changed_params(), + Binding::MultiShard(ref mut servers, _) => servers + .iter_mut() + .for_each(|server| server.reset_changed_params()), + _ => (), + } + } + pub(super) fn dirty(&mut self) { match self { Binding::Direct(Some(ref mut server)) => server.mark_dirty(true), @@ -381,12 +418,41 @@ impl Binding { } } - pub fn copy_mode(&self) -> bool { + pub fn is_multishard(&self) -> bool { + match self { + Binding::MultiShard(ref servers, _) => !servers.is_empty(), + _ => false, + } + } + + pub fn is_direct(&self) -> bool { + matches!(self, Binding::Direct(Some(_))) + } + + pub fn in_copy_mode(&self) -> bool { match self { Binding::Admin(_) => false, - Binding::MultiShard(ref servers, _state) => servers.iter().all(|s| s.copy_mode()), - Binding::Direct(Some(ref server)) => server.copy_mode(), + Binding::MultiShard(ref servers, _state) => servers.iter().all(|s| s.in_copy_mode()), + Binding::Direct(Some(ref server)) => server.in_copy_mode(), _ => false, } } + + /// Number of connected shards. + pub fn shards(&self) -> Result { + Ok(match self { + Binding::Admin(_) => 1, + Binding::Direct(Some(_)) => 1, + Binding::MultiShard(ref servers, _) => { + if servers.is_empty() { + return Err(Error::MultiShardNotConnected); + } else { + servers.len() + } + } + _ => { + return Err(Error::NotConnected); + } + }) + } } diff --git a/pgdog/src/backend/pool/connection/binding_test.rs b/pgdog/src/backend/pool/connection/binding_test.rs index fa53ebf1e..b028b4872 100644 --- a/pgdog/src/backend/pool/connection/binding_test.rs +++ b/pgdog/src/backend/pool/connection/binding_test.rs @@ -4,13 +4,15 @@ mod tests { use crate::{ backend::{ - pool::Pool, - pool::{connection::binding::Binding, PoolConfig}, + pool::{connection::binding::Binding, Pool, PoolConfig}, server::test::test_server, }, frontend::{ client::query_engine::TwoPcPhase, - router::{parser::Shard, Route}, + router::{ + parser::{Shard, ShardWithPriority}, + Route, + }, }, }; @@ -46,7 +48,7 @@ mod tests { crate::backend::pool::Guard::new(pool3, server3, now), ]; - let route = Route::write(Shard::All); + let route = Route::write(ShardWithPriority::new_default_unset(Shard::All)); let multishard = MultiShard::new(3, &route); let mut binding = Binding::MultiShard(guards, Box::new(multishard)); diff --git a/pgdog/src/backend/pool/connection/buffer.rs b/pgdog/src/backend/pool/connection/buffer.rs index d5f017af7..db755c5e5 100644 --- a/pgdog/src/backend/pool/connection/buffer.rs +++ b/pgdog/src/backend/pool/connection/buffer.rs @@ -6,7 +6,10 @@ use std::{ }; use crate::{ - frontend::router::parser::{Aggregate, DistinctBy, DistinctColumn, OrderBy, RewritePlan}, + frontend::router::parser::{ + rewrite::statement::aggregate::AggregateRewritePlan, Aggregate, DistinctBy, DistinctColumn, + OrderBy, + }, net::{ messages::{DataRow, FromBytes, Message, Protocol, ToBytes, Vector}, Decoder, @@ -136,7 +139,7 @@ impl Buffer { &mut self, aggregate: &Aggregate, decoder: &Decoder, - plan: &RewritePlan, + plan: &AggregateRewritePlan, ) -> Result<(), super::Error> { let buffer: VecDeque = std::mem::take(&mut self.buffer); let mut rows = if aggregate.is_empty() { @@ -158,7 +161,7 @@ impl Buffer { Ok(()) } - fn drop_helper_columns(rows: &mut VecDeque, plan: &RewritePlan) { + fn drop_helper_columns(rows: &mut VecDeque, plan: &AggregateRewritePlan) { if plan.drop_columns().is_empty() { return; } @@ -274,7 +277,7 @@ mod test { buf.add(dr.message().unwrap()).unwrap(); } - buf.aggregate(&agg, &Decoder::from(&rd), &RewritePlan::new()) + buf.aggregate(&agg, &Decoder::from(&rd), &AggregateRewritePlan::default()) .unwrap(); buf.full(); @@ -301,7 +304,7 @@ mod test { } } - buf.aggregate(&agg, &Decoder::from(&rd), &RewritePlan::new()) + buf.aggregate(&agg, &Decoder::from(&rd), &AggregateRewritePlan::default()) .unwrap(); buf.full(); diff --git a/pgdog/src/backend/pool/connection/lazy.rs b/pgdog/src/backend/pool/connection/lazy.rs new file mode 100644 index 000000000..8c156fa30 --- /dev/null +++ b/pgdog/src/backend/pool/connection/lazy.rs @@ -0,0 +1,3 @@ +//! Lazy connection guard. +//! +//! Handles server synchronization and lazy connection creation. diff --git a/pgdog/src/backend/pool/connection/mirror/handler.rs b/pgdog/src/backend/pool/connection/mirror/handler.rs index 92c2e406c..6b53ed0cd 100644 --- a/pgdog/src/backend/pool/connection/mirror/handler.rs +++ b/pgdog/src/backend/pool/connection/mirror/handler.rs @@ -68,7 +68,7 @@ impl MirrorHandler { } MirrorHandlerState::Idle => { let roll = if self.exposure < 1.0 { - thread_rng().gen_range(0.0..1.0) + rng().random_range(0.0..1.0) } else { 0.99 }; diff --git a/pgdog/src/backend/pool/connection/mirror/mod.rs b/pgdog/src/backend/pool/connection/mirror/mod.rs index 949289cf2..a162bbe1a 100644 --- a/pgdog/src/backend/pool/connection/mirror/mod.rs +++ b/pgdog/src/backend/pool/connection/mirror/mod.rs @@ -2,7 +2,7 @@ use std::time::Duration; -use rand::{thread_rng, Rng}; +use rand::{rng, Rng}; use tokio::select; use tokio::time::{sleep, Instant}; use tokio::{spawn, sync::mpsc::*}; @@ -13,9 +13,8 @@ use crate::config::{config, ConfigAndUsers}; use crate::frontend::client::query_engine::{QueryEngine, QueryEngineContext}; use crate::frontend::client::timeouts::Timeouts; use crate::frontend::client::TransactionType; -use crate::frontend::comms::comms; -use crate::frontend::PreparedStatements; -use crate::net::{Parameter, Parameters, Stream}; +use crate::frontend::{ClientComms, PreparedStatements}; +use crate::net::{BackendKeyData, Parameter, Parameters, Stream}; use crate::frontend::ClientRequest; @@ -33,6 +32,8 @@ pub use request::*; /// to PgDog. #[derive(Debug)] pub struct Mirror { + /// Random identifier for this mirror connection. + pub id: BackendKeyData, /// Mirror's prepared statements. Should be similar /// to client's statements, if exposure is high. pub prepared_statements: PreparedStatements, @@ -51,10 +52,11 @@ pub struct Mirror { impl Mirror { fn new(params: &Parameters, config: &ConfigAndUsers) -> Self { Self { + id: BackendKeyData::new(), prepared_statements: PreparedStatements::new(), params: params.clone(), timeouts: Timeouts::from_config(&config.config.general), - stream: Stream::DevNull, + stream: Stream::dev_null(), transaction: None, cross_shard_disabled: config.config.general.cross_shard_disabled, } @@ -90,7 +92,17 @@ impl Mirror { ]); // Same query engine as the client, except with a potentially different database config. - let mut query_engine = QueryEngine::new(¶ms, &comms(), false, &None)?; + let mut query_engine = QueryEngine::new( + ¶ms, + &ClientComms::new(&BackendKeyData::new()), + false, + &None, + )?; + + // Mirror must read server responses to keep the connection synchronized, + // so disable test_mode which skips reading responses. + #[cfg(test)] + query_engine.set_test_mode(false); // Mirror traffic handler. let mut mirror = Self::new(¶ms, &config); @@ -210,7 +222,7 @@ mod test { #[tokio::test] async fn test_mirror() { config::load_test(); - let cluster = Cluster::new_test(); + let cluster = Cluster::new_test(&config()); cluster.launch(); let mut mirror = Mirror::spawn("pgdog", &cluster, None).unwrap(); let mut conn = cluster.primary(0, &Request::default()).await.unwrap(); @@ -260,7 +272,7 @@ mod test { #[tokio::test] async fn test_mirror_stats_tracking() { config::load_test(); - let cluster = Cluster::new_test(); + let cluster = Cluster::new_test(&config()); cluster.launch(); // Get initial stats diff --git a/pgdog/src/backend/pool/connection/mod.rs b/pgdog/src/backend/pool/connection/mod.rs index 20a36d4c8..7ea427941 100644 --- a/pgdog/src/backend/pool/connection/mod.rs +++ b/pgdog/src/backend/pool/connection/mod.rs @@ -8,7 +8,7 @@ use crate::{ admin::server::AdminServer, backend::{ databases::{self, databases}, - reload_notify, PubSubClient, + pool, reload_notify, PubSubClient, }, config::{config, PoolerMode, User}, frontend::{ @@ -100,10 +100,7 @@ impl Connection { debug!("detected configuration reload, reloading cluster"); // Wait to reload pools until they are ready. - if let Some(wait) = reload_notify::ready() { - wait.await; - } - self.reload()?; + self.safe_reload().await?; return self.try_conn(request, route).await; } Err(err) => { @@ -203,18 +200,23 @@ impl Connection { match &self.binding { Binding::Admin(_) => Ok(ParameterStatus::fake()), _ => { - // Try a replica. If not, try the primary. - if self.connect(request, &Route::read(Some(0))).await.is_err() { - self.connect(request, &Route::write(Some(0))).await?; - }; - let mut params = vec![]; - for param in self.server()?.params().iter() { - if let Some(value) = param.1.as_str() { - params.push(ParameterStatus::from((param.0.as_str(), value))); + // Get params from the first database that answers. + // Parameters are cached on the pool. + for shard in self.cluster()?.shards() { + for pool in shard.pools() { + if let Ok(params) = pool.params(request).await { + let mut result = vec![]; + for param in params.iter() { + if let Some(value) = param.1.as_str() { + result.push(ParameterStatus::from((param.0.as_str(), value))); + } + } + + return Ok(result); + } } } - self.disconnect(); - Ok(params) + Err(Error::Pool(pool::Error::AllReplicasDown)) } } } @@ -224,12 +226,14 @@ impl Connection { /// Only await this future inside a `select!`. One of the conditions /// suspends this loop indefinitely and expects another `select!` branch /// to cancel it. + /// pub(crate) async fn read(&mut self) -> Result { select! { notification = self.pub_sub.recv() => { Ok(notification.ok_or(Error::ProtocolOutOfSync)?.message()?) } + // This is cancel-safe. message = self.binding.read() => { message } @@ -289,7 +293,7 @@ impl Connection { router: &mut Router, streaming: bool, ) -> Result<(), Error> { - if client_request.copy() && !streaming { + if client_request.is_copy() && !streaming { let rows = router .copy_data(client_request) .map_err(|e| Error::Router(e.to_string()))?; @@ -307,8 +311,17 @@ impl Connection { Ok(()) } + /// Reload synchronized with partial config changes. + pub async fn safe_reload(&mut self) -> Result<(), Error> { + if let Some(wait) = reload_notify::ready() { + wait.await; + } + + self.reload() + } + /// Fetch the cluster from the global database store. - pub(crate) fn reload(&mut self) -> Result<(), Error> { + fn reload(&mut self) -> Result<(), Error> { match self.binding { Binding::Direct(_) | Binding::MultiShard(_, _) => { let user = (self.user.as_str(), self.database.as_str()); @@ -373,7 +386,7 @@ impl Connection { } /// Get connected servers addresses. - pub(crate) fn addr(&mut self) -> Result, Error> { + pub(crate) fn addr(&self) -> Result, Error> { Ok(match self.binding { Binding::Direct(Some(ref server)) => vec![server.addr()], Binding::MultiShard(ref servers, _) => servers.iter().map(|s| s.addr()).collect(), @@ -383,36 +396,23 @@ impl Connection { }) } - /// Get a connected server, if any. If multi-shard, get the first one. - #[inline] - fn server(&mut self) -> Result<&mut Guard, Error> { - Ok(match self.binding { - Binding::Direct(ref mut server) => server.as_mut().ok_or(Error::NotConnected)?, - Binding::MultiShard(ref mut servers, _) => { - servers.first_mut().ok_or(Error::NotConnected)? - } - _ => return Err(Error::NotConnected), - }) - } - /// Get cluster if any. #[inline] pub(crate) fn cluster(&self) -> Result<&Cluster, Error> { - self.cluster.as_ref().ok_or(Error::NotConnected) + self.cluster.as_ref().ok_or(Error::ClusterNotConnected) } - /// Transaction mode pooling. + /// Pooler is in session mode. #[inline] - pub(crate) fn transaction_mode(&self) -> bool { + pub(crate) fn session_mode(&self) -> bool { self.cluster() - .map(|c| c.pooler_mode() == PoolerMode::Transaction) + .map(|c| c.pooler_mode() == PoolerMode::Session) .unwrap_or(true) } - /// Pooler is in session mod #[inline] - pub(crate) fn session_mode(&self) -> bool { - !self.transaction_mode() + pub(crate) fn pooler_mode(&self) -> PoolerMode { + self.cluster().map(|c| c.pooler_mode()).unwrap_or_default() } /// This is an admin DB connection. diff --git a/pgdog/src/backend/pool/connection/multi_shard/error.rs b/pgdog/src/backend/pool/connection/multi_shard/error.rs index 60d9fca83..e42660b20 100644 --- a/pgdog/src/backend/pool/connection/multi_shard/error.rs +++ b/pgdog/src/backend/pool/connection/multi_shard/error.rs @@ -28,6 +28,9 @@ pub enum Error { #[error("net error: {0}")] Net(#[from] crate::net::Error), + + #[error("unsupported aggregation {function}: {reason}")] + UnsupportedAggregation { function: String, reason: String }, } impl From for Error { diff --git a/pgdog/src/backend/pool/connection/multi_shard/mod.rs b/pgdog/src/backend/pool/connection/multi_shard/mod.rs index 12266431c..36c71f8d9 100644 --- a/pgdog/src/backend/pool/connection/multi_shard/mod.rs +++ b/pgdog/src/backend/pool/connection/multi_shard/mod.rs @@ -39,6 +39,9 @@ struct Counters { bind_complete: usize, command_complete: Option, transaction_error: bool, + copy_done: usize, + copy_out: usize, + copy_data: usize, } /// Multi-shard state. @@ -99,7 +102,7 @@ impl MultiShard { self.counters.transaction_error = true; } - forward = if self.counters.ready_for_query % self.shards == 0 { + forward = if self.counters.ready_for_query.is_multiple_of(self.shards) { if self.counters.transaction_error { Some(ReadyForQuery::error().message()?) } else { @@ -124,7 +127,11 @@ impl MultiShard { }; self.counters.command_complete_count += 1; - if self.counters.command_complete_count % self.shards == 0 { + if self + .counters + .command_complete_count + .is_multiple_of(self.shards) + { self.buffer.full(); if !self.buffer.is_empty() { @@ -132,7 +139,7 @@ impl MultiShard { .aggregate( self.route.aggregate(), &self.decoder, - self.route.rewrite_plan(), + self.route.aggregate_rewrite_plan(), ) .map_err(Error::from)?; @@ -169,7 +176,7 @@ impl MultiShard { if self.counters.row_description == self.shards { // Only send it to the client once all shards sent it, // so we don't get early requests from clients. - let plan = self.route.rewrite_plan(); + let plan = self.route.aggregate_rewrite_plan(); if plan.drop_columns().is_empty() { forward = Some(message); } else { @@ -181,7 +188,11 @@ impl MultiShard { 'I' => { self.counters.empty_query_response += 1; - if self.counters.empty_query_response % self.shards == 0 { + if self + .counters + .empty_query_response + .is_multiple_of(self.shards) + { forward = Some(message); } } @@ -193,7 +204,9 @@ impl MultiShard { self.validator.validate_data_row(&data_row)?; } - if !self.should_buffer() && self.counters.row_description % self.shards == 0 { + if !self.should_buffer() + && self.counters.row_description.is_multiple_of(self.shards) + { forward = Some(message); } else { self.buffer.add(message).map_err(Error::from)?; @@ -202,28 +215,28 @@ impl MultiShard { 'G' => { self.counters.copy_in += 1; - if self.counters.copy_in % self.shards == 0 { + if self.counters.copy_in.is_multiple_of(self.shards) { forward = Some(message); } } 'n' => { self.counters.no_data += 1; - if self.counters.no_data % self.shards == 0 { + if self.counters.no_data.is_multiple_of(self.shards) { forward = Some(message); } } '1' => { self.counters.parse_complete += 1; - if self.counters.parse_complete % self.shards == 0 { + if self.counters.parse_complete.is_multiple_of(self.shards) { forward = Some(message); } } '3' => { self.counters.close_complete += 1; - if self.counters.close_complete % self.shards == 0 { + if self.counters.close_complete.is_multiple_of(self.shards) { forward = Some(message); } } @@ -231,14 +244,37 @@ impl MultiShard { '2' => { self.counters.bind_complete += 1; - if self.counters.bind_complete % self.shards == 0 { + if self.counters.bind_complete.is_multiple_of(self.shards) { + forward = Some(message); + } + } + + 'c' => { + self.counters.copy_done += 1; + if self.counters.copy_done.is_multiple_of(self.shards) { + forward = Some(message); + } + } + + 'd' => { + self.counters.copy_data += 1; + forward = Some(message); + } + + 'H' => { + self.counters.copy_out += 1; + if self.counters.copy_out.is_multiple_of(self.shards) { forward = Some(message); } } 't' => { self.counters.parameter_description += 1; - if self.counters.parameter_description % self.shards == 0 { + if self + .counters + .parameter_description + .is_multiple_of(self.shards) + { forward = Some(message); } } diff --git a/pgdog/src/backend/pool/connection/multi_shard/test.rs b/pgdog/src/backend/pool/connection/multi_shard/test.rs index 36cc9df6a..e96115313 100644 --- a/pgdog/src/backend/pool/connection/multi_shard/test.rs +++ b/pgdog/src/backend/pool/connection/multi_shard/test.rs @@ -1,4 +1,7 @@ -use crate::net::{DataRow, Field}; +use crate::{ + frontend::router::parser::{Shard, ShardWithPriority}, + net::{DataRow, Field}, +}; use super::*; @@ -59,7 +62,10 @@ fn test_inconsistent_data_rows() { #[test] fn test_rd_before_dr() { - let mut multi_shard = MultiShard::new(3, &Route::read(None)); + let mut multi_shard = MultiShard::new( + 3, + &Route::read(ShardWithPriority::new_default_unset(Shard::All)), + ); let rd = RowDescription::new(&[Field::bigint("id")]); let mut dr = DataRow::new(); dr.add(1i64); diff --git a/pgdog/src/backend/pool/error.rs b/pgdog/src/backend/pool/error.rs index c9b392175..c1fddccb8 100644 --- a/pgdog/src/backend/pool/error.rs +++ b/pgdog/src/backend/pool/error.rs @@ -1,6 +1,8 @@ //! Connection pool errors. use thiserror::Error; +use crate::net::BackendKeyData; + #[derive(Debug, Error, PartialEq, Clone, Copy)] pub enum Error { #[error("checkout timeout")] @@ -59,4 +61,16 @@ pub enum Error { #[error("pub/sub disabled")] PubSubDisabled, + + #[error("pool {0} has no health target")] + PoolNoHealthTarget(u64), + + #[error("pool is not healthy")] + PoolUnhealthy, + + #[error("checked in untracked connection: {0}")] + UntrackedConnCheckin(BackendKeyData), + + #[error("mapping missing: {0}")] + MappingMissing(usize), } diff --git a/pgdog/src/backend/pool/guard.rs b/pgdog/src/backend/pool/guard.rs index 24cb4b13a..1bd30b738 100644 --- a/pgdog/src/backend/pool/guard.rs +++ b/pgdog/src/backend/pool/guard.rs @@ -2,14 +2,14 @@ use std::ops::{Deref, DerefMut}; +use pgdog_config::pooling::ConnectionRecovery; use tokio::time::timeout; use tokio::{spawn, time::Instant}; use tracing::{debug, error}; -use crate::backend::Server; +use crate::backend::{Error, Server}; use crate::state::State; -use super::Error; use super::{cleanup::Cleanup, Pool}; /// Connection guard. @@ -59,24 +59,39 @@ impl Guard { let sync_prepared = server.sync_prepared(); let needs_drain = server.needs_drain(); let force_close = server.force_close(); + let needs_cleanup = rollback || reset || sync_prepared || needs_drain; server.reset_changed_params(); // No need to delay checkin unless we have to. - if (rollback || reset || sync_prepared || needs_drain) && !force_close { - let rollback_timeout = pool.inner().config.rollback_timeout(); + if needs_cleanup && !force_close { + let rollback_timeout = pool.inner().config.rollback_timeout; + let conn_recovery = pool.inner().config.connection_recovery; + let addr = self.pool.addr().clone(); + spawn(async move { - if timeout( + match timeout( rollback_timeout, - Self::cleanup_internal(&mut server, cleanup), + Self::cleanup_internal(&mut server, cleanup, conn_recovery), ) .await - .is_err() { - error!("rollback timeout [{}]", server.addr()); - }; + Ok(Ok(_)) => (), + Err(_) => { + error!("server cleanup timed out [{}]", server.addr()); + server.stats_mut().state(State::ForceClose); + } + Ok(Err(err)) => { + error!("server cleanup failed: {} [{}]", err, server.addr()); + if !server.error() { + server.stats_mut().state(State::ForceClose); + } + } + } - pool.checkin(server); + if let Err(err) = pool.checkin(server) { + error!("pool checkin error: {} [{}]", err, addr); + } }); } else { debug!( @@ -84,64 +99,74 @@ impl Guard { server.stats().state, server.addr(), ); - pool.checkin(server); + if let Err(err) = pool.checkin(server) { + error!("pool checkin error: {} [{}]", err, self.pool.addr()); + } } } } - async fn cleanup_internal(server: &mut Box, cleanup: Cleanup) -> Result<(), Error> { + async fn cleanup_internal( + server: &mut Box, + cleanup: Cleanup, + conn_recovery: ConnectionRecovery, + ) -> Result<(), Error> { let schema_changed = server.schema_changed(); let sync_prepared = server.sync_prepared(); let needs_drain = server.needs_drain(); if needs_drain { - // Receive whatever data the client left before disconnecting. - debug!( - "[cleanup] draining data from \"{}\" server [{}]", - server.stats().state, - server.addr() - ); - server.drain().await; + if conn_recovery.can_recover() { + // Receive whatever data the client left before disconnecting. + debug!( + "[cleanup] draining data from \"{}\" server [{}]", + server.stats().state, + server.addr() + ); + server.drain().await?; + } else { + server.stats_mut().state(State::ForceClose); + return Ok(()); + } } + let rollback = server.in_transaction(); // Rollback any unfinished transactions, // but only if the server is in sync (protocol-wise). if rollback { - debug!( - "[cleanup] rolling back server transaction, in \"{}\" state [{}]", - server.stats().state, - server.addr(), - ); - server.rollback().await; + if conn_recovery.can_rollback() { + debug!( + "[cleanup] rolling back server transaction, in \"{}\" state [{}]", + server.stats().state, + server.addr(), + ); + server.rollback().await?; + } else { + server.stats_mut().state(State::ForceClose); + return Ok(()); + } } if cleanup.needed() { debug!( - "[cleanup] {}, server in \"{}\" state [{}]", - cleanup, + "[cleanup] running {} cleanup queries, server in \"{}\" state [{}]", + cleanup.len(), server.stats().state, server.addr() ); - match server.execute_batch(cleanup.queries()).await { - Err(_) => { - error!("server reset error [{}]", server.addr()); - } - Ok(_) => { - if cleanup.is_deallocate() { - server.prepared_statements_mut().clear(); - } - server.cleaned(); - } - } + server.execute_batch(cleanup.queries()).await?; - match server.close_many(cleanup.close()).await { - Ok(_) => (), - Err(err) => { - server.stats_mut().state(State::Error); - error!("server close error: {} [{}]", err, server.addr()); - } + if cleanup.is_deallocate() { + server.prepared_statements_mut().clear(); } + server.cleaned(); + + debug!( + "[cleanup] closing {} prepared statements", + cleanup.close().len() + ); + server.close_many(cleanup.close()).await?; } if schema_changed { @@ -158,13 +183,7 @@ impl Guard { server.stats().state, server.addr() ); - if let Err(err) = server.sync_prepared_statements().await { - error!( - "prepared statements sync error: {:?} [{}]", - err, - server.addr() - ); - } + server.sync_prepared_statements().await?; } Ok(()) @@ -193,9 +212,19 @@ impl Drop for Guard { #[cfg(test)] mod test { + use std::time::Duration; + + use pgdog_config::pooling::ConnectionRecovery; + use tokio::time::{sleep, timeout, Instant}; + use crate::{ - backend::pool::{test::pool, Request}, - net::{Describe, Flush, Parse, Protocol, Query, Sync}, + backend::{ + pool::{ + cleanup::Cleanup, test::pool, Address, Config, Guard, Pool, PoolConfig, Request, + }, + server::test::test_server, + }, + net::{Describe, Flush, Parse, Protocol, ProtocolMessage, Query, Sync}, }; #[tokio::test] @@ -279,4 +308,477 @@ mod test { let guard = pool.get(&Request::default()).await.unwrap(); assert!(guard.prepared_statements().is_empty()); } + + #[tokio::test] + async fn test_rollback_timeout() { + crate::logger(); + + let config = Config { + max: 1, + min: 0, + rollback_timeout: Duration::from_millis(100), + ..Default::default() + }; + + let pool = Pool::new(&PoolConfig { + address: Address::new_test(), + config, + }); + pool.launch(); + + { + let mut guard = pool.get(&Request::default()).await.unwrap(); + + guard.execute("BEGIN").await.unwrap(); + assert!(guard.in_transaction()); + + guard + .send(&vec![Query::new("SELECT pg_sleep(1)").into()].into()) + .await + .unwrap(); + } + + sleep(Duration::from_millis(500)).await; + + { + let state = pool.lock(); + assert_eq!(state.errors, 0); + assert_eq!(state.idle(), 0); + assert_eq!(state.total(), 0); + assert_eq!(state.force_close, 1); + } + + // Will create new connection. + let mut server = pool.get(&Request::default()).await.unwrap(); + let one: Vec = server.fetch_all("SELECT 1").await.unwrap(); + assert_eq!(one[0], 1); + } + + #[tokio::test] + async fn test_cleanup_close_drain() { + crate::logger(); + + let mut server = Guard::new( + Pool::new_test(), + Box::new(test_server().await), + Instant::now(), + ); + server.prepared_statements_mut().set_capacity(1); + + for i in 0..5 { + server + .send( + &vec![ + ProtocolMessage::from(Parse::named(format!("test_{}", i), "SELECT 1")), + Flush.into(), + ] + .into(), + ) + .await + .unwrap(); + + let ok = server.read().await.unwrap(); + assert_eq!(ok.code(), '1'); + assert!(server.done()); + } + assert_eq!(server.prepared_statements().len(), 5); + server + .send(&vec![Query::new("SHOW prepared_statements").into()].into()) + .await + .unwrap(); + let mut guard = server; + let mut server = guard.server.take().unwrap(); + let cleanup = Cleanup::new(&guard, &mut server); + assert_eq!(cleanup.close().len(), 4); + assert!(server.needs_drain()); + + Guard::cleanup_internal(&mut server, cleanup, ConnectionRecovery::Recover) + .await + .unwrap(); + + assert!(server.done()); + assert!(!server.needs_drain()); + + let one: Vec = server.fetch_all("SELECT 1").await.unwrap(); + assert_eq!(one[0], 1); + } + + #[tokio::test] + async fn test_cancel_safety_partial_send() { + let mut server = test_server().await; + let select = (0..50_000_000).into_iter().map(|_| 'b').collect::(); + let select = Query::new(format!("SELECT '{}'", select)); + let res = timeout( + Duration::from_millis(1), + server.send(&vec![select.into()].into()), + ) + .await; + assert!(res.is_err()); + assert!(server.force_close()); + assert!(server.io_in_progress()) + } + + #[tokio::test] + async fn test_conn_recovery_recover_with_needs_drain() { + crate::logger(); + + let mut server = Guard::new( + Pool::new_test(), + Box::new(test_server().await), + Instant::now(), + ); + server.prepared_statements_mut().set_capacity(1); + + server + .send( + &vec![ + ProtocolMessage::from(Parse::named("test_0", "SELECT 1")), + Flush.into(), + ] + .into(), + ) + .await + .unwrap(); + + let ok = server.read().await.unwrap(); + assert_eq!(ok.code(), '1'); + assert!(server.done()); + + server + .send(&vec![Query::new("SHOW prepared_statements").into()].into()) + .await + .unwrap(); + + let mut guard = server; + let mut server = guard.server.take().unwrap(); + let cleanup = Cleanup::new(&guard, &mut server); + + assert!(server.needs_drain()); + + Guard::cleanup_internal(&mut server, cleanup, ConnectionRecovery::Recover) + .await + .unwrap(); + + assert!(server.done()); + assert!(!server.needs_drain()); + + let one: Vec = server.fetch_all("SELECT 1").await.unwrap(); + assert_eq!(one[0], 1); + } + + #[tokio::test] + async fn test_conn_recovery_rollback_only_with_needs_drain() { + crate::logger(); + + let mut server = Guard::new( + Pool::new_test(), + Box::new(test_server().await), + Instant::now(), + ); + server.prepared_statements_mut().set_capacity(1); + + server + .send( + &vec![ + ProtocolMessage::from(Parse::named("test_0", "SELECT 1")), + Flush.into(), + ] + .into(), + ) + .await + .unwrap(); + + let ok = server.read().await.unwrap(); + assert_eq!(ok.code(), '1'); + assert!(server.done()); + + server + .send(&vec![Query::new("SHOW prepared_statements").into()].into()) + .await + .unwrap(); + + let mut guard = server; + let mut server = guard.server.take().unwrap(); + let cleanup = Cleanup::new(&guard, &mut server); + + assert!(server.needs_drain()); + + Guard::cleanup_internal(&mut server, cleanup, ConnectionRecovery::RollbackOnly) + .await + .unwrap(); + + use crate::state::State; + assert_eq!(server.stats().state, State::ForceClose); + assert!(server.needs_drain()); + } + + #[tokio::test] + async fn test_conn_recovery_drop_with_needs_drain() { + crate::logger(); + + let mut server = Guard::new( + Pool::new_test(), + Box::new(test_server().await), + Instant::now(), + ); + server.prepared_statements_mut().set_capacity(1); + + server + .send( + &vec![ + ProtocolMessage::from(Parse::named("test_0", "SELECT 1")), + Flush.into(), + ] + .into(), + ) + .await + .unwrap(); + + let ok = server.read().await.unwrap(); + assert_eq!(ok.code(), '1'); + assert!(server.done()); + + server + .send(&vec![Query::new("SHOW prepared_statements").into()].into()) + .await + .unwrap(); + + let mut guard = server; + let mut server = guard.server.take().unwrap(); + let cleanup = Cleanup::new(&guard, &mut server); + + assert!(server.needs_drain()); + + Guard::cleanup_internal(&mut server, cleanup, ConnectionRecovery::Drop) + .await + .unwrap(); + + use crate::state::State; + assert_eq!(server.stats().state, State::ForceClose); + assert!(server.needs_drain()); + } + + #[tokio::test] + async fn test_conn_recovery_recover_with_rollback() { + crate::logger(); + + let mut server = Guard::new( + Pool::new_test(), + Box::new(test_server().await), + Instant::now(), + ); + + server + .send(&vec![Query::new("BEGIN").into()].into()) + .await + .unwrap(); + + loop { + let msg = server.read().await.unwrap(); + if msg.code() == 'Z' { + break; + } + } + + assert!(server.in_transaction()); + + let mut guard = server; + let mut server = guard.server.take().unwrap(); + let cleanup = Cleanup::new(&guard, &mut server); + + Guard::cleanup_internal(&mut server, cleanup, ConnectionRecovery::Recover) + .await + .unwrap(); + + assert!(!server.in_transaction()); + + let one: Vec = server.fetch_all("SELECT 1").await.unwrap(); + assert_eq!(one[0], 1); + } + + #[tokio::test] + async fn test_conn_recovery_rollback_only_with_rollback() { + crate::logger(); + + let mut server = Guard::new( + Pool::new_test(), + Box::new(test_server().await), + Instant::now(), + ); + + server + .send(&vec![Query::new("BEGIN").into()].into()) + .await + .unwrap(); + + loop { + let msg = server.read().await.unwrap(); + if msg.code() == 'Z' { + break; + } + } + + assert!(server.in_transaction()); + + let mut guard = server; + let mut server = guard.server.take().unwrap(); + let cleanup = Cleanup::new(&guard, &mut server); + + Guard::cleanup_internal(&mut server, cleanup, ConnectionRecovery::RollbackOnly) + .await + .unwrap(); + + assert!(!server.in_transaction()); + + let one: Vec = server.fetch_all("SELECT 1").await.unwrap(); + assert_eq!(one[0], 1); + } + + #[tokio::test] + async fn test_conn_recovery_drop_with_rollback() { + crate::logger(); + + let mut server = Guard::new( + Pool::new_test(), + Box::new(test_server().await), + Instant::now(), + ); + + server + .send(&vec![Query::new("BEGIN").into()].into()) + .await + .unwrap(); + + loop { + let msg = server.read().await.unwrap(); + if msg.code() == 'Z' { + break; + } + } + + assert!(server.in_transaction()); + + let mut guard = server; + let mut server = guard.server.take().unwrap(); + let cleanup = Cleanup::new(&guard, &mut server); + + Guard::cleanup_internal(&mut server, cleanup, ConnectionRecovery::Drop) + .await + .unwrap(); + + use crate::state::State; + assert_eq!(server.stats().state, State::ForceClose); + assert!(server.in_transaction()); + } + + #[tokio::test] + async fn test_conn_recovery_drop_with_needs_drain_and_rollback() { + crate::logger(); + + let mut server = Guard::new( + Pool::new_test(), + Box::new(test_server().await), + Instant::now(), + ); + server.prepared_statements_mut().set_capacity(1); + + server + .send( + &vec![ + ProtocolMessage::from(Parse::named("test_0", "SELECT 1")), + Flush.into(), + ] + .into(), + ) + .await + .unwrap(); + + let ok = server.read().await.unwrap(); + assert_eq!(ok.code(), '1'); + assert!(server.done()); + + server + .send(&vec![Query::new("BEGIN").into()].into()) + .await + .unwrap(); + + loop { + let msg = server.read().await.unwrap(); + if msg.code() == 'Z' { + break; + } + } + + assert!(server.in_transaction()); + + server + .send(&vec![Query::new("SHOW prepared_statements").into()].into()) + .await + .unwrap(); + + let mut guard = server; + let mut server = guard.server.take().unwrap(); + let cleanup = Cleanup::new(&guard, &mut server); + + assert!(server.needs_drain()); + assert!(server.in_transaction()); + + Guard::cleanup_internal(&mut server, cleanup, ConnectionRecovery::Drop) + .await + .unwrap(); + + use crate::state::State; + assert_eq!(server.stats().state, State::ForceClose); + assert!(server.needs_drain()); + assert!(server.in_transaction()); + } + + #[tokio::test] + async fn test_cleanup_syncs_prepared_statements() { + crate::logger(); + + let mut server = Guard::new( + Pool::new_test(), + Box::new(test_server().await), + Instant::now(), + ); + + assert!(!server.sync_prepared()); + + server + .send(&vec![Query::new("PREPARE test_stmt AS SELECT $1::bigint").into()].into()) + .await + .unwrap(); + + for c in ['C', 'Z'] { + let msg = server.read().await.unwrap(); + assert_eq!(msg.code(), c); + } + + assert!( + server.sync_prepared(), + "sync_prepared flag should be set after PREPARE command" + ); + + let mut guard = server; + let mut server = guard.server.take().unwrap(); + let cleanup = Cleanup::new(&guard, &mut server); + + Guard::cleanup_internal(&mut server, cleanup, ConnectionRecovery::Recover) + .await + .unwrap(); + + assert!( + !server.sync_prepared(), + "sync_prepared flag should be cleared after cleanup" + ); + + assert!( + server.prepared_statements_mut().contains("test_stmt"), + "Statement should be in local cache after sync" + ); + + let one: Vec = server.fetch_all("SELECT 1").await.unwrap(); + assert_eq!(one[0], 1); + } } diff --git a/pgdog/src/backend/pool/healthcheck.rs b/pgdog/src/backend/pool/healthcheck.rs index a9c45c7ec..03d2034d3 100644 --- a/pgdog/src/backend/pool/healthcheck.rs +++ b/pgdog/src/backend/pool/healthcheck.rs @@ -49,17 +49,17 @@ impl<'a> Healtcheck<'a> { /// Perform the healtcheck if it's required. pub async fn healthcheck(&mut self) -> Result<(), Error> { - let healtcheck_age = self.conn.healthcheck_age(self.now); + let health_check_age = self.conn.healthcheck_age(self.now); - if healtcheck_age < self.healthcheck_interval { + if health_check_age < self.healthcheck_interval { return Ok(()); } match timeout(self.healthcheck_timeout, self.conn.healthcheck(";")).await { Ok(Ok(())) => Ok(()), Ok(Err(err)) => { - error!("server error: {} [{}]", err, self.pool.addr()); - Err(Error::ServerError) + error!("health check server error: {} [{}]", err, self.pool.addr()); + Err(Error::HealthcheckError) } Err(_) => Err(Error::HealthcheckError), } diff --git a/pgdog/src/backend/pool/inner.rs b/pgdog/src/backend/pool/inner.rs index 0b48fcbe2..f85458b48 100644 --- a/pgdog/src/backend/pool/inner.rs +++ b/pgdog/src/backend/pool/inner.rs @@ -2,13 +2,16 @@ use std::cmp::max; use std::collections::VecDeque; +use std::fmt::Display; +use std::time::Duration; use crate::backend::{stats::Counts as BackendCounts, Server}; +use crate::backend::{ConnectReason, DisconnectReason}; use crate::net::messages::BackendKeyData; use tokio::time::Instant; -use super::{Ban, Config, Error, Mapping, Oids, Pool, Request, Stats, Taken, Waiter}; +use super::{Config, Error, Mapping, Oids, Pool, Request, Stats, Taken, Waiter}; /// Pool internals protected by a mutex. #[derive(Default)] @@ -22,8 +25,6 @@ pub(super) struct Inner { pub(super) config: Config, /// Number of clients waiting for a connection. pub(super) waiting: VecDeque, - /// Pool ban status. - pub(super) ban: Option, /// Pool is online and available to clients. pub(super) online: bool, /// Pool is paused. @@ -44,8 +45,10 @@ pub(super) struct Inner { /// The pool has been changed and connections should be returned /// to the new pool. moved: Option, + /// Unique pool identifier. id: u64, - pub(super) replica_lag: ReplicaLag, + /// Replica lag. + pub(super) replica_lag: Duration, } impl std::fmt::Debug for Inner { @@ -68,7 +71,6 @@ impl Inner { taken: Taken::default(), config, waiting: VecDeque::new(), - ban: None, online: false, paused: false, force_close: 0, @@ -79,7 +81,7 @@ impl Inner { oids: None, moved: None, id, - replica_lag: ReplicaLag::default(), + replica_lag: Duration::ZERO, } } /// Total number of connections managed by the pool. @@ -88,12 +90,24 @@ impl Inner { self.idle() + self.checked_out() } + /// The pool is full and will not + /// create any more connections. + #[inline] + pub(super) fn full(&self) -> bool { + self.total() >= self.max() + } + /// Number of idle connections in the pool. #[inline] pub(super) fn idle(&self) -> usize { self.idle_connections.len() } + #[cfg(test)] + pub(super) fn idle_conns(&self) -> &[Box] { + &self.idle_connections + } + /// Number of connections checked out of the pool /// by clients. #[inline] @@ -131,7 +145,7 @@ impl Inner { /// The pool should create more connections now. #[inline] - pub(super) fn should_create(&self) -> bool { + pub(super) fn should_create(&self) -> ShouldCreate { let below_min = self.total() < self.min(); let below_max = self.total() < self.max(); let maintain_min = below_min && below_max; @@ -139,26 +153,24 @@ impl Inner { below_max && !self.waiting.is_empty() && self.idle_connections.is_empty(); let maintenance_on = self.online && !self.paused; - !self.banned() && (client_needs || maintenance_on && maintain_min) - } - - /// Check if the pool ban should be removed. - #[inline] - pub(super) fn check_ban(&mut self, now: Instant) -> bool { - if self.ban.is_none() { - return false; - } + // Clients from banned pools won't be able to request connections + // unless it's a primary. + let reason = if client_needs { + ConnectReason::ClientWaiting + } else if maintenance_on && maintain_min { + ConnectReason::BelowMin + } else { + return ShouldCreate::No; + }; - let mut unbanned = false; - if let Some(ban) = self.ban.take() { - if !ban.expired(now) { - self.ban = Some(ban); - } else { - unbanned = true; - } + ShouldCreate::Yes { + reason, + min: self.min(), + max: self.max(), + idle: self.idle(), + taken: self.checked_out(), + waiting: self.waiting.len(), } - - unbanned } /// Close connections that have exceeded the max age. @@ -167,12 +179,15 @@ impl Inner { let max_age = self.config.max_age; let mut removed = 0; - self.idle_connections.retain(|c| { + self.idle_connections.retain_mut(|c| { let age = c.age(now); let keep = age < max_age; if !keep { removed += 1; } + if !keep { + c.disconnect_reason(DisconnectReason::Old); + } keep }); @@ -186,16 +201,22 @@ impl Inner { let (mut remove, mut removed) = (self.can_remove(), 0); let idle_timeout = self.config.idle_timeout; - self.idle_connections.retain(|c| { + self.idle_connections.retain_mut(|c| { let idle_for = c.idle_for(now); - if remove > 0 && idle_for >= idle_timeout { + let keep = if remove > 0 && idle_for >= idle_timeout { remove -= 1; removed += 1; false } else { true + }; + + if !keep { + c.disconnect_reason(DisconnectReason::Idle); } + + keep }); removed @@ -209,39 +230,43 @@ impl Inner { /// Take connection from the idle pool. #[inline(always)] - pub(super) fn take(&mut self, request: &Request) -> Option> { + pub(super) fn take(&mut self, request: &Request) -> Result>, Error> { if let Some(conn) = self.idle_connections.pop() { self.taken.take(&Mapping { client: request.id, server: *(conn.id()), - }); + })?; - Some(conn) + Ok(Some(conn)) } else { - None + Ok(None) } } /// Place connection back into the pool /// or give it to a waiting client. #[inline] - pub(super) fn put(&mut self, conn: Box, now: Instant) { + pub(super) fn put(&mut self, mut conn: Box, now: Instant) -> Result<(), Error> { // Try to give it to a client that's been waiting, if any. let id = *conn.id(); - if let Some(waiter) = self.waiting.pop_front() { - if let Err(conn) = waiter.tx.send(Ok(conn)) { - self.idle_connections.push(conn.unwrap()); + while let Some(waiter) = self.waiting.pop_front() { + if let Err(conn_ret) = waiter.tx.send(Ok(conn)) { + conn = conn_ret.unwrap(); // SAFETY: We sent Ok(conn), we'll get back Ok(conn) if channel is closed. } else { self.taken.take(&Mapping { server: id, client: waiter.request.id, - }); + })?; self.stats.counts.server_assignment_count += 1; self.stats.counts.wait_time += now.duration_since(waiter.request.created_at); + return Ok(()); } - } else { - self.idle_connections.push(conn); } + + // No waiters, put connection in idle list. + self.idle_connections.push(conn); + + Ok(()) } #[inline] @@ -252,36 +277,41 @@ impl Inner { /// Dump all idle connections. #[inline] pub(super) fn dump_idle(&mut self) { + for conn in &mut self.idle_connections { + conn.disconnect_reason(DisconnectReason::Offline); + } self.idle_connections.clear(); } /// Take all idle connections and tell active ones to /// be returned to a different pool instance. #[inline] - #[allow(clippy::vec_box)] // Server is a very large struct, reading it when moving between contains is expensive. + #[allow(clippy::vec_box)] // Server is a very large struct, reading it when moving between containers is expensive. pub(super) fn move_conns_to(&mut self, destination: &Pool) -> (Vec>, Taken) { self.moved = Some(destination.clone()); - let idle = std::mem::take(&mut self.idle_connections) - .into_iter() - .collect(); + let mut idle = std::mem::take(&mut self.idle_connections); let taken = std::mem::take(&mut self.taken); + for conn in idle.iter_mut() { + conn.stats_mut().pool_id = destination.id(); + } + (idle, taken) } - #[inline(always)] /// Check a connection back into the pool if it's ok to do so. /// Otherwise, drop the connection and close it. /// /// Return: true if the pool should be banned, false otherwise. + #[inline(always)] pub(super) fn maybe_check_in( &mut self, mut server: Box, now: Instant, stats: BackendCounts, - ) -> CheckInResult { + ) -> Result { let mut result = CheckInResult { - banned: false, + server_error: false, replenish: true, }; @@ -289,12 +319,14 @@ impl Inner { result.replenish = false; // Prevents deadlocks. if moved.id() != self.id { - moved.lock().maybe_check_in(server, now, stats); - return result; + server.stats_mut().pool_id = moved.id(); + server.stats_mut().update(); + moved.lock().maybe_check_in(server, now, stats)?; + return Ok(result); } } - self.taken.check_in(server.id()); + self.taken.check_in(server.id())?; // Update stats self.stats.counts = self.stats.counts + stats; @@ -302,31 +334,36 @@ impl Inner { // Ban the pool from serving more clients. if server.error() { self.errors += 1; - result.banned = self.maybe_ban(now, Error::ServerError); - return result; + result.server_error = true; + server.disconnect_reason(DisconnectReason::Error); + + return Ok(result); } // Pool is offline or paused, connection should be closed. if !self.online || self.paused { result.replenish = false; - return result; + return Ok(result); } // Close connections exceeding max age. if server.age(now) >= self.config.max_age { - return result; + server.disconnect_reason(DisconnectReason::Old); + return Ok(result); } // Force close the connection. if server.force_close() { self.force_close += 1; - return result; + server.disconnect_reason(DisconnectReason::ForceClose); + return Ok(result); } // Close connections in replication mode, // they are generally not re-usable. if server.replication_mode() { - return result; + server.disconnect_reason(DisconnectReason::ReplicationMode); + return Ok(result); } if server.re_synced() { @@ -337,14 +374,20 @@ impl Inner { // Finally, if the server is ok, // place the connection back into the idle list. if server.can_check_in() { - self.put(server, now); + self.put(server, now)?; + result.replenish = false; } else { self.out_of_sync += 1; + server.disconnect_reason(DisconnectReason::OutOfSync); } - result + Ok(result) } + /// Remove waiter from the queue. + /// + /// This happens if the waiter timed out, e.g. checkout timeout, + /// or the caller got cancelled. #[inline] pub(super) fn remove_waiter(&mut self, id: &BackendKeyData) { if let Some(waiter) = self.waiting.pop_front() { @@ -364,138 +407,62 @@ impl Inner { } } - /// Ban the pool from serving traffic if that's allowed per configuration. - #[inline] - pub fn maybe_ban(&mut self, now: Instant, reason: Error) -> bool { - if self.config.bannable || reason == Error::ManualBan { - let ban = Ban { - created_at: now, - reason, - ban_timeout: self.config.ban_timeout(), - }; - self.ban = Some(ban); - - // Tell every waiting client that this pool is busted. - self.close_waiters(Error::Banned); - - // Clear the idle connection pool. - self.idle_connections.clear(); - - true - } else { - false - } - } - #[inline] pub(super) fn close_waiters(&mut self, err: Error) { for waiter in self.waiting.drain(..) { let _ = waiter.tx.send(Err(err)); } } - - /// Remove the pool ban unless it' been manually banned. - #[inline(always)] - pub fn maybe_unban(&mut self) -> bool { - let mut unbanned = false; - if let Some(ban) = self.ban.take() { - if ban.reason == Error::ManualBan { - self.ban = Some(ban); - } else { - unbanned = true; - } - } - - unbanned - } - - pub fn unban(&mut self) -> bool { - self.ban.take().is_some() - } - - #[inline(always)] - pub fn banned(&self) -> bool { - self.ban.is_some() - } - - #[inline(always)] - #[allow(dead_code)] - pub fn manually_banned(&self) -> bool { - self.ban.map(|ban| ban.manual()).unwrap_or(false) - } } +/// Result of connection check into the pool. #[derive(Debug, Copy, Clone)] pub(super) struct CheckInResult { - pub(super) banned: bool, + pub(super) server_error: bool, pub(super) replenish: bool, } -// ------------------------------------------------------------------------------------------------- -// ----- ReplicaLag -------------------------------------------------------------------------------- - -#[derive(Clone, Copy, Debug)] -pub enum ReplicaLag { - NonApplicable, - Duration(std::time::Duration), - Bytes(u64), - Unknown, +#[derive(Debug, Copy, Clone, PartialEq)] +pub(super) enum ShouldCreate { + No, + Yes { + reason: ConnectReason, + min: usize, + max: usize, + idle: usize, + taken: usize, + waiting: usize, + }, } -impl ReplicaLag { - pub fn simple_display(&self) -> String { - match self { - Self::NonApplicable => "n/a".to_string(), - Self::Duration(d) => { - let total_secs = d.as_secs(); - let minutes = total_secs / 60; - let seconds = total_secs % 60; - - if minutes > 0 { - return if seconds > 0 { - format!("{}m{}s", minutes, seconds) - } else { - format!("{}m", minutes) - }; - } - - if total_secs > 0 { - return format!("{}s", total_secs); - } - - let millis = d.as_millis(); - if millis > 0 { - return format!("{}ms", millis); - } - - "<1ms".to_string() - } - Self::Bytes(b) => format!("{}B", b), - Self::Unknown => "unknown".to_string(), - } - } -} - -impl Default for ReplicaLag { - fn default() -> Self { - Self::Unknown +impl ShouldCreate { + pub(super) fn yes(&self) -> bool { + matches!(self, Self::Yes { .. }) } } -impl std::fmt::Display for ReplicaLag { +impl Display for ShouldCreate { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { - Self::NonApplicable => write!(f, "NonApplicable"), - Self::Duration(d) => write!(f, "Duration({:?})", d), - Self::Bytes(b) => write!(f, "Bytes({})", b), - Self::Unknown => write!(f, "Unknown"), + Self::No => write!(f, "no"), + Self::Yes { + reason, + min, + max, + idle, + taken, + waiting, + } => { + write!( + f, + "reason={}, min={}, max={}, idle={}, taken={}, waiting={}", + reason, min, max, idle, taken, waiting + ) + } } } } -// ------------------------------------------------------------------------------------------------- -// ------------------------------------------------------------------------------------------------- - #[cfg(test)] mod test { use std::time::Duration; @@ -507,150 +474,681 @@ mod test { use super::*; #[test] - fn test_invariants() { - let mut inner = Inner::default(); + fn test_default_state() { + let inner = Inner::default(); - // Defaults. - assert!(!inner.banned()); assert_eq!(inner.idle(), 0); + assert_eq!(inner.checked_out(), 0); + assert_eq!(inner.total(), 0); assert!(!inner.online); assert!(!inner.paused); + } - inner.idle_connections.push(Box::new(Server::default())); - inner.idle_connections.push(Box::new(Server::default())); - inner.idle_connections.push(Box::new(Server::default())); - assert_eq!(inner.idle(), 3); - - // The ban list. bans clear idle connections. - let banned = inner.maybe_ban(Instant::now(), Error::CheckoutTimeout); - assert!(banned); - assert_eq!(inner.idle(), 0); + #[test] + fn test_offline_pool_behavior() { + let mut inner = Inner::default(); - let unbanned = inner.check_ban(Instant::now() + Duration::from_secs(100)); - assert!(!unbanned); - assert!(inner.banned()); - let unbanned = inner.check_ban(Instant::now() + Duration::from_secs(301)); - assert!(unbanned); - assert!(!inner.banned()); - let unbanned = inner.maybe_unban(); - assert!(!unbanned); - assert!(!inner.banned()); - let banned = inner.maybe_ban(Instant::now(), Error::ManualBan); - assert!(banned); - assert!(!inner.maybe_unban()); - assert!(inner.banned()); - let banned = inner.maybe_ban(Instant::now(), Error::ServerError); - assert!(banned); - - // Testing check-in server. - let result = inner.maybe_check_in( - Box::new(Server::default()), - Instant::now(), - BackendCounts::default(), - ); - assert!(!result.banned); - assert_eq!(inner.idle(), 0); // pool offline + let server = Box::new(Server::default()); + let server_id = *server.id(); + inner + .taken + .take(&Mapping { + client: BackendKeyData::new(), + server: server_id, + }) + .unwrap(); + + let result = inner + .maybe_check_in(server, Instant::now(), BackendCounts::default()) + .unwrap(); + + assert!(!result.server_error); + assert_eq!(inner.idle(), 0); // pool offline, connection not added + assert_eq!(inner.total(), 0); + } + #[test] + fn test_paused_pool_behavior() { + let mut inner = Inner::default(); inner.online = true; inner.paused = true; - inner.maybe_check_in( - Box::new(Server::default()), - Instant::now(), - BackendCounts::default(), - ); - assert_eq!(inner.total(), 0); // pool paused; + + let server = Box::new(Server::default()); + let server_id = *server.id(); + inner + .taken + .take(&Mapping { + client: BackendKeyData::new(), + server: server_id, + }) + .unwrap(); + + inner + .maybe_check_in(server, Instant::now(), BackendCounts::default()) + .unwrap(); + + assert_eq!(inner.total(), 0); // pool paused, connection not added + } + + #[test] + fn test_online_pool_accepts_connections() { + let mut inner = Inner::default(); + inner.online = true; inner.paused = false; - assert!( - !inner - .maybe_check_in( - Box::new(Server::default()), - Instant::now(), - BackendCounts::default() - ) - .banned - ); - assert!(inner.idle() > 0); + + let server = Box::new(Server::default()); + let server_id = *server.id(); + inner + .taken + .take(&Mapping { + client: BackendKeyData::new(), + server: server_id, + }) + .unwrap(); + + let result = inner + .maybe_check_in(server, Instant::now(), BackendCounts::default()) + .unwrap(); + + assert!(!result.server_error); assert_eq!(inner.idle(), 1); + assert_eq!(inner.total(), 1); + } - let server = Box::new(Server::new_error()); + #[test] + fn test_server_error_handling() { + let mut inner = Inner::default(); + inner.online = true; - assert_eq!(inner.checked_out(), 0); - inner.taken.take(&Mapping { - client: BackendKeyData::new(), - server: *server.id(), - }); + let server = Box::new(Server::new_error()); + let server_id = *server.id(); + + // Simulate server being checked out + inner + .taken + .take(&Mapping { + client: BackendKeyData::new(), + server: server_id, + }) + .unwrap(); assert_eq!(inner.checked_out(), 1); - let result = inner.maybe_check_in(server, Instant::now(), BackendCounts::default()); - assert!(result.banned); - assert_eq!(inner.ban.unwrap().reason, Error::ServerError); - assert!(inner.taken.is_empty()); - inner.ban = None; + let result = inner + .maybe_check_in(server, Instant::now(), BackendCounts::default()) + .unwrap(); + assert!(result.server_error); + assert!(inner.taken.is_empty()); // Error server removed from taken + assert_eq!(inner.idle(), 0); // Error server not added to idle + } + + #[test] + fn test_should_create_with_waiting_clients() { + let mut inner = Inner::default(); + inner.online = true; inner.config.max = 5; + inner.config.min = 1; + inner.waiting.push_back(Waiter { request: Request::default(), tx: channel().0, }); - assert_eq!(inner.config.min, 1); + assert_eq!(inner.idle(), 0); - assert!(inner.should_create()); + assert!(matches!( + inner.should_create(), + ShouldCreate::Yes { + reason: ConnectReason::ClientWaiting, + min: 1, + max: 5, + idle: 0, + taken: 0, + waiting: 1, + } + )); + } + #[test] + fn test_should_create_below_minimum() { + let mut inner = Inner::default(); + inner.online = true; inner.config.min = 2; - assert_eq!(inner.config.max, 5); + inner.config.max = 5; + assert!(inner.total() < inner.min()); assert!(inner.total() < inner.max()); - assert!(!inner.banned() && inner.online); - assert!(inner.should_create()); - - inner.config.max = 1; - assert!(inner.should_create()); + assert!(matches!( + inner.should_create(), + ShouldCreate::Yes { + reason: ConnectReason::BelowMin, + min: 2, + max: 5, + idle: 0, + taken: 0, + waiting: 0, + } + )); + } + #[test] + fn test_should_not_create_at_max() { + let mut inner = Inner::default(); + inner.online = true; inner.config.max = 3; - assert!(inner.should_create()); + assert!(!inner.full()); + // Add 2 idle connections and 1 checked out connection to reach max + inner.idle_connections.push(Box::new(Server::default())); + inner.idle_connections.push(Box::new(Server::default())); + inner + .taken + .take(&Mapping { + client: BackendKeyData::new(), + server: BackendKeyData::new(), + }) + .unwrap(); + + assert_eq!(inner.idle(), 2); + assert_eq!(inner.checked_out(), 1); + assert_eq!(inner.total(), inner.config.max); + assert!(inner.full()); + assert_eq!(inner.should_create(), ShouldCreate::No); + } + + #[test] + fn test_close_idle_respects_minimum() { + let mut inner = Inner::default(); + inner.config.min = 2; + inner.config.max = 3; + inner.config.idle_timeout = Duration::from_millis(5_000); + + // Add connections to max inner.idle_connections.push(Box::new(Server::default())); inner.idle_connections.push(Box::new(Server::default())); inner.idle_connections.push(Box::new(Server::default())); - assert!(!inner.should_create()); - // Close idle connections. - inner.config.idle_timeout = Duration::from_millis(5_000); // 5 seconds. + // Close idle connections - shouldn't close any initially inner.close_idle(Instant::now()); - assert_eq!(inner.idle(), inner.config.max); // Didn't close any. + assert_eq!(inner.idle(), inner.config.max); + + // Close after timeout - should respect minimum for _ in 0..10 { inner.close_idle(Instant::now() + Duration::from_secs(6)); } assert_eq!(inner.idle(), inner.config.min); + + // Further closing should still respect minimum inner.config.min = 1; inner.close_idle(Instant::now() + Duration::from_secs(6)); assert_eq!(inner.idle(), inner.config.min); + } - // Close old connections. + #[test] + fn test_close_old_ignores_minimum() { + let mut inner = Inner::default(); + inner.online = true; + inner.config.min = 1; inner.config.max_age = Duration::from_millis(60_000); + + // Add a connection + let server = Box::new(Server::default()); + let server_id = *server.id(); + inner + .taken + .take(&Mapping { + client: BackendKeyData::new(), + server: server_id, + }) + .unwrap(); + + inner + .maybe_check_in(server, Instant::now(), BackendCounts::default()) + .unwrap(); + assert_eq!(inner.idle(), 1); + + // Close old connections before max age - should keep connection inner.close_old(Instant::now() + Duration::from_secs(59)); assert_eq!(inner.idle(), 1); + + // Close old connections after max age - ignores minimum inner.close_old(Instant::now() + Duration::from_secs(61)); - assert_eq!(inner.idle(), 0); // This ignores the min setting! + assert_eq!(inner.idle(), 0); + } - assert!(inner.should_create()); + #[test] + fn test_connection_lifecycle() { + let mut inner = Inner::default(); assert_eq!(inner.total(), 0); - inner.taken.take(&Mapping::default()); + + // Simulate taking a connection + inner.taken.take(&Mapping::default()).unwrap(); assert_eq!(inner.total(), 1); + assert_eq!(inner.checked_out(), 1); + + // Clear taken connections inner.taken.clear(); assert_eq!(inner.total(), 0); + assert_eq!(inner.checked_out(), 0); + } + + #[test] + fn test_max_age_enforcement_on_checkin() { + let mut inner = Inner::default(); + inner.online = true; + inner.config.max_age = Duration::from_millis(60_000); let server = Box::new(Server::default()); - let result = inner.maybe_check_in( - server, - Instant::now() + Duration::from_secs(61), - BackendCounts::default(), - ); - - assert!(!result.banned); - // Not checked in because of max age. - assert_eq!(inner.total(), 0); + let server_id = *server.id(); + inner + .taken + .take(&Mapping { + client: BackendKeyData::new(), + server: server_id, + }) + .unwrap(); + + inner + .maybe_check_in( + server, + Instant::now() + Duration::from_secs(61), // Exceeds max age + BackendCounts::default(), + ) + .unwrap(); + + assert_eq!(inner.total(), 0); // Connection not added due to max age + } + + #[test] + fn test_peer_lookup() { + let mut inner = Inner::default(); + let client_id = BackendKeyData::new(); + let server_id = BackendKeyData::new(); + + assert_eq!(inner.peer(&client_id), None); + + inner + .taken + .take(&Mapping { + client: client_id, + server: server_id, + }) + .unwrap(); + + assert_eq!(inner.peer(&client_id), Some(server_id)); + } + + #[test] + fn test_taken_server_returns_server_when_mapped() { + let mut taken = Taken::default(); + let client_id = BackendKeyData::new(); + let server_id = BackendKeyData::new(); + + // No mapping yet + assert_eq!(taken.server(&client_id), None); + + // Add mapping + taken + .take(&Mapping { + client: client_id, + server: server_id, + }) + .unwrap(); + + // Server should be returned for mapped client + assert_eq!(taken.server(&client_id), Some(server_id)); + + // Different client should return None + let other_client = BackendKeyData::new(); + assert_eq!(taken.server(&other_client), None); + } + + #[test] + fn test_can_remove() { + let mut inner = Inner::default(); + inner.config.min = 2; + inner.config.max = 5; + + assert_eq!(inner.can_remove(), 0); // total=0, min=2 + + inner.idle_connections.push(Box::new(Server::default())); + assert_eq!(inner.can_remove(), 0); // total=1, min=2 + + inner.idle_connections.push(Box::new(Server::default())); + assert_eq!(inner.can_remove(), 0); // total=2, min=2 + + inner.idle_connections.push(Box::new(Server::default())); + assert_eq!(inner.can_remove(), 1); // total=3, min=2 + } + + #[test] + fn test_take_connection() { + let mut inner = Inner::default(); + let request = Request::default(); + + assert!(inner.take(&request).unwrap().is_none()); + + inner.idle_connections.push(Box::new(Server::default())); + let server = inner.take(&request); + assert!(server.unwrap().is_some()); + assert_eq!(inner.idle(), 0); + assert_eq!(inner.checked_out(), 1); + } + + #[test] + fn test_put_connection_with_waiter() { + let mut inner = Inner::default(); + let (tx, mut rx) = channel(); + let waiter_request = Request::default(); + + inner.waiting.push_back(Waiter { + request: waiter_request, + tx, + }); + + let server = Box::new(Server::default()); + inner.put(server, Instant::now()).unwrap(); + + assert_eq!(inner.idle(), 0); // Connection given to waiter, not idle + assert_eq!(inner.checked_out(), 1); // Connection now checked out to waiter + assert!(inner.waiting.is_empty()); // Waiter was served + + // Verify waiter received the connection + assert!(rx.try_recv().is_ok()); + } + + #[test] + fn test_put_connection_no_waiters() { + let mut inner = Inner::default(); + let server = Box::new(Server::default()); + + inner.put(server, Instant::now()).unwrap(); + + assert_eq!(inner.idle(), 1); // Connection added to idle pool + assert_eq!(inner.checked_out(), 0); + assert!(inner.waiting.is_empty()); + } + + #[test] + fn test_dump_idle() { + let mut inner = Inner::default(); + inner.idle_connections.push(Box::new(Server::default())); + inner.idle_connections.push(Box::new(Server::default())); + + assert_eq!(inner.idle(), 2); + inner.dump_idle(); + assert_eq!(inner.idle(), 0); + } + + #[test] + fn test_remove_waiter() { + let mut inner = Inner::default(); + let (tx1, _) = channel(); + let (tx2, _) = channel(); + let (tx3, _) = channel(); + + let req1 = Request::default(); + let req2 = Request::default(); + let req3 = Request::default(); + let target_id = req2.id; + + inner.waiting.push_back(Waiter { + request: req1, + tx: tx1, + }); + inner.waiting.push_back(Waiter { + request: req2, + tx: tx2, + }); + inner.waiting.push_back(Waiter { + request: req3, + tx: tx3, + }); + + assert_eq!(inner.waiting.len(), 3); + inner.remove_waiter(&target_id); + assert_eq!(inner.waiting.len(), 2); + } + + #[test] + fn test_close_waiters() { + let mut inner = Inner::default(); + let (tx1, mut rx1) = channel(); + let (tx2, mut rx2) = channel(); + + inner.waiting.push_back(Waiter { + request: Request::default(), + tx: tx1, + }); + inner.waiting.push_back(Waiter { + request: Request::default(), + tx: tx2, + }); + + assert_eq!(inner.waiting.len(), 2); + inner.close_waiters(Error::CheckoutTimeout); + assert_eq!(inner.waiting.len(), 0); + + // Verify waiters received the correct error + assert_eq!(rx1.try_recv().unwrap().unwrap_err(), Error::CheckoutTimeout); + assert_eq!(rx2.try_recv().unwrap().unwrap_err(), Error::CheckoutTimeout); + } + + #[test] + fn test_should_create_for_waiting_clients_even_above_minimum() { + let mut inner = Inner::default(); + inner.online = true; + inner.config.min = 1; + inner.config.max = 5; + + // Add connections above minimum but all are checked out (no idle) + inner + .taken + .take(&Mapping { + client: BackendKeyData::new(), + server: BackendKeyData::new(), + }) + .unwrap(); + inner + .taken + .take(&Mapping { + client: BackendKeyData::new(), + server: BackendKeyData::new(), + }) + .unwrap(); + + // Add a waiting client + inner.waiting.push_back(Waiter { + request: Request::default(), + tx: channel().0, + }); + + assert!(inner.total() > inner.min()); // Above minimum + assert!(inner.total() < inner.max()); // Below maximum + assert_eq!(inner.idle(), 0); // No idle connections + assert!(!inner.waiting.is_empty()); // Has waiting clients + assert!(matches!( + inner.should_create(), + ShouldCreate::Yes { + reason: ConnectReason::ClientWaiting, + min: 1, + max: 5, + idle: 0, + taken: 2, + waiting: 1, + } + )); + } + + #[test] + fn test_should_not_create_offline() { + let mut inner = Inner::default(); + inner.online = false; + inner.config.min = 2; + + assert!(inner.total() < inner.min()); + assert_eq!(inner.should_create(), ShouldCreate::No); + } + + #[test] + fn test_set_taken() { + let mut inner = Inner::default(); + let mapping = Mapping { + client: BackendKeyData::new(), + server: BackendKeyData::new(), + }; + + assert_eq!(inner.checked_out(), 0); + + let mut taken = Taken::default(); + taken.take(&mapping).unwrap(); + + inner.set_taken(taken); + assert_eq!(inner.checked_out(), 1); + } + + #[test] + fn test_put_connection_skips_dropped_waiters() { + let mut inner = Inner::default(); + let (tx1, _rx1) = channel(); // Will be dropped + let (tx2, _rx2) = channel(); // Will be dropped + let (tx3, mut rx3) = channel(); // Will remain active + + let req1 = Request::default(); + let req2 = Request::default(); + let req3 = Request::default(); + + // Add three waiters to the queue + inner.waiting.push_back(Waiter { + request: req1, + tx: tx1, + }); + inner.waiting.push_back(Waiter { + request: req2, + tx: tx2, + }); + inner.waiting.push_back(Waiter { + request: req3, + tx: tx3, + }); + + // Drop the first two receivers to simulate cancelled waiters + drop(_rx1); + drop(_rx2); + + assert_eq!(inner.waiting.len(), 3); + + let server = Box::new(Server::default()); + inner.put(server, Instant::now()).unwrap(); + + // All waiters should be removed from queue since we tried each one + assert_eq!(inner.waiting.len(), 0); + // Connection should be given to the third waiter (the only one still listening) + assert_eq!(inner.checked_out(), 1); + assert_eq!(inner.idle(), 0); + + // Verify the third waiter received the connection + assert!(rx3.try_recv().is_ok()); + } + + #[test] + fn test_put_connection_all_waiters_dropped() { + let mut inner = Inner::default(); + let (tx1, _rx1) = channel(); + let (tx2, _rx2) = channel(); + + let req1 = Request::default(); + let req2 = Request::default(); + + inner.waiting.push_back(Waiter { + request: req1, + tx: tx1, + }); + inner.waiting.push_back(Waiter { + request: req2, + tx: tx2, + }); + + // Drop all receivers + drop(_rx1); + drop(_rx2); + + assert_eq!(inner.waiting.len(), 2); + + let server = Box::new(Server::default()); + inner.put(server, Instant::now()).unwrap(); + + // All waiters should be removed since they were all dropped + assert_eq!(inner.waiting.len(), 0); + // Connection should go to idle pool since no waiters could receive it + assert_eq!(inner.idle(), 1); + assert_eq!(inner.checked_out(), 0); + } + + #[test] + fn test_same_client_checks_out_two_connections() { + let mut inner = Inner::default(); + inner.online = true; + inner.config.max = 2; + inner.config.min = 0; + + // Add two idle connections to the pool + let server1 = Box::new(Server::default()); + let server1_id = *server1.id(); + let server2 = Box::new(Server::default()); + let server2_id = *server2.id(); + inner.idle_connections.push(server1); + inner.idle_connections.push(server2); + + assert_eq!(inner.idle(), 2); + assert_eq!(inner.checked_out(), 0); + assert_eq!(inner.total(), 2); + + // Same client ID for both requests + let client_id = BackendKeyData::new(); + let request = Request::new(client_id); + + // Check out first connection + let conn1 = inner + .take(&request) + .unwrap() + .expect("should get connection"); + assert_eq!(inner.idle(), 1); + assert_eq!(inner.checked_out(), 1); + assert_eq!(inner.total(), 2); + + // Check out second connection with the same client ID + let conn2 = inner + .take(&request) + .unwrap() + .expect("should get connection"); + assert_eq!(inner.idle(), 0); + assert_eq!(inner.checked_out(), 2); + assert_eq!(inner.total(), 2); + + // Verify the connections are different + assert_ne!(conn1.id(), conn2.id()); + + // Check in both connections + let now = Instant::now(); + inner + .maybe_check_in(conn1, now, BackendCounts::default()) + .unwrap(); + assert_eq!(inner.idle(), 1); + assert_eq!(inner.checked_out(), 1); + assert_eq!(inner.total(), 2); + + inner + .maybe_check_in(conn2, now, BackendCounts::default()) + .unwrap(); + assert_eq!(inner.idle(), 2); + assert_eq!(inner.checked_out(), 0); + assert_eq!(inner.total(), 2); + + // Verify the specific servers are back in the idle pool + let idle_ids: Vec<_> = inner.idle_conns().iter().map(|s| *s.id()).collect(); + assert!(idle_ids.contains(&server1_id)); + assert!(idle_ids.contains(&server2_id)); } } diff --git a/pgdog/src/backend/pool/lb/ban.rs b/pgdog/src/backend/pool/lb/ban.rs new file mode 100644 index 000000000..eddf30cf7 --- /dev/null +++ b/pgdog/src/backend/pool/lb/ban.rs @@ -0,0 +1,396 @@ +use super::*; +use parking_lot::RwLock; +use std::time::Instant; + +use tracing::{error, warn}; + +/// Load balancer target ban. +#[derive(Clone, Debug)] +pub struct Ban { + inner: Arc>, + pool: Pool, +} + +impl Ban { + /// Create new ban handler. + pub(super) fn new(pool: &Pool) -> Self { + Self { + inner: Arc::new(RwLock::new(BanInner { ban: None })), + pool: pool.clone(), + } + } + + /// Check if the database is banned. + pub fn banned(&self) -> bool { + self.inner.read().ban.is_some() + } + + /// Get ban error, if any. + pub fn error(&self) -> Option { + self.inner.read().ban.as_ref().map(|b| b.error) + } + + /// Unban the database. + pub fn unban(&self, manual_check: bool) { + let mut guard = self.inner.upgradable_read(); + if let Some(ref ban) = guard.ban { + if ban.error != Error::ManualBan || !manual_check { + guard.with_upgraded(|guard| { + guard.ban = None; + }); + } + warn!("resuming read queries [{}]", self.pool.addr()); + } + } + + /// Get reference to the connection pool. + pub fn pool(&self) -> &Pool { + &self.pool + } + + /// Ban the database for the ban_timeout duration. + pub fn ban(&self, error: Error, ban_timeout: Duration) -> bool { + let created_at = Instant::now(); + let mut guard = self.inner.upgradable_read(); + + if guard.ban.is_none() { + guard.with_upgraded(|guard| { + guard.ban = Some(BanEntry { + created_at, + error, + ban_timeout, + }); + self.pool.lock().dump_idle(); + }); + drop(guard); + error!("read queries banned: {} [{}]", error, self.pool.addr()); + true + } else { + false + } + } + + /// Remove ban if it has expired. + pub(super) fn unban_if_expired(&self, now: Instant) -> bool { + let mut guard = self.inner.upgradable_read(); + let unbanned = if guard + .ban + .as_ref() + .map(|b| b.expired(now) && b.error != Error::ManualBan) + .unwrap_or(false) + { + let healthy = self.pool.healthy(); + if !healthy { + warn!( + "not resuming read queries, pool is unhealthy [{}]", + self.pool.addr() + ); + return false; + } + guard.with_upgraded(|guard| { + guard.ban = None; + }); + + true + } else { + false + }; + drop(guard); + if unbanned { + warn!("resuming read queries [{}]", self.pool.addr()); + } + unbanned + } +} + +#[derive(Debug)] +struct BanEntry { + created_at: Instant, + error: Error, + ban_timeout: Duration, +} + +#[derive(Debug)] +pub(super) struct BanInner { + ban: Option, +} + +impl BanEntry { + fn expired(&self, now: Instant) -> bool { + now.duration_since(self.created_at) >= self.ban_timeout + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::thread; + + #[test] + fn test_new_ban_is_not_banned() { + let pool = Pool::new_test(); + let ban = Ban::new(&pool); + assert!(!ban.banned()); + assert!(ban.error().is_none()); + } + + #[test] + fn test_ban_sets_banned_state() { + let pool = Pool::new_test(); + let ban = Ban::new(&pool); + let result = ban.ban(Error::ServerError, Duration::from_secs(1)); + assert!(result); + assert!(ban.banned()); + } + + #[test] + fn test_ban_returns_error() { + let pool = Pool::new_test(); + let ban = Ban::new(&pool); + ban.ban(Error::ConnectTimeout, Duration::from_secs(1)); + assert_eq!(ban.error(), Some(Error::ConnectTimeout)); + } + + #[test] + fn test_ban_twice_returns_false() { + let pool = Pool::new_test(); + let ban = Ban::new(&pool); + let first = ban.ban(Error::ServerError, Duration::from_secs(1)); + let second = ban.ban(Error::ConnectTimeout, Duration::from_secs(1)); + assert!(first); + assert!(!second); + assert_eq!(ban.error(), Some(Error::ServerError)); + } + + #[test] + fn test_unban_clears_state() { + let pool = Pool::new_test(); + let ban = Ban::new(&pool); + ban.ban(Error::ServerError, Duration::from_secs(1)); + ban.unban(false); + assert!(!ban.banned()); + assert!(ban.error().is_none()); + } + + #[test] + fn test_unban_manual_check_does_not_clear_manual_ban() { + let pool = Pool::new_test(); + let ban = Ban::new(&pool); + ban.ban(Error::ManualBan, Duration::from_secs(1)); + ban.unban(true); + assert!(ban.banned()); + assert_eq!(ban.error(), Some(Error::ManualBan)); + } + + #[test] + fn test_unban_manual_check_clears_non_manual_ban() { + let pool = Pool::new_test(); + let ban = Ban::new(&pool); + ban.ban(Error::ServerError, Duration::from_secs(1)); + ban.unban(true); + assert!(!ban.banned()); + assert!(ban.error().is_none()); + } + + #[test] + fn test_unban_if_expired_before_expiration() { + let pool = Pool::new_test(); + let ban = Ban::new(&pool); + let now = Instant::now(); + ban.ban(Error::ServerError, Duration::from_secs(10)); + let unbanned = ban.unban_if_expired(now); + assert!(!unbanned); + assert!(ban.banned()); + } + + #[test] + fn test_unban_if_expired_after_expiration() { + let pool = Pool::new_test(); + let ban = Ban::new(&pool); + let now = Instant::now(); + ban.ban(Error::ServerError, Duration::from_millis(1)); + let future = now + Duration::from_millis(10); + let unbanned = ban.unban_if_expired(future); + assert!(unbanned); + assert!(!ban.banned()); + } + + #[test] + fn test_unban_if_expired_not_banned() { + let pool = Pool::new_test(); + let ban = Ban::new(&pool); + let now = Instant::now(); + let unbanned = ban.unban_if_expired(now); + assert!(!unbanned); + assert!(!ban.banned()); + } + + #[test] + fn test_pool_reference() { + let pool = Pool::new_test(); + let ban = Ban::new(&pool); + assert_eq!(ban.pool().id(), pool.id()); + } + + #[test] + fn test_ban_race_condition_multiple_threads() { + let pool = Pool::new_test(); + let ban = Ban::new(&pool); + let ban_clone1 = ban.clone(); + let ban_clone2 = ban.clone(); + let ban_clone3 = ban.clone(); + + let success_count = Arc::new(AtomicUsize::new(0)); + let success1 = success_count.clone(); + let success2 = success_count.clone(); + let success3 = success_count.clone(); + + let h1 = thread::spawn(move || { + if ban_clone1.ban(Error::ServerError, Duration::from_secs(1)) { + success1.fetch_add(1, Ordering::SeqCst); + } + }); + + let h2 = thread::spawn(move || { + if ban_clone2.ban(Error::ConnectTimeout, Duration::from_secs(1)) { + success2.fetch_add(1, Ordering::SeqCst); + } + }); + + let h3 = thread::spawn(move || { + if ban_clone3.ban(Error::CheckoutTimeout, Duration::from_secs(1)) { + success3.fetch_add(1, Ordering::SeqCst); + } + }); + + h1.join().unwrap(); + h2.join().unwrap(); + h3.join().unwrap(); + + assert_eq!(success_count.load(Ordering::SeqCst), 1); + assert!(ban.banned()); + } + + #[test] + fn test_concurrent_ban_and_check() { + let pool = Pool::new_test(); + let ban = Ban::new(&pool); + let ban_clone = ban.clone(); + + let h1 = thread::spawn(move || { + for _ in 0..1000 { + let _ = ban_clone.banned(); + } + }); + + for _ in 0..100 { + ban.ban(Error::ServerError, Duration::from_secs(1)); + ban.unban(false); + } + + h1.join().unwrap(); + } + + #[test] + fn test_concurrent_unban_if_expired() { + let pool = Pool::new_test(); + let ban = Ban::new(&pool); + ban.ban(Error::ServerError, Duration::from_millis(1)); + thread::sleep(Duration::from_millis(10)); + + let ban_clone1 = ban.clone(); + let ban_clone2 = ban.clone(); + let ban_clone3 = ban.clone(); + + let unban_count = Arc::new(AtomicUsize::new(0)); + let count1 = unban_count.clone(); + let count2 = unban_count.clone(); + let count3 = unban_count.clone(); + + let h1 = thread::spawn(move || { + let now = Instant::now(); + if ban_clone1.unban_if_expired(now) { + count1.fetch_add(1, Ordering::SeqCst); + } + }); + + let h2 = thread::spawn(move || { + let now = Instant::now(); + if ban_clone2.unban_if_expired(now) { + count2.fetch_add(1, Ordering::SeqCst); + } + }); + + let h3 = thread::spawn(move || { + let now = Instant::now(); + if ban_clone3.unban_if_expired(now) { + count3.fetch_add(1, Ordering::SeqCst); + } + }); + + h1.join().unwrap(); + h2.join().unwrap(); + h3.join().unwrap(); + + assert_eq!(unban_count.load(Ordering::SeqCst), 1); + assert!(!ban.banned()); + } + + #[test] + fn test_ban_preserves_first_error() { + let pool = Pool::new_test(); + let ban = Ban::new(&pool); + + ban.ban(Error::ServerError, Duration::from_secs(1)); + ban.ban(Error::ConnectTimeout, Duration::from_secs(2)); + ban.ban(Error::CheckoutTimeout, Duration::from_secs(3)); + + assert_eq!(ban.error(), Some(Error::ServerError)); + } + + #[test] + fn test_multiple_ban_unban_cycles() { + let pool = Pool::new_test(); + let ban = Ban::new(&pool); + + for _ in 0..10 { + assert!(ban.ban(Error::ServerError, Duration::from_secs(1))); + assert!(ban.banned()); + ban.unban(false); + assert!(!ban.banned()); + } + } + + #[test] + fn test_unban_if_expired_does_not_unban_unhealthy_pool() { + let pool = Pool::new_test(); + let ban = Ban::new(&pool); + let now = Instant::now(); + + ban.ban(Error::ServerError, Duration::from_millis(1)); + pool.inner().health.toggle(false); + + let future = now + Duration::from_millis(10); + let unbanned = ban.unban_if_expired(future); + + assert!(!unbanned); + assert!(ban.banned()); + } + + #[test] + fn test_unban_if_expired_does_not_unban_manual_ban() { + let pool = Pool::new_test(); + let ban = Ban::new(&pool); + let now = Instant::now(); + + ban.ban(Error::ManualBan, Duration::from_millis(1)); + + let future = now + Duration::from_millis(10); + let unbanned = ban.unban_if_expired(future); + + assert!(!unbanned); + assert!(ban.banned()); + assert_eq!(ban.error(), Some(Error::ManualBan)); + } +} diff --git a/pgdog/src/backend/pool/lb/mod.rs b/pgdog/src/backend/pool/lb/mod.rs new file mode 100644 index 000000000..fabd54fa4 --- /dev/null +++ b/pgdog/src/backend/pool/lb/mod.rs @@ -0,0 +1,328 @@ +//! Load balanced connection pool. + +use std::{ + sync::{ + atomic::{AtomicBool, AtomicUsize, Ordering}, + Arc, + }, + time::Duration, +}; + +use rand::seq::SliceRandom; +use tokio::{ + sync::Notify, + time::{timeout, Instant}, +}; +use tracing::warn; + +use crate::config::{LoadBalancingStrategy, ReadWriteSplit, Role}; +use crate::net::messages::BackendKeyData; + +use super::{Error, Guard, Pool, PoolConfig, Request}; + +pub mod ban; +pub mod monitor; +pub mod target_health; + +use ban::Ban; +use monitor::*; +pub use target_health::*; + +#[cfg(test)] +mod test; + +/// Read query load balancer target. +#[derive(Clone, Debug)] +pub struct Target { + pub pool: Pool, + pub ban: Ban, + replica: Arc, + pub health: TargetHealth, +} + +impl Target { + pub(super) fn new(pool: Pool, role: Role) -> Self { + let ban = Ban::new(&pool); + Self { + ban, + replica: Arc::new(AtomicBool::new(role == Role::Replica)), + health: pool.inner().health.clone(), + pool, + } + } + + /// Get role. + pub(super) fn role(&self) -> Role { + if self.replica.load(Ordering::Relaxed) { + Role::Replica + } else { + Role::Primary + } + } + + /// Set role. + pub(super) fn set_role(&self, role: Role) -> bool { + let value = role == Role::Replica; + let old = self.replica.swap(value, Ordering::Relaxed); + value != old + } +} + +/// Load balancer. +#[derive(Clone, Default, Debug)] +pub struct LoadBalancer { + /// Read/write targets. + pub(super) targets: Vec, + /// Checkout timeout. + pub(super) checkout_timeout: Duration, + /// Round robin atomic counter. + pub(super) round_robin: Arc, + /// Chosen load balancing strategy. + pub(super) lb_strategy: LoadBalancingStrategy, + /// Maintenance. notification. + pub(super) maintenance: Arc, + /// Read/write split. + pub(super) rw_split: ReadWriteSplit, +} + +impl LoadBalancer { + /// Create new replicas pools. + pub fn new( + primary: &Option, + addrs: &[PoolConfig], + lb_strategy: LoadBalancingStrategy, + rw_split: ReadWriteSplit, + ) -> LoadBalancer { + let checkout_timeout = primary + .as_ref() + .map(|primary| primary.config().checkout_timeout) + .unwrap_or(Duration::ZERO) + + addrs + .iter() + .map(|c| c.config.checkout_timeout) + .sum::(); + + let mut targets: Vec<_> = addrs + .iter() + .map(|config| Target::new(Pool::new(config), Role::Replica)) + .collect(); + + let primary_target = primary + .as_ref() + .map(|pool| Target::new(pool.clone(), Role::Primary)); + + if let Some(primary) = primary_target { + targets.push(primary); + } + + Self { + targets, + checkout_timeout, + round_robin: Arc::new(AtomicUsize::new(0)), + lb_strategy, + maintenance: Arc::new(Notify::new()), + rw_split, + } + } + + /// Get the primary pool, if configured. + pub fn primary(&self) -> Option<&Pool> { + self.primary_target().map(|target| &target.pool) + } + + /// Get the primary read target containing the pool, ban state, and health. + /// + /// Unlike [`primary()`], this returns the full target struct which allows + /// access to ban and health state for monitoring and testing purposes. + pub fn primary_target(&self) -> Option<&Target> { + self.targets + .iter() + .rev() // If there is a primary, it's likely to be last. + .find(|target| target.role() == Role::Primary) + } + + /// Detect database roles from pg_is_in_recovery() and + /// return new primary (if any), and replicas. + pub fn redetect_roles(&self) -> bool { + let mut promoted = false; + + let mut targets = self + .targets + .clone() + .into_iter() + .map(|target| (target.pool.lsn_stats(), target)) + .collect::>(); + + // Pick primary by latest data. The one with the most + // up-to-date lsn number and pg_is_in_recovery() = false + // is the new primary. + // + // The old primary is still part of the config and will be demoted + // to replica. If it's down, it will be banned from serving traffic. + // + let now = Instant::now(); + targets.sort_by_cached_key(|target| target.0.lsn_age(now)); + + let primary = targets + .iter() + .position(|target| !target.0.replica && target.0.valid()); + + if let Some(primary) = primary { + promoted = targets[primary].1.set_role(Role::Primary); + + if promoted { + warn!("new primary chosen: {}", targets[primary].1.pool.addr()); + + // Demote everyone else to replicas. + targets + .iter() + .enumerate() + .filter(|(i, _)| *i != primary) + .for_each(|(_, target)| { + target.1.set_role(Role::Replica); + }); + } + } + + promoted + } + + /// Launch replica pools and start the monitor. + pub fn launch(&self) { + self.targets.iter().for_each(|target| target.pool.launch()); + Monitor::spawn(self); + } + + /// Get a live connection from the pool. + pub async fn get(&self, request: &Request) -> Result { + match timeout(self.checkout_timeout, self.get_internal(request)).await { + Ok(Ok(conn)) => Ok(conn), + Ok(Err(err)) => Err(err), + Err(_) => Err(Error::ReplicaCheckoutTimeout), + } + } + + /// Move connections from this replica set to another. + pub fn move_conns_to(&self, destination: &LoadBalancer) -> Result<(), Error> { + assert_eq!(self.targets.len(), destination.targets.len()); + + for (from, to) in self.targets.iter().zip(destination.targets.iter()) { + from.pool.move_conns_to(&to.pool)?; + } + + Ok(()) + } + + /// The two replica sets are referring to the same databases. + pub fn can_move_conns_to(&self, destination: &LoadBalancer) -> bool { + self.targets.len() == destination.targets.len() + && self + .targets + .iter() + .zip(destination.targets.iter()) + .all(|(a, b)| a.pool.can_move_conns_to(&b.pool)) + } + + /// There are no replicas. + pub fn has_replicas(&self) -> bool { + self.targets + .iter() + .any(|target| target.role() == Role::Replica) + } + + /// Cancel a query if one is running. + pub async fn cancel(&self, id: &BackendKeyData) -> Result<(), super::super::Error> { + for target in &self.targets { + target.pool.cancel(id).await?; + } + + Ok(()) + } + + /// Replica pools handle. + pub fn pools(&self) -> Vec<&Pool> { + self.targets.iter().map(|target| &target.pool).collect() + } + + /// Collect all connection pools used for read queries. + pub fn pools_with_roles_and_bans(&self) -> Vec<(Role, Ban, Pool)> { + let result: Vec<_> = self + .targets + .iter() + .map(|target| (target.role(), target.ban.clone(), target.pool.clone())) + .collect(); + + result + } + + async fn get_internal(&self, request: &Request) -> Result { + use LoadBalancingStrategy::*; + use ReadWriteSplit::*; + + let mut candidates: Vec<&Target> = self.targets.iter().collect(); + + let primary_reads = match self.rw_split { + IncludePrimary => true, + IncludePrimaryIfReplicaBanned => candidates.iter().any(|target| target.ban.banned()), + ExcludePrimary => false, + }; + + if !primary_reads { + candidates.retain(|target| target.role() == Role::Replica); + } + + if candidates.is_empty() { + return Err(Error::AllReplicasDown); + } + + match self.lb_strategy { + Random => candidates.shuffle(&mut rand::rng()), + RoundRobin => { + let first = self.round_robin.fetch_add(1, Ordering::Relaxed) % candidates.len(); + let mut reshuffled = vec![]; + reshuffled.extend_from_slice(&candidates[first..]); + reshuffled.extend_from_slice(&candidates[..first]); + candidates = reshuffled; + } + LeastActiveConnections => { + candidates.sort_by_cached_key(|target| target.pool.lock().idle()); + } + } + + // Only ban a candidate pool if there are more than one + // and we have alternates. + let bannable = candidates.len() > 1; + + for target in &candidates { + if target.ban.banned() { + continue; + } + match target.pool.get(request).await { + Ok(conn) => return Ok(conn), + Err(Error::Offline) => { + continue; + } + Err(err) => { + if bannable { + target.ban.ban(err, target.pool.config().ban_timeout); + } + } + } + } + + candidates.iter().for_each(|target| target.ban.unban(true)); + + Err(Error::AllReplicasDown) + } + + /// Shutdown replica pools. + /// + /// N.B. The primary pool is managed by `super::Shard`. + pub fn shutdown(&self) { + for target in &self.targets { + target.pool.shutdown(); + } + + self.maintenance.notify_waiters(); + } +} diff --git a/pgdog/src/backend/pool/lb/monitor.rs b/pgdog/src/backend/pool/lb/monitor.rs new file mode 100644 index 000000000..83bb9a7c6 --- /dev/null +++ b/pgdog/src/backend/pool/lb/monitor.rs @@ -0,0 +1,100 @@ +use std::time::Instant; + +use super::*; + +use tokio::{select, spawn, task::JoinHandle, time::interval}; +use tracing::debug; + +static MAINTENANCE: Duration = Duration::from_millis(333); + +#[derive(Clone, Debug)] +pub(super) struct Monitor { + replicas: LoadBalancer, +} + +impl Monitor { + /// Create new replica targets monitor. + pub(super) fn spawn(replicas: &LoadBalancer) -> JoinHandle<()> { + let monitor = Self { + replicas: replicas.clone(), + }; + + spawn(async move { + monitor.run().await; + }) + } + + async fn run(&self) { + let mut interval = interval(MAINTENANCE); + + debug!("replicas monitor running"); + + let targets = &self.replicas.targets; + + loop { + let mut check_offline = false; + let mut ban_targets = Vec::new(); + + select! { + _ = interval.tick() => {} + _ = self.replicas.maintenance.notified() => { + check_offline = true; + } + } + + if check_offline { + let offline = self + .replicas + .targets + .iter() + .all(|target| !target.pool.lock().online); + + if offline { + break; + } + } + + let now = Instant::now(); + let mut banned = 0; + + for (i, target) in targets.iter().enumerate() { + let healthy = target.health.healthy(); + // Clear expired bans. + if healthy { + target.ban.unban_if_expired(now); + } + + let bannable = + targets.len() > 1 && target.pool.config().ban_timeout > Duration::ZERO; + + // Check health and ban if unhealthy. + if !healthy && bannable { + let already_banned = target.ban.banned(); + if already_banned || !healthy { + banned += 1; + } + if !healthy { + ban_targets.push(i); + } + } + } + + // Clear all bans if all targets are unhealthy. + if targets.len() == banned { + targets.iter().for_each(|target| { + target.ban.unban(false); + }); + } else { + for i in ban_targets { + targets.get(i).map(|target| { + target + .ban + .ban(Error::PoolUnhealthy, target.pool.config().ban_timeout) + }); + } + } + } + + debug!("replicas monitor shut down"); + } +} diff --git a/pgdog/src/backend/pool/lb/target_health.rs b/pgdog/src/backend/pool/lb/target_health.rs new file mode 100644 index 000000000..7e89d76b6 --- /dev/null +++ b/pgdog/src/backend/pool/lb/target_health.rs @@ -0,0 +1,30 @@ +//! Keep a record of each pool's health. + +use std::sync::{ + atomic::{AtomicBool, Ordering}, + Arc, +}; + +#[derive(Clone, Debug)] +pub struct TargetHealth { + #[allow(dead_code)] + pub(super) id: u64, + pub(super) healthy: Arc, +} + +impl TargetHealth { + pub fn new(id: u64) -> Self { + Self { + id, + healthy: Arc::new(AtomicBool::new(true)), + } + } + + pub fn toggle(&self, healthy: bool) { + self.healthy.swap(healthy, Ordering::SeqCst); + } + + pub fn healthy(&self) -> bool { + self.healthy.load(Ordering::Relaxed) + } +} diff --git a/pgdog/src/backend/pool/lb/test.rs b/pgdog/src/backend/pool/lb/test.rs new file mode 100644 index 000000000..a9523acd3 --- /dev/null +++ b/pgdog/src/backend/pool/lb/test.rs @@ -0,0 +1,1041 @@ +use std::collections::HashSet; +use std::time::Duration; +use tokio::time::sleep; + +use crate::backend::pool::{Address, Config, Error, PoolConfig, Request}; +use crate::config::LoadBalancingStrategy; + +use super::*; +use monitor::Monitor; + +fn create_test_pool_config(host: &str, port: u16) -> PoolConfig { + PoolConfig { + address: Address { + host: host.into(), + port, + user: "pgdog".into(), + password: "pgdog".into(), + database_name: "pgdog".into(), + ..Default::default() + }, + config: Config { + max: 1, + checkout_timeout: Duration::from_millis(1000), + ban_timeout: Duration::from_millis(100), + ..Default::default() + }, + ..Default::default() + } +} + +fn setup_test_replicas() -> LoadBalancer { + let pool_config1 = create_test_pool_config("127.0.0.1", 5432); + let pool_config2 = create_test_pool_config("localhost", 5432); + + let replicas = LoadBalancer::new( + &None, + &[pool_config1, pool_config2], + LoadBalancingStrategy::Random, + ReadWriteSplit::IncludePrimary, + ); + replicas.launch(); + replicas +} + +#[tokio::test] +async fn test_replica_ban_recovery_after_timeout() { + let replicas = setup_test_replicas(); + + // Ban the first replica with very short timeout + let ban = &replicas.targets[0].ban; + ban.ban(Error::ServerError, Duration::from_millis(50)); + + assert!(ban.banned()); + + // Wait for ban to expire + sleep(Duration::from_millis(60)).await; + + // Check if ban would be removed (simulate monitor behavior) + let now = std::time::Instant::now(); + let unbanned = ban.unban_if_expired(now); + + assert!(unbanned); + assert!(!ban.banned()); + + replicas.shutdown(); +} + +#[tokio::test] +async fn test_replica_manual_unban() { + let replicas = setup_test_replicas(); + + // Ban the first replica + let ban = &replicas.targets[0].ban; + ban.ban(Error::ServerError, Duration::from_millis(1000)); + + assert!(ban.banned()); + + // Manually unban + ban.unban(false); + + assert!(!ban.banned()); + + replicas.shutdown(); +} + +#[tokio::test] +async fn test_replica_ban_error_retrieval() { + let replicas = setup_test_replicas(); + + let ban = &replicas.targets[0].ban; + + // No error initially + assert!(ban.error().is_none()); + + // Ban with specific error + ban.ban(Error::ServerError, Duration::from_millis(100)); + + // Should return the ban error + let error = ban.error().unwrap(); + assert!(matches!(error, Error::ServerError)); + + replicas.shutdown(); +} + +#[tokio::test] +async fn test_multiple_replica_banning() { + let replicas = setup_test_replicas(); + + // Ban both replicas + for i in 0..2 { + let ban = &replicas.targets[i].ban; + ban.ban(Error::ServerError, Duration::from_millis(100)); + + assert!(ban.banned()); + } + + // Both should be banned + assert_eq!( + replicas.targets.iter().filter(|r| r.ban.banned()).count(), + 2 + ); + + replicas.shutdown(); +} + +#[tokio::test] +async fn test_replica_ban_idempotency() { + let replicas = setup_test_replicas(); + + let ban = &replicas.targets[0].ban; + + // First ban should succeed + let first_ban = ban.ban(Error::ServerError, Duration::from_millis(100)); + assert!(first_ban); + assert!(ban.banned()); + + // Second ban of same replica should not create new ban + let second_ban = ban.ban(Error::ConnectTimeout, Duration::from_millis(200)); + assert!(!second_ban); + assert!(ban.banned()); + + // Error should still be the original one + assert!(matches!(ban.error().unwrap(), Error::ServerError)); + + replicas.shutdown(); +} + +#[tokio::test] +async fn test_pools_with_roles_and_bans() { + let replicas = setup_test_replicas(); + + let pools_info = replicas.pools_with_roles_and_bans(); + + // Should have 2 replica pools (no primary in this test) + assert_eq!(pools_info.len(), 2); + + // All should be replica role + for (role, _ban, _pool) in &pools_info { + assert!(matches!(role, crate::config::Role::Replica)); + } + + replicas.shutdown(); +} + +#[tokio::test] +async fn test_primary_pool_banning() { + let primary_config = create_test_pool_config("127.0.0.1", 5432); + let primary_pool = Pool::new(&primary_config); + primary_pool.launch(); + + let replica_configs = [create_test_pool_config("localhost", 5432)]; + + let replicas = LoadBalancer::new( + &Some(primary_pool), + &replica_configs, + LoadBalancingStrategy::Random, + ReadWriteSplit::IncludePrimary, + ); + replicas.launch(); + + // Test primary ban exists + assert!(replicas.primary_target().is_some()); + + let primary_ban = &replicas.primary_target().unwrap().ban; + + // Ban primary for reads + primary_ban.ban(Error::ServerError, Duration::from_millis(100)); + + assert!(primary_ban.banned()); + + // Check pools with roles includes primary + let pools_info = replicas.pools_with_roles_and_bans(); + assert_eq!(pools_info.len(), 2); // 1 replica + 1 primary + + let has_primary = pools_info + .iter() + .any(|(role, _ban, _pool)| matches!(role, crate::config::Role::Primary)); + assert!(has_primary); + + // Shutdown both primary and replicas + replicas.shutdown(); +} + +#[tokio::test] +async fn test_ban_timeout_not_expired() { + let replicas = setup_test_replicas(); + + let ban = &replicas.targets[0].ban; + ban.ban(Error::ServerError, Duration::from_millis(1000)); // Long timeout + + assert!(ban.banned()); + + // Check immediately - should not be expired + let now = std::time::Instant::now(); + let unbanned = ban.unban_if_expired(now); + + assert!(!unbanned); + assert!(ban.banned()); + + replicas.shutdown(); +} + +#[tokio::test] +async fn test_unban_if_expired_checks_pool_health() { + let replicas = setup_test_replicas(); + + let ban = &replicas.targets[0].ban; + let pool = &replicas.targets[0].pool; + + ban.ban(Error::ServerError, Duration::from_millis(50)); + assert!(ban.banned()); + + pool.inner().health.toggle(false); + + sleep(Duration::from_millis(60)).await; + + let now = std::time::Instant::now(); + let unbanned = ban.unban_if_expired(now); + + assert!(!unbanned); + assert!(ban.banned()); + + replicas.shutdown(); +} + +#[tokio::test] +async fn test_replica_ban_clears_idle_connections() { + let replicas = setup_test_replicas(); + + // Get a connection and return it to create idle connections + let request = Request::default(); + let conn = replicas.pools()[0] + .get(&request) + .await + .expect("Should be able to get connection from launched pool"); + + // Verify we have a valid connection + assert!(!conn.error()); + + drop(conn); // Return to pool as idle + + // Give a moment for the connection to be properly returned to idle state + sleep(Duration::from_millis(10)).await; + + // Check that we have idle connections before banning + let idle_before = replicas.pools()[0].lock().idle(); + assert!( + idle_before > 0, + "Should have idle connections before banning, but found {}", + idle_before + ); + + let ban = &replicas.targets[0].ban; + + // Ban should trigger dump_idle() on the pool + ban.ban(Error::ServerError, Duration::from_millis(100)); + + // Verify the ban was applied + assert!(ban.banned()); + + // Verify that idle connections were cleared + let idle_after = replicas.pools()[0].lock().idle(); + assert_eq!( + idle_after, 0, + "Idle connections should be cleared after banning" + ); + + replicas.shutdown(); +} + +#[tokio::test] +async fn test_monitor_automatic_ban_expiration() { + let replicas = setup_test_replicas(); + + // Ban the first replica with very short timeout + let ban = &replicas.targets[0].ban; + ban.ban(Error::ServerError, Duration::from_millis(100)); + + assert!(ban.banned()); + + // Wait longer than the ban timeout to allow monitor to process + // The monitor runs every 333ms, so we wait for at least one cycle + sleep(Duration::from_millis(400)).await; + + // The monitor should have automatically unbanned the replica + // Note: Since the monitor runs in a background task spawned during Replicas::new(), + // and we can't easily control its timing in tests, we check that the ban + // can be expired when checked + let now = std::time::Instant::now(); + let would_be_unbanned = ban.unban_if_expired(now); + + // Either it was already unbanned by the monitor, or it would be unbanned now + assert!(!ban.banned() || would_be_unbanned); + + replicas.shutdown(); +} + +#[tokio::test] +async fn test_read_write_split_exclude_primary() { + let primary_config = create_test_pool_config("127.0.0.1", 5432); + let primary_pool = Pool::new(&primary_config); + primary_pool.launch(); + + let replica_configs = [ + create_test_pool_config("localhost", 5432), + create_test_pool_config("127.0.0.1", 5432), + ]; + + let replicas = LoadBalancer::new( + &Some(primary_pool), + &replica_configs, + LoadBalancingStrategy::Random, + ReadWriteSplit::ExcludePrimary, + ); + replicas.launch(); + + let request = Request::default(); + + // Try getting connections multiple times and verify primary is never used + let mut replica_ids = HashSet::new(); + for _ in 0..100 { + let conn = replicas.get(&request).await.unwrap(); + replica_ids.insert(conn.pool.id()); + } + + // Should only use replica pools, not primary + assert_eq!(replica_ids.len(), 2); + + // Verify primary pool ID is not in the set of used pools + let primary_id = replicas.primary().unwrap().id(); + assert!(!replica_ids.contains(&primary_id)); + + // Shutdown both primary and replicas + replicas.shutdown(); +} + +#[tokio::test] +async fn test_read_write_split_include_primary() { + let primary_config = create_test_pool_config("127.0.0.1", 5432); + let primary_pool = Pool::new(&primary_config); + primary_pool.launch(); + + let replica_configs = [create_test_pool_config("localhost", 5432)]; + + let replicas = LoadBalancer::new( + &Some(primary_pool), + &replica_configs, + LoadBalancingStrategy::Random, + ReadWriteSplit::IncludePrimary, + ); + replicas.launch(); + + let request = Request::default(); + + // Try getting connections multiple times and verify both primary and replica can be used + let mut used_pool_ids = HashSet::new(); + for _ in 0..20 { + let conn = replicas.get(&request).await.unwrap(); + used_pool_ids.insert(conn.pool.id()); + } + + // Should use both primary and replica pools + assert_eq!(used_pool_ids.len(), 2); + + // Verify primary pool ID is in the set of used pools + let primary_id = replicas.primary().unwrap().id(); + assert!(used_pool_ids.contains(&primary_id)); + + // Shutdown both primary and replicas + replicas.shutdown(); +} + +#[tokio::test] +async fn test_read_write_split_exclude_primary_no_primary() { + // Test exclude primary setting when no primary exists + let replica_configs = [ + create_test_pool_config("localhost", 5432), + create_test_pool_config("127.0.0.1", 5432), + ]; + + let replicas = LoadBalancer::new( + &None, + &replica_configs, + LoadBalancingStrategy::Random, + ReadWriteSplit::ExcludePrimary, + ); + replicas.launch(); + + let request = Request::default(); + + // Should work normally with just replicas + let mut replica_ids = HashSet::new(); + for _ in 0..10 { + let conn = replicas.get(&request).await.unwrap(); + replica_ids.insert(conn.pool.id()); + } + + assert_eq!(replica_ids.len(), 2); + + replicas.shutdown(); +} + +#[tokio::test] +async fn test_read_write_split_include_primary_no_primary() { + // Test include primary setting when no primary exists + let replica_configs = [ + create_test_pool_config("localhost", 5432), + create_test_pool_config("127.0.0.1", 5432), + ]; + + let replicas = LoadBalancer::new( + &None, + &replica_configs, + LoadBalancingStrategy::Random, + ReadWriteSplit::IncludePrimary, + ); + replicas.launch(); + + let request = Request::default(); + + // Should work normally with just replicas + let mut replica_ids = HashSet::new(); + for _ in 0..10 { + let conn = replicas.get(&request).await.unwrap(); + replica_ids.insert(conn.pool.id()); + } + + // Should use both replica pools + assert_eq!(replica_ids.len(), 2); + + replicas.shutdown(); +} + +#[tokio::test] +async fn test_read_write_split_with_banned_primary() { + let primary_config = create_test_pool_config("127.0.0.1", 5432); + let primary_pool = Pool::new(&primary_config); + primary_pool.launch(); + + let replica_configs = [create_test_pool_config("localhost", 5432)]; + + let replicas = LoadBalancer::new( + &Some(primary_pool), + &replica_configs, + LoadBalancingStrategy::Random, + ReadWriteSplit::IncludePrimary, + ); + replicas.launch(); + + // Ban the primary + let primary_ban = &replicas.targets.last().unwrap().ban; + primary_ban.ban(Error::ServerError, Duration::from_millis(1000)); + + let request = Request::default(); + + // Should only use replica even though primary inclusion is enabled + let mut used_pool_ids = HashSet::new(); + for _ in 0..10 { + let conn = replicas.get(&request).await.unwrap(); + used_pool_ids.insert(conn.pool.id()); + } + + // Should only use replica pool since primary is banned + assert_eq!(used_pool_ids.len(), 1); + + // Verify primary pool ID is not in the set of used pools + let primary_id = replicas.targets.last().unwrap().pool.id(); + assert!(!used_pool_ids.contains(&primary_id)); + + // Shutdown both primary and replicas + replicas.shutdown(); +} + +#[tokio::test] +async fn test_read_write_split_with_banned_replicas() { + let primary_config = create_test_pool_config("127.0.0.1", 5432); + let primary_pool = Pool::new(&primary_config); + primary_pool.launch(); + + let replica_configs = [create_test_pool_config("localhost", 5432)]; + + let replicas = LoadBalancer::new( + &Some(primary_pool), + &replica_configs, + LoadBalancingStrategy::Random, + ReadWriteSplit::IncludePrimary, + ); + replicas.launch(); + + // Ban the replica + let replica_ban = &replicas.targets[0].ban; + replica_ban.ban(Error::ServerError, Duration::from_millis(1000)); + + let request = Request::default(); + + // Should only use primary since replica is banned + let mut used_pool_ids = HashSet::new(); + for _ in 0..10 { + let conn = replicas.get(&request).await.unwrap(); + used_pool_ids.insert(conn.pool.id()); + } + + // Should only use primary pool since replica is banned + assert_eq!(used_pool_ids.len(), 1); + + // Verify primary pool ID is in the set of used pools + let primary_id = replicas.targets.last().unwrap().pool.id(); + assert!(used_pool_ids.contains(&primary_id)); + + // Shutdown both primary and replicas + replicas.shutdown(); +} + +#[tokio::test] +async fn test_read_write_split_exclude_primary_with_round_robin() { + let primary_config = create_test_pool_config("127.0.0.1", 5432); + let primary_pool = Pool::new(&primary_config); + primary_pool.launch(); + + let replica_configs = [ + create_test_pool_config("localhost", 5432), + create_test_pool_config("127.0.0.1", 5432), + ]; + + let replicas = LoadBalancer::new( + &Some(primary_pool), + &replica_configs, + LoadBalancingStrategy::RoundRobin, + ReadWriteSplit::ExcludePrimary, + ); + replicas.launch(); + + let request = Request::default(); + + // Collect pool IDs from multiple requests to verify round-robin behavior + let mut pool_sequence = Vec::new(); + for _ in 0..8 { + let conn = replicas.get(&request).await.unwrap(); + pool_sequence.push(conn.pool.id()); + } + + // Should use both replicas (round-robin) + let unique_ids: HashSet<_> = pool_sequence.iter().collect(); + assert_eq!(unique_ids.len(), 2); + + // Verify primary is never used + let primary_id = replicas.targets.last().unwrap().pool.id(); + assert!(!pool_sequence.contains(&primary_id)); + + // Verify round-robin pattern: each pool should be different from the previous one + for i in 1..pool_sequence.len() { + assert_ne!( + pool_sequence[i], + pool_sequence[i - 1], + "Round-robin pattern broken: consecutive pools are the same at positions {} and {}", + i - 1, + i + ); + } + + // Shutdown both primary and replicas + replicas.shutdown(); +} + +#[tokio::test] +async fn test_monitor_shuts_down_on_notify() { + let pool_config1 = create_test_pool_config("127.0.0.1", 5432); + let pool_config2 = create_test_pool_config("localhost", 5432); + + let replicas = LoadBalancer::new( + &None, + &[pool_config1, pool_config2], + LoadBalancingStrategy::Random, + ReadWriteSplit::IncludePrimary, + ); + + replicas + .targets + .iter() + .for_each(|target| target.pool.launch()); + let monitor_handle = Monitor::spawn(&replicas); + + // Give monitor time to start and register notified() future + sleep(Duration::from_millis(10)).await; + + replicas.shutdown(); + + let result = tokio::time::timeout(Duration::from_secs(1), monitor_handle).await; + + assert!( + result.is_ok(), + "Monitor should shut down within timeout after notify" + ); + assert!( + result.unwrap().is_ok(), + "Monitor task should complete successfully" + ); +} + +#[tokio::test] +async fn test_monitor_bans_unhealthy_target() { + let replicas = setup_test_replicas(); + + replicas.targets[0].health.toggle(false); + + sleep(Duration::from_millis(400)).await; + + assert!(replicas.targets[0].ban.banned()); + + replicas.shutdown(); +} + +#[tokio::test] +async fn test_monitor_clears_expired_bans() { + let replicas = setup_test_replicas(); + + replicas.targets[0] + .ban + .ban(Error::ServerError, Duration::from_millis(50)); + + sleep(Duration::from_millis(400)).await; + + assert!(!replicas.targets[0].ban.banned()); + + replicas.shutdown(); +} + +#[tokio::test] +async fn test_monitor_does_not_ban_single_target() { + let pool_config = create_test_pool_config("127.0.0.1", 5432); + + let replicas = LoadBalancer::new( + &None, + &[pool_config], + LoadBalancingStrategy::Random, + ReadWriteSplit::IncludePrimary, + ); + replicas.launch(); + + replicas.targets[0].health.toggle(false); + + sleep(Duration::from_millis(400)).await; + + assert!(!replicas.targets[0].ban.banned()); + + replicas.shutdown(); +} + +#[tokio::test] +async fn test_monitor_unbans_all_when_all_unhealthy() { + let replicas = setup_test_replicas(); + + replicas.targets[0].health.toggle(false); + replicas.targets[1].health.toggle(false); + + sleep(Duration::from_millis(400)).await; + + assert!(!replicas.targets[0].ban.banned()); + assert!(!replicas.targets[1].ban.banned()); + + replicas.shutdown(); +} + +#[tokio::test] +async fn test_monitor_does_not_ban_with_zero_ban_timeout() { + let pool_config1 = PoolConfig { + address: Address { + host: "127.0.0.1".into(), + port: 5432, + user: "pgdog".into(), + password: "pgdog".into(), + database_name: "pgdog".into(), + ..Default::default() + }, + config: Config { + max: 1, + checkout_timeout: Duration::from_millis(1000), + ban_timeout: Duration::ZERO, + ..Default::default() + }, + ..Default::default() + }; + + let pool_config2 = PoolConfig { + address: Address { + host: "localhost".into(), + port: 5432, + user: "pgdog".into(), + password: "pgdog".into(), + database_name: "pgdog".into(), + ..Default::default() + }, + config: Config { + max: 1, + checkout_timeout: Duration::from_millis(1000), + ban_timeout: Duration::ZERO, + ..Default::default() + }, + ..Default::default() + }; + + let replicas = LoadBalancer::new( + &None, + &[pool_config1, pool_config2], + LoadBalancingStrategy::Random, + ReadWriteSplit::IncludePrimary, + ); + replicas.launch(); + + replicas.targets[0].health.toggle(false); + + sleep(Duration::from_millis(400)).await; + + assert!(!replicas.targets[0].ban.banned()); + + replicas.shutdown(); +} + +#[tokio::test] +async fn test_monitor_health_state_race() { + use tokio::spawn; + + let replicas = setup_test_replicas(); + let target = replicas.targets[0].clone(); + + let toggle_task = spawn(async move { + for _ in 0..50 { + target.health.toggle(false); + sleep(Duration::from_micros(100)).await; + target.health.toggle(true); + sleep(Duration::from_micros(100)).await; + } + }); + + sleep(Duration::from_millis(500)).await; + + toggle_task.await.unwrap(); + + let banned = replicas.targets[0].ban.banned(); + let healthy = replicas.targets[0].health.healthy(); + + assert!( + !banned || !healthy, + "Pool should not be banned if healthy after race" + ); + + replicas.shutdown(); +} + +#[tokio::test] +async fn test_include_primary_if_replica_banned_no_bans() { + let primary_config = create_test_pool_config("127.0.0.1", 5432); + let primary_pool = Pool::new(&primary_config); + primary_pool.launch(); + + let replica_configs = [create_test_pool_config("localhost", 5432)]; + + let replicas = LoadBalancer::new( + &Some(primary_pool), + &replica_configs, + LoadBalancingStrategy::Random, + ReadWriteSplit::IncludePrimaryIfReplicaBanned, + ); + replicas.launch(); + + let request = Request::default(); + + // When no replicas are banned, primary should NOT be used + let mut used_pool_ids = HashSet::new(); + for _ in 0..20 { + let conn = replicas.get(&request).await.unwrap(); + used_pool_ids.insert(conn.pool.id()); + } + + // Should only use replica pool + assert_eq!(used_pool_ids.len(), 1); + + // Verify primary pool ID is not in the set of used pools + let primary_id = replicas.primary().unwrap().id(); + assert!(!used_pool_ids.contains(&primary_id)); + + // Shutdown both primary and replicas + replicas.shutdown(); +} + +#[tokio::test] +async fn test_include_primary_if_replica_banned_with_ban() { + let primary_config = create_test_pool_config("127.0.0.1", 5432); + let primary_pool = Pool::new(&primary_config); + primary_pool.launch(); + + let replica_configs = [create_test_pool_config("localhost", 5432)]; + + let replicas = LoadBalancer::new( + &Some(primary_pool), + &replica_configs, + LoadBalancingStrategy::Random, + ReadWriteSplit::IncludePrimaryIfReplicaBanned, + ); + replicas.launch(); + + // Ban the replica + let replica_ban = &replicas.targets[0].ban; + replica_ban.ban(Error::ServerError, Duration::from_millis(1000)); + + let request = Request::default(); + + // When replica is banned, primary SHOULD be used + let mut used_pool_ids = HashSet::new(); + for _ in 0..20 { + let conn = replicas.get(&request).await.unwrap(); + used_pool_ids.insert(conn.pool.id()); + } + + // Should only use primary pool since replica is banned + assert_eq!(used_pool_ids.len(), 1); + + // Verify primary pool ID is in the set of used pools + let primary_id = replicas.primary().unwrap().id(); + assert!(used_pool_ids.contains(&primary_id)); + + // Shutdown both primary and replicas + replicas.shutdown(); +} + +#[tokio::test] +async fn test_has_replicas_with_replicas() { + let replicas = setup_test_replicas(); + + assert!(replicas.has_replicas()); + + replicas.shutdown(); +} + +#[tokio::test] +async fn test_has_replicas_with_primary_and_replicas() { + let primary_config = create_test_pool_config("127.0.0.1", 5432); + let primary_pool = Pool::new(&primary_config); + primary_pool.launch(); + + let replica_configs = [create_test_pool_config("localhost", 5432)]; + + let lb = LoadBalancer::new( + &Some(primary_pool), + &replica_configs, + LoadBalancingStrategy::Random, + ReadWriteSplit::IncludePrimary, + ); + lb.launch(); + + assert!(lb.has_replicas()); + + lb.shutdown(); +} + +#[tokio::test] +async fn test_has_replicas_primary_only() { + let primary_config = create_test_pool_config("127.0.0.1", 5432); + let primary_pool = Pool::new(&primary_config); + primary_pool.launch(); + + let lb = LoadBalancer::new( + &Some(primary_pool), + &[], + LoadBalancingStrategy::Random, + ReadWriteSplit::IncludePrimary, + ); + lb.launch(); + + assert!(!lb.has_replicas()); + + lb.shutdown(); +} + +#[tokio::test] +async fn test_has_replicas_empty() { + let lb = LoadBalancer::new( + &None, + &[], + LoadBalancingStrategy::Random, + ReadWriteSplit::IncludePrimary, + ); + + assert!(!lb.has_replicas()); +} + +#[tokio::test] +async fn test_set_role() { + let replicas = setup_test_replicas(); + + // Initially all targets are replicas + assert_eq!(replicas.targets[0].role(), Role::Replica); + assert_eq!(replicas.targets[1].role(), Role::Replica); + + // Setting replica to replica returns false (no change) + let changed = replicas.targets[0].set_role(Role::Replica); + assert!(!changed); + assert_eq!(replicas.targets[0].role(), Role::Replica); + + // Setting replica to primary returns true (changed) + let changed = replicas.targets[0].set_role(Role::Primary); + assert!(changed); + assert_eq!(replicas.targets[0].role(), Role::Primary); + + // Setting primary to primary returns false (no change) + let changed = replicas.targets[0].set_role(Role::Primary); + assert!(!changed); + assert_eq!(replicas.targets[0].role(), Role::Primary); + + // Setting primary to replica returns true (changed) + let changed = replicas.targets[0].set_role(Role::Replica); + assert!(changed); + assert_eq!(replicas.targets[0].role(), Role::Replica); + + replicas.shutdown(); +} + +#[tokio::test] +async fn test_can_move_conns_to_same_config() { + let pool_config1 = create_test_pool_config("127.0.0.1", 5432); + let pool_config2 = create_test_pool_config("localhost", 5432); + + let lb1 = LoadBalancer::new( + &None, + &[pool_config1.clone(), pool_config2.clone()], + LoadBalancingStrategy::Random, + ReadWriteSplit::IncludePrimary, + ); + + let lb2 = LoadBalancer::new( + &None, + &[pool_config1, pool_config2], + LoadBalancingStrategy::Random, + ReadWriteSplit::IncludePrimary, + ); + + assert!(lb1.can_move_conns_to(&lb2)); +} + +#[tokio::test] +async fn test_can_move_conns_to_different_count() { + let pool_config1 = create_test_pool_config("127.0.0.1", 5432); + let pool_config2 = create_test_pool_config("localhost", 5432); + + let lb1 = LoadBalancer::new( + &None, + &[pool_config1.clone(), pool_config2], + LoadBalancingStrategy::Random, + ReadWriteSplit::IncludePrimary, + ); + + let lb2 = LoadBalancer::new( + &None, + &[pool_config1], + LoadBalancingStrategy::Random, + ReadWriteSplit::IncludePrimary, + ); + + assert!(!lb1.can_move_conns_to(&lb2)); +} + +#[tokio::test] +async fn test_can_move_conns_to_different_addresses() { + let pool_config1 = create_test_pool_config("127.0.0.1", 5432); + let pool_config2 = create_test_pool_config("localhost", 5432); + let pool_config3 = create_test_pool_config("127.0.0.1", 5433); + + let lb1 = LoadBalancer::new( + &None, + &[pool_config1, pool_config2], + LoadBalancingStrategy::Random, + ReadWriteSplit::IncludePrimary, + ); + + let lb2 = LoadBalancer::new( + &None, + &[pool_config3.clone(), pool_config3], + LoadBalancingStrategy::Random, + ReadWriteSplit::IncludePrimary, + ); + + assert!(!lb1.can_move_conns_to(&lb2)); +} + +#[tokio::test] +async fn test_monitor_unbans_all_when_second_target_becomes_unhealthy_after_first_banned() { + let replicas = setup_test_replicas(); + + // First target becomes unhealthy + replicas.targets[0].health.toggle(false); + + // Wait for monitor to ban the first target + sleep(Duration::from_millis(400)).await; + + assert!( + replicas.targets[0].ban.banned(), + "First target should be banned" + ); + assert!( + !replicas.targets[1].ban.banned(), + "Second target should not be banned yet" + ); + + // Now second target becomes unhealthy (first is already banned) + replicas.targets[1].health.toggle(false); + + // Wait for monitor to process - should unban all since all are unhealthy + sleep(Duration::from_millis(400)).await; + + // Both should be unbanned because all targets are unhealthy + assert!( + !replicas.targets[0].ban.banned(), + "First target should be unbanned when all targets are unhealthy" + ); + assert!( + !replicas.targets[1].ban.banned(), + "Second target should be unbanned when all targets are unhealthy" + ); + + replicas.shutdown(); +} diff --git a/pgdog/src/backend/pool/lsn_monitor.rs b/pgdog/src/backend/pool/lsn_monitor.rs new file mode 100644 index 000000000..c01054011 --- /dev/null +++ b/pgdog/src/backend/pool/lsn_monitor.rs @@ -0,0 +1,252 @@ +use std::time::Duration; + +use tokio::{ + select, spawn, + time::{interval, sleep, timeout, Instant}, +}; +use tracing::{debug, error, trace}; + +use crate::{ + backend::replication::publisher::Lsn, + net::{DataRow, Format, TimestampTz}, +}; + +use super::*; + +static AURORA_DETECTION_QUERY: &str = "SELECT aurora_version()"; + +static LSN_QUERY: &str = " +SELECT + pg_is_in_recovery() AS replica, + CASE + WHEN pg_is_in_recovery() THEN + COALESCE( + pg_last_wal_replay_lsn(), + pg_last_wal_receive_lsn() + ) + ELSE + pg_current_wal_lsn() + END AS lsn, + CASE + WHEN pg_is_in_recovery() THEN + COALESCE( + pg_last_wal_replay_lsn(), + pg_last_wal_receive_lsn() + ) - '0/0'::pg_lsn + ELSE + pg_current_wal_lsn() - '0/0'::pg_lsn + END AS offset_bytes, + CASE + WHEN pg_is_in_recovery() THEN + COALESCE(pg_last_xact_replay_timestamp(), now()) + ELSE + now() + END AS timestamp +"; + +static AURORA_LSN_QUERY: &str = " +SELECT + pg_is_in_recovery() AS replica, + '0/0'::pg_lsn AS lsn, + 0::bigint AS offset_bytes, + now() AS timestamp +"; + +/// LSN information. +#[derive(Debug, Clone, Copy)] +pub struct LsnStats { + /// pg_is_in_recovery() + pub replica: bool, + /// Replay LSN on replica, current LSN on primary. + pub lsn: Lsn, + /// LSN position in bytes from 0. + pub offset_bytes: i64, + /// Server timestamp. + pub timestamp: TimestampTz, + /// Our timestamp. + pub fetched: Instant, + /// Running on Aurora. + pub aurora: bool, +} + +impl Default for LsnStats { + fn default() -> Self { + Self { + replica: true, // Replica unless proven otherwise. + lsn: Lsn::default(), + offset_bytes: 0, + timestamp: TimestampTz::default(), + fetched: Instant::now(), + aurora: false, + } + } +} + +impl LsnStats { + /// How old the stats are. + pub fn lsn_age(&self, now: Instant) -> Duration { + now.duration_since(self.fetched) + } + + /// Stats contain real data. + pub fn valid(&self) -> bool { + self.aurora || self.lsn.lsn > 0 + } +} +impl LsnStats { + fn from_row(value: DataRow, aurora: bool) -> Self { + Self { + replica: value.get(0, Format::Text).unwrap_or_default(), + lsn: value.get(1, Format::Text).unwrap_or_default(), + offset_bytes: value.get(2, Format::Text).unwrap_or_default(), + timestamp: value.get(3, Format::Text).unwrap_or_default(), + fetched: Instant::now(), + aurora, + } + } +} + +/// LSN monitor loop. +pub(super) struct LsnMonitor { + pool: Pool, +} + +impl LsnMonitor { + pub(super) fn run(pool: &Pool) { + let monitor = Self { pool: pool.clone() }; + + spawn(async move { + monitor.spawn().await; + }); + } + + async fn run_query(&self, conn: &mut Guard, query: &str) -> Option { + match timeout(self.pool.config().lsn_check_timeout, conn.fetch_all(query)).await { + Ok(Ok(rows)) => rows.into_iter().next(), + Ok(Err(err)) => { + error!("lsn monitor query error: {} [{}]", err, self.pool.addr()); + None + } + Err(_) => { + error!("lsn monitor query timeout [{}]", self.pool.addr()); + None + } + } + } + + async fn detect_aurora(&self, conn: &mut Guard) -> Option { + match timeout( + self.pool.config().lsn_check_timeout, + conn.fetch_all::(AURORA_DETECTION_QUERY), + ) + .await + { + Ok(Ok(_)) => { + debug!("aurora detected [{}]", self.pool.addr()); + Some(true) + } + Ok(Err(crate::backend::Error::ExecutionError(_))) => Some(false), + Ok(Err(err)) => { + error!( + "lsn monitor aurora detection error: {} [{}]", + err, + self.pool.addr() + ); + None + } + Err(_) => { + error!( + "lsn monitor aurora detection timeout [{}]", + self.pool.addr() + ); + None + } + } + } + + async fn spawn(&self) { + select! { + _ = sleep(self.pool.config().lsn_check_delay) => {}, + _ = self.pool.comms().shutdown.notified() => { return; } + } + + debug!("lsn monitor loop is running [{}]", self.pool.addr()); + + let mut aurora_detected: Option = None; + let mut interval = interval(self.pool.config().lsn_check_interval); + + loop { + select! { + _ = interval.tick() => {}, + _ = self.pool.comms().shutdown.notified() => { break; } + } + + let mut conn = match self.pool.get(&Request::default()).await { + Ok(conn) => conn, + Err(Error::Offline) => break, + Err(err) => { + error!("lsn monitor checkout error: {} [{}]", err, self.pool.addr()); + continue; + } + }; + + if aurora_detected.is_none() { + aurora_detected = self.detect_aurora(&mut conn).await; + } + + let Some(aurora) = aurora_detected else { + continue; + }; + + let query = if aurora { AURORA_LSN_QUERY } else { LSN_QUERY }; + + if let Some(row) = self.run_query(&mut conn, query).await { + drop(conn); + let stats = LsnStats::from_row(row, aurora); + *self.pool.inner().lsn_stats.write() = stats; + trace!("lsn monitor stats updated [{}]", self.pool.addr()); + } + } + + debug!("lsn monitor shutdown [{}]", self.pool.addr()); + } +} + +#[cfg(test)] +mod test { + use super::*; + + #[test] + fn test_aurora_stats_valid_with_zero_lsn() { + let stats = LsnStats { + replica: true, + lsn: Lsn::default(), + offset_bytes: 0, + timestamp: TimestampTz::default(), + fetched: Instant::now(), + aurora: true, + }; + + assert!( + stats.valid(), + "Aurora stats should be valid even with zero LSN" + ); + } + + #[test] + fn test_non_aurora_stats_invalid_with_zero_lsn() { + let stats = LsnStats { + replica: true, + lsn: Lsn::default(), + offset_bytes: 0, + timestamp: TimestampTz::default(), + fetched: Instant::now(), + aurora: false, + }; + + assert!( + !stats.valid(), + "Non-Aurora stats should be invalid with zero LSN" + ); + } +} diff --git a/pgdog/src/backend/pool/mod.rs b/pgdog/src/backend/pool/mod.rs index 17f93181f..2f2d79b37 100644 --- a/pgdog/src/backend/pool/mod.rs +++ b/pgdog/src/backend/pool/mod.rs @@ -1,7 +1,6 @@ //! Manage connections to the servers. pub mod address; -pub mod ban; pub mod cleanup; pub mod cluster; pub mod comms; @@ -12,12 +11,13 @@ pub mod error; pub mod guard; pub mod healthcheck; pub mod inner; +pub mod lb; +pub mod lsn_monitor; pub mod mapping; pub mod mirror_stats; pub mod monitor; pub mod oids; pub mod pool_impl; -pub mod replicas; pub mod request; pub mod shard; pub mod state; @@ -32,20 +32,21 @@ pub use connection::Connection; pub use error::Error; pub use guard::Guard; pub use healthcheck::Healtcheck; +pub use lb::LoadBalancer; +pub use lsn_monitor::LsnStats; pub use mirror_stats::MirrorStats; -use monitor::Monitor; +pub use monitor::Monitor; pub use oids::Oids; pub use pool_impl::Pool; -pub use replicas::Replicas; pub use request::Request; pub use shard::Shard; pub use state::State; pub use stats::Stats; -use ban::Ban; use comms::Comms; use inner::Inner; use mapping::Mapping; +use shard::ShardConfig; use taken::Taken; use waiting::{Waiter, Waiting}; diff --git a/pgdog/src/backend/pool/monitor.rs b/pgdog/src/backend/pool/monitor.rs index de781d870..06405b4ae 100644 --- a/pgdog/src/backend/pool/monitor.rs +++ b/pgdog/src/backend/pool/monitor.rs @@ -35,7 +35,8 @@ use std::time::Duration; use super::{Error, Guard, Healtcheck, Oids, Pool, Request}; -use crate::backend::Server; +use crate::backend::pool::inner::ShouldCreate; +use crate::backend::{ConnectReason, DisconnectReason, Server}; use tokio::time::{interval, sleep, timeout, Instant}; use tokio::{select, task::spawn}; @@ -49,7 +50,7 @@ static MAINTENANCE: Duration = Duration::from_millis(333); /// /// See [`crate::backend::pool::monitor`] module documentation /// for more details. -pub(super) struct Monitor { +pub struct Monitor { pool: Pool, } @@ -75,7 +76,7 @@ impl Monitor { let pool = self.pool.clone(); spawn(async move { Self::stats(pool).await }); - // Delay starting healthchecks to give + // Delay starting health checks to give // time for the pool to spin up. let pool = self.pool.clone(); let (delay, replication_mode) = { @@ -119,10 +120,17 @@ impl Monitor { break; } - if should_create { - let ok = self.replenish().await; + if let ShouldCreate::Yes { reason, .. } = should_create { + info!("new connection requested: {} [{}]", should_create, self.pool.addr()); + let ok = match self.replenish(reason).await { + Ok(ok) => ok, + Err(err) => { + error!("monitor error: {}", err); + false + } + }; if !ok { - self.pool.ban(Error::ServerError); + self.pool.inner().health.toggle(false); } } } @@ -137,17 +145,16 @@ impl Monitor { debug!("maintenance loop is shut down [{}]", self.pool.addr()); } - /// The healthcheck loop. + /// The health check loop. /// - /// Runs regularly and ensures the pool triggers healthchecks on idle connections. + /// Runs regularly and ensures the pool triggers health checks on idle connections. async fn healthchecks(pool: Pool) { let mut tick = interval(pool.lock().config().idle_healthcheck_interval()); let comms = pool.comms(); - debug!("healthchecks running [{}]", pool.addr()); + debug!("health checks running [{}]", pool.addr()); loop { - let mut unbanned = false; select! { _ = tick.tick() => { { @@ -158,29 +165,21 @@ impl Monitor { break; } - // Pool is paused, skip healtcheck. + // Pool is paused, skip health check. if guard.paused { continue; } - } - // If the server is okay, remove the ban if it had one. - if let Ok(true) = Self::healthcheck(&pool).await { - unbanned = pool.lock().maybe_unban(); - } + let _ = Self::healthcheck(&pool).await; } _ = comms.shutdown.notified() => break, } - - if unbanned { - info!("pool unbanned due to healtcheck [{}]", pool.addr()); - } } - debug!("healthchecks stopped [{}]", pool.addr()); + debug!("health checks stopped [{}]", pool.addr()); } /// Perform maintenance on the pool periodically. @@ -204,7 +203,7 @@ impl Monitor { // If a client is waiting already, // create it a connection. - if guard.should_create() { + if guard.should_create().yes() { comms.request.notify_one(); } @@ -215,11 +214,6 @@ impl Monitor { guard.close_idle(now); guard.close_old(now); - let unbanned = guard.check_ban(now); - - if unbanned { - info!("pool unbanned due to maintenance [{}]", pool.addr()); - } } _ = comms.shutdown.notified() => break, @@ -230,14 +224,17 @@ impl Monitor { } /// Replenish pool with one new connection. - async fn replenish(&self) -> bool { - if let Ok(conn) = Self::create_connection(&self.pool).await { + async fn replenish(&self, reason: ConnectReason) -> Result { + if let Ok(conn) = Self::create_connection(&self.pool, reason).await { + let now = Instant::now(); let server = Box::new(conn); let mut guard = self.pool.lock(); - guard.put(server, Instant::now()); - true + if guard.online { + guard.put(server, now)?; + } + Ok(true) } else { - false + Ok(false) } } @@ -255,14 +252,28 @@ impl Monitor { Ok(()) } + pub async fn healthcheck(pool: &Pool) -> Result { + match Self::healthcheck_internal(pool).await { + Ok(result) => { + pool.inner().health.toggle(result); + Ok(result) + } + + Err(err) => { + pool.inner().health.toggle(false); + Err(err) + } + } + } + /// Perform a periodic healthcheck on the pool. - async fn healthcheck(pool: &Pool) -> Result { + async fn healthcheck_internal(pool: &Pool) -> Result { let conn = { let mut guard = pool.lock(); - if !guard.online || guard.banned() { + if !guard.online { return Ok(false); } - guard.take(&Request::default()) + guard.take(&Request::default())? }; let healthcheck_timeout = pool.config().healthcheck_timeout; @@ -276,26 +287,26 @@ impl Monitor { ) .healthcheck() .await?; - - Ok(true) } else { // Create a new one and close it. info!("creating new healthcheck connection [{}]", pool.addr()); - let mut server = Self::create_connection(pool) + let mut server = Self::create_connection(pool, ConnectReason::Healthcheck) .await .map_err(|_| Error::HealthcheckError)?; + server.disconnect_reason(DisconnectReason::Healthcheck); + Healtcheck::mandatory(&mut server, pool, healthcheck_timeout) .healthcheck() .await?; - - Ok(true) } + + Ok(true) } async fn stats(pool: Pool) { - let duration = Duration::from_secs(15); + let duration = pool.config().stats_period; let comms = pool.comms(); loop { @@ -314,22 +325,34 @@ impl Monitor { } } - pub(super) async fn create_connection(pool: &Pool) -> Result { + pub(super) async fn create_connection( + pool: &Pool, + reason: ConnectReason, + ) -> Result { let connect_timeout = pool.config().connect_timeout; let connect_attempts = pool.config().connect_attempts; let connect_attempt_delay = pool.config().connect_attempt_delay; let options = pool.server_options(); let mut error = Error::ServerError; + let now = Instant::now(); for attempt in 0..connect_attempts { match timeout( connect_timeout, - Server::connect(pool.addr(), options.clone()), + Server::connect(pool.addr(), options.clone(), reason), ) .await { - Ok(Ok(conn)) => return Ok(conn), + Ok(Ok(conn)) => { + let elapsed = now.elapsed(); + { + let mut guard = pool.lock(); + guard.stats.counts.connect_count += 1; + guard.stats.counts.connect_time += elapsed; + } + return Ok(conn); + } Ok(Err(err)) => { error!( @@ -369,6 +392,7 @@ impl Monitor { #[cfg(test)] mod test { use crate::backend::pool::test::pool; + use crate::backend::pool::{Address, Config, PoolConfig}; use super::*; @@ -378,9 +402,104 @@ mod test { let pool = pool(); let ok = Monitor::healthcheck(&pool).await.unwrap(); assert!(ok); + } - pool.ban(Error::ManualBan); - let ok = Monitor::healthcheck(&pool).await.unwrap(); - assert!(!ok); + #[tokio::test] + async fn test_healthcheck_sets_health_true_on_success() { + crate::logger(); + let pool = pool(); + + pool.inner().health.toggle(false); + assert!(!pool.inner().health.healthy()); + + let result = Monitor::healthcheck(&pool).await.unwrap(); + + assert!(result); + assert!(pool.inner().health.healthy()); + } + + #[tokio::test] + async fn test_healthcheck_sets_health_false_on_offline_pool() { + crate::logger(); + let pool = pool(); + + pool.inner().health.toggle(true); + assert!(pool.inner().health.healthy()); + + pool.shutdown(); + + let result = Monitor::healthcheck(&pool).await.unwrap(); + + assert!(!result); + assert!(!pool.inner().health.healthy()); + } + + #[tokio::test] + async fn test_healthcheck_sets_health_false_on_connection_error() { + crate::logger(); + + let config = Config { + max: 1, + min: 1, + healthcheck_timeout: Duration::from_millis(10), + ..Default::default() + }; + + let pool = Pool::new(&PoolConfig { + address: Address { + host: "127.0.0.1".into(), + port: 1, + database_name: "pgdog".into(), + user: "pgdog".into(), + password: "pgdog".into(), + ..Default::default() + }, + config, + }); + pool.launch(); + + pool.inner().health.toggle(true); + assert!(pool.inner().health.healthy()); + + let result = Monitor::healthcheck(&pool).await; + + assert!(result.is_err()); + assert!(!pool.inner().health.healthy()); + } + + #[tokio::test] + async fn test_replenish_only_when_pool_is_online() { + crate::logger(); + + let config = Config { + max: 5, + min: 2, + ..Default::default() + }; + + let pool = Pool::new(&PoolConfig { + address: Address { + host: "127.0.0.1".into(), + port: 5432, + database_name: "pgdog".into(), + user: "pgdog".into(), + password: "pgdog".into(), + ..Default::default() + }, + config, + }); + pool.launch(); + + let initial_total = pool.lock().total(); + + pool.shutdown(); + assert!(!pool.lock().online); + + let monitor = Monitor { pool: pool.clone() }; + let ok = monitor.replenish(ConnectReason::Other).await.unwrap(); + + assert!(ok); + assert_eq!(pool.lock().total(), initial_total); + assert!(!pool.lock().online); } } diff --git a/pgdog/src/backend/pool/pool_impl.rs b/pgdog/src/backend/pool/pool_impl.rs index 9e065a5cb..fc645c7b8 100644 --- a/pgdog/src/backend/pool/pool_impl.rs +++ b/pgdog/src/backend/pool/pool_impl.rs @@ -4,21 +4,22 @@ use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::Arc; use std::time::Duration; -use once_cell::sync::Lazy; +use once_cell::sync::{Lazy, OnceCell}; +use parking_lot::RwLock; use parking_lot::{lock_api::MutexGuard, Mutex, RawMutex}; -use tokio::time::Instant; -use tracing::{error, info}; +use tokio::time::{timeout, Instant}; +use tracing::error; -use crate::backend::{Server, ServerOptions}; +use crate::backend::pool::LsnStats; +use crate::backend::{ConnectReason, DisconnectReason, Server, ServerOptions}; use crate::config::PoolerMode; -use crate::net::messages::{BackendKeyData, DataRow, Format}; -use crate::net::Parameter; +use crate::net::messages::BackendKeyData; +use crate::net::{Parameter, Parameters}; use super::inner::CheckInResult; -use super::inner::ReplicaLag; use super::{ - Address, Comms, Config, Error, Guard, Healtcheck, Inner, Monitor, Oids, PoolConfig, Request, - State, Waiting, + lb::TargetHealth, lsn_monitor::LsnMonitor, Address, Comms, Config, Error, Guard, Healtcheck, + Inner, Monitor, Oids, PoolConfig, Request, State, Waiting, }; static ID_COUNTER: Lazy> = Lazy::new(|| Arc::new(AtomicU64::new(0))); @@ -38,6 +39,9 @@ pub(crate) struct InnerSync { pub(super) inner: Mutex, pub(super) id: u64, pub(super) config: Config, + pub(super) health: TargetHealth, + pub(super) params: OnceCell, + pub(super) lsn_stats: RwLock, } impl std::fmt::Debug for Pool { @@ -59,6 +63,9 @@ impl Pool { inner: Mutex::new(Inner::new(config.config, id)), id, config: config.config, + health: TargetHealth::new(id), + params: OnceCell::new(), + lsn_stats: RwLock::new(LsnStats::default()), }), } } @@ -78,85 +85,104 @@ impl Pool { &self.inner } + pub fn healthy(&self) -> bool { + self.inner.health.healthy() + } + /// Launch the maintenance loop, bringing the pool online. pub fn launch(&self) { let mut guard = self.lock(); if !guard.online { guard.online = true; Monitor::run(self); + LsnMonitor::run(self); } } - pub async fn get_forced(&self, request: &Request) -> Result { - self.get_internal(request, true).await - } - pub async fn get(&self, request: &Request) -> Result { - self.get_internal(request, false).await - } - - /// Get a connection from the pool. - async fn get_internal(&self, request: &Request, unban: bool) -> Result { - let pool = self.clone(); - - // Fast path, idle connection probably available. - let (server, granted_at, paused) = { - // Ask for time before we acquire the lock - // and only if we actually waited for a connection. - let granted_at = request.created_at; - let elapsed = granted_at.saturating_duration_since(request.created_at); - let mut guard = self.lock(); - - if !guard.online { - return Err(Error::Offline); + match timeout(self.config().checkout_timeout, self.get_internal(request)).await { + Ok(Ok(conn)) => Ok(conn), + Err(_) => { + self.inner.health.toggle(false); + Err(Error::CheckoutTimeout) } - - // Try this only once. If the pool still - // has an error after a checkout attempt, - // return error. - if unban && guard.banned() { - guard.maybe_unban(); + Ok(Err(err)) => { + self.inner.health.toggle(false); + Err(err) } + } + } - if guard.banned() { - return Err(Error::Banned); + /// Get a connection from the pool. + async fn get_internal(&self, request: &Request) -> Result { + loop { + let pool = self.clone(); + + // Fast path, idle connection probably available. + let (server, granted_at, paused) = { + // Ask for time before we acquire the lock + // and only if we actually waited for a connection. + let granted_at = request.created_at; + let elapsed = granted_at.saturating_duration_since(request.created_at); + let mut guard = self.lock(); + + if !guard.online { + return Err(Error::Offline); + } + + let conn = guard.take(request)?; + + if conn.is_some() { + guard.stats.counts.wait_time += elapsed; + guard.stats.counts.server_assignment_count += 1; + } + + (conn, granted_at, guard.paused) + }; + + if paused { + self.comms().ready.notified().await; } - let conn = guard.take(request); - - if conn.is_some() { - guard.stats.counts.wait_time += elapsed; - guard.stats.counts.server_assignment_count += 1; + let (server, granted_at) = if let Some(mut server) = server { + server + .prepared_statements_mut() + .set_capacity(self.inner.config.prepared_statements_limit); + server.set_pooler_mode(self.inner.config.pooler_mode); + (Guard::new(pool, server, granted_at), granted_at) + } else { + // Slow path, pool is empty, will create new connection + // or wait for one to be returned if the pool is maxed out. + let mut waiting = Waiting::new(pool, request)?; + waiting.wait().await? + }; + + match self + .maybe_healthcheck( + server, + self.inner.config.healthcheck_timeout, + self.inner.config.healthcheck_interval, + granted_at, + ) + .await + { + Ok(conn) => return Ok(conn), + // Try another connection. + Err(Error::HealthcheckError) => continue, + Err(err) => return Err(err), } - - (conn, granted_at, guard.paused) - }; - - if paused { - self.comms().ready.notified().await; } + } - let (server, granted_at) = if let Some(mut server) = server { - server - .prepared_statements_mut() - .set_capacity(self.inner.config.prepared_statements_limit); - server.set_pooler_mode(self.inner.config.pooler_mode); - (Guard::new(pool, server, granted_at), granted_at) + /// Get server parameters, fetch them if necessary. + pub async fn params(&self, request: &Request) -> Result<&Parameters, Error> { + if let Some(params) = self.inner.params.get() { + Ok(params) } else { - // Slow path, pool is empty, will create new connection - // or wait for one to be returned if the pool is maxed out. - let waiting = Waiting::new(pool, request)?; - waiting.wait().await? - }; - - return self - .maybe_healthcheck( - server, - self.inner.config.healthcheck_timeout, - self.inner.config.healthcheck_interval, - granted_at, - ) - .await; + let conn = self.get(request).await?; + let params = conn.params().clone(); + Ok(self.inner.params.get_or_init(|| params)) + } } /// Perform a health check on the connection if one is needed. @@ -176,24 +202,19 @@ impl Pool { ); if let Err(err) = healthcheck.healthcheck().await { + conn.disconnect_reason(DisconnectReason::Unhealthy); drop(conn); - self.ban(Error::HealthcheckError); + self.inner.health.toggle(false); return Err(err); + } else if !self.inner.health.healthy() { + self.inner.health.toggle(true); } Ok(conn) } - /// Create new identical connection pool. - pub fn duplicate(&self) -> Pool { - Pool::new(&PoolConfig { - address: self.addr().clone(), - config: *self.lock().config(), - }) - } - /// Check the connection back into the pool. - pub(super) fn checkin(&self, mut server: Box) { + pub(super) fn checkin(&self, mut server: Box) -> Result<(), Error> { // Server is checked in right after transaction finished // in transaction mode but can be checked in anytime in session mode. let now = if server.pooler_mode() == &PoolerMode::Session { @@ -202,19 +223,25 @@ impl Pool { server.stats().last_used }; - let counts = server.stats_mut().reset_last_checkout(); + let counts = { + let stats = server.stats_mut(); + stats.client_id = None; + stats.reset_last_checkout() + }; // Check everything and maybe check the connection // into the idle pool. - let CheckInResult { banned, replenish } = - { self.lock().maybe_check_in(server, now, counts) }; + let CheckInResult { + server_error, + replenish, + } = { self.lock().maybe_check_in(server, now, counts)? }; - if banned { + if server_error { error!( - "pool banned on check in: {} [{}]", - Error::ServerError, + "pool received broken server connection, closing [{}]", self.addr() ); + self.inner.health.toggle(false); } // Notify maintenance that we need a new connection because @@ -222,6 +249,8 @@ impl Pool { if replenish { self.comms().request.notify_one(); } + + Ok(()) } /// Server connection used by the client. @@ -238,43 +267,14 @@ impl Pool { Ok(()) } - /// Is this pool banned? - pub fn banned(&self) -> bool { - self.lock().banned() - } - /// Pool is available to serve connections. pub fn available(&self) -> bool { let guard = self.lock(); !guard.paused && guard.online } - /// Ban this connection pool from serving traffic. - pub fn ban(&self, reason: Error) { - let now = Instant::now(); - let banned = self.lock().maybe_ban(now, reason); - - if banned { - error!("pool banned explicitly: {} [{}]", reason, self.addr()); - } - } - - /// Unban this pool from serving traffic, unless manually banned. - #[allow(dead_code)] - pub fn maybe_unban(&self) { - let unbanned = self.lock().maybe_unban(); - if unbanned { - info!("pool unbanned [{}]", self.addr()); - } - } - - pub fn unban(&self) { - if self.lock().unban() { - info!("pool unbanned [{}]", self.addr()); - } - } - /// Connection pool unique identifier. + #[inline] pub(crate) fn id(&self) -> u64 { self.inner.id } @@ -283,7 +283,7 @@ impl Pool { /// to a new instance of the pool. /// /// This shuts down the pool. - pub(crate) fn move_conns_to(&self, destination: &Pool) { + pub(crate) fn move_conns_to(&self, destination: &Pool) -> Result<(), Error> { // Ensure no deadlock. assert!(self.inner.id != destination.id()); let now = Instant::now(); @@ -295,13 +295,15 @@ impl Pool { from_guard.online = false; let (idle, taken) = from_guard.move_conns_to(destination); for server in idle { - to_guard.put(server, now); + to_guard.put(server, now)?; } to_guard.set_taken(taken); } destination.launch(); self.shutdown(); + + Ok(()) } /// The two pools refer to the same database. @@ -322,15 +324,14 @@ impl Pool { { let mut guard = self.lock(); guard.paused = false; - guard.ban = None; } self.comms().ready.notify_waiters(); } /// Create a connection to the pool, untracked by the logic here. - pub async fn standalone(&self) -> Result { - Monitor::create_connection(self).await + pub async fn standalone(&self, reason: ConnectReason) -> Result { + Monitor::create_connection(self, reason).await } /// Shutdown the pool. @@ -386,7 +387,7 @@ impl Pool { if let Some(statement_timeout) = config.statement_timeout { params.push(Parameter { name: "statement_timeout".into(), - value: statement_timeout.as_millis().to_string(), + value: statement_timeout.as_millis().to_string().into(), }); } @@ -404,7 +405,10 @@ impl Pool { }); } - ServerOptions { params } + ServerOptions { + params, + pool_id: self.id(), + } } /// Pool state. @@ -412,6 +416,11 @@ impl Pool { State::get(self) } + /// LSN stats + pub fn lsn_stats(&self) -> LsnStats { + *self.inner().lsn_stats.read() + } + /// Update pool configuration used in internals. #[cfg(test)] pub(crate) fn update_config(&self, config: Config) { @@ -422,68 +431,4 @@ impl Pool { pub fn oids(&self) -> Option { self.lock().oids } - - /// `pg_current_wal_flush_lsn()` on the primary. - pub async fn wal_flush_lsn(&self) -> Result { - let mut guard = self.get(&Request::default()).await?; - - let rows: Vec = guard - .fetch_all("SELECT pg_current_wal_flush_lsn()") - .await - .map_err(|_| Error::PrimaryLsnQueryFailed)?; - - let lsn = rows - .first() - .map(|r| r.get::(0, Format::Text).unwrap_or_default()) - .unwrap_or_default(); - - parse_pg_lsn(&lsn).map_err(|_| Error::ReplicaLsnQueryFailed) - } - - /// `pg_last_wal_replay_lsn()` on a replica. - pub async fn wal_replay_lsn(&self) -> Result { - let mut guard = self.get(&Request::default()).await?; - - let rows: Vec = guard - .fetch_all("SELECT pg_last_wal_replay_lsn()") - .await - .map_err(|_| Error::ReplicaLsnQueryFailed)?; - - let lsn = rows - .first() - .map(|r| r.get::(0, Format::Text).unwrap_or_default()) - .unwrap_or_default(); - - parse_pg_lsn(&lsn).map_err(|_| Error::ReplicaLsnQueryFailed) - } - - pub fn set_replica_lag(&self, replica_lag: ReplicaLag) { - self.lock().replica_lag = replica_lag; - } -} - -// ------------------------------------------------------------------------------------------------- -// ----- Utils :: Parse LSN ------------------------------------------------------------------------ - -#[derive(Debug)] -enum ParseLsnError { - MissingSlash, - InvalidHex, -} - -/// Parse PostgreSQL LSN string to u64 bytes. -/// See spec: https://www.postgresql.org/docs/current/datatype-pg-lsn.html -fn parse_pg_lsn(s: &str) -> Result { - let (hi_str, lo_str) = s.split_once('/').ok_or(ParseLsnError::MissingSlash)?; - - let hi = u32::from_str_radix(hi_str, 16).map_err(|_| ParseLsnError::InvalidHex)? as u64; - let lo = u32::from_str_radix(lo_str, 16).map_err(|_| ParseLsnError::InvalidHex)? as u64; - - let shifted_hi = hi << 32; - let lsn_value = shifted_hi | lo; - - Ok(lsn_value) } - -// ------------------------------------------------------------------------------------------------- -// ------------------------------------------------------------------------------------------------- diff --git a/pgdog/src/backend/pool/replicas.rs b/pgdog/src/backend/pool/replicas.rs deleted file mode 100644 index 399197626..000000000 --- a/pgdog/src/backend/pool/replicas.rs +++ /dev/null @@ -1,168 +0,0 @@ -//! Replicas pool. - -use std::{ - sync::{ - atomic::{AtomicUsize, Ordering}, - Arc, - }, - time::Duration, -}; - -use rand::seq::SliceRandom; -use tokio::time::timeout; -use tracing::error; - -use crate::config::LoadBalancingStrategy; -use crate::net::messages::BackendKeyData; - -use super::{Error, Guard, Pool, PoolConfig, Request}; - -/// Replicas pools. -#[derive(Clone, Default, Debug)] -pub struct Replicas { - /// Connection pools. - pub(super) pools: Vec, - /// Checkout timeout. - pub(super) checkout_timeout: Duration, - /// Round robin atomic counter. - pub(super) round_robin: Arc, - /// Chosen load balancing strategy. - pub(super) lb_strategy: LoadBalancingStrategy, -} - -impl Replicas { - /// Create new replicas pools. - pub fn new(addrs: &[PoolConfig], lb_strategy: LoadBalancingStrategy) -> Replicas { - let checkout_timeout = addrs - .iter() - .map(|c| c.config.checkout_timeout()) - .sum::(); - Self { - pools: addrs.iter().map(Pool::new).collect(), - checkout_timeout, - round_robin: Arc::new(AtomicUsize::new(0)), - lb_strategy, - } - } - - /// Get a live connection from the pool. - pub async fn get(&self, request: &Request, primary: &Option) -> Result { - match timeout(self.checkout_timeout, self.get_internal(request, primary)).await { - Ok(Ok(conn)) => Ok(conn), - Ok(Err(err)) => Err(err), - Err(_) => Err(Error::ReplicaCheckoutTimeout), - } - } - - /// Move connections from this replica set to another. - pub fn move_conns_to(&self, destination: &Replicas) { - assert_eq!(self.pools.len(), destination.pools.len()); - - for (from, to) in self.pools.iter().zip(destination.pools.iter()) { - from.move_conns_to(to); - } - } - - /// The two replica sets are referring to the same databases. - pub fn can_move_conns_to(&self, destination: &Replicas) -> bool { - self.pools.len() == destination.pools.len() - && self - .pools - .iter() - .zip(destination.pools.iter()) - .all(|(a, b)| a.can_move_conns_to(b)) - } - - /// How many replicas we are connected to. - pub fn len(&self) -> usize { - self.pools.len() - } - - /// There are no replicas. - pub fn is_empty(&self) -> bool { - self.len() == 0 - } - - /// Create new identical replica pool. - pub fn duplicate(&self) -> Replicas { - Self { - pools: self.pools.iter().map(|p| p.duplicate()).collect(), - checkout_timeout: self.checkout_timeout, - round_robin: Arc::new(AtomicUsize::new(0)), - lb_strategy: self.lb_strategy, - } - } - - /// Cancel a query if one is running. - pub async fn cancel(&self, id: &BackendKeyData) -> Result<(), super::super::Error> { - for pool in &self.pools { - pool.cancel(id).await?; - } - - Ok(()) - } - - /// Pools handle. - pub fn pools(&self) -> &[Pool] { - &self.pools - } - - async fn get_internal( - &self, - request: &Request, - primary: &Option, - ) -> Result { - let mut unbanned = false; - loop { - let mut candidates = self.pools.iter().collect::>(); - - if let Some(primary) = primary { - candidates.push(primary); - } - - use LoadBalancingStrategy::*; - - match self.lb_strategy { - Random => candidates.shuffle(&mut rand::thread_rng()), - RoundRobin => { - let first = self.round_robin.fetch_add(1, Ordering::Relaxed) % candidates.len(); - let mut reshuffled = vec![]; - reshuffled.extend_from_slice(&candidates[first..]); - reshuffled.extend_from_slice(&candidates[..first]); - candidates = reshuffled; - } - LeastActiveConnections => { - candidates.sort_by_cached_key(|pool| pool.lock().idle()); - } - } - - let mut banned = 0; - - for candidate in &candidates { - match candidate.get(request).await { - Ok(conn) => return Ok(conn), - Err(Error::Offline) => continue, - Err(Error::Banned) => { - banned += 1; - continue; - } - Err(err) => { - error!("{} [{}]", err, candidate.addr()); - } - } - } - - // All replicas are banned, unban everyone. - if banned == candidates.len() && !unbanned { - candidates - .iter() - .for_each(|candidate| candidate.maybe_unban()); - unbanned = true; - } else { - break; - } - } - - Err(Error::AllReplicasDown) - } -} diff --git a/pgdog/src/backend/pool/shard.rs b/pgdog/src/backend/pool/shard.rs deleted file mode 100644 index 3d4722870..000000000 --- a/pgdog/src/backend/pool/shard.rs +++ /dev/null @@ -1,507 +0,0 @@ -//! A shard is a collection of replicas and an optional primary. - -use std::ops::Deref; -use std::sync::Arc; -use std::time::Duration; -use tokio::sync::broadcast; -use tokio::time::{interval, sleep}; -use tokio::{join, select, spawn, sync::Notify}; -use tracing::{debug, error}; - -use crate::backend::PubSubListener; -use crate::config::{config, LoadBalancingStrategy, ReadWriteSplit, Role}; -use crate::net::messages::BackendKeyData; -use crate::net::NotificationResponse; - -use super::inner::ReplicaLag; -use super::{Error, Guard, Pool, PoolConfig, Replicas, Request}; - -// ------------------------------------------------------------------------------------------------- -// ----- Public Interface -------------------------------------------------------------------------- - -#[derive(Clone, Debug)] -pub struct Shard { - inner: Arc, -} - -impl Shard { - pub fn new( - primary: &Option, - replicas: &[PoolConfig], - lb_strategy: LoadBalancingStrategy, - rw_split: ReadWriteSplit, - ) -> Self { - Self { - inner: Arc::new(ShardInner::new(primary, replicas, lb_strategy, rw_split)), - } - } - - /// Get connection to primary database. - pub async fn primary(&self, request: &Request) -> Result { - self.primary - .as_ref() - .ok_or(Error::NoPrimary)? - .get_forced(request) - .await - } - - /// Get connection to one of the replica databases, using the configured - /// load balancing algorithm. - pub async fn replica(&self, request: &Request) -> Result { - if self.replicas.is_empty() { - self.primary - .as_ref() - .ok_or(Error::NoDatabases)? - .get(request) - .await - } else { - use ReadWriteSplit::*; - - let primary = match self.rw_split { - IncludePrimary => &self.primary, - ExcludePrimary => &None, - }; - - self.replicas.get(request, primary).await - } - } - - /// Get connection to primary if configured, otherwise replica. - pub async fn primary_or_replica(&self, request: &Request) -> Result { - if self.primary.is_some() { - self.primary(request).await - } else { - self.replica(request).await - } - } - - pub fn move_conns_to(&self, destination: &Shard) { - if let Some(ref primary) = self.primary { - if let Some(ref other) = destination.primary { - primary.move_conns_to(other); - } - } - - self.replicas.move_conns_to(&destination.replicas); - } - - pub(crate) fn can_move_conns_to(&self, other: &Shard) -> bool { - if let Some(ref primary) = self.primary { - if let Some(ref other) = other.primary { - if !primary.can_move_conns_to(other) { - return false; - } - } else { - return false; - } - } - - self.replicas.can_move_conns_to(&other.replicas) - } - - /// Listen for notifications on channel. - pub async fn listen( - &self, - channel: &str, - ) -> Result, Error> { - if let Some(ref listener) = self.pub_sub { - listener.listen(channel).await - } else { - Err(Error::PubSubDisabled) - } - } - - /// Notify channel with optional payload (payload can be empty string). - pub async fn notify(&self, channel: &str, payload: &str) -> Result<(), Error> { - if let Some(ref listener) = self.pub_sub { - listener.notify(channel, payload).await - } else { - Err(Error::PubSubDisabled) - } - } - - /// Clone pools but keep them independent. - pub fn duplicate(&self) -> Self { - let primary = self - .inner - .primary - .as_ref() - .map(|primary| primary.duplicate()); - let pub_sub = if self.pub_sub.is_some() { - primary.as_ref().map(PubSubListener::new) - } else { - None - }; - - Self { - inner: Arc::new(ShardInner { - primary, - pub_sub, - replicas: self.inner.replicas.duplicate(), - rw_split: self.inner.rw_split, - comms: ShardComms::default(), // Create new comms instead of duplicating - }), - } - } - - /// Bring every pool online. - pub fn launch(&self) { - self.pools().iter().for_each(|pool| pool.launch()); - ShardMonitor::run(self); - if let Some(ref listener) = self.pub_sub { - listener.launch(); - } - } - - pub fn has_primary(&self) -> bool { - self.primary.is_some() - } - - pub fn has_replicas(&self) -> bool { - !self.replicas.is_empty() - } - - pub async fn cancel(&self, id: &BackendKeyData) -> Result<(), super::super::Error> { - if let Some(ref primary) = self.primary { - primary.cancel(id).await?; - } - self.replicas.cancel(id).await?; - - Ok(()) - } - - pub fn pools(&self) -> Vec { - self.pools_with_roles() - .into_iter() - .map(|(_, pool)| pool) - .collect() - } - - pub fn pools_with_roles(&self) -> Vec<(Role, Pool)> { - let mut pools = vec![]; - if let Some(primary) = self.primary.clone() { - pools.push((Role::Primary, primary)); - } - - pools.extend( - self.replicas - .pools() - .iter() - .map(|p| (Role::Replica, p.clone())), - ); - - pools - } - - /// Shutdown every pool. - pub fn shutdown(&self) { - self.comms.shutdown.notify_waiters(); - self.pools().iter().for_each(|pool| pool.shutdown()); - if let Some(ref listener) = self.pub_sub { - listener.shutdown(); - } - } - - fn comms(&self) -> &ShardComms { - &self.comms - } -} - -impl Deref for Shard { - type Target = ShardInner; - #[inline] - fn deref(&self) -> &Self::Target { - &self.inner - } -} - -// ------------------------------------------------------------------------------------------------- -// ----- Private Implementation -------------------------------------------------------------------- - -#[derive(Default, Debug)] -pub struct ShardInner { - primary: Option, - replicas: Replicas, - rw_split: ReadWriteSplit, - comms: ShardComms, - pub_sub: Option, -} - -impl ShardInner { - fn new( - primary: &Option, - replicas: &[PoolConfig], - lb_strategy: LoadBalancingStrategy, - rw_split: ReadWriteSplit, - ) -> Self { - let primary = primary.as_ref().map(Pool::new); - let replicas = Replicas::new(replicas, lb_strategy); - let comms = ShardComms { - shutdown: Notify::new(), - }; - let pub_sub = if config().pub_sub_enabled() { - primary.as_ref().map(PubSubListener::new) - } else { - None - }; - - Self { - primary, - replicas, - rw_split, - comms, - pub_sub, - } - } -} - -// ------------------------------------------------------------------------------------------------- -// ----- Comms ------------------------------------------------------------------------------------- - -#[derive(Debug)] -struct ShardComms { - pub shutdown: Notify, -} - -impl Default for ShardComms { - fn default() -> Self { - Self { - shutdown: Notify::new(), - } - } -} - -// ------------------------------------------------------------------------------------------------- -// ----- Monitoring -------------------------------------------------------------------------------- - -struct ShardMonitor {} - -impl ShardMonitor { - pub fn run(shard: &Shard) { - let shard = shard.clone(); - spawn(async move { Self::monitor_replicas(shard).await }); - } -} - -impl ShardMonitor { - async fn monitor_replicas(shard: Shard) { - let cfg_handle = config(); - let Some(rl_config) = cfg_handle.config.replica_lag.as_ref() else { - return; - }; - - let mut tick = interval(rl_config.check_interval); - - debug!("replica monitoring running"); - let comms = shard.comms(); - - loop { - select! { - _ = tick.tick() => { - Self::process_replicas(&shard, rl_config.max_age).await; - } - _ = comms.shutdown.notified() => break, - } - } - - debug!("replica monitoring stopped"); - } - - async fn process_replicas(shard: &Shard, max_age: Duration) { - let Some(primary) = shard.primary.as_ref() else { - return; - }; - - primary.set_replica_lag(ReplicaLag::NonApplicable); - - let lsn_metrics = match collect_lsn_metrics(primary, max_age).await { - Some(m) => m, - None => { - error!("failed to collect LSN metrics"); - return; - } - }; - - for replica in shard.replicas.pools() { - Self::process_single_replica( - replica, - lsn_metrics.max_lsn, - lsn_metrics.average_bytes_per_sec, - max_age, - ) - .await; - } - } - - // TODO -> [process_single_replica] - // - The current query logic prevents pools from executing any query once banned. - // - This includes running their own LSN queries. - // - For this reason, we cannot ban replicas for lagging just yet - // - For now, we simply tracing::error!() it for now. - // - It's sensible to ban replicas from making user queries when it's lagging too much, but... - // unexposed PgDog admin queries should be allowed on "banned" replicas. - // - TLDR; We need a way to distinguish between user and admin queries, and let admin... - // queries run on "banned" replicas. - - async fn process_single_replica( - replica: &Pool, - primary_lsn: u64, - lsn_throughput: f64, - _max_age: Duration, // used to make banning decisions when it's supported later - ) { - if replica.banned() { - replica.set_replica_lag(ReplicaLag::Unknown); - return; - }; - - let replay_lsn = match replica.wal_replay_lsn().await { - Ok(lsn) => lsn, - Err(e) => { - error!("replica {} LSN query failed: {}", replica.id(), e); - return; - } - }; - - let bytes_behind = primary_lsn.saturating_sub(replay_lsn); - - let mut lag = ReplicaLag::Bytes(bytes_behind); - if lsn_throughput > 0.0 { - let duration = Duration::from_secs_f64(bytes_behind as f64 / lsn_throughput); - lag = ReplicaLag::Duration(duration); - } - if bytes_behind == 0 { - lag = ReplicaLag::Duration(Duration::ZERO); - } - - replica.set_replica_lag(lag); - } -} - -// ------------------------------------------------------------------------------------------------- -// ----- Utils :: Primary LSN Metrics -------------------------------------------------------------- - -struct LsnMetrics { - pub average_bytes_per_sec: f64, - pub max_lsn: u64, -} - -/// Sample WAL LSN at 0, half, and full window; compute avg rate and max LSN. -async fn collect_lsn_metrics(primary: &Pool, window: Duration) -> Option { - if window.is_zero() { - return None; - } - - let half_window = window / 2; - let half_window_in_seconds = half_window.as_secs_f64(); - - // fire three futures at once - let f0 = primary.wal_flush_lsn(); - let f1 = async { - sleep(half_window).await; - primary.wal_flush_lsn().await - }; - let f2 = async { - sleep(window).await; - primary.wal_flush_lsn().await - }; - - // collect results concurrently - let (r0, r1, r2) = join!(f0, f1, f2); - - let lsn_initial = r0.ok()?; - let lsn_half = r1.ok()?; - let lsn_full = r2.ok()?; - - let rate1 = (lsn_half.saturating_sub(lsn_initial)) as f64 / half_window_in_seconds; - let rate2 = (lsn_full.saturating_sub(lsn_half)) as f64 / half_window_in_seconds; - let average_rate = (rate1 + rate2) / 2.0; - - let max_lsn = lsn_initial.max(lsn_half).max(lsn_full); - - let metrics = LsnMetrics { - average_bytes_per_sec: average_rate, - max_lsn, - }; - - Some(metrics) -} - -// ------------------------------------------------------------------------------------------------- -// ----- Tests ------------------------------------------------------------------------------------- - -#[cfg(test)] -mod test { - use std::collections::BTreeSet; - - use crate::backend::pool::{Address, Config}; - - use super::*; - - #[tokio::test] - async fn test_exclude_primary() { - crate::logger(); - - let primary = &Some(PoolConfig { - address: Address::new_test(), - config: Config::default(), - }); - - let replicas = &[PoolConfig { - address: Address::new_test(), - config: Config::default(), - }]; - - let shard = Shard::new( - primary, - replicas, - LoadBalancingStrategy::Random, - ReadWriteSplit::ExcludePrimary, - ); - shard.launch(); - - for _ in 0..25 { - let replica_id = shard.replicas.pools[0].id(); - - let conn = shard.replica(&Request::default()).await.unwrap(); - assert_eq!(conn.pool.id(), replica_id); - } - - shard.shutdown(); - } - - #[tokio::test] - async fn test_include_primary() { - crate::logger(); - - let primary = &Some(PoolConfig { - address: Address::new_test(), - config: Config::default(), - }); - - let replicas = &[PoolConfig { - address: Address::new_test(), - config: Config::default(), - }]; - - let shard = Shard::new( - primary, - replicas, - LoadBalancingStrategy::Random, - ReadWriteSplit::IncludePrimary, - ); - shard.launch(); - let mut ids = BTreeSet::new(); - - for _ in 0..25 { - let conn = shard.replica(&Request::default()).await.unwrap(); - ids.insert(conn.pool.id()); - } - - shard.shutdown(); - - assert_eq!(ids.len(), 2); - } -} - -// ------------------------------------------------------------------------------------------------- -// ------------------------------------------------------------------------------------------------- diff --git a/pgdog/src/backend/pool/shard/mod.rs b/pgdog/src/backend/pool/shard/mod.rs new file mode 100644 index 000000000..9429b431d --- /dev/null +++ b/pgdog/src/backend/pool/shard/mod.rs @@ -0,0 +1,352 @@ +//! A shard is a collection of replicas and an optional primary. + +use std::ops::Deref; +use std::sync::Arc; +use std::time::Duration; +use tokio::sync::broadcast; +use tokio::{select, spawn, sync::Notify}; +use tracing::debug; + +use crate::backend::databases::User; +use crate::backend::pool::lb::ban::Ban; +use crate::backend::PubSubListener; +use crate::config::{config, LoadBalancingStrategy, ReadWriteSplit, Role}; +use crate::net::messages::BackendKeyData; +use crate::net::NotificationResponse; + +use super::{Error, Guard, LoadBalancer, Pool, PoolConfig, Request}; + +pub mod monitor; +pub mod role_detector; + +use monitor::*; +use role_detector::*; + +pub(super) struct ShardConfig<'a> { + /// Shard number. + pub(super) number: usize, + /// Shard primary database, if any. + pub(super) primary: &'a Option, + /// Shard replica databases. + pub(super) replicas: &'a [PoolConfig], + /// Load balancing strategy for replicas. + pub(super) lb_strategy: LoadBalancingStrategy, + /// Primary/replica read/write split strategy. + pub(super) rw_split: ReadWriteSplit, + /// Cluster identifier (user/password). + pub(super) identifier: Arc, + /// LSN check interval + pub(super) lsn_check_interval: Duration, +} + +/// Connection pools for a single database shard. +/// +/// Includes a primary and replicas. +#[derive(Clone, Debug)] +pub struct Shard { + inner: Arc, +} + +impl Shard { + /// Create new shard connection pools from configuration. + /// + /// # Arguments + /// + /// * `primary`: Primary configuration, if any. Primary databases are optional. + /// * `replica`: List of replica database configurations. + /// * `lb_strategy`: Query load balancing strategy, e.g., random, round robin, etc. + /// * `rw_split`: Read/write traffic splitting strategy. + /// + pub(super) fn new(config: ShardConfig<'_>) -> Self { + Self { + inner: Arc::new(ShardInner::new(config)), + } + } + + /// Get connection to the primary database. + pub async fn primary(&self, request: &Request) -> Result { + self.lb + .primary() + .ok_or(Error::NoPrimary)? + .get(request) + .await + } + + /// Get connection to one of the replica databases, using the configured + /// load balancing algorithm. + pub async fn replica(&self, request: &Request) -> Result { + self.lb.get(request).await + } + + /// Get connection to primary if configured, otherwise replica. + pub async fn primary_or_replica(&self, request: &Request) -> Result { + if let Ok(primary) = self.primary(request).await { + Ok(primary) + } else { + self.replica(request).await + } + } + + /// Move connections from this shard to another shard, preserving them. + /// + /// This is done during configuration reloading, if no significant changes are made to + /// the configuration. + pub fn move_conns_to(&self, destination: &Shard) -> Result<(), Error> { + self.lb.move_conns_to(&destination.lb)?; + + Ok(()) + } + + /// Checks if the connection pools from this shard are compatible + /// with the other shard. If yes, they can be moved without closing them. + pub(crate) fn can_move_conns_to(&self, other: &Shard) -> bool { + self.lb.can_move_conns_to(&other.lb) + } + + /// Listen for notifications on channel. + pub async fn listen( + &self, + channel: &str, + ) -> Result, Error> { + if let Some(ref listener) = self.pub_sub { + listener.listen(channel).await + } else { + Err(Error::PubSubDisabled) + } + } + + /// Notify channel with optional payload (payload can be empty string). + pub async fn notify(&self, channel: &str, payload: &str) -> Result<(), Error> { + if let Some(ref listener) = self.pub_sub { + listener.notify(channel, payload).await + } else { + Err(Error::PubSubDisabled) + } + } + + /// Bring every pool online. + pub fn launch(&self) { + self.lb.launch(); + ShardMonitor::run(self); + if let Some(ref listener) = self.pub_sub { + listener.launch(); + } + } + + /// Returns true if the shard has a primary database. + pub fn has_primary(&self) -> bool { + self.lb.primary().is_some() + } + + /// Returns true if the shard has any replica databases. + pub fn has_replicas(&self) -> bool { + self.lb.has_replicas() + } + + /// Request a query to be cancelled on any of the servers in the connection pools + /// in this shard. + /// + /// # Arguments + /// + /// * `id`: Client unique identifier. Clients can execute one query at a time. + /// + /// If these connection pools aren't running the query sent by this client, this is a no-op. + /// + pub async fn cancel(&self, id: &BackendKeyData) -> Result<(), super::super::Error> { + self.lb.cancel(id).await?; + + Ok(()) + } + + /// Get all connection pools. + pub fn pools(&self) -> Vec { + self.pools_with_roles() + .into_iter() + .map(|(_, pool)| pool) + .collect() + } + + /// Get all connection pools along with their roles (i.e., primary or replica). + pub fn pools_with_roles(&self) -> Vec<(Role, Pool)> { + let mut pools = vec![]; + + pools.extend( + self.lb + .targets + .iter() + .map(|target| (target.role(), target.pool.clone())), + ); + + pools + } + + /// Get all connection pools with bans and their role in the shard. + pub fn pools_with_roles_and_bans(&self) -> Vec<(Role, Ban, Pool)> { + self.lb.pools_with_roles_and_bans() + } + + /// Shutdown every pool and maintenance task in this shard. + pub fn shutdown(&self) { + self.comms.shutdown.notify_waiters(); + if let Some(ref listener) = self.pub_sub { + listener.shutdown(); + } + self.lb.shutdown(); + } + + fn comms(&self) -> &ShardComms { + &self.comms + } + + pub fn number(&self) -> usize { + self.number + } + + pub fn identifier(&self) -> &User { + &self.identifier + } + + /// Re-detect primary/replica roles and re-build + /// the shard routing logic. + pub fn redetect_roles(&self) -> bool { + self.lb.redetect_roles() + } +} + +impl Deref for Shard { + type Target = ShardInner; + #[inline] + fn deref(&self) -> &Self::Target { + &self.inner + } +} + +/// Shard connection pools +/// and internal state. +#[derive(Default, Debug, Clone)] +pub struct ShardInner { + number: usize, + lb: LoadBalancer, + comms: Arc, + pub_sub: Option, + identifier: Arc, +} + +impl ShardInner { + fn new(shard: ShardConfig) -> Self { + let ShardConfig { + number, + primary, + replicas, + lb_strategy, + rw_split, + identifier, + lsn_check_interval, + } = shard; + let primary = primary.as_ref().map(Pool::new); + let lb = LoadBalancer::new(&primary, replicas, lb_strategy, rw_split); + let comms = Arc::new(ShardComms { + shutdown: Notify::new(), + lsn_check_interval, + }); + let pub_sub = if config().pub_sub_enabled() { + primary.as_ref().map(PubSubListener::new) + } else { + None + }; + + Self { + number, + lb, + comms, + pub_sub, + identifier, + } + } +} + +#[cfg(test)] +mod test { + use std::collections::BTreeSet; + + use crate::backend::pool::Address; + + use super::*; + + #[tokio::test] + async fn test_exclude_primary() { + crate::logger(); + + let primary = &Some(PoolConfig { + address: Address::new_test(), + ..Default::default() + }); + + let replicas = &[PoolConfig { + address: Address::new_test(), + ..Default::default() + }]; + + let shard = Shard::new(ShardConfig { + number: 0, + primary, + replicas, + lb_strategy: LoadBalancingStrategy::Random, + rw_split: ReadWriteSplit::ExcludePrimary, + identifier: Arc::new(User { + user: "pgdog".into(), + database: "pgdog".into(), + }), + lsn_check_interval: Duration::MAX, + }); + shard.launch(); + + for _ in 0..25 { + let replica_id = shard.lb.targets[0].pool.id(); + + let conn = shard.replica(&Request::default()).await.unwrap(); + assert_eq!(conn.pool.id(), replica_id); + } + + shard.shutdown(); + } + + #[tokio::test] + async fn test_include_primary() { + crate::logger(); + + let primary = &Some(PoolConfig { + address: Address::new_test(), + ..Default::default() + }); + + let replicas = &[PoolConfig { + address: Address::new_test(), + ..Default::default() + }]; + + let shard = Shard::new(ShardConfig { + number: 0, + primary, + replicas, + lb_strategy: LoadBalancingStrategy::Random, + rw_split: ReadWriteSplit::IncludePrimary, + identifier: Arc::new(User { + user: "pgdog".into(), + database: "pgdog".into(), + }), + lsn_check_interval: Duration::MAX, + }); + shard.launch(); + let mut ids = BTreeSet::new(); + + for _ in 0..25 { + let conn = shard.replica(&Request::default()).await.unwrap(); + ids.insert(conn.pool.id()); + } + + shard.shutdown(); + + assert_eq!(ids.len(), 2); + } +} diff --git a/pgdog/src/backend/pool/shard/monitor.rs b/pgdog/src/backend/pool/shard/monitor.rs new file mode 100644 index 000000000..99cbab83a --- /dev/null +++ b/pgdog/src/backend/pool/shard/monitor.rs @@ -0,0 +1,118 @@ +use super::*; + +use tokio::time::interval; +use tracing::{info, warn}; + +/// Shard communication primitives. +#[derive(Debug)] +pub(super) struct ShardComms { + pub(super) shutdown: Notify, + pub(super) lsn_check_interval: Duration, +} + +impl Default for ShardComms { + fn default() -> Self { + Self { + shutdown: Notify::new(), + lsn_check_interval: Duration::MAX, + } + } +} + +/// Monitor shard connection pools for various stats. +/// +/// Currently, only checking for replica lag, if any replicas are configured +/// and this is enabled. +pub(super) struct ShardMonitor { + shard: Shard, +} + +impl ShardMonitor { + /// Run the shard monitor. + pub(super) fn run(shard: &Shard) { + let monitor = Self { + shard: shard.clone(), + }; + + spawn(async move { monitor.spawn().await }); + } +} + +impl ShardMonitor { + async fn spawn(&self) { + let maintenance = (self.shard.comms().lsn_check_interval / 2) + .clamp(Duration::from_millis(333), Duration::MAX); + let mut maintenance = interval(maintenance); + + debug!( + "shard {} monitor running [{}]", + self.shard.number(), + self.shard.identifier() + ); + + let mut detector = RoleDetector::new(&self.shard); + + if detector.enabled() { + info!( + "automatic database role detection is enabled for shard {} [{}]", + self.shard.number(), + self.shard.identifier() + ); + } + + loop { + select! { + _ = maintenance.tick() => {}, + _ = self.shard.comms().shutdown.notified() => { + break; + }, + } + + if detector.changed() { + warn!( + "database role changed in shard {} [{}]", + self.shard.number(), + self.shard.identifier() + ); + } + + let pool_with_stats = self + .shard + .pools() + .iter() + .map(|pool| (pool.clone(), pool.lsn_stats())) + .collect::>(); + + let primary = pool_with_stats.iter().find(|pair| !pair.1.replica); + + // There is a primary. If not, replica lag cannot be + // calculated. + if let Some(primary) = primary { + let replicas = pool_with_stats.iter().filter(|pair| pair.1.replica); + for replica in replicas { + // Primary is ahead, there is replica lag. + let lag = if primary.1.lsn.lsn > replica.1.lsn.lsn { + // Assuming databases use the same timezone, + // since they are primary & replicas and database + // clocks are correctly synchronized. + let lag_ms = (primary.1.timestamp.to_naive_datetime() + - replica.1.timestamp.to_naive_datetime()) + .num_milliseconds() + .clamp(0, i64::MAX); + Duration::from_millis(lag_ms as u64) + } else { + Duration::ZERO + }; + replica.0.lock().replica_lag = lag; + } + primary.0.lock().replica_lag = Duration::ZERO; + } + } + + debug!( + "shard {} monitor shutdown [{}]", + self.shard.number(), + self.shard.identifier() + ); + } +} diff --git a/pgdog/src/backend/pool/shard/role_detector.rs b/pgdog/src/backend/pool/shard/role_detector.rs new file mode 100644 index 000000000..849fd57d6 --- /dev/null +++ b/pgdog/src/backend/pool/shard/role_detector.rs @@ -0,0 +1,182 @@ +use super::Shard; + +pub(super) struct RoleDetector { + enabled: bool, + shard: Shard, +} + +impl RoleDetector { + /// Create new role change detector. + pub(super) fn new(shard: &Shard) -> Self { + Self { + enabled: shard + .pools() + .iter() + .all(|pool| pool.config().role_detection), + shard: shard.clone(), + } + } + + /// Detect role change in the shard. + pub(super) fn changed(&mut self) -> bool { + if self.enabled() { + self.shard.redetect_roles() + } else { + false + } + } + + /// Role detector is enabled. + pub(super) fn enabled(&self) -> bool { + self.enabled + } +} + +#[cfg(test)] +mod test { + use std::sync::Arc; + use std::time::Duration; + + use tokio::time::Instant; + + use crate::backend::databases::User; + use crate::backend::pool::lsn_monitor::LsnStats; + use crate::backend::pool::{Address, Config, PoolConfig}; + use crate::backend::replication::publisher::Lsn; + use crate::config::{LoadBalancingStrategy, ReadWriteSplit}; + + use super::super::ShardConfig; + use super::*; + + fn create_test_pool_config(host: &str, port: u16, role_detection: bool) -> PoolConfig { + PoolConfig { + address: Address { + host: host.into(), + port, + user: "pgdog".into(), + password: "pgdog".into(), + database_name: "pgdog".into(), + ..Default::default() + }, + config: Config { + role_detection, + ..Default::default() + }, + } + } + + fn create_test_shard(primary: &Option, replicas: &[PoolConfig]) -> Shard { + Shard::new(ShardConfig { + number: 0, + primary, + replicas, + lb_strategy: LoadBalancingStrategy::Random, + rw_split: ReadWriteSplit::ExcludePrimary, + identifier: Arc::new(User { + user: "pgdog".into(), + database: "pgdog".into(), + }), + lsn_check_interval: Duration::MAX, + }) + } + + fn set_lsn_stats(shard: &Shard, index: usize, replica: bool, lsn: i64) { + let pools = shard.pools(); + let stats = LsnStats { + replica, + lsn: Lsn::from_i64(lsn), + offset_bytes: lsn, + fetched: Instant::now(), + ..Default::default() + }; + *pools[index].inner().lsn_stats.write() = stats; + } + + #[test] + fn test_changed_returns_false_when_lsn_stats_invalid() { + let primary = Some(create_test_pool_config("127.0.0.1", 5432, true)); + let replicas = [create_test_pool_config("localhost", 5432, true)]; + let shard = create_test_shard(&primary, &replicas); + + let mut detector = RoleDetector::new(&shard); + + assert!(detector.enabled()); + assert!(!detector.changed()); + } + + #[test] + fn test_changed_returns_false_when_roles_unchanged() { + let primary = Some(create_test_pool_config("127.0.0.1", 5432, true)); + let replicas = [create_test_pool_config("localhost", 5432, true)]; + let shard = create_test_shard(&primary, &replicas); + + set_lsn_stats(&shard, 0, true, 100); + set_lsn_stats(&shard, 1, false, 200); + + let mut detector = RoleDetector::new(&shard); + + assert!(detector.enabled()); + assert!(!detector.changed()); + } + + #[test] + fn test_changed_returns_true_on_failover() { + let primary = Some(create_test_pool_config("127.0.0.1", 5432, true)); + let replicas = [create_test_pool_config("localhost", 5432, true)]; + let shard = create_test_shard(&primary, &replicas); + + set_lsn_stats(&shard, 0, true, 100); + set_lsn_stats(&shard, 1, false, 200); + + let mut detector = RoleDetector::new(&shard); + + assert!(detector.enabled()); + assert!(!detector.changed()); + + set_lsn_stats(&shard, 0, false, 300); + set_lsn_stats(&shard, 1, true, 200); + + assert!(detector.changed()); + } + + #[test] + fn test_changed_returns_false_after_roles_stabilize() { + let primary = Some(create_test_pool_config("127.0.0.1", 5432, true)); + let replicas = [create_test_pool_config("localhost", 5432, true)]; + let shard = create_test_shard(&primary, &replicas); + + set_lsn_stats(&shard, 0, true, 100); + set_lsn_stats(&shard, 1, false, 200); + + let mut detector = RoleDetector::new(&shard); + assert!(detector.enabled()); + assert!(!detector.changed()); + + set_lsn_stats(&shard, 0, false, 300); + set_lsn_stats(&shard, 1, true, 200); + + assert!(detector.changed()); + + assert!(!detector.changed()); + } + + #[test] + fn test_disabled_when_not_all_roles_auto() { + let primary = Some(create_test_pool_config("127.0.0.1", 5432, false)); + let replicas = [create_test_pool_config("localhost", 5432, true)]; + let shard = create_test_shard(&primary, &replicas); + + set_lsn_stats(&shard, 0, true, 100); + set_lsn_stats(&shard, 1, false, 200); + + let mut detector = RoleDetector::new(&shard); + + assert!(!detector.enabled()); + assert!(!detector.changed()); + + set_lsn_stats(&shard, 0, false, 300); + set_lsn_stats(&shard, 1, true, 200); + + assert!(!detector.changed()); + } +} diff --git a/pgdog/src/backend/pool/state.rs b/pgdog/src/backend/pool/state.rs index e67520d78..d3e614600 100644 --- a/pgdog/src/backend/pool/state.rs +++ b/pgdog/src/backend/pool/state.rs @@ -3,7 +3,7 @@ use std::time::Duration; use crate::config::PoolerMode; use tokio::time::Instant; -use super::{inner::ReplicaLag, Ban, Config, Pool, Stats}; +use super::{Config, LsnStats, Pool, Stats}; /// Pool state. #[derive(Debug)] @@ -24,10 +24,6 @@ pub struct State { pub paused: bool, /// Number of clients waiting for a connection. pub waiting: usize, - /// Pool ban. - pub ban: Option, - /// Pool is banned. - pub banned: bool, /// Errors. pub errors: usize, /// Out of sync @@ -41,12 +37,17 @@ pub struct State { /// Pool mode pub pooler_mode: PoolerMode, /// Lag - pub replica_lag: ReplicaLag, + pub replica_lag: Duration, + /// Force closed. + pub force_close: usize, + /// LSN stats. + pub lsn_stats: LsnStats, } impl State { pub(super) fn get(pool: &Pool) -> Self { let now = Instant::now(); + let lsn_stats = pool.lsn_stats(); let guard = pool.lock(); State { @@ -58,8 +59,6 @@ impl State { config: guard.config, paused: guard.paused, waiting: guard.waiting.len(), - ban: guard.ban, - banned: guard.ban.is_some(), errors: guard.errors, out_of_sync: guard.out_of_sync, re_synced: guard.re_synced, @@ -72,6 +71,8 @@ impl State { .unwrap_or(Duration::ZERO), pooler_mode: guard.config().pooler_mode, replica_lag: guard.replica_lag, + force_close: guard.force_close, + lsn_stats, } } } diff --git a/pgdog/src/backend/pool/stats.rs b/pgdog/src/backend/pool/stats.rs index 529a378f7..da63a0417 100644 --- a/pgdog/src/backend/pool/stats.rs +++ b/pgdog/src/backend/pool/stats.rs @@ -1,6 +1,11 @@ -//! Pool stats. +//! +//! Pool statistics. +//! +//! Used in SHOW POOLS admin database command +//! and in Prometheus metrics. +//! -use crate::backend::stats::Counts as BackendCounts; +use crate::{backend::stats::Counts as BackendCounts, config::Memory, net::MessageBufferStats}; use std::{ iter::Sum, @@ -8,22 +13,54 @@ use std::{ time::Duration, }; +/// Pool statistics. +/// +/// These are updated after each connection check-in. +/// #[derive(Debug, Clone, Default, Copy)] pub struct Counts { + /// Number of committed transactions. pub xact_count: usize, + /// Number of transactions committed with 2-phase commit. pub xact_2pc_count: usize, + /// Number of executed queries. pub query_count: usize, + /// How many times a server has been given to a client. + /// In transaction mode, this equals to `xact_count`. pub server_assignment_count: usize, + /// Number of bytes received by server connections. pub received: usize, + /// Number of bytes sent to server connections. pub sent: usize, + /// Total duration of all transactions. pub xact_time: Duration, + /// Total time spent idling inside transactions. + pub idle_xact_time: Duration, + /// Total time spent executing queries. pub query_time: Duration, + /// Total time clients spent waiting for a connection from the pool. pub wait_time: Duration, + /// Total count of Parse messages sent to server connections. pub parse_count: usize, + /// Total count of Bind messages sent to server connections. pub bind_count: usize, + /// Number of times the pool had to rollback unfinished transactions. pub rollbacks: usize, + /// Number of times the pool sent the health check query. pub healthchecks: usize, + /// Total count of Close messages sent to server connections. pub close: usize, + /// Total number of network-related errors detected on server connections. + pub errors: usize, + /// Total number of server connections that were cleaned after a dirty session. + pub cleaned: usize, + /// Total number of times servers had to synchronize prepared statements from Postgres' + /// pg_prepared_statements view. + pub prepared_sync: usize, + /// Total time spent creating server connections. + pub connect_time: Duration, + /// Total number of times the pool attempted to create server connections. + pub connect_count: usize, } impl Sub for Counts { @@ -40,13 +77,19 @@ impl Sub for Counts { received: self.received.saturating_sub(rhs.received), sent: self.sent.saturating_sub(rhs.sent), xact_time: self.xact_time.saturating_sub(rhs.xact_time), + idle_xact_time: self.idle_xact_time.saturating_sub(rhs.idle_xact_time), query_time: self.query_time.saturating_sub(rhs.query_time), wait_time: self.wait_time.saturating_sub(rhs.wait_time), parse_count: self.parse_count.saturating_sub(rhs.parse_count), - bind_count: self.parse_count.saturating_sub(rhs.bind_count), + bind_count: self.bind_count.saturating_sub(rhs.bind_count), rollbacks: self.rollbacks.saturating_sub(rhs.rollbacks), - healthchecks: self.healthchecks.saturating_add(rhs.healthchecks), - close: self.close.saturating_add(rhs.close), + healthchecks: self.healthchecks.saturating_sub(rhs.healthchecks), + close: self.close.saturating_sub(rhs.close), + errors: self.errors.saturating_sub(rhs.errors), + cleaned: self.cleaned.saturating_sub(rhs.cleaned), + prepared_sync: self.prepared_sync.saturating_sub(rhs.prepared_sync), + connect_time: self.connect_time.saturating_sub(rhs.connect_time), + connect_count: self.connect_count.saturating_sub(rhs.connect_count), } } } @@ -56,20 +99,32 @@ impl Div for Counts { fn div(self, rhs: usize) -> Self::Output { Self { - xact_count: self.xact_count.saturating_div(rhs), - xact_2pc_count: self.xact_2pc_count.saturating_div(rhs), - query_count: self.query_count.saturating_div(rhs), - server_assignment_count: self.server_assignment_count.saturating_div(rhs), - received: self.received.saturating_div(rhs), - sent: self.sent.saturating_div(rhs), + xact_count: self.xact_count.checked_div(rhs).unwrap_or(0), + xact_2pc_count: self.xact_2pc_count.checked_div(rhs).unwrap_or(0), + query_count: self.query_count.checked_div(rhs).unwrap_or(0), + server_assignment_count: self.server_assignment_count.checked_div(rhs).unwrap_or(0), + received: self.received.checked_div(rhs).unwrap_or(0), + sent: self.sent.checked_div(rhs).unwrap_or(0), xact_time: self.xact_time.checked_div(rhs as u32).unwrap_or_default(), + idle_xact_time: self + .idle_xact_time + .checked_div(rhs as u32) + .unwrap_or_default(), query_time: self.query_time.checked_div(rhs as u32).unwrap_or_default(), wait_time: self.wait_time.checked_div(rhs as u32).unwrap_or_default(), - parse_count: self.parse_count.saturating_div(rhs), - bind_count: self.parse_count.saturating_div(rhs), - rollbacks: self.rollbacks.saturating_div(rhs), - healthchecks: self.healthchecks.saturating_div(rhs), - close: self.close.saturating_div(rhs), + parse_count: self.parse_count.checked_div(rhs).unwrap_or(0), + bind_count: self.bind_count.checked_div(rhs).unwrap_or(0), + rollbacks: self.rollbacks.checked_div(rhs).unwrap_or(0), + healthchecks: self.healthchecks.checked_div(rhs).unwrap_or(0), + close: self.close.checked_div(rhs).unwrap_or(0), + errors: self.errors.checked_div(rhs).unwrap_or(0), + cleaned: self.cleaned.checked_div(rhs).unwrap_or(0), + prepared_sync: self.prepared_sync.checked_div(rhs).unwrap_or(0), + connect_time: self + .connect_time + .checked_div(rhs as u32) + .unwrap_or_default(), + connect_count: self.connect_count.checked_div(rhs).unwrap_or(0), } } } @@ -87,12 +142,18 @@ impl Add for Counts { sent: self.sent + rhs.bytes_sent, query_time: self.query_time + rhs.query_time, xact_time: self.xact_time + rhs.transaction_time, + idle_xact_time: self.idle_xact_time + rhs.idle_in_transaction_time, wait_time: self.wait_time, parse_count: self.parse_count + rhs.parse, bind_count: self.bind_count + rhs.bind, rollbacks: self.rollbacks + rhs.rollbacks, healthchecks: self.healthchecks + rhs.healthchecks, close: self.close + rhs.close, + errors: self.errors + rhs.errors, + cleaned: self.cleaned + rhs.cleaned, + prepared_sync: self.prepared_sync + rhs.prepared_sync, + connect_count: self.connect_count, + connect_time: self.connect_time, } } } @@ -122,13 +183,19 @@ impl Add for Counts { received: self.received.saturating_add(rhs.received), sent: self.sent.saturating_add(rhs.sent), xact_time: self.xact_time.saturating_add(rhs.xact_time), + idle_xact_time: self.idle_xact_time.saturating_add(rhs.idle_xact_time), query_time: self.query_time.saturating_add(rhs.query_time), wait_time: self.wait_time.saturating_add(rhs.wait_time), parse_count: self.parse_count.saturating_add(rhs.parse_count), - bind_count: self.parse_count.saturating_add(rhs.bind_count), + bind_count: self.bind_count.saturating_add(rhs.bind_count), rollbacks: self.rollbacks.saturating_add(rhs.rollbacks), healthchecks: self.healthchecks.saturating_add(rhs.healthchecks), close: self.close.saturating_add(rhs.close), + errors: self.errors.saturating_add(rhs.errors), + cleaned: self.cleaned.saturating_add(rhs.cleaned), + prepared_sync: self.prepared_sync.saturating_add(rhs.prepared_sync), + connect_count: self.connect_count.saturating_add(rhs.connect_count), + connect_time: self.connect_time.saturating_add(rhs.connect_time), } } } @@ -137,6 +204,7 @@ impl Add for Counts { pub struct Stats { // Total counts. pub counts: Counts, + /// Counts since last average calculation. last_counts: Counts, // Average counts. pub averages: Counts, @@ -147,8 +215,419 @@ impl Stats { pub fn calc_averages(&mut self, time: Duration) { let secs = time.as_secs() as usize; if secs > 0 { - self.averages = (self.counts - self.last_counts) / secs; + let diff = self.counts - self.last_counts; + self.averages = diff / secs; + self.averages.query_time = diff.query_time + / diff + .query_count + .try_into() + .unwrap_or(u32::MAX) + .clamp(1, u32::MAX); + self.averages.xact_time = diff.xact_time + / diff + .xact_count + .try_into() + .unwrap_or(u32::MAX) + .clamp(1, u32::MAX); + self.averages.wait_time = diff.wait_time + / diff + .server_assignment_count + .try_into() + .unwrap_or(u32::MAX) + .clamp(1, u32::MAX); + self.averages.connect_time = diff.connect_time + / diff + .connect_count + .try_into() + .unwrap_or(u32::MAX) + .clamp(1, u32::MAX); + let queries_in_xact = diff + .query_count + .wrapping_sub(diff.xact_count) + .clamp(1, u32::MAX as usize); + self.averages.idle_xact_time = + diff.idle_xact_time / queries_in_xact.try_into().unwrap_or(u32::MAX); + self.last_counts = self.counts; } } } + +/// Statistics calculated for the network buffer used +/// by clients and servers. +#[derive(Debug, Clone, Default, Copy)] +pub struct MemoryStats { + /// Message buffer stats. + pub buffer: MessageBufferStats, + /// Memory used by prepared statements. + pub prepared_statements: usize, + /// Memory used by the network stream buffer. + pub stream: usize, +} + +impl MemoryStats { + /// Create new memory stats tracker. + pub fn new(config: &Memory) -> Self { + Self { + buffer: MessageBufferStats { + bytes_alloc: config.message_buffer, + ..Default::default() + }, + prepared_statements: 0, + stream: config.net_buffer, + } + } + + /// Calculate total memory usage. + pub fn total(&self) -> usize { + self.buffer.bytes_alloc + self.prepared_statements + self.stream + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_add_trait() { + let a = Counts { + xact_count: 10, + xact_2pc_count: 5, + query_count: 20, + server_assignment_count: 3, + received: 1000, + sent: 2000, + xact_time: Duration::from_secs(5), + idle_xact_time: Duration::from_secs(3), + query_time: Duration::from_secs(3), + wait_time: Duration::from_secs(2), + parse_count: 15, + bind_count: 18, + rollbacks: 2, + healthchecks: 1, + close: 4, + errors: 3, + cleaned: 6, + prepared_sync: 7, + connect_time: Duration::from_secs(1), + connect_count: 8, + }; + + let b = Counts { + xact_count: 5, + xact_2pc_count: 3, + query_count: 10, + server_assignment_count: 2, + received: 500, + sent: 1000, + xact_time: Duration::from_secs(2), + idle_xact_time: Duration::from_secs(5), + query_time: Duration::from_secs(1), + wait_time: Duration::from_secs(1), + parse_count: 8, + bind_count: 9, + rollbacks: 1, + healthchecks: 2, + close: 3, + errors: 1, + cleaned: 2, + prepared_sync: 3, + connect_time: Duration::from_secs(2), + connect_count: 4, + }; + + let result = a + b; + + assert_eq!(result.xact_count, 15); + assert_eq!(result.xact_2pc_count, 8); + assert_eq!(result.query_count, 30); + assert_eq!(result.server_assignment_count, 5); + assert_eq!(result.received, 1500); + assert_eq!(result.sent, 3000); + assert_eq!(result.xact_time, Duration::from_secs(7)); + assert_eq!(result.idle_xact_time, Duration::from_secs(8)); + assert_eq!(result.query_time, Duration::from_secs(4)); + assert_eq!(result.wait_time, Duration::from_secs(3)); + assert_eq!(result.parse_count, 23); + assert_eq!(result.bind_count, 27); + assert_eq!(result.rollbacks, 3); + assert_eq!(result.healthchecks, 3); + assert_eq!(result.close, 7); + assert_eq!(result.errors, 4); + assert_eq!(result.cleaned, 8); + assert_eq!(result.prepared_sync, 10); + assert_eq!(result.connect_time, Duration::from_secs(3)); + assert_eq!(result.connect_count, 12); + } + + #[test] + fn test_sub_trait() { + let a = Counts { + xact_count: 10, + xact_2pc_count: 5, + query_count: 20, + server_assignment_count: 3, + received: 1000, + sent: 2000, + xact_time: Duration::from_secs(5), + idle_xact_time: Duration::from_secs(3), + query_time: Duration::from_secs(3), + wait_time: Duration::from_secs(2), + parse_count: 15, + bind_count: 18, + rollbacks: 2, + healthchecks: 4, + close: 5, + errors: 3, + cleaned: 6, + prepared_sync: 7, + connect_time: Duration::from_secs(3), + connect_count: 8, + }; + + let b = Counts { + xact_count: 5, + xact_2pc_count: 3, + query_count: 10, + server_assignment_count: 2, + received: 500, + sent: 1000, + xact_time: Duration::from_secs(2), + idle_xact_time: Duration::from_secs(2), + query_time: Duration::from_secs(1), + wait_time: Duration::from_secs(1), + parse_count: 8, + bind_count: 9, + rollbacks: 1, + healthchecks: 2, + close: 3, + errors: 1, + cleaned: 2, + prepared_sync: 3, + connect_time: Duration::from_secs(1), + connect_count: 4, + }; + + let result = a - b; + + assert_eq!(result.xact_count, 5); + assert_eq!(result.xact_2pc_count, 2); + assert_eq!(result.query_count, 10); + assert_eq!(result.server_assignment_count, 1); + assert_eq!(result.received, 500); + assert_eq!(result.sent, 1000); + assert_eq!(result.xact_time, Duration::from_secs(3)); + assert_eq!(result.idle_xact_time, Duration::from_secs(1)); + assert_eq!(result.query_time, Duration::from_secs(2)); + assert_eq!(result.wait_time, Duration::from_secs(1)); + assert_eq!(result.parse_count, 7); + assert_eq!(result.bind_count, 9); + assert_eq!(result.rollbacks, 1); + assert_eq!(result.healthchecks, 2); + assert_eq!(result.close, 2); + assert_eq!(result.errors, 2); + assert_eq!(result.cleaned, 4); + assert_eq!(result.prepared_sync, 4); + assert_eq!(result.connect_time, Duration::from_secs(2)); + assert_eq!(result.connect_count, 4); + } + + #[test] + fn test_sub_trait_saturating() { + let a = Counts { + xact_count: 5, + bind_count: 3, + ..Default::default() + }; + + let b = Counts { + xact_count: 10, + bind_count: 5, + ..Default::default() + }; + + let result = a - b; + + assert_eq!(result.xact_count, 0); + assert_eq!(result.bind_count, 0); + } + + #[test] + fn test_div_trait() { + let a = Counts { + xact_count: 10, + xact_2pc_count: 6, + query_count: 20, + server_assignment_count: 4, + received: 1000, + sent: 2000, + xact_time: Duration::from_secs(10), + idle_xact_time: Duration::from_secs(20), + query_time: Duration::from_secs(6), + wait_time: Duration::from_secs(4), + parse_count: 15, + bind_count: 18, + rollbacks: 2, + healthchecks: 8, + close: 12, + errors: 3, + cleaned: 6, + prepared_sync: 9, + connect_time: Duration::from_secs(8), + connect_count: 4, + }; + + let result = a / 2; + + assert_eq!(result.xact_count, 5); + assert_eq!(result.xact_2pc_count, 3); + assert_eq!(result.query_count, 10); + assert_eq!(result.server_assignment_count, 2); + assert_eq!(result.received, 500); + assert_eq!(result.sent, 1000); + assert_eq!(result.xact_time, Duration::from_secs(5)); + assert_eq!(result.idle_xact_time, Duration::from_secs(10)); + assert_eq!(result.query_time, Duration::from_secs(3)); + assert_eq!(result.wait_time, Duration::from_secs(2)); + assert_eq!(result.parse_count, 7); + assert_eq!(result.bind_count, 9); + assert_eq!(result.rollbacks, 1); + assert_eq!(result.healthchecks, 4); + assert_eq!(result.close, 6); + assert_eq!(result.errors, 1); + assert_eq!(result.cleaned, 3); + assert_eq!(result.prepared_sync, 4); + assert_eq!(result.connect_time, Duration::from_secs(4)); + assert_eq!(result.connect_count, 2); + } + + #[test] + fn test_div_by_zero() { + let a = Counts { + xact_count: 10, + xact_time: Duration::from_secs(10), + ..Default::default() + }; + + let result = a / 0; + + assert_eq!(result.xact_count, 0); + assert_eq!(result.xact_time, Duration::ZERO); + } + + #[test] + fn test_add_backend_counts() { + let pool_counts = Counts { + xact_count: 10, + xact_2pc_count: 5, + query_count: 20, + server_assignment_count: 3, + received: 1000, + sent: 2000, + xact_time: Duration::from_secs(5), + idle_xact_time: Duration::from_secs(10), + query_time: Duration::from_secs(3), + wait_time: Duration::from_secs(2), + parse_count: 15, + bind_count: 18, + rollbacks: 2, + healthchecks: 1, + close: 4, + errors: 3, + cleaned: 6, + prepared_sync: 7, + connect_time: Duration::from_secs(1), + connect_count: 8, + }; + + let backend_counts = BackendCounts { + bytes_sent: 500, + bytes_received: 300, + transactions: 5, + transactions_2pc: 2, + queries: 10, + rollbacks: 1, + errors: 2, + prepared_statements: 0, + query_time: Duration::from_secs(2), + transaction_time: Duration::from_secs(3), + idle_in_transaction_time: Duration::from_secs(5), + parse: 7, + bind: 8, + healthchecks: 3, + close: 2, + cleaned: 4, + prepared_sync: 5, + }; + + let result = pool_counts + backend_counts; + + assert_eq!(result.xact_count, 15); + assert_eq!(result.xact_2pc_count, 7); + assert_eq!(result.query_count, 30); + assert_eq!(result.server_assignment_count, 3); + assert_eq!(result.received, 1300); + assert_eq!(result.sent, 2500); + assert_eq!(result.xact_time, Duration::from_secs(8)); + assert_eq!(result.idle_xact_time, Duration::from_secs(15)); + assert_eq!(result.query_time, Duration::from_secs(5)); + assert_eq!(result.wait_time, Duration::from_secs(2)); + assert_eq!(result.parse_count, 22); + assert_eq!(result.bind_count, 26); + assert_eq!(result.rollbacks, 3); + assert_eq!(result.healthchecks, 4); + assert_eq!(result.close, 6); + assert_eq!(result.errors, 5); + assert_eq!(result.cleaned, 10); + assert_eq!(result.prepared_sync, 12); + assert_eq!(result.connect_count, 8); + assert_eq!(result.connect_time, Duration::from_secs(1)); + } + + #[test] + fn test_calc_averages() { + let mut stats = Stats::default(); + + stats.counts.query_count = 10; + stats.counts.query_time = Duration::from_millis(500); + stats.counts.xact_count = 5; + stats.counts.xact_time = Duration::from_millis(1000); + stats.counts.server_assignment_count = 4; + stats.counts.wait_time = Duration::from_millis(200); + stats.counts.connect_count = 2; + stats.counts.connect_time = Duration::from_millis(100); + + stats.counts.idle_xact_time = Duration::from_millis(250); + + stats.calc_averages(Duration::from_secs(1)); + + assert_eq!(stats.averages.query_time, Duration::from_millis(50)); + assert_eq!(stats.averages.xact_time, Duration::from_millis(200)); + assert_eq!(stats.averages.wait_time, Duration::from_millis(50)); + assert_eq!(stats.averages.connect_time, Duration::from_millis(50)); + // idle_xact_time is divided by (query_count - xact_count) = 10 - 5 = 5 + assert_eq!(stats.averages.idle_xact_time, Duration::from_millis(50)); + } + + #[test] + fn test_calc_averages_division_by_zero() { + let mut stats = Stats::default(); + + // Set some time values but leave counts at zero + stats.counts.query_time = Duration::from_millis(100); + stats.counts.xact_time = Duration::from_millis(200); + stats.counts.wait_time = Duration::from_millis(50); + stats.counts.connect_time = Duration::from_millis(25); + stats.counts.idle_xact_time = Duration::from_millis(75); + + // Should not panic - clamp(1, u32::MAX) prevents division by zero + stats.calc_averages(Duration::from_secs(1)); + + // When counts are zero, times are divided by 1 (the clamped minimum) + assert_eq!(stats.averages.query_time, Duration::from_millis(100)); + assert_eq!(stats.averages.xact_time, Duration::from_millis(200)); + assert_eq!(stats.averages.wait_time, Duration::from_millis(50)); + assert_eq!(stats.averages.connect_time, Duration::from_millis(25)); + assert_eq!(stats.averages.idle_xact_time, Duration::from_millis(75)); + } +} diff --git a/pgdog/src/backend/pool/taken.rs b/pgdog/src/backend/pool/taken.rs index 0d13e48c9..62c66740b 100644 --- a/pgdog/src/backend/pool/taken.rs +++ b/pgdog/src/backend/pool/taken.rs @@ -2,32 +2,50 @@ use fnv::FnvHashMap as HashMap; use crate::net::BackendKeyData; -use super::Mapping; +use super::{Error, Mapping}; #[derive(Default, Clone, Debug)] pub(super) struct Taken { + /// Guaranteed to be unique per client/server connection. + taken: HashMap, + /// Guaranteed to be unique because servers can only be mapped + /// to one client at a time. + server_client: HashMap, + /// Not unique, but will contain the server that's actively executing a query + /// for that client. client_server: HashMap, - server_client: HashMap, + /// Counter that guarantees uniqueness. Wraparound happens after a gazillion billion transactions. + counter: usize, } impl Taken { #[inline] - pub(super) fn take(&mut self, mapping: &Mapping) { + pub(super) fn take(&mut self, mapping: &Mapping) -> Result<(), Error> { + self.taken.insert(self.counter, *mapping); + self.server_client.insert(mapping.server, self.counter); self.client_server.insert(mapping.client, mapping.server); - self.server_client.insert(mapping.server, mapping.client); + self.counter = self.counter.wrapping_add(1); + Ok(()) } #[inline] - pub(super) fn check_in(&mut self, server: &BackendKeyData) { - let client = self.server_client.remove(server); - if let Some(client) = client { - self.client_server.remove(&client); - } + pub(super) fn check_in(&mut self, server: &BackendKeyData) -> Result<(), Error> { + let counter = self + .server_client + .remove(server) + .ok_or(Error::UntrackedConnCheckin(*server))?; + let mapping = self + .taken + .remove(&counter) + .ok_or(Error::MappingMissing(counter))?; + self.client_server.remove(&mapping.client); + + Ok(()) } #[inline] pub(super) fn len(&self) -> usize { - self.client_server.len() // Both should always be the same length. + self.taken.len() } #[allow(dead_code)] @@ -37,17 +55,11 @@ impl Taken { #[inline] pub(super) fn server(&self, client: &BackendKeyData) -> Option { - self.client_server.get(client).cloned() - } - - #[allow(dead_code)] - pub(super) fn client(&self, server: &BackendKeyData) -> Option { - self.server_client.get(server).cloned() + self.client_server.get(client).copied() } #[cfg(test)] pub(super) fn clear(&mut self) { - self.client_server.clear(); - self.server_client.clear(); + self.taken.clear(); } } diff --git a/pgdog/src/backend/pool/test/mod.rs b/pgdog/src/backend/pool/test/mod.rs index 07d4f2215..61f48db13 100644 --- a/pgdog/src/backend/pool/test/mod.rs +++ b/pgdog/src/backend/pool/test/mod.rs @@ -2,12 +2,12 @@ use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; -use std::time::{Duration, Instant}; +use std::time::Duration; use rand::Rng; use tokio::spawn; use tokio::task::yield_now; -use tokio::time::{sleep, timeout}; +use tokio::time::{sleep, timeout, Instant}; use tokio_util::task::TaskTracker; use crate::net::ProtocolMessage; @@ -16,8 +16,6 @@ use crate::state::State; use super::*; -mod replica; - pub fn pool() -> Pool { let config = Config { max: 1, @@ -78,7 +76,7 @@ async fn test_pool_checkout() { assert_eq!(pool.lock().idle(), 0); assert_eq!(pool.lock().total(), 1); - assert!(!pool.lock().should_create()); + assert_eq!(pool.lock().should_create(), inner::ShouldCreate::No); let err = timeout(Duration::from_millis(100), pool.get(&Request::default())).await; @@ -100,7 +98,7 @@ async fn test_concurrency() { let pool = pool.clone(); tracker.spawn(async move { let _conn = pool.get(&Request::default()).await.unwrap(); - let duration = rand::thread_rng().gen_range(0..10); + let duration = rand::rng().random_range(0..10); sleep(Duration::from_millis(duration)).await; }); } @@ -133,7 +131,7 @@ async fn test_concurrency_with_gas() { let pool = pool.clone(); tracker.spawn(async move { let _conn = pool.get(&Request::default()).await.unwrap(); - let duration = rand::thread_rng().gen_range(0..10); + let duration = rand::rng().random_range(0..10); assert!(pool.lock().checked_out() > 0); assert!(pool.lock().total() <= 10); sleep(Duration::from_millis(duration)).await; @@ -146,21 +144,6 @@ async fn test_concurrency_with_gas() { assert_eq!(pool.lock().total(), 10); } -#[tokio::test] -async fn test_bans() { - let pool = pool(); - let mut config = *pool.lock().config(); - config.checkout_timeout = Duration::from_millis(100); - pool.update_config(config); - - pool.ban(Error::CheckoutTimeout); - assert!(pool.banned()); - - // Will timeout getting a connection from a banned pool. - let conn = pool.get(&Request::default()).await; - assert!(conn.is_err()); -} - #[tokio::test] async fn test_offline() { let pool = pool(); @@ -168,7 +151,6 @@ async fn test_offline() { pool.shutdown(); assert!(!pool.lock().online); - assert!(!pool.banned()); // Cannot get a connection from the pool. let err = pool.get(&Request::default()).await; @@ -190,7 +172,6 @@ async fn test_pause() { pool.get(&Request::default()) .await .expect_err("checkout timeout"); - pool.maybe_unban(); drop(hold); // Make sure we're not blocked still. drop(pool.get(&Request::default()).await.unwrap()); @@ -285,6 +266,7 @@ async fn test_incomplete_request_recovery() { for query in ["SELECT 1", "BEGIN"] { let mut conn = pool.get(&Request::default()).await.unwrap(); + let conn_id = *(conn.id()); conn.send(&vec![ProtocolMessage::from(Query::new(query))].into()) .await @@ -301,6 +283,10 @@ async fn test_incomplete_request_recovery() { } else { assert_eq!(state.stats.counts.rollbacks, 0); } + + // Verify the same connection is reused + let conn = pool.get(&Request::default()).await.unwrap(); + assert_eq!(conn.id(), &conn_id); } } @@ -315,6 +301,48 @@ async fn test_force_close() { assert_eq!(pool.lock().force_close, 1); } +#[tokio::test] +async fn test_server_force_close_discards_connection() { + crate::logger(); + + let config = Config { + max: 1, + min: 0, + ..Default::default() + }; + + let pool = Pool::new(&PoolConfig { + address: Address { + host: "127.0.0.1".into(), + port: 5432, + database_name: "pgdog".into(), + user: "pgdog".into(), + password: "pgdog".into(), + ..Default::default() + }, + config, + }); + pool.launch(); + + let mut conn = pool.get(&Request::default()).await.unwrap(); + conn.execute("BEGIN").await.unwrap(); + assert!(conn.in_transaction()); + + assert_eq!(pool.lock().total(), 1); + assert_eq!(pool.lock().idle(), 0); + assert_eq!(pool.lock().checked_out(), 1); + + conn.stats_mut().state(State::ForceClose); + drop(conn); + + sleep(Duration::from_millis(100)).await; + + let state = pool.state(); + assert_eq!(state.force_close, 1); + assert_eq!(state.idle, 0); + assert_eq!(state.total, 0); +} + #[tokio::test] async fn test_query_stats() { let pool = pool(); @@ -439,3 +467,189 @@ async fn test_prepared_statements_limit() { assert_eq!(guard.prepared_statements_mut().len(), 100); assert_eq!(guard.stats().total.prepared_statements, 100); // stats are accurate. } + +#[tokio::test] +async fn test_idle_healthcheck_loop() { + crate::logger(); + + let config = Config { + max: 1, + min: 1, + idle_healthcheck_interval: Duration::from_millis(100), + idle_healthcheck_delay: Duration::from_millis(10), + ..Default::default() + }; + + let pool = Pool::new(&PoolConfig { + address: Address { + host: "127.0.0.1".into(), + port: 5432, + database_name: "pgdog".into(), + user: "pgdog".into(), + password: "pgdog".into(), + ..Default::default() + }, + config, + }); + pool.launch(); + + let initial_healthchecks = pool.state().stats.counts.healthchecks; + + sleep(Duration::from_millis(350)).await; + + let after_healthchecks = pool.state().stats.counts.healthchecks; + + assert!( + after_healthchecks > initial_healthchecks, + "Expected healthchecks to increase from {} but got {}", + initial_healthchecks, + after_healthchecks + ); + assert!( + after_healthchecks >= initial_healthchecks + 2, + "Expected at least 2 healthchecks to run in 350ms with 100ms interval, got {} (increase of {})", + after_healthchecks, + after_healthchecks - initial_healthchecks + ); +} + +#[tokio::test] +async fn test_checkout_timeout() { + crate::logger(); + + let config = Config { + max: 1, + min: 1, + checkout_timeout: Duration::from_millis(100), + ..Default::default() + }; + + let pool = Pool::new(&PoolConfig { + address: Address::new_test(), + config, + }); + pool.launch(); + + // Hold the only connection + let _conn = pool.get(&Request::default()).await.unwrap(); + + // Try to get another connection - should timeout + let result = pool.get(&Request::default()).await; + assert!(result.is_err()); + assert_eq!(result.unwrap_err(), Error::CheckoutTimeout); + assert!(pool.lock().waiting.is_empty()); +} + +#[tokio::test] +async fn test_move_conns_to() { + crate::logger(); + + let config = Config { + max: 3, + min: 0, + ..Default::default() + }; + + let source = Pool::new(&PoolConfig { + address: Address { + host: "127.0.0.1".into(), + port: 5432, + database_name: "pgdog".into(), + user: "pgdog".into(), + password: "pgdog".into(), + ..Default::default() + }, + config, + }); + source.launch(); + + let destination = Pool::new(&PoolConfig { + address: Address { + host: "127.0.0.1".into(), + port: 5432, + database_name: "pgdog".into(), + user: "pgdog".into(), + password: "pgdog".into(), + ..Default::default() + }, + config, + }); + + let conn1 = source.get(&Request::default()).await.unwrap(); + let conn2 = source.get(&Request::default()).await.unwrap(); + + drop(conn1); + + sleep(Duration::from_millis(50)).await; + + assert_eq!(source.lock().idle(), 1); + assert_eq!(source.lock().checked_out(), 1); + assert_eq!(source.lock().total(), 2); + assert_eq!(destination.lock().total(), 0); + assert!(!destination.lock().online); + + source.move_conns_to(&destination).unwrap(); + + assert!(!source.lock().online); + assert!(destination.lock().online); + assert_eq!(destination.lock().total(), 2); + assert_eq!(source.lock().total(), 0); + let new_pool_id = destination.id(); + for conn in destination.lock().idle_conns() { + assert_eq!(conn.stats().pool_id, new_pool_id); + } + + drop(conn2); + + sleep(Duration::from_millis(50)).await; + + assert_eq!(destination.lock().idle(), 2); + assert_eq!(destination.lock().checked_out(), 0); +} + +#[tokio::test] +async fn test_lsn_monitor() { + crate::logger(); + + let config = Config { + max: 1, + min: 1, + lsn_check_delay: Duration::from_millis(10), + lsn_check_interval: Duration::from_millis(50), + lsn_check_timeout: Duration::from_millis(5_000), + ..Default::default() + }; + + let pool = Pool::new(&PoolConfig { + address: Address::new_test(), + config, + }); + + let initial_stats = pool.lsn_stats(); + assert!(!initial_stats.valid()); + + pool.launch(); + + sleep(Duration::from_millis(200)).await; + + let stats = pool.lsn_stats(); + assert!( + stats.valid(), + "LSN stats should be valid after monitor runs" + ); + assert!(!stats.replica, "Local PostgreSQL should not be a replica"); + assert!(stats.lsn.lsn > 0, "LSN should be greater than 0"); + assert!( + stats.offset_bytes > 0, + "Offset bytes should be greater than 0" + ); + + let age = stats.lsn_age(Instant::now()); + assert!( + age < Duration::from_millis(300), + "LSN stats age should be recent, got {:?}", + age + ); + + pool.shutdown(); +} diff --git a/pgdog/src/backend/pool/test/replica.rs b/pgdog/src/backend/pool/test/replica.rs deleted file mode 100644 index 732a11193..000000000 --- a/pgdog/src/backend/pool/test/replica.rs +++ /dev/null @@ -1,65 +0,0 @@ -use crate::config::LoadBalancingStrategy; - -// use super::pool; -use super::*; -use tokio::spawn; - -fn replicas() -> Replicas { - let one = PoolConfig { - address: Address { - host: "127.0.0.1".into(), - port: 5432, - user: "pgdog".into(), - password: "pgdog".into(), - database_name: "pgdog".into(), - ..Default::default() - }, - config: Config { - max: 1, - checkout_timeout: Duration::from_millis(1000), - ..Default::default() - }, - }; - let mut two = one.clone(); - two.address.host = "localhost".into(); - let replicas = Replicas::new(&[one, two], LoadBalancingStrategy::Random); - replicas.pools().iter().for_each(|p| p.launch()); - replicas -} - -#[tokio::test] -async fn test_replicas() { - let replicas = replicas(); - - for pool in 0..2 { - let mut tasks = vec![]; - replicas.pools[pool].ban(Error::CheckoutTimeout); - - for _ in 0..10000 { - let replicas = replicas.clone(); - let other = if pool == 0 { 1 } else { 0 }; - tasks.push(spawn(async move { - assert!(replicas.pools[pool].banned()); - assert!(!replicas.pools[other].banned()); - let conn = replicas.get(&Request::default(), &None).await.unwrap(); - assert_eq!(conn.addr(), replicas.pools[other].addr()); - assert!(replicas.pools[pool].banned()); - assert!(!replicas.pools[other].banned()); - })); - } - - for task in tasks { - task.await.unwrap(); - } - - replicas.pools[pool].maybe_unban(); - } - - replicas.pools[0].ban(Error::CheckoutTimeout); - replicas.pools[1].ban(Error::CheckoutTimeout); - - // All replicas banned, unban everyone. - assert!(replicas.pools.iter().all(|pool| pool.banned())); - replicas.get(&Request::default(), &None).await.unwrap(); - assert!(replicas.pools.iter().all(|pool| !pool.banned())); -} diff --git a/pgdog/src/backend/pool/waiting.rs b/pgdog/src/backend/pool/waiting.rs index e3477f69c..7417f5951 100644 --- a/pgdog/src/backend/pool/waiting.rs +++ b/pgdog/src/backend/pool/waiting.rs @@ -1,60 +1,77 @@ use crate::backend::Server; use super::{Error, Guard, Pool, Request}; -use tokio::{ - sync::oneshot::*, - time::{timeout, Instant}, -}; +use tokio::{sync::oneshot::*, time::Instant}; pub(super) struct Waiting { pool: Pool, - rx: Receiver, Error>>, + rx: Option, Error>>>, request: Request, + waiting: bool, +} + +impl Drop for Waiting { + fn drop(&mut self) { + if self.waiting { + self.pool.lock().remove_waiter(&self.request.id); + } + } } impl Waiting { + /// Create new waiter. + /// + /// N.B. You must call and await `Waiting::wait`, otherwise you'll leak waiters. + /// pub(super) fn new(pool: Pool, request: &Request) -> Result { let request = *request; let (tx, rx) = channel(); - { + let full = { let mut guard = pool.lock(); if !guard.online { return Err(Error::Offline); } - guard.waiting.push_back(Waiter { request, tx }) - } + guard.waiting.push_back(Waiter { request, tx }); + guard.full() + }; // Tell maintenance we are in line waiting for a connection. - pool.comms().request.notify_one(); + if !full { + pool.comms().request.notify_one(); + } - Ok(Self { pool, rx, request }) + Ok(Self { + pool, + rx: Some(rx), + request, + waiting: true, + }) } - pub(super) async fn wait(self) -> Result<(Guard, Instant), Error> { - let checkout_timeout = self.pool.inner().config.checkout_timeout; - let server = timeout(checkout_timeout, self.rx).await; + /// Wait for connection from the pool. + pub(super) async fn wait(&mut self) -> Result<(Guard, Instant), Error> { + let rx = self.rx.take().expect("waiter rx taken"); + + // Can be cancelled. Drop will remove the waiter from the queue. + let server = rx.await; + + // Disarm the guard. We can't be cancelled beyond this point. + self.waiting = false; let now = Instant::now(); match server { - Ok(Ok(server)) => { + Ok(server) => { let server = server?; Ok((Guard::new(self.pool.clone(), server, now), now)) } - Err(_err) => { - let mut guard = self.pool.lock(); - if !guard.banned() { - guard.maybe_ban(now, Error::CheckoutTimeout); - } - guard.remove_waiter(&self.request.id); + Err(_) => { + // Should not be possible. + // This means someone else removed my waiter from the wait queue, + // indicating a bug in the pool. Err(Error::CheckoutTimeout) } - - // Should not be possible. - // This means someone removed my waiter from the wait queue, - // indicating a bug in the pool. - Ok(Err(_)) => Err(Error::CheckoutTimeout), } } } @@ -64,3 +81,99 @@ pub(super) struct Waiter { pub(super) request: Request, pub(super) tx: Sender, Error>>, } + +#[cfg(test)] +mod tests { + use super::*; + use crate::backend::pool::Pool; + use crate::net::messages::BackendKeyData; + use tokio::time::{sleep, timeout, Duration}; + + #[tokio::test] + async fn test_cancellation_safety() { + let pool = Pool::new_test(); + pool.launch(); + + let num_tasks = 10; + let mut wait_tasks = Vec::new(); + + for i in 0..num_tasks { + let pool_clone = pool.clone(); + let request = Request::new(BackendKeyData::new()); + let mut waiting = Waiting::new(pool_clone, &request).unwrap(); + + let wait_task = tokio::spawn(async move { waiting.wait().await }); + + wait_tasks.push((wait_task, i)); + } + + { + let pool_guard = pool.lock(); + assert_eq!( + pool_guard.waiting.len(), + num_tasks, + "All waiters should be in queue" + ); + } + + sleep(Duration::from_millis(5)).await; + + for (wait_task, i) in wait_tasks { + if i % 2 == 0 { + sleep(Duration::from_millis(1)).await; + } + wait_task.abort(); + } + + sleep(Duration::from_millis(10)).await; + + let pool_guard = pool.lock(); + assert!( + pool_guard.waiting.is_empty(), + "All waiters should be removed from queue on cancellation" + ); + } + + #[tokio::test] + async fn test_timeout_removes_waiter() { + let config = crate::backend::pool::Config { + max: 1, + min: 1, + checkout_timeout: Duration::from_millis(10), + ..Default::default() + }; + + let pool = Pool::new(&crate::backend::pool::PoolConfig { + address: crate::backend::pool::Address { + host: "127.0.0.1".into(), + port: 5432, + database_name: "pgdog".into(), + user: "pgdog".into(), + password: "pgdog".into(), + ..Default::default() + }, + config, + }); + pool.launch(); + + sleep(Duration::from_millis(100)).await; + + let _conn = pool.get(&Request::default()).await.unwrap(); + + let request = Request::new(BackendKeyData::new()); + let waiter_pool = pool.clone(); + let get_conn = async move { + let mut waiting = Waiting::new(waiter_pool.clone(), &request).unwrap(); + waiting.wait().await + }; + let result = timeout(Duration::from_millis(100), get_conn).await; + + assert!(result.is_err()); + + let pool_guard = pool.lock(); + assert!( + pool_guard.waiting.is_empty(), + "Waiter should be removed on timeout" + ); + } +} diff --git a/pgdog/src/backend/prepared_statements.rs b/pgdog/src/backend/prepared_statements.rs index e0f1ca4be..0ee5a481e 100644 --- a/pgdog/src/backend/prepared_statements.rs +++ b/pgdog/src/backend/prepared_statements.rs @@ -95,8 +95,8 @@ impl PreparedStatements { let message = self.check_prepared(bind.statement())?; match message { Some(message) => { - self.state.add_ignore('1', bind.statement()); - self.prepared(bind.statement()); + self.state.add_ignore('1'); + self.parses.push_back(bind.statement().to_string()); self.state.add('2'); return Ok(HandleResult::Prepend(message)); } @@ -115,8 +115,8 @@ impl PreparedStatements { match message { Some(message) => { - self.state.add_ignore('1', describe.statement()); - self.prepared(describe.statement()); + self.state.add_ignore('1'); + self.parses.push_back(describe.statement().to_string()); self.state.add(ExecutionCode::DescriptionOrNothing); // t self.state.add(ExecutionCode::DescriptionOrNothing); // T return Ok(HandleResult::Prepend(message)); @@ -156,7 +156,6 @@ impl PreparedStatements { self.state.add_simulated(ParseComplete.message()?); return Ok(HandleResult::Drop); } else { - self.prepared(parse.name()); self.state.add('1'); self.parses.push_back(parse.name().to_string()); } @@ -175,7 +174,16 @@ impl PreparedStatements { self.state.add('3'); } } - ProtocolMessage::Prepare { .. } => (), + ProtocolMessage::Prepare { name, .. } => { + if self.contains(name) { + return Ok(HandleResult::Drop); + } else { + self.parses.push_back(name.clone()); + self.state.add_ignore('C'); + self.state.add_ignore('Z'); + return Ok(HandleResult::Forward); + } + } ProtocolMessage::CopyDone(_) => { self.state.action('c')?; } @@ -198,14 +206,11 @@ impl PreparedStatements { // Cleanup prepared statements state. match code { 'E' => { - let parse = self.parses.pop_front(); - let describe = self.describes.pop_front(); - if let Some(parse) = parse { - self.remove(&parse); - } - if let Some(describe) = describe { - self.remove(&describe); - } + // Backend ignored any subsequent extended commands. + // These prepared statements have not been prepared, even if they + // are syntactically valid. + self.describes.clear(); + self.parses.clear(); } 'T' => { @@ -222,16 +227,18 @@ impl PreparedStatements { self.describes.pop_front(); } - '1' => { - self.parses.pop_front(); + '1' | 'C' => { + if let Some(name) = self.parses.pop_front() { + self.prepared(&name); + } } 'G' => { self.state.prepend('G'); // Next thing we'll see is a CopyFail or CopyDone. } - 'c' | 'f' => { - // Backend told us the copy failed or succeeded. + // Backend told us the copy is done. + 'c' => { self.state.action(code)?; } @@ -240,12 +247,6 @@ impl PreparedStatements { match action { Action::Ignore => Ok(false), - Action::ForwardAndRemove(names) => { - for name in names { - self.remove(&name); - } - Ok(true) - } Action::Forward => Ok(true), } } @@ -255,16 +256,19 @@ impl PreparedStatements { self.state.done() && self.parses.is_empty() && self.describes.is_empty() } + /// The server connection has more messages to send + /// to the client. pub(crate) fn has_more_messages(&self) -> bool { self.state.has_more_messages() } - pub(crate) fn copy_mode(&self) -> bool { - self.state.copy_mode() + /// The server connection is in COPY mode. + pub(crate) fn in_copy_mode(&self) -> bool { + self.state.in_copy_mode() } fn check_prepared(&mut self, name: &str) -> Result, Error> { - if !self.contains(name) { + if !self.contains(name) && !self.parses.iter().any(|s| s == name) { let parse = self.parse(name); if let Some(parse) = parse { Ok(Some(ProtocolMessage::Parse(parse))) @@ -295,7 +299,7 @@ impl PreparedStatements { /// Get the Parse message stored in the global prepared statements /// cache for this statement. pub(crate) fn parse(&self, name: &str) -> Option { - self.global_cache.read().parse(name) + self.global_cache.read().rewritten_parse(name) } /// Get the globally stored RowDescription for this prepared statement, diff --git a/pgdog/src/backend/protocol/state.rs b/pgdog/src/backend/protocol/state.rs index c069bf226..b8505e5d6 100644 --- a/pgdog/src/backend/protocol/state.rs +++ b/pgdog/src/backend/protocol/state.rs @@ -10,7 +10,6 @@ use std::{collections::VecDeque, fmt::Debug}; pub enum Action { Forward, Ignore, - ForwardAndRemove(VecDeque), } #[derive(Debug, Copy, Clone, PartialEq)] @@ -71,7 +70,6 @@ impl MemoryUsage for ExecutionItem { #[derive(Debug, Clone, Default)] pub struct ProtocolState { queue: VecDeque, - names: VecDeque, simulated: VecDeque, extended: bool, out_of_sync: bool, @@ -81,7 +79,6 @@ impl MemoryUsage for ProtocolState { #[inline] fn memory_usage(&self) -> usize { self.queue.memory_usage() - + self.names.memory_usage() + self.simulated.memory_usage() + self.extended.memory_usage() + self.out_of_sync.memory_usage() @@ -95,11 +92,10 @@ impl ProtocolState { /// This is used for preparing statements that the client expects to be there /// but the server connection doesn't have yet. /// - pub(crate) fn add_ignore(&mut self, code: impl Into, name: &str) { + pub(crate) fn add_ignore(&mut self, code: impl Into) { let code = code.into(); self.extended = self.extended || code.extended(); self.queue.push_back(ExecutionItem::Ignore(code)); - self.names.push_back(name.to_owned()); } /// Add a message to the execution queue. We expect this message @@ -186,11 +182,8 @@ impl ProtocolState { // Used for preparing statements that the client expects to be there. ExecutionItem::Ignore(in_queue) => { - self.names.pop_front().ok_or(Error::ProtocolOutOfSync)?; if code == in_queue { Ok(Action::Ignore) - } else if code == ExecutionCode::Error { - Ok(Action::ForwardAndRemove(std::mem::take(&mut self.names))) } else { Err(Error::ProtocolOutOfSync) } @@ -198,7 +191,7 @@ impl ProtocolState { } } - pub(crate) fn copy_mode(&self) -> bool { + pub(crate) fn in_copy_mode(&self) -> bool { self.queue.front() == Some(&ExecutionItem::Code(ExecutionCode::Copy)) } @@ -215,6 +208,11 @@ impl ProtocolState { &self.queue } + #[cfg(test)] + pub(crate) fn queue_mut(&mut self) -> &mut VecDeque { + &mut self.queue + } + pub(crate) fn done(&self) -> bool { self.is_empty() && !self.out_of_sync } @@ -233,10 +231,615 @@ impl ProtocolState { mod test { use super::*; + // ======================================== + // Simple Query Protocol Tests + // ======================================== + + #[test] + fn test_simple_query_with_results() { + let mut state = ProtocolState::default(); + // Simple query: SELECT * FROM users + // Expected: RowDescription -> DataRow(s) -> CommandComplete -> ReadyForQuery + state.add('T'); // RowDescription + state.add('C'); // CommandComplete + state.add('Z'); // ReadyForQuery + + assert_eq!(state.action('T').unwrap(), Action::Forward); + // DataRows are not tracked, they come between T and C + assert_eq!(state.action('C').unwrap(), Action::Forward); + assert_eq!(state.action('Z').unwrap(), Action::Forward); + assert!(state.is_empty()); + } + + #[test] + fn test_simple_query_no_results() { + let mut state = ProtocolState::default(); + // Simple query: INSERT/UPDATE/DELETE + // Expected: CommandComplete -> ReadyForQuery + state.add('C'); // CommandComplete + state.add('Z'); // ReadyForQuery + + assert_eq!(state.action('C').unwrap(), Action::Forward); + assert_eq!(state.action('Z').unwrap(), Action::Forward); + assert!(state.is_empty()); + } + + #[test] + fn test_simple_query_empty() { + let mut state = ProtocolState::default(); + // Empty query + // Expected: EmptyQueryResponse -> ReadyForQuery + state.add('I'); // EmptyQueryResponse + state.add('Z'); // ReadyForQuery + + assert_eq!(state.action('I').unwrap(), Action::Forward); + assert_eq!(state.action('Z').unwrap(), Action::Forward); + assert!(state.is_empty()); + } + + #[test] + fn test_simple_query_error() { + let mut state = ProtocolState::default(); + // Query with syntax error + // Expected: ErrorResponse -> ReadyForQuery + state.add('C'); // CommandComplete (expected but won't arrive) + state.add('Z'); // ReadyForQuery + + // Error clears the queue except ReadyForQuery + assert_eq!(state.action('E').unwrap(), Action::Forward); + assert_eq!(state.len(), 1); // Only ReadyForQuery remains + assert_eq!(state.action('Z').unwrap(), Action::Forward); + assert!(state.is_empty()); + } + + #[test] + fn test_simple_query_multiple_results() { + let mut state = ProtocolState::default(); + // Multiple SELECT statements in one query + // Expected: T->C->T->C->Z + state.add('T'); // RowDescription for first query + state.add('C'); // CommandComplete for first query + state.add('T'); // RowDescription for second query + state.add('C'); // CommandComplete for second query + state.add('Z'); // ReadyForQuery + + assert_eq!(state.action('T').unwrap(), Action::Forward); + assert_eq!(state.action('C').unwrap(), Action::Forward); + assert_eq!(state.action('T').unwrap(), Action::Forward); + assert_eq!(state.action('C').unwrap(), Action::Forward); + assert_eq!(state.action('Z').unwrap(), Action::Forward); + assert!(state.is_empty()); + } + + // ======================================== + // Extended Query Protocol Tests + // ======================================== + + #[test] + fn test_extended_parse_bind_execute_sync() { + let mut state = ProtocolState::default(); + // Basic extended query: Parse -> Bind -> Execute -> Sync + state.add('1'); // ParseComplete + state.add('2'); // BindComplete + state.add('C'); // CommandComplete + state.add('Z'); // ReadyForQuery + + assert_eq!(state.action('1').unwrap(), Action::Forward); + assert_eq!(state.action('2').unwrap(), Action::Forward); + assert_eq!(state.action('C').unwrap(), Action::Forward); + assert_eq!(state.action('Z').unwrap(), Action::Forward); + assert!(state.is_empty()); + assert!(state.extended); + } + + #[test] + fn test_extended_parse_bind_describe_execute_sync() { + let mut state = ProtocolState::default(); + // Extended query with Describe: Parse -> Bind -> Describe -> Execute -> Sync + state.add('1'); // ParseComplete + state.add('2'); // BindComplete + state.add('T'); // RowDescription from Describe + state.add('C'); // CommandComplete + state.add('Z'); // ReadyForQuery + + assert_eq!(state.action('1').unwrap(), Action::Forward); + assert_eq!(state.action('2').unwrap(), Action::Forward); + assert_eq!(state.action('T').unwrap(), Action::Forward); + assert_eq!(state.action('C').unwrap(), Action::Forward); + assert_eq!(state.action('Z').unwrap(), Action::Forward); + assert!(state.is_empty()); + } + + #[test] + fn test_extended_describe_statement_returns_nodata() { + let mut state = ProtocolState::default(); + // Describe a statement that doesn't return data (e.g., INSERT) + state.add('1'); // ParseComplete + state.add('n'); // NoData from Describe + state.add('2'); // BindComplete + state.add('C'); // CommandComplete + state.add('Z'); // ReadyForQuery + + assert_eq!(state.action('1').unwrap(), Action::Forward); + assert_eq!(state.action('n').unwrap(), Action::Forward); + assert_eq!(state.action('2').unwrap(), Action::Forward); + assert_eq!(state.action('C').unwrap(), Action::Forward); + assert_eq!(state.action('Z').unwrap(), Action::Forward); + assert!(state.is_empty()); + } + + #[test] + fn test_extended_portal_suspended() { + let mut state = ProtocolState::default(); + // Execute with row limit, returns PortalSuspended + state.add('1'); // ParseComplete + state.add('2'); // BindComplete + state.add('s'); // PortalSuspended (partial results) + state.add('Z'); // ReadyForQuery + + assert_eq!(state.action('1').unwrap(), Action::Forward); + assert_eq!(state.action('2').unwrap(), Action::Forward); + assert_eq!(state.action('s').unwrap(), Action::Forward); + assert_eq!(state.action('Z').unwrap(), Action::Forward); + assert!(state.is_empty()); + } + + #[test] + fn test_extended_close_statement() { + let mut state = ProtocolState::default(); + // Close a prepared statement + state.add('3'); // CloseComplete + state.add('Z'); // ReadyForQuery + + assert_eq!(state.action('3').unwrap(), Action::Forward); + assert_eq!(state.action('Z').unwrap(), Action::Forward); + assert!(state.is_empty()); + } + + #[test] + fn test_extended_pipelined_queries() { + let mut state = ProtocolState::default(); + // Multiple queries in one pipeline: Parse->Bind->Execute->Parse->Bind->Execute->Sync + state.add('1'); // ParseComplete #1 + state.add('2'); // BindComplete #1 + state.add('C'); // CommandComplete #1 + state.add('1'); // ParseComplete #2 + state.add('2'); // BindComplete #2 + state.add('C'); // CommandComplete #2 + state.add('Z'); // ReadyForQuery + + assert_eq!(state.action('1').unwrap(), Action::Forward); + assert_eq!(state.action('2').unwrap(), Action::Forward); + assert_eq!(state.action('C').unwrap(), Action::Forward); + assert_eq!(state.action('1').unwrap(), Action::Forward); + assert_eq!(state.action('2').unwrap(), Action::Forward); + assert_eq!(state.action('C').unwrap(), Action::Forward); + assert_eq!(state.action('Z').unwrap(), Action::Forward); + assert!(state.is_empty()); + } + + // ======================================== + // Error Handling Tests + // ======================================== + + #[test] + fn test_extended_parse_error() { + let mut state = ProtocolState::default(); + // Parse fails (syntax error) + state.add('1'); // ParseComplete (expected but won't arrive) + state.add('2'); // BindComplete (won't be reached) + state.add('Z'); // ReadyForQuery + + // Error clears queue except ReadyForQuery + assert_eq!(state.action('E').unwrap(), Action::Forward); + assert!(state.out_of_sync); + assert_eq!(state.len(), 1); // Only ReadyForQuery remains + assert_eq!(state.action('Z').unwrap(), Action::Forward); + assert!(!state.out_of_sync); + assert!(state.is_empty()); + } + + #[test] + fn test_extended_bind_error() { + let mut state = ProtocolState::default(); + // Parse succeeds, Bind fails + state.add('1'); // ParseComplete + state.add('2'); // BindComplete (expected but won't arrive) + state.add('C'); // CommandComplete (won't be reached) + state.add('Z'); // ReadyForQuery + + assert_eq!(state.action('1').unwrap(), Action::Forward); + assert_eq!(state.action('E').unwrap(), Action::Forward); + assert!(state.out_of_sync); + assert_eq!(state.len(), 1); + assert_eq!(state.action('Z').unwrap(), Action::Forward); + assert!(!state.out_of_sync); + } + + #[test] + fn test_extended_execute_error() { + let mut state = ProtocolState::default(); + // Parse and Bind succeed, Execute fails + state.add('1'); // ParseComplete + state.add('2'); // BindComplete + state.add('C'); // CommandComplete (expected but won't arrive) + state.add('Z'); // ReadyForQuery + + assert_eq!(state.action('1').unwrap(), Action::Forward); + assert_eq!(state.action('2').unwrap(), Action::Forward); + assert_eq!(state.action('E').unwrap(), Action::Forward); + assert!(state.out_of_sync); + assert_eq!(state.action('Z').unwrap(), Action::Forward); + assert!(!state.out_of_sync); + } + + #[test] + fn test_simple_query_error_no_out_of_sync() { + let mut state = ProtocolState::default(); + // Simple query error should NOT set out_of_sync + state.add('C'); // CommandComplete (expected but won't arrive) + state.add('Z'); // ReadyForQuery + + assert!(!state.extended); + assert_eq!(state.action('E').unwrap(), Action::Forward); + assert!(!state.out_of_sync); // Simple query doesn't set out_of_sync + assert_eq!(state.action('Z').unwrap(), Action::Forward); + } + + #[test] + fn test_extended_error_in_pipeline() { + let mut state = ProtocolState::default(); + // Pipeline with error in middle: P->B->E->P->B->E->Sync + // If first Execute fails, rest of pipeline still processes + state.add('1'); // ParseComplete #1 + state.add('2'); // BindComplete #1 + state.add('C'); // CommandComplete #1 (won't arrive) + state.add('1'); // ParseComplete #2 + state.add('2'); // BindComplete #2 + state.add('C'); // CommandComplete #2 + state.add('Z'); // ReadyForQuery + + assert_eq!(state.action('1').unwrap(), Action::Forward); + assert_eq!(state.action('2').unwrap(), Action::Forward); + assert_eq!(state.action('E').unwrap(), Action::Forward); + assert!(state.out_of_sync); + // After error in extended protocol, we're out of sync + // Server still sends remaining responses but we're waiting for ReadyForQuery + assert_eq!(state.action('Z').unwrap(), Action::Forward); + assert!(!state.out_of_sync); + } + + // ======================================== + // COPY Protocol Tests + // ======================================== + + #[test] + fn test_copy_in_success() { + let mut state = ProtocolState::default(); + // COPY FROM STDIN + state.add('G'); // CopyInResponse + state.add('C'); // CommandComplete (after CopyDone or CopyFail) + state.add('Z'); // ReadyForQuery + + // Check copy_mode before consuming the message + assert!(state.in_copy_mode()); + assert_eq!(state.action('G').unwrap(), Action::Forward); + // After consuming 'G', we're no longer in copy mode (it's popped from queue) + assert!(!state.in_copy_mode()); + // CopyData messages ('d') would be sent here but aren't tracked + assert_eq!(state.action('C').unwrap(), Action::Forward); + assert_eq!(state.action('Z').unwrap(), Action::Forward); + assert!(state.is_empty()); + } + + #[test] + fn test_copy_fail() { + let mut state = ProtocolState::default(); + // COPY that fails + state.add('G'); // CopyInResponse + state.add('C'); // CommandComplete (won't arrive due to error) + state.add('Z'); // ReadyForQuery + + assert_eq!(state.action('G').unwrap(), Action::Forward); + // Client sends CopyFail ('f') + // Server responds with ErrorResponse + assert_eq!(state.action('E').unwrap(), Action::Forward); + assert_eq!(state.action('Z').unwrap(), Action::Forward); + } + + // ======================================== + // Ignore Tests (for statement preparation) + // ======================================== + + #[test] + fn test_ignore_parse_complete() { + let mut state = ProtocolState::default(); + state.add_ignore('1'); + assert_eq!(state.action('1').unwrap(), Action::Ignore); + assert!(state.is_empty()); + } + + #[test] + fn test_ignore_bind_complete() { + let mut state = ProtocolState::default(); + state.add_ignore('2'); + assert_eq!(state.action('2').unwrap(), Action::Ignore); + assert!(state.is_empty()); + } + + #[test] + fn test_ignore_error_behavior() { + let mut state = ProtocolState::default(); + state.add_ignore('1'); + state.add_ignore('2'); + + // When we get an error with Ignore items in queue, + // the Error arm is triggered first (before checking queue items) + // so it clears the queue and returns Forward, not ForwardAndRemove + let result = state.action('E').unwrap(); + assert_eq!(result, Action::Forward); + // Queue should be empty after error + assert!(state.is_empty()); + // Note: The ForwardAndRemove logic in the Ignore arm (line 192-193) + // is unreachable because Error is handled at the top of action() + // This may be dead code or a bug in the implementation. + } + + #[test] + fn test_ignore_wrong_code_is_out_of_sync() { + let mut state = ProtocolState::default(); + state.add_ignore('1'); + // We expect ParseComplete but get BindComplete + assert!(state.action('2').is_err()); + } + + // ======================================== + // Simulated Messages Tests + // ======================================== + + #[test] + fn test_simulated_message() { + let mut state = ProtocolState::default(); + // Create a simulated CloseComplete message + let message = Message::new(bytes::Bytes::from(vec![b'3', 0, 0, 0, 4])); + state.add_simulated(message.clone()); + + assert_eq!(state.len(), 1); + let retrieved = state.get_simulated(); + assert!(retrieved.is_some()); + assert_eq!(retrieved.unwrap().code(), '3'); + assert!(state.is_empty()); + } + + #[test] + fn test_simulated_message_wrong_position() { + let mut state = ProtocolState::default(); + let message = Message::new(bytes::Bytes::from(vec![b'3', 0, 0, 0, 4])); + state.add('1'); // ParseComplete expected first + state.add_simulated(message); + + // get_simulated should return None because CloseComplete is not at front + let retrieved = state.get_simulated(); + assert!(retrieved.is_none()); + assert_eq!(state.len(), 2); + } + + // ======================================== + // State Management Tests + // ======================================== + + #[test] + fn test_prepend() { + let mut state = ProtocolState::default(); + state.add('C'); // CommandComplete + state.add('Z'); // ReadyForQuery + state.prepend('T'); // RowDescription should come first + + assert_eq!(state.action('T').unwrap(), Action::Forward); + assert_eq!(state.action('C').unwrap(), Action::Forward); + assert_eq!(state.action('Z').unwrap(), Action::Forward); + assert!(state.is_empty()); + } + + #[test] + fn test_untracked_messages_always_forward() { + let mut state = ProtocolState::default(); + state.add('C'); // CommandComplete + + // Untracked messages (like DataRow 'D', NoticeResponse 'N', etc.) should always forward + // even if they're not in the queue + assert_eq!(state.action('D').unwrap(), Action::Forward); + assert_eq!(state.action('N').unwrap(), Action::Forward); + assert_eq!(state.action('S').unwrap(), Action::Forward); // ParameterStatus + + // Original queue should be unchanged + assert_eq!(state.len(), 1); + assert_eq!(state.action('C').unwrap(), Action::Forward); + } + + #[test] + fn test_ready_for_query_when_expecting_other() { + let mut state = ProtocolState::default(); + state.add('T'); // RowDescription + state.add('Z'); // ReadyForQuery + + // If we receive ReadyForQuery but we're expecting RowDescription first: + // - The code sets out_of_sync = false + // - Pops 'T' from queue + // - Checks: is received code NOT RFQ AND expected code IS RFQ? + // - No, received IS RFQ, so we don't push back + // - We consume 'T' and move on, 'Z' remains in queue + let result = state.action('Z'); + assert!(result.is_ok()); + assert_eq!(state.len(), 1); // Only 'Z' remains + assert_eq!(state.action('Z').unwrap(), Action::Forward); + assert!(state.is_empty()); + } + + #[test] + fn test_done_when_empty_and_in_sync() { + let mut state = ProtocolState::default(); + assert!(state.done()); + + state.add('Z'); + assert!(!state.done()); // Has messages + + assert_eq!(state.action('Z').unwrap(), Action::Forward); + assert!(state.done()); // Empty and in sync + } + #[test] - fn test_state() { + fn test_not_done_when_out_of_sync() { let mut state = ProtocolState::default(); - state.add_ignore('1', "test"); + state.add('1'); // ParseComplete + state.add('Z'); // ReadyForQuery + + assert_eq!(state.action('E').unwrap(), Action::Forward); + assert!(state.out_of_sync); + assert!(!state.done()); // Out of sync, not done even though has messages + + assert_eq!(state.action('Z').unwrap(), Action::Forward); + assert!(state.done()); // Back in sync and empty + } + + #[test] + fn test_out_of_sync_empty_queue() { + let mut state = ProtocolState::default(); + // Error with no pending ReadyForQuery + let result = state.action('E'); + assert!(result.is_ok()); + assert!(state.is_empty()); // Queue is empty + } + + // ======================================== + // Edge Cases and Complex Scenarios + // ======================================== + + #[test] + fn test_multiple_untracked_between_tracked() { + let mut state = ProtocolState::default(); + state.add('T'); // RowDescription + state.add('C'); // CommandComplete + state.add('Z'); // ReadyForQuery + + assert_eq!(state.action('T').unwrap(), Action::Forward); + // Multiple DataRows (untracked) + assert_eq!(state.action('D').unwrap(), Action::Forward); + assert_eq!(state.action('D').unwrap(), Action::Forward); + assert_eq!(state.action('D').unwrap(), Action::Forward); + // NoticeResponse (untracked) + assert_eq!(state.action('N').unwrap(), Action::Forward); + assert_eq!(state.action('C').unwrap(), Action::Forward); + assert_eq!(state.action('Z').unwrap(), Action::Forward); + assert!(state.is_empty()); + } + + #[test] + fn test_parameter_description() { + let mut state = ProtocolState::default(); + // Describe a prepared statement's parameters + state.add('1'); // ParseComplete + state.add('t'); // ParameterDescription + state.add('Z'); // ReadyForQuery + + assert_eq!(state.action('1').unwrap(), Action::Forward); + assert_eq!(state.action('t').unwrap(), Action::Forward); + assert_eq!(state.action('Z').unwrap(), Action::Forward); + assert!(state.is_empty()); + } + + #[test] + fn test_empty_queue_non_error_message() { + let mut state = ProtocolState::default(); + // Receiving a tracked message when queue is empty should be OutOfSync + let result = state.action('C'); + assert!(result.is_err()); + } + + #[test] + fn test_mixed_simple_and_extended() { + let mut state = ProtocolState::default(); + // This shouldn't happen in practice, but test state tracking + // Simple query + state.add('C'); // CommandComplete + state.add('Z'); // ReadyForQuery + assert_eq!(state.action('C').unwrap(), Action::Forward); + assert_eq!(state.action('Z').unwrap(), Action::Forward); + + // Now extended query + state.add('1'); // ParseComplete + state.add('2'); // BindComplete + state.add('C'); // CommandComplete + state.add('Z'); // ReadyForQuery + + assert_eq!(state.action('1').unwrap(), Action::Forward); + assert!(state.extended); // Now marked as extended + assert_eq!(state.action('2').unwrap(), Action::Forward); + assert_eq!(state.action('C').unwrap(), Action::Forward); + assert_eq!(state.action('Z').unwrap(), Action::Forward); + } + + #[test] + fn test_copy_mode_detection() { + let mut state = ProtocolState::default(); + assert!(!state.in_copy_mode()); + + state.add('G'); // CopyInResponse + state.add('C'); // CommandComplete + assert!(state.in_copy_mode()); + + assert_eq!(state.action('G').unwrap(), Action::Forward); + assert!(!state.in_copy_mode()); // No longer at front + + assert_eq!(state.action('C').unwrap(), Action::Forward); + assert!(!state.in_copy_mode()); + } + + #[test] + fn test_has_more_messages() { + let mut state = ProtocolState::default(); + assert!(!state.has_more_messages()); + + state.add('C'); + assert!(state.has_more_messages()); + + assert_eq!(state.action('C').unwrap(), Action::Forward); + assert!(!state.has_more_messages()); + } + + #[test] + fn test_names_cleared_on_error() { + // This test verifies that when an error occurs, both queue AND names + // are cleared to maintain the invariant that they stay synchronized. + + let mut state = ProtocolState::default(); + state.add_ignore('1'); + state.add_ignore('2'); + state.add_ignore('3'); + state.add('Z'); // ReadyForQuery + + // Error should clear both queue (except RFQ) and names + assert_eq!(state.action('E').unwrap(), Action::Forward); + + // Consume the ReadyForQuery + assert_eq!(state.action('Z').unwrap(), Action::Forward); + assert!(state.is_empty()); // Queue is empty + + // Now if we add a new ignore item, it should work correctly + // because names was also cleared + state.add_ignore('1'); assert_eq!(state.action('1').unwrap(), Action::Ignore); + assert!(state.is_empty()); // Both queue and names should be empty + + // Verify we can continue using the state normally + state.add_ignore('2'); + state.add('C'); + state.add('Z'); + + // Process normally + assert_eq!(state.action('2').unwrap(), Action::Ignore); + assert_eq!(state.action('C').unwrap(), Action::Forward); + assert_eq!(state.action('Z').unwrap(), Action::Forward); + assert!(state.is_empty()); } } diff --git a/pgdog/src/backend/pub_sub/listener.rs b/pgdog/src/backend/pub_sub/listener.rs index bff7d3d98..7bb5de35b 100644 --- a/pgdog/src/backend/pub_sub/listener.rs +++ b/pgdog/src/backend/pub_sub/listener.rs @@ -15,11 +15,11 @@ use tokio::{ use tracing::{debug, error, info}; use crate::{ - backend::{self, pool::Error, Pool}, + backend::{self, pool::Error, ConnectReason, DisconnectReason, Pool}, config::config, net::{ - FromBytes, NotificationResponse, Parameter, Parameters, Protocol, ProtocolMessage, Query, - ToBytes, + BackendKeyData, FromBytes, NotificationResponse, Parameter, Parameters, Protocol, + ProtocolMessage, Query, ToBytes, }, }; @@ -160,12 +160,17 @@ impl PubSubListener { ) -> Result<(), backend::Error> { info!("pub/sub started [{}]", pool.addr()); - let mut server = pool.standalone().await?; + let mut server = pool.standalone(ConnectReason::PubSub).await?; + server - .link_client(&Parameters::from(vec![Parameter { - name: "application_name".into(), - value: "PgDog Pub/Sub Listener".into(), - }])) + .link_client( + &BackendKeyData::new(), + &Parameters::from(vec![Parameter { + name: "application_name".into(), + value: "PgDog Pub/Sub Listener".into(), + }]), + None, + ) .await?; // Re-listen on all channels when re-starting the task. @@ -175,7 +180,10 @@ impl PubSubListener { .keys() .map(|channel| Request::Subscribe(channel.to_string()).into()) .collect::>(); - server.send(&resub.into()).await?; + + if !resub.is_empty() { + server.send(&resub.into()).await?; + } loop { select! { @@ -210,6 +218,7 @@ impl PubSubListener { debug!("pub/sub request {:?}", req); server.send(&vec![req.into()].into()).await?; } else { + server.disconnect_reason(DisconnectReason::Offline); break; } } diff --git a/pgdog/src/backend/replication/logical/copy_statement.rs b/pgdog/src/backend/replication/logical/copy_statement.rs index 5f9cdecc1..35b872247 100644 --- a/pgdog/src/backend/replication/logical/copy_statement.rs +++ b/pgdog/src/backend/replication/logical/copy_statement.rs @@ -2,11 +2,12 @@ //! Generate COPY statement for table synchronization. //! +use super::publisher::PublicationTable; + /// COPY statement generator. #[derive(Debug, Clone)] pub struct CopyStatement { - schema: String, - table: String, + table: PublicationTable, columns: Vec, } @@ -19,10 +20,9 @@ impl CopyStatement { /// * `table`: Name of the table. /// * `columns`: Table column names. /// - pub fn new(schema: &str, table: &str, columns: &[String]) -> CopyStatement { + pub fn new(table: &PublicationTable, columns: &[String]) -> CopyStatement { CopyStatement { - schema: schema.to_owned(), - table: table.to_owned(), + table: table.clone(), columns: columns.to_vec(), } } @@ -37,12 +37,28 @@ impl CopyStatement { self.copy(false) } + fn schema_name(&self, out: bool) -> &str { + if out || self.table.parent_schema.is_empty() { + &self.table.schema + } else { + &self.table.parent_schema + } + } + + fn table_name(&self, out: bool) -> &str { + if out || self.table.parent_name.is_empty() { + &self.table.name + } else { + &self.table.parent_name + } + } + // Generate the statement. fn copy(&self, out: bool) -> String { format!( r#"COPY "{}"."{}" ({}) {} WITH (FORMAT binary)"#, - self.schema, - self.table, + self.schema_name(out), + self.table_name(out), self.columns .iter() .map(|c| format!(r#""{}""#, c)) @@ -59,16 +75,42 @@ mod test { #[test] fn test_copy_stmt() { - let copy = CopyStatement::new("public", "test", &["id".into(), "email".into()]).copy_in(); + let table = PublicationTable { + schema: "public".into(), + name: "test".into(), + ..Default::default() + }; + + let copy = CopyStatement::new(&table, &["id".into(), "email".into()]); + let copy_in = copy.copy_in(); assert_eq!( - copy.to_string(), + copy_in, r#"COPY "public"."test" ("id", "email") FROM STDIN WITH (FORMAT binary)"# ); - let copy = CopyStatement::new("public", "test", &["id".into(), "email".into()]).copy_out(); assert_eq!( - copy.to_string(), + copy.copy_out(), r#"COPY "public"."test" ("id", "email") TO STDOUT WITH (FORMAT binary)"# ); + + let table = PublicationTable { + schema: "public".into(), + name: "test_0".into(), + parent_name: "test".into(), + parent_schema: "public".into(), + ..Default::default() + }; + + let copy = CopyStatement::new(&table, &["id".into(), "email".into()]); + let copy_in = copy.copy_in(); + assert_eq!( + copy_in, + r#"COPY "public"."test" ("id", "email") FROM STDIN WITH (FORMAT binary)"# + ); + + assert_eq!( + copy.copy_out(), + r#"COPY "public"."test_0" ("id", "email") TO STDOUT WITH (FORMAT binary)"# + ); } } diff --git a/pgdog/src/backend/replication/logical/error.rs b/pgdog/src/backend/replication/logical/error.rs index 3f6a03706..02e91ccb9 100644 --- a/pgdog/src/backend/replication/logical/error.rs +++ b/pgdog/src/backend/replication/logical/error.rs @@ -2,17 +2,20 @@ use std::num::ParseIntError; use thiserror::Error; -use crate::net::ErrorResponse; +use crate::{backend::replication::publisher::PublicationTable, net::ErrorResponse}; #[derive(Debug, Error)] pub enum Error { - #[error("{0}")] + #[error("backend: {0}")] Backend(#[from] crate::backend::Error), - #[error("{0}")] + #[error("pool: {0}")] Pool(#[from] crate::backend::pool::Error), - #[error("{0}")] + #[error("router: {0}")] + Router(#[from] crate::frontend::router::Error), + + #[error("net: {0}")] Net(#[from] crate::net::Error), #[error("transaction not started")] @@ -21,6 +24,12 @@ pub enum Error { #[error("out of sync, got {0}")] OutOfSync(char), + #[error("out of sync during commit, got {0}")] + CommitOutOfSync(char), + + #[error("out of sync during relation prepare, got {0}")] + RelationOutOfSync(char), + #[error("missing data")] MissingData, @@ -30,7 +39,7 @@ pub enum Error { #[error("copy error")] Copy, - #[error("{0}")] + #[error("pg_error: {0}")] PgError(Box), #[error("table \"{0}\".\"{1}\" has no replica identity")] @@ -39,13 +48,16 @@ pub enum Error { #[error("lsn decode")] LsnDecode, + #[error("replication slot \"{0}\" doesn't exist, but it should")] + MissingReplicationSlot(String), + #[error("parse int")] ParseInt(#[from] ParseIntError), #[error("shard has no primary")] NoPrimary, - #[error("{0}")] + #[error("parser: {0}")] Parser(#[from] crate::frontend::router::parser::Error), #[error("not connected")] @@ -65,6 +77,12 @@ pub enum Error { #[error("no replicas available for table sync")] NoReplicas, + + #[error("table {0} doesn't have a primary key")] + NoPrimaryKey(PublicationTable), + + #[error("router returned incorrect command")] + IncorrectCommand, } impl From for Error { diff --git a/pgdog/src/backend/replication/logical/publisher/copy.rs b/pgdog/src/backend/replication/logical/publisher/copy.rs index 2994742da..5f2134353 100644 --- a/pgdog/src/backend/replication/logical/publisher/copy.rs +++ b/pgdog/src/backend/replication/logical/publisher/copy.rs @@ -2,7 +2,7 @@ use crate::{ backend::Server, net::{CopyData, ErrorResponse, FromBytes, Protocol, Query, ToBytes}, }; -use tracing::trace; +use tracing::{debug, trace}; use super::{ super::{CopyStatement, Error}, @@ -17,8 +17,7 @@ pub struct Copy { impl Copy { pub fn new(table: &Table) -> Self { let stmt = CopyStatement::new( - &table.table.schema, - &table.table.name, + &table.table, &table .columns .iter() @@ -34,9 +33,10 @@ impl Copy { return Err(Error::TransactionNotStarted); } - server - .send(&vec![Query::new(self.stmt.copy_out()).into()].into()) - .await?; + let query = Query::new(self.stmt.copy_out()); + debug!("{} [{}]", query.query(), server.addr()); + + server.send(&vec![query.into()].into()).await?; let result = server.read().await?; match result.code() { 'E' => return Err(ErrorResponse::from_bytes(result.to_bytes()?)?.into()), diff --git a/pgdog/src/backend/replication/logical/publisher/publisher_impl.rs b/pgdog/src/backend/replication/logical/publisher/publisher_impl.rs index d7b97cdf9..d20aff42d 100644 --- a/pgdog/src/backend/replication/logical/publisher/publisher_impl.rs +++ b/pgdog/src/backend/replication/logical/publisher/publisher_impl.rs @@ -1,6 +1,7 @@ use std::collections::HashMap; use std::time::Duration; +use pgdog_config::QueryParserEngine; use tokio::{select, spawn}; use tracing::{debug, error, info}; @@ -27,15 +28,22 @@ pub struct Publisher { tables: HashMap>, /// Replication slots. slots: HashMap, + /// Query parser engine. + query_parser_engine: QueryParserEngine, } impl Publisher { - pub fn new(cluster: &Cluster, publication: &str) -> Self { + pub fn new( + cluster: &Cluster, + publication: &str, + query_parser_engine: QueryParserEngine, + ) -> Self { Self { cluster: cluster.clone(), publication: publication.to_string(), tables: HashMap::new(), slots: HashMap::new(), + query_parser_engine, } } @@ -44,7 +52,8 @@ impl Publisher { for (number, shard) in self.cluster.shards().iter().enumerate() { // Load tables from publication. let mut primary = shard.primary(&Request::default()).await?; - let tables = Table::load(&self.publication, &mut primary).await?; + let tables = + Table::load(&self.publication, &mut primary, self.query_parser_engine).await?; self.tables.insert(number, tables); } @@ -59,11 +68,12 @@ impl Publisher { /// If you're doing a cross-shard transaction, parts of it can be lost. /// /// TODO: Add support for 2-phase commit. - async fn create_slots(&mut self) -> Result<(), Error> { + async fn create_slots(&mut self, slot_name: Option) -> Result<(), Error> { for (number, shard) in self.cluster.shards().iter().enumerate() { let addr = shard.primary(&Request::default()).await?.addr().clone(); - let mut slot = ReplicationSlot::replication(&self.publication, &addr); + let mut slot = + ReplicationSlot::replication(&self.publication, &addr, slot_name.clone()); slot.create_slot().await?; self.slots.insert(number, slot); @@ -76,7 +86,11 @@ impl Publisher { /// /// This uses a dedicated replication slot which will survive crashes and reboots. /// N.B.: The slot needs to be manually dropped! - pub async fn replicate(&mut self, dest: &Cluster) -> Result<(), Error> { + pub async fn replicate( + &mut self, + dest: &Cluster, + slot_name: Option, + ) -> Result<(), Error> { // Replicate shards in parallel. let mut streams = vec![]; @@ -87,7 +101,7 @@ impl Publisher { // Create replication slots if we haven't already. if self.slots.is_empty() { - self.create_slots().await?; + self.create_slots(slot_name).await?; } for (number, _) in self.cluster.shards().iter().enumerate() { @@ -98,7 +112,7 @@ impl Publisher { .get(&number) .ok_or(Error::NoReplicationTables(number))?; // Handles the logical replication stream messages. - let mut stream = StreamSubscriber::new(dest, tables); + let mut stream = StreamSubscriber::new(dest, tables, self.query_parser_engine); // Take ownership of the slot for replication. let mut slot = self @@ -166,13 +180,19 @@ impl Publisher { /// re-sharding the cluster in the process. /// /// TODO: Parallelize shard syncs. - pub async fn data_sync(&mut self, dest: &Cluster) -> Result<(), Error> { + pub async fn data_sync( + &mut self, + dest: &Cluster, + replicate: bool, + slot_name: Option, + ) -> Result<(), Error> { // Create replication slots. - self.create_slots().await?; + self.create_slots(slot_name.clone()).await?; for (number, shard) in self.cluster.shards().iter().enumerate() { let mut primary = shard.primary(&Request::default()).await?; - let tables = Table::load(&self.publication, &mut primary).await?; + let tables = + Table::load(&self.publication, &mut primary, self.query_parser_engine).await?; let include_primary = !shard.has_replicas(); let replicas = shard @@ -181,6 +201,7 @@ impl Publisher { .filter(|(r, _)| match *r { Role::Replica => true, Role::Primary => include_primary, + Role::Auto => false, }) .map(|(_, p)| p) .collect::>(); @@ -199,8 +220,10 @@ impl Publisher { self.tables.insert(number, tables); } - // Replicate changes. - self.replicate(dest).await?; + if replicate { + // Replicate changes. + self.replicate(dest, slot_name).await?; + } Ok(()) } diff --git a/pgdog/src/backend/replication/logical/publisher/queries.rs b/pgdog/src/backend/replication/logical/publisher/queries.rs index 4ce094e8e..fef0e5c9b 100644 --- a/pgdog/src/backend/replication/logical/publisher/queries.rs +++ b/pgdog/src/backend/replication/logical/publisher/queries.rs @@ -3,6 +3,8 @@ //! TODO: I think these are Postgres-version specific, so we need to handle that //! later. These were fetched from CREATE SUBSCRIPTION ran on Postgres 17. //! +use std::fmt::Display; + use crate::{ backend::Server, net::{DataRow, Format}, @@ -11,21 +13,39 @@ use crate::{ use super::super::Error; /// Get list of tables in publication. -static TABLES: &str = "SELECT DISTINCT n.nspname, c.relname, gpt.attrs +static TABLES: &str = "SELECT DISTINCT + n.nspname, + c.relname, + gpt.attrs, + COALESCE(pn.nspname::text, '') AS parent_schema, + COALESCE(p.relname::text, '') AS parent_table FROM pg_class c JOIN pg_namespace n ON n.oid = c.relnamespace -JOIN ( SELECT (pg_get_publication_tables(VARIADIC array_agg(pubname::text))).* - FROM pg_publication - WHERE pubname IN ($1)) AS gpt - ON gpt.relid = c.oid -ORDER BY n.nspname, c.relname"; +JOIN ( + SELECT (pg_get_publication_tables(VARIADIC array_agg(pubname::text))).* + FROM pg_publication + WHERE pubname IN ($1) +) AS gpt + ON gpt.relid = c.oid +LEFT JOIN pg_inherits i ON i.inhrelid = c.oid -- only present if c is a child partition +LEFT JOIN pg_class p ON p.oid = i.inhparent -- immediate parent partitioned table +LEFT JOIN pg_namespace pn ON pn.oid = p.relnamespace +ORDER BY n.nspname, c.relname;"; /// Table included in a publication. -#[derive(Debug, Clone, PartialEq)] +#[derive(Debug, Clone, PartialEq, Default)] pub struct PublicationTable { pub schema: String, pub name: String, pub attributes: String, + pub parent_schema: String, + pub parent_name: String, +} + +impl Display for PublicationTable { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "\"{}\".\"{}\"", self.schema, self.name) + } } impl PublicationTable { @@ -37,6 +57,22 @@ impl PublicationTable { .fetch_all(TABLES.replace("$1", &format!("'{}'", publication))) .await?) } + + pub fn destination_name(&self) -> &str { + if self.parent_name.is_empty() { + &self.name + } else { + &self.parent_name + } + } + + pub fn destination_schema(&self) -> &str { + if self.parent_schema.is_empty() { + &self.schema + } else { + &self.parent_schema + } + } } impl From for PublicationTable { @@ -45,6 +81,8 @@ impl From for PublicationTable { schema: value.get(0, Format::Text).unwrap_or_default(), name: value.get(1, Format::Text).unwrap_or_default(), attributes: value.get(2, Format::Text).unwrap_or_default(), + parent_schema: value.get(3, Format::Text).unwrap_or_default(), + parent_name: value.get(4, Format::Text).unwrap_or_default(), } } } diff --git a/pgdog/src/backend/replication/logical/publisher/slot.rs b/pgdog/src/backend/replication/logical/publisher/slot.rs index 8667c3d28..6ecae929d 100644 --- a/pgdog/src/backend/replication/logical/publisher/slot.rs +++ b/pgdog/src/backend/replication/logical/publisher/slot.rs @@ -1,17 +1,18 @@ use super::super::Error; use crate::{ - backend::{pool::Address, Server, ServerOptions}, + backend::{self, pool::Address, ConnectReason, Server, ServerOptions}, net::{ replication::StatusUpdate, CopyData, CopyDone, DataRow, ErrorResponse, Format, FromBytes, - Protocol, Query, ToBytes, + FromDataType, Protocol, Query, ToBytes, }, util::random_string, }; +use bytes::Bytes; use std::{fmt::Display, str::FromStr, time::Duration}; use tokio::time::timeout; use tracing::{debug, trace}; -#[derive(Debug, Clone, Default, Copy)] +#[derive(Debug, Clone, Default, Copy, Eq, PartialEq, Ord, PartialOrd)] pub struct Lsn { pub high: i64, pub low: i64, @@ -27,6 +28,20 @@ impl Lsn { } } +impl FromDataType for Lsn { + fn decode(bytes: &[u8], encoding: Format) -> Result { + let val = String::decode(bytes, encoding)?; + Self::from_str(&val).map_err(|_| crate::net::Error::NotPgLsn) + } + + fn encode(&self, encoding: Format) -> Result { + match encoding { + Format::Text => Ok(Bytes::from(self.to_string())), + Format::Binary => todo!(), + } + } +} + impl FromStr for Lsn { type Err = Error; @@ -91,8 +106,8 @@ pub struct ReplicationSlot { impl ReplicationSlot { /// Create replication slot used for streaming the WAL. - pub fn replication(publication: &str, address: &Address) -> Self { - let name = format!("__pgdog_repl_{}", random_string(19).to_lowercase()); + pub fn replication(publication: &str, address: &Address, name: Option) -> Self { + let name = name.unwrap_or(format!("__pgdog_repl_{}", random_string(19).to_lowercase())); Self { address: address.clone(), @@ -124,7 +139,14 @@ impl ReplicationSlot { /// Connect to database using replication mode. pub async fn connect(&mut self) -> Result<(), Error> { - self.server = Some(Server::connect(&self.address, ServerOptions::new_replication()).await?); + self.server = Some( + Server::connect( + &self.address, + ServerOptions::new_replication(), + ConnectReason::Replication, + ) + .await?, + ); Ok(()) } @@ -156,26 +178,69 @@ impl ReplicationSlot { self.snapshot ); - let result = self + let existing_slot = format!( + " + SELECT + slot_name, + restart_lsn, + confirmed_flush_lsn + FROM + pg_replication_slots + WHERE slot_name = '{}'", + self.name + ); + + match self .server()? .fetch_all::(&start_replication) - .await? - .pop() - .ok_or(Error::MissingData)?; - - let lsn = result - .get::(1, Format::Text) - .ok_or(Error::MissingData)?; - - let lsn = Lsn::from_str(&lsn)?; - self.lsn = lsn; + .await + { + Ok(mut result) => { + let result = result.pop().ok_or(Error::MissingData)?; + let lsn = result + .get::(1, Format::Text) + .ok_or(Error::MissingData)?; + let lsn = Lsn::from_str(&lsn)?; + self.lsn = lsn; + + debug!( + "replication slot \"{}\" at lsn {} created [{}]", + self.name, self.lsn, self.address, + ); + + Ok(lsn) + } - debug!( - "replication slot \"{}\" at lsn {} created [{}]", - self.name, self.lsn, self.address, - ); + Err(err) => match err { + backend::Error::ExecutionError(err) => { + // duplicate object. + if err.code == "42710" { + let exists: Option = + self.server()?.fetch_all(existing_slot).await?.pop(); + + if let Some(lsn) = + exists.and_then(|slot| slot.get::(2, Format::Text)) + { + let lsn = Lsn::from_str(&lsn)?; + self.lsn = lsn; + + debug!( + "using existing replication slot \"{}\" at lsn {} [{}]", + self.name, self.lsn, self.address, + ); + + Ok(lsn) + } else { + Err(Error::MissingReplicationSlot(self.name.clone())) + } + } else { + Err(backend::Error::ExecutionError(err).into()) + } + } - Ok(lsn) + err => Err(err.into()), + }, + } } /// Drop the slot. diff --git a/pgdog/src/backend/replication/logical/publisher/table.rs b/pgdog/src/backend/replication/logical/publisher/table.rs index 585e900b1..c109b43e5 100644 --- a/pgdog/src/backend/replication/logical/publisher/table.rs +++ b/pgdog/src/backend/replication/logical/publisher/table.rs @@ -2,6 +2,8 @@ use std::time::Duration; +use pgdog_config::QueryParserEngine; + use crate::backend::pool::Address; use crate::backend::replication::publisher::progress::Progress; use crate::backend::replication::publisher::Lsn; @@ -26,10 +28,16 @@ pub struct Table { pub columns: Vec, /// Table data as of this LSN. pub lsn: Lsn, + /// Query parser engine. + pub query_parser_engine: QueryParserEngine, } impl Table { - pub async fn load(publication: &str, server: &mut Server) -> Result, Error> { + pub async fn load( + publication: &str, + server: &mut Server, + query_parser_engine: QueryParserEngine, + ) -> Result, Error> { let tables = PublicationTable::load(publication, server).await?; let mut results = vec![]; @@ -43,12 +51,22 @@ impl Table { identity, columns, lsn: Lsn::default(), + query_parser_engine, }); } Ok(results) } + /// Check that the table supports replication. + pub fn valid(&self) -> Result<(), Error> { + if !self.columns.iter().any(|c| c.identity) { + return Err(Error::NoPrimaryKey(self.table.clone())); + } + + Ok(()) + } + /// Upsert record into table. pub fn insert(&self, upsert: bool) -> String { let names = format!( @@ -91,7 +109,59 @@ impl Table { format!( "INSERT INTO \"{}\".\"{}\" {} {} {}", - self.table.schema, self.table.name, names, values, on_conflict + self.table.destination_schema(), + self.table.destination_name(), + names, + values, + on_conflict + ) + } + + /// Update record in table. + pub fn update(&self) -> String { + let set_clause = self + .columns + .iter() + .enumerate() + .filter(|(_, c)| !c.identity) + .map(|(i, c)| format!("\"{}\" = ${}", c.name, i + 1)) + .collect::>() + .join(", "); + + let where_clause = self + .columns + .iter() + .enumerate() + .filter(|(_, c)| c.identity) + .map(|(i, c)| format!("\"{}\" = ${}", c.name, i + 1)) + .collect::>() + .join(" AND "); + + format!( + "UPDATE \"{}\".\"{}\" SET {} WHERE {}", + self.table.destination_schema(), + self.table.destination_name(), + set_clause, + where_clause + ) + } + + /// Delete record from table. + pub fn delete(&self) -> String { + let where_clause = self + .columns + .iter() + .enumerate() + .filter(|(_, c)| c.identity) + .map(|(i, c)| format!("\"{}\" = ${}", c.name, i + 1)) + .collect::>() + .join(" AND "); + + format!( + "DELETE FROM \"{}\".\"{}\" WHERE {}", + self.table.destination_schema(), + self.table.destination_name(), + where_clause ) } @@ -120,7 +190,7 @@ impl Table { // Create new standalone connection for the copy. // let mut server = Server::connect(source, ServerOptions::new_replication()).await?; - let mut copy_sub = CopySubscriber::new(copy.statement(), dest)?; + let mut copy_sub = CopySubscriber::new(copy.statement(), dest, self.query_parser_engine)?; copy_sub.connect().await?; // Create sync slot. @@ -171,6 +241,7 @@ impl Table { mod test { use crate::backend::replication::logical::publisher::test::setup_publication; + use crate::config::config; use super::*; @@ -179,15 +250,25 @@ mod test { crate::logger(); let mut publication = setup_publication().await; - let tables = Table::load("publication_test", &mut publication.server) - .await - .unwrap(); + let tables = Table::load( + "publication_test", + &mut publication.server, + config().config.general.query_parser_engine, + ) + .await + .unwrap(); assert_eq!(tables.len(), 2); for table in tables { let upsert = table.insert(true); assert!(pg_query::parse(&upsert).is_ok()); + + let update = table.update(); + assert!(pg_query::parse(&update).is_ok()); + + let delete = table.delete(); + assert!(pg_query::parse(&delete).is_ok()); } publication.cleanup().await; diff --git a/pgdog/src/backend/replication/logical/subscriber/context.rs b/pgdog/src/backend/replication/logical/subscriber/context.rs new file mode 100644 index 000000000..405f60e0c --- /dev/null +++ b/pgdog/src/backend/replication/logical/subscriber/context.rs @@ -0,0 +1,62 @@ +use lazy_static::lazy_static; + +use super::super::Error; +use crate::{ + backend::Cluster, + frontend::{ + client::Sticky, router::parser::Shard, ClientRequest, Command, Router, RouterContext, + }, + net::{replication::TupleData, Bind, Parameters, Parse}, +}; + +#[derive(Debug)] +pub struct StreamContext<'a> { + request: ClientRequest, + cluster: &'a Cluster, + bind: Bind, +} + +impl<'a> StreamContext<'a> { + /// Construct new stream context. + pub fn new(cluster: &'a Cluster, tuple: &TupleData, stmt: &Parse) -> Self { + let bind = tuple.to_bind(stmt.name()); + let request = ClientRequest::from(vec![stmt.clone().into(), bind.clone().into()]); + + Self { + request, + cluster, + bind, + } + } + + pub fn shard(&'a mut self) -> Result { + let router_context = self.router_context()?; + let mut router = Router::new(); + let route = router.query(router_context)?; + + if let Command::Query(route) = route { + Ok(route.shard().clone()) + } else { + Err(Error::IncorrectCommand) + } + } + + /// Get Bind message. + pub fn bind(&self) -> &Bind { + &self.bind + } + + /// Construct router context. + pub fn router_context(&'a mut self) -> Result, Error> { + lazy_static! { + static ref PARAMS: Parameters = Parameters::default(); + } + Ok(RouterContext::new( + &self.request, + self.cluster, + &PARAMS, + None, + Sticky::new(), + )?) + } +} diff --git a/pgdog/src/backend/replication/logical/subscriber/copy.rs b/pgdog/src/backend/replication/logical/subscriber/copy.rs index 8b3c50595..63e1b498b 100644 --- a/pgdog/src/backend/replication/logical/subscriber/copy.rs +++ b/pgdog/src/backend/replication/logical/subscriber/copy.rs @@ -1,10 +1,12 @@ //! Shard COPY stream from one source //! between N shards. -use pg_query::NodeEnum; +use pg_query::{parse_raw, NodeEnum}; +use pgdog_config::QueryParserEngine; +use tracing::debug; use crate::{ - backend::{replication::subscriber::ParallelConnection, Cluster}, + backend::{replication::subscriber::ParallelConnection, Cluster, ConnectReason}, config::Role, frontend::router::parser::{CopyParser, Shard}, net::{CopyData, CopyDone, ErrorResponse, FromBytes, Protocol, Query, ToBytes}, @@ -33,8 +35,17 @@ impl CopySubscriber { /// 1. What kind of encoding we use. /// 2. Which column is used for sharding. /// - pub fn new(copy_stmt: &CopyStatement, cluster: &Cluster) -> Result { - let stmt = pg_query::parse(copy_stmt.clone().copy_in().as_str())?; + pub fn new( + copy_stmt: &CopyStatement, + cluster: &Cluster, + query_parser_engine: QueryParserEngine, + ) -> Result { + let stmt = match query_parser_engine { + QueryParserEngine::PgQueryProtobuf => { + pg_query::parse(copy_stmt.clone().copy_in().as_str()) + } + QueryParserEngine::PgQueryRaw => parse_raw(copy_stmt.clone().copy_in().as_str()), + }?; let stmt = stmt .protobuf .stmts @@ -47,9 +58,7 @@ impl CopySubscriber { .as_ref() .ok_or(Error::MissingData)?; let copy = if let NodeEnum::CopyStmt(stmt) = stmt { - CopyParser::new(stmt, cluster) - .map_err(|_| Error::MissingData)? - .ok_or(Error::MissingData)? + CopyParser::new(stmt, cluster).map_err(|_| Error::MissingData)? } else { return Err(Error::MissingData); }; @@ -74,7 +83,7 @@ impl CopySubscriber { .find(|(role, _)| role == &Role::Primary) .ok_or(Error::NoPrimary)? .1 - .standalone() + .standalone(ConnectReason::Replication) .await?; servers.push(ParallelConnection::new(primary)?); } @@ -102,6 +111,8 @@ impl CopySubscriber { } for server in &mut self.connections { + debug!("{} [{}]", stmt.query(), server.addr()); + server.send_one(&stmt.clone().into()).await?; server.flush().await?; @@ -196,7 +207,8 @@ mod test { use bytes::Bytes; use crate::{ - backend::pool::Request, + backend::{pool::Request, replication::publisher::PublicationTable}, + config::config, frontend::router::parser::binary::{header::Header, Data, Tuple}, }; @@ -206,8 +218,14 @@ mod test { async fn test_subscriber() { crate::logger(); - let copy = CopyStatement::new("pgdog", "sharded", &["id".into(), "value".into()]); - let cluster = Cluster::new_test(); + let table = PublicationTable { + schema: "pgdog".into(), + name: "sharded".into(), + ..Default::default() + }; + + let copy = CopyStatement::new(&table, &["id".into(), "value".into()]); + let cluster = Cluster::new_test(&config()); cluster.launch(); cluster @@ -220,7 +238,9 @@ mod test { .await .unwrap(); - let mut subscriber = CopySubscriber::new(©, &cluster).unwrap(); + let mut subscriber = + CopySubscriber::new(©, &cluster, config().config.general.query_parser_engine) + .unwrap(); subscriber.start_copy().await.unwrap(); let header = CopyData::new(&Header::new().to_bytes().unwrap()); diff --git a/pgdog/src/backend/replication/logical/subscriber/mod.rs b/pgdog/src/backend/replication/logical/subscriber/mod.rs index b7a8c0ccc..a83b3cff1 100644 --- a/pgdog/src/backend/replication/logical/subscriber/mod.rs +++ b/pgdog/src/backend/replication/logical/subscriber/mod.rs @@ -1,6 +1,9 @@ +pub mod context; pub mod copy; pub mod parallel_connection; pub mod stream; + +pub use context::StreamContext; pub use copy::CopySubscriber; pub use parallel_connection::ParallelConnection; pub use stream::StreamSubscriber; diff --git a/pgdog/src/backend/replication/logical/subscriber/parallel_connection.rs b/pgdog/src/backend/replication/logical/subscriber/parallel_connection.rs index 8e48f7ba4..4c4b15713 100644 --- a/pgdog/src/backend/replication/logical/subscriber/parallel_connection.rs +++ b/pgdog/src/backend/replication/logical/subscriber/parallel_connection.rs @@ -10,6 +10,7 @@ use tokio::sync::{ Notify, }; +use crate::backend::pool::Address; use crate::{ backend::Server, frontend::ClientRequest, @@ -43,6 +44,7 @@ pub struct ParallelConnection { tx: Sender, rx: Receiver, stop: Arc, + address: Address, } impl ParallelConnection { @@ -91,6 +93,11 @@ impl ParallelConnection { Ok(()) } + /// Server address. + pub fn addr(&self) -> &Address { + &self.address + } + // Move server connection into its own Tokio task. pub fn new(server: Server) -> Result { // Ideally we don't hardcode these. PgDog @@ -98,6 +105,7 @@ impl ParallelConnection { let (tx1, rx1) = channel(4096); let (tx2, rx2) = channel(4096); let stop = Arc::new(Notify::new()); + let address = server.addr().clone(); let listener = Listener { stop: stop.clone(), @@ -113,6 +121,7 @@ impl ParallelConnection { }); Ok(Self { + address, tx: tx1, rx: rx2, stop, diff --git a/pgdog/src/backend/replication/logical/subscriber/stream.rs b/pgdog/src/backend/replication/logical/subscriber/stream.rs index e0a05fcdc..5eb11a039 100644 --- a/pgdog/src/backend/replication/logical/subscriber/stream.rs +++ b/pgdog/src/backend/replication/logical/subscriber/stream.rs @@ -4,26 +4,30 @@ //! into idempotent prepared statements. //! use std::{ - collections::HashMap, + collections::{HashMap, HashSet}, + fmt::Display, sync::atomic::{AtomicUsize, Ordering}, }; use once_cell::sync::Lazy; use pg_query::{ + parse_raw, protobuf::{InsertStmt, ParseResult}, NodeEnum, }; -use tracing::trace; +use pgdog_config::QueryParserEngine; +use tracing::{debug, trace}; use super::super::{publisher::Table, Error}; +use super::StreamContext; use crate::{ - backend::{Cluster, Server, ShardingSchema}, + backend::{Cluster, ConnectReason, Server}, config::Role, - frontend::router::parser::{Insert, Shard}, + frontend::router::parser::Shard, net::{ replication::{ - xlog_data::XLogPayload, Commit as XLogCommit, Insert as XLogInsert, Relation, - StatusUpdate, + xlog_data::XLogPayload, Commit as XLogCommit, Delete as XLogDelete, + Insert as XLogInsert, Relation, StatusUpdate, Update as XLogUpdate, }, Bind, CopyData, ErrorResponse, Execute, Flush, FromBytes, Parse, Protocol, Sync, ToBytes, }, @@ -46,37 +50,47 @@ struct Key { name: String, } +impl Display for Key { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, r#""{}"."{}""#, self.schema, self.name) + } +} + #[derive(Default, Debug, Clone)] struct Statements { - #[allow(dead_code)] insert: Statement, - upsert: Statement, #[allow(dead_code)] + upsert: Statement, update: Statement, + delete: Statement, } #[derive(Default, Debug, Clone)] struct Statement { - name: String, - query: String, ast: ParseResult, + parse: Parse, } impl Statement { - fn parse(&self) -> Parse { - Parse::named(&self.name, &self.query) + fn parse(&self) -> &Parse { + &self.parse } - fn new(query: &str) -> Result { - let ast = pg_query::parse(query)?.protobuf; + fn new(query: &str, query_parser_engine: QueryParserEngine) -> Result { + let ast = match query_parser_engine { + QueryParserEngine::PgQueryProtobuf => pg_query::parse(query), + QueryParserEngine::PgQueryRaw => parse_raw(query), + }? + .protobuf; + let name = statement_name(); Ok(Self { - name: statement_name(), - query: query.to_string(), ast, + parse: Parse::named(name, query.to_string()), }) } #[allow(clippy::borrowed_box)] + #[allow(dead_code)] fn insert(&self) -> Option<&Box> { self.ast .stmts @@ -95,16 +109,16 @@ impl Statement { .flatten() .flatten() } -} + fn query(&self) -> &str { + self.parse.query() + } +} #[derive(Debug)] pub struct StreamSubscriber { /// Destination cluster. cluster: Cluster, - /// Sharding schema. - sharding_schema: ShardingSchema, - // Relation markers sent by the publisher. // Happens once per connection. relations: HashMap, @@ -115,6 +129,9 @@ pub struct StreamSubscriber { // Statements statements: HashMap, + // Partitioned tables dedup. + partitioned_dedup: HashSet, + // LSNs for each table table_lsns: HashMap, @@ -127,15 +144,23 @@ pub struct StreamSubscriber { // Bytes sharded bytes_sharded: usize, + + // Query parser engine. + query_parser_engine: QueryParserEngine, } impl StreamSubscriber { - pub fn new(cluster: &Cluster, tables: &[Table]) -> Self { + pub fn new( + cluster: &Cluster, + tables: &[Table], + query_parser_engine: QueryParserEngine, + ) -> Self { + let cluster = cluster.logical_stream(); Self { - cluster: cluster.clone(), - sharding_schema: cluster.sharding_schema(), + cluster, relations: HashMap::new(), statements: HashMap::new(), + partitioned_dedup: HashSet::new(), table_lsns: HashMap::new(), tables: tables .iter() @@ -153,6 +178,7 @@ impl StreamSubscriber { lsn: 0, // Unknown, bytes_sharded: 0, lsn_changed: true, + query_parser_engine, } } @@ -167,7 +193,7 @@ impl StreamSubscriber { .find(|(r, _)| r == &Role::Primary) .ok_or(Error::NoPrimary)? .1 - .standalone() + .standalone(ConnectReason::Replication) .await?; conns.push(primary); } @@ -231,33 +257,101 @@ impl StreamSubscriber { // Convert Insert into an idempotent "upsert" and apply it to // the right shard(s). async fn insert(&mut self, insert: XLogInsert) -> Result<(), Error> { - if let Some(table_lsn) = self.table_lsns.get(&insert.oid) { - // Don't apply change if table is ahead. - if self.lsn < *table_lsn { - return Ok(()); - } + if self.lsn_applied(&insert.oid) { + return Ok(()); } if let Some(statements) = self.statements.get(&insert.oid) { // Convert TupleData into a Bind message. We can now insert that tuple // using a prepared statement. - let bind = insert.tuple_data.to_bind(&statements.upsert.name); - - // Upserts are idempotent. Even if we rewind the stream, - // we are able to replay changes we already applied safely. - if let Some(upsert) = statements.upsert.insert() { - let upsert = Insert::new(upsert); - let val = upsert.shard(&self.sharding_schema, Some(&bind))?; - self.send(&val, &bind).await?; + let mut context = + StreamContext::new(&self.cluster, &insert.tuple_data, statements.insert.parse()); + let bind = context.bind().clone(); + let shard = context.shard()?; + + self.send(&shard, &bind).await?; + } + + // Update table LSN. + self.table_lsns.insert(insert.oid, self.lsn); + + Ok(()) + } + + async fn update(&mut self, update: XLogUpdate) -> Result<(), Error> { + if self.lsn_applied(&update.oid) { + return Ok(()); + } + + if let Some(statements) = self.statements.get(&update.oid) { + // Primary key update. + // + // Convert it into a delete of the old row + // and an insert with the new row. + if let Some(key) = update.key { + let delete = XLogDelete { + key: Some(key), + oid: update.oid, + old: None, + }; + let insert = XLogInsert { + xid: None, + oid: update.oid, + tuple_data: update.new, + }; + self.delete(delete).await?; + self.insert(insert).await?; + } else { + let mut context = + StreamContext::new(&self.cluster, &update.new, statements.update.parse()); + let bind = context.bind().clone(); + let shard = context.shard()?; + + self.send(&shard, &bind).await?; } + } - // Update table LSN. - self.table_lsns.insert(insert.oid, self.lsn); + // Update table LSN. + self.table_lsns.insert(update.oid, self.lsn); + + Ok(()) + } + + async fn delete(&mut self, delete: XLogDelete) -> Result<(), Error> { + if self.lsn_applied(&delete.oid) { + return Ok(()); + } + + if let Some(statements) = self.statements.get(&delete.oid) { + // Convert TupleData into a Bind message. We can now insert that tuple + // using a prepared statement. + if let Some(key) = delete.key_non_null() { + let mut context = + StreamContext::new(&self.cluster, &key, statements.delete.parse()); + let bind = context.bind().clone(); + let shard = context.shard()?; + + self.send(&shard, &bind).await?; + } } + // Update table LSN. + self.table_lsns.insert(delete.oid, self.lsn); + Ok(()) } + fn lsn_applied(&self, oid: &i32) -> bool { + if let Some(table_lsn) = self.table_lsns.get(oid) { + // Don't apply change if table is ahead. + if self.lsn < *table_lsn { + return true; + } + } + + false + } + // Handle Commit message. // // Send Sync to all shards, ensuring they close the transaction. @@ -280,7 +374,7 @@ impl StreamSubscriber { } 'Z' => break, '2' | 'C' => continue, - c => return Err(Error::OutOfSync(c)), + c => return Err(Error::CommitOutOfSync(c)), } } } @@ -303,12 +397,40 @@ impl StreamSubscriber { if let Some(table) = table { // Prepare queries for this table. Prepared statements // are much faster. - let insert = Statement::new(&table.insert(false))?; - let upsert = Statement::new(&table.insert(true))?; + + table.valid()?; + + let dest_key = Key { + schema: table.table.destination_schema().to_string(), + name: table.table.destination_name().to_string(), + }; + + if self.partitioned_dedup.contains(&dest_key) { + debug!("queries for table {} already prepared", dest_key); + return Ok(()); + } + + let insert = Statement::new(&table.insert(false), self.query_parser_engine)?; + let upsert = Statement::new(&table.insert(true), self.query_parser_engine)?; + let update = Statement::new(&table.update(), self.query_parser_engine)?; + let delete = Statement::new(&table.delete(), self.query_parser_engine)?; for server in &mut self.connections { + for stmt in &[&insert, &upsert, &update, &delete] { + debug!("preparing \"{}\" [{}]", stmt.query(), server.addr()); + } + server - .send(&vec![insert.parse().into(), upsert.parse().into(), Sync.into()].into()) + .send( + &vec![ + insert.parse().clone().into(), + upsert.parse().clone().into(), + update.parse().clone().into(), + delete.parse().clone().into(), + Sync.into(), + ] + .into(), + ) .await?; } @@ -325,7 +447,7 @@ impl StreamSubscriber { } 'Z' => break, '1' => continue, - c => return Err(Error::OutOfSync(c)), + c => return Err(Error::RelationOutOfSync(c)), } } } @@ -335,13 +457,15 @@ impl StreamSubscriber { Statements { insert, upsert, - update: Statement::default(), + update, + delete, }, ); // Only record tables we expect to stream changes for. self.table_lsns.insert(relation.oid, table.lsn.lsn); self.relations.insert(relation.oid, relation); + self.partitioned_dedup.insert(dest_key); } Ok(()) @@ -362,6 +486,8 @@ impl StreamSubscriber { if let Some(payload) = xlog.payload() { match payload { XLogPayload::Insert(insert) => self.insert(insert).await?, + XLogPayload::Update(update) => self.update(update).await?, + XLogPayload::Delete(delete) => self.delete(delete).await?, XLogPayload::Commit(commit) => { self.commit(commit).await?; status_update = Some(self.status_update()); diff --git a/pgdog/src/backend/replication/mod.rs b/pgdog/src/backend/replication/mod.rs index 430c7d457..0dae9d135 100644 --- a/pgdog/src/backend/replication/mod.rs +++ b/pgdog/src/backend/replication/mod.rs @@ -2,10 +2,12 @@ pub mod buffer; pub mod config; pub mod error; pub mod logical; +pub mod sharded_schema; pub mod sharded_tables; pub use buffer::Buffer; pub use config::ReplicationConfig; pub use error::Error; pub use logical::*; +pub use sharded_schema::*; pub use sharded_tables::{ShardedColumn, ShardedTables}; diff --git a/pgdog/src/backend/replication/sharded_schema.rs b/pgdog/src/backend/replication/sharded_schema.rs new file mode 100644 index 000000000..272d7f99b --- /dev/null +++ b/pgdog/src/backend/replication/sharded_schema.rs @@ -0,0 +1,65 @@ +use pgdog_config::sharding::ShardedSchema; +use std::{collections::HashMap, ops::Deref, sync::Arc}; + +use crate::frontend::router::parser::Schema; + +#[derive(Debug, Clone)] +pub struct ShardedSchemas { + inner: Arc, +} + +#[derive(Debug)] +struct Inner { + schemas: HashMap, + default_mapping: Option, +} + +impl Inner { + fn new(schemas: Vec) -> Self { + let without_default = schemas + .iter() + .filter(|schema| !schema.is_default()) + .cloned(); + let default_mapping = schemas.iter().find(|schema| schema.is_default()).cloned(); + + Self { + schemas: without_default + .into_iter() + .map(|schema| (schema.name().to_string(), schema)) + .collect(), + default_mapping, + } + } +} + +impl Deref for ShardedSchemas { + type Target = HashMap; + + fn deref(&self) -> &Self::Target { + &self.inner.schemas + } +} + +impl ShardedSchemas { + pub fn get<'a>(&self, schema: Option>) -> Option<&ShardedSchema> { + if let Some(schema) = schema { + if let Some(schema) = self.inner.schemas.get(schema.name) { + return Some(schema); + } + } + + self.inner.default_mapping.as_ref() + } + + pub fn new(schemas: Vec) -> Self { + Self { + inner: Arc::new(Inner::new(schemas)), + } + } +} + +impl Default for ShardedSchemas { + fn default() -> Self { + Self::new(vec![]) + } +} diff --git a/pgdog/src/backend/replication/sharded_tables.rs b/pgdog/src/backend/replication/sharded_tables.rs index 805f50c54..33680f296 100644 --- a/pgdog/src/backend/replication/sharded_tables.rs +++ b/pgdog/src/backend/replication/sharded_tables.rs @@ -1,15 +1,20 @@ //! Tables sharded in the database. +use pgdog_config::OmnishardedTable; + use crate::{ config::{DataType, ShardedTable}, - frontend::router::sharding::Mapping, + frontend::router::{parser::Column, sharding::Mapping}, net::messages::Vector, }; -use std::{collections::HashSet, sync::Arc}; +use std::{ + collections::{HashMap, HashSet}, + sync::Arc, +}; #[derive(Default, Debug)] struct Inner { tables: Vec, - omnisharded: HashSet, + omnisharded: HashMap, // Name <-> sticky routing /// This is set only if we have the same sharding scheme /// across all tables, i.e., 3 tables with the same data type /// and list/range/hash function. @@ -46,7 +51,7 @@ impl From<&[ShardedTable]> for ShardedTables { } impl ShardedTables { - pub fn new(tables: Vec, omnisharded_tables: Vec) -> Self { + pub fn new(tables: Vec, omnisharded_tables: Vec) -> Self { let mut common_mapping = HashSet::new(); for table in &tables { common_mapping.insert(( @@ -69,7 +74,10 @@ impl ShardedTables { Self { inner: Arc::new(Inner { tables, - omnisharded: omnisharded_tables.into_iter().collect(), + omnisharded: omnisharded_tables + .into_iter() + .map(|table| (table.name, table.sticky_routing)) + .collect(), common_mapping, }), } @@ -79,7 +87,7 @@ impl ShardedTables { &self.inner.tables } - pub fn omnishards(&self) -> &HashSet { + pub fn omnishards(&self) -> &HashMap { &self.inner.omnisharded } @@ -95,6 +103,35 @@ impl ShardedTables { .find(|t| t.name.as_deref() == Some(name)) } + /// Determine if the column is sharded and return its data type, + /// as declared in the schema. + pub fn get_table(&self, column: Column<'_>) -> Option<&ShardedTable> { + // Only fully-qualified columns can be matched. + let table = column.table()?; + + for candidate in &self.inner.tables { + if let Some(table_name) = candidate.name.as_ref() { + if !table.name_match(table_name) { + continue; + } + } + + if let Some(schema_name) = candidate.schema.as_ref() { + if let Some(schema) = table.schema() { + if schema.name != schema_name { + continue; + } + } + } + + if column.name == candidate.column { + return Some(candidate); + } + } + + None + } + /// Find out which column (if any) is sharded in the given table. pub fn sharded_column(&self, table: &str, columns: &[&str]) -> Option { let with_names = self diff --git a/pgdog/src/backend/schema/invalid_index.sql b/pgdog/src/backend/schema/invalid_index.sql new file mode 100644 index 000000000..b7e5688c3 --- /dev/null +++ b/pgdog/src/backend/schema/invalid_index.sql @@ -0,0 +1,28 @@ +SELECT + n.nspname AS schema_name, + c.relname AS table_name, + i.relname AS index_name, + idx.indisvalid + FROM + pg_index idx + JOIN pg_class i ON i.oid = idx.indexrelid + JOIN pg_class c ON c.oid = idx.indrelid + JOIN pg_namespace n ON n.oid = c.relnamespace + WHERE + i.relname = $1; + + Or if you want to include the schema name: + + SELECT + n.nspname AS schema_name, + c.relname AS table_name, + i.relname AS index_name, + idx.indisvalid + FROM + pg_index idx + JOIN pg_class i ON i.oid = idx.indexrelid + JOIN pg_class c ON c.oid = idx.indrelid + JOIN pg_namespace n ON n.oid = c.relnamespace + WHERE + i.relname = $1 + AND n.nspname = $2; diff --git a/pgdog/src/backend/schema/mod.rs b/pgdog/src/backend/schema/mod.rs index a148562c1..b8ad33eb1 100644 --- a/pgdog/src/backend/schema/mod.rs +++ b/pgdog/src/backend/schema/mod.rs @@ -58,6 +58,19 @@ impl Schema { }) } + #[cfg(test)] + pub(crate) fn from_parts( + search_path: Vec, + relations: HashMap<(String, String), Relation>, + ) -> Self { + Self { + inner: Arc::new(Inner { + search_path, + relations, + }), + } + } + /// Load schema from primary database. pub async fn from_cluster(cluster: &Cluster, shard: usize) -> Result { let mut primary = cluster.primary(shard, &Request::default()).await?; diff --git a/pgdog/src/backend/schema/relation.rs b/pgdog/src/backend/schema/relation.rs index ac3b5dc44..9f914a308 100644 --- a/pgdog/src/backend/schema/relation.rs +++ b/pgdog/src/backend/schema/relation.rs @@ -98,6 +98,24 @@ impl Relation { } } +#[cfg(test)] +impl Relation { + pub(crate) fn test_table(schema: &str, name: &str, columns: HashMap) -> Self { + Self { + schema: schema.into(), + name: name.into(), + type_: "table".into(), + owner: String::new(), + persistence: String::new(), + access_method: String::new(), + size: 0, + description: String::new(), + oid: 0, + columns, + } + } +} + #[cfg(test)] mod test { use crate::backend::pool::{test::pool, Request}; diff --git a/pgdog/src/backend/schema/setup.sql b/pgdog/src/backend/schema/setup.sql index ad6c5bc26..37dd2d33b 100644 --- a/pgdog/src/backend/schema/setup.sql +++ b/pgdog/src/backend/schema/setup.sql @@ -7,7 +7,8 @@ GRANT USAGE ON SCHEMA pgdog TO PUBLIC; CREATE TABLE IF NOT EXISTS pgdog.config ( shard INTEGER NOT NULL, shards INTEGER NOT NULL, - created_at TIMESTAMPTZ NOT NULL DEFAULT NOW() + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + PRIMARY KEY(shard, shards) ); CREATE OR REPLACE FUNCTION pgdog.config_trigger() RETURNS trigger AS $body$ diff --git a/pgdog/src/backend/schema/sync/error.rs b/pgdog/src/backend/schema/sync/error.rs index 9b2d4cae8..8b81b32d7 100644 --- a/pgdog/src/backend/schema/sync/error.rs +++ b/pgdog/src/backend/schema/sync/error.rs @@ -34,4 +34,7 @@ pub enum Error { #[error("missing entity in dump")] MissingEntity, + + #[error("publication \"{0}\" has no tables")] + PublicationNoTables(String), } diff --git a/pgdog/src/backend/schema/sync/pg_dump.rs b/pgdog/src/backend/schema/sync/pg_dump.rs index 1ca582137..fd94006cf 100644 --- a/pgdog/src/backend/schema/sync/pg_dump.rs +++ b/pgdog/src/backend/schema/sync/pg_dump.rs @@ -4,24 +4,34 @@ use std::{ops::Deref, str::from_utf8}; use lazy_static::lazy_static; use pg_query::{ - protobuf::{AlterTableType, ConstrType, ParseResult}, + protobuf::{AlterTableType, ConstrType, ObjectType, ParseResult}, NodeEnum, }; +use pgdog_config::QueryParserEngine; use regex::Regex; use tracing::{info, trace, warn}; use super::{progress::Progress, Error}; use crate::{ - backend::{ - pool::{Address, Request}, - replication::publisher::PublicationTable, - schema::sync::progress::Item, - Cluster, - }, + backend::{self, pool::Request, replication::publisher::PublicationTable, Cluster}, config::config, frontend::router::parser::{sequence::Sequence, Column, Table}, }; +fn deparse_node(node: NodeEnum) -> Result { + match config().config.general.query_parser_engine { + QueryParserEngine::PgQueryProtobuf => node.deparse(), + QueryParserEngine::PgQueryRaw => node.deparse_raw(), + } +} + +fn parse(query: &str) -> Result { + match config().config.general.query_parser_engine { + QueryParserEngine::PgQueryProtobuf => pg_query::parse(query), + QueryParserEngine::PgQueryRaw => pg_query::parse_raw(query), + } +} + use tokio::process::Command; #[derive(Debug, Clone)] @@ -38,16 +48,26 @@ impl PgDump { } } + fn clean(source: &str) -> String { + lazy_static! { + static ref CLEANUP_RE: Regex = Regex::new(r"(?m)^\\(?:un)?restrict.*\n?").unwrap(); + } + let cleaned = CLEANUP_RE.replace_all(source, ""); + + cleaned.to_string() + } + /// Dump schema from source cluster. - pub async fn dump(&self) -> Result, Error> { + pub async fn dump(&self) -> Result { let mut comparison: Vec = vec![]; let addr = self .source .shards() .first() .ok_or(Error::NoDatabases)? - .primary_or_replica(&Request::default()) - .await? + .pools() + .first() + .ok_or(Error::NoDatabases)? .addr() .clone(); @@ -70,58 +90,15 @@ impl PgDump { server.addr(), self.source.name() ); - continue; } } - let mut result = vec![]; - info!( - "dumping schema for {} tables [{}, {}]", - comparison.len(), - addr, - self.source.name() - ); - - let progress = Progress::new(comparison.len()); - - for table in comparison { - let cmd = PgDumpCommand { - table: table.name.clone(), - schema: table.schema.clone(), - address: addr.clone(), - }; - progress.next(Item::TableDump { - schema: table.schema.clone(), - name: table.name.clone(), - }); - - let dump = cmd.execute().await?; - result.push(dump); - } - - progress.done(); - - Ok(result) - } -} - -struct PgDumpCommand { - table: String, - schema: String, - address: Address, -} - -impl PgDumpCommand { - fn clean(source: &str) -> String { - lazy_static! { - static ref CLEANUP_RE: Regex = Regex::new(r"(?m)^\\(?:un)?restrict.*\n?").unwrap(); + if comparison.is_empty() { + return Err(Error::PublicationNoTables(self.publication.clone())); } - let cleaned = CLEANUP_RE.replace_all(source, ""); - cleaned.to_string() - } + info!("dumping schema [{}, {}]", comparison.len(), addr,); - async fn execute(&self) -> Result { let config = config(); let pg_dump_path = config .config @@ -131,20 +108,16 @@ impl PgDumpCommand { .unwrap_or("pg_dump"); let output = Command::new(pg_dump_path) - .arg("-t") - .arg(&self.table) - .arg("-n") - .arg(&self.schema) .arg("--schema-only") .arg("-h") - .arg(&self.address.host) + .arg(&addr.host) .arg("-p") - .arg(self.address.port.to_string()) + .arg(addr.port.to_string()) .arg("-U") - .arg(&self.address.user) - .env("PGPASSWORD", &self.address.password) + .arg(&addr.user) + .env("PGPASSWORD", &addr.password) .arg("-d") - .arg(&self.address.database_name) + .arg(&addr.database_name) .output() .await?; @@ -159,13 +132,11 @@ impl PgDumpCommand { let cleaned = Self::clean(&original); trace!("[pg_dump (clean)] {}", cleaned); - let stmts = pg_query::parse(&cleaned)?.protobuf; + let stmts = parse(&cleaned)?.protobuf; Ok(PgDumpOutput { stmts, original: cleaned, - table: self.table.clone(), - schema: self.schema.clone(), }) } } @@ -174,8 +145,6 @@ impl PgDumpCommand { pub struct PgDumpOutput { stmts: ParseResult, original: String, - pub table: String, - pub schema: String, } #[derive(Debug, Copy, Clone, PartialEq)] @@ -190,16 +159,17 @@ pub enum Statement<'a> { Index { table: Table<'a>, name: &'a str, - sql: &'a str, + sql: String, }, Table { table: Table<'a>, - sql: &'a str, + sql: String, }, Other { - sql: &'a str, + sql: String, + idempotent: bool, }, SequenceOwner { @@ -221,7 +191,7 @@ impl<'a> Deref for Statement<'a> { Self::Index { sql, .. } => sql, Self::Table { sql, .. } => sql, Self::SequenceOwner { sql, .. } => sql, - Self::Other { sql } => sql, + Self::Other { sql, .. } => sql, Self::SequenceSetMax { sql, .. } => sql.as_str(), } } @@ -229,7 +199,19 @@ impl<'a> Deref for Statement<'a> { impl<'a> From<&'a str> for Statement<'a> { fn from(value: &'a str) -> Self { - Self::Other { sql: value } + Self::Other { + sql: value.to_string(), + idempotent: true, + } + } +} + +impl<'a> From for Statement<'a> { + fn from(value: String) -> Self { + Self::Other { + sql: value, + idempotent: true, + } } } @@ -252,21 +234,44 @@ impl PgDumpOutput { if let Some(ref node) = node.node { match node { NodeEnum::CreateStmt(stmt) => { + let sql = { + let mut stmt = stmt.clone(); + stmt.if_not_exists = true; + deparse_node(NodeEnum::CreateStmt(stmt))? + }; if state == SyncState::PreData { // CREATE TABLE is always good. let table = stmt.relation.as_ref().map(Table::from).unwrap_or_default(); - result.push(Statement::Table { - table, - sql: original, - }); + result.push(Statement::Table { table, sql }); } } - NodeEnum::CreateSeqStmt(_) => { + NodeEnum::CreateSeqStmt(stmt) => { + let mut stmt = stmt.clone(); + stmt.if_not_exists = true; + let sql = deparse_node(NodeEnum::CreateSeqStmt(stmt))?; if state == SyncState::PreData { // Bring sequences over. - result.push(original.into()); + result.push(sql.into()); + } + } + + NodeEnum::CreateExtensionStmt(stmt) => { + let mut stmt = stmt.clone(); + stmt.if_not_exists = true; + let sql = deparse_node(NodeEnum::CreateExtensionStmt(stmt))?; + if state == SyncState::PreData { + result.push(sql.into()); + } + } + + NodeEnum::CreateSchemaStmt(stmt) => { + let mut stmt = stmt.clone(); + stmt.if_not_exists = true; + let sql = deparse_node(NodeEnum::CreateSchemaStmt(stmt))?; + if state == SyncState::PreData { + result.push(sql.into()); } } @@ -287,24 +292,115 @@ impl PgDumpOutput { | ConstrType::ConstrNull ) { if state == SyncState::PreData { - result.push(original.into()); + result.push(Statement::Other { + sql: original.to_string(), + idempotent: false, + }); } } else if state == SyncState::PostData { - result.push(original.into()); + result.push(Statement::Other { + sql: original.to_string(), + idempotent: false, + }); + } + } + } + } + AlterTableType::AtAttachPartition => { + // Index partitions need to be attached to indexes, + // which we create in the post-data step. + match stmt.objtype() { + ObjectType::ObjectIndex => { + if state == SyncState::PostData { + result.push(Statement::Other { + sql: original.to_string(), + idempotent: false, + }); + } + } + + _ => { + if state == SyncState::PreData { + result.push(Statement::Other { + sql: original.to_string(), + idempotent: false, + }); } } } } + AlterTableType::AtColumnDefault => { if state == SyncState::PreData { result.push(original.into()) } } - AlterTableType::AtChangeOwner => { - continue; // Don't change owners, for now. + + AlterTableType::AtAddIdentity => { + if state == SyncState::Cutover { + if let Some(ref node) = cmd.def { + if let Some(NodeEnum::Constraint( + ref constraint, + )) = node.node + { + for option in &constraint.options { + if let Some(NodeEnum::DefElem( + ref elem, + )) = option.node + { + if elem.defname == "sequence_name" { + if let Some(ref node) = elem.arg + { + if let Some( + NodeEnum::List( + ref list, + ), + ) = node.node + { + let table = Table::try_from(list).map_err(|_| Error::MissingEntity)?; + let sequence = + Sequence::from( + table, + ); + let table = stmt + .relation + .as_ref() + .map(Table::from); + let schema = table + .and_then( + |table| { + table.schema + }, + ); + let table = table.map( + |table| table.name, + ); + let column = Column { + table, + name: cmd + .name + .as_str(), + schema, + }; + let sql = sequence.setval_from_column(&column).map_err(|_| Error::MissingEntity)?; + + result.push(Statement::SequenceSetMax { sequence, sql }); + } + } + } + } + } + } + } + } else if state == SyncState::PreData { + result.push(original.into()); + } } + // AlterTableType::AtChangeOwner => { + // continue; // Don't change owners, for now. + // } _ => { - if state == SyncState::PostData { + if state == SyncState::PreData { result.push(original.into()); } } @@ -313,6 +409,21 @@ impl PgDumpOutput { } } + NodeEnum::CreateTrigStmt(stmt) => { + let mut stmt = stmt.clone(); + stmt.replace = true; + + if state == SyncState::PreData { + result.push(deparse_node(NodeEnum::CreateTrigStmt(stmt))?.into()); + } + } + + // Skip these. + NodeEnum::CreatePublicationStmt(_) + | NodeEnum::CreateSubscriptionStmt(_) + | NodeEnum::AlterPublicationStmt(_) + | NodeEnum::AlterSubscriptionStmt(_) => (), + NodeEnum::AlterSeqStmt(stmt) => { if matches!(state, SyncState::PreData | SyncState::Cutover) { let sequence = stmt @@ -331,7 +442,7 @@ impl PgDumpOutput { sequence, sql: original, }); - } else { + } else if state == SyncState::Cutover { let sql = sequence .setval_from_column(&column) .map_err(|_| Error::MissingEntity)?; @@ -342,21 +453,90 @@ impl PgDumpOutput { NodeEnum::IndexStmt(stmt) => { if state == SyncState::PostData { + let sql = { + let mut stmt = stmt.clone(); + stmt.concurrent = stmt + .relation + .as_ref() + .map(|relation| relation.inh) // ONLY used for partitioned tables, which can't be created concurrently. + .unwrap_or(false); + stmt.if_not_exists = true; + deparse_node(NodeEnum::IndexStmt(stmt))? + }; + let table = stmt.relation.as_ref().map(Table::from).unwrap_or_default(); + result.push(Statement::Index { table, name: stmt.idxname.as_str(), - sql: original, + sql, + }); + } + } + + NodeEnum::ViewStmt(stmt) => { + let mut stmt = stmt.clone(); + stmt.replace = true; + + if state == SyncState::PreData { + result.push(Statement::Other { + sql: deparse_node(NodeEnum::ViewStmt(stmt))?, + idempotent: true, + }); + } + } + + NodeEnum::CreateTableAsStmt(stmt) => { + let mut stmt = stmt.clone(); + stmt.if_not_exists = true; + + if state == SyncState::PreData { + result.push(Statement::Other { + sql: deparse_node(NodeEnum::CreateTableAsStmt(stmt))?, + idempotent: true, + }); + } + } + + NodeEnum::CreateFunctionStmt(stmt) => { + let mut stmt = stmt.clone(); + stmt.replace = true; + + if state == SyncState::PreData { + result.push(Statement::Other { + sql: deparse_node(NodeEnum::CreateFunctionStmt(stmt))?, + idempotent: true, + }); + } + } + + NodeEnum::AlterOwnerStmt(stmt) => { + if stmt.object_type() != ObjectType::ObjectPublication + && state == SyncState::PreData + { + result.push(Statement::Other { + sql: original.to_string(), + idempotent: true, + }); + } + } + + NodeEnum::CreateEnumStmt(_) + | NodeEnum::CreateDomainStmt(_) + | NodeEnum::CompositeTypeStmt(_) => { + if state == SyncState::PreData { + result.push(Statement::Other { + sql: original.to_owned(), + idempotent: false, }); } } NodeEnum::VariableSetStmt(_) => continue, NodeEnum::SelectStmt(_) => continue, - _ => { - if state == SyncState::PostData { + if state == SyncState::PreData { result.push(original.into()); } } @@ -381,19 +561,37 @@ impl PgDumpOutput { let mut primary = shard.primary(&Request::default()).await?; info!( - "syncing schema for \"{}\".\"{}\" into shard {} [{}, {}]", - self.schema, - self.table, + "syncing schema into shard {} [{}, {}]", num, primary.addr(), dest.name() ); - let progress = Progress::new(stmts.len()); + let mut progress = Progress::new(stmts.len()); for stmt in &stmts { progress.next(stmt); if let Err(err) = primary.execute(stmt.deref()).await { + if let backend::Error::ExecutionError(ref err) = err { + let code = &err.code; + + if let Statement::Other { idempotent, .. } = stmt { + if !idempotent { + if matches!(code.as_str(), "42P16" | "42710" | "42809" | "42P07") { + warn!("entity already exists, skipping"); + continue; + } else if !ignore_errors { + return Err(Error::Backend(backend::Error::ExecutionError( + err.clone(), + ))); + } else { + warn!("skipping: {}", err); + } + } + } + } else { + return Err(err.into()); + } if ignore_errors { warn!("skipping: {}", err); } else { @@ -410,98 +608,12 @@ impl PgDumpOutput { #[cfg(test)] mod test { - use crate::backend::server::test::test_server; - use super::*; #[tokio::test] async fn test_pg_dump_execute() { - let mut server = test_server().await; - - let queries = vec![ - "DROP PUBLICATION IF EXISTS test_pg_dump_execute", - "CREATE TABLE IF NOT EXISTS test_pg_dump_execute(id BIGSERIAL PRIMARY KEY, email VARCHAR UNIQUE, created_at TIMESTAMPTZ)", - "CREATE INDEX ON test_pg_dump_execute USING btree(created_at)", - "CREATE TABLE IF NOT EXISTS test_pg_dump_execute_fk(fk BIGINT NOT NULL REFERENCES test_pg_dump_execute(id), meta JSONB)", - "CREATE PUBLICATION test_pg_dump_execute FOR TABLE test_pg_dump_execute, test_pg_dump_execute_fk" - ]; - - for query in queries { - server.execute(query).await.unwrap(); - } - - let output = PgDumpCommand { - table: "test_pg_dump_execute".into(), - schema: "pgdog".into(), - address: server.addr().clone(), - } - .execute() - .await - .unwrap(); - - let output_pre = output.statements(SyncState::PreData).unwrap(); - let output_post = output.statements(SyncState::PostData).unwrap(); - let output_cutover = output.statements(SyncState::Cutover).unwrap(); - - let mut dest = test_server().await; - dest.execute("DROP SCHEMA IF EXISTS test_pg_dump_execute_dest CASCADE") - .await - .unwrap(); - - dest.execute("CREATE SCHEMA test_pg_dump_execute_dest") - .await - .unwrap(); - dest.execute("SET search_path TO test_pg_dump_execute_dest, public") - .await - .unwrap(); - - for stmt in output_pre { - // Hack around us using the same database as destination. - // I know, not very elegant. - let stmt = stmt.replace("pgdog.", "test_pg_dump_execute_dest."); - dest.execute(stmt).await.unwrap(); - } - - for i in 0..5 { - let id = dest.fetch_all::("INSERT INTO test_pg_dump_execute_dest.test_pg_dump_execute VALUES (DEFAULT, 'test@test', NOW()) RETURNING id") - .await - .unwrap(); - assert_eq!(id[0], i + 1); // Sequence has made it over. - - // Unique index didn't make it over. - } - - dest.execute("DELETE FROM test_pg_dump_execute_dest.test_pg_dump_execute") - .await - .unwrap(); - - for stmt in output_post { - let stmt = stmt.replace("pgdog.", "test_pg_dump_execute_dest."); - dest.execute(stmt).await.unwrap(); - } - - let q = "INSERT INTO test_pg_dump_execute_dest.test_pg_dump_execute VALUES (DEFAULT, 'test@test', NOW()) RETURNING id"; - assert!(dest.execute(q).await.is_ok()); - let err = dest.execute(q).await.err().unwrap(); - assert!(err.to_string().contains( - r#"duplicate key value violates unique constraint "test_pg_dump_execute_email_key""# - )); // Unique index made it over. - - assert_eq!(output_cutover.len(), 1); - for stmt in output_cutover { - let stmt = stmt.replace("pgdog.", "test_pg_dump_execute_dest."); - assert!(stmt.starts_with("SELECT setval('")); - dest.execute(stmt).await.unwrap(); - } - - dest.execute("DROP SCHEMA test_pg_dump_execute_dest CASCADE") - .await - .unwrap(); - - server - .execute("DROP TABLE test_pg_dump_execute CASCADE") - .await - .unwrap(); + let cluster = Cluster::new_test_single_shard(&config()); + let _pg_dump = PgDump::new(&cluster, "test_pg_dump_execute"); } #[test] @@ -557,6 +669,43 @@ ALTER TABLE ONLY public.users \unrestrict nu6jB5ogH2xGMn2dB3dMyMbSZ2PsVDqB2IaWK6zZVjngeba0UrnmxMy6s63SwzR "#; - let _parse = pg_query::parse(&PgDumpCommand::clean(&dump)).unwrap(); + let _parse = pg_query::parse(&PgDump::clean(&dump)).unwrap(); + } + + #[test] + fn test_generated_identity() { + let q = "ALTER TABLE public.users ALTER COLUMN id ADD GENERATED ALWAYS AS IDENTITY ( + SEQUENCE NAME public.users_id_seq + START WITH 1 + INCREMENT BY 1 + NO MINVALUE + NO MAXVALUE + CACHE 1 + );"; + let parse = pg_query::parse(q).unwrap(); + let output = PgDumpOutput { + stmts: parse.protobuf, + original: q.to_string(), + }; + let statements = output.statements(SyncState::Cutover).unwrap(); + match statements.first() { + Some(Statement::SequenceSetMax { sequence, sql }) => { + assert_eq!(sequence.table.name, "users_id_seq"); + assert_eq!( + sequence.table.schema().map(|schema| schema.name), + Some("public") + ); + assert_eq!( + sql, + r#"SELECT setval('"public"."users_id_seq"', COALESCE((SELECT MAX("id") FROM "public"."users"), 1), true);"# + ); + } + + _ => panic!("not a set sequence max"), + } + let statements = output.statements(SyncState::PreData).unwrap(); + assert!(!statements.is_empty()); + let statements = output.statements(SyncState::PostData).unwrap(); + assert!(statements.is_empty()); } } diff --git a/pgdog/src/backend/schema/sync/progress.rs b/pgdog/src/backend/schema/sync/progress.rs index ec47a1221..351a45943 100644 --- a/pgdog/src/backend/schema/sync/progress.rs +++ b/pgdog/src/backend/schema/sync/progress.rs @@ -1,12 +1,10 @@ use std::fmt::Display; -use std::sync::Arc; -use tokio::sync::Notify; +use tokio::sync::mpsc; use tokio::time::Instant; use tokio::{select, spawn}; use tracing::info; use super::Statement; -use parking_lot::Mutex; #[derive(Clone)] pub enum Item { @@ -23,17 +21,11 @@ pub enum Item { schema: String, name: String, }, - // SequenceOwner { - // sequence: String, - // owner: String, - // }, Other { sql: String, }, } -// Remove pg_dump comments. -// Only for displaying purposes! Don't use for executing queries. fn no_comments(sql: &str) -> String { let mut output = String::new(); for line in sql.lines() { @@ -44,7 +36,23 @@ fn no_comments(sql: &str) -> String { output.push('\n'); } - output.trim().to_string() + let result = output.trim().replace("\n", " "); + + let mut prev_space = false; + let mut collapsed = String::new(); + for c in result.chars() { + if c == ' ' { + if !prev_space { + collapsed.push(c); + prev_space = true; + } + } else { + collapsed.push(c); + prev_space = false; + } + } + + collapsed } impl Default for Item { @@ -87,16 +95,12 @@ impl Item { impl From<&Statement<'_>> for Item { fn from(value: &Statement<'_>) -> Self { match value { - Statement::Index { table, name, .. } => Item::Index { - schema: table.schema.as_ref().unwrap_or(&"").to_string(), - table: table.name.to_string(), - name: name.to_string(), - }, + Statement::Index { sql, .. } => Item::Other { sql: sql.clone() }, Statement::Table { table, .. } => Item::Table { schema: table.schema.as_ref().unwrap_or(&"").to_string(), name: table.name.to_string(), }, - Statement::Other { sql } => Item::Other { + Statement::Other { sql, .. } => Item::Other { sql: sql.to_string(), }, Statement::SequenceOwner { sql, .. } => Item::Other { @@ -109,85 +113,54 @@ impl From<&Statement<'_>> for Item { } } -struct ItemTracker { - item: Item, - timer: Instant, -} - -impl ItemTracker { - fn new() -> Self { - Self { - item: Item::default(), - timer: Instant::now(), - } - } -} - -struct Comms { - updated: Notify, - shutdown: Notify, -} -impl Comms { - fn new() -> Self { - Self { - updated: Notify::new(), - shutdown: Notify::new(), - } - } +enum Message { + Next(Item), + Shutdown, } #[derive(Clone)] pub struct Progress { - item: Arc>, - comms: Arc, - total: usize, + tx: mpsc::UnboundedSender, + timer: Instant, } impl Progress { pub fn new(total: usize) -> Self { - let me = Self { - item: Arc::new(Mutex::new(ItemTracker::new())), - comms: Arc::new(Comms::new()), - total, - }; - - let task = me.clone(); + let (tx, rx) = mpsc::unbounded_channel(); + let timer = Instant::now(); spawn(async move { - task.listen().await; + Self::listen(rx, total).await; }); - me + Self { tx, timer } } - pub fn done(&self) { - let elapsed = self.item.lock().timer.elapsed(); - + pub fn done(&mut self) { + let elapsed = self.timer.elapsed(); info!("finished in {:.3}s", elapsed.as_secs_f64()); + self.timer = Instant::now(); } pub fn next(&self, item: impl Into) { - { - let mut guard = self.item.lock(); - guard.item = item.into().clone(); - guard.timer = Instant::now(); - } - self.comms.updated.notify_one(); + let _ = self.tx.send(Message::Next(item.into())); } - async fn listen(&self) { + async fn listen(mut rx: mpsc::UnboundedReceiver, total: usize) { let mut counter = 1; loop { select! { - _ = self.comms.updated.notified() => { - let item = self.item.lock().item.clone(); - info!("[{}/{}] {} {}", counter, self.total, item.action(), item); - counter += 1; - } - - _ = self.comms.shutdown.notified() => { - break; + msg = rx.recv() => { + match msg { + Some(Message::Next(item)) => { + info!("[{}/{}] {} {}", counter, total, item.action(), item); + counter += 1; + } + Some(Message::Shutdown) | None => { + break; + } + } } } } @@ -196,6 +169,6 @@ impl Progress { impl Drop for Progress { fn drop(&mut self) { - self.comms.shutdown.notify_one(); + let _ = self.tx.send(Message::Shutdown); } } diff --git a/pgdog/src/backend/server.rs b/pgdog/src/backend/server.rs index 1a27b721c..4ed0d5a1f 100644 --- a/pgdog/src/backend/server.rs +++ b/pgdog/src/backend/server.rs @@ -1,4 +1,5 @@ //! PostgreSQL server connection. + use std::time::Duration; use bytes::{BufMut, BytesMut}; @@ -12,18 +13,20 @@ use tokio::{ use tracing::{debug, error, info, trace, warn}; use super::{ - pool::Address, prepared_statements::HandleResult, Error, PreparedStatements, ServerOptions, - Stats, + pool::Address, prepared_statements::HandleResult, ConnectReason, DisconnectReason, Error, + PreparedStatements, ServerOptions, Stats, }; use crate::{ auth::{md5, scram::Client}, + backend::pool::stats::MemoryStats, + config::AuthType, frontend::ClientRequest, net::{ messages::{ hello::SslReply, Authentication, BackendKeyData, ErrorResponse, FromBytes, Message, ParameterStatus, Password, Protocol, Query, ReadyForQuery, Startup, Terminate, ToBytes, }, - Close, Parameter, ProtocolMessage, Sync, + Close, MessageBuffer, Parameter, ProtocolMessage, Sync, }, stats::memory::MemoryUsage, }; @@ -45,7 +48,6 @@ pub struct Server { stream: Option, id: BackendKeyData, params: Parameters, - startup_options: ServerOptions, changed_params: Parameters, client_params: Parameters, stats: Stats, @@ -58,7 +60,8 @@ pub struct Server { re_synced: bool, replication_mode: bool, pooler_mode: PoolerMode, - stream_buffer: BytesMut, + stream_buffer: MessageBuffer, + disconnect_reason: Option, } impl MemoryUsage for Server { @@ -78,14 +81,18 @@ impl MemoryUsage for Server { impl Server { /// Create new PostgreSQL server connection. - pub async fn connect(addr: &Address, options: ServerOptions) -> Result { + pub async fn connect( + addr: &Address, + options: ServerOptions, + connect_reason: ConnectReason, + ) -> Result { debug!("=> {}", addr); let stream = TcpStream::connect(addr.addr().await?).await?; tweak(&stream)?; + let cfg = config(); - let mut stream = Stream::plain(stream); + let mut stream = Stream::plain(stream, cfg.config.memory.net_buffer); - let cfg = config(); let tls_mode = cfg.config.general.tls_verify; // Only attempt TLS if not in Disabled mode @@ -119,7 +126,7 @@ impl Server { Ok(tls_stream) => { debug!("TLS handshake successful with {}", addr.host); let cipher = tokio_rustls::TlsStream::Client(tls_stream); - stream = Stream::tls(cipher); + stream = Stream::tls(cipher, cfg.config.memory.net_buffer); } Err(e) => { error!("TLS handshake failed with {:?} [{}]", e, addr); @@ -156,6 +163,7 @@ impl Server { // Perform authentication. let mut scram = Client::new(&addr.user, &addr.password); + let mut auth_type = AuthType::Trust; loop { let message = stream.read().await?; @@ -174,6 +182,7 @@ impl Server { stream.send_flush(&password).await?; } Authentication::Sasl(_) => { + auth_type = AuthType::Scram; let initial = Password::sasl_initial(&scram.first()?); stream.send_flush(&initial).await?; } @@ -188,6 +197,7 @@ impl Server { scram.server_last(&data)?; } Authentication::Md5(salt) => { + auth_type = AuthType::Md5; let client = md5::Client::new_salt(&addr.user, &addr.password, &salt)?; stream.send_flush(&client.response()).await?; } @@ -235,15 +245,20 @@ impl Server { let id = key_data.ok_or(Error::NoBackendKeyData)?; let params: Parameters = params.into(); - info!("new server connection [{}]", addr); + info!( + "new server connection [{}, auth: {}, reason: {}] {}", + addr, + auth_type, + connect_reason, + if stream.is_tls() { "πŸ”“" } else { "" }, + ); let mut server = Server { addr: addr.clone(), stream: Some(stream), id, - stats: Stats::connect(id, addr, ¶ms), + stats: Stats::connect(id, addr, ¶ms, &options, &cfg.config.memory), replication_mode: options.replication_mode(), - startup_options: options, params, changed_params: Parameters::default(), client_params: Parameters::default(), @@ -255,11 +270,11 @@ impl Server { in_transaction: false, re_synced: false, pooler_mode: PoolerMode::Transaction, - stream_buffer: BytesMut::with_capacity(1024), + stream_buffer: MessageBuffer::new(cfg.config.memory.message_buffer), + disconnect_reason: None, }; - server.stats.memory_used(server.memory_usage()); // Stream capacity. - server.stats().update(); + server.stats.memory_used(server.memory_stats()); // Stream capacity. Ok(server) } @@ -308,6 +323,10 @@ impl Server { HandleResult::Forward => [Some(message), None], }; + for message in queue.iter().flatten() { + trace!("{:#?} >>> [{}]", message, self.addr()); + } + for message in queue.into_iter().flatten() { match self.stream().send(message).await { Ok(sent) => self.stats.send(sent), @@ -333,18 +352,17 @@ impl Server { } /// Read a single message from the server. + /// + /// # Cancellation safety + /// + /// This method is cancel-safe. + /// pub async fn read(&mut self) -> Result { let message = loop { if let Some(message) = self.prepared_statements.state_mut().get_simulated() { return Ok(message.backend()); } - match self - .stream - .as_mut() - .unwrap() - .read_buf(&mut self.stream_buffer) - .await - { + match self.stream_buffer.read(self.stream.as_mut().unwrap()).await { Ok(message) => { let message = message.stream(self.streaming).backend(); match self.prepared_statements.forward(&message) { @@ -354,11 +372,16 @@ impl Server { } } Err(err) => { + if let Error::ProtocolOutOfSync = err { + // conservatively, we do not know for sure if this is recoverable + self.stats.state(State::Error); + } error!( - "{:?} got: {}, extended buffer: {:?}", + "{:?} got: {}, extended buffer: {:?}, state: {}", err, message.code(), self.prepared_statements.state(), + self.stats.state, ); return Err(err); } @@ -377,11 +400,11 @@ impl Server { match message.code() { 'Z' => { let now = Instant::now(); - self.stats.query(now); - self.stats.memory_used(self.memory_usage()); - let rfq = ReadyForQuery::from_bytes(message.payload())?; + self.stats.query(now, rfq.status == 'T'); + self.stats.memory_used(self.memory_stats()); + match rfq.status { 'I' => { self.in_transaction = false; @@ -397,7 +420,7 @@ impl Server { return Err(Error::UnexpectedTransactionStatus(status)); } } - + self.stream_buffer.shrink_to_fit(); self.streaming = false; } 'E' => { @@ -432,36 +455,91 @@ impl Server { _ => (), } + trace!("{:#?} <<< [{}]", message, self.addr()); + Ok(message) } /// Synchronize parameters between client and server. - pub async fn link_client(&mut self, params: &Parameters) -> Result { + pub async fn link_client( + &mut self, + id: &BackendKeyData, + params: &Parameters, + start_transaction: Option<&str>, + ) -> Result { // Sync application_name parameter // and update it in the stats. - let default_name = "PgDog"; - let server_name = self - .client_params - .get_default("application_name", default_name); - let client_name = params.get_default("application_name", default_name); - self.stats.link_client(client_name, server_name); + let server_name = self.client_params.get_default("application_name", "PgDog"); + let client_name = params.get_default("application_name", "PgDog"); + + self.stats.link_client(client_name, server_name, id); // Clear any params previously tracked by SET. self.changed_params.clear(); + let mut clear_params = false; // Compare client and server params. - if !params.identical(&self.client_params) { + let mut executed = if !params.identical(&self.client_params) { + // Construct client parameter SET queries. let tracked = params.tracked(); + // Construct RESET queries to reset any current params + // to their default values. let mut queries = self.client_params.reset_queries(); - queries.extend(tracked.set_queries()); + + // Combine both to create a new, fresh session state + // on this connection. + queries.extend(tracked.set_queries(false)); + + // Set state on the connection only if + // there are any params to change. if !queries.is_empty() { debug!("syncing {} params", queries.len()); + self.execute_batch(&queries).await?; + clear_params = true; } + + // Update params on this connection. self.client_params = tracked; - Ok(queries.len()) + + queries.len() + } else { + 0 + }; + + // Sync any parameters set inside the transaction. These will + // need to be revered on rollback or commited on commit. + if let Some(start_transaction) = start_transaction { + self.execute(start_transaction).await?; + let transaction_sets = params.set_queries(true); + + if !transaction_sets.is_empty() { + debug!("syncing {} in-transaction params", transaction_sets.len()); + + self.execute_batch(&transaction_sets).await?; + clear_params = true; + + self.client_params.copy_in_transaction(params); + } + + executed += transaction_sets.len(); + } + + if clear_params { + // We can receive ParameterStatus messages here, + // but we should ignore them since we are managing the session state. + self.changed_params.clear(); + } + + Ok(executed) + } + + // Handle COMMIT/ROLLBACK for in-transaction params tracking. + pub fn transaction_params_hook(&mut self, rollback: bool) { + if rollback { + self.client_params.rollback(); } else { - Ok(0) + self.client_params.commit(); } } @@ -481,18 +559,26 @@ impl Server { self.prepared_statements.done() && !self.in_transaction() } + /// Server hasn't finished sending or receiving a complete message. + pub fn io_in_progress(&self) -> bool { + self.stream + .as_ref() + .map(|stream| stream.io_in_progress()) + .unwrap_or(false) + } + /// Server can execute a query. pub fn in_sync(&self) -> bool { matches!( - self.stats.state, + self.stats().state, State::Idle | State::IdleInTransaction | State::TransactionError ) } - /// Server is done executing all queries and is + /// Server is done executing all queries and isz /// not inside a transaction. pub fn can_check_in(&self) -> bool { - self.stats.state == State::Idle + self.stats().state == State::Idle } /// Server hasn't sent all messages yet. @@ -500,8 +586,8 @@ impl Server { self.prepared_statements.has_more_messages() || self.streaming } - pub fn copy_mode(&self) -> bool { - self.prepared_statements.copy_mode() + pub fn in_copy_mode(&self) -> bool { + self.prepared_statements.in_copy_mode() } /// Server is still inside a transaction. @@ -513,7 +599,7 @@ impl Server { /// The server connection permanently failed. #[inline] pub fn error(&self) -> bool { - self.stats.state == State::Error + self.stats().state == State::Error } /// Did the schema change and prepared statements are broken. @@ -534,7 +620,7 @@ impl Server { /// Close the connection, don't do any recovery. pub fn force_close(&self) -> bool { - self.stats.state == State::ForceClose + self.stats().state == State::ForceClose || self.io_in_progress() } /// Server parameters. @@ -633,54 +719,50 @@ impl Server { debug!("running healthcheck \"{}\" [{}]", query, self.addr); self.execute(query).await?; + self.stats.healthcheck(); Ok(()) } /// Attempt to rollback the transaction on this server, if any has been started. - pub async fn rollback(&mut self) { + pub(super) async fn rollback(&mut self) -> Result<(), Error> { if self.in_transaction() { - if let Err(_err) = self.execute("ROLLBACK").await { - self.stats.state(State::Error); - } + self.execute("ROLLBACK").await?; self.stats.rollback(); } + // Reset any params left over. + self.transaction_params_hook(true); + if !self.done() { - self.stats.state(State::Error); + Err(Error::RollbackFailed) + } else { + Ok(()) } } - pub async fn drain(&mut self) { + /// Drain any remaining messages on the server connection, + /// attempting to return the connection into a synchronized state. + pub(super) async fn drain(&mut self) -> Result<(), Error> { while self.has_more_messages() { - if self.read().await.is_err() { - self.stats.state(State::Error); - break; - } + self.read().await?; } if !self.in_sync() { - if self - .send(&vec![ProtocolMessage::Sync(Sync)].into()) - .await - .is_err() - { - self.stats.state(State::Error); - return; - } + self.send(&vec![ProtocolMessage::Sync(Sync)].into()).await?; + while !self.in_sync() { - if self.read().await.is_err() { - self.stats.state(State::Error); - break; - } + self.read().await?; } - - self.re_synced = true; } + + self.re_synced = true; + Ok(()) } - pub async fn sync_prepared_statements(&mut self) -> Result<(), Error> { + /// Synchronize prepared statements from Postgres. + pub(super) async fn sync_prepared_statements(&mut self) -> Result<(), Error> { let names = self .fetch_all::("SELECT name FROM pg_prepared_statements") .await?; @@ -692,13 +774,14 @@ impl Server { debug!("prepared statements synchronized [{}]", self.addr()); let count = self.prepared_statements.len(); - self.stats_mut().set_prepared_statements(count); + self.stats.set_prepared_statements(count); + self.sync_prepared = false; Ok(()) } /// Close any prepared statements that exceed cache capacity. - pub fn ensure_prepared_capacity(&mut self) -> Vec { + pub(super) fn ensure_prepared_capacity(&mut self) -> Vec { let close = self.prepared_statements.ensure_capacity(); self.stats .close_many(close.len(), self.prepared_statements.len()); @@ -706,7 +789,7 @@ impl Server { } /// Close multiple prepared statements. - pub async fn close_many(&mut self, close: &[Close]) -> Result<(), Error> { + pub(super) async fn close_many(&mut self, close: &[Close]) -> Result<(), Error> { if close.is_empty() { return Ok(()); } @@ -749,10 +832,6 @@ impl Server { Ok(()) } - pub async fn reconnect(&self) -> Result { - Self::connect(&self.addr, self.startup_options.clone()).await - } - /// Reset error state caused by schema change. #[inline] pub fn reset_schema_changed(&mut self) { @@ -775,6 +854,13 @@ impl Server { self.re_synced } + #[inline] + pub fn disconnect_reason(&mut self, reason: DisconnectReason) { + if self.disconnect_reason.is_none() { + self.disconnect_reason = Some(reason); + } + } + /// Server connection unique identifier. #[inline] pub fn id(&self) -> &BackendKeyData { @@ -784,19 +870,19 @@ impl Server { /// How old this connection is. #[inline] pub fn age(&self, instant: Instant) -> Duration { - instant.duration_since(self.stats.created_at) + instant.duration_since(self.stats().created_at) } /// How long this connection has been idle. #[inline] pub fn idle_for(&self, instant: Instant) -> Duration { - instant.duration_since(self.stats.last_used) + instant.duration_since(self.stats().last_used) } /// How long has it been since the last connection healthcheck. #[inline] pub fn healthcheck_age(&self, instant: Instant) -> Duration { - if let Some(last_healthcheck) = self.stats.last_healthcheck { + if let Some(last_healthcheck) = self.stats().last_healthcheck { instant.duration_since(last_healthcheck) } else { Duration::MAX @@ -830,6 +916,7 @@ impl Server { #[inline] pub(super) fn cleaned(&mut self) { self.dirty = false; + self.stats.cleaned(); } /// Server is streaming data. @@ -872,21 +959,31 @@ impl Server { pub fn prepared_statements(&self) -> &PreparedStatements { &self.prepared_statements } + + #[inline] + pub fn memory_stats(&self) -> MemoryStats { + MemoryStats { + buffer: *self.stream_buffer.stats(), + prepared_statements: self.prepared_statements.memory_used(), + stream: self + .stream + .as_ref() + .map(|s| s.memory_usage()) + .unwrap_or_default(), + } + } } impl Drop for Server { fn drop(&mut self) { - self.stats.disconnect(); + self.stats().disconnect(); if let Some(mut stream) = self.stream.take() { - // If you see a lot of these, tell your clients - // to not send queries unless they are willing to stick - // around for results. - let out_of_sync = if self.done() { - " ".into() - } else { - format!(" {} ", self.stats.state) - }; - info!("closing{}server connection [{}]", out_of_sync, self.addr,); + info!( + "closing server connection [{}, state: {}, reason: {}]", + self.addr, + self.stats.state, + self.disconnect_reason.take().unwrap_or_default(), + ); spawn(async move { stream.write_all(&Terminate.to_bytes()?).await?; @@ -900,9 +997,9 @@ impl Drop for Server { // Used for testing. #[cfg(test)] pub mod test { - use crate::{frontend::PreparedStatements, net::*}; + use crate::{config::Memory, frontend::PreparedStatements, net::*}; - use super::*; + use super::{Error, *}; impl Default for Server { fn default() -> Self { @@ -911,11 +1008,16 @@ pub mod test { Self { stream: None, id, - startup_options: ServerOptions::default(), params: Parameters::default(), changed_params: Parameters::default(), client_params: Parameters::default(), - stats: Stats::connect(id, &addr, &Parameters::default()), + stats: Stats::connect( + id, + &addr, + &Parameters::default(), + &ServerOptions::default(), + &Memory::default(), + ), prepared_statements: super::PreparedStatements::new(), addr, dirty: false, @@ -926,7 +1028,8 @@ pub mod test { re_synced: false, replication_mode: false, pooler_mode: PoolerMode::Transaction, - stream_buffer: BytesMut::with_capacity(1024), + stream_buffer: MessageBuffer::new(4096), + disconnect_reason: None, } } } @@ -941,15 +1044,23 @@ pub mod test { } pub async fn test_server() -> Server { - Server::connect(&Address::new_test(), ServerOptions::default()) - .await - .unwrap() + Server::connect( + &Address::new_test(), + ServerOptions::default(), + ConnectReason::Other, + ) + .await + .unwrap() } pub async fn test_replication_server() -> Server { - Server::connect(&Address::new_test(), ServerOptions::new_replication()) - .await - .unwrap() + Server::connect( + &Address::new_test(), + ServerOptions::new_replication(), + ConnectReason::Replication, + ) + .await + .unwrap() } #[tokio::test] @@ -1035,7 +1146,7 @@ pub mod test { "", &[Parameter { len: 1, - data: "1".as_bytes().to_vec(), + data: "1".as_bytes().into(), }], &[Format::Text], ); @@ -1058,7 +1169,11 @@ pub mod test { assert_eq!(msg.code(), c); } - assert!(server.done()) + assert!(server.done()); + assert_eq!( + server.stats().total.idle_in_transaction_time, + Duration::ZERO + ); } } @@ -1080,7 +1195,7 @@ pub mod test { &name, &[Parameter { len: 1, - data: "1".as_bytes().to_vec(), + data: "1".as_bytes().into(), }], ); @@ -1166,7 +1281,7 @@ pub mod test { "__pgdog_1", &[Parameter { len: 1, - data: "1".as_bytes().to_vec(), + data: "1".as_bytes().into(), }], )), Execute::new().into(), @@ -1190,11 +1305,15 @@ pub mod test { async fn test_bad_parse() { let mut server = test_server().await; for _ in 0..25 { + let good_parse = Parse::named("test2", "SELECT $1"); let parse = Parse::named("test", "SELECT bad syntax;"); + let good_parse_2 = Parse::named("test3", "SELECT $2"); // Will be ignored due to error above. server .send( &vec![ + good_parse.into(), ProtocolMessage::from(parse), + good_parse_2.into(), Describe::new_statement("test").into(), Sync {}.into(), ] @@ -1202,12 +1321,22 @@ pub mod test { ) .await .unwrap(); - for c in ['E', 'Z'] { + for c in ['1', 'E', 'Z'] { let msg = server.read().await.unwrap(); assert_eq!(msg.code(), c); } assert!(server.done()); - assert!(server.prepared_statements.is_empty()); + assert_eq!(server.prepared_statements.len(), 1); + assert!(server.prepared_statements.contains("test2")); + server + .send(&vec![Query::new("SELECT 1").into()].into()) + .await + .unwrap(); + + for c in ['T', 'D', 'C', 'Z'] { + let msg = server.read().await.unwrap(); + assert_eq!(msg.code(), c); + } } } @@ -1332,6 +1461,9 @@ pub mod test { .unwrap(); for c in ['T', 'D', 'C', 'T', 'D', 'C', 'Z'] { let msg = server.read().await.unwrap(); + if c != 'Z' { + assert!(!server.done()); + } assert_eq!(c, msg.code()); } } @@ -1348,7 +1480,7 @@ pub mod test { "test_1", &[crate::net::bind::Parameter { len: 1, - data: "1".as_bytes().to_vec(), + data: "1".as_bytes().into(), }], ) .into(), @@ -1409,7 +1541,7 @@ pub mod test { "test", &[crate::net::bind::Parameter { len: 1, - data: "1".as_bytes().to_vec(), + data: "1".as_bytes().into(), }], ) .into(), @@ -1474,7 +1606,7 @@ pub mod test { "test", &[crate::net::bind::Parameter { len: 1, - data: "1".as_bytes().to_vec(), + data: "1".as_bytes().into(), }], ) .into(), @@ -1624,8 +1756,10 @@ pub mod test { let mut server = test_server().await; let mut params = Parameters::default(); params.insert("application_name", "test_sync_params"); - println!("server state: {}", server.stats().state); - let changed = server.link_client(¶ms).await.unwrap(); + let changed = server + .link_client(&BackendKeyData::new(), ¶ms, None) + .await + .unwrap(); assert_eq!(changed, 1); let app_name = server @@ -1634,7 +1768,10 @@ pub mod test { .unwrap(); assert_eq!(app_name[0], "test_sync_params"); - let changed = server.link_client(¶ms).await.unwrap(); + let changed = server + .link_client(&BackendKeyData::new(), ¶ms, None) + .await + .unwrap(); assert_eq!(changed, 0); } @@ -1683,7 +1820,7 @@ pub mod test { assert!(server.needs_drain()); assert!(!server.has_more_messages()); assert!(server.done()); - server.drain().await; + server.drain().await.unwrap(); assert!(server.in_sync()); assert!(server.done()); assert!(!server.needs_drain()); @@ -1705,9 +1842,9 @@ pub mod test { assert!(!server.needs_drain()); assert!(server.in_transaction()); - server.drain().await; // Nothing will be done. + server.drain().await.unwrap(); // Nothing will be done. assert!(server.in_transaction()); - server.rollback().await; + server.rollback().await.unwrap(); assert!(server.in_sync()); assert!(!server.in_transaction()); @@ -1721,20 +1858,28 @@ pub mod test { let mut server = test_server().await; - let changed = server.link_client(¶ms).await?; + let changed = server + .link_client(&BackendKeyData::new(), ¶ms, None) + .await?; assert_eq!(changed, 1); - let changed = server.link_client(¶ms).await?; + let changed = server + .link_client(&BackendKeyData::new(), ¶ms, None) + .await?; assert_eq!(changed, 0); for i in 0..25 { let value = format!("apples_{}", i); params.insert("application_name", value); - let changed = server.link_client(¶ms).await?; + let changed = server + .link_client(&BackendKeyData::new(), ¶ms, None) + .await?; assert_eq!(changed, 2); // RESET, SET. - let changed = server.link_client(¶ms).await?; + let changed = server + .link_client(&BackendKeyData::new(), ¶ms, None) + .await?; assert_eq!(changed, 0); } @@ -1828,6 +1973,36 @@ pub mod test { assert!(server.in_sync()); } + #[tokio::test] + async fn test_copy_client_fail() { + let mut server = test_server().await; + server.execute("BEGIN").await.unwrap(); + server + .execute("CREATE TABLE IF NOT EXISTS test_copy_t (id BIGINT)") + .await + .unwrap(); + server + .send( + &vec![ + Query::new("COPY test_copy_t(id) FROM STDIN CSV").into(), + CopyData::new(b"1\n").into(), + CopyFail::new("something went wrong").into(), + ] + .into(), + ) + .await + .unwrap(); + for c in ['G', 'E', 'Z'] { + if c != 'Z' { + assert!(!server.done()); + assert!(server.has_more_messages()); + } + assert_eq!(server.read().await.unwrap().code(), c); + } + + server.execute("ROLLBACK").await.unwrap(); + } + #[tokio::test] async fn test_query_stats() { let mut server = test_server().await; @@ -1901,7 +2076,7 @@ pub mod test { "", &[Parameter { len: 4, - data: "1234".as_bytes().to_vec(), + data: "1234".as_bytes().into(), }], ) .into(), @@ -2057,4 +2232,277 @@ pub mod test { assert_eq!(server.prepared_statements.len(), 0); assert!(server.done()); } + + #[tokio::test] + async fn test_drain_chaos() { + use crate::net::bind::Parameter; + use rand::Rng; + + let mut server = test_server().await; + let mut rng = rand::rng(); + + for iteration in 0..1000 { + let name = format!("chaos_test_{}", iteration); + let use_sync = rng.random_bool(0.5); + + if rng.random_bool(0.2) { + let bad_parse = Parse::named(&name, "SELECT invalid syntax"); + server + .send( + &vec![ + ProtocolMessage::from(bad_parse), + ProtocolMessage::from(Bind::new_params(&name, &[])), + ProtocolMessage::from(Execute::new()), + ProtocolMessage::from(Flush), + ] + .into(), + ) + .await + .unwrap(); + + let messages_to_read = rng.random_range(0..=1); + for _ in 0..messages_to_read { + let _ = server.read().await; + } + } else { + let parse = Parse::named(&name, "SELECT $1, $2"); + let bind = Bind::new_params( + &name, + &[ + Parameter { + len: 1, + data: "1".as_bytes().into(), + }, + Parameter { + len: 1, + data: "2".as_bytes().into(), + }, + ], + ); + let execute = Execute::new(); + + let messages = if use_sync { + vec![ + ProtocolMessage::from(parse), + ProtocolMessage::from(bind), + ProtocolMessage::from(execute), + ProtocolMessage::from(Sync), + ] + } else { + vec![ + ProtocolMessage::from(parse), + ProtocolMessage::from(bind), + ProtocolMessage::from(execute), + ProtocolMessage::from(Flush), + ] + }; + + server.send(&messages.into()).await.unwrap(); + + let expected_messages = if use_sync { + vec!['1', '2', 'D', 'C', 'Z'] + } else { + vec!['1', '2', 'D', 'C'] + }; + let messages_to_read = rng.random_range(0..expected_messages.len()); + + for i in 0..messages_to_read { + let msg = server.read().await.unwrap(); + assert_eq!(msg.code(), expected_messages[i]); + } + } + + server.drain().await.unwrap(); + + assert!(server.in_sync(), "Server should be in sync after drain"); + assert!(!server.error(), "Server should not be in error state"); + assert!(server.done(), "Server should be done after drain"); + } + } + + #[tokio::test] + async fn test_idle_in_transaction() { + let mut server = test_server().await; + + server.execute("SELECT 1").await.unwrap(); + assert_eq!( + server.stats().total.idle_in_transaction_time, + Duration::ZERO, + ); + assert_eq!( + server.stats().last_checkout.idle_in_transaction_time, + Duration::ZERO, + ); + + server.execute("BEGIN").await.unwrap(); + assert_eq!( + server.stats().total.idle_in_transaction_time, + Duration::ZERO, + ); + + server.execute("SELECT 1").await.unwrap(); + tokio::time::sleep(Duration::from_millis(50)).await; + + server.execute("SELECT 2").await.unwrap(); + + let idle_time = server.stats().total.idle_in_transaction_time; + assert!( + idle_time >= Duration::from_millis(50), + "Expected idle time >= 50ms, got {:?}", + idle_time + ); + assert!( + idle_time < Duration::from_millis(200), + "Expected idle time < 200ms, got {:?}", + idle_time + ); + + tokio::time::sleep(Duration::from_millis(100)).await; + + server.execute("COMMIT").await.unwrap(); + + let final_idle_time = server.stats().total.idle_in_transaction_time; + assert!( + final_idle_time >= Duration::from_millis(150), + "Expected final idle time >= 150ms, got {:?}", + final_idle_time + ); + assert!( + final_idle_time < Duration::from_millis(400), + "Expected final idle time < 400ms, got {:?}", + final_idle_time + ); + + server.execute("SELECT 3").await.unwrap(); + assert_eq!( + server.stats().total.idle_in_transaction_time, + final_idle_time, + ); + } + + #[tokio::test] + async fn test_prepare_forces_sync_prepared_flag() { + let mut server = test_server().await; + + assert!(!server.sync_prepared()); + + server + .send(&vec![Query::new("PREPARE test_stmt AS SELECT $1::bigint").into()].into()) + .await + .unwrap(); + + for c in ['C', 'Z'] { + let msg = server.read().await.unwrap(); + assert_eq!(msg.code(), c); + } + + assert!( + server.sync_prepared(), + "sync_prepared flag should be set after PREPARE command" + ); + + server.sync_prepared_statements().await.unwrap(); + + assert!( + !server.sync_prepared(), + "sync_prepared flag should be cleared after sync_prepared_statements()" + ); + + server.execute("SELECT 1").await.unwrap(); + + assert!( + !server.sync_prepared(), + "sync_prepared flag should remain false after regular queries" + ); + } + + #[tokio::test] + async fn test_protocol_out_of_sync_sets_error_state() { + let mut server = test_server().await; + + server + .send(&vec![Query::new("SELECT 1").into()].into()) + .await + .unwrap(); + + for c in ['T', 'D'] { + let msg = server.read().await.unwrap(); + assert_eq!(msg.code(), c); + } + + // simulate an unlikely, but existent out-of-sync state + server + .prepared_statements_mut() + .state_mut() + .queue_mut() + .clear(); + + let res = server.read().await; + assert!( + matches!(res, Err(Error::ProtocolOutOfSync)), + "protocol should be out of sync" + ); + assert!( + server.stats().state == State::Error, + "state should be Error after detecting desync" + ) + } + + #[tokio::test] + async fn test_reset_clears_client_params() { + let mut server = test_server().await; + let mut params = Parameters::default(); + params.insert("application_name", "test_reset"); + + // Sync params to server + let changed = server + .link_client(&BackendKeyData::new(), ¶ms, None) + .await + .unwrap(); + assert_eq!(changed, 1); + + // Same params should not need re-sync + let changed = server + .link_client(&BackendKeyData::new(), ¶ms, None) + .await + .unwrap(); + assert_eq!(changed, 0); + + // Execute RESET ALL which clears client_params + server.execute("RESET ALL").await.unwrap(); + + // Now link_client should need to re-sync because client_params was cleared + let changed = server + .link_client(&BackendKeyData::new(), ¶ms, None) + .await + .unwrap(); + assert!( + changed > 0, + "expected re-sync after RESET ALL cleared client_params" + ); + } + + #[tokio::test] + async fn test_error_decoding() { + let mut server = test_server().await; + let err = server + .execute(Query::new("SELECT * FROM test_error_decoding")) + .await + .expect_err("expected this query to fail"); + assert!( + matches!(err, Error::ExecutionError(_)), + "expected execution error" + ); + if let Error::ExecutionError(err) = err { + assert_eq!( + err.message, + "relation \"test_error_decoding\" does not exist" + ); + assert_eq!(err.severity, "ERROR"); + assert_eq!(err.code, "42P01"); + assert_eq!(err.context, None); + assert_eq!(err.routine, Some("parserOpenTable".into())); // Might break in the future. + assert_eq!(err.detail, None); + } + } } diff --git a/pgdog/src/backend/server_options.rs b/pgdog/src/backend/server_options.rs index a3e25de17..7242abf90 100644 --- a/pgdog/src/backend/server_options.rs +++ b/pgdog/src/backend/server_options.rs @@ -1,15 +1,20 @@ -use crate::net::Parameter; +use crate::net::{parameter::ParameterValue, Parameter}; #[derive(Debug, Clone, Default)] pub struct ServerOptions { pub params: Vec, + pub pool_id: u64, } impl ServerOptions { pub fn replication_mode(&self) -> bool { - self.params - .iter() - .any(|p| p.name == "replication" && p.value == "database") + self.params.iter().any(|p| { + p.name == "replication" + && match p.value { + ParameterValue::String(ref value) => value == "database", + _ => false, + } + }) } pub fn new_replication() -> Self { @@ -18,6 +23,7 @@ impl ServerOptions { name: "replication".into(), value: "database".into(), }], + pool_id: 0, } } } diff --git a/pgdog/src/backend/stats.rs b/pgdog/src/backend/stats.rs index dfce7075c..3541e2a6d 100644 --- a/pgdog/src/backend/stats.rs +++ b/pgdog/src/backend/stats.rs @@ -11,6 +11,8 @@ use parking_lot::Mutex; use tokio::time::Instant; use crate::{ + backend::{pool::stats::MemoryStats, Pool, ServerOptions}, + config::Memory, net::{messages::BackendKeyData, Parameters}, state::State, }; @@ -25,6 +27,17 @@ pub fn stats() -> HashMap { STATS.lock().clone() } +/// Get idle-in-transaction server connections for connection pool. +pub fn idle_in_transaction(pool: &Pool) -> usize { + STATS + .lock() + .values() + .filter(|stat| { + stat.stats.pool_id == pool.id() && stat.stats.state == State::IdleInTransaction + }) + .count() +} + /// Update stats to latest version. fn update(id: BackendKeyData, stats: Stats) { let mut guard = STATS.lock(); @@ -60,11 +73,13 @@ pub struct Counts { pub prepared_statements: usize, pub query_time: Duration, pub transaction_time: Duration, + pub idle_in_transaction_time: Duration, pub parse: usize, pub bind: usize, pub healthchecks: usize, pub close: usize, - pub memory_used: usize, + pub cleaned: usize, + pub prepared_sync: usize, } impl Add for Counts { @@ -79,16 +94,18 @@ impl Add for Counts { queries: self.queries.saturating_add(rhs.queries), rollbacks: self.rollbacks.saturating_add(rhs.rollbacks), errors: self.errors.saturating_add(rhs.errors), - prepared_statements: self - .prepared_statements - .saturating_add(rhs.prepared_statements), + prepared_statements: rhs.prepared_statements, // It's a gauge. query_time: self.query_time.saturating_add(rhs.query_time), transaction_time: self.query_time.saturating_add(rhs.transaction_time), + idle_in_transaction_time: self + .idle_in_transaction_time + .saturating_add(rhs.idle_in_transaction_time), parse: self.parse.saturating_add(rhs.parse), bind: self.bind.saturating_add(rhs.bind), healthchecks: self.healthchecks.saturating_add(rhs.healthchecks), close: self.close.saturating_add(rhs.close), - memory_used: self.memory_used, // It's a gauge. + cleaned: self.cleaned.saturating_add(rhs.cleaned), + prepared_sync: self.prepared_sync.saturating_add(rhs.prepared_sync), } } } @@ -104,13 +121,23 @@ pub struct Stats { pub created_at_time: SystemTime, pub total: Counts, pub last_checkout: Counts, + pub pool_id: u64, + pub client_id: Option, + pub memory: MemoryStats, query_timer: Option, transaction_timer: Option, + idle_in_transaction_timer: Option, } impl Stats { /// Register new server with statistics. - pub fn connect(id: BackendKeyData, addr: &Address, params: &Parameters) -> Self { + pub fn connect( + id: BackendKeyData, + addr: &Address, + params: &Parameters, + options: &ServerOptions, + config: &Memory, + ) -> Self { let now = Instant::now(); let stats = Stats { id, @@ -123,6 +150,10 @@ impl Stats { last_checkout: Counts::default(), query_timer: None, transaction_timer: None, + pool_id: options.pool_id, + client_id: None, + memory: MemoryStats::new(config), + idle_in_transaction_timer: None, }; STATS.lock().insert( @@ -151,7 +182,8 @@ impl Stats { self.update(); } - pub fn link_client(&mut self, client_name: &str, server_server: &str) { + pub fn link_client(&mut self, client_name: &str, server_server: &str, id: &BackendKeyData) { + self.client_id = Some(*id); if client_name != server_server { let mut guard = STATS.lock(); if let Some(entry) = guard.get_mut(&self.id) { @@ -172,6 +204,8 @@ impl Stats { /// for stats. pub fn set_prepared_statements(&mut self, size: usize) { self.total.prepared_statements = size; + self.total.prepared_sync += 1; + self.last_checkout.prepared_sync += 1; self.update(); } @@ -214,9 +248,14 @@ impl Stats { } /// A query has been completed. - pub fn query(&mut self, now: Instant) { + pub fn query(&mut self, now: Instant, idle_in_transaction: bool) { self.total.queries += 1; self.last_checkout.queries += 1; + + if idle_in_transaction { + self.idle_in_transaction_timer = Some(now); + } + if let Some(query_timer) = self.query_timer.take() { let duration = now.duration_since(query_timer); self.total.query_time += duration; @@ -240,6 +279,7 @@ impl Stats { } fn activate(&mut self) { + // Client started a query/transaction. if self.state == State::Active { let now = Instant::now(); if self.transaction_timer.is_none() { @@ -248,6 +288,11 @@ impl Stats { if self.query_timer.is_none() { self.query_timer = Some(now); } + if let Some(idle_in_transaction_timer) = self.idle_in_transaction_timer.take() { + let elapsed = now.duration_since(idle_in_transaction_timer); + self.last_checkout.idle_in_transaction_time += elapsed; + self.total.idle_in_transaction_time += elapsed; + } } } @@ -272,9 +317,14 @@ impl Stats { } #[inline] - pub fn memory_used(&mut self, memory: usize) { - self.total.memory_used = memory; - self.last_checkout.memory_used = memory; + pub fn memory_used(&mut self, stats: MemoryStats) { + self.memory = stats; + } + + #[inline] + pub fn cleaned(&mut self) { + self.last_checkout.cleaned += 1; + self.total.cleaned += 1; } /// Track rollbacks. diff --git a/pgdog/src/cli.rs b/pgdog/src/cli.rs index 05bbbbbe5..11212c43c 100644 --- a/pgdog/src/cli.rs +++ b/pgdog/src/cli.rs @@ -5,12 +5,13 @@ use clap::{Parser, Subcommand}; use std::fs::read_to_string; use thiserror::Error; use tokio::{select, signal::ctrl_c}; -use tracing::error; +use tracing::{error, info}; use crate::backend::schema::sync::config::ShardConfig; use crate::backend::schema::sync::pg_dump::{PgDump, SyncState}; use crate::backend::{databases::databases, replication::logical::Publisher}; -use crate::config::{Config, Users}; +use crate::config::{config, Config, Users}; +use crate::frontend::router::cli::RouterCli; /// PgDog is a PostgreSQL pooler, proxy, load balancer and query router. #[derive(Parser, Debug)] @@ -55,6 +56,21 @@ pub enum Commands { path: Option, }, + /// Execute the router on the queries. + Route { + /// User in users.toml. + #[arg(short, long)] + user: String, + + /// Database in pgdog.toml. + #[arg(short, long)] + database: String, + + /// Path to the file containing the queries. + #[arg(short, long)] + file: PathBuf, + }, + /// Check configuration files for errors. Configcheck, @@ -64,9 +80,7 @@ pub enum Commands { /// Source database name. #[arg(long)] from_database: String, - /// Source user name. - #[arg(long)] - from_user: String, + /// Publication name. #[arg(long)] publication: String, @@ -74,13 +88,18 @@ pub enum Commands { /// Destination database. #[arg(long)] to_database: String, - /// Destination user name. - #[arg(long)] - to_user: String, /// Replicate or copy data over. #[arg(long, default_value = "false")] - replicate: bool, + replicate_only: bool, + + /// Replicate or copy data over. + #[arg(long, default_value = "false")] + sync_only: bool, + + /// Name of the replication slot to create/use. + #[arg(long)] + replication_slot: Option, }, /// Copy schema from source to destination cluster. @@ -107,6 +126,10 @@ pub enum Commands { /// Data sync has been complete. #[arg(long)] data_sync_complete: bool, + + /// Execute cutover statements. + #[arg(long)] + cutover: bool, }, /// Perform cluster configuration steps @@ -204,31 +227,43 @@ pub fn config_check( } pub async fn data_sync(commands: Commands) -> Result<(), Box> { - let (source, destination, publication, replicate) = if let Commands::DataSync { - from_database, - from_user, - to_database, - to_user, - publication, - replicate, - } = commands - { - let source = databases().cluster((from_user.as_str(), from_database.as_str()))?; - let dest = databases().cluster((to_user.as_str(), to_database.as_str()))?; + let (source, destination, publication, replicate_only, sync_only, replication_slot) = + if let Commands::DataSync { + from_database, + to_database, + publication, + replicate_only, + sync_only, + replication_slot, + } = commands + { + let source = databases().schema_owner(&from_database)?; + let dest = databases().schema_owner(&to_database)?; - (source, dest, publication, replicate) - } else { - return Ok(()); - }; + ( + source, + dest, + publication, + replicate_only, + sync_only, + replication_slot, + ) + } else { + return Ok(()); + }; - let mut publication = Publisher::new(&source, &publication); - if replicate { - if let Err(err) = publication.replicate(&destination).await { + let mut publication = Publisher::new( + &source, + &publication, + config().config.general.query_parser_engine, + ); + if replicate_only { + if let Err(err) = publication.replicate(&destination, replication_slot).await { error!("{}", err); } } else { select! { - result = publication.data_sync(&destination) => { + result = publication.data_sync(&destination, !sync_only, replication_slot) => { if let Err(err) = result { error!("{}", err); } @@ -244,7 +279,7 @@ pub async fn data_sync(commands: Commands) -> Result<(), Box Result<(), Box> { - let (source, destination, publication, dry_run, ignore_errors, data_sync_complete) = + let (source, destination, publication, dry_run, ignore_errors, data_sync_complete, cutover) = if let Commands::SchemaSync { from_database, to_database, @@ -252,6 +287,7 @@ pub async fn schema_sync(commands: Commands) -> Result<(), Box Result<(), Box Result<(), Box Result<(), Box Result<(), Box> { Ok(()) } + +pub async fn route(commands: Commands) -> Result<(), Box> { + if let Commands::Route { + user, + database, + file, + } = commands + { + let cli = RouterCli::new(&database, &user, file).await?; + let cmds = cli.run()?; + + for cmd in cmds { + info!("{:?}", cmd); + } + } + + Ok(()) +} diff --git a/pgdog/src/config/auth.rs b/pgdog/src/config/auth.rs index b2780751e..dc5c0db19 100644 --- a/pgdog/src/config/auth.rs +++ b/pgdog/src/config/auth.rs @@ -1,47 +1 @@ -use serde::{Deserialize, Serialize}; -use std::str::FromStr; - -#[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq)] -#[serde(rename_all = "snake_case")] -pub enum PassthoughAuth { - #[default] - Disabled, - Enabled, - EnabledPlain, -} - -#[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq)] -#[serde(rename_all = "snake_case")] -pub enum AuthType { - Md5, - #[default] - Scram, - Trust, -} - -impl AuthType { - pub fn md5(&self) -> bool { - matches!(self, Self::Md5) - } - - pub fn scram(&self) -> bool { - matches!(self, Self::Scram) - } - - pub fn trust(&self) -> bool { - matches!(self, Self::Trust) - } -} - -impl FromStr for AuthType { - type Err = String; - - fn from_str(s: &str) -> Result { - match s.to_lowercase().as_str() { - "md5" => Ok(Self::Md5), - "scram" => Ok(Self::Scram), - "trust" => Ok(Self::Trust), - _ => Err(format!("Invalid auth type: {}", s)), - } - } -} +pub use pgdog_config::auth::*; diff --git a/pgdog/src/config/convert.rs b/pgdog/src/config/convert.rs index 75b7f21b5..14c4f8c3c 100644 --- a/pgdog/src/config/convert.rs +++ b/pgdog/src/config/convert.rs @@ -3,31 +3,48 @@ use crate::net::{Parameters, Password}; use super::User; -impl User { - pub(crate) fn from_params(params: &Parameters, password: &Password) -> Result { - let user = params - .get("user") - .ok_or(Error::IncompleteStartup)? - .as_str() - .ok_or(Error::IncompleteStartup)?; - let database = params.get_default("database", user); - let password = password.password().ok_or(Error::IncompleteStartup)?; +pub fn user_from_params(params: &Parameters, password: &Password) -> Result { + let user = params + .get("user") + .ok_or(Error::IncompleteStartup)? + .as_str() + .ok_or(Error::IncompleteStartup)?; + let database = params.get_default("database", user); + let password = password.password().ok_or(Error::IncompleteStartup)?; - Ok(Self { - name: user.to_owned(), - database: database.to_owned(), - password: Some(password.to_owned()), - ..Default::default() - }) - } - - /// New user from user, password and database. - pub(crate) fn new(user: &str, password: &str, database: &str) -> Self { - Self { - name: user.to_owned(), - database: database.to_owned(), - password: Some(password.to_owned()), - ..Default::default() - } - } + Ok(User { + name: user.to_owned(), + database: database.to_owned(), + password: Some(password.to_owned()), + ..Default::default() + }) } + +// impl User { +// pub(crate) fn from_params(params: &Parameters, password: &Password) -> Result { +// let user = params +// .get("user") +// .ok_or(Error::IncompleteStartup)? +// .as_str() +// .ok_or(Error::IncompleteStartup)?; +// let database = params.get_default("database", user); +// let password = password.password().ok_or(Error::IncompleteStartup)?; + +// Ok(Self { +// name: user.to_owned(), +// database: database.to_owned(), +// password: Some(password.to_owned()), +// ..Default::default() +// }) +// } + +// /// New user from user, password and database. +// pub(crate) fn new(user: &str, password: &str, database: &str) -> Self { +// Self { +// name: user.to_owned(), +// database: database.to_owned(), +// password: Some(password.to_owned()), +// ..Default::default() +// } +// } +// } diff --git a/pgdog/src/config/core.rs b/pgdog/src/config/core.rs index 362776207..95a0f3832 100644 --- a/pgdog/src/config/core.rs +++ b/pgdog/src/config/core.rs @@ -1,466 +1 @@ -use serde::{Deserialize, Serialize}; -use std::collections::{HashMap, HashSet}; -use std::fs::read_to_string; -use std::path::PathBuf; -use tracing::{info, warn}; - -use super::database::Database; -use super::error::Error; -use super::general::General; -use super::networking::{MultiTenant, Tcp}; -use super::pooling::{PoolerMode, Stats}; -use super::replication::{MirrorConfig, Mirroring, ReplicaLag, Replication}; -use super::sharding::{ManualQuery, OmnishardedTables, ShardedMapping, ShardedTable}; -use super::users::{Admin, Plugin, Users}; - -#[derive(Debug, Clone)] -pub struct ConfigAndUsers { - /// pgdog.toml - pub config: Config, - /// users.toml - pub users: Users, - /// Path to pgdog.toml. - pub config_path: PathBuf, - /// Path to users.toml. - pub users_path: PathBuf, -} - -impl ConfigAndUsers { - /// Load configuration from disk or use defaults. - pub fn load(config_path: &PathBuf, users_path: &PathBuf) -> Result { - let config: Config = if let Ok(config) = read_to_string(config_path) { - let config = match toml::from_str(&config) { - Ok(config) => config, - Err(err) => return Err(Error::config(&config, err)), - }; - info!("loaded \"{}\"", config_path.display()); - config - } else { - warn!( - "\"{}\" doesn't exist, loading defaults instead", - config_path.display() - ); - Config::default() - }; - - if config.admin.random() { - #[cfg(debug_assertions)] - info!("[debug only] admin password: {}", config.admin.password); - #[cfg(not(debug_assertions))] - warn!("admin password has been randomly generated"); - } - - if config.multi_tenant.is_some() { - info!("multi-tenant protection enabled"); - } - - let users: Users = if let Ok(users) = read_to_string(users_path) { - let mut users: Users = toml::from_str(&users)?; - users.check(&config); - info!("loaded \"{}\"", users_path.display()); - users - } else { - warn!( - "\"{}\" doesn't exist, loading defaults instead", - users_path.display() - ); - Users::default() - }; - - Ok(ConfigAndUsers { - config, - users, - config_path: config_path.to_owned(), - users_path: users_path.to_owned(), - }) - } - - /// Prepared statements are enabled. - pub fn prepared_statements(&self) -> bool { - // Disable prepared statements automatically in session mode - if self.config.general.pooler_mode == PoolerMode::Session { - false - } else { - self.config.general.prepared_statements.enabled() - } - } - - /// Prepared statements are in "full" mode (used for query parser decision). - pub fn prepared_statements_full(&self) -> bool { - self.config.general.prepared_statements.full() - } - - pub fn pub_sub_enabled(&self) -> bool { - self.config.general.pub_sub_channel_size > 0 - } -} - -impl Default for ConfigAndUsers { - fn default() -> Self { - Self { - config: Config::default(), - users: Users::default(), - config_path: PathBuf::from("pgdog.toml"), - users_path: PathBuf::from("users.toml"), - } - } -} - -/// Configuration. -#[derive(Serialize, Deserialize, Debug, Clone, Default)] -#[serde(deny_unknown_fields)] -pub struct Config { - /// General configuration. - #[serde(default)] - pub general: General, - - /// Statistics. - #[serde(default)] - pub stats: Stats, - - /// TCP settings - #[serde(default)] - pub tcp: Tcp, - - /// Multi-tenant - pub multi_tenant: Option, - - /// Servers. - #[serde(default)] - pub databases: Vec, - - #[serde(default)] - pub plugins: Vec, - - #[serde(default)] - pub admin: Admin, - - /// List of sharded tables. - #[serde(default)] - pub sharded_tables: Vec, - - /// Queries routed manually to a single shard. - #[serde(default)] - pub manual_queries: Vec, - - /// List of omnisharded tables. - #[serde(default)] - pub omnisharded_tables: Vec, - - /// Explicit sharding key mappings. - #[serde(default)] - pub sharded_mappings: Vec, - - /// Replica lag configuration. - #[serde(default, deserialize_with = "ReplicaLag::deserialize_optional")] - pub replica_lag: Option, - - /// Replication config. - #[serde(default)] - pub replication: Replication, - - /// Mirroring configurations. - #[serde(default)] - pub mirroring: Vec, -} - -impl Config { - /// Organize all databases by name for quicker retrieval. - pub fn databases(&self) -> HashMap>> { - let mut databases = HashMap::new(); - for database in &self.databases { - let entry = databases - .entry(database.name.clone()) - .or_insert_with(Vec::new); - while entry.len() <= database.shard { - entry.push(vec![]); - } - entry - .get_mut(database.shard) - .unwrap() - .push(database.clone()); - } - databases - } - - /// Organize sharded tables by database name. - pub fn sharded_tables(&self) -> HashMap> { - let mut tables = HashMap::new(); - - for table in &self.sharded_tables { - let entry = tables - .entry(table.database.clone()) - .or_insert_with(Vec::new); - entry.push(table.clone()); - } - - tables - } - - pub fn omnisharded_tables(&self) -> HashMap> { - let mut tables = HashMap::new(); - - for table in &self.omnisharded_tables { - let entry = tables - .entry(table.database.clone()) - .or_insert_with(Vec::new); - for t in &table.tables { - entry.push(t.clone()); - } - } - - tables - } - - /// Manual queries. - pub fn manual_queries(&self) -> HashMap { - let mut queries = HashMap::new(); - - for query in &self.manual_queries { - queries.insert(query.fingerprint.clone(), query.clone()); - } - - queries - } - - /// Sharded mappings. - pub fn sharded_mappings( - &self, - ) -> HashMap<(String, String, Option), Vec> { - let mut mappings = HashMap::new(); - - for mapping in &self.sharded_mappings { - let mapping = mapping.clone(); - let entry = mappings - .entry(( - mapping.database.clone(), - mapping.column.clone(), - mapping.table.clone(), - )) - .or_insert_with(Vec::new); - entry.push(mapping); - } - - mappings - } - - pub fn check(&self) { - // Check databases. - let mut duplicate_primaries = HashSet::new(); - for database in self.databases.clone() { - let id = ( - database.name.clone(), - database.role, - database.shard, - database.port, - ); - let new = duplicate_primaries.insert(id); - if !new { - warn!( - "database \"{}\" (shard={}) has a duplicate {}", - database.name, database.shard, database.role, - ); - } - } - } - - /// Multi-tenanncy is enabled. - pub fn multi_tenant(&self) -> &Option { - &self.multi_tenant - } - - /// Get mirroring configuration for a specific source/destination pair. - pub fn get_mirroring_config( - &self, - source_db: &str, - destination_db: &str, - ) -> Option { - self.mirroring - .iter() - .find(|m| m.source_db == source_db && m.destination_db == destination_db) - .map(|m| MirrorConfig { - queue_length: m.queue_length.unwrap_or(self.general.mirror_queue), - exposure: m.exposure.unwrap_or(self.general.mirror_exposure), - }) - } - - /// Get all mirroring configurations mapped by source database. - pub fn mirroring_by_source(&self) -> HashMap> { - let mut result = HashMap::new(); - - for mirror in &self.mirroring { - let config = MirrorConfig { - queue_length: mirror.queue_length.unwrap_or(self.general.mirror_queue), - exposure: mirror.exposure.unwrap_or(self.general.mirror_exposure), - }; - - result - .entry(mirror.source_db.clone()) - .or_insert_with(Vec::new) - .push((mirror.destination_db.clone(), config)); - } - - result - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::config::{PoolerMode, PreparedStatements}; - use std::time::Duration; - - #[test] - fn test_basic() { - let source = r#" -[general] -host = "0.0.0.0" -port = 6432 -default_pool_size = 15 -pooler_mode = "transaction" - -[[databases]] -name = "production" -role = "primary" -host = "127.0.0.1" -port = 5432 -database_name = "postgres" - -[tcp] -keepalive = true -interval = 5000 -time = 1000 -user_timeout = 1000 -retries = 5 - -[[plugins]] -name = "pgdog_routing" - -[multi_tenant] -column = "tenant_id" -"#; - - let config: Config = toml::from_str(source).unwrap(); - assert_eq!(config.databases[0].name, "production"); - assert_eq!(config.plugins[0].name, "pgdog_routing"); - assert!(config.tcp.keepalive()); - assert_eq!(config.tcp.interval().unwrap(), Duration::from_millis(5000)); - assert_eq!( - config.tcp.user_timeout().unwrap(), - Duration::from_millis(1000) - ); - assert_eq!(config.tcp.time().unwrap(), Duration::from_millis(1000)); - assert_eq!(config.tcp.retries().unwrap(), 5); - assert_eq!(config.multi_tenant.unwrap().column, "tenant_id"); - } - - #[test] - fn test_prepared_statements_disabled_in_session_mode() { - let mut config = ConfigAndUsers::default(); - - // Test transaction mode (default) - prepared statements should be enabled - config.config.general.pooler_mode = PoolerMode::Transaction; - config.config.general.prepared_statements = PreparedStatements::Extended; - assert!( - config.prepared_statements(), - "Prepared statements should be enabled in transaction mode" - ); - - // Test session mode - prepared statements should be disabled - config.config.general.pooler_mode = PoolerMode::Session; - config.config.general.prepared_statements = PreparedStatements::Extended; - assert!( - !config.prepared_statements(), - "Prepared statements should be disabled in session mode" - ); - - // Test session mode with full prepared statements - should still be disabled - config.config.general.pooler_mode = PoolerMode::Session; - config.config.general.prepared_statements = PreparedStatements::Full; - assert!( - !config.prepared_statements(), - "Prepared statements should be disabled in session mode even when set to Full" - ); - - // Test transaction mode with disabled prepared statements - should remain disabled - config.config.general.pooler_mode = PoolerMode::Transaction; - config.config.general.prepared_statements = PreparedStatements::Disabled; - assert!(!config.prepared_statements(), "Prepared statements should remain disabled when explicitly set to Disabled in transaction mode"); - } - - #[test] - fn test_mirroring_config() { - let source = r#" -[general] -host = "0.0.0.0" -port = 6432 -mirror_queue = 128 -mirror_exposure = 1.0 - -[[databases]] -name = "source_db" -host = "127.0.0.1" -port = 5432 - -[[databases]] -name = "destination_db1" -host = "127.0.0.1" -port = 5433 - -[[databases]] -name = "destination_db2" -host = "127.0.0.1" -port = 5434 - -[[mirroring]] -source_db = "source_db" -destination_db = "destination_db1" -queue_length = 256 -exposure = 0.5 - -[[mirroring]] -source_db = "source_db" -destination_db = "destination_db2" -exposure = 0.75 -"#; - - let config: Config = toml::from_str(source).unwrap(); - - // Verify we have 2 mirroring configurations - assert_eq!(config.mirroring.len(), 2); - - // Check first mirroring config - assert_eq!(config.mirroring[0].source_db, "source_db"); - assert_eq!(config.mirroring[0].destination_db, "destination_db1"); - assert_eq!(config.mirroring[0].queue_length, Some(256)); - assert_eq!(config.mirroring[0].exposure, Some(0.5)); - - // Check second mirroring config - assert_eq!(config.mirroring[1].source_db, "source_db"); - assert_eq!(config.mirroring[1].destination_db, "destination_db2"); - assert_eq!(config.mirroring[1].queue_length, None); // Should use global default - assert_eq!(config.mirroring[1].exposure, Some(0.75)); - - // Verify global defaults are still set - assert_eq!(config.general.mirror_queue, 128); - assert_eq!(config.general.mirror_exposure, 1.0); - - // Test get_mirroring_config method - let mirror_config = config - .get_mirroring_config("source_db", "destination_db1") - .unwrap(); - assert_eq!(mirror_config.queue_length, 256); - assert_eq!(mirror_config.exposure, 0.5); - - let mirror_config2 = config - .get_mirroring_config("source_db", "destination_db2") - .unwrap(); - assert_eq!(mirror_config2.queue_length, 128); // Uses global default - assert_eq!(mirror_config2.exposure, 0.75); - - // Non-existent mirror config should return None - assert!(config - .get_mirroring_config("source_db", "non_existent") - .is_none()); - } -} +pub use pgdog_config::core::*; diff --git a/pgdog/src/config/database.rs b/pgdog/src/config/database.rs index ea98ab8dd..d82c9ae1f 100644 --- a/pgdog/src/config/database.rs +++ b/pgdog/src/config/database.rs @@ -1,138 +1 @@ -use serde::{Deserialize, Serialize}; -use std::str::FromStr; - -use super::pooling::PoolerMode; - -#[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq, Copy)] -#[serde(rename_all = "snake_case")] -pub enum ReadWriteStrategy { - #[default] - Conservative, - Aggressive, -} - -impl FromStr for ReadWriteStrategy { - type Err = String; - - fn from_str(s: &str) -> Result { - match s.to_lowercase().as_str() { - "conservative" => Ok(Self::Conservative), - "aggressive" => Ok(Self::Aggressive), - _ => Err(format!("Invalid read-write strategy: {}", s)), - } - } -} - -#[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq, Copy)] -#[serde(rename_all = "snake_case")] -pub enum LoadBalancingStrategy { - #[default] - Random, - RoundRobin, - LeastActiveConnections, -} - -impl FromStr for LoadBalancingStrategy { - type Err = String; - - fn from_str(s: &str) -> Result { - match s.to_lowercase().replace(['_', '-'], "").as_str() { - "random" => Ok(Self::Random), - "roundrobin" => Ok(Self::RoundRobin), - "leastactiveconnections" => Ok(Self::LeastActiveConnections), - _ => Err(format!("Invalid load balancing strategy: {}", s)), - } - } -} - -#[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq, Copy)] -#[serde(rename_all = "snake_case")] -pub enum ReadWriteSplit { - #[default] - IncludePrimary, - ExcludePrimary, -} - -impl FromStr for ReadWriteSplit { - type Err = String; - - fn from_str(s: &str) -> Result { - match s.to_lowercase().replace(['_', '-'], "").as_str() { - "includeprimary" => Ok(Self::IncludePrimary), - "excludeprimary" => Ok(Self::ExcludePrimary), - _ => Err(format!("Invalid read-write split: {}", s)), - } - } -} - -/// Database server proxied by pgDog. -#[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq, Ord, PartialOrd, Eq)] -#[serde(deny_unknown_fields)] -pub struct Database { - /// Database name visible to the clients. - pub name: String, - /// Database role, e.g. primary. - #[serde(default)] - pub role: Role, - /// Database host or IP address, e.g. 127.0.0.1. - pub host: String, - /// Database port, e.g. 5432. - #[serde(default = "Database::port")] - pub port: u16, - /// Shard. - #[serde(default)] - pub shard: usize, - /// PostgreSQL database name, e.g. "postgres". - pub database_name: Option, - /// Use this user to connect to the database, overriding the userlist. - pub user: Option, - /// Use this password to login, overriding the userlist. - pub password: Option, - // Maximum number of connections to this database from this pooler. - // #[serde(default = "Database::max_connections")] - // pub max_connections: usize, - /// Pool size for this database pools, overriding `default_pool_size`. - pub pool_size: Option, - /// Minimum pool size for this database pools, overriding `min_pool_size`. - pub min_pool_size: Option, - /// Pooler mode. - pub pooler_mode: Option, - /// Statement timeout. - pub statement_timeout: Option, - /// Idle timeout. - pub idle_timeout: Option, - /// Read-only mode. - pub read_only: Option, - /// Server lifetime. - pub server_lifetime: Option, -} - -impl Database { - #[allow(dead_code)] - fn max_connections() -> usize { - usize::MAX - } - - fn port() -> u16 { - 5432 - } -} - -#[derive( - Serialize, Deserialize, Debug, Clone, Default, PartialEq, Ord, PartialOrd, Eq, Hash, Copy, -)] -#[serde(rename_all = "snake_case")] -pub enum Role { - #[default] - Primary, - Replica, -} - -impl std::fmt::Display for Role { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - Self::Primary => write!(f, "primary"), - Self::Replica => write!(f, "replica"), - } - } -} +pub use pgdog_config::database::*; diff --git a/pgdog/src/config/error.rs b/pgdog/src/config/error.rs index aced2f5ab..8d7cf6c8a 100644 --- a/pgdog/src/config/error.rs +++ b/pgdog/src/config/error.rs @@ -1,64 +1,3 @@ //! Configuration errors. -use thiserror::Error; - -/// Configuration error. -#[derive(Debug, Error)] -pub enum Error { - #[error("{0}")] - Io(#[from] std::io::Error), - - #[error("{0}")] - Deser(#[from] toml::de::Error), - - #[error("{0}, line {1}")] - MissingField(String, usize), - - #[error("{0}")] - Url(#[from] url::ParseError), - - #[error("{0}")] - Json(#[from] serde_json::Error), - - #[error("incomplete startup")] - IncompleteStartup, - - #[error("no database urls in environment")] - NoDbsInEnv, -} - -impl Error { - pub fn config(source: &str, err: toml::de::Error) -> Self { - let span = err.span(); - let message = err.message(); - - let span = if let Some(span) = span { - span - } else { - return Self::MissingField(message.into(), 0); - }; - - let mut lines = vec![]; - let mut line = 1; - for (i, c) in source.chars().enumerate() { - if c == '\n' { - lines.push((line, i)); - line += 1; - } - } - - let mut lines = lines.into_iter().peekable(); - - while let Some(line) = lines.next() { - if span.start < line.1 { - if let Some(next) = lines.peek() { - if next.1 > span.start { - return Self::MissingField(message.into(), line.0); - } - } - } - } - - Self::MissingField(message.into(), 0) - } -} +pub use pgdog_config::Error; diff --git a/pgdog/src/config/general.rs b/pgdog/src/config/general.rs index 08ed6cbba..942bbabe6 100644 --- a/pgdog/src/config/general.rs +++ b/pgdog/src/config/general.rs @@ -1,845 +1 @@ -use serde::{Deserialize, Serialize}; -use std::env; -use std::net::Ipv4Addr; -use std::path::PathBuf; -use std::time::Duration; - -use super::auth::{AuthType, PassthoughAuth}; -use super::database::{LoadBalancingStrategy, ReadWriteSplit, ReadWriteStrategy}; -use super::networking::TlsVerifyMode; -use super::pooling::{PoolerMode, PreparedStatements}; - -#[derive(Serialize, Deserialize, Debug, Clone)] -pub struct General { - /// Run on this address. - #[serde(default = "General::host")] - pub host: String, - /// Run on this port. - #[serde(default = "General::port")] - pub port: u16, - /// Spawn this many Tokio threads. - #[serde(default = "General::workers")] - pub workers: usize, - /// Default pool size, e.g. 10. - #[serde(default = "General::default_pool_size")] - pub default_pool_size: usize, - /// Minimum number of connections to maintain in the pool. - #[serde(default = "General::min_pool_size")] - pub min_pool_size: usize, - /// Pooler mode, e.g. transaction. - #[serde(default)] - pub pooler_mode: PoolerMode, - /// How often to check a connection. - #[serde(default = "General::healthcheck_interval")] - pub healthcheck_interval: u64, - /// How often to issue a healthcheck via an idle connection. - #[serde(default = "General::idle_healthcheck_interval")] - pub idle_healthcheck_interval: u64, - /// Delay idle healthchecks by this time at startup. - #[serde(default = "General::idle_healthcheck_delay")] - pub idle_healthcheck_delay: u64, - /// Healthcheck timeout. - #[serde(default = "General::healthcheck_timeout")] - pub healthcheck_timeout: u64, - /// HTTP health check port. - pub healthcheck_port: Option, - /// Maximum duration of a ban. - #[serde(default = "General::ban_timeout")] - pub ban_timeout: u64, - /// Rollback timeout. - #[serde(default = "General::rollback_timeout")] - pub rollback_timeout: u64, - /// Load balancing strategy. - #[serde(default = "General::load_balancing_strategy")] - pub load_balancing_strategy: LoadBalancingStrategy, - /// How aggressive should the query parser be in determining reads. - #[serde(default)] - pub read_write_strategy: ReadWriteStrategy, - /// Read write split. - #[serde(default)] - pub read_write_split: ReadWriteSplit, - /// TLS certificate. - pub tls_certificate: Option, - /// TLS private key. - pub tls_private_key: Option, - /// TLS verification mode (for connecting to servers) - #[serde(default = "General::default_tls_verify")] - pub tls_verify: TlsVerifyMode, - /// TLS CA certificate (for connecting to servers). - pub tls_server_ca_certificate: Option, - /// Shutdown timeout. - #[serde(default = "General::default_shutdown_timeout")] - pub shutdown_timeout: u64, - /// Broadcast IP. - pub broadcast_address: Option, - /// Broadcast port. - #[serde(default = "General::broadcast_port")] - pub broadcast_port: u16, - /// Load queries to file (warning: slow, don't use in production). - #[serde(default)] - pub query_log: Option, - /// Enable OpenMetrics server on this port. - pub openmetrics_port: Option, - /// OpenMetrics prefix. - pub openmetrics_namespace: Option, - - /// Prepared statatements support. - #[serde(default)] - pub prepared_statements: PreparedStatements, - /// Limit on the number of prepared statements in the server cache. - #[serde(default = "General::prepared_statements_limit")] - pub prepared_statements_limit: usize, - #[serde(default = "General::query_cache_limit")] - pub query_cache_limit: usize, - /// Automatically add connection pools for user/database pairs we don't have. - #[serde(default = "General::default_passthrough_auth")] - pub passthrough_auth: PassthoughAuth, - /// Server connect timeout. - #[serde(default = "General::default_connect_timeout")] - pub connect_timeout: u64, - /// Attempt connections multiple times on bad networks. - #[serde(default = "General::connect_attempts")] - pub connect_attempts: u64, - /// How long to wait between connection attempts. - #[serde(default = "General::default_connect_attempt_delay")] - pub connect_attempt_delay: u64, - /// How long to wait for a query to return the result before aborting. Dangerous: don't use unless your network is bad. - #[serde(default = "General::default_query_timeout")] - pub query_timeout: u64, - /// Checkout timeout. - #[serde(default = "General::checkout_timeout")] - pub checkout_timeout: u64, - /// Dry run for sharding. Parse the query, route to shard 0. - #[serde(default)] - pub dry_run: bool, - /// Idle timeout. - #[serde(default = "General::idle_timeout")] - pub idle_timeout: u64, - /// Client idle timeout. - #[serde(default = "General::default_client_idle_timeout")] - pub client_idle_timeout: u64, - /// Server lifetime. - #[serde(default = "General::server_lifetime")] - pub server_lifetime: u64, - /// Mirror queue size. - #[serde(default = "General::mirror_queue")] - pub mirror_queue: usize, - /// Mirror exposure - #[serde(default = "General::mirror_exposure")] - pub mirror_exposure: f32, - #[serde(default)] - pub auth_type: AuthType, - /// Disable cross-shard queries. - #[serde(default)] - pub cross_shard_disabled: bool, - /// How often to refresh DNS entries, in ms. - #[serde(default)] - pub dns_ttl: Option, - /// LISTEN/NOTIFY channel size. - #[serde(default)] - pub pub_sub_channel_size: usize, - /// Log client connections. - #[serde(default = "General::log_connections")] - pub log_connections: bool, - /// Log client disconnections. - #[serde(default = "General::log_disconnections")] - pub log_disconnections: bool, - /// Two-phase commit. - #[serde(default)] - pub two_phase_commit: bool, - /// Two-phase commit automatic transactions. - #[serde(default)] - pub two_phase_commit_auto: Option, -} - -impl Default for General { - fn default() -> Self { - Self { - host: Self::host(), - port: Self::port(), - workers: Self::workers(), - default_pool_size: Self::default_pool_size(), - min_pool_size: Self::min_pool_size(), - pooler_mode: Self::pooler_mode(), - healthcheck_interval: Self::healthcheck_interval(), - idle_healthcheck_interval: Self::idle_healthcheck_interval(), - idle_healthcheck_delay: Self::idle_healthcheck_delay(), - healthcheck_timeout: Self::healthcheck_timeout(), - healthcheck_port: Self::healthcheck_port(), - ban_timeout: Self::ban_timeout(), - rollback_timeout: Self::rollback_timeout(), - load_balancing_strategy: Self::load_balancing_strategy(), - read_write_strategy: Self::read_write_strategy(), - read_write_split: Self::read_write_split(), - tls_certificate: Self::tls_certificate(), - tls_private_key: Self::tls_private_key(), - tls_verify: Self::default_tls_verify(), - tls_server_ca_certificate: Self::tls_server_ca_certificate(), - shutdown_timeout: Self::default_shutdown_timeout(), - broadcast_address: Self::broadcast_address(), - broadcast_port: Self::broadcast_port(), - query_log: Self::query_log(), - openmetrics_port: Self::openmetrics_port(), - openmetrics_namespace: Self::openmetrics_namespace(), - prepared_statements: Self::prepared_statements(), - prepared_statements_limit: Self::prepared_statements_limit(), - query_cache_limit: Self::query_cache_limit(), - passthrough_auth: Self::default_passthrough_auth(), - connect_timeout: Self::default_connect_timeout(), - connect_attempt_delay: Self::default_connect_attempt_delay(), - connect_attempts: Self::connect_attempts(), - query_timeout: Self::default_query_timeout(), - checkout_timeout: Self::checkout_timeout(), - dry_run: Self::dry_run(), - idle_timeout: Self::idle_timeout(), - client_idle_timeout: Self::default_client_idle_timeout(), - mirror_queue: Self::mirror_queue(), - mirror_exposure: Self::mirror_exposure(), - auth_type: Self::auth_type(), - cross_shard_disabled: Self::cross_shard_disabled(), - dns_ttl: Self::default_dns_ttl(), - pub_sub_channel_size: Self::pub_sub_channel_size(), - log_connections: Self::log_connections(), - log_disconnections: Self::log_disconnections(), - two_phase_commit: bool::default(), - two_phase_commit_auto: None, - server_lifetime: Self::server_lifetime(), - } - } -} - -impl General { - fn env_or_default(env_var: &str, default: T) -> T { - env::var(env_var) - .ok() - .and_then(|v| v.parse().ok()) - .unwrap_or(default) - } - - fn env_string_or_default(env_var: &str, default: &str) -> String { - env::var(env_var).unwrap_or_else(|_| default.to_string()) - } - - fn env_bool_or_default(env_var: &str, default: bool) -> bool { - env::var(env_var) - .ok() - .and_then(|v| match v.to_lowercase().as_str() { - "true" | "1" | "yes" | "on" => Some(true), - "false" | "0" | "no" | "off" => Some(false), - _ => None, - }) - .unwrap_or(default) - } - - fn env_option(env_var: &str) -> Option { - env::var(env_var).ok().and_then(|v| v.parse().ok()) - } - - fn env_option_string(env_var: &str) -> Option { - env::var(env_var).ok().filter(|s| !s.is_empty()) - } - - fn env_enum_or_default(env_var: &str) -> T { - env::var(env_var) - .ok() - .and_then(|v| v.parse().ok()) - .unwrap_or_default() - } - - fn host() -> String { - Self::env_string_or_default("PGDOG_HOST", "0.0.0.0") - } - - pub fn port() -> u16 { - Self::env_or_default("PGDOG_PORT", 6432) - } - - fn workers() -> usize { - Self::env_or_default("PGDOG_WORKERS", 2) - } - - fn default_pool_size() -> usize { - Self::env_or_default("PGDOG_DEFAULT_POOL_SIZE", 10) - } - - fn min_pool_size() -> usize { - Self::env_or_default("PGDOG_MIN_POOL_SIZE", 1) - } - - fn healthcheck_interval() -> u64 { - Self::env_or_default("PGDOG_HEALTHCHECK_INTERVAL", 30_000) - } - - fn idle_healthcheck_interval() -> u64 { - Self::env_or_default("PGDOG_IDLE_HEALTHCHECK_INTERVAL", 30_000) - } - - fn idle_healthcheck_delay() -> u64 { - Self::env_or_default("PGDOG_IDLE_HEALTHCHECK_DELAY", 5_000) - } - - fn healthcheck_port() -> Option { - Self::env_option("PGDOG_HEALTHCHECK_PORT") - } - - fn ban_timeout() -> u64 { - Self::env_or_default( - "PGDOG_BAN_TIMEOUT", - Duration::from_secs(300).as_millis() as u64, - ) - } - - fn rollback_timeout() -> u64 { - Self::env_or_default("PGDOG_ROLLBACK_TIMEOUT", 5_000) - } - - fn idle_timeout() -> u64 { - Self::env_or_default( - "PGDOG_IDLE_TIMEOUT", - Duration::from_secs(60).as_millis() as u64, - ) - } - - fn default_client_idle_timeout() -> u64 { - Self::env_or_default( - "PGDOG_CLIENT_IDLE_TIMEOUT", - Duration::MAX.as_millis() as u64, - ) - } - - fn default_query_timeout() -> u64 { - Self::env_or_default("PGDOG_QUERY_TIMEOUT", Duration::MAX.as_millis() as u64) - } - - pub(crate) fn query_timeout(&self) -> Duration { - Duration::from_millis(self.query_timeout) - } - - pub fn dns_ttl(&self) -> Option { - self.dns_ttl.map(Duration::from_millis) - } - - pub(crate) fn client_idle_timeout(&self) -> Duration { - Duration::from_millis(self.client_idle_timeout) - } - - pub(crate) fn connect_attempt_delay(&self) -> Duration { - Duration::from_millis(self.connect_attempt_delay) - } - - fn load_balancing_strategy() -> LoadBalancingStrategy { - Self::env_enum_or_default("PGDOG_LOAD_BALANCING_STRATEGY") - } - - fn default_tls_verify() -> TlsVerifyMode { - env::var("PGDOG_TLS_VERIFY") - .ok() - .and_then(|v| v.parse().ok()) - .unwrap_or(TlsVerifyMode::Prefer) - } - - fn default_shutdown_timeout() -> u64 { - Self::env_or_default("PGDOG_SHUTDOWN_TIMEOUT", 60_000) - } - - fn default_connect_timeout() -> u64 { - Self::env_or_default("PGDOG_CONNECT_TIMEOUT", 5_000) - } - - fn default_connect_attempt_delay() -> u64 { - Self::env_or_default("PGDOG_CONNECT_ATTEMPT_DELAY", 0) - } - - fn connect_attempts() -> u64 { - Self::env_or_default("PGDOG_CONNECT_ATTEMPTS", 1) - } - - fn pooler_mode() -> PoolerMode { - Self::env_enum_or_default("PGDOG_POOLER_MODE") - } - - fn read_write_strategy() -> ReadWriteStrategy { - Self::env_enum_or_default("PGDOG_READ_WRITE_STRATEGY") - } - - fn read_write_split() -> ReadWriteSplit { - Self::env_enum_or_default("PGDOG_READ_WRITE_SPLIT") - } - - fn prepared_statements() -> PreparedStatements { - Self::env_enum_or_default("PGDOG_PREPARED_STATEMENTS") - } - - fn auth_type() -> AuthType { - Self::env_enum_or_default("PGDOG_AUTH_TYPE") - } - - fn tls_certificate() -> Option { - Self::env_option_string("PGDOG_TLS_CERTIFICATE").map(PathBuf::from) - } - - fn tls_private_key() -> Option { - Self::env_option_string("PGDOG_TLS_PRIVATE_KEY").map(PathBuf::from) - } - - fn tls_server_ca_certificate() -> Option { - Self::env_option_string("PGDOG_TLS_SERVER_CA_CERTIFICATE").map(PathBuf::from) - } - - fn query_log() -> Option { - Self::env_option_string("PGDOG_QUERY_LOG").map(PathBuf::from) - } - - pub fn openmetrics_port() -> Option { - Self::env_option("PGDOG_OPENMETRICS_PORT") - } - - pub fn openmetrics_namespace() -> Option { - Self::env_option_string("PGDOG_OPENMETRICS_NAMESPACE") - } - - fn default_dns_ttl() -> Option { - Self::env_option("PGDOG_DNS_TTL") - } - - pub fn pub_sub_channel_size() -> usize { - Self::env_or_default("PGDOG_PUB_SUB_CHANNEL_SIZE", 0) - } - - pub fn dry_run() -> bool { - Self::env_bool_or_default("PGDOG_DRY_RUN", false) - } - - pub fn cross_shard_disabled() -> bool { - Self::env_bool_or_default("PGDOG_CROSS_SHARD_DISABLED", false) - } - - pub fn broadcast_address() -> Option { - Self::env_option("PGDOG_BROADCAST_ADDRESS") - } - - pub fn broadcast_port() -> u16 { - Self::env_or_default("PGDOG_BROADCAST_PORT", Self::port() + 1) - } - - fn healthcheck_timeout() -> u64 { - Self::env_or_default( - "PGDOG_HEALTHCHECK_TIMEOUT", - Duration::from_secs(5).as_millis() as u64, - ) - } - - fn checkout_timeout() -> u64 { - Self::env_or_default( - "PGDOG_CHECKOUT_TIMEOUT", - Duration::from_secs(5).as_millis() as u64, - ) - } - - pub fn mirror_queue() -> usize { - Self::env_or_default("PGDOG_MIRROR_QUEUE", 128) - } - - pub fn mirror_exposure() -> f32 { - Self::env_or_default("PGDOG_MIRROR_EXPOSURE", 1.0) - } - - pub fn prepared_statements_limit() -> usize { - Self::env_or_default("PGDOG_PREPARED_STATEMENTS_LIMIT", usize::MAX) - } - - pub fn query_cache_limit() -> usize { - Self::env_or_default("PGDOG_QUERY_CACHE_LIMIT", 50_000) - } - - pub fn log_connections() -> bool { - Self::env_bool_or_default("PGDOG_LOG_CONNECTIONS", true) - } - - pub fn log_disconnections() -> bool { - Self::env_bool_or_default("PGDOG_LOG_DISCONNECTIONS", true) - } - - pub fn server_lifetime() -> u64 { - Self::env_or_default( - "PGDOG_SERVER_LIFETIME", - Duration::from_secs(3600 * 24).as_millis() as u64, - ) - } - - fn default_passthrough_auth() -> PassthoughAuth { - if let Ok(auth) = env::var("PGDOG_PASSTHROUGH_AUTH") { - // TODO: figure out why toml::from_str doesn't work. - match auth.as_str() { - "enabled" => PassthoughAuth::Enabled, - "disabled" => PassthoughAuth::Disabled, - "enabled_plain" => PassthoughAuth::EnabledPlain, - _ => PassthoughAuth::default(), - } - } else { - PassthoughAuth::default() - } - } - - /// Get shutdown timeout as a duration. - pub fn shutdown_timeout(&self) -> Duration { - Duration::from_millis(self.shutdown_timeout) - } - - /// Get TLS config, if any. - pub fn tls(&self) -> Option<(&PathBuf, &PathBuf)> { - if let Some(cert) = &self.tls_certificate { - if let Some(key) = &self.tls_private_key { - return Some((cert, key)); - } - } - - None - } - - pub fn passthrough_auth(&self) -> bool { - self.tls().is_some() && self.passthrough_auth == PassthoughAuth::Enabled - || self.passthrough_auth == PassthoughAuth::EnabledPlain - } - - /// Support for LISTEN/NOTIFY. - pub fn pub_sub_enabled(&self) -> bool { - self.pub_sub_channel_size > 0 - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_env_workers() { - env::set_var("PGDOG_WORKERS", "8"); - assert_eq!(General::workers(), 8); - env::remove_var("PGDOG_WORKERS"); - assert_eq!(General::workers(), 2); - } - - #[test] - fn test_env_pool_sizes() { - env::set_var("PGDOG_DEFAULT_POOL_SIZE", "50"); - env::set_var("PGDOG_MIN_POOL_SIZE", "5"); - - assert_eq!(General::default_pool_size(), 50); - assert_eq!(General::min_pool_size(), 5); - - env::remove_var("PGDOG_DEFAULT_POOL_SIZE"); - env::remove_var("PGDOG_MIN_POOL_SIZE"); - - assert_eq!(General::default_pool_size(), 10); - assert_eq!(General::min_pool_size(), 1); - } - - #[test] - fn test_env_timeouts() { - env::set_var("PGDOG_HEALTHCHECK_INTERVAL", "60000"); - env::set_var("PGDOG_HEALTHCHECK_TIMEOUT", "10000"); - env::set_var("PGDOG_CONNECT_TIMEOUT", "10000"); - env::set_var("PGDOG_CHECKOUT_TIMEOUT", "15000"); - env::set_var("PGDOG_IDLE_TIMEOUT", "120000"); - - assert_eq!(General::healthcheck_interval(), 60000); - assert_eq!(General::healthcheck_timeout(), 10000); - assert_eq!(General::default_connect_timeout(), 10000); - assert_eq!(General::checkout_timeout(), 15000); - assert_eq!(General::idle_timeout(), 120000); - - env::remove_var("PGDOG_HEALTHCHECK_INTERVAL"); - env::remove_var("PGDOG_HEALTHCHECK_TIMEOUT"); - env::remove_var("PGDOG_CONNECT_TIMEOUT"); - env::remove_var("PGDOG_CHECKOUT_TIMEOUT"); - env::remove_var("PGDOG_IDLE_TIMEOUT"); - - assert_eq!(General::healthcheck_interval(), 30000); - assert_eq!(General::healthcheck_timeout(), 5000); - assert_eq!(General::default_connect_timeout(), 5000); - assert_eq!(General::checkout_timeout(), 5000); - assert_eq!(General::idle_timeout(), 60000); - } - - #[test] - fn test_env_invalid_values() { - env::set_var("PGDOG_WORKERS", "invalid"); - env::set_var("PGDOG_DEFAULT_POOL_SIZE", "not_a_number"); - - assert_eq!(General::workers(), 2); - assert_eq!(General::default_pool_size(), 10); - - env::remove_var("PGDOG_WORKERS"); - env::remove_var("PGDOG_DEFAULT_POOL_SIZE"); - } - - #[test] - fn test_env_host_port() { - // Test existing env var functionality - env::set_var("PGDOG_HOST", "192.168.1.1"); - env::set_var("PGDOG_PORT", "8432"); - - assert_eq!(General::host(), "192.168.1.1"); - assert_eq!(General::port(), 8432); - - env::remove_var("PGDOG_HOST"); - env::remove_var("PGDOG_PORT"); - - assert_eq!(General::host(), "0.0.0.0"); - assert_eq!(General::port(), 6432); - } - - #[test] - fn test_env_enum_fields() { - // Test pooler mode - env::set_var("PGDOG_POOLER_MODE", "session"); - assert_eq!(General::pooler_mode(), PoolerMode::Session); - env::remove_var("PGDOG_POOLER_MODE"); - assert_eq!(General::pooler_mode(), PoolerMode::Transaction); - - // Test load balancing strategy - env::set_var("PGDOG_LOAD_BALANCING_STRATEGY", "round_robin"); - assert_eq!( - General::load_balancing_strategy(), - LoadBalancingStrategy::RoundRobin - ); - env::remove_var("PGDOG_LOAD_BALANCING_STRATEGY"); - assert_eq!( - General::load_balancing_strategy(), - LoadBalancingStrategy::Random - ); - - // Test read-write strategy - env::set_var("PGDOG_READ_WRITE_STRATEGY", "aggressive"); - assert_eq!( - General::read_write_strategy(), - ReadWriteStrategy::Aggressive - ); - env::remove_var("PGDOG_READ_WRITE_STRATEGY"); - assert_eq!( - General::read_write_strategy(), - ReadWriteStrategy::Conservative - ); - - // Test read-write split - env::set_var("PGDOG_READ_WRITE_SPLIT", "exclude_primary"); - assert_eq!(General::read_write_split(), ReadWriteSplit::ExcludePrimary); - env::remove_var("PGDOG_READ_WRITE_SPLIT"); - assert_eq!(General::read_write_split(), ReadWriteSplit::IncludePrimary); - - // Test TLS verify mode - env::set_var("PGDOG_TLS_VERIFY", "verify_full"); - assert_eq!(General::default_tls_verify(), TlsVerifyMode::VerifyFull); - env::remove_var("PGDOG_TLS_VERIFY"); - assert_eq!(General::default_tls_verify(), TlsVerifyMode::Prefer); - - // Test prepared statements - env::set_var("PGDOG_PREPARED_STATEMENTS", "full"); - assert_eq!(General::prepared_statements(), PreparedStatements::Full); - env::remove_var("PGDOG_PREPARED_STATEMENTS"); - assert_eq!(General::prepared_statements(), PreparedStatements::Extended); - - // Test auth type - env::set_var("PGDOG_AUTH_TYPE", "md5"); - assert_eq!(General::auth_type(), AuthType::Md5); - env::remove_var("PGDOG_AUTH_TYPE"); - assert_eq!(General::auth_type(), AuthType::Scram); - } - - #[test] - fn test_env_additional_timeouts() { - env::set_var("PGDOG_IDLE_HEALTHCHECK_INTERVAL", "45000"); - env::set_var("PGDOG_IDLE_HEALTHCHECK_DELAY", "10000"); - env::set_var("PGDOG_BAN_TIMEOUT", "600000"); - env::set_var("PGDOG_ROLLBACK_TIMEOUT", "10000"); - env::set_var("PGDOG_SHUTDOWN_TIMEOUT", "120000"); - env::set_var("PGDOG_CONNECT_ATTEMPT_DELAY", "1000"); - env::set_var("PGDOG_QUERY_TIMEOUT", "30000"); - env::set_var("PGDOG_CLIENT_IDLE_TIMEOUT", "3600000"); - - assert_eq!(General::idle_healthcheck_interval(), 45000); - assert_eq!(General::idle_healthcheck_delay(), 10000); - assert_eq!(General::ban_timeout(), 600000); - assert_eq!(General::rollback_timeout(), 10000); - assert_eq!(General::default_shutdown_timeout(), 120000); - assert_eq!(General::default_connect_attempt_delay(), 1000); - assert_eq!(General::default_query_timeout(), 30000); - assert_eq!(General::default_client_idle_timeout(), 3600000); - - env::remove_var("PGDOG_IDLE_HEALTHCHECK_INTERVAL"); - env::remove_var("PGDOG_IDLE_HEALTHCHECK_DELAY"); - env::remove_var("PGDOG_BAN_TIMEOUT"); - env::remove_var("PGDOG_ROLLBACK_TIMEOUT"); - env::remove_var("PGDOG_SHUTDOWN_TIMEOUT"); - env::remove_var("PGDOG_CONNECT_ATTEMPT_DELAY"); - env::remove_var("PGDOG_QUERY_TIMEOUT"); - env::remove_var("PGDOG_CLIENT_IDLE_TIMEOUT"); - - assert_eq!(General::idle_healthcheck_interval(), 30000); - assert_eq!(General::idle_healthcheck_delay(), 5000); - assert_eq!(General::ban_timeout(), 300000); - assert_eq!(General::rollback_timeout(), 5000); - assert_eq!(General::default_shutdown_timeout(), 60000); - assert_eq!(General::default_connect_attempt_delay(), 0); - } - - #[test] - fn test_env_path_fields() { - env::set_var("PGDOG_TLS_CERTIFICATE", "/path/to/cert.pem"); - env::set_var("PGDOG_TLS_PRIVATE_KEY", "/path/to/key.pem"); - env::set_var("PGDOG_TLS_SERVER_CA_CERTIFICATE", "/path/to/ca.pem"); - env::set_var("PGDOG_QUERY_LOG", "/var/log/pgdog/queries.log"); - - assert_eq!( - General::tls_certificate(), - Some(PathBuf::from("/path/to/cert.pem")) - ); - assert_eq!( - General::tls_private_key(), - Some(PathBuf::from("/path/to/key.pem")) - ); - assert_eq!( - General::tls_server_ca_certificate(), - Some(PathBuf::from("/path/to/ca.pem")) - ); - assert_eq!( - General::query_log(), - Some(PathBuf::from("/var/log/pgdog/queries.log")) - ); - - env::remove_var("PGDOG_TLS_CERTIFICATE"); - env::remove_var("PGDOG_TLS_PRIVATE_KEY"); - env::remove_var("PGDOG_TLS_SERVER_CA_CERTIFICATE"); - env::remove_var("PGDOG_QUERY_LOG"); - - assert_eq!(General::tls_certificate(), None); - assert_eq!(General::tls_private_key(), None); - assert_eq!(General::tls_server_ca_certificate(), None); - assert_eq!(General::query_log(), None); - } - - #[test] - fn test_env_numeric_fields() { - env::set_var("PGDOG_BROADCAST_PORT", "7432"); - env::set_var("PGDOG_OPENMETRICS_PORT", "9090"); - env::set_var("PGDOG_PREPARED_STATEMENTS_LIMIT", "1000"); - env::set_var("PGDOG_QUERY_CACHE_LIMIT", "500"); - env::set_var("PGDOG_CONNECT_ATTEMPTS", "3"); - env::set_var("PGDOG_MIRROR_QUEUE", "256"); - env::set_var("PGDOG_MIRROR_EXPOSURE", "0.5"); - env::set_var("PGDOG_DNS_TTL", "60000"); - env::set_var("PGDOG_PUB_SUB_CHANNEL_SIZE", "100"); - - assert_eq!(General::broadcast_port(), 7432); - assert_eq!(General::openmetrics_port(), Some(9090)); - assert_eq!(General::prepared_statements_limit(), 1000); - assert_eq!(General::query_cache_limit(), 500); - assert_eq!(General::connect_attempts(), 3); - assert_eq!(General::mirror_queue(), 256); - assert_eq!(General::mirror_exposure(), 0.5); - assert_eq!(General::default_dns_ttl(), Some(60000)); - assert_eq!(General::pub_sub_channel_size(), 100); - - env::remove_var("PGDOG_BROADCAST_PORT"); - env::remove_var("PGDOG_OPENMETRICS_PORT"); - env::remove_var("PGDOG_PREPARED_STATEMENTS_LIMIT"); - env::remove_var("PGDOG_QUERY_CACHE_LIMIT"); - env::remove_var("PGDOG_CONNECT_ATTEMPTS"); - env::remove_var("PGDOG_MIRROR_QUEUE"); - env::remove_var("PGDOG_MIRROR_EXPOSURE"); - env::remove_var("PGDOG_DNS_TTL"); - env::remove_var("PGDOG_PUB_SUB_CHANNEL_SIZE"); - - assert_eq!(General::broadcast_port(), General::port() + 1); - assert_eq!(General::openmetrics_port(), None); - assert_eq!(General::prepared_statements_limit(), usize::MAX); - assert_eq!(General::query_cache_limit(), 50_000); - assert_eq!(General::connect_attempts(), 1); - assert_eq!(General::mirror_queue(), 128); - assert_eq!(General::mirror_exposure(), 1.0); - assert_eq!(General::default_dns_ttl(), None); - assert_eq!(General::pub_sub_channel_size(), 0); - } - - #[test] - fn test_env_boolean_fields() { - env::set_var("PGDOG_DRY_RUN", "true"); - env::set_var("PGDOG_CROSS_SHARD_DISABLED", "yes"); - env::set_var("PGDOG_LOG_CONNECTIONS", "false"); - env::set_var("PGDOG_LOG_DISCONNECTIONS", "0"); - - assert_eq!(General::dry_run(), true); - assert_eq!(General::cross_shard_disabled(), true); - assert_eq!(General::log_connections(), false); - assert_eq!(General::log_disconnections(), false); - - env::remove_var("PGDOG_DRY_RUN"); - env::remove_var("PGDOG_CROSS_SHARD_DISABLED"); - env::remove_var("PGDOG_LOG_CONNECTIONS"); - env::remove_var("PGDOG_LOG_DISCONNECTIONS"); - - assert_eq!(General::dry_run(), false); - assert_eq!(General::cross_shard_disabled(), false); - assert_eq!(General::log_connections(), true); - assert_eq!(General::log_disconnections(), true); - } - - #[test] - fn test_env_other_fields() { - env::set_var("PGDOG_BROADCAST_ADDRESS", "192.168.1.100"); - env::set_var("PGDOG_OPENMETRICS_NAMESPACE", "pgdog_metrics"); - - assert_eq!( - General::broadcast_address(), - Some("192.168.1.100".parse().unwrap()) - ); - assert_eq!( - General::openmetrics_namespace(), - Some("pgdog_metrics".to_string()) - ); - - env::remove_var("PGDOG_BROADCAST_ADDRESS"); - env::remove_var("PGDOG_OPENMETRICS_NAMESPACE"); - - assert_eq!(General::broadcast_address(), None); - assert_eq!(General::openmetrics_namespace(), None); - } - - #[test] - fn test_env_invalid_enum_values() { - env::set_var("PGDOG_POOLER_MODE", "invalid_mode"); - env::set_var("PGDOG_AUTH_TYPE", "not_an_auth"); - env::set_var("PGDOG_TLS_VERIFY", "bad_verify"); - - // Should fall back to defaults for invalid values - assert_eq!(General::pooler_mode(), PoolerMode::Transaction); - assert_eq!(General::auth_type(), AuthType::Scram); - assert_eq!(General::default_tls_verify(), TlsVerifyMode::Prefer); - - env::remove_var("PGDOG_POOLER_MODE"); - env::remove_var("PGDOG_AUTH_TYPE"); - env::remove_var("PGDOG_TLS_VERIFY"); - } - - #[test] - fn test_general_default_uses_env_vars() { - // Set some environment variables - env::set_var("PGDOG_WORKERS", "8"); - env::set_var("PGDOG_POOLER_MODE", "session"); - env::set_var("PGDOG_AUTH_TYPE", "trust"); - env::set_var("PGDOG_DRY_RUN", "true"); - - let general = General::default(); - - assert_eq!(general.workers, 8); - assert_eq!(general.pooler_mode, PoolerMode::Session); - assert_eq!(general.auth_type, AuthType::Trust); - assert_eq!(general.dry_run, true); - - env::remove_var("PGDOG_WORKERS"); - env::remove_var("PGDOG_POOLER_MODE"); - env::remove_var("PGDOG_AUTH_TYPE"); - env::remove_var("PGDOG_DRY_RUN"); - } -} +pub use pgdog_config::general::General; diff --git a/pgdog/src/config/memory.rs b/pgdog/src/config/memory.rs new file mode 100644 index 000000000..95f0b563e --- /dev/null +++ b/pgdog/src/config/memory.rs @@ -0,0 +1 @@ +pub use pgdog_config::memory::*; diff --git a/pgdog/src/config/mod.rs b/pgdog/src/config/mod.rs index 3efcd93c1..75c729e81 100644 --- a/pgdog/src/config/mod.rs +++ b/pgdog/src/config/mod.rs @@ -1,45 +1,31 @@ //! Configuration. // Submodules -pub mod auth; pub mod convert; pub mod core; pub mod database; pub mod error; pub mod general; +pub mod memory; pub mod networking; pub mod overrides; pub mod pooling; pub mod replication; +pub mod rewrite; pub mod sharding; -pub mod url; pub mod users; -// Re-export from error module -pub use error::Error; - -// Re-export from overrides module -pub use overrides::Overrides; - -// Re-export core configuration types pub use core::{Config, ConfigAndUsers}; - -// Re-export from general module +pub use database::{Database, Role}; +pub use error::Error; pub use general::General; - -// Re-export from auth module -pub use auth::{AuthType, PassthoughAuth}; - -// Re-export from pooling module -pub use pooling::{PoolerMode, PreparedStatements, Stats}; - -// Re-export from database module -pub use database::{Database, LoadBalancingStrategy, ReadWriteSplit, ReadWriteStrategy, Role}; - -// Re-export from networking module +pub use memory::*; pub use networking::{MultiTenant, Tcp, TlsVerifyMode}; - -// Re-export from users module +pub use overrides::Overrides; +pub use pgdog_config::auth::{AuthType, PassthoughAuth}; +pub use pgdog_config::{LoadBalancingStrategy, ReadWriteSplit, ReadWriteStrategy}; +pub use pooling::{ConnectionRecovery, PoolerMode, PreparedStatements, Stats}; +pub use rewrite::{Rewrite, RewriteMode}; pub use users::{Admin, Plugin, User, Users}; // Re-export from sharding module @@ -95,6 +81,8 @@ pub fn from_urls(urls: &[String]) -> Result { /// Extract all database URLs from the environment and /// create the config. pub fn from_env() -> Result { + let _lock = LOCK.lock(); + let mut urls = vec![]; let mut index = 1; while let Ok(url) = env::var(format!("PGDOG_DATABASE_URL_{}", index)) { @@ -103,10 +91,26 @@ pub fn from_env() -> Result { } if urls.is_empty() { - Err(Error::NoDbsInEnv) - } else { - from_urls(&urls) + return Err(Error::NoDbsInEnv); } + + let mut config = (*config()).clone(); + config = config.databases_from_urls(&urls)?; + + // Extract mirroring configuration + let mut mirror_strs = vec![]; + let mut index = 1; + while let Ok(mirror_str) = env::var(format!("PGDOG_MIRRORING_{}", index)) { + mirror_strs.push(mirror_str); + index += 1; + } + + if !mirror_strs.is_empty() { + config = config.mirroring_from_strings(&mirror_strs)?; + } + + CONFIG.store(Arc::new(config.clone())); + Ok(config) } /// Override some settings. @@ -140,6 +144,11 @@ pub fn overrides(overrides: Overrides) { // Test helper functions #[cfg(test)] pub fn load_test() { + load_test_with_pooler_mode(PoolerMode::Transaction) +} + +#[cfg(test)] +pub fn load_test_with_pooler_mode(pooler_mode: PoolerMode) { use crate::backend::databases::init; let mut config = ConfigAndUsers::default(); @@ -147,17 +156,19 @@ pub fn load_test() { name: "pgdog".into(), host: "127.0.0.1".into(), port: 5432, + pooler_mode: Some(pooler_mode), ..Default::default() }]; config.users.users = vec![User { name: "pgdog".into(), database: "pgdog".into(), password: Some("pgdog".into()), + pooler_mode: Some(pooler_mode), ..Default::default() }]; set(config).unwrap(); - init(); + init().unwrap(); } #[cfg(test)] @@ -191,5 +202,110 @@ pub fn load_test_replicas() { }]; set(config).unwrap(); - init(); + init().unwrap(); +} + +#[cfg(test)] +pub fn load_test_sharded() { + use pgdog_config::ShardedSchema; + + use crate::backend::databases::init; + + let mut config = ConfigAndUsers::default(); + config.config.general.min_pool_size = 0; + config.config.databases = vec![ + Database { + name: "pgdog".into(), + host: "127.0.0.1".into(), + port: 5432, + role: Role::Primary, + database_name: Some("shard_0".into()), + shard: 0, + ..Default::default() + }, + Database { + name: "pgdog".into(), + host: "127.0.0.1".into(), + port: 5432, + role: Role::Replica, + read_only: Some(true), + database_name: Some("shard_0".into()), + shard: 0, + ..Default::default() + }, + Database { + name: "pgdog".into(), + host: "127.0.0.1".into(), + port: 5432, + role: Role::Primary, + database_name: Some("shard_1".into()), + shard: 1, + ..Default::default() + }, + Database { + name: "pgdog".into(), + host: "127.0.0.1".into(), + port: 5432, + role: Role::Replica, + read_only: Some(true), + database_name: Some("shard_1".into()), + shard: 1, + ..Default::default() + }, + ]; + config.config.sharded_tables = vec![ + ShardedTable { + database: "pgdog".into(), + name: Some("sharded".into()), + column: "id".into(), + ..Default::default() + }, + ShardedTable { + database: "pgdog".into(), + name: Some("sharded_varchar".into()), + column: "id_varchar".into(), + data_type: DataType::Varchar, + ..Default::default() + }, + ShardedTable { + database: "pgdog".into(), + name: Some("sharded_uuid".into()), + column: "id_uuid".into(), + data_type: DataType::Uuid, + ..Default::default() + }, + ]; + config.config.sharded_schemas = vec![ + ShardedSchema { + database: "pgdog".into(), + name: Some("acustomer".into()), + shard: 0, + ..Default::default() + }, + ShardedSchema { + database: "pgdog".into(), + name: Some("bcustomer".into()), + shard: 1, + ..Default::default() + }, + ShardedSchema { + database: "pgdog".into(), + name: Some("all".into()), + all: true, + ..Default::default() + }, + ]; + config.config.rewrite.enabled = true; + config.config.rewrite.split_inserts = RewriteMode::Rewrite; + config.config.rewrite.shard_key = RewriteMode::Rewrite; + config.config.general.load_balancing_strategy = LoadBalancingStrategy::RoundRobin; + config.users.users = vec![User { + name: "pgdog".into(), + database: "pgdog".into(), + password: Some("pgdog".into()), + ..Default::default() + }]; + + set(config).unwrap(); + init().unwrap(); } diff --git a/pgdog/src/config/networking.rs b/pgdog/src/config/networking.rs index 48354073b..34ff28084 100644 --- a/pgdog/src/config/networking.rs +++ b/pgdog/src/config/networking.rs @@ -1,112 +1 @@ -use serde::{Deserialize, Serialize}; -use std::str::FromStr; -use std::time::Duration; - -use crate::util::human_duration_optional; - -#[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq, Copy)] -#[serde(rename_all = "snake_case")] -pub enum TlsVerifyMode { - #[default] - Disabled, - Prefer, - VerifyCa, - VerifyFull, -} - -impl FromStr for TlsVerifyMode { - type Err = String; - - fn from_str(s: &str) -> Result { - match s.to_lowercase().replace(['_', '-'], "").as_str() { - "disabled" => Ok(Self::Disabled), - "prefer" => Ok(Self::Prefer), - "verifyca" => Ok(Self::VerifyCa), - "verifyfull" => Ok(Self::VerifyFull), - _ => Err(format!("Invalid TLS verify mode: {}", s)), - } - } -} - -#[derive(Serialize, Deserialize, PartialEq, Debug, Clone)] -pub struct Tcp { - #[serde(default = "Tcp::default_keepalive")] - keepalive: bool, - user_timeout: Option, - time: Option, - interval: Option, - retries: Option, - congestion_control: Option, -} - -impl std::fmt::Display for Tcp { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!( - f, - "keepalive={} user_timeout={} time={} interval={}, retries={}, congestion_control={}", - self.keepalive(), - human_duration_optional(self.user_timeout()), - human_duration_optional(self.time()), - human_duration_optional(self.interval()), - if let Some(retries) = self.retries() { - retries.to_string() - } else { - "default".into() - }, - if let Some(ref c) = self.congestion_control { - c.as_str() - } else { - "" - }, - ) - } -} - -impl Default for Tcp { - fn default() -> Self { - Self { - keepalive: Self::default_keepalive(), - user_timeout: None, - time: None, - interval: None, - retries: None, - congestion_control: None, - } - } -} - -impl Tcp { - fn default_keepalive() -> bool { - true - } - - pub fn keepalive(&self) -> bool { - self.keepalive - } - - pub fn time(&self) -> Option { - self.time.map(Duration::from_millis) - } - - pub fn interval(&self) -> Option { - self.interval.map(Duration::from_millis) - } - - pub fn user_timeout(&self) -> Option { - self.user_timeout.map(Duration::from_millis) - } - - pub fn retries(&self) -> Option { - self.retries - } - - pub fn congestion_control(&self) -> &Option { - &self.congestion_control - } -} - -#[derive(Serialize, Deserialize, PartialEq, Debug, Clone)] -#[serde(rename_all = "snake_case")] -pub struct MultiTenant { - pub column: String, -} +pub use pgdog_config::{MultiTenant, Tcp, TlsVerifyMode}; diff --git a/pgdog/src/config/overrides.rs b/pgdog/src/config/overrides.rs index 1769713dc..d1b7edbc6 100644 --- a/pgdog/src/config/overrides.rs +++ b/pgdog/src/config/overrides.rs @@ -1,6 +1 @@ -#[derive(Debug, Clone, Default)] -pub struct Overrides { - pub default_pool_size: Option, - pub min_pool_size: Option, - pub session_mode: Option, -} +pub use pgdog_config::Overrides; diff --git a/pgdog/src/config/pooling.rs b/pgdog/src/config/pooling.rs index 5cb3f7cf5..1dde71b21 100644 --- a/pgdog/src/config/pooling.rs +++ b/pgdog/src/config/pooling.rs @@ -1,67 +1 @@ -use serde::{Deserialize, Serialize}; -use std::str::FromStr; - -#[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq)] -#[serde(rename_all = "snake_case")] -pub enum PreparedStatements { - Disabled, - #[default] - Extended, - Full, -} - -impl PreparedStatements { - pub fn full(&self) -> bool { - matches!(self, PreparedStatements::Full) - } - - pub fn enabled(&self) -> bool { - !matches!(self, PreparedStatements::Disabled) - } -} - -impl FromStr for PreparedStatements { - type Err = String; - - fn from_str(s: &str) -> Result { - match s.to_lowercase().as_str() { - "disabled" => Ok(Self::Disabled), - "extended" => Ok(Self::Extended), - "full" => Ok(Self::Full), - _ => Err(format!("Invalid prepared statements mode: {}", s)), - } - } -} - -/// Empty struct for stats -#[derive(Serialize, Deserialize, Debug, Clone, Default)] -pub struct Stats {} - -#[derive(Serialize, Deserialize, Debug, Clone, Copy, Default, PartialEq, Eq, Ord, PartialOrd)] -#[serde(rename_all = "snake_case")] -pub enum PoolerMode { - #[default] - Transaction, - Session, -} - -impl std::fmt::Display for PoolerMode { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - Self::Transaction => write!(f, "transaction"), - Self::Session => write!(f, "session"), - } - } -} - -impl FromStr for PoolerMode { - type Err = String; - - fn from_str(s: &str) -> Result { - match s.to_lowercase().as_str() { - "transaction" => Ok(Self::Transaction), - "session" => Ok(Self::Session), - _ => Err(format!("Invalid pooler mode: {}", s)), - } - } -} +pub use pgdog_config::{pooling::ConnectionRecovery, PoolerMode, PreparedStatements, Stats}; diff --git a/pgdog/src/config/replication.rs b/pgdog/src/config/replication.rs index fd24f6b42..a9c790a15 100644 --- a/pgdog/src/config/replication.rs +++ b/pgdog/src/config/replication.rs @@ -1,141 +1 @@ -use serde::{Deserialize, Serialize}; -use std::path::PathBuf; -use std::time::Duration; - -#[derive(Deserialize)] -struct RawReplicaLag { - #[serde(default)] - check_interval: Option, - #[serde(default)] - max_age: Option, -} - -#[derive(Debug, Clone)] -pub struct ReplicaLag { - pub check_interval: Duration, - pub max_age: Duration, -} - -impl ReplicaLag { - fn default_max_age() -> Duration { - Duration::from_millis(25) - } - - fn default_check_interval() -> Duration { - Duration::from_millis(1000) - } - - /// Custom "all-or-none" deserializer that returns Option. - pub fn deserialize_optional<'de, D>(de: D) -> Result, D::Error> - where - D: serde::de::Deserializer<'de>, - { - let maybe: Option = Option::deserialize(de)?; - - Ok(match maybe { - None => None, - - Some(RawReplicaLag { - check_interval: None, - max_age: None, - }) => None, - - Some(RawReplicaLag { - check_interval: Some(ci_u64), - max_age: Some(ma_u64), - }) => Some(ReplicaLag { - check_interval: Duration::from_millis(ci_u64), - max_age: Duration::from_millis(ma_u64), - }), - - Some(RawReplicaLag { - check_interval: None, - max_age: Some(ma_u64), - }) => Some(ReplicaLag { - check_interval: Self::default_check_interval(), - max_age: Duration::from_millis(ma_u64), - }), - - _ => { - return Err(serde::de::Error::custom( - "replica_lag: cannot set check_interval without max_age", - )) - } - }) - } -} - -// NOTE: serialize and deserialize are not inverses. -// - Normally you'd expect ser <-> deser to round-trip, but here deser applies defaults... -// for missing fields -// - Serializes takes those applied defaults into account so that ReplicaLag always reflects... -// the actual effective values. -// - This ensures pgdog.admin sees the true config that is applied, not just what was configured. - -impl Serialize for ReplicaLag { - fn serialize(&self, serializer: S) -> Result - where - S: serde::Serializer, - { - use serde::ser::SerializeStruct; - let mut state = serializer.serialize_struct("ReplicaLag", 2)?; - state.serialize_field("check_interval", &(self.check_interval.as_millis() as u64))?; - state.serialize_field("max_age", &(self.max_age.as_millis() as u64))?; - state.end() - } -} - -impl Default for ReplicaLag { - fn default() -> Self { - Self { - check_interval: Self::default_check_interval(), - max_age: Self::default_max_age(), - } - } -} - -/// Replication configuration. -#[derive(Serialize, Deserialize, PartialEq, Debug, Clone)] -#[serde(rename_all = "snake_case")] -pub struct Replication { - /// Path to the pg_dump executable. - #[serde(default = "Replication::pg_dump_path")] - pub pg_dump_path: PathBuf, -} - -impl Replication { - fn pg_dump_path() -> PathBuf { - PathBuf::from("pg_dump") - } -} - -impl Default for Replication { - fn default() -> Self { - Self { - pg_dump_path: Self::pg_dump_path(), - } - } -} - -/// Mirroring configuration. -#[derive(Serialize, Deserialize, Debug, Clone, Default)] -#[serde(deny_unknown_fields)] -pub struct Mirroring { - /// Source database name to mirror from. - pub source_db: String, - /// Destination database name to mirror to. - pub destination_db: String, - /// Queue length for this mirror (overrides global mirror_queue). - pub queue_length: Option, - /// Exposure for this mirror (overrides global mirror_exposure). - pub exposure: Option, -} - -/// Runtime mirror configuration with resolved values. -#[derive(Debug, Clone)] -pub struct MirrorConfig { - /// Queue length for this mirror. - pub queue_length: usize, - /// Exposure for this mirror. - pub exposure: f32, -} +pub use pgdog_config::replication::*; diff --git a/pgdog/src/config/rewrite.rs b/pgdog/src/config/rewrite.rs new file mode 100644 index 000000000..e273e5978 --- /dev/null +++ b/pgdog/src/config/rewrite.rs @@ -0,0 +1 @@ +pub use pgdog_config::rewrite::*; diff --git a/pgdog/src/config/sharding.rs b/pgdog/src/config/sharding.rs index 534e610f9..c88567f54 100644 --- a/pgdog/src/config/sharding.rs +++ b/pgdog/src/config/sharding.rs @@ -1,174 +1,4 @@ -use serde::{Deserialize, Serialize}; -use std::collections::{hash_map::DefaultHasher, HashSet}; -use std::hash::{Hash, Hasher as StdHasher}; -use std::path::PathBuf; -use tracing::{info, warn}; - -use super::error::Error; -use crate::frontend::router::sharding::Mapping; -use crate::net::messages::Vector; - -/// Sharded table. -#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Default)] -#[serde(rename_all = "snake_case", deny_unknown_fields)] -pub struct ShardedTable { - /// Database this table belongs to. - pub database: String, - /// Table name. If none specified, all tables with the specified - /// column are considered sharded. - pub name: Option, - /// Table sharded on this column. - #[serde(default)] - pub column: String, - /// This table is the primary sharding anchor (e.g. "users"). - #[serde(default)] - pub primary: bool, - /// Centroids for vector sharding. - #[serde(default)] - pub centroids: Vec, - #[serde(default)] - pub centroids_path: Option, - /// Data type of the column. - #[serde(default)] - pub data_type: DataType, - /// How many centroids to probe. - #[serde(default)] - pub centroid_probes: usize, - /// Hasher function. - #[serde(default)] - pub hasher: Hasher, - /// Explicit routing rules. - #[serde(skip, default)] - pub mapping: Option, -} - -impl ShardedTable { - /// Load centroids from file, if provided. - /// - /// Centroids can be very large vectors (1000+ columns). - /// Hardcoding them in pgdog.toml is then impractical. - pub fn load_centroids(&mut self) -> Result<(), Error> { - if let Some(centroids_path) = &self.centroids_path { - if let Ok(f) = std::fs::read_to_string(centroids_path) { - let centroids: Vec = serde_json::from_str(&f)?; - self.centroids = centroids; - info!("loaded {} centroids", self.centroids.len()); - } else { - warn!( - "centroids at path \"{}\" not found", - centroids_path.display() - ); - } - } - - if self.centroid_probes < 1 { - self.centroid_probes = (self.centroids.len() as f32).sqrt().ceil() as usize; - if self.centroid_probes > 0 { - info!("setting centroid probes to {}", self.centroid_probes); - } - } - - Ok(()) - } -} - -#[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq)] -#[serde(rename_all = "snake_case")] -pub enum Hasher { - #[default] - Postgres, - Sha1, -} - -#[derive(Serialize, Deserialize, PartialEq, Debug, Clone, Default, Copy, Eq, Hash)] -#[serde(rename_all = "snake_case")] -pub enum DataType { - #[default] - Bigint, - Uuid, - Vector, - Varchar, -} - -#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Default, Eq)] -#[serde(rename_all = "snake_case", deny_unknown_fields)] -pub struct ShardedMapping { - pub database: String, - pub column: String, - pub table: Option, - pub kind: ShardedMappingKind, - pub start: Option, - pub end: Option, - #[serde(default)] - pub values: HashSet, - pub shard: usize, -} - -impl Hash for ShardedMapping { - fn hash(&self, state: &mut H) { - self.database.hash(state); - self.column.hash(state); - self.table.hash(state); - self.kind.hash(state); - self.start.hash(state); - self.end.hash(state); - - // Hash the values in a deterministic way by XORing their individual hashes - let mut values_hash = 0u64; - for value in &self.values { - let mut hasher = DefaultHasher::new(); - value.hash(&mut hasher); - values_hash ^= hasher.finish(); - } - values_hash.hash(state); - - self.shard.hash(state); - } -} - -#[derive(Serialize, Deserialize, PartialEq, Debug, Clone, Default, Hash, Eq)] -#[serde(rename_all = "snake_case", deny_unknown_fields)] -pub enum ShardedMappingKind { - #[default] - List, - Range, -} - -#[derive(Serialize, Deserialize, PartialEq, Debug, Clone, Eq, Hash)] -#[serde(untagged)] -pub enum FlexibleType { - Integer(i64), - Uuid(uuid::Uuid), - String(String), -} - -impl From for FlexibleType { - fn from(value: i64) -> Self { - Self::Integer(value) - } -} - -impl From for FlexibleType { - fn from(value: uuid::Uuid) -> Self { - Self::Uuid(value) - } -} - -impl From for FlexibleType { - fn from(value: String) -> Self { - Self::String(value) - } -} - -#[derive(Serialize, Deserialize, PartialEq, Debug, Clone, Default)] -#[serde(rename_all = "snake_case", deny_unknown_fields)] -pub struct OmnishardedTables { - pub database: String, - pub tables: Vec, -} - -/// Queries with manual routing rules. -#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, Hash)] -pub struct ManualQuery { - pub fingerprint: String, -} +pub use pgdog_config::sharding::{ + DataType, FlexibleType, Hasher, ManualQuery, OmnishardedTables, QueryParserLevel, + ShardedMapping, ShardedMappingKey, ShardedMappingKind, ShardedTable, +}; diff --git a/pgdog/src/config/url.rs b/pgdog/src/config/url.rs deleted file mode 100644 index a780f4ae3..000000000 --- a/pgdog/src/config/url.rs +++ /dev/null @@ -1,88 +0,0 @@ -//! Parse URL and convert to config struct. -use std::{collections::BTreeSet, env::var}; -use url::Url; - -use super::{ConfigAndUsers, Database, Error, User, Users}; - -fn database_name(url: &Url) -> String { - let database = url.path().chars().skip(1).collect::(); - if database.is_empty() { - "postgres".into() - } else { - database - } -} - -impl From<&Url> for Database { - fn from(value: &Url) -> Self { - let host = value - .host() - .map(|host| host.to_string()) - .unwrap_or("127.0.0.1".into()); - let port = value.port().unwrap_or(5432); - - Database { - name: database_name(value), - host, - port, - ..Default::default() - } - } -} - -impl From<&Url> for User { - fn from(value: &Url) -> Self { - let user = value.username(); - let user = if user.is_empty() { - var("USER").unwrap_or("postgres".into()) - } else { - user.to_string() - }; - let password = value.password().unwrap_or("postgres").to_owned(); - - User { - name: user, - password: Some(password), - database: database_name(value), - ..Default::default() - } - } -} - -impl ConfigAndUsers { - /// Load from database URLs. - pub fn databases_from_urls(mut self, urls: &[String]) -> Result { - let urls = urls - .iter() - .map(|url| Url::parse(url)) - .collect::, url::ParseError>>()?; - let databases = urls - .iter() - .map(Database::from) - .collect::>() // Make sure we only have unique entries. - .into_iter() - .collect::>(); - let users = urls - .iter() - .map(User::from) - .collect::>() // Make sure we only have unique entries. - .into_iter() - .collect::>(); - - self.users = Users { users }; - self.config.databases = databases; - - Ok(self) - } -} - -#[cfg(test)] -mod test { - use super::*; - - #[test] - fn test_url() { - let url = Url::parse("postgres://user:password@host:5432/name").unwrap(); - println!("{:#?}", url); - } -} diff --git a/pgdog/src/config/users.rs b/pgdog/src/config/users.rs index 2f3dfa3a6..e3f34933c 100644 --- a/pgdog/src/config/users.rs +++ b/pgdog/src/config/users.rs @@ -1,167 +1 @@ -use serde::{Deserialize, Serialize}; -use std::collections::HashMap; -use std::env; -use tracing::warn; - -use super::core::Config; -use super::pooling::PoolerMode; -use crate::util::random_string; - -/// pgDog plugin. -#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)] -pub struct Plugin { - /// Plugin name. - pub name: String, -} - -/// Users and passwords. -#[derive(Serialize, Deserialize, Debug, Clone, Default)] -#[serde(deny_unknown_fields)] -pub struct Users { - /// Users and passwords. - #[serde(default)] - pub users: Vec, -} - -impl Users { - /// Organize users by database name. - pub fn users(&self) -> HashMap> { - let mut users = HashMap::new(); - - for user in &self.users { - let entry = users.entry(user.database.clone()).or_insert_with(Vec::new); - entry.push(user.clone()); - } - - users - } - - pub fn check(&mut self, config: &Config) { - for user in &mut self.users { - if user.password().is_empty() { - if !config.general.passthrough_auth() { - warn!( - "user \"{}\" doesn't have a password and passthrough auth is disabled", - user.name - ); - } - - if let Some(min_pool_size) = user.min_pool_size { - if min_pool_size > 0 { - warn!("user \"{}\" (database \"{}\") doesn't have a password configured, \ - so we can't connect to the server to maintain min_pool_size of {}; setting it to 0", user.name, user.database, min_pool_size); - user.min_pool_size = Some(0); - } - } - } - } - } -} - -/// User allowed to connect to pgDog. -#[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq, Eq, Ord, PartialOrd)] -#[serde(deny_unknown_fields)] -pub struct User { - /// User name. - pub name: String, - /// Database name, from pgdog.toml. - pub database: String, - /// User's password. - pub password: Option, - /// Pool size for this user pool, overriding `default_pool_size`. - pub pool_size: Option, - /// Minimum pool size for this user pool, overriding `min_pool_size`. - pub min_pool_size: Option, - /// Pooler mode. - pub pooler_mode: Option, - /// Server username. - pub server_user: Option, - /// Server password. - pub server_password: Option, - /// Statement timeout. - pub statement_timeout: Option, - /// Relication mode. - #[serde(default)] - pub replication_mode: bool, - /// Sharding into this database. - pub replication_sharding: Option, - /// Idle timeout. - pub idle_timeout: Option, - /// Read-only mode. - pub read_only: Option, - /// Schema owner. - #[serde(default)] - pub schema_admin: bool, - /// Disable cross-shard queries for this user. - pub cross_shard_disabled: Option, - /// Two-pc. - pub two_phase_commit: Option, - /// Automatic transactions. - pub two_phase_commit_auto: Option, - /// Server lifetime. - pub server_lifetime: Option, -} - -impl User { - pub fn password(&self) -> &str { - if let Some(ref s) = self.password { - s.as_str() - } else { - "" - } - } -} - -/// Admin database settings. -#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)] -#[serde(deny_unknown_fields)] -pub struct Admin { - /// Admin database name. - #[serde(default = "Admin::name")] - pub name: String, - /// Admin user name. - #[serde(default = "Admin::user")] - pub user: String, - /// Admin user's password. - #[serde(default = "Admin::password")] - pub password: String, -} - -impl Default for Admin { - fn default() -> Self { - Self { - name: Self::name(), - user: Self::user(), - password: admin_password(), - } - } -} - -impl Admin { - fn name() -> String { - "admin".into() - } - - fn user() -> String { - "admin".into() - } - - fn password() -> String { - admin_password() - } - - /// The password has been randomly generated. - pub fn random(&self) -> bool { - let prefix = "_pgdog_"; - self.password.starts_with(prefix) && self.password.len() == prefix.len() + 12 - } -} - -fn admin_password() -> String { - if let Ok(password) = env::var("PGDOG_ADMIN_PASSWORD") { - password - } else { - let pw = random_string(12); - format!("_pgdog_{}", pw) - } -} +pub use pgdog_config::users::*; diff --git a/pgdog/src/frontend/buffered_query.rs b/pgdog/src/frontend/buffered_query.rs index b76a6c5c3..0ef537490 100644 --- a/pgdog/src/frontend/buffered_query.rs +++ b/pgdog/src/frontend/buffered_query.rs @@ -20,6 +20,14 @@ impl BufferedQuery { matches!(self, Self::Prepared(_)) } + pub fn prepared(&self) -> bool { + if let Self::Prepared(ref parse) = self { + !parse.anonymous() + } else { + false + } + } + pub fn simple(&self) -> bool { matches!(self, Self::Query(_)) } diff --git a/pgdog/src/frontend/client/mod.rs b/pgdog/src/frontend/client/mod.rs index 928abbbbf..aafe4ca8e 100644 --- a/pgdog/src/frontend/client/mod.rs +++ b/pgdog/src/frontend/client/mod.rs @@ -1,38 +1,42 @@ //! Frontend client. use std::net::SocketAddr; -use std::time::Instant; +use std::sync::Arc; +use std::time::{Duration, Instant}; -use bytes::BytesMut; use timeouts::Timeouts; -use tokio::time::timeout; -use tokio::{select, spawn}; +use tokio::{select, spawn, time::timeout}; use tracing::{debug, enabled, error, info, trace, Level as LogLevel}; -use super::{ClientRequest, Comms, Error, PreparedStatements}; +use super::{ClientRequest, Error, PreparedStatements}; use crate::auth::{md5, scram::Server}; use crate::backend::maintenance_mode; +use crate::backend::pool::stats::MemoryStats; use crate::backend::{ databases, pool::{Connection, Request}, }; -use crate::config::{self, config, AuthType}; +use crate::config::convert::user_from_params; +use crate::config::{self, config, AuthType, ConfigAndUsers}; use crate::frontend::client::query_engine::{QueryEngine, QueryEngineContext}; +use crate::frontend::ClientComms; use crate::net::messages::{ Authentication, BackendKeyData, ErrorResponse, FromBytes, Message, Password, Protocol, ReadyForQuery, ToBytes, }; -use crate::net::ProtocolMessage; -use crate::net::{parameter::Parameters, Stream}; +use crate::net::{parameter::Parameters, MessageBuffer, ProtocolMessage, Stream}; use crate::state::State; use crate::stats::memory::MemoryUsage; use crate::util::user_database_from_params; -// pub mod counter; pub mod query_engine; +pub mod sticky; pub mod timeouts; +pub(crate) use sticky::Sticky; + /// Frontend client. +#[derive(Debug)] pub struct Client { addr: SocketAddr, stream: Stream, @@ -40,7 +44,7 @@ pub struct Client { #[allow(dead_code)] connect_params: Parameters, params: Parameters, - comms: Comms, + comms: ClientComms, admin: bool, streaming: bool, shutdown: bool, @@ -48,8 +52,9 @@ pub struct Client { transaction: Option, timeouts: Timeouts, client_request: ClientRequest, - stream_buffer: BytesMut, + stream_buffer: MessageBuffer, passthrough_password: Option, + sticky: Sticky, } #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] @@ -57,7 +62,8 @@ pub enum TransactionType { ReadOnly, #[default] ReadWrite, - Error, + ErrorReadWrite, + ErrorReadOnly, } impl TransactionType { @@ -70,7 +76,7 @@ impl TransactionType { } pub fn error(&self) -> bool { - matches!(self, Self::Error) + matches!(self, Self::ErrorReadWrite | Self::ErrorReadOnly) } } @@ -82,11 +88,11 @@ impl MemoryUsage for Client { + std::mem::size_of::() + self.connect_params.memory_usage() + self.params.memory_usage() - + std::mem::size_of::() + + std::mem::size_of::() + std::mem::size_of::() * 5 + self.prepared_statements.memory_used() + std::mem::size_of::() - + self.stream_buffer.memory_usage() + + self.stream_buffer.capacity() + self.client_request.memory_usage() + self .passthrough_password @@ -99,19 +105,55 @@ impl MemoryUsage for Client { impl Client { /// Create new frontend client from the given TCP stream. pub async fn spawn( - mut stream: Stream, + stream: Stream, params: Parameters, addr: SocketAddr, - mut comms: Comms, + config: Arc, ) -> Result<(), Error> { - let (user, database) = user_database_from_params(¶ms); - let config = config::config(); + let login_timeout = Duration::from_millis(config.config.general.client_login_timeout); + + match timeout(login_timeout, Self::login(stream, params, addr, config)).await { + Ok(Ok(Some(mut client))) => { + if client.admin { + // Admin clients are not waited on during shutdown. + spawn(async move { + client.spawn_internal().await; + }); + } else { + client.spawn_internal().await; + } + + Ok(()) + } + Err(_) => { + error!("client login timeout [{}]", addr); + Ok(()) + } + Ok(Ok(None)) => Ok(()), + Ok(Err(err)) => Err(err), + } + } + + /// Create new frontend client from the given TCP stream. + async fn login( + mut stream: Stream, + params: Parameters, + addr: SocketAddr, + config: Arc, + ) -> Result, Error> { + // Bail immediately if TLS is required but the connection isn't using it. + if config.config.general.tls_client_required && !stream.is_tls() { + stream.fatal(ErrorResponse::tls_required()).await?; + return Ok(None); + } + let (user, database) = user_database_from_params(¶ms); let admin = database == config.config.admin.name && config.config.admin.user == user; let admin_password = &config.config.admin.password; let auth_type = &config.config.general.auth_type; - let id = BackendKeyData::new(); + let id = BackendKeyData::new_client(); + let comms = ClientComms::new(&id); // Auto database. let exists = databases::databases().exists((user, database)); @@ -131,7 +173,7 @@ impl Client { }; if !exists { - let user = config::User::from_params(¶ms, &password).ok(); + let user = user_from_params(¶ms, &password).ok(); if let Some(user) = user { databases::add(user); } @@ -146,7 +188,7 @@ impl Client { Ok(conn) => conn, Err(_) => { stream.fatal(ErrorResponse::auth(user, database)).await?; - return Ok(()); + return Ok(None); } }; @@ -156,36 +198,55 @@ impl Client { conn.cluster()?.password() }; + let mut auth_ok = false; + + if let Some(ref passthrough_password) = passthrough_password { + if passthrough_password != password && auth_type != &AuthType::Trust { + stream.fatal(ErrorResponse::auth(user, database)).await?; + return Ok(None); + } else { + auth_ok = true; + } + } + let auth_type = &config.config.general.auth_type; - let auth_ok = match (auth_type, stream.is_tls()) { - // TODO: SCRAM doesn't work with TLS currently because of - // lack of support for channel binding in our scram library. - // Defaulting to MD5. - (AuthType::Scram, true) | (AuthType::Md5, _) => { - let md5 = md5::Client::new(user, password); - stream.send_flush(&md5.challenge()).await?; - let password = Password::from_bytes(stream.read().await?.to_bytes()?)?; - if let Password::PasswordMessage { response } = password { - md5.check(&response) - } else { - false + if !auth_ok { + auth_ok = match auth_type { + AuthType::Md5 => { + let md5 = md5::Client::new(user, password); + stream.send_flush(&md5.challenge()).await?; + let password = Password::from_bytes(stream.read().await?.to_bytes()?)?; + if let Password::PasswordMessage { response } = password { + md5.check(&response) + } else { + false + } } - } - (AuthType::Scram, false) => { - stream.send_flush(&Authentication::scram()).await?; + AuthType::Scram => { + stream.send_flush(&Authentication::scram()).await?; - let scram = Server::new(password); - let res = scram.handle(&mut stream).await; - matches!(res, Ok(true)) - } + let scram = Server::new(password); + let res = scram.handle(&mut stream).await; + matches!(res, Ok(true)) + } + + AuthType::Plain => { + stream + .send_flush(&Authentication::ClearTextPassword) + .await?; + let response = stream.read().await?; + let response = Password::from_bytes(response.to_bytes()?)?; + response.password() == Some(password) + } - (AuthType::Trust, _) => true, + AuthType::Trust => true, + } }; if !auth_ok { stream.fatal(ErrorResponse::auth(user, database)).await?; - return Ok(()); + return Ok(None); } else { stream.send(&Authentication::Ok).await?; } @@ -193,16 +254,21 @@ impl Client { // Check if the pooler is shutting down. if comms.offline() && !admin { stream.fatal(ErrorResponse::shutting_down()).await?; - return Ok(()); + return Ok(None); } let server_params = match conn.parameters(&Request::new(id)).await { Ok(params) => params, Err(err) => { if err.no_server() { - error!("connection pool is down"); - stream.fatal(ErrorResponse::connection()).await?; - return Ok(()); + error!( + "aborting new client connection, connection pool is down [{}]", + addr + ); + stream + .fatal(ErrorResponse::connection(user, database)) + .await?; + return Ok(None); } else { return Err(err.into()); } @@ -215,16 +281,29 @@ impl Client { stream.send(&id).await?; stream.send_flush(&ReadyForQuery::idle()).await?; - comms.connect(&id, addr, ¶ms); + comms.connect(addr, ¶ms); if config.config.general.log_connections { info!( - r#"client "{}" connected to database "{}" [{}]"#, - user, database, addr + r#"client "{}" connected to database "{}" [{}, auth: {}] {}"#, + user, + database, + addr, + if passthrough_password.is_some() { + "passthrough".into() + } else { + auth_type.to_string() + }, + if stream.is_tls() { "πŸ”“" } else { "" } ); } - let mut client = Self { + debug!( + "client \"{}\" startup parameters: {} [{}]", + user, params, addr + ); + + Ok(Some(Self { addr, stream, id, @@ -232,55 +311,48 @@ impl Client { admin, streaming: false, params: params.clone(), - connect_params: params, prepared_statements: PreparedStatements::new(), transaction: None, timeouts: Timeouts::from_config(&config.config.general), client_request: ClientRequest::new(), - stream_buffer: BytesMut::new(), + stream_buffer: MessageBuffer::new(config.config.memory.message_buffer), shutdown: false, - passthrough_password, - }; - - drop(conn); - - if client.admin { - // Admin clients are not waited on during shutdown. - spawn(async move { - client.spawn_internal().await; - }); - } else { - client.spawn_internal().await; - } - - Ok(()) + sticky: Sticky::from_params(¶ms), + connect_params: params, + })) } #[cfg(test)] - pub fn new_test(stream: Stream, addr: SocketAddr) -> Self { - use crate::{config::config, frontend::comms::comms}; + pub fn new_test(stream: Stream, params: Parameters) -> Self { + use crate::config::config; let mut connect_params = Parameters::default(); connect_params.insert("user", "pgdog"); connect_params.insert("database", "pgdog"); + connect_params.merge(params); + + let id = BackendKeyData::new(); + let mut prepared_statements = PreparedStatements::new(); + prepared_statements.level = config().config.general.prepared_statements; Self { stream, - addr, - id: BackendKeyData::new(), - comms: comms(), + addr: SocketAddr::from(([127, 0, 0, 1], 1234)), + id, + comms: ClientComms::new(&id), streaming: false, - prepared_statements: PreparedStatements::new(), + prepared_statements, connect_params: connect_params.clone(), - params: connect_params, admin: false, transaction: None, timeouts: Timeouts::from_config(&config().config.general), client_request: ClientRequest::new(), - stream_buffer: BytesMut::new(), + stream_buffer: MessageBuffer::new(4096), shutdown: false, passthrough_password: None, + sticky: Sticky::from_params(&connect_params), + params: connect_params, } } @@ -381,7 +453,9 @@ impl Client { message: Message, ) -> Result<(), Error> { let mut context = QueryEngineContext::new(self); - query_engine.server_message(&mut context, message).await?; + query_engine + .process_server_message(&mut context, message) + .await?; self.transaction = context.transaction(); Ok(()) @@ -416,6 +490,9 @@ impl Client { } } + // Check buffer size once per request. + self.stream_buffer.shrink_to_fit(); + Ok(()) } @@ -424,7 +501,7 @@ impl Client { /// This ensures we don't check out a connection from the pool until the client /// sent a complete request. async fn buffer(&mut self, state: State) -> Result { - self.client_request.messages.clear(); + self.client_request.clear(); // Only start timer once we receive the first message. let mut timer = None; @@ -432,20 +509,19 @@ impl Client { // Check config once per request. let config = config::config(); // Configure prepared statements cache. - self.prepared_statements.enabled = config.prepared_statements(); - self.prepared_statements.capacity = config.config.general.prepared_statements_limit; + self.prepared_statements.level = config.prepared_statements(); self.timeouts = Timeouts::from_config(&config.config.general); - while !self.client_request.full() { + while !self.client_request.is_complete() { let idle_timeout = self .timeouts .client_idle_timeout(&state, &self.client_request); let message = - match timeout(idle_timeout, self.stream.read_buf(&mut self.stream_buffer)).await { + match timeout(idle_timeout, self.stream_buffer.read(&mut self.stream)).await { Err(_) => { self.stream - .fatal(ErrorResponse::client_idle_timeout(idle_timeout)) + .fatal(ErrorResponse::client_idle_timeout(idle_timeout, &state)) .await?; return Ok(BufferEvent::DisconnectAbrupt); } @@ -492,6 +568,15 @@ impl Client { pub fn in_transaction(&self) -> bool { self.transaction.is_some() } + + /// Get client memory stats. + pub fn memory_stats(&self) -> MemoryStats { + MemoryStats { + buffer: *self.stream_buffer.stats(), + prepared_statements: self.prepared_statements.memory_used(), + stream: self.stream.memory_usage(), + } + } } impl Drop for Client { diff --git a/pgdog/src/frontend/client/query_engine/connect.rs b/pgdog/src/frontend/client/query_engine/connect.rs index 38b81e6a6..a122e3d11 100644 --- a/pgdog/src/frontend/client/query_engine/connect.rs +++ b/pgdog/src/frontend/client/query_engine/connect.rs @@ -1,5 +1,7 @@ use tokio::time::timeout; +use crate::frontend::router::parser::ShardWithPriority; + use super::*; use tracing::{error, trace}; @@ -8,46 +10,60 @@ impl QueryEngine { /// Connect to backend, if necessary. /// /// Return true if connected, false otherwise. + /// + /// # Arguments + /// + /// - context: Query engine context. + /// - connect_route: Override which route to use for connecting to backend(s). + /// Used to connect to all shards for an explicit cross-shard transaction + /// started with `BEGIN`. + /// pub(super) async fn connect( &mut self, context: &mut QueryEngineContext<'_>, - route: &Route, + connect_route: Option<&Route>, ) -> Result { if self.backend.connected() { + self.debug_connected(context, true); return Ok(true); } - let request = Request::new(self.client_id); + let request = Request::new(*context.id); self.stats.waiting(request.created_at); - self.comms.stats(self.stats); + self.comms.update_stats(self.stats); + + let connect_route = connect_route.unwrap_or(context.client_request.route()); - let connected = match self.backend.connect(&request, route).await { + let connected = match self.backend.connect(&request, connect_route).await { Ok(_) => { self.stats.connected(); - self.stats.locked(route.lock_session()); + self.stats + .locked(context.client_request.route().is_lock_session()); // This connection will be locked to this client // until they disconnect. // // Used in case the client runs an advisory lock // or another leaky transaction mode abstraction. - self.backend.lock(route.lock_session()); - - if let Ok(addr) = self.backend.addr() { - debug!( - "client paired with [{}] using route [{}] [{:.4}ms]", - addr.into_iter() - .map(|a| a.to_string()) - .collect::>() - .join(","), - route, - self.stats.wait_time.as_secs_f64() * 1000.0 - ); - } + self.backend + .lock(context.client_request.route().is_lock_session()); + + self.debug_connected(context, false); let query_timeout = context.timeouts.query_timeout(&self.stats.state); + + let begin_stmt = self.begin_stmt.take(); + // We may need to sync params with the server and that reads from the socket. - timeout(query_timeout, self.backend.link_client(context.params)).await??; + timeout( + query_timeout, + self.backend.link_client( + context.id, + context.params, + begin_stmt.as_ref().map(|stmt| stmt.query()), + ), + ) + .await??; true } @@ -57,10 +73,16 @@ impl QueryEngine { if err.no_server() { error!("{} [{:?}]", err, context.stream.peer_addr()); + + let error = ErrorResponse::from_err(&err); + + self.hooks.on_engine_error(context, &error)?; + let bytes_sent = context .stream - .error(ErrorResponse::from_err(&err), context.in_transaction()) + .error(error, context.in_transaction()) .await?; + self.stats.sent(bytes_sent); self.backend.disconnect(); self.router.reset(); @@ -72,7 +94,7 @@ impl QueryEngine { } }; - self.comms.stats(self.stats); + self.comms.update_stats(self.stats); Ok(connected) } @@ -81,24 +103,53 @@ impl QueryEngine { pub(super) async fn connect_transaction( &mut self, context: &mut QueryEngineContext<'_>, - route: &Route, ) -> Result { debug!("connecting to backend(s) to serve transaction"); - let route = self.transaction_route(route)?; + let route = self.transaction_route(context.client_request.route())?; trace!("transaction routing to {:#?}", route); - self.connect(context, &route).await + self.connect(context, Some(&route)).await } pub(super) fn transaction_route(&mut self, route: &Route) -> Result { let cluster = self.backend.cluster()?; if cluster.shards().len() == 1 { - Ok(Route::write(Shard::Direct(0)).set_read(route.is_read())) + Ok( + Route::write(ShardWithPriority::new_override_transaction(Shard::Direct( + 0, + ))) + .with_read(route.is_read()), + ) + } else if route.is_search_path_driven() { + // Schema-based routing will only go to one shard. + Ok(route.clone()) } else { - Ok(Route::write(Shard::All).set_read(route.is_read())) + Ok( + Route::write(ShardWithPriority::new_override_transaction(Shard::All)) + .with_read(route.is_read()), + ) + } + } + + fn debug_connected(&self, context: &QueryEngineContext<'_>, connected: bool) { + if let Ok(addr) = self.backend.addr() { + debug!( + "{} [{}] using route [{}] [{:.4}ms]", + if connected { + "already connected to" + } else { + "client paired with" + }, + addr.into_iter() + .map(|a| a.to_string()) + .collect::>() + .join(","), + context.client_request.route(), + self.stats.wait_time.as_secs_f64() * 1000.0 + ); } } } diff --git a/pgdog/src/frontend/client/query_engine/context.rs b/pgdog/src/frontend/client/query_engine/context.rs index 13e413a7a..b54751a35 100644 --- a/pgdog/src/frontend/client/query_engine/context.rs +++ b/pgdog/src/frontend/client/query_engine/context.rs @@ -1,15 +1,18 @@ use crate::{ - backend::pool::connection::mirror::Mirror, + backend::pool::{connection::mirror::Mirror, stats::MemoryStats}, frontend::{ - client::{timeouts::Timeouts, TransactionType}, + client::{timeouts::Timeouts, Sticky, TransactionType}, + router::parser::rewrite::statement::plan::RewriteResult, Client, ClientRequest, PreparedStatements, }, - net::{Parameters, Stream}, - stats::memory::MemoryUsage, + net::{BackendKeyData, Parameters, Stream}, }; +#[allow(dead_code)] /// Context passed to the query engine to execute a query. pub struct QueryEngineContext<'a> { + /// Client ID running the query. + pub(super) id: &'a BackendKeyData, /// Prepared statements cache. pub(super) prepared_statements: &'a mut PreparedStatements, /// Client session parameters. @@ -27,18 +30,23 @@ pub struct QueryEngineContext<'a> { /// Cross shard queries are disabled. pub(super) cross_shard_disabled: Option, /// Client memory usage. - pub(super) memory_usage: usize, + pub(super) memory_stats: MemoryStats, /// Is the client an admin. pub(super) admin: bool, /// Executing rollback statement. pub(super) rollback: bool, + /// Sticky config: + pub(super) sticky: Sticky, + /// Rewrite result. + pub(super) rewrite_result: Option, } impl<'a> QueryEngineContext<'a> { pub fn new(client: &'a mut Client) -> Self { - let memory_usage = client.memory_usage(); + let memory_stats = client.memory_stats(); Self { + id: &client.id, prepared_statements: &mut client.prepared_statements, params: &mut client.params, client_request: &mut client.client_request, @@ -46,10 +54,12 @@ impl<'a> QueryEngineContext<'a> { transaction: client.transaction, timeouts: client.timeouts, cross_shard_disabled: None, - memory_usage, + memory_stats, admin: client.admin, requests_left: 0, rollback: false, + sticky: client.sticky, + rewrite_result: None, } } @@ -62,6 +72,7 @@ impl<'a> QueryEngineContext<'a> { /// Create context from mirror. pub fn new_mirror(mirror: &'a mut Mirror, buffer: &'a mut ClientRequest) -> Self { Self { + id: &mirror.id, prepared_statements: &mut mirror.prepared_statements, params: &mut mirror.params, client_request: buffer, @@ -69,10 +80,12 @@ impl<'a> QueryEngineContext<'a> { transaction: mirror.transaction, timeouts: mirror.timeouts, cross_shard_disabled: None, - memory_usage: 0, + memory_stats: MemoryStats::default(), admin: false, requests_left: 0, rollback: false, + sticky: Sticky::new(), + rewrite_result: None, } } diff --git a/pgdog/src/frontend/client/query_engine/end_transaction.rs b/pgdog/src/frontend/client/query_engine/end_transaction.rs index ce447fc17..207e2e6f2 100644 --- a/pgdog/src/frontend/client/query_engine/end_transaction.rs +++ b/pgdog/src/frontend/client/query_engine/end_transaction.rs @@ -43,10 +43,10 @@ impl QueryEngine { pub(super) async fn end_connected( &mut self, context: &mut QueryEngineContext<'_>, - route: &Route, rollback: bool, extended: bool, ) -> Result<(), Error> { + self.backend.transaction_params_hook(rollback); let cluster = self.backend.cluster()?; // If we experienced an error and client @@ -70,7 +70,7 @@ impl QueryEngine { // 2pc is used only for writes and is not needed for rollbacks. let two_pc = cluster.two_pc_enabled() - && route.is_write() + && context.client_request.route().is_write() && !rollback && context.transaction().map(|t| t.write()).unwrap_or(false); @@ -91,7 +91,7 @@ impl QueryEngine { self.notify_buffer.clear(); } context.rollback = rollback; - self.execute(context, route).await?; + self.execute(context).await?; } Ok(()) @@ -130,20 +130,21 @@ impl QueryEngine { #[cfg(test)] mod tests { use super::*; + use crate::config::load_test; use crate::frontend::client::TransactionType; use crate::net::Stream; #[tokio::test] async fn test_transaction_state_not_cleared() { + load_test(); + // Create a test client with DevNull stream (doesn't require real I/O) - let mut client = crate::frontend::Client::new_test( - Stream::DevNull, - std::net::SocketAddr::from(([127, 0, 0, 1], 1234)), - ); + let mut client = + crate::frontend::Client::new_test(Stream::dev_null(), Parameters::default()); client.transaction = Some(TransactionType::ReadWrite); // Create a default query engine (avoids backend connection) - let mut engine = QueryEngine::default(); + let mut engine = QueryEngine::from_client(&client).unwrap(); // state copied from client let mut context = QueryEngineContext::new(&mut client); let result = engine.end_not_connected(&mut context, false, false).await; diff --git a/pgdog/src/frontend/client/query_engine/fake.rs b/pgdog/src/frontend/client/query_engine/fake.rs new file mode 100644 index 000000000..94d33d525 --- /dev/null +++ b/pgdog/src/frontend/client/query_engine/fake.rs @@ -0,0 +1,59 @@ +use tokio::io::AsyncWriteExt; + +use crate::net::{ + BindComplete, CommandComplete, NoData, ParameterDescription, ParseComplete, ProtocolMessage, + ReadyForQuery, RowDescription, +}; + +use super::*; + +impl QueryEngine { + /// Respond to a command sent by the client + /// in a way that won't make it suspicious. + pub(crate) async fn fake_command_response( + &mut self, + context: &mut QueryEngineContext<'_>, + command: &str, + ) -> Result<(), Error> { + let mut sent = 0; + for message in context.client_request.iter() { + sent += match message { + ProtocolMessage::Parse(_) => context.stream.send(&ParseComplete).await?, + ProtocolMessage::Bind(_) => context.stream.send(&BindComplete).await?, + ProtocolMessage::Describe(describe) => { + if describe.is_statement() { + context + .stream + .send(&ParameterDescription::default()) + .await? + + context.stream.send(&RowDescription::default()).await? + } else { + context.stream.send(&NoData).await? + } + } + ProtocolMessage::Execute(_) => { + context.stream.send(&CommandComplete::new(command)).await? + } + ProtocolMessage::Sync(_) => { + context + .stream + .send(&ReadyForQuery::in_transaction(context.in_transaction())) + .await? + } + ProtocolMessage::Query(_) => { + context.stream.send(&CommandComplete::new(command)).await? + + context + .stream + .send(&ReadyForQuery::in_transaction(context.in_transaction())) + .await? + } + + _ => 0, + } + } + context.stream.flush().await?; + self.stats.sent(sent); + + Ok(()) + } +} diff --git a/pgdog/src/frontend/client/query_engine/hooks/mod.rs b/pgdog/src/frontend/client/query_engine/hooks/mod.rs new file mode 100644 index 000000000..cbcbbc62f --- /dev/null +++ b/pgdog/src/frontend/client/query_engine/hooks/mod.rs @@ -0,0 +1,56 @@ +//! Query hooks. +#![allow(unused_variables, dead_code)] +use super::*; + +#[derive(Debug)] +pub struct QueryEngineHooks; + +impl Default for QueryEngineHooks { + fn default() -> Self { + Self::new() + } +} + +impl QueryEngineHooks { + pub(super) fn new() -> Self { + Self {} + } + + pub(super) fn before_execution( + &mut self, + context: &mut QueryEngineContext<'_>, + ) -> Result<(), Error> { + Ok(()) + } + + pub(super) fn after_connected( + &mut self, + context: &mut QueryEngineContext<'_>, + backend: &Connection, + ) -> Result<(), Error> { + Ok(()) + } + + pub(super) fn after_execution( + &mut self, + context: &mut QueryEngineContext<'_>, + ) -> Result<(), Error> { + Ok(()) + } + + pub(super) fn on_server_message( + &mut self, + context: &mut QueryEngineContext<'_>, + message: &Message, + ) -> Result<(), Error> { + Ok(()) + } + + pub(super) fn on_engine_error( + &mut self, + context: &mut QueryEngineContext<'_>, + error: &ErrorResponse, + ) -> Result<(), Error> { + Ok(()) + } +} diff --git a/pgdog/src/frontend/client/query_engine/internal_values.rs b/pgdog/src/frontend/client/query_engine/internal_values.rs new file mode 100644 index 000000000..4721d30b4 --- /dev/null +++ b/pgdog/src/frontend/client/query_engine/internal_values.rs @@ -0,0 +1,50 @@ +use crate::{ + net::{CommandComplete, DataRow, Field, Protocol, ReadyForQuery, RowDescription}, + unique_id, +}; + +use super::*; + +impl QueryEngine { + /// SHOW pgdog.shards. + pub(super) async fn show_internal_value( + &mut self, + context: &mut QueryEngineContext<'_>, + field: String, + value: String, + ) -> Result<(), Error> { + let bytes_sent = context + .stream + .send_many(&[ + RowDescription::new(&[Field::text(&field)]).message()?, + DataRow::from_columns(vec![value]).message()?, + CommandComplete::from_str("SHOW").message()?, + ReadyForQuery::in_transaction(context.in_transaction()).message()?, + ]) + .await?; + + self.stats.sent(bytes_sent); + + Ok(()) + } + + pub(super) async fn unique_id( + &mut self, + context: &mut QueryEngineContext<'_>, + ) -> Result<(), Error> { + let id = unique_id::UniqueId::generator()?.next_id(); + let bytes_sent = context + .stream + .send_many(&[ + RowDescription::new(&[Field::bigint("unique_id")]).message()?, + DataRow::from_columns(vec![id.to_string()]).message()?, + CommandComplete::from_str("SHOW").message()?, + ReadyForQuery::in_transaction(context.in_transaction()).message()?, + ]) + .await?; + + self.stats.sent(bytes_sent); + + Ok(()) + } +} diff --git a/pgdog/src/frontend/client/query_engine/mod.rs b/pgdog/src/frontend/client/query_engine/mod.rs index f04779828..bb2de8353 100644 --- a/pgdog/src/frontend/client/query_engine/mod.rs +++ b/pgdog/src/frontend/client/query_engine/mod.rs @@ -1,10 +1,12 @@ use crate::{ backend::pool::{Connection, Request}, + config::config, frontend::{ + client::query_engine::{hooks::QueryEngineHooks, route_query::ClusterCheck}, router::{parser::Shard, Route}, - BufferedQuery, Client, Command, Comms, Error, Router, RouterContext, Stats, + BufferedQuery, Client, ClientComms, Command, Error, Router, RouterContext, Stats, }, - net::{BackendKeyData, ErrorResponse, Message, Parameters}, + net::{ErrorResponse, Message, Parameters}, state::State, }; @@ -15,46 +17,74 @@ pub mod context; pub mod deallocate; pub mod discard; pub mod end_transaction; +pub mod fake; +pub mod hooks; pub mod incomplete_requests; +pub mod internal_values; +pub mod multi_step; pub mod notify_buffer; -pub mod prepared_statements; pub mod pub_sub; pub mod query; +pub mod rewrite; pub mod route_query; pub mod set; -pub mod show_shards; +pub mod shard_key_rewrite; pub mod start_transaction; -pub mod two_pc; -pub mod unknown_command; - +#[cfg(test)] +mod test; #[cfg(test)] mod testing; +pub mod two_pc; +pub mod unknown_command; +use self::query::ExplainResponseState; pub use context::QueryEngineContext; use notify_buffer::NotifyBuffer; pub use two_pc::phase::TwoPcPhase; use two_pc::TwoPc; -#[derive(Default, Debug)] +#[derive(Debug)] +pub struct TestMode { + pub enabled: bool, +} + +impl Default for TestMode { + fn default() -> Self { + Self::new() + } +} + +impl TestMode { + pub fn new() -> Self { + Self { + #[cfg(test)] + enabled: true, + #[cfg(not(test))] + enabled: false, + } + } +} + +#[derive(Debug)] pub struct QueryEngine { begin_stmt: Option, router: Router, - comms: Comms, + comms: ClientComms, stats: Stats, backend: Connection, streaming: bool, - client_id: BackendKeyData, - test_mode: bool, - set_route: Option, + test_mode: TestMode, two_pc: TwoPc, notify_buffer: NotifyBuffer, + pending_explain: Option, + hooks: QueryEngineHooks, } impl QueryEngine { /// Create new query engine. pub fn new( params: &Parameters, - comms: &Comms, + comms: &ClientComms, admin: bool, passthrough_password: &Option, ) -> Result { @@ -65,13 +95,16 @@ impl QueryEngine { Ok(Self { backend, - client_id: comms.client_id(), comms: comms.clone(), - #[cfg(test)] - test_mode: true, - #[cfg(not(test))] - test_mode: false, - ..Default::default() + hooks: QueryEngineHooks::new(), + test_mode: TestMode::new(), + stats: Stats::default(), + streaming: bool::default(), + two_pc: TwoPc::default(), + notify_buffer: NotifyBuffer::default(), + pending_explain: None, + begin_stmt: None, + router: Router::default(), }) } @@ -103,10 +136,20 @@ impl QueryEngine { pub async fn handle(&mut self, context: &mut QueryEngineContext<'_>) -> Result<(), Error> { self.stats .received(context.client_request.total_message_len()); + self.set_state(State::Active); // Client is active. // Rewrite prepared statements. self.rewrite_extended(context)?; + if let ClusterCheck::Offline = self.cluster_check(context).await? { + return Ok(()); + } + + // Rewrite statement if necessary. + if !self.parse_and_rewrite(context).await? { + return Ok(()); + } + // Intercept commands we don't have to forward to a server. if self.intercept_incomplete(context).await? { self.update_stats(context); @@ -114,29 +157,40 @@ impl QueryEngine { } // Route transaction to the right servers. - if !self.route_transaction(context).await? { + if !self.route_query(context).await? { self.update_stats(context); - debug!("transaction has nowhere to go"); + debug!("query has nowhere to go"); return Ok(()); } + self.hooks.before_execution(context)?; + // Queue up request to mirrors, if any. // Do this before sending query to actual server // to have accurate timings between queries. self.backend.mirror(context.client_request); + self.pending_explain = None; + let command = self.router.command(); - let route = if let Some(ref route) = self.set_route { - route.clone() - } else { - command.route().clone() - }; - // FIXME, we should not to copy route twice. - context.client_request.route = Some(route.clone()); + if let Some(trace) = context + .client_request + .route // Admin commands don't have a route. + .as_mut() + .and_then(|route| route.take_explain()) + { + if config().config.general.expanded_explain { + self.pending_explain = Some(ExplainResponseState::new(trace)); + } + } match command { - Command::Shards(shards) => self.show_shards(context, *shards).await?, + Command::InternalField { name, value } => { + self.show_internal_value(context, name.clone(), value.clone()) + .await? + } + Command::UniqueId => self.unique_id(context).await?, Command::StartTransaction { query, transaction_type, @@ -146,34 +200,36 @@ impl QueryEngine { .await? } Command::CommitTransaction { extended } => { - self.set_route = None; - if self.backend.connected() || *extended { let extended = *extended; - let transaction_route = self.transaction_route(&route)?; + let transaction_route = + self.transaction_route(context.client_request.route())?; context.client_request.route = Some(transaction_route.clone()); context.cross_shard_disabled = Some(false); - self.end_connected(context, &transaction_route, false, extended) - .await?; + self.end_connected(context, false, extended).await?; } else { self.end_not_connected(context, false, *extended).await? } + + if context.params.commit() { + self.comms.update_params(context.params); + } } Command::RollbackTransaction { extended } => { - self.set_route = None; - if self.backend.connected() || *extended { let extended = *extended; - let transaction_route = self.transaction_route(&route)?; + let transaction_route = + self.transaction_route(context.client_request.route())?; context.client_request.route = Some(transaction_route.clone()); context.cross_shard_disabled = Some(false); - self.end_connected(context, &transaction_route, true, extended) - .await?; + self.end_connected(context, true, extended).await?; } else { self.end_not_connected(context, true, *extended).await? } + + context.params.rollback(); } - Command::Query(_) => self.execute(context, &route).await?, + Command::Query(_) => self.execute(context).await?, Command::Listen { channel, shard } => { self.listen(context, &channel.clone(), shard.clone()) .await? @@ -187,26 +243,20 @@ impl QueryEngine { .await? } Command::Unlisten(channel) => self.unlisten(context, &channel.clone()).await?, - Command::Set { name, value } => { - if self.backend.connected() { - self.execute(context, &route).await? - } else { - self.set(context, name.clone(), value.clone()).await? - } - } - Command::SetRoute(route) => { - self.set_route(context, route.clone()).await?; - } - Command::Copy(_) => self.execute(context, &route).await?, - Command::Rewrite(query) => { - context.client_request.rewrite(query)?; - self.execute(context, &route).await?; + Command::Set { + name, value, local, .. + } => { + self.set(context, name.clone(), value.clone(), *local) + .await?; } + Command::Copy(_) => self.execute(context).await?, Command::Deallocate => self.deallocate(context).await?, Command::Discard { extended } => self.discard(context, *extended).await?, command => self.unknown_command(context, command.clone()).await?, } + self.hooks.after_execution(context)?; + if context.in_error() { self.backend.mirror_clear(); self.notify_buffer.clear(); @@ -234,14 +284,14 @@ impl QueryEngine { self.stats .prepared_statements(context.prepared_statements.len_local()); - self.stats.memory_used(context.memory_usage); + self.stats.memory_used(context.memory_stats); - self.comms.stats(self.stats); + self.comms.update_stats(self.stats); } pub fn set_state(&mut self, state: State) { self.stats.state = state; - self.comms.stats(self.stats); + self.comms.update_stats(self.stats); } pub fn get_state(&self) -> State { diff --git a/pgdog/src/frontend/client/query_engine/multi_step/error.rs b/pgdog/src/frontend/client/query_engine/multi_step/error.rs new file mode 100644 index 000000000..793436ce1 --- /dev/null +++ b/pgdog/src/frontend/client/query_engine/multi_step/error.rs @@ -0,0 +1,48 @@ +use thiserror::Error; + +use crate::net::ErrorResponse; + +#[derive(Debug, Error)] +pub enum Error { + #[error("{0}")] + Update(#[from] UpdateError), + + #[error("frontend: {0}")] + Frontend(Box), + + #[error("backend: {0}")] + Backend(#[from] crate::backend::Error), + + #[error("rewrite: {0}")] + Rewrite(#[from] crate::frontend::router::parser::rewrite::statement::Error), + + #[error("router: {0}")] + Router(#[from] crate::frontend::router::Error), + + #[error("{0}")] + Execution(ErrorResponse), + + #[error("net: {0}")] + Net(#[from] crate::net::Error), +} + +#[derive(Debug, Error)] +pub enum UpdateError { + #[error("sharding key updates are forbidden")] + Disabled, + + #[error("sharding key update must be executed inside a transaction")] + TransactionRequired, + + #[error("sharding key update intermediate query has no route")] + NoRoute, + + #[error("sharding key update changes more than one row ({0})")] + TooManyRows(usize), +} + +impl From for Error { + fn from(value: crate::frontend::Error) -> Self { + Self::Frontend(Box::new(value)) + } +} diff --git a/pgdog/src/frontend/client/query_engine/multi_step/forward_check.rs b/pgdog/src/frontend/client/query_engine/multi_step/forward_check.rs new file mode 100644 index 000000000..a79dbf429 --- /dev/null +++ b/pgdog/src/frontend/client/query_engine/multi_step/forward_check.rs @@ -0,0 +1,40 @@ +use fnv::FnvHashSet as HashSet; + +use crate::{frontend::ClientRequest, net::Protocol}; + +#[derive(Debug, Clone)] +pub(crate) struct ForwardCheck { + codes: HashSet, + sent: HashSet, + describe: bool, +} + +impl ForwardCheck { + /// Create new forward checker from a client request. + /// + /// Will construct a mapping to allow only the messages the client expects through + /// + pub(crate) fn new(request: &ClientRequest) -> Self { + Self { + codes: request.iter().map(|m| m.code()).collect(), + describe: request.iter().any(|m| m.code() == 'D'), + sent: HashSet::default(), + } + } + + /// Check if we should forward a particular message to the client. + pub(crate) fn forward(&mut self, code: char) -> bool { + let forward = match code { + '1' => self.codes.contains(&'P'), // ParseComplete + '2' => self.codes.contains(&'B'), // BindComplete + 'D' | 'E' => true, // DataRow + 'T' => self.describe && !self.sent.contains(&'T') || self.codes.contains(&'Q'), + 't' => self.describe && !self.sent.contains(&'t'), + _ => false, + }; + + self.sent.insert(code); + + forward + } +} diff --git a/pgdog/src/frontend/client/query_engine/multi_step/insert.rs b/pgdog/src/frontend/client/query_engine/multi_step/insert.rs new file mode 100644 index 000000000..e82a94783 --- /dev/null +++ b/pgdog/src/frontend/client/query_engine/multi_step/insert.rs @@ -0,0 +1,89 @@ +use super::{CommandType, MultiServerState}; +use crate::{ + frontend::{ + client::query_engine::{QueryEngine, QueryEngineContext}, + ClientRequest, Command, Router, RouterContext, + }, + net::Protocol, +}; + +use super::super::Error; + +#[derive(Debug)] +pub(crate) struct InsertMulti<'a> { + /// Requests split by the rewrite engine. + requests: Vec, + /// Execution state. + state: MultiServerState, + /// Query engine. + engine: &'a mut QueryEngine, +} + +impl<'a> InsertMulti<'a> { + /// Create multi-shard INSERT handler + /// from query engine and a set of routed requests. + pub(crate) fn from_engine(engine: &'a mut QueryEngine, requests: Vec) -> Self { + Self { + state: MultiServerState::new(requests.len()), + requests, + engine, + } + } + + /// Execute the multi-shard INSERT. + pub(crate) async fn execute( + &'a mut self, + context: &mut QueryEngineContext<'_>, + ) -> Result { + let cluster = self.engine.backend.cluster()?; + for request in self.requests.iter_mut() { + let context = RouterContext::new( + request, + cluster, + context.params, + context.transaction(), + context.sticky, + )?; + let mut router = Router::new(); + let command = router.query(context)?; + if let Command::Query(route) = command { + request.route = Some(route.clone()); + } else { + return Err(Error::NoRoute); + } + } + + if !self.engine.backend.is_multishard() { + return Err(Error::MultiShardRequired); + } + + for request in self.requests.iter() { + self.engine + .backend + .handle_client_request(request, &mut self.engine.router, self.engine.streaming) + .await?; + + while self.engine.backend.has_more_messages() { + let message = self.engine.read_server_message(context).await?; + + if self.state.forward(&message)? { + self.engine.process_server_message(context, message).await?; + } + } + } + + if let Some(cc) = self.state.command_complete(CommandType::Insert) { + self.engine + .process_server_message(context, cc.message()?) + .await?; + } + + if let Some(rfq) = self.state.ready_for_query(context.in_transaction()) { + self.engine + .process_server_message(context, rfq.message()?) + .await?; + } + + Ok(self.state.error()) + } +} diff --git a/pgdog/src/frontend/client/query_engine/multi_step/mod.rs b/pgdog/src/frontend/client/query_engine/multi_step/mod.rs new file mode 100644 index 000000000..7b80478f0 --- /dev/null +++ b/pgdog/src/frontend/client/query_engine/multi_step/mod.rs @@ -0,0 +1,14 @@ +pub(crate) mod error; +pub mod forward_check; +pub mod insert; +pub mod state; +pub mod update; + +pub(crate) use error::{Error, UpdateError}; +pub(crate) use forward_check::*; +pub(crate) use insert::InsertMulti; +pub use state::{CommandType, MultiServerState}; +pub(crate) use update::UpdateMulti; + +#[cfg(test)] +mod test; diff --git a/pgdog/src/frontend/client/query_engine/multi_step/state.rs b/pgdog/src/frontend/client/query_engine/multi_step/state.rs new file mode 100644 index 000000000..cff6af113 --- /dev/null +++ b/pgdog/src/frontend/client/query_engine/multi_step/state.rs @@ -0,0 +1,87 @@ +use fnv::FnvHashMap as HashMap; + +use super::super::Error; +use crate::net::{CommandComplete, FromBytes, Message, Protocol, ReadyForQuery, ToBytes}; + +#[derive(Debug, Clone)] +pub enum CommandType { + Insert, + Update, + Delete, +} + +#[derive(Debug, Clone)] +pub struct MultiServerState { + servers: usize, + rows: usize, + counters: HashMap, +} + +impl MultiServerState { + /// New multi-server execution state. + pub fn new(servers: usize) -> Self { + Self { + servers, + rows: 0, + counters: HashMap::default(), + } + } + + /// Should the message be forwarded to the client. + pub fn forward(&mut self, message: &Message) -> Result { + let code = message.code(); + let count = self.counters.entry(code).or_default(); + *count += 1; + + Ok(match code { + 'T' | '1' | '2' | '3' | 't' => *count == 1, + 'C' => { + let command_complete = CommandComplete::from_bytes(message.to_bytes()?)?; + self.rows += command_complete.rows()?.unwrap_or(0); + false + } + 'Z' => false, + 'n' => *count == self.servers && !self.counters.contains_key(&'D'), + 'I' => *count == self.servers && !self.counters.contains_key(&'C'), + _ => true, + }) + } + + /// Number of rows returned. + pub fn rows(&self) -> usize { + self.rows + } + + /// Error happened. + pub fn error(&self) -> bool { + self.counters.contains_key(&'E') + } + + /// Create CommandComplete (C) message. + pub fn command_complete(&self, command_type: CommandType) -> Option { + if !self.counters.contains_key(&'C') || self.error() { + return None; + } + + let name = match command_type { + CommandType::Delete => "DELETE", + CommandType::Update => "UPDATE", + CommandType::Insert => "INSERT 0", + }; + + Some(CommandComplete::new(format!("{} {}", name, self.rows()))) + } + + /// Create ReadyForQuery (C) message. + pub fn ready_for_query(&self, in_transaction: bool) -> Option { + if !self.counters.contains_key(&'Z') { + return None; + } + + if self.error() { + Some(ReadyForQuery::error()) + } else { + Some(ReadyForQuery::in_transaction(in_transaction)) + } + } +} diff --git a/pgdog/src/frontend/client/query_engine/multi_step/test/mod.rs b/pgdog/src/frontend/client/query_engine/multi_step/test/mod.rs new file mode 100644 index 000000000..ffd1c8085 --- /dev/null +++ b/pgdog/src/frontend/client/query_engine/multi_step/test/mod.rs @@ -0,0 +1,20 @@ +use tokio::{io::AsyncWriteExt, net::TcpStream}; + +use crate::{ + frontend::client::test::read_messages, + net::{Query, ToBytes}, +}; + +pub mod prepared; +pub mod simple; +pub mod update; + +async fn truncate_table(table: &str, stream: &mut TcpStream) { + let query = Query::new(format!("TRUNCATE {}", table)) + .to_bytes() + .unwrap(); + stream.write_all(&query).await.unwrap(); + stream.flush().await.unwrap(); + + read_messages(stream, &['C', 'Z']).await; +} diff --git a/pgdog/src/frontend/client/query_engine/multi_step/test/prepared.rs b/pgdog/src/frontend/client/query_engine/multi_step/test/prepared.rs new file mode 100644 index 000000000..947850250 --- /dev/null +++ b/pgdog/src/frontend/client/query_engine/multi_step/test/prepared.rs @@ -0,0 +1,83 @@ +mod insert { + use tokio::{io::AsyncWriteExt, spawn}; + + use crate::{ + frontend::client::{ + query_engine::multi_step::test::truncate_table, + test::{read_messages, test_client_sharded}, + }, + net::{ + bind::Parameter, Bind, CommandComplete, DataRow, Describe, Execute, Flush, Format, + FromBytes, Parse, Sync, ToBytes, + }, + }; + + #[tokio::test] + async fn test_prepared() { + crate::logger(); + + let (mut stream, mut client) = test_client_sharded().await; + spawn(async move { client.run().await.unwrap() }); + + truncate_table("sharded", &mut stream).await; + + let stmt = Parse::named( + "test_multi", + "INSERT INTO sharded (id, value) VALUES ($1, $2), ($3, $4), ($5, $6), ($7, 'test_value_4') RETURNING *", + ); + let desc = Describe::new_statement("test_multi"); + let flush = Flush; + + let params = Bind::new_params( + "test_multi", + &[ + Parameter::new("123423425245".as_bytes()), + Parameter::new("test_value_1".as_bytes()), + Parameter::new("123423425246".as_bytes()), + Parameter::new("test_value_2".as_bytes()), + Parameter::new("123423425247".as_bytes()), + Parameter::new("test_value_3".as_bytes()), + Parameter::new("12342342524823424".as_bytes()), + ], + ); + let exec = Execute::new(); + let sync = Sync; + + stream.write_all(&stmt.to_bytes().unwrap()).await.unwrap(); + stream.write_all(&desc.to_bytes().unwrap()).await.unwrap(); + stream.write_all(&flush.to_bytes().unwrap()).await.unwrap(); + stream.flush().await.unwrap(); + + let _ = read_messages(&mut stream, &['1', 't', 'T']).await; + + stream.write_all(¶ms.to_bytes().unwrap()).await.unwrap(); + stream.write_all(&exec.to_bytes().unwrap()).await.unwrap(); + stream.write_all(&sync.to_bytes().unwrap()).await.unwrap(); + stream.flush().await.unwrap(); + + let messages = read_messages(&mut stream, &['2', 'D', 'D', 'D', 'D', 'C', 'Z']).await; + + // Assert DataRow values (messages[1..5] are the 4 DataRow messages) + let row1 = DataRow::from_bytes(messages[1].to_bytes().unwrap()).unwrap(); + assert_eq!(row1.get::(0, Format::Text), Some(123423425245)); + assert_eq!(row1.get_text(1), Some("test_value_1".to_string())); + + let row2 = DataRow::from_bytes(messages[2].to_bytes().unwrap()).unwrap(); + assert_eq!(row2.get::(0, Format::Text), Some(123423425246)); + assert_eq!(row2.get_text(1), Some("test_value_2".to_string())); + + let row3 = DataRow::from_bytes(messages[3].to_bytes().unwrap()).unwrap(); + assert_eq!(row3.get::(0, Format::Text), Some(123423425247)); + assert_eq!(row3.get_text(1), Some("test_value_3".to_string())); + + let row4 = DataRow::from_bytes(messages[4].to_bytes().unwrap()).unwrap(); + assert_eq!(row4.get::(0, Format::Text), Some(12342342524823424)); + assert_eq!(row4.get_text(1), Some("test_value_4".to_string())); + + // Assert CommandComplete returns 4 rows + let cc = CommandComplete::from_bytes(messages[5].to_bytes().unwrap()).unwrap(); + assert_eq!(cc.rows().unwrap(), Some(4)); + + truncate_table("sharded", &mut stream).await; + } +} diff --git a/pgdog/src/frontend/client/query_engine/multi_step/test/simple.rs b/pgdog/src/frontend/client/query_engine/multi_step/test/simple.rs new file mode 100644 index 000000000..b3d291f92 --- /dev/null +++ b/pgdog/src/frontend/client/query_engine/multi_step/test/simple.rs @@ -0,0 +1,52 @@ +mod insert { + use rand::{rng, Rng}; + use tokio::{io::AsyncWriteExt, spawn}; + + use crate::{ + frontend::client::{ + query_engine::multi_step::test::truncate_table, + test::{read_messages, test_client_sharded}, + }, + net::{CommandComplete, FromBytes, Query, ToBytes}, + }; + + #[tokio::test] + async fn test_simple() { + crate::logger(); + + let (mut stream, mut client) = test_client_sharded().await; + + spawn(async move { + client.run().await.unwrap(); + }); + + let values = (0..5) + .into_iter() + .map(|_| { + let val = rng().random::(); + (format!("'val_{}'", val), val) + }) + .map(|tuple| format!("({}, {})", tuple.1, tuple.0)) + .collect::>() + .join(", "); + + stream + .write_all( + &Query::new(format!("INSERT INTO sharded (id, value) VALUES {}", values)) + .to_bytes() + .unwrap(), + ) + .await + .unwrap(); + stream.flush().await.unwrap(); + + let messages = read_messages(&mut stream, &['C', 'Z']).await; + assert_eq!(messages.len(), 2); + + let command_complete = + CommandComplete::from_bytes(messages[0].to_bytes().unwrap()).unwrap(); + assert_eq!(command_complete.rows().unwrap().unwrap(), 5); + + truncate_table("sharded", &mut stream).await; + } +} diff --git a/pgdog/src/frontend/client/query_engine/multi_step/test/update.rs b/pgdog/src/frontend/client/query_engine/multi_step/test/update.rs new file mode 100644 index 000000000..89765a662 --- /dev/null +++ b/pgdog/src/frontend/client/query_engine/multi_step/test/update.rs @@ -0,0 +1,571 @@ +use rand::{rng, Rng}; + +use crate::{ + expect_message, + frontend::{ + client::{ + query_engine::{multi_step::UpdateMulti, QueryEngineContext}, + test::TestClient, + }, + ClientRequest, + }, + net::{ + bind::Parameter, Bind, CommandComplete, DataRow, Describe, ErrorResponse, Execute, Flush, + Format, Parameters, Parse, Protocol, Query, ReadyForQuery, RowDescription, Sync, + TransactionState, + }, +}; + +use super::super::super::Error; + +async fn same_shard_check(request: ClientRequest) -> Result<(), Error> { + let mut client = TestClient::new_rewrites(Parameters::default()).await; + client.client().client_request.extend(request.messages); + + let mut context = QueryEngineContext::new(&mut client.client); + client.engine.parse_and_rewrite(&mut context).await?; + client.engine.route_query(&mut context).await?; + + assert!( + context.client_request.route().shard().is_direct(), + "UPDATE stmt should be using direct-to-shard routing" + ); + + client.engine.connect(&mut context, None).await?; + + assert!( + client.engine.backend.is_direct(), + "backend should be connected with Binding::Direct" + ); + + let rewrite = context + .client_request + .ast + .as_ref() + .expect("ast to exist") + .rewrite_plan + .clone() + .sharding_key_update + .clone() + .expect("sharding key update to exist"); + + let mut update = UpdateMulti::new(&mut client.engine, rewrite); + assert!( + update.is_same_shard(&context).unwrap(), + "query should not trigger multi-shard update" + ); + + // Won't error out because the query goes to the same shard + // as the old shard. + update.execute(&mut context).await?; + + Ok(()) +} + +#[tokio::test] +async fn test_update_check_simple() { + same_shard_check( + vec![Query::new("UPDATE sharded SET id = 1 WHERE id = 1 AND value = 'test'").into()].into(), + ) + .await + .unwrap(); +} + +#[tokio::test] +async fn test_update_check_extended() { + same_shard_check( + vec![ + Parse::new_anonymous("UPDATE sharded SET id = $1 WHERE id = $1 AND value = $2").into(), + Bind::new_params( + "", + &[ + Parameter::new("1234".as_bytes()), + Parameter::new("test".as_bytes()), + ], + ) + .into(), + Execute::new().into(), + Sync.into(), + ] + .into(), + ) + .await + .unwrap(); + + same_shard_check( + vec![ + Parse::new_anonymous( + "UPDATE sharded SET id = $1, value = $2 WHERE id = $3 AND value = $4", + ) + .into(), + Bind::new_params( + "", + &[ + Parameter::new("1234".as_bytes()), + Parameter::new("test".as_bytes()), + Parameter::new("1234".as_bytes()), + Parameter::new("test2".as_bytes()), + ], + ) + .into(), + Execute::new().into(), + Sync.into(), + ] + .into(), + ) + .await + .unwrap(); +} + +#[tokio::test] +async fn test_row_same_shard_no_transaction() { + crate::logger(); + let mut client = TestClient::new_rewrites(Parameters::default()).await; + + let shard_0 = client.random_id_for_shard(0); + let shard_0_1 = client.random_id_for_shard(0); + + client + .send_simple(Query::new(format!( + "INSERT INTO sharded (id, value) VALUES ({}, 'test value')", + shard_0 + ))) + .await; + client.read_until('Z').await.unwrap(); + + client.client.client_request = ClientRequest::from(vec![Query::new(format!( + "UPDATE sharded SET id = {} WHERE value = 'test value' AND id = {}", + shard_0_1, shard_0 + )) + .into()]); + + let mut context = QueryEngineContext::new(&mut client.client); + + client.engine.parse_and_rewrite(&mut context).await.unwrap(); + + assert!( + context + .client_request + .ast + .as_ref() + .expect("ast to exist") + .rewrite_plan + .sharding_key_update + .is_some(), + "sharding key update should exist on the request" + ); + + client.engine.route_query(&mut context).await.unwrap(); + client.engine.execute(&mut context).await.unwrap(); + + let cmd = client.read().await; + + assert_eq!( + CommandComplete::try_from(cmd).unwrap().command(), + "UPDATE 1" + ); + + expect_message!(client.read().await, ReadyForQuery); +} + +#[tokio::test] +async fn test_no_rows_updated() { + let mut client = TestClient::new_rewrites(Parameters::default()).await; + let id = rng().random::(); + + // Transaction not required because + // it'll check for existing row first (on the same shard). + client + .send_simple(Query::new(format!( + "UPDATE sharded SET id = {} WHERE id = {}", + id, + id + 1 + ))) + .await; + let cc = client.read().await; + expect_message!(cc.clone(), CommandComplete); + assert_eq!(CommandComplete::try_from(cc).unwrap().command(), "UPDATE 0"); + expect_message!(client.read().await, ReadyForQuery); +} + +#[tokio::test] +async fn test_transaction_required() { + let mut client = TestClient::new_rewrites(Parameters::default()).await; + + let shard_0 = client.random_id_for_shard(0); + let shard_1 = client.random_id_for_shard(1); + + client + .send_simple(Query::new(format!( + "INSERT INTO sharded (id) VALUES ({}) ON CONFLICT(id) DO NOTHING", + shard_0 + ))) + .await; + client.read_until('Z').await.unwrap(); + + client + .send_simple(Query::new(format!( + "UPDATE sharded SET id = {} WHERE id = {}", + shard_1, shard_0 + ))) + .await; + let err = ErrorResponse::try_from(client.read().await).expect("expected error"); + assert_eq!( + err.message, + "sharding key update must be executed inside a transaction" + ); + // Connection still good. + client.send_simple(Query::new("SELECT 1")).await; + client.read_until('Z').await.unwrap(); +} + +#[tokio::test] +async fn test_move_rows_simple() { + let mut client = TestClient::new_rewrites(Parameters::default()).await; + + let shard_0_id = client.random_id_for_shard(0); + let shard_1_id = client.random_id_for_shard(1); + + client + .send_simple(Query::new(format!( + "INSERT INTO sharded (id) VALUES ({}) ON CONFLICT(id) DO NOTHING", + shard_0_id + ))) + .await; + client.read_until('Z').await.unwrap(); + + client.send_simple(Query::new("BEGIN")).await; + client.read_until('Z').await.unwrap(); + + client + .try_send_simple(Query::new(format!( + "UPDATE sharded SET id = {} WHERE id = {} RETURNING id", + shard_1_id, shard_0_id + ))) + .await + .unwrap(); + + let reply = client.read_until('Z').await.unwrap(); + + let shard_1_id_str = shard_1_id.to_string(); + reply + .into_iter() + .zip(['T', 'D', 'C', 'Z']) + .for_each(|(message, code)| { + assert_eq!(message.code(), code); + match code { + 'C' => assert_eq!( + CommandComplete::try_from(message).unwrap().command(), + "UPDATE 1" + ), + 'Z' => assert!( + ReadyForQuery::try_from(message).unwrap().state().unwrap() + == TransactionState::InTrasaction + ), + 'T' => assert_eq!( + RowDescription::try_from(message) + .unwrap() + .field(0) + .unwrap() + .name, + "id" + ), + 'D' => assert_eq!( + DataRow::try_from(message).unwrap().column(0).unwrap(), + shard_1_id_str.as_bytes() + ), + _ => unreachable!(), + } + }); +} + +#[tokio::test] +async fn test_move_rows_extended() { + let mut client = TestClient::new_rewrites(Parameters::default()).await; + + let shard_0_id = client.random_id_for_shard(0); + let shard_1_id = client.random_id_for_shard(1); + + client + .send_simple(Query::new(format!( + "INSERT INTO sharded (id) VALUES ({}) ON CONFLICT(id) DO NOTHING", + shard_0_id + ))) + .await; + client.read_until('Z').await.unwrap(); + + client.send_simple(Query::new("BEGIN")).await; + client.read_until('Z').await.unwrap(); + + client + .send(Parse::new_anonymous( + "UPDATE sharded SET id = $2 WHERE id = $1 RETURNING id", + )) + .await; + client + .send(Bind::new_params( + "", + &[ + Parameter::new(shard_0_id.to_string().as_bytes()), + Parameter::new(shard_1_id.to_string().as_bytes()), + ], + )) + .await; + client.send(Execute::new()).await; + client.send(Sync).await; + client.try_process().await.unwrap(); + + let reply = client.read_until('Z').await.unwrap(); + + let shard_1_id_str = shard_1_id.to_string(); + reply + .into_iter() + .zip(['1', '2', 'D', 'C', 'Z']) + .for_each(|(message, code)| { + assert_eq!(message.code(), code); + match code { + 'C' => assert_eq!( + CommandComplete::try_from(message).unwrap().command(), + "UPDATE 1" + ), + 'Z' => assert!( + ReadyForQuery::try_from(message).unwrap().state().unwrap() + == TransactionState::InTrasaction + ), + 'D' => assert_eq!( + DataRow::try_from(message).unwrap().column(0).unwrap(), + shard_1_id_str.as_bytes() + ), + '1' | '2' => (), + _ => unreachable!(), + } + }); +} + +#[tokio::test] +async fn test_move_rows_prepared() { + crate::logger(); + let mut client = TestClient::new_rewrites(Parameters::default()).await; + + let shard_0_id = client.random_id_for_shard(0); + let shard_1_id = client.random_id_for_shard(1); + + client + .send_simple(Query::new(format!( + "INSERT INTO sharded (id) VALUES ({}) ON CONFLICT(id) DO NOTHING", + shard_0_id + ))) + .await; + client.read_until('Z').await.unwrap(); + + client.send_simple(Query::new("BEGIN")).await; + client.read_until('Z').await.unwrap(); + + client + .send(Parse::named( + "__test_1", + "UPDATE sharded SET id = $2 WHERE id = $1 RETURNING id", + )) + .await; + client.send(Describe::new_statement("__test_1")).await; + client.send(Flush).await; + client.try_process().await.unwrap(); + + let reply = client.read_until('T').await.unwrap(); + + reply + .into_iter() + .zip(['1', 't', 'T']) + .for_each(|(message, code)| { + assert_eq!(message.code(), code); + + match code { + 'T' => assert_eq!( + RowDescription::try_from(message) + .unwrap() + .field(0) + .unwrap() + .name, + "id" + ), + + 't' | '1' => (), + _ => unreachable!(), + } + }); + + client + .send(Bind::new_params( + "__test_1", + &[ + Parameter::new(shard_0_id.to_string().as_bytes()), + Parameter::new(shard_1_id.to_string().as_bytes()), + ], + )) + .await; + client.send(Execute::new()).await; + client.send(Sync).await; + client.try_process().await.unwrap(); + + let reply = client.read_until('Z').await.unwrap(); + + let shard_1_id_str = shard_1_id.to_string(); + reply + .into_iter() + .zip(['2', 'D', 'C', 'Z']) + .for_each(|(message, code)| { + assert_eq!(message.code(), code); + match code { + 'C' => assert_eq!( + CommandComplete::try_from(message).unwrap().command(), + "UPDATE 1" + ), + 'Z' => assert!( + ReadyForQuery::try_from(message).unwrap().state().unwrap() + == TransactionState::InTrasaction + ), + 'D' => assert_eq!( + DataRow::try_from(message).unwrap().column(0).unwrap(), + shard_1_id_str.as_bytes() + ), + '1' | '2' => (), + _ => unreachable!(), + } + }); +} + +#[tokio::test] +async fn test_same_shard_binary() { + let mut client = TestClient::new_rewrites(Parameters::default()).await; + let id = client.random_id_for_shard(0); + client + .send_simple(Query::new(format!( + "INSERT INTO sharded (id) VALUES ({})", + id + ))) + .await; + client.read_until('Z').await.unwrap(); + let id_2 = client.random_id_for_shard(0); + client + .send(Parse::new_anonymous( + "UPDATE sharded SET id = $1 WHERE id = $2 RETURNING *", + )) + .await; + client + .send(Bind::new_params_codes( + "", + &[ + Parameter::new(&id_2.to_be_bytes()), + Parameter::new(&id.to_be_bytes()), + ], + &[Format::Binary], + )) + .await; + client.send(Execute::new()).await; + client.send(Sync).await; + client.try_process().await.unwrap(); + let messages = client.read_until('Z').await.unwrap(); + + messages + .into_iter() + .zip(['1', '2', 'D', 'C', 'Z']) + .for_each(|(message, code)| { + assert_eq!(message.code(), code); + if message.code() == 'C' { + assert_eq!( + CommandComplete::try_from(message).unwrap().command(), + "UPDATE 1" + ); + } + }); +} + +#[tokio::test] +async fn test_update_with_expr() { + // Test that UPDATE with expression columns (not simple values) works correctly. + // This validates the bind parameter alignment fix where expression columns + // don't consume bind parameter slots. + // + // Note: Expressions that reference the original row's columns (like COALESCE(value, 'default')) + // won't work because they're inserted literally into the INSERT statement where those + // columns don't exist. Only standalone expressions like 'prefix' || 'suffix' work. + let mut client = TestClient::new_rewrites(Parameters::default()).await; + + // Use random IDs to avoid conflicts with other tests + let shard_0_id = client.random_id_for_shard(0); + let shard_1_id = client.random_id_for_shard(1); + + // Insert a row into shard 0 + client + .send_simple(Query::new(format!( + "INSERT INTO sharded (id, value) VALUES ({}, 'original') ON CONFLICT(id) DO UPDATE SET value = 'original'", + shard_0_id + ))) + .await; + client.read_until('Z').await.unwrap(); + + client.send_simple(Query::new("BEGIN")).await; + client.read_until('Z').await.unwrap(); + + // UPDATE that moves row to different shard with an expression column. + // Use a standalone expression that doesn't reference any columns. + client + .try_send_simple(Query::new(format!( + "UPDATE sharded SET id = {}, value = 'prefix' || '_suffix' WHERE id = {} RETURNING id, value", + shard_1_id, shard_0_id + ))) + .await + .unwrap(); + + let reply = client.read_until('Z').await.unwrap(); + + let shard_1_id_str = shard_1_id.to_string(); + reply + .into_iter() + .zip(['T', 'D', 'C', 'Z']) + .for_each(|(message, code)| { + assert_eq!(message.code(), code); + match code { + 'C' => assert_eq!( + CommandComplete::try_from(message).unwrap().command(), + "UPDATE 1" + ), + 'Z' => assert!( + ReadyForQuery::try_from(message).unwrap().state().unwrap() + == TransactionState::InTrasaction + ), + 'T' => { + let rd = RowDescription::try_from(message).unwrap(); + assert_eq!(rd.field(0).unwrap().name, "id"); + assert_eq!(rd.field(1).unwrap().name, "value"); + } + 'D' => { + let dr = DataRow::try_from(message).unwrap(); + assert_eq!(dr.column(0).unwrap(), shard_1_id_str.as_bytes()); + // The value should be 'prefix_suffix' from the expression + assert_eq!(dr.column(1).unwrap(), "prefix_suffix".as_bytes()); + } + _ => unreachable!(), + } + }); + + client.send_simple(Query::new("COMMIT")).await; + client.read_until('Z').await.unwrap(); + + // Verify the row was actually moved to the new shard with correct values + client + .send_simple(Query::new(format!( + "SELECT id, value FROM sharded WHERE id = {}", + shard_1_id + ))) + .await; + let reply = client.read_until('Z').await.unwrap(); + + let data_row = reply + .iter() + .find(|m| m.code() == 'D') + .expect("should have data row"); + let dr = DataRow::try_from(data_row.clone()).unwrap(); + assert_eq!(dr.column(0).unwrap(), shard_1_id_str.as_bytes()); + assert_eq!(dr.column(1).unwrap(), "prefix_suffix".as_bytes()); +} diff --git a/pgdog/src/frontend/client/query_engine/multi_step/update.rs b/pgdog/src/frontend/client/query_engine/multi_step/update.rs new file mode 100644 index 000000000..2d92b1fed --- /dev/null +++ b/pgdog/src/frontend/client/query_engine/multi_step/update.rs @@ -0,0 +1,300 @@ +use pgdog_config::RewriteMode; +use tracing::debug; + +use crate::{ + frontend::{ + client::query_engine::{QueryEngine, QueryEngineContext}, + router::parser::rewrite::statement::ShardingKeyUpdate, + ClientRequest, Command, Router, RouterContext, + }, + net::{CommandComplete, DataRow, ErrorResponse, Protocol, ReadyForQuery, RowDescription}, +}; + +use super::{Error, ForwardCheck, UpdateError}; + +#[derive(Debug, Clone, Default)] +pub(super) struct Row { + data_row: DataRow, + row_description: RowDescription, +} + +#[derive(Debug)] +pub(crate) struct UpdateMulti<'a> { + pub(super) rewrite: ShardingKeyUpdate, + pub(super) engine: &'a mut QueryEngine, +} + +impl<'a> UpdateMulti<'a> { + /// Create new sharding key update handler. + pub(crate) fn new(engine: &'a mut QueryEngine, rewrite: ShardingKeyUpdate) -> Self { + Self { rewrite, engine } + } + + /// Execute sharding key update, if needed. + pub(crate) async fn execute( + &mut self, + context: &mut QueryEngineContext<'_>, + ) -> Result<(), Error> { + match self.execute_internal(context).await { + Ok(()) => Ok(()), + Err(err) => { + // These are recoverable with a ROLLBACK. + if matches!(err, Error::Update(_) | Error::Execution(_)) { + self.engine + .error_response(context, ErrorResponse::from_err(&err)) + .await?; + Ok(()) + } else { + // These are bad, disconnecting the client. + Err(err) + } + } + } + } + + /// Execute sharding key update, if needed. + pub(super) async fn execute_internal( + &mut self, + context: &mut QueryEngineContext<'_>, + ) -> Result<(), Error> { + let mut check = self.rewrite.check.build_request(context.client_request)?; + self.route(&mut check, context)?; + + // The new row is on the same shard as the old row + // and we know this from the statement itself, e.g. + // + // UPDATE my_table SET shard_key = $1 WHERE shard_key = $2 + // + // This is very likely if the number of shards is low or + // you're using an ORM that puts all record columns + // into the SET clause. + // + if self.is_same_shard(context)? { + // Serve original request as-is. + debug!("[update] row is on the same shard"); + self.execute_original(context).await?; + + return Ok(()); + } + + // Fetch the old row from whatever shard it is on. + let row = self.fetch_row(context).await?; + + if let Some(row) = row { + self.insert_row(context, row).await?; + } else { + // This happens, but the UPDATE's WHERE clause + // doesn't match any rows, so this whole thing is a no-op. + self.engine + .fake_command_response(context, "UPDATE 0") + .await?; + } + + Ok(()) + } + + /// Create row. + pub(super) async fn insert_row( + &mut self, + context: &mut QueryEngineContext<'_>, + row: Row, + ) -> Result<(), Error> { + let mut request = self.rewrite.insert.build_request( + context.client_request, + &row.row_description, + &row.data_row, + )?; + self.route(&mut request, context)?; + + let original_shard = context.client_request.route().shard(); + let new_shard = request.route().shard(); + + // The new row maps to the same shard as the old row. + // We don't need to do the multi-step UPDATE anymore. + // Forward the original request as-is. + if original_shard.is_direct() && new_shard == original_shard { + debug!("[update] selected row is on the same shard"); + self.execute_original(context).await + } else { + debug!("[update] executing multi-shard insert/delete"); + + // Check if we are allowed to do this operation by the config. + if self.engine.backend.cluster()?.rewrite().shard_key == RewriteMode::Error { + self.engine + .error_response(context, ErrorResponse::from_err(&UpdateError::Disabled)) + .await?; + return Ok(()); + } + + if !context.in_transaction() && !self.engine.backend.is_multishard() + // Do this check at the last possible moment. + // Just in case we change how transactions are + // routed in the future. + { + return Err(UpdateError::TransactionRequired.into()); + } + + self.delete_row(context).await?; + self.execute_request_internal( + context, + &mut request, + self.rewrite.insert.is_returning(), + ) + .await?; + + self.engine + .process_server_message(context, CommandComplete::new("UPDATE 1").message()?) // We only allow to update one row at a time. + .await?; + self.engine + .process_server_message( + context, + ReadyForQuery::in_transaction(context.in_transaction()).message()?, + ) + .await?; + + Ok(()) + } + } + + /// Execute request and return messages to the client if forward_reply is true. + async fn execute_request_internal( + &mut self, + context: &mut QueryEngineContext<'_>, + request: &mut ClientRequest, + forward_reply: bool, + ) -> Result<(), Error> { + self.engine + .backend + .handle_client_request(request, &mut Router::default(), false) + .await?; + + let mut checker = ForwardCheck::new(context.client_request); + + while self.engine.backend.has_more_messages() { + let message = self.engine.read_server_message(context).await?; + let code = message.code(); + + if code == 'E' { + return Err(Error::Execution(ErrorResponse::try_from(message)?)); + } + + if forward_reply && checker.forward(code) { + self.engine.process_server_message(context, message).await?; + } + } + + Ok(()) + } + + async fn execute_original( + &mut self, + context: &mut QueryEngineContext<'_>, + ) -> Result<(), Error> { + // Serve original request as-is. + self.engine + .backend + .handle_client_request( + context.client_request, + &mut self.engine.router, + self.engine.streaming, + ) + .await?; + + while self.engine.backend.has_more_messages() { + let message = self.engine.read_server_message(context).await?; + self.engine.process_server_message(context, message).await?; + } + + Ok(()) + } + + pub(super) async fn delete_row( + &mut self, + context: &mut QueryEngineContext<'_>, + ) -> Result<(), Error> { + let mut request = self.rewrite.delete.build_request(context.client_request)?; + self.route(&mut request, context)?; + + self.execute_request_internal(context, &mut request, false) + .await + } + + pub(super) async fn fetch_row( + &mut self, + context: &mut QueryEngineContext<'_>, + ) -> Result, Error> { + let mut request = self.rewrite.select.build_request(context.client_request)?; + self.route(&mut request, context)?; + + self.engine + .backend + .handle_client_request(&mut request, &mut Router::default(), false) + .await?; + + let mut row = Row::default(); + let mut rows = 0; + + while self.engine.backend.has_more_messages() { + let message = self.engine.read_server_message(context).await?; + match message.code() { + 'D' => { + row.data_row = DataRow::try_from(message)?; + rows += 1; + } + 'T' => row.row_description = RowDescription::try_from(message)?, + 'E' => return Err(Error::Execution(ErrorResponse::try_from(message)?)), + _ => (), + } + } + + match rows { + 0 => return Ok(None), + 1 => (), + n => return Err(UpdateError::TooManyRows(n).into()), + } + + Ok(Some(row)) + } + + /// Returns true if the new sharding key resides on the same shard + /// as the old sharding key. + /// + /// This is an optimization to avoid doing a multi-shard UPDATE when + /// we don't have to. + pub(super) fn is_same_shard(&self, context: &QueryEngineContext<'_>) -> Result { + let mut check = self.rewrite.check.build_request(context.client_request)?; + self.route(&mut check, context)?; + + let new_shard = check.route().shard(); + let old_shard = context.client_request.route().shard(); + + // The sharding key isn't actually being changed + // or it maps to the same shard as before. + Ok(new_shard == old_shard) + } + + fn route( + &self, + request: &mut ClientRequest, + context: &QueryEngineContext<'_>, + ) -> Result<(), Error> { + let cluster = self.engine.backend.cluster()?; + + let context = RouterContext::new( + request, + cluster, + context.params, + context.transaction(), + context.sticky, + )?; + let mut router = Router::new(); + let command = router.query(context)?; + if let Command::Query(route) = command { + request.route = Some(route.clone()); + } else { + return Err(UpdateError::NoRoute.into()); + } + + Ok(()) + } +} diff --git a/pgdog/src/frontend/client/query_engine/prepared_statements.rs b/pgdog/src/frontend/client/query_engine/prepared_statements.rs deleted file mode 100644 index 68b38845e..000000000 --- a/pgdog/src/frontend/client/query_engine/prepared_statements.rs +++ /dev/null @@ -1,16 +0,0 @@ -use super::*; - -impl QueryEngine { - /// Rewrite extended protocol messages. - pub(super) fn rewrite_extended( - &mut self, - context: &mut QueryEngineContext<'_>, - ) -> Result<(), Error> { - for message in context.client_request.iter_mut() { - if message.extended() && context.prepared_statements.enabled { - context.prepared_statements.maybe_rewrite(message)?; - } - } - Ok(()) - } -} diff --git a/pgdog/src/frontend/client/query_engine/pub_sub.rs b/pgdog/src/frontend/client/query_engine/pub_sub.rs index 9d1240913..104f6316d 100644 --- a/pgdog/src/frontend/client/query_engine/pub_sub.rs +++ b/pgdog/src/frontend/client/query_engine/pub_sub.rs @@ -61,7 +61,7 @@ impl QueryEngine { let bytes_sent = context .stream .send_many(&[ - CommandComplete::new(command).message()?, + CommandComplete::new(command).message()?.backend(), ReadyForQuery::in_transaction(context.in_transaction()).message()?, ]) .await?; diff --git a/pgdog/src/frontend/client/query_engine/query.rs b/pgdog/src/frontend/client/query_engine/query.rs index 6d0efa024..3335bb485 100644 --- a/pgdog/src/frontend/client/query_engine/query.rs +++ b/pgdog/src/frontend/client/query_engine/query.rs @@ -1,15 +1,19 @@ use tokio::time::timeout; +use tracing::trace; use crate::{ - frontend::client::TransactionType, + frontend::{ + client::TransactionType, + router::parser::{explain_trace::ExplainTrace, rewrite::statement::plan::RewriteResult}, + }, net::{ - FromBytes, Message, Protocol, ProtocolMessage, Query, ReadyForQuery, ToBytes, - TransactionState, + DataRow, FromBytes, Message, Protocol, ProtocolMessage, Query, ReadyForQuery, + RowDescription, ToBytes, TransactionState, }, state::State, }; -use tracing::debug; +use tracing::{debug, error}; use super::*; @@ -18,7 +22,6 @@ impl QueryEngine { pub(super) async fn execute( &mut self, context: &mut QueryEngineContext<'_>, - route: &Route, ) -> Result<(), Error> { // Check that we're not in a transaction error state. if !self.transaction_error_check(context).await? { @@ -27,39 +30,25 @@ impl QueryEngine { // Check if we need to do 2pc automatically // for single-statement writes. - self.two_pc_check(context, route); + self.two_pc_check(context); // We need to run a query now. - if let Some(begin_stmt) = self.begin_stmt.take() { + if context.in_transaction() { // Connect to one shard if not sharded or to all shards // for a cross-shard tranasction. - if !self.connect_transaction(context, route).await? { + if !self.connect_transaction(context).await? { return Ok(()); } - - self.backend.execute(begin_stmt.query()).await?; - } else if !self.connect(context, route).await? { + } else if !self.connect(context, None).await? { return Ok(()); } // Check we can run this query. - if !self.cross_shard_check(context, route).await? { + if !self.cross_shard_check(context).await? { return Ok(()); } - if let Some(sql) = route.rewritten_sql() { - match context.client_request.rewrite(sql) { - Ok(()) => (), - Err(crate::net::Error::OnlySimpleForRewrites) => { - context.client_request.rewrite_prepared( - sql, - context.prepared_statements, - route.rewrite_plan(), - ); - } - Err(err) => return Err(err.into()), - } - } + self.hooks.after_connected(context, &self.backend)?; // Set response format. for msg in context.client_request.messages.iter() { @@ -68,27 +57,61 @@ impl QueryEngine { } } - self.backend - .handle_client_request(context.client_request, &mut self.router, self.streaming) - .await?; + match context.rewrite_result.take() { + Some(RewriteResult::InsertSplit(requests)) => { + multi_step::InsertMulti::from_engine(self, requests) + .execute(context) + .await?; + } - while self.backend.has_more_messages() - && !self.backend.copy_mode() - && !self.streaming - && !self.test_mode - { - let message = timeout( - context.timeouts.query_timeout(&State::Active), - self.backend.read(), - ) - .await??; - self.server_message(context, message).await?; + Some(RewriteResult::InPlace) | None => { + self.backend + .handle_client_request(context.client_request, &mut self.router, self.streaming) + .await?; + + while self.backend.has_more_messages() + && !self.backend.in_copy_mode() + && !self.streaming + && !self.test_mode.enabled + { + let message = self.read_server_message(context).await?; + self.process_server_message(context, message).await?; + } + } + + Some(RewriteResult::ShardingKeyUpdate(sharding_key_update)) => { + multi_step::UpdateMulti::new(self, sharding_key_update) + .execute(context) + .await?; + } } Ok(()) } - pub async fn server_message( + pub async fn read_server_message( + &mut self, + context: &mut QueryEngineContext<'_>, + ) -> Result { + let message = match timeout( + context.timeouts.query_timeout(&State::Active), + self.backend.read(), + ) + .await + { + Ok(response) => response?, + Err(err) => { + // Close the conn, it could be stuck executing a query + // or dead. + self.backend.force_close(); + return Err(err.into()); + } + }; + + Ok(message) + } + + pub async fn process_server_message( &mut self, context: &mut QueryEngineContext<'_>, message: Message, @@ -96,9 +119,35 @@ impl QueryEngine { self.streaming = message.streaming(); let code = message.code(); + let payload = if code == 'T' { + Some(message.payload()) + } else { + None + }; let mut message = message.backend(); let has_more_messages = self.backend.has_more_messages(); + if let Some(bytes) = payload { + if let Some(state) = self.pending_explain.as_mut() { + if let Ok(row_description) = RowDescription::from_bytes(bytes) { + state.capture_row_description(row_description); + } else { + state.annotated = true; + } + } + } + + if code == 'C' { + self.emit_explain_rows(context).await?; + } + + if code == 'E' { + if let Some(state) = self.pending_explain.as_mut() { + state.annotated = true; + } + self.pending_explain = None; + } + // Messages that we need to send to the client immediately. // ReadyForQuery (B) | CopyInResponse (B) | ErrorResponse(B) | NoticeResponse(B) | NotificationResponse (B) let flush = matches!(code, 'Z' | 'G' | 'E' | 'N' | 'A') @@ -115,7 +164,12 @@ impl QueryEngine { match state { TransactionState::Error => { - context.transaction = Some(TransactionType::Error); + let error_state = match context.transaction { + Some(TransactionType::ReadOnly) => Some(TransactionType::ErrorReadOnly), + Some(TransactionType::ReadWrite) => Some(TransactionType::ErrorReadWrite), + _ => None, + }; + context.transaction = error_state; if self.two_pc.auto() { self.end_two_pc(true).await?; // TODO: this records a 2pc transaction in client @@ -133,10 +187,23 @@ impl QueryEngine { self.end_two_pc(false).await?; two_pc_auto = true; } - if context.transaction.is_none() { + match context.transaction { // Query parser is disabled, so the server is responsible for telling us // we started a transaction. - context.transaction = Some(TransactionType::ReadWrite); + None => { + context.transaction = Some(TransactionType::ReadWrite); + } + + // Restore transaction state after rollback to savepoint. + Some(TransactionType::ErrorReadOnly) => { + context.transaction = Some(TransactionType::ReadOnly); + } + + Some(TransactionType::ErrorReadWrite) => { + context.transaction = Some(TransactionType::ReadWrite); + } + + _ => (), } } } @@ -161,12 +228,47 @@ impl QueryEngine { // Do this before flushing, because flushing can take time. self.cleanup_backend(context); + trace!("{:#?} >>> {:?}", message, context.stream.peer_addr()); + if flush { context.stream.send_flush(&message).await?; } else { context.stream.send(&message).await?; } + if code == 'Z' { + self.pending_explain = None; + } + self.hooks.on_server_message(context, &message)?; + + Ok(()) + } + + async fn emit_explain_rows( + &mut self, + context: &mut QueryEngineContext<'_>, + ) -> Result<(), Error> { + if let Some(state) = self.pending_explain.as_mut() { + if !state.should_emit() { + return Ok(()); + } + + if state.row_description.is_none() { + return Ok(()); + } + + for line in state.lines.clone() { + let mut row = DataRow::new(); + row.add(line); + let message = row.message()?; + let len = message.len(); + context.stream.send(&message).await?; + self.stats.sent(len); + } + + state.annotated = true; + } + Ok(()) } @@ -176,7 +278,7 @@ impl QueryEngine { // Release the connection back into the pool before flushing data to client. // Flushing can take a minute and we don't want to block the connection from being reused. - if self.backend.transaction_mode() && context.requests_left == 0 { + if !self.backend.session_mode() && context.requests_left == 0 { self.backend.disconnect(); } @@ -203,7 +305,6 @@ impl QueryEngine { async fn cross_shard_check( &mut self, context: &mut QueryEngineContext<'_>, - route: &Route, ) -> Result { // Check for cross-shard queries. if context.cross_shard_disabled.is_none() { @@ -220,29 +321,26 @@ impl QueryEngine { debug!("cross-shard queries disabled: {}", cross_shard_disabled,); if cross_shard_disabled - && route.is_cross_shard() + && context.client_request.route().is_cross_shard() && !context.admin - && context.client_request.executable() + && context.client_request.is_executable() { - let bytes_sent = context - .stream - .error( - ErrorResponse::cross_shard_disabled(), - context.in_transaction(), - ) - .await?; - self.stats.sent(bytes_sent); + let query = context.client_request.query()?; + let error = ErrorResponse::cross_shard_disabled(query.as_ref().map(|q| q.query())); + + self.error_response(context, error).await?; if self.backend.connected() && self.backend.done() { self.backend.disconnect(); } + Ok(false) } else { Ok(true) } } - fn two_pc_check(&mut self, context: &mut QueryEngineContext<'_>, route: &Route) { + fn two_pc_check(&mut self, context: &mut QueryEngineContext<'_>) { let enabled = self .backend .cluster() @@ -250,9 +348,9 @@ impl QueryEngine { .unwrap_or_default(); if enabled - && route.should_2pc() + && context.client_request.route().should_2pc() && self.begin_stmt.is_none() - && context.client_request.executable() + && context.client_request.is_executable() && !context.in_transaction() { debug!("[2pc] enabling automatic transaction"); @@ -265,19 +363,84 @@ impl QueryEngine { &mut self, context: &mut QueryEngineContext<'_>, ) -> Result { - if context.in_error() && !context.rollback && context.client_request.executable() { - let bytes_sent = context - .stream - .error( - ErrorResponse::in_failed_transaction(), - context.in_transaction(), - ) - .await?; - self.stats.sent(bytes_sent); + let shards = if let Ok(shards) = self.backend.shards() { + shards + } else { + return Ok(true); + }; + if shards > 1 // This check only matters for cross-shard queries + && context.in_error() + && !context.rollback + && context.client_request.is_executable() + && !context.client_request.route().rollback_savepoint() + { + let error = ErrorResponse::in_failed_transaction(); + + self.error_response(context, error).await?; Ok(false) } else { Ok(true) } } + + pub(super) async fn error_response( + &mut self, + context: &mut QueryEngineContext<'_>, + mut error: ErrorResponse, + ) -> Result<(), Error> { + error!("{:?} [{:?}]", error.message, context.stream.peer_addr()); + + // Attach query context. + if error.detail.is_none() { + let query = context + .client_request + .query()? + .map(|q| q.query().to_owned()); + error.detail = Some(query.unwrap_or_default()); + } + + self.hooks.on_engine_error(context, &error)?; + + let bytes_sent = context + .stream + .error(error, context.in_transaction()) + .await?; + self.stats.sent(bytes_sent); + + Ok(()) + } +} + +#[derive(Debug, Default, Clone)] +pub(super) struct ExplainResponseState { + lines: Vec, + row_description: Option, + annotated: bool, + supported: bool, +} + +impl ExplainResponseState { + pub fn new(trace: ExplainTrace) -> Self { + Self { + lines: trace.render_lines(), + row_description: None, + annotated: false, + supported: false, + } + } + + pub fn capture_row_description(&mut self, row_description: RowDescription) { + self.supported = row_description.fields.len() == 1 + && matches!(row_description.field(0).map(|f| f.type_oid), Some(25)); + if self.supported { + self.row_description = Some(row_description); + } else { + self.annotated = true; + } + } + + pub fn should_emit(&self) -> bool { + self.supported && !self.annotated + } } diff --git a/pgdog/src/frontend/client/query_engine/rewrite.rs b/pgdog/src/frontend/client/query_engine/rewrite.rs new file mode 100644 index 000000000..a36e28eb1 --- /dev/null +++ b/pgdog/src/frontend/client/query_engine/rewrite.rs @@ -0,0 +1,70 @@ +use crate::{config::PreparedStatements, frontend::router::parser::Cache}; + +use super::*; + +impl QueryEngine { + /// Rewrite extended protocol messages. + pub(super) fn rewrite_extended( + &mut self, + context: &mut QueryEngineContext<'_>, + ) -> Result<(), Error> { + for message in context.client_request.iter_mut() { + if message.extended() { + let level = context.prepared_statements.level; + match (level, message.anonymous()) { + (PreparedStatements::ExtendedAnonymous, _) + | (PreparedStatements::Extended, false) => { + context.prepared_statements.maybe_rewrite(message)? + } + _ => (), + } + } + } + Ok(()) + } + + /// Parse client request and rewrite it, if necessary. + pub(super) async fn parse_and_rewrite( + &mut self, + context: &mut QueryEngineContext<'_>, + ) -> Result { + let use_parser = self + .backend + .cluster() + .map(|cluster| cluster.use_query_parser()) + .unwrap_or(false); + + if !use_parser { + return Ok(true); + } + + let query = context.client_request.query()?; + if let Some(query) = query { + let ast = match Cache::get().query( + &query, + &self.backend.cluster()?.sharding_schema(), + context.prepared_statements, + ) { + Ok(ast) => ast, + Err(err) => { + self.error_response(context, ErrorResponse::syntax(err.to_string().as_str())) + .await?; + return Ok(false); + } + }; + context.client_request.ast = Some(ast); + } + + let plan = context + .client_request + .ast + .as_ref() + .map(|ast| ast.rewrite_plan.clone()); + + if let Some(plan) = plan { + context.rewrite_result = Some(plan.apply(context.client_request)?); + } + + Ok(true) + } +} diff --git a/pgdog/src/frontend/client/query_engine/route_query.rs b/pgdog/src/frontend/client/query_engine/route_query.rs index e8369680a..ad24bf924 100644 --- a/pgdog/src/frontend/client/query_engine/route_query.rs +++ b/pgdog/src/frontend/client/query_engine/route_query.rs @@ -1,13 +1,63 @@ -use tracing::{error, trace}; +use pgdog_config::PoolerMode; +use tracing::trace; use super::*; +#[derive(Debug, Clone)] +pub enum ClusterCheck { + Ok, + Offline, +} + impl QueryEngine { - pub(super) async fn route_transaction( + /// Get mutable reference to the backend connection. + pub fn backend(&mut self) -> &mut Connection { + &mut self.backend + } + + /// Check that the cluster is still valid and online. + pub async fn cluster_check( &mut self, context: &mut QueryEngineContext<'_>, - ) -> Result { + ) -> Result { // Admin doesn't have a cluster. + if let Ok(cluster) = self.backend.cluster() { + if !context.in_transaction() && !cluster.online() { + let identifier = cluster.identifier(); + + // Reload cluster config. + self.backend.safe_reload().await?; + + if self.backend.cluster().is_ok() { + Ok(ClusterCheck::Ok) + } else { + self.error_response( + context, + ErrorResponse::connection(&identifier.user, &identifier.database), + ) + .await?; + Ok(ClusterCheck::Offline) + } + } else { + Ok(ClusterCheck::Ok) + } + } else { + Ok(ClusterCheck::Ok) + } + } + + pub(super) async fn route_query( + &mut self, + context: &mut QueryEngineContext<'_>, + ) -> Result { + // Check that we can route this transaction at all. + if self.backend.pooler_mode() == PoolerMode::Statement && context.client_request.is_begin() + { + self.error_response(context, ErrorResponse::transaction_statement_mode()) + .await?; + return Ok(false); + } + let cluster = if let Ok(cluster) = self.backend.cluster() { cluster } else { @@ -17,28 +67,23 @@ impl QueryEngine { let router_context = RouterContext::new( context.client_request, cluster, - context.prepared_statements, context.params, context.transaction, + context.sticky, )?; match self.router.query(router_context) { - Ok(cmd) => { + Ok(command) => { + context.client_request.route = Some(command.route().clone()); trace!( "routing {:#?} to {:#?}", context.client_request.messages, - cmd + command, ); } Err(err) => { - error!("{:?} [{:?}]", err, context.stream.peer_addr()); - let bytes_sent = context - .stream - .error( - ErrorResponse::syntax(err.to_string().as_str()), - context.in_transaction(), - ) + self.error_response(context, ErrorResponse::syntax(err.to_string().as_str())) .await?; - self.stats.sent(bytes_sent); + return Ok(false); } } diff --git a/pgdog/src/frontend/client/query_engine/set.rs b/pgdog/src/frontend/client/query_engine/set.rs index d6ba8041e..8ba36e57d 100644 --- a/pgdog/src/frontend/client/query_engine/set.rs +++ b/pgdog/src/frontend/client/query_engine/set.rs @@ -1,4 +1,4 @@ -use crate::net::{parameter::ParameterValue, CommandComplete, Protocol, ReadyForQuery}; +use crate::net::parameter::ParameterValue; use super::*; @@ -8,43 +8,22 @@ impl QueryEngine { context: &mut QueryEngineContext<'_>, name: String, value: ParameterValue, + local: bool, ) -> Result<(), Error> { - context.params.insert(name, value); - self.comms.update_params(context.params); - - let bytes_sent = context - .stream - .send_many(&[ - CommandComplete::from_str("SET").message()?.backend(), - ReadyForQuery::in_transaction(context.in_transaction()) - .message()? - .backend(), - ]) - .await?; - - self.stats.sent(bytes_sent); - - Ok(()) - } - - pub(crate) async fn set_route( - &mut self, - context: &mut QueryEngineContext<'_>, - route: Route, - ) -> Result<(), Error> { - self.set_route = Some(route); - - let bytes_sent = context - .stream - .send_many(&[ - CommandComplete::from_str("SET").message()?.backend(), - ReadyForQuery::in_transaction(context.in_transaction()) - .message()? - .backend(), - ]) - .await?; - - self.stats.sent(bytes_sent); + if context.in_transaction() { + context + .params + .insert_transaction(name, value.clone(), local); + } else { + context.params.insert(name, value.clone()); + self.comms.update_params(context.params); + } + + if self.backend.connected() { + self.execute(context).await?; + } else { + self.fake_command_response(context, "SET").await?; + } Ok(()) } diff --git a/pgdog/src/frontend/client/query_engine/shard_key_rewrite.rs b/pgdog/src/frontend/client/query_engine/shard_key_rewrite.rs new file mode 100644 index 000000000..0b1666874 --- /dev/null +++ b/pgdog/src/frontend/client/query_engine/shard_key_rewrite.rs @@ -0,0 +1,255 @@ +// use std::collections::HashMap; + +// use super::*; +// use crate::{ +// backend::pool::{connection::binding::Binding, Guard, Request}, +// frontend::router::{ +// self as router, +// parser::{ +// self as parser, +// rewrite::{AssignmentValue, ShardKeyRewritePlan}, +// Shard, +// }, +// }, +// net::messages::Protocol, +// net::{ +// messages::{ +// bind::Format, command_complete::CommandComplete, Bind, DataRow, FromBytes, Message, +// RowDescription, ToBytes, +// }, +// ErrorResponse, ReadyForQuery, +// }, +// util::escape_identifier, +// }; +// use pgdog_plugin::pg_query::NodeEnum; +// use tracing::warn; + +// #[cfg(test)] +// mod tests { +// use super::*; +// use crate::frontend::router::{ +// parser::{rewrite::Assignment, route::Shard, table::OwnedTable, ShardWithPriority}, +// Route, +// }; +// use crate::{ +// backend::{ +// databases::{self, databases, lock, User as DbUser}, +// pool::{cluster::Cluster, Request}, +// }, +// config::{ +// self, +// core::ConfigAndUsers, +// database::Database, +// sharding::{DataType, FlexibleType, ShardedMapping, ShardedMappingKind, ShardedTable}, +// users::User as ConfigUser, +// RewriteMode, +// }, +// frontend::Client, +// net::{Query, Stream}, +// }; +// use std::collections::HashSet; + +// async fn configure_cluster(two_pc_enabled: bool) -> Cluster { +// let mut cfg = ConfigAndUsers::default(); +// cfg.config.general.two_phase_commit = two_pc_enabled; +// cfg.config.general.two_phase_commit_auto = Some(false); +// cfg.config.rewrite.enabled = true; +// cfg.config.rewrite.shard_key = RewriteMode::Rewrite; + +// cfg.config.databases = vec![ +// Database { +// name: "pgdog_sharded".into(), +// host: "127.0.0.1".into(), +// port: 5432, +// database_name: Some("pgdog".into()), +// shard: 0, +// ..Default::default() +// }, +// Database { +// name: "pgdog_sharded".into(), +// host: "127.0.0.1".into(), +// port: 5432, +// database_name: Some("pgdog".into()), +// shard: 1, +// ..Default::default() +// }, +// ]; + +// cfg.config.sharded_tables = vec![ShardedTable { +// database: "pgdog_sharded".into(), +// name: Some("sharded".into()), +// column: "id".into(), +// data_type: DataType::Bigint, +// primary: true, +// ..Default::default() +// }]; + +// let shard0_values = HashSet::from([ +// FlexibleType::Integer(1), +// FlexibleType::Integer(2), +// FlexibleType::Integer(3), +// FlexibleType::Integer(4), +// ]); +// let shard1_values = HashSet::from([ +// FlexibleType::Integer(5), +// FlexibleType::Integer(6), +// FlexibleType::Integer(7), +// ]); + +// cfg.config.sharded_mappings = vec![ +// ShardedMapping { +// database: "pgdog_sharded".into(), +// table: Some("sharded".into()), +// column: "id".into(), +// kind: ShardedMappingKind::List, +// values: shard0_values, +// shard: 0, +// ..Default::default() +// }, +// ShardedMapping { +// database: "pgdog_sharded".into(), +// table: Some("sharded".into()), +// column: "id".into(), +// kind: ShardedMappingKind::List, +// values: shard1_values, +// shard: 1, +// ..Default::default() +// }, +// ]; + +// cfg.users.users = vec![ConfigUser { +// name: "pgdog".into(), +// database: "pgdog_sharded".into(), +// password: Some("pgdog".into()), +// two_phase_commit: Some(two_pc_enabled), +// two_phase_commit_auto: Some(false), +// ..Default::default() +// }]; + +// config::set(cfg).unwrap(); +// databases::init().unwrap(); + +// let user = DbUser { +// user: "pgdog".into(), +// database: "pgdog_sharded".into(), +// }; + +// databases() +// .all() +// .get(&user) +// .expect("cluster missing") +// .clone() +// } + +// async fn prepare_table(cluster: &Cluster) { +// let request = Request::default(); +// let mut primary = cluster.primary(0, &request).await.unwrap(); +// primary +// .execute("CREATE TABLE IF NOT EXISTS sharded (id BIGINT PRIMARY KEY, value TEXT)") +// .await +// .unwrap(); +// primary.execute("TRUNCATE TABLE sharded").await.unwrap(); +// primary +// .execute("INSERT INTO sharded (id, value) VALUES (1, 'old')") +// .await +// .unwrap(); +// } + +// async fn table_state(cluster: &Cluster) -> (i64, i64) { +// let request = Request::default(); +// let mut primary = cluster.primary(0, &request).await.unwrap(); +// let old_id = primary +// .fetch_all::("SELECT COUNT(*)::bigint FROM sharded WHERE id = 1") +// .await +// .unwrap()[0]; +// let new_id = primary +// .fetch_all::("SELECT COUNT(*)::bigint FROM sharded WHERE id = 5") +// .await +// .unwrap()[0]; +// (old_id, new_id) +// } + +// fn new_client() -> Client { +// let stream = Stream::dev_null(); +// let mut client = Client::new_test(stream, Parameters::default()); +// client.params.insert("database", "pgdog_sharded"); +// client.connect_params.insert("database", "pgdog_sharded"); +// client +// } + +// #[tokio::test] +// async fn shard_key_rewrite_moves_row_between_shards() { +// crate::logger(); +// let _lock = lock(); + +// let cluster = configure_cluster(true).await; +// prepare_table(&cluster).await; + +// let mut client = new_client(); +// client +// .client_request +// .messages +// .push(Query::new("UPDATE sharded SET id = 5 WHERE id = 1").into()); + +// let mut engine = QueryEngine::from_client(&client).unwrap(); +// let mut context = QueryEngineContext::new(&mut client); + +// engine.handle(&mut context).await.unwrap(); + +// let (old_count, new_count) = table_state(&cluster).await; +// assert_eq!(old_count, 0, "old row must be removed"); +// assert_eq!( +// new_count, 1, +// "new row must be inserted on destination shard" +// ); + +// databases::shutdown(); +// config::load_test(); +// } + +// #[test] +// fn build_delete_sql_requires_where_clause() { +// let parsed = pgdog_plugin::pg_query::parse("UPDATE sharded SET id = 5") +// .expect("parse update without where"); +// let stmt = parsed +// .protobuf +// .stmts +// .first() +// .and_then(|node| node.stmt.as_ref()) +// .and_then(|node| node.node.as_ref()) +// .expect("statement node"); + +// let mut update_stmt = match stmt { +// NodeEnum::UpdateStmt(update) => (**update).clone(), +// _ => panic!("expected update statement"), +// }; + +// update_stmt.where_clause = None; + +// let plan = ShardKeyRewritePlan::new( +// OwnedTable { +// name: "sharded".into(), +// schema: None, +// alias: None, +// }, +// Route::write(ShardWithPriority::new_default_unset(Shard::Direct(0))), +// Some(1), +// update_stmt, +// vec![Assignment::new("id".into(), AssignmentValue::Integer(5))], +// ); + +// let err = build_delete_sql(&plan).expect_err("expected invariant error"); +// match err { +// Error::Router(router::Error::Parser(parser::Error::ShardKeyRewriteInvariant { +// reason, +// })) => { +// assert!( +// reason.contains("without WHERE clause"), +// "unexpected reason: {}", +// reason +// ); +// } +// other => panic!("unexpected error variant: {other:?}"), +// } +// } +// } diff --git a/pgdog/src/frontend/client/query_engine/show_shards.rs b/pgdog/src/frontend/client/query_engine/show_shards.rs deleted file mode 100644 index 3ca403bfc..000000000 --- a/pgdog/src/frontend/client/query_engine/show_shards.rs +++ /dev/null @@ -1,26 +0,0 @@ -use crate::net::{CommandComplete, DataRow, Field, Protocol, ReadyForQuery, RowDescription}; - -use super::*; - -impl QueryEngine { - /// SHOW pgdog.shards. - pub(super) async fn show_shards( - &mut self, - context: &mut QueryEngineContext<'_>, - shards: usize, - ) -> Result<(), Error> { - let bytes_sent = context - .stream - .send_many(&[ - RowDescription::new(&[Field::bigint("shards")]).message()?, - DataRow::from_columns(vec![shards]).message()?, - CommandComplete::from_str("SHOW").message()?, - ReadyForQuery::in_transaction(context.in_transaction()).message()?, - ]) - .await?; - - self.stats.sent(bytes_sent); - - Ok(()) - } -} diff --git a/pgdog/src/frontend/client/query_engine/test/mod.rs b/pgdog/src/frontend/client/query_engine/test/mod.rs new file mode 100644 index 000000000..ebeeb25fe --- /dev/null +++ b/pgdog/src/frontend/client/query_engine/test/mod.rs @@ -0,0 +1,32 @@ +use pgdog_config::General; + +use crate::{ + backend::databases::reload_from_existing, + config::{config, load_test, load_test_sharded, set}, + frontend::Client, + net::{Parameters, Stream}, +}; + +pub mod prelude; +mod rewrite_extended; +mod rewrite_insert_split; +mod rewrite_simple_prepared; +mod set; +mod set_schema_sharding; + +pub(super) fn test_client() -> Client { + load_test(); + Client::new_test(Stream::dev_null(), Parameters::default()) +} + +pub(super) fn test_sharded_client() -> Client { + load_test_sharded(); + Client::new_test(Stream::dev_null(), Parameters::default()) +} + +pub(super) fn change_config(f: impl FnOnce(&mut General)) { + let mut config = (*config()).clone(); + f(&mut config.config.general); + set(config).unwrap(); + reload_from_existing().unwrap(); +} diff --git a/pgdog/src/frontend/client/query_engine/test/prelude.rs b/pgdog/src/frontend/client/query_engine/test/prelude.rs new file mode 100644 index 000000000..4f616b686 --- /dev/null +++ b/pgdog/src/frontend/client/query_engine/test/prelude.rs @@ -0,0 +1,14 @@ +pub use crate::{ + frontend::{ + client::{ + query_engine::{QueryEngine, QueryEngineContext}, + test::TestClient, + Client, + }, + ClientRequest, + }, + net::{ + bind::Parameter, Bind, Execute, Flush, Parameters, Parse, Protocol, ProtocolMessage, Query, + Stream, Sync, + }, +}; diff --git a/pgdog/src/frontend/client/query_engine/test/rewrite_extended.rs b/pgdog/src/frontend/client/query_engine/test/rewrite_extended.rs new file mode 100644 index 000000000..5f154b925 --- /dev/null +++ b/pgdog/src/frontend/client/query_engine/test/rewrite_extended.rs @@ -0,0 +1,82 @@ +use super::prelude::*; +use super::{test_client, test_sharded_client}; + +async fn run_test(messages: Vec) -> Vec { + let mut results = vec![]; + + for mut client in [test_client(), test_sharded_client()] { + client.client_request = messages.clone().into(); + + let mut engine = QueryEngine::from_client(&client).unwrap(); + let mut context = QueryEngineContext::new(&mut client); + + engine.rewrite_extended(&mut context).unwrap(); + + results.push(client.client_request.messages.clone()); + } + + assert_eq!( + results[0], results[1], + "expected rewrite_extended to work the same for sharded and unsharded databases" + ); + + results.pop().unwrap() +} + +#[tokio::test] +async fn test_rewrite_prepared() { + let messages = run_test(vec![ + ProtocolMessage::from(Parse::named("__test_1", "SELECT $1, $2, $3")), + ProtocolMessage::from(Bind::new_statement("__test_1")), + ProtocolMessage::from(Execute::new()), + ProtocolMessage::from(Sync), + ]) + .await; + + assert!( + matches!(messages[0].clone(), ProtocolMessage::Parse(parse) if parse.name() == "__pgdog_1"), + "parse should of been renamed to __pgdog_1", + ); + + assert!( + matches!(messages[1].clone(), ProtocolMessage::Bind(bind) if bind.statement() =="__pgdog_1"), + "bind should of been renamed to __pgdog_1", + ); + + assert_eq!(messages.len(), 4); +} + +#[tokio::test] +async fn test_rewrite_extended() { + let messages = run_test(vec![ + ProtocolMessage::from(Parse::named("", "SELECT $1, $2, $3")), + ProtocolMessage::from(Bind::new_statement("")), + ProtocolMessage::from(Execute::new()), + ProtocolMessage::from(Sync), + ]) + .await; + + assert!( + matches!(messages[0].clone(), ProtocolMessage::Parse(parse) if parse.name() == ""), + "parse should not be renamed", + ); + + assert!( + matches!(messages[1].clone(), ProtocolMessage::Bind(bind) if bind.statement() ==""), + "bind should not be renamed", + ); + + assert_eq!(messages.len(), 4); +} + +#[tokio::test] +async fn test_rewrite_simple() { + let messages = run_test(vec![ProtocolMessage::from(Query::new("SELECT 1, 2, 3"))]).await; + + assert!( + matches!(messages[0].clone(), ProtocolMessage::Query(query) if query.query() == "SELECT 1, 2, 3"), + "query should not be altered", + ); + + assert_eq!(messages.len(), 1); +} diff --git a/pgdog/src/frontend/client/query_engine/test/rewrite_insert_split.rs b/pgdog/src/frontend/client/query_engine/test/rewrite_insert_split.rs new file mode 100644 index 000000000..db04585a5 --- /dev/null +++ b/pgdog/src/frontend/client/query_engine/test/rewrite_insert_split.rs @@ -0,0 +1,153 @@ +use crate::frontend::router::parser::rewrite::statement::plan::RewriteResult; + +use super::prelude::*; + +use super::{test_client, test_sharded_client}; + +async fn run_test(messages: Vec) -> Vec { + let mut client = test_sharded_client(); + + client.client_request = ClientRequest::from(messages); + + let mut engine = QueryEngine::from_client(&client).unwrap(); + let mut context = QueryEngineContext::new(&mut client); + + engine.rewrite_extended(&mut context).unwrap(); + engine.parse_and_rewrite(&mut context).await.unwrap(); + + assert!( + matches!(context.rewrite_result, Some(RewriteResult::InsertSplit(_))), + "expected rewrite insert split" + ); + + match context.rewrite_result.unwrap() { + RewriteResult::InsertSplit(requests) => requests, + _ => unreachable!(), + } +} + +#[tokio::test] +async fn test_insert_split() { + let requests = run_test(vec![ + ProtocolMessage::Parse(Parse::new_anonymous( + "INSERT INTO test (id, email) VALUES ($1, $2), ($3, $4)", + )), + ProtocolMessage::Bind(Bind::new_params( + "", + &[ + Parameter::new("1".as_bytes()), + Parameter::new("test@test.com".as_bytes()), + Parameter::new("1234567890102334".as_bytes()), + Parameter::new("test2@test.com".as_bytes()), + ], + )), + ProtocolMessage::Execute(Execute::new()), + ProtocolMessage::Sync(Sync), + ]) + .await; + + assert_eq!( + requests.len(), + 2, + "expected rewrite split to contain two requests" + ); + + for (request, (id, email)) in requests.iter().zip(vec![ + ("1".as_bytes(), "test@test.com".as_bytes()), + ("1234567890102334".as_bytes(), "test2@test.com".as_bytes()), + ]) { + assert!( + matches!(request[0].clone(), ProtocolMessage::Parse(parse) if parse.query() == "INSERT INTO test (id, email) VALUES ($1, $2)" && parse.anonymous()), + "expected single tuple insert with no name" + ); + match request[1].clone() { + ProtocolMessage::Bind(bind) => { + assert_eq!(bind.params_raw().get(0).unwrap().data, id); + assert_eq!(bind.params_raw().get(1).unwrap().data, email); + assert!(bind.anonymous()); + } + _ => panic!("expected bind"), + } + } +} +#[tokio::test] +async fn test_insert_split_prepared() { + let requests = run_test(vec![ + ProtocolMessage::Parse(Parse::named( + "__test_1", + "INSERT INTO test (id, email) VALUES ($1, $2), ($3, $4)", + )), + ProtocolMessage::Bind(Bind::new_params( + "__test_1", + &[ + Parameter::new("1".as_bytes()), + Parameter::new("test@test.com".as_bytes()), + Parameter::new("1234567890102334".as_bytes()), + Parameter::new("test2@test.com".as_bytes()), + ], + )), + ]) + .await; + + assert_eq!(requests.len(), 2); + + for (request, (id, email)) in requests.iter().zip(vec![ + ("1".as_bytes(), "test@test.com".as_bytes()), + ("1234567890102334".as_bytes(), "test2@test.com".as_bytes()), + ]) { + assert!( + matches!(request[0].clone(), ProtocolMessage::Parse(parse) if parse.query() == "INSERT INTO test (id, email) VALUES ($1, $2)" && parse.name() == "__pgdog_2"), + "expected single tuple insert" + ); + match request[1].clone() { + ProtocolMessage::Bind(bind) => { + assert_eq!(bind.params_raw().get(0).unwrap().data, id); + assert_eq!(bind.params_raw().get(1).unwrap().data, email); + assert_eq!(bind.statement(), "__pgdog_2"); + } + _ => panic!("expected bind"), + } + } +} + +#[tokio::test] +async fn test_insert_split_simple() { + let requests = run_test(vec![ProtocolMessage::Query(Query::new( + "INSERT INTO test (id, email) VALUES (1, 'test@test.com'), (2, 'test2@test.com') RETURNING *", + ))]) + .await; + + assert_eq!(requests.len(), 2); + + match requests[0][0].clone() { + ProtocolMessage::Query(query) => assert_eq!( + query.query(), + "INSERT INTO test (id, email) VALUES (1, 'test@test.com') RETURNING *" + ), + _ => panic!("not a query"), + } + + match requests[1][0].clone() { + ProtocolMessage::Query(query) => assert_eq!( + query.query(), + "INSERT INTO test (id, email) VALUES (2, 'test2@test.com') RETURNING *" + ), + _ => panic!("not a query"), + } +} + +#[tokio::test] +async fn test_insert_split_not_sharded() { + let mut client = test_client(); + client.client_request = ClientRequest::from(vec![ + ProtocolMessage::Parse(Parse::new_anonymous( + "INSERT INTO test (id, email) VALUES ($1, $2), ($3, $4)", + )), + ProtocolMessage::Other(Flush.message().unwrap()), + ]); + let mut engine = QueryEngine::from_client(&client).unwrap(); + let mut context = QueryEngineContext::new(&mut client); + engine.parse_and_rewrite(&mut context).await.unwrap(); + + assert!(context.rewrite_result.is_none()); +} diff --git a/pgdog/src/frontend/client/query_engine/test/rewrite_simple_prepared.rs b/pgdog/src/frontend/client/query_engine/test/rewrite_simple_prepared.rs new file mode 100644 index 000000000..9541aa21c --- /dev/null +++ b/pgdog/src/frontend/client/query_engine/test/rewrite_simple_prepared.rs @@ -0,0 +1,53 @@ +use pgdog_config::PreparedStatements; + +use crate::config::load_test; + +use super::change_config; +use super::prelude::*; + +async fn run_test(client: &mut Client, messages: &[ProtocolMessage]) -> Vec { + client.client_request = ClientRequest::from(messages.to_vec()); + let mut engine = QueryEngine::from_client(&client).unwrap(); + let mut context = QueryEngineContext::new(client); + + assert!(engine.parse_and_rewrite(&mut context).await.unwrap()); + + client.client_request.messages.clone() +} + +#[tokio::test] +async fn test_rewrite_prepare() { + load_test(); + + change_config(|general| { + general.prepared_statements = PreparedStatements::Full; + }); + + let mut client = Client::new_test(Stream::dev_null(), Parameters::default()); + + let messages = run_test( + &mut client, + &[Query::new("PREPARE __test_1 AS SELECT $1, $2, $3").into()], + ) + .await; + + assert!( + matches!(messages[0].clone(), ProtocolMessage::Query(query) if query.query() == "PREPARE __pgdog_1 AS SELECT $1, $2, $3"), + "expected rewritten prepared statement" + ); + + let messages = run_test( + &mut client, + &[Query::new("EXECUTE __test_1(1, 2, 3)").into()], + ) + .await; + + assert!( + matches!(messages[0].clone(), ProtocolMessage::Prepare { name, statement } if name == "__pgdog_1" && statement == "SELECT $1, $2, $3") + ); + + assert!( + matches!(messages[1].clone(), ProtocolMessage::Query(query) if query.query() == "EXECUTE __pgdog_1(1, 2, 3)"), + "expected rewritten prepared statement" + ); +} diff --git a/pgdog/src/frontend/client/query_engine/test/set.rs b/pgdog/src/frontend/client/query_engine/test/set.rs new file mode 100644 index 000000000..9f2645c98 --- /dev/null +++ b/pgdog/src/frontend/client/query_engine/test/set.rs @@ -0,0 +1,144 @@ +use crate::{ + expect_message, + net::{parameter::ParameterValue, CommandComplete, ReadyForQuery}, +}; + +use super::prelude::*; + +#[tokio::test] +async fn test_set() { + let mut test_client = TestClient::new_sharded(Parameters::default()).await; + + test_client + .send_simple(Query::new("SET application_name TO 'test_set'")) + .await; + + assert_eq!( + expect_message!(test_client.read().await, CommandComplete).command(), + "SET" + ); + assert_eq!( + expect_message!(test_client.read().await, ReadyForQuery).status, + 'I' + ); + + assert_eq!( + test_client.client().params.get("application_name").unwrap(), + &ParameterValue::String("test_set".into()), + ); +} + +#[tokio::test] +async fn test_set_search_path() { + let mut test_client = TestClient::new_sharded(Parameters::default()).await; + + test_client + .send_simple(Query::new( + "SET search_path TO \"$user\", public, acustomer", + )) + .await; + + assert_eq!( + expect_message!(test_client.read().await, CommandComplete).command(), + "SET" + ); + assert_eq!( + expect_message!(test_client.read().await, ReadyForQuery).status, + 'I' + ); + + assert_eq!( + test_client.client().params.get("search_path").unwrap(), + &ParameterValue::Tuple(vec!["$user".into(), "public".into(), "acustomer".into()]), + ); +} + +#[tokio::test] +async fn test_set_inside_transaction() { + let mut test_client = TestClient::new_sharded(Parameters::default()).await; + + test_client.send_simple(Query::new("BEGIN")).await; + + assert_eq!( + expect_message!(test_client.read().await, CommandComplete).command(), + "BEGIN" + ); + assert_eq!( + expect_message!(test_client.read().await, ReadyForQuery).status, + 'T' + ); + + test_client + .send_simple(Query::new("SET search_path TO acustomer, public")) + .await; + + assert_eq!( + expect_message!(test_client.read().await, CommandComplete).command(), + "SET" + ); + assert_eq!( + expect_message!(test_client.read().await, ReadyForQuery).status, + 'T' + ); + + test_client.send_simple(Query::new("COMMIT")).await; + + assert_eq!( + expect_message!(test_client.read().await, CommandComplete).command(), + "COMMIT" + ); + assert_eq!( + expect_message!(test_client.read().await, ReadyForQuery).status, + 'I' + ); + + assert_eq!( + test_client.client().params.get("search_path").unwrap(), + &ParameterValue::Tuple(vec!["acustomer".into(), "public".into()]), + ); +} + +#[tokio::test] +async fn test_set_inside_transaction_rollback() { + let mut test_client = TestClient::new_sharded(Parameters::default()).await; + + test_client.send_simple(Query::new("BEGIN")).await; + + assert_eq!( + expect_message!(test_client.read().await, CommandComplete).command(), + "BEGIN" + ); + assert_eq!( + expect_message!(test_client.read().await, ReadyForQuery).status, + 'T' + ); + + test_client + .send_simple(Query::new("SET search_path TO acustomer, public")) + .await; + + assert_eq!( + expect_message!(test_client.read().await, CommandComplete).command(), + "SET" + ); + assert_eq!( + expect_message!(test_client.read().await, ReadyForQuery).status, + 'T' + ); + + test_client.send_simple(Query::new("ROLLBACK")).await; + + assert_eq!( + expect_message!(test_client.read().await, CommandComplete).command(), + "ROLLBACK" + ); + assert_eq!( + expect_message!(test_client.read().await, ReadyForQuery).status, + 'I' + ); + + assert!( + test_client.client().params.get("search_path").is_none(), + "search_path should not be set", + ); +} diff --git a/pgdog/src/frontend/client/query_engine/test/set_schema_sharding.rs b/pgdog/src/frontend/client/query_engine/test/set_schema_sharding.rs new file mode 100644 index 000000000..671c7f977 --- /dev/null +++ b/pgdog/src/frontend/client/query_engine/test/set_schema_sharding.rs @@ -0,0 +1,39 @@ +use super::prelude::*; + +use crate::net::{DataRow, Parameter}; + +#[tokio::test] +async fn test_set_works_cross_shard_disabled() { + let mut client = TestClient::new_cross_shard_disabled(Parameters::from(vec![Parameter { + name: "search_path".into(), + value: "bcustomer".into(), + }])) + .await; + + client.send_simple(Query::new("BEGIN")).await; + let reply = client.read_until('Z').await.unwrap(); + assert_eq!(reply.len(), 2); + + client.send_simple(Query::new("SELECT 1")).await; + let reply = client.read_until('Z').await.unwrap(); + assert_eq!(reply.len(), 4); + + client + .send_simple(Query::new("SET statement_timeout TO '12345s'")) + .await; + let reply = client.read_until('Z').await.unwrap(); + assert_eq!(reply.len(), 2); + + client + .send_simple(Query::new("SHOW statement_timeout")) + .await; + let reply = client.read_until('Z').await.unwrap(); + assert_eq!(reply.len(), 4); + + let row = DataRow::try_from(reply[1].clone()).unwrap(); + assert_eq!(row.get_text(0).unwrap(), "12345s"); + + client.send_simple(Query::new("COMMIT")).await; + let reply = client.read_until('Z').await.unwrap(); + assert_eq!(reply.len(), 2); +} diff --git a/pgdog/src/frontend/client/query_engine/testing.rs b/pgdog/src/frontend/client/query_engine/testing.rs index ab0f0bb68..ee80f0a93 100644 --- a/pgdog/src/frontend/client/query_engine/testing.rs +++ b/pgdog/src/frontend/client/query_engine/testing.rs @@ -1,10 +1,6 @@ use super::*; impl QueryEngine { - pub fn backend(&mut self) -> &mut Connection { - &mut self.backend - } - pub fn router(&mut self) -> &mut Router { &mut self.router } @@ -12,4 +8,8 @@ impl QueryEngine { pub fn stats(&mut self) -> &mut Stats { &mut self.stats } + + pub fn set_test_mode(&mut self, test_mode: bool) { + self.test_mode.enabled = test_mode; + } } diff --git a/pgdog/src/frontend/client/query_engine/two_pc/manager.rs b/pgdog/src/frontend/client/query_engine/two_pc/manager.rs index 8c2062250..4dae8a587 100644 --- a/pgdog/src/frontend/client/query_engine/two_pc/manager.rs +++ b/pgdog/src/frontend/client/query_engine/two_pc/manager.rs @@ -23,7 +23,10 @@ use crate::{ two_pc::{TwoPcGuard, TwoPcTransaction}, TwoPcPhase, }, - router::{parser::Shard, Route}, + router::{ + parser::{Shard, ShardWithPriority}, + Route, + }, }, }; @@ -195,7 +198,10 @@ impl Manager { }; connection - .connect(&Request::default(), &Route::write(Shard::All)) + .connect( + &Request::default(), + &Route::write(ShardWithPriority::new_override_transaction(Shard::All)), + ) .await?; connection.two_pc(&transaction.to_string(), phase).await?; connection.disconnect(); diff --git a/pgdog/src/frontend/client/query_engine/two_pc/test.rs b/pgdog/src/frontend/client/query_engine/two_pc/test.rs index 2b98bb9d4..e98bbd878 100644 --- a/pgdog/src/frontend/client/query_engine/two_pc/test.rs +++ b/pgdog/src/frontend/client/query_engine/two_pc/test.rs @@ -4,7 +4,10 @@ use crate::{ pool::{Connection, Request}, }, config, - frontend::router::{parser::Shard, Route}, + frontend::router::{ + parser::{Shard, ShardWithPriority}, + Route, + }, logger, net::Protocol, }; @@ -20,9 +23,12 @@ async fn test_cleanup_transaction_phase_one() { let transaction = two_pc.transaction(); let mut conn = Connection::new(cluster.user(), cluster.name(), false, &None).unwrap(); - conn.connect(&Request::default(), &Route::write(Shard::All)) - .await - .unwrap(); + conn.connect( + &Request::default(), + &Route::write(ShardWithPriority::new_default_unset(Shard::All)), + ) + .await + .unwrap(); conn.execute("BEGIN").await.unwrap(); conn.execute("CREATE TABLE test_cleanup_transaction_phase_one(id BIGINT)") @@ -53,9 +59,12 @@ async fn test_cleanup_transaction_phase_one() { let transactions = Manager::get().transactions(); assert!(transactions.is_empty()); - conn.connect(&Request::default(), &Route::write(Shard::All)) - .await - .unwrap(); + conn.connect( + &Request::default(), + &Route::write(ShardWithPriority::new_default_unset(Shard::All)), + ) + .await + .unwrap(); let two_pc = conn .execute("SELECT * FROM pg_prepared_xacts") @@ -84,9 +93,12 @@ async fn test_cleanup_transaction_phase_two() { let transaction = two_pc.transaction(); let mut conn = Connection::new(cluster.user(), cluster.name(), false, &None).unwrap(); - conn.connect(&Request::default(), &Route::write(Shard::All)) - .await - .unwrap(); + conn.connect( + &Request::default(), + &Route::write(ShardWithPriority::new_default_unset(Shard::All)), + ) + .await + .unwrap(); conn.execute("BEGIN").await.unwrap(); conn.execute("CREATE TABLE test_cleanup_transaction_phase_two(id BIGINT)") @@ -122,9 +134,12 @@ async fn test_cleanup_transaction_phase_two() { let transactions = Manager::get().transactions(); assert!(transactions.is_empty()); - conn.connect(&Request::default(), &Route::write(Shard::All)) - .await - .unwrap(); + conn.connect( + &Request::default(), + &Route::write(ShardWithPriority::new_default_unset(Shard::All)), + ) + .await + .unwrap(); let two_pc = conn .execute("SELECT * FROM pg_prepared_xacts") diff --git a/pgdog/src/frontend/client/query_engine/two_pc/transaction.rs b/pgdog/src/frontend/client/query_engine/two_pc/transaction.rs index ae79e1b3c..f6f9042a6 100644 --- a/pgdog/src/frontend/client/query_engine/two_pc/transaction.rs +++ b/pgdog/src/frontend/client/query_engine/two_pc/transaction.rs @@ -1,4 +1,4 @@ -use rand::{thread_rng, Rng}; +use rand::{rng, Rng}; use std::fmt::Display; #[derive(Debug, Clone, Copy, PartialEq, Hash, Eq)] @@ -8,7 +8,7 @@ impl TwoPcTransaction { pub(crate) fn new() -> Self { // Transactions have random identifiers, // so multiple instances of PgDog don't create an identical transaction. - Self(thread_rng().gen()) + Self(rng().random_range(0..usize::MAX)) } } diff --git a/pgdog/src/frontend/client/sticky.rs b/pgdog/src/frontend/client/sticky.rs new file mode 100644 index 000000000..1b274e7a0 --- /dev/null +++ b/pgdog/src/frontend/client/sticky.rs @@ -0,0 +1,78 @@ +//! Sticky settings for clients that override +//! default routing behavior determined by the query parser. + +use pgdog_config::Role; +use rand::{rng, Rng}; + +use crate::net::{parameter::ParameterValue, Parameters}; + +#[derive(Debug, Clone, Copy)] +pub struct Sticky { + /// Which shard to use for omnisharded queries, making them + /// stick to only one database. + pub omni_index: usize, + + /// Desired database role. This comes from `target_session_attrs` + /// provided by the client. + pub role: Option, +} + +impl Default for Sticky { + fn default() -> Self { + Self::new() + } +} + +impl Sticky { + /// Create new sticky config. + pub fn new() -> Self { + Self::from_params(&Parameters::default()) + } + + #[cfg(test)] + pub fn new_test() -> Self { + Self { + omni_index: 1, + role: None, + } + } + + /// Create Sticky from params. + pub fn from_params(params: &Parameters) -> Self { + let role = params.get("pgdog.role").and_then(|value| match value { + ParameterValue::String(value) => match value.as_str() { + "primary" => Some(Role::Primary), + "replica" => Some(Role::Replica), + _ => None, + }, + _ => None, + }); + + Self { + omni_index: rng().random_range(1..usize::MAX), + role, + } + } +} + +#[cfg(test)] +mod test { + use super::*; + + #[test] + fn test_sticky() { + let params = Parameters::default(); + assert!(Sticky::from_params(¶ms).role.is_none()); + + for (attr, role) in [ + ("primary", Some(Role::Primary)), + ("replica", Some(Role::Replica)), + ("random", None), + ] { + let mut params = Parameters::default(); + params.insert("pgdog.role", attr); + let sticky = Sticky::from_params(¶ms); + assert_eq!(sticky.role, role); + } + } +} diff --git a/pgdog/src/frontend/client/test/mod.rs b/pgdog/src/frontend/client/test/mod.rs index d7bd15a5c..2ee4cc405 100644 --- a/pgdog/src/frontend/client/test/mod.rs +++ b/pgdog/src/frontend/client/test/mod.rs @@ -1,7 +1,8 @@ use std::time::{Duration, Instant}; +use pgdog_config::PoolerMode; use tokio::{ - io::{AsyncReadExt, AsyncWriteExt, BufStream}, + io::{AsyncRead, AsyncReadExt, AsyncWriteExt}, net::{TcpListener, TcpStream}, time::timeout, }; @@ -10,21 +11,28 @@ use bytes::{Buf, BufMut, BytesMut}; use crate::{ backend::databases::databases, - config::{config, load_test, load_test_replicas, set, Role}, + config::{ + config, load_test, load_test_replicas, load_test_sharded, load_test_with_pooler_mode, set, + PreparedStatements, Role, + }, frontend::{ client::{BufferEvent, QueryEngine}, - Client, + prepared_statements, Client, }, net::{ bind::Parameter, Bind, Close, CommandComplete, DataRow, Describe, ErrorResponse, Execute, - Field, Flush, Format, FromBytes, Parse, Protocol, Query, ReadyForQuery, RowDescription, - Sync, Terminate, ToBytes, + Field, Flush, Format, FromBytes, Message, Parameters, Parse, Protocol, Query, + ReadyForQuery, RowDescription, Sync, Terminate, ToBytes, }, state::State, }; use super::Stream; +pub mod target_session_attrs; +pub mod test_client; +pub use test_client::TestClient; + // // cargo nextest runs these in separate processes. // That's important otherwise I'm not sure what would happen. @@ -40,18 +48,37 @@ pub async fn test_client(replicas: bool) -> (TcpStream, Client) { parallel_test_client().await } +/// Load test client with 4 databases (2 shards, 1 replica each). +pub async fn test_client_sharded() -> (TcpStream, Client) { + load_test_sharded(); + + parallel_test_client().await +} + +pub async fn test_client_with_params(params: Parameters, replicas: bool) -> (TcpStream, Client) { + if replicas { + load_test_replicas(); + } else { + load_test(); + } + + parallel_test_client_with_params(params).await +} + pub async fn parallel_test_client() -> (TcpStream, Client) { + parallel_test_client_with_params(Parameters::default()).await +} + +pub async fn parallel_test_client_with_params(params: Parameters) -> (TcpStream, Client) { let addr = "127.0.0.1:0".to_string(); let conn_addr = addr.clone(); let stream = TcpListener::bind(&conn_addr).await.unwrap(); let port = stream.local_addr().unwrap().port(); let connect_handle = tokio::spawn(async move { - let (stream, addr) = stream.accept().await.unwrap(); - - let stream = BufStream::new(stream); - let stream = Stream::Plain(stream); + let (stream, _) = stream.accept().await.unwrap(); + let stream = Stream::plain(stream, 4096); - Client::new_test(stream, addr) + Client::new_test(stream, params) }); let conn = TcpStream::connect(&format!("127.0.0.1:{}", port)) @@ -72,6 +99,55 @@ macro_rules! new_client { }}; } +pub fn buffer(messages: &[impl ToBytes]) -> BytesMut { + let mut buf = BytesMut::new(); + for message in messages { + buf.put(message.to_bytes().unwrap()); + } + buf +} + +/// Read a series of messages from the stream and make sure +/// they arrive in the right order. +pub async fn read_messages(stream: &mut (impl AsyncRead + Unpin), codes: &[char]) -> Vec { + let mut result = vec![]; + let mut error = false; + + for code in codes { + let c = stream.read_u8().await.unwrap(); + + if c as char != *code { + if c as char == 'E' { + error = true; + } else { + panic!("got message code {}, expected {}", c as char, *code); + } + } + + let len = stream.read_i32().await.unwrap(); + let mut data = vec![0u8; len as usize - 4]; + stream.read_exact(&mut data).await.unwrap(); + + let mut message = BytesMut::new(); + message.put_u8(c); + message.put_i32(len); + message.put_slice(&data); + + let message = Message::new(message.freeze()); + + if error { + panic!( + "{:#}", + ErrorResponse::from_bytes(message.to_bytes().unwrap()).unwrap() + ); + } else { + result.push(message); + } + } + + result +} + macro_rules! buffer { ( $( $msg:block ),* ) => {{ let mut buf = BytesMut::new(); @@ -307,8 +383,8 @@ async fn test_client_with_replicas() { client.run().await.unwrap(); }); let buf = buffer!( - { Parse::new_anonymous("SELECT * FROM test_client_with_replicas") }, - { Bind::new_statement("") }, + { Parse::named("test", "SELECT * FROM test_client_with_replicas") }, + { Bind::new_statement("test") }, { Execute::new() }, { Sync } ); @@ -380,6 +456,7 @@ async fn test_client_with_replicas() { assert!(state.stats.counts.healthchecks <= idle + 1); // TODO: same pool_sent -= (healthcheck_len_sent * state.stats.counts.healthchecks) as isize; } + Role::Auto => unreachable!("role auto"), } } @@ -447,7 +524,6 @@ async fn test_transaction_state() { assert!(client.transaction.is_some()); assert!(engine.router().route().is_write()); - assert!(engine.router().in_transaction()); conn.write_all(&buffer!( { Parse::named("test", "SELECT $1") }, @@ -462,7 +538,6 @@ async fn test_transaction_state() { assert!(client.transaction.is_some()); assert!(engine.router().route().is_write()); - assert!(engine.router().in_transaction()); for c in ['1', 't', 'T', 'Z'] { let msg = engine.backend().read().await.unwrap(); @@ -479,7 +554,7 @@ async fn test_transaction_state() { "test", &[Parameter { len: 1, - data: "1".as_bytes().to_vec(), + data: "1".as_bytes().into(), }], ) }, @@ -503,7 +578,6 @@ async fn test_transaction_state() { assert!(client.transaction.is_some()); assert!(engine.router().route().is_write()); - assert!(engine.router().in_transaction()); conn.write_all(&buffer!({ Query::new("COMMIT") })) .await @@ -526,6 +600,8 @@ async fn test_transaction_state() { #[tokio::test] async fn test_close_parse() { + crate::logger(); + let (mut conn, mut client, mut engine) = new_client!(true); conn.write_all(&buffer!({ Close::named("test") }, { Sync })) @@ -700,3 +776,186 @@ async fn test_parse_describe_flush_bind_execute_close_sync() { handle.await.unwrap(); } + +#[tokio::test] +async fn test_client_login_timeout() { + use crate::config::config; + use tokio::time::sleep; + + crate::logger(); + load_test(); + + let mut config = (*config()).clone(); + config.config.general.client_login_timeout = 100; + set(config).unwrap(); + + let addr = "127.0.0.1:0".to_string(); + let stream = TcpListener::bind(&addr).await.unwrap(); + let port = stream.local_addr().unwrap().port(); + + let handle = tokio::spawn(async move { + let (stream, addr) = stream.accept().await.unwrap(); + let stream = Stream::plain(stream, 4096); + + let mut params = crate::net::parameter::Parameters::default(); + params.insert("user", "pgdog"); + params.insert("database", "pgdog"); + + Client::spawn(stream, params, addr, crate::config::config()).await + }); + + let conn = TcpStream::connect(&format!("127.0.0.1:{}", port)) + .await + .unwrap(); + + sleep(Duration::from_millis(150)).await; + + drop(conn); + + handle.await.unwrap().unwrap(); +} + +#[tokio::test] +async fn test_statement_mode() { + crate::logger(); + + load_test_with_pooler_mode(PoolerMode::Statement); + let (mut conn, mut client) = parallel_test_client().await; + + let _ = tokio::spawn(async move { + client.run().await.unwrap(); + }); + + let req = buffer!({ Query::new("BEGIN") }); + conn.write_all(&req).await.unwrap(); + + let msgs = read!(conn, ['E', 'Z']); + let error = ErrorResponse::from_bytes(msgs[0].clone().freeze()).unwrap(); + assert_eq!(error.code, "58000"); +} + +#[tokio::test] +async fn test_client_login_timeout_does_not_affect_queries() { + crate::logger(); + load_test(); + + let mut config = (*config()).clone(); + config.config.general.client_login_timeout = 100; + set(config).unwrap(); + + let (mut conn, mut client, _) = new_client!(false); + + let handle = tokio::spawn(async move { + client.run().await.unwrap(); + }); + + tokio::time::sleep(Duration::from_millis(150)).await; + + let buf = buffer!({ Query::new("SELECT pg_sleep(0.2)") }); + conn.write_all(&buf).await.unwrap(); + + let msgs = read!(conn, ['T', 'D', 'C', 'Z']); + assert_eq!(msgs[0][0] as char, 'T'); + assert_eq!(msgs[3][0] as char, 'Z'); + + conn.write_all(&buffer!({ Terminate })).await.unwrap(); + handle.await.unwrap(); +} + +#[tokio::test] +async fn test_anon_prepared_statements() { + crate::logger(); + load_test(); + + let (mut conn, mut client, _) = new_client!(false); + + let mut c = (*config()).clone(); + c.config.general.prepared_statements = PreparedStatements::ExtendedAnonymous; + set(c).unwrap(); + + let handle = tokio::spawn(async move { + client.run().await.unwrap(); + }); + + conn.write_all(&buffer!( + { Parse::new_anonymous("SELECT 1") }, + { Bind::new_params("", &[]) }, + { Execute::new() }, + { Sync } + )) + .await + .unwrap(); + + let _ = read!(conn, ['1', '2', 'D', 'C', 'Z']); + + { + let cache = prepared_statements::PreparedStatements::global(); + let read = cache.read(); + assert!(!read.is_empty()); + } + + conn.write_all(&buffer!({ Terminate })).await.unwrap(); + handle.await.unwrap(); +} + +#[tokio::test] +async fn test_anon_prepared_statements_extended() { + crate::logger(); + load_test(); + + let (mut conn, mut client, _) = new_client!(false); + + let mut c = (*config()).clone(); + c.config.general.prepared_statements = PreparedStatements::Extended; + set(c).unwrap(); + + let handle = tokio::spawn(async move { + client.run().await.unwrap(); + }); + + conn.write_all(&buffer!( + { Parse::new_anonymous("SELECT 1") }, + { Bind::new_params("", &[]) }, + { Execute::new() }, + { Sync } + )) + .await + .unwrap(); + + let _ = read!(conn, ['1', '2', 'D', 'C', 'Z']); + + { + let cache = prepared_statements::PreparedStatements::global(); + let read = cache.read(); + assert!(read.is_empty()); + } + + conn.write_all(&buffer!({ Terminate })).await.unwrap(); + handle.await.unwrap(); +} + +#[tokio::test] +async fn test_query_timeout() { + crate::logger(); + load_test(); + + let (mut conn, mut client, mut engine) = new_client!(false); + + let mut c = (*config()).clone(); + c.config.general.query_timeout = 50; + set(c).unwrap(); + + engine.set_test_mode(false); + + let buf = buffer!({ Query::new("SELECT pg_sleep(0.2)") }); + conn.write_all(&buf).await.unwrap(); + + client.buffer(State::Idle).await.unwrap(); + let result = client.client_messages(&mut engine).await; + + assert!(result.is_err()); + + let pools = databases().cluster(("pgdog", "pgdog")).unwrap().shards()[0].pools(); + let state = pools[0].state(); + assert_eq!(state.force_close, 1); +} diff --git a/pgdog/src/frontend/client/test/target_session_attrs.rs b/pgdog/src/frontend/client/test/target_session_attrs.rs new file mode 100644 index 000000000..471ff8c50 --- /dev/null +++ b/pgdog/src/frontend/client/test/target_session_attrs.rs @@ -0,0 +1,58 @@ +use super::*; + +async fn run_target_session_test(property: &str, query: &str) -> Message { + let mut params = Parameters::default(); + params.insert("pgdog.role", property); + + let (mut stream, mut client) = test_client_with_params(params, true).await; + let mut engine = QueryEngine::from_client(&client).unwrap(); + + let expected = if property == "primary" { + Role::Primary + } else if property == "replica" { + Role::Replica + } else { + panic!("unexpected property: {}", property); + }; + assert_eq!(client.sticky.role, Some(expected)); + + stream + .write_all(&Query::new(query).to_bytes().unwrap()) + .await + .unwrap(); + stream.flush().await.unwrap(); + + client.buffer(State::Idle).await.unwrap(); + client.client_messages(&mut engine).await.unwrap(); + + let reply = engine.backend().read().await.unwrap(); + + reply +} + +#[tokio::test] +async fn test_target_session_attrs_standby() { + let reply = run_target_session_test( + "replica", + "CREATE TABLE test_target_session_attrs_standby(id BIGINT)", + ) + .await; + assert_eq!(reply.code(), 'E'); + let error = ErrorResponse::from_bytes(reply.to_bytes().unwrap()).unwrap(); + assert_eq!( + error.message, + "cannot execute CREATE TABLE in a read-only transaction" + ); +} + +#[tokio::test] +async fn test_target_session_attrs_primary() { + for _ in 0..5 { + let reply = run_target_session_test( + "primary", + "CREATE TABLE IF NOT EXISTS test_target_session_attrs_primary(id BIGINT)", + ) + .await; + assert_ne!(reply.code(), 'E'); + } +} diff --git a/pgdog/src/frontend/client/test/test_client.rs b/pgdog/src/frontend/client/test/test_client.rs new file mode 100644 index 000000000..c828e5053 --- /dev/null +++ b/pgdog/src/frontend/client/test/test_client.rs @@ -0,0 +1,222 @@ +use std::{fmt::Debug, ops::Deref}; + +use bytes::{BufMut, Bytes, BytesMut}; +use pgdog_config::RewriteMode; +use rand::{rng, Rng}; +use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + net::{TcpListener, TcpStream}, +}; + +use crate::{ + backend::databases::{reload_from_existing, shutdown}, + config::{config, load_test_replicas, load_test_sharded, set}, + frontend::{ + client::query_engine::QueryEngine, + router::{parser::Shard, sharding::ContextBuilder}, + Client, + }, + net::{ErrorResponse, Message, Parameters, Protocol, Stream}, +}; + +/// Try to convert a Message to the specified type. +/// If conversion fails and the message is an ErrorResponse, panic with its contents. +#[cfg(test)] +#[macro_export] +macro_rules! expect_message { + ($message:expr, $ty:ty) => {{ + use crate::net::Protocol; + let message: crate::net::Message = $message; + match <$ty as TryFrom>::try_from(message.clone()) { + Ok(val) => val, + Err(_) => { + match >::try_from( + message.clone(), + ) { + Ok(err) => panic!("expected {}, got ErrorResponse: {:?}", stringify!($ty), err), + Err(_) => panic!( + "expected {}, got message with code '{}'", + stringify!($ty), + message.code() + ), + } + } + } + }}; +} + +/// Test client. +#[derive(Debug)] +pub struct TestClient { + pub(crate) client: Client, + pub(crate) engine: QueryEngine, + pub(crate) conn: TcpStream, +} + +impl TestClient { + /// Create new test client after the login phase + /// is complete. + /// + /// Config needs to be loaded. + /// + async fn new(params: Parameters) -> Self { + let addr = "127.0.0.1:0".to_string(); + let conn_addr = addr.clone(); + let stream = TcpListener::bind(&conn_addr).await.unwrap(); + let port = stream.local_addr().unwrap().port(); + let connect_handle = tokio::spawn(async move { + let (stream, _) = stream.accept().await.unwrap(); + let stream = Stream::plain(stream, 4096); + + Client::new_test(stream, params) + }); + + let conn = TcpStream::connect(&format!("127.0.0.1:{}", port)) + .await + .unwrap(); + let client = connect_handle.await.unwrap(); + + Self { + conn, + engine: QueryEngine::from_client(&client).expect("create query engine from client"), + client, + } + } + + /// New sharded client with parameters. + pub(crate) async fn new_sharded(params: Parameters) -> Self { + load_test_sharded(); + Self::new(params).await + } + + /// New client with replicas but not sharded. + #[allow(dead_code)] + pub(crate) async fn new_replicas(params: Parameters) -> Self { + load_test_replicas(); + Self::new(params).await + } + + /// New client with cross-shard-queries disabled. + pub(crate) async fn new_cross_shard_disabled(params: Parameters) -> Self { + load_test_sharded(); + + let mut config = config().deref().clone(); + config.config.general.cross_shard_disabled = true; + set(config).unwrap(); + reload_from_existing().unwrap(); + + Self::new(params).await + } + + /// Create client that will rewrite all queries. + pub(crate) async fn new_rewrites(params: Parameters) -> Self { + load_test_sharded(); + + let mut config = config().deref().clone(); + config.config.rewrite.enabled = true; + config.config.rewrite.shard_key = RewriteMode::Rewrite; + config.config.rewrite.split_inserts = RewriteMode::Rewrite; + + set(config).unwrap(); + reload_from_existing().unwrap(); + + Self::new(params).await + } + + /// Send message to client. + pub(crate) async fn send(&mut self, message: impl Protocol) { + let message = message.to_bytes().expect("message to convert to bytes"); + self.conn.write_all(&message).await.expect("write_all"); + self.conn.flush().await.expect("flush"); + } + + /// Send a simple query and panic on any errors. + pub(crate) async fn send_simple(&mut self, message: impl Protocol) { + self.try_send_simple(message).await.unwrap() + } + + /// Try to send a simple query and return the error, if any. + pub(crate) async fn try_send_simple( + &mut self, + message: impl Protocol, + ) -> Result<(), Box> { + self.send(message).await; + self.try_process().await + } + + /// Read a message received from the servers. + pub(crate) async fn read(&mut self) -> Message { + let code = self.conn.read_u8().await.expect("code"); + let len = self.conn.read_i32().await.expect("len"); + let mut rest = vec![0u8; len as usize - 4]; + self.conn.read_exact(&mut rest).await.expect("read_exact"); + + let mut payload = BytesMut::new(); + payload.put_u8(code); + payload.put_i32(len); + payload.put(Bytes::from(rest)); + + Message::new(payload.freeze()).backend() + } + + /// Inspect client state. + pub(crate) fn client(&mut self) -> &mut Client { + &mut self.client + } + + /// Process a request. + pub(crate) async fn try_process(&mut self) -> Result<(), Box> { + self.engine.set_test_mode(false); + self.client.buffer(self.engine.stats().state).await?; + self.client.client_messages(&mut self.engine).await?; + self.engine.set_test_mode(true); + + Ok(()) + } + + /// Read all messages until an expected last message. + pub(crate) async fn read_until(&mut self, code: char) -> Result, ErrorResponse> { + let mut result = vec![]; + loop { + let message = self.read().await; + result.push(message.clone()); + + if message.code() == code { + break; + } + + if message.code() == 'E' && code != 'E' { + let error = ErrorResponse::try_from(message).unwrap(); + return Err(error); + } + } + + Ok(result) + } + + /// Generate a random ID for a given shard. + pub(crate) fn random_id_for_shard(&mut self, shard: usize) -> i64 { + let cluster = self.engine.backend().cluster().unwrap().clone(); + + loop { + let id: i64 = rng().random(); + let calc = ContextBuilder::new(cluster.sharded_tables().first().unwrap()) + .data(id) + .shards(cluster.shards().len()) + .build() + .unwrap() + .apply() + .unwrap(); + + if calc == Shard::Direct(shard) { + return id; + } + } + } +} + +impl Drop for TestClient { + fn drop(&mut self) { + shutdown(); + } +} diff --git a/pgdog/src/frontend/client/timeouts.rs b/pgdog/src/frontend/client/timeouts.rs index dea4962f2..59896f0fe 100644 --- a/pgdog/src/frontend/client/timeouts.rs +++ b/pgdog/src/frontend/client/timeouts.rs @@ -6,6 +6,7 @@ use crate::{config::General, frontend::ClientRequest, state::State}; pub struct Timeouts { pub(super) query_timeout: Duration, pub(super) client_idle_timeout: Duration, + pub(super) idle_in_transaction_timeout: Duration, } impl Default for Timeouts { @@ -13,6 +14,7 @@ impl Default for Timeouts { Self { query_timeout: Duration::MAX, client_idle_timeout: Duration::MAX, + idle_in_transaction_timeout: Duration::MAX, } } } @@ -22,6 +24,7 @@ impl Timeouts { Self { query_timeout: general.query_timeout(), client_idle_timeout: general.client_idle_timeout(), + idle_in_transaction_timeout: general.client_idle_in_transaction_timeout(), } } @@ -48,7 +51,40 @@ impl Timeouts { Duration::MAX } } + State::IdleInTransaction => { + // Client is sending the request, don't fire. + if !client_request.messages.is_empty() { + Duration::MAX + } else { + self.idle_in_transaction_timeout + } + } + _ => Duration::MAX, } } } + +#[cfg(test)] +mod test { + + use crate::{config::config, net::Query}; + + use super::*; + + #[test] + fn test_idle_in_transaction_timeout() { + let config = config(); // Will be default. + let timeout = Timeouts::from_config(&config.config.general); + + let actual = timeout.client_idle_timeout(&State::IdleInTransaction, &ClientRequest::new()); + assert_eq!(actual, timeout.idle_in_transaction_timeout); + assert_eq!(actual.as_millis(), u64::MAX.into()); + + let actual = timeout.client_idle_timeout( + &State::IdleInTransaction, + &ClientRequest::from(vec![Query::new("SELECT 1").into()]), + ); + assert_eq!(actual, Duration::MAX); + } +} diff --git a/pgdog/src/frontend/client_request.rs b/pgdog/src/frontend/client_request.rs index 69051a84c..19e580f7d 100644 --- a/pgdog/src/frontend/client_request.rs +++ b/pgdog/src/frontend/client_request.rs @@ -1,12 +1,16 @@ //! ClientRequest (messages buffer). +//! +//! Contains exactly one request. +//! use std::ops::{Deref, DerefMut}; use lazy_static::lazy_static; +use regex::Regex; use crate::{ - frontend::router::parser::RewritePlan, + frontend::router::Ast, net::{ - messages::{Bind, CopyData, Protocol, Query}, + messages::{Bind, CopyData, Protocol}, Error, Flush, ProtocolMessage, }, stats::memory::MemoryUsage, @@ -16,11 +20,17 @@ use super::{router::Route, PreparedStatements}; pub use super::BufferedQuery; -/// Message buffer. +/// Client request, containing exactly one query. #[derive(Debug, Clone)] pub struct ClientRequest { + /// Messages, e.g. Query, or Parse, Bind, Execute, etc. pub messages: Vec, + /// The route this request will take in our query engine. + /// When the request is created, this is not known yet. + /// The QueryEngine will set the route once it handles the request. pub route: Option, + /// The statement AST, if we parsed the request with our query parser. + pub ast: Option, } impl MemoryUsage for ClientRequest { @@ -28,6 +38,7 @@ impl MemoryUsage for ClientRequest { fn memory_usage(&self) -> usize { // ProtocolMessage uses memory allocated by BytesMut (mostly). self.messages.capacity() * std::mem::size_of::() + + std::mem::size_of::>() } } @@ -43,12 +54,20 @@ impl ClientRequest { Self { messages: Vec::with_capacity(5), route: None, + ast: None, } } - /// The buffer is full and the client won't send any more messages - /// until it gets a reply, or we don't want to buffer the data in memory. - pub fn full(&self) -> bool { + /// Remove any saved state from the request. + pub fn clear(&mut self) { + self.messages.clear(); + self.route = None; + self.ast = None; + } + + /// We received a complete request and we are ready to + /// send it to the query engine. + pub fn is_complete(&self) -> bool { if let Some(message) = self.messages.last() { // Flush (F) | Sync (F) | Query (F) | CopyDone (F) | CopyFail (F) if matches!(message.code(), 'H' | 'S' | 'Q' | 'c' | 'f') { @@ -108,6 +127,28 @@ impl ClientRequest { Ok(None) } + /// Quick and cheap way to find out if this + /// request is starting a transaction. + pub fn is_begin(&self) -> bool { + lazy_static! { + static ref BEGIN: Regex = Regex::new("(?i)^BEGIN").unwrap(); + } + + for message in &self.messages { + let query = match message { + ProtocolMessage::Parse(parse) => parse.query(), + ProtocolMessage::Query(query) => query.query(), + _ => continue, + }; + + if BEGIN.is_match(query) { + return true; + } + } + + false + } + /// If this buffer contains bound parameters, retrieve them. pub fn parameters(&self) -> Result, Error> { for message in &self.messages { @@ -139,11 +180,12 @@ impl ClientRequest { Self { messages, route: self.route.clone(), + ast: self.ast.clone(), } } /// The buffer has COPY messages. - pub fn copy(&self) -> bool { + pub fn is_copy(&self) -> bool { self.messages .last() .map(|m| m.code() == 'd' || m.code() == 'c') @@ -152,46 +194,22 @@ impl ClientRequest { /// The client is setting state on the connection /// which we can no longer ignore. - pub(crate) fn executable(&self) -> bool { + pub(crate) fn is_executable(&self) -> bool { self.messages .iter() .any(|m| ['E', 'Q', 'B'].contains(&m.code())) } /// Rewrite query in buffer. - pub fn rewrite(&mut self, query: &str) -> Result<(), Error> { + pub fn rewrite(&mut self, request: &[ProtocolMessage]) -> Result<(), Error> { if self.messages.iter().any(|c| c.code() != 'Q') { return Err(Error::OnlySimpleForRewrites); } self.messages.clear(); - self.messages.push(Query::new(query).into()); + self.messages.extend(request.to_vec()); Ok(()) } - /// Rewrite prepared statement SQL before sending it to the backend. - pub fn rewrite_prepared( - &mut self, - query: &str, - prepared: &mut PreparedStatements, - plan: &RewritePlan, - ) -> bool { - let mut updated = false; - - for message in self.messages.iter_mut() { - if let ProtocolMessage::Parse(parse) = message { - parse.set_query(query); - let name = parse.name().to_owned(); - let _ = prepared.update_query(&name, query); - if !plan.is_noop() { - prepared.set_rewrite_plan(&name, plan.clone()); - } - updated = true; - } - } - - updated - } - /// Get the route for this client request. pub fn route(&self) -> &Route { lazy_static! { @@ -279,6 +297,7 @@ impl From> for ClientRequest { ClientRequest { messages, route: None, + ast: None, } } } @@ -299,7 +318,7 @@ impl DerefMut for ClientRequest { #[cfg(test)] mod test { - use crate::net::{Describe, Execute, Parse, Sync}; + use crate::net::{Describe, Execute, Parse, Query, Sync}; use super::*; @@ -466,4 +485,16 @@ mod test { assert_eq!(second_slice[2].code(), 'H'); // Flush assert_eq!(second_slice[3].code(), 'S'); // Sync } + + #[test] + fn test_detect_begin() { + for query in [ + ProtocolMessage::Query(Query::new("begin")), + ProtocolMessage::Query(Query::new("BEGIN WORK REPEATABLE READ")), + ProtocolMessage::Parse(Parse::new_anonymous("BEGIN")), + ] { + let req = ClientRequest::from(vec![query]); + assert!(req.is_begin()); + } + } } diff --git a/pgdog/src/frontend/comms.rs b/pgdog/src/frontend/comms.rs index b9ff153e8..2887e2d16 100644 --- a/pgdog/src/frontend/comms.rs +++ b/pgdog/src/frontend/comms.rs @@ -1,6 +1,7 @@ //! Communication to/from connected clients. use std::net::SocketAddr; +use std::ops::Deref; use std::sync::{ atomic::{AtomicBool, Ordering}, Arc, @@ -40,7 +41,6 @@ struct Global { #[derive(Clone, Debug)] pub struct Comms { global: Arc, - id: Option, } impl Default for Comms { @@ -59,7 +59,6 @@ impl Comms { clients: Mutex::new(HashMap::default()), tracker: TaskTracker::new(), }), - id: None, } } @@ -73,15 +72,6 @@ impl Comms { self.global.clients.lock().len() } - pub fn clients_memory(&self) -> usize { - self.global - .clients - .lock() - .values() - .map(|v| v.stats.memory_used) - .sum::() - } - pub fn tracker(&self) -> &TaskTracker { &self.global.tracker } @@ -97,39 +87,31 @@ impl Comms { } /// New client connected. - pub fn connect(&mut self, id: &BackendKeyData, addr: SocketAddr, params: &Parameters) -> Self { + pub fn connect(&self, id: &BackendKeyData, addr: SocketAddr, params: &Parameters) { self.global .clients .lock() - .insert(*id, ConnectedClient::new(addr, params)); - self.id = Some(*id); - self.clone() + .insert(*id, ConnectedClient::new(id, addr, params)); } /// Update client parameters. - pub fn update_params(&self, params: &Parameters) { - if let Some(id) = self.id { - let mut guard = self.global.clients.lock(); - if let Some(entry) = guard.get_mut(&id) { - entry.paramters = params.clone(); - } + pub fn update_params(&self, id: &BackendKeyData, params: Parameters) { + let mut guard = self.global.clients.lock(); + if let Some(entry) = guard.get_mut(id) { + entry.paramters = params; } } /// Client disconnected. - pub fn disconnect(&mut self) { - if let Some(id) = self.id.take() { - self.global.clients.lock().remove(&id); - } + pub fn disconnect(&self, id: &BackendKeyData) { + self.global.clients.lock().remove(id); } /// Update stats. - pub fn stats(&self, stats: Stats) { - if let Some(ref id) = self.id { - let mut guard = self.global.clients.lock(); - if let Some(entry) = guard.get_mut(id) { - entry.stats = stats; - } + pub fn update_stats(&self, id: &BackendKeyData, stats: Stats) { + let mut guard = self.global.clients.lock(); + if let Some(entry) = guard.get_mut(id) { + entry.stats = stats; } } @@ -149,8 +131,43 @@ impl Comms { pub fn offline(&self) -> bool { self.global.offline.load(Ordering::Relaxed) } +} + +#[derive(Debug, Clone)] +pub struct ClientComms { + comms: Comms, + id: BackendKeyData, +} - pub fn client_id(&self) -> BackendKeyData { - self.id.unwrap_or_default() +impl Deref for ClientComms { + type Target = Comms; + + fn deref(&self) -> &Self::Target { + &self.comms + } +} + +impl ClientComms { + pub fn disconnect(&self) { + self.comms.disconnect(&self.id); + } + + pub fn update_stats(&self, stats: Stats) { + self.comms.update_stats(&self.id, stats); + } + + pub fn new(id: &BackendKeyData) -> Self { + Self { + id: *id, + comms: comms(), + } + } + + pub fn connect(&self, addr: SocketAddr, params: &Parameters) { + self.comms.connect(&self.id, addr, params) + } + + pub fn update_params(&self, params: &Parameters) { + self.comms.update_params(&self.id, params.clone()); } } diff --git a/pgdog/src/frontend/connected_client.rs b/pgdog/src/frontend/connected_client.rs index ce77d346f..dd7318e32 100644 --- a/pgdog/src/frontend/connected_client.rs +++ b/pgdog/src/frontend/connected_client.rs @@ -1,7 +1,7 @@ use chrono::{DateTime, Local}; use std::net::SocketAddr; -use crate::net::Parameters; +use crate::net::{BackendKeyData, Parameters}; use super::Stats; @@ -16,12 +16,15 @@ pub struct ConnectedClient { pub connected_at: DateTime, /// Client connection parameters. pub paramters: Parameters, + /// Identifier. + pub id: BackendKeyData, } impl ConnectedClient { /// New connected client. - pub fn new(addr: SocketAddr, params: &Parameters) -> Self { + pub fn new(id: &BackendKeyData, addr: SocketAddr, params: &Parameters) -> Self { Self { + id: *id, stats: Stats::new(), addr, connected_at: Local::now(), diff --git a/pgdog/src/frontend/error.rs b/pgdog/src/frontend/error.rs index fbbc4b8e4..bfa1cac74 100644 --- a/pgdog/src/frontend/error.rs +++ b/pgdog/src/frontend/error.rs @@ -4,6 +4,8 @@ use std::io::ErrorKind; use thiserror::Error; +use crate::unique_id; + /// Frontend error. #[derive(Debug, Error)] pub enum Error { @@ -37,7 +39,7 @@ pub enum Error { #[error("{0}")] PreparedStatements(#[from] super::prepared_statements::Error), - #[error("prepared staatement \"{0}\" is missing")] + #[error("prepared statement \"{0}\" is missing")] MissingPreparedStatement(String), #[error("query timeout")] @@ -45,6 +47,29 @@ pub enum Error { #[error("join error")] Join(#[from] tokio::task::JoinError), + + #[error("unique id: {0}")] + UniqueId(#[from] unique_id::Error), + + #[error("parser: {0}")] + Parser(#[from] crate::frontend::router::parser::Error), + + #[error("rewrite: {0}")] + Rewrite(#[from] crate::frontend::router::parser::rewrite::statement::Error), + + #[error("query has no route")] + NoRoute, + + #[error("multi-tuple insert requires multi-shard binding")] + MultiShardRequired, + + #[error("sharding key updates are forbidden")] + ShardingKeyUpdateForbidden, + + // FIXME: layer errors better so we don't have + // to reach so deep into a module. + #[error("{0}")] + Multi(#[from] crate::frontend::client::query_engine::multi_step::error::Error), } impl Error { diff --git a/pgdog/src/frontend/listener.rs b/pgdog/src/frontend/listener.rs index 2bb907eee..250ec6533 100644 --- a/pgdog/src/frontend/listener.rs +++ b/pgdog/src/frontend/listener.rs @@ -20,10 +20,7 @@ use tokio::{select, spawn}; use tracing::{error, info, warn}; -use super::{ - comms::{comms, Comms}, - Client, Error, -}; +use super::{comms::comms, Client, Error}; /// Client connections listener and handler. #[derive(Debug, Clone)] @@ -45,21 +42,18 @@ impl Listener { pub async fn listen(&mut self) -> Result<(), Error> { info!("πŸ• PgDog listening on {}", self.addr); let listener = TcpListener::bind(&self.addr).await?; - let comms = comms(); - let shutdown_signal = comms.shutting_down(); + let shutdown_signal = comms().shutting_down(); let mut sighup = Sighup::new()?; loop { - let comms = comms.clone(); - select! { connection = listener.accept() => { + let comms = comms(); let (stream, addr) = connection?; let offline = comms.offline(); - let client_comms = comms.clone(); let future = async move { - match Self::handle_client(stream, addr, client_comms).await { + match Self::handle_client(stream, addr).await { Ok(_) => (), Err(err) => if !err.disconnect() { error!("client crashed: {:?}", err); @@ -117,8 +111,9 @@ impl Listener { let shutdown_timeout = config().config.general.shutdown_timeout(); info!( - "waiting up to {:.3}s for clients to finish transactions", - shutdown_timeout.as_secs_f64() + "waiting up to {:.3}s for {} clients to finish transactions", + shutdown_timeout.as_secs_f64(), + comms().tracker().len(), ); let comms = comms(); @@ -131,15 +126,39 @@ impl Listener { "terminating {} client connections due to shutdown timeout", comms.tracker().len() ); + + // If a shutdown termination timeout is configured, enforce it here. + // This will ensure that we don't wait indefinitely for databases to respond. + if let Some(termination_timeout) = + config().config.general.shutdown_termination_timeout() + { + // Shutdown timeout elapsed; cancel any still-running queries before tearing pools down. + let cancel_futures = comms.clients().into_keys().map(|id| async move { + if let Err(err) = databases().cancel(&id).await { + error!(?id, "cancel request failed during shutdown: {err}"); + } + }); + let cancel_all = futures::future::join_all(cancel_futures); + + if timeout(termination_timeout, cancel_all).await.is_err() { + error!( + "forced shutdown: abandoning {} outstanding cancel requests after waiting {:.3}s" , + comms.clients().len(), + termination_timeout.as_secs_f64() + ); + } + } } self.shutdown.notify_waiters(); } - async fn handle_client(stream: TcpStream, addr: SocketAddr, comms: Comms) -> Result<(), Error> { + async fn handle_client(stream: TcpStream, addr: SocketAddr) -> Result<(), Error> { tweak(&stream)?; + let config = config(); + + let mut stream = Stream::plain(stream, config.config.memory.net_buffer); - let mut stream = Stream::plain(stream); let tls = acceptor(); loop { @@ -159,11 +178,14 @@ impl Listener { match startup { Startup::Ssl => { - if let Some(tls) = tls { + if let Some(tls) = tls.as_ref() { stream.send_flush(&SslReply::Yes).await?; let plain = stream.take()?; let cipher = tls.accept(plain).await?; - stream = Stream::tls(tokio_rustls::TlsStream::Server(cipher)); + stream = Stream::tls( + tokio_rustls::TlsStream::Server(cipher), + config.config.memory.net_buffer, + ); } else { stream.send_flush(&SslReply::No).await?; } @@ -175,7 +197,7 @@ impl Listener { } Startup::Startup { params } => { - Client::spawn(stream, params, addr, comms).await?; + Client::spawn(stream, params, addr, config).await?; break; } diff --git a/pgdog/src/frontend/mod.rs b/pgdog/src/frontend/mod.rs index bff961d8c..1dbd952e4 100644 --- a/pgdog/src/frontend/mod.rs +++ b/pgdog/src/frontend/mod.rs @@ -18,9 +18,9 @@ pub mod stats; pub use buffered_query::BufferedQuery; pub use client::Client; pub use client_request::ClientRequest; -pub use comms::Comms; +pub use comms::{ClientComms, Comms}; pub use connected_client::ConnectedClient; -pub use error::Error; +pub(crate) use error::Error; pub use prepared_statements::{PreparedStatements, Rewrite}; #[cfg(debug_assertions)] pub use query_logger::QueryLogger; diff --git a/pgdog/src/frontend/prepared_statements/global_cache.rs b/pgdog/src/frontend/prepared_statements/global_cache.rs index e2e545810..76ee0331f 100644 --- a/pgdog/src/frontend/prepared_statements/global_cache.rs +++ b/pgdog/src/frontend/prepared_statements/global_cache.rs @@ -1,24 +1,28 @@ use bytes::Bytes; use crate::{ - frontend::router::parser::RewritePlan, net::messages::{Parse, RowDescription}, stats::memory::MemoryUsage, }; use std::{collections::hash_map::HashMap, str::from_utf8}; +use fnv::FnvHashSet as HashSet; + // Format the globally unique prepared statement // name based on the counter. fn global_name(counter: usize) -> String { format!("__pgdog_{}", counter) } -#[derive(Debug, Clone)] +#[derive(Debug, Clone, Default)] pub struct Statement { parse: Parse, + rewrite: Option, row_description: Option, + #[allow(dead_code)] version: usize, - rewrite_plan: Option, + cache_key: CacheKey, + evict_on_close: bool, } impl MemoryUsage for Statement { @@ -30,8 +34,8 @@ impl MemoryUsage for Statement { } else { 0 } - // Rewrite plans are small; treat as zero-cost for now. - + 0 + + self.cache_key.memory_usage() + + self.evict_on_close.memory_usage() } } @@ -41,11 +45,7 @@ impl Statement { } fn cache_key(&self) -> CacheKey { - CacheKey { - query: self.parse.query_ref(), - data_types: self.parse.data_types_ref(), - version: self.version, - } + self.cache_key.clone() } } @@ -56,7 +56,7 @@ impl Statement { /// with different data types, we can't re-use it and /// need to plan a new one. /// -#[derive(Debug, Clone, PartialEq, Hash, Eq)] +#[derive(Debug, Clone, PartialEq, Hash, Eq, Default)] pub struct CacheKey { pub query: Bytes, pub data_types: Bytes, @@ -112,6 +112,7 @@ impl CachedStmt { pub struct GlobalCache { statements: HashMap, names: HashMap, + unused: HashSet, counter: usize, versions: usize, } @@ -123,6 +124,7 @@ impl MemoryUsage for GlobalCache { + self.names.memory_usage() + self.counter.memory_usage() + self.versions.memory_usage() + + self.unused.len() * std::mem::size_of::() } } @@ -140,6 +142,9 @@ impl GlobalCache { }; if let Some(entry) = self.statements.get_mut(&parse_key) { + if entry.used == 0 { + self.unused.remove(&entry.counter); + } entry.used += 1; (false, global_name(entry.counter)) } else { @@ -147,14 +152,14 @@ impl GlobalCache { let name = global_name(self.counter); let parse = parse.rename(&name); - let parse_key = CacheKey { + let cache_key = CacheKey { query: parse.query_ref(), data_types: parse.data_types_ref(), version: 0, }; self.statements.insert( - parse_key, + cache_key.clone(), CachedStmt { counter: self.counter, used: 1, @@ -165,9 +170,8 @@ impl GlobalCache { name.clone(), Statement { parse, - row_description: None, - version: 0, - rewrite_plan: None, + cache_key, + ..Default::default() }, ); @@ -202,15 +206,22 @@ impl GlobalCache { name.clone(), Statement { parse, - row_description: None, version: self.versions, - rewrite_plan: None, + cache_key: key, + ..Default::default() }, ); name } + /// Rewrite prepared statement in the global cache. + pub fn rewrite(&mut self, parse: &Parse) { + if let Some(stmt) = self.names.get_mut(parse.name()) { + stmt.rewrite = Some(parse.clone()); + } + } + /// Client sent a Describe for a prepared statement and received a RowDescription. /// We record the RowDescription for later use by the results decoder. pub fn insert_row_description(&mut self, name: &str, row_description: &RowDescription) { @@ -221,35 +232,11 @@ impl GlobalCache { } } - pub fn update_query(&mut self, name: &str, sql: &str) -> bool { - if let Some(statement) = self.names.get_mut(name) { - let old_key = statement.cache_key(); - let cached = self.statements.remove(&old_key); - statement.parse.set_query(sql); - let new_key = statement.cache_key(); - if let Some(entry) = cached { - self.statements.insert(new_key, entry); - } - true - } else { - false - } - } - - pub fn set_rewrite_plan(&mut self, name: &str, plan: RewritePlan) { - if let Some(statement) = self.names.get_mut(name) { - statement.rewrite_plan = Some(plan); - } - } - - pub fn rewrite_plan(&self, name: &str) -> Option { - self.names.get(name).and_then(|s| s.rewrite_plan.clone()) - } - - #[cfg(test)] + /// Clear the global cache. pub fn reset(&mut self) { self.statements.clear(); self.names.clear(); + self.unused.clear(); self.counter = 0; self.versions = 0; } @@ -270,6 +257,25 @@ impl GlobalCache { self.names.get(name).map(|p| p.parse.clone()) } + /// Get the rewritten Parse statement. + /// + /// Used for preparing this statement on a server connection. + /// + pub fn rewritten_parse(&self, name: &str) -> Option { + self.names + .get(name) + .map(|p| p.rewrite.clone().unwrap_or(p.parse.clone())) + } + + /// Returns true if this prepared statement has been + /// rewritten by the rewrite engine. + pub fn is_rewritten(&self, name: &str) -> bool { + self.names + .get(name) + .map(|p| p.rewrite.is_some()) + .unwrap_or_default() + } + /// Get the RowDescription message for the prepared statement. /// /// It can be used to decode results received from executing the prepared @@ -289,46 +295,38 @@ impl GlobalCache { } /// Close prepared statement. - pub fn close(&mut self, name: &str, capacity: usize) -> bool { - let used = if let Some(stmt) = self.names.get(name) { - if let Some(stmt) = self.statements.get_mut(&stmt.cache_key()) { - stmt.used = stmt.used.saturating_sub(1); - stmt.used > 0 - } else { - false + pub fn close(&mut self, name: &str) { + if let Some(statement) = self.names.get(name) { + let key = statement.cache_key(); + + if let Some(entry) = self.statements.get_mut(&key) { + entry.used = entry.used.saturating_sub(1); + if entry.used == 0 && statement.evict_on_close { + self.remove(name); + } else if entry.used == 0 { + self.unused.insert(entry.counter); + } } - } else { - false - }; - - if !used && self.len() > capacity { - self.remove(name); - true - } else { - false } } /// Close all unused statements exceeding capacity. pub fn close_unused(&mut self, capacity: usize) -> usize { - let mut remove = self.statements.len() as i64 - capacity as i64; - let mut to_remove = vec![]; - for stmt in self.statements.values() { - if remove <= 0 { - break; - } - - if stmt.used == 0 { - to_remove.push(stmt.name()); - remove -= 1; - } + if capacity == 0 { + let removed = self.len(); + self.reset(); + return removed; } - for name in &to_remove { - self.remove(name); + let over = self.len().saturating_sub(capacity); + let remove = self.unused.iter().take(over).copied().collect::>(); + + for counter in &remove { + self.unused.remove(counter); + self.remove(&global_name(*counter)); } - to_remove.len() + remove.len() } /// Remove statement from global cache. @@ -343,6 +341,9 @@ impl GlobalCache { if let Some(stmt) = self.names.get(name) { if let Some(stmt) = self.statements.get_mut(&stmt.cache_key()) { stmt.used = stmt.used.saturating_sub(1); + if stmt.used == 0 { + self.unused.insert(stmt.counter); + } } } } @@ -380,21 +381,21 @@ mod test { assert_eq!(entry.used, 26); for _ in 0..25 { - cache.close("__pgdog_1", 0); + cache.close("__pgdog_1"); } let entry = cache.statements.get(&stmt.cache_key()).unwrap(); assert_eq!(entry.used, 1); + assert!(cache.unused.is_empty()); - cache.close("__pgdog_1", 0); - assert!(cache.statements.is_empty()); - assert!(cache.names.is_empty()); + cache.close("__pgdog_1"); + let entry = cache.statements.get(&stmt.cache_key()).unwrap(); + assert_eq!(entry.used, 0); + assert!(cache.unused.contains(&1)); // __pgdog_1 let name = cache.insert_anyway(&parse); - cache.close(&name, 0); - - assert!(cache.names.is_empty()); - assert!(cache.statements.is_empty()); + cache.close(&name); + assert!(cache.unused.contains(&2)); // __pgdog_2 } #[test] @@ -409,11 +410,8 @@ mod test { names.push(name); } - assert_eq!(cache.close_unused(0), 0); - for name in &names[0..5] { - assert!(!cache.close(name, 25)); // Won't close because - // capacity is enough to keep unused around. + cache.close(name); } assert_eq!(cache.close_unused(26), 0); @@ -422,4 +420,196 @@ mod test { assert_eq!(cache.close_unused(19), 0); assert_eq!(cache.len(), 20); } + + #[test] + fn test_reuse_statement_after_becomes_unused() { + let mut cache = GlobalCache::default(); + let parse = Parse::named("test", "SELECT $1"); + + let (new, name) = cache.insert(&parse); + assert!(new); + assert_eq!(cache.len(), 1); + + cache.close(&name); + let stmt = cache.names.get(&name).unwrap().clone(); + let entry = cache.statements.get(&stmt.cache_key()).unwrap(); + assert_eq!(entry.used, 0); + assert!(cache.unused.contains(&1)); + + let (new_again, name_again) = cache.insert(&parse); + assert!(!new_again); + assert_eq!(name, name_again); + assert!(!cache.unused.contains(&1)); + + let entry = cache.statements.get(&stmt.cache_key()).unwrap(); + assert_eq!(entry.used, 1); + } + + #[test] + fn test_close_nonexistent_statement() { + let mut cache = GlobalCache::default(); + let parse = Parse::named("test", "SELECT 1"); + cache.insert(&parse); + + cache.close("__pgdog_999"); + assert_eq!(cache.len(), 1); + assert!(cache.unused.is_empty()); + } + + #[test] + fn test_close_unused_with_capacity_zero() { + let mut cache = GlobalCache::default(); + + for i in 0..10 { + let parse = Parse::named("test", format!("SELECT {}", i)); + let (_, name) = cache.insert(&parse); + cache.close(&name); + } + + assert_eq!(cache.len(), 10); + assert_eq!(cache.unused.len(), 10); + + let removed = cache.close_unused(0); + assert_eq!(removed, 10); + assert_eq!(cache.len(), 0); + assert!(cache.unused.is_empty()); + assert!(cache.names.is_empty()); + assert!(cache.statements.is_empty()); + } + + #[test] + fn test_close_unused_when_nothing_unused() { + let mut cache = GlobalCache::default(); + + for i in 0..10 { + let parse = Parse::named("test", format!("SELECT {}", i)); + cache.insert(&parse); + } + + assert_eq!(cache.len(), 10); + assert!(cache.unused.is_empty()); + + let removed = cache.close_unused(5); + assert_eq!(removed, 0); + assert_eq!(cache.len(), 10); + } + + #[test] + fn test_decrement_marks_as_unused() { + let mut cache = GlobalCache::default(); + let parse = Parse::named("test", "SELECT 1"); + + let (_, name) = cache.insert(&parse); + cache.insert(&parse); + cache.insert(&parse); + + let stmt = cache.names.get(&name).unwrap().clone(); + let entry = cache.statements.get(&stmt.cache_key()).unwrap(); + assert_eq!(entry.used, 3); + + cache.decrement(&name); + let entry = cache.statements.get(&stmt.cache_key()).unwrap(); + assert_eq!(entry.used, 2); + assert!(cache.unused.is_empty()); + + cache.decrement(&name); + cache.decrement(&name); + let entry = cache.statements.get(&stmt.cache_key()).unwrap(); + assert_eq!(entry.used, 0); + assert!(cache.unused.contains(&1)); + + cache.decrement(&name); + let entry = cache.statements.get(&stmt.cache_key()).unwrap(); + assert_eq!(entry.used, 0); + } + + #[test] + fn test_both_maps_cleaned_up_on_removal() { + let mut cache = GlobalCache::default(); + let mut names = vec![]; + + for i in 0..5 { + let parse = Parse::named("test", format!("SELECT {}", i)); + let (_, name) = cache.insert(&parse); + names.push(name); + } + + assert_eq!(cache.len(), 5); + assert_eq!(cache.statements.len(), 5); + assert_eq!(cache.names.len(), 5); + + for name in &names { + cache.close(name); + } + + assert_eq!(cache.unused.len(), 5); + + cache.close_unused(0); + + assert_eq!(cache.len(), 0); + assert_eq!(cache.statements.len(), 0); + assert_eq!(cache.names.len(), 0); + assert_eq!(cache.unused.len(), 0); + + for name in &names { + assert!(cache.parse(name).is_none()); + assert!(cache.query(name).is_none()); + } + } + + #[test] + fn test_complex_interleaved_operations() { + let mut cache = GlobalCache::default(); + + let parse1 = Parse::named("test", "SELECT 1"); + let parse2 = Parse::named("test", "SELECT 2"); + let parse3 = Parse::named("test", "SELECT 3"); + + let (_, name1) = cache.insert(&parse1); + let (_, name2) = cache.insert(&parse2); + let (_, name3) = cache.insert(&parse3); + + cache.insert(&parse1); + cache.insert(&parse1); + + assert_eq!(cache.len(), 3); + + cache.close(&name1); + cache.close(&name2); + cache.close(&name3); + + assert_eq!(cache.unused.len(), 2); + assert!(cache.unused.contains(&2)); + assert!(cache.unused.contains(&3)); + assert!(!cache.unused.contains(&1)); + + cache.close(&name1); + cache.close(&name1); + assert_eq!(cache.unused.len(), 3); + assert!(cache.unused.contains(&1)); + + cache.close_unused(2); + assert_eq!(cache.len(), 2); + + let parse_exists = cache.parse(&name1).is_some(); + let parse_new = Parse::named("test", "SELECT 99"); + let (is_new, new_name) = cache.insert(&parse_new); + assert!(is_new); + + cache.close(&new_name); + assert_eq!(cache.unused.len(), 3); + + cache.close_unused(1); + assert_eq!(cache.len(), 1); + + if parse_exists { + assert!(cache.parse(&name1).is_some()); + } + + cache.close_unused(0); + assert_eq!(cache.len(), 0); + assert!(cache.statements.is_empty()); + assert!(cache.names.is_empty()); + assert!(cache.unused.is_empty()); + } } diff --git a/pgdog/src/frontend/prepared_statements/mod.rs b/pgdog/src/frontend/prepared_statements/mod.rs index 0df616134..f48513c4d 100644 --- a/pgdog/src/frontend/prepared_statements/mod.rs +++ b/pgdog/src/frontend/prepared_statements/mod.rs @@ -1,12 +1,14 @@ //! Prepared statements cache. -use std::{collections::HashMap, sync::Arc}; +use std::{collections::HashMap, sync::Arc, time::Duration}; use once_cell::sync::Lazy; use parking_lot::RwLock; +use tokio::{spawn, time::sleep}; +use tracing::debug; use crate::{ - frontend::router::parser::RewritePlan, + config::{config, PreparedStatements as PreparedStatementsLevel}, net::{Parse, ProtocolMessage}, stats::memory::MemoryUsage, }; @@ -17,7 +19,6 @@ pub mod rewrite; pub use error::Error; pub use global_cache::GlobalCache; - pub use rewrite::Rewrite; static CACHE: Lazy = Lazy::new(PreparedStatements::default); @@ -26,8 +27,7 @@ static CACHE: Lazy = Lazy::new(PreparedStatements::default); pub struct PreparedStatements { pub(super) global: Arc>, pub(super) local: HashMap, - pub(super) enabled: bool, - pub(super) capacity: usize, + pub(super) level: PreparedStatementsLevel, pub(super) memory_used: usize, } @@ -35,8 +35,7 @@ impl MemoryUsage for PreparedStatements { #[inline] fn memory_usage(&self) -> usize { self.local.memory_usage() - + self.enabled.memory_usage() - + self.capacity.memory_usage() + + std::mem::size_of::() + std::mem::size_of::>>() } } @@ -46,8 +45,7 @@ impl Default for PreparedStatements { Self { global: Arc::new(RwLock::new(GlobalCache::default())), local: HashMap::default(), - enabled: true, - capacity: usize::MAX, + level: PreparedStatementsLevel::Extended, memory_used: 0, } } @@ -97,26 +95,18 @@ impl PreparedStatements { parse.rename_fast(&name) } - /// Update stored SQL for a prepared statement after a rewrite. - pub fn update_query(&mut self, name: &str, sql: &str) -> bool { - self.global.write().update_query(name, sql) - } - - /// Store rewrite plan metadata for a prepared statement. - pub fn set_rewrite_plan(&mut self, name: &str, plan: RewritePlan) { - self.global.write().set_rewrite_plan(name, plan); - } - - /// Retrieve stored rewrite plan for a prepared statement, if any. - pub fn rewrite_plan(&self, name: &str) -> Option { - self.global.read().rewrite_plan(name) - } - /// Get global statement counter. pub fn name(&self, name: &str) -> Option<&String> { self.local.get(name) } + /// Get globally-prepared statement by local name. + pub fn parse(&self, name: &str) -> Option { + self.local + .get(name) + .and_then(|name| self.global.read().parse(name)) + } + /// Number of prepared statements in the local cache. pub fn len_local(&self) -> usize { self.local.len() @@ -131,7 +121,7 @@ impl PreparedStatements { pub fn close(&mut self, name: &str) { if let Some(global_name) = self.local.remove(name) { { - self.global.write().close(&global_name, self.capacity); + self.global.write().close(&global_name); } self.memory_used = self.memory_usage(); } @@ -143,7 +133,7 @@ impl PreparedStatements { let mut global = self.global.write(); for global_name in self.local.values() { - global.close(global_name, self.capacity); + global.close(global_name); } } @@ -155,6 +145,31 @@ impl PreparedStatements { pub fn memory_used(&self) -> usize { self.memory_used } + + /// Set the prepared statements level. + #[cfg(test)] + pub fn set_level(&mut self, level: PreparedStatementsLevel) { + self.level = level; + } +} + +/// Run prepared statements maintenance task +/// every second. +pub fn start_maintenance() { + spawn(async move { + debug!("prepared statements cache maintenance started"); + loop { + sleep(Duration::from_secs(1)).await; + run_maintenance(); + } + }); +} + +/// Check prepared statements cache for overflows +/// and remove any unused statements exceeding the limit. +pub fn run_maintenance() { + let capacity = config().config.general.prepared_statements_limit; + PreparedStatements::global().write().close_unused(capacity); } #[cfg(test)] @@ -166,7 +181,6 @@ mod test { #[test] fn test_maybe_rewrite() { let mut statements = PreparedStatements::default(); - statements.capacity = 0; let mut messages = vec![ ProtocolMessage::from(Parse::named("__sqlx_1", "SELECT 1")), @@ -183,7 +197,6 @@ mod test { statements.close_all(); assert!(statements.local.is_empty()); - assert!(statements.global.read().names().is_empty()); let mut messages = vec![ ProtocolMessage::from(Parse::named("__sqlx_1", "SELECT 1")), @@ -200,7 +213,6 @@ mod test { statements.close("__sqlx_1"); assert!(statements.local.is_empty()); - assert!(statements.global.read().names().is_empty()); } #[test] diff --git a/pgdog/src/frontend/router/cli.rs b/pgdog/src/frontend/router/cli.rs new file mode 100644 index 000000000..54accbd84 --- /dev/null +++ b/pgdog/src/frontend/router/cli.rs @@ -0,0 +1,60 @@ +//! Invoke the router on query. + +use std::path::Path; + +use tokio::fs::read_to_string; + +use super::Error; +use crate::{ + backend::databases::databases, + frontend::{client::Sticky, router::QueryParser, Command, RouterContext}, + net::{Parameters, ProtocolMessage, Query}, +}; + +#[derive(Debug, Clone)] +pub struct RouterCli { + database: String, + user: String, + queries: Vec, +} + +impl RouterCli { + pub async fn new( + database: impl ToString, + user: impl ToString, + file: impl AsRef, + ) -> Result { + let queries = read_to_string(file).await?; + let queries = queries + .split(";") + .filter(|q| !q.trim().is_empty()) + .map(|s| s.to_string()) + .collect(); + + Ok(Self { + database: database.to_string(), + user: user.to_string(), + queries, + }) + } + + pub fn run(&self) -> Result, Error> { + let mut result = vec![]; + let cluster = databases().cluster((self.user.as_str(), self.database.as_str()))?; + + for query in &self.queries { + let mut qp = QueryParser::default(); + let req = vec![ProtocolMessage::from(Query::new(query))]; + let cmd = qp.parse(RouterContext::new( + &req.into(), + &cluster, + &Parameters::default(), + None, + Sticky::new(), + )?)?; + result.push(cmd); + } + + Ok(result) + } +} diff --git a/pgdog/src/frontend/router/context.rs b/pgdog/src/frontend/router/context.rs index fa58fdb34..935522d0d 100644 --- a/pgdog/src/frontend/router/context.rs +++ b/pgdog/src/frontend/router/context.rs @@ -1,14 +1,16 @@ -use super::Error; +use super::{Error, ParameterHints}; use crate::{ backend::Cluster, - frontend::{client::TransactionType, BufferedQuery, ClientRequest, PreparedStatements}, + frontend::{ + client::{Sticky, TransactionType}, + router::Ast, + BufferedQuery, ClientRequest, + }, net::{Bind, Parameters}, }; #[derive(Debug)] pub struct RouterContext<'a> { - /// Prepared statements. - pub prepared_statements: &'a mut PreparedStatements, /// Bound parameters to the query. pub bind: Option<&'a Bind>, /// Query we're looking it. @@ -16,7 +18,7 @@ pub struct RouterContext<'a> { /// Cluster configuration. pub cluster: &'a Cluster, /// Client parameters, e.g. search_path. - pub params: &'a Parameters, + pub parameter_hints: ParameterHints<'a>, /// Client inside transaction, pub transaction: Option, /// Currently executing COPY statement. @@ -25,30 +27,38 @@ pub struct RouterContext<'a> { pub executable: bool, /// Two-pc enabled pub two_pc: bool, + /// Sticky omnisharded index. + pub sticky: Sticky, + /// Extended protocol. + pub extended: bool, + /// AST. + pub ast: Option, } impl<'a> RouterContext<'a> { pub fn new( buffer: &'a ClientRequest, cluster: &'a Cluster, - stmt: &'a mut PreparedStatements, params: &'a Parameters, transaction: Option, + sticky: Sticky, ) -> Result { let query = buffer.query()?; let bind = buffer.parameters()?; - let copy_mode = buffer.copy(); + let copy_mode = buffer.is_copy(); Ok(Self { - query, bind, - params, - prepared_statements: stmt, + parameter_hints: ParameterHints::from(params), cluster, transaction, copy_mode, - executable: buffer.executable(), + executable: buffer.is_executable(), two_pc: cluster.two_pc_enabled(), + sticky, + extended: matches!(query, Some(BufferedQuery::Prepared(_))) || bind.is_some(), + query, + ast: buffer.ast.clone(), }) } diff --git a/pgdog/src/frontend/router/mod.rs b/pgdog/src/frontend/router/mod.rs index 888c95164..b50c3b44a 100644 --- a/pgdog/src/frontend/router/mod.rs +++ b/pgdog/src/frontend/router/mod.rs @@ -1,8 +1,10 @@ //! Query router. +pub mod cli; pub mod context; pub mod copy; pub mod error; +pub mod parameter_hints; pub mod parser; pub mod round_robin; pub mod search_path; @@ -12,10 +14,13 @@ pub use copy::CopyRow; pub use error::Error; use lazy_static::lazy_static; use parser::Shard; -pub use parser::{Command, QueryParser, Route}; +pub use parser::{Ast, Command, QueryParser, Route}; + +use crate::frontend::router::parser::ShardWithPriority; use super::ClientRequest; pub use context::RouterContext; +pub use parameter_hints::ParameterHints; pub use search_path::SearchPath; pub use sharding::{Lists, Ranges}; @@ -73,7 +78,8 @@ impl Router { /// Get current route. pub fn route(&self) -> &Route { lazy_static! { - static ref DEFAULT_ROUTE: Route = Route::write(Shard::All); + static ref DEFAULT_ROUTE: Route = + Route::write(ShardWithPriority::new_default_unset(Shard::All)); } match self.command() { @@ -88,11 +94,6 @@ impl Router { self.latest_command = Command::default(); } - /// Query parser is inside a transaction. - pub fn in_transaction(&self) -> bool { - self.query_parser.in_transaction() - } - /// Get last commmand computed by the query parser. pub fn command(&self) -> &Command { &self.latest_command diff --git a/pgdog/src/frontend/router/parameter_hints.rs b/pgdog/src/frontend/router/parameter_hints.rs new file mode 100644 index 000000000..41b32e3b2 --- /dev/null +++ b/pgdog/src/frontend/router/parameter_hints.rs @@ -0,0 +1,164 @@ +use pgdog_config::Role; + +use super::parser::Error; +use crate::{ + backend::ShardingSchema, + frontend::router::{ + parser::{Schema, Shard, ShardWithPriority, ShardsWithPriority}, + sharding::{ContextBuilder, SchemaSharder}, + }, + net::{parameter::ParameterValue, Parameters}, +}; + +#[derive(Debug, Clone)] +pub struct ParameterHints<'a> { + pub search_path: Option<&'a ParameterValue>, + pub pgdog_shard: Option<&'a ParameterValue>, + pub pgdog_sharding_key: Option<&'a ParameterValue>, + pub pgdog_role: Option<&'a ParameterValue>, +} + +impl<'a> From<&'a Parameters> for ParameterHints<'a> { + fn from(value: &'a Parameters) -> Self { + Self { + search_path: value.search_path(), + pgdog_shard: value.get("pgdog.shard"), + pgdog_role: value.get("pgdog.role"), + pgdog_sharding_key: value.get("pgdog.sharding_key"), + } + } +} + +impl ParameterHints<'_> { + /// Compute shard from parameters. + pub(crate) fn compute_shard( + &self, + shards: &mut ShardsWithPriority, + sharding_schema: &ShardingSchema, + ) -> Result<(), Error> { + let mut schema_sharder = SchemaSharder::default(); + + if let Some(ParameterValue::Integer(val)) = self.pgdog_shard { + shards.push(ShardWithPriority::new_set(Shard::Direct(*val as usize))); + } + if let Some(ParameterValue::String(val)) = self.pgdog_shard { + if let Ok(shard) = val.parse() { + shards.push(ShardWithPriority::new_set(Shard::Direct(shard))); + } + } + if let Some(ParameterValue::String(val)) = self.pgdog_sharding_key { + if sharding_schema.schemas.is_empty() { + let ctx = + ContextBuilder::infer_from_from_and_config(val.as_str(), sharding_schema)? + .shards(sharding_schema.shards) + .build()?; + let shard = ctx.apply()?; + shards.push(ShardWithPriority::new_set(shard)); + } else { + schema_sharder.resolve(Some(Schema::from(val.as_str())), &sharding_schema.schemas); + + if let Some((shard, _)) = schema_sharder.get() { + shards.push(ShardWithPriority::new_set(shard.clone())); + } + } + } + if let Some(search_path) = self.search_path { + match search_path { + ParameterValue::String(search_path) => { + let schema = Schema::from(search_path.as_str()); + schema_sharder.resolve(Some(schema), &sharding_schema.schemas); + } + ParameterValue::Tuple(search_paths) => { + for schema in search_paths { + let schema = Schema::from(schema.as_str()); + schema_sharder.resolve(Some(schema), &sharding_schema.schemas); + } + } + + _ => (), + } + + if let Some((shard, schema)) = schema_sharder.get() { + shards.push(ShardWithPriority::new_search_path(shard.clone(), schema)); + } + } + Ok(()) + } + + /// Compute role from parameter value. + pub(crate) fn compute_role(&self) -> Option { + match self.pgdog_role { + Some(ParameterValue::String(val)) => match val.as_str() { + "replica" => Some(Role::Replica), + "primary" => Some(Role::Primary), + _ => None, + }, + + _ => None, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::backend::replication::ShardedSchemas; + use pgdog_config::sharding::ShardedSchema; + + fn make_sharding_schema(schemas: &[(&str, usize)]) -> ShardingSchema { + let sharded_schemas: Vec = schemas + .iter() + .map(|(name, shard)| ShardedSchema { + database: "test".to_string(), + name: Some(name.to_string()), + shard: *shard, + all: false, + }) + .collect(); + + ShardingSchema { + shards: schemas.len(), + schemas: ShardedSchemas::new(sharded_schemas), + ..Default::default() + } + } + + #[test] + fn test_sharding_key_with_schema_name() { + let sharding_schema = make_sharding_schema(&[("sales", 1)]); + + let sharding_key = ParameterValue::String("sales".to_string()); + let hints = ParameterHints { + search_path: None, + pgdog_shard: None, + pgdog_sharding_key: Some(&sharding_key), + pgdog_role: None, + }; + + let mut shards = ShardsWithPriority::default(); + hints.compute_shard(&mut shards, &sharding_schema).unwrap(); + + let result = shards.shard(); + assert_eq!(*result, Shard::Direct(1)); + } + + #[test] + fn test_sharding_key_takes_priority_over_search_path() { + let sharding_schema = make_sharding_schema(&[("sales", 0), ("inventory", 1)]); + + let sharding_key = ParameterValue::String("sales".to_string()); + let search_path = ParameterValue::String("inventory".to_string()); + let hints = ParameterHints { + search_path: Some(&search_path), + pgdog_shard: None, + pgdog_sharding_key: Some(&sharding_key), + pgdog_role: None, + }; + + let mut shards = ShardsWithPriority::default(); + hints.compute_shard(&mut shards, &sharding_schema).unwrap(); + + let result = shards.shard(); + assert_eq!(*result, Shard::Direct(0)); + } +} diff --git a/pgdog/src/frontend/router/parser/aggregate.rs b/pgdog/src/frontend/router/parser/aggregate.rs index 19ccc2e95..759728af6 100644 --- a/pgdog/src/frontend/router/parser/aggregate.rs +++ b/pgdog/src/frontend/router/parser/aggregate.rs @@ -1,11 +1,9 @@ use pg_query::protobuf::Integer; -use pg_query::protobuf::{a_const::Val, SelectStmt}; +use pg_query::protobuf::{a_const::Val, Node, SelectStmt, String as PgQueryString}; use pg_query::NodeEnum; use crate::frontend::router::parser::{ExpressionRegistry, Function}; -use super::Error; - #[derive(Debug, Clone, PartialEq)] pub struct AggregateTarget { column: usize, @@ -39,6 +37,26 @@ pub enum AggregateFunction { Min, Avg, Sum, + StddevPop, + StddevSamp, + VarPop, + VarSamp, +} + +impl AggregateFunction { + pub fn as_str(&self) -> &'static str { + match self { + AggregateFunction::Count => "count", + AggregateFunction::Max => "max", + AggregateFunction::Min => "min", + AggregateFunction::Avg => "avg", + AggregateFunction::Sum => "sum", + AggregateFunction::StddevPop => "stddev_pop", + AggregateFunction::StddevSamp => "stddev_samp", + AggregateFunction::VarPop => "var_pop", + AggregateFunction::VarSamp => "var_samp", + } + } } #[derive(Debug, Clone, PartialEq, Default)] @@ -47,9 +65,63 @@ pub struct Aggregate { group_by: Vec, } +fn target_list_to_index(stmt: &SelectStmt, column_names: Vec<&String>) -> Option { + for (idx, node) in stmt.target_list.iter().enumerate() { + if let Some(NodeEnum::ResTarget(res_target_box)) = node.node.as_ref() { + let res_target = res_target_box.as_ref(); + if let Some(node_box) = res_target.val.as_ref() { + if let Some(NodeEnum::ColumnRef(column_ref)) = node_box.node.as_ref() { + let select_names: Vec<&String> = column_ref + .fields + .iter() + .filter_map(|field_node| { + if let Some(node_box) = field_node.node.as_ref() { + match node_box { + NodeEnum::String(PgQueryString { + sval: found_column_name, + .. + }) => Some(found_column_name), + _ => None, + } + } else { + None + } + }) + .collect(); + + if select_names.is_empty() { + continue; + } + + if columns_match(&column_names, &select_names) { + return Some(idx); + } + } + } + } + } + None +} + +fn columns_match(group_by_names: &[&String], select_names: &[&String]) -> bool { + if group_by_names == select_names { + return true; + } + + if group_by_names.len() == 1 && select_names.len() == 2 { + return select_names[1] == group_by_names[0]; + } + + if group_by_names.len() == 2 && select_names.len() == 1 { + return group_by_names[1] == select_names[0]; + } + + false +} + impl Aggregate { /// Figure out what aggregates are present and which ones PgDog supports. - pub fn parse(stmt: &SelectStmt) -> Result { + pub fn parse(stmt: &SelectStmt) -> Self { let mut targets = vec![]; let mut registry = ExpressionRegistry::new(); let group_by = stmt @@ -61,6 +133,20 @@ impl Aggregate { Val::Ival(Integer { ival }) => Some(*ival as usize - 1), // We use 0-indexed arrays, Postgres uses 1-indexed. _ => None, }), + NodeEnum::ColumnRef(column_ref) => { + let column_names: Vec<&String> = column_ref + .fields + .iter() + .filter_map(|node| match node { + Node { + node: + Some(NodeEnum::String(PgQueryString { sval: column_name })), + } => Some(column_name), + _ => None, + }) + .collect(); + Some(target_list_to_index(stmt, column_names)) + } _ => None, }) }) @@ -78,6 +164,10 @@ impl Aggregate { "min" => Some(AggregateFunction::Min), "sum" => Some(AggregateFunction::Sum), "avg" => Some(AggregateFunction::Avg), + "stddev" | "stddev_samp" => Some(AggregateFunction::StddevSamp), + "stddev_pop" => Some(AggregateFunction::StddevPop), + "variance" | "var_samp" => Some(AggregateFunction::VarSamp), + "var_pop" => Some(AggregateFunction::VarPop), _ => None, }; @@ -106,7 +196,7 @@ impl Aggregate { } } - Ok(Self { targets, group_by }) + Self { targets, group_by } } pub fn targets(&self) -> &[AggregateTarget] { @@ -165,7 +255,7 @@ mod test { .unwrap(); match query.stmt.unwrap().node.unwrap() { NodeEnum::SelectStmt(stmt) => { - let aggr = Aggregate::parse(&stmt).unwrap(); + let aggr = Aggregate::parse(&stmt); assert_eq!( aggr.targets().first().unwrap().function, AggregateFunction::Count @@ -187,7 +277,7 @@ mod test { .unwrap(); match query.stmt.unwrap().node.unwrap() { NodeEnum::SelectStmt(stmt) => { - let aggr = Aggregate::parse(&stmt).unwrap(); + let aggr = Aggregate::parse(&stmt); assert_eq!(aggr.targets().len(), 2); let count = &aggr.targets()[0]; let avg = &aggr.targets()[1]; @@ -212,7 +302,7 @@ mod test { .unwrap(); match query.stmt.unwrap().node.unwrap() { NodeEnum::SelectStmt(stmt) => { - let aggr = Aggregate::parse(&stmt).unwrap(); + let aggr = Aggregate::parse(&stmt); assert_eq!(aggr.targets().len(), 2); let count = &aggr.targets()[0]; let avg = &aggr.targets()[1]; @@ -237,7 +327,7 @@ mod test { .unwrap(); match query.stmt.unwrap().node.unwrap() { NodeEnum::SelectStmt(stmt) => { - let aggr = Aggregate::parse(&stmt).unwrap(); + let aggr = Aggregate::parse(&stmt); assert_eq!(aggr.targets().len(), 2); let count = &aggr.targets()[0]; let avg = &aggr.targets()[1]; @@ -250,4 +340,319 @@ mod test { _ => panic!("not a select"), } } + + #[test] + fn test_parse_stddev_variants() { + let query = pg_query::parse( + "SELECT STDDEV(price), STDDEV_SAMP(price), STDDEV_POP(price) FROM menu", + ) + .unwrap() + .protobuf + .stmts + .first() + .cloned() + .unwrap(); + match query.stmt.unwrap().node.unwrap() { + NodeEnum::SelectStmt(stmt) => { + let aggr = Aggregate::parse(&stmt); + assert_eq!(aggr.targets().len(), 3); + assert!(matches!( + aggr.targets()[0].function(), + AggregateFunction::StddevSamp + )); + assert!(matches!( + aggr.targets()[1].function(), + AggregateFunction::StddevSamp + )); + assert!(matches!( + aggr.targets()[2].function(), + AggregateFunction::StddevPop + )); + } + _ => panic!("not a select"), + } + } + + #[test] + fn test_parse_variance_variants() { + let query = + pg_query::parse("SELECT VARIANCE(price), VAR_SAMP(price), VAR_POP(price) FROM menu") + .unwrap() + .protobuf + .stmts + .first() + .cloned() + .unwrap(); + match query.stmt.unwrap().node.unwrap() { + NodeEnum::SelectStmt(stmt) => { + let aggr = Aggregate::parse(&stmt); + assert_eq!(aggr.targets().len(), 3); + assert!(matches!( + aggr.targets()[0].function(), + AggregateFunction::VarSamp + )); + assert!(matches!( + aggr.targets()[1].function(), + AggregateFunction::VarSamp + )); + assert!(matches!( + aggr.targets()[2].function(), + AggregateFunction::VarPop + )); + } + _ => panic!("not a select"), + } + } + + #[test] + fn test_parse_group_by_ordinals() { + let query = + pg_query::parse("SELECT price, category_id, SUM(quantity) FROM menu GROUP BY 1, 2") + .unwrap() + .protobuf + .stmts + .first() + .cloned() + .unwrap(); + match query.stmt.unwrap().node.unwrap() { + NodeEnum::SelectStmt(stmt) => { + let aggr = Aggregate::parse(&stmt); + assert_eq!(aggr.group_by(), &[0, 1]); + assert_eq!(aggr.targets().len(), 1); + let target = &aggr.targets()[0]; + assert!(matches!(target.function(), AggregateFunction::Sum)); + assert_eq!(target.column(), 2); + } + _ => panic!("not a select"), + } + } + + #[test] + fn test_parse_sum_distinct_sets_flag() { + let query = pg_query::parse("SELECT SUM(DISTINCT price) FROM menu") + .unwrap() + .protobuf + .stmts + .first() + .cloned() + .unwrap(); + match query.stmt.unwrap().node.unwrap() { + NodeEnum::SelectStmt(stmt) => { + let aggr = Aggregate::parse(&stmt); + assert_eq!(aggr.targets().len(), 1); + let target = &aggr.targets()[0]; + assert!(matches!(target.function(), AggregateFunction::Sum)); + assert!(target.is_distinct()); + } + _ => panic!("not a select"), + } + } + + #[test] + fn test_parse_group_by_column_name_single() { + let query = pg_query::parse("SELECT user_id, COUNT(1) FROM example GROUP BY user_id") + .unwrap() + .protobuf + .stmts + .first() + .cloned() + .unwrap(); + match query.stmt.unwrap().node.unwrap() { + NodeEnum::SelectStmt(stmt) => { + let aggr = Aggregate::parse(&stmt); + assert_eq!(aggr.group_by(), &[0]); + assert_eq!(aggr.targets().len(), 1); + let target = &aggr.targets()[0]; + assert!(matches!(target.function(), AggregateFunction::Count)); + assert_eq!(target.column(), 1); + } + _ => panic!("not a select"), + } + } + + #[test] + fn test_parse_group_by_column_name_multiple() { + let query = pg_query::parse( + "SELECT COUNT(*), user_id, category_id FROM example GROUP BY user_id, category_id", + ) + .unwrap() + .protobuf + .stmts + .first() + .cloned() + .unwrap(); + match query.stmt.unwrap().node.unwrap() { + NodeEnum::SelectStmt(stmt) => { + let aggr = Aggregate::parse(&stmt); + assert_eq!(aggr.group_by(), &[1, 2]); + assert_eq!(aggr.targets().len(), 1); + let target = &aggr.targets()[0]; + assert!(matches!(target.function(), AggregateFunction::Count)); + assert_eq!(target.column(), 0); + } + _ => panic!("not a select"), + } + } + + #[test] + fn test_parse_group_by_qualified_column_name() { + let query = pg_query::parse( + "SELECT COUNT(1), example.user_id FROM example GROUP BY example.user_id", + ) + .unwrap() + .protobuf + .stmts + .first() + .cloned() + .unwrap(); + match query.stmt.unwrap().node.unwrap() { + NodeEnum::SelectStmt(stmt) => { + let aggr = Aggregate::parse(&stmt); + assert_eq!(aggr.group_by(), &[1]); + assert_eq!(aggr.targets().len(), 1); + let target = &aggr.targets()[0]; + assert!(matches!(target.function(), AggregateFunction::Count)); + assert_eq!(target.column(), 0); + } + _ => panic!("not a select"), + } + } + + #[test] + fn test_parse_group_by_mixed_ordinal_and_column_name() { + let query = pg_query::parse( + "SELECT user_id, category_id, SUM(quantity) FROM example GROUP BY user_id, 2", + ) + .unwrap() + .protobuf + .stmts + .first() + .cloned() + .unwrap(); + match query.stmt.unwrap().node.unwrap() { + NodeEnum::SelectStmt(stmt) => { + let aggr = Aggregate::parse(&stmt); + assert_eq!(aggr.group_by(), &[0, 1]); + assert_eq!(aggr.targets().len(), 1); + let target = &aggr.targets()[0]; + assert!(matches!(target.function(), AggregateFunction::Sum)); + assert_eq!(target.column(), 2); + } + _ => panic!("not a select"), + } + } + + #[test] + fn test_parse_group_by_column_not_in_select() { + let query = pg_query::parse("SELECT COUNT(*) FROM example GROUP BY user_id") + .unwrap() + .protobuf + .stmts + .first() + .cloned() + .unwrap(); + match query.stmt.unwrap().node.unwrap() { + NodeEnum::SelectStmt(stmt) => { + let aggr = Aggregate::parse(&stmt); + let empty: Vec = vec![]; + assert_eq!(aggr.group_by(), empty.as_slice()); + assert_eq!(aggr.targets().len(), 1); + } + _ => panic!("not a select"), + } + } + + #[test] + fn test_parse_group_by_with_multiple_aggregates() { + let query = pg_query::parse( + "SELECT COUNT(*), SUM(price), user_id, AVG(price) FROM example GROUP BY user_id", + ) + .unwrap() + .protobuf + .stmts + .first() + .cloned() + .unwrap(); + match query.stmt.unwrap().node.unwrap() { + NodeEnum::SelectStmt(stmt) => { + let aggr = Aggregate::parse(&stmt); + assert_eq!(aggr.group_by(), &[2]); + assert_eq!(aggr.targets().len(), 3); + assert!(matches!( + aggr.targets()[0].function(), + AggregateFunction::Count + )); + assert!(matches!( + aggr.targets()[1].function(), + AggregateFunction::Sum + )); + assert!(matches!( + aggr.targets()[2].function(), + AggregateFunction::Avg + )); + } + _ => panic!("not a select"), + } + } + + #[test] + fn test_parse_group_by_qualified_matches_select_unqualified() { + let query = + pg_query::parse("SELECT user_id, COUNT(1) FROM example GROUP BY example.user_id") + .unwrap() + .protobuf + .stmts + .first() + .cloned() + .unwrap(); + match query.stmt.unwrap().node.unwrap() { + NodeEnum::SelectStmt(stmt) => { + let aggr = Aggregate::parse(&stmt); + assert_eq!(aggr.group_by(), &[0]); + assert_eq!(aggr.targets().len(), 1); + } + _ => panic!("not a select"), + } + } + + #[test] + fn test_parse_group_by_unqualified_matches_select_qualified() { + let query = + pg_query::parse("SELECT example.user_id, COUNT(1) FROM example GROUP BY user_id") + .unwrap() + .protobuf + .stmts + .first() + .cloned() + .unwrap(); + match query.stmt.unwrap().node.unwrap() { + NodeEnum::SelectStmt(stmt) => { + let aggr = Aggregate::parse(&stmt); + assert_eq!(aggr.group_by(), &[0]); + assert_eq!(aggr.targets().len(), 1); + } + _ => panic!("not a select"), + } + } + + #[test] + fn test_parse_group_by_both_qualified_order_matters() { + let query = pg_query::parse( + "SELECT example.user_id, COUNT(1) FROM example GROUP BY example.user_id", + ) + .unwrap() + .protobuf + .stmts + .first() + .cloned() + .unwrap(); + match query.stmt.unwrap().node.unwrap() { + NodeEnum::SelectStmt(stmt) => { + let aggr = Aggregate::parse(&stmt); + assert_eq!(aggr.group_by(), &[0]); + assert_eq!(aggr.targets().len(), 1); + } + _ => panic!("not a select"), + } + } } diff --git a/pgdog/src/frontend/router/parser/cache.rs b/pgdog/src/frontend/router/parser/cache.rs deleted file mode 100644 index af8688fe6..000000000 --- a/pgdog/src/frontend/router/parser/cache.rs +++ /dev/null @@ -1,348 +0,0 @@ -//! AST cache. -//! -//! Shared between all clients and databases. - -use lru::LruCache; -use once_cell::sync::Lazy; -use pg_query::*; -use std::collections::HashMap; - -use parking_lot::Mutex; -use std::sync::Arc; -use tracing::debug; - -use crate::{ - backend::ShardingSchema, - config::Role, - frontend::router::parser::{comment::comment, Shard}, -}; - -use super::Route; - -static CACHE: Lazy = Lazy::new(Cache::new); - -#[derive(Default, Debug, Clone, Copy)] -pub struct Stats { - /// Cache hits. - pub hits: usize, - /// Cache misses (new queries). - pub misses: usize, - /// Direct shard queries. - pub direct: usize, - /// Multi-shard queries. - pub multi: usize, -} - -/// Abstract syntax tree (query) cache entry, -/// with statistics. -#[derive(Debug, Clone)] -pub struct CachedAst { - /// pg_query-produced AST. - pub ast: Arc, - /// Statistics. Use a separate Mutex to avoid - /// contention when updating them. - pub stats: Arc>, - /// Was this entry cached? - pub cached: bool, - /// Shard. - pub shard: Shard, - /// Role. - pub role: Option, -} - -impl CachedAst { - /// Create new cache entry from pg_query's AST. - fn new(query: &str, schema: &ShardingSchema) -> std::result::Result { - let ast = parse(query).map_err(super::Error::PgQuery)?; - let (shard, role) = comment(query, schema)?; - - Ok(Self { - cached: true, - ast: Arc::new(ast), - shard, - role, - stats: Arc::new(Mutex::new(Stats { - hits: 1, - ..Default::default() - })), - }) - } - - /// Get the reference to the AST. - pub fn ast(&self) -> &ParseResult { - &self.ast - } - - /// Update stats for this statement, given the route - /// calculated by the query parser. - pub fn update_stats(&self, route: &Route) { - let mut guard = self.stats.lock(); - - if route.is_cross_shard() { - guard.multi += 1; - } else { - guard.direct += 1; - } - } -} - -/// Mutex-protected query cache. -#[derive(Debug)] -struct Inner { - /// Least-recently-used cache. - queries: LruCache, - /// Cache global stats. - stats: Stats, -} - -/// AST cache. -#[derive(Clone, Debug)] -pub struct Cache { - inner: Arc>, -} - -impl Cache { - /// Create new cache. Should only be done once at pooler startup. - fn new() -> Self { - Self { - inner: Arc::new(Mutex::new(Inner { - queries: LruCache::unbounded(), - stats: Stats::default(), - })), - } - } - - /// Resize cache to capacity, evicting any statements exceeding the capacity. - /// - /// Minimum capacity is 1. - pub fn resize(capacity: usize) { - let capacity = if capacity == 0 { 1 } else { capacity }; - - CACHE - .inner - .lock() - .queries - .resize(capacity.try_into().unwrap()); - - debug!("ast cache size set to {}", capacity); - } - - /// Parse a statement by either getting it from cache - /// or using pg_query parser. - /// - /// N.B. There is a race here that allows multiple threads to - /// parse the same query. That's better imo than locking the data structure - /// while we parse the query. - pub fn parse( - &self, - query: &str, - schema: &ShardingSchema, - ) -> std::result::Result { - { - let mut guard = self.inner.lock(); - let ast = guard.queries.get_mut(query).map(|entry| { - entry.stats.lock().hits += 1; // No contention on this. - entry.clone() - }); - if let Some(ast) = ast { - guard.stats.hits += 1; - return Ok(ast); - } - } - - // Parse query without holding lock. - let entry = CachedAst::new(query, schema)?; - - let mut guard = self.inner.lock(); - guard.queries.put(query.to_owned(), entry.clone()); - guard.stats.misses += 1; - - Ok(entry) - } - - /// Parse a statement but do not store it in the cache. - pub fn parse_uncached( - &self, - query: &str, - schema: &ShardingSchema, - ) -> std::result::Result { - let mut entry = CachedAst::new(query, schema)?; - entry.cached = false; - Ok(entry) - } - - /// Record a query sent over the simple protocol, while removing parameters. - pub fn record_normalized( - &self, - query: &str, - route: &Route, - schema: &ShardingSchema, - ) -> std::result::Result<(), super::Error> { - let normalized = pg_query::normalize(query).map_err(super::Error::PgQuery)?; - - { - let mut guard = self.inner.lock(); - if let Some(entry) = guard.queries.get(&normalized) { - entry.update_stats(route); - guard.stats.hits += 1; - return Ok(()); - } - } - - let entry = CachedAst::new(query, schema)?; - entry.update_stats(route); - - let mut guard = self.inner.lock(); - guard.queries.put(normalized, entry); - guard.stats.misses += 1; - - Ok(()) - } - - /// Get global cache instance. - pub fn get() -> Self { - CACHE.clone() - } - - /// Get cache stats. - pub fn stats() -> (Stats, usize) { - let cache = Self::get(); - let (len, query_stats, mut stats) = { - let guard = cache.inner.lock(); - ( - guard.queries.len(), - guard - .queries - .iter() - .map(|c| c.1.stats.clone()) - .collect::>(), - guard.stats.clone(), - ) - }; - for stat in query_stats { - let guard = stat.lock(); - stats.direct += guard.direct; - stats.multi += guard.multi; - } - (stats, len) - } - - /// Get a copy of all queries stored in the cache. - pub fn queries() -> HashMap { - Self::get() - .inner - .lock() - .queries - .iter() - .map(|i| (i.0.clone(), i.1.clone())) - .collect() - } - - /// Reset cache, removing all statements - /// and setting stats to 0. - pub fn reset() { - let cache = Self::get(); - let mut guard = cache.inner.lock(); - guard.queries.clear(); - guard.stats.hits = 0; - guard.stats.misses = 0; - } -} - -#[cfg(test)] -mod test { - use tokio::spawn; - - use super::*; - use std::time::{Duration, Instant}; - - #[tokio::test(flavor = "multi_thread")] - async fn bench_ast_cache() { - let query = "SELECT - u.username, - p.product_name, - SUM(oi.quantity * oi.price) AS total_revenue, - AVG(r.rating) AS average_rating, - COUNT(DISTINCT c.country) AS countries_purchased_from - FROM users u - INNER JOIN orders o ON u.user_id = o.user_id - INNER JOIN order_items oi ON o.order_id = oi.order_id - INNER JOIN products p ON oi.product_id = p.product_id - LEFT JOIN reviews r ON o.order_id = r.order_id - LEFT JOIN customer_addresses c ON o.shipping_address_id = c.address_id - WHERE - o.order_date BETWEEN '2023-01-01' AND '2023-12-31' - AND p.category IN ('Electronics', 'Clothing') - AND (r.rating > 4 OR r.rating IS NULL) - GROUP BY u.username, p.product_name - HAVING COUNT(DISTINCT c.country) > 2 - ORDER BY total_revenue DESC; -"; - - let times = 10_000; - let threads = 5; - - let mut tasks = vec![]; - for _ in 0..threads { - let handle = spawn(async move { - let mut parse_time = Duration::ZERO; - for _ in 0..(times / threads) { - let start = Instant::now(); - parse(query).unwrap(); - parse_time += start.elapsed(); - } - - parse_time - }); - tasks.push(handle); - } - - let mut parse_time = Duration::ZERO; - for task in tasks { - parse_time += task.await.unwrap(); - } - - println!("[bench_ast_cache]: parse time: {:?}", parse_time); - - // Simulate lock contention. - let mut tasks = vec![]; - - for _ in 0..threads { - let handle = spawn(async move { - let mut cached_time = Duration::ZERO; - for _ in 0..(times / threads) { - let start = Instant::now(); - Cache::get() - .parse(query, &ShardingSchema::default()) - .unwrap(); - cached_time += start.elapsed(); - } - - cached_time - }); - tasks.push(handle); - } - - let mut cached_time = Duration::ZERO; - for task in tasks { - cached_time += task.await.unwrap(); - } - - println!("[bench_ast_cache]: cached time: {:?}", cached_time); - - let faster = parse_time.as_micros() as f64 / cached_time.as_micros() as f64; - println!( - "[bench_ast_cache]: cached is {:.4} times faster than parsed", - faster - ); // 32x on my M1 - - assert!(faster > 10.0); - } - - #[test] - fn test_normalize() { - let q = "SELECT * FROM users WHERE id = 1"; - let normalized = normalize(q).unwrap(); - assert_eq!(normalized, "SELECT * FROM users WHERE id = $1"); - } -} diff --git a/pgdog/src/frontend/router/parser/cache/ast.rs b/pgdog/src/frontend/router/parser/cache/ast.rs new file mode 100644 index 000000000..d04a8546c --- /dev/null +++ b/pgdog/src/frontend/router/parser/cache/ast.rs @@ -0,0 +1,187 @@ +use pg_query::{parse, parse_raw, protobuf::ObjectType, NodeEnum, NodeRef, ParseResult}; +use pgdog_config::QueryParserEngine; +use std::fmt::Debug; +use std::{collections::HashSet, ops::Deref}; + +use parking_lot::Mutex; +use std::sync::Arc; + +use super::super::{ + comment::comment, Error, Route, Shard, StatementRewrite, StatementRewriteContext, Table, +}; +use super::{Fingerprint, Stats}; +use crate::frontend::router::parser::rewrite::statement::RewritePlan; +use crate::frontend::{BufferedQuery, PreparedStatements}; +use crate::{backend::ShardingSchema, config::Role}; + +/// Abstract syntax tree (query) cache entry, +/// with statistics. +#[derive(Debug, Clone)] +pub struct Ast { + /// Was this entry cached? + pub cached: bool, + /// Inner sync. + inner: Arc, +} + +#[derive(Debug)] +pub struct AstInner { + /// Cached AST. + pub ast: ParseResult, + /// AST stats. + pub stats: Mutex, + /// Shard. + pub comment_shard: Option, + /// Role. + pub comment_role: Option, + /// Rewrite plan. + pub rewrite_plan: RewritePlan, + /// Fingerprint. + pub fingerprint: Fingerprint, +} + +impl AstInner { + /// Create new AST record, with no rewrite or comment routing. + pub fn new(ast: ParseResult) -> Self { + Self { + ast, + stats: Mutex::new(Stats::new()), + comment_role: None, + comment_shard: None, + rewrite_plan: RewritePlan::default(), + fingerprint: Fingerprint::default(), + } + } +} + +impl Deref for Ast { + type Target = AstInner; + + fn deref(&self) -> &Self::Target { + &self.inner + } +} + +impl Ast { + /// Parse statement and run the rewrite engine, if necessary. + pub fn new( + query: &BufferedQuery, + schema: &ShardingSchema, + prepared_statements: &mut PreparedStatements, + ) -> Result { + let mut ast = match schema.query_parser_engine { + QueryParserEngine::PgQueryProtobuf => parse(query), + QueryParserEngine::PgQueryRaw => parse_raw(query), + } + .map_err(Error::PgQuery)?; + let (comment_shard, comment_role) = comment(query, schema)?; + let fingerprint = + Fingerprint::new(query, schema.query_parser_engine).map_err(Error::PgQuery)?; + + // Don't rewrite statements that will be + // sent to a direct shard. + let rewrite_plan = if comment_shard.is_none() { + StatementRewrite::new(StatementRewriteContext { + stmt: &mut ast.protobuf, + extended: query.extended(), + prepared: query.prepared(), + prepared_statements, + schema, + }) + .maybe_rewrite()? + } else { + RewritePlan::default() + }; + + Ok(Self { + cached: true, + inner: Arc::new(AstInner { + stats: Mutex::new(Stats::new()), + comment_shard, + comment_role, + ast, + rewrite_plan, + fingerprint, + }), + }) + } + + /// Record new AST entry, without rewriting or comment-routing. + pub fn new_record(query: &str, query_parser_engine: QueryParserEngine) -> Result { + let ast = match query_parser_engine { + QueryParserEngine::PgQueryProtobuf => parse(query), + QueryParserEngine::PgQueryRaw => parse_raw(query), + } + .map_err(Error::PgQuery)?; + + Ok(Self { + cached: true, + inner: Arc::new(AstInner::new(ast)), + }) + } + + /// Create new AST from a parse result. + pub fn from_parse_result(parse_result: ParseResult) -> Self { + Self { + cached: true, + inner: Arc::new(AstInner::new(parse_result)), + } + } + + /// Get the reference to the AST. + pub fn parse_result(&self) -> &ParseResult { + &self.ast + } + + /// Get a list of tables referenced by the query. + /// + /// This is better than pg_query's version because we + /// also handle `NodeRef::CreateStmt` and we handle identifiers correctly. + /// + pub fn tables<'a>(&'a self) -> Vec> { + let mut tables = HashSet::new(); + + for node in self.ast.protobuf.nodes() { + match node.0 { + NodeRef::RangeVar(table) => { + let table = Table::from(table); + tables.insert(table); + } + + NodeRef::CreateStmt(stmt) => { + if let Some(ref stmt) = stmt.relation { + tables.insert(Table::from(stmt)); + } + } + + NodeRef::DropStmt(stmt) => { + if stmt.remove_type() == ObjectType::ObjectTable { + for object in &stmt.objects { + if let Some(NodeEnum::List(ref list)) = object.node { + if let Ok(table) = Table::try_from(list) { + tables.insert(table); + } + } + } + } + } + + _ => (), + } + } + + tables.into_iter().collect() + } + + /// Update stats for this statement, given the route + /// calculated by the query parser. + pub fn update_stats(&self, route: &Route) { + let mut guard = self.stats.lock(); + + if route.is_cross_shard() { + guard.multi += 1; + } else { + guard.direct += 1; + } + } +} diff --git a/pgdog/src/frontend/router/parser/cache/cache_impl.rs b/pgdog/src/frontend/router/parser/cache/cache_impl.rs new file mode 100644 index 000000000..a08f9a51b --- /dev/null +++ b/pgdog/src/frontend/router/parser/cache/cache_impl.rs @@ -0,0 +1,221 @@ +use lru::LruCache; +use once_cell::sync::Lazy; +use pg_query::normalize; +use pgdog_config::QueryParserEngine; +use std::collections::HashMap; + +use parking_lot::Mutex; +use std::sync::Arc; +use tracing::debug; + +use super::super::{Error, Route}; +use super::Ast; +use crate::backend::ShardingSchema; +use crate::frontend::{BufferedQuery, PreparedStatements}; + +static CACHE: Lazy = Lazy::new(Cache::new); + +/// Cache statistics. +#[derive(Default, Debug, Clone, Copy)] +pub struct Stats { + /// Cache hits. + pub hits: usize, + /// Cache misses (new queries). + pub misses: usize, + /// Direct shard queries. + pub direct: usize, + /// Multi-shard queries. + pub multi: usize, +} + +impl Stats { + /// Create new statistics record for an AST entry. + pub fn new() -> Self { + Self { + hits: 1, + ..Default::default() + } + } +} + +/// Mutex-protected query cache. +#[derive(Debug)] +struct Inner { + /// Least-recently-used cache. + queries: LruCache, + /// Cache global stats. + stats: Stats, +} + +/// AST cache. +#[derive(Clone, Debug)] +pub struct Cache { + inner: Arc>, +} + +impl Cache { + /// Create new cache. Should only be done once at pooler startup. + fn new() -> Self { + Self { + inner: Arc::new(Mutex::new(Inner { + queries: LruCache::unbounded(), + stats: Stats::default(), + })), + } + } + + /// Resize cache to capacity, evicting any statements exceeding the capacity. + /// + /// Minimum capacity is 1. + pub fn resize(capacity: usize) { + let capacity = if capacity == 0 { 1 } else { capacity }; + + CACHE + .inner + .lock() + .queries + .resize(capacity.try_into().unwrap()); + + debug!("ast cache size set to {}", capacity); + } + + /// Handle parsing a query. + pub fn query( + &self, + query: &BufferedQuery, + schema: &ShardingSchema, + prepared_statements: &mut PreparedStatements, + ) -> Result { + match query { + BufferedQuery::Prepared(_) => self.parse(query, schema, prepared_statements), + BufferedQuery::Query(_) => self.simple(query, schema, prepared_statements), + } + } + + /// Parse a statement by either getting it from cache + /// or using pg_query parser. + /// + /// N.B. There is a race here that allows multiple threads to + /// parse the same query. That's better imo than locking the data structure + /// while we parse the query. + fn parse( + &self, + query: &BufferedQuery, + schema: &ShardingSchema, + prepared_statements: &mut PreparedStatements, + ) -> Result { + { + let mut guard = self.inner.lock(); + let ast = guard.queries.get_mut(query.query()).map(|entry| { + entry.stats.lock().hits += 1; // No contention on this. + entry.clone() + }); + if let Some(ast) = ast { + guard.stats.hits += 1; + return Ok(ast); + } + } + + // Parse query without holding lock. + let entry = Ast::new(query, schema, prepared_statements)?; + + let mut guard = self.inner.lock(); + guard.queries.put(query.query().to_string(), entry.clone()); + guard.stats.misses += 1; + + Ok(entry) + } + + /// Parse and rewrite a statement but do not store it in the cache, + /// because it may contain parameter values. + fn simple( + &self, + query: &BufferedQuery, + schema: &ShardingSchema, + prepared_statements: &mut PreparedStatements, + ) -> Result { + let mut entry = Ast::new(query, schema, prepared_statements)?; + entry.cached = false; + Ok(entry) + } + + /// Record a query sent over the simple protocol, while removing parameters. + /// + /// Used by dry run mode to keep stats on what queries are routed correctly, + /// and which are not. + /// + pub fn record_normalized( + &self, + query: &str, + route: &Route, + query_parser_engine: QueryParserEngine, + ) -> Result<(), Error> { + let normalized = normalize(query).map_err(Error::PgQuery)?; + + { + let mut guard = self.inner.lock(); + if let Some(entry) = guard.queries.get(&normalized) { + entry.update_stats(route); + guard.stats.hits += 1; + return Ok(()); + } + } + + let entry = Ast::new_record(&normalized, query_parser_engine)?; + entry.update_stats(route); + + let mut guard = self.inner.lock(); + guard.queries.put(normalized, entry); + guard.stats.misses += 1; + + Ok(()) + } + + /// Get global cache instance. + pub fn get() -> Self { + CACHE.clone() + } + + /// Get cache stats. + pub fn stats() -> (Stats, usize) { + let cache = Self::get(); + let (len, query_stats, mut stats) = { + let guard = cache.inner.lock(); + ( + guard.queries.len(), + guard + .queries + .iter() + .map(|c| *c.1.stats.lock()) + .collect::>(), + guard.stats, + ) + }; + for stat in query_stats { + stats.direct += stat.direct; + stats.multi += stat.multi; + } + (stats, len) + } + + /// Get a copy of all queries stored in the cache. + pub fn queries() -> HashMap { + Self::get() + .inner + .lock() + .queries + .iter() + .map(|i| (i.0.clone(), i.1.clone())) + .collect() + } + + /// Reset cache, removing all statements + /// and setting stats to 0. + pub fn reset() { + let cache = Self::get(); + let mut guard = cache.inner.lock(); + guard.queries.clear(); + guard.stats.hits = 0; + guard.stats.misses = 0; + } +} diff --git a/pgdog/src/frontend/router/parser/cache/fingerprint.rs b/pgdog/src/frontend/router/parser/cache/fingerprint.rs new file mode 100644 index 000000000..7b7907a28 --- /dev/null +++ b/pgdog/src/frontend/router/parser/cache/fingerprint.rs @@ -0,0 +1,49 @@ +use std::{fmt::Debug, ops::Deref}; + +use pg_query::{fingerprint, fingerprint_raw}; +use pgdog_config::QueryParserEngine; + +/// Query fingerprint. +pub struct Fingerprint { + fingerprint: pg_query::Fingerprint, +} + +impl Fingerprint { + /// Fingerprint a query. + pub(crate) fn new(query: &str, engine: QueryParserEngine) -> Result { + Ok(Self { + fingerprint: match engine { + QueryParserEngine::PgQueryProtobuf => fingerprint(query), + QueryParserEngine::PgQueryRaw => fingerprint_raw(query), + }?, + }) + } +} + +impl Debug for Fingerprint { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Fingerprint") + .field("value", &self.fingerprint.value) + .field("hex", &self.fingerprint.hex) + .finish() + } +} + +impl Default for Fingerprint { + fn default() -> Self { + Self { + fingerprint: pg_query::Fingerprint { + value: 0, + hex: "".into(), + }, + } + } +} + +impl Deref for Fingerprint { + type Target = pg_query::Fingerprint; + + fn deref(&self) -> &Self::Target { + &self.fingerprint + } +} diff --git a/pgdog/src/frontend/router/parser/cache/mod.rs b/pgdog/src/frontend/router/parser/cache/mod.rs new file mode 100644 index 000000000..867467ede --- /dev/null +++ b/pgdog/src/frontend/router/parser/cache/mod.rs @@ -0,0 +1,14 @@ +//! AST cache. +//! +//! Shared between all clients and databases. +//! +pub mod ast; +pub mod cache_impl; +pub mod fingerprint; + +pub use ast::*; +pub use cache_impl::*; +pub use fingerprint::*; + +#[cfg(test)] +pub mod test; diff --git a/pgdog/src/frontend/router/parser/cache/test.rs b/pgdog/src/frontend/router/parser/cache/test.rs new file mode 100644 index 000000000..e3ecb05a9 --- /dev/null +++ b/pgdog/src/frontend/router/parser/cache/test.rs @@ -0,0 +1,122 @@ +use pg_query::{normalize, parse}; +use tokio::spawn; + +use crate::{ + backend::ShardingSchema, + frontend::{BufferedQuery, PreparedStatements}, + net::{Parse, Query}, +}; + +use super::*; +use std::time::{Duration, Instant}; + +#[tokio::test(flavor = "multi_thread")] +async fn bench_ast_cache() { + let query = "SELECT + u.username, + p.product_name, + SUM(oi.quantity * oi.price) AS total_revenue, + AVG(r.rating) AS average_rating, + COUNT(DISTINCT c.country) AS countries_purchased_from + FROM users u + INNER JOIN orders o ON u.user_id = o.user_id + INNER JOIN order_items oi ON o.order_id = oi.order_id + INNER JOIN products p ON oi.product_id = p.product_id + LEFT JOIN reviews r ON o.order_id = r.order_id + LEFT JOIN customer_addresses c ON o.shipping_address_id = c.address_id + WHERE + o.order_date BETWEEN '2023-01-01' AND '2023-12-31' + AND p.category IN ('Electronics', 'Clothing') + AND (r.rating > 4 OR r.rating IS NULL) + GROUP BY u.username, p.product_name + HAVING COUNT(DISTINCT c.country) > 2 + ORDER BY total_revenue DESC; +"; + + let times = 10_000; + let threads = 5; + + let mut tasks = vec![]; + for _ in 0..threads { + let handle = spawn(async move { + let mut parse_time = Duration::ZERO; + for _ in 0..(times / threads) { + let start = Instant::now(); + parse(query).unwrap(); + parse_time += start.elapsed(); + } + + parse_time + }); + tasks.push(handle); + } + + let mut parse_time = Duration::ZERO; + for task in tasks { + parse_time += task.await.unwrap(); + } + + println!("[bench_ast_cache]: parse time: {:?}", parse_time); + + // Simulate lock contention. + let mut tasks = vec![]; + + for _ in 0..threads { + let handle = spawn(async move { + let mut cached_time = Duration::ZERO; + let mut prepared_statements = PreparedStatements::default(); + for _ in 0..(times / threads) { + let start = Instant::now(); + Cache::get() + .query( + &BufferedQuery::Prepared(Parse::new_anonymous(query)), + &ShardingSchema::default(), + &mut prepared_statements, + ) + .unwrap(); + cached_time += start.elapsed(); + } + + cached_time + }); + tasks.push(handle); + } + + let mut cached_time = Duration::ZERO; + for task in tasks { + cached_time += task.await.unwrap(); + } + + println!("[bench_ast_cache]: cached time: {:?}", cached_time); + + let faster = parse_time.as_micros() as f64 / cached_time.as_micros() as f64; + println!( + "[bench_ast_cache]: cached is {:.4} times faster than parsed", + faster + ); // 32x on my M1 + + assert!(faster > 10.0); +} + +#[test] +fn test_normalize() { + let q = "SELECT * FROM users WHERE id = 1"; + let normalized = normalize(q).unwrap(); + assert_eq!(normalized, "SELECT * FROM users WHERE id = $1"); +} + +#[test] +fn test_tables_list() { + let mut prepared_statements = PreparedStatements::default(); + for q in [ + "CREATE TABLE private_schema.test (id BIGINT)", + "SELECT * FROM private_schema.test a INNER JOIN public_schema.test b ON a.id = b.id LIMIT 5", + "INSERT INTO public_schema.test VALUES ($1, $2)", + "DELETE FROM private_schema.test", + "DROP TABLE private_schema.test", + ] { + let ast = Ast::new(&BufferedQuery::Query(Query::new(q)), &ShardingSchema::default(), &mut prepared_statements).unwrap(); + let tables = ast.tables(); + println!("{:?}", tables); + } +} diff --git a/pgdog/src/frontend/router/parser/column.rs b/pgdog/src/frontend/router/parser/column.rs index 66875b464..605660650 100644 --- a/pgdog/src/frontend/router/parser/column.rs +++ b/pgdog/src/frontend/router/parser/column.rs @@ -6,7 +6,7 @@ use pg_query::{ }; use std::fmt::{Display, Formatter, Result as FmtResult}; -use super::Table; +use super::{Error, Table}; use crate::util::escape_identifier; /// Column name extracted from a query. @@ -44,9 +44,7 @@ impl<'a> Column<'a> { pub fn to_owned(&self) -> OwnedColumn { OwnedColumn::from(*self) } -} -impl<'a> Column<'a> { pub fn from_string(string: &'a Node) -> Result { match &string.node { Some(NodeEnum::String(protobuf::String { sval })) => Ok(Self { @@ -57,6 +55,14 @@ impl<'a> Column<'a> { _ => Err(()), } } + + /// Fully-qualify this column with a table. + pub fn qualify(&mut self, table: Table<'a>) { + if self.table.is_none() { + self.table = Some(table.name); + self.schema = table.schema; + } + } } impl<'a> Display for Column<'a> { @@ -114,7 +120,7 @@ impl<'a> From<&'a OwnedColumn> for Column<'a> { } impl<'a> TryFrom<&'a Node> for Column<'a> { - type Error = (); + type Error = Error; fn try_from(value: &'a Node) -> Result { Column::try_from(&value.node) @@ -122,7 +128,7 @@ impl<'a> TryFrom<&'a Node> for Column<'a> { } impl<'a> TryFrom<&'a Option> for Column<'a> { - type Error = (); + type Error = Error; fn try_from(value: &'a Option) -> Result { fn from_node(node: &Node) -> Option<&str> { @@ -133,12 +139,15 @@ impl<'a> TryFrom<&'a Option> for Column<'a> { } } - fn from_slice<'a>(nodes: &'a [Node]) -> Result, ()> { + fn from_slice<'a>(nodes: &'a [Node]) -> Result, Error> { match nodes.len() { 3 => { let schema = nodes.first().and_then(from_node); let table = nodes.get(1).and_then(from_node); - let name = nodes.get(2).and_then(from_node).ok_or(())?; + let name = nodes + .get(2) + .and_then(from_node) + .ok_or(Error::ColumnDecode)?; Ok(Column { schema, @@ -149,7 +158,10 @@ impl<'a> TryFrom<&'a Option> for Column<'a> { 2 => { let table = nodes.first().and_then(from_node); - let name = nodes.get(1).and_then(from_node).ok_or(())?; + let name = nodes + .get(1) + .and_then(from_node) + .ok_or(Error::ColumnDecode)?; Ok(Column { schema: None, @@ -159,7 +171,10 @@ impl<'a> TryFrom<&'a Option> for Column<'a> { } 1 => { - let name = nodes.first().and_then(from_node).ok_or(())?; + let name = nodes + .first() + .and_then(from_node) + .ok_or(Error::ColumnDecode)?; Ok(Column { name, @@ -167,7 +182,7 @@ impl<'a> TryFrom<&'a Option> for Column<'a> { }) } - _ => Err(()), + _ => Err(Error::ColumnDecode), } } @@ -186,26 +201,36 @@ impl<'a> TryFrom<&'a Option> for Column<'a> { if let Some(ref node) = list.arg { Ok(Column::try_from(&node.node)?) } else { - Err(()) + Err(Error::ColumnDecode) } } else { - Err(()) + Err(Error::ColumnDecode) } } - _ => Err(()), + _ => Err(Error::ColumnDecode), } } } impl<'a> TryFrom<&Option<&'a Node>> for Column<'a> { - type Error = (); + type Error = Error; fn try_from(value: &Option<&'a Node>) -> Result { if let Some(value) = value { (*value).try_into() } else { - Err(()) + Err(Error::ColumnDecode) + } + } +} + +impl<'a> From<&'a str> for Column<'a> { + fn from(value: &'a str) -> Self { + Column { + name: value, + table: None, + schema: None, } } } @@ -214,7 +239,7 @@ impl<'a> TryFrom<&Option<&'a Node>> for Column<'a> { mod test { use pg_query::{parse, NodeEnum}; - use super::Column; + use super::{Column, Error}; #[test] fn test_column() { @@ -226,7 +251,7 @@ mod test { .cols .iter() .map(Column::try_from) - .collect::, ()>>() + .collect::, Error>>() .unwrap(); assert_eq!( columns, diff --git a/pgdog/src/frontend/router/parser/command.rs b/pgdog/src/frontend/router/parser/command.rs index 038d930e2..d384f543c 100644 --- a/pgdog/src/frontend/router/parser/command.rs +++ b/pgdog/src/frontend/router/parser/command.rs @@ -24,10 +24,14 @@ pub enum Command { Set { name: String, value: ParameterValue, + local: bool, + route: Route, }, PreparedStatement(Prepare), - Rewrite(String), - Shards(usize), + InternalField { + name: String, + value: String, + }, Deallocate, Discard { extended: bool, @@ -42,17 +46,19 @@ pub enum Command { shard: Shard, }, Unlisten(String), - SetRoute(Route), + UniqueId, } impl Command { pub fn route(&self) -> &Route { lazy_static! { - static ref DEFAULT_ROUTE: Route = Route::write(Shard::All); + static ref DEFAULT_ROUTE: Route = + Route::write(ShardWithPriority::new_default_unset(Shard::All)); } match self { Self::Query(route) => route, + Self::Set { route, .. } => route, _ => &DEFAULT_ROUTE, } } @@ -60,7 +66,9 @@ impl Command { impl Default for Command { fn default() -> Self { - Command::Query(Route::write(Shard::All)) + Command::Query(Route::write(ShardWithPriority::new_default_unset( + Shard::All, + ))) } } @@ -103,11 +111,13 @@ impl Command { pub(crate) fn dry_run(self) -> Self { match self { Command::Query(mut query) => { - query.set_shard_mut(0); + query.set_shard_mut(ShardWithPriority::new_override_dry_run(Shard::Direct(0))); Command::Query(query) } - Command::Copy(_) => Command::Query(Route::write(Some(0))), + Command::Copy(_) => Command::Query(Route::write( + ShardWithPriority::new_override_dry_run(Shard::Direct(0)), + )), _ => self, } } diff --git a/pgdog/src/frontend/router/parser/comment.rs b/pgdog/src/frontend/router/parser/comment.rs index fa2de70d6..0dffacf38 100644 --- a/pgdog/src/frontend/router/parser/comment.rs +++ b/pgdog/src/frontend/router/parser/comment.rs @@ -1,5 +1,7 @@ use once_cell::sync::Lazy; +use pg_query::scan_raw; use pg_query::{protobuf::Token, scan}; +use pgdog_config::QueryParserEngine; use regex::Regex; use crate::backend::ShardingSchema; @@ -29,8 +31,15 @@ fn get_matched_value<'a>(caps: &'a regex::Captures<'a>) -> Option<&'a str> { /// /// See [`SHARD`] and [`SHARDING_KEY`] for the style of comment we expect. /// -pub fn comment(query: &str, schema: &ShardingSchema) -> Result<(Shard, Option), Error> { - let tokens = scan(query).map_err(Error::PgQuery)?; +pub fn comment( + query: &str, + schema: &ShardingSchema, +) -> Result<(Option, Option), Error> { + let tokens = match schema.query_parser_engine { + QueryParserEngine::PgQueryProtobuf => scan(query), + QueryParserEngine::PgQueryRaw => scan_raw(query), + } + .map_err(Error::PgQuery)?; let mut role = None; for token in tokens.tokens.iter() { @@ -47,21 +56,26 @@ pub fn comment(query: &str, schema: &ShardingSchema) -> Result<(Shard, Option() - .ok() - .map(Shard::Direct) - .unwrap_or(Shard::All), + Some( + shard + .as_str() + .parse::() + .ok() + .map(Shard::Direct) + .unwrap_or(Shard::All), + ), role, )); } @@ -69,7 +83,7 @@ pub fn comment(query: &str, schema: &ShardingSchema) -> Result<(Shard, Option { pub(super) router_context: RouterContext<'a>, /// How aggressively we want to send reads to replicas. pub(super) rw_strategy: &'a ReadWriteStrategy, - /// Are we re-writing prepared statements sent over the simple protocol? - pub(super) full_prepared_statements: bool, /// Do we need the router at all? Shortcut to bypass this for unsharded /// clusters with databases that only read or write. pub(super) router_needed: bool, - /// Do we have support for LISTEN/NOTIFY enabled? - pub(super) pub_sub_enabled: bool, /// Are we running multi-tenant checks? pub(super) multi_tenant: &'a Option, /// Dry run enabled? pub(super) dry_run: bool, + /// Expanded EXPLAIN annotations enabled? + pub(super) expanded_explain: bool, + /// Shards calculator. + pub(super) shards_calculator: ShardsWithPriority, } impl<'a> QueryParserContext<'a> { /// Create query parser context from router context. - pub fn new(router_context: RouterContext<'a>) -> Self { - let config = config(); - Self { + pub fn new(router_context: RouterContext<'a>) -> Result { + let mut shards_calculator = ShardsWithPriority::default(); + let sharding_schema = router_context.cluster.sharding_schema(); + + router_context + .parameter_hints + .compute_shard(&mut shards_calculator, &sharding_schema)?; + + Ok(Self { read_only: router_context.cluster.read_only(), write_only: router_context.cluster.write_only(), shards: router_context.cluster.shards().len(), - sharding_schema: router_context.cluster.sharding_schema(), + sharding_schema, rw_strategy: router_context.cluster.read_write_strategy(), - full_prepared_statements: config.prepared_statements_full(), router_needed: router_context.cluster.router_needed(), - pub_sub_enabled: config.config.general.pub_sub_enabled(), multi_tenant: router_context.cluster.multi_tenant(), - dry_run: config.config.general.dry_run, + dry_run: router_context.cluster.dry_run(), + expanded_explain: router_context.cluster.expanded_explain(), router_context, - } + shards_calculator, + }) } /// Write override enabled? @@ -71,6 +79,7 @@ impl<'a> QueryParserContext<'a> { self.router_context.transaction(), Some(TransactionType::ReadWrite) ) && self.rw_conservative() + || self.router_context.parameter_hints.compute_role() == Some(Role::Primary) } /// Are we using the conservative read/write separation strategy? @@ -82,11 +91,7 @@ impl<'a> QueryParserContext<'a> { /// /// Shortcut to avoid the overhead if we can. pub(super) fn use_parser(&self) -> bool { - self.full_prepared_statements - || self.router_needed - || self.pub_sub_enabled - || self.multi_tenant().is_some() - || self.dry_run + self.router_context.cluster.use_query_parser() } /// Get the query we're parsing, if any. @@ -94,11 +99,6 @@ impl<'a> QueryParserContext<'a> { self.router_context.query.as_ref().ok_or(Error::EmptyQuery) } - /// Mutable reference to client's prepared statements cache. - pub(super) fn prepared_statements(&mut self) -> &mut PreparedStatements { - self.router_context.prepared_statements - } - /// Multi-tenant checks. pub(super) fn multi_tenant(&self) -> &Option { self.multi_tenant @@ -136,4 +136,8 @@ impl<'a> QueryParserContext<'a> { params, } } + + pub(super) fn expanded_explain(&self) -> bool { + self.expanded_explain + } } diff --git a/pgdog/src/frontend/router/parser/copy.rs b/pgdog/src/frontend/router/parser/copy.rs index b5c0bc5da..f2875e614 100644 --- a/pgdog/src/frontend/router/parser/copy.rs +++ b/pgdog/src/frontend/router/parser/copy.rs @@ -60,7 +60,7 @@ pub struct CopyParser { delimiter: Option, /// Number of columns columns: usize, - /// This is a COPY coming from the server. + /// This is a COPY coming from the client. is_from: bool, /// Stream parser. stream: CopyStream, @@ -70,6 +70,8 @@ pub struct CopyParser { sharded_table: Option, /// The sharding column is in this position in each row. sharded_column: usize, + /// Schema shard. + schema_shard: Option, } impl Default for CopyParser { @@ -83,13 +85,14 @@ impl Default for CopyParser { sharding_schema: ShardingSchema::default(), sharded_table: None, sharded_column: 0, + schema_shard: None, } } } impl CopyParser { /// Create new copy parser from a COPY statement. - pub fn new(stmt: &CopyStmt, cluster: &Cluster) -> Result, Error> { + pub fn new(stmt: &CopyStmt, cluster: &Cluster) -> Result { let mut parser = Self { is_from: stmt.is_from, ..Default::default() @@ -109,6 +112,13 @@ impl CopyParser { let table = Table::from(rel); + // The CopyParser is used for replicating + // data during data-sync. This will ensure all rows + // are sent to the right schema-based shard. + if let Some(schema) = cluster.sharding_schema().schemas.get(table.schema()) { + parser.schema_shard = Some(schema.shard().into()); + } + if let Some(key) = Tables::new(&cluster.sharding_schema()).key(table, &columns) { parser.sharded_table = Some(key.table.clone()); parser.sharded_column = key.position; @@ -178,7 +188,7 @@ impl CopyParser { }; parser.sharding_schema = cluster.sharding_schema(); - Ok(Some(parser)) + Ok(parser) } #[inline] @@ -202,7 +212,10 @@ impl CopyParser { if self.headers && self.is_from { let headers = stream.headers()?; if let Some(headers) = headers { - rows.push(CopyRow::new(headers.to_string().as_bytes(), Shard::All)); + rows.push(CopyRow::new( + headers.to_string().as_bytes(), + self.schema_shard.clone().unwrap_or(Shard::All), + )); } self.headers = false; } @@ -222,6 +235,8 @@ impl CopyParser { .build()?; ctx.apply()? + } else if let Some(schema_shard) = self.schema_shard.clone() { + schema_shard } else { Shard::All }; @@ -233,7 +248,10 @@ impl CopyParser { CopyStream::Binary(stream) => { if self.headers { let header = stream.header()?; - rows.push(CopyRow::new(&header.to_bytes()?, Shard::All)); + rows.push(CopyRow::new( + &header.to_bytes()?, + self.schema_shard.clone().unwrap_or(Shard::All), + )); self.headers = false; } @@ -241,7 +259,10 @@ impl CopyParser { let tuple = tuple?; if tuple.end() { let terminator = (-1_i16).to_be_bytes(); - rows.push(CopyRow::new(&terminator, Shard::All)); + rows.push(CopyRow::new( + &terminator, + self.schema_shard.clone().unwrap_or(Shard::All), + )); break; } let shard = if let Some(table) = &self.sharded_table { @@ -258,6 +279,8 @@ impl CopyParser { } else { Shard::All } + } else if let Some(schema_shard) = self.schema_shard.clone() { + schema_shard } else { Shard::All }; @@ -276,6 +299,8 @@ impl CopyParser { mod test { use pg_query::parse; + use crate::config::config; + use super::*; #[test] @@ -288,9 +313,7 @@ mod test { _ => panic!("not a copy"), }; - let mut copy = CopyParser::new(©, &Cluster::default()) - .unwrap() - .unwrap(); + let mut copy = CopyParser::new(©, &Cluster::default()).unwrap(); assert_eq!(copy.delimiter(), '\t'); assert!(!copy.headers); @@ -312,9 +335,7 @@ mod test { _ => panic!("not a copy"), }; - let mut copy = CopyParser::new(©, &Cluster::default()) - .unwrap() - .unwrap(); + let mut copy = CopyParser::new(©, &Cluster::default()).unwrap(); assert!(copy.is_from); assert_eq!(copy.delimiter(), ','); @@ -353,9 +374,7 @@ mod test { _ => panic!("not a copy"), }; - let mut copy = CopyParser::new(©, &Cluster::new_test()) - .unwrap() - .unwrap(); + let mut copy = CopyParser::new(©, &Cluster::new_test(&config())).unwrap(); let rows = copy.shard(&[copy_data]).unwrap(); assert_eq!(rows.len(), 3); @@ -377,9 +396,7 @@ mod test { _ => panic!("not a copy"), }; - let mut copy = CopyParser::new(©, &Cluster::default()) - .unwrap() - .unwrap(); + let mut copy = CopyParser::new(©, &Cluster::default()).unwrap(); assert_eq!(copy.delimiter(), ','); assert!(!copy.headers); @@ -403,9 +420,7 @@ mod test { _ => panic!("not a copy"), }; - let mut copy = CopyParser::new(©, &Cluster::default()) - .unwrap() - .unwrap(); + let mut copy = CopyParser::new(©, &Cluster::default()).unwrap(); assert!(copy.is_from); assert!(copy.headers); let mut data = b"PGCOPY".to_vec(); diff --git a/pgdog/src/frontend/router/parser/error.rs b/pgdog/src/frontend/router/parser/error.rs index c4f9e422b..06e148bda 100644 --- a/pgdog/src/frontend/router/parser/error.rs +++ b/pgdog/src/frontend/router/parser/error.rs @@ -2,6 +2,7 @@ use thiserror::Error; +use super::rewrite::statement::Error as RewriteError; use crate::frontend::router::sharding; #[derive(Debug, Error)] @@ -68,4 +69,28 @@ pub enum Error { #[error("regex error")] RegexError, + + #[error("cross-shard truncate not supported when schema-sharding is used")] + CrossShardTruncateSchemaSharding, + + #[error("prepared statement \"{0}\" doesn't exist")] + PreparedStatementDoesntExist(String), + + #[error("column decode error")] + ColumnDecode, + + #[error("table decode error")] + TableDecode, + + #[error("parameter ${0} not in bind")] + BindParameterMissing(i32), + + #[error("statement is not a SELECT")] + NotASelect, + + #[error("rewrite: {0}")] + Rewrite(#[from] RewriteError), + + #[error("sharded databases require the query parser to be enabled")] + QueryParserRequired, } diff --git a/pgdog/src/frontend/router/parser/explain_trace.rs b/pgdog/src/frontend/router/parser/explain_trace.rs new file mode 100644 index 000000000..f9ea25dff --- /dev/null +++ b/pgdog/src/frontend/router/parser/explain_trace.rs @@ -0,0 +1,236 @@ +use pgdog_config::Role; + +use crate::frontend::router::parser::route::Shard; + +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct ExplainTrace { + summary: ExplainSummary, + steps: Vec, +} + +impl ExplainTrace { + pub fn new(summary: ExplainSummary, steps: Vec) -> Self { + Self { summary, steps } + } + + pub fn summary(&self) -> &ExplainSummary { + &self.summary + } + + pub fn steps(&self) -> &[ExplainEntry] { + &self.steps + } + + pub fn render_lines(&self) -> Vec { + let mut lines = vec![String::new(), "PgDog Routing:".to_string()]; + lines.push(format!( + " Summary: shard={} role={}", + self.summary.shard, + if self.summary.read { + "replica" + } else { + "primary" + } + )); + + for entry in &self.steps { + lines.push(entry.render_line()); + } + + lines + } +} + +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct ExplainSummary { + pub shard: Shard, + pub read: bool, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ExplainEntry { + pub shard: Option, + pub description: String, +} + +impl ExplainEntry { + pub fn new(shard: Option, description: impl Into) -> Self { + Self { + shard, + description: description.into(), + } + } + + fn render_line(&self) -> String { + match &self.shard { + Some(Shard::Direct(shard)) => { + format!(" Shard {}: {}", shard, self.description) + } + Some(Shard::Multi(shards)) => { + format!(" Shards {:?}: {}", shards, self.description) + } + Some(Shard::All) => format!(" All shards: {}", self.description), + None => format!(" Note: {}", self.description), + } + } +} + +/// EXPLAIN recorder. +#[derive(Debug, Default)] +pub struct ExplainRecorder { + entries: Vec, + comment: Option, + plugin: Option, +} + +impl ExplainRecorder { + pub fn new() -> Self { + Self::default() + } + + pub fn record_entry(&mut self, shard: Option, description: impl Into) { + self.entries.push(ExplainEntry::new(shard, description)); + } + + pub fn clear(&mut self) { + self.entries.clear(); + self.comment = None; + self.plugin = None; + } + + pub fn record_comment_override(&mut self, shard: Shard, role: Option) { + let mut description = match shard { + Shard::Direct(_) | Shard::Multi(_) | Shard::All => { + format!("manual override to shard={}", shard) + } + }; + + if let Some(role) = role { + description.push_str(&format!(" role={}", role)); + } + + self.comment = Some(ExplainEntry::new(Some(shard), description)); + } + + pub fn record_plugin_override( + &mut self, + plugin: impl Into, + shard: Option, + read: Option, + ) { + let mut description = format!("plugin {} adjusted routing", plugin.into()); + if let Some(shard) = &shard { + description.push_str(&format!(" shard={}", shard)); + } + if let Some(read) = read { + description.push_str(&format!( + " role={}", + if read { "replica" } else { "primary" } + )); + } + self.plugin = Some(ExplainEntry::new(shard, description)); + } + + pub fn finalize(mut self, summary: ExplainSummary) -> ExplainTrace { + if let Some(comment) = self.comment.take() { + self.entries.insert(0, comment); + } + + if let Some(plugin) = self.plugin.take() { + self.entries.push(plugin); + } + + if self.entries.is_empty() { + let description = match summary.shard { + Shard::All => "no sharding key matched; broadcasting".to_string(), + Shard::Multi(ref shards) if !shards.is_empty() => { + format!("multiple shards matched: {:?}", shards) + } + Shard::Multi(_) => "multiple shards matched".to_string(), + Shard::Direct(_) => "direct routing without recorded hints".to_string(), + }; + self.entries.push(ExplainEntry::new(None, description)); + } + + ExplainTrace::new(summary, self.entries) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn render_lines_formats_summary_and_entries() { + let trace = ExplainTrace::new( + ExplainSummary { + shard: Shard::Direct(2), + read: true, + }, + vec![ExplainEntry::new( + Some(Shard::Direct(2)), + "matched sharding key", + )], + ); + + let lines = trace.render_lines(); + assert_eq!(lines[0], ""); + assert_eq!(lines[1], "PgDog Routing:"); + assert_eq!(lines[2], " Summary: shard=2 role=replica"); + assert_eq!(lines[3], " Shard 2: matched sharding key"); + } + + #[test] + fn finalize_inserts_comment_and_plugin_entries() { + let mut recorder = ExplainRecorder::new(); + recorder.record_entry(Some(Shard::Direct(7)), "matched sharding key"); + recorder.record_comment_override(Shard::Direct(3), Some(Role::Primary)); + recorder.record_plugin_override("test_plugin", Some(Shard::Direct(9)), Some(true)); + + let trace = recorder.finalize(ExplainSummary { + shard: Shard::Direct(9), + read: true, + }); + + let descriptions: Vec<&str> = trace + .steps() + .iter() + .map(|entry| entry.description.as_str()) + .collect(); + + assert_eq!(descriptions[0], "manual override to shard=3 role=primary"); + assert_eq!(descriptions[1], "matched sharding key"); + assert_eq!( + descriptions[2], + "plugin test_plugin adjusted routing shard=9 role=replica" + ); + } + + #[test] + fn finalize_injects_fallback_when_no_entries() { + let trace = ExplainRecorder::new().finalize(ExplainSummary { + shard: Shard::All, + read: false, + }); + + assert_eq!(trace.steps().len(), 1); + assert_eq!( + trace.steps()[0].description, + "no sharding key matched; broadcasting" + ); + assert!(trace.steps()[0].shard.is_none()); + } + + #[test] + fn finalize_reports_multiple_shards() { + let trace = ExplainRecorder::new().finalize(ExplainSummary { + shard: Shard::Multi(vec![1, 5]), + read: true, + }); + + assert_eq!( + trace.steps()[0].description, + "multiple shards matched: [1, 5]" + ); + } +} diff --git a/pgdog/src/frontend/router/parser/expression.rs b/pgdog/src/frontend/router/parser/expression.rs index fd87a0b95..d7fc5851c 100644 --- a/pgdog/src/frontend/router/parser/expression.rs +++ b/pgdog/src/frontend/router/parser/expression.rs @@ -83,10 +83,6 @@ impl CanonicalExpr { .collect::>(); if func.agg_distinct { - // DISTINCT aggregate arguments form a set, so sort them to - // ensure canonical equality across permutations. Non-DISTINCT - // functions must retain argument order because most built-ins - // are not commutative. args.sort(); } @@ -249,14 +245,4 @@ mod tests { let regular = registry.intern(&targets[1]); assert_ne!(distinct, regular); } - - #[test] - fn registry_preserves_argument_order_for_functions() { - let mut registry = ExpressionRegistry::new(); - let targets = extract_targets("SELECT substr(name, 1, 2), substr(name, 2, 1) FROM menu"); - assert_eq!(targets.len(), 2); - let first = registry.intern(&targets[0]); - let second = registry.intern(&targets[1]); - assert_ne!(first, second); - } } diff --git a/pgdog/src/frontend/router/parser/from_clause.rs b/pgdog/src/frontend/router/parser/from_clause.rs index 9b80849ad..bf0fbf074 100644 --- a/pgdog/src/frontend/router/parser/from_clause.rs +++ b/pgdog/src/frontend/router/parser/from_clause.rs @@ -3,7 +3,7 @@ use pg_query::{Node, NodeEnum}; use super::*; /// Handle FROM clause. -#[derive(Copy, Clone)] +#[derive(Copy, Clone, Debug)] pub struct FromClause<'a> { nodes: &'a [Node], } @@ -64,4 +64,37 @@ impl<'a> FromClause<'a> { None } + + /// Get all tables from the FROM clause. + pub fn tables(&'a self) -> Vec> { + let mut tables = vec![]; + + fn tables_recursive(node: &Node) -> Vec> { + let mut tables = vec![]; + match node.node { + Some(NodeEnum::RangeVar(ref range_var)) => { + tables.push(Table::from(range_var)); + } + + Some(NodeEnum::JoinExpr(ref join)) => { + if let Some(ref node) = join.larg { + tables.extend(tables_recursive(node)); + } + if let Some(ref node) = join.rarg { + tables.extend(tables_recursive(node)); + } + } + + _ => (), + } + + tables + } + + for node in self.nodes { + tables.extend(tables_recursive(node)); + } + + tables + } } diff --git a/pgdog/src/frontend/router/parser/insert.rs b/pgdog/src/frontend/router/parser/insert.rs index 02942b0eb..0fe0b3908 100644 --- a/pgdog/src/frontend/router/parser/insert.rs +++ b/pgdog/src/frontend/router/parser/insert.rs @@ -31,7 +31,7 @@ impl<'a> Insert<'a> { .cols .iter() .map(Column::try_from) - .collect::>, ()>>() + .collect::>, Error>>() .ok() .unwrap_or(vec![]) } @@ -57,6 +57,17 @@ impl<'a> Insert<'a> { vec![] } + /// Calculate the number of tuples in the statement. + pub fn num_tuples(&self) -> usize { + if let Some(select) = &self.stmt.select_stmt { + if let Some(NodeEnum::SelectStmt(stmt)) = &select.node { + return stmt.values_lists.len(); + } + } + + 0 + } + /// Get the sharding key for the statement. pub fn shard( &'a self, @@ -66,29 +77,39 @@ impl<'a> Insert<'a> { let tables = Tables::new(schema); let columns = self.columns(); let table = self.table(); + let tuples = self.tuples(); + let key = table.and_then(|table| tables.key(table, &columns)); + if let Some(table) = table { + // Schema-based routing. + if let Some(schema) = schema.schemas.get(table.schema()) { + return Ok(schema.shard().into()); + } + } + + if self.num_tuples() != 1 { + debug!("multiple tuples in an INSERT statement"); + return Ok(Shard::All); + } + if let Some(key) = key { if let Some(bind) = bind { if let Ok(Some(param)) = bind.parameter(key.position) { - // Arrays not supported as sharding keys at the moment. - let value = ShardingValue::from_param(¶m, key.table.data_type)?; - let ctx = ContextBuilder::new(key.table) - .value(value) - .shards(schema.shards) - .build()?; - return Ok(ctx.apply()?); + if param.is_null() { + return Ok(Shard::All); + } else { + // Arrays not supported as sharding keys at the moment. + let value = ShardingValue::from_param(¶m, key.table.data_type)?; + let ctx = ContextBuilder::new(key.table) + .value(value) + .shards(schema.shards) + .build()?; + return Ok(ctx.apply()?); + } } } - let tuples = self.tuples(); - - // TODO: support rewriting INSERTs to run against multiple shards. - if tuples.len() != 1 { - debug!("multiple tuples in an INSERT statement"); - return Ok(Shard::All); - } - if let Some(value) = tuples.first().and_then(|tuple| tuple.get(key.position)) { match value { Value::Integer(int) => { @@ -130,6 +151,7 @@ mod test { use crate::config::ShardedTable; use crate::net::bind::Parameter; use crate::net::Format; + use bytes::Bytes; use super::super::Value; use super::*; @@ -231,6 +253,7 @@ mod test { ], vec![], ), + ..Default::default() }; match &select.node { @@ -243,7 +266,7 @@ mod test { "", &[Parameter { len: 1, - data: "3".as_bytes().to_vec(), + data: "3".as_bytes().into(), }], ); @@ -254,7 +277,7 @@ mod test { "", &[Parameter { len: 8, - data: 234_i64.to_be_bytes().to_vec(), + data: Bytes::copy_from_slice(&234_i64.to_be_bytes()), }], &[Format::Binary], ); @@ -306,4 +329,32 @@ mod test { _ => panic!("not a select"), } } + + #[test] + fn test_null_sharding_key_routes_to_all() { + let query = parse("INSERT INTO sharded (id, value) VALUES ($1, 'test')").unwrap(); + let select = query.protobuf.stmts.first().unwrap().stmt.as_ref().unwrap(); + let schema = ShardingSchema { + shards: 3, + tables: ShardedTables::new( + vec![ShardedTable { + name: Some("sharded".into()), + column: "id".into(), + ..Default::default() + }], + vec![], + ), + ..Default::default() + }; + + match &select.node { + Some(NodeEnum::InsertStmt(stmt)) => { + let insert = Insert::new(stmt); + let bind = Bind::new_params("", &[Parameter::new_null()]); + let shard = insert.shard(&schema, Some(&bind)).unwrap(); + assert!(matches!(shard, Shard::All)); + } + _ => panic!("not an insert"), + } + } } diff --git a/pgdog/src/frontend/router/parser/mod.rs b/pgdog/src/frontend/router/parser/mod.rs index c4ab263d1..8fad8c44e 100644 --- a/pgdog/src/frontend/router/parser/mod.rs +++ b/pgdog/src/frontend/router/parser/mod.rs @@ -11,6 +11,7 @@ pub mod copy; pub mod csv; pub mod distinct; pub mod error; +pub mod explain_trace; mod expression; pub mod from_clause; pub mod function; @@ -22,10 +23,10 @@ pub mod order_by; pub mod prepare; pub mod query; pub mod rewrite; -pub mod rewrite_engine; -pub mod rewrite_plan; pub mod route; +pub mod schema; pub mod sequence; +pub mod statement; pub mod table; pub mod tuple; pub mod value; @@ -35,7 +36,7 @@ pub use expression::ExpressionRegistry; pub use aggregate::{Aggregate, AggregateFunction, AggregateTarget}; pub use binary::BinaryStream; -pub use cache::Cache; +pub use cache::{Ast, Cache}; pub use column::{Column, OwnedColumn}; pub use command::Command; pub use context::QueryParserContext; @@ -51,10 +52,13 @@ pub use limit::{Limit, LimitClause}; pub use order_by::OrderBy; pub use prepare::Prepare; pub use query::QueryParser; -pub use rewrite_engine::RewriteEngine; -pub use rewrite_plan::{HelperMapping, QueryRewriter, RewriteOutput, RewritePlan}; -pub use route::{Route, Shard}; +pub use rewrite::{ + Assignment, AssignmentValue, ShardKeyRewritePlan, StatementRewrite, StatementRewriteContext, +}; +pub use route::{Route, Shard, ShardWithPriority, ShardsWithPriority}; +pub use schema::Schema; pub use sequence::{OwnedSequence, Sequence}; +pub use statement::StatementParser; pub use table::{OwnedTable, Table}; pub use tuple::Tuple; pub use value::Value; diff --git a/pgdog/src/frontend/router/parser/multi_tenant.rs b/pgdog/src/frontend/router/parser/multi_tenant.rs index 8ce29f749..d9f7ef736 100644 --- a/pgdog/src/frontend/router/parser/multi_tenant.rs +++ b/pgdog/src/frontend/router/parser/multi_tenant.rs @@ -8,7 +8,7 @@ use crate::{ router::parser::{where_clause::TablesSource, Table, WhereClause}, SearchPath, }, - net::Parameters, + net::parameter::ParameterValue, }; pub struct MultiTenantCheck<'a> { @@ -16,7 +16,7 @@ pub struct MultiTenantCheck<'a> { config: &'a MultiTenant, schema: Schema, ast: &'a ParseResult, - parameters: &'a Parameters, + search_path: Option<&'a ParameterValue>, } impl<'a> MultiTenantCheck<'a> { @@ -25,13 +25,13 @@ impl<'a> MultiTenantCheck<'a> { config: &'a MultiTenant, schema: Schema, ast: &'a ParseResult, - parameters: &'a Parameters, + search_path: Option<&'a ParameterValue>, ) -> Self { Self { config, schema, ast, - parameters, + search_path, user, } } @@ -79,7 +79,7 @@ impl<'a> MultiTenantCheck<'a> { } fn check(&self, table: Table, where_clause: Option) -> Result<(), Error> { - let search_path = SearchPath::new(self.user, self.parameters, &self.schema); + let search_path = SearchPath::new(self.user, self.search_path, &self.schema); let schemas = search_path.resolve(); for schema in schemas { @@ -106,3 +106,63 @@ impl<'a> MultiTenantCheck<'a> { Ok(()) } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::backend::schema::{columns::Column, Relation, Schema}; + use std::collections::HashMap; + + fn schema_with_tenant_column(column: &str) -> Schema { + let mut columns = HashMap::new(); + columns.insert( + column.to_string(), + Column { + table_catalog: "catalog".into(), + table_schema: "public".into(), + table_name: "accounts".into(), + column_name: column.into(), + column_default: String::new(), + is_nullable: false, + data_type: "bigint".into(), + }, + ); + + let relation = Relation::test_table("public", "accounts", columns); + let mut relations = HashMap::new(); + relations.insert(("public".into(), "accounts".into()), relation); + + Schema::from_parts(vec!["$user".into(), "public".into()], relations) + } + + #[test] + fn multi_tenant_check_passes_with_matching_filter() { + let schema = schema_with_tenant_column("tenant_id"); + let ast = pg_query::parse("SELECT * FROM accounts WHERE tenant_id = 1") + .expect("parse select statement"); + let config = MultiTenant { + column: "tenant_id".into(), + }; + + let check = MultiTenantCheck::new("alice", &config, schema, &ast, None); + assert!(check.run().is_ok()); + } + + #[test] + fn multi_tenant_check_requires_tenant_column_in_filter() { + let schema = schema_with_tenant_column("tenant_id"); + let ast = pg_query::parse("SELECT * FROM accounts WHERE other_id = 1") + .expect("parse select statement"); + let config = MultiTenant { + column: "tenant_id".into(), + }; + + let check = MultiTenantCheck::new("alice", &config, schema, &ast, None); + let err = check + .run() + .expect_err("expected tenant id validation error"); + matches!(err, Error::MultiTenantId) + .then_some(()) + .expect("should return multi-tenant id error"); + } +} diff --git a/pgdog/src/frontend/router/parser/prepare.rs b/pgdog/src/frontend/router/parser/prepare.rs index 7feaff225..f727677e5 100644 --- a/pgdog/src/frontend/router/parser/prepare.rs +++ b/pgdog/src/frontend/router/parser/prepare.rs @@ -1,5 +1,7 @@ -use super::Error; use pg_query::protobuf::PrepareStmt; +use pgdog_config::QueryParserEngine; + +use super::Error; #[derive(Debug, Clone, PartialEq)] pub struct Prepare { @@ -7,16 +9,17 @@ pub struct Prepare { statement: String, } -impl TryFrom<&PrepareStmt> for Prepare { - type Error = super::Error; - - fn try_from(value: &PrepareStmt) -> Result { - let statement = value - .query - .as_ref() - .ok_or(Error::EmptyQuery)? - .deparse() - .map_err(|_| Error::EmptyQuery)?; +impl Prepare { + pub fn from_stmt( + value: &PrepareStmt, + query_parser_engine: QueryParserEngine, + ) -> Result { + let query = value.query.as_ref().ok_or(Error::EmptyQuery)?; + let statement = match query_parser_engine { + QueryParserEngine::PgQueryProtobuf => query.deparse(), + QueryParserEngine::PgQueryRaw => query.deparse_raw(), + } + .map_err(|_| Error::EmptyQuery)?; Ok(Self { name: value.name.to_string(), @@ -44,7 +47,8 @@ mod test { .unwrap(); match ast.node.unwrap() { NodeEnum::PrepareStmt(stmt) => { - let prepare = Prepare::try_from(stmt.as_ref()).unwrap(); + let prepare = + Prepare::from_stmt(stmt.as_ref(), QueryParserEngine::PgQueryProtobuf).unwrap(); assert_eq!(prepare.name, "test"); assert_eq!(prepare.statement, "SELECT $1, $2"); } diff --git a/pgdog/src/frontend/router/parser/query/ddl.rs b/pgdog/src/frontend/router/parser/query/ddl.rs new file mode 100644 index 000000000..23ed84991 --- /dev/null +++ b/pgdog/src/frontend/router/parser/query/ddl.rs @@ -0,0 +1,635 @@ +use pg_query::parse; + +use super::*; + +impl QueryParser { + /// Handle DDL, e.g. CREATE, DROP, ALTER, etc. + pub(super) fn ddl( + &mut self, + node: &Option, + context: &mut QueryParserContext<'_>, + ) -> Result { + let command = Self::shard_ddl( + node, + &context.sharding_schema, + &mut context.shards_calculator, + )?; + + Ok(command) + } + + pub(super) fn shard_ddl( + node: &Option, + schema: &ShardingSchema, + calculator: &mut ShardsWithPriority, + ) -> Result { + let mut shard = Shard::All; + + match node { + Some(NodeEnum::CreateStmt(stmt)) => { + shard = Self::shard_ddl_table(&stmt.relation, schema)?.unwrap_or(Shard::All); + } + + Some(NodeEnum::CreateSeqStmt(stmt)) => { + shard = Self::shard_ddl_table(&stmt.sequence, schema)?.unwrap_or(Shard::All); + } + + Some(NodeEnum::DropStmt(stmt)) => match stmt.remove_type() { + ObjectType::ObjectTable + | ObjectType::ObjectIndex + | ObjectType::ObjectView + | ObjectType::ObjectSequence => { + let table = Table::try_from(&stmt.objects).ok(); + if let Some(table) = table { + if let Some(schema) = schema.schemas.get(table.schema()) { + shard = schema.shard().into(); + } + } + } + + ObjectType::ObjectSchema => { + if let Some(Node { + node: Some(NodeEnum::String(string)), + }) = stmt.objects.first() + { + if let Some(schema) = schema.schemas.get(Some(string.sval.as_str().into())) + { + shard = schema.shard().into(); + } + } + } + + _ => (), + }, + + Some(NodeEnum::CreateSchemaStmt(stmt)) => { + if let Some(schema) = schema.schemas.get(Some(stmt.schemaname.as_str().into())) { + shard = schema.shard().into(); + } + } + + Some(NodeEnum::IndexStmt(stmt)) => { + shard = Self::shard_ddl_table(&stmt.relation, schema)?.unwrap_or(Shard::All); + } + + Some(NodeEnum::ViewStmt(stmt)) => { + shard = Self::shard_ddl_table(&stmt.view, schema)?.unwrap_or(Shard::All); + } + + Some(NodeEnum::CreateTableAsStmt(stmt)) => { + if let Some(into) = &stmt.into { + shard = Self::shard_ddl_table(&into.rel, schema)?.unwrap_or(Shard::All); + } + } + + Some(NodeEnum::CreateFunctionStmt(stmt)) => { + let table = Table::try_from(&stmt.funcname).ok(); + if let Some(table) = table { + shard = schema + .schemas + .get(table.schema()) + .map(|schema| schema.shard().into()) + .unwrap_or(Shard::All); + } + } + + Some(NodeEnum::CreateEnumStmt(stmt)) => { + let table = Table::try_from(&stmt.type_name).ok(); + if let Some(table) = table { + shard = schema + .schemas + .get(table.schema()) + .map(|schema| schema.shard().into()) + .unwrap_or(Shard::All); + } + } + + Some(NodeEnum::AlterOwnerStmt(stmt)) => { + shard = Self::shard_ddl_table(&stmt.relation, schema)?.unwrap_or(Shard::All); + } + + Some(NodeEnum::RenameStmt(stmt)) => { + shard = Self::shard_ddl_table(&stmt.relation, schema)?.unwrap_or(Shard::All); + } + + Some(NodeEnum::AlterTableStmt(stmt)) => { + shard = Self::shard_ddl_table(&stmt.relation, schema)?.unwrap_or(Shard::All); + } + + Some(NodeEnum::AlterSeqStmt(stmt)) => { + shard = Self::shard_ddl_table(&stmt.sequence, schema)?.unwrap_or(Shard::All); + } + + Some(NodeEnum::LockStmt(stmt)) => { + if let Some(node) = stmt.relations.first() { + if let Some(NodeEnum::RangeVar(ref table)) = node.node { + let table = Table::from(table); + shard = schema + .schemas + .get(table.schema()) + .map(|schema| schema.shard().into()) + .unwrap_or(Shard::All); + } + } + } + + Some(NodeEnum::VacuumStmt(stmt)) => { + for rel in &stmt.rels { + if let Some(NodeEnum::VacuumRelation(ref stmt)) = rel.node { + shard = + Self::shard_ddl_table(&stmt.relation, schema)?.unwrap_or(Shard::All); + } + } + } + + Some(NodeEnum::VacuumRelation(stmt)) => { + shard = Self::shard_ddl_table(&stmt.relation, schema)?.unwrap_or(Shard::All); + } + + // DO $$ BEGIN ... END + Some(NodeEnum::DoStmt(stmt)) => { + if let Some(inner) = stmt.args.first() { + if let Some(NodeEnum::DefElem(ref elem)) = inner.node { + if let Some(ref arg) = elem.arg { + if let Some(NodeEnum::String(ref string)) = arg.node { + // Parse each statement individually. + // The first DDL statement to return a direct shard will be used. + // TODO: handle non-DDL statements in here as well, + // need a full recursive call back to QueryParser::query basically, but that requires a refactor. + for stmt in string.sval.lines() { + if let Ok(stmt) = parse(stmt) { + if let Some(node) = stmt + .protobuf + .stmts + .first() + .map(|stmt| &stmt.stmt) + .cloned() + .flatten() + { + // Use a fresh calculator for each inner statement + // to avoid pollution from statements that don't match + // any DDL pattern (like BEGIN, END, etc.) + let mut inner_calculator = + ShardsWithPriority::default(); + let command = Self::shard_ddl( + &node.node, + schema, + &mut inner_calculator, + )?; + if let Command::Query(query) = command { + if !query.is_cross_shard() { + shard = query.shard().clone(); + break; + } + } + } + } + } + } + } + } + } + } + + Some(NodeEnum::TruncateStmt(stmt)) => { + let mut shards = HashSet::new(); + for relation in &stmt.relations { + if let Some(NodeEnum::RangeVar(ref relation)) = relation.node { + shards.insert( + Self::shard_ddl_table(&Some(relation.clone()), schema)? + .unwrap_or(Shard::All), + ); + } + } + + match shards.len() { + 0 => (), + 1 => { + shard = shards.iter().next().unwrap().clone(); + } + _ => return Err(Error::CrossShardTruncateSchemaSharding), + } + } + + // All others are not handled. + // They are sent to all shards concurrently. + _ => (), + }; + + calculator.push(ShardWithPriority::new_table(shard)); + + Ok(Command::Query(Route::write(calculator.shard()))) + } + + pub(super) fn shard_ddl_table( + range_var: &Option, + schema: &ShardingSchema, + ) -> Result, Error> { + let table = range_var.as_ref().map(Table::from); + if let Some(table) = table { + if let Some(sharded_schema) = schema.schemas.get(table.schema()) { + return Ok(Some(sharded_schema.shard().into())); + } + } + + Ok(None) + } +} + +#[cfg(test)] +mod test { + use pg_query::{parse, NodeEnum}; + use pgdog_config::ShardedSchema; + + use crate::{ + backend::{replication::ShardedSchemas, ShardingSchema}, + frontend::router::{ + parser::{Shard, ShardsWithPriority}, + QueryParser, + }, + }; + + fn test_schema() -> ShardingSchema { + ShardingSchema { + shards: 2, + schemas: ShardedSchemas::new(vec![ + ShardedSchema { + name: Some("shard_0".into()), + shard: 0, + ..Default::default() + }, + ShardedSchema { + name: Some("shard_1".into()), + shard: 1, + ..Default::default() + }, + ]), + ..Default::default() + } + } + + fn parse_stmt(query: &str) -> Option { + let root = parse(query) + .unwrap() + .protobuf + .stmts + .first() + .unwrap() + .clone() + .stmt + .unwrap() + .node; + root + } + + #[test] + fn test_create_table_sharded_schema() { + let root = parse_stmt("CREATE TABLE shard_0.test (id BIGINT)"); + let mut calculator = ShardsWithPriority::default(); + let command = QueryParser::shard_ddl(&root, &test_schema(), &mut calculator).unwrap(); + assert_eq!(command.route().shard(), &Shard::Direct(0)); + } + + #[test] + fn test_create_table_unsharded_schema() { + let root = parse_stmt("CREATE TABLE unsharded.test (id BIGINT)"); + let mut calculator = ShardsWithPriority::default(); + let command = QueryParser::shard_ddl(&root, &test_schema(), &mut calculator).unwrap(); + assert_eq!(command.route().shard(), &Shard::All); + } + + #[test] + fn test_create_table_no_schema() { + let root = parse_stmt("CREATE TABLE test (id BIGINT)"); + let mut calculator = ShardsWithPriority::default(); + let command = QueryParser::shard_ddl(&root, &test_schema(), &mut calculator).unwrap(); + assert_eq!(command.route().shard(), &Shard::All); + } + + #[test] + fn test_create_sequence_sharded() { + let root = parse_stmt("CREATE SEQUENCE shard_1.test_seq"); + let mut calculator = ShardsWithPriority::default(); + let command = QueryParser::shard_ddl(&root, &test_schema(), &mut calculator).unwrap(); + assert_eq!(command.route().shard(), &Shard::Direct(1)); + } + + #[test] + fn test_create_sequence_unsharded() { + let root = parse_stmt("CREATE SEQUENCE public.test_seq"); + let mut calculator = ShardsWithPriority::default(); + let command = QueryParser::shard_ddl(&root, &test_schema(), &mut calculator).unwrap(); + assert_eq!(command.route().shard(), &Shard::All); + } + + #[test] + fn test_drop_table_sharded() { + let root = parse_stmt("DROP TABLE shard_0.test"); + let mut calculator = ShardsWithPriority::default(); + let command = QueryParser::shard_ddl(&root, &test_schema(), &mut calculator).unwrap(); + assert_eq!(command.route().shard(), &Shard::Direct(0)); + } + + #[test] + fn test_drop_table_unsharded() { + let root = parse_stmt("DROP TABLE public.test"); + let mut calculator = ShardsWithPriority::default(); + let command = QueryParser::shard_ddl(&root, &test_schema(), &mut calculator).unwrap(); + assert_eq!(command.route().shard(), &Shard::All); + } + + #[test] + fn test_drop_index_sharded() { + let root = parse_stmt("DROP INDEX shard_1.test_idx"); + let mut calculator = ShardsWithPriority::default(); + let command = QueryParser::shard_ddl(&root, &test_schema(), &mut calculator).unwrap(); + assert_eq!(command.route().shard(), &Shard::Direct(1)); + } + + #[test] + fn test_drop_view_sharded() { + let root = parse_stmt("DROP VIEW shard_0.test_view"); + let mut calculator = ShardsWithPriority::default(); + let command = QueryParser::shard_ddl(&root, &test_schema(), &mut calculator).unwrap(); + assert_eq!(command.route().shard(), &Shard::Direct(0)); + } + + #[test] + fn test_drop_sequence_sharded() { + let root = parse_stmt("DROP SEQUENCE shard_1.test_seq"); + let mut calculator = ShardsWithPriority::default(); + let command = QueryParser::shard_ddl(&root, &test_schema(), &mut calculator).unwrap(); + assert_eq!(command.route().shard(), &Shard::Direct(1)); + } + + #[test] + fn test_drop_schema_sharded() { + let root = parse_stmt("DROP SCHEMA shard_0"); + let mut calculator = ShardsWithPriority::default(); + let command = QueryParser::shard_ddl(&root, &test_schema(), &mut calculator).unwrap(); + assert_eq!(command.route().shard(), &Shard::Direct(0)); + } + + #[test] + fn test_drop_schema_unsharded() { + let root = parse_stmt("DROP SCHEMA public"); + let mut calculator = ShardsWithPriority::default(); + let command = QueryParser::shard_ddl(&root, &test_schema(), &mut calculator).unwrap(); + assert_eq!(command.route().shard(), &Shard::All); + } + + #[test] + fn test_create_schema_sharded() { + let root = parse_stmt("CREATE SCHEMA shard_0"); + let mut calculator = ShardsWithPriority::default(); + let command = QueryParser::shard_ddl(&root, &test_schema(), &mut calculator).unwrap(); + assert_eq!(command.route().shard(), &Shard::Direct(0)); + } + + #[test] + fn test_create_schema_unsharded() { + let root = parse_stmt("CREATE SCHEMA new_schema"); + let mut calculator = ShardsWithPriority::default(); + let command = QueryParser::shard_ddl(&root, &test_schema(), &mut calculator).unwrap(); + assert_eq!(command.route().shard(), &Shard::All); + } + + #[test] + fn test_create_index_sharded() { + let root = parse_stmt("CREATE INDEX test_idx ON shard_1.test (id)"); + let mut calculator = ShardsWithPriority::default(); + let command = QueryParser::shard_ddl(&root, &test_schema(), &mut calculator).unwrap(); + assert_eq!(command.route().shard(), &Shard::Direct(1)); + + let root = parse_stmt("CREATE UNIQUE INDEX test_idx ON shard_1.test (id)"); + let mut calculator = ShardsWithPriority::default(); + let command = QueryParser::shard_ddl(&root, &test_schema(), &mut calculator).unwrap(); + assert_eq!(command.route().shard(), &Shard::Direct(1)); + } + + #[test] + fn test_do_begin() { + let root = parse_stmt( + r#"DO $$ BEGIN + ALTER TABLE "shard_1"."foo" ADD CONSTRAINT "foo_id_foo2_id_fk" FOREIGN KEY ("id") REFERENCES "shard_1"."foo2"("id") ON DELETE cascade ON UPDATE cascade; + EXCEPTION + WHEN duplicate_object THEN null; + END $$;"#, + ); + let mut calculator = ShardsWithPriority::default(); + let command = QueryParser::shard_ddl(&root, &test_schema(), &mut calculator).unwrap(); + assert_eq!(command.route().shard(), &Shard::Direct(1)); + } + + #[test] + fn test_create_index_unsharded() { + let root = parse_stmt("CREATE INDEX test_idx ON public.test (id)"); + let mut calculator = ShardsWithPriority::default(); + let command = QueryParser::shard_ddl(&root, &test_schema(), &mut calculator).unwrap(); + assert_eq!(command.route().shard(), &Shard::All); + } + + #[test] + fn test_create_view_sharded() { + let root = parse_stmt("CREATE VIEW shard_0.test_view AS SELECT 1"); + let mut calculator = ShardsWithPriority::default(); + let command = QueryParser::shard_ddl(&root, &test_schema(), &mut calculator).unwrap(); + assert_eq!(command.route().shard(), &Shard::Direct(0)); + } + + #[test] + fn test_create_view_unsharded() { + let root = parse_stmt("CREATE VIEW public.test_view AS SELECT 1"); + let mut calculator = ShardsWithPriority::default(); + let command = QueryParser::shard_ddl(&root, &test_schema(), &mut calculator).unwrap(); + assert_eq!(command.route().shard(), &Shard::All); + } + + #[test] + fn test_create_table_as_sharded() { + let root = parse_stmt("CREATE TABLE shard_1.new_table AS SELECT * FROM other"); + let mut calculator = ShardsWithPriority::default(); + let command = QueryParser::shard_ddl(&root, &test_schema(), &mut calculator).unwrap(); + assert_eq!(command.route().shard(), &Shard::Direct(1)); + } + + #[test] + fn test_create_table_as_unsharded() { + let root = parse_stmt("CREATE TABLE public.new_table AS SELECT * FROM other"); + let mut calculator = ShardsWithPriority::default(); + let command = QueryParser::shard_ddl(&root, &test_schema(), &mut calculator).unwrap(); + assert_eq!(command.route().shard(), &Shard::All); + } + + #[test] + fn test_lock_table() { + let root = parse_stmt(r#"LOCK TABLE "shard_1"."__migrations_table""#); + let mut calculator = ShardsWithPriority::default(); + let command = QueryParser::shard_ddl(&root, &test_schema(), &mut calculator).unwrap(); + assert_eq!(command.route().shard(), &Shard::Direct(1)); + } + + #[test] + fn test_create_function_sharded() { + let root = parse_stmt( + "CREATE FUNCTION shard_0.test_func() RETURNS void AS $$ BEGIN END; $$ LANGUAGE plpgsql", + ); + let mut calculator = ShardsWithPriority::default(); + let command = QueryParser::shard_ddl(&root, &test_schema(), &mut calculator).unwrap(); + assert_eq!(command.route().shard(), &Shard::Direct(0)); + } + + #[test] + fn test_create_function_unsharded() { + let root = parse_stmt( + "CREATE FUNCTION public.test_func() RETURNS void AS $$ BEGIN END; $$ LANGUAGE plpgsql", + ); + let mut calculator = ShardsWithPriority::default(); + let command = QueryParser::shard_ddl(&root, &test_schema(), &mut calculator).unwrap(); + assert_eq!(command.route().shard(), &Shard::All); + } + + #[test] + fn test_create_enum_sharded() { + let root = parse_stmt("CREATE TYPE shard_1.mood AS ENUM ('sad', 'ok', 'happy')"); + let mut calculator = ShardsWithPriority::default(); + let command = QueryParser::shard_ddl(&root, &test_schema(), &mut calculator).unwrap(); + assert_eq!(command.route().shard(), &Shard::Direct(1)); + } + + #[test] + fn test_create_enum_unsharded() { + let root = parse_stmt("CREATE TYPE public.mood AS ENUM ('sad', 'ok', 'happy')"); + let mut calculator = ShardsWithPriority::default(); + let command = QueryParser::shard_ddl(&root, &test_schema(), &mut calculator).unwrap(); + assert_eq!(command.route().shard(), &Shard::All); + } + + #[test] + fn test_alter_owner_sharded() { + let root = parse_stmt("ALTER TABLE shard_0.test OWNER TO new_owner"); + let mut calculator = ShardsWithPriority::default(); + let command = QueryParser::shard_ddl(&root, &test_schema(), &mut calculator).unwrap(); + assert_eq!(command.route().shard(), &Shard::Direct(0)); + } + + #[test] + fn test_alter_owner_unsharded() { + let root = parse_stmt("ALTER TABLE public.test OWNER TO new_owner"); + let mut calculator = ShardsWithPriority::default(); + let command = QueryParser::shard_ddl(&root, &test_schema(), &mut calculator).unwrap(); + assert_eq!(command.route().shard(), &Shard::All); + } + + #[test] + fn test_rename_table_sharded() { + let root = parse_stmt("ALTER TABLE shard_1.test RENAME TO new_test"); + let mut calculator = ShardsWithPriority::default(); + let command = QueryParser::shard_ddl(&root, &test_schema(), &mut calculator).unwrap(); + assert_eq!(command.route().shard(), &Shard::Direct(1)); + } + + #[test] + fn test_rename_table_unsharded() { + let root = parse_stmt("ALTER TABLE public.test RENAME TO new_test"); + let mut calculator = ShardsWithPriority::default(); + let command = QueryParser::shard_ddl(&root, &test_schema(), &mut calculator).unwrap(); + assert_eq!(command.route().shard(), &Shard::All); + } + + #[test] + fn test_alter_table_sharded() { + let root = parse_stmt("ALTER TABLE shard_0.test ADD COLUMN new_col INT"); + let mut calculator = ShardsWithPriority::default(); + let command = QueryParser::shard_ddl(&root, &test_schema(), &mut calculator).unwrap(); + assert_eq!(command.route().shard(), &Shard::Direct(0)); + } + + #[test] + fn test_alter_table_unsharded() { + let root = parse_stmt("ALTER TABLE public.test ADD COLUMN new_col INT"); + let mut calculator = ShardsWithPriority::default(); + let command = QueryParser::shard_ddl(&root, &test_schema(), &mut calculator).unwrap(); + assert_eq!(command.route().shard(), &Shard::All); + } + + #[test] + fn test_alter_sequence_sharded() { + let root = parse_stmt("ALTER SEQUENCE shard_1.test_seq RESTART WITH 100"); + let mut calculator = ShardsWithPriority::default(); + let command = QueryParser::shard_ddl(&root, &test_schema(), &mut calculator).unwrap(); + assert_eq!(command.route().shard(), &Shard::Direct(1)); + } + + #[test] + fn test_alter_sequence_unsharded() { + let root = parse_stmt("ALTER SEQUENCE public.test_seq RESTART WITH 100"); + let mut calculator = ShardsWithPriority::default(); + let command = QueryParser::shard_ddl(&root, &test_schema(), &mut calculator).unwrap(); + assert_eq!(command.route().shard(), &Shard::All); + } + + #[test] + fn test_vacuum_sharded() { + let root = parse_stmt("VACUUM shard_0.test"); + let mut calculator = ShardsWithPriority::default(); + let command = QueryParser::shard_ddl(&root, &test_schema(), &mut calculator).unwrap(); + assert_eq!(command.route().shard(), &Shard::Direct(0)); + } + + #[test] + fn test_vacuum_unsharded() { + let root = parse_stmt("VACUUM public.test"); + let mut calculator = ShardsWithPriority::default(); + let command = QueryParser::shard_ddl(&root, &test_schema(), &mut calculator).unwrap(); + assert_eq!(command.route().shard(), &Shard::All); + } + + #[test] + fn test_vacuum_no_table() { + let root = parse_stmt("VACUUM"); + let mut calculator = ShardsWithPriority::default(); + let command = QueryParser::shard_ddl(&root, &test_schema(), &mut calculator).unwrap(); + assert_eq!(command.route().shard(), &Shard::All); + } + + #[test] + fn test_truncate_single_table_sharded() { + let root = parse_stmt("TRUNCATE shard_0.test"); + let mut calculator = ShardsWithPriority::default(); + let command = QueryParser::shard_ddl(&root, &test_schema(), &mut calculator).unwrap(); + assert_eq!(command.route().shard(), &Shard::Direct(0)); + } + + #[test] + fn test_truncate_single_table_unsharded() { + let root = parse_stmt("TRUNCATE public.test"); + let mut calculator = ShardsWithPriority::default(); + let command = QueryParser::shard_ddl(&root, &test_schema(), &mut calculator).unwrap(); + assert_eq!(command.route().shard(), &Shard::All); + } + + #[test] + fn test_truncate_multiple_tables_same_shard() { + let root = parse_stmt("TRUNCATE shard_0.test1, shard_0.test2"); + let mut calculator = ShardsWithPriority::default(); + let command = QueryParser::shard_ddl(&root, &test_schema(), &mut calculator).unwrap(); + assert_eq!(command.route().shard(), &Shard::Direct(0)); + } + + #[test] + fn test_truncate_cross_shard_error() { + let root = parse_stmt("TRUNCATE shard_0.test1, shard_1.test2"); + let mut calculator = ShardsWithPriority::default(); + let result = QueryParser::shard_ddl(&root, &test_schema(), &mut calculator); + assert!(result.is_err()); + } + + #[test] + fn test_unhandled_ddl_defaults_to_all() { + let root = parse_stmt("COMMENT ON TABLE public.test IS 'test comment'"); + let mut calculator = ShardsWithPriority::default(); + let command = QueryParser::shard_ddl(&root, &test_schema(), &mut calculator).unwrap(); + assert_eq!(command.route().shard(), &Shard::All); + } +} diff --git a/pgdog/src/frontend/router/parser/query/delete.rs b/pgdog/src/frontend/router/parser/query/delete.rs index a1856aa77..f6163130f 100644 --- a/pgdog/src/frontend/router/parser/query/delete.rs +++ b/pgdog/src/frontend/router/parser/query/delete.rs @@ -1,28 +1,44 @@ -use crate::frontend::router::parser::where_clause::TablesSource; - +use super::StatementParser; use super::*; impl QueryParser { pub(super) fn delete( + &mut self, stmt: &DeleteStmt, - context: &QueryParserContext, + context: &mut QueryParserContext, ) -> Result { - let table = stmt.relation.as_ref().map(Table::from); - - if let Some(table) = table { - let source = TablesSource::from(table); - let where_clause = WhereClause::new(&source, &stmt.where_clause); + let shard = StatementParser::from_delete( + stmt, + context.router_context.bind, + &context.sharding_schema, + self.recorder_mut(), + ) + .shard()?; - if let Some(where_clause) = where_clause { - let shards = Self::where_clause( - &context.sharding_schema, - &where_clause, - context.router_context.bind, - )?; - return Ok(Command::Query(Route::write(Self::converge(shards)))); + let shard = match shard { + Some(shard) => { + if let Some(recorder) = self.recorder_mut() { + recorder.record_entry( + Some(shard.clone()), + "DELETE matched WHERE clause for sharding key", + ); + } + shard } - } + None => { + if let Some(recorder) = self.recorder_mut() { + recorder.record_entry(None, "DELETE fell back to broadcast"); + } + Shard::default() + } + }; + + context + .shards_calculator + .push(ShardWithPriority::new_table(shard)); - Ok(Command::Query(Route::write(None))) + Ok(Command::Query(Route::write( + context.shards_calculator.shard(), + ))) } } diff --git a/pgdog/src/frontend/router/parser/query/explain.rs b/pgdog/src/frontend/router/parser/query/explain.rs index 36b0ca9ce..60ac14fba 100644 --- a/pgdog/src/frontend/router/parser/query/explain.rs +++ b/pgdog/src/frontend/router/parser/query/explain.rs @@ -3,21 +3,43 @@ use super::*; impl QueryParser { pub(super) fn explain( &mut self, + cached_ast: &Ast, stmt: &ExplainStmt, context: &mut QueryParserContext, ) -> Result { let query = stmt.query.as_ref().ok_or(Error::EmptyQuery)?; let node = query.node.as_ref().ok_or(Error::EmptyQuery)?; - match node { - NodeEnum::SelectStmt(ref stmt) => self.select(stmt, context), - NodeEnum::InsertStmt(ref stmt) => Self::insert(stmt, context), - NodeEnum::UpdateStmt(ref stmt) => Self::update(stmt, context), - NodeEnum::DeleteStmt(ref stmt) => Self::delete(stmt, context), + if context.expanded_explain() { + if self.explain_recorder.is_none() { + self.explain_recorder = Some(ExplainRecorder::new()); + } + } else { + self.explain_recorder = None; + } + + let result = match node { + NodeEnum::SelectStmt(ref stmt) => self.select(cached_ast, stmt, context), + NodeEnum::InsertStmt(ref stmt) => self.insert(stmt, context), + NodeEnum::UpdateStmt(ref stmt) => self.update(stmt, context), + NodeEnum::DeleteStmt(ref stmt) => self.delete(stmt, context), _ => { // For other statement types, route to all shards - Ok(Command::Query(Route::write(None))) + context + .shards_calculator + .push(ShardWithPriority::new_table(Shard::All)); + Ok(Command::Query(Route::write( + context.shards_calculator.shard(), + ))) + } + }; + + match result { + Ok(command) => Ok(command), + Err(err) => { + self.explain_recorder = None; + Err(err) } } } @@ -28,19 +50,43 @@ mod tests { use super::*; use crate::backend::Cluster; - use crate::frontend::{ClientRequest, PreparedStatements, RouterContext}; - use crate::net::messages::{Bind, Parameter, Parse, Query}; - use crate::net::Parameters; + use crate::config::{self, config}; + use crate::frontend::client::Sticky; + use crate::frontend::router::Ast; + use crate::frontend::{BufferedQuery, ClientRequest, PreparedStatements, RouterContext}; + use crate::net::{ + messages::{Bind, Parameter, Parse, Query}, + Parameters, + }; + use bytes::Bytes; + use std::sync::Once; + + fn enable_expanded_explain() { + static INIT: Once = Once::new(); + INIT.call_once(|| { + let mut cfg = config().as_ref().clone(); + cfg.config.general.expanded_explain = true; + config::set(cfg).unwrap(); + }); + } // Helper function to route a plain SQL statement and return its `Route`. fn route(sql: &str) -> Route { - let buffer = ClientRequest::from(vec![Query::new(sql).into()]); - - let cluster = Cluster::new_test(); + enable_expanded_explain(); + let cluster = Cluster::new_test(&config()); let mut stmts = PreparedStatements::default(); - let params = Parameters::default(); - let ctx = RouterContext::new(&buffer, &cluster, &mut stmts, ¶ms, None).unwrap(); + let ast = Ast::new( + &BufferedQuery::Query(Query::new(sql)), + &cluster.sharding_schema(), + &mut stmts, + ) + .unwrap(); + let mut buffer = ClientRequest::from(vec![Query::new(sql).into()]); + buffer.ast = Some(ast); + + let params = Parameters::default(); + let ctx = RouterContext::new(&buffer, &cluster, ¶ms, None, Sticky::new()).unwrap(); match QueryParser::default().parse(ctx).unwrap().clone() { Command::Query(route) => route, @@ -50,23 +96,32 @@ mod tests { // Helper function to route a parameterized SQL statement and return its `Route`. fn route_parameterized(sql: &str, values: &[&[u8]]) -> Route { + enable_expanded_explain(); let parse_msg = Parse::new_anonymous(sql); let parameters = values .iter() .map(|v| Parameter { len: v.len() as i32, - data: v.to_vec(), + data: Bytes::copy_from_slice(v), }) .collect::>(); let bind = Bind::new_params("", ¶meters); - let buffer: ClientRequest = vec![parse_msg.into(), bind.into()].into(); - let cluster = Cluster::new_test(); + let cluster = Cluster::new_test(&config()); let mut stmts = PreparedStatements::default(); - let params = Parameters::default(); - let ctx = RouterContext::new(&buffer, &cluster, &mut stmts, ¶ms, None).unwrap(); + let ast = Ast::new( + &BufferedQuery::Prepared(Parse::new_anonymous(sql)), + &cluster.sharding_schema(), + &mut stmts, + ) + .unwrap(); + let mut buffer: ClientRequest = vec![parse_msg.into(), bind.into()].into(); + buffer.ast = Some(ast); + + let params = Parameters::default(); + let ctx = RouterContext::new(&buffer, &cluster, ¶ms, None, Sticky::new()).unwrap(); match QueryParser::default().parse(ctx).unwrap().clone() { Command::Query(route) => route, @@ -94,10 +149,16 @@ mod tests { let r = route("EXPLAIN SELECT * FROM sharded WHERE id = 1"); assert!(matches!(r.shard(), Shard::Direct(_))); assert!(r.is_read()); + let lines = r.explain().unwrap().render_lines(); + assert!(lines + .iter() + .any(|line| line.contains("matched sharding key"))); let r = route_parameterized("EXPLAIN SELECT * FROM sharded WHERE id = $1", &[b"11"]); assert!(matches!(r.shard(), Shard::Direct(_))); assert!(r.is_read()); + let lines = r.explain().unwrap().render_lines(); + assert!(lines.iter().any(|line| line.contains("parameter"))); } #[test] @@ -105,6 +166,8 @@ mod tests { let r = route("EXPLAIN SELECT * FROM sharded"); assert_eq!(r.shard(), &Shard::All); assert!(r.is_read()); + let lines = r.explain().unwrap().render_lines(); + assert!(lines.iter().any(|line| line.contains("broadcast"))); } #[test] @@ -115,6 +178,10 @@ mod tests { ); assert!(matches!(r.shard(), Shard::Direct(_))); assert!(r.is_write()); + let lines = r.explain().unwrap().render_lines(); + assert!(lines + .iter() + .any(|line| line.contains("INSERT matched sharding key"))); } #[test] @@ -125,10 +192,18 @@ mod tests { ); assert!(matches!(r.shard(), Shard::Direct(_))); assert!(r.is_write()); + let lines = r.explain().unwrap().render_lines(); + assert!(lines + .iter() + .any(|line| line.contains("UPDATE matched WHERE clause"))); let r = route("EXPLAIN UPDATE sharded SET active = true"); assert_eq!(r.shard(), &Shard::All); assert!(r.is_write()); + let lines = r.explain().unwrap().render_lines(); + assert!(lines + .iter() + .any(|line| line.contains("UPDATE fell back to broadcast"))); } #[test] @@ -136,10 +211,18 @@ mod tests { let r = route_parameterized("EXPLAIN DELETE FROM sharded WHERE id = $1", &[b"11"]); assert!(matches!(r.shard(), Shard::Direct(_))); assert!(r.is_write()); + let lines = r.explain().unwrap().render_lines(); + assert!(lines + .iter() + .any(|line| line.contains("DELETE matched WHERE clause"))); let r = route("EXPLAIN DELETE FROM sharded"); assert_eq!(r.shard(), &Shard::All); assert!(r.is_write()); + let lines = r.explain().unwrap().render_lines(); + assert!(lines + .iter() + .any(|line| line.contains("DELETE fell back to broadcast"))); } #[test] @@ -157,6 +240,27 @@ mod tests { fn test_explain_with_comment_override() { let r = route("/* pgdog_shard: 5 */ EXPLAIN SELECT * FROM sharded"); assert_eq!(r.shard(), &Shard::Direct(5)); + let lines = r.explain().unwrap().render_lines(); + assert_eq!(lines[3], " Shard 5: manual override to shard=5"); + } + + #[test] + fn test_explain_select_broadcast_entry() { + let lines = route("EXPLAIN SELECT * FROM sharded") + .explain() + .unwrap() + .render_lines(); + assert_eq!(lines[3], " Note: no sharding key matched; broadcasting"); + } + + #[test] + fn test_explain_select_parameter_entry() { + let lines = route_parameterized("EXPLAIN SELECT * FROM sharded WHERE id = $1", &[b"22"]) + .explain() + .unwrap() + .render_lines(); + + assert!(lines.iter().any(|line| line.contains("parameter"))); } #[test] diff --git a/pgdog/src/frontend/router/parser/query/mod.rs b/pgdog/src/frontend/router/parser/query/mod.rs index 77870913c..e953b4aeb 100644 --- a/pgdog/src/frontend/router/parser/query/mod.rs +++ b/pgdog/src/frontend/router/parser/query/mod.rs @@ -1,17 +1,14 @@ //! Route queries to correct shards. -use std::collections::HashSet; +use std::{collections::HashSet, ops::Deref}; use crate::{ backend::{databases::databases, ShardingSchema}, config::Role, - frontend::{ - router::{ - context::RouterContext, - parser::{rewrite::Rewrite, OrderBy, Shard}, - round_robin, - sharding::{Centroids, ContextBuilder, Value as ShardingValue}, - }, - BufferedQuery, + frontend::router::{ + context::RouterContext, + parser::{OrderBy, Shard}, + round_robin, + sharding::{Centroids, ContextBuilder}, }, net::{ messages::{Bind, Vector}, @@ -20,10 +17,15 @@ use crate::{ plugin::plugins, }; -use super::*; +use super::{ + explain_trace::{ExplainRecorder, ExplainSummary}, + *, +}; +mod ddl; mod delete; mod explain; mod plugins; +mod schema_sharding; mod select; mod set; mod shared; @@ -33,7 +35,6 @@ mod update; use multi_tenant::MultiTenantCheck; use pgdog_plugin::pg_query::{ - fingerprint, protobuf::{a_const::Val, *}, NodeEnum, }; @@ -47,64 +48,131 @@ use tracing::{debug, trace}; /// /// 1. Which shard it should go to /// 2. Is it a read or a write -/// 3. Does it need to be re-rewritten to something else, e.g. prepared statement. /// /// It's re-created for each query we process. Struct variables are used /// to store intermediate state or to store external context for the duration /// of the parsing. /// -#[derive(Debug)] +#[derive(Debug, Default)] pub struct QueryParser { - // The statement is executed inside a transaction. - in_transaction: bool, // No matter what query is executed, we'll send it to the primary. write_override: bool, - // Currently calculated shard. - shard: Shard, // Plugin read override. plugin_output: PluginOutput, + // Record explain output. + explain_recorder: Option, } -impl Default for QueryParser { - fn default() -> Self { - Self { - in_transaction: false, - write_override: false, - shard: Shard::All, - plugin_output: PluginOutput::default(), +impl QueryParser { + fn recorder_mut(&mut self) -> Option<&mut ExplainRecorder> { + self.explain_recorder.as_mut() + } + + fn ensure_explain_recorder( + &mut self, + ast: &pg_query::ParseResult, + context: &QueryParserContext, + ) { + if self.explain_recorder.is_some() || !context.expanded_explain() { + return; + } + + if let Some(root) = ast.protobuf.stmts.first() { + if let Some(node) = root.stmt.as_ref().and_then(|stmt| stmt.node.as_ref()) { + if matches!(node, NodeEnum::ExplainStmt(_)) { + self.explain_recorder = Some(ExplainRecorder::new()); + } + } } } -} -impl QueryParser { - /// Indicates we are in a transaction. - pub fn in_transaction(&self) -> bool { - self.in_transaction + fn attach_explain(&mut self, command: &mut Command) { + if let (Some(recorder), Command::Query(route)) = (self.explain_recorder.take(), command) { + let summary = ExplainSummary { + shard: route.shard().clone(), + read: route.is_read(), + }; + route.set_explain(recorder.finalize(summary)); + } } /// Parse a query and return a command. pub fn parse(&mut self, context: RouterContext) -> Result { - let mut qp_context = QueryParserContext::new(context); + let mut context = QueryParserContext::new(context)?; - let mut command = if qp_context.query().is_ok() { - self.in_transaction = qp_context.router_context.in_transaction(); - self.write_override = qp_context.write_override(); + let mut command = if context.query().is_ok() { + self.write_override = context.write_override(); - self.query(&mut qp_context)? + self.query(&mut context)? } else { Command::default() }; - // If the cluster only has one shard, use direct-to-shard queries. - if let Command::Query(ref mut query) = command { - if !matches!(query.shard(), Shard::Direct(_)) && qp_context.shards == 1 { - query.set_shard_mut(0); + if let Command::Query(route) = &mut command { + if route.is_cross_shard() && context.shards == 1 { + context + .shards_calculator + .push(ShardWithPriority::new_override_only_one_shard( + Shard::Direct(0), + )); + route.set_shard_mut(context.shards_calculator.shard()); + } + + route.set_search_path_driven_mut(context.shards_calculator.is_search_path()); + + if let Some(role) = context.router_context.sticky.role { + match role { + Role::Primary => route.set_read(false), + _ => route.set_read(true), + } } } + debug!("query router decision: {:#?}", command); + + self.attach_explain(&mut command); + Ok(command) } + /// Bypass the query parser if we can. + fn query_parser_bypass(context: &mut QueryParserContext) -> Option { + let shard = context.shards_calculator.shard(); + + if !shard.is_direct() && context.shards > 1 { + return None; + } + + if !shard.is_direct() { + context + .shards_calculator + .push(ShardWithPriority::new_override_parser_disabled( + Shard::Direct(0), + )); + } + + let shard = context.shards_calculator.shard(); + + // Cluster is read-only and only has one shard. + if context.read_only { + Some(Route::read(shard)) + } + // Cluster doesn't have replicas and has only one shard. + else if context.write_only { + Some(Route::write(shard)) + + // The role is specified in the connection parameter (pgdog.role). + } else if let Some(role) = context.router_context.parameter_hints.compute_role() { + Some(match role { + Role::Replica => Route::read(shard), + Role::Primary | Role::Auto => Route::write(shard), + }) + // Default to primary. + } else { + Some(Route::write(shard)) + } + } + /// Parse a query and return a command that tells us what to do with it. /// /// # Arguments @@ -124,54 +192,48 @@ impl QueryParser { ); if !use_parser { - // Cluster is read-only and only has one shard. - if context.read_only { - return Ok(Command::Query(Route::read(Shard::Direct(0)))); - } - // Cluster doesn't have replicas and has only one shard. - if context.write_only { - return Ok(Command::Query(Route::write(Shard::Direct(0)))); + // Try to figure out where we can send the query without + // parsing SQL. + if let Some(route) = Self::query_parser_bypass(context) { + return Ok(Command::Query(route)); + } else { + return Err(Error::QueryParserRequired); } } - let cache = Cache::get(); + let statement = context + .router_context + .ast + .clone() + .ok_or(Error::EmptyQuery)?; - // Get the AST from cache or parse the statement live. - let statement = match context.query()? { - // Only prepared statements (or just extended) are cached. - BufferedQuery::Prepared(query) => { - cache.parse(query.query(), &context.sharding_schema)? - } - // Don't cache simple queries. - // - // They contain parameter values, which makes the cache - // too large to be practical. - // - // Make your clients use prepared statements - // or at least send statements with placeholders using the - // extended protocol. - BufferedQuery::Query(query) => { - cache.parse_uncached(query.query(), &context.sharding_schema)? - } - }; + self.ensure_explain_recorder(statement.parse_result(), context); // Parse hardcoded shard from a query comment. if context.router_needed || context.dry_run { - self.shard = statement.shard.clone(); - if let Some(role) = statement.role { + if let Some(comment_shard) = statement.comment_shard.clone() { + context + .shards_calculator + .push(ShardWithPriority::new_comment(comment_shard)); + } + + let role_override = statement.comment_role; + if let Some(role) = role_override { self.write_override = role == Role::Primary; } + + if statement.comment_shard.is_some() || role_override.is_some() { + let shard = context.shards_calculator.shard(); + + if let Some(recorder) = self.recorder_mut() { + recorder.record_comment_override(shard.deref().clone(), role_override); + } + } } debug!("{}", context.query()?.query()); trace!("{:#?}", statement); - let rewrite = Rewrite::new(statement.ast()); - if rewrite.needs_rewrite() { - debug!("rewrite needed"); - return rewrite.rewrite(context.prepared_statements()); - } - if let Some(multi_tenant) = context.multi_tenant() { debug!("running multi-tenant check"); @@ -179,8 +241,8 @@ impl QueryParser { context.router_context.cluster.user(), multi_tenant, context.router_context.cluster.schema(), - statement.ast(), - context.router_context.params, + statement.parse_result(), + context.router_context.parameter_hints.search_path, ) .run()?; } @@ -191,15 +253,20 @@ impl QueryParser { // We don't expect clients to send multiple queries. If they do // only the first one is used for routing. // - let root = statement.ast().protobuf.stmts.first(); + let root = statement.parse_result().protobuf.stmts.first(); let root = if let Some(root) = root { root.stmt.as_ref().ok_or(Error::EmptyQuery)? } else { + context + .shards_calculator + .push(ShardWithPriority::new_rr_empty_query(Shard::Direct( + round_robin::next() % context.shards, + ))); // Send empty query to any shard. - return Ok(Command::Query(Route::read(Shard::Direct( - round_robin::next() % context.shards, - )))); + return Ok(Command::Query(Route::read( + context.shards_calculator.shard(), + ))); }; let mut command = match root.node { @@ -212,15 +279,15 @@ impl QueryParser { return Ok(Command::Deallocate); } // SELECT statements. - Some(NodeEnum::SelectStmt(ref stmt)) => self.select(stmt, context), + Some(NodeEnum::SelectStmt(ref stmt)) => self.select(&statement, stmt, context), // COPY statements. Some(NodeEnum::CopyStmt(ref stmt)) => Self::copy(stmt, context), // INSERT statements. - Some(NodeEnum::InsertStmt(ref stmt)) => Self::insert(stmt, context), + Some(NodeEnum::InsertStmt(ref stmt)) => self.insert(stmt, context), // UPDATE statements. - Some(NodeEnum::UpdateStmt(ref stmt)) => Self::update(stmt, context), + Some(NodeEnum::UpdateStmt(ref stmt)) => self.update(stmt, context), // DELETE statements. - Some(NodeEnum::DeleteStmt(ref stmt)) => Self::delete(stmt, context), + Some(NodeEnum::DeleteStmt(ref stmt)) => self.delete(stmt, context), // Transaction control statements, // e.g. BEGIN, COMMIT, etc. Some(NodeEnum::TransactionStmt(ref stmt)) => match self.transaction(stmt, context)? { @@ -258,12 +325,7 @@ impl QueryParser { return Ok(Command::Unlisten(stmt.conditionname.clone())); } - Some(NodeEnum::ExplainStmt(ref stmt)) => self.explain(stmt, context), - - // VACUUM. - Some(NodeEnum::VacuumRelation(_)) | Some(NodeEnum::VacuumStmt(_)) => { - Ok(Command::Query(Route::write(None).set_maintenace())) - } + Some(NodeEnum::ExplainStmt(ref stmt)) => self.explain(&statement, stmt, context), Some(NodeEnum::DiscardStmt { .. }) => { return Ok(Command::Discard { @@ -271,17 +333,27 @@ impl QueryParser { }) } - // All others are not handled. - // They are sent to all shards concurrently. - _ => Ok(Command::Query(Route::write(None))), + _ => self.ddl(&root.node, context), }?; // e.g. Parse, Describe, Flush-style flow. if !context.router_context.executable { - if let Command::Query(query) = command { - return Ok(Command::Query( - query.set_shard(round_robin::next() % context.shards), - )); + if let Command::Query(ref query) = command { + if query.is_cross_shard() && statement.rewrite_plan.insert_split.is_empty() { + context + .shards_calculator + .push(ShardWithPriority::new_rr_not_executable(Shard::Direct( + round_robin::next() % context.shards, + ))); + + // Since this query isn't executable and we decided + // to route it to any shard, we can early return here. + return Ok(Command::Query( + query + .clone() + .with_shard(context.shards_calculator.shard().clone()), + )); + } } } @@ -295,9 +367,10 @@ impl QueryParser { }, )?; - // Overwrite shard using shard we got from a comment, if any. - if let Shard::Direct(shard) = self.shard { - if let Command::Query(ref mut route) = command { + // Set shard on route, if we're ready. + if let Command::Query(ref mut route) = command { + let shard = context.shards_calculator.shard(); + if shard.is_direct() { route.set_shard_mut(shard); } } @@ -306,11 +379,14 @@ impl QueryParser { // Plugins override what we calculated above. if let Command::Query(ref mut route) = command { if let Some(read) = self.plugin_output.read { - route.set_read_mut(read); + route.set_read(read); } if let Some(ref shard) = self.plugin_output.shard { - route.set_shard_raw_mut(shard); + context + .shards_calculator + .push(ShardWithPriority::new_plugin(shard.clone())); + route.set_shard_raw_mut(context.shards_calculator.shard()); } } @@ -322,7 +398,12 @@ impl QueryParser { // if context.shards == 1 && !context.dry_run { if let Command::Query(ref mut route) = command { - route.set_shard_mut(0); + context + .shards_calculator + .push(ShardWithPriority::new_override_only_one_shard( + Shard::Direct(0), + )); + route.set_shard_mut(context.shards_calculator.shard()); } } @@ -332,35 +413,36 @@ impl QueryParser { // Looking through manual queries to see if we have any // with the fingerprint. // - if route.shard().all() { + if route.shard().is_all() { let databases = databases(); // Only fingerprint the query if some manual queries are configured. // Otherwise, we're wasting time parsing SQL. if !databases.manual_queries().is_empty() { - let fingerprint = - fingerprint(context.query()?.query()).map_err(Error::PgQuery)?; - debug!("fingerprint: {}", fingerprint.hex); - let manual_route = databases.manual_query(&fingerprint.hex).cloned(); + let fingerprint = &statement.fingerprint.hex; + debug!("fingerprint: {}", fingerprint); + let manual_route = databases.manual_query(fingerprint).cloned(); // TODO: check routing logic required by config. if manual_route.is_some() { - route.set_shard_mut(round_robin::next() % context.shards); + context.shards_calculator.push(ShardWithPriority::new_table( + Shard::Direct(round_robin::next() % context.shards), + )); + route.set_shard_mut(context.shards_calculator.shard().clone()); } } } } - debug!("query router decision: {:#?}", command); - statement.update_stats(command.route()); if context.dry_run { // Record statement in cache with normalized parameters. if !statement.cached { - cache.record_normalized( - context.query()?.query(), + let query = context.query()?.query(); + Cache::get().record_normalized( + query, command.route(), - &context.sharding_schema, + context.sharding_schema.query_parser_engine, )?; } Ok(command.dry_run()) @@ -370,12 +452,46 @@ impl QueryParser { } /// Handle COPY command. - fn copy(stmt: &CopyStmt, context: &QueryParserContext) -> Result { + fn copy(stmt: &CopyStmt, context: &mut QueryParserContext) -> Result { + // Schema-based routing. + // + // We do this here as well because COPY
TO STDOUT + // doesn't use the CopyParser (doesn't need to, normally), + // so we need to handle this case here. + // + // The CopyParser itself has handling for schema-based sharding, + // but that's only used for logical replication during the first + // phase of data-sync. + // + let table = stmt.relation.as_ref().map(Table::from); + if let Some(table) = table { + if let Some(schema) = context.sharding_schema.schemas.get(table.schema()) { + let shard: Shard = schema.shard().into(); + context + .shards_calculator + .push(ShardWithPriority::new_table(shard)); + if !stmt.is_from { + return Ok(Command::Query(Route::read( + context.shards_calculator.shard(), + ))); + } else { + return Ok(Command::Query(Route::write( + context.shards_calculator.shard(), + ))); + } + } + } + let parser = CopyParser::new(stmt, context.router_context.cluster)?; - if let Some(parser) = parser { - Ok(Command::Copy(Box::new(parser))) + if !stmt.is_from { + context + .shards_calculator + .push(ShardWithPriority::new_table(Shard::All)); + Ok(Command::Query(Route::read( + context.shards_calculator.shard(), + ))) } else { - Ok(Command::Query(Route::write(None))) + Ok(Command::Copy(Box::new(parser))) } } @@ -386,9 +502,29 @@ impl QueryParser { /// * `stmt`: INSERT statement from pg_query. /// * `context`: Query parser context. /// - fn insert(stmt: &InsertStmt, context: &QueryParserContext) -> Result { + fn insert( + &mut self, + stmt: &InsertStmt, + context: &mut QueryParserContext, + ) -> Result { let insert = Insert::new(stmt); - let shard = insert.shard(&context.sharding_schema, context.router_context.bind)?; + context.shards_calculator.push(ShardWithPriority::new_table( + insert.shard(&context.sharding_schema, context.router_context.bind)?, + )); + let shard = context.shards_calculator.shard(); + + if let Some(recorder) = self.recorder_mut() { + match shard.deref() { + Shard::Direct(_) => recorder + .record_entry(Some(shard.deref().clone()), "INSERT matched sharding key"), + Shard::Multi(_) => recorder.record_entry( + Some(shard.deref().clone()), + "INSERT targeted multiple shards", + ), + Shard::All => recorder.record_entry(None, "INSERT broadcasted"), + }; + } + Ok(Command::Query(Route::write(shard))) } } diff --git a/pgdog/src/frontend/router/parser/query/plugins.rs b/pgdog/src/frontend/router/parser/query/plugins.rs index eecdd9bd4..14e992ba9 100644 --- a/pgdog/src/frontend/router/parser/query/plugins.rs +++ b/pgdog/src/frontend/router/parser/query/plugins.rs @@ -1,5 +1,6 @@ -use crate::frontend::router::parser::cache::CachedAst; +use crate::frontend::router::parser::cache::Ast; use pgdog_plugin::{ReadWrite, Shard as PdShard}; +use std::string::String as StdString; use super::*; @@ -8,6 +9,7 @@ use super::*; pub(super) struct PluginOutput { pub(super) shard: Option, pub(super) read: Option, + pub(super) plugin_name: Option, } impl PluginOutput { @@ -21,7 +23,7 @@ impl QueryParser { pub(super) fn plugins( &mut self, context: &QueryParserContext, - statement: &CachedAst, + statement: &Ast, read: bool, ) -> Result<(), Error> { // Don't run plugins on Parse only. @@ -43,12 +45,15 @@ impl QueryParser { // The first plugin to returns something, wins. debug!("executing {} router plugins", plugins.len()); - let mut context = - context.plugin_context(&statement.ast().protobuf, &context.router_context.bind); + let mut context = context.plugin_context( + &statement.parse_result().protobuf, + &context.router_context.bind, + ); context.write_override = if self.write_override || !read { 1 } else { 0 }; for plugin in plugins { if let Some(route) = plugin.route(context) { + let plugin_name = plugin.name().to_owned(); match route.shard.try_into() { Ok(shard) => match shard { PdShard::All => self.plugin_output.shard = Some(Shard::All), @@ -69,10 +74,21 @@ impl QueryParser { _ => self.plugin_output.read = None, } + self.plugin_output.plugin_name = Some(plugin_name.clone()); + if self.plugin_output.provided() { + let shard_override = self.plugin_output.shard.clone(); + let read_override = self.plugin_output.read; + if let Some(recorder) = self.recorder_mut() { + recorder.record_plugin_override( + plugin_name.clone(), + shard_override, + read_override, + ); + } debug!( "plugin \"{}\" returned route [{}, {}]", - plugin.name(), + plugin_name, match self.plugin_output.shard.as_ref() { Some(shard) => format!("shard={}", shard), None => "shard=unknown".to_string(), diff --git a/pgdog/src/frontend/router/parser/query/schema_sharding.rs b/pgdog/src/frontend/router/parser/query/schema_sharding.rs new file mode 100644 index 000000000..bdbc5417c --- /dev/null +++ b/pgdog/src/frontend/router/parser/query/schema_sharding.rs @@ -0,0 +1,56 @@ +use super::*; + +impl QueryParser { + // pub(super) fn check_search_path_for_shard( + // &mut self, + // context: &QueryParserContext<'_>, + // ) -> Result, Error> { + // // Shortcut. + // if context.sharding_schema.schemas.is_empty() { + // return Ok(None); + // } + + // // Check search_path for schema. + // let search_path = context.router_context.search_path; + // let mut schema_sharder = SchemaSharder::default(); + + // match search_path { + // Some(ParameterValue::String(search_path)) => { + // let schema = Schema::from(search_path.as_str()); + // schema_sharder.resolve(Some(schema), &context.sharding_schema.schemas); + // if let Some((shard, schema)) = schema_sharder.get() { + // if let Some(recorder) = self.recorder_mut() { + // recorder.clear(); + // recorder.record_entry( + // Some(shard.clone()), + // format!("matched schema {} in search_path", schema), + // ); + // } + // return Ok(Some(shard)); + // } + // } + + // Some(ParameterValue::Tuple(search_paths)) => { + // for schema in search_paths { + // let schema = Schema::from(schema.as_str()); + // schema_sharder.resolve(Some(schema), &context.sharding_schema.schemas); + // } + + // if let Some((shard, schema)) = schema_sharder.get() { + // if let Some(recorder) = self.recorder_mut() { + // recorder.clear(); + // recorder.record_entry( + // Some(shard.clone()), + // format!("matched schema {} in search_path", schema), + // ); + // } + // return Ok(Some(shard)); + // } + // } + + // _ => (), + // } + + // Ok(None) + // } +} diff --git a/pgdog/src/frontend/router/parser/query/select.rs b/pgdog/src/frontend/router/parser/query/select.rs index b1f11c92e..a69a40968 100644 --- a/pgdog/src/frontend/router/parser/query/select.rs +++ b/pgdog/src/frontend/router/parser/query/select.rs @@ -1,6 +1,9 @@ -use crate::frontend::router::parser::{from_clause::FromClause, where_clause::TablesSource}; +use crate::frontend::router::parser::{ + cache::Ast, from_clause::FromClause, where_clause::TablesSource, +}; use super::*; +use shared::ConvergeAlgorithm; impl QueryParser { /// Handle SELECT statement. @@ -12,6 +15,7 @@ impl QueryParser { /// pub(super) fn select( &mut self, + cached_ast: &Ast, stmt: &SelectStmt, context: &mut QueryParserContext, ) -> Result { @@ -27,31 +31,42 @@ impl QueryParser { writes.writes = true; } - if matches!(self.shard, Shard::Direct(_)) { + // Early return for any direct-to-shard queries. + if context.shards_calculator.shard().is_direct() { return Ok(Command::Query( - Route::read(self.shard.clone()).set_write(writes), + Route::read(context.shards_calculator.shard().clone()).with_write(writes), )); } + let mut shards = HashSet::new(); + + let shard = StatementParser::from_select( + stmt, + context.router_context.bind, + &context.sharding_schema, + self.recorder_mut(), + ) + .shard()?; + + if let Some(shard) = shard { + shards.insert(shard); + } + // `SELECT NOW()`, `SELECT 1`, etc. - if stmt.from_clause.is_empty() { + if shards.is_empty() && stmt.from_clause.is_empty() { + context + .shards_calculator + .push(ShardWithPriority::new_rr_no_table(Shard::Direct( + round_robin::next() % context.shards, + ))); + return Ok(Command::Query( - Route::read(Some(round_robin::next() % context.shards)).set_write(writes), + Route::read(context.shards_calculator.shard().clone()).with_write(writes), )); } let order_by = Self::select_sort(&stmt.sort_clause, context.router_context.bind); - let mut shards = HashSet::new(); let from_clause = TablesSource::from(FromClause::new(&stmt.from_clause)); - let where_clause = WhereClause::new(&from_clause, &stmt.where_clause); - - if let Some(ref where_clause) = where_clause { - shards = Self::where_clause( - &context.sharding_schema, - &where_clause, - context.router_context.bind, - )?; - } // Shard by vector in ORDER BY clause. for order in &order_by { @@ -62,66 +77,90 @@ impl QueryParser { || table.name.as_deref() == from_clause.table_name()) { let centroids = Centroids::from(&table.centroids); - shards.insert(centroids.shard( - vector, - context.shards, - table.centroid_probes, - )); + let shard: Shard = centroids + .shard(vector, context.shards, table.centroid_probes) + .into(); + if let Some(recorder) = self.recorder_mut() { + recorder.record_entry( + Some(shard.clone()), + format!("ORDER BY vector distance on {}", column_name), + ); + } + shards.insert(shard); } } } } - let shard = Self::converge(shards); - let aggregates = Aggregate::parse(stmt)?; + let shard = Self::converge(&shards, ConvergeAlgorithm::default()); + let aggregates = Aggregate::parse(stmt); let limit = LimitClause::new(stmt, context.router_context.bind).limit_offset()?; let distinct = Distinct::new(stmt).distinct()?; - let mut query = Route::select(shard, order_by, aggregates, limit, distinct); + context + .shards_calculator + .push(ShardWithPriority::new_table(shard)); - let mut omni = false; + let mut query = Route::select( + context.shards_calculator.shard().clone(), + order_by, + aggregates, + limit, + distinct, + ); + + // Omnisharded tables check. if query.is_all_shards() { - if let Some(name) = from_clause.table_name() { - omni = context.sharding_schema.tables.omnishards().contains(name); - } - } + let tables = from_clause.tables(); + let mut sticky = false; + let omni = tables.iter().all(|table| { + let is_sticky = context.sharding_schema.tables.omnishards().get(table.name); - if omni { - query.set_shard_mut(round_robin::next() % context.shards); + if let Some(is_sticky) = is_sticky { + if *is_sticky { + sticky = true; + } + true + } else { + false + } + }); + + if omni { + let shard = if sticky { + context.router_context.sticky.omni_index + } else { + round_robin::next() + } % context.shards; + + context + .shards_calculator + .push(ShardWithPriority::new_rr_omni(Shard::Direct(shard))); + + query.set_shard_mut(context.shards_calculator.shard().clone()); + + if let Some(recorder) = self.recorder_mut() { + recorder.record_entry( + Some(shard.into()), + format!( + "SELECT matched omnisharded tables: {}", + tables + .iter() + .map(|table| table.name) + .collect::>() + .join(", ") + ), + ); + } + } } // Only rewrite if query is cross-shard. if query.is_cross_shard() && context.shards > 1 { - if let Some(buffered_query) = context.router_context.query.as_ref() { - let rewrite = - RewriteEngine::new().rewrite_select(buffered_query.query(), query.aggregate()); - if !rewrite.plan.is_noop() { - if let BufferedQuery::Prepared(parse) = buffered_query { - let name = parse.name().to_owned(); - { - let prepared = context.prepared_statements(); - prepared.set_rewrite_plan(&name, rewrite.plan.clone()); - prepared.update_query(&name, &rewrite.sql); - } - } - query.set_rewrite(rewrite.plan, rewrite.sql); - } else if let BufferedQuery::Prepared(parse) = buffered_query { - let name = parse.name().to_owned(); - let stored_plan = { - let prepared = context.prepared_statements(); - prepared.rewrite_plan(&name) - }; - if let Some(plan) = stored_plan { - if !plan.is_noop() { - query.clear_rewrite(); - *query.rewrite_plan_mut() = plan; - } - } - } - } + query.with_aggregate_rewrite_plan_mut(cached_ast.rewrite_plan.aggregates.clone()); } - Ok(Command::Query(query.set_write(writes))) + Ok(Command::Query(query.with_write(writes))) } /// Handle the `ORDER BY` clause of a `SELECT` statement. @@ -195,7 +234,7 @@ impl QueryParser { Value::Vector(vec) => vector = Some(vec), _ => (), } - }; + } if let Ok(col) = Column::try_from(&e.node) { column = Some(col.name.to_owned()); diff --git a/pgdog/src/frontend/router/parser/query/set.rs b/pgdog/src/frontend/router/parser/query/set.rs index 0b469ef2a..353a136c6 100644 --- a/pgdog/src/frontend/router/parser/query/set.rs +++ b/pgdog/src/frontend/router/parser/query/set.rs @@ -13,122 +13,58 @@ impl QueryParser { stmt: &VariableSetStmt, context: &QueryParserContext, ) -> Result { - match stmt.name.as_str() { - "pgdog.shard" => { - if self.in_transaction { - let node = stmt - .args - .first() - .ok_or(Error::SetShard)? - .node - .as_ref() - .ok_or(Error::SetShard)?; - if let NodeEnum::AConst(AConst { - val: Some(a_const::Val::Ival(Integer { ival })), - .. - }) = node - { - return Ok(Command::SetRoute( - Route::write(Some(*ival as usize)).set_read(context.read_only), - )); - } - } else { - return Err(Error::RequiresTransaction); - } - } - - "pgdog.sharding_key" => { - if self.in_transaction { - let node = stmt - .args - .first() - .ok_or(Error::SetShard)? - .node - .as_ref() - .ok_or(Error::SetShard)?; + let transaction_state = stmt.name.starts_with("TRANSACTION"); + let value = Self::parse_set_value(stmt)?; - if let NodeEnum::AConst(AConst { - val: Some(Val::Sval(String { sval })), - .. - }) = node - { - let shard = if context.sharding_schema.shards > 1 { - let ctx = ContextBuilder::infer_from_from_and_config( - sval.as_str(), - &context.sharding_schema, - )? - .shards(context.shards) - .build()?; - ctx.apply()? - } else { - Shard::Direct(0) - }; - return Ok(Command::SetRoute( - Route::write(shard).set_read(context.read_only), - )); - } - } else { - return Err(Error::RequiresTransaction); - } + if let Some(value) = value { + if !transaction_state { + return Ok(Command::Set { + name: stmt.name.to_string(), + value, + local: stmt.is_local, + route: Route::write(context.shards_calculator.shard()), + }); } + } - // TODO: Handle SET commands for updating client - // params without touching the server. - name => { - if !self.in_transaction { - let mut value = vec![]; - - for node in &stmt.args { - if let Some(NodeEnum::AConst(AConst { val: Some(val), .. })) = &node.node { - match val { - Val::Sval(String { sval }) => { - value.push(sval.to_string()); - } + Ok(Command::Query( + Route::write(context.shards_calculator.shard().clone()).with_read(context.read_only), + )) + } - Val::Ival(Integer { ival }) => { - value.push(ival.to_string()); - } + pub fn parse_set_value(stmt: &VariableSetStmt) -> Result, Error> { + let mut value = vec![]; - Val::Fval(Float { fval }) => { - value.push(fval.to_string()); - } + for node in &stmt.args { + if let Some(NodeEnum::AConst(AConst { val: Some(val), .. })) = &node.node { + match val { + Val::Sval(String { sval }) => { + value.push(sval.to_string()); + } - Val::Boolval(Boolean { boolval }) => { - value.push(boolval.to_string()); - } + Val::Ival(Integer { ival }) => { + value.push(ival.to_string()); + } - _ => (), - } - } + Val::Fval(Float { fval }) => { + value.push(fval.to_string()); } - match value.len() { - 0 => (), - 1 => { - return Ok(Command::Set { - name: name.to_string(), - value: ParameterValue::String(value.pop().unwrap()), - }) - } - _ => { - return Ok(Command::Set { - name: name.to_string(), - value: ParameterValue::Tuple(value), - }) - } + Val::Boolval(Boolean { boolval }) => { + value.push(boolval.to_string()); } + + _ => (), } } } - let shard = if let Shard::Direct(_) = self.shard { - self.shard.clone() - } else { - Shard::All + let value = match value.len() { + 0 => None, + 1 => Some(ParameterValue::String(value.pop().unwrap())), + _ => Some(ParameterValue::Tuple(value)), }; - Ok(Command::Query( - Route::write(shard).set_read(context.read_only), - )) + Ok(value) } } diff --git a/pgdog/src/frontend/router/parser/query/shared.rs b/pgdog/src/frontend/router/parser/query/shared.rs index ebed8093a..433b72689 100644 --- a/pgdog/src/frontend/router/parser/query/shared.rs +++ b/pgdog/src/frontend/router/parser/query/shared.rs @@ -1,83 +1,173 @@ use super::*; +#[derive(Debug, Clone, Default, Copy, PartialEq)] +pub(super) enum ConvergeAlgorithm { + // Take the first direct shard we find + FirstDirect, + // If All is present, make it cross-shard. + // If multiple shards are present, make it multi. + // Else, make it direct. + #[default] + AllFirstElseMulti, +} + impl QueryParser { /// Converge to a single route given multiple shards. - pub(super) fn converge(shards: HashSet) -> Shard { + pub(super) fn converge(shards: &HashSet, algorithm: ConvergeAlgorithm) -> Shard { let shard = if shards.len() == 1 { shards.iter().next().cloned().unwrap() } else { - let mut multi = vec![]; + let mut multi = HashSet::new(); let mut all = false; - for shard in &shards { + for shard in shards { match shard { Shard::All => { all = true; break; } - Shard::Direct(v) => multi.push(*v), - Shard::Multi(m) => multi.extend(m), + Shard::Direct(v) => { + multi.insert(*v); + } + Shard::Multi(m) => multi.extend(m.iter()), }; } + + if algorithm == ConvergeAlgorithm::FirstDirect { + let direct = shards.iter().find(|shard| shard.is_direct()); + if let Some(direct) = direct { + return direct.clone(); + } + } + if all || shards.is_empty() { Shard::All } else { - Shard::Multi(multi) + Shard::Multi(multi.into_iter().collect()) } }; shard } +} - /// Handle WHERRE clause in SELECT, UPDATE an DELETE statements. - pub(super) fn where_clause( - sharding_schema: &ShardingSchema, - where_clause: &WhereClause, - params: Option<&Bind>, - ) -> Result, Error> { - let mut shards = HashSet::new(); - // Complexity: O(number of sharded tables * number of columns in the query) - for table in sharding_schema.tables().tables() { - let table_name = table.name.as_deref(); - let keys = where_clause.keys(table_name, &table.column); - for key in keys { - match key { - Key::Constant { value, array } => { - if array { - shards.insert(Shard::All); - break; - } - - let ctx = ContextBuilder::new(table) - .data(value.as_str()) - .shards(sharding_schema.shards) - .build()?; - shards.insert(ctx.apply()?); - } +#[cfg(test)] +mod tests { + use super::*; + use std::collections::HashSet; - Key::Parameter { pos, array } => { - // Don't hash individual values yet. - // The odds are high this will go to all shards anyway. - if array { - shards.insert(Shard::All); - break; - } else if let Some(params) = params { - if let Some(param) = params.parameter(pos)? { - let value = ShardingValue::from_param(¶m, table.data_type)?; - let ctx = ContextBuilder::new(table) - .value(value) - .shards(sharding_schema.shards) - .build()?; - shards.insert(ctx.apply()?); - } - } - } + #[test] + fn single_direct_returns_itself() { + let shards = HashSet::from([Shard::Direct(5)]); - // Null doesn't help. - Key::Null => (), - } + let result = QueryParser::converge(&shards, ConvergeAlgorithm::AllFirstElseMulti); + assert_eq!(result, Shard::Direct(5)); + + let result = QueryParser::converge(&shards, ConvergeAlgorithm::FirstDirect); + assert_eq!(result, Shard::Direct(5)); + } + + #[test] + fn single_all_returns_itself() { + let shards = HashSet::from([Shard::All]); + + let result = QueryParser::converge(&shards, ConvergeAlgorithm::AllFirstElseMulti); + assert_eq!(result, Shard::All); + + let result = QueryParser::converge(&shards, ConvergeAlgorithm::FirstDirect); + assert_eq!(result, Shard::All); + } + + #[test] + fn single_multi_returns_itself() { + let shards = HashSet::from([Shard::Multi(vec![1, 2, 3])]); + + let result = QueryParser::converge(&shards, ConvergeAlgorithm::AllFirstElseMulti); + assert_eq!(result, Shard::Multi(vec![1, 2, 3])); + + let result = QueryParser::converge(&shards, ConvergeAlgorithm::FirstDirect); + assert_eq!(result, Shard::Multi(vec![1, 2, 3])); + } + + #[test] + fn multiple_direct_all_first_else_multi_returns_multi() { + let shards = HashSet::from([Shard::Direct(1), Shard::Direct(2)]); + + let result = QueryParser::converge(&shards, ConvergeAlgorithm::AllFirstElseMulti); + match result { + Shard::Multi(mut v) => { + v.sort(); + assert_eq!(v, vec![1, 2]); } + other => panic!("expected Multi, got {:?}", other), } + } + + #[test] + fn multiple_direct_first_direct_returns_one_direct() { + let shards = HashSet::from([Shard::Direct(1), Shard::Direct(2)]); + + let result = QueryParser::converge(&shards, ConvergeAlgorithm::FirstDirect); + assert!( + matches!(result, Shard::Direct(1) | Shard::Direct(2)), + "expected Direct(1) or Direct(2), got {:?}", + result + ); + } + + #[test] + fn all_present_all_first_else_multi_returns_all() { + let shards = HashSet::from([Shard::All, Shard::Direct(1)]); + + let result = QueryParser::converge(&shards, ConvergeAlgorithm::AllFirstElseMulti); + assert_eq!(result, Shard::All); + } + + #[test] + fn all_present_first_direct_returns_direct() { + let shards = HashSet::from([Shard::All, Shard::Direct(1)]); + + let result = QueryParser::converge(&shards, ConvergeAlgorithm::FirstDirect); + assert_eq!(result, Shard::Direct(1)); + } + + #[test] + fn empty_set_returns_all() { + let shards = HashSet::new(); + + let result = QueryParser::converge(&shards, ConvergeAlgorithm::AllFirstElseMulti); + assert_eq!(result, Shard::All); + + let result = QueryParser::converge(&shards, ConvergeAlgorithm::FirstDirect); + assert_eq!(result, Shard::All); + } + + #[test] + fn multi_and_direct_merge_into_multi() { + let shards = HashSet::from([Shard::Multi(vec![1, 2]), Shard::Direct(3)]); + + let result = QueryParser::converge(&shards, ConvergeAlgorithm::AllFirstElseMulti); + match result { + Shard::Multi(mut v) => { + v.sort(); + assert_eq!(v, vec![1, 2, 3]); + } + other => panic!("expected Multi, got {:?}", other), + } + } + + #[test] + fn multi_and_direct_first_direct_returns_direct() { + let shards = HashSet::from([Shard::Multi(vec![1, 2]), Shard::Direct(3)]); + + let result = QueryParser::converge(&shards, ConvergeAlgorithm::FirstDirect); + assert_eq!(result, Shard::Direct(3)); + } + + #[test] + fn all_with_multi_first_direct_no_direct_returns_all() { + let shards = HashSet::from([Shard::All, Shard::Multi(vec![1, 2])]); - Ok(shards) + let result = QueryParser::converge(&shards, ConvergeAlgorithm::FirstDirect); + assert_eq!(result, Shard::All); } } diff --git a/pgdog/src/frontend/router/parser/query/show.rs b/pgdog/src/frontend/router/parser/query/show.rs index 1e34bc90a..a3338b02b 100644 --- a/pgdog/src/frontend/router/parser/query/show.rs +++ b/pgdog/src/frontend/router/parser/query/show.rs @@ -6,13 +6,22 @@ impl QueryParser { pub(super) fn show( &mut self, stmt: &VariableShowStmt, - context: &QueryParserContext, + context: &mut QueryParserContext, ) -> Result { match stmt.name.as_str() { - "pgdog.shards" => Ok(Command::Shards(context.shards)), + "pgdog.shards" => Ok(Command::InternalField { + name: "shards".into(), + value: context.shards.to_string(), + }), + "pgdog.unique_id" => Ok(Command::UniqueId), _ => { - let shard = Shard::Direct(round_robin::next() % context.shards); - let route = Route::write(shard).set_read(context.read_only); + context + .shards_calculator + .push(ShardWithPriority::new_rr_no_table(Shard::Direct( + round_robin::next() % context.shards, + ))); + let route = Route::write(context.shards_calculator.shard().clone()) + .with_read(context.read_only); Ok(Command::Query(route)) } } @@ -22,23 +31,32 @@ impl QueryParser { #[cfg(test)] mod test_show { use crate::backend::Cluster; + use crate::config::config; + use crate::frontend::client::Sticky; use crate::frontend::router::parser::Shard; - use crate::frontend::router::QueryParser; - use crate::frontend::{ClientRequest, PreparedStatements, RouterContext}; + use crate::frontend::router::{Ast, QueryParser}; + use crate::frontend::{BufferedQuery, ClientRequest, PreparedStatements, RouterContext}; use crate::net::messages::Query; use crate::net::Parameters; #[test] fn show_runs_on_a_direct_shard_round_robin() { - let mut ps = PreparedStatements::default(); - let c = Cluster::new_test(); - let p = Parameters::default(); + let c = Cluster::new_test(&config()); let mut parser = QueryParser::default(); // First call let query = "SHOW TRANSACTION ISOLATION LEVEL"; - let buffer = ClientRequest::from(vec![Query::new(query).into()]); - let context = RouterContext::new(&buffer, &c, &mut ps, &p, None).unwrap(); + let mut ast = Ast::new( + &BufferedQuery::Query(Query::new(query)), + &c.sharding_schema(), + &mut PreparedStatements::default(), + ) + .unwrap(); + ast.cached = false; + let mut buffer = ClientRequest::from(vec![Query::new(query).into()]); + buffer.ast = Some(ast); + let params = Parameters::default(); + let context = RouterContext::new(&buffer, &c, ¶ms, None, Sticky::new()).unwrap(); let first = parser.parse(context).unwrap().clone(); let first_shard = first.route().shard(); @@ -46,8 +64,17 @@ mod test_show { // Second call let query = "SHOW TRANSACTION ISOLATION LEVEL"; - let buffer = ClientRequest::from(vec![Query::new(query).into()]); - let context = RouterContext::new(&buffer, &c, &mut ps, &p, None).unwrap(); + let mut ast = Ast::new( + &BufferedQuery::Query(Query::new(query)), + &c.sharding_schema(), + &mut PreparedStatements::default(), + ) + .unwrap(); + ast.cached = false; + let mut buffer = ClientRequest::from(vec![Query::new(query).into()]); + buffer.ast = Some(ast); + let params = Parameters::default(); + let context = RouterContext::new(&buffer, &c, ¶ms, None, Sticky::new()).unwrap(); let second = parser.parse(context).unwrap().clone(); let second_shard = second.route().shard(); diff --git a/pgdog/src/frontend/router/parser/query/test.rs b/pgdog/src/frontend/router/parser/query/test/mod.rs similarity index 56% rename from pgdog/src/frontend/router/parser/query/test.rs rename to pgdog/src/frontend/router/parser/query/test/mod.rs index 9e06aad6e..b2fd43219 100644 --- a/pgdog/src/frontend/router/parser/query/test.rs +++ b/pgdog/src/frontend/router/parser/query/test/mod.rs @@ -2,19 +2,61 @@ use std::ops::Deref; use crate::{ config::{self, config}, - frontend::client::TransactionType, net::{ messages::{parse::Parse, Parameter}, - Close, Format, Sync, + Close, Format, Parameters, Sync, }, }; +use bytes::Bytes; use super::{super::Shard, *}; use crate::backend::Cluster; use crate::config::ReadWriteStrategy; -use crate::frontend::{ClientRequest, PreparedStatements, RouterContext}; +use crate::frontend::{ + client::{Sticky, TransactionType}, + router::Ast, + BufferedQuery, ClientRequest, PreparedStatements, RouterContext, +}; use crate::net::messages::Query; -use crate::net::Parameters; + +pub mod setup; + +pub mod test_bypass; +pub mod test_comments; +pub mod test_ddl; +pub mod test_delete; +pub mod test_dml; +pub mod test_explain; +pub mod test_functions; +pub mod test_insert; +pub mod test_rr; +pub mod test_schema_sharding; +pub mod test_search_path; +pub mod test_select; +pub mod test_set; +pub mod test_sharding; +pub mod test_special; +pub mod test_subqueries; +pub mod test_transaction; + +fn parse_query(query: &str) -> Command { + let mut query_parser = QueryParser::default(); + let cluster = Cluster::new_test(&config()); + let ast = Ast::new( + &BufferedQuery::Query(Query::new(query)), + &cluster.sharding_schema(), + &mut PreparedStatements::default(), + ) + .unwrap(); + let mut client_request = ClientRequest::from(vec![Query::new(query).into()]); + client_request.ast = Some(ast); + + let params = Parameters::default(); + let context = + RouterContext::new(&client_request, &cluster, ¶ms, None, Sticky::new_test()).unwrap(); + let command = query_parser.parse(context).unwrap().clone(); + command +} macro_rules! command { ($query:expr) => {{ @@ -24,17 +66,30 @@ macro_rules! command { ($query:expr, $in_transaction:expr) => {{ let query = $query; let mut query_parser = QueryParser::default(); - let client_request = ClientRequest::from(vec![Query::new(query).into()]); - let cluster = Cluster::new_test(); - let mut stmt = PreparedStatements::default(); - let params = Parameters::default(); + let cluster = Cluster::new_test(&crate::config::config()); + let mut ast = Ast::new( + &BufferedQuery::Query(Query::new($query)), + &cluster.sharding_schema(), + &mut PreparedStatements::default(), + ) + .unwrap(); + ast.cached = false; // Simple protocol queries are not cached + let mut client_request = ClientRequest::from(vec![Query::new(query).into()]); + client_request.ast = Some(ast); let transaction = if $in_transaction { Some(TransactionType::ReadWrite) } else { None }; - let context = - RouterContext::new(&client_request, &cluster, &mut stmt, ¶ms, transaction).unwrap(); + let params = Parameters::default(); + let context = RouterContext::new( + &client_request, + &cluster, + ¶ms, + transaction, + Sticky::new(), + ) + .unwrap(); let command = query_parser.parse(context).unwrap().clone(); (command, query_parser) @@ -61,9 +116,19 @@ macro_rules! query { macro_rules! query_parser { ($qp:expr, $query:expr, $in_transaction:expr, $cluster:expr) => {{ let cluster = $cluster; + let mut client_request: ClientRequest = vec![$query.into()].into(); + + let buffered_query = client_request + .query() + .expect("parsing query") + .expect("no query in client request"); + let mut prep_stmts = PreparedStatements::default(); - let params = Parameters::default(); - let client_request: ClientRequest = vec![$query.into()].into(); + + let mut ast = + Ast::new(&buffered_query, &cluster.sharding_schema(), &mut prep_stmts).unwrap(); + ast.cached = false; // Dry run test needs this. + client_request.ast = Some(ast); let maybe_transaction = if $in_transaction { Some(TransactionType::ReadWrite) @@ -71,12 +136,13 @@ macro_rules! query_parser { None }; + let params = Parameters::default(); let router_context = RouterContext::new( &client_request, &cluster, - &mut prep_stmts, ¶ms, maybe_transaction, + Sticky::new(), ) .unwrap(); @@ -84,7 +150,12 @@ macro_rules! query_parser { }}; ($qp:expr, $query:expr, $in_transaction:expr) => { - query_parser!($qp, $query, $in_transaction, Cluster::new_test()) + query_parser!( + $qp, + $query, + $in_transaction, + Cluster::new_test(&crate::config::config()) + ) }; } @@ -99,18 +170,28 @@ macro_rules! parse { .into_iter() .map(|p| Parameter { len: p.len() as i32, - data: p.to_vec(), + data: Bytes::copy_from_slice(&p), }) .collect::>(); let bind = Bind::new_params_codes($name, ¶ms, $codes); + let cluster = Cluster::new_test(&crate::config::config()); + let ast = Ast::new( + &BufferedQuery::Prepared(Parse::new_anonymous($query)), + &cluster.sharding_schema(), + &mut PreparedStatements::default(), + ) + .unwrap(); + let mut client_request = ClientRequest::from(vec![parse.into(), bind.into()]); + client_request.ast = Some(ast); + let client_params = Parameters::default(); let route = QueryParser::default() .parse( RouterContext::new( - &ClientRequest::from(vec![parse.into(), bind.into()]), - &Cluster::new_test(), - &mut PreparedStatements::default(), - &Parameters::default(), + &client_request, + &cluster, + &client_params, None, + Sticky::new(), ) .unwrap(), ) @@ -135,7 +216,7 @@ fn test_insert() { "INSERT INTO sharded (id, email) VALUES ($1, $2)", ["11".as_bytes(), "test@test.com".as_bytes()] ); - assert_eq!(route.shard(), &Shard::direct(1)); + assert_eq!(route.shard(), &Shard::new_direct(1)); } #[test] @@ -194,98 +275,131 @@ fn test_select_for_update() { } #[test] -fn test_prepared_avg_rewrite_plan() { - let route = parse!( - "avg_test", - "SELECT AVG(price) FROM menu", - Vec::>::new() - ); +fn test_omni() { + let mut omni_round_robin = HashSet::new(); + let q = "SELECT sharded_omni.* FROM sharded_omni WHERE sharded_omni.id = 1"; - assert!(!route.rewrite_plan().is_noop()); - assert_eq!(route.rewrite_plan().drop_columns(), &[1]); - let rewritten = route - .rewritten_sql() - .expect("rewrite should produce SQL for prepared average"); - assert!( - rewritten.to_lowercase().contains("count"), - "helper COUNT should be injected" - ); -} + for _ in 0..10 { + let command = parse_query(q); + match command { + Command::Query(query) => { + assert!(matches!(query.shard(), Shard::Direct(_))); + omni_round_robin.insert(query.shard().clone()); + } -#[test] -fn test_omni() { - let q = "SELECT sharded_omni.* FROM sharded_omni WHERE sharded_omni.id = $1"; - let route = query!(q); - assert!(matches!(route.shard(), Shard::Direct(_))); - let (_, qp) = command!(q); - assert!(!qp.in_transaction); -} + _ => {} + } + } -#[test] -fn test_set() { - let (command, _) = command!(r#"SET "pgdog.shard" TO 1"#, true); - match command { - Command::SetRoute(route) => assert_eq!(route.shard(), &Shard::Direct(1)), - _ => panic!("not a set route"), + assert_eq!(omni_round_robin.len(), 2); + + // Test sticky routing + let mut omni_sticky = HashSet::new(); + let q = + "SELECT sharded_omni_sticky.* FROM sharded_omni_sticky WHERE sharded_omni_sticky.id = $1"; + + for _ in 0..10 { + let command = parse_query(q); + match command { + Command::Query(query) => { + assert!(matches!(query.shard(), Shard::Direct(_))); + omni_sticky.insert(query.shard().clone()); + } + + _ => {} + } } - let (_, qp) = command!(r#"SET "pgdog.shard" TO 1"#, true); - assert!(qp.in_transaction); - let (command, _) = command!(r#"SET "pgdog.sharding_key" TO '11'"#, true); - match command { - Command::SetRoute(route) => assert_eq!(route.shard(), &Shard::Direct(1)), - _ => panic!("not a set route"), + assert_eq!(omni_sticky.len(), 1); + + // Test that sharded tables take priority. + let q = " + SELECT + sharded_omni.*, + sharded.* + FROM + sharded_omni + INNER JOIN + sharded + ON sharded_omni.id = sharded.i + WHERE sharded.id = 5"; + + let route = query!(q); + let shard = route.shard().clone(); + + for _ in 0..5 { + let route = query!(q); + // Test that shard doesn't change (i.e. not round robin) + assert_eq!(&shard, route.shard()); + assert!(matches!(shard, Shard::Direct(_))); } - let (_, qp) = command!(r#"SET "pgdog.sharding_key" TO '11'"#, true); - assert!(qp.in_transaction); - for (command, qp) in [ + // Test that all tables have to be omnisharded. + let q = "SELECT * FROM sharded_omni INNER JOIN not_sharded ON sharded_omni.id = not_sharded.id WHERE sharded_omni = $1"; + let route = query!(q); + assert!(matches!(route.shard(), Shard::All)); +} + +#[test] +fn test_set() { + // let (command, _) = command!(r#"SET "pgdog.shard" TO 1"#, true); + // match command { + // Command::SetRoute(route) => assert_eq!(route.shard(), &Shard::Direct(1)), + // _ => panic!("not a set route"), + // } + // let (_, _) = command!(r#"SET "pgdog.shard" TO 1"#, true); + + // let (command, _) = command!(r#"SET "pgdog.sharding_key" TO '11'"#, true); + // match command { + // Command::SetRoute(route) => assert_eq!(route.shard(), &Shard::Direct(1)), + // _ => panic!("not a set route"), + // } + let (_, _) = command!(r#"SET "pgdog.sharding_key" TO '11'"#, true); + + for (command, _) in [ command!("SET TimeZone TO 'UTC'"), command!("SET TIME ZONE 'UTC'"), ] { match command { - Command::Set { name, value } => { + Command::Set { name, value, .. } => { assert_eq!(name, "timezone"); assert_eq!(value, ParameterValue::from("UTC")); } _ => panic!("not a set"), }; - assert!(!qp.in_transaction); } - let (command, qp) = command!("SET statement_timeout TO 3000"); + let (command, _) = command!("SET statement_timeout TO 3000"); match command { - Command::Set { name, value } => { + Command::Set { name, value, .. } => { assert_eq!(name, "statement_timeout"); assert_eq!(value, ParameterValue::from("3000")); } _ => panic!("not a set"), }; - assert!(!qp.in_transaction); // TODO: user shouldn't be able to set these. // The server will report an error on synchronization. - let (command, qp) = command!("SET is_superuser TO true"); + let (command, _) = command!("SET is_superuser TO true"); match command { - Command::Set { name, value } => { + Command::Set { name, value, .. } => { assert_eq!(name, "is_superuser"); assert_eq!(value, ParameterValue::from("true")); } _ => panic!("not a set"), }; - assert!(!qp.in_transaction); let (_, mut qp) = command!("BEGIN"); assert!(qp.write_override); let command = query_parser!(qp, Query::new(r#"SET statement_timeout TO 3000"#), true); - match command { - Command::Query(q) => assert!(q.is_write()), - _ => panic!("set should trigger binding"), - } + assert!( + matches!(command, Command::Set { .. }), + "set must be intercepted inside transactions" + ); let (command, _) = command!("SET search_path TO \"$user\", public, \"APPLES\""); match command { - Command::Set { name, value } => { + Command::Set { name, value, .. } => { assert_eq!(name, "search_path"); assert_eq!( value, @@ -295,29 +409,37 @@ fn test_set() { _ => panic!("search path"), } - let buffer: ClientRequest = vec![Query::new(r#"SET statement_timeout TO 1"#).into()].into(); - let cluster = Cluster::new_test(); + let query_str = r#"SET statement_timeout TO 1"#; + let cluster = Cluster::new_test(&config()); let mut prep_stmts = PreparedStatements::default(); - let params = Parameters::default(); + let buffered_query = BufferedQuery::Query(Query::new(query_str)); + let mut ast = Ast::new(&buffered_query, &cluster.sharding_schema(), &mut prep_stmts).unwrap(); + ast.cached = false; + let mut buffer: ClientRequest = vec![Query::new(query_str).into()].into(); + buffer.ast = Some(ast); let transaction = Some(TransactionType::ReadWrite); + let params = Parameters::default(); let router_context = - RouterContext::new(&buffer, &cluster, &mut prep_stmts, ¶ms, transaction).unwrap(); - let mut context = QueryParserContext::new(router_context); + RouterContext::new(&buffer, &cluster, ¶ms, transaction, Sticky::new()).unwrap(); + let mut context = QueryParserContext::new(router_context).unwrap(); for read_only in [true, false] { context.read_only = read_only; // Overriding context above. let mut qp = QueryParser::default(); - qp.in_transaction = true; let route = qp.query(&mut context).unwrap(); match route { - Command::Query(route) => { - assert_eq!(route.is_read(), read_only); - } - cmd => panic!("not a query: {:?}", cmd), + Command::Set { .. } => {} + _ => panic!("set must be intercepted"), } } + + let command = command!("SET LOCAL work_mem TO 1"); + match command.0 { + Command::Set { local, .. } => assert!(local), + _ => panic!("not a set"), + } } #[test] @@ -328,7 +450,6 @@ fn test_transaction() { _ => panic!("not a query"), }; - assert!(qp.in_transaction); assert!(qp.write_override); let route = query_parser!(qp, Parse::named("test", "SELECT $1"), true); @@ -337,7 +458,7 @@ fn test_transaction() { _ => panic!("not a select"), } - let mut cluster = Cluster::new_test(); + let mut cluster = Cluster::new_test(&config()); cluster.set_read_write_strategy(ReadWriteStrategy::Aggressive); let command = query_parser!( QueryParser::default(), @@ -346,9 +467,7 @@ fn test_transaction() { cluster.clone() ); assert!(matches!(command, Command::StartTransaction { .. })); - assert!(qp.in_transaction); - qp.in_transaction = true; let route = query_parser!( qp, Query::new("SET application_name TO 'test'"), @@ -356,9 +475,13 @@ fn test_transaction() { cluster.clone() ); match route { - Command::Query(q) => { - assert!(q.is_write()); + Command::Set { + name, value, local, .. + } => { + assert_eq!(name, "application_name"); + assert_eq!(value.as_str().unwrap(), "test"); assert!(!cluster.read_only()); + assert!(!local); } _ => panic!("not a query"), @@ -382,23 +505,22 @@ fn test_begin_extended() { #[test] fn test_show_shards() { - let (cmd, qp) = command!("SHOW pgdog.shards"); - assert!(matches!(cmd, Command::Shards(2))); - assert!(!qp.in_transaction); + let (cmd, _) = command!("SHOW pgdog.shards"); + assert!(matches!(cmd, Command::InternalField { .. })); } #[test] fn test_write_functions() { let route = query!("SELECT pg_advisory_lock($1)"); assert!(route.is_write()); - assert!(route.lock_session()); + assert!(route.is_lock_session()); } #[test] fn test_write_nolock() { let route = query!("SELECT nextval('234')"); assert!(route.is_write()); - assert!(!route.lock_session()); + assert!(!route.is_lock_session()); } #[test] @@ -414,12 +536,9 @@ fn test_cte() { fn test_function_begin() { let (cmd, mut qp) = command!("BEGIN"); assert!(matches!(cmd, Command::StartTransaction { .. })); - assert!(qp.in_transaction); - let cluster = Cluster::new_test(); + let cluster = Cluster::new_test(&config()); let mut prep_stmts = PreparedStatements::default(); - let params = Parameters::default(); - let buffer: ClientRequest = vec![Query::new( - "SELECT + let query_str = "SELECT ROW(t1.*) AS tt1, ROW(t2.*) AS tt2 FROM t1 @@ -432,20 +551,22 @@ WHERE t2.account = ( WHERE t2.id = $1 ) - ", - ) - .into()] - .into(); + "; + let buffered_query = BufferedQuery::Query(Query::new(query_str)); + let mut ast = Ast::new(&buffered_query, &cluster.sharding_schema(), &mut prep_stmts).unwrap(); + ast.cached = false; + let mut buffer: ClientRequest = vec![Query::new(query_str).into()].into(); + buffer.ast = Some(ast); let transaction = Some(TransactionType::ReadWrite); + let params = Parameters::default(); let router_context = - RouterContext::new(&buffer, &cluster, &mut prep_stmts, ¶ms, transaction).unwrap(); - let mut context = QueryParserContext::new(router_context); + RouterContext::new(&buffer, &cluster, ¶ms, transaction, Sticky::new()).unwrap(); + let mut context = QueryParserContext::new(router_context).unwrap(); let route = qp.query(&mut context).unwrap(); match route { Command::Query(query) => assert!(query.is_write()), _ => panic!("not a select"), } - assert!(qp.in_transaction); } #[test] @@ -469,7 +590,7 @@ fn test_comment() { ); match command { - Command::Query(query) => assert_eq!(query.shard(), &Shard::Direct(1)), // Round-robin because it's only a parse + Command::Query(query) => assert_eq!(query.shard(), &Shard::Direct(1234)), _ => panic!("not a query"), } } @@ -491,15 +612,14 @@ fn test_limit_offset() { #[test] fn test_close_direct_one_shard() { - let cluster = Cluster::new_test_single_shard(); + let cluster = Cluster::new_test_single_shard(&config()); let mut qp = QueryParser::default(); let buf: ClientRequest = vec![Close::named("test").into(), Sync.into()].into(); - let mut pp = PreparedStatements::default(); - let params = Parameters::default(); let transaction = None; + let params = Parameters::default(); - let context = RouterContext::new(&buf, &cluster, &mut pp, ¶ms, transaction).unwrap(); + let context = RouterContext::new(&buf, &cluster, ¶ms, transaction, Sticky::new()).unwrap(); let cmd = qp.parse(context).unwrap(); @@ -549,9 +669,9 @@ fn test_commit_prepared() { fn test_dry_run_simple() { let mut config = config().deref().clone(); config.config.general.dry_run = true; - config::set(config).unwrap(); + config::set(config.clone()).unwrap(); - let cluster = Cluster::new_test_single_shard(); + let cluster = Cluster::new_test_single_shard(&config); let command = query_parser!( QueryParser::default(), Query::new("/* pgdog_sharding_key: 1234 */ SELECT * FROM sharded"), @@ -567,17 +687,52 @@ fn test_dry_run_simple() { #[test] fn test_set_comments() { - let command = query_parser!( - QueryParser::default(), - Query::new("/* pgdog_sharding_key: 1234 */ SET statement_timeout TO 1"), - true - ); - assert_eq!(command.route().shard(), &Shard::Direct(0)); + // let command = query_parser!( + // QueryParser::default(), + // Query::new("/* pgdog_sharding_key: 1234 */ SET statement_timeout TO 1"), + // true + // ); + // assert_eq!(command.route().shard(), &Shard::Direct(0)); let command = query_parser!( QueryParser::default(), Query::new("SET statement_timeout TO 1"), true ); - assert_eq!(command.route().shard(), &Shard::All); + assert!(matches!(command, Command::Set { .. })); +} + +#[test] +fn test_subqueries() { + println!( + "{:#?}", + pg_query::parse(r#" + SELECT + count(*) AS "count" + FROM + ( + SELECT "companies".* FROM + ( + SELECT "companies".* + FROM "companies" + INNER JOIN "organizations_relevant_companies" ON + ("organizations_relevant_companies"."company_id" = "companies"."id") + WHERE + ( + ("organizations_relevant_companies"."org_id" = 1) + AND NOT + ( + EXISTS + ( + SELECT * FROM "hidden_globals" + WHERE + ( + ("hidden_globals"."org_id" = 1) + AND ("hidden_globals"."global_company_id" = "organizations_relevant_companies"."company_id") + ) + ) + ) + ) + ) AS "companies" OFFSET 0) AS "t1" LIMIT 1"#).unwrap() + ); } diff --git a/pgdog/src/frontend/router/parser/query/test/setup.rs b/pgdog/src/frontend/router/parser/query/test/setup.rs new file mode 100644 index 000000000..1df2493a4 --- /dev/null +++ b/pgdog/src/frontend/router/parser/query/test/setup.rs @@ -0,0 +1,163 @@ +use std::ops::Deref; + +use pgdog_config::ConfigAndUsers; + +use crate::{ + backend::Cluster, + config::{self, config, ReadWriteStrategy}, + frontend::{ + client::{Sticky, TransactionType}, + router::{ + parser::{Cache, Error}, + QueryParser, + }, + ClientRequest, Command, PreparedStatements, RouterContext, + }, + net::{parameter::ParameterValue, Parameters, ProtocolMessage}, +}; + +pub(super) use crate::net::*; + +/// Test harness for QueryParser that uses builder pattern for configuration. +pub(crate) struct QueryParserTest { + cluster: Cluster, + params: Parameters, + transaction: Option, + sticky: Sticky, + prepared: PreparedStatements, + pub(crate) parser: QueryParser, + last_parse: Option, +} + +impl QueryParserTest { + /// Create a new test with default settings (no transaction, default cluster). + pub(crate) fn new() -> Self { + Self::new_with_config(&config()) + } + + /// Create a test with a single-shard cluster. + pub(crate) fn new_single_shard(config: &ConfigAndUsers) -> Self { + let cluster = Cluster::new_test_single_shard(config); + + Self { + cluster, + params: Parameters::default(), + transaction: None, + sticky: Sticky::default(), + parser: QueryParser::default(), + prepared: PreparedStatements::new(), + last_parse: None, + } + } + + /// Create new test with specific general settings. + pub(crate) fn new_with_config(config: &ConfigAndUsers) -> Self { + let cluster = Cluster::new_test(config); + + Self { + cluster, + params: Parameters::default(), + transaction: None, + sticky: Sticky::default(), + parser: QueryParser::default(), + prepared: PreparedStatements::new(), + last_parse: None, + } + } + + /// Set whether we're in a transaction. + pub(crate) fn in_transaction(mut self, in_tx: bool) -> Self { + self.transaction = if in_tx { + Some(TransactionType::ReadWrite) + } else { + None + }; + self + } + + /// Set the read/write strategy on the cluster. + pub(crate) fn with_read_write_strategy(mut self, strategy: ReadWriteStrategy) -> Self { + self.cluster.set_read_write_strategy(strategy); + self + } + + /// Enable dry run mode for this test. + pub(crate) fn with_dry_run(mut self) -> Self { + let mut updated = config().deref().clone(); + updated.config.general.dry_run = true; + config::set(updated).unwrap(); + // Recreate cluster with the new config + self.cluster = Cluster::new_test(&config()); + self + } + + /// Set a parameter value. + pub(crate) fn with_param( + mut self, + name: impl ToString, + value: impl Into, + ) -> Self { + self.params.insert(name, value); + self + } + + /// Startup parameters. + + /// Execute a request and return the command (panics on error). + pub(crate) fn execute(&mut self, request: Vec) -> Command { + self.try_execute(request).expect("execute failed") + } + + /// Execute a request and return Result (for testing error conditions). + pub(crate) fn try_execute(&mut self, request: Vec) -> Result { + let mut request: ClientRequest = request.into(); + + for message in request.iter_mut() { + if let ProtocolMessage::Parse(parse) = message { + let (_, name) = PreparedStatements::global().write().insert(parse); + self.last_parse = Some(name); + } + + if let ProtocolMessage::Bind(bind) = message { + if let Some(ref name) = self.last_parse { + bind.rename(name); + } + } + + if let ProtocolMessage::Describe(desc) = message { + if let Some(ref name) = self.last_parse { + desc.rename(name); + } + } + } + + // Some requests (like Close) don't have a query + if let Ok(Some(buffered_query)) = request.query() { + let ast = Cache::get() + .query( + &buffered_query, + &self.cluster.sharding_schema(), + &mut self.prepared, + ) + .unwrap(); + request.ast = Some(ast); + } + + let router_ctx = RouterContext::new( + &request, + &self.cluster, + &self.params, + self.transaction, + self.sticky, + ) + .unwrap(); + + let command = self.parser.parse(router_ctx)?; + Ok(command.clone()) + } + + /// Get access to the cluster (for assertions). + pub(crate) fn cluster(&self) -> &Cluster { + &self.cluster + } +} diff --git a/pgdog/src/frontend/router/parser/query/test/test_bypass.rs b/pgdog/src/frontend/router/parser/query/test/test_bypass.rs new file mode 100644 index 000000000..8c7909ab9 --- /dev/null +++ b/pgdog/src/frontend/router/parser/query/test/test_bypass.rs @@ -0,0 +1,102 @@ +//! Tests that test what the query parser is disabled +//! and we have only one shard (but we have replicas). +//! +//! QueryParser::query_parser_bypass. +//! +use pgdog_config::QueryParserLevel; + +use crate::{ + config::config, + frontend::router::parser::{Error, Shard}, + net::Query, +}; + +use super::setup::QueryParserTest; + +fn setup() -> QueryParserTest { + let mut config = (*config()).clone(); + config.config.general.query_parser = QueryParserLevel::Off; + QueryParserTest::new_single_shard(&config) +} + +fn setup_sharded() -> QueryParserTest { + let mut config = (*config()).clone(); + config.config.general.query_parser = QueryParserLevel::Off; + QueryParserTest::new_with_config(&config) +} + +const QUERIES: &[&str] = &[ + "SELECT 1", + "CREATE TABLE test (id BIGINT)", + "SELECT * FROM test", + "INSERT INTO test (id) VALUES (1)", +]; + +#[tokio::test] +async fn test_replica() { + let mut test = setup().with_param("pgdog.role", "replica"); + + for query in QUERIES { + let result = test.try_execute(vec![Query::new(query).into()]).unwrap(); + assert!(result.route().is_read()); + assert_eq!(result.route().shard(), &Shard::Direct(0)) + } +} + +#[tokio::test] +async fn test_primary() { + let mut test = setup().with_param("pgdog.role", "primary"); + + for query in QUERIES { + let result = test.try_execute(vec![Query::new(query).into()]).unwrap(); + assert!(result.route().is_write()); + assert_eq!(result.route().shard(), &Shard::Direct(0)) + } +} + +#[tokio::test] +async fn test_no_hints() { + let mut test = setup(); + + for query in QUERIES { + let result = test.try_execute(vec![Query::new(query).into()]).unwrap(); + assert!(result.route().is_write()); + assert_eq!(result.route().shard(), &Shard::Direct(0)) + } +} + +#[tokio::test] +async fn test_sharded_with_shard() { + let mut test = setup_sharded().with_param("pgdog.shard", "1"); + + for query in QUERIES { + let result = test.try_execute(vec![Query::new(query).into()]).unwrap(); + assert!(result.route().is_write()); + assert_eq!(result.route().shard(), &Shard::Direct(1)) + } +} + +#[tokio::test] +async fn test_sharded_with_shard_and_replica() { + let mut test = setup_sharded() + .with_param("pgdog.shard", "1") + .with_param("pgdog.role", "replica"); + + for query in QUERIES { + let result = test.try_execute(vec![Query::new(query).into()]).unwrap(); + assert!(result.route().is_read()); + assert_eq!(result.route().shard(), &Shard::Direct(1)) + } +} + +#[tokio::test] +async fn test_sharded_no_hints() { + let mut test = setup_sharded(); + + for query in QUERIES { + let result = test + .try_execute(vec![Query::new(query).into()]) + .unwrap_err(); + assert!(matches!(result, Error::QueryParserRequired)); + } +} diff --git a/pgdog/src/frontend/router/parser/query/test/test_comments.rs b/pgdog/src/frontend/router/parser/query/test/test_comments.rs new file mode 100644 index 000000000..83ff6a72d --- /dev/null +++ b/pgdog/src/frontend/router/parser/query/test/test_comments.rs @@ -0,0 +1,41 @@ +use crate::frontend::router::parser::Shard; + +use super::setup::*; + +#[test] +fn test_comment_pgdog_role_primary() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new("/* pgdog_role: primary */ SELECT 1").into()]); + + assert!(command.route().is_write()); +} + +#[test] +fn test_comment_pgdog_shard() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![ + Query::new("/* pgdog_shard: 1234 */ SELECT 1234").into() + ]); + + assert_eq!(command.route().shard(), &Shard::Direct(1234)); +} + +#[test] +fn test_comment_pgdog_shard_extended() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![ + Parse::named( + "__test_comment", + "/* pgdog_shard: 1234 */ SELECT * FROM sharded WHERE id = $1", + ) + .into(), + Bind::new_statement("__test_comment").into(), + Execute::new().into(), + Sync.into(), + ]); + + assert_eq!(command.route().shard(), &Shard::Direct(1234)); +} diff --git a/pgdog/src/frontend/router/parser/query/test/test_ddl.rs b/pgdog/src/frontend/router/parser/query/test/test_ddl.rs new file mode 100644 index 000000000..bf2307d10 --- /dev/null +++ b/pgdog/src/frontend/router/parser/query/test/test_ddl.rs @@ -0,0 +1,332 @@ +use crate::frontend::router::parser::Shard; +use crate::frontend::Command; + +use super::setup::{QueryParserTest, *}; + +#[test] +fn test_create_table() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new( + "CREATE TABLE test_table (id SERIAL PRIMARY KEY, name TEXT)", + ) + .into()]); + + assert!(command.route().is_write()); + assert_eq!(command.route().shard(), &Shard::All); +} + +#[test] +fn test_drop_table() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new("DROP TABLE IF EXISTS test_table").into()]); + + assert!(command.route().is_write()); + assert_eq!(command.route().shard(), &Shard::All); +} + +#[test] +fn test_alter_table() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new( + "ALTER TABLE sharded ADD COLUMN new_col TEXT", + ) + .into()]); + + assert!(command.route().is_write()); + assert_eq!(command.route().shard(), &Shard::All); +} + +#[test] +fn test_create_index() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![ + Query::new("CREATE INDEX idx_test ON sharded (email)").into() + ]); + + assert!(command.route().is_write()); + assert_eq!(command.route().shard(), &Shard::All); +} + +#[test] +fn test_drop_index() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new("DROP INDEX IF EXISTS idx_test").into()]); + + assert!(command.route().is_write()); + assert_eq!(command.route().shard(), &Shard::All); +} + +#[test] +fn test_truncate() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new("TRUNCATE TABLE sharded").into()]); + + assert!(command.route().is_write()); + assert_eq!(command.route().shard(), &Shard::All); +} + +#[test] +fn test_create_sequence() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new("CREATE SEQUENCE test_seq").into()]); + + assert!(command.route().is_write()); + assert_eq!(command.route().shard(), &Shard::All); +} + +#[test] +fn test_vacuum() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new("VACUUM sharded").into()]); + + assert!(command.route().is_write()); +} + +#[test] +fn test_analyze() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new("ANALYZE sharded").into()]); + + assert!(command.route().is_write()); +} + +#[test] +fn test_commit() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new("COMMIT").into()]); + + match command { + Command::CommitTransaction { .. } => {} + _ => panic!("expected CommitTransaction, got {command:?}"), + } +} + +#[test] +fn test_rollback() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new("ROLLBACK").into()]); + + match command { + Command::RollbackTransaction { .. } => {} + _ => panic!("expected RollbackTransaction, got {command:?}"), + } +} + +// --- Schema-based sharding tests --- + +#[test] +fn test_create_table_shard_0() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new( + "CREATE TABLE shard_0.test_table (id SERIAL PRIMARY KEY, name TEXT)", + ) + .into()]); + + assert!(command.route().is_write()); + assert_eq!(command.route().shard(), &Shard::Direct(0)); +} + +#[test] +fn test_create_table_shard_1() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new( + "CREATE TABLE shard_1.test_table (id SERIAL PRIMARY KEY, name TEXT)", + ) + .into()]); + + assert!(command.route().is_write()); + assert_eq!(command.route().shard(), &Shard::Direct(1)); +} + +#[test] +fn test_drop_table_shard_0() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![ + Query::new("DROP TABLE IF EXISTS shard_0.test_table").into() + ]); + + assert!(command.route().is_write()); + assert_eq!(command.route().shard(), &Shard::Direct(0)); +} + +#[test] +fn test_drop_table_shard_1() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![ + Query::new("DROP TABLE IF EXISTS shard_1.test_table").into() + ]); + + assert!(command.route().is_write()); + assert_eq!(command.route().shard(), &Shard::Direct(1)); +} + +#[test] +fn test_alter_table_shard_0() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new( + "ALTER TABLE shard_0.sharded ADD COLUMN new_col TEXT", + ) + .into()]); + + assert!(command.route().is_write()); + assert_eq!(command.route().shard(), &Shard::Direct(0)); +} + +#[test] +fn test_alter_table_shard_1() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new( + "ALTER TABLE shard_1.sharded ADD COLUMN new_col TEXT", + ) + .into()]); + + assert!(command.route().is_write()); + assert_eq!(command.route().shard(), &Shard::Direct(1)); +} + +#[test] +fn test_create_index_shard_0() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new( + "CREATE INDEX idx_test ON shard_0.sharded (email)", + ) + .into()]); + + assert!(command.route().is_write()); + assert_eq!(command.route().shard(), &Shard::Direct(0)); +} + +#[test] +fn test_create_index_shard_1() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new( + "CREATE INDEX idx_test ON shard_1.sharded (email)", + ) + .into()]); + + assert!(command.route().is_write()); + assert_eq!(command.route().shard(), &Shard::Direct(1)); +} + +#[test] +fn test_drop_index_shard_0() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![ + Query::new("DROP INDEX IF EXISTS shard_0.idx_test").into() + ]); + + assert!(command.route().is_write()); + assert_eq!(command.route().shard(), &Shard::Direct(0)); +} + +#[test] +fn test_drop_index_shard_1() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![ + Query::new("DROP INDEX IF EXISTS shard_1.idx_test").into() + ]); + + assert!(command.route().is_write()); + assert_eq!(command.route().shard(), &Shard::Direct(1)); +} + +#[test] +fn test_truncate_shard_0() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new("TRUNCATE TABLE shard_0.sharded").into()]); + + assert!(command.route().is_write()); + assert_eq!(command.route().shard(), &Shard::Direct(0)); +} + +#[test] +fn test_truncate_shard_1() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new("TRUNCATE TABLE shard_1.sharded").into()]); + + assert!(command.route().is_write()); + assert_eq!(command.route().shard(), &Shard::Direct(1)); +} + +#[test] +fn test_create_sequence_shard_0() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new("CREATE SEQUENCE shard_0.test_seq").into()]); + + assert!(command.route().is_write()); + assert_eq!(command.route().shard(), &Shard::Direct(0)); +} + +#[test] +fn test_create_sequence_shard_1() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new("CREATE SEQUENCE shard_1.test_seq").into()]); + + assert!(command.route().is_write()); + assert_eq!(command.route().shard(), &Shard::Direct(1)); +} + +#[test] +fn test_vacuum_shard_0() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new("VACUUM shard_0.sharded").into()]); + + assert!(command.route().is_write()); + assert_eq!(command.route().shard(), &Shard::Direct(0)); +} + +#[test] +fn test_vacuum_shard_1() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new("VACUUM shard_1.sharded").into()]); + + assert!(command.route().is_write()); + assert_eq!(command.route().shard(), &Shard::Direct(1)); +} + +#[test] +fn test_analyze_shard_0() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new("ANALYZE shard_0.sharded").into()]); + + assert!(command.route().is_write()); + assert_eq!(command.route().shard(), &Shard::Direct(0)); +} + +#[test] +fn test_analyze_shard_1() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new("ANALYZE shard_1.sharded").into()]); + + assert!(command.route().is_write()); + assert_eq!(command.route().shard(), &Shard::Direct(1)); +} diff --git a/pgdog/src/frontend/router/parser/query/test/test_delete.rs b/pgdog/src/frontend/router/parser/query/test/test_delete.rs new file mode 100644 index 000000000..168e194c9 --- /dev/null +++ b/pgdog/src/frontend/router/parser/query/test/test_delete.rs @@ -0,0 +1,86 @@ +use crate::frontend::router::parser::Shard; +use crate::net::messages::Parameter; + +use super::setup::{QueryParserTest, *}; + +#[test] +fn test_delete_with_sharding_key() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![ + Parse::named("__test_delete", "DELETE FROM sharded WHERE id = $1").into(), + Bind::new_params("__test_delete", &[Parameter::new(b"5")]).into(), + Execute::new().into(), + Sync.into(), + ]); + + assert!(command.route().is_write()); + assert!(matches!(command.route().shard(), Shard::Direct(_))); +} + +#[test] +fn test_delete_without_sharding_key() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![ + Query::new("DELETE FROM sharded WHERE email = 'test'").into() + ]); + + assert!(command.route().is_write()); + assert_eq!(command.route().shard(), &Shard::All); +} + +#[test] +fn test_delete_all_rows() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new("DELETE FROM sharded").into()]); + + assert!(command.route().is_write()); + assert_eq!(command.route().shard(), &Shard::All); +} + +#[test] +fn test_delete_with_returning() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![ + Parse::named( + "__test_del_ret", + "DELETE FROM sharded WHERE id = $1 RETURNING *", + ) + .into(), + Bind::new_params("__test_del_ret", &[Parameter::new(b"7")]).into(), + Execute::new().into(), + Sync.into(), + ]); + + assert!(command.route().is_write()); + assert!(matches!(command.route().shard(), Shard::Direct(_))); +} + +#[test] +fn test_delete_with_subquery() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new( + "DELETE FROM sharded WHERE id IN (SELECT id FROM other_table WHERE status = 'inactive')", + ) + .into()]); + + assert!(command.route().is_write()); + assert_eq!(command.route().shard(), &Shard::All); +} + +#[test] +fn test_delete_using_join() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new( + "DELETE FROM sharded USING other_table WHERE sharded.id = other_table.sharded_id", + ) + .into()]); + + assert!(command.route().is_write()); + assert_eq!(command.route().shard(), &Shard::All); +} diff --git a/pgdog/src/frontend/router/parser/query/test/test_dml.rs b/pgdog/src/frontend/router/parser/query/test/test_dml.rs new file mode 100644 index 000000000..cb085b75e --- /dev/null +++ b/pgdog/src/frontend/router/parser/query/test/test_dml.rs @@ -0,0 +1,67 @@ +use crate::frontend::router::parser::Shard; +use crate::net::messages::Parameter; + +use super::setup::*; + +#[test] +fn test_insert_with_params() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![ + Parse::named( + "__test_insert", + "INSERT INTO sharded (id, email) VALUES ($1, $2)", + ) + .into(), + Bind::new_params( + "__test_insert", + &[Parameter::new(b"11"), Parameter::new(b"test@test.com")], + ) + .into(), + Describe::new_statement("__test_insert").into(), + Execute::new().into(), + Sync.into(), + ]); + + assert_eq!(command.route().shard(), &Shard::Direct(1)); +} + +#[test] +fn test_insert_returning() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new( + "INSERT INTO foo (id) VALUES ($1::UUID) ON CONFLICT (id) DO UPDATE SET id = excluded.id RETURNING id", + ) + .into()]); + + assert!(command.route().is_write()); +} + +#[test] +fn test_select_for_update() { + let mut test = QueryParserTest::new(); + + // Without parameters - goes to all shards + let command = test.execute(vec![Query::new( + "SELECT * FROM sharded WHERE id = $1 FOR UPDATE", + ) + .into()]); + assert!(command.route().is_write()); + assert!(matches!(command.route().shard(), Shard::All)); + + // With parameter - goes to specific shard + let mut test = QueryParserTest::new(); + let command = test.execute(vec![ + Parse::named( + "__test_sfu", + "SELECT * FROM sharded WHERE id = $1 FOR UPDATE", + ) + .into(), + Bind::new_params("__test_sfu", &[Parameter::new(b"1")]).into(), + Execute::new().into(), + Sync.into(), + ]); + assert!(matches!(command.route().shard(), Shard::Direct(_))); + assert!(command.route().is_write()); +} diff --git a/pgdog/src/frontend/router/parser/query/test/test_explain.rs b/pgdog/src/frontend/router/parser/query/test/test_explain.rs new file mode 100644 index 000000000..7036f9fae --- /dev/null +++ b/pgdog/src/frontend/router/parser/query/test/test_explain.rs @@ -0,0 +1,82 @@ +use crate::frontend::router::parser::Shard; +use crate::net::messages::Parameter; + +use super::setup::{QueryParserTest, *}; + +#[test] +fn test_explain_select() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new( + "EXPLAIN SELECT * FROM sharded WHERE id = 1", + ) + .into()]); + + assert!(command.route().is_read()); + assert!(matches!(command.route().shard(), Shard::Direct(_))); +} + +#[test] +fn test_explain_analyze_select() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new( + "EXPLAIN ANALYZE SELECT * FROM sharded WHERE id = 1", + ) + .into()]); + + // EXPLAIN ANALYZE actually runs the query, so it should be treated as write + assert!(command.route().is_read()); +} + +#[test] +fn test_explain_insert() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new( + "EXPLAIN INSERT INTO sharded (id, email) VALUES (1, 'test')", + ) + .into()]); + + assert!(command.route().is_write()); +} + +#[test] +fn test_explain_with_params() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![ + Parse::named( + "__test_explain", + "EXPLAIN SELECT * FROM sharded WHERE id = $1", + ) + .into(), + Bind::new_params("__test_explain", &[Parameter::new(b"5")]).into(), + Execute::new().into(), + Sync.into(), + ]); + + assert!(command.route().is_read()); + assert!(matches!(command.route().shard(), Shard::Direct(_))); +} + +#[test] +fn test_explain_all_shards() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new("EXPLAIN SELECT * FROM sharded").into()]); + + assert_eq!(command.route().shard(), &Shard::All); +} + +#[test] +fn test_explain_verbose() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new( + "EXPLAIN (VERBOSE, COSTS) SELECT * FROM sharded WHERE id = 1", + ) + .into()]); + + assert!(command.route().is_read()); +} diff --git a/pgdog/src/frontend/router/parser/query/test/test_functions.rs b/pgdog/src/frontend/router/parser/query/test/test_functions.rs new file mode 100644 index 000000000..0b37fab37 --- /dev/null +++ b/pgdog/src/frontend/router/parser/query/test/test_functions.rs @@ -0,0 +1,21 @@ +use super::setup::*; + +#[test] +fn test_write_function_advisory_lock() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new("SELECT pg_advisory_lock($1)").into()]); + + assert!(command.route().is_write()); + assert!(command.route().is_lock_session()); +} + +#[test] +fn test_write_function_nextval() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new("SELECT nextval('234')").into()]); + + assert!(command.route().is_write()); + assert!(!command.route().is_lock_session()); +} diff --git a/pgdog/src/frontend/router/parser/query/test/test_insert.rs b/pgdog/src/frontend/router/parser/query/test/test_insert.rs new file mode 100644 index 000000000..ff3a76d5d --- /dev/null +++ b/pgdog/src/frontend/router/parser/query/test/test_insert.rs @@ -0,0 +1,89 @@ +use crate::frontend::router::parser::Shard; +use crate::net::messages::Parameter; + +use super::setup::*; + +#[test] +fn test_insert_numeric() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new( + "INSERT INTO sharded (id, sample_numeric) VALUES (2, -987654321.123456789::NUMERIC)", + ) + .into()]); + + assert!(command.route().is_write()); + assert!(matches!(command.route().shard(), Shard::Direct(_))); +} + +#[test] +fn test_insert_negative_sharding_key() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![ + Query::new("INSERT INTO sharded (id) VALUES (-5)").into() + ]); + + assert!(command.route().is_write()); + assert!(matches!(command.route().shard(), Shard::Direct(_))); +} + +#[test] +fn test_insert_with_cast_on_sharding_key() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new( + "INSERT INTO sharded (id, value) VALUES (42::BIGINT, 'test')", + ) + .into()]); + + assert!(command.route().is_write()); + assert!(matches!(command.route().shard(), Shard::Direct(_))); +} + +#[test] +fn test_insert_multi_row() { + let mut test = QueryParserTest::new(); + + // Multi-row inserts go to all shards (split by query engine later) + let command = test.execute(vec![ + Parse::named( + "__test_multi", + "INSERT INTO sharded (id, value) VALUES ($1, 'a'), ($2, 'b')", + ) + .into(), + Bind::new_params( + "__test_multi", + &[Parameter::new(b"0"), Parameter::new(b"2")], + ) + .into(), + Execute::new().into(), + Sync.into(), + ]); + + assert!(command.route().is_write()); + assert!(matches!(command.route().shard(), Shard::All)); +} + +#[test] +fn test_insert_select() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new( + "INSERT INTO sharded (id, value) SELECT id, value FROM other_table WHERE id = 1", + ) + .into()]); + + assert!(command.route().is_write()); + assert!(command.route().is_all_shards()); +} + +#[test] +fn test_insert_default_values() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new("INSERT INTO sharded DEFAULT VALUES").into()]); + + assert!(command.route().is_write()); + assert!(command.route().is_all_shards()); +} diff --git a/pgdog/src/frontend/router/parser/query/test/test_rr.rs b/pgdog/src/frontend/router/parser/query/test/test_rr.rs new file mode 100644 index 000000000..53f42e7c0 --- /dev/null +++ b/pgdog/src/frontend/router/parser/query/test/test_rr.rs @@ -0,0 +1,34 @@ +use crate::frontend::router::parser::{ + route::{RoundRobinReason, ShardSource}, + Shard, +}; + +use super::setup::*; + +#[test] +fn test_rr_executable() { + let mut test = QueryParserTest::new(); + let command = test.execute(vec![ + Parse::named( + "__test_1", + "INSERT INTO some_table (id, value) VALUES ($1, $2)", + ) + .into(), + Describe::new_statement("__test_1").into(), + Flush.into(), + ]); + + assert!(matches!(command.route().shard(), Shard::Direct(_))); + assert_eq!( + command.route().shard_with_priority().source(), + &ShardSource::RoundRobin(RoundRobinReason::NotExecutable) + ); + + let command = test.execute(vec![ + Bind::new_params("__test_1", &[]).into(), + Execute::new().into(), + Sync.into(), + ]); + + assert!(matches!(command.route().shard(), Shard::All)); +} diff --git a/pgdog/src/frontend/router/parser/query/test/test_schema_sharding.rs b/pgdog/src/frontend/router/parser/query/test/test_schema_sharding.rs new file mode 100644 index 000000000..50d018541 --- /dev/null +++ b/pgdog/src/frontend/router/parser/query/test/test_schema_sharding.rs @@ -0,0 +1,293 @@ +use crate::frontend::router::parser::{route::ShardSource, Shard}; +use crate::net::parameter::ParameterValue; + +use super::setup::{QueryParserTest, *}; + +// --- SELECT queries with schema-qualified tables --- + +#[test] +fn test_select_from_shard_0_schema() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![ + Query::new("SELECT * FROM shard_0.users WHERE id = 1").into() + ]); + + assert!(command.route().is_read()); + assert_eq!(command.route().shard(), &Shard::Direct(0)); +} + +#[test] +fn test_select_from_shard_1_schema() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![ + Query::new("SELECT * FROM shard_1.users WHERE id = 1").into() + ]); + + assert!(command.route().is_read()); + assert_eq!(command.route().shard(), &Shard::Direct(1)); +} + +#[test] +fn test_select_from_unsharded_schema_goes_to_all() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![ + Query::new("SELECT * FROM public.users WHERE id = 1").into() + ]); + + // Unknown schema goes to all shards + assert_eq!(command.route().shard(), &Shard::All); +} + +// --- INSERT queries with schema-qualified tables --- + +#[test] +fn test_insert_into_shard_0_schema() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new( + "INSERT INTO shard_0.users (id, name) VALUES (1, 'test')", + ) + .into()]); + + assert!(command.route().is_write()); + assert_eq!(command.route().shard(), &Shard::Direct(0)); +} + +#[test] +fn test_insert_into_shard_1_schema() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new( + "INSERT INTO shard_1.users (id, name) VALUES (1, 'test')", + ) + .into()]); + + assert!(command.route().is_write()); + assert_eq!(command.route().shard(), &Shard::Direct(1)); +} + +// --- UPDATE queries with schema-qualified tables --- + +#[test] +fn test_update_shard_0_schema() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new( + "UPDATE shard_0.users SET name = 'updated' WHERE id = 1", + ) + .into()]); + + assert!(command.route().is_write()); + assert_eq!(command.route().shard(), &Shard::Direct(0)); +} + +#[test] +fn test_update_shard_1_schema() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new( + "UPDATE shard_1.users SET name = 'updated' WHERE id = 1", + ) + .into()]); + + assert!(command.route().is_write()); + assert_eq!(command.route().shard(), &Shard::Direct(1)); +} + +// --- DELETE queries with schema-qualified tables --- + +#[test] +fn test_delete_from_shard_0_schema() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![ + Query::new("DELETE FROM shard_0.users WHERE id = 1").into() + ]); + + assert!(command.route().is_write()); + assert_eq!(command.route().shard(), &Shard::Direct(0)); +} + +#[test] +fn test_delete_from_shard_1_schema() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![ + Query::new("DELETE FROM shard_1.users WHERE id = 1").into() + ]); + + assert!(command.route().is_write()); + assert_eq!(command.route().shard(), &Shard::Direct(1)); +} + +// --- JOIN queries with schema-qualified tables --- + +#[test] +fn test_join_same_schema() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new( + "SELECT * FROM shard_0.users u JOIN shard_0.orders o ON u.id = o.user_id", + ) + .into()]); + + assert!(command.route().is_read()); + assert_eq!(command.route().shard(), &Shard::Direct(0)); +} + +#[test] +fn test_join_different_schemas_uses_first_found() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new( + "SELECT * FROM shard_0.users u JOIN shard_1.orders o ON u.id = o.user_id", + ) + .into()]); + + // Cross-schema join uses the first schema found in the query + assert!(matches!(command.route().shard(), Shard::Direct(_))); +} + +// --- DDL with schema-qualified tables --- + +#[test] +fn test_create_table_in_shard_0_schema() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new( + "CREATE TABLE shard_0.new_table (id BIGINT PRIMARY KEY)", + ) + .into()]); + + assert!(command.route().is_write()); + assert_eq!(command.route().shard(), &Shard::Direct(0)); +} + +#[test] +fn test_create_table_in_shard_1_schema() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new( + "CREATE TABLE shard_1.new_table (id BIGINT PRIMARY KEY)", + ) + .into()]); + + assert!(command.route().is_write()); + assert_eq!(command.route().shard(), &Shard::Direct(1)); +} + +#[test] +fn test_drop_table_in_sharded_schema() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![ + Query::new("DROP TABLE IF EXISTS shard_0.old_table").into() + ]); + + assert!(command.route().is_write()); + assert_eq!(command.route().shard(), &Shard::Direct(0)); +} + +#[test] +fn test_alter_table_in_sharded_schema() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new( + "ALTER TABLE shard_1.users ADD COLUMN email TEXT", + ) + .into()]); + + assert!(command.route().is_write()); + assert_eq!(command.route().shard(), &Shard::Direct(1)); +} + +// --- Index operations with schema-qualified tables --- + +#[test] +fn test_create_index_on_sharded_schema() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new( + "CREATE INDEX idx_users_name ON shard_0.users (name)", + ) + .into()]); + + assert!(command.route().is_write()); + assert_eq!(command.route().shard(), &Shard::Direct(0)); +} + +// --- Schema takes priority over table-based sharding --- + +#[test] +#[ignore = "this is actually not the case currently"] +fn test_schema_sharding_takes_priority_over_table_sharding() { + let mut test = QueryParserTest::new(); + + // "sharded" table is configured for table-based sharding with column "id" + // id=1 would normally hash to some shard, but schema prefix should override + let command = test.execute(vec![Query::new( + "SELECT * FROM shard_1.sharded WHERE id = 1", + ) + .into()]); + + // Schema-based routing should take priority: shard_1 -> shard 1 + assert_eq!(command.route().shard(), &Shard::Direct(1)); +} + +#[test] +fn test_schema_sharding_priority_on_insert() { + let mut test = QueryParserTest::new(); + + // Even though "sharded" table has sharding key "id", schema should win + let command = test.execute(vec![Query::new( + "INSERT INTO shard_0.sharded (id, email) VALUES (999, 'test@test.com')", + ) + .into()]); + + assert_eq!(command.route().shard(), &Shard::Direct(0)); +} + +#[test] +#[ignore = "this is not currently how it works, but it should"] +fn test_schema_sharding_priority_on_update() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new( + "UPDATE shard_1.sharded SET email = 'new@test.com' WHERE id = 1", + ) + .into()]); + + assert_eq!(command.route().shard(), &Shard::Direct(1)); +} + +#[test] +#[ignore = "this is not currently how it works, but it should"] +fn test_schema_sharding_priority_on_delete() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new( + "DELETE FROM shard_0.sharded WHERE id = 999", + ) + .into()]); + + assert_eq!(command.route().shard(), &Shard::Direct(0)); +} + +// --- SET commands with schema sharding via search_path --- + +#[test] +fn test_set_routes_to_shard_from_search_path() { + let mut test = + QueryParserTest::new().with_param("search_path", ParameterValue::String("shard_0".into())); + + let command = test.execute(vec![Query::new("SET statement_timeout TO 1000").into()]); + + assert_eq!(command.route().shard(), &Shard::Direct(0)); + assert_eq!( + command.route().shard_with_priority().source(), + &ShardSource::SearchPath("shard_0".into()) + ); +} diff --git a/pgdog/src/frontend/router/parser/query/test/test_search_path.rs b/pgdog/src/frontend/router/parser/query/test/test_search_path.rs new file mode 100644 index 000000000..bfd1a53e5 --- /dev/null +++ b/pgdog/src/frontend/router/parser/query/test/test_search_path.rs @@ -0,0 +1,302 @@ +use crate::frontend::router::parser::Shard; +use crate::net::parameter::ParameterValue; + +use super::setup::{QueryParserTest, *}; + +// --- search_path routing for DML --- + +#[test] +fn test_search_path_shard_0_routes_select() { + let mut test = QueryParserTest::new().with_param( + "search_path", + ParameterValue::Tuple(vec!["$user".into(), "shard_0".into(), "public".into()]), + ); + + let command = test.execute(vec![Query::new("SELECT * FROM users WHERE id = 1").into()]); + + assert!(command.route().is_read()); + assert_eq!(command.route().shard(), &Shard::Direct(0)); +} + +#[test] +fn test_search_path_shard_1_routes_select() { + let mut test = QueryParserTest::new().with_param( + "search_path", + ParameterValue::Tuple(vec!["public".into(), "shard_1".into(), "$user".into()]), + ); + + let command = test.execute(vec![Query::new("SELECT * FROM users WHERE id = 1").into()]); + + assert!(command.route().is_read()); + assert_eq!(command.route().shard(), &Shard::Direct(1)); +} + +#[test] +fn test_search_path_shard_0_routes_insert() { + let mut test = QueryParserTest::new().with_param( + "search_path", + ParameterValue::Tuple(vec!["pg_catalog".into(), "shard_0".into()]), + ); + + let command = test.execute(vec![Query::new( + "INSERT INTO users (id, name) VALUES (1, 'test')", + ) + .into()]); + + assert!(command.route().is_write()); + assert_eq!(command.route().shard(), &Shard::Direct(0)); +} + +#[test] +fn test_search_path_shard_1_routes_insert() { + let mut test = QueryParserTest::new().with_param( + "search_path", + ParameterValue::Tuple(vec![ + "$user".into(), + "public".into(), + "pg_temp".into(), + "shard_1".into(), + ]), + ); + + let command = test.execute(vec![Query::new( + "INSERT INTO users (id, name) VALUES (1, 'test')", + ) + .into()]); + + assert!(command.route().is_write()); + assert_eq!(command.route().shard(), &Shard::Direct(1)); +} + +#[test] +fn test_search_path_shard_0_routes_update() { + let mut test = QueryParserTest::new().with_param( + "search_path", + ParameterValue::Tuple(vec!["shard_0".into(), "pg_catalog".into(), "public".into()]), + ); + + let command = test.execute(vec![Query::new( + "UPDATE users SET name = 'updated' WHERE id = 1", + ) + .into()]); + + assert!(command.route().is_write()); + assert_eq!(command.route().shard(), &Shard::Direct(0)); +} + +#[test] +fn test_search_path_shard_1_routes_update() { + let mut test = QueryParserTest::new().with_param( + "search_path", + ParameterValue::Tuple(vec!["information_schema".into(), "shard_1".into()]), + ); + + let command = test.execute(vec![Query::new( + "UPDATE users SET name = 'updated' WHERE id = 1", + ) + .into()]); + + assert!(command.route().is_write()); + assert_eq!(command.route().shard(), &Shard::Direct(1)); +} + +#[test] +fn test_search_path_shard_0_routes_delete() { + let mut test = QueryParserTest::new().with_param( + "search_path", + ParameterValue::Tuple(vec!["public".into(), "$user".into(), "shard_0".into()]), + ); + + let command = test.execute(vec![Query::new("DELETE FROM users WHERE id = 1").into()]); + + assert!(command.route().is_write()); + assert_eq!(command.route().shard(), &Shard::Direct(0)); +} + +#[test] +fn test_search_path_shard_1_routes_delete() { + let mut test = QueryParserTest::new().with_param( + "search_path", + ParameterValue::Tuple(vec!["shard_1".into(), "public".into()]), + ); + + let command = test.execute(vec![Query::new("DELETE FROM users WHERE id = 1").into()]); + + assert!(command.route().is_write()); + assert_eq!(command.route().shard(), &Shard::Direct(1)); +} + +// --- search_path routing priority (first sharded schema wins) --- + +#[test] +fn test_search_path_first_shard_wins_shard_0() { + let mut test = QueryParserTest::new().with_param( + "search_path", + ParameterValue::Tuple(vec![ + "public".into(), + "shard_0".into(), + "shard_1".into(), + "$user".into(), + ]), + ); + + let command = test.execute(vec![Query::new("SELECT * FROM users WHERE id = 1").into()]); + + assert!(command.route().is_read()); + assert_eq!(command.route().shard(), &Shard::Direct(0)); +} + +#[test] +fn test_search_path_first_shard_wins_shard_1() { + let mut test = QueryParserTest::new().with_param( + "search_path", + ParameterValue::Tuple(vec![ + "$user".into(), + "pg_catalog".into(), + "shard_1".into(), + "shard_0".into(), + ]), + ); + + let command = test.execute(vec![Query::new("SELECT * FROM users WHERE id = 1").into()]); + + assert!(command.route().is_read()); + assert_eq!(command.route().shard(), &Shard::Direct(1)); +} + +#[test] +fn test_search_path_shard_at_end_still_matches() { + let mut test = QueryParserTest::new().with_param( + "search_path", + ParameterValue::Tuple(vec![ + "$user".into(), + "public".into(), + "pg_catalog".into(), + "information_schema".into(), + "shard_1".into(), + ]), + ); + + let command = test.execute(vec![Query::new("SELECT * FROM users WHERE id = 1").into()]); + + assert!(command.route().is_read()); + assert_eq!(command.route().shard(), &Shard::Direct(1)); +} + +// --- search_path with no sharded schema routes to all --- + +#[test] +fn test_search_path_no_sharded_schema_routes_to_all() { + let mut test = QueryParserTest::new().with_param( + "search_path", + ParameterValue::Tuple(vec!["$user".into(), "public".into(), "pg_catalog".into()]), + ); + + let command = test.execute(vec![Query::new("SELECT * FROM users WHERE id = 1").into()]); + + assert_eq!(command.route().shard(), &Shard::All); +} + +#[test] +fn test_search_path_only_system_schemas_routes_to_all() { + let mut test = QueryParserTest::new().with_param( + "search_path", + ParameterValue::Tuple(vec![ + "pg_catalog".into(), + "information_schema".into(), + "pg_temp".into(), + ]), + ); + + let command = test.execute(vec![Query::new("SELECT * FROM users WHERE id = 1").into()]); + + assert_eq!(command.route().shard(), &Shard::All); +} + +// --- search_path routing for DDL --- + +#[test] +fn test_search_path_shard_0_routes_create_table() { + let mut test = QueryParserTest::new().with_param( + "search_path", + ParameterValue::Tuple(vec!["$user".into(), "shard_0".into(), "public".into()]), + ); + + let command = test.execute(vec![Query::new( + "CREATE TABLE new_table (id SERIAL PRIMARY KEY)", + ) + .into()]); + + assert!(command.route().is_write()); + assert_eq!(command.route().shard(), &Shard::Direct(0)); +} + +#[test] +fn test_search_path_shard_1_routes_create_table() { + let mut test = QueryParserTest::new().with_param( + "search_path", + ParameterValue::Tuple(vec!["public".into(), "pg_catalog".into(), "shard_1".into()]), + ); + + let command = test.execute(vec![Query::new( + "CREATE TABLE new_table (id SERIAL PRIMARY KEY)", + ) + .into()]); + + assert!(command.route().is_write()); + assert_eq!(command.route().shard(), &Shard::Direct(1)); +} + +#[test] +fn test_search_path_shard_0_routes_drop_table() { + let mut test = QueryParserTest::new().with_param( + "search_path", + ParameterValue::Tuple(vec!["pg_temp".into(), "$user".into(), "shard_0".into()]), + ); + + let command = test.execute(vec![Query::new("DROP TABLE IF EXISTS old_table").into()]); + + assert!(command.route().is_write()); + assert_eq!(command.route().shard(), &Shard::Direct(0)); +} + +#[test] +fn test_search_path_shard_1_routes_alter_table() { + let mut test = QueryParserTest::new().with_param( + "search_path", + ParameterValue::Tuple(vec!["shard_1".into(), "$user".into(), "public".into()]), + ); + + let command = test.execute(vec![ + Query::new("ALTER TABLE users ADD COLUMN email TEXT").into() + ]); + + assert!(command.route().is_write()); + assert_eq!(command.route().shard(), &Shard::Direct(1)); +} + +// --- search_path driven flag --- + +#[test] +fn test_search_path_driven_flag_is_set() { + let mut test = QueryParserTest::new().with_param( + "search_path", + ParameterValue::Tuple(vec!["$user".into(), "public".into(), "shard_0".into()]), + ); + + let command = test.execute(vec![Query::new("SELECT * FROM users WHERE id = 1").into()]); + + assert!(command.route().is_search_path_driven()); +} + +#[test] +fn test_search_path_driven_flag_not_set_for_no_sharded_schema() { + let mut test = QueryParserTest::new().with_param( + "search_path", + ParameterValue::Tuple(vec!["$user".into(), "public".into(), "pg_catalog".into()]), + ); + + let command = test.execute(vec![Query::new("SELECT * FROM users WHERE id = 1").into()]); + + assert!(!command.route().is_search_path_driven()); +} diff --git a/pgdog/src/frontend/router/parser/query/test/test_select.rs b/pgdog/src/frontend/router/parser/query/test/test_select.rs new file mode 100644 index 000000000..d2e405b9f --- /dev/null +++ b/pgdog/src/frontend/router/parser/query/test/test_select.rs @@ -0,0 +1,149 @@ +use crate::frontend::router::parser::{DistinctBy, DistinctColumn, Shard}; +use crate::net::messages::Parameter; + +use super::setup::*; + +#[test] +fn test_order_by_vector_simple() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new( + "SELECT * FROM embeddings ORDER BY embedding <-> '[1,2,3]'", + ) + .into()]); + + let route = command.route(); + let order_by = route.order_by().first().unwrap(); + assert!(order_by.asc()); +} + +#[test] +fn test_order_by_vector_with_params() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![ + Parse::named( + "__test_order", + "SELECT * FROM embeddings ORDER BY embedding <-> $1", + ) + .into(), + Bind::new_params("__test_order", &[Parameter::new(b"[4.0,5.0,6.0]")]).into(), + Execute::new().into(), + Sync.into(), + ]); + + let route = command.route(); + let order_by = route.order_by().first().unwrap(); + assert!(order_by.asc()); +} + +#[test] +fn test_limit_offset_simple() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![ + Query::new("SELECT * FROM users LIMIT 25 OFFSET 5").into() + ]); + + let route = command.route(); + assert_eq!(route.limit().offset, Some(5)); + assert_eq!(route.limit().limit, Some(25)); +} + +#[test] +fn test_limit_offset_with_params() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![ + Parse::named("__test_limit", "SELECT * FROM users LIMIT $1 OFFSET $2").into(), + Bind::new_params( + "__test_limit", + &[Parameter::new(b"1"), Parameter::new(b"25")], + ) + .into(), + Execute::new().into(), + Sync.into(), + ]); + + let route = command.route(); + assert_eq!(route.limit().limit, Some(1)); + assert_eq!(route.limit().offset, Some(25)); +} + +#[test] +fn test_distinct_row() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new("SELECT DISTINCT * FROM users").into()]); + + let route = command.route(); + let distinct = route.distinct().as_ref().unwrap(); + assert_eq!(distinct, &DistinctBy::Row); +} + +#[test] +fn test_distinct_on_columns() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new( + "SELECT DISTINCT ON(1, email) * FROM users", + ) + .into()]); + + let route = command.route(); + let distinct = route.distinct().as_ref().unwrap(); + assert_eq!( + distinct, + &DistinctBy::Columns(vec![ + DistinctColumn::Index(0), + DistinctColumn::Name(String::from("email")) + ]) + ); +} + +#[test] +fn test_any_literal() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new( + "SELECT * FROM sharded WHERE id = ANY('{1, 2, 3}')", + ) + .into()]); + + assert_eq!(command.route().shard(), &Shard::All); +} + +#[test] +fn test_any_with_param() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![ + Parse::named("__test_any", "SELECT * FROM sharded WHERE id = ANY($1)").into(), + Bind::new_params("__test_any", &[Parameter::new(b"{1, 2, 3}")]).into(), + Execute::new().into(), + Sync.into(), + ]); + + assert_eq!(command.route().shard(), &Shard::All); +} + +#[test] +fn test_cte_read() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new("WITH s AS (SELECT 1) SELECT 2").into()]); + + assert!(command.route().is_read()); +} + +#[test] +fn test_cte_write() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new( + "WITH s AS (SELECT 1), s2 AS (INSERT INTO test VALUES ($1) RETURNING *), s3 AS (SELECT 123) SELECT * FROM s", + ) + .into()]); + + assert!(command.route().is_write()); +} diff --git a/pgdog/src/frontend/router/parser/query/test/test_set.rs b/pgdog/src/frontend/router/parser/query/test/test_set.rs new file mode 100644 index 000000000..0cc8a3f1a --- /dev/null +++ b/pgdog/src/frontend/router/parser/query/test/test_set.rs @@ -0,0 +1,39 @@ +use crate::{frontend::Command, net::parameter::ParameterValue}; + +use super::setup::*; + +#[test] +fn test_set_comment() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new( + "/* pgdog_sharding_key: 1234 */ SET statement_timeout TO 1", + ) + .into()]); + + assert!( + matches!(command.clone(), Command::Set { name, value, local, route } if name == "statement_timeout" && !local && value == ParameterValue::String("1".into()) && route.shard().is_direct()), + "expected Command::Set, got {:#?}", + command, + ); +} + +#[test] +fn test_set_transaction_level() { + let mut test = QueryParserTest::new(); + + for query in [ + "SET TRANSACTION SNAPSHOT '00000003-0000001B-1'", + "SET TRANSACTION ISOLATION LEVEL REPEATABLE READ", + "set transaction isolation level repeatable read", + "set transaction snapshot '00000003-0000001B-1'", + "SET SESSION CHARACTERISTICS AS TRANSACTION ISOLATION LEVEL SERIALIZABLE", + ] { + let command = test.execute(vec![Query::new(query).into()]); + assert!( + matches!(command.clone(), Command::Query(_)), + "{:#?}", + command + ); + } +} diff --git a/pgdog/src/frontend/router/parser/query/test/test_sharding.rs b/pgdog/src/frontend/router/parser/query/test/test_sharding.rs new file mode 100644 index 000000000..57064a5c5 --- /dev/null +++ b/pgdog/src/frontend/router/parser/query/test/test_sharding.rs @@ -0,0 +1,123 @@ +use std::collections::HashSet; + +use crate::config::config; +use crate::frontend::router::parser::{Cache, Shard}; +use crate::frontend::Command; + +use super::setup::{QueryParserTest, *}; + +#[test] +fn test_show_shards() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new("SHOW pgdog.shards").into()]); + + assert!(matches!(command, Command::InternalField { .. })); +} + +#[test] +fn test_close_direct_single_shard() { + let mut test = QueryParserTest::new_single_shard(&config()); + + let command = test.execute(vec![Close::named("test").into(), Sync.into()]); + + match command { + Command::Query(route) => assert_eq!(route.shard(), &Shard::Direct(0), "{:?}", route), + _ => panic!("expected Query, got {command:?}"), + } +} + +#[test] +fn test_dry_run_simple() { + let mut test = QueryParserTest::new_single_shard(&config()).with_dry_run(); + + let command = test.execute(vec![Query::new( + "/* pgdog_sharding_key: 1234 */ SELECT * FROM sharded", + ) + .into()]); + + let cache = Cache::queries(); + let stmt = cache.values().next().unwrap(); + assert_eq!(stmt.stats.lock().direct, 1); + assert_eq!(stmt.stats.lock().multi, 0); + assert_eq!(command.route().shard(), &Shard::Direct(0)); +} + +#[test] +fn test_omni_round_robin() { + let mut omni_round_robin = HashSet::new(); + let q = "SELECT sharded_omni.* FROM sharded_omni WHERE sharded_omni.id = 1"; + + for _ in 0..10 { + let mut test = QueryParserTest::new(); + let command = test.execute(vec![Query::new(q).into()]); + + match command { + Command::Query(query) => { + assert!(matches!(query.shard(), Shard::Direct(_))); + omni_round_robin.insert(query.shard().clone()); + } + _ => {} + } + } + + assert_eq!(omni_round_robin.len(), 2); +} + +#[test] +fn test_omni_sticky() { + let mut omni_sticky = HashSet::new(); + let q = + "SELECT sharded_omni_sticky.* FROM sharded_omni_sticky WHERE sharded_omni_sticky.id = $1"; + + // Use a single test instance with consistent Sticky across all iterations + let mut test = QueryParserTest::new(); + for _ in 0..10 { + let command = test.execute(vec![Query::new(q).into()]); + + match command { + Command::Query(query) => { + assert!(matches!(query.shard(), Shard::Direct(_))); + omni_sticky.insert(query.shard().clone()); + } + _ => {} + } + } + + assert_eq!(omni_sticky.len(), 1); +} + +#[test] +fn test_omni_sharded_table_takes_priority() { + let q = " + SELECT + sharded_omni.*, + sharded.* + FROM + sharded_omni + INNER JOIN + sharded + ON sharded_omni.id = sharded.i + WHERE sharded.id = 5"; + + let mut test = QueryParserTest::new(); + let command = test.execute(vec![Query::new(q).into()]); + let first_shard = command.route().shard().clone(); + + for _ in 0..5 { + let mut test = QueryParserTest::new(); + let command = test.execute(vec![Query::new(q).into()]); + assert_eq!(&first_shard, command.route().shard()); + assert!(matches!(first_shard, Shard::Direct(_))); + } +} + +#[test] +fn test_omni_all_tables_must_be_omnisharded() { + let q = "SELECT * FROM sharded_omni INNER JOIN not_sharded ON sharded_omni.id = not_sharded.id WHERE sharded_omni = $1"; + + let mut test = QueryParserTest::new(); + let command = test.execute(vec![Query::new(q).into()]); + + assert!(matches!(command.route().shard(), Shard::All)); +} diff --git a/pgdog/src/frontend/router/parser/query/test/test_special.rs b/pgdog/src/frontend/router/parser/query/test/test_special.rs new file mode 100644 index 000000000..d3b314d0f --- /dev/null +++ b/pgdog/src/frontend/router/parser/query/test/test_special.rs @@ -0,0 +1,329 @@ +use crate::frontend::router::parser::Shard; +use crate::frontend::Command; +use crate::net::messages::Parameter; + +use super::setup::{QueryParserTest, *}; + +// --- NULL handling --- + +#[test] +fn test_where_is_null() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![ + Query::new("SELECT * FROM sharded WHERE id IS NULL").into() + ]); + + // NULL can be on any shard + assert_eq!(command.route().shard(), &Shard::All); +} + +#[test] +fn test_where_is_not_null() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new( + "SELECT * FROM sharded WHERE id IS NOT NULL", + ) + .into()]); + + assert_eq!(command.route().shard(), &Shard::All); +} + +// --- OR conditions --- + +#[test] +fn test_where_or_same_key() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new( + "SELECT * FROM sharded WHERE id = 1 OR id = 2", + ) + .into()]); + + // OR with different values goes to all shards + assert_eq!(command.route().shard(), &Shard::All); +} + +#[test] +fn test_where_in_list() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new( + "SELECT * FROM sharded WHERE id IN (1, 2, 3)", + ) + .into()]); + + // IN list with multiple values routes to Multi shard covering those specific shards + assert!(matches!(command.route().shard(), Shard::Multi(_))); +} + +// --- CASE expressions --- + +#[test] +fn test_case_expression() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new( + "SELECT CASE WHEN id = 1 THEN 'one' ELSE 'other' END FROM sharded WHERE id = 1", + ) + .into()]); + + assert!(matches!(command.route().shard(), Shard::Direct(_))); +} + +// --- UNION queries --- + +#[test] +fn test_union() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new( + "SELECT id FROM sharded WHERE id = 1 UNION SELECT id FROM sharded WHERE id = 2", + ) + .into()]); + + assert!(command.route().is_read()); + // UNION finds the first matching shard key - both sides have literals so it picks one + assert!(matches!(command.route().shard(), Shard::Direct(_))); +} + +#[test] +fn test_union_all() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new( + "SELECT id FROM sharded WHERE id = 1 UNION ALL SELECT id FROM sharded WHERE id = 2", + ) + .into()]); + + assert!(command.route().is_read()); +} + +// --- Aggregate functions --- + +#[test] +fn test_count_all() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new("SELECT COUNT(*) FROM sharded").into()]); + + assert!(command.route().is_read()); + assert_eq!(command.route().shard(), &Shard::All); +} + +#[test] +fn test_sum_with_group_by() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new( + "SELECT email, SUM(amount) FROM sharded GROUP BY email", + ) + .into()]); + + assert!(command.route().is_read()); + assert_eq!(command.route().shard(), &Shard::All); +} + +// --- Window functions --- + +#[test] +fn test_window_function() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new( + "SELECT id, ROW_NUMBER() OVER (ORDER BY id) FROM sharded WHERE id = 1", + ) + .into()]); + + assert!(command.route().is_read()); +} + +// --- RETURNING clause --- + +#[test] +fn test_insert_returning() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![ + Parse::named( + "__test_ins_ret", + "INSERT INTO sharded (id, email) VALUES ($1, $2) RETURNING id, email", + ) + .into(), + Bind::new_params( + "__test_ins_ret", + &[Parameter::new(b"5"), Parameter::new(b"test@test.com")], + ) + .into(), + Execute::new().into(), + Sync.into(), + ]); + + assert!(command.route().is_write()); +} + +#[test] +fn test_update_returning() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![ + Parse::named( + "__test_upd_ret", + "UPDATE sharded SET email = $1 WHERE id = $2 RETURNING *", + ) + .into(), + Bind::new_params( + "__test_upd_ret", + &[Parameter::new(b"new@test.com"), Parameter::new(b"5")], + ) + .into(), + Execute::new().into(), + Sync.into(), + ]); + + assert!(command.route().is_write()); +} + +// --- COALESCE and NULL functions --- + +#[test] +fn test_coalesce_in_where() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new( + "SELECT * FROM sharded WHERE COALESCE(id, 0) = 1", + ) + .into()]); + + assert!(command.route().is_read()); +} + +// --- BETWEEN --- + +#[test] +fn test_between_on_sharding_key() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new( + "SELECT * FROM sharded WHERE id BETWEEN 1 AND 10", + ) + .into()]); + + // BETWEEN can span multiple shards + assert_eq!(command.route().shard(), &Shard::All); +} + +// --- LIKE/ILIKE --- + +#[test] +fn test_like_on_non_sharding_column() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new( + "SELECT * FROM sharded WHERE email LIKE '%test%'", + ) + .into()]); + + assert_eq!(command.route().shard(), &Shard::All); +} + +// --- SHOW commands --- + +#[test] +fn test_show_tables() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new("SHOW TABLES").into()]); + + match command { + Command::Query(_) | Command::InternalField { .. } => {} + _ => panic!("expected Query or InternalField, got {command:?}"), + } +} + +#[test] +fn test_show_server_version() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new("SHOW server_version").into()]); + + match command { + Command::Query(_) | Command::InternalField { .. } => {} + _ => panic!("expected Query or InternalField, got {command:?}"), + } +} + +// --- Locking --- + +#[test] +fn test_select_for_share() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![ + Parse::named( + "__test_share", + "SELECT * FROM sharded WHERE id = $1 FOR SHARE", + ) + .into(), + Bind::new_params("__test_share", &[Parameter::new(b"1")]).into(), + Execute::new().into(), + Sync.into(), + ]); + + assert!(command.route().is_write()); +} + +#[test] +fn test_select_for_no_key_update() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![ + Parse::named( + "__test_nku", + "SELECT * FROM sharded WHERE id = $1 FOR NO KEY UPDATE", + ) + .into(), + Bind::new_params("__test_nku", &[Parameter::new(b"1")]).into(), + Execute::new().into(), + Sync.into(), + ]); + + assert!(command.route().is_write()); +} + +// --- Cast expressions --- + +#[test] +fn test_cast_sharding_key() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![ + Parse::named( + "__test_cast", + "SELECT * FROM sharded WHERE id = $1::INTEGER", + ) + .into(), + Bind::new_params("__test_cast", &[Parameter::new(b"5")]).into(), + Execute::new().into(), + Sync.into(), + ]); + + assert!(matches!(command.route().shard(), Shard::Direct(_))); +} + +// --- Empty/edge parameter values --- + +#[test] +fn test_empty_string_param() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![ + Parse::named("__test_empty", "SELECT * FROM sharded WHERE email = $1").into(), + Bind::new_params("__test_empty", &[Parameter::new(b"")]).into(), + Execute::new().into(), + Sync.into(), + ]); + + assert!(command.route().is_read()); +} diff --git a/pgdog/src/frontend/router/parser/query/test/test_subqueries.rs b/pgdog/src/frontend/router/parser/query/test/test_subqueries.rs new file mode 100644 index 000000000..db9ef1509 --- /dev/null +++ b/pgdog/src/frontend/router/parser/query/test/test_subqueries.rs @@ -0,0 +1,134 @@ +use crate::frontend::router::parser::Shard; +use crate::net::messages::Parameter; + +use super::setup::{QueryParserTest, *}; + +#[test] +fn test_subquery_in_where() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new( + "SELECT * FROM sharded WHERE id IN (SELECT id FROM other_table WHERE status = 'active')", + ) + .into()]); + + assert!(command.route().is_read()); + assert_eq!(command.route().shard(), &Shard::All); +} + +#[test] +fn test_subquery_in_from() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new( + "SELECT * FROM (SELECT id, email FROM sharded WHERE id = 1) AS sub", + ) + .into()]); + + assert!(command.route().is_read()); +} + +#[test] +fn test_correlated_subquery() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new( + "SELECT * FROM sharded s1 WHERE EXISTS (SELECT 1 FROM sharded s2 WHERE s2.id = s1.id)", + ) + .into()]); + + assert!(command.route().is_read()); + assert_eq!(command.route().shard(), &Shard::All); +} + +#[test] +fn test_scalar_subquery() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new( + "SELECT *, (SELECT COUNT(*) FROM other_table) AS total FROM sharded WHERE id = 1", + ) + .into()]); + + assert!(command.route().is_read()); +} + +#[test] +fn test_subquery_with_sharding_key() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![ + Parse::named( + "__test_sub", + "SELECT * FROM sharded WHERE id = (SELECT id FROM other WHERE pk = $1)", + ) + .into(), + Bind::new_params("__test_sub", &[Parameter::new(b"5")]).into(), + Execute::new().into(), + Sync.into(), + ]); + + // Can't route to specific shard because we don't know the subquery result + assert_eq!(command.route().shard(), &Shard::All); +} + +#[test] +fn test_nested_subqueries() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new( + "SELECT * FROM sharded WHERE id IN (SELECT id FROM other WHERE status IN (SELECT status FROM statuses))", + ) + .into()]); + + assert!(command.route().is_read()); + assert_eq!(command.route().shard(), &Shard::All); +} + +#[test] +fn test_cte_with_insert() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new( + "WITH new_rows AS (INSERT INTO sharded (id, email) VALUES (1, 'test') RETURNING *) SELECT * FROM new_rows", + ) + .into()]); + + assert!(command.route().is_write()); +} + +#[test] +fn test_cte_with_delete() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new( + "WITH deleted AS (DELETE FROM sharded WHERE id = 1 RETURNING *) SELECT * FROM deleted", + ) + .into()]); + + assert!(command.route().is_write()); +} + +#[test] +fn test_cte_with_update() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new( + "WITH updated AS (UPDATE sharded SET email = 'new' WHERE id = 1 RETURNING *) SELECT * FROM updated", + ) + .into()]); + + assert!(command.route().is_write()); +} + +#[test] +fn test_recursive_cte() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new( + "WITH RECURSIVE nums AS (SELECT 1 AS n UNION ALL SELECT n + 1 FROM nums WHERE n < 10) SELECT * FROM nums", + ) + .into()]); + + assert!(command.route().is_read()); +} diff --git a/pgdog/src/frontend/router/parser/query/test/test_transaction.rs b/pgdog/src/frontend/router/parser/query/test/test_transaction.rs new file mode 100644 index 000000000..0551260fb --- /dev/null +++ b/pgdog/src/frontend/router/parser/query/test/test_transaction.rs @@ -0,0 +1,100 @@ +use crate::config::ReadWriteStrategy; +use crate::frontend::Command; +use crate::net::parameter::ParameterValue; + +use super::setup::*; + +#[test] +fn test_begin_simple() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new("BEGIN").into()]); + + match command { + Command::StartTransaction { query: q, .. } => assert_eq!(q.query(), "BEGIN"), + _ => panic!("expected StartTransaction, got {command:?}"), + } +} + +#[test] +fn test_begin_extended() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![ + Parse::new_anonymous("BEGIN").into(), + Bind::new_statement("").into(), + Execute::new().into(), + Sync.into(), + ]); + + match command { + Command::StartTransaction { extended, .. } => assert!(extended), + _ => panic!("expected StartTransaction, got {command:?}"), + } +} + +#[test] +fn test_begin_sets_write_override() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new("BEGIN").into()]); + assert!(matches!(command, Command::StartTransaction { .. })); + assert!(test.parser.write_override); +} + +// NOTE: The test for "SELECT after BEGIN is treated as write" requires passing +// in_transaction=true to the router context for the subsequent query, plus reusing +// the same QueryParser instance. This is better tested by the original +// test_transaction test in mod.rs using the query_parser! macro. + +#[test] +fn test_begin_with_aggressive_strategy() { + let mut test = QueryParserTest::new().with_read_write_strategy(ReadWriteStrategy::Aggressive); + + let command = test.execute(vec![Query::new("BEGIN").into()]); + + assert!(matches!(command, Command::StartTransaction { .. })); +} + +#[test] +fn test_set_in_transaction() { + let mut test = QueryParserTest::new().in_transaction(true); + + let command = test.execute(vec![Query::new("SET statement_timeout TO 3000").into()]); + + match command { + Command::Set { name, value, .. } => { + assert_eq!(name, "statement_timeout"); + assert_eq!(value, ParameterValue::from("3000")); + } + _ => panic!("expected Set, got {command:?}"), + } +} + +#[test] +fn test_set_application_name_in_transaction() { + let mut test = QueryParserTest::new() + .in_transaction(true) + .with_read_write_strategy(ReadWriteStrategy::Aggressive); + + let command = test.execute(vec![Query::new("SET application_name TO 'test'").into()]); + + match command.clone() { + Command::Set { + name, + value, + local, + route, + } => { + assert_eq!(name, "application_name"); + assert_eq!(value.as_str().unwrap(), "test"); + assert!(!local); + assert!(!test.cluster().read_only()); + assert!(route.is_write()); + assert!(route.shard().is_all()); + } + _ => panic!("expected Set, got {command:?}"), + } + + assert!(command.route().shard().is_all()); +} diff --git a/pgdog/src/frontend/router/parser/query/transaction.rs b/pgdog/src/frontend/router/parser/query/transaction.rs index c2e855a31..ef9e6379a 100644 --- a/pgdog/src/frontend/router/parser/query/transaction.rs +++ b/pgdog/src/frontend/router/parser/query/transaction.rs @@ -13,9 +13,10 @@ impl QueryParser { pub(super) fn transaction( &mut self, stmt: &TransactionStmt, - context: &QueryParserContext, + context: &mut QueryParserContext, ) -> Result { let extended = !context.query()?.simple(); + let mut rollback_savepoint = false; if context.rw_conservative() && !context.read_only { self.write_override = true; @@ -29,7 +30,6 @@ impl QueryParser { return Ok(Command::RollbackTransaction { extended }) } TransactionStmtKind::TransStmtBegin | TransactionStmtKind::TransStmtStart => { - self.in_transaction = true; let transaction_type = Self::transaction_type(&stmt.options).unwrap_or_default(); return Ok(Command::StartTransaction { query: context.query()?.clone(), @@ -37,6 +37,7 @@ impl QueryParser { extended, }); } + TransactionStmtKind::TransStmtRollbackTo => rollback_savepoint = true, TransactionStmtKind::TransStmtPrepare | TransactionStmtKind::TransStmtCommitPrepared | TransactionStmtKind::TransStmtRollbackPrepared => { @@ -47,7 +48,14 @@ impl QueryParser { _ => (), } - Ok(Command::Query(Route::write(None))) + context + .shards_calculator + .push(ShardWithPriority::new_table(Shard::All)); + + Ok(Command::Query( + Route::write(context.shards_calculator.shard()) + .with_rollback_savepoint(rollback_savepoint), + )) } #[inline] diff --git a/pgdog/src/frontend/router/parser/query/update.rs b/pgdog/src/frontend/router/parser/query/update.rs index 02f95c9f3..eb2a9dab7 100644 --- a/pgdog/src/frontend/router/parser/query/update.rs +++ b/pgdog/src/frontend/router/parser/query/update.rs @@ -1,28 +1,123 @@ -use crate::frontend::router::parser::where_clause::TablesSource; - use super::*; impl QueryParser { pub(super) fn update( + &mut self, stmt: &UpdateStmt, - context: &QueryParserContext, + context: &mut QueryParserContext, ) -> Result { - let table = stmt.relation.as_ref().map(Table::from); - - if let Some(table) = table { - let source = TablesSource::from(table); - let where_clause = WhereClause::new(&source, &stmt.where_clause); - - if let Some(where_clause) = where_clause { - let shards = Self::where_clause( - &context.sharding_schema, - &where_clause, - context.router_context.bind, - )?; - return Ok(Command::Query(Route::write(Self::converge(shards)))); + let mut parser = StatementParser::from_update( + stmt, + context.router_context.bind, + &context.sharding_schema, + self.recorder_mut(), + ); + let shard = parser.shard()?; + if let Some(shard) = shard { + if let Some(recorder) = self.recorder_mut() { + recorder.record_entry( + Some(shard.clone()), + "UPDATE matched WHERE clause for sharding key", + ); + } + context + .shards_calculator + .push(ShardWithPriority::new_table(shard)); + } else if let Some(recorder) = self.recorder_mut() { + recorder.record_entry(None, "UPDATE fell back to broadcast"); + } + + Ok(Command::Query(Route::write( + context.shards_calculator.shard(), + ))) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn update_preserves_decimal_values() { + let parsed = pgdog_plugin::pg_query::parse( + "UPDATE transactions SET amount = 50.00, status = 'completed' WHERE id = 1", + ) + .expect("parse"); + + let stmt = parsed + .protobuf + .stmts + .first() + .and_then(|node| node.stmt.as_ref()) + .and_then(|node| node.node.as_ref()) + .expect("statement node"); + + let update = match stmt { + NodeEnum::UpdateStmt(update) => update, + _ => panic!("expected update stmt"), + }; + + // Check that we can extract assignment values including decimals + let mut found_decimal = false; + let mut found_string = false; + + for target in &update.target_list { + if let Some(NodeEnum::ResTarget(res)) = &target.node { + if let Some(val) = &res.val { + let value = Value::try_from(&val.node).unwrap(); + match value { + Value::Float(f) => { + assert_eq!(f, 50.0); + found_decimal = true; + } + Value::String(s) => { + assert_eq!(s, "completed"); + found_string = true; + } + _ => {} + } + } } } + assert!(found_decimal, "Should have found decimal value"); + assert!(found_string, "Should have found string value"); + } + + #[test] + fn update_with_quoted_decimal() { + let parsed = + pgdog_plugin::pg_query::parse("UPDATE transactions SET amount = '50.00' WHERE id = 1") + .expect("parse"); - Ok(Command::Query(Route::write(Shard::All))) + let stmt = parsed + .protobuf + .stmts + .first() + .and_then(|node| node.stmt.as_ref()) + .and_then(|node| node.node.as_ref()) + .expect("statement node"); + + let update = match stmt { + NodeEnum::UpdateStmt(update) => update, + _ => panic!("expected update stmt"), + }; + + // Quoted decimals should be treated as strings + let mut found_string = false; + for target in &update.target_list { + if let Some(NodeEnum::ResTarget(res)) = &target.node { + if let Some(val) = &res.val { + let value = Value::try_from(&val.node).unwrap(); + match value { + Value::String(s) => { + assert_eq!(s, "50.00"); + found_string = true; + } + _ => {} + } + } + } + } + assert!(found_string, "Should have found string value"); } } diff --git a/pgdog/src/frontend/router/parser/rewrite/mod.rs b/pgdog/src/frontend/router/parser/rewrite/mod.rs index c7574299a..17688a336 100644 --- a/pgdog/src/frontend/router/parser/rewrite/mod.rs +++ b/pgdog/src/frontend/router/parser/rewrite/mod.rs @@ -1,102 +1,5 @@ -use pg_query::{NodeEnum, ParseResult}; +mod shard_key; +pub mod statement; +pub use statement::{StatementRewrite, StatementRewriteContext}; -use super::{Command, Error}; -use crate::frontend::PreparedStatements; -use crate::net::Parse; - -#[derive(Debug, Clone)] -pub struct Rewrite<'a> { - ast: &'a ParseResult, -} - -impl<'a> Rewrite<'a> { - pub fn new(ast: &'a ParseResult) -> Self { - Self { ast } - } - - /// Statement needs to be rewritten. - pub fn needs_rewrite(&self) -> bool { - for stmt in &self.ast.protobuf.stmts { - if let Some(ref stmt) = stmt.stmt { - if let Some(ref node) = stmt.node { - match node { - NodeEnum::PrepareStmt(_) => return true, - NodeEnum::ExecuteStmt(_) => return true, - NodeEnum::DeallocateStmt(_) => return true, - _ => (), - } - } - } - } - - false - } - - pub fn rewrite(&self, prepared_statements: &mut PreparedStatements) -> Result { - let mut ast = self.ast.protobuf.clone(); - - for stmt in &mut ast.stmts { - if let Some(ref mut stmt) = stmt.stmt { - if let Some(ref mut node) = stmt.node { - match node { - NodeEnum::PrepareStmt(ref mut stmt) => { - let statement = stmt.query.as_ref().ok_or(Error::EmptyQuery)?; - let statement = statement.deparse().map_err(|_| Error::EmptyQuery)?; - let mut parse = Parse::named(&stmt.name, &statement); - prepared_statements.insert_anyway(&mut parse); - stmt.name = parse.name().to_string(); - } - - NodeEnum::ExecuteStmt(ref mut stmt) => { - let name = prepared_statements.name(&stmt.name); - if let Some(name) = name { - stmt.name = name.to_string(); - } - } - - NodeEnum::DeallocateStmt(_) => return Ok(Command::Deallocate), - - _ => (), - } - } - } - } - - Ok(Command::Rewrite( - ast.deparse().map_err(|_| Error::EmptyQuery)?, - )) - } -} - -#[cfg(test)] -mod test { - use std::sync::Arc; - - use super::*; - - #[test] - fn test_rewrite_prepared() { - let ast = pg_query::parse("BEGIN; PREPARE test AS SELECT $1, $2, $3; PREPARE test2 AS SELECT * FROM my_table WHERE id = $1; COMMIT;").unwrap(); - let rewrite = Rewrite::new(&ast); - assert!(rewrite.needs_rewrite()); - let mut prepared_statements = PreparedStatements::new(); - let queries = rewrite.rewrite(&mut prepared_statements).unwrap(); - match queries { - Command::Rewrite(queries) => assert_eq!(queries, "BEGIN; PREPARE __pgdog_1 AS SELECT $1, $2, $3; PREPARE __pgdog_2 AS SELECT * FROM my_table WHERE id = $1; COMMIT"), - _ => panic!("not a rewrite"), - } - } - - #[test] - fn test_deallocate() { - for q in ["DEALLOCATE ALL", "DEALLOCATE test"] { - let ast = pg_query::parse(q).unwrap(); - let ast = Arc::new(ast); - let rewrite = Rewrite::new(&ast) - .rewrite(&mut PreparedStatements::new()) - .unwrap(); - - assert!(matches!(rewrite, Command::Deallocate)); - } - } -} +pub use shard_key::{Assignment, AssignmentValue, ShardKeyRewritePlan}; diff --git a/pgdog/src/frontend/router/parser/rewrite/shard_key.rs b/pgdog/src/frontend/router/parser/rewrite/shard_key.rs new file mode 100644 index 000000000..75ca77c98 --- /dev/null +++ b/pgdog/src/frontend/router/parser/rewrite/shard_key.rs @@ -0,0 +1,97 @@ +use pgdog_plugin::pg_query::protobuf::UpdateStmt; + +use crate::frontend::router::Route; + +use super::super::table::OwnedTable; + +#[derive(Debug, Clone, PartialEq)] +pub enum AssignmentValue { + Parameter(i32), + Integer(i64), + Float(String), + String(String), + Boolean(bool), + Null, + Column(String), +} + +impl AssignmentValue { + pub fn as_parameter(&self) -> Option { + if let Self::Parameter(param) = self { + Some(*param) + } else { + None + } + } +} + +#[derive(Debug, Clone, PartialEq)] +pub struct Assignment { + column: String, + value: AssignmentValue, +} + +impl Assignment { + pub fn new(column: String, value: AssignmentValue) -> Self { + Self { column, value } + } + + pub fn column(&self) -> &str { + &self.column + } + + pub fn value(&self) -> &AssignmentValue { + &self.value + } +} + +#[derive(Debug, Clone)] +pub struct ShardKeyRewritePlan { + table: OwnedTable, + route: Route, + new_shard: Option, + statement: UpdateStmt, + assignments: Vec, +} + +impl ShardKeyRewritePlan { + pub fn new( + table: OwnedTable, + route: Route, + new_shard: Option, + statement: UpdateStmt, + assignments: Vec, + ) -> Self { + Self { + table, + route, + new_shard, + statement, + assignments, + } + } + + pub fn table(&self) -> &OwnedTable { + &self.table + } + + pub fn route(&self) -> &Route { + &self.route + } + + pub fn new_shard(&self) -> Option { + self.new_shard + } + + pub fn statement(&self) -> &UpdateStmt { + &self.statement + } + + pub fn assignments(&self) -> &[Assignment] { + &self.assignments + } + + pub fn set_new_shard(&mut self, shard: usize) { + self.new_shard = Some(shard); + } +} diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/aggregate/engine.rs b/pgdog/src/frontend/router/parser/rewrite/statement/aggregate/engine.rs new file mode 100644 index 000000000..3b6180df3 --- /dev/null +++ b/pgdog/src/frontend/router/parser/rewrite/statement/aggregate/engine.rs @@ -0,0 +1,445 @@ +use crate::frontend::router::parser::aggregate::{Aggregate, AggregateFunction}; +use pg_query::protobuf::{ + a_const::Val, AConst, FuncCall, Integer, Node, ParseResult, ResTarget, String as PgString, +}; +use pg_query::NodeEnum; + +use super::{AggregateRewritePlan, HelperKind, HelperMapping, RewriteOutput}; + +/// Query rewrite engine. Currently supports injecting helper aggregates for AVG and +/// variance-related functions that require additional helper aggregates when run +/// across shards. +#[derive(Default)] +pub struct AggregatesRewrite; + +impl AggregatesRewrite { + /// Rewrite a SELECT query in-place, adding helper aggregates when necessary. + pub fn rewrite_select(&self, ast: &mut ParseResult, aggregate: &Aggregate) -> RewriteOutput { + self.rewrite_parsed(ast, aggregate) + } + + fn rewrite_parsed(&self, parsed: &mut ParseResult, aggregate: &Aggregate) -> RewriteOutput { + let Some(raw_stmt) = parsed.stmts.first() else { + return RewriteOutput::default(); + }; + + let Some(stmt) = raw_stmt.stmt.as_ref() else { + return RewriteOutput::default(); + }; + + let Some(NodeEnum::SelectStmt(select)) = stmt.node.as_ref() else { + return RewriteOutput::default(); + }; + + let mut plan = AggregateRewritePlan::new(); + let mut helper_nodes: Vec = Vec::new(); + let mut planned_aliases: Vec = Vec::new(); + let base_len = select.target_list.len(); + + for target in aggregate.targets() { + if aggregate.targets().iter().any(|other| { + matches!(other.function(), AggregateFunction::Count) + && other.expr_id() == target.expr_id() + && other.is_distinct() == target.is_distinct() + }) && matches!(target.function(), AggregateFunction::Avg) + { + continue; + } + + let Some(node) = select.target_list.get(target.column()) else { + continue; + }; + + let Some((location, helper_specs)) = ({ + if let Some(NodeEnum::ResTarget(res_target)) = node.node.as_ref() { + if let Some(original_value) = res_target.val.as_ref() { + if let Some(func_call) = Self::extract_func_call(original_value) { + let specs = Self::helper_specs( + func_call, + target.function(), + target.is_distinct(), + ); + Some((res_target.location, specs)) + } else { + None + } + } else { + None + } + } else { + None + } + }) else { + continue; + }; + + if helper_specs.is_empty() { + continue; + } + + for helper in helper_specs { + let HelperSpec { func, kind } = helper; + + let helper_alias = format!( + "__pgdog_{}_expr{}_col{}", + kind.alias_suffix(), + target.expr_id(), + target.column() + ); + + let alias_exists = select.target_list.iter().any(|existing| { + if let Some(NodeEnum::ResTarget(res)) = existing.node.as_ref() { + res.name == helper_alias + } else { + false + } + }) || planned_aliases + .iter() + .any(|existing| existing == &helper_alias); + + if alias_exists { + continue; + } + + let helper_column = base_len + helper_nodes.len(); + + let helper_res = ResTarget { + name: helper_alias.clone(), + indirection: vec![], + val: Some(Box::new(Node { + node: Some(NodeEnum::FuncCall(Box::new(func))), + })), + location, + }; + + helper_nodes.push(Node { + node: Some(NodeEnum::ResTarget(Box::new(helper_res))), + }); + planned_aliases.push(helper_alias.clone()); + + plan.add_drop_column(helper_column); + plan.add_helper(HelperMapping { + target_column: target.column(), + helper_column, + expr_id: target.expr_id(), + distinct: target.is_distinct(), + kind, + alias: helper_alias, + }); + } + } + + if helper_nodes.is_empty() { + return RewriteOutput::default(); + } + + let Some(raw_stmt) = parsed.stmts.first_mut() else { + return RewriteOutput::default(); + }; + + let Some(stmt) = raw_stmt.stmt.as_mut() else { + return RewriteOutput::default(); + }; + + let Some(NodeEnum::SelectStmt(select)) = stmt.node.as_mut() else { + return RewriteOutput::default(); + }; + + select.target_list.extend(helper_nodes); + + RewriteOutput::new(plan) + } + + fn extract_func_call(node: &Node) -> Option<&FuncCall> { + match node.node.as_ref()? { + NodeEnum::FuncCall(func) => Some(func), + NodeEnum::TypeCast(cast) => cast + .arg + .as_deref() + .and_then(|inner| Self::extract_func_call(inner)), + NodeEnum::CollateClause(collate) => collate + .arg + .as_deref() + .and_then(|inner| Self::extract_func_call(inner)), + NodeEnum::CoerceToDomain(coerce) => coerce + .arg + .as_deref() + .and_then(|inner| Self::extract_func_call(inner)), + NodeEnum::ResTarget(res) => res + .val + .as_deref() + .and_then(|inner| Self::extract_func_call(inner)), + _ => None, + } + } + + fn build_count_func(original: &FuncCall, distinct: bool) -> FuncCall { + FuncCall { + funcname: vec![Node { + node: Some(NodeEnum::String(PgString { + sval: "count".into(), + })), + }], + args: original.args.clone(), + agg_order: original.agg_order.clone(), + agg_filter: original.agg_filter.clone(), + over: original.over.clone(), + agg_within_group: original.agg_within_group, + agg_star: original.agg_star, + agg_distinct: distinct, + func_variadic: original.func_variadic, + funcformat: original.funcformat, + location: original.location, + } + } + + fn build_sum_func(original: &FuncCall, distinct: bool) -> FuncCall { + FuncCall { + funcname: vec![Node { + node: Some(NodeEnum::String(PgString { sval: "sum".into() })), + }], + args: original.args.clone(), + agg_order: original.agg_order.clone(), + agg_filter: original.agg_filter.clone(), + over: original.over.clone(), + agg_within_group: original.agg_within_group, + agg_star: original.agg_star, + agg_distinct: distinct, + func_variadic: original.func_variadic, + funcformat: original.funcformat, + location: original.location, + } + } + + fn build_sum_of_squares_func(original: &FuncCall, distinct: bool) -> Option { + let arg = original.args.first()?.clone(); + + let two = Node { + node: Some(NodeEnum::AConst(AConst { + val: Some(Val::Ival(Integer { ival: 2 })), + location: original.location, + isnull: false, + })), + }; + + let power = FuncCall { + funcname: vec![Node { + node: Some(NodeEnum::String(PgString { + sval: "power".into(), + })), + }], + args: vec![arg, two], + agg_order: vec![], + agg_filter: None, + over: None, + agg_within_group: false, + agg_star: false, + agg_distinct: false, + func_variadic: false, + funcformat: original.funcformat, + location: original.location, + }; + + Some(FuncCall { + funcname: vec![Node { + node: Some(NodeEnum::String(PgString { sval: "sum".into() })), + }], + args: vec![Node { + node: Some(NodeEnum::FuncCall(Box::new(power))), + }], + agg_order: original.agg_order.clone(), + agg_filter: original.agg_filter.clone(), + over: original.over.clone(), + agg_within_group: original.agg_within_group, + agg_star: false, + agg_distinct: distinct, + func_variadic: original.func_variadic, + funcformat: original.funcformat, + location: original.location, + }) + } + + fn helper_specs( + func_call: &FuncCall, + function: &AggregateFunction, + distinct: bool, + ) -> Vec { + match function { + AggregateFunction::Avg => vec![HelperSpec { + func: Self::build_count_func(func_call, distinct), + kind: HelperKind::Count, + }], + AggregateFunction::StddevSamp + | AggregateFunction::StddevPop + | AggregateFunction::VarSamp + | AggregateFunction::VarPop => { + let mut helpers = vec![ + HelperSpec { + func: Self::build_count_func(func_call, distinct), + kind: HelperKind::Count, + }, + HelperSpec { + func: Self::build_sum_func(func_call, distinct), + kind: HelperKind::Sum, + }, + ]; + + if let Some(sum_sq) = Self::build_sum_of_squares_func(func_call, distinct) { + helpers.push(HelperSpec { + func: sum_sq, + kind: HelperKind::SumSquares, + }); + } + + helpers + } + _ => vec![], + } + } +} + +#[derive(Debug, Clone)] +struct HelperSpec { + func: FuncCall, + kind: HelperKind, +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::frontend::router::parser::aggregate::Aggregate; + + fn select(ast: &ParseResult) -> &pg_query::protobuf::SelectStmt { + match ast + .stmts + .first() + .and_then(|stmt| stmt.stmt.as_ref()) + .and_then(|stmt| stmt.node.as_ref()) + { + Some(NodeEnum::SelectStmt(select)) => select, + _ => panic!("not a select"), + } + } + + fn rewrite(sql: &str) -> (ParseResult, RewriteOutput) { + let mut parsed = pg_query::parse(sql).unwrap().protobuf; + let aggregate = { + let stmt = select(&parsed); + Aggregate::parse(stmt) + }; + let output = AggregatesRewrite::default().rewrite_select(&mut parsed, &aggregate); + (parsed, output) + } + + #[test] + fn rewrite_engine_noop() { + let (ast, output) = rewrite("SELECT COUNT(price) FROM menu"); + assert!(output.plan.is_noop()); + assert_eq!(select(&ast).target_list.len(), 1); + } + + #[test] + fn rewrite_engine_adds_helper() { + let (ast, output) = rewrite("SELECT AVG(price) FROM menu"); + assert!(!output.plan.is_noop()); + assert_eq!(output.plan.drop_columns(), &[1]); + assert_eq!(output.plan.helpers().len(), 1); + let helper = &output.plan.helpers()[0]; + assert_eq!(helper.target_column, 0); + assert_eq!(helper.helper_column, 1); + assert_eq!(helper.expr_id, 0); + assert!(!helper.distinct); + assert!(matches!(helper.kind, HelperKind::Count)); + + let aggregate = Aggregate::parse(select(&ast)); + assert_eq!(aggregate.targets().len(), 2); + assert!(aggregate + .targets() + .iter() + .any(|target| matches!(target.function(), AggregateFunction::Count))); + } + + #[test] + fn rewrite_engine_skips_when_count_exists() { + let sql = "SELECT COUNT(price), AVG(price) FROM menu"; + let (ast, output) = rewrite(sql); + assert!(output.plan.is_noop()); + assert_eq!(select(&ast).target_list.len(), 2); + } + + #[test] + fn rewrite_engine_handles_mismatched_pair() { + let (ast, output) = rewrite("SELECT COUNT(price::numeric), AVG(price) FROM menu"); + assert_eq!(output.plan.drop_columns(), &[2]); + assert_eq!(output.plan.helpers().len(), 1); + let helper = &output.plan.helpers()[0]; + assert_eq!(helper.target_column, 1); + assert_eq!(helper.helper_column, 2); + assert!(!helper.distinct); + assert!(matches!(helper.kind, HelperKind::Count)); + + let aggregate = Aggregate::parse(select(&ast)); + assert_eq!(aggregate.targets().len(), 3); + assert!( + aggregate + .targets() + .iter() + .filter(|target| matches!(target.function(), AggregateFunction::Count)) + .count() + >= 2 + ); + } + + #[test] + fn rewrite_engine_multiple_avg_helpers() { + let (ast, output) = rewrite("SELECT AVG(price), AVG(discount) FROM menu"); + assert_eq!(output.plan.drop_columns(), &[2, 3]); + assert_eq!(output.plan.helpers().len(), 2); + + let helper_price = &output.plan.helpers()[0]; + assert_eq!(helper_price.target_column, 0); + assert_eq!(helper_price.helper_column, 2); + assert!(matches!(helper_price.kind, HelperKind::Count)); + + let helper_discount = &output.plan.helpers()[1]; + assert_eq!(helper_discount.target_column, 1); + assert_eq!(helper_discount.helper_column, 3); + assert!(matches!(helper_discount.kind, HelperKind::Count)); + + let aggregate = Aggregate::parse(select(&ast)); + assert_eq!(aggregate.targets().len(), 4); + assert_eq!( + aggregate + .targets() + .iter() + .filter(|target| matches!(target.function(), AggregateFunction::Count)) + .count(), + 2 + ); + } + + #[test] + fn rewrite_engine_stddev_helpers() { + let (ast, output) = rewrite("SELECT STDDEV(price) FROM menu"); + assert!(!output.plan.is_noop()); + assert_eq!(output.plan.drop_columns(), &[1, 2, 3]); + assert_eq!(output.plan.helpers().len(), 3); + + let kinds: Vec = output + .plan + .helpers() + .iter() + .map(|helper| { + assert_eq!(helper.target_column, 0); + helper.kind + }) + .collect(); + + assert!(kinds.contains(&HelperKind::Count)); + assert!(kinds.contains(&HelperKind::Sum)); + assert!(kinds.contains(&HelperKind::SumSquares)); + + // Expect original STDDEV plus three helpers. + assert_eq!(select(&ast).target_list.len(), 4); + } +} diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/aggregate/mod.rs b/pgdog/src/frontend/router/parser/rewrite/statement/aggregate/mod.rs new file mode 100644 index 000000000..0676031f0 --- /dev/null +++ b/pgdog/src/frontend/router/parser/rewrite/statement/aggregate/mod.rs @@ -0,0 +1,44 @@ +pub mod engine; +pub mod plan; + +use super::{Error, RewritePlan, StatementRewrite}; +use crate::frontend::router::parser::aggregate::Aggregate; +use pg_query::NodeEnum; + +pub use engine::AggregatesRewrite; +pub use plan::{AggregateRewritePlan, HelperKind, HelperMapping, RewriteOutput}; + +impl StatementRewrite<'_> { + /// Add missing COUNT(*) and other helps when using aggregates. + pub(super) fn rewrite_aggregates(&mut self, plan: &mut RewritePlan) -> Result<(), Error> { + if self.schema.shards == 1 { + return Ok(()); + } + + let Some(raw_stmt) = self.stmt.stmts.first() else { + return Ok(()); + }; + + let Some(stmt) = raw_stmt.stmt.as_ref() else { + return Ok(()); + }; + + let Some(NodeEnum::SelectStmt(select)) = stmt.node.as_ref() else { + return Ok(()); + }; + + let aggregate = Aggregate::parse(select); + if aggregate.is_empty() { + return Ok(()); + } + + let output = AggregatesRewrite.rewrite_select(self.stmt, &aggregate); + if output.plan.is_noop() { + return Ok(()); + } + + plan.aggregates = output.plan; + self.rewritten = true; + Ok(()) + } +} diff --git a/pgdog/src/frontend/router/parser/rewrite_plan.rs b/pgdog/src/frontend/router/parser/rewrite/statement/aggregate/plan.rs similarity index 60% rename from pgdog/src/frontend/router/parser/rewrite_plan.rs rename to pgdog/src/frontend/router/parser/rewrite/statement/aggregate/plan.rs index 7dac29cf1..78b349b6b 100644 --- a/pgdog/src/frontend/router/parser/rewrite_plan.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/aggregate/plan.rs @@ -1,20 +1,46 @@ +/// Type of aggregate function added to the result set. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum HelperKind { + /// COUNT(*) or COUNT(name) + Count, + /// SUM(column) + Sum, + /// SUM(POWER(column, 2)) + SumSquares, +} + +impl HelperKind { + /// Suffix for the aggregate function. + pub fn alias_suffix(&self) -> &'static str { + match self { + HelperKind::Count => "count", + HelperKind::Sum => "sum", + HelperKind::SumSquares => "sumsq", + } + } +} + +/// Context on the aggregate function column added to the result set. #[derive(Debug, Clone, PartialEq)] pub struct HelperMapping { - pub avg_column: usize, + pub target_column: usize, pub helper_column: usize, pub expr_id: usize, pub distinct: bool, + pub kind: HelperKind, + pub alias: String, } /// Plan describing how the proxy rewrites a query and its results. #[derive(Debug, Clone, Default, PartialEq)] -pub struct RewritePlan { +pub struct AggregateRewritePlan { /// Column indexes (0-based) to drop from the row description/results after execution. drop_columns: Vec, helpers: Vec, } -impl RewritePlan { +impl AggregateRewritePlan { + /// Create new no-op aggregate rewrite plan. pub fn new() -> Self { Self { drop_columns: Vec::new(), @@ -22,6 +48,7 @@ impl RewritePlan { } } + /// Is this plan a no-op? Doesn't do anything. pub fn is_noop(&self) -> bool { self.drop_columns.is_empty() && self.helpers.is_empty() } @@ -47,27 +74,22 @@ impl RewritePlan { #[derive(Debug, Default, Clone)] pub struct RewriteOutput { - pub sql: String, - pub plan: RewritePlan, + pub plan: AggregateRewritePlan, } impl RewriteOutput { - pub fn new(sql: String, plan: RewritePlan) -> Self { - Self { sql, plan } + pub fn new(plan: AggregateRewritePlan) -> Self { + Self { plan } } } -pub trait QueryRewriter { - fn rewrite(&self, sql: &str, route: &crate::frontend::router::Route) -> RewriteOutput; -} - #[cfg(test)] mod tests { use super::*; #[test] fn rewrite_plan_noop() { - let plan = RewritePlan::new(); + let plan = AggregateRewritePlan::new(); assert!(plan.is_noop()); assert!(plan.drop_columns().is_empty()); assert!(plan.helpers().is_empty()); @@ -75,7 +97,7 @@ mod tests { #[test] fn rewrite_plan_drop_columns() { - let mut plan = RewritePlan::new(); + let mut plan = AggregateRewritePlan::new(); plan.add_drop_column(1); plan.add_drop_column(4); assert_eq!(plan.drop_columns(), &[1, 4]); @@ -83,25 +105,28 @@ mod tests { #[test] fn rewrite_plan_helpers() { - let mut plan = RewritePlan::new(); + let mut plan = AggregateRewritePlan::new(); plan.add_helper(HelperMapping { - avg_column: 0, + target_column: 0, helper_column: 1, expr_id: 7, distinct: false, + kind: HelperKind::Count, + alias: "__pgdog_count_expr7_col0".into(), }); assert_eq!(plan.helpers().len(), 1); let helper = &plan.helpers()[0]; - assert_eq!(helper.avg_column, 0); + assert_eq!(helper.target_column, 0); assert_eq!(helper.helper_column, 1); assert_eq!(helper.expr_id, 7); assert!(!helper.distinct); + assert!(matches!(helper.kind, HelperKind::Count)); + assert_eq!(helper.alias, "__pgdog_count_expr7_col0"); } #[test] fn rewrite_output_defaults() { let output = RewriteOutput::default(); assert!(output.plan.is_noop()); - assert!(output.sql.is_empty()); } } diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/error.rs b/pgdog/src/frontend/router/parser/rewrite/statement/error.rs new file mode 100644 index 000000000..0125957bb --- /dev/null +++ b/pgdog/src/frontend/router/parser/rewrite/statement/error.rs @@ -0,0 +1,31 @@ +use thiserror::Error; + +#[derive(Debug, Error)] +pub enum Error { + #[error("unique_id generation failed: {0}")] + UniqueId(#[from] crate::unique_id::Error), + + #[error("pg_query: {0}")] + PgQuery(#[from] pg_query::Error), + + #[error("cache: {0}")] + Cache(String), + + #[error("sharding key assignment unsupported: {0}")] + UnsupportedShardingKeyUpdate(String), + + #[error("net: {0}")] + Net(#[from] crate::net::Error), + + #[error("missing parameter: ${0}")] + MissingParameter(u16), + + #[error("empty query")] + EmptyQuery, + + #[error("missing column: ${0}")] + MissingColumn(usize), + + #[error("WHERE clause is required")] + WhereClauseMissing, +} diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/insert.rs b/pgdog/src/frontend/router/parser/rewrite/statement/insert.rs new file mode 100644 index 000000000..d1c4077c7 --- /dev/null +++ b/pgdog/src/frontend/router/parser/rewrite/statement/insert.rs @@ -0,0 +1,527 @@ +use pg_query::{Node, NodeEnum}; +use pgdog_config::{QueryParserEngine, RewriteMode}; + +use crate::frontend::router::parser::Cache; +use crate::frontend::router::Ast; +use crate::frontend::{BufferedQuery, ClientRequest}; +use crate::net::messages::bind::{Format, Parameter}; +use crate::net::{Bind, Parse, ProtocolMessage, Query}; + +use super::{Error, RewritePlan, StatementRewrite}; + +#[derive(Debug, Clone)] +pub struct InsertSplit { + /// Parameter positions in the original Bind message + /// that should be used to build the Bind message specific to this + /// insert statement. + params: Vec, + + /// The split up INSERT statement with parameters and/or values. + stmt: String, + + /// The statement AST. + ast: Ast, + + /// The global prepared statement name for this split. + /// Only set when the original statement was a named prepared statement. + statement_name: Option, +} + +impl InsertSplit { + /// Get the parameter positions from the original Bind message. + pub fn params(&self) -> &[u16] { + &self.params + } + + /// Get the SQL statement. + pub fn stmt(&self) -> &str { + &self.stmt + } + + /// Get the AST. + pub fn ast(&self) -> &Ast { + &self.ast + } + + /// Get the global prepared statement name, if this split was registered. + pub fn statement_name(&self) -> Option<&str> { + self.statement_name.as_deref() + } + + /// Build a ClientRequest from this split and the original request. + pub fn build_request(&self, request: &ClientRequest) -> ClientRequest { + let mut new_request = ClientRequest::default(); + + for message in &request.messages { + let new_message = match message { + ProtocolMessage::Parse(parse) => { + let mut new_parse = parse.clone(); + new_parse.set_query(&self.stmt); + + if let Some(name) = self.statement_name() { + new_parse.rename_fast(name); + } + + ProtocolMessage::Parse(new_parse) + } + ProtocolMessage::Query(query) => { + let mut new_query = query.clone(); + new_query.set_query(&self.stmt); + ProtocolMessage::Query(new_query) + } + ProtocolMessage::Bind(bind) => { + let new_bind = self.extract_bind_params(bind); + ProtocolMessage::Bind(new_bind) + } + other => other.clone(), + }; + new_request.messages.push(new_message); + new_request.ast = Some(self.ast.clone()); + } + + new_request + } + + /// Extract specific parameters from a Bind message based on this split's param indices. + fn extract_bind_params(&self, bind: &Bind) -> Bind { + let params: Vec = self + .params + .iter() + .filter_map(|&idx| bind.params_raw().get(idx as usize).cloned()) + .collect(); + + let codes: Vec = if bind.format_codes_raw().len() == 1 { + // Uniform format: keep it + bind.format_codes_raw().clone() + } else if bind.format_codes_raw().len() == bind.params_raw().len() { + // One-to-one mapping: extract corresponding codes + self.params + .iter() + .filter_map(|&idx| bind.format_codes_raw().get(idx as usize).copied()) + .collect() + } else { + // No codes (all text) + Vec::new() + }; + + // Use the split's registered statement name if available, + // otherwise fall back to the original bind's statement name. + let statement_name = self + .statement_name + .as_deref() + .unwrap_or_else(|| bind.statement()); + + Bind::new_params_codes(statement_name, ¶ms, &codes) + } +} + +/// Build separate ClientRequests for each insert split. +pub fn build_split_requests(splits: &[InsertSplit], request: &ClientRequest) -> Vec { + splits + .iter() + .map(|split| split.build_request(request)) + .collect() +} + +impl StatementRewrite<'_> { + /// Split up multi-tuple INSERT statements into separate single-tuple statements + /// for individual execution. + /// + /// # Example + /// + /// ```sql + /// INSERT INTO my_table (id, value) VALUES ($1, $2), ($3, $4) + /// ``` + /// + /// becomes + /// + /// ```sql + /// INSERT INTO my_table (id, value) VALUES ($1, $2) + /// INSERT INTO my_table (id, value) VALUES ($1, $2) -- These are copied from params $3 and $4 + /// ``` + /// + pub(super) fn split_insert(&mut self, plan: &mut RewritePlan) -> Result<(), Error> { + // Don't rewrite INSERTs in unsharded databases. + if self.schema.shards == 1 || self.schema.rewrite.split_inserts != RewriteMode::Rewrite { + return Ok(()); + } + + let splits: Vec<(Vec, String)> = { + let values_lists = match self.get_insert_values_lists() { + Some(lists) if lists.len() > 1 => lists, + _ => return Ok(()), + }; + + values_lists + .iter() + .map(|values_list| self.build_single_tuple_insert(values_list)) + .collect::, _>>()? + }; + + // Now create Ast for each split (needs mutable borrow of prepared_statements) + let cache = Cache::get(); + for (params, stmt) in splits { + let query = if self.extended { + BufferedQuery::Prepared(Parse::named("", &stmt)) + } else { + BufferedQuery::Query(Query::new(&stmt)) + }; + let ast = cache + .query(&query, self.schema, self.prepared_statements) + .map_err(|e| Error::Cache(e.to_string()))?; + + // If this is a named prepared statement, register the split in the global cache + // and store the assigned name for use in Bind messages. + let statement_name = if self.prepared { + // Name will be assigned by `insert`. + let mut parse = Parse::named("", &stmt); + self.prepared_statements.insert(&mut parse); + Some(parse.name().to_owned()) + } else { + None + }; + + plan.insert_split.push(InsertSplit { + params, + stmt, + ast, + statement_name, + }); + } + + Ok(()) + } + + /// Get the values_lists from an INSERT statement, if present. + fn get_insert_values_lists(&self) -> Option<&[Node]> { + let stmt = self.stmt.stmts.first()?; + let node = stmt.stmt.as_ref()?; + + if let NodeEnum::InsertStmt(insert) = node.node.as_ref()? { + let select = insert.select_stmt.as_ref()?; + if let NodeEnum::SelectStmt(select_stmt) = select.node.as_ref()? { + if !select_stmt.values_lists.is_empty() { + return Some(&select_stmt.values_lists); + } + } + } + None + } + + /// Build a single-tuple INSERT from the original statement with just one values_list. + /// Returns the parameter positions (0-indexed) and the SQL string. + fn build_single_tuple_insert(&self, values_list: &Node) -> Result<(Vec, String), Error> { + let mut ast = self.stmt.clone(); + let mut params = Vec::new(); + + // Collect parameter references from this values_list + Self::collect_params(values_list, &mut params); + + // Renumber parameters to start from $1 + let mut new_values_list = values_list.clone(); + Self::renumber_params(&mut new_values_list, ¶ms); + + // Replace the values_lists with just this one tuple + if let Some(stmt) = ast.stmts.first_mut() { + if let Some(node) = stmt.stmt.as_mut() { + if let Some(NodeEnum::InsertStmt(insert)) = node.node.as_mut() { + if let Some(select) = insert.select_stmt.as_mut() { + if let Some(NodeEnum::SelectStmt(select_stmt)) = select.node.as_mut() { + select_stmt.values_lists = vec![new_values_list]; + } + } + } + } + } + + let stmt = match self.schema.query_parser_engine { + QueryParserEngine::PgQueryProtobuf => ast.deparse(), + QueryParserEngine::PgQueryRaw => ast.deparse_raw(), + }?; + + Ok((params, stmt)) + } + + /// Collect all parameter references from a node tree. + fn collect_params(node: &Node, params: &mut Vec) { + if let Some(node_enum) = &node.node { + match node_enum { + NodeEnum::ParamRef(param) => { + if param.number > 0 { + params.push((param.number - 1) as u16); + } + } + NodeEnum::List(list) => { + for item in &list.items { + Self::collect_params(item, params); + } + } + NodeEnum::TypeCast(cast) => { + if let Some(arg) = &cast.arg { + Self::collect_params(arg, params); + } + } + _ => {} + } + } + } + + /// Renumber parameters in a node tree based on their position in the params list. + fn renumber_params(node: &mut Node, params: &[u16]) { + if let Some(node_enum) = &mut node.node { + match node_enum { + NodeEnum::ParamRef(param) => { + if param.number > 0 { + let old_pos = (param.number - 1) as u16; + if let Some(new_pos) = params.iter().position(|&p| p == old_pos) { + param.number = (new_pos + 1) as i32; + } + } + } + NodeEnum::List(list) => { + for item in &mut list.items { + Self::renumber_params(item, params); + } + } + NodeEnum::TypeCast(cast) => { + if let Some(arg) = &mut cast.arg { + Self::renumber_params(arg, params); + } + } + _ => {} + } + } + } +} + +#[cfg(test)] +mod tests { + use pgdog_config::Rewrite; + + use super::*; + use crate::backend::ShardingSchema; + use crate::frontend::router::parser::StatementRewriteContext; + use crate::frontend::PreparedStatements; + + fn default_schema() -> ShardingSchema { + ShardingSchema { + shards: 2, + rewrite: Rewrite { + enabled: true, + split_inserts: RewriteMode::Rewrite, + ..Default::default() + }, + ..Default::default() + } + } + + fn parse_and_split(sql: &str) -> Vec { + let mut ast = pg_query::parse(sql).unwrap().protobuf; + let mut prepared = PreparedStatements::default(); + let schema = default_schema(); + let mut rewriter = StatementRewrite::new(StatementRewriteContext { + stmt: &mut ast, + extended: false, + prepared: false, + prepared_statements: &mut prepared, + schema: &schema, + }); + let mut plan = RewritePlan::default(); + rewriter.split_insert(&mut plan).unwrap(); + plan.insert_split + } + + #[test] + fn test_split_insert_with_params() { + let splits = parse_and_split("INSERT INTO my_table (id, value) VALUES ($1, $2), ($3, $4)"); + + assert_eq!(splits.len(), 2); + + // First tuple uses params 0 and 1 (original $1, $2) + assert_eq!(splits[0].params(), &[0, 1]); + assert_eq!( + splits[0].stmt(), + "INSERT INTO my_table (id, value) VALUES ($1, $2)" + ); + + // Second tuple uses params 2 and 3 (original $3, $4), renumbered to $1, $2 + assert_eq!(splits[1].params(), &[2, 3]); + assert_eq!( + splits[1].stmt(), + "INSERT INTO my_table (id, value) VALUES ($1, $2)" + ); + } + + #[test] + fn test_split_insert_single_tuple_no_split() { + let splits = parse_and_split("INSERT INTO my_table (id, value) VALUES ($1, $2)"); + + // Single tuple should not be split + assert!(splits.is_empty()); + } + + #[test] + fn test_split_insert_literal_values() { + let splits = parse_and_split("INSERT INTO my_table (id, value) VALUES (1, 'a'), (2, 'b')"); + + assert_eq!(splits.len(), 2); + + // No params for literal values + assert!(splits[0].params().is_empty()); + assert_eq!( + splits[0].stmt(), + "INSERT INTO my_table (id, value) VALUES (1, 'a')" + ); + + assert!(splits[1].params().is_empty()); + assert_eq!( + splits[1].stmt(), + "INSERT INTO my_table (id, value) VALUES (2, 'b')" + ); + } + + #[test] + fn test_split_insert_mixed_params_and_literals() { + let splits = + parse_and_split("INSERT INTO my_table (id, value) VALUES ($1, 'a'), ($2, 'b')"); + + assert_eq!(splits.len(), 2); + + assert_eq!(splits[0].params(), &[0]); + assert_eq!( + splits[0].stmt(), + "INSERT INTO my_table (id, value) VALUES ($1, 'a')" + ); + + assert_eq!(splits[1].params(), &[1]); + assert_eq!( + splits[1].stmt(), + "INSERT INTO my_table (id, value) VALUES ($1, 'b')" + ); + } + + #[test] + fn test_extract_bind_params() { + let splits = parse_and_split("INSERT INTO t (a, b) VALUES ($1, $2), ($3, $4)"); + let bind = Bind::new_params( + "test", + &[ + Parameter::new(b"p0"), + Parameter::new(b"p1"), + Parameter::new(b"p2"), + Parameter::new(b"p3"), + ], + ); + + // First split uses params 0 and 1 + let extracted = splits[0].extract_bind_params(&bind); + assert_eq!(extracted.params_raw().len(), 2); + assert_eq!(extracted.params_raw()[0].data.as_ref(), b"p0"); + assert_eq!(extracted.params_raw()[1].data.as_ref(), b"p1"); + + // Second split uses params 2 and 3 + let extracted = splits[1].extract_bind_params(&bind); + assert_eq!(extracted.params_raw().len(), 2); + assert_eq!(extracted.params_raw()[0].data.as_ref(), b"p2"); + assert_eq!(extracted.params_raw()[1].data.as_ref(), b"p3"); + } + + #[test] + fn test_extract_bind_params_with_format_codes() { + let splits = parse_and_split("INSERT INTO t (a, b) VALUES ($1, $2), ($3, $4)"); + let bind = Bind::new_params_codes( + "test", + &[ + Parameter::new(b"p0"), + Parameter::new(b"p1"), + Parameter::new(b"p2"), + Parameter::new(b"p3"), + ], + &[Format::Text, Format::Binary, Format::Text, Format::Binary], + ); + + // Second split uses params 2 and 3 (Text, Binary) + let extracted = splits[1].extract_bind_params(&bind); + assert_eq!(extracted.params_raw().len(), 2); + assert_eq!(extracted.params_raw()[0].data.as_ref(), b"p2"); + assert_eq!(extracted.params_raw()[1].data.as_ref(), b"p3"); + assert_eq!(extracted.format_codes_raw().len(), 2); + assert_eq!(extracted.format_codes_raw()[0], Format::Text); + assert_eq!(extracted.format_codes_raw()[1], Format::Binary); + } + + #[test] + fn test_extract_bind_params_uniform_format() { + let splits = parse_and_split("INSERT INTO t (a) VALUES ($1), ($2)"); + let bind = Bind::new_params_codes( + "test", + &[Parameter::new(b"p0"), Parameter::new(b"p1")], + &[Format::Binary], // Uniform format + ); + + let extracted = splits[0].extract_bind_params(&bind); + assert_eq!(extracted.params_raw().len(), 1); + assert_eq!(extracted.format_codes_raw().len(), 1); + assert_eq!(extracted.format_codes_raw()[0], Format::Binary); + } + + #[test] + fn test_extract_bind_params_mixed_params_and_literals() { + let splits = parse_and_split("INSERT INTO t (a, b) VALUES ($1, 'lit1'), ($2, 'lit2')"); + let bind = Bind::new_params( + "test", + &[ + Parameter::new(b"value_for_param1"), + Parameter::new(b"value_for_param2"), + ], + ); + + assert_eq!(splits.len(), 2); + + // First split: statement uses $1 with literal, bind extracts param 0 + assert_eq!(splits[0].stmt(), "INSERT INTO t (a, b) VALUES ($1, 'lit1')"); + let extracted = splits[0].extract_bind_params(&bind); + assert_eq!(extracted.params_raw().len(), 1); + assert_eq!(extracted.params_raw()[0].data.as_ref(), b"value_for_param1"); + + // Second split: statement uses $1 (renumbered from $2) with literal, bind extracts param 1 + assert_eq!(splits[1].stmt(), "INSERT INTO t (a, b) VALUES ($1, 'lit2')"); + let extracted = splits[1].extract_bind_params(&bind); + assert_eq!(extracted.params_raw().len(), 1); + assert_eq!(extracted.params_raw()[0].data.as_ref(), b"value_for_param2"); + } + + #[test] + fn test_extract_bind_params_varying_param_counts() { + // First tuple has 2 params, second tuple has 1 param and 1 literal + let splits = parse_and_split("INSERT INTO t (a, b) VALUES ($1, $2), ($3, 'literal')"); + let bind = Bind::new_params( + "test", + &[ + Parameter::new(b"p1"), + Parameter::new(b"p2"), + Parameter::new(b"p3"), + ], + ); + + assert_eq!(splits.len(), 2); + + // First split: uses params 0 and 1 (original $1, $2) + assert_eq!(splits[0].stmt(), "INSERT INTO t (a, b) VALUES ($1, $2)"); + assert_eq!(splits[0].params(), &[0, 1]); + let extracted = splits[0].extract_bind_params(&bind); + assert_eq!(extracted.params_raw().len(), 2); + assert_eq!(extracted.params_raw()[0].data.as_ref(), b"p1"); + assert_eq!(extracted.params_raw()[1].data.as_ref(), b"p2"); + + // Second split: uses param 2 (original $3), renumbered to $1 + assert_eq!( + splits[1].stmt(), + "INSERT INTO t (a, b) VALUES ($1, 'literal')" + ); + assert_eq!(splits[1].params(), &[2]); + let extracted = splits[1].extract_bind_params(&bind); + assert_eq!(extracted.params_raw().len(), 1); + assert_eq!(extracted.params_raw()[0].data.as_ref(), b"p3"); + } +} diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs b/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs new file mode 100644 index 000000000..2081643a6 --- /dev/null +++ b/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs @@ -0,0 +1,122 @@ +//! Statement rewriter. + +use pg_query::protobuf::ParseResult; +use pg_query::Node; +use pgdog_config::QueryParserEngine; + +use crate::backend::ShardingSchema; +use crate::frontend::PreparedStatements; + +pub mod aggregate; +pub mod error; +pub mod insert; +pub mod plan; +pub mod simple_prepared; +pub mod unique_id; +pub mod update; +pub mod visitor; + +pub use error::Error; +pub use insert::InsertSplit; +pub(crate) use plan::RewritePlan; +pub use simple_prepared::SimplePreparedResult; +pub(crate) use update::*; + +/// Statement rewrite engine context. +#[derive(Debug)] +pub struct StatementRewriteContext<'a> { + /// The AST of the statement we are rewriting. + pub stmt: &'a mut ParseResult, + /// The statement is using the extended protocol with placeholders. + pub extended: bool, + /// The statement is named, so we need to save any derivatives into the global + /// statement cache. + pub prepared: bool, + /// Reference to global prepared stmt cache. + pub prepared_statements: &'a mut PreparedStatements, + /// Sharding schema. + pub schema: &'a ShardingSchema, +} + +#[derive(Debug)] +pub struct StatementRewrite<'a> { + /// SQL statement. + stmt: &'a mut ParseResult, + /// The statement was rewritten. + rewritten: bool, + /// Statement is using the extended protocol, so + /// we need to rewrite function calls with parameters + /// and not actual values. + extended: bool, + /// The statement is named (prepared), so we need to save + /// any derivatives into the global statement cache. + prepared: bool, + /// Prepared statements cache for name mapping. + prepared_statements: &'a mut PreparedStatements, + /// Sharding schema for cache lookups. + schema: &'a ShardingSchema, +} + +impl<'a> StatementRewrite<'a> { + /// Create new statement rewriter. + /// + /// More often than not, it won't do anything. + /// + pub fn new(ctx: StatementRewriteContext<'a>) -> Self { + Self { + stmt: ctx.stmt, + rewritten: false, + extended: ctx.extended, + prepared: ctx.prepared, + prepared_statements: ctx.prepared_statements, + schema: ctx.schema, + } + } + + /// Maybe rewrite the statement and produce a rewrite plan + /// we can apply to Bind messages. + pub fn maybe_rewrite(&mut self) -> Result { + let params = visitor::count_params(self.stmt); + let mut plan = RewritePlan { + params, + ..Default::default() + }; + + // Handle top-level PREPARE/EXECUTE statements. + let prepared_result = self.rewrite_simple_prepared()?; + if prepared_result.rewritten { + self.rewritten = true; + plan.prepares = prepared_result.prepares; + } + + // Track the next parameter number to use + let mut next_param = plan.params as i32 + 1; + + // if self.schema.rewrite.enabled { + let extended = self.extended; + visitor::visit_and_mutate_nodes(self.stmt, |node| -> Result, Error> { + match Self::rewrite_unique_id(node, extended, &mut next_param)? { + Some(replacement) => { + plan.unique_ids += 1; + self.rewritten = true; + Ok(Some(replacement)) + } + None => Ok(None), + } + })?; + + self.rewrite_aggregates(&mut plan)?; + + if self.rewritten { + plan.stmt = Some(match self.schema.query_parser_engine { + QueryParserEngine::PgQueryProtobuf => self.stmt.deparse(), + QueryParserEngine::PgQueryRaw => self.stmt.deparse_raw(), + }?); + } + + self.split_insert(&mut plan)?; + self.sharding_key_update(&mut plan)?; + + Ok(plan) + } +} diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/plan.rs b/pgdog/src/frontend/router/parser/rewrite/statement/plan.rs new file mode 100644 index 000000000..c41f25949 --- /dev/null +++ b/pgdog/src/frontend/router/parser/rewrite/statement/plan.rs @@ -0,0 +1,265 @@ +use crate::frontend::{ClientRequest, PreparedStatements}; +use crate::net::messages::bind::{Format, Parameter}; +use crate::net::{Bind, Parse, ProtocolMessage, Query}; +use crate::unique_id::UniqueId; + +use super::insert::build_split_requests; +use super::{aggregate::AggregateRewritePlan, Error, InsertSplit, ShardingKeyUpdate}; + +/// Statement rewrite plan. +/// +/// Executed in order of fields in this struct. +/// +#[derive(Default, Clone, Debug)] +pub struct RewritePlan { + /// Number of parameters ($1, $2, etc.) in + /// the original statement. This is calculated first, + /// and $params+n parameters are added to the statement to + /// substitute values we are rewriting. + pub(crate) params: u16, + + /// Number of unique IDs to append to the Bind message. + pub(crate) unique_ids: u16, + + /// Rewritten SQL statement. + pub(crate) stmt: Option, + + /// Prepared statements to prepend to the client request. + /// Each tuple contains (name, statement) for ProtocolMessage::Prepare. + pub(crate) prepares: Vec<(String, String)>, + + /// Splitting of multi-tuple INSERT statements into + /// multiple queries. + pub(crate) insert_split: Vec, + + /// Position in the result where the count(*) or count(name) + /// functions are added. + pub(crate) aggregates: AggregateRewritePlan, + + /// Sharding key is being updated, we need to execute + /// a multi-step plan. + pub(crate) sharding_key_update: Option, +} + +#[derive(Debug, Clone)] +pub(crate) enum RewriteResult { + InPlace, + InsertSplit(Vec), + ShardingKeyUpdate(ShardingKeyUpdate), +} + +impl RewritePlan { + /// Apply the rewrite plan to a Bind message by appending generated unique IDs. + pub(crate) fn apply_bind(&self, bind: &mut Bind) -> Result<(), Error> { + let format = bind.default_param_format(); + + for _ in 0..self.unique_ids { + let generator = UniqueId::generator()?; + let id = generator.next_id(); + let param = match format { + Format::Binary => Parameter::new(&id.to_be_bytes()), + Format::Text => Parameter::new(id.to_string().as_bytes()), + }; + bind.push_param(param, format); + } + Ok(()) + } + + /// Apply the rewrite plan to a Parse message by updating the SQL. + pub(crate) fn apply_parse(&self, parse: &mut Parse) { + if let Some(ref stmt) = self.stmt { + parse.set_query(stmt); + if !parse.anonymous() { + PreparedStatements::global().write().rewrite(parse); + } + } + } + + /// Apply the rewrite plan to a Query message by updating the SQL. + pub(crate) fn apply_query(&self, query: &mut Query) { + if let Some(ref stmt) = self.stmt { + query.set_query(stmt); + } + } + + /// Apply the rewrite plan to a ClientRequest. + pub(crate) fn apply(&self, request: &mut ClientRequest) -> Result { + // Prepend any required Prepare messages for EXECUTE statements. + if !self.prepares.is_empty() { + let prepends: Vec = self + .prepares + .iter() + .map(|(name, statement)| ProtocolMessage::Prepare { + name: name.clone(), + statement: statement.clone(), + }) + .collect(); + request.messages.splice(0..0, prepends); + } + + for message in request.messages.iter_mut() { + match message { + ProtocolMessage::Parse(parse) => self.apply_parse(parse), + ProtocolMessage::Query(query) => self.apply_query(query), + ProtocolMessage::Bind(bind) => self.apply_bind(bind)?, + _ => {} + } + } + + if !self.insert_split.is_empty() { + let requests = build_split_requests(&self.insert_split, request); + return Ok(RewriteResult::InsertSplit(requests)); + } + + if let Some(sharding_key_update) = &self.sharding_key_update { + if request.is_executable() { + return Ok(RewriteResult::ShardingKeyUpdate( + sharding_key_update.clone(), + )); + } + } + + Ok(RewriteResult::InPlace) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::collections::HashSet; + + #[test] + fn test_apply_bind_no_unique_ids() { + unsafe { + std::env::set_var("NODE_ID", "pgdog-1"); + } + let plan = RewritePlan::default(); + let mut bind = Bind::default(); + plan.apply_bind(&mut bind).unwrap(); + assert_eq!(bind.params_raw().len(), 0); + } + + #[test] + fn test_apply_bind_text_format() { + unsafe { + std::env::set_var("NODE_ID", "pgdog-1"); + } + let plan = RewritePlan { + unique_ids: 1, + ..Default::default() + }; + let mut bind = Bind::default(); + plan.apply_bind(&mut bind).unwrap(); + assert_eq!(bind.params_raw().len(), 1); + + // Default format is Text, so data should be a string + let param = &bind.params_raw()[0]; + let text = std::str::from_utf8(¶m.data).unwrap(); + let _id: i64 = text.parse().expect("should be valid i64 text"); + + // No format codes needed for all-text + assert_eq!(bind.format_codes_raw().len(), 0); + } + + #[test] + fn test_apply_bind_binary_format_uniform() { + unsafe { + std::env::set_var("NODE_ID", "pgdog-1"); + } + let plan = RewritePlan { + params: 1, + unique_ids: 1, + ..Default::default() + }; + // Create bind with uniform binary format (1 code applies to all) + let mut bind = + Bind::new_params_codes("test", &[Parameter::new(b"existing")], &[Format::Binary]); + plan.apply_bind(&mut bind).unwrap(); + assert_eq!(bind.params_raw().len(), 2); + + // Should use binary format: 8 bytes big-endian + let param = &bind.params_raw()[1]; + assert_eq!(param.data.len(), 8, "binary bigint should be 8 bytes"); + let id = i64::from_be_bytes(param.data[..].try_into().unwrap()); + assert!(id > 0, "ID should be positive"); + + // Uniform format preserved (still 1 code) + assert_eq!(bind.format_codes_raw().len(), 1); + assert_eq!(bind.format_codes_raw()[0], Format::Binary); + } + + #[test] + fn test_apply_bind_binary_format_one_to_one() { + unsafe { + std::env::set_var("NODE_ID", "pgdog-1"); + } + let plan = RewritePlan { + params: 2, + unique_ids: 1, + ..Default::default() + }; + // Create bind with one-to-one format codes + let mut bind = Bind::new_params_codes( + "test", + &[Parameter::new(b"a"), Parameter::new(b"b")], + &[Format::Binary, Format::Binary], + ); + plan.apply_bind(&mut bind).unwrap(); + assert_eq!(bind.params_raw().len(), 3); + + // New param should be text (default for one-to-one) + let param = &bind.params_raw()[2]; + let text = std::str::from_utf8(¶m.data).unwrap(); + let _: i64 = text.parse().expect("should be valid i64 text"); + + // Format code added for new param + assert_eq!(bind.format_codes_raw().len(), 3); + assert_eq!(bind.format_codes_raw()[2], Format::Text); + } + + #[test] + fn test_apply_bind_multiple_unique_ids() { + unsafe { + std::env::set_var("NODE_ID", "pgdog-1"); + } + let plan = RewritePlan { + unique_ids: 3, + ..Default::default() + }; + let mut bind = Bind::default(); + plan.apply_bind(&mut bind).unwrap(); + assert_eq!(bind.params_raw().len(), 3); + + let mut ids = HashSet::new(); + for param in bind.params_raw() { + let text = std::str::from_utf8(¶m.data).unwrap(); + let id: i64 = text.parse().expect("should be valid i64"); + ids.insert(id); + } + assert_eq!(ids.len(), 3, "all IDs should be unique"); + } + + #[test] + fn test_apply_bind_appends_to_existing_params() { + unsafe { + std::env::set_var("NODE_ID", "pgdog-1"); + } + let plan = RewritePlan { + params: 2, + unique_ids: 2, + ..Default::default() + }; + let mut bind = Bind::new_params( + "test", + &[Parameter::new(b"existing1"), Parameter::new(b"existing2")], + ); + plan.apply_bind(&mut bind).unwrap(); + assert_eq!(bind.params_raw().len(), 4); + + assert_eq!(bind.params_raw()[0].data.as_ref(), b"existing1"); + assert_eq!(bind.params_raw()[1].data.as_ref(), b"existing2"); + + let text = std::str::from_utf8(&bind.params_raw()[2].data).unwrap(); + let _: i64 = text.parse().expect("should be valid i64"); + } +} diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/simple_prepared.rs b/pgdog/src/frontend/router/parser/rewrite/statement/simple_prepared.rs new file mode 100644 index 000000000..54b306a34 --- /dev/null +++ b/pgdog/src/frontend/router/parser/rewrite/statement/simple_prepared.rs @@ -0,0 +1,237 @@ +use pg_query::{Error as PgQueryError, NodeEnum}; +use pgdog_config::QueryParserEngine; + +use crate::net::Parse; +use crate::{backend::ShardingSchema, frontend::PreparedStatements}; + +use super::{Error, StatementRewrite}; + +/// Result of rewriting all PREPARE/EXECUTE statements in a query. +#[derive(Debug, Clone, Default)] +pub struct SimplePreparedResult { + /// Whether any statement was rewritten. + pub rewritten: bool, + /// Prepared statements to prepend (name, statement) for EXECUTE rewrites. + pub prepares: Vec<(String, String)>, +} + +/// Result of rewriting a single PREPARE or EXECUTE SQL command. +#[derive(Debug, Clone)] +enum SimplePreparedRewrite { + /// Node was not a PREPARE or EXECUTE statement. + None, + /// PREPARE statement was rewritten. + Prepared, + /// EXECUTE statement was rewritten. Contains the global name and statement + /// needed to prepend a ProtocolMessage::Prepare. + Executed { name: String, statement: String }, +} + +impl StatementRewrite<'_> { + /// Rewrites all top-level `PREPARE` and `EXECUTE` SQL commands. + /// + /// # More details + /// + /// `PREPARE __stmt_1 AS SELECT $1` is rewritten as `PREPARE __pgdog_1 AS SELECT $1` and + /// `SELECT $1` is stored in the global cache using `insert_anyway`. + /// + /// `EXECUTE __stmt_1(1)` is rewritten to `EXECUTE __pgdog_1(1)`. Additionally, the caller + /// should prepend `ProtocolMessage::Prepare` to the client request using the returned + /// name and statement. + /// + pub(super) fn rewrite_simple_prepared(&mut self) -> Result { + let mut result = SimplePreparedResult::default(); + + if !self.prepared_statements.level.full() { + return Ok(result); + } + + for stmt in &mut self.stmt.stmts { + if let Some(ref mut node) = stmt.stmt { + if let Some(ref mut inner) = node.node { + match rewrite_single_prepared(inner, self.prepared_statements, self.schema)? { + SimplePreparedRewrite::Prepared => { + result.rewritten = true; + } + SimplePreparedRewrite::Executed { name, statement } => { + result.prepares.push((name, statement)); + result.rewritten = true; + } + SimplePreparedRewrite::None => {} + } + } + } + } + + Ok(result) + } +} + +/// Rewrites a single `PREPARE` or `EXECUTE` node. +fn rewrite_single_prepared( + node: &mut NodeEnum, + prepared_statements: &mut PreparedStatements, + schema: &ShardingSchema, +) -> Result { + match node { + NodeEnum::PrepareStmt(stmt) => { + let query = stmt + .query + .as_ref() + .ok_or(Error::PgQuery(PgQueryError::Parse( + "missing query in PREPARE".into(), + )))?; + let query = match schema.query_parser_engine { + QueryParserEngine::PgQueryProtobuf => query.deparse(), + QueryParserEngine::PgQueryRaw => query.deparse_raw(), + } + .map_err(Error::PgQuery)?; + + let mut parse = Parse::named(&stmt.name, &query); + prepared_statements.insert_anyway(&mut parse); + stmt.name = parse.name().to_string(); + + Ok(SimplePreparedRewrite::Prepared) + } + + NodeEnum::ExecuteStmt(stmt) => { + let parse = prepared_statements.parse(&stmt.name); + if let Some(parse) = parse { + let global_name = parse.name().to_string(); + let statement = parse.query().to_string(); + stmt.name = global_name.clone(); + + Ok(SimplePreparedRewrite::Executed { + name: global_name, + statement, + }) + } else { + Err(Error::PgQuery(PgQueryError::Parse(format!( + "prepared statement '{}' does not exist", + stmt.name + )))) + } + } + + _ => Ok(SimplePreparedRewrite::None), + } +} + +#[cfg(test)] +mod tests { + use super::super::{RewritePlan, StatementRewrite, StatementRewriteContext}; + use super::*; + use crate::backend::ShardingSchema; + use crate::config::PreparedStatements as PreparedStatementsLevel; + use pg_query::parse; + use pg_query::protobuf::ParseResult; + use pgdog_config::Rewrite; + + struct TestContext { + ps: PreparedStatements, + schema: ShardingSchema, + } + + impl TestContext { + fn new() -> Self { + let mut ps = PreparedStatements::default(); + ps.set_level(PreparedStatementsLevel::Full); + Self { + ps, + schema: ShardingSchema { + shards: 1, + rewrite: Rewrite { + enabled: true, + ..Default::default() + }, + ..Default::default() + }, + } + } + + fn rewrite(&mut self, sql: &str) -> Result<(ParseResult, RewritePlan), Error> { + let mut ast = parse(sql).unwrap().protobuf; + let mut rewrite = StatementRewrite::new(StatementRewriteContext { + stmt: &mut ast, + extended: false, + prepared: false, + prepared_statements: &mut self.ps, + schema: &self.schema, + }); + let plan = rewrite.maybe_rewrite()?; + Ok((ast, plan)) + } + } + + #[test] + fn test_rewrite_prepare() { + let mut ctx = TestContext::new(); + let (ast, plan) = ctx.rewrite("PREPARE test_stmt AS SELECT $1, $2").unwrap(); + + let sql = ast.deparse().unwrap(); + assert!( + sql.contains("__pgdog_"), + "PREPARE should be renamed to __pgdog_N, got: {sql}" + ); + assert!( + !sql.contains("test_stmt"), + "original name should be replaced: {sql}" + ); + assert!(plan.prepares.is_empty()); + assert!(plan.stmt.is_some()); + } + + #[test] + fn test_rewrite_execute() { + let mut ctx = TestContext::new(); + ctx.rewrite("PREPARE test_stmt AS SELECT 1").unwrap(); + let (ast, plan) = ctx.rewrite("EXECUTE test_stmt").unwrap(); + + let sql = ast.deparse().unwrap(); + assert!( + sql.contains("__pgdog_"), + "EXECUTE should use global name, got: {sql}" + ); + assert_eq!(plan.prepares.len(), 1); + + let (name, statement) = &plan.prepares[0]; + assert!(name.starts_with("__pgdog_")); + assert_eq!(statement, "SELECT 1"); + } + + #[test] + fn test_rewrite_execute_with_params() { + let mut ctx = TestContext::new(); + ctx.rewrite("PREPARE test_stmt AS SELECT $1, $2").unwrap(); + let (ast, plan) = ctx.rewrite("EXECUTE test_stmt(1, 'hello')").unwrap(); + + let sql = ast.deparse().unwrap(); + assert!( + sql.contains("__pgdog_"), + "EXECUTE should use global name, got: {sql}" + ); + assert!( + sql.contains("(1, 'hello')"), + "EXECUTE params should be preserved, got: {sql}" + ); + assert_eq!(plan.prepares.len(), 1); + } + + #[test] + fn test_execute_nonexistent_fails() { + let mut ctx = TestContext::new(); + let result = ctx.rewrite("EXECUTE nonexistent_stmt"); + assert!(result.is_err()); + } + + #[test] + fn test_no_rewrite_for_regular_select() { + let mut ctx = TestContext::new(); + let (ast, plan) = ctx.rewrite("SELECT 1, 2, 3").unwrap(); + + let sql = ast.deparse().unwrap(); + assert_eq!(sql, "SELECT 1, 2, 3"); + assert!(plan.prepares.is_empty()); + assert!(plan.stmt.is_none()); + } +} diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/unique_id.rs b/pgdog/src/frontend/router/parser/rewrite/statement/unique_id.rs new file mode 100644 index 000000000..054d72fa2 --- /dev/null +++ b/pgdog/src/frontend/router/parser/rewrite/statement/unique_id.rs @@ -0,0 +1,478 @@ +use pg_query::protobuf::{a_const::Val, AConst, ParamRef, String as PgString, TypeCast, TypeName}; +use pg_query::{Node, NodeEnum}; + +use super::StatementRewrite; + +impl StatementRewrite<'_> { + /// Attempt to rewrite a pgdog.unique_id() call. + /// + /// Returns `Ok(Some(replacement_node))` if the node is a unique_id call, + /// `Ok(None)` otherwise. Increments `next_param` when in extended mode. + pub(super) fn rewrite_unique_id( + node: &Node, + extended: bool, + next_param: &mut i32, + ) -> Result, super::Error> { + if !Self::is_unique_id(node) { + return Ok(None); + } + + let replacement = if extended { + let param_num = *next_param; + *next_param += 1; + Self::param_bigint(param_num) + } else { + let unique_id = crate::unique_id::UniqueId::generator()?.next_id(); + Self::literal_bigint(unique_id) + }; + + Ok(Some(replacement)) + } + + /// Create a parameter reference cast to bigint: $N::bigint + fn param_bigint(number: i32) -> Node { + let param_ref = Node { + node: Some(NodeEnum::ParamRef(ParamRef { + number, + ..Default::default() + })), + }; + + Node { + node: Some(NodeEnum::TypeCast(Box::new(TypeCast { + arg: Some(Box::new(param_ref)), + type_name: Some(Self::bigint_type()), + ..Default::default() + }))), + } + } + + /// Create a literal value cast to bigint: ::bigint + fn literal_bigint(value: i64) -> Node { + let literal = Node { + node: Some(NodeEnum::AConst(AConst { + val: Some(Val::Sval(PgString { + sval: value.to_string(), + })), + ..Default::default() + })), + }; + + Node { + node: Some(NodeEnum::TypeCast(Box::new(TypeCast { + arg: Some(Box::new(literal)), + type_name: Some(Self::bigint_type()), + ..Default::default() + }))), + } + } + + /// Create a TypeName for bigint (int8). + fn bigint_type() -> TypeName { + TypeName { + names: vec![ + Node { + node: Some(NodeEnum::String(PgString { + sval: "pg_catalog".to_string(), + })), + }, + Node { + node: Some(NodeEnum::String(PgString { + sval: "int8".to_string(), + })), + }, + ], + ..Default::default() + } + } + + /// Check if a node is a function call to pgdog.unique_id(). + fn is_unique_id(node: &Node) -> bool { + let Some(NodeEnum::FuncCall(func)) = &node.node else { + return false; + }; + + // Must have exactly 2 parts: schema "pgdog" and function "unique_id" + if func.funcname.len() != 2 { + return false; + } + + let schema = func.funcname.first().and_then(|n| match &n.node { + Some(NodeEnum::String(s)) => Some(s.sval.as_str()), + _ => None, + }); + + let name = func.funcname.get(1).and_then(|n| match &n.node { + Some(NodeEnum::String(s)) => Some(s.sval.as_str()), + _ => None, + }); + + matches!((schema, name), (Some("pgdog"), Some("unique_id"))) + } +} + +#[cfg(test)] +mod tests { + use pgdog_config::Rewrite; + + use super::*; + use crate::backend::ShardingSchema; + use crate::frontend::router::parser::StatementRewriteContext; + use crate::frontend::PreparedStatements; + + fn default_schema() -> ShardingSchema { + ShardingSchema { + shards: 1, + rewrite: Rewrite { + enabled: true, + ..Default::default() + }, + ..Default::default() + } + } + + fn parse_first_target(sql: &str) -> Node { + let ast = pg_query::parse(sql).unwrap(); + let stmt = ast.protobuf.stmts.first().unwrap().stmt.as_ref().unwrap(); + match &stmt.node { + Some(NodeEnum::SelectStmt(select)) => { + let res_target = select.target_list.first().unwrap(); + match &res_target.node { + Some(NodeEnum::ResTarget(res)) => *res.val.as_ref().unwrap().clone(), + _ => panic!("expected ResTarget"), + } + } + _ => panic!("expected SelectStmt"), + } + } + + #[test] + fn test_is_unique_id_qualified() { + let node = parse_first_target("SELECT pgdog.unique_id()"); + assert!(StatementRewrite::is_unique_id(&node)); + } + + #[test] + fn test_is_unique_id_unqualified() { + let node = parse_first_target("SELECT unique_id()"); + assert!(!StatementRewrite::is_unique_id(&node)); + } + + #[test] + fn test_is_unique_id_wrong_schema() { + let node = parse_first_target("SELECT other.unique_id()"); + assert!(!StatementRewrite::is_unique_id(&node)); + } + + #[test] + fn test_is_unique_id_wrong_function() { + let node = parse_first_target("SELECT pgdog.other_func()"); + assert!(!StatementRewrite::is_unique_id(&node)); + } + + #[test] + fn test_is_unique_id_not_function() { + let node = parse_first_target("SELECT 1"); + assert!(!StatementRewrite::is_unique_id(&node)); + } + + #[test] + fn test_rewrite_select_extended_single() { + let mut ast = pg_query::parse("SELECT pgdog.unique_id()") + .unwrap() + .protobuf; + let mut ps = PreparedStatements::default(); + let schema = default_schema(); + let mut rewrite = StatementRewrite::new(StatementRewriteContext { + stmt: &mut ast, + extended: true, + prepared: false, + prepared_statements: &mut ps, + schema: &schema, + }); + let plan = rewrite.maybe_rewrite().unwrap(); + + let sql = ast.deparse().unwrap(); + assert_eq!(sql, "SELECT $1::bigint"); + assert_eq!(plan.params, 0); + assert_eq!(plan.unique_ids, 1); + } + + #[test] + fn test_rewrite_select_extended_with_existing_params() { + let mut ast = pg_query::parse("SELECT pgdog.unique_id(), $1, $2") + .unwrap() + .protobuf; + let mut ps = PreparedStatements::default(); + let schema = default_schema(); + let mut rewrite = StatementRewrite::new(StatementRewriteContext { + stmt: &mut ast, + extended: true, + prepared: false, + prepared_statements: &mut ps, + schema: &schema, + }); + let plan = rewrite.maybe_rewrite().unwrap(); + + let sql = ast.deparse().unwrap(); + assert_eq!(sql, "SELECT $3::bigint, $1, $2"); + assert_eq!(plan.params, 2); + assert_eq!(plan.unique_ids, 1); + } + + #[test] + fn test_rewrite_select_extended_multiple_unique_ids() { + let mut ast = pg_query::parse("SELECT pgdog.unique_id(), pgdog.unique_id()") + .unwrap() + .protobuf; + let mut ps = PreparedStatements::default(); + let schema = default_schema(); + let mut rewrite = StatementRewrite::new(StatementRewriteContext { + stmt: &mut ast, + extended: true, + prepared: false, + prepared_statements: &mut ps, + schema: &schema, + }); + let plan = rewrite.maybe_rewrite().unwrap(); + + let sql = ast.deparse().unwrap(); + assert_eq!(sql, "SELECT $1::bigint, $2::bigint"); + assert_eq!(plan.params, 0); + assert_eq!(plan.unique_ids, 2); + } + + #[test] + fn test_rewrite_select_simple() { + unsafe { + std::env::set_var("NODE_ID", "pgdog-1"); + } + let mut ast = pg_query::parse("SELECT pgdog.unique_id()") + .unwrap() + .protobuf; + let mut ps = PreparedStatements::default(); + let schema = default_schema(); + let mut rewrite = StatementRewrite::new(StatementRewriteContext { + stmt: &mut ast, + extended: false, + prepared: false, + prepared_statements: &mut ps, + schema: &schema, + }); + let plan = rewrite.maybe_rewrite().unwrap(); + + let sql = ast.deparse().unwrap(); + // Value should be a bigint literal cast + assert!( + sql.contains("::bigint"), + "Expected ::bigint cast, got: {sql}" + ); + assert!( + !sql.contains("pgdog.unique_id"), + "Function should be replaced: {sql}" + ); + assert_eq!(plan.params, 0); + assert_eq!(plan.unique_ids, 1); + } + + #[test] + fn test_rewrite_select_simple_multiple_unique_ids() { + unsafe { + std::env::set_var("NODE_ID", "pgdog-1"); + } + let mut ast = pg_query::parse("SELECT pgdog.unique_id(), pgdog.unique_id()") + .unwrap() + .protobuf; + let mut ps = PreparedStatements::default(); + let schema = default_schema(); + let mut rewrite = StatementRewrite::new(StatementRewriteContext { + stmt: &mut ast, + extended: false, + prepared: false, + prepared_statements: &mut ps, + schema: &schema, + }); + let plan = rewrite.maybe_rewrite().unwrap(); + + let sql = ast.deparse().unwrap(); + // Each unique_id call should get a different value + assert!( + !sql.contains("pgdog.unique_id"), + "Functions should be replaced: {sql}" + ); + assert_eq!(plan.unique_ids, 2); + } + + #[test] + fn test_rewrite_no_unique_id() { + let mut ast = pg_query::parse("SELECT 1, 2, 3").unwrap().protobuf; + let mut ps = PreparedStatements::default(); + let schema = default_schema(); + let mut rewrite = StatementRewrite::new(StatementRewriteContext { + stmt: &mut ast, + extended: true, + prepared: false, + prepared_statements: &mut ps, + schema: &schema, + }); + let plan = rewrite.maybe_rewrite().unwrap(); + + let sql = ast.deparse().unwrap(); + assert_eq!(sql, "SELECT 1, 2, 3"); + assert_eq!(plan.unique_ids, 0); + } + + #[test] + fn test_rewrite_insert_values() { + let mut ast = + pg_query::parse("INSERT INTO t (id, name) VALUES (pgdog.unique_id(), 'test')") + .unwrap() + .protobuf; + let mut ps = PreparedStatements::default(); + let schema = default_schema(); + let mut rewrite = StatementRewrite::new(StatementRewriteContext { + stmt: &mut ast, + extended: true, + prepared: false, + prepared_statements: &mut ps, + schema: &schema, + }); + let plan = rewrite.maybe_rewrite().unwrap(); + + let sql = ast.deparse().unwrap(); + assert_eq!(sql, "INSERT INTO t (id, name) VALUES ($1::bigint, 'test')"); + assert_eq!(plan.unique_ids, 1); + } + + #[test] + fn test_rewrite_insert_multiple_rows() { + let mut ast = + pg_query::parse("INSERT INTO t (id) VALUES (pgdog.unique_id()), (pgdog.unique_id())") + .unwrap() + .protobuf; + let mut ps = PreparedStatements::default(); + let schema = default_schema(); + let mut rewrite = StatementRewrite::new(StatementRewriteContext { + stmt: &mut ast, + extended: true, + prepared: false, + prepared_statements: &mut ps, + schema: &schema, + }); + let plan = rewrite.maybe_rewrite().unwrap(); + + let sql = ast.deparse().unwrap(); + assert_eq!(sql, "INSERT INTO t (id) VALUES ($1::bigint), ($2::bigint)"); + assert_eq!(plan.unique_ids, 2); + } + + #[test] + fn test_rewrite_insert_select() { + let mut ast = pg_query::parse("INSERT INTO t (id) SELECT pgdog.unique_id() FROM s") + .unwrap() + .protobuf; + let mut ps = PreparedStatements::default(); + let schema = default_schema(); + let mut rewrite = StatementRewrite::new(StatementRewriteContext { + stmt: &mut ast, + extended: true, + prepared: false, + prepared_statements: &mut ps, + schema: &schema, + }); + let plan = rewrite.maybe_rewrite().unwrap(); + + let sql = ast.deparse().unwrap(); + assert_eq!(sql, "INSERT INTO t (id) SELECT $1::bigint FROM s"); + assert_eq!(plan.unique_ids, 1); + } + + #[test] + fn test_rewrite_update_set() { + let mut ast = pg_query::parse("UPDATE t SET id = pgdog.unique_id() WHERE name = 'test'") + .unwrap() + .protobuf; + let mut ps = PreparedStatements::default(); + let schema = default_schema(); + let mut rewrite = StatementRewrite::new(StatementRewriteContext { + stmt: &mut ast, + extended: true, + prepared: false, + prepared_statements: &mut ps, + schema: &schema, + }); + let plan = rewrite.maybe_rewrite().unwrap(); + + let sql = ast.deparse().unwrap(); + assert_eq!(sql, "UPDATE t SET id = $1::bigint WHERE name = 'test'"); + assert_eq!(plan.unique_ids, 1); + } + + #[test] + fn test_rewrite_update_where() { + let mut ast = pg_query::parse("UPDATE t SET name = 'new' WHERE id = pgdog.unique_id()") + .unwrap() + .protobuf; + let mut ps = PreparedStatements::default(); + let schema = default_schema(); + let mut rewrite = StatementRewrite::new(StatementRewriteContext { + stmt: &mut ast, + extended: true, + prepared: false, + prepared_statements: &mut ps, + schema: &schema, + }); + let plan = rewrite.maybe_rewrite().unwrap(); + + let sql = ast.deparse().unwrap(); + assert_eq!(sql, "UPDATE t SET name = 'new' WHERE id = $1::bigint"); + assert_eq!(plan.unique_ids, 1); + } + + #[test] + fn test_rewrite_delete_where() { + let mut ast = pg_query::parse("DELETE FROM t WHERE id = pgdog.unique_id()") + .unwrap() + .protobuf; + let mut ps = PreparedStatements::default(); + let schema = default_schema(); + let mut rewrite = StatementRewrite::new(StatementRewriteContext { + stmt: &mut ast, + extended: true, + prepared: false, + prepared_statements: &mut ps, + schema: &schema, + }); + let plan = rewrite.maybe_rewrite().unwrap(); + + let sql = ast.deparse().unwrap(); + assert_eq!(sql, "DELETE FROM t WHERE id = $1::bigint"); + assert_eq!(plan.unique_ids, 1); + } + + #[test] + fn test_rewrite_insert_returning() { + let mut ast = pg_query::parse( + "INSERT INTO t (id) VALUES (pgdog.unique_id()) RETURNING pgdog.unique_id()", + ) + .unwrap() + .protobuf; + let mut ps = PreparedStatements::default(); + let schema = default_schema(); + let mut rewrite = StatementRewrite::new(StatementRewriteContext { + stmt: &mut ast, + extended: true, + prepared: false, + prepared_statements: &mut ps, + schema: &schema, + }); + let plan = rewrite.maybe_rewrite().unwrap(); + + let sql = ast.deparse().unwrap(); + assert_eq!( + sql, + "INSERT INTO t (id) VALUES ($1::bigint) RETURNING $2::bigint" + ); + assert_eq!(plan.unique_ids, 2); + } +} diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/update.rs b/pgdog/src/frontend/router/parser/rewrite/statement/update.rs new file mode 100644 index 000000000..c4fda855c --- /dev/null +++ b/pgdog/src/frontend/router/parser/rewrite/statement/update.rs @@ -0,0 +1,1190 @@ +use std::{collections::HashMap, ops::Deref, sync::Arc}; + +use pg_query::{ + protobuf::{ + AExpr, AExprKind, AStar, ColumnRef, DeleteStmt, InsertStmt, LimitOption, List, + OverridingKind, ParamRef, ParseResult, RangeVar, RawStmt, ResTarget, SelectStmt, + SetOperation, String as PgString, UpdateStmt, + }, + Node, NodeEnum, +}; +use pgdog_config::{QueryParserEngine, RewriteMode}; + +use crate::{ + frontend::{ + router::{ + parser::{rewrite::statement::visitor::visit_and_mutate_nodes, Column, Table, Value}, + Ast, + }, + BufferedQuery, ClientRequest, + }, + net::{ + bind::Parameter, Bind, DataRow, Describe, Execute, Flush, Format, FromDataType, Parse, + ProtocolMessage, Query, RowDescription, Sync, + }, +}; + +use super::*; + +#[derive(Debug, Clone)] +pub(crate) struct Statement { + pub(crate) ast: Ast, + pub(crate) stmt: String, + pub(crate) params: Vec, +} + +impl Statement { + /// Create new Bind message for the statement from original Bind. + pub(crate) fn rewrite_bind(&self, bind: &Bind) -> Result { + let mut new = Bind::new_statement(""); // We use anonymous prepared + // statements for execution. + for param in &self.params { + let param = bind + .parameter(*param as usize - 1)? + .ok_or(Error::MissingParameter(*param))?; + new.push_param(param.parameter().clone(), param.format()); + } + + Ok(new) + } + + /// Build request from statement. + /// + /// Use the same protocol as the original statement. + /// + pub(crate) fn build_request(&self, request: &ClientRequest) -> Result { + let query = request.query()?.ok_or(Error::EmptyQuery)?; + let params = request.parameters()?; + + let mut request = ClientRequest::new(); + + match query { + BufferedQuery::Query(_) => { + request.push(Query::new(self.stmt.clone()).into()); + } + BufferedQuery::Prepared(_) => { + request.push(Parse::new_anonymous(&self.stmt).into()); + request.push(Describe::new_statement("").into()); + if let Some(params) = params { + request.push(self.rewrite_bind(params)?.into()); + request.push(Execute::new().into()); + request.push(Sync.into()); + } else { + // This shouldn't really happen since we don't rewrite + // non-executable requests. + request.push(Flush.into()); + } + } + } + + request.ast = Some(self.ast.clone()); + + Ok(request) + } +} + +#[derive(Debug, Clone)] +pub(crate) struct ShardingKeyUpdate { + inner: Arc, +} + +impl Deref for ShardingKeyUpdate { + type Target = Inner; + + fn deref(&self) -> &Self::Target { + &self.inner + } +} + +#[derive(Debug, Clone)] +pub(crate) struct Inner { + /// Fetch the whole old row. + pub(crate) select: Statement, + /// Check that the row actually moves shards. + pub(crate) check: Statement, + /// Delete old row from shard. + pub(crate) delete: Statement, + /// Partial insert statement. + pub(crate) insert: Insert, +} + +/// Partially built INSERT statement. +#[derive(Debug, Clone)] +pub(crate) struct Insert { + pub(super) table: Option, + /// Mapping of column name to `column name = value` from + /// the original UPDATE statement. + pub(super) mapping: HashMap, + /// Return columns. + pub(super) returning_list: Vec, + /// Returning list deparsed. + pub(super) returnin_list_deparsed: Option, +} + +impl Insert { + /// Build an INSERT statement built from an existing + /// UPDATE statement and a row returned by a SELECT statement. + /// + pub(crate) fn build_request( + &self, + request: &ClientRequest, + row_description: &RowDescription, + data_row: &DataRow, + ) -> Result { + let params = request.parameters()?; + + let mut bind = Bind::new_statement(""); + let mut columns = vec![]; + let mut values = vec![]; + let mut columns_str = vec![]; + let mut values_str = vec![]; + + let mut bind_idx = 0; + for (row_idx, field) in row_description.iter().enumerate() { + columns_str.push(format!(r#""{}""#, field.name.replace("\"", "\"\""))); // Escape " + + if let Some(value) = self.mapping.get(&field.name) { + let value = match value { + UpdateValue::Value(value) => { + values_str.push(format!("${}", bind_idx + 1)); + Value::try_from(value).unwrap() // SAFETY: We check that the value is valid. + } + UpdateValue::Expr(expr) => { + values_str.push(expr.clone()); + continue; + } + }; + + match value { + Value::Placeholder(number) => { + let param = params + .as_ref() + .expect("param") + .parameter(number as usize - 1)? + .ok_or(Error::MissingParameter(number as u16))?; + bind.push_param(param.parameter().clone(), param.format()) + } + + Value::Integer(int) => { + bind.push_param(Parameter::new(int.to_string().as_bytes()), Format::Text) + } + + Value::String(s) => bind.push_param(Parameter::new(s.as_bytes()), Format::Text), + + Value::Float(f) => { + bind.push_param(Parameter::new(f.to_string().as_bytes()), Format::Text) + } + + Value::Boolean(b) => bind.push_param( + Parameter::new(if b { "t".as_bytes() } else { "f".as_bytes() }), + Format::Text, + ), + + Value::Vector(vec) => { + bind.push_param(Parameter::new(&vec.encode(Format::Text)?), Format::Text) + } + + Value::Null => bind.push_param(Parameter::new_null(), Format::Text), + } + } else { + let value = data_row + .get_raw(row_idx) + .ok_or(Error::MissingColumn(row_idx))?; + + if value.is_null() { + bind.push_param(Parameter::new_null(), Format::Text); + } else { + bind.push_param(Parameter::new(value), Format::Text); + } + + values_str.push(format!("${}", bind_idx + 1)); + } + + columns.push(Node { + node: Some(NodeEnum::ResTarget(Box::new(ResTarget { + name: field.name.clone(), + ..Default::default() + }))), + }); + + values.push(Node { + node: Some(NodeEnum::ParamRef(ParamRef { + number: bind_idx + 1, + ..Default::default() + })), + }); + + bind_idx += 1; + } + + let insert = InsertStmt { + relation: self.table.clone(), + cols: columns, + select_stmt: Some(Box::new(Node { + node: Some(NodeEnum::SelectStmt(Box::new(SelectStmt { + target_list: vec![], + from_clause: vec![], + limit_option: LimitOption::Default.into(), + where_clause: None, + op: SetOperation::SetopNone.into(), + values_lists: vec![Node { + node: Some(NodeEnum::List(List { items: values })), + }], + ..Default::default() + }))), + })), + returning_list: self.returning_list.clone(), + r#override: OverridingKind::OverridingNotSet.into(), + ..Default::default() + }; + + let table = self.table.as_ref().map(Table::from).unwrap(); // SAFETY: We check that UPDATE has a table. + + // This is probably one of the few places in the code where + // we shouldn't use the parser. It's quicker to concatenate strings + // than to call pg_query::deparse because of the Protobuf (de)ser. + // + // TODO: Replace protobuf (de)ser with native mappings and use the + // parser again. + // + let stmt = format!( + "INSERT INTO {} ({}) VALUES ({}){}", + table, + columns_str.join(", "), + values_str.join(", "), + if let Some(ref returning_list) = self.returnin_list_deparsed { + format!("RETURNING {}", returning_list) + } else { + "".into() + } + ); + + // Build the AST to be used with the router. + // It's identical to the string-generated statement above. + let insert = parse_result(NodeEnum::InsertStmt(Box::new(insert))); + let insert = pg_query::ParseResult::new(insert, "".into()); + + let ast = Ast::from_parse_result(insert); + + let mut req = ClientRequest::from(vec![ + ProtocolMessage::from(Parse::new_anonymous(&stmt)), + Describe::new_statement("").into(), // So we get both T and t, + bind.into(), + Execute::new().into(), + Sync.into(), + ]); + req.ast = Some(ast); + Ok(req) + } + + /// Do we have to return the rows to the client? + pub(crate) fn is_returning(&self) -> bool { + !self.returning_list.is_empty() && self.returnin_list_deparsed.is_some() + } +} + +impl<'a> StatementRewrite<'a> { + /// Create a plan for shardking key updates, if we suspect there is one + /// in the query. + pub(super) fn sharding_key_update(&mut self, plan: &mut RewritePlan) -> Result<(), Error> { + if self.schema.shards == 1 || self.schema.rewrite.shard_key == RewriteMode::Ignore { + return Ok(()); + } + + let stmt = self + .stmt + .stmts + .first() + .and_then(|stmt| stmt.stmt.as_ref().map(|stmt| stmt.node.as_ref())) + .flatten(); + + let stmt = if let Some(NodeEnum::UpdateStmt(stmt)) = stmt { + stmt + } else { + // TODO: Handle EXPLAIN ANALYZE which needs to execute. + // We could return a combined plan for all 3 queries + // we need to execute. + return Ok(()); + }; + + if let Some(value) = self.sharding_key_update_check(stmt)? { + // Without a WHERE clause, this is a huge + // cross-shard rewrite. + if stmt.where_clause.is_none() { + return Err(Error::WhereClauseMissing); + } + plan.sharding_key_update = + Some(create_stmts(stmt, value, self.schema.query_parser_engine)?); + } + + Ok(()) + } + + /// Check if the sharding key could be updated. + fn sharding_key_update_check( + &'a self, + stmt: &'a UpdateStmt, + ) -> Result>, Error> { + let table = if let Some(table) = stmt.relation.as_ref().map(Table::from) { + table + } else { + return Ok(None); + }; + + Ok(stmt + .target_list + .iter() + .filter(|column| { + if let Ok(mut column) = Column::try_from(&column.node) { + column.qualify(table); + self.schema.tables().get_table(column).is_some() + } else { + false + } + }) + .map(|column| { + if let Some(NodeEnum::ResTarget(res)) = &column.node { + // Check that it's a value assignment and not something like + // id = id + 1 + let supported = res + .val + .as_ref() + .map(|node| Value::try_from(&node.node)) + .transpose() + .is_ok(); + + if supported { + Ok(Some(res)) + } else { + // FIXME: + // + // We can technically support this. We can inject this into + // the `SELECT` statement we use to pull the existing row + // and use the computed value for assignment. + // + let expr = res + .val + .as_ref() + .map(|node| deparse_expr(node, self.schema.query_parser_engine)) + .transpose()? + .unwrap_or_else(|| "".to_string()); + Err(Error::UnsupportedShardingKeyUpdate(format!( + "\"{}\" = {}", + res.name, expr + ))) + } + } else { + Ok(None) + } + }) + .next() + .transpose()? + .flatten()) + } +} + +/// Visit all ParamRef nodes in a ParseResult and renumber them sequentially. +/// Returns a sorted list of the original parameter numbers. +fn rewrite_params(parse_result: &mut ParseResult) -> Result, Error> { + let mut params = HashMap::new(); + + visit_and_mutate_nodes(parse_result, |node| -> Result, Error> { + if let Some(NodeEnum::ParamRef(ref mut param)) = node.node { + if let Some(existing) = params.get(¶m.number) { + param.number = *existing; + } else { + let number = params.len() as i32 + 1; + params.insert(param.number, number); + param.number = number; + } + } + + Ok(None) + })?; + + let mut params: Vec<(i32, i32)> = params.into_iter().collect(); + params.sort_by(|a, b| a.1.cmp(&b.1)); + + Ok(params + .into_iter() + .map(|(original, _)| original as u16) + .collect()) +} + +#[derive(Debug, Clone)] +pub(super) enum UpdateValue { + Value(Node), + Expr(String), // We deparse the expression because we can't handle it yet. +} + +/// # Example +/// +/// ```ignore +/// UPDATE sharded SET id = $1, email = $2 WHERE id = $3 AND user_id = $4 +/// ``` +/// +/// ```ignore +/// [ +/// ("id", (id, $1)), +/// ("email", (email, $2)) +/// ] +/// ``` +/// +/// This allows us to build a partial INSERT statement. +/// +fn res_targets_to_insert_res_targets( + stmt: &UpdateStmt, + query_parser_engine: QueryParserEngine, +) -> Result, Error> { + let mut result = HashMap::new(); + for target in &stmt.target_list { + if let Some(ref node) = target.node { + if let NodeEnum::ResTarget(ref target) = node { + let valid = target + .val + .as_ref() + .map(|value| Value::try_from(&value.node).is_ok()) + .unwrap_or_default(); + let value = if valid { + UpdateValue::Value(*target.val.clone().unwrap()) + } else { + UpdateValue::Expr(deparse_expr( + target.val.as_ref().unwrap(), + query_parser_engine, + )?) + }; + result.insert(target.name.clone(), value); + } + } + } + + Ok(result) +} + +/// Convert a ResTarget (from UPDATE SET clause) to an AExpr equality expression. +/// +/// Transforms `SET column = value` into `column = value` expression +/// for use in shard routing validation. +fn res_target_to_a_expr(res_target: &ResTarget) -> AExpr { + let column_ref = ColumnRef { + fields: vec![Node { + node: Some(NodeEnum::String(PgString { + sval: res_target.name.clone(), + })), + }], + location: res_target.location, + }; + + AExpr { + kind: AExprKind::AexprOp.into(), + name: vec![Node { + node: Some(NodeEnum::String(PgString { sval: "=".into() })), + }], + lexpr: Some(Box::new(Node { + node: Some(NodeEnum::ColumnRef(column_ref)), + })), + rexpr: res_target.val.clone(), + ..Default::default() + } +} + +fn select_star() -> Vec { + vec![Node { + node: Some(NodeEnum::ResTarget(Box::new(ResTarget { + name: "".into(), + val: Some(Box::new(Node { + node: Some(NodeEnum::ColumnRef(ColumnRef { + fields: vec![Node { + node: Some(NodeEnum::AStar(AStar {})), + }], + ..Default::default() + })), + })), + ..Default::default() + }))), + }] +} + +fn parse_result(node: NodeEnum) -> ParseResult { + ParseResult { + version: pg_query::PG_VERSION_NUM as i32, + stmts: vec![RawStmt { + stmt: Some(Box::new(Node { + node: Some(node), + ..Default::default() + })), + ..Default::default() + }], + ..Default::default() + } +} + +/// Deparse an expression node by wrapping it in a SELECT statement. +fn deparse_expr(node: &Node, query_parser_engine: QueryParserEngine) -> Result { + Ok(deparse_list( + &[Node { + node: Some(NodeEnum::ResTarget(Box::new(ResTarget { + val: Some(Box::new(node.clone())), + ..Default::default() + }))), + }], + query_parser_engine, + )? + .unwrap()) // SAFETY: we are not passing in an empty list. +} + +/// Deparse a list of expressions by wrapping them into a SELECT statement. +fn deparse_list( + list: &[Node], + query_parser_engine: QueryParserEngine, +) -> Result, Error> { + if list.is_empty() { + return Ok(None); + } + + let stmt = SelectStmt { + target_list: list.to_vec(), + limit_option: LimitOption::Default.into(), + op: SetOperation::SetopNone.into(), + ..Default::default() + }; + let result = parse_result(NodeEnum::SelectStmt(Box::new(stmt))); + let string = match query_parser_engine { + QueryParserEngine::PgQueryProtobuf => result.deparse()?, + QueryParserEngine::PgQueryRaw => result.deparse_raw()?, + } + .strip_prefix("SELECT ") + .unwrap_or_default() + .to_string(); + + Ok(Some(string)) +} + +fn create_stmts( + stmt: &UpdateStmt, + new_value: &ResTarget, + query_parser_engine: QueryParserEngine, +) -> Result { + let select = SelectStmt { + target_list: select_star(), + from_clause: vec![Node { + node: Some(NodeEnum::RangeVar(stmt.relation.clone().unwrap())), // SAFETY: we checked the UPDATE stmt has a table name. + }], + limit_option: LimitOption::Default.into(), + where_clause: stmt.where_clause.clone(), + op: SetOperation::SetopNone.into(), + ..Default::default() + }; + + let mut select = parse_result(NodeEnum::SelectStmt(Box::new(select))); + + let params = rewrite_params(&mut select)?; + let select = pg_query::ParseResult::new(select, "".into()); + + let select = Statement { + stmt: match query_parser_engine { + QueryParserEngine::PgQueryProtobuf => select.deparse()?, + QueryParserEngine::PgQueryRaw => select.deparse_raw()?, + }, + ast: Ast::from_parse_result(select), + params, + }; + + let delete = DeleteStmt { + relation: stmt.relation.clone(), + where_clause: stmt.where_clause.clone(), + ..Default::default() + }; + + let mut delete = parse_result(NodeEnum::DeleteStmt(Box::new(delete))); + + let params = rewrite_params(&mut delete)?; + + let delete = pg_query::ParseResult::new(delete, "".into()); + + let delete = Statement { + stmt: match query_parser_engine { + QueryParserEngine::PgQueryProtobuf => delete.deparse()?, + QueryParserEngine::PgQueryRaw => delete.deparse_raw()?, + }, + ast: Ast::from_parse_result(delete), + params, + }; + + let check = SelectStmt { + target_list: select_star(), + from_clause: vec![Node { + node: Some(NodeEnum::RangeVar(stmt.relation.clone().unwrap())), // SAFETY: we checked the UPDATE stmt has a table name. + }], + limit_option: LimitOption::Default.into(), + where_clause: Some(Box::new(Node { + node: Some(NodeEnum::AExpr(Box::new(res_target_to_a_expr(new_value)))), + })), + op: SetOperation::SetopNone.into(), + ..Default::default() + }; + + let mut check = parse_result(NodeEnum::SelectStmt(Box::new(check))); + let params = rewrite_params(&mut check)?; + let check = pg_query::ParseResult::new(check, "".into()); + + let check = Statement { + stmt: match query_parser_engine { + QueryParserEngine::PgQueryProtobuf => check.deparse()?, + QueryParserEngine::PgQueryRaw => check.deparse_raw()?, + }, + ast: Ast::from_parse_result(check), + params, + }; + + Ok(ShardingKeyUpdate { + inner: Arc::new(Inner { + select, + delete, + check, + insert: Insert { + table: stmt.relation.clone(), + mapping: res_targets_to_insert_res_targets(stmt, query_parser_engine)?, + returning_list: stmt.returning_list.clone(), + returnin_list_deparsed: deparse_list(&stmt.returning_list, query_parser_engine)?, + }, + }), + }) +} + +#[cfg(test)] +mod test { + use pg_query::parse; + use pgdog_config::{Rewrite, ShardedTable}; + + use crate::backend::{replication::ShardedSchemas, ShardedTables}; + use crate::net::messages::row_description::Field; + + use super::*; + + fn default_schema() -> ShardingSchema { + ShardingSchema { + shards: 2, + tables: ShardedTables::new( + vec![ShardedTable { + database: "pgdog".into(), + name: Some("sharded".into()), + column: "id".into(), + ..Default::default() + }], + vec![], + ), + schemas: ShardedSchemas::new(vec![]), + rewrite: Rewrite { + enabled: true, + shard_key: RewriteMode::Rewrite, + ..Default::default() + }, + ..Default::default() + } + } + + fn run_test(query: &str) -> Result, Error> { + let mut stmt = parse(query)?; + let schema = default_schema(); + let mut stmts = PreparedStatements::new(); + + let ctx = StatementRewriteContext { + stmt: &mut stmt.protobuf, + schema: &schema, + extended: true, + prepared: false, + prepared_statements: &mut stmts, + }; + let mut plan = RewritePlan::default(); + StatementRewrite::new(ctx).sharding_key_update(&mut plan)?; + Ok(plan.sharding_key_update) + } + + #[test] + fn test_select_basic_where_param() { + let result = run_test("UPDATE sharded SET id = $1 WHERE email = $2") + .unwrap() + .unwrap(); + + // SELECT should have WHERE clause with param renumbered to $1 + assert_eq!(result.select.stmt, "SELECT * FROM sharded WHERE email = $1"); + assert_eq!(result.select.params, vec![2]); + } + + #[test] + fn test_select_multiple_where_params() { + let result = run_test("UPDATE sharded SET id = $1 WHERE email = $2 AND name = $3") + .unwrap() + .unwrap(); + + assert_eq!( + result.select.stmt, + "SELECT * FROM sharded WHERE email = $1 AND name = $2" + ); + assert_eq!(result.select.params, vec![2, 3]); + assert!(!result.insert.is_returning()); + } + + #[test] + fn test_select_non_sequential_params() { + // Params in WHERE are $3 and $5, should be renumbered to $1 and $2 + let result = run_test( + "UPDATE sharded SET id = $1, value = $2, other = $4 WHERE email = $3 AND name = $5", + ) + .unwrap() + .unwrap(); + + assert_eq!( + result.select.stmt, + "SELECT * FROM sharded WHERE email = $1 AND name = $2" + ); + assert_eq!(result.select.params, vec![3, 5]); + } + + #[test] + fn test_select_single_where_param() { + let result = run_test("UPDATE sharded SET id = $1 WHERE email = $2") + .unwrap() + .unwrap(); + + assert_eq!(result.select.stmt, "SELECT * FROM sharded WHERE email = $1"); + assert_eq!(result.select.params, vec![2]); + } + + #[test] + fn test_delete_basic() { + let result = run_test("UPDATE sharded SET id = $1 WHERE email = $2") + .unwrap() + .unwrap(); + + assert_eq!(result.delete.stmt, "DELETE FROM sharded WHERE email = $1"); + } + + #[test] + fn test_delete_multiple_where_params() { + let result = run_test("UPDATE sharded SET id = $1 WHERE email = $2 AND name = $3") + .unwrap() + .unwrap(); + + assert_eq!( + result.delete.stmt, + "DELETE FROM sharded WHERE email = $1 AND name = $2" + ); + } + + #[test] + fn test_no_params_in_where() { + let result = run_test("UPDATE sharded SET id = $1 WHERE email = 'test@example.com'") + .unwrap() + .unwrap(); + + assert_eq!( + result.select.stmt, + "SELECT * FROM sharded WHERE email = 'test@example.com'" + ); + assert_eq!(result.select.params, Vec::::new()); + } + + #[test] + fn test_where_with_in_clause() { + let result = run_test("UPDATE sharded SET id = $1 WHERE email IN ($2, $3, $4)") + .unwrap() + .unwrap(); + + assert_eq!( + result.select.stmt, + "SELECT * FROM sharded WHERE email IN ($1, $2, $3)" + ); + assert_eq!(result.select.params, vec![2, 3, 4]); + } + + #[test] + fn test_where_with_comparison_operators() { + let result = run_test("UPDATE sharded SET id = $1 WHERE count > $2 AND count < $3") + .unwrap() + .unwrap(); + + assert_eq!( + result.select.stmt, + "SELECT * FROM sharded WHERE count > $1 AND count < $2" + ); + assert_eq!(result.select.params, vec![2, 3]); + } + + #[test] + fn test_where_with_or_condition() { + let result = run_test("UPDATE sharded SET id = $1 WHERE email = $2 OR name = $3") + .unwrap() + .unwrap(); + + assert_eq!( + result.select.stmt, + "SELECT * FROM sharded WHERE email = $1 OR name = $2" + ); + assert_eq!(result.select.params, vec![2, 3]); + } + + #[test] + fn test_high_param_numbers() { + let result = run_test("UPDATE sharded SET id = $10 WHERE email = $20 AND name = $30") + .unwrap() + .unwrap(); + + assert_eq!( + result.select.stmt, + "SELECT * FROM sharded WHERE email = $1 AND name = $2" + ); + assert_eq!(result.select.params, vec![20, 30]); + } + + #[test] + fn test_non_sharding_key_update_returns_none() { + // Updating a non-sharding column should return None + let result = run_test("UPDATE sharded SET email = $1 WHERE id = $2").unwrap(); + assert!(result.is_none()); + } + + #[test] + fn test_where_with_like() { + let result = run_test("UPDATE sharded SET id = $1 WHERE email LIKE $2") + .unwrap() + .unwrap(); + + assert_eq!( + result.select.stmt, + "SELECT * FROM sharded WHERE email LIKE $1" + ); + assert_eq!(result.select.params, vec![2]); + } + + #[test] + fn test_where_with_is_null() { + let result = run_test("UPDATE sharded SET id = $1 WHERE email = $2 AND deleted_at IS NULL") + .unwrap() + .unwrap(); + + assert_eq!( + result.select.stmt, + "SELECT * FROM sharded WHERE email = $1 AND deleted_at IS NULL" + ); + assert_eq!(result.select.params, vec![2]); + } + + #[test] + fn test_where_with_between() { + let result = run_test("UPDATE sharded SET id = $1 WHERE created_at BETWEEN $2 AND $3") + .unwrap() + .unwrap(); + + assert_eq!( + result.select.stmt, + "SELECT * FROM sharded WHERE created_at BETWEEN $1 AND $2" + ); + assert_eq!(result.select.params, vec![2, 3]); + } + + #[test] + fn test_same_param_used_twice() { + // Same parameter $2 used twice in WHERE clause + let result = run_test("UPDATE sharded SET id = $1 WHERE email = $2 OR name = $2") + .unwrap() + .unwrap(); + + // Both occurrences should be renumbered to $1 + assert_eq!( + result.select.stmt, + "SELECT * FROM sharded WHERE email = $1 OR name = $1" + ); + // Only one unique param in the mapping + assert_eq!(result.select.params, vec![2]); + } + + #[test] + fn test_same_param_used_multiple_times() { + // $2 used three times + let result = run_test("UPDATE sharded SET id = $1 WHERE a = $2 AND b = $2 AND c = $2") + .unwrap() + .unwrap(); + + assert_eq!( + result.select.stmt, + "SELECT * FROM sharded WHERE a = $1 AND b = $1 AND c = $1" + ); + assert_eq!(result.select.params, vec![2]); + } + + #[test] + fn test_mixed_repeated_and_unique_params() { + // $2 used twice, $3 used once + let result = run_test("UPDATE sharded SET id = $1 WHERE a = $2 AND b = $3 AND c = $2") + .unwrap() + .unwrap(); + + assert_eq!( + result.select.stmt, + "SELECT * FROM sharded WHERE a = $1 AND b = $2 AND c = $1" + ); + assert_eq!(result.select.params, vec![2, 3]); + } + + #[test] + fn test_repeated_params_in_in_clause() { + // Same param repeated in IN clause (unusual but valid) + let result = run_test("UPDATE sharded SET id = $1 WHERE email IN ($2, $3, $2)") + .unwrap() + .unwrap(); + + assert_eq!( + result.select.stmt, + "SELECT * FROM sharded WHERE email IN ($1, $2, $1)" + ); + assert_eq!(result.select.params, vec![2, 3]); + } + + #[test] + fn test_delete_with_repeated_params() { + let result = run_test("UPDATE sharded SET id = $1 WHERE email = $2 OR name = $2") + .unwrap() + .unwrap(); + + assert_eq!( + result.delete.stmt, + "DELETE FROM sharded WHERE email = $1 OR name = $1" + ); + assert_eq!(result.delete.params, vec![2]); + } + + #[test] + fn test_sharding_key_not_changed() { + let result = run_test("UPDATE sharded SET id = $1 WHERE id = $1 AND email = $2") + .unwrap() + .unwrap(); + assert_eq!(result.check.stmt, "SELECT * FROM sharded WHERE id = $1"); + assert_eq!(result.check.params, vec![1]); + } + + #[test] + fn test_unsupported_assignment() { + let result = run_test("UPDATE sharded SET id = random() WHERE id = $1"); + assert!(matches!( + result, + Err(Error::UnsupportedShardingKeyUpdate(msg)) if msg == "\"id\" = random()" + )); + } + + #[test] + fn test_unsupported_assignment_arithmetic_add() { + let result = run_test("UPDATE sharded SET id = id + 1 WHERE id = $1"); + assert!(matches!( + result, + Err(Error::UnsupportedShardingKeyUpdate(msg)) if msg == "\"id\" = id + 1" + )); + } + + #[test] + fn test_unsupported_assignment_arithmetic_multiply() { + let result = run_test("UPDATE sharded SET id = id * 2 WHERE id = $1"); + assert!(matches!( + result, + Err(Error::UnsupportedShardingKeyUpdate(msg)) if msg == "\"id\" = id * 2" + )); + } + + #[test] + fn test_unsupported_assignment_arithmetic_with_param() { + let result = run_test("UPDATE sharded SET id = id + $2 WHERE id = $1"); + assert!(matches!( + result, + Err(Error::UnsupportedShardingKeyUpdate(msg)) if msg == "\"id\" = id + $2" + )); + } + + #[test] + fn test_unsupported_assignment_now() { + let result = run_test("UPDATE sharded SET id = now() WHERE id = $1"); + assert!(matches!( + result, + Err(Error::UnsupportedShardingKeyUpdate(msg)) if msg == "\"id\" = now()" + )); + } + + #[test] + fn test_unsupported_assignment_coalesce() { + let result = run_test("UPDATE sharded SET id = coalesce(id, 0) WHERE id = $1"); + assert!(matches!( + result, + Err(Error::UnsupportedShardingKeyUpdate(msg)) if msg == "\"id\" = COALESCE(id, 0)" + )); + } + + #[test] + fn test_unsupported_assignment_case() { + let result = + run_test("UPDATE sharded SET id = CASE WHEN id > 0 THEN 1 ELSE 0 END WHERE id = $1"); + assert!(matches!( + result, + Err(Error::UnsupportedShardingKeyUpdate(msg)) if msg == "\"id\" = CASE WHEN id > 0 THEN 1 ELSE 0 END" + )); + } + + #[test] + fn test_unsupported_assignment_subquery() { + let result = + run_test("UPDATE sharded SET id = (SELECT max(id) FROM sharded) WHERE id = $1"); + assert!(matches!( + result, + Err(Error::UnsupportedShardingKeyUpdate(msg)) if msg == "\"id\" = (SELECT max(id) FROM sharded)" + )); + } + + #[test] + fn test_unsupported_assignment_column_reference() { + let result = run_test("UPDATE sharded SET id = other_column WHERE id = $1"); + assert!(matches!( + result, + Err(Error::UnsupportedShardingKeyUpdate(msg)) if msg == "\"id\" = other_column" + )); + } + + #[test] + fn test_unsupported_assignment_concat() { + let result = run_test("UPDATE sharded SET id = id || '_suffix' WHERE id = $1"); + assert!(matches!( + result, + Err(Error::UnsupportedShardingKeyUpdate(msg)) if msg == "\"id\" = id || '_suffix'" + )); + } + + #[test] + fn test_unsupported_assignment_negation() { + let result = run_test("UPDATE sharded SET id = -id WHERE id = $1"); + assert!(matches!( + result, + Err(Error::UnsupportedShardingKeyUpdate(msg)) if msg == "\"id\" = - id" + )); + } + + #[test] + fn test_return_rows() { + let result = run_test("UPDATE sharded SET id = $1 WHERE id = $2 RETURNING *") + .unwrap() + .unwrap(); + assert_eq!(result.insert.returnin_list_deparsed, Some("*".into())); + + let result = + run_test("UPDATE sharded SET id = $1 WHERE id = $2 RETURNING id, email, random()") + .unwrap() + .unwrap(); + assert_eq!( + result.insert.returnin_list_deparsed, + Some("id, email, random()".into()) + ); + } + + #[test] + fn test_res_targets_to_insert_res_targets_expr_branch() { + // Test that expression assignments (non-simple values) are deparsed correctly + // and stored as UpdateValue::Expr in the insert mapping. + let result = run_test("UPDATE sharded SET id = $1, email = random() WHERE id = $2") + .unwrap() + .unwrap(); + + // The id column should be UpdateValue::Value (simple parameter) + let id_value = result.insert.mapping.get("id").unwrap(); + assert!(matches!(id_value, UpdateValue::Value(_))); + + // The email column should be UpdateValue::Expr with the deparsed expression + let email_value = result.insert.mapping.get("email").unwrap(); + match email_value { + UpdateValue::Expr(expr) => assert_eq!(expr, "random()"), + _ => panic!("Expected UpdateValue::Expr for email"), + } + } + + #[test] + fn test_res_targets_to_insert_res_targets_expr_arithmetic() { + // Test arithmetic expressions are deparsed correctly + let result = run_test("UPDATE sharded SET id = $1, counter = counter + 1 WHERE id = $2") + .unwrap() + .unwrap(); + + let counter_value = result.insert.mapping.get("counter").unwrap(); + match counter_value { + UpdateValue::Expr(expr) => assert_eq!(expr, "counter + 1"), + _ => panic!("Expected UpdateValue::Expr for counter"), + } + } + + #[test] + fn test_res_targets_to_insert_res_targets_expr_coalesce() { + // Test COALESCE expressions are deparsed correctly + let result = + run_test("UPDATE sharded SET id = $1, name = COALESCE(name, 'default') WHERE id = $2") + .unwrap() + .unwrap(); + + let name_value = result.insert.mapping.get("name").unwrap(); + match name_value { + UpdateValue::Expr(expr) => assert_eq!(expr, "COALESCE(name, 'default')"), + _ => panic!("Expected UpdateValue::Expr for name"), + } + } + + #[test] + fn test_insert_build_request_with_expr_column() { + // Test that INSERT statement is built correctly when there are expression columns. + // The expression should appear directly in the VALUES clause. + // Use literal values (not placeholders) to avoid needing bind parameters. + let result = run_test("UPDATE sharded SET id = 42, email = random() WHERE id = 1") + .unwrap() + .unwrap(); + + // Create a mock row description matching the SELECT * result + let row_description = RowDescription::new(&[ + Field::bigint("id"), + Field::text("email"), + Field::text("other_col"), + ]); + + // Create a mock data row with values for columns not in the UPDATE SET clause + let mut data_row = DataRow::new(); + data_row.add("1"); // id - will be overwritten by mapping + data_row.add("old@example.com"); // email - will be overwritten by mapping + data_row.add("other_value"); // other_col - from existing row + + // Create a simple query request (not prepared statement) + let request = ClientRequest::from(vec![ProtocolMessage::from(Query::new( + "UPDATE sharded SET id = 42, email = random() WHERE id = 1", + ))]); + + let insert_request = result + .insert + .build_request(&request, &row_description, &data_row) + .unwrap(); + + // Get the query from the request to verify the INSERT statement + let query = insert_request.query().unwrap().unwrap(); + let stmt = query.query(); + + // The INSERT should contain the expression random() directly in VALUES + assert!( + stmt.contains("random()"), + "INSERT statement should contain the expression: {}", + stmt + ); + // Verify it's an INSERT statement + assert!( + stmt.starts_with("INSERT INTO"), + "Should be an INSERT statement: {}", + stmt + ); + // Verify parameter numbering is correct: $1 for id, random() for email, $2 for other_col + // (not $3, which would be wrong if we used row index instead of bind param index) + assert!( + stmt.contains("$1") && stmt.contains("$2") && !stmt.contains("$3"), + "Parameter numbering should be sequential without gaps: {}", + stmt + ); + } +} diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/visitor.rs b/pgdog/src/frontend/router/parser/rewrite/statement/visitor.rs new file mode 100644 index 000000000..55087e515 --- /dev/null +++ b/pgdog/src/frontend/router/parser/rewrite/statement/visitor.rs @@ -0,0 +1,306 @@ +//! AST visitor utilities for statement rewriting. + +use pg_query::protobuf::ParseResult; +use pg_query::{Node, NodeEnum}; + +/// Count the maximum parameter number ($1, $2, etc.) in the parse result. +pub fn count_params(ast: &mut ParseResult) -> u16 { + let mut max_param = 0i32; + let _: Result<(), std::convert::Infallible> = visit_and_mutate_nodes(ast, |node| { + if let Some(NodeEnum::ParamRef(param)) = &node.node { + max_param = max_param.max(param.number); + } + Ok(None) + }); + max_param.max(0) as u16 +} + +/// Recursively visit and potentially mutate all nodes in the AST. +/// The callback returns Ok(Some(new_node)) to replace, Ok(None) to keep, or Err to abort. +pub(super) fn visit_and_mutate_nodes(ast: &mut ParseResult, mut callback: F) -> Result<(), E> +where + F: FnMut(&mut Node) -> Result, E>, +{ + for stmt in &mut ast.stmts { + if let Some(ref mut node) = stmt.stmt { + visit_and_mutate_node(node, &mut callback)?; + } + } + Ok(()) +} + +pub(super) fn visit_and_mutate_node(node: &mut Node, callback: &mut F) -> Result<(), E> +where + F: FnMut(&mut Node) -> Result, E>, +{ + // Try to replace this node + if let Some(replacement) = callback(node)? { + *node = replacement; + return Ok(()); + } + + // Otherwise, recurse into children + let Some(inner) = &mut node.node else { + return Ok(()); + }; + + visit_and_mutate_children(inner, callback) +} + +pub(super) fn visit_and_mutate_children( + node: &mut NodeEnum, + callback: &mut F, +) -> Result<(), E> +where + F: FnMut(&mut Node) -> Result, E>, +{ + match node { + NodeEnum::SelectStmt(stmt) => { + for target in &mut stmt.target_list { + visit_and_mutate_node(target, callback)?; + } + for from in &mut stmt.from_clause { + visit_and_mutate_node(from, callback)?; + } + if let Some(where_clause) = &mut stmt.where_clause { + visit_and_mutate_node(where_clause, callback)?; + } + if let Some(having) = &mut stmt.having_clause { + visit_and_mutate_node(having, callback)?; + } + for group in &mut stmt.group_clause { + visit_and_mutate_node(group, callback)?; + } + for order in &mut stmt.sort_clause { + visit_and_mutate_node(order, callback)?; + } + if let Some(limit) = &mut stmt.limit_count { + visit_and_mutate_node(limit, callback)?; + } + if let Some(offset) = &mut stmt.limit_offset { + visit_and_mutate_node(offset, callback)?; + } + for cte in stmt.with_clause.iter_mut().flat_map(|w| &mut w.ctes) { + visit_and_mutate_node(cte, callback)?; + } + for values in &mut stmt.values_lists { + visit_and_mutate_node(values, callback)?; + } + } + + NodeEnum::InsertStmt(stmt) => { + if let Some(select) = &mut stmt.select_stmt { + visit_and_mutate_node(select, callback)?; + } + for returning in &mut stmt.returning_list { + visit_and_mutate_node(returning, callback)?; + } + for cte in stmt.with_clause.iter_mut().flat_map(|w| &mut w.ctes) { + visit_and_mutate_node(cte, callback)?; + } + } + + NodeEnum::UpdateStmt(stmt) => { + for target in &mut stmt.target_list { + visit_and_mutate_node(target, callback)?; + } + if let Some(where_clause) = &mut stmt.where_clause { + visit_and_mutate_node(where_clause, callback)?; + } + for from in &mut stmt.from_clause { + visit_and_mutate_node(from, callback)?; + } + for returning in &mut stmt.returning_list { + visit_and_mutate_node(returning, callback)?; + } + for cte in stmt.with_clause.iter_mut().flat_map(|w| &mut w.ctes) { + visit_and_mutate_node(cte, callback)?; + } + } + + NodeEnum::DeleteStmt(stmt) => { + if let Some(where_clause) = &mut stmt.where_clause { + visit_and_mutate_node(where_clause, callback)?; + } + for using in &mut stmt.using_clause { + visit_and_mutate_node(using, callback)?; + } + for returning in &mut stmt.returning_list { + visit_and_mutate_node(returning, callback)?; + } + for cte in stmt.with_clause.iter_mut().flat_map(|w| &mut w.ctes) { + visit_and_mutate_node(cte, callback)?; + } + } + + NodeEnum::ResTarget(res) => { + if let Some(val) = &mut res.val { + visit_and_mutate_node(val, callback)?; + } + } + + NodeEnum::AExpr(expr) => { + if let Some(lexpr) = &mut expr.lexpr { + visit_and_mutate_node(lexpr, callback)?; + } + if let Some(rexpr) = &mut expr.rexpr { + visit_and_mutate_node(rexpr, callback)?; + } + } + + NodeEnum::FuncCall(func) => { + for arg in &mut func.args { + visit_and_mutate_node(arg, callback)?; + } + } + + NodeEnum::TypeCast(cast) => { + if let Some(arg) = &mut cast.arg { + visit_and_mutate_node(arg, callback)?; + } + } + + NodeEnum::SubLink(sub) => { + if let Some(subselect) = &mut sub.subselect { + visit_and_mutate_node(subselect, callback)?; + } + if let Some(testexpr) = &mut sub.testexpr { + visit_and_mutate_node(testexpr, callback)?; + } + } + + NodeEnum::BoolExpr(bool_expr) => { + for arg in &mut bool_expr.args { + visit_and_mutate_node(arg, callback)?; + } + } + + NodeEnum::RangeSubselect(range) => { + if let Some(subquery) = &mut range.subquery { + visit_and_mutate_node(subquery, callback)?; + } + } + + NodeEnum::JoinExpr(join) => { + if let Some(larg) = &mut join.larg { + visit_and_mutate_node(larg, callback)?; + } + if let Some(rarg) = &mut join.rarg { + visit_and_mutate_node(rarg, callback)?; + } + if let Some(quals) = &mut join.quals { + visit_and_mutate_node(quals, callback)?; + } + } + + NodeEnum::CommonTableExpr(cte) => { + if let Some(query) = &mut cte.ctequery { + visit_and_mutate_node(query, callback)?; + } + } + + NodeEnum::List(list) => { + for item in &mut list.items { + visit_and_mutate_node(item, callback)?; + } + } + + NodeEnum::SortBy(sort) => { + if let Some(node) = &mut sort.node { + visit_and_mutate_node(node, callback)?; + } + } + + NodeEnum::CoalesceExpr(coalesce) => { + for arg in &mut coalesce.args { + visit_and_mutate_node(arg, callback)?; + } + } + + NodeEnum::CaseExpr(case) => { + if let Some(arg) = &mut case.arg { + visit_and_mutate_node(arg, callback)?; + } + for when in &mut case.args { + visit_and_mutate_node(when, callback)?; + } + if let Some(defresult) = &mut case.defresult { + visit_and_mutate_node(defresult, callback)?; + } + } + + NodeEnum::CaseWhen(when) => { + if let Some(expr) = &mut when.expr { + visit_and_mutate_node(expr, callback)?; + } + if let Some(result) = &mut when.result { + visit_and_mutate_node(result, callback)?; + } + } + + NodeEnum::NullTest(test) => { + if let Some(arg) = &mut test.arg { + visit_and_mutate_node(arg, callback)?; + } + } + + NodeEnum::RowExpr(row) => { + for arg in &mut row.args { + visit_and_mutate_node(arg, callback)?; + } + } + + NodeEnum::ArrayExpr(arr) => { + for elem in &mut arr.elements { + visit_and_mutate_node(elem, callback)?; + } + } + + _ => (), + } + + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_count_params_none() { + let mut ast = pg_query::parse("SELECT 1").unwrap(); + assert_eq!(count_params(&mut ast.protobuf), 0); + } + + #[test] + fn test_count_params_single() { + let mut ast = pg_query::parse("SELECT $1").unwrap(); + assert_eq!(count_params(&mut ast.protobuf), 1); + } + + #[test] + fn test_count_params_multiple() { + let mut ast = pg_query::parse("SELECT $1, $2, $3").unwrap(); + assert_eq!(count_params(&mut ast.protobuf), 3); + } + + #[test] + fn test_count_params_out_of_order() { + let mut ast = pg_query::parse("SELECT $3, $1, $5").unwrap(); + assert_eq!(count_params(&mut ast.protobuf), 5); + } + + #[test] + fn test_count_params_in_where() { + let mut ast = pg_query::parse("SELECT * FROM t WHERE id = $1 AND name = $2").unwrap(); + assert_eq!(count_params(&mut ast.protobuf), 2); + } + + #[test] + fn test_count_params_in_subquery() { + let mut ast = + pg_query::parse("SELECT * FROM t WHERE id IN (SELECT id FROM s WHERE val = $1)") + .unwrap(); + assert_eq!(count_params(&mut ast.protobuf), 1); + } +} diff --git a/pgdog/src/frontend/router/parser/rewrite_engine.rs b/pgdog/src/frontend/router/parser/rewrite_engine.rs deleted file mode 100644 index ccbbd45b4..000000000 --- a/pgdog/src/frontend/router/parser/rewrite_engine.rs +++ /dev/null @@ -1,328 +0,0 @@ -use super::{Aggregate, AggregateFunction, HelperMapping, RewriteOutput, RewritePlan}; -use pg_query::protobuf::{FuncCall, Node, ResTarget, String as PgString}; -use pg_query::{NodeEnum, ParseResult}; -use tracing::debug; - -/// Query rewrite engine. Currently supports injecting helper aggregates for AVG. -#[derive(Default)] -pub struct RewriteEngine; - -impl RewriteEngine { - pub fn new() -> Self { - Self::default() - } - - /// Rewrite a SELECT query, adding helper aggregates when necessary. - pub fn rewrite_select(&self, sql: &str, aggregate: &Aggregate) -> RewriteOutput { - match pg_query::parse(sql) { - Ok(parsed) => self.rewrite_parsed(parsed, aggregate, sql), - Err(err) => { - debug!("rewrite failed to parse SELECT: {}", err); - RewriteOutput::new(sql.to_string(), RewritePlan::new()) - } - } - } - - fn rewrite_parsed( - &self, - parsed: ParseResult, - aggregate: &Aggregate, - original_sql: &str, - ) -> RewriteOutput { - let mut ast = parsed.protobuf.clone(); - let Some(raw_stmt) = ast.stmts.first_mut() else { - return RewriteOutput::new(original_sql.to_string(), RewritePlan::new()); - }; - - let Some(stmt) = raw_stmt.stmt.as_mut() else { - return RewriteOutput::new(original_sql.to_string(), RewritePlan::new()); - }; - - let Some(NodeEnum::SelectStmt(select)) = stmt.node.as_mut() else { - return RewriteOutput::new(original_sql.to_string(), RewritePlan::new()); - }; - - let mut plan = RewritePlan::new(); - let mut modified = false; - - for target in aggregate - .targets() - .iter() - .filter(|t| matches!(t.function(), AggregateFunction::Avg)) - { - if aggregate.targets().iter().any(|other| { - matches!(other.function(), AggregateFunction::Count) - && other.expr_id() == target.expr_id() - && other.is_distinct() == target.is_distinct() - }) { - continue; - } - - let Some(node) = select.target_list.get(target.column()) else { - continue; - }; - let Some(NodeEnum::ResTarget(res_target)) = node.node.as_ref() else { - continue; - }; - let Some(original_value) = res_target.val.as_ref() else { - continue; - }; - let Some(func_call) = Self::extract_func_call(original_value) else { - continue; - }; - - let helper_index = select.target_list.len(); - let helper_alias = format!("__pgdog_count_{}", helper_index); - - let helper_func = Self::build_count_func(func_call, target.is_distinct()); - let helper_res = ResTarget { - name: helper_alias, - indirection: vec![], - val: Some(Box::new(Node { - node: Some(NodeEnum::FuncCall(Box::new(helper_func))), - })), - location: res_target.location, - }; - - select.target_list.push(Node { - node: Some(NodeEnum::ResTarget(Box::new(helper_res))), - }); - - plan.add_drop_column(helper_index); - plan.add_helper(HelperMapping { - avg_column: target.column(), - helper_column: helper_index, - expr_id: target.expr_id(), - distinct: target.is_distinct(), - }); - - modified = true; - } - - if !modified { - return RewriteOutput::new(original_sql.to_string(), RewritePlan::new()); - } - - let rewritten_sql = ast.deparse().unwrap_or_else(|_| original_sql.to_string()); - RewriteOutput::new(rewritten_sql, plan) - } - - fn extract_func_call(node: &Node) -> Option<&FuncCall> { - match node.node.as_ref()? { - NodeEnum::FuncCall(func) => Some(func), - NodeEnum::TypeCast(cast) => cast - .arg - .as_deref() - .and_then(|inner| Self::extract_func_call(inner)), - NodeEnum::CollateClause(collate) => collate - .arg - .as_deref() - .and_then(|inner| Self::extract_func_call(inner)), - NodeEnum::CoerceToDomain(coerce) => coerce - .arg - .as_deref() - .and_then(|inner| Self::extract_func_call(inner)), - NodeEnum::ResTarget(res) => res - .val - .as_deref() - .and_then(|inner| Self::extract_func_call(inner)), - _ => None, - } - } - - fn build_count_func(original: &FuncCall, distinct: bool) -> FuncCall { - FuncCall { - funcname: vec![Node { - node: Some(NodeEnum::String(PgString { - sval: "count".into(), - })), - }], - args: original.args.clone(), - agg_order: original.agg_order.clone(), - agg_filter: original.agg_filter.clone(), - over: original.over.clone(), - agg_within_group: original.agg_within_group, - agg_star: original.agg_star, - agg_distinct: distinct, - func_variadic: original.func_variadic, - funcformat: original.funcformat, - location: original.location, - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::frontend::router::parser::aggregate::Aggregate; - - fn rewrite(sql: &str) -> RewriteOutput { - let ast = pg_query::parse(sql).unwrap(); - let stmt = match ast - .protobuf - .stmts - .first() - .and_then(|stmt| stmt.stmt.as_ref()) - { - Some(raw) => match raw.node.as_ref().unwrap() { - NodeEnum::SelectStmt(select) => select, - _ => panic!("not a select"), - }, - None => panic!("empty"), - }; - let aggregate = Aggregate::parse(stmt).unwrap(); - RewriteEngine::new().rewrite_select(sql, &aggregate) - } - - #[test] - fn rewrite_engine_noop() { - let output = rewrite("SELECT COUNT(price) FROM menu"); - assert!(output.plan.is_noop()); - } - - #[test] - fn rewrite_engine_adds_helper() { - let output = rewrite("SELECT AVG(price) FROM menu"); - assert!(!output.plan.is_noop()); - assert_eq!(output.plan.drop_columns(), &[1]); - assert_eq!(output.plan.helpers().len(), 1); - let helper = &output.plan.helpers()[0]; - assert_eq!(helper.avg_column, 0); - assert_eq!(helper.helper_column, 1); - assert_eq!(helper.expr_id, 0); - assert!(!helper.distinct); - - let parsed = pg_query::parse(&output.sql).unwrap(); - let stmt = match parsed - .protobuf - .stmts - .first() - .and_then(|stmt| stmt.stmt.as_ref()) - { - Some(raw) => match raw.node.as_ref().unwrap() { - NodeEnum::SelectStmt(select) => select, - _ => panic!("not select"), - }, - None => panic!("empty"), - }; - let aggregate = Aggregate::parse(stmt).unwrap(); - assert_eq!(aggregate.targets().len(), 2); - assert!(aggregate - .targets() - .iter() - .any(|target| matches!(target.function(), AggregateFunction::Count))); - } - - #[test] - fn rewrite_engine_handles_avg_with_casts() { - let output = rewrite("SELECT AVG(amount)::numeric::bigint::real FROM sharded"); - assert_eq!(output.plan.drop_columns(), &[1]); - assert_eq!(output.plan.helpers().len(), 1); - let helper = &output.plan.helpers()[0]; - assert_eq!(helper.avg_column, 0); - assert_eq!(helper.helper_column, 1); - - let parsed = pg_query::parse(&output.sql).unwrap(); - let stmt = match parsed - .protobuf - .stmts - .first() - .and_then(|stmt| stmt.stmt.as_ref()) - { - Some(raw) => match raw.node.as_ref().unwrap() { - NodeEnum::SelectStmt(select) => select, - _ => panic!("not select"), - }, - None => panic!("empty"), - }; - let aggregate = Aggregate::parse(stmt).unwrap(); - assert_eq!(aggregate.targets().len(), 2); - assert!(aggregate - .targets() - .iter() - .any(|target| matches!(target.function(), AggregateFunction::Count))); - } - - #[test] - fn rewrite_engine_skips_when_count_exists() { - let sql = "SELECT COUNT(price), AVG(price) FROM menu"; - let output = rewrite(sql); - assert!(output.plan.is_noop()); - assert_eq!(output.sql, sql); - } - - #[test] - fn rewrite_engine_handles_mismatched_pair() { - let output = rewrite("SELECT COUNT(price::numeric), AVG(price) FROM menu"); - assert_eq!(output.plan.drop_columns(), &[2]); - assert_eq!(output.plan.helpers().len(), 1); - let helper = &output.plan.helpers()[0]; - assert_eq!(helper.avg_column, 1); - assert_eq!(helper.helper_column, 2); - assert!(!helper.distinct); - - // Ensure the rewritten SQL now contains both AVG and helper COUNT for the AVG target. - let parsed = pg_query::parse(&output.sql).unwrap(); - let stmt = match parsed - .protobuf - .stmts - .first() - .and_then(|stmt| stmt.stmt.as_ref()) - { - Some(raw) => match raw.node.as_ref().unwrap() { - NodeEnum::SelectStmt(select) => select, - _ => panic!("not select"), - }, - None => panic!("empty"), - }; - let aggregate = Aggregate::parse(stmt).unwrap(); - assert_eq!(aggregate.targets().len(), 3); - assert!( - aggregate - .targets() - .iter() - .filter(|target| matches!(target.function(), AggregateFunction::Count)) - .count() - >= 2 - ); - } - - #[test] - fn rewrite_engine_multiple_avg_helpers() { - let output = rewrite("SELECT AVG(price), AVG(discount) FROM menu"); - assert_eq!(output.plan.drop_columns(), &[2, 3]); - assert_eq!(output.plan.helpers().len(), 2); - - let helper_price = &output.plan.helpers()[0]; - assert_eq!(helper_price.avg_column, 0); - assert_eq!(helper_price.helper_column, 2); - - let helper_discount = &output.plan.helpers()[1]; - assert_eq!(helper_discount.avg_column, 1); - assert_eq!(helper_discount.helper_column, 3); - - let parsed = pg_query::parse(&output.sql).unwrap(); - let stmt = match parsed - .protobuf - .stmts - .first() - .and_then(|stmt| stmt.stmt.as_ref()) - { - Some(raw) => match raw.node.as_ref().unwrap() { - NodeEnum::SelectStmt(select) => select, - _ => panic!("not select"), - }, - None => panic!("empty"), - }; - let aggregate = Aggregate::parse(stmt).unwrap(); - assert_eq!(aggregate.targets().len(), 4); - assert_eq!( - aggregate - .targets() - .iter() - .filter(|target| matches!(target.function(), AggregateFunction::Count)) - .count(), - 2 - ); - } -} diff --git a/pgdog/src/frontend/router/parser/route.rs b/pgdog/src/frontend/router/parser/route.rs index ceea4514f..0a3673ca2 100644 --- a/pgdog/src/frontend/router/parser/route.rs +++ b/pgdog/src/frontend/router/parser/route.rs @@ -1,13 +1,20 @@ -use std::fmt::Display; +use std::{fmt::Display, ops::Deref}; + +use lazy_static::lazy_static; use super::{ - Aggregate, DistinctBy, FunctionBehavior, Limit, LockingBehavior, OrderBy, RewritePlan, + explain_trace::ExplainTrace, rewrite::statement::aggregate::AggregateRewritePlan, Aggregate, + DistinctBy, FunctionBehavior, Limit, LockingBehavior, OrderBy, }; +/// The shard destination for a statement. #[derive(Debug, Clone, PartialEq, PartialOrd, Ord, Eq, Hash, Default)] pub enum Shard { + /// Direct-to-shard number. Direct(usize), + /// Multiple shards, enumerated. Multi(Vec), + /// All shards. #[default] All, } @@ -27,13 +34,20 @@ impl Display for Shard { } impl Shard { - pub fn all(&self) -> bool { + /// Returns true if this is an all-shard query. + pub fn is_all(&self) -> bool { matches!(self, Shard::All) } - pub fn direct(shard: usize) -> Self { + /// Create new direct-to-shard mapping. + pub fn new_direct(shard: usize) -> Self { Self::Direct(shard) } + + /// Returns true if this is a direct-to-shard mapping. + pub fn is_direct(&self) -> bool { + matches!(self, Self::Direct(_)) + } } impl From> for Shard { @@ -46,11 +60,23 @@ impl From> for Shard { } } +impl From for Shard { + fn from(value: usize) -> Self { + Shard::Direct(value) + } +} + +impl From> for Shard { + fn from(value: Vec) -> Self { + Shard::Multi(value) + } +} + /// Path a query should take and any transformations -/// that should be applied along the way. -#[derive(Debug, Clone, Default, PartialEq)] +/// that should be applied to the response. +#[derive(Debug, Clone, Default, PartialEq, derive_builder::Builder)] pub struct Route { - shard: Shard, + shard: ShardWithPriority, read: bool, order_by: Vec, aggregate: Aggregate, @@ -58,8 +84,11 @@ pub struct Route { lock_session: bool, distinct: Option, maintenance: bool, - rewrite_plan: RewritePlan, + rewrite_plan: AggregateRewritePlan, rewritten_sql: Option, + explain: Option, + rollback_savepoint: bool, + search_path_driven: bool, } impl Display for Route { @@ -67,16 +96,16 @@ impl Display for Route { write!( f, "shard={}, role={}", - self.shard, + self.shard.deref(), if self.read { "replica" } else { "primary" } ) } } impl Route { - /// SELECT query. + /// Create new route for a `SELECT` query. pub fn select( - shard: Shard, + shard: ShardWithPriority, order_by: Vec, aggregate: Aggregate, limit: Limit, @@ -94,26 +123,30 @@ impl Route { } /// A query that should go to a replica. - pub fn read(shard: impl Into) -> Self { + pub fn read(shard: ShardWithPriority) -> Self { Self { - shard: shard.into(), + shard, read: true, ..Default::default() } } /// A write query. - pub fn write(shard: impl Into) -> Self { + pub fn write(shard: ShardWithPriority) -> Self { Self { - shard: shard.into(), + shard, ..Default::default() } } + /// Returns true if this is a query that + /// can be sent to a replica. pub fn is_read(&self) -> bool { self.read } + /// Returns true if this query can only be sent + /// to a primary. pub fn is_write(&self) -> bool { !self.is_read() } @@ -123,15 +156,23 @@ impl Route { &self.shard } - /// Should this query go to all shards? + pub fn shard_with_priority(&self) -> &ShardWithPriority { + &self.shard + } + + /// Returns true if this query should go to all shards. pub fn is_all_shards(&self) -> bool { - matches!(self.shard, Shard::All) + matches!(*self.shard, Shard::All) } + /// Returns true if this query should be sent to multiple + /// but not all shards. pub fn is_multi_shard(&self) -> bool { - matches!(self.shard, Shard::Multi(_)) + matches!(*self.shard, Shard::Multi(_)) } + /// Returns true if this query should be sent to + /// more than one shard. pub fn is_cross_shard(&self) -> bool { self.is_all_shards() || self.is_multi_shard() } @@ -148,26 +189,29 @@ impl Route { &mut self.aggregate } - pub fn set_shard_mut(&mut self, shard: usize) { - self.shard = Shard::Direct(shard); + pub fn set_shard_mut(&mut self, shard: ShardWithPriority) { + self.shard = shard; } - pub fn set_shard(mut self, shard: usize) -> Self { + pub fn with_shard(mut self, shard: ShardWithPriority) -> Self { self.set_shard_mut(shard); self } - pub fn set_maintenace(mut self) -> Self { - self.maintenance = true; - self + pub fn set_search_path_driven_mut(&mut self, schema_driven: bool) { + self.search_path_driven = schema_driven; + } + + pub fn is_search_path_driven(&self) -> bool { + self.search_path_driven } pub fn is_maintenance(&self) -> bool { self.maintenance } - pub fn set_shard_raw_mut(&mut self, shard: &Shard) { - self.shard = shard.clone(); + pub fn set_shard_raw_mut(&mut self, shard: ShardWithPriority) { + self.shard = shard; } pub fn should_buffer(&self) -> bool { @@ -178,21 +222,42 @@ impl Route { &self.limit } - pub fn set_read(mut self, read: bool) -> Self { - self.set_read_mut(read); + pub fn with_read(mut self, read: bool) -> Self { + self.set_read(read); self } - pub fn set_read_mut(&mut self, read: bool) { + pub fn set_read(&mut self, read: bool) { self.read = read; } - pub fn set_write(mut self, write: FunctionBehavior) -> Self { - self.set_write_mut(write); + pub fn explain(&self) -> Option<&ExplainTrace> { + self.explain.as_ref() + } + + pub fn set_explain(&mut self, trace: ExplainTrace) { + self.explain = Some(trace); + } + + pub fn take_explain(&mut self) -> Option { + self.explain.take() + } + + pub fn with_rollback_savepoint(mut self, rollback: bool) -> Self { + self.rollback_savepoint = rollback; + self + } + + pub fn rollback_savepoint(&self) -> bool { + self.rollback_savepoint + } + + pub fn with_write(mut self, write: FunctionBehavior) -> Self { + self.set_write(write); self } - pub fn set_write_mut(&mut self, write: FunctionBehavior) { + pub fn set_write(&mut self, write: FunctionBehavior) { let FunctionBehavior { writes, locking_behavior, @@ -201,12 +266,7 @@ impl Route { self.lock_session = matches!(locking_behavior, LockingBehavior::Lock); } - pub fn set_lock_session(mut self) -> Self { - self.lock_session = true; - self - } - - pub fn lock_session(&self) -> bool { + pub fn is_lock_session(&self) -> bool { self.lock_session } @@ -218,29 +278,293 @@ impl Route { self.is_cross_shard() && self.is_write() && !self.is_maintenance() } - pub fn rewrite_plan(&self) -> &RewritePlan { + pub fn aggregate_rewrite_plan(&self) -> &AggregateRewritePlan { &self.rewrite_plan } - pub fn rewrite_plan_mut(&mut self) -> &mut RewritePlan { - &mut self.rewrite_plan + pub fn with_aggregate_rewrite_plan_mut(&mut self, plan: AggregateRewritePlan) { + self.rewrite_plan = plan; } +} - pub fn set_rewrite(&mut self, plan: RewritePlan, sql: String) { - self.rewrite_plan = plan; - self.rewritten_sql = Some(sql); +/// Shard source. +/// +/// N.B. Ordering here matters. Don't move these around, +/// unless you're changing the algorithm. +/// +/// These are ranked from least priority to highest +/// priority. +#[derive(Debug, Clone, PartialEq, Eq, Ord, PartialOrd, Default)] +pub enum ShardSource { + #[default] + DefaultUnset, + Table, + RoundRobin(RoundRobinReason), + SearchPath(String), + Set, + Comment, + Plugin, + Override(OverrideReason), +} + +#[derive(Debug, Clone, PartialEq, Eq, Ord, PartialOrd)] +pub enum RoundRobinReason { + PrimaryShardedTableInsert, + Omni, + NotExecutable, + NoTable, + EmptyQuery, +} + +#[derive(Debug, Clone, PartialEq, Eq, Ord, PartialOrd)] +pub enum OverrideReason { + DryRun, + ParserDisabled, + Transaction, + OnlyOneShard, + RewriteUpdate, +} + +#[derive(Debug, Clone, PartialEq, Eq, Ord, PartialOrd, Default)] +pub struct ShardWithPriority { + source: ShardSource, + shard: Shard, +} + +impl ShardWithPriority { + /// Create new shard with comment-level priority. + pub fn new_comment(shard: Shard) -> Self { + Self { + shard, + source: ShardSource::Comment, + } + } + + pub fn new_plugin(shard: Shard) -> Self { + Self { + shard, + source: ShardSource::Plugin, + } + } + + /// Create new shard with table-level priority. + pub fn new_table(shard: Shard) -> Self { + Self { + shard, + source: ShardSource::Table, + } + } + + /// Create new shard with highest priority. + pub fn new_override_parser_disabled(shard: Shard) -> Self { + Self { + shard, + source: ShardSource::Override(OverrideReason::ParserDisabled), + } } - pub fn clear_rewrite(&mut self) { - self.rewrite_plan = RewritePlan::new(); - self.rewritten_sql = None; + pub fn new_override_rewrite_update(shard: Shard) -> Self { + Self { + shard, + source: ShardSource::Override(OverrideReason::RewriteUpdate), + } } - pub fn rewritten_sql(&self) -> Option<&str> { - self.rewritten_sql.as_deref() + pub fn new_override_dry_run(shard: Shard) -> Self { + Self { + shard, + source: ShardSource::Override(OverrideReason::DryRun), + } } - pub fn take_rewritten_sql(&mut self) -> Option { - self.rewritten_sql.take() + pub fn new_override_transaction(shard: Shard) -> Self { + Self { + shard, + source: ShardSource::Override(OverrideReason::Transaction), + } + } + + pub fn new_override_only_one_shard(shard: Shard) -> Self { + Self { + shard, + source: ShardSource::Override(OverrideReason::OnlyOneShard), + } + } + + pub fn new_default_unset(shard: Shard) -> Self { + Self { + shard, + source: ShardSource::DefaultUnset, + } + } + + pub fn new_rr_omni(shard: Shard) -> Self { + Self { + shard, + source: ShardSource::RoundRobin(RoundRobinReason::Omni), + } + } + + pub fn new_rr_not_executable(shard: Shard) -> Self { + Self { + shard, + source: ShardSource::RoundRobin(RoundRobinReason::NotExecutable), + } + } + + pub fn new_rr_primary_insert(shard: Shard) -> Self { + Self { + shard, + source: ShardSource::RoundRobin(RoundRobinReason::PrimaryShardedTableInsert), + } + } + + pub fn new_rr_no_table(shard: Shard) -> Self { + Self { + shard, + source: ShardSource::RoundRobin(RoundRobinReason::NoTable), + } + } + + pub fn new_rr_empty_query(shard: Shard) -> Self { + Self { + shard, + source: ShardSource::RoundRobin(RoundRobinReason::EmptyQuery), + } + } + + /// New SET-based routing. + pub fn new_set(shard: Shard) -> Self { + Self { + shard, + source: ShardSource::Set, + } + } + + /// New search_path-based shard. + pub fn new_search_path(shard: Shard, schema: &str) -> Self { + Self { + shard, + source: ShardSource::SearchPath(schema.to_string()), + } + } + + #[cfg(test)] + pub(crate) fn source(&self) -> &ShardSource { + &self.source + } +} + +impl Deref for ShardWithPriority { + type Target = Shard; + + fn deref(&self) -> &Self::Target { + &self.shard + } +} + +/// Ordered collection of set shards. +#[derive(Default, Debug, Clone)] +pub struct ShardsWithPriority { + max: Option, +} + +impl ShardsWithPriority { + /// Get currently computed shard. + pub(crate) fn shard(&self) -> ShardWithPriority { + lazy_static! { + static ref DEFAULT_SHARD: ShardWithPriority = ShardWithPriority { + shard: Shard::All, + source: ShardSource::DefaultUnset, + }; + } + + self.peek().cloned().unwrap_or(DEFAULT_SHARD.clone()) + } + + pub(crate) fn push(&mut self, shard: ShardWithPriority) { + if let Some(ref max) = self.max { + if max < &shard { + self.max = Some(shard); + } + } else { + self.max = Some(shard); + } + } + + pub(crate) fn peek(&self) -> Option<&ShardWithPriority> { + self.max.as_ref() + } + + /// Schema-path based routing priority is used. + pub(crate) fn is_search_path(&self) -> bool { + self.peek() + .map(|shard| matches!(shard.source, ShardSource::SearchPath(_))) + .unwrap_or_default() + } +} + +#[cfg(test)] +mod test { + use super::*; + + #[test] + fn test_shard_ord() { + assert!(Shard::Direct(0) < Shard::All); + assert!(Shard::Multi(vec![]) < Shard::All); + } + + #[test] + fn test_source_ord() { + assert!(ShardSource::Table < ShardSource::RoundRobin(RoundRobinReason::NotExecutable)); + assert!(ShardSource::Table < ShardSource::SearchPath(String::new())); + assert!(ShardSource::SearchPath(String::new()) < ShardSource::Set); + assert!(ShardSource::Set < ShardSource::Comment); + assert!(ShardSource::Comment < ShardSource::Override(OverrideReason::OnlyOneShard)); + } + + #[test] + fn test_shard_with_priority_ord() { + let shard = Shard::Direct(0); + + assert!( + ShardWithPriority::new_table(shard.clone()) + < ShardWithPriority::new_rr_omni(shard.clone()) + ); + assert!( + ShardWithPriority::new_table(shard.clone()) + < ShardWithPriority::new_search_path(shard.clone(), "schema") + ); + assert!( + ShardWithPriority::new_search_path(shard.clone(), "schema") + < ShardWithPriority::new_set(shard.clone()) + ); + assert!( + ShardWithPriority::new_set(shard.clone()) + < ShardWithPriority::new_comment(shard.clone()) + ); + assert!( + ShardWithPriority::new_comment(shard.clone()) + < ShardWithPriority::new_override_dry_run(shard.clone()) + ); + } + + #[test] + fn test_comment_override_set() { + let mut shards = ShardsWithPriority::default(); + + shards.push(ShardWithPriority::new_set(Shard::Direct(1))); + assert_eq!(shards.shard().deref(), &Shard::Direct(1)); + + shards.push(ShardWithPriority::new_comment(Shard::Direct(2))); + assert_eq!(shards.shard().deref(), &Shard::Direct(2)); + + let mut shards = ShardsWithPriority::default(); + + shards.push(ShardWithPriority::new_comment(Shard::Direct(3))); + assert_eq!(shards.shard().deref(), &Shard::Direct(3)); + + shards.push(ShardWithPriority::new_set(Shard::Direct(4))); + assert_eq!(shards.shard().deref(), &Shard::Direct(3)); } } diff --git a/pgdog/src/frontend/router/parser/schema.rs b/pgdog/src/frontend/router/parser/schema.rs new file mode 100644 index 000000000..a2e1f7388 --- /dev/null +++ b/pgdog/src/frontend/router/parser/schema.rs @@ -0,0 +1,9 @@ +pub struct Schema<'a> { + pub name: &'a str, +} + +impl<'a> From<&'a str> for Schema<'a> { + fn from(value: &'a str) -> Self { + Schema { name: value } + } +} diff --git a/pgdog/src/frontend/router/parser/statement.rs b/pgdog/src/frontend/router/parser/statement.rs new file mode 100644 index 000000000..da3276333 --- /dev/null +++ b/pgdog/src/frontend/router/parser/statement.rs @@ -0,0 +1,1925 @@ +use std::collections::{HashMap, HashSet}; + +use pg_query::{ + protobuf::{ + AExprKind, BoolExprType, DeleteStmt, InsertStmt, RangeVar, RawStmt, SelectStmt, UpdateStmt, + }, + Node, NodeEnum, +}; + +use super::{ + super::sharding::Value as ShardingValue, explain_trace::ExplainRecorder, Column, Error, Table, + Value, +}; +use crate::{ + backend::ShardingSchema, + frontend::router::{parser::Shard, sharding::ContextBuilder, sharding::SchemaSharder}, + net::Bind, +}; + +/// Context for searching a SELECT statement, tracking table aliases. +#[derive(Debug, Default, Clone)] +struct SearchContext<'a> { + /// Maps alias -> full Table (including schema) + aliases: HashMap<&'a str, Table<'a>>, + /// The primary table from the FROM clause (if simple) + table: Option>, +} + +impl<'a> SearchContext<'a> { + /// Build context from a FROM clause, extracting table aliases. + fn from_from_clause(nodes: &'a [Node]) -> Self { + let mut ctx = Self::default(); + ctx.extract_aliases(nodes); + + // Try to get the primary table for simple queries + if nodes.len() == 1 { + if let Some(table) = nodes.first().and_then(|n| Table::try_from(n).ok()) { + ctx.table = Some(table); + } + } + + ctx + } + + fn extract_aliases(&mut self, nodes: &'a [Node]) { + for node in nodes { + self.extract_alias_from_node(node); + } + } + + fn extract_alias_from_node(&mut self, node: &'a Node) { + match &node.node { + Some(NodeEnum::RangeVar(range_var)) => { + if let Some(ref alias) = range_var.alias { + let table = Table::from(range_var); + self.aliases.insert(alias.aliasname.as_str(), table); + } + } + Some(NodeEnum::JoinExpr(join)) => { + if let Some(ref larg) = join.larg { + self.extract_alias_from_node(larg); + } + if let Some(ref rarg) = join.rarg { + self.extract_alias_from_node(rarg); + } + } + Some(NodeEnum::RangeSubselect(subselect)) => { + if let Some(ref alias) = subselect.alias { + // For subselects, we don't have a real table name + // but we record the alias anyway for future use + self.aliases.insert( + alias.aliasname.as_str(), + Table { + name: alias.aliasname.as_str(), + schema: None, + alias: None, + }, + ); + } + } + _ => {} + } + } + + /// Resolve a table reference (which may be an alias) to the actual Table. + fn resolve_table(&self, name: &str) -> Option> { + self.aliases.get(name).copied() + } +} + +#[derive(Debug)] +enum SearchResult<'a> { + Column(Column<'a>), + Value(Value<'a>), + Values(Vec>), + Match(Shard), + Matches(Vec), + None, +} + +struct ValueIterator<'a, 'b> { + source: &'b SearchResult<'a>, + pos: usize, +} + +impl<'a, 'b> Iterator for ValueIterator<'a, 'b> { + type Item = &'b Value<'a>; + + fn next(&mut self) -> Option { + let next = match self.source { + SearchResult::Value(ref val) => { + if self.pos == 0 { + Some(val) + } else { + None + } + } + SearchResult::Values(values) => values.get(self.pos), + _ => None, + }; + + self.pos += 1; + + next + } +} + +impl<'a> SearchResult<'a> { + fn is_none(&self) -> bool { + matches!(self, Self::None) + } + + fn is_match(&self) -> bool { + matches!(self, Self::Match(_) | Self::Matches(_)) + } + + fn merge(self, other: Self) -> Self { + match (self, other) { + (Self::Match(first), Self::Match(second)) => Self::Matches(vec![first, second]), + (Self::Match(shard), Self::Matches(mut shards)) + | (Self::Matches(mut shards), Self::Match(shard)) => Self::Matches({ + shards.push(shard); + shards + }), + (Self::None, other) | (other, Self::None) => other, + _ => Self::None, + } + } + + fn iter<'b>(&'b self) -> ValueIterator<'a, 'b> { + ValueIterator { + source: self, + pos: 0, + } + } +} + +enum Statement<'a> { + Select(&'a SelectStmt), + Update(&'a UpdateStmt), + Delete(&'a DeleteStmt), + Insert(&'a InsertStmt), +} + +pub struct StatementParser<'a, 'b, 'c> { + stmt: Statement<'a>, + bind: Option<&'b Bind>, + schema: &'b ShardingSchema, + recorder: Option<&'c mut ExplainRecorder>, +} + +impl<'a, 'b, 'c> StatementParser<'a, 'b, 'c> { + fn new( + stmt: Statement<'a>, + bind: Option<&'b Bind>, + schema: &'b ShardingSchema, + recorder: Option<&'c mut ExplainRecorder>, + ) -> Self { + Self { + stmt, + bind, + schema, + recorder, + } + } + + pub fn from_select( + stmt: &'a SelectStmt, + bind: Option<&'b Bind>, + schema: &'b ShardingSchema, + recorder: Option<&'c mut ExplainRecorder>, + ) -> Self { + Self::new(Statement::Select(stmt), bind, schema, recorder) + } + + pub fn from_update( + stmt: &'a UpdateStmt, + bind: Option<&'b Bind>, + schema: &'b ShardingSchema, + recorder: Option<&'c mut ExplainRecorder>, + ) -> Self { + Self::new(Statement::Update(stmt), bind, schema, recorder) + } + + pub fn from_delete( + stmt: &'a DeleteStmt, + bind: Option<&'b Bind>, + schema: &'b ShardingSchema, + recorder: Option<&'c mut ExplainRecorder>, + ) -> Self { + Self::new(Statement::Delete(stmt), bind, schema, recorder) + } + + pub fn from_insert( + stmt: &'a InsertStmt, + bind: Option<&'b Bind>, + schema: &'b ShardingSchema, + recorder: Option<&'c mut ExplainRecorder>, + ) -> Self { + Self::new(Statement::Insert(stmt), bind, schema, recorder) + } + + /// Record a sharding key match. + fn record_sharding_key(&mut self, shard: &Shard, column: Column<'_>, value: &Value<'_>) { + if let Some(recorder) = self.recorder.as_mut() { + let col_str = if let Some(table) = column.table { + format!("{}.{}", table, column.name) + } else { + column.name.to_string() + }; + let description = match value { + Value::Placeholder(pos) => { + format!("matched sharding key {} using parameter ${}", col_str, pos) + } + _ => format!("matched sharding key {} using constant", col_str), + }; + recorder.record_entry(Some(shard.clone()), description); + } + } + + pub fn from_raw( + raw: &'a RawStmt, + bind: Option<&'b Bind>, + schema: &'b ShardingSchema, + recorder: Option<&'c mut ExplainRecorder>, + ) -> Result { + match raw.stmt.as_ref().and_then(|n| n.node.as_ref()) { + Some(NodeEnum::SelectStmt(stmt)) => Ok(Self::from_select(stmt, bind, schema, recorder)), + Some(NodeEnum::UpdateStmt(stmt)) => Ok(Self::from_update(stmt, bind, schema, recorder)), + Some(NodeEnum::DeleteStmt(stmt)) => Ok(Self::from_delete(stmt, bind, schema, recorder)), + Some(NodeEnum::InsertStmt(stmt)) => Ok(Self::from_insert(stmt, bind, schema, recorder)), + _ => Err(Error::NotASelect), + } + } + + pub fn shard(&mut self) -> Result, Error> { + let result = match self.stmt { + Statement::Select(stmt) => self.shard_select(stmt), + Statement::Update(stmt) => self.shard_update(stmt), + Statement::Delete(stmt) => self.shard_delete(stmt), + Statement::Insert(stmt) => self.shard_insert(stmt), + }?; + + // Key-based sharding succeeded + if result.is_some() { + return Ok(result); + } else if self.schema.schemas.is_empty() { + return Ok(None); + } + + // Fallback to schema-based sharding + let tables = self.extract_tables(); + let mut schema_sharder = SchemaSharder::default(); + for table in &tables { + schema_sharder.resolve(table.schema(), &self.schema.schemas); + } + + if let Some((shard, schema_name)) = schema_sharder.get() { + if let Some(recorder) = self.recorder.as_mut() { + recorder.record_entry( + Some(shard.clone()), + format!("matched schema {}", schema_name), + ); + } + return Ok(Some(shard)); + } + + Ok(None) + } + + /// Extract all tables referenced in the statement. + fn extract_tables(&self) -> Vec> { + let mut tables = Vec::new(); + match self.stmt { + Statement::Select(stmt) => self.extract_tables_from_select(stmt, &mut tables), + Statement::Update(stmt) => self.extract_tables_from_update(stmt, &mut tables), + Statement::Delete(stmt) => self.extract_tables_from_delete(stmt, &mut tables), + Statement::Insert(stmt) => self.extract_tables_from_insert(stmt, &mut tables), + } + tables + } + + fn extract_tables_from_select(&self, stmt: &'a SelectStmt, tables: &mut Vec>) { + // Handle UNION/INTERSECT/EXCEPT + if let Some(ref larg) = stmt.larg { + self.extract_tables_from_select(larg, tables); + } + if let Some(ref rarg) = stmt.rarg { + self.extract_tables_from_select(rarg, tables); + } + + // FROM clause + for node in &stmt.from_clause { + self.extract_tables_from_node(node, tables); + } + + // WITH clause (CTEs) + if let Some(ref with_clause) = stmt.with_clause { + for cte in &with_clause.ctes { + if let Some(NodeEnum::CommonTableExpr(ref cte_expr)) = cte.node { + if let Some(ref ctequery) = cte_expr.ctequery { + if let Some(NodeEnum::SelectStmt(ref inner_select)) = ctequery.node { + self.extract_tables_from_select(inner_select, tables); + } + } + } + } + } + + // WHERE clause subqueries + if let Some(ref where_clause) = stmt.where_clause { + self.extract_tables_from_node(where_clause, tables); + } + } + + fn extract_tables_from_update(&self, stmt: &'a UpdateStmt, tables: &mut Vec>) { + // Main relation + if let Some(ref relation) = stmt.relation { + tables.push(Table::from(relation)); + } + + // FROM clause + for node in &stmt.from_clause { + self.extract_tables_from_node(node, tables); + } + + // WITH clause (CTEs) + if let Some(ref with_clause) = stmt.with_clause { + for cte in &with_clause.ctes { + if let Some(NodeEnum::CommonTableExpr(ref cte_expr)) = cte.node { + if let Some(ref ctequery) = cte_expr.ctequery { + if let Some(NodeEnum::SelectStmt(ref inner_select)) = ctequery.node { + self.extract_tables_from_select(inner_select, tables); + } + } + } + } + } + + // WHERE clause subqueries + if let Some(ref where_clause) = stmt.where_clause { + self.extract_tables_from_node(where_clause, tables); + } + } + + fn extract_tables_from_delete(&self, stmt: &'a DeleteStmt, tables: &mut Vec>) { + // Main relation + if let Some(ref relation) = stmt.relation { + tables.push(Table::from(relation)); + } + + // USING clause + for node in &stmt.using_clause { + self.extract_tables_from_node(node, tables); + } + + // WITH clause (CTEs) + if let Some(ref with_clause) = stmt.with_clause { + for cte in &with_clause.ctes { + if let Some(NodeEnum::CommonTableExpr(ref cte_expr)) = cte.node { + if let Some(ref ctequery) = cte_expr.ctequery { + if let Some(NodeEnum::SelectStmt(ref inner_select)) = ctequery.node { + self.extract_tables_from_select(inner_select, tables); + } + } + } + } + } + + // WHERE clause subqueries + if let Some(ref where_clause) = stmt.where_clause { + self.extract_tables_from_node(where_clause, tables); + } + } + + fn extract_tables_from_insert(&self, stmt: &'a InsertStmt, tables: &mut Vec>) { + // Main relation + if let Some(ref relation) = stmt.relation { + tables.push(Table::from(relation)); + } + + // WITH clause (CTEs) + if let Some(ref with_clause) = stmt.with_clause { + for cte in &with_clause.ctes { + if let Some(NodeEnum::CommonTableExpr(ref cte_expr)) = cte.node { + if let Some(ref ctequery) = cte_expr.ctequery { + if let Some(NodeEnum::SelectStmt(ref inner_select)) = ctequery.node { + self.extract_tables_from_select(inner_select, tables); + } + } + } + } + } + + // SELECT part of INSERT ... SELECT + if let Some(ref select_stmt) = stmt.select_stmt { + if let Some(NodeEnum::SelectStmt(ref inner_select)) = select_stmt.node { + self.extract_tables_from_select(inner_select, tables); + } + } + } + + fn extract_tables_from_node(&self, node: &'a Node, tables: &mut Vec>) { + match &node.node { + Some(NodeEnum::RangeVar(range_var)) => { + tables.push(Table::from(range_var)); + } + Some(NodeEnum::JoinExpr(join)) => { + if let Some(ref larg) = join.larg { + self.extract_tables_from_node(larg, tables); + } + if let Some(ref rarg) = join.rarg { + self.extract_tables_from_node(rarg, tables); + } + } + Some(NodeEnum::RangeSubselect(subselect)) => { + if let Some(ref subquery) = subselect.subquery { + if let Some(NodeEnum::SelectStmt(ref inner_select)) = subquery.node { + self.extract_tables_from_select(inner_select, tables); + } + } + } + Some(NodeEnum::SubLink(sublink)) => { + if let Some(ref subselect) = sublink.subselect { + if let Some(NodeEnum::SelectStmt(ref inner_select)) = subselect.node { + self.extract_tables_from_select(inner_select, tables); + } + } + } + Some(NodeEnum::SelectStmt(inner_select)) => { + self.extract_tables_from_select(inner_select, tables); + } + Some(NodeEnum::BoolExpr(bool_expr)) => { + for arg in &bool_expr.args { + self.extract_tables_from_node(arg, tables); + } + } + Some(NodeEnum::AExpr(a_expr)) => { + if let Some(ref lexpr) = a_expr.lexpr { + self.extract_tables_from_node(lexpr, tables); + } + if let Some(ref rexpr) = a_expr.rexpr { + self.extract_tables_from_node(rexpr, tables); + } + } + _ => {} + } + } + + fn shard_select(&mut self, stmt: &'a SelectStmt) -> Result, Error> { + let ctx = SearchContext::from_from_clause(&stmt.from_clause); + let result = self.search_select_stmt(stmt, &ctx)?; + + match result { + SearchResult::Match(shard) => Ok(Some(shard)), + SearchResult::Matches(shards) => Ok(Self::converge(&shards)), + _ => Ok(None), + } + } + + fn shard_update(&mut self, stmt: &'a UpdateStmt) -> Result, Error> { + let ctx = self.context_from_relation(&stmt.relation); + let result = self.search_update_stmt(stmt, &ctx)?; + + match result { + SearchResult::Match(shard) => Ok(Some(shard)), + SearchResult::Matches(shards) => Ok(Self::converge(&shards)), + _ => Ok(None), + } + } + + fn shard_delete(&mut self, stmt: &'a DeleteStmt) -> Result, Error> { + let ctx = self.context_from_relation(&stmt.relation); + let result = self.search_delete_stmt(stmt, &ctx)?; + + match result { + SearchResult::Match(shard) => Ok(Some(shard)), + SearchResult::Matches(shards) => Ok(Self::converge(&shards)), + _ => Ok(None), + } + } + + fn shard_insert(&mut self, stmt: &'a InsertStmt) -> Result, Error> { + let ctx = self.context_from_relation(&stmt.relation); + let result = self.search_insert_stmt(stmt, &ctx)?; + + match result { + SearchResult::Match(shard) => Ok(Some(shard)), + SearchResult::Matches(shards) => Ok(Self::converge(&shards)), + _ => Ok(None), + } + } + + fn context_from_relation(&self, relation: &'a Option) -> SearchContext<'a> { + let mut ctx = SearchContext::default(); + if let Some(ref range_var) = relation { + let table = Table::from(range_var); + ctx.table = Some(table); + if let Some(ref alias) = range_var.alias { + ctx.aliases.insert(alias.aliasname.as_str(), table); + } + } + ctx + } + + fn converge(shards: &[Shard]) -> Option { + let shards: HashSet = shards.iter().cloned().collect(); + match shards.len() { + 0 => None, + 1 => shards.into_iter().next(), + _ => { + let mut multi = vec![]; + for shard in shards.into_iter() { + match shard { + Shard::All => return Some(Shard::All), + Shard::Direct(direct) => multi.push(direct), + Shard::Multi(many) => multi.extend(many), + } + } + Some(Shard::Multi(multi)) + } + } + } + + fn compute_shard( + &mut self, + column: Column<'a>, + value: Value<'a>, + ) -> Result, Error> { + if let Some(table) = self.schema.tables().get_table(column) { + let context = ContextBuilder::new(table); + let shard = match value { + Value::Placeholder(pos) => { + let param = self + .bind + .map(|bind| bind.parameter(pos as usize - 1)) + .transpose()? + .flatten(); + // Expect params to be accurate. + let param = if let Some(param) = param { + param + } else { + return Ok(None); + }; + let value = ShardingValue::from_param(¶m, table.data_type)?; + Some( + context + .value(value) + .shards(self.schema.shards) + .build()? + .apply()?, + ) + } + + Value::String(val) => Some( + context + .data(val) + .shards(self.schema.shards) + .build()? + .apply()?, + ), + + Value::Integer(val) => Some( + context + .data(val) + .shards(self.schema.shards) + .build()? + .apply()?, + ), + Value::Null => None, + _ => None, + }; + + Ok(shard) + } else { + Ok(None) + } + } + + fn select_search( + &mut self, + node: &'a Node, + ctx: &SearchContext<'a>, + ) -> Result, Error> { + match node.node { + // Value types - these are leaf nodes representing actual values + Some(NodeEnum::AConst(_)) + | Some(NodeEnum::ParamRef(_)) + | Some(NodeEnum::FuncCall(_)) => { + if let Ok(value) = Value::try_from(&node.node) { + return Ok(SearchResult::Value(value)); + } + Ok(SearchResult::None) + } + + Some(NodeEnum::TypeCast(ref cast)) => { + if let Some(ref arg) = cast.arg { + return self.select_search(arg, ctx); + } + Ok(SearchResult::None) + } + + Some(NodeEnum::SelectStmt(ref stmt)) => { + // Build context with aliases from the FROM clause + let ctx = SearchContext::from_from_clause(&stmt.from_clause); + self.search_select_stmt(stmt, &ctx) + } + + Some(NodeEnum::RangeSubselect(ref subselect)) => { + if let Some(ref node) = subselect.subquery { + self.select_search(node, ctx) + } else { + Ok(SearchResult::None) + } + } + + Some(NodeEnum::ColumnRef(_)) => { + let mut column = Column::try_from(&node.node)?; + + // If column has no table, qualify with context table + if column.table().is_none() { + if let Some(ref table) = ctx.table { + column.qualify(*table); + } + } + + Ok(SearchResult::Column(column)) + } + + Some(NodeEnum::AExpr(ref expr)) => { + let kind = expr.kind(); + let mut supported = false; + + if matches!( + kind, + AExprKind::AexprOp | AExprKind::AexprIn | AExprKind::AexprOpAny + ) { + supported = expr + .name + .first() + .map(|node| match node.node { + Some(NodeEnum::String(ref string)) => string.sval.as_str(), + _ => "", + }) + .unwrap_or_default() + == "="; + } + + if !supported { + return Ok(SearchResult::None); + } + + let is_any = matches!(kind, AExprKind::AexprOpAny); + + let mut results = vec![]; + + if let Some(ref left) = expr.lexpr { + results.push(self.select_search(left, ctx)?); + } + + if let Some(ref right) = expr.rexpr { + results.push(self.select_search(right, ctx)?); + } + + if results.len() != 2 { + Ok(SearchResult::None) + } else { + let right = results.pop().unwrap(); + let left = results.pop().unwrap(); + + // If either side is already a match (from subquery), return it + if right.is_match() { + return Ok(right); + } + if left.is_match() { + return Ok(left); + } + + match (right, left) { + (SearchResult::Column(column), values) + | (values, SearchResult::Column(column)) => { + // For ANY expressions with sharding columns, we can't reliably + // parse array literals or parameters, so route to all shards. + if is_any + && matches!(values, SearchResult::Value(_)) + && self.schema.tables().get_table(column).is_some() + { + return Ok(SearchResult::Match(Shard::All)); + } + + let mut shards = HashSet::new(); + for value in values.iter() { + if let Some(shard) = + self.compute_shard_with_ctx(column, value.clone(), ctx)? + { + shards.insert(shard); + } + } + + match shards.len() { + 0 => Ok(SearchResult::None), + 1 => Ok(SearchResult::Match(shards.into_iter().next().unwrap())), + _ => Ok(SearchResult::Matches(shards.into_iter().collect())), + } + } + _ => Ok(SearchResult::None), + } + } + } + + Some(NodeEnum::List(ref list)) => { + let mut values = vec![]; + + for value in &list.items { + if let Ok(value) = Value::try_from(&value.node) { + values.push(value); + } + } + + Ok(SearchResult::Values(values)) + } + + Some(NodeEnum::WithClause(ref with_clause)) => { + for cte in &with_clause.ctes { + let result = self.select_search(cte, ctx)?; + if !result.is_none() { + return Ok(result); + } + } + + Ok(SearchResult::None) + } + + Some(NodeEnum::JoinExpr(ref join)) => { + let mut results = vec![]; + + if let Some(ref left) = join.larg { + results.push(self.select_search(left, ctx)?); + } + if let Some(ref right) = join.rarg { + results.push(self.select_search(right, ctx)?); + } + + results.retain(|result| result.is_match()); + + let result = results + .into_iter() + .fold(SearchResult::None, |acc, x| acc.merge(x)); + + Ok(result) + } + + Some(NodeEnum::BoolExpr(ref expr)) => { + // Only AND expressions can determine a shard. + // OR expressions could route to multiple shards. + if expr.boolop() != BoolExprType::AndExpr { + return Ok(SearchResult::None); + } + + for arg in &expr.args { + let result = self.select_search(arg, ctx)?; + if result.is_match() { + return Ok(result); + } + } + + Ok(SearchResult::None) + } + + Some(NodeEnum::SubLink(ref sublink)) => { + if let Some(ref subselect) = sublink.subselect { + return self.select_search(subselect, ctx); + } + Ok(SearchResult::None) + } + + Some(NodeEnum::CommonTableExpr(ref cte)) => { + if let Some(ref ctequery) = cte.ctequery { + return self.select_search(ctequery, ctx); + } + Ok(SearchResult::None) + } + + _ => Ok(SearchResult::None), + } + } + + /// Search a SELECT statement with its own context. + fn search_select_stmt( + &mut self, + stmt: &'a SelectStmt, + ctx: &SearchContext<'a>, + ) -> Result, Error> { + // Handle UNION/INTERSECT/EXCEPT (set operations) + // These have larg and rarg instead of a regular SELECT structure + if let Some(ref larg) = stmt.larg { + let larg_ctx = SearchContext::from_from_clause(&larg.from_clause); + let result = self.search_select_stmt(larg, &larg_ctx)?; + if !result.is_none() { + return Ok(result); + } + } + if let Some(ref rarg) = stmt.rarg { + let rarg_ctx = SearchContext::from_from_clause(&rarg.from_clause); + let result = self.search_select_stmt(rarg, &rarg_ctx)?; + if !result.is_none() { + return Ok(result); + } + } + + if let Some(ref with_clause) = stmt.with_clause { + for cte in &with_clause.ctes { + let result = self.select_search(cte, ctx)?; + if !result.is_none() { + return Ok(result); + } + } + } + + // Search WHERE clause + if let Some(ref where_clause) = stmt.where_clause { + let result = self.select_search(where_clause, ctx)?; + if !result.is_none() { + return Ok(result); + } + } + + for from_ in &stmt.from_clause { + let result = self.select_search(from_, ctx)?; + if !result.is_none() { + return Ok(result); + } + } + + Ok(SearchResult::None) + } + + /// Compute shard with alias resolution from context. + fn compute_shard_with_ctx( + &mut self, + column: Column<'a>, + value: Value<'a>, + ctx: &SearchContext<'a>, + ) -> Result, Error> { + // Resolve table alias if present + let resolved_column = if let Some(table_ref) = column.table() { + if let Some(resolved) = ctx.resolve_table(table_ref.name) { + Column { + name: column.name, + table: Some(resolved.name), + schema: resolved.schema, + } + } else { + column + } + } else { + column + }; + + let shard = self.compute_shard(resolved_column, value.clone())?; + if let Some(ref shard) = shard { + self.record_sharding_key(shard, resolved_column, &value); + } + Ok(shard) + } + + /// Search an UPDATE statement for sharding keys. + fn search_update_stmt( + &mut self, + stmt: &'a UpdateStmt, + ctx: &SearchContext<'a>, + ) -> Result, Error> { + // Handle CTEs (WITH clause) + if let Some(ref with_clause) = stmt.with_clause { + for cte in &with_clause.ctes { + let result = self.select_search(cte, ctx)?; + if !result.is_none() { + return Ok(result); + } + } + } + + // Search WHERE clause + if let Some(ref where_clause) = stmt.where_clause { + let result = self.select_search(where_clause, ctx)?; + if !result.is_none() { + return Ok(result); + } + } + + // Search FROM clause (UPDATE ... FROM ...) + for from_ in &stmt.from_clause { + let result = self.select_search(from_, ctx)?; + if !result.is_none() { + return Ok(result); + } + } + + Ok(SearchResult::None) + } + + /// Search a DELETE statement for sharding keys. + fn search_delete_stmt( + &mut self, + stmt: &'a DeleteStmt, + ctx: &SearchContext<'a>, + ) -> Result, Error> { + // Handle CTEs (WITH clause) + if let Some(ref with_clause) = stmt.with_clause { + for cte in &with_clause.ctes { + let result = self.select_search(cte, ctx)?; + if !result.is_none() { + return Ok(result); + } + } + } + + // Search WHERE clause + if let Some(ref where_clause) = stmt.where_clause { + let result = self.select_search(where_clause, ctx)?; + if !result.is_none() { + return Ok(result); + } + } + + // Search USING clause (DELETE ... USING ...) + for using_ in &stmt.using_clause { + let result = self.select_search(using_, ctx)?; + if !result.is_none() { + return Ok(result); + } + } + + Ok(SearchResult::None) + } + + /// Search an INSERT statement for sharding keys. + fn search_insert_stmt( + &mut self, + stmt: &'a InsertStmt, + ctx: &SearchContext<'a>, + ) -> Result, Error> { + // Handle CTEs (WITH clause) + if let Some(ref with_clause) = stmt.with_clause { + for cte in &with_clause.ctes { + let result = self.select_search(cte, ctx)?; + if !result.is_none() { + return Ok(result); + } + } + } + + // Get the column names from INSERT INTO table (col1, col2, ...) + let columns: Vec<&str> = stmt + .cols + .iter() + .filter_map(|node| match &node.node { + Some(NodeEnum::ResTarget(target)) => Some(target.name.as_str()), + _ => None, + }) + .collect(); + + // The select_stmt field contains either VALUES or a SELECT subquery + if let Some(ref select_node) = stmt.select_stmt { + if let Some(NodeEnum::SelectStmt(ref select_stmt)) = select_node.node { + // Check if this is VALUES (has values_lists) - need special handling + // to match column positions with sharding keys + if !select_stmt.values_lists.is_empty() { + for values_list in &select_stmt.values_lists { + if let Some(NodeEnum::List(ref list)) = values_list.node { + for (pos, value_node) in list.items.iter().enumerate() { + // Check if this position corresponds to a sharding key column + if let Some(column_name) = columns.get(pos) { + let column = Column { + name: column_name, + table: ctx.table.map(|t| t.name), + schema: ctx.table.and_then(|t| t.schema), + }; + + if self.schema.tables().get_table(column).is_some() { + // Try to extract the value directly + if let Ok(value) = Value::try_from(value_node) { + if let Some(shard) = + self.compute_shard_with_ctx(column, value, ctx)? + { + return Ok(SearchResult::Match(shard)); + } + } + } + } + + // Search subqueries in values recursively + let result = self.select_search(value_node, ctx)?; + if result.is_match() { + return Ok(result); + } + } + } + } + } + + // Handle INSERT ... SELECT by recursively searching the SelectStmt + let result = self.select_search(select_node, ctx)?; + if !result.is_none() { + return Ok(result); + } + } + } + + Ok(SearchResult::None) + } +} + +#[cfg(test)] +mod test { + use pgdog_config::{FlexibleType, Mapping, ShardedMapping, ShardedMappingKind, ShardedTable}; + + use crate::backend::ShardedTables; + use crate::net::messages::{Bind, Parameter}; + + use super::*; + + fn run_test(stmt: &str, bind: Option<&Bind>) -> Result, Error> { + let schema = ShardingSchema { + shards: 3, + tables: ShardedTables::new( + vec![ + ShardedTable { + column: "id".into(), + name: Some("sharded".into()), + ..Default::default() + }, + ShardedTable { + column: "sharded_id".into(), + ..Default::default() + }, + ShardedTable { + column: "list_id".into(), + mapping: Mapping::new(&[ShardedMapping { + kind: ShardedMappingKind::List, + values: vec![FlexibleType::Integer(1), FlexibleType::Integer(2)] + .into_iter() + .collect(), + ..Default::default() + }]), + ..Default::default() + }, + // Schema-qualified sharded table with different column name + ShardedTable { + column: "tenant_id".into(), + name: Some("schema_sharded".into()), + schema: Some("myschema".into()), + ..Default::default() + }, + ], + vec![], + ), + ..Default::default() + }; + let raw = pg_query::parse(stmt) + .unwrap() + .protobuf + .stmts + .first() + .cloned() + .unwrap(); + let mut parser = StatementParser::from_raw(&raw, bind, &schema, None)?; + parser.shard() + } + + #[test] + fn test_simple_select() { + let result = run_test("SELECT * FROM sharded WHERE id = 1", None); + assert!(result.unwrap().is_some()); + } + + #[test] + fn test_select_with_and() { + let result = run_test("SELECT * FROM sharded WHERE id = 1 AND name = 'foo'", None).unwrap(); + assert!(result.is_some()); + } + + #[test] + fn test_select_with_or_returns_none() { + // OR expressions can't determine a single shard + let result = run_test("SELECT * FROM sharded WHERE id = 1 OR id = 2", None).unwrap(); + assert!(result.is_none()); + } + + #[test] + fn test_select_with_subquery() { + let result = run_test( + "SELECT * FROM sharded WHERE id IN (SELECT sharded_id FROM other WHERE sharded_id = 1)", + None, + ) + .unwrap(); + assert!(result.is_some()); + } + + #[test] + fn test_select_with_cte() { + let result = run_test( + "WITH cte AS (SELECT * FROM sharded WHERE id = 1) SELECT * FROM cte", + None, + ) + .unwrap(); + assert!(result.is_some()); + } + + #[test] + fn test_select_with_join() { + let result = run_test( + "SELECT * FROM sharded s JOIN other o ON s.id = o.sharded_id WHERE s.id = 1", + None, + ) + .unwrap(); + assert!(result.is_some()); + } + + #[test] + fn test_select_with_type_cast() { + let result = run_test("SELECT * FROM sharded WHERE id = '1'::int", None).unwrap(); + assert!(result.is_some()); + } + + #[test] + fn test_select_from_subquery() { + let result = run_test( + "SELECT * FROM (SELECT * FROM sharded WHERE id = 1) AS sub", + None, + ) + .unwrap(); + assert!(result.is_some()); + } + + #[test] + fn test_select_with_nested_cte() { + let result = run_test( + "WITH cte1 AS (SELECT * FROM sharded WHERE id = 1), \ + cte2 AS (SELECT * FROM cte1) \ + SELECT * FROM cte2", + None, + ) + .unwrap(); + assert!(result.is_some()); + } + + #[test] + fn test_select_no_where_returns_none() { + let result = run_test("SELECT * FROM sharded", None).unwrap(); + assert!(result.is_none()); + } + + #[test] + fn test_select_with_in_list() { + let result = run_test("SELECT * FROM sharded WHERE id IN (1, 2, 3)", None).unwrap(); + // IN with multiple values should return a shard match + assert!(result.is_some()); + } + + #[test] + fn test_select_with_not_equals_returns_none() { + // != operator is not supported for sharding + let result = run_test("SELECT * FROM sharded WHERE id != 1", None).unwrap(); + assert!(result.is_none()); + } + + #[test] + fn test_select_with_greater_than_returns_none() { + // > operator is not supported for sharding + let result = run_test("SELECT * FROM sharded WHERE id > 1", None).unwrap(); + assert!(result.is_none()); + } + + #[test] + fn test_select_with_complex_and() { + let result = run_test( + "SELECT * FROM sharded WHERE id = 1 AND status = 'active' AND created_at > now()", + None, + ) + .unwrap(); + assert!(result.is_some()); + } + + #[test] + fn test_select_with_left_join() { + let result = run_test( + "SELECT * FROM sharded s LEFT JOIN other o ON s.id = o.sharded_id WHERE s.id = 1", + None, + ) + .unwrap(); + assert!(result.is_some()); + } + + #[test] + fn test_select_with_multiple_joins() { + let result = run_test( + "SELECT * FROM sharded s \ + JOIN other o ON s.id = o.sharded_id \ + JOIN third t ON o.id = t.other_id \ + WHERE s.id = 1", + None, + ) + .unwrap(); + assert!(result.is_some()); + } + + #[test] + fn test_select_with_exists_subquery() { + let result = run_test( + "SELECT * FROM sharded WHERE EXISTS (SELECT 1 FROM other WHERE sharded_id = 1)", + None, + ) + .unwrap(); + // EXISTS subquery should find the shard condition inside + assert!(result.is_some()); + } + + #[test] + fn test_select_with_scalar_subquery() { + // Scalar subquery where shard is determined by the subquery's WHERE clause + let result = run_test( + "SELECT * FROM sharded WHERE id = (SELECT sharded_id FROM other WHERE sharded_id = 1 LIMIT 1)", + None, + ) + .unwrap(); + // The subquery's sharded_id = 1 should be found + assert!(result.is_some()); + } + + #[test] + fn test_select_with_recursive_cte() { + // Recursive CTEs have UNION - we look at the base case + let result = run_test( + "WITH RECURSIVE cte AS ( \ + SELECT * FROM sharded WHERE id = 1 \ + UNION ALL \ + SELECT s.* FROM sharded s JOIN cte c ON s.parent_id = c.id \ + ) SELECT * FROM cte", + None, + ) + .unwrap(); + // The base case has id = 1 + assert!(result.is_some()); + } + + #[test] + fn test_select_with_union() { + let result = run_test( + "SELECT * FROM sharded WHERE id = 1 UNION SELECT * FROM sharded WHERE id = 2", + None, + ) + .unwrap(); + // UNION queries should find at least one shard + assert!(result.is_some()); + } + + #[test] + fn test_select_with_nested_subselects() { + let result = run_test( + "SELECT * FROM sharded WHERE id IN ( \ + SELECT * FROM other WHERE x IN ( \ + SELECT y FROM third WHERE sharded_id = 1 \ + ) \ + )", + None, + ) + .unwrap(); + // The innermost subquery has sharded_id = 1 + assert!(result.is_some()); + } + + // Bound parameter tests + + #[test] + fn test_bound_simple_select() { + let bind = Bind::new_params("", &[Parameter::new(b"1")]); + let result = run_test("SELECT * FROM sharded WHERE id = $1", Some(&bind)).unwrap(); + assert!(result.is_some()); + } + + #[test] + fn test_bound_select_with_and() { + let bind = Bind::new_params("", &[Parameter::new(b"1")]); + let result = run_test( + "SELECT * FROM sharded WHERE id = $1 AND name = 'foo'", + Some(&bind), + ) + .unwrap(); + assert!(result.is_some()); + } + + #[test] + fn test_bound_select_with_or_returns_none() { + let bind = Bind::new_params("", &[Parameter::new(b"1"), Parameter::new(b"2")]); + let result = run_test( + "SELECT * FROM sharded WHERE id = $1 OR id = $2", + Some(&bind), + ) + .unwrap(); + assert!(result.is_none()); + } + + #[test] + fn test_bound_select_with_subquery() { + let bind = Bind::new_params("", &[Parameter::new(b"1")]); + let result = run_test( + "SELECT * FROM sharded WHERE id IN (SELECT sharded_id FROM other WHERE sharded_id = $1)", + Some(&bind), + ) + .unwrap(); + assert!(result.is_some()); + } + + #[test] + fn test_bound_select_with_cte() { + let bind = Bind::new_params("", &[Parameter::new(b"1")]); + let result = run_test( + "WITH cte AS (SELECT * FROM sharded WHERE id = $1) SELECT * FROM cte", + Some(&bind), + ) + .unwrap(); + assert!(result.is_some()); + } + + #[test] + fn test_bound_select_with_join() { + let bind = Bind::new_params("", &[Parameter::new(b"1")]); + let result = run_test( + "SELECT * FROM sharded s JOIN other o ON s.id = o.sharded_id WHERE s.id = $1", + Some(&bind), + ) + .unwrap(); + assert!(result.is_some()); + } + + #[test] + fn test_bound_select_with_type_cast() { + let bind = Bind::new_params("", &[Parameter::new(b"1")]); + let result = run_test("SELECT * FROM sharded WHERE id = $1::int", Some(&bind)).unwrap(); + assert!(result.is_some()); + } + + #[test] + fn test_bound_select_from_subquery() { + let bind = Bind::new_params("", &[Parameter::new(b"1")]); + let result = run_test( + "SELECT * FROM (SELECT * FROM sharded WHERE id = $1) AS sub", + Some(&bind), + ) + .unwrap(); + assert!(result.is_some()); + } + + #[test] + fn test_bound_select_with_nested_cte() { + let bind = Bind::new_params("", &[Parameter::new(b"1")]); + let result = run_test( + "WITH cte1 AS (SELECT * FROM sharded WHERE id = $1), \ + cte2 AS (SELECT * FROM cte1) \ + SELECT * FROM cte2", + Some(&bind), + ) + .unwrap(); + assert!(result.is_some()); + } + + #[test] + fn test_bound_select_with_in_list() { + let bind = Bind::new_params( + "", + &[ + Parameter::new(b"1"), + Parameter::new(b"2"), + Parameter::new(b"3"), + ], + ); + let result = run_test( + "SELECT * FROM sharded WHERE id IN ($1, $2, $3)", + Some(&bind), + ) + .unwrap(); + assert!(result.is_some()); + } + + #[test] + fn test_bound_select_with_any_array() { + // ANY($1) with an array parameter - $1 is a single array value like '{1,2,3}' + // Array parameters route to all shards since we can't reliably parse them + let bind = Bind::new_params("", &[Parameter::new(b"{1,2,3}")]); + let result = run_test("SELECT * FROM sharded WHERE id = ANY($1)", Some(&bind)).unwrap(); + assert_eq!(result, Some(Shard::All)); + } + + #[test] + fn test_bound_select_with_not_equals_returns_none() { + let bind = Bind::new_params("", &[Parameter::new(b"1")]); + let result = run_test("SELECT * FROM sharded WHERE id != $1", Some(&bind)).unwrap(); + assert!(result.is_none()); + } + + #[test] + fn test_bound_select_with_greater_than_returns_none() { + let bind = Bind::new_params("", &[Parameter::new(b"1")]); + let result = run_test("SELECT * FROM sharded WHERE id > $1", Some(&bind)).unwrap(); + assert!(result.is_none()); + } + + #[test] + fn test_bound_select_with_complex_and() { + let bind = Bind::new_params("", &[Parameter::new(b"1")]); + let result = run_test( + "SELECT * FROM sharded WHERE id = $1 AND status = 'active' AND created_at > now()", + Some(&bind), + ) + .unwrap(); + assert!(result.is_some()); + } + + #[test] + fn test_bound_select_with_left_join() { + let bind = Bind::new_params("", &[Parameter::new(b"1")]); + let result = run_test( + "SELECT * FROM sharded s LEFT JOIN other o ON s.id = o.sharded_id WHERE s.id = $1", + Some(&bind), + ) + .unwrap(); + assert!(result.is_some()); + } + + #[test] + fn test_bound_select_with_multiple_joins() { + let bind = Bind::new_params("", &[Parameter::new(b"1")]); + let result = run_test( + "SELECT * FROM sharded s \ + JOIN other o ON s.id = o.sharded_id \ + JOIN third t ON o.id = t.other_id \ + WHERE s.id = $1", + Some(&bind), + ) + .unwrap(); + assert!(result.is_some()); + } + + #[test] + fn test_bound_select_with_exists_subquery() { + let bind = Bind::new_params("", &[Parameter::new(b"1")]); + let result = run_test( + "SELECT * FROM sharded WHERE EXISTS (SELECT 1 FROM other WHERE sharded_id = $1)", + Some(&bind), + ) + .unwrap(); + assert!(result.is_some()); + } + + #[test] + fn test_bound_select_with_scalar_subquery() { + let bind = Bind::new_params("", &[Parameter::new(b"1")]); + let result = run_test( + "SELECT * FROM sharded WHERE id = (SELECT sharded_id FROM other WHERE sharded_id = $1 LIMIT 1)", + Some(&bind), + ) + .unwrap(); + assert!(result.is_some()); + } + + #[test] + fn test_bound_select_with_recursive_cte() { + let bind = Bind::new_params("", &[Parameter::new(b"1")]); + let result = run_test( + "WITH RECURSIVE cte AS ( \ + SELECT * FROM sharded WHERE id = $1 \ + UNION ALL \ + SELECT s.* FROM sharded s JOIN cte c ON s.parent_id = c.id \ + ) SELECT * FROM cte", + Some(&bind), + ) + .unwrap(); + assert!(result.is_some()); + } + + #[test] + fn test_bound_select_with_union() { + let bind = Bind::new_params("", &[Parameter::new(b"1"), Parameter::new(b"2")]); + let result = run_test( + "SELECT * FROM sharded WHERE id = $1 UNION SELECT * FROM sharded WHERE id = $2", + Some(&bind), + ) + .unwrap(); + assert!(result.is_some()); + } + + #[test] + fn test_bound_select_with_nested_subselects() { + let bind = Bind::new_params("", &[Parameter::new(b"1")]); + let result = run_test( + "SELECT * FROM sharded WHERE id IN ( \ + SELECT * FROM other WHERE x IN ( \ + SELECT y FROM third WHERE sharded_id = $1 \ + ) \ + )", + Some(&bind), + ) + .unwrap(); + assert!(result.is_some()); + } + + #[test] + fn test_bound_select_with_cte_and_subquery() { + let bind = Bind::new_params("", &[Parameter::new(b"1")]); + let result = run_test( + "WITH cte AS (SELECT * FROM sharded WHERE id = $1) \ + SELECT * FROM cte WHERE id IN (SELECT sharded_id FROM other)", + Some(&bind), + ) + .unwrap(); + assert!(result.is_some()); + } + + #[test] + fn test_bound_select_with_multiple_ctes_and_subquery() { + let bind = Bind::new_params("", &[Parameter::new(b"1")]); + let result = run_test( + "WITH cte1 AS (SELECT * FROM sharded WHERE id = $1), \ + cte2 AS (SELECT * FROM other WHERE sharded_id IN (SELECT id FROM cte1)) \ + SELECT * FROM cte2", + Some(&bind), + ) + .unwrap(); + assert!(result.is_some()); + } + + #[test] + fn test_bound_select_with_cte_subquery_and_join() { + let bind = Bind::new_params("", &[Parameter::new(b"1")]); + let result = run_test( + "WITH cte AS (SELECT * FROM sharded WHERE id = $1) \ + SELECT c.*, o.* FROM cte c \ + JOIN other o ON c.id = o.sharded_id \ + WHERE o.x IN (SELECT y FROM third)", + Some(&bind), + ) + .unwrap(); + assert!(result.is_some()); + } + + // Schema-qualified table tests + + #[test] + fn test_select_with_schema_qualified_table() { + let result = run_test( + "SELECT * FROM myschema.schema_sharded WHERE tenant_id = 1", + None, + ) + .unwrap(); + assert!(result.is_some()); + } + + #[test] + fn test_select_with_schema_qualified_alias() { + let result = run_test( + "SELECT * FROM myschema.schema_sharded s WHERE s.tenant_id = 1", + None, + ) + .unwrap(); + assert!(result.is_some()); + } + + #[test] + fn test_bound_select_with_schema_qualified_alias() { + let bind = Bind::new_params("", &[Parameter::new(b"1")]); + let result = run_test( + "SELECT * FROM myschema.schema_sharded s WHERE s.tenant_id = $1", + Some(&bind), + ) + .unwrap(); + assert!(result.is_some()); + } + + #[test] + fn test_select_with_schema_qualified_join() { + let result = run_test( + "SELECT * FROM myschema.schema_sharded s \ + JOIN other o ON s.id = o.sharded_id WHERE s.tenant_id = 1", + None, + ) + .unwrap(); + assert!(result.is_some()); + } + + #[test] + fn test_select_wrong_schema_returns_none() { + let result = run_test( + "SELECT * FROM otherschema.schema_sharded WHERE tenant_id = 1", + None, + ) + .unwrap(); + assert!(result.is_none()); + } + + #[test] + fn test_select_wrong_schema_alias_returns_none() { + let result = run_test( + "SELECT * FROM otherschema.schema_sharded s WHERE s.tenant_id = 1", + None, + ) + .unwrap(); + assert!(result.is_none()); + } + + #[test] + fn test_select_with_any_array_literal() { + let result = run_test("SELECT * FROM sharded WHERE id = ANY('{1, 2, 3}')", None).unwrap(); + // ANY with array literal routes to all shards + assert_eq!(result, Some(Shard::All)); + } + + // UPDATE statement tests + + #[test] + fn test_simple_update() { + let result = run_test("UPDATE sharded SET name = 'foo' WHERE id = 1", None); + assert!(result.unwrap().is_some()); + } + + #[test] + fn test_update_with_bound_param() { + let bind = Bind::new_params("", &[Parameter::new(b"1")]); + let result = run_test("UPDATE sharded SET name = 'foo' WHERE id = $1", Some(&bind)); + assert!(result.unwrap().is_some()); + } + + #[test] + fn test_update_with_and() { + let result = run_test( + "UPDATE sharded SET name = 'foo' WHERE id = 1 AND status = 'active'", + None, + ); + assert!(result.unwrap().is_some()); + } + + #[test] + fn test_update_no_where_returns_none() { + let result = run_test("UPDATE sharded SET name = 'foo'", None); + assert!(result.unwrap().is_none()); + } + + #[test] + fn test_update_with_subquery() { + let result = run_test( + "UPDATE sharded SET name = 'foo' WHERE id IN (SELECT sharded_id FROM other WHERE sharded_id = 1)", + None, + ); + assert!(result.unwrap().is_some()); + } + + // DELETE statement tests + + #[test] + fn test_simple_delete() { + let result = run_test("DELETE FROM sharded WHERE id = 1", None); + assert!(result.unwrap().is_some()); + } + + #[test] + fn test_delete_with_bound_param() { + let bind = Bind::new_params("", &[Parameter::new(b"1")]); + let result = run_test("DELETE FROM sharded WHERE id = $1", Some(&bind)); + assert!(result.unwrap().is_some()); + } + + #[test] + fn test_delete_with_and() { + let result = run_test( + "DELETE FROM sharded WHERE id = 1 AND status = 'active'", + None, + ); + assert!(result.unwrap().is_some()); + } + + #[test] + fn test_delete_no_where_returns_none() { + let result = run_test("DELETE FROM sharded", None); + assert!(result.unwrap().is_none()); + } + + #[test] + fn test_delete_with_subquery() { + let result = run_test( + "DELETE FROM sharded WHERE id IN (SELECT sharded_id FROM other WHERE sharded_id = 1)", + None, + ); + assert!(result.unwrap().is_some()); + } + + #[test] + fn test_delete_with_cte() { + let result = run_test( + "WITH to_delete AS (SELECT id FROM sharded WHERE id = 1) DELETE FROM sharded WHERE id IN (SELECT id FROM to_delete)", + None, + ); + assert!(result.unwrap().is_some()); + } + + // INSERT statement tests + + #[test] + fn test_simple_insert_with_value() { + let result = run_test("INSERT INTO sharded (id, name) VALUES (1, 'foo')", None); + assert!(result.unwrap().is_some()); + } + + #[test] + fn test_insert_with_bound_param() { + let bind = Bind::new_params("", &[Parameter::new(b"1"), Parameter::new(b"foo")]); + let result = run_test( + "INSERT INTO sharded (id, name) VALUES ($1, $2)", + Some(&bind), + ); + assert!(result.unwrap().is_some()); + } + + #[test] + fn test_insert_with_subquery_in_values() { + let result = run_test( + "INSERT INTO sharded (id, name) VALUES ((SELECT sharded_id FROM other WHERE sharded_id = 1), 'foo')", + None, + ); + assert!(result.unwrap().is_some()); + } + + #[test] + fn test_insert_with_subquery_in_values_param() { + let bind = Bind::new_params("", &[Parameter::new(b"1")]); + let result = run_test( + "INSERT INTO sharded (id, name) VALUES ((SELECT sharded_id FROM other WHERE sharded_id = $1), 'foo')", + Some(&bind), + ); + assert!(result.unwrap().is_some()); + } + + #[test] + fn test_insert_select() { + let result = run_test( + "INSERT INTO sharded (id, name) SELECT sharded_id, name FROM other WHERE sharded_id = 1", + None, + ); + assert!(result.unwrap().is_some()); + } + + #[test] + fn test_insert_select_with_param() { + let bind = Bind::new_params("", &[Parameter::new(b"1")]); + let result = run_test( + "INSERT INTO sharded (id, name) SELECT sharded_id, name FROM other WHERE sharded_id = $1", + Some(&bind), + ); + assert!(result.unwrap().is_some()); + } + + #[test] + fn test_insert_with_cte() { + let result = run_test( + "WITH src AS (SELECT id, name FROM sharded WHERE id = 1) INSERT INTO sharded (id, name) SELECT id, name FROM src", + None, + ); + assert!(result.unwrap().is_some()); + } + + #[test] + fn test_insert_no_sharding_key_returns_none() { + let result = run_test("INSERT INTO sharded (name) VALUES ('foo')", None); + assert!(result.unwrap().is_none()); + } + + // Schema-based sharding fallback tests + use crate::backend::replication::ShardedSchemas; + use pgdog_config::sharding::ShardedSchema; + + fn run_test_with_schemas(stmt: &str, bind: Option<&Bind>) -> Result, Error> { + let schema = ShardingSchema { + shards: 3, + tables: ShardedTables::new( + vec![ShardedTable { + column: "id".into(), + name: Some("sharded".into()), + ..Default::default() + }], + vec![], + ), + schemas: ShardedSchemas::new(vec![ + ShardedSchema { + database: "test".to_string(), + name: Some("sales".to_string()), + shard: 1, + all: false, + }, + ShardedSchema { + database: "test".to_string(), + name: Some("inventory".to_string()), + shard: 2, + all: false, + }, + ]), + ..Default::default() + }; + let raw = pg_query::parse(stmt) + .unwrap() + .protobuf + .stmts + .first() + .cloned() + .unwrap(); + let mut parser = StatementParser::from_raw(&raw, bind, &schema, None)?; + parser.shard() + } + + #[test] + fn test_schema_sharding_select_fallback() { + // No sharding key in WHERE clause, but table is in a sharded schema + let result = run_test_with_schemas("SELECT * FROM sales.products", None).unwrap(); + assert_eq!(result, Some(Shard::Direct(1))); + } + + #[test] + fn test_schema_sharding_select_with_join() { + // JOIN between tables in the same sharded schema + let result = run_test_with_schemas( + "SELECT * FROM sales.products p JOIN sales.orders o ON p.id = o.product_id", + None, + ) + .unwrap(); + assert_eq!(result, Some(Shard::Direct(1))); + } + + #[test] + fn test_schema_sharding_update_fallback() { + // No sharding key in WHERE clause, but table is in a sharded schema + let result = run_test_with_schemas("UPDATE sales.products SET name = 'foo'", None).unwrap(); + assert_eq!(result, Some(Shard::Direct(1))); + } + + #[test] + fn test_schema_sharding_delete_fallback() { + // No sharding key in WHERE clause, but table is in a sharded schema + let result = run_test_with_schemas("DELETE FROM sales.products", None).unwrap(); + assert_eq!(result, Some(Shard::Direct(1))); + } + + #[test] + fn test_schema_sharding_insert_fallback() { + // No sharding key in values, but table is in a sharded schema + let result = + run_test_with_schemas("INSERT INTO sales.products (name) VALUES ('foo')", None) + .unwrap(); + assert_eq!(result, Some(Shard::Direct(1))); + } + + #[test] + fn test_schema_sharding_with_subquery() { + // Subquery references table in a sharded schema + let result = run_test_with_schemas( + "SELECT * FROM unsharded WHERE id IN (SELECT id FROM sales.products)", + None, + ) + .unwrap(); + assert_eq!(result, Some(Shard::Direct(1))); + } + + #[test] + fn test_schema_sharding_with_cte() { + // CTE references table in a sharded schema + let result = run_test_with_schemas( + "WITH cte AS (SELECT * FROM sales.products) SELECT * FROM cte", + None, + ) + .unwrap(); + assert_eq!(result, Some(Shard::Direct(1))); + } + + #[test] + fn test_key_sharding_takes_precedence() { + // Both key-based and schema-based sharding could match, + // but key-based should take precedence + let result = run_test_with_schemas("SELECT * FROM sharded WHERE id = 1", None).unwrap(); + // Key-based sharding returns a shard (not necessarily shard 1) + assert!(result.is_some()); + } + + #[test] + fn test_no_schema_no_key_returns_none() { + // Table not in sharded schema and no sharding key + let result = run_test_with_schemas("SELECT * FROM public.unknown", None).unwrap(); + assert!(result.is_none()); + } + + #[test] + fn test_schema_sharding_different_schemas() { + // Different sharded schemas route to different shards + let result1 = run_test_with_schemas("SELECT * FROM sales.products", None).unwrap(); + let result2 = run_test_with_schemas("SELECT * FROM inventory.items", None).unwrap(); + assert_eq!(result1, Some(Shard::Direct(1))); + assert_eq!(result2, Some(Shard::Direct(2))); + } +} diff --git a/pgdog/src/frontend/router/parser/table.rs b/pgdog/src/frontend/router/parser/table.rs index 85c4fbc10..cda74e4ab 100644 --- a/pgdog/src/frontend/router/parser/table.rs +++ b/pgdog/src/frontend/router/parser/table.rs @@ -1,11 +1,15 @@ use std::fmt::Display; -use pg_query::{protobuf::RangeVar, Node, NodeEnum}; +use pg_query::{ + protobuf::{List, RangeVar}, + Node, NodeEnum, +}; +use super::{Error, Schema}; use crate::util::escape_identifier; /// Table name in a query. -#[derive(Debug, Clone, Copy, PartialEq, Default)] +#[derive(Debug, Clone, Copy, PartialEq, Default, Hash, Eq)] pub struct Table<'a> { /// Table name. pub name: &'a str, @@ -50,6 +54,10 @@ impl<'a> Table<'a> { pub fn name_match(&self, name: &str) -> bool { Some(name) == self.alias || name == self.name } + + pub fn schema(&self) -> Option> { + self.schema.map(|s| s.into()) + } } impl Display for OwnedTable { @@ -92,20 +100,56 @@ impl<'a> TryFrom<&'a Node> for Table<'a> { } impl<'a> TryFrom<&'a Vec> for Table<'a> { - type Error = (); + type Error = Error; fn try_from(value: &'a Vec) -> Result { - let table = value - .first() - .and_then(|node| { - node.node.as_ref().map(|node| match node { - NodeEnum::RangeVar(var) => Some(Table::from(var)), - _ => None, - }) - }) - .flatten() - .ok_or(())?; - Ok(table) + match value.len() { + 1 => { + let table = value + .first() + .and_then(|node| { + node.node.as_ref().map(|node| match node { + NodeEnum::RangeVar(var) => Some(Ok(Table::from(var))), + NodeEnum::List(list) => Some(Table::try_from(list)), + NodeEnum::String(str) => Some(Ok(Table::from(str.sval.as_str()))), + _ => None, + }) + }) + .flatten() + .ok_or(Error::TableDecode)?; + return table; + } + + 2 => { + let schema = value.iter().next().unwrap().node.as_ref().and_then(|node| { + if let NodeEnum::String(sval) = node { + Some(sval.sval.as_str()) + } else { + None + } + }); + let table = value.iter().last().unwrap().node.as_ref().and_then(|node| { + if let NodeEnum::String(sval) = node { + Some(sval.sval.as_str()) + } else { + None + } + }); + if let Some(schema) = schema { + if let Some(table) = table { + return Ok(Table { + name: table, + schema: Some(schema), + alias: None, + }); + } + } + } + + _ => (), + } + + Err(Error::TableDecode) } } @@ -127,3 +171,49 @@ impl<'a> From<&'a RangeVar> for Table<'a> { } } } + +impl<'a> TryFrom<&'a List> for Table<'a> { + type Error = Error; + + fn try_from(value: &'a List) -> Result { + fn str_value(list: &List, pos: usize) -> Option<&str> { + if let Some(NodeEnum::String(ref schema)) = list.items.get(pos).unwrap().node { + Some(schema.sval.as_str()) + } else { + None + } + } + + match value.items.len() { + 2 => { + let schema = str_value(value, 0); + let name = str_value(value, 1).ok_or(Error::TableDecode)?; + Ok(Table { + schema, + name, + alias: None, + }) + } + + 1 => { + let name = str_value(value, 0).ok_or(Error::TableDecode)?; + Ok(Table { + schema: None, + name, + alias: None, + }) + } + + _ => Err(Error::TableDecode), + } + } +} + +impl<'a> From<&'a str> for Table<'a> { + fn from(value: &'a str) -> Self { + Table { + name: value, + ..Default::default() + } + } +} diff --git a/pgdog/src/frontend/router/parser/tuple.rs b/pgdog/src/frontend/router/parser/tuple.rs index 40c5dc6f0..70c9e2b30 100644 --- a/pgdog/src/frontend/router/parser/tuple.rs +++ b/pgdog/src/frontend/router/parser/tuple.rs @@ -20,8 +20,23 @@ impl<'a> TryFrom<&'a List> for Tuple<'a> { let mut values = vec![]; for value in &value.items { - let value = value.try_into()?; - values.push(value); + if let Ok(value) = Value::try_from(value) { + values.push(value); + } else { + // FIXME: + // + // This basically makes all values we can't parse NULL. + // Normally, the result of that is the query is sent to all + // shards, quietly. + // + // I think the right thing here is to throw an error, + // but more likely it'll be a value we don't actually need for sharding. + // + // We should check if its value we actually need and only then + // throw an error. + // + values.push(Value::Null); + } } Ok(Self { values }) diff --git a/pgdog/src/frontend/router/parser/value.rs b/pgdog/src/frontend/router/parser/value.rs index af9249dcb..8f46759c2 100644 --- a/pgdog/src/frontend/router/parser/value.rs +++ b/pgdog/src/frontend/router/parser/value.rs @@ -1,22 +1,45 @@ //! Value extracted from a query. +use std::fmt::Display; + use pg_query::{ protobuf::{a_const::Val, *}, NodeEnum, }; -use crate::net::messages::Vector; +use crate::net::{messages::Vector, vector::str_to_vector}; /// A value extracted from a query. #[derive(Debug, Clone, PartialEq)] pub enum Value<'a> { String(&'a str), Integer(i64), + Float(f64), Boolean(bool), Null, Placeholder(i32), Vector(Vector), - Function(&'a str), +} + +impl Display for Value<'_> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::String(s) => write!(f, "'{}'", s.replace("'", "''")), + Self::Integer(i) => write!(f, "{}", i), + Self::Float(s) => write!(f, "{}", s), + Self::Null => write!(f, "NULL"), + Self::Boolean(b) => write!(f, "{}", if *b { "true" } else { "false" }), + Self::Vector(v) => write!( + f, + "{}", + v.iter() + .map(|v| v.to_string()) + .collect::>() + .join(",") + ), + Self::Placeholder(p) => write!(f, "${}", p), + } + } } impl Value<'_> { @@ -39,7 +62,7 @@ impl<'a> From<&'a AConst> for Value<'a> { match value.val.as_ref() { Some(Val::Sval(s)) => { if s.sval.starts_with('[') && s.sval.ends_with(']') { - if let Ok(vector) = Vector::try_from(s.sval.as_str()) { + if let Ok(vector) = str_to_vector(s.sval.as_str()) { Value::Vector(vector) } else { Value::String(s.sval.as_str()) @@ -54,13 +77,21 @@ impl<'a> From<&'a AConst> for Value<'a> { Some(Val::Boolval(b)) => Value::Boolean(b.boolval), Some(Val::Ival(i)) => Value::Integer(i.ival as i64), Some(Val::Fval(Float { fval })) => { - if let Ok(val) = fval.parse() { - Value::Integer(val) + if fval.contains(".") { + if let Ok(float) = fval.parse() { + Value::Float(float) + } else { + Value::String(fval.as_str()) + } } else { - Value::Null + match fval.parse::() { + Ok(i) => Value::Integer(i), // Integers over 2.2B and under -2.2B are sent as "floats" + Err(_) => Value::String(fval.as_str()), + } } } - _ => Value::Null, + Some(Val::Bsval(bsval)) => Value::String(bsval.bsval.as_str()), + None => Value::Null, } } } @@ -77,28 +108,39 @@ impl<'a> TryFrom<&'a Option> for Value<'a> { type Error = (); fn try_from(value: &'a Option) -> Result { - match value { - Some(NodeEnum::AConst(a_const)) => Ok(a_const.into()), - Some(NodeEnum::ParamRef(param_ref)) => Ok(Value::Placeholder(param_ref.number)), - Some(NodeEnum::FuncCall(func)) => { - if let Some(Node { - node: Some(NodeEnum::String(sval)), - }) = func.funcname.first() - { - Ok(Value::Function(&sval.sval)) - } else { - Ok(Value::Null) - } - } + Ok(match value { + Some(NodeEnum::AConst(a_const)) => a_const.into(), + Some(NodeEnum::ParamRef(param_ref)) => Value::Placeholder(param_ref.number), Some(NodeEnum::TypeCast(cast)) => { if let Some(ref arg) = cast.arg { - Value::try_from(&arg.node) + Value::try_from(&arg.node)? } else { - Ok(Value::Null) + Value::Null } } - _ => Ok(Value::Null), - } + + Some(NodeEnum::AExpr(expr)) => { + if expr.kind() == AExprKind::AexprOp { + if let Some(Node { + node: Some(NodeEnum::String(pg_query::protobuf::String { sval })), + }) = expr.name.first() + { + if sval == "-" { + if let Some(ref node) = expr.rexpr { + let value = Value::try_from(&node.node)?; + if let Value::Float(float) = value { + return Ok(Value::Float(-float)); + } + } + } + } + } + + return Err(()); + } + + _ => return Err(()), + }) } } @@ -121,4 +163,34 @@ mod test { let vector = Value::try_from(&node).unwrap(); assert_eq!(vector.vector().unwrap()[0], 1.0.into()); } + + #[test] + fn test_negative_numeric_with_cast() { + let stmt = + pg_query::parse("INSERT INTO t (id, val) VALUES (2, -987654321.123456789::NUMERIC)") + .unwrap(); + + let insert = match stmt.protobuf.stmts[0].stmt.as_ref().unwrap().node.as_ref() { + Some(NodeEnum::InsertStmt(insert)) => insert, + _ => panic!("expected InsertStmt"), + }; + + let select = insert.select_stmt.as_ref().unwrap(); + let values = match select.node.as_ref() { + Some(NodeEnum::SelectStmt(s)) => &s.values_lists, + _ => panic!("expected SelectStmt"), + }; + + // values_lists[0] is a List node containing the tuple items + let tuple = match values[0].node.as_ref() { + Some(NodeEnum::List(list)) => &list.items, + _ => panic!("expected List"), + }; + + // Second value in the VALUES tuple is our negative numeric + let neg_numeric_node = &tuple[1]; + let value = Value::try_from(&neg_numeric_node.node).unwrap(); + + assert_eq!(value, Value::Float(-987654321.123456789)); + } } diff --git a/pgdog/src/frontend/router/parser/where_clause.rs b/pgdog/src/frontend/router/parser/where_clause.rs index 0485eb3ca..970ce3a8e 100644 --- a/pgdog/src/frontend/router/parser/where_clause.rs +++ b/pgdog/src/frontend/router/parser/where_clause.rs @@ -10,7 +10,7 @@ use crate::frontend::router::parser::{from_clause::FromClause, Table}; use super::Key; -#[derive(Copy, Clone)] +#[derive(Copy, Clone, Debug)] pub enum TablesSource<'a> { Table(Table<'a>), FromClause(FromClause<'a>), @@ -54,6 +54,13 @@ impl<'a> TablesSource<'a> { Self::FromClause(fc) => fc.table_name(), } } + + pub fn tables(&'a self) -> Vec> { + match self { + Self::Table(table) => vec![*table], + Self::FromClause(from_clause) => from_clause.tables(), + } + } } #[derive(Debug)] diff --git a/pgdog/src/frontend/router/search_path.rs b/pgdog/src/frontend/router/search_path.rs index 9234d9261..741362945 100644 --- a/pgdog/src/frontend/router/search_path.rs +++ b/pgdog/src/frontend/router/search_path.rs @@ -1,7 +1,4 @@ -use crate::{ - backend::Schema, - net::{parameter::ParameterValue, Parameters}, -}; +use crate::{backend::Schema, net::parameter::ParameterValue}; #[derive(Debug)] pub struct SearchPath<'a> { @@ -24,9 +21,13 @@ impl<'a> SearchPath<'a> { schemas } - pub(crate) fn new(user: &'a str, params: &'a Parameters, schema: &'a Schema) -> Self { + pub(crate) fn new( + user: &'a str, + search_path: Option<&'a ParameterValue>, + schema: &'a Schema, + ) -> Self { let default_path = schema.search_path(); - let search_path = match params.get("search_path") { + let search_path = match search_path { Some(ParameterValue::Tuple(overriden)) => overriden.as_slice(), _ => default_path, }; diff --git a/pgdog/src/frontend/router/sharding/context.rs b/pgdog/src/frontend/router/sharding/context.rs index 6b62d905d..a7b0ac777 100644 --- a/pgdog/src/frontend/router/sharding/context.rs +++ b/pgdog/src/frontend/router/sharding/context.rs @@ -1,5 +1,5 @@ use crate::frontend::router::parser::Shard; -use tracing::debug; +use tracing::trace; use super::{Error, Hasher, Operator, Value}; @@ -14,7 +14,7 @@ impl Context<'_> { pub fn apply(&self) -> Result { match &self.operator { Operator::Shards(shards) => { - debug!("sharding using hash"); + trace!("sharding using hash"); if let Some(hash) = self.value.hash(self.hasher)? { return Ok(Shard::Direct(hash as usize % shards)); } @@ -25,19 +25,19 @@ impl Context<'_> { probes, centroids, } => { - debug!("sharding using k-means"); + trace!("sharding using k-means"); if let Some(vector) = self.value.vector()? { - return Ok(centroids.shard(&vector, *shards, *probes)); + return Ok(centroids.shard(&vector, *shards, *probes).into()); } } Operator::Range(ranges) => { - debug!("sharding using range"); + trace!("sharding using range"); return ranges.shard(&self.value); } Operator::List(lists) => { - debug!("sharding using lists"); + trace!("sharding using lists"); return lists.shard(&self.value); } } diff --git a/pgdog/src/frontend/router/sharding/context_builder.rs b/pgdog/src/frontend/router/sharding/context_builder.rs index 7f9450184..d52edcec3 100644 --- a/pgdog/src/frontend/router/sharding/context_builder.rs +++ b/pgdog/src/frontend/router/sharding/context_builder.rs @@ -1,3 +1,8 @@ +//! Context builder for sharding a value. +//! +//! Manages mapping a value (integer, string, etc.) +//! to a shard number, given a sharded mapping in pgdog.toml. +//! use crate::{ backend::ShardingSchema, config::{DataType, Hasher as HasherConfig, ShardedTable}, @@ -6,6 +11,7 @@ use crate::{ use super::{Centroids, Context, Data, Error, Hasher, Lists, Operator, Ranges, Value}; +/// Sharding context builder. #[derive(Debug)] pub struct ContextBuilder<'a> { data_type: DataType, @@ -21,6 +27,8 @@ pub struct ContextBuilder<'a> { } impl<'a> ContextBuilder<'a> { + /// Create new context builder from a sharded table + /// in the config. pub fn new(table: &'a ShardedTable) -> Self { Self { data_type: table.data_type, @@ -42,7 +50,8 @@ impl<'a> ContextBuilder<'a> { } } - /// Infer sharding function from config. + /// Infer sharding function from config, iff + /// only one sharding function is configured. pub fn infer_from_from_and_config( value: &'a str, sharding_schema: &'a ShardingSchema, @@ -141,6 +150,7 @@ impl<'a> ContextBuilder<'a> { } } + /// Set the number of shards in the configuration. pub fn shards(mut self, shards: usize) -> Self { if let Some(centroids) = self.centroids.take() { self.operator = Some(Operator::Centroids { @@ -196,12 +206,13 @@ mod test { shards: 2, tables: ShardedTables::new( vec![ShardedTable { - data_type: DataType::Varchar, mapping: None, + data_type: DataType::Varchar, ..Default::default() }], vec![], ), + ..Default::default() }; let ctx = ContextBuilder::infer_from_from_and_config("test_value", &schema) @@ -231,6 +242,7 @@ mod test { }], vec![], ), + ..Default::default() }; let ctx = ContextBuilder::infer_from_from_and_config("15", &schema) @@ -261,6 +273,7 @@ mod test { }], vec![], ), + ..Default::default() }; let builder = ContextBuilder::infer_from_from_and_config("1", &schema).unwrap(); diff --git a/pgdog/src/frontend/router/sharding/distance_simd_rust.rs b/pgdog/src/frontend/router/sharding/distance_simd_rust.rs index f9e1d9792..812cd46c5 100644 --- a/pgdog/src/frontend/router/sharding/distance_simd_rust.rs +++ b/pgdog/src/frontend/router/sharding/distance_simd_rust.rs @@ -1,242 +1,10 @@ -use crate::net::messages::data_types::Float; - -#[cfg(target_arch = "x86_64")] -use std::arch::x86_64::*; - -#[cfg(target_arch = "aarch64")] -use std::arch::aarch64::*; - -/// Helper function to convert Float slice to f32 slice -/// SAFETY: Float is a newtype wrapper around f32, so the memory layout is identical -#[inline(always)] -unsafe fn float_slice_to_f32(floats: &[Float]) -> &[f32] { - // This is safe because Float is a transparent wrapper around f32 - std::slice::from_raw_parts(floats.as_ptr() as *const f32, floats.len()) -} - -/// Scalar reference implementation - no allocations -#[inline] -pub fn euclidean_distance_scalar(p: &[Float], q: &[Float]) -> f32 { - debug_assert_eq!(p.len(), q.len()); - - let mut sum = 0.0f32; - for i in 0..p.len() { - let diff = q[i].0 - p[i].0; - sum += diff * diff; - } - sum.sqrt() -} - -/// SIMD implementation for x86_64 with SSE - optimized version -#[cfg(all(target_arch = "x86_64", target_feature = "sse"))] -#[inline] -pub fn euclidean_distance_sse(p: &[Float], q: &[Float]) -> f32 { - debug_assert_eq!(p.len(), q.len()); - - unsafe { - // Convert Float slices to f32 slices to avoid temporary arrays - let p_f32 = float_slice_to_f32(p); - let q_f32 = float_slice_to_f32(q); - - let mut sum1 = _mm_setzero_ps(); - let mut sum2 = _mm_setzero_ps(); - let chunks = p.len() / 8; // Process 8 at a time (2 SSE vectors) - - // Unroll loop to process 8 floats per iteration - for i in 0..chunks { - let idx = i * 8; - - // First 4 floats - direct load from slice - let p_vec1 = _mm_loadu_ps(p_f32.as_ptr().add(idx)); - let q_vec1 = _mm_loadu_ps(q_f32.as_ptr().add(idx)); - let diff1 = _mm_sub_ps(q_vec1, p_vec1); - sum1 = _mm_add_ps(sum1, _mm_mul_ps(diff1, diff1)); - - // Second 4 floats - let p_vec2 = _mm_loadu_ps(p_f32.as_ptr().add(idx + 4)); - let q_vec2 = _mm_loadu_ps(q_f32.as_ptr().add(idx + 4)); - let diff2 = _mm_sub_ps(q_vec2, p_vec2); - sum2 = _mm_add_ps(sum2, _mm_mul_ps(diff2, diff2)); - } - - // Combine accumulators - let sum = _mm_add_ps(sum1, sum2); - - // More efficient horizontal sum using shuffles - let shuf = _mm_shuffle_ps(sum, sum, 0b01_00_11_10); // [2,3,0,1] - let sums = _mm_add_ps(sum, shuf); - let shuf = _mm_shuffle_ps(sums, sums, 0b10_11_00_01); // [1,0,3,2] - let result = _mm_add_ps(sums, shuf); - let mut total = _mm_cvtss_f32(result); - - // Handle remaining elements - for i in (chunks * 8)..p.len() { - let diff = q[i].0 - p[i].0; - total += diff * diff; - } - - total.sqrt() - } -} - -/// SIMD implementation for x86_64 with AVX2 - optimized version -#[cfg(all(target_arch = "x86_64", target_feature = "avx2"))] -#[inline] -pub fn euclidean_distance_avx2(p: &[Float], q: &[Float]) -> f32 { - debug_assert_eq!(p.len(), q.len()); - - unsafe { - // Convert Float slices to f32 slices to avoid temporary arrays - let p_f32 = float_slice_to_f32(p); - let q_f32 = float_slice_to_f32(q); - - let mut sum1 = _mm256_setzero_ps(); - let mut sum2 = _mm256_setzero_ps(); - let chunks = p.len() / 16; // Process 16 at a time (2 AVX2 vectors) - - // Unroll loop to process 16 floats per iteration - for i in 0..chunks { - let idx = i * 16; - - // First 8 floats - direct load from slice - let p_vec1 = _mm256_loadu_ps(p_f32.as_ptr().add(idx)); - let q_vec1 = _mm256_loadu_ps(q_f32.as_ptr().add(idx)); - let diff1 = _mm256_sub_ps(q_vec1, p_vec1); - - #[cfg(target_feature = "fma")] - { - sum1 = _mm256_fmadd_ps(diff1, diff1, sum1); - } - #[cfg(not(target_feature = "fma"))] - { - sum1 = _mm256_add_ps(sum1, _mm256_mul_ps(diff1, diff1)); - } - - // Second 8 floats - let p_vec2 = _mm256_loadu_ps(p_f32.as_ptr().add(idx + 8)); - let q_vec2 = _mm256_loadu_ps(q_f32.as_ptr().add(idx + 8)); - let diff2 = _mm256_sub_ps(q_vec2, p_vec2); - - #[cfg(target_feature = "fma")] - { - sum2 = _mm256_fmadd_ps(diff2, diff2, sum2); - } - #[cfg(not(target_feature = "fma"))] - { - sum2 = _mm256_add_ps(sum2, _mm256_mul_ps(diff2, diff2)); - } - } - - // Combine accumulators - let sum = _mm256_add_ps(sum1, sum2); - - // More efficient horizontal sum using extract and SSE - let high = _mm256_extractf128_ps(sum, 1); - let low = _mm256_castps256_ps128(sum); - let sum128 = _mm_add_ps(low, high); - - // Use SSE horizontal sum (more efficient than storing to array) - let shuf = _mm_shuffle_ps(sum128, sum128, 0b01_00_11_10); - let sums = _mm_add_ps(sum128, shuf); - let shuf = _mm_shuffle_ps(sums, sums, 0b10_11_00_01); - let result = _mm_add_ps(sums, shuf); - let mut total = _mm_cvtss_f32(result); - - // Handle remaining elements - for i in (chunks * 16)..p.len() { - let diff = q[i].0 - p[i].0; - total += diff * diff; - } - - total.sqrt() - } -} - -/// SIMD implementation for ARM NEON - optimized version -#[cfg(target_arch = "aarch64")] -#[inline] -pub fn euclidean_distance_neon(p: &[Float], q: &[Float]) -> f32 { - debug_assert_eq!(p.len(), q.len()); - - unsafe { - // Convert Float slices to f32 slices to avoid temporary arrays - let p_f32 = float_slice_to_f32(p); - let q_f32 = float_slice_to_f32(q); - - let mut sum1 = vdupq_n_f32(0.0); - let mut sum2 = vdupq_n_f32(0.0); - let chunks = p.len() / 8; // Process 8 at a time (2 vectors) - - // Unroll loop to process 8 floats per iteration (2x4) - for i in 0..chunks { - let idx = i * 8; - - // First 4 floats - direct load from slice - let p_vec1 = vld1q_f32(p_f32.as_ptr().add(idx)); - let q_vec1 = vld1q_f32(q_f32.as_ptr().add(idx)); - let diff1 = vsubq_f32(q_vec1, p_vec1); - sum1 = vfmaq_f32(sum1, diff1, diff1); - - // Second 4 floats - let p_vec2 = vld1q_f32(p_f32.as_ptr().add(idx + 4)); - let q_vec2 = vld1q_f32(q_f32.as_ptr().add(idx + 4)); - let diff2 = vsubq_f32(q_vec2, p_vec2); - sum2 = vfmaq_f32(sum2, diff2, diff2); - } - - // Combine the two accumulators - let sum = vaddq_f32(sum1, sum2); - - // Horizontal sum - more efficient version - let sum_pairs = vpaddq_f32(sum, sum); - let sum_final = vpaddq_f32(sum_pairs, sum_pairs); - let mut total = vgetq_lane_f32(sum_final, 0); - - // Handle remaining elements - for i in (chunks * 8)..p.len() { - let diff = q[i].0 - p[i].0; - total += diff * diff; - } - - total.sqrt() - } -} - -/// Auto-select best implementation based on CPU features -#[inline] -pub fn euclidean_distance(p: &[Float], q: &[Float]) -> f32 { - #[cfg(all(target_arch = "x86_64", target_feature = "avx2"))] - { - euclidean_distance_avx2(p, q) - } - - #[cfg(all( - target_arch = "x86_64", - target_feature = "sse", - not(target_feature = "avx2") - ))] - { - euclidean_distance_sse(p, q) - } - - #[cfg(target_arch = "aarch64")] - { - euclidean_distance_neon(p, q) - } - - #[cfg(not(any( - all(target_arch = "x86_64", target_feature = "sse"), - target_arch = "aarch64" - )))] - { - euclidean_distance_scalar(p, q) - } -} +pub use pgdog_vector::distance_simd_rust::*; #[cfg(test)] mod tests { - use super::*; - use crate::net::messages::Vector; + + use pgdog_vector::distance_simd_rust::*; + use pgdog_vector::{Float, Vector}; #[test] fn test_no_allocations() { diff --git a/pgdog/src/frontend/router/sharding/error.rs b/pgdog/src/frontend/router/sharding/error.rs index eae69ba34..6973a0908 100644 --- a/pgdog/src/frontend/router/sharding/error.rs +++ b/pgdog/src/frontend/router/sharding/error.rs @@ -1,11 +1,11 @@ -use std::{array::TryFromSliceError, ffi::NulError, num::ParseIntError}; +use std::{array::TryFromSliceError, ffi::NulError}; use thiserror::Error; #[derive(Debug, Error)] pub enum Error { #[error("{0}")] - Parse(#[from] ParseIntError), + ParseInt(String), #[error("{0}")] Size(#[from] TryFromSliceError), @@ -39,4 +39,7 @@ pub enum Error { #[error("sharding key value isn't valid")] InvalidValue, + + #[error("config error: {0}")] + ConfigError(#[from] pgdog_config::Error), } diff --git a/pgdog/src/frontend/router/sharding/list.rs b/pgdog/src/frontend/router/sharding/list.rs index 7fd4ae935..8cd15ee63 100644 --- a/pgdog/src/frontend/router/sharding/list.rs +++ b/pgdog/src/frontend/router/sharding/list.rs @@ -1,9 +1,8 @@ -use std::collections::{hash_map::DefaultHasher, HashMap}; -use std::hash::{Hash, Hasher}; - use super::{Error, Mapping, Shard, Value}; -use crate::config::{FlexibleType, ShardedMapping, ShardedMappingKind}; +use crate::config::FlexibleType; + +pub use pgdog_config::sharding::ListShards; #[derive(Debug, PartialEq, Eq)] pub struct Lists<'a> { @@ -25,59 +24,14 @@ impl<'a> Lists<'a> { let uuid = value.uuid()?; if let Some(integer) = integer { - self.list.shard(&FlexibleType::Integer(integer)) + Ok(self.list.shard(&FlexibleType::Integer(integer))?.into()) } else if let Some(uuid) = uuid { - self.list.shard(&FlexibleType::Uuid(uuid)) + Ok(self.list.shard(&FlexibleType::Uuid(uuid))?.into()) } else if let Some(varchar) = varchar { - self.list.shard(&FlexibleType::String(varchar.to_string())) - } else { - Ok(Shard::All) - } - } -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct ListShards { - mapping: HashMap, -} - -impl Hash for ListShards { - fn hash(&self, state: &mut H) { - // Hash the mapping in a deterministic way by XORing individual key-value hashes - let mut mapping_hash = 0u64; - for (key, value) in &self.mapping { - let mut hasher = DefaultHasher::new(); - key.hash(&mut hasher); - value.hash(&mut hasher); - mapping_hash ^= hasher.finish(); - } - mapping_hash.hash(state); - } -} - -impl ListShards { - pub fn is_empty(&self) -> bool { - self.mapping.is_empty() - } - - pub fn new(mappings: &[ShardedMapping]) -> Self { - let mut mapping = HashMap::new(); - - for map in mappings - .iter() - .filter(|m| m.kind == ShardedMappingKind::List) - { - for value in &map.values { - mapping.insert(value.clone(), map.shard); - } - } - - Self { mapping } - } - - pub fn shard(&self, value: &FlexibleType) -> Result { - if let Some(shard) = self.mapping.get(value) { - Ok(Shard::Direct(*shard)) + Ok(self + .list + .shard(&FlexibleType::String(varchar.to_string()))? + .into()) } else { Ok(Shard::All) } diff --git a/pgdog/src/frontend/router/sharding/mapping.rs b/pgdog/src/frontend/router/sharding/mapping.rs index 894c34cf2..b7b51d67c 100644 --- a/pgdog/src/frontend/router/sharding/mapping.rs +++ b/pgdog/src/frontend/router/sharding/mapping.rs @@ -1,39 +1,12 @@ -use super::ListShards; -use crate::{ - config::{ShardedMapping, ShardedMappingKind}, - frontend::router::Ranges, -}; +pub use pgdog_config::sharding::Mapping; -#[derive(Debug, Clone, Eq, PartialEq, Hash)] -pub enum Mapping { - Range(Vec), // TODO: optimize with a BTreeMap. - List(ListShards), // Optimized. -} - -impl Mapping { - pub fn new(mappings: &[ShardedMapping]) -> Option { - let range = mappings - .iter() - .filter(|m| m.kind == ShardedMappingKind::Range) - .cloned() - .collect::>(); - let list = mappings.iter().any(|m| m.kind == ShardedMappingKind::List); - - if !range.is_empty() { - Some(Self::Range(range)) - } else if list { - Some(Self::List(ListShards::new(mappings))) - } else { - None - } - } +use crate::frontend::router::Ranges; - pub fn valid(&self) -> bool { - match self { - Self::Range(_) => Ranges::new(&Some(self.clone())) - .map(|r| r.valid()) - .unwrap_or(false), - Self::List(_) => true, - } +pub fn mapping_valid(mapping: &Mapping) -> bool { + match mapping { + Mapping::Range(_) => Ranges::new(&Some(mapping.clone())) + .map(|r| r.valid()) + .unwrap_or(false), + Mapping::List(_) => true, } } diff --git a/pgdog/src/frontend/router/sharding/mod.rs b/pgdog/src/frontend/router/sharding/mod.rs index 3c36b4857..e4ddd271a 100644 --- a/pgdog/src/frontend/router/sharding/mod.rs +++ b/pgdog/src/frontend/router/sharding/mod.rs @@ -3,7 +3,10 @@ use uuid::Uuid; use crate::{ backend::ShardingSchema, config::{DataType, ShardedTable}, - net::messages::{Format, FromDataType, ParameterWithFormat, Vector}, + net::{ + messages::{Format, FromDataType, ParameterWithFormat, Vector}, + vector::str_to_vector, + }, }; // pub mod context; @@ -19,6 +22,7 @@ pub mod list; pub mod mapping; pub mod operator; pub mod range; +pub mod schema; pub mod tables; #[cfg(test)] pub mod test; @@ -30,6 +34,7 @@ pub use context_builder::*; pub use error::Error; pub use hasher::Hasher; pub use operator::*; +pub use schema::SchemaSharder; pub use tables::*; pub use value::*; pub use vector::{Centroids, Distance}; @@ -102,9 +107,13 @@ pub(crate) fn shard_value( .ok() .map(Shard::Direct) .unwrap_or(Shard::All), - DataType::Vector => Vector::try_from(value) + DataType::Vector => str_to_vector(value) .ok() - .map(|v| Centroids::from(centroids).shard(&v, shards, centroid_probes)) + .map(|v| { + Centroids::from(centroids) + .shard(&v, shards, centroid_probes) + .into() + }) .unwrap_or(Shard::All), DataType::Varchar => Shard::Direct(varchar(value.as_bytes()) as usize % shards), } @@ -120,15 +129,19 @@ pub(crate) fn shard_binary( match data_type { DataType::Bigint => i64::decode(bytes, Format::Binary) .ok() - .map(|i| Shard::direct(bigint(i) as usize % shards)) + .map(|i| Shard::new_direct(bigint(i) as usize % shards)) .unwrap_or(Shard::All), DataType::Uuid => Uuid::decode(bytes, Format::Binary) .ok() - .map(|u| Shard::direct(uuid(u) as usize % shards)) + .map(|u| Shard::new_direct(uuid(u) as usize % shards)) .unwrap_or(Shard::All), DataType::Vector => Vector::decode(bytes, Format::Binary) .ok() - .map(|v| Centroids::from(centroids).shard(&v, shards, centroid_probes)) + .map(|v| { + Centroids::from(centroids) + .shard(&v, shards, centroid_probes) + .into() + }) .unwrap_or(Shard::All), DataType::Varchar => Shard::Direct(varchar(bytes) as usize % shards), } diff --git a/pgdog/src/frontend/router/sharding/schema.rs b/pgdog/src/frontend/router/sharding/schema.rs new file mode 100644 index 000000000..05ee4e1ce --- /dev/null +++ b/pgdog/src/frontend/router/sharding/schema.rs @@ -0,0 +1,151 @@ +use crate::{ + backend::replication::ShardedSchemas, + frontend::router::parser::{Schema, Shard}, + net::parameter::ParameterValue, +}; + +#[derive(Debug, Default, Clone)] +pub struct SchemaSharder { + catch_all: bool, + current: Option, + schema: Option, +} + +impl SchemaSharder { + /// Resolve current schema. + pub fn resolve(&mut self, schema: Option>, schemas: &ShardedSchemas) { + if schemas.is_empty() { + return; + } + + let check = schemas.get(schema); + if let Some(schema) = check { + let catch_all = schema.is_default(); + let set = + catch_all && self.current.is_none() || self.catch_all || self.current.is_none(); + if set { + self.current = Some(schema.shard().into()); + self.catch_all = catch_all; + self.schema = Some(schema.name().to_owned()); + } + } + } + + /// Resolve current schema from connection parameter. + pub fn resolve_parameter(&mut self, parameter: &ParameterValue, schemas: &ShardedSchemas) { + if schemas.is_empty() { + return; + } + + match parameter { + ParameterValue::String(search_path) => { + let schema = Schema::from(search_path.as_str()); + self.resolve(Some(schema), schemas) + } + + ParameterValue::Tuple(search_paths) => { + for schema in search_paths { + let schema = Schema::from(schema.as_str()); + self.resolve(Some(schema), schemas); + } + } + + _ => (), + } + } + + pub fn get(&self) -> Option<(Shard, &str)> { + if let Some(current) = self.current.as_ref() { + if let Some(schema) = self.schema.as_ref() { + return Some((current.clone(), schema.as_str())); + } + } + + None + } +} + +#[cfg(test)] +mod tests { + use super::*; + use pgdog_config::sharding::ShardedSchema; + + #[test] + fn test_catch_all_chosen_only_when_no_specific_match() { + // Create a catch-all schema (no name) that routes to shard 0 + let catch_all = ShardedSchema { + database: "test".to_string(), + name: None, // This makes it a catch-all/default + shard: 0, + all: false, + }; + + // Create a specific schema "sales" that routes to shard 1 + let sales_schema = ShardedSchema { + database: "test".to_string(), + name: Some("sales".to_string()), + shard: 1, + all: false, + }; + + let schemas = ShardedSchemas::new(vec![catch_all, sales_schema]); + + // Test 1: When we resolve "sales", we should get shard 1 (specific match, not catch-all) + let mut sharder = SchemaSharder::default(); + sharder.resolve(Some(Schema { name: "sales" }), &schemas); + let result = sharder.get(); + assert_eq!( + result.as_ref().map(|(s, _)| s.clone()), + Some(Shard::Direct(1)) + ); + assert_eq!(result.as_ref().map(|(_, name)| *name), Some("sales")); + assert!(!sharder.catch_all); + + // Test 2: When we resolve "unknown", we should get shard 0 (catch-all) + let mut sharder = SchemaSharder::default(); + sharder.resolve(Some(Schema { name: "unknown" }), &schemas); + let result = sharder.get(); + assert_eq!( + result.as_ref().map(|(s, _)| s.clone()), + Some(Shard::Direct(0)) + ); + assert_eq!(result.as_ref().map(|(_, name)| *name), Some("*")); + assert!(sharder.catch_all); + + // Test 3: When we resolve None (no schema specified), we should get shard 0 (catch-all) + let mut sharder = SchemaSharder::default(); + sharder.resolve(None, &schemas); + let result = sharder.get(); + assert_eq!( + result.as_ref().map(|(s, _)| s.clone()), + Some(Shard::Direct(0)) + ); + assert_eq!(result.as_ref().map(|(_, name)| *name), Some("*")); + assert!(sharder.catch_all); + + // Test 4: When catch-all is resolved first, then specific match, + // the specific match should override the catch-all + let mut sharder = SchemaSharder::default(); + sharder.resolve(Some(Schema { name: "unknown" }), &schemas); // catch-all first + sharder.resolve(Some(Schema { name: "sales" }), &schemas); // specific second + let result = sharder.get(); + assert_eq!( + result.as_ref().map(|(s, _)| s.clone()), + Some(Shard::Direct(1)) + ); + assert_eq!(result.as_ref().map(|(_, name)| *name), Some("sales")); + assert!(!sharder.catch_all); + + // Test 5: When specific match is resolved first, catch-all should NOT override + let mut sharder = SchemaSharder::default(); + sharder.resolve(Some(Schema { name: "sales" }), &schemas); // specific first + sharder.resolve(Some(Schema { name: "unknown" }), &schemas); // catch-all second + let result = sharder.get(); + assert_eq!( + result.as_ref().map(|(s, _)| s.clone()), + Some(Shard::Direct(1)) + ); + assert_eq!(result.as_ref().map(|(_, name)| *name), Some("sales")); + assert!(!sharder.catch_all); + } +} diff --git a/pgdog/src/frontend/router/sharding/test/mod.rs b/pgdog/src/frontend/router/sharding/test/mod.rs index 25c71c027..a9194475a 100644 --- a/pgdog/src/frontend/router/sharding/test/mod.rs +++ b/pgdog/src/frontend/router/sharding/test/mod.rs @@ -1,6 +1,6 @@ use std::{collections::HashSet, str::from_utf8}; -use rand::{seq::SliceRandom, thread_rng}; +use rand::seq::SliceRandom; use crate::{ backend::server::test::test_server, @@ -17,7 +17,7 @@ async fn test_shard_varchar() { let mut server = test_server().await; let inserts = (0..100) .map(|i| { - words.shuffle(&mut thread_rng()); + words.shuffle(&mut rand::rng()); let word = words.first().unwrap(); Query::new(format!( @@ -124,7 +124,7 @@ async fn test_binary_encoding() { "", &[Parameter { len: 5, - data: "test1".as_bytes().to_vec(), + data: "test1".as_bytes().into(), }], &[Format::Binary], &[1], diff --git a/pgdog/src/frontend/router/sharding/value.rs b/pgdog/src/frontend/router/sharding/value.rs index a59bc0a19..7a77a628a 100644 --- a/pgdog/src/frontend/router/sharding/value.rs +++ b/pgdog/src/frontend/router/sharding/value.rs @@ -54,6 +54,8 @@ impl<'a> Value<'a> { } } + /// Convert parameter to value, given the data type + /// and known encoding. pub fn from_param( param: &'a ParameterWithFormat<'a>, data_type: DataType, @@ -109,7 +111,10 @@ impl<'a> Value<'a> { if self.data_type == DataType::Bigint { match self.data { Data::Integer(int) => Ok(Some(int)), - Data::Text(text) => Ok(Some(text.parse()?)), + Data::Text(text) => Ok(Some( + text.parse() + .map_err(|_| Error::ParseInt(text.to_string()))?, + )), Data::Binary(data) => match data.len() { 2 => Ok(Some(i16::from_be_bytes(data.try_into()?) as i64)), 4 => Ok(Some(i32::from_be_bytes(data.try_into()?) as i64)), @@ -151,7 +156,12 @@ impl<'a> Value<'a> { pub fn hash(&self, hasher: Hasher) -> Result, Error> { match self.data_type { DataType::Bigint => match self.data { - Data::Text(text) => Ok(Some(hasher.bigint(text.parse()?))), + Data::Text(text) => Ok(Some( + hasher.bigint( + text.parse() + .map_err(|_| Error::ParseInt(text.to_string()))?, + ), + )), Data::Binary(data) => Ok(Some(hasher.bigint(match data.len() { 2 => i16::from_be_bytes(data.try_into()?) as i64, 4 => i32::from_be_bytes(data.try_into()?) as i64, diff --git a/pgdog/src/frontend/router/sharding/vector.rs b/pgdog/src/frontend/router/sharding/vector.rs index d73f3d964..138ac0278 100644 --- a/pgdog/src/frontend/router/sharding/vector.rs +++ b/pgdog/src/frontend/router/sharding/vector.rs @@ -1,65 +1,4 @@ -use crate::{frontend::router::parser::Shard, net::messages::Vector}; - -// Use the SIMD module from the parent sharding module -use super::distance_simd_rust; - -pub enum Distance<'a> { - Euclidean(&'a Vector, &'a Vector), -} - -impl Distance<'_> { - pub fn distance(&self) -> f32 { - match self { - Self::Euclidean(p, q) => { - assert_eq!(p.len(), q.len()); - // Avoids temporary array allocations by working directly with the Float slices - distance_simd_rust::euclidean_distance(p, q) - } - } - } - - // Fallback implementation for testing - pub fn distance_scalar(&self) -> f32 { - match self { - Self::Euclidean(p, q) => { - assert_eq!(p.len(), q.len()); - // Use scalar implementation - distance_simd_rust::euclidean_distance_scalar(p, q) - } - } - } -} - -#[derive(Debug)] -pub struct Centroids<'a> { - centroids: &'a [Vector], -} - -impl Centroids<'_> { - /// Find the shards with the closest centroids, - /// according to the number of probes. - pub fn shard(&self, vector: &Vector, shards: usize, probes: usize) -> Shard { - let mut selected = vec![]; - let mut centroids = self.centroids.iter().enumerate().collect::>(); - centroids.sort_by(|(_, a), (_, b)| { - a.distance_l2(vector) - .partial_cmp(&b.distance_l2(vector)) - .unwrap_or(std::cmp::Ordering::Equal) - }); - let centroids = centroids.into_iter().take(probes); - for (i, _) in centroids { - selected.push(i % shards); - } - - Shard::Multi(selected) - } -} - -impl<'a> From<&'a Vec> for Centroids<'a> { - fn from(centroids: &'a Vec) -> Self { - Centroids { centroids } - } -} +pub use pgdog_vector::{Centroids, Distance}; #[cfg(test)] mod test { diff --git a/pgdog/src/frontend/stats.rs b/pgdog/src/frontend/stats.rs index edb0c63db..5ab262bcc 100644 --- a/pgdog/src/frontend/stats.rs +++ b/pgdog/src/frontend/stats.rs @@ -3,7 +3,7 @@ use std::time::{Duration, SystemTime}; use tokio::time::Instant; -use crate::state::State; +use crate::{backend::pool::stats::MemoryStats, state::State}; /// Client statistics. #[derive(Copy, Clone, Debug)] @@ -37,7 +37,7 @@ pub struct Stats { pub last_request: SystemTime, /// Number of bytes used by the stream buffer, where all the messages /// are stored until they are processed. - pub memory_used: usize, + pub memory_stats: MemoryStats, /// Number of prepared statements in the local cache. pub prepared_statements: usize, /// Client is locked to a particular server. @@ -69,7 +69,7 @@ impl Stats { query_timer: now, wait_timer: now, last_request: SystemTime::now(), - memory_used: 0, + memory_stats: MemoryStats::default(), prepared_statements: 0, locked: false, } @@ -127,8 +127,8 @@ impl Stats { self.bytes_sent += bytes; } - pub(super) fn memory_used(&mut self, memory: usize) { - self.memory_used = memory; + pub(super) fn memory_used(&mut self, memory: MemoryStats) { + self.memory_stats = memory; } pub(super) fn idle(&mut self, in_transaction: bool) { diff --git a/pgdog/src/healthcheck.rs b/pgdog/src/healthcheck.rs index 384c159da..171c9f70f 100644 --- a/pgdog/src/healthcheck.rs +++ b/pgdog/src/healthcheck.rs @@ -45,7 +45,7 @@ async fn healthcheck( .into_iter() .flatten() .collect::>(); - pools.iter().all(|p| p.banned()) + pools.iter().all(|p| !p.healthy()) }); let response = if broken { "down" } else { "up" }; diff --git a/pgdog/src/lib.rs b/pgdog/src/lib.rs index 2a252d300..6ebf419db 100644 --- a/pgdog/src/lib.rs +++ b/pgdog/src/lib.rs @@ -16,13 +16,25 @@ pub mod state; pub mod stats; #[cfg(feature = "tui")] pub mod tui; +pub mod unique_id; pub mod util; use tracing::level_filters::LevelFilter; use tracing_subscriber::{fmt, prelude::*, EnvFilter}; +#[cfg(test)] +use std::alloc::System; use std::io::IsTerminal; +#[cfg(not(test))] +#[cfg(not(target_env = "msvc"))] +#[global_allocator] +static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc; + +#[cfg(test)] +#[global_allocator] +static GLOBAL: &stats_alloc::StatsAlloc = &stats_alloc::INSTRUMENTED_SYSTEM; + /// Setup the logger, so `info!`, `debug!` /// and other macros actually output something. /// diff --git a/pgdog/src/main.rs b/pgdog/src/main.rs index d970dd878..b31327147 100644 --- a/pgdog/src/main.rs +++ b/pgdog/src/main.rs @@ -6,6 +6,7 @@ use pgdog::backend::pool::dns_cache::DnsCache; use pgdog::cli::{self, Commands}; use pgdog::config::{self, config}; use pgdog::frontend::listener::Listener; +use pgdog::frontend::prepared_statements; use pgdog::plugin; use pgdog::stats; use pgdog::util::pgdog_version; @@ -15,13 +16,6 @@ use tracing::{error, info}; use std::process::exit; -#[cfg(not(target_env = "msvc"))] -use tikv_jemallocator::Jemalloc; - -#[cfg(not(target_env = "msvc"))] -#[global_allocator] -static GLOBAL: Jemalloc = Jemalloc; - fn main() -> Result<(), Box> { let args = cli::Cli::parse(); @@ -83,18 +77,28 @@ fn main() -> Result<(), Box> { let runtime = match config.config.general.workers { 0 => { let mut binding = Builder::new_current_thread(); - binding.enable_all(); + binding + .enable_all() + .thread_stack_size(config.config.memory.stack_size); binding } workers => { - info!("spawning {} workers", workers); let mut builder = Builder::new_multi_thread(); - builder.worker_threads(workers).enable_all(); + builder + .worker_threads(workers) + .enable_all() + .thread_stack_size(config.config.memory.stack_size); builder } } .build()?; + info!( + "spawning {} threads (stack size: {}MiB)", + config.config.general.workers, + config.config.memory.stack_size / 1024 / 1024 + ); + runtime.block_on(async move { pgdog(args.command).await })?; Ok(()) @@ -106,7 +110,7 @@ async fn pgdog(command: Option) -> Result<(), Box) -> Result<(), Box) -> Result<(), Box { if let Commands::DataSync { .. } = command { info!("πŸ”„ entering data sync mode"); - cli::data_sync(command.clone()).await?; + if let Err(err) = cli::data_sync(command.clone()).await { + error!("{}", err); + } } if let Commands::SchemaSync { .. } = command { info!("πŸ”„ entering schema sync mode"); - cli::schema_sync(command.clone()).await?; + if let Err(err) = cli::schema_sync(command.clone()).await { + error!("{}", err); + } } if let Commands::Setup { database } = command { info!("πŸ”„ entering setup mode"); cli::setup(database).await?; } + + if let Commands::Route { .. } = command { + if let Err(err) = cli::route(command.clone()).await { + error!("{}", err); + } + } } } diff --git a/pgdog/src/net/discovery/listener.rs b/pgdog/src/net/discovery/listener.rs index faa872fca..2f05b0917 100644 --- a/pgdog/src/net/discovery/listener.rs +++ b/pgdog/src/net/discovery/listener.rs @@ -42,7 +42,7 @@ impl Listener { /// Create new listener. fn new() -> Self { Self { - id: rand::thread_rng().gen(), + id: rand::rng().random(), inner: Arc::new(Mutex::new(Inner { peers: HashMap::new(), })), diff --git a/pgdog/src/net/error.rs b/pgdog/src/net/error.rs index 1bf121cba..d28bf04ee 100644 --- a/pgdog/src/net/error.rs +++ b/pgdog/src/net/error.rs @@ -96,4 +96,7 @@ pub enum Error { #[error("not a boolean")] NotBoolean, + + #[error("not a pg_lsn")] + NotPgLsn, } diff --git a/pgdog/src/net/messages/backend_key.rs b/pgdog/src/net/messages/backend_key.rs index 4514b02a8..9f07d903c 100644 --- a/pgdog/src/net/messages/backend_key.rs +++ b/pgdog/src/net/messages/backend_key.rs @@ -1,9 +1,21 @@ //! BackendKeyData (B) message. +use std::fmt::Display; +use std::sync::atomic::AtomicI32; +use std::sync::atomic::Ordering; + use crate::net::messages::code; use crate::net::messages::prelude::*; +use once_cell::sync::Lazy; use rand::Rng; +static COUNTER: Lazy = Lazy::new(|| AtomicI32::new(0)); + +// This wraps around. +fn next_counter() -> i32 { + COUNTER.fetch_add(1, Ordering::SeqCst) +} + /// BackendKeyData (B) #[derive(Copy, Clone, Debug, Hash, PartialEq, Eq)] pub struct BackendKeyData { @@ -13,6 +25,12 @@ pub struct BackendKeyData { pub secret: i32, } +impl Display for BackendKeyData { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "pid={}, secret={}", self.pid, self.secret) + } +} + impl Default for BackendKeyData { fn default() -> Self { Self::new() @@ -23,8 +41,18 @@ impl BackendKeyData { /// Create new random BackendKeyData (B) message. pub fn new() -> Self { Self { - pid: rand::thread_rng().gen(), - secret: rand::thread_rng().gen(), + pid: rand::rng().random(), + secret: rand::rng().random(), + } + } + + /// Create new BackendKeyData for a connected client. + /// + /// This counts client IDs incrementally. + pub fn new_client() -> Self { + Self { + pid: next_counter(), + secret: rand::rng().random(), } } } diff --git a/pgdog/src/net/messages/bind.rs b/pgdog/src/net/messages/bind.rs index 101e2c087..a8cb2652a 100644 --- a/pgdog/src/net/messages/bind.rs +++ b/pgdog/src/net/messages/bind.rs @@ -7,6 +7,7 @@ use super::prelude::*; use super::Error; use super::FromDataType; use super::Vector; +use bytes::BytesMut; use std::fmt::Debug; use std::str::from_utf8; @@ -34,7 +35,7 @@ pub struct Parameter { /// Parameter data length. pub len: i32, /// Parameter data. - pub data: Vec, + pub data: Bytes, } impl Debug for Parameter { @@ -59,14 +60,14 @@ impl Parameter { pub fn new_null() -> Self { Self { len: -1, - data: vec![], + data: Bytes::new(), } } pub fn new(data: &[u8]) -> Self { Self { len: data.len() as i32, - data: data.to_vec(), + data: Bytes::copy_from_slice(data), } } } @@ -111,6 +112,14 @@ impl<'a> ParameterWithFormat<'a> { pub fn data(&'a self) -> &'a [u8] { &self.parameter.data } + + pub fn is_null(&self) -> bool { + self.parameter.len < 0 + } + + pub fn parameter(&self) -> &Parameter { + self.parameter + } } /// Bind (F) message. @@ -124,8 +133,8 @@ pub struct Bind { codes: Vec, /// Parameters. params: Vec, - /// Results format. - results: Vec, + /// Results format (raw bytes, 2 bytes per i16). + results: Bytes, /// Original payload. original: Option, } @@ -137,7 +146,7 @@ impl Default for Bind { statement: Bytes::from("\0"), codes: vec![], params: vec![], - results: vec![], + results: Bytes::new(), original: None, } } @@ -149,7 +158,7 @@ impl Bind { + self.statement.len() + self.codes.len() * std::mem::size_of::() + 2 // num codes + self.params.iter().map(|p| p.len()).sum::() + 2 // num params - + self.results.len() * std::mem::size_of::() + 2 // num results + + self.results.len() + 2 // num results (results already stores raw bytes) + 4 // len + 1 // code } @@ -237,7 +246,11 @@ impl Bind { results: &[i16], ) -> Self { let mut me = Self::new_params_codes(name, params, codes); - me.results = results.to_vec(); + let mut buf = BytesMut::with_capacity(results.len() * 2); + for result in results { + buf.put_i16(*result); + } + me.results = buf.freeze(); me } @@ -249,6 +262,47 @@ impl Bind { pub fn format_codes_raw(&self) -> &Vec { &self.codes } + + /// Push a parameter to the end of the parameter list with the given format. + /// + /// Handles format codes correctly per PostgreSQL semantics: + /// - If codes.len() == 0: all parameters use Text + /// - If codes.len() == 1: all parameters use that one format (uniform) + /// - If codes.len() == params.len() (and > 1): one-to-one mapping + pub fn push_param(&mut self, param: Parameter, format: Format) { + if self.codes.len() == 1 { + // Uniform format: if new format differs, expand to one-to-one + if self.codes[0] != format { + let existing_format = self.codes[0]; + self.codes = vec![existing_format; self.params.len()]; + self.codes.push(format); + } + // If format matches, keep uniform (no change to codes) + } else if self.codes.len() > 1 && self.codes.len() == self.params.len() { + // One-to-one mapping: add the new format + self.codes.push(format); + } else if self.codes.is_empty() && format == Format::Binary { + // No codes (all text): if adding binary, need to expand + self.codes = vec![Format::Text; self.params.len()]; + self.codes.push(Format::Binary); + } + // If codes.len() == 0 and format is Text, no codes needed + + self.params.push(param); + self.original = None; + } + + /// Get the effective format for new parameters. + pub fn default_param_format(&self) -> Format { + if self.codes.len() == 1 { + self.codes[0] + } else if self.codes.is_empty() { + Format::Text + } else { + // One-to-one mapping: default to Text for new params + Format::Text + } + } } impl FromBytes for Bind { @@ -262,29 +316,31 @@ impl FromBytes for Bind { from_utf8(&portal[0..portal.len() - 1])?; from_utf8(&statement[0..statement.len() - 1])?; - let num_codes = bytes.get_i16(); - let codes = (0..num_codes) + let num_codes = bytes.get_u16(); + let codes = (0..num_codes as usize) .map(|_| match bytes.get_i16() { 0 => Format::Text, _ => Format::Binary, }) .collect(); - let num_params = bytes.get_i16(); - let params = (0..num_params) + let num_params = bytes.get_u16(); + let params = (0..num_params as usize) .map(|_| { let len = bytes.get_i32(); let data = if len >= 0 { - let mut data = Vec::with_capacity(len as usize); - (0..len).for_each(|_| data.push(bytes.get_u8())); - data + bytes.split_to(len as usize) } else { - vec![] + Bytes::new() }; Parameter { len, data } }) .collect(); let num_results = bytes.get_i16(); - let results = (0..num_results).map(|_| bytes.get_i16()).collect(); + let results = if num_results > 0 { + bytes.split_to((num_results * 2) as usize) + } else { + Bytes::new() + }; Ok(Self { portal, @@ -309,22 +365,20 @@ impl ToBytes for Bind { payload.put(self.portal.clone()); payload.put(self.statement.clone()); - payload.put_i16(self.codes.len() as i16); + payload.put_u16(self.codes.len() as u16); for code in &self.codes { payload.put_i16(match code { Format::Text => 0, Format::Binary => 1, }); } - payload.put_i16(self.params.len() as i16); + payload.put_u16(self.params.len() as u16); for param in &self.params { payload.put_i32(param.len); payload.put(¶m.data[..]); } - payload.put_i16(self.results.len() as i16); - for result in &self.results { - payload.put_i16(*result); - } + payload.put_i16((self.results.len() / 2) as i16); + payload.put(self.results.clone()); Ok(payload.freeze()) } } @@ -358,14 +412,18 @@ mod test { params: vec![ Parameter { len: 2, - data: vec![0, 1], + data: Bytes::copy_from_slice(&[0, 1]), }, Parameter { len: 4, - data: "test".as_bytes().to_vec(), + data: Bytes::from("test"), }, ], - results: vec![0], + results: { + let mut buf = BytesMut::with_capacity(2); + buf.put_i16(0); + buf.freeze() + }, }; let bytes = bind.to_bytes().unwrap(); let mut original = Bind::from_bytes(bytes.clone()).unwrap(); @@ -400,7 +458,7 @@ mod test { statement: "test\0".into(), codes: vec![Format::Binary], params: vec![Parameter { - data: jsonb.as_bytes().to_vec(), + data: Bytes::copy_from_slice(jsonb.as_bytes()), len: jsonb.len() as i32, }], ..Default::default() @@ -434,4 +492,19 @@ mod test { assert_eq!(msg.code(), c); } } + + #[test] + fn test_large_parameter_count_round_trip() { + let count = 35_000; + let params: Vec = (0..count).map(|_| Parameter::new_null()).collect(); + let bind = Bind::new_params("__pgdog_large", ¶ms); + + let bytes = bind.to_bytes().unwrap(); + let decoded = Bind::from_bytes(bytes.clone()).unwrap(); + + assert_eq!(decoded.params_raw().len(), count); + assert_eq!(decoded.codes().len(), 0); + assert_eq!(decoded.statement(), "__pgdog_large"); + assert_eq!(bytes.len(), decoded.len()); + } } diff --git a/pgdog/src/net/messages/buffer.rs b/pgdog/src/net/messages/buffer.rs new file mode 100644 index 000000000..35a414874 --- /dev/null +++ b/pgdog/src/net/messages/buffer.rs @@ -0,0 +1,372 @@ +//! Cancel-safe and memory-efficient +//! read buffer for Postgres messages. + +use std::{io::Cursor, ops::Add}; + +use bytes::{Buf, BytesMut}; +use tokio::io::AsyncReadExt; + +use crate::net::stream::eof; + +use super::{Error, Message}; + +const HEADER_SIZE: usize = 5; + +#[derive(Default, Debug, Clone)] +pub struct MessageBuffer { + buffer: BytesMut, + capacity: usize, + stats: MessageBufferStats, +} + +impl MessageBuffer { + /// Create new cancel-safe + /// message buffer. + pub fn new(capacity: usize) -> Self { + Self { + buffer: BytesMut::with_capacity(capacity), + capacity, + stats: MessageBufferStats { + bytes_alloc: capacity, + ..Default::default() + }, + } + } + + /// Buffer capacity. + pub fn capacity(&self) -> usize { + self.buffer.capacity() + } + + async fn read_internal( + &mut self, + stream: &mut (impl Unpin + AsyncReadExt), + ) -> Result { + loop { + if let Some(size) = self.message_size() { + if self.have_message() { + return Ok(Message::new(self.buffer.split_to(size).freeze())); + } + + self.ensure_capacity(size); // Reserve at least enough space for the whole message. + } + + // Ensure there is 1/4 of the buffer + // available at all times. This clears memory usage + // frequently. + self.ensure_capacity(self.capacity / 4); + + let read = eof(stream.read_buf(&mut self.buffer).await)?; + self.stats.bytes_used += read; + + if read == 0 { + return Err(Error::UnexpectedEof); + } + } + } + + // This may or may not allocate memory, depending on how big of + // a message we are receiving. + fn ensure_capacity(&mut self, amount: usize) { + if self.buffer.try_reclaim(amount) { + // I know this isn't exactly right, we could be reclaiming more. + // But undercounting is better than overcounting. + self.stats.bytes_used = self.stats.bytes_used.saturating_sub(amount); + self.stats.reclaims += 1; + } else { + self.buffer.reserve(amount); + // Possibly undercounting. + self.stats.bytes_alloc += amount; + } + } + + fn have_message(&self) -> bool { + self.message_size() + .map(|len| self.buffer.len() >= len) + .unwrap_or(false) + } + + fn message_size(&self) -> Option { + if self.buffer.len() >= HEADER_SIZE { + let mut cur = Cursor::new(&self.buffer); + let _code = cur.get_u8(); + let len = cur.get_i32() as usize + 1; + Some(len) + } else { + None + } + } + + /// Re-allcoate buffer if it exceeds capacity. + pub fn shrink_to_fit(&mut self) -> bool { + // Re-allocate the buffer to save on memory. + if self.stats.bytes_alloc > self.capacity * 2 { + // Create new buffer and copy contents. + let mut buffer = BytesMut::with_capacity(self.capacity); + buffer.extend_from_slice(&self.buffer); + + // Update stats. + self.stats.bytes_used = self.buffer.len(); + self.buffer = buffer; + self.stats.reallocs += 1; + self.stats.bytes_alloc = self.capacity; // Possibly undercounting. + true + } else { + false + } + } + + /// Get buffer stats. + pub fn stats(&self) -> &MessageBufferStats { + &self.stats + } + + /// Read a Postgres message off of a stream. + /// + /// # Cancellation safety + /// + /// This method is cancel-safe. + /// + pub async fn read( + &mut self, + stream: &mut (impl Unpin + AsyncReadExt), + ) -> Result { + self.read_internal(stream).await + } +} + +#[derive(Debug, Copy, Clone, Default)] +pub struct MessageBufferStats { + pub reallocs: usize, + pub reclaims: usize, + pub bytes_used: usize, + pub bytes_alloc: usize, +} + +impl Add for MessageBufferStats { + type Output = Self; + + fn add(self, rhs: Self) -> Self::Output { + Self { + reallocs: rhs.reallocs + self.reallocs, + reclaims: rhs.reclaims + self.reclaims, + bytes_used: rhs.bytes_used + self.bytes_used, + bytes_alloc: rhs.bytes_alloc + self.bytes_alloc, + } + } +} + +#[cfg(test)] +mod test { + use super::*; + use crate::net::{FromBytes, Parse, Protocol, Sync, ToBytes}; + use bytes::BufMut; + use std::time::Duration; + use tokio::{ + io::AsyncWriteExt, + net::{TcpListener, TcpStream}, + spawn, + sync::mpsc, + time::interval, + }; + + #[tokio::test(flavor = "multi_thread")] + async fn test_read() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let (tx, mut rx) = mpsc::channel(1); + + spawn(async move { + let mut conn = TcpStream::connect(addr).await.unwrap(); + use rand::{rngs::StdRng, Rng, SeedableRng}; + let mut rng = StdRng::from_os_rng(); + + for i in 0..5000 { + let msg = Sync.to_bytes().unwrap(); + conn.write_all(&msg).await.unwrap(); + + let query_len = rng.random_range(10..=1000); + let query: String = (0..query_len) + .map(|_| rng.sample(rand::distr::Alphanumeric) as char) + .collect(); + + let msg = Parse::named(format!("test_{}", i), &query) + .to_bytes() + .unwrap(); + conn.write_all(&msg).await.unwrap(); + conn.flush().await.unwrap(); + } + rx.recv().await; + }); + + let (mut conn, _) = listener.accept().await.unwrap(); + let mut buf = MessageBuffer::default(); + + let mut counter = 0; + let mut interrupted = 0; + let mut interval = interval(Duration::from_millis(1)); + + while counter < 10000 { + let msg = tokio::select! { + msg = buf.read(&mut conn) => { + msg.unwrap() + } + + _ = interval.tick() => { + interrupted += 1; + continue; + } + }; + + if counter % 2 == 0 { + assert_eq!(msg.code(), 'S'); + } else { + assert_eq!(msg.code(), 'P'); + let parse = Parse::from_bytes(msg.to_bytes().unwrap()).unwrap(); + assert_eq!(parse.name(), format!("test_{}", counter / 2)); + } + + counter += 1; + } + + tx.send(0).await.unwrap(); + + assert!(interrupted > 0, "no cancellations"); + assert_eq!(counter, 10000, "didnt receive all messages"); + assert!(matches!( + buf.read(&mut conn).await.err(), + Some(Error::UnexpectedEof) + )); + assert!(buf.capacity() > 0); + } + + #[test] + fn test_bytes_mut() { + let region = stats_alloc::Region::new(crate::GLOBAL); + + let mut original = BytesMut::with_capacity(5 * 1000); + assert_eq!(original.capacity(), 5 * 1000); + assert_eq!(original.len(), 0); + + for _ in 0..(5 * 25 * 1000) { + original.put_u8('S' as u8); + original.put_i32(4); + + let sync = original.split_to(5); + assert_eq!(sync.capacity(), 5); + assert_eq!(sync.len(), 5); + + // Removes it from the buffer, giving that space back. + drop(sync); + } + + assert_eq!(region.change().allocations, 2); + assert!(region.change().bytes_allocated < 6000); // Depends on the allocator, but it will never be more. + } + + #[tokio::test] + async fn test_shrink_to_fit() { + use std::io::Cursor; + + let mut stream_data = Vec::new(); + + // Create a large message (10KB query) + let large_query = "SELECT * FROM ".to_string() + &"x".repeat(10_000); + let large_msg = Parse::named("large", &large_query).to_bytes().unwrap(); + stream_data.extend_from_slice(&large_msg); + + // Create a small message + let small_msg = Sync.to_bytes().unwrap(); + stream_data.extend_from_slice(&small_msg); + + let mut cursor = Cursor::new(stream_data); + let mut buf = MessageBuffer::new(4096); + + // Read the large message + let msg = buf.read(&mut cursor).await.unwrap(); + assert_eq!(msg.code(), 'P'); + + // At this point, bytes_used should be > BUFFER_SIZE + let bytes_used_before = buf.stats.bytes_used; + assert!(bytes_used_before > 4096); + + // Shrink the buffer + assert!(buf.shrink_to_fit()); + + // After shrinking, we should have reset to BUFFER_SIZE capacity + assert_eq!(buf.buffer.capacity(), 4096); + + // Should still be able to read the next message + let msg = buf.read(&mut cursor).await.unwrap(); + assert_eq!(msg.code(), 'S'); + } + + #[tokio::test] + async fn test_shrink_to_fit_preserves_partial_data() { + use bytes::BufMut; + + let mut buf = MessageBuffer::new(4096); + + // Simulate having allocated memory for a large message + buf.stats.bytes_alloc = 4096 * 3; + buf.stats.bytes_used = 4096 * 2; + + // Put some partial message data in the buffer (incomplete header) + buf.buffer.put_u8(b'P'); + buf.buffer.put_u8(0); + buf.buffer.put_u8(0); + + let data_before = buf.buffer.clone(); + + // Shrink should preserve the partial data + assert!(buf.shrink_to_fit()); + + assert_eq!(buf.stats().bytes_alloc, 4096); + assert_eq!(buf.buffer.len(), data_before.len()); + assert_eq!(buf.buffer[..], data_before[..]); + assert_eq!(buf.buffer.capacity(), 4096); + } + + #[tokio::test] + async fn test_shrink_to_fit_no_realloc_when_under_capacity() { + use std::io::Cursor; + + let mut stream_data = Vec::new(); + + // Create several small messages that won't exceed BUFFER_SIZE + for i in 0..10 { + let query = format!("SELECT {}", i); + let msg = Parse::named(format!("stmt_{}", i), &query) + .to_bytes() + .unwrap(); + stream_data.extend_from_slice(&msg); + } + + let mut cursor = Cursor::new(stream_data); + let mut buf = MessageBuffer::new(4096); + + // Read all small messages + for _ in 0..10 { + let msg = buf.read(&mut cursor).await.unwrap(); + assert_eq!(msg.code(), 'P'); + } + + // At this point, bytes_used should be below BUFFER_SIZE + let bytes_used = buf.stats.bytes_used; + assert!(bytes_used <= 4096); + + let capacity_before = buf.buffer.capacity(); + let reallocs_before = buf.stats.reallocs; + let bytes_alloc_before = buf.stats.bytes_alloc; + let frees_before = buf.stats.reclaims; + + // Should not reallocate since we haven't exceeded BUFFER_SIZE + assert!(!buf.shrink_to_fit()); + + // Verify no reallocation occurred and stats remain unchanged + assert_eq!(buf.buffer.capacity(), capacity_before); + assert_eq!(buf.stats.reallocs, reallocs_before); + assert_eq!(buf.stats.bytes_alloc, bytes_alloc_before); + assert_eq!(buf.stats.reclaims, frees_before); + } +} diff --git a/pgdog/src/net/messages/close.rs b/pgdog/src/net/messages/close.rs index 24a7fa2f2..70fa067a6 100644 --- a/pgdog/src/net/messages/close.rs +++ b/pgdog/src/net/messages/close.rs @@ -6,7 +6,7 @@ use std::str::from_utf8_unchecked; use super::code; use super::prelude::*; -#[derive(Clone)] +#[derive(Clone, PartialEq)] pub struct Close { payload: Bytes, } diff --git a/pgdog/src/net/messages/copy_done.rs b/pgdog/src/net/messages/copy_done.rs index b2fcf9ff5..1aa3b188b 100644 --- a/pgdog/src/net/messages/copy_done.rs +++ b/pgdog/src/net/messages/copy_done.rs @@ -1,6 +1,6 @@ use super::{code, prelude::*}; -#[derive(Debug, Clone)] +#[derive(Debug, Clone, PartialEq)] pub struct CopyDone; impl FromBytes for CopyDone { diff --git a/pgdog/src/net/messages/copy_fail.rs b/pgdog/src/net/messages/copy_fail.rs index 2bc3f9af7..9b74cde9e 100644 --- a/pgdog/src/net/messages/copy_fail.rs +++ b/pgdog/src/net/messages/copy_fail.rs @@ -1,10 +1,19 @@ use super::{code, prelude::*}; -#[derive(Debug, Clone)] +#[derive(Debug, Clone, PartialEq)] pub struct CopyFail { error: Bytes, } +impl CopyFail { + pub fn new(error: impl ToString) -> Self { + let error = error.to_string(); + Self { + error: Bytes::from(format!("{}\0", error)), + } + } +} + impl FromBytes for CopyFail { fn from_bytes(mut bytes: Bytes) -> Result { code!(bytes, 'f'); diff --git a/pgdog/src/net/messages/data_row.rs b/pgdog/src/net/messages/data_row.rs index b25f663bf..16b987669 100644 --- a/pgdog/src/net/messages/data_row.rs +++ b/pgdog/src/net/messages/data_row.rs @@ -52,6 +52,10 @@ impl Data { is_null: true, } } + + pub(crate) fn is_null(&self) -> bool { + self.is_null + } } /// DataRow message. @@ -102,6 +106,15 @@ impl ToDataRowColumn for i64 { } } +impl ToDataRowColumn for Option { + fn to_data_row_column(&self) -> Data { + match self { + Some(value) => ToDataRowColumn::to_data_row_column(value), + None => Data::null(), + } + } +} + impl ToDataRowColumn for usize { fn to_data_row_column(&self) -> Data { Bytes::copy_from_slice(self.to_string().as_bytes()).into() @@ -245,6 +258,11 @@ impl DataRow { .and_then(|col| T::decode(&col, format).ok()) } + /// Get raw column data. + pub(crate) fn get_raw(&self, index: usize) -> Option<&Data> { + self.columns.get(index) + } + /// Get column at index given row description. pub fn get_column<'a>( &self, diff --git a/pgdog/src/net/messages/data_types/float.rs b/pgdog/src/net/messages/data_types/float.rs index 4e4ef10a9..b613da1cc 100644 --- a/pgdog/src/net/messages/data_types/float.rs +++ b/pgdog/src/net/messages/data_types/float.rs @@ -1,47 +1,10 @@ use bytes::{Buf, BufMut, BytesMut}; -use serde::{Deserialize, Serialize}; -use std::cmp::Ordering; -use std::fmt::Display; -use std::hash::{Hash, Hasher}; use crate::net::messages::data_row::Data; use super::*; -/// Wrapper type for f32 that implements Ord for PostgreSQL compatibility -#[derive(Debug, Clone, Copy, Serialize, Deserialize)] -pub struct Float(pub f32); - -impl PartialOrd for Float { - fn partial_cmp(&self, other: &Self) -> Option { - Some(self.cmp(other)) - } -} - -impl Ord for Float { - fn cmp(&self, other: &Self) -> Ordering { - // PostgreSQL ordering: NaN is greater than all other values - match (self.0.is_nan(), other.0.is_nan()) { - (true, true) => Ordering::Equal, - (true, false) => Ordering::Greater, - (false, true) => Ordering::Less, - (false, false) => self.0.partial_cmp(&other.0).unwrap_or(Ordering::Equal), - } - } -} - -impl PartialEq for Float { - fn eq(&self, other: &Self) -> bool { - // PostgreSQL treats NaN as equal to NaN for indexing purposes - if self.0.is_nan() && other.0.is_nan() { - true - } else { - self.0 == other.0 - } - } -} - -impl Eq for Float {} +pub use pgdog_vector::Float; impl FromDataType for Float { fn decode(bytes: &[u8], encoding: Format) -> Result { @@ -100,49 +63,10 @@ impl ToDataRowColumn for Float { } } -impl From for Float { - fn from(value: f32) -> Self { - Float(value) - } -} - -impl From for f32 { - fn from(value: Float) -> Self { - value.0 - } -} - -impl Display for Float { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - if self.0.is_nan() { - write!(f, "NaN") - } else if self.0.is_infinite() { - if self.0.is_sign_positive() { - write!(f, "Infinity") - } else { - write!(f, "-Infinity") - } - } else { - write!(f, "{}", self.0) - } - } -} - -impl Hash for Float { - fn hash(&self, state: &mut H) { - if self.0.is_nan() { - // All NaN values hash to the same value - 0u8.hash(state); - } else { - // Use bit representation for consistent hashing - self.0.to_bits().hash(state); - } - } -} - #[cfg(test)] mod tests { use super::*; + use std::hash::{Hash, Hasher}; #[test] fn test_float_nan_handling() { diff --git a/pgdog/src/net/messages/data_types/timestamp.rs b/pgdog/src/net/messages/data_types/timestamp.rs index dabe30e0e..d9e4a4089 100644 --- a/pgdog/src/net/messages/data_types/timestamp.rs +++ b/pgdog/src/net/messages/data_types/timestamp.rs @@ -84,6 +84,21 @@ macro_rules! assign { } impl Timestamp { + /// Convert Postgres timestamp to timestamp without timezone. + pub fn to_naive_datetime(&self) -> NaiveDateTime { + NaiveDateTime::new( + NaiveDate::from_ymd_opt(self.year as i32, self.month as u32, self.day as u32) + .unwrap_or_default(), + NaiveTime::from_hms_micro_opt( + self.hour as u32, + self.minute as u32, + self.second as u32, + self.micros as u32, + ) + .unwrap_or_default(), + ) + } + /// Create a timestamp representing positive infinity pub fn infinity() -> Self { Self { diff --git a/pgdog/src/net/messages/data_types/timestamptz.rs b/pgdog/src/net/messages/data_types/timestamptz.rs index aa993d3bf..dac79b484 100644 --- a/pgdog/src/net/messages/data_types/timestamptz.rs +++ b/pgdog/src/net/messages/data_types/timestamptz.rs @@ -44,12 +44,15 @@ impl ToDataRowColumn for TimestampTz { #[cfg(test)] mod test { + use chrono::{Datelike, Timelike}; + use super::*; #[test] fn test_timestamptz() { let ts = "2025-03-05 14:55:02.436109-08".as_bytes(); let ts = TimestampTz::decode(ts, Format::Text).unwrap(); + let time = ts.to_naive_datetime(); assert_eq!(ts.year, 2025); assert_eq!(ts.month, 3); @@ -59,5 +62,12 @@ mod test { assert_eq!(ts.second, 2); assert_eq!(ts.micros, 436109); assert_eq!(ts.offset, Some(-8)); + + assert_eq!(time.date().year(), 2025); + assert_eq!(time.date().month(), 3); + assert_eq!(time.date().day(), 5); + assert_eq!(time.time().hour(), 14); + assert_eq!(time.time().minute(), 55); + assert_eq!(time.time().second(), 2); } } diff --git a/pgdog/src/net/messages/data_types/vector.rs b/pgdog/src/net/messages/data_types/vector.rs index 3689d6a55..6f38c75d3 100644 --- a/pgdog/src/net/messages/data_types/vector.rs +++ b/pgdog/src/net/messages/data_types/vector.rs @@ -1,46 +1,15 @@ -use crate::{ - frontend::router::sharding::vector::Distance, - net::{ - messages::{Format, ToDataRowColumn}, - Error, - }, +use std::str::from_utf8; + +use crate::net::{ + messages::{Format, ToDataRowColumn}, + Error, }; use bytes::{Buf, BufMut, Bytes, BytesMut}; -use serde::{ - de::{self, Visitor}, - ser::SerializeSeq, - Deserialize, Serialize, -}; -use std::{fmt::Debug, ops::Deref, str::from_utf8}; -use super::{Datum, Float, FromDataType}; - -#[derive(Clone, PartialEq, PartialOrd, Ord, Eq, Hash)] -#[repr(C)] -pub struct Vector { - values: Vec, -} +use super::{Datum, FromDataType}; +use pgdog_vector::Float; -impl Debug for Vector { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - if self.values.len() > 3 { - f.debug_struct("Vector") - .field( - "values", - &format!( - "[{}..{}]", - self.values[0], - self.values[self.values.len() - 1] - ), - ) - .finish() - } else { - f.debug_struct("Vector") - .field("values", &self.values) - .finish() - } - } -} +pub use pgdog_vector::Vector; impl FromDataType for Vector { fn decode(mut bytes: &[u8], encoding: Format) -> Result { @@ -92,75 +61,8 @@ impl ToDataRowColumn for Vector { } } -impl Vector { - /// Length of the vector. - pub fn len(&self) -> usize { - self.values.len() - } - - /// Is the vector empty? - pub fn is_empty(&self) -> bool { - self.len() == 0 - } - - /// Compute L2 distance between the vectors. - pub fn distance_l2(&self, other: &Self) -> f32 { - Distance::Euclidean(self, other).distance() - } -} - -impl Deref for Vector { - type Target = Vec; - - fn deref(&self) -> &Self::Target { - &self.values - } -} - -impl From<&[f64]> for Vector { - fn from(value: &[f64]) -> Self { - Self { - values: value.iter().map(|v| Float(*v as f32)).collect(), - } - } -} - -impl From<&[f32]> for Vector { - fn from(value: &[f32]) -> Self { - Self { - values: value.iter().map(|v| Float(*v)).collect(), - } - } -} - -impl From> for Vector { - fn from(value: Vec) -> Self { - Self { - values: value.into_iter().map(Float::from).collect(), - } - } -} - -impl From> for Vector { - fn from(value: Vec) -> Self { - Self { - values: value.into_iter().map(|v| Float(v as f32)).collect(), - } - } -} - -impl From> for Vector { - fn from(value: Vec) -> Self { - Self { values: value } - } -} - -impl TryFrom<&str> for Vector { - type Error = Error; - - fn try_from(value: &str) -> Result { - Self::decode(value.as_bytes(), Format::Text) - } +pub fn str_to_vector(value: &str) -> Result { + FromDataType::decode(value.as_bytes(), Format::Text) } impl From for Datum { @@ -181,50 +83,6 @@ impl TryFrom for Vector { } } -struct VectorVisitor; - -impl<'de> Visitor<'de> for VectorVisitor { - type Value = Vector; - - fn visit_seq(self, mut seq: A) -> Result - where - A: de::SeqAccess<'de>, - { - let mut results = vec![]; - while let Some(n) = seq.next_element::()? { - results.push(n); - } - - Ok(Vector::from(results.as_slice())) - } - - fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { - formatter.write_str("expected a list of floating points") - } -} - -impl<'de> Deserialize<'de> for Vector { - fn deserialize(deserializer: D) -> Result - where - D: de::Deserializer<'de>, - { - deserializer.deserialize_seq(VectorVisitor) - } -} - -impl Serialize for Vector { - fn serialize(&self, serializer: S) -> Result - where - S: serde::Serializer, - { - let mut seq = serializer.serialize_seq(Some(self.len()))?; - for v in &self.values { - seq.serialize_element(v)?; - } - seq.end() - } -} - #[cfg(test)] mod test { use super::*; @@ -282,24 +140,4 @@ mod test { assert!(encoded_str.contains("Infinity")); assert!(encoded_str.contains("-Infinity")); } - - #[test] - fn test_vector_distance_with_special_values() { - // Test distance calculation with normal values - let v1 = Vector::from(&[3.0, 4.0][..]); - let v2 = Vector::from(&[0.0, 0.0][..]); - let distance = v1.distance_l2(&v2); - assert_eq!(distance, 5.0); // 3-4-5 triangle - - // Test distance with NaN - let v_nan = Vector::from(vec![Float(f32::NAN), Float(1.0)]); - let v_normal = Vector::from(&[1.0, 1.0][..]); - let distance_nan = v_nan.distance_l2(&v_normal); - assert!(distance_nan.is_nan()); - - // Test distance with Infinity - let v_inf = Vector::from(vec![Float(f32::INFINITY), Float(1.0)]); - let distance_inf = v_inf.distance_l2(&v_normal); - assert!(distance_inf.is_infinite()); - } } diff --git a/pgdog/src/net/messages/describe.rs b/pgdog/src/net/messages/describe.rs index d5a9c2a28..0f956316b 100644 --- a/pgdog/src/net/messages/describe.rs +++ b/pgdog/src/net/messages/describe.rs @@ -7,7 +7,7 @@ use super::code; use super::prelude::*; /// Describe (F) message. -#[derive(Clone)] +#[derive(Clone, PartialEq)] pub struct Describe { payload: Bytes, original: Option, diff --git a/pgdog/src/net/messages/error_response.rs b/pgdog/src/net/messages/error_response.rs index f9107aabf..f782b3e75 100644 --- a/pgdog/src/net/messages/error_response.rs +++ b/pgdog/src/net/messages/error_response.rs @@ -3,14 +3,16 @@ use std::fmt::Display; use std::time::Duration; -use crate::net::c_string_buf; - use super::prelude::*; +use crate::{ + net::{c_string_buf, code}, + state::State, +}; /// ErrorResponse (B) message. -#[derive(Debug)] +#[derive(Debug, Clone)] pub struct ErrorResponse { - severity: String, + pub severity: String, pub code: String, pub message: String, pub detail: Option, @@ -50,25 +52,65 @@ impl ErrorResponse { } } - pub fn cross_shard_disabled() -> ErrorResponse { + pub fn client_login_timeout(timeout: Duration) -> ErrorResponse { + let mut error = Self::client_idle_timeout(timeout, &State::Active); + error.message = "client login timeout".into(); + error.detail = Some(format!( + "client_login_timeout of {}ms expired", + timeout.as_millis() + )); + + error + } + + pub fn cross_shard_disabled(query: Option<&str>) -> ErrorResponse { ErrorResponse { severity: "ERROR".into(), code: "58000".into(), message: "cross-shard queries are disabled".into(), - detail: Some("query doesn't have a sharding key".into()), + detail: Some(format!( + "query doesn't have a sharding key{}", + if let Some(query) = query { + format!(": {}", query) + } else { + "".into() + } + )), context: None, file: None, routine: None, } } - pub fn client_idle_timeout(duration: Duration) -> ErrorResponse { + pub fn transaction_statement_mode() -> ErrorResponse { + ErrorResponse { + severity: "ERROR".into(), + code: "58000".into(), + message: "transaction control statements are not supported in statement pooler mode" + .into(), + ..Default::default() + } + } + + pub fn client_idle_timeout(duration: Duration, state: &State) -> ErrorResponse { ErrorResponse { severity: "FATAL".into(), code: "57P05".into(), - message: "disconnecting idle client".into(), + message: format!( + "disconnecting {} client", + if state == &State::IdleInTransaction { + "idle in transaction" + } else { + "idle" + } + ), detail: Some(format!( - "client_idle_timeout of {}ms expired", + "{} of {}ms expired", + if state == &State::IdleInTransaction { + "client_idle_in_transaction_timeout" + } else { + "client_idle_timeout" + }, duration.as_millis() )), context: None, @@ -78,11 +120,14 @@ impl ErrorResponse { } /// Connection error. - pub fn connection() -> ErrorResponse { + pub fn connection(user: &str, database: &str) -> ErrorResponse { ErrorResponse { severity: "ERROR".into(), code: "58000".into(), - message: "connection pool is down".into(), + message: format!( + r#"connection pool for user "{}" and database "{}" is down"#, + user, database + ), detail: None, context: None, file: None, @@ -115,6 +160,18 @@ impl ErrorResponse { } } + pub fn tls_required() -> ErrorResponse { + Self { + severity: "FATAL".into(), + code: "08004".into(), + message: "only TLS connections are allowed".into(), + detail: None, + context: None, + file: None, + routine: None, + } + } + pub fn from_err(err: &impl std::error::Error) -> Self { let message = err.to_string(); Self { @@ -163,6 +220,8 @@ impl Display for ErrorResponse { impl FromBytes for ErrorResponse { fn from_bytes(mut bytes: Bytes) -> Result { + code!(bytes, 'E'); + let _len = bytes.get_i32(); let mut error_response = ErrorResponse::default(); diff --git a/pgdog/src/net/messages/execute.rs b/pgdog/src/net/messages/execute.rs index dace25ba0..1be68080a 100644 --- a/pgdog/src/net/messages/execute.rs +++ b/pgdog/src/net/messages/execute.rs @@ -6,7 +6,7 @@ use crate::net::c_string_buf_len; use super::code; use super::prelude::*; -#[derive(Clone)] +#[derive(Clone, PartialEq)] pub struct Execute { payload: Bytes, portal_len: usize, diff --git a/pgdog/src/net/messages/hello.rs b/pgdog/src/net/messages/hello.rs index 906418ad2..84f901b06 100644 --- a/pgdog/src/net/messages/hello.rs +++ b/pgdog/src/net/messages/hello.rs @@ -54,7 +54,10 @@ impl Startup { let value = c_string(stream).await?; - if name == "options" { + if name == "search_path" { + let value = search_path(&value); + params.insert(name, value); + } else if name == "options" { let kvs = value.split("-c"); for kv in kvs { let mut nvs = kv.split("="); @@ -66,6 +69,11 @@ impl Startup { let name = name.trim().to_string(); let value = value.trim().to_string(); if !name.is_empty() && !value.is_empty() { + let value = if name == "search_path" { + search_path(&value) + } else { + ParameterValue::from(value) + }; params.insert(name, value); } } @@ -233,6 +241,14 @@ impl FromBytes for SslReply { } } +fn search_path(value: &str) -> ParameterValue { + let value = value + .split(",") + .map(|value| value.to_string()) + .collect::>(); + ParameterValue::Tuple(value) +} + #[cfg(test)] mod test { use crate::net::messages::ToBytes; diff --git a/pgdog/src/net/messages/mod.rs b/pgdog/src/net/messages/mod.rs index 66651b23c..f080ac706 100644 --- a/pgdog/src/net/messages/mod.rs +++ b/pgdog/src/net/messages/mod.rs @@ -3,6 +3,7 @@ pub mod auth; pub mod backend_key; pub mod bind; pub mod bind_complete; +pub mod buffer; pub mod close; pub mod close_complete; pub mod command_complete; @@ -17,6 +18,7 @@ pub mod error_response; pub mod execute; pub mod flush; pub mod hello; +pub mod no_data; pub mod notice_response; pub mod notification_response; pub mod parameter_description; @@ -36,6 +38,7 @@ pub use auth::{Authentication, Password}; pub use backend_key::BackendKeyData; pub use bind::{Bind, Format, Parameter, ParameterWithFormat}; pub use bind_complete::BindComplete; +pub use buffer::{MessageBuffer, MessageBufferStats}; pub use close::Close; pub use close_complete::CloseComplete; pub use command_complete::CommandComplete; @@ -50,6 +53,7 @@ pub use error_response::ErrorResponse; pub use execute::Execute; pub use flush::Flush; pub use hello::Startup; +pub use no_data::NoData; pub use notice_response::NoticeResponse; pub use notification_response::NotificationResponse; pub use parameter_description::ParameterDescription; @@ -257,3 +261,44 @@ macro_rules! code { } pub(crate) use code; + +macro_rules! from_message { + ($ty:tt) => { + impl TryFrom for $ty { + type Error = crate::net::Error; + + fn try_from(message: Message) -> Result<$ty, Self::Error> { + <$ty as FromBytes>::from_bytes(message.to_bytes()?) + } + } + }; +} + +from_message!(Authentication); +from_message!(BackendKeyData); +from_message!(Bind); +from_message!(BindComplete); +from_message!(Close); +from_message!(CloseComplete); +from_message!(CommandComplete); +from_message!(CopyData); +from_message!(CopyDone); +from_message!(CopyFail); +from_message!(DataRow); +from_message!(Describe); +from_message!(EmptyQueryResponse); +from_message!(ErrorResponse); +from_message!(Execute); +from_message!(Flush); +from_message!(NoData); +from_message!(NoticeResponse); +from_message!(NotificationResponse); +from_message!(ParameterDescription); +from_message!(ParameterStatus); +from_message!(Parse); +from_message!(ParseComplete); +from_message!(Query); +from_message!(ReadyForQuery); +from_message!(RowDescription); +from_message!(Sync); +from_message!(Terminate); diff --git a/pgdog/src/net/messages/no_data.rs b/pgdog/src/net/messages/no_data.rs new file mode 100644 index 000000000..61bc6424d --- /dev/null +++ b/pgdog/src/net/messages/no_data.rs @@ -0,0 +1,24 @@ +use super::{code, prelude::*}; + +#[derive(Debug, Clone, Default)] +pub struct NoData; + +impl FromBytes for NoData { + fn from_bytes(mut bytes: Bytes) -> Result { + code!(bytes, 'n'); + Ok(Self) + } +} + +impl ToBytes for NoData { + fn to_bytes(&self) -> Result { + let payload = Payload::named(self.code()); + Ok(payload.freeze()) + } +} + +impl Protocol for NoData { + fn code(&self) -> char { + 'n' + } +} diff --git a/pgdog/src/net/messages/parameter_description.rs b/pgdog/src/net/messages/parameter_description.rs index aa27d431f..729abf7ea 100644 --- a/pgdog/src/net/messages/parameter_description.rs +++ b/pgdog/src/net/messages/parameter_description.rs @@ -1,7 +1,7 @@ use super::code; use super::prelude::*; -#[derive(Debug, Clone)] +#[derive(Debug, Clone, Default)] pub struct ParameterDescription { params: Vec, } @@ -10,9 +10,9 @@ impl FromBytes for ParameterDescription { fn from_bytes(mut bytes: Bytes) -> Result { code!(bytes, 't'); let _len = bytes.get_i32(); - let num_params = bytes.get_i16(); - let mut params = vec![]; - for _ in 0..num_params { + let num_params = bytes.get_u16(); + let mut params = Vec::with_capacity(num_params as usize); + for _ in 0..num_params as usize { params.push(bytes.get_i32()); } Ok(Self { params }) @@ -22,7 +22,7 @@ impl FromBytes for ParameterDescription { impl ToBytes for ParameterDescription { fn to_bytes(&self) -> Result { let mut payload = Payload::named(self.code()); - payload.put_i16(self.params.len() as i16); + payload.put_u16(self.params.len() as u16); for param in &self.params { payload.put_i32(*param); } @@ -36,3 +36,59 @@ impl Protocol for ParameterDescription { 't' } } + +impl ParameterDescription { + /// Create an empty parameter description. + pub fn empty() -> Self { + Self { params: Vec::new() } + } + + /// Create a parameter description from a list of type OIDs. + pub fn from_params(params: Vec) -> Self { + Self { params } + } +} + +#[cfg(test)] +mod test { + use super::*; + + #[test] + fn parameter_description_round_trip_small() { + let params = vec![23, 42, 87]; + let description = ParameterDescription { + params: params.clone(), + }; + + let bytes = description.to_bytes().unwrap(); + let mut buf = bytes.clone(); + assert_eq!(buf.get_u8(), b't'); + let len = buf.get_i32(); + assert_eq!(len as usize, bytes.len() - 1); // message length excludes the leading code + assert_eq!(buf.get_u16(), params.len() as u16); + + let decoded = ParameterDescription::from_bytes(bytes).unwrap(); + assert_eq!(decoded.params, params); + } + + #[test] + fn parameter_description_round_trip_max_parameter_count() { + let count = u16::MAX as usize; + let params: Vec = (0..count).map(|i| i as i32).collect(); + let description = ParameterDescription { + params: params.clone(), + }; + + let bytes = description.to_bytes().unwrap(); + let mut buf = bytes.clone(); + assert_eq!(buf.get_u8(), b't'); + let len = buf.get_i32(); + assert_eq!(len as usize, bytes.len() - 1); + assert_eq!(buf.get_u16(), count as u16); + + let decoded = ParameterDescription::from_bytes(bytes).unwrap(); + assert_eq!(decoded.params.len(), count); + assert_eq!(decoded.params[0], 0); + assert_eq!(decoded.params[count - 1], (count - 1) as i32); + } +} diff --git a/pgdog/src/net/messages/parameter_status.rs b/pgdog/src/net/messages/parameter_status.rs index 7cfafde98..15c573989 100644 --- a/pgdog/src/net/messages/parameter_status.rs +++ b/pgdog/src/net/messages/parameter_status.rs @@ -3,6 +3,7 @@ use crate::net::{ c_string_buf, messages::{code, prelude::*}, + parameter::ParameterValue, Parameter, }; @@ -12,7 +13,7 @@ pub struct ParameterStatus { /// Parameter name, e.g. `client_encoding`. pub name: String, /// Parameter value, e.g. `UTF8`. - pub value: String, + pub value: ParameterValue, } impl From for ParameterStatus { @@ -28,7 +29,7 @@ impl From<(T, T)> for ParameterStatus { fn from(value: (T, T)) -> Self { Self { name: value.0.to_string(), - value: value.1.to_string(), + value: ParameterValue::String(value.1.to_string()), } } } @@ -61,7 +62,7 @@ impl ParameterStatus { }, ParameterStatus { name: "server_version".into(), - value: env!("CARGO_PKG_VERSION").to_string() + " (PgDog)", + value: (env!("CARGO_PKG_VERSION").to_string() + " (PgDog)").into(), }, ParameterStatus { name: "standard_conforming_strings".into(), @@ -76,7 +77,7 @@ impl ToBytes for ParameterStatus { let mut payload = Payload::named(self.code()); payload.put_string(&self.name); - payload.put_string(&self.value); + payload.put(self.value.to_bytes()?); Ok(payload.freeze()) } @@ -89,7 +90,15 @@ impl FromBytes for ParameterStatus { let _len = bytes.get_i32(); let name = c_string_buf(&mut bytes); - let value = c_string_buf(&mut bytes); + let mut value = c_string_buf(&mut bytes) + .split(",") + .map(|s| s.trim().to_string()) + .collect::>(); + let value = if value.len() == 1 { + ParameterValue::String(value.pop().unwrap()) + } else { + ParameterValue::Tuple(value) + }; Ok(Self { name, value }) } @@ -100,3 +109,106 @@ impl Protocol for ParameterStatus { 'S' } } + +#[cfg(test)] +mod test { + use super::*; + use bytes::BytesMut; + + fn make_parameter_status_bytes(name: &str, value: &str) -> Bytes { + let mut buf = BytesMut::new(); + buf.put_u8(b'S'); + // Length: 4 bytes for len + name + null + value + null + let len = 4 + name.len() + 1 + value.len() + 1; + buf.put_i32(len as i32); + buf.put_slice(name.as_bytes()); + buf.put_u8(0); + buf.put_slice(value.as_bytes()); + buf.put_u8(0); + buf.freeze() + } + + #[test] + fn test_from_bytes_single_value() { + let bytes = make_parameter_status_bytes("client_encoding", "UTF8"); + let status = ParameterStatus::from_bytes(bytes).unwrap(); + + assert_eq!(status.name, "client_encoding"); + assert_eq!(status.value, ParameterValue::String("UTF8".into())); + } + + #[test] + fn test_from_bytes_comma_separated_tuple() { + let bytes = make_parameter_status_bytes("search_path", "$user, public"); + let status = ParameterStatus::from_bytes(bytes).unwrap(); + + assert_eq!(status.name, "search_path"); + assert_eq!( + status.value, + ParameterValue::Tuple(vec!["$user".into(), "public".into()]) + ); + } + + #[test] + fn test_from_bytes_comma_separated_with_extra_spaces() { + let bytes = make_parameter_status_bytes("search_path", " $user , public "); + let status = ParameterStatus::from_bytes(bytes).unwrap(); + + assert_eq!(status.name, "search_path"); + assert_eq!( + status.value, + ParameterValue::Tuple(vec!["$user".into(), "public".into()]) + ); + } + + #[test] + fn test_to_bytes_roundtrip_string() { + let original = ParameterStatus { + name: "client_encoding".into(), + value: ParameterValue::String("UTF8".into()), + }; + + let bytes = original.to_bytes().unwrap(); + let parsed = ParameterStatus::from_bytes(bytes).unwrap(); + + assert_eq!(parsed.name, original.name); + assert_eq!(parsed.value, original.value); + } + + #[test] + fn test_from_tuple() { + let status: ParameterStatus = ("test_name", "test_value").into(); + + assert_eq!(status.name, "test_name"); + assert_eq!(status.value, ParameterValue::String("test_value".into())); + } + + #[test] + fn test_from_parameter() { + let param = Parameter { + name: "search_path".into(), + value: ParameterValue::Tuple(vec!["$user".into(), "public".into()]), + }; + + let status: ParameterStatus = param.into(); + + assert_eq!(status.name, "search_path"); + assert_eq!( + status.value, + ParameterValue::Tuple(vec!["$user".into(), "public".into()]) + ); + } + + #[test] + fn test_to_parameter() { + let status = ParameterStatus { + name: "timezone".into(), + value: ParameterValue::String("UTC".into()), + }; + + let param: Parameter = status.into(); + + assert_eq!(param.name, "timezone"); + assert_eq!(param.value, ParameterValue::String("UTC".into())); + } +} diff --git a/pgdog/src/net/messages/parse.rs b/pgdog/src/net/messages/parse.rs index 06e2be6a7..f6a95ab8a 100644 --- a/pgdog/src/net/messages/parse.rs +++ b/pgdog/src/net/messages/parse.rs @@ -33,12 +33,12 @@ impl Debug for Parse { } impl Parse { + /// Length in bytes of the message. pub fn len(&self) -> usize { self.name.len() + self.query.len() + self.data_types.len() + 5 } /// New anonymous prepared statement. - #[cfg(test)] pub fn new_anonymous(query: &str) -> Self { Self { name: Bytes::from("\0"), diff --git a/pgdog/src/net/messages/query.rs b/pgdog/src/net/messages/query.rs index 20529699f..72188804b 100644 --- a/pgdog/src/net/messages/query.rs +++ b/pgdog/src/net/messages/query.rs @@ -6,7 +6,7 @@ use bytes::Bytes; use std::str::{from_utf8, from_utf8_unchecked}; /// Query (F) message. -#[derive(Clone)] +#[derive(Clone, PartialEq)] pub struct Query { /// Query string. pub payload: Bytes, @@ -39,6 +39,13 @@ impl Query { // Don't read the trailing null byte. unsafe { from_utf8_unchecked(&self.payload[5..self.payload.len() - 1]) } } + + /// Update the SQL query. + pub fn set_query(&mut self, query: &str) { + let mut payload = Payload::named('Q'); + payload.put_string(query); + self.payload = payload.freeze(); + } } impl FromBytes for Query { diff --git a/pgdog/src/net/messages/replication/hot_standby_feedback.rs b/pgdog/src/net/messages/replication/hot_standby_feedback.rs index e7bcea243..3896c5f41 100644 --- a/pgdog/src/net/messages/replication/hot_standby_feedback.rs +++ b/pgdog/src/net/messages/replication/hot_standby_feedback.rs @@ -39,3 +39,28 @@ impl ToBytes for HotStandbyFeedback { Ok(payload.freeze()) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn hot_standby_feedback_roundtrip() { + let feedback = HotStandbyFeedback { + system_clock: 1234, + global_xmin: 42, + epoch: 7, + catalog_min: 11, + epoch_catalog_min: 13, + }; + + let bytes = feedback.to_bytes().expect("serialize hot standby feedback"); + let decoded = HotStandbyFeedback::from_bytes(bytes).expect("decode hot standby feedback"); + + assert_eq!(decoded.system_clock, 1234); + assert_eq!(decoded.global_xmin, 42); + assert_eq!(decoded.epoch, 7); + assert_eq!(decoded.catalog_min, 11); + assert_eq!(decoded.epoch_catalog_min, 13); + } +} diff --git a/pgdog/src/net/messages/replication/keep_alive.rs b/pgdog/src/net/messages/replication/keep_alive.rs index 5ef938f22..81e191dc0 100644 --- a/pgdog/src/net/messages/replication/keep_alive.rs +++ b/pgdog/src/net/messages/replication/keep_alive.rs @@ -46,3 +46,33 @@ impl ToBytes for KeepAlive { Ok(payload.freeze()) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn keep_alive_roundtrip_and_reply_flag() { + let ka = KeepAlive { + wal_end: 9876, + system_clock: 5432, + reply: 1, + }; + + assert!(ka.reply()); + + let bytes = ka.to_bytes().expect("serialize keepalive"); + let decoded = KeepAlive::from_bytes(bytes).expect("decode keepalive"); + + assert_eq!(decoded.wal_end, 9876); + assert_eq!(decoded.system_clock, 5432); + assert_eq!(decoded.reply, 1); + assert!(decoded.reply()); + + let wrapped = decoded.wrapped().expect("wrap keepalive copydata"); + let meta = wrapped.replication_meta().expect("decode replication meta"); + matches!(meta, ReplicationMeta::KeepAlive(_)) + .then_some(()) + .expect("should be keepalive meta"); + } +} diff --git a/pgdog/src/net/messages/replication/logical/delete.rs b/pgdog/src/net/messages/replication/logical/delete.rs index d722c54d9..ef1f8633f 100644 --- a/pgdog/src/net/messages/replication/logical/delete.rs +++ b/pgdog/src/net/messages/replication/logical/delete.rs @@ -1,3 +1,5 @@ +use crate::net::replication::logical::tuple_data::Identifier; + use super::super::super::code; use super::super::super::prelude::*; use super::tuple_data::TupleData; @@ -9,6 +11,23 @@ pub struct Delete { pub old: Option, } +impl Delete { + pub fn key_non_null(&self) -> Option { + if let Some(ref key) = self.key { + let columns = key + .columns + .clone() + .into_iter() + .filter(|column| column.identifier != Identifier::Null) + .collect(); + + Some(TupleData { columns }) + } else { + None + } + } +} + impl FromBytes for Delete { fn from_bytes(mut bytes: Bytes) -> Result { code!(bytes, 'D'); diff --git a/pgdog/src/net/messages/replication/logical/tuple_data.rs b/pgdog/src/net/messages/replication/logical/tuple_data.rs index e90778487..71b1bcfc5 100644 --- a/pgdog/src/net/messages/replication/logical/tuple_data.rs +++ b/pgdog/src/net/messages/replication/logical/tuple_data.rs @@ -74,7 +74,7 @@ impl TupleData { .columns .iter() .map(|c| { - if c.data.is_empty() { + if c.identifier == Identifier::Null { Parameter::new_null() } else { Parameter::new(&c.data) @@ -87,7 +87,7 @@ impl TupleData { } /// Explains what's inside the column. -#[derive(Debug, Clone)] +#[derive(Debug, Clone, PartialEq)] pub enum Identifier { Format(Format), Null, @@ -128,3 +128,32 @@ impl FromBytes for TupleData { Self::from_buffer(&mut bytes) } } + +#[cfg(test)] +mod test { + use super::*; + + #[test] + fn test_null_conversion() { + let data = TupleData { + columns: vec![ + Column { + identifier: Identifier::Null, + len: 0, + data: Bytes::new(), + }, + Column { + identifier: Identifier::Format(Format::Text), + len: 4, + data: Bytes::from(String::from("1234")), + }, + ], + }; + + let bind = data.to_bind("__pgdog_1"); + assert_eq!(bind.statement(), "__pgdog_1"); + assert!(bind.parameter(0).unwrap().unwrap().is_null()); + assert!(!bind.parameter(1).unwrap().unwrap().is_null()); + assert_eq!(bind.parameter(1).unwrap().unwrap().bigint().unwrap(), 1234); + } +} diff --git a/pgdog/src/net/messages/replication/mod.rs b/pgdog/src/net/messages/replication/mod.rs index 09a17d9ea..1f04ead99 100644 --- a/pgdog/src/net/messages/replication/mod.rs +++ b/pgdog/src/net/messages/replication/mod.rs @@ -47,3 +47,61 @@ impl ToBytes for ReplicationMeta { } } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn replication_meta_roundtrip_variants() { + let feedback = HotStandbyFeedback { + system_clock: 1, + global_xmin: 2, + epoch: 3, + catalog_min: 4, + epoch_catalog_min: 5, + }; + let keepalive = KeepAlive { + wal_end: 6, + system_clock: 7, + reply: 0, + }; + let status = StatusUpdate { + last_written: 8, + last_flushed: 9, + last_applied: 10, + system_clock: 11, + reply: 1, + }; + + for meta in [ + ReplicationMeta::HotStandbyFeedback(feedback.clone()), + ReplicationMeta::KeepAlive(keepalive.clone()), + ReplicationMeta::StatusUpdate(status.clone()), + ] { + let bytes = meta.to_bytes().expect("serialize replication meta"); + let decoded = ReplicationMeta::from_bytes(bytes).expect("decode replication meta"); + match (meta, decoded) { + ( + ReplicationMeta::HotStandbyFeedback(expected), + ReplicationMeta::HotStandbyFeedback(actual), + ) => { + assert_eq!(actual.system_clock, expected.system_clock); + assert_eq!(actual.global_xmin, expected.global_xmin); + } + (ReplicationMeta::KeepAlive(expected), ReplicationMeta::KeepAlive(actual)) => { + assert_eq!(actual.wal_end, expected.wal_end); + assert_eq!(actual.system_clock, expected.system_clock); + } + ( + ReplicationMeta::StatusUpdate(expected), + ReplicationMeta::StatusUpdate(actual), + ) => { + assert_eq!(actual.last_written, expected.last_written); + assert_eq!(actual.reply, expected.reply); + } + _ => panic!("replication meta variant mismatch"), + } + } + } +} diff --git a/pgdog/src/net/messages/replication/status_update.rs b/pgdog/src/net/messages/replication/status_update.rs index 496bc3139..9f5e0fb33 100644 --- a/pgdog/src/net/messages/replication/status_update.rs +++ b/pgdog/src/net/messages/replication/status_update.rs @@ -110,4 +110,29 @@ mod test { _ => panic!("not a status update"), } } + + #[test] + fn status_update_from_keepalive_inherits_wal_end() { + let keepalive = KeepAlive { + wal_end: 123, + system_clock: 456, + reply: 0, + }; + + let update: StatusUpdate = keepalive.into(); + assert_eq!(update.last_written, 123); + assert_eq!(update.last_flushed, 123); + assert_eq!(update.last_applied, 123); + assert_eq!(update.reply, 0); + } + + #[test] + fn status_update_new_reply_sets_reply_flag() { + let lsn = Lsn::from_i64(999); + let update = StatusUpdate::new_reply(lsn); + assert_eq!(update.last_written, 999); + assert_eq!(update.last_flushed, 999); + assert_eq!(update.last_applied, 999); + assert_eq!(update.reply, 1); + } } diff --git a/pgdog/src/net/messages/sync.rs b/pgdog/src/net/messages/sync.rs index a5f1a61b5..d537d1102 100644 --- a/pgdog/src/net/messages/sync.rs +++ b/pgdog/src/net/messages/sync.rs @@ -1,7 +1,7 @@ use super::code; use super::prelude::*; -#[derive(Debug, Clone)] +#[derive(Debug, Clone, PartialEq)] pub struct Sync; impl Default for Sync { diff --git a/pgdog/src/net/parameter.rs b/pgdog/src/net/parameter.rs index 398e7f59c..a9db66c6d 100644 --- a/pgdog/src/net/parameter.rs +++ b/pgdog/src/net/parameter.rs @@ -1,7 +1,9 @@ //! Startup parameter. +use bytes::{BufMut, Bytes, BytesMut}; +use tracing::debug; use std::{ - collections::BTreeMap, + collections::{btree_map, BTreeMap}, fmt::Display, hash::{DefaultHasher, Hash, Hasher}, ops::{Deref, DerefMut}, @@ -9,7 +11,7 @@ use std::{ use once_cell::sync::Lazy; -use crate::stats::memory::MemoryUsage; +use crate::{net::ToBytes, stats::memory::MemoryUsage}; use super::{messages::Query, Error}; @@ -19,25 +21,26 @@ static IMMUTABLE_PARAMS: Lazy> = Lazy::new(|| { String::from("user"), String::from("client_encoding"), String::from("replication"), + String::from("pgdog.role"), + String::from("pgdog.shard"), + String::from("pgdog.sharding_key"), ]) }); -// static IMMUTABLE_PARAMS: &[&str] = &["database", "user", "client_encoding"]; - /// Startup parameter. #[derive(Debug, Clone, PartialEq)] pub struct Parameter { /// Parameter name. pub name: String, /// Parameter value. - pub value: String, + pub value: ParameterValue, } impl From<(T, T)> for Parameter { fn from(value: (T, T)) -> Self { Self { name: value.0.to_string(), - value: value.1.to_string(), + value: ParameterValue::String(value.1.to_string()), } } } @@ -52,6 +55,28 @@ pub struct MergeResult { pub enum ParameterValue { String(String), Tuple(Vec), + Integer(i32), +} + +impl ToBytes for ParameterValue { + fn to_bytes(&self) -> Result { + let mut bytes = BytesMut::new(); + match self { + Self::String(string) => bytes.put_slice(string.as_bytes()), + Self::Tuple(ref values) => { + let values = values + .iter() + .map(|value| value.as_bytes().to_vec()) + .collect::>() + .join(", ".as_bytes()); + bytes.put(Bytes::from(values)); + } + Self::Integer(integer) => bytes.put_slice(integer.to_string().as_bytes()), + } + bytes.put_u8(0); + + Ok(bytes.freeze()) + } } impl MemoryUsage for ParameterValue { @@ -60,22 +85,40 @@ impl MemoryUsage for ParameterValue { match self { Self::String(v) => v.memory_usage(), Self::Tuple(vals) => vals.memory_usage(), + Self::Integer(int) => int.memory_usage(), } } } impl Display for ParameterValue { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + fn quote(value: &str) -> String { + let value = if value.starts_with("\"") && value.ends_with("\"") { + let mut value = value.to_string(); + value.remove(0); + value.pop(); + value + } else { + value.to_string() + }; + + if value.is_empty() || value.contains("\"") { + format!("'{}'", value.replace("'", "''")) + } else { + format!(r#""{}""#, value.replace("\"", "\"\"")) + } + } match self { - Self::String(s) => write!(f, "'{}'", s), + Self::String(s) => write!(f, "{}", quote(s)), Self::Tuple(t) => write!( f, "{}", t.iter() - .map(|s| format!("'{}'", s)) + .map(|s| quote(s).to_string()) .collect::>() .join(", ") ), + Self::Integer(int) => write!(f, "{}", int), } } } @@ -104,10 +147,34 @@ impl ParameterValue { /// List of parameters. #[derive(Default, Debug, Clone, PartialEq)] pub struct Parameters { + /// Save parameters set at connection startup & set with `SET` command + /// outside a transaction. params: BTreeMap, + /// Save parameters set with `SET` inside a transaction. These will + /// need to be rolled back or saved depending on if the transaction is + /// rolled back or not. + transaction_params: BTreeMap, + /// Parameters set with `SET LOCAL`. These need to be thrown away no matter + /// what but we need to intercept them for databases that have cross shard + /// queries disabled. + transaction_local_params: BTreeMap, + /// Hash of `params` to avoid syncing params between clients and servers + /// when they are the same. hash: u64, } +impl Display for Parameters { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let output = self + .params + .iter() + .map(|(k, v)| format!("{}={}", k, v)) + .collect::>() + .join(", "); + write!(f, "{}", output) + } +} + impl MemoryUsage for Parameters { #[inline] fn memory_usage(&self) -> usize { @@ -121,6 +188,8 @@ impl From> for Parameters { Self { params: value, hash, + transaction_params: BTreeMap::new(), + transaction_local_params: BTreeMap::new(), } } } @@ -140,6 +209,67 @@ impl Parameters { result } + /// Recompute hash when params are cleared. + pub fn clear(&mut self) { + self.params.clear(); + self.hash = Self::compute_hash(&self.params); + } + + /// Get parameter. + pub fn get(&self, name: &str) -> Option<&ParameterValue> { + if let Some(param) = self.transaction_local_params.get(name) { + Some(param) + } else if let Some(param) = self.transaction_params.get(name) { + Some(param) + } else { + self.params.get(name) + } + } + + /// Get an iterator over in-transaction params. + pub fn in_transaction_iter(&self) -> btree_map::Iter<'_, String, ParameterValue> { + self.transaction_params.iter() + } + + /// Insert a parameter, but only for the duration of the transaction. + pub fn insert_transaction( + &mut self, + name: impl ToString, + value: impl Into, + local: bool, + ) -> Option { + let name = name.to_string().to_lowercase(); + if local { + self.transaction_local_params.insert(name, value.into()) + } else { + self.transaction_params.insert(name, value.into()) + } + } + + /// Commit params we saved during the transaction. + pub fn commit(&mut self) -> bool { + debug!( + "saved {} in-transaction params", + self.transaction_params.len() + ); + let changed = !self.transaction_params.is_empty(); + + self.params + .extend(std::mem::take(&mut self.transaction_params)); + self.transaction_local_params.clear(); + + if changed { + self.hash = Self::compute_hash(&self.params); + } + changed + } + + /// Remove any params we saved during the transaction. + pub fn rollback(&mut self) { + self.transaction_params.clear(); + self.transaction_local_params.clear(); + } + fn compute_hash(params: &BTreeMap) -> u64 { let mut hasher = DefaultHasher::new(); let mut entries = 0; @@ -176,11 +306,38 @@ impl Parameters { self.hash == other.hash } - pub fn set_queries(&self) -> Vec { - self.params - .iter() - .map(|(name, value)| Query::new(format!(r#"SET "{}" TO {}"#, name, value))) - .collect() + /// Generate SET queries to change server state. + /// + /// # Arguments + /// + /// * `transaction`: Generate `SET` statements from in-transaction params only. + /// + pub fn set_queries(&self, transaction_only: bool) -> Vec { + fn query(name: &str, value: &ParameterValue, local: bool) -> Query { + let set = if local { "SET LOCAL" } else { "SET" }; + Query::new(format!(r#"{} "{}" TO {}"#, set, name, value)) + } + + if transaction_only { + let mut sets = self + .transaction_params + .iter() + .map(|(key, value)| query(key, value, false)) + .collect::>(); + + sets.extend( + self.transaction_local_params + .iter() + .map(|(key, value)| query(key, value, true)), + ); + + sets + } else { + self.params + .iter() + .map(|(key, value)| query(key, value, false)) + .collect() + } } pub fn reset_queries(&self) -> Vec { @@ -218,6 +375,36 @@ impl Parameters { self.get(name) .map_or(default_value, |p| p.as_str().unwrap_or(default_value)) } + + /// Merge other into self. + pub fn merge(&mut self, other: Self) { + self.params.extend(other.params); + self.transaction_params.extend(other.transaction_params); + self.transaction_local_params + .extend(other.transaction_local_params); + Self::compute_hash(&self.params); + } + + /// Copy params set inside the transaction. + pub fn copy_in_transaction(&mut self, other: &Self) { + self.transaction_params.extend( + other + .transaction_params + .iter() + .map(|(key, value)| (key.clone(), value.clone())), + ); + self.transaction_local_params.extend( + other + .transaction_local_params + .iter() + .map(|(key, value)| (key.clone(), value.clone())), + ); + } + + /// Get search_path, if set. + pub fn search_path(&self) -> Option<&ParameterValue> { + self.get("search_path") + } } impl Deref for Parameters { @@ -238,10 +425,15 @@ impl From> for Parameters { fn from(value: Vec) -> Self { let params = value .into_iter() - .map(|p| (p.name, ParameterValue::String(p.value))) + .map(|p| (p.name, p.value)) .collect::>(); let hash = Self::compute_hash(¶ms); - Self { params, hash } + Self { + params, + hash, + transaction_params: BTreeMap::new(), + transaction_local_params: BTreeMap::new(), + } } } @@ -251,7 +443,7 @@ impl From<&Parameters> for Vec { for (key, value) in &val.params { result.push(Parameter { name: key.to_string(), - value: value.to_string(), + value: value.clone(), }); } @@ -261,7 +453,9 @@ impl From<&Parameters> for Vec { #[cfg(test)] mod test { + use crate::backend::server::test::test_server; use crate::net::parameter::ParameterValue; + use crate::net::ToBytes; use super::Parameters; @@ -283,4 +477,271 @@ mod test { assert!(Parameters::default().identical(&Parameters::default())); } + + #[test] + fn test_insert_transaction_non_local() { + let mut params = Parameters::default(); + params.insert("application_name", "test"); + params.insert_transaction("search_path", "public", false); + + // Transaction param should be accessible via get + assert_eq!( + params.get("search_path"), + Some(&ParameterValue::String("public".into())) + ); + + // Regular param should still be accessible + assert_eq!( + params.get("application_name"), + Some(&ParameterValue::String("test".into())) + ); + } + + #[test] + fn test_insert_transaction_local() { + let mut params = Parameters::default(); + params.insert_transaction("search_path", "public", true); + + // Local param should be accessible via get + assert_eq!( + params.get("search_path"), + Some(&ParameterValue::String("public".into())) + ); + } + + #[test] + fn test_get_priority_local_over_transaction() { + let mut params = Parameters::default(); + params.insert("search_path", "base"); + params.insert_transaction("search_path", "transaction", false); + params.insert_transaction("search_path", "local", true); + + // Local should take priority + assert_eq!( + params.get("search_path"), + Some(&ParameterValue::String("local".into())) + ); + } + + #[test] + fn test_get_priority_transaction_over_regular() { + let mut params = Parameters::default(); + params.insert("search_path", "base"); + params.insert_transaction("search_path", "transaction", false); + + // Transaction should take priority over regular + assert_eq!( + params.get("search_path"), + Some(&ParameterValue::String("transaction".into())) + ); + } + + #[test] + fn test_commit_clears_local_params() { + let mut params = Parameters::default(); + params.insert_transaction("search_path", "transaction", false); + params.insert_transaction("timezone", "local_tz", true); + + assert!(params.commit()); + + // Transaction param should be committed to regular params + assert_eq!( + params.get("search_path"), + Some(&ParameterValue::String("transaction".into())) + ); + + // Local param should be cleared (not committed) + assert_eq!(params.get("timezone"), None); + } + + #[test] + fn test_rollback_clears_both_transaction_and_local() { + let mut params = Parameters::default(); + params.insert("base", "value"); + params.insert_transaction("search_path", "transaction", false); + params.insert_transaction("timezone", "local_tz", true); + + params.rollback(); + + // Both transaction and local params should be cleared + assert_eq!(params.get("search_path"), None); + assert_eq!(params.get("timezone"), None); + + // Base param should remain + assert_eq!( + params.get("base"), + Some(&ParameterValue::String("value".into())) + ); + } + + #[test] + fn test_set_queries_transaction_only_includes_set_local() { + let mut params = Parameters::default(); + params.insert_transaction("search_path", "public", false); + params.insert_transaction("timezone", "UTC", true); + + let queries = params.set_queries(true); + + assert_eq!(queries.len(), 2); + + // Check that we have both SET and SET LOCAL queries + let query_strings: Vec = queries.iter().map(|q| q.query().to_string()).collect(); + + assert!(query_strings + .iter() + .any(|q| q.contains("SET \"search_path\"") && !q.contains("SET LOCAL"))); + assert!(query_strings.iter().any(|q| q.contains("SET LOCAL"))); + } + + #[test] + fn test_copy_in_transaction() { + let mut source = Parameters::default(); + source.insert_transaction("search_path", "public", false); + source.insert_transaction("timezone", "UTC", true); + + let mut dest = Parameters::default(); + dest.copy_in_transaction(&source); + + // Both transaction and local params should be copied + assert_eq!( + dest.get("search_path"), + Some(&ParameterValue::String("public".into())) + ); + assert_eq!( + dest.get("timezone"), + Some(&ParameterValue::String("UTC".into())) + ); + } + + #[test] + fn test_parameter_value_to_bytes_string() { + let value = ParameterValue::String("test".into()); + let bytes = value.to_bytes().unwrap(); + + assert_eq!(&bytes[..], b"test\0"); + } + + #[test] + fn test_parameter_value_to_bytes_tuple() { + let value = ParameterValue::Tuple(vec!["a".into(), "b".into()]); + let bytes = value.to_bytes().unwrap(); + + assert_eq!(&bytes[..], b"a, b\0"); + } + + #[test] + fn test_parameter_value_display_string() { + let value = ParameterValue::String("test".into()); + assert_eq!(format!("{}", value), r#""test""#); + } + + #[test] + fn test_parameter_value_display_tuple() { + let value = ParameterValue::Tuple(vec!["$user".into(), "public".into()]); + assert_eq!(format!("{}", value), r#""$user", "public""#); + } + + #[test] + fn test_parameter_value_display_already_quoted() { + // If value is already quoted, it should strip quotes and re-quote + let value = ParameterValue::String(r#""already quoted""#.into()); + assert_eq!(format!("{}", value), r#""already quoted""#); + } + + #[test] + fn test_merge_includes_local_params() { + let mut params1 = Parameters::default(); + params1.insert("base", "value"); + + let mut params2 = Parameters::default(); + params2.insert_transaction("search_path", "public", false); + params2.insert_transaction("timezone", "UTC", true); + + params1.merge(params2); + + // All params should be merged + assert_eq!( + params1.get("base"), + Some(&ParameterValue::String("value".into())) + ); + assert_eq!( + params1.get("search_path"), + Some(&ParameterValue::String("public".into())) + ); + assert_eq!( + params1.get("timezone"), + Some(&ParameterValue::String("UTC".into())) + ); + } + + #[test] + fn test_json_parameter_value() { + assert_eq!( + ParameterValue::String(r#"{"sampling_state":"1","span_id":"2a9abb846bb02bfe","trace_id":"6b9e798174650d2f6e8262ec175f241f"}"#.into()).to_string(), + r#"'{"sampling_state":"1","span_id":"2a9abb846bb02bfe","trace_id":"6b9e798174650d2f6e8262ec175f241f"}'"# + ); + } + + #[test] + fn test_empty_parameter_value() { + assert_eq!(ParameterValue::String("".into()).to_string(), "''"); + } + + #[test] + fn test_clear_resets_hash() { + let mut params = Parameters::default(); + params.insert("application_name", "test_app"); + params.insert("TimeZone", "UTC"); + + // Verify params are not identical to empty (hash differs) + assert!(!params.identical(&Parameters::default())); + + params.clear(); + + // After clear, hash should be reset to match empty Parameters + assert!(params.identical(&Parameters::default())); + } + + #[tokio::test] + async fn test_escape_chars() { + let mut server = test_server().await; + + for quote in ["'", "\""] { + let base = r#"my_app_nameQUOTE;CREATE/**/TABLE/**/poc_table_two/**/(dummy_column/**/INTEGER);SET/**/application_name/**/TO/**/QUOTEyour_app_name"#; + let param = base.replace("QUOTE", quote); + // Postgres truncates identifiers. + let truncated = + r#"my_app_nameQUOTE;CREATE/**/TABLE/**/poc_table_two/**/(dummy_column/"# + .replace("QUOTE", quote); + + let mut params = Parameters::default(); + params.insert("application_name", param.clone()); + + let query = params.set_queries(false).first().unwrap().clone(); + + assert!(query.query().contains(¶m)); + + server.execute(query).await.unwrap(); + + let param: Vec = server.fetch_all("SHOW application_name").await.unwrap(); + let param = param.first().unwrap().clone(); + + assert_eq!(param, truncated); + } + } + + #[tokio::test] + async fn test_set_with_server() { + let mut server = test_server().await; + + let mut params = Parameters::default(); + params.insert("application_name", "test_set_with_server"); + + let query = params.set_queries(false).first().unwrap().clone(); + server.execute(query).await.unwrap(); + + let param: Vec = server.fetch_all("SHOW application_name").await.unwrap(); + let param = param.first().unwrap(); + assert_eq!(param, "test_set_with_server"); + } } diff --git a/pgdog/src/net/protocol_message.rs b/pgdog/src/net/protocol_message.rs index 6460e61b8..f76f5ff47 100644 --- a/pgdog/src/net/protocol_message.rs +++ b/pgdog/src/net/protocol_message.rs @@ -6,7 +6,7 @@ use super::{ Protocol, Query, Sync, ToBytes, }; -#[derive(Debug, Clone)] +#[derive(Debug, Clone, PartialEq)] pub enum ProtocolMessage { Bind(Bind), Parse(Parse), @@ -31,6 +31,17 @@ impl ProtocolMessage { ) } + pub fn anonymous(&self) -> bool { + use ProtocolMessage::*; + + match self { + Bind(bind) => bind.anonymous(), + Parse(parse) => parse.anonymous(), + Describe(describe) => describe.anonymous(), + _ => false, + } + } + pub fn len(&self) -> usize { match self { Self::Bind(bind) => bind.len(), @@ -93,7 +104,9 @@ impl ToBytes for ProtocolMessage { Self::Bind(bind) => bind.to_bytes(), Self::Parse(parse) => parse.to_bytes(), Self::Describe(describe) => describe.to_bytes(), - Self::Prepare { statement, .. } => Query::new(statement).to_bytes(), + Self::Prepare { statement, name } => { + Query::new(format!("PREPARE {} AS {}", name, statement)).to_bytes() + } Self::Execute(execute) => execute.to_bytes(), Self::Close(close) => close.to_bytes(), Self::Query(query) => query.to_bytes(), diff --git a/pgdog/src/net/stream.rs b/pgdog/src/net/stream.rs index fb59259e0..0a39eb1ff 100644 --- a/pgdog/src/net/stream.rs +++ b/pgdog/src/net/stream.rs @@ -4,7 +4,7 @@ use bytes::{BufMut, BytesMut}; use pin_project::pin_project; use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, BufStream, ReadBuf}; use tokio::net::TcpStream; -use tracing::{debug, enabled, trace, Level}; +use tracing::trace; use std::io::{Error, ErrorKind}; use std::net::SocketAddr; @@ -14,16 +14,26 @@ use std::task::Context; use super::messages::{ErrorResponse, Message, Protocol, ReadyForQuery, Terminate}; -/// A network socket. -#[pin_project(project = StreamProjection)] +/// Inner stream types. +#[pin_project(project = StreamInnerProjection)] #[derive(Debug)] #[allow(clippy::large_enum_variant)] -pub enum Stream { +enum StreamInner { Plain(#[pin] BufStream), Tls(#[pin] BufStream>), DevNull, } +/// A network socket. +#[pin_project] +#[derive(Debug)] +pub struct Stream { + #[pin] + inner: StreamInner, + io_in_progress: bool, + capacity: usize, +} + impl AsyncRead for Stream { fn poll_read( self: Pin<&mut Self>, @@ -31,10 +41,10 @@ impl AsyncRead for Stream { buf: &mut ReadBuf<'_>, ) -> std::task::Poll> { let project = self.project(); - match project { - StreamProjection::Plain(stream) => stream.poll_read(cx, buf), - StreamProjection::Tls(stream) => stream.poll_read(cx, buf), - StreamProjection::DevNull => std::task::Poll::Ready(Ok(())), + match project.inner.project() { + StreamInnerProjection::Plain(stream) => stream.poll_read(cx, buf), + StreamInnerProjection::Tls(stream) => stream.poll_read(cx, buf), + StreamInnerProjection::DevNull => std::task::Poll::Ready(Ok(())), } } } @@ -46,10 +56,10 @@ impl AsyncWrite for Stream { buf: &[u8], ) -> std::task::Poll> { let project = self.project(); - match project { - StreamProjection::Plain(stream) => stream.poll_write(cx, buf), - StreamProjection::Tls(stream) => stream.poll_write(cx, buf), - StreamProjection::DevNull => std::task::Poll::Ready(Ok(buf.len())), + match project.inner.project() { + StreamInnerProjection::Plain(stream) => stream.poll_write(cx, buf), + StreamInnerProjection::Tls(stream) => stream.poll_write(cx, buf), + StreamInnerProjection::DevNull => std::task::Poll::Ready(Ok(buf.len())), } } @@ -58,10 +68,10 @@ impl AsyncWrite for Stream { cx: &mut Context<'_>, ) -> std::task::Poll> { let project = self.project(); - match project { - StreamProjection::Plain(stream) => stream.poll_flush(cx), - StreamProjection::Tls(stream) => stream.poll_flush(cx), - StreamProjection::DevNull => std::task::Poll::Ready(Ok(())), + match project.inner.project() { + StreamInnerProjection::Plain(stream) => stream.poll_flush(cx), + StreamInnerProjection::Tls(stream) => stream.poll_flush(cx), + StreamInnerProjection::DevNull => std::task::Poll::Ready(Ok(())), } } @@ -70,52 +80,79 @@ impl AsyncWrite for Stream { cx: &mut Context<'_>, ) -> std::task::Poll> { let project = self.project(); - match project { - StreamProjection::Plain(stream) => stream.poll_shutdown(cx), - StreamProjection::Tls(stream) => stream.poll_shutdown(cx), - StreamProjection::DevNull => std::task::Poll::Ready(Ok(())), + match project.inner.project() { + StreamInnerProjection::Plain(stream) => stream.poll_shutdown(cx), + StreamInnerProjection::Tls(stream) => stream.poll_shutdown(cx), + StreamInnerProjection::DevNull => std::task::Poll::Ready(Ok(())), } } } impl Stream { + /// Memory used by the stream buffers. + pub fn memory_usage(&self) -> usize { + self.capacity * 2 + } + /// Wrap an unencrypted TCP stream. - pub fn plain(stream: TcpStream) -> Self { - Self::Plain(BufStream::with_capacity(9126, 9126, stream)) + pub fn plain(stream: TcpStream, capacity: usize) -> Self { + Self { + inner: StreamInner::Plain(BufStream::with_capacity(capacity, capacity, stream)), + io_in_progress: false, + capacity, + } } /// Wrap an encrypted TCP stream. - pub fn tls(stream: tokio_rustls::TlsStream) -> Self { - Self::Tls(BufStream::with_capacity(9126, 9126, stream)) + pub fn tls(stream: tokio_rustls::TlsStream, capacity: usize) -> Self { + Self { + inner: StreamInner::Tls(BufStream::with_capacity(capacity, capacity, stream)), + io_in_progress: false, + capacity, + } + } + + /// Create a dev null stream that discards all data. + pub fn dev_null() -> Self { + Self { + inner: StreamInner::DevNull, + io_in_progress: false, + capacity: 0, + } } /// This is a TLS stream. pub fn is_tls(&self) -> bool { - matches!(self, Self::Tls(_)) + matches!(self.inner, StreamInner::Tls(_)) } /// Get peer address if any. We're not using UNIX sockets (yet) /// so the peer address should always be available. pub fn peer_addr(&self) -> PeerAddr { - match self { - Self::Plain(stream) => stream.get_ref().peer_addr().ok().into(), - Self::Tls(stream) => stream.get_ref().get_ref().0.peer_addr().ok().into(), - Self::DevNull => PeerAddr { addr: None }, + match &self.inner { + StreamInner::Plain(stream) => stream.get_ref().peer_addr().ok().into(), + StreamInner::Tls(stream) => stream.get_ref().get_ref().0.peer_addr().ok().into(), + StreamInner::DevNull => PeerAddr { addr: None }, } } /// Check socket is okay while we wait for something else. pub async fn check(&mut self) -> Result<(), crate::net::Error> { let mut buf = [0u8; 1]; - match self { - Self::Plain(plain) => eof(plain.get_mut().peek(&mut buf).await)?, - Self::Tls(tls) => eof(tls.get_mut().get_mut().0.peek(&mut buf).await)?, - Self::DevNull => 0, + match &mut self.inner { + StreamInner::Plain(plain) => eof(plain.get_mut().peek(&mut buf).await)?, + StreamInner::Tls(tls) => eof(tls.get_mut().get_mut().0.peek(&mut buf).await)?, + StreamInner::DevNull => 0, }; Ok(()) } + /// Get the current io_in_progress state. + pub fn io_in_progress(&self) -> bool { + self.io_in_progress + } + /// Send data via the stream. /// /// # Performance @@ -123,34 +160,34 @@ impl Stream { /// This is fast because the stream is buffered. Make sure to call [`Stream::send_flush`] /// for the last message in the exchange. pub async fn send(&mut self, message: &impl Protocol) -> Result { - let bytes = message.to_bytes()?; - - match self { - Stream::Plain(ref mut stream) => eof(stream.write_all(&bytes).await)?, - Stream::Tls(ref mut stream) => eof(stream.write_all(&bytes).await)?, - Self::DevNull => (), - } - - if !enabled!(Level::TRACE) { - debug!("{:?} <-- {}", self.peer_addr(), message.code()); - } else { - trace!("{:?} <-- {:#?}", self.peer_addr(), message); - } + self.io_in_progress = true; + let result = async { + let bytes = message.to_bytes()?; + + match &mut self.inner { + StreamInner::Plain(ref mut stream) => eof(stream.write_all(&bytes).await)?, + StreamInner::Tls(ref mut stream) => eof(stream.write_all(&bytes).await)?, + StreamInner::DevNull => (), + } - #[cfg(debug_assertions)] - { - use crate::net::messages::FromBytes; - use tracing::error; + #[cfg(debug_assertions)] + { + use crate::net::messages::FromBytes; + use tracing::error; - if message.code() == 'E' { - let error = ErrorResponse::from_bytes(bytes.clone())?; - if !error.message.is_empty() { - error!("{:?} <-- {}", self.peer_addr(), error) + if message.code() == 'E' { + let error = ErrorResponse::from_bytes(bytes.clone())?; + if !error.message.is_empty() { + error!("{:?} <-- {}", self.peer_addr(), error) + } } } - } - Ok(bytes.len()) + Ok(bytes.len()) + } + .await; + self.io_in_progress = false; + result } /// Send data via the stream and flush the buffer, @@ -199,30 +236,35 @@ impl Stream { /// Read data into a buffer, avoiding unnecessary allocations. pub async fn read_buf(&mut self, bytes: &mut BytesMut) -> Result { - let code = eof(self.read_u8().await)?; - let len = eof(self.read_i32().await)?; - - bytes.put_u8(code); - bytes.put_i32(len); - - // Length must be at least 4 bytes. - if len < 4 { - return Err(crate::net::Error::UnexpectedEof); - } + let result = async { + let code = eof(self.read_u8().await)?; + self.io_in_progress = true; + bytes.put_u8(code); + let len = eof(self.read_i32().await)?; + bytes.put_i32(len); + + // Length must be at least 4 bytes. + if len < 4 { + return Err(crate::net::Error::UnexpectedEof); + } - let capacity = len as usize + 1; - bytes.reserve(capacity); // self + 1 byte for the message code - unsafe { - // SAFETY: We reserved the memory above, so it's there. - // It contains garbage but we're about to write to it. - bytes.set_len(capacity); - } + let capacity = len as usize + 1; + bytes.reserve(capacity); // self + 1 byte for the message code + unsafe { + // SAFETY: We reserved the memory above, so it's there. + // It contains garbage but we're about to write to it. + bytes.set_len(capacity); + } - eof(self.read_exact(&mut bytes[5..capacity]).await)?; + eof(self.read_exact(&mut bytes[5..capacity]).await)?; - let message = Message::new(bytes.split().freeze()); + let message = Message::new(bytes.split().freeze()); - Ok(message) + Ok(message) + } + .await; + self.io_in_progress = false; + result } /// Send an error to the client and disconnect gracefully. @@ -254,14 +296,14 @@ impl Stream { /// Get the wrapped TCP stream back. pub(crate) fn take(self) -> Result { - match self { - Self::Plain(stream) => Ok(stream.into_inner()), + match self.inner { + StreamInner::Plain(stream) => Ok(stream.into_inner()), _ => Err(crate::net::Error::UnexpectedTlsRequest), } } } -fn eof(result: std::io::Result) -> Result { +pub fn eof(result: std::io::Result) -> Result { match result { Ok(val) => Ok(val), Err(err) => { @@ -303,3 +345,27 @@ impl std::fmt::Debug for PeerAddr { } } } + +#[cfg(test)] +mod tests { + use super::*; + use tokio::net::TcpListener; + + #[tokio::test] + async fn test_io_in_progress_initially_false() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + + let client = tokio::spawn(async move { TcpStream::connect(addr).await.unwrap() }); + + let (server_stream, _) = listener.accept().await.unwrap(); + let stream = Stream::plain(server_stream, 4096); + + assert!( + !stream.io_in_progress(), + "io_in_progress should be false initially" + ); + + client.await.unwrap(); + } +} diff --git a/pgdog/src/net/tls.rs b/pgdog/src/net/tls.rs index 57a268795..b347aa11c 100644 --- a/pgdog/src/net/tls.rs +++ b/pgdog/src/net/tls.rs @@ -1,9 +1,15 @@ //! TLS configuration. -use std::{path::PathBuf, sync::Arc}; +use std::{ + path::{Path, PathBuf}, + sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, + }, +}; use crate::config::TlsVerifyMode; -use once_cell::sync::OnceCell; +use arc_swap::ArcSwapOption; use rustls::pki_types::{CertificateDer, PrivateKeyDer}; use tokio_rustls::rustls::{ self, @@ -12,154 +18,142 @@ use tokio_rustls::rustls::{ ClientConfig, }; use tokio_rustls::{TlsAcceptor, TlsConnector}; -use tracing::{debug, info}; +use tracing::{debug, info, warn}; use crate::config::config; use super::Error; -static ACCEPTOR: OnceCell> = OnceCell::new(); -static CONNECTOR: OnceCell = OnceCell::new(); +static ACCEPTOR: ArcSwapOption = ArcSwapOption::const_empty(); +static ACCEPTOR_BUILD_COUNT: AtomicUsize = AtomicUsize::new(0); -/// Get preloaded TLS acceptor. -pub fn acceptor() -> Option<&'static TlsAcceptor> { - if let Some(Some(acceptor)) = ACCEPTOR.get() { - return Some(acceptor); - } +static CONNECTOR: ArcSwapOption = ArcSwapOption::const_empty(); - None +#[derive(Clone, Debug, PartialEq)] +struct ConnectorConfigKey { + mode: TlsVerifyMode, + ca_path: Option, } -/// Create a new TLS acceptor from the cert and key. -/// -/// This is not atomic, so call it on startup only. -pub fn load_acceptor(cert: &PathBuf, key: &PathBuf) -> Result, Error> { - if let Some(acceptor) = ACCEPTOR.get() { - return Ok(acceptor.clone()); +impl ConnectorConfigKey { + fn new(mode: TlsVerifyMode, ca_path: Option<&PathBuf>) -> Self { + Self { + mode, + ca_path: ca_path.cloned(), + } } +} - let pem = if let Ok(pem) = CertificateDer::from_pem_file(cert) { - pem - } else { - let _ = ACCEPTOR.set(None); - return Ok(None); - }; +struct ConnectorCacheEntry { + key: ConnectorConfigKey, + config: Arc, +} - let key = if let Ok(key) = PrivateKeyDer::from_pem_file(key) { - key - } else { - let _ = ACCEPTOR.set(None); - return Ok(None); - }; +impl ConnectorCacheEntry { + fn new(key: ConnectorConfigKey, config: Arc) -> Arc { + Arc::new(Self { key, config }) + } - let config = rustls::ServerConfig::builder() - .with_no_client_auth() - .with_single_cert(vec![pem], key)?; + fn connector(&self) -> TlsConnector { + TlsConnector::from(self.config.clone()) + } +} - let acceptor = tokio_rustls::TlsAcceptor::from(Arc::new(config)); +#[cfg(test)] +static CONNECTOR_BUILD_COUNT: AtomicUsize = AtomicUsize::new(0); - info!("πŸ”‘ TLS on"); +#[cfg(test)] +fn increment_connector_build_count() { + CONNECTOR_BUILD_COUNT.fetch_add(1, Ordering::SeqCst); +} - // A bit of a race, but it's not a big deal unless this is called - // with different certificate/secret key. - let _ = ACCEPTOR.set(Some(acceptor.clone())); +#[cfg(not(test))] +fn increment_connector_build_count() {} - Ok(Some(acceptor)) +/// Get the current TLS acceptor snapshot, if TLS is enabled. +pub fn acceptor() -> Option> { + ACCEPTOR.load_full() } -/// Create new TLS connector using the default configuration. +/// Create new TLS connector using the current configuration. pub fn connector() -> Result { - if let Some(connector) = CONNECTOR.get() { - return Ok(connector.clone()); - } - let config = config(); - let connector = connector_with_verify_mode( + connector_with_verify_mode( config.config.general.tls_verify, config.config.general.tls_server_ca_certificate.as_ref(), - )?; - - let _ = CONNECTOR.set(connector.clone()); - - Ok(connector) + ) } /// Preload TLS at startup. pub fn load() -> Result<(), Error> { + reload() +} + +/// Rebuild TLS primitives according to the current configuration. +/// +/// This validates the new settings and swaps them in atomically. If validation +/// fails, the existing TLS acceptor remains active. +pub fn reload() -> Result<(), Error> { + debug!("reloading TLS configuration"); + let config = config(); + let general = &config.config.general; - if let Some((cert, key)) = config.config.general.tls() { - load_acceptor(cert, key)?; - } + // Always validate upstream TLS settings so we surface CA issues early. + let _ = connector_with_verify_mode( + general.tls_verify, + general.tls_server_ca_certificate.as_ref(), + )?; + + let tls_paths = general.tls(); + let new_acceptor = tls_paths + .map(|(cert, key)| build_acceptor(cert, key)) + .transpose()?; - connector()?; + match (new_acceptor, tls_paths) { + (Some(acceptor), Some((cert, _))) => { + let acceptor = Arc::new(acceptor); + let previous = ACCEPTOR.swap(Some(acceptor)); + + if previous.is_none() { + info!(cert = %cert.display(), "πŸ”‘ TLS enabled"); + } else { + info!(cert = %cert.display(), "πŸ” TLS certificate reloaded"); + } + } + (None, _) => { + let previous = ACCEPTOR.swap(None); + if previous.is_some() { + info!("πŸ”“ TLS disabled"); + } + } + // This state should be unreachable because `new_acceptor` is `Some` + // iff `tls_paths` is `Some`. + (Some(_), None) => { + warn!("TLS acceptor built without configuration; keeping previous value"); + } + } Ok(()) } -#[derive(Debug)] -struct AllowAllVerifier; +fn build_acceptor(cert: &Path, key: &Path) -> Result { + let pem = CertificateDer::from_pem_file(cert)?; + let key = PrivateKeyDer::from_pem_file(key)?; -impl ServerCertVerifier for AllowAllVerifier { - fn verify_server_cert( - &self, - _end_entity: &CertificateDer<'_>, - _intermediates: &[CertificateDer<'_>], - _server_name: &rustls::pki_types::ServerName<'_>, - _ocsp_response: &[u8], - _now: rustls::pki_types::UnixTime, - ) -> Result { - // Accept self-signed certs or certs signed by any CA. - // Doesn't protect against MITM attacks. - Ok(ServerCertVerified::assertion()) - } - - fn verify_tls12_signature( - &self, - _message: &[u8], - _cert: &CertificateDer<'_>, - _dss: &rustls::DigitallySignedStruct, - ) -> Result { - Ok(rustls::client::danger::HandshakeSignatureValid::assertion()) - } + let config = rustls::ServerConfig::builder() + .with_no_client_auth() + .with_single_cert(vec![pem], key)?; - fn verify_tls13_signature( - &self, - _message: &[u8], - _cert: &CertificateDer<'_>, - _dss: &rustls::DigitallySignedStruct, - ) -> Result { - Ok(rustls::client::danger::HandshakeSignatureValid::assertion()) - } + ACCEPTOR_BUILD_COUNT.fetch_add(1, Ordering::SeqCst); - fn supported_verify_schemes(&self) -> Vec { - vec![ - rustls::SignatureScheme::RSA_PKCS1_SHA1, - rustls::SignatureScheme::RSA_PKCS1_SHA256, - rustls::SignatureScheme::RSA_PKCS1_SHA384, - rustls::SignatureScheme::RSA_PKCS1_SHA512, - rustls::SignatureScheme::RSA_PSS_SHA256, - rustls::SignatureScheme::RSA_PSS_SHA384, - rustls::SignatureScheme::RSA_PSS_SHA512, - rustls::SignatureScheme::ECDSA_NISTP256_SHA256, - rustls::SignatureScheme::ECDSA_NISTP384_SHA384, - rustls::SignatureScheme::ECDSA_NISTP521_SHA512, - rustls::SignatureScheme::ED25519, - rustls::SignatureScheme::ED448, - ] - } + Ok(TlsAcceptor::from(Arc::new(config))) } -/// Create a TLS connector with the specified verification mode. -pub fn connector_with_verify_mode( - mode: TlsVerifyMode, - ca_cert_path: Option<&PathBuf>, -) -> Result { - // Load root certificates +fn build_connector(config_key: &ConnectorConfigKey) -> Result, Error> { let mut roots = rustls::RootCertStore::empty(); - // If a custom CA certificate is provided, load it - if let Some(ca_path) = ca_cert_path { + if let Some(ca_path) = config_key.ca_path.as_ref() { debug!("loading CA certificate from: {}", ca_path.display()); let certs = CertificateDer::pem_file_iter(ca_path) @@ -193,9 +187,10 @@ pub fn connector_with_verify_mode( "No valid certificates could be added from CA file", ))); } - } else if mode == TlsVerifyMode::VerifyCa || mode == TlsVerifyMode::VerifyFull { - // For Certificate and Full modes, we need CA certificates - // Load system native certificates as fallback + } else if matches!( + config_key.mode, + TlsVerifyMode::VerifyCa | TlsVerifyMode::VerifyFull + ) { debug!("no custom CA certificate provided, loading system certificates"); let result = rustls_native_certs::load_native_certs(); for cert in result.certs { @@ -210,15 +205,10 @@ pub fn connector_with_verify_mode( debug!("loaded {} system CA certificates", roots.len()); } - // Create the appropriate config based on the verification mode - let config = match mode { - TlsVerifyMode::Disabled => { - // For Disabled mode, we still create a connector but it won't be used - // The server connection logic should skip TLS entirely - ClientConfig::builder() - .with_root_certificates(roots) - .with_no_client_auth() - } + let config = match config_key.mode { + TlsVerifyMode::Disabled => ClientConfig::builder() + .with_root_certificates(roots) + .with_no_client_auth(), TlsVerifyMode::Prefer => { let verifier = AllowAllVerifier; ClientConfig::builder() @@ -243,7 +233,108 @@ pub fn connector_with_verify_mode( .with_no_client_auth(), }; - Ok(TlsConnector::from(Arc::new(config))) + increment_connector_build_count(); + + Ok(Arc::new(config)) +} + +#[cfg_attr(not(test), allow(dead_code))] +#[doc(hidden)] +pub fn test_acceptor_build_count() -> usize { + ACCEPTOR_BUILD_COUNT.load(Ordering::SeqCst) +} + +#[cfg_attr(not(test), allow(dead_code))] +#[doc(hidden)] +pub fn test_reset_acceptor() { + ACCEPTOR.store(None); + ACCEPTOR_BUILD_COUNT.store(0, Ordering::SeqCst); +} + +#[cfg(test)] +#[doc(hidden)] +pub fn test_connector_build_count() -> usize { + CONNECTOR_BUILD_COUNT.load(Ordering::SeqCst) +} + +#[cfg(test)] +#[doc(hidden)] +pub fn test_reset_connector() { + CONNECTOR.store(None); + CONNECTOR_BUILD_COUNT.store(0, Ordering::SeqCst); +} + +#[derive(Debug)] +struct AllowAllVerifier; + +impl ServerCertVerifier for AllowAllVerifier { + fn verify_server_cert( + &self, + _end_entity: &CertificateDer<'_>, + _intermediates: &[CertificateDer<'_>], + _server_name: &rustls::pki_types::ServerName<'_>, + _ocsp_response: &[u8], + _now: rustls::pki_types::UnixTime, + ) -> Result { + // Accept self-signed certs or certs signed by any CA. + // Doesn't protect against MITM attacks. + Ok(ServerCertVerified::assertion()) + } + + fn verify_tls12_signature( + &self, + _message: &[u8], + _cert: &CertificateDer<'_>, + _dss: &rustls::DigitallySignedStruct, + ) -> Result { + Ok(rustls::client::danger::HandshakeSignatureValid::assertion()) + } + + fn verify_tls13_signature( + &self, + _message: &[u8], + _cert: &CertificateDer<'_>, + _dss: &rustls::DigitallySignedStruct, + ) -> Result { + Ok(rustls::client::danger::HandshakeSignatureValid::assertion()) + } + + fn supported_verify_schemes(&self) -> Vec { + vec![ + rustls::SignatureScheme::RSA_PKCS1_SHA1, + rustls::SignatureScheme::RSA_PKCS1_SHA256, + rustls::SignatureScheme::RSA_PKCS1_SHA384, + rustls::SignatureScheme::RSA_PKCS1_SHA512, + rustls::SignatureScheme::RSA_PSS_SHA256, + rustls::SignatureScheme::RSA_PSS_SHA384, + rustls::SignatureScheme::RSA_PSS_SHA512, + rustls::SignatureScheme::ECDSA_NISTP256_SHA256, + rustls::SignatureScheme::ECDSA_NISTP384_SHA384, + rustls::SignatureScheme::ECDSA_NISTP521_SHA512, + rustls::SignatureScheme::ED25519, + rustls::SignatureScheme::ED448, + ] + } +} + +/// Create a TLS connector with the specified verification mode. +pub fn connector_with_verify_mode( + mode: TlsVerifyMode, + ca_cert_path: Option<&PathBuf>, +) -> Result { + let config_key = ConnectorConfigKey::new(mode, ca_cert_path); + + if let Some(entry) = CONNECTOR.load_full() { + if entry.key == config_key { + return Ok(entry.connector()); + } + } + + let client_config = build_connector(&config_key)?; + let connector = TlsConnector::from(client_config.clone()); + CONNECTOR.store(Some(ConnectorCacheEntry::new(config_key, client_config))); + + Ok(connector) } /// Certificate verifier that validates certificates but skips hostname verification @@ -332,6 +423,41 @@ impl ServerCertVerifier for NoHostnameVerifier { mod tests { use super::*; use crate::config::TlsVerifyMode; + use std::sync::Arc; + + #[test] + fn acceptor_reuse_snapshot() { + crate::logger(); + + super::test_reset_acceptor(); + + let cert = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tests/tls/cert.pem"); + let key = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tests/tls/key.pem"); + + let mut cfg = crate::config::ConfigAndUsers::default(); + cfg.config.general.tls_certificate = Some(cert.clone()); + cfg.config.general.tls_private_key = Some(key.clone()); + + crate::config::set(cfg).unwrap(); + + super::reload().unwrap(); + + assert_eq!(super::test_acceptor_build_count(), 1, "acceptor built once"); + + let first = super::acceptor().expect("acceptor initialized"); + let second = super::acceptor().expect("acceptor initialized"); + + assert!(Arc::ptr_eq(&first, &second), "cached acceptor reused"); + assert_eq!( + super::test_acceptor_build_count(), + 1, + "no additional builds" + ); + + super::test_reset_acceptor(); + + crate::config::set(crate::config::ConfigAndUsers::default()).unwrap(); + } #[tokio::test] async fn test_connector_with_verify_mode() { @@ -347,6 +473,92 @@ mod tests { assert!(full.is_ok()); } + #[tokio::test] + async fn connector_reuses_cached_config() { + crate::logger(); + + super::test_reset_acceptor(); + super::test_reset_connector(); + + let cert_path = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tests/tls/cert.pem"); + let key_path = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tests/tls/key.pem"); + let ca_path = cert_path.clone(); + + let mut cfg = crate::config::ConfigAndUsers::default(); + cfg.config.general.tls_certificate = Some(cert_path.clone()); + cfg.config.general.tls_private_key = Some(key_path.clone()); + cfg.config.general.tls_server_ca_certificate = Some(ca_path.clone()); + cfg.config.general.tls_verify = TlsVerifyMode::VerifyFull; + + crate::config::set(cfg).unwrap(); + + let _first = super::connector_with_verify_mode(TlsVerifyMode::VerifyFull, Some(&ca_path)) + .expect("first connector builds"); + let first_cache = super::CONNECTOR + .load_full() + .expect("connector cached after first build"); + + let _second = super::connector_with_verify_mode(TlsVerifyMode::VerifyFull, Some(&ca_path)) + .expect("second connector reuses cache"); + let second_cache = super::CONNECTOR + .load_full() + .expect("connector cached after second build"); + + assert_eq!( + super::test_connector_build_count(), + 1, + "connector built once" + ); + assert!( + Arc::ptr_eq(&first_cache, &second_cache), + "cache entry reused" + ); + assert!( + Arc::ptr_eq(&first_cache.config, &second_cache.config), + "client config reused" + ); + + super::reload().expect("reload succeeds"); + + let post_reload_cache = super::CONNECTOR + .load_full() + .expect("connector cached after reload"); + + assert_eq!( + super::test_connector_build_count(), + 1, + "reload does not rebuild connector" + ); + assert!( + Arc::ptr_eq(&second_cache, &post_reload_cache), + "reload retains cache entry" + ); + assert!( + Arc::ptr_eq(&second_cache.config, &post_reload_cache.config), + "reload retains client config" + ); + + let _third = super::connector_with_verify_mode(TlsVerifyMode::VerifyFull, Some(&ca_path)) + .expect("third connector still reuses cache"); + let third_cache = super::CONNECTOR + .load_full() + .expect("connector cached after third build"); + + assert_eq!( + super::test_connector_build_count(), + 1, + "additional calls reuse existing connector" + ); + assert!( + Arc::ptr_eq(&post_reload_cache, &third_cache), + "cache entry unchanged" + ); + + super::test_reset_connector(); + super::test_reset_acceptor(); + crate::config::set(crate::config::ConfigAndUsers::default()).unwrap(); + } + #[tokio::test] async fn test_connector_with_verify_mode_missing_ca_file() { crate::logger(); diff --git a/pgdog/src/stats/clients.rs b/pgdog/src/stats/clients.rs index 45745002c..4372fb70e 100644 --- a/pgdog/src/stats/clients.rs +++ b/pgdog/src/stats/clients.rs @@ -34,7 +34,10 @@ impl OpenMetric for Clients { #[cfg(test)] mod test { - use crate::stats::Metric; + use crate::{ + config::{self, ConfigAndUsers}, + stats::Metric, + }; use super::*; @@ -51,4 +54,19 @@ mod test { ); assert_eq!(lines.next().unwrap(), "clients 25"); } + + #[test] + fn clients_load_uses_global_state() { + config::set(ConfigAndUsers::default()).unwrap(); + + let metric = Clients::load(); + let rendered = metric.to_string(); + let lines: Vec<&str> = rendered.lines().collect(); + assert_eq!(lines[0], "# TYPE clients gauge"); + assert_eq!( + lines[1], + "# HELP clients Total number of connected clients." + ); + assert_eq!(lines.last().copied(), Some("clients 0")); + } } diff --git a/pgdog/src/stats/open_metric.rs b/pgdog/src/stats/open_metric.rs index 6b24d706a..a97bf64a9 100644 --- a/pgdog/src/stats/open_metric.rs +++ b/pgdog/src/stats/open_metric.rs @@ -162,4 +162,29 @@ mod test { assert_eq!(render.lines().next().unwrap(), "# TYPE pgdog.test gauge"); assert_eq!(render.lines().last().unwrap(), "pgdog.test 5"); } + + #[test] + fn measurement_render_formats_labels() { + let measurement = Measurement { + labels: vec![ + ("role".into(), "primary".into()), + ("shard".into(), "0".into()), + ], + measurement: MeasurementType::Integer(42), + }; + + let rendered = measurement.render("pool_clients"); + assert_eq!(rendered, "pool_clients{role=\"primary\",shard=\"0\"} 42"); + } + + #[test] + fn measurement_render_rounds_floats() { + let measurement = Measurement { + labels: vec![], + measurement: MeasurementType::Float(1.23456), + }; + + let rendered = measurement.render("query_latency_seconds"); + assert_eq!(rendered, "query_latency_seconds 1.235"); + } } diff --git a/pgdog/src/stats/pools.rs b/pgdog/src/stats/pools.rs index 85e9fd879..39fd6240d 100644 --- a/pgdog/src/stats/pools.rs +++ b/pgdog/src/stats/pools.rs @@ -1,4 +1,4 @@ -use crate::backend::databases::databases; +use crate::backend::{self, databases::databases}; use super::{Measurement, Metric, OpenMetric}; @@ -61,10 +61,24 @@ impl Pools { let mut avg_received = vec![]; let mut total_xact_time = vec![]; let mut avg_xact_time = vec![]; + let mut total_idle_xact_time = vec![]; + let mut avg_idle_xact_time = vec![]; let mut total_query_time = vec![]; let mut avg_query_time = vec![]; let mut total_close = vec![]; let mut avg_close = vec![]; + let mut total_server_errors = vec![]; + let mut avg_server_errors = vec![]; + let mut total_cleaned = vec![]; + let mut avg_cleaned = vec![]; + let mut total_rollbacks = vec![]; + let mut avg_rollbacks = vec![]; + let mut total_connect_time = vec![]; + let mut avg_connect_time = vec![]; + let mut total_connect_count = vec![]; + let mut avg_connect_count = vec![]; + let mut total_sv_xact_idle = vec![]; + for (user, cluster) in databases().all() { for (shard_num, shard) in cluster.shards().iter().enumerate() { for (role, pool) in shard.pools_with_roles() { @@ -172,6 +186,16 @@ impl Pools { measurement: averages.xact_time.as_millis().into(), }); + total_idle_xact_time.push(Measurement { + labels: labels.clone(), + measurement: totals.idle_xact_time.as_millis().into(), + }); + + avg_idle_xact_time.push(Measurement { + labels: labels.clone(), + measurement: averages.idle_xact_time.as_millis().into(), + }); + total_query_time.push(Measurement { labels: labels.clone(), measurement: totals.query_time.as_millis().into(), @@ -191,6 +215,61 @@ impl Pools { labels: labels.clone(), measurement: averages.close.into(), }); + + total_server_errors.push(Measurement { + labels: labels.clone(), + measurement: totals.errors.into(), + }); + + avg_server_errors.push(Measurement { + labels: labels.clone(), + measurement: averages.errors.into(), + }); + + total_cleaned.push(Measurement { + labels: labels.clone(), + measurement: totals.cleaned.into(), + }); + + avg_cleaned.push(Measurement { + labels: labels.clone(), + measurement: averages.cleaned.into(), + }); + + total_rollbacks.push(Measurement { + labels: labels.clone(), + measurement: totals.rollbacks.into(), + }); + + avg_rollbacks.push(Measurement { + labels: labels.clone(), + measurement: averages.rollbacks.into(), + }); + + total_connect_time.push(Measurement { + labels: labels.clone(), + measurement: totals.connect_time.as_millis().into(), + }); + + avg_connect_time.push(Measurement { + labels: labels.clone(), + measurement: averages.connect_time.as_millis().into(), + }); + + total_connect_count.push(Measurement { + labels: labels.clone(), + measurement: totals.connect_count.into(), + }); + + avg_connect_count.push(Measurement { + labels: labels.clone(), + measurement: averages.connect_count.into(), + }); + + total_sv_xact_idle.push(Measurement { + labels: labels.clone(), + measurement: backend::stats::idle_in_transaction(&pool).into(), + }); } } } @@ -340,6 +419,22 @@ impl Pools { metric_type: None, })); + metrics.push(Metric::new(PoolMetric { + name: "total_idle_xact_time".into(), + measurements: total_idle_xact_time, + help: "Total time spent idling inside transactions.".into(), + unit: None, + metric_type: Some("counter".into()), + })); + + metrics.push(Metric::new(PoolMetric { + name: "avg_idle_xact_time".into(), + measurements: avg_idle_xact_time, + help: "Average time spent idling inside transactions.".into(), + unit: None, + metric_type: None, + })); + metrics.push(Metric::new(PoolMetric { name: "total_query_time".into(), measurements: total_query_time, @@ -372,10 +467,155 @@ impl Pools { metric_type: None, })); + metrics.push(Metric::new(PoolMetric { + name: "total_server_errors".into(), + measurements: total_server_errors, + help: "Total number of errors returned by server connections.".into(), + unit: None, + metric_type: Some("counter".into()), + })); + + metrics.push(Metric::new(PoolMetric { + name: "avg_server_errors".into(), + measurements: avg_server_errors, + help: "Average number of errors returned by server connections.".into(), + unit: None, + metric_type: None, + })); + + metrics.push(Metric::new(PoolMetric { + name: "total_cleaned".into(), + measurements: total_cleaned, + help: "Total number of times server connections were cleaned from client parameters." + .into(), + unit: None, + metric_type: Some("counter".into()), + })); + + metrics.push(Metric::new(PoolMetric { + name: "avg_cleaned".into(), + measurements: avg_cleaned, + help: "Average number of times server connections were cleaned from client parameters." + .into(), + unit: None, + metric_type: None, + })); + + metrics.push(Metric::new(PoolMetric { + name: "total_rollbacks".into(), + measurements: total_rollbacks, + help: + "Total number of abandoned transactions that had to be rolled back automatically." + .into(), + unit: None, + metric_type: Some("counter".into()), + })); + + metrics.push(Metric::new(PoolMetric { + name: "avg_rollbacks".into(), + measurements: avg_rollbacks, + help: + "Average number of abandoned transactions that had to be rolled back automatically." + .into(), + unit: None, + metric_type: None, + })); + + metrics.push(Metric::new(PoolMetric { + name: "total_connect_time".into(), + measurements: total_connect_time, + help: "Total time spent connecting to servers.".into(), + unit: None, + metric_type: Some("counter".into()), + })); + + metrics.push(Metric::new(PoolMetric { + name: "avg_connect_time".into(), + measurements: avg_connect_time, + help: "Average time spent connecting to servers.".into(), + unit: None, + metric_type: None, + })); + + metrics.push(Metric::new(PoolMetric { + name: "total_connect_count".into(), + measurements: total_connect_count, + help: "Total number of connections established to servers.".into(), + unit: None, + metric_type: Some("counter".into()), + })); + + metrics.push(Metric::new(PoolMetric { + name: "avg_connect_count".into(), + measurements: avg_connect_count, + help: "Average number of connections established to servers.".into(), + unit: None, + metric_type: None, + })); + + metrics.push(Metric::new(PoolMetric { + name: "sv_idle_xact".into(), + measurements: total_sv_xact_idle, + help: "Servers currently idle in transaction.".into(), + unit: None, + metric_type: None, + })); + Pools { metrics } } } +#[cfg(test)] +mod tests { + use crate::config::{self, ConfigAndUsers}; + + use super::*; + + #[test] + fn pool_metric_defaults_to_gauge() { + let metric = PoolMetric { + name: "cl_waiting".into(), + measurements: vec![Measurement { + labels: vec![], + measurement: 3usize.into(), + }], + help: "Waiting clients per pool".into(), + unit: None, + metric_type: None, + }; + + assert_eq!(metric.metric_type(), "gauge"); + assert!(metric.unit().is_none()); + assert_eq!(metric.help(), Some("Waiting clients per pool".into())); + } + + #[test] + fn pool_metric_renders_labels_and_unit() { + config::set(ConfigAndUsers::default()).unwrap(); + + let metric = PoolMetric { + name: "sv_active".into(), + measurements: vec![Measurement { + labels: vec![ + ("user".into(), "alice".into()), + ("database".into(), "app".into()), + ], + measurement: 5usize.into(), + }], + help: "Active servers per pool".into(), + unit: Some("connections".into()), + metric_type: Some("gauge".into()), + }; + + let rendered = Metric::new(metric).to_string(); + let lines: Vec<&str> = rendered.lines().collect(); + assert_eq!(lines[0], "# TYPE sv_active gauge"); + assert_eq!(lines[1], "# UNIT sv_active connections"); + assert_eq!(lines[2], "# HELP sv_active Active servers per pool"); + assert_eq!(lines[3], "sv_active{user=\"alice\",database=\"app\"} 5"); + } +} + impl std::fmt::Display for Pools { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { for pool in &self.metrics { diff --git a/pgdog/src/stats/query_cache.rs b/pgdog/src/stats/query_cache.rs index 85692bd86..7939f8afb 100644 --- a/pgdog/src/stats/query_cache.rs +++ b/pgdog/src/stats/query_cache.rs @@ -112,3 +112,84 @@ impl OpenMetric for QueryCacheMetric { }] } } + +#[cfg(test)] +mod tests { + use crate::config::{self, ConfigAndUsers}; + + use super::*; + + #[test] + fn query_cache_metric_renders_counter_and_gauge() { + config::set(ConfigAndUsers::default()).unwrap(); + + let counter_metric = QueryCacheMetric { + name: "query_cache_hits".into(), + help: "Hits".into(), + value: 7, + gauge: false, + }; + let counter_render = Metric::new(counter_metric).to_string(); + let counter_lines: Vec<&str> = counter_render.lines().collect(); + assert_eq!(counter_lines[0], "# TYPE query_cache_hits counter"); + assert_eq!(counter_lines[2], "query_cache_hits 7"); + + let gauge_metric = QueryCacheMetric { + name: "query_cache_size".into(), + help: "Size".into(), + value: 3, + gauge: true, + }; + let gauge_render = Metric::new(gauge_metric).to_string(); + let gauge_lines: Vec<&str> = gauge_render.lines().collect(); + assert_eq!(gauge_lines[0], "# TYPE query_cache_size gauge"); + assert_eq!(gauge_lines[2], "query_cache_size 3"); + } + + #[test] + fn query_cache_metrics_expose_all_counters() { + let cache = QueryCache { + stats: Stats { + hits: 1, + misses: 2, + direct: 3, + multi: 4, + }, + len: 5, + prepared_statements: 6, + prepared_statements_memory: 7, + }; + + let metrics = cache.metrics(); + let metric_names: Vec = metrics.iter().map(|metric| metric.name()).collect(); + assert_eq!( + metric_names, + vec![ + "query_cache_hits".to_string(), + "query_cache_misses".to_string(), + "query_cache_direct".to_string(), + "query_cache_cross".to_string(), + "query_cache_size".to_string(), + "prepared_statements".to_string(), + "prepared_statements_memory_used".to_string(), + ] + ); + + let hits_metric = &metrics[0]; + let hits_value = hits_metric + .measurements() + .first() + .unwrap() + .measurement + .clone(); + match hits_value { + MeasurementType::Integer(value) => assert_eq!(value, 1), + other => panic!("expected integer measurement, got {:?}", other), + } + + let memory_metric = metrics.last().unwrap(); + assert_eq!(memory_metric.metric_type(), "gauge"); + let rendered = memory_metric.to_string(); + assert!(rendered.contains("prepared_statements_memory_used 7")); + } +} diff --git a/pgdog/src/unique_id.rs b/pgdog/src/unique_id.rs new file mode 100644 index 000000000..47d11a71e --- /dev/null +++ b/pgdog/src/unique_id.rs @@ -0,0 +1,278 @@ +//! Globally unique 64-bit (i64) ID generator. +//! +//! Relies on 2 invariants: +//! +//! 1. Each instance of PgDog must have a unique, numeric NODE_ID, +//! not exceeding 1023. +//! 2. Each instance of PgDog has reasonably accurate and synchronized +//! clock, so `std::time::SystemTime` returns a good value. +//! +use std::thread; +use std::time::UNIX_EPOCH; +use std::time::{Duration, SystemTime}; + +use once_cell::sync::OnceCell; +use parking_lot::Mutex; +use thiserror::Error; + +use crate::config::config; +use crate::util::{instance_id, node_id}; + +const NODE_BITS: u64 = 10; // Max 1023 nodes +const SEQUENCE_BITS: u64 = 12; +const TIMESTAMP_BITS: u64 = 41; // 41 bits = ~69 years, keeps i64 sign bit clear +const MAX_NODE_ID: u64 = (1 << NODE_BITS) - 1; // 1023 +const MAX_SEQUENCE: u64 = (1 << SEQUENCE_BITS) - 1; // 4095 +const MAX_TIMESTAMP: u64 = (1 << TIMESTAMP_BITS) - 1; +const PGDOG_EPOCH: u64 = 1764184395000; // Wednesday, November 26, 2025 11:13:15 AM GMT-08:00 +const NODE_SHIFT: u8 = SEQUENCE_BITS as u8; // 12 +const TIMESTAMP_SHIFT: u8 = (SEQUENCE_BITS + NODE_BITS) as u8; // 22 + // Maximum offset to ensure base_id + offset doesn't overflow i64 +const MAX_OFFSET: u64 = i64::MAX as u64 + - ((MAX_TIMESTAMP << TIMESTAMP_SHIFT) | (MAX_NODE_ID << NODE_SHIFT) | MAX_SEQUENCE); + +static UNIQUE_ID: OnceCell = OnceCell::new(); + +#[derive(Debug, Default)] +struct State { + last_timestamp_ms: u64, + sequence: u64, +} + +impl State { + // Generate next unique ID in a distributed sequence. + // The `node_id` argument must be globally unique. + fn next_id(&mut self, node_id: u64, id_offset: u64) -> u64 { + let mut now = wait_until(self.last_timestamp_ms); + + if now == self.last_timestamp_ms { + self.sequence = (self.sequence + 1) & MAX_SEQUENCE; + // Wraparound. + if self.sequence == 0 { + now = wait_until(now + 1); + } + } else { + // Reset sequence to zero once we reach next ms. + self.sequence = 0; + } + + self.last_timestamp_ms = now; + + let elapsed = self.last_timestamp_ms - PGDOG_EPOCH; + assert!( + elapsed <= MAX_TIMESTAMP, + "unique_id timestamp overflow: {elapsed} > {MAX_TIMESTAMP}" + ); + let timestamp_part = (elapsed & MAX_TIMESTAMP) << TIMESTAMP_SHIFT; + let node_part = node_id << NODE_SHIFT; + let sequence_part = self.sequence; + + let base_id = timestamp_part | node_part | sequence_part; + base_id + id_offset + } +} + +// Get current time in ms. +fn now_ms() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .expect("SystemTime is before UNIX_EPOCH") + .as_millis() as u64 +} + +// Get a monotonically increasing timestamp in ms. +// Protects against clock drift. +fn wait_until(target_ms: u64) -> u64 { + loop { + let now = now_ms(); + if now >= target_ms { + return now; + } + thread::sleep(Duration::from_millis(1)); + } +} + +#[derive(Debug, Error)] +pub enum Error { + #[error("invalid node identifier: {0}")] + InvalidNodeId(String), + + #[error("node ID exceeding maximum (1023): {0}")] + NodeIdTooLarge(u64), + + #[error("id_offset too large, would overflow i64: {0}")] + OffsetTooLarge(u64), +} + +#[derive(Debug)] +pub struct UniqueId { + node_id: u64, + id_offset: u64, + inner: Mutex, +} + +impl UniqueId { + /// Initialize the UniqueId generator. + fn new() -> Result { + let node_id = node_id().map_err(|_| Error::InvalidNodeId(instance_id().to_string()))?; + + let min_id = config().config.general.unique_id_min; + + if node_id > MAX_NODE_ID { + return Err(Error::NodeIdTooLarge(node_id)); + } + + if min_id > MAX_OFFSET { + return Err(Error::OffsetTooLarge(min_id)); + } + + Ok(Self { + node_id, + id_offset: min_id, + inner: Mutex::new(State::default()), + }) + } + + /// Get (and initialize, if necessary) the unique ID generator. + pub fn generator() -> Result<&'static UniqueId, Error> { + UNIQUE_ID.get_or_try_init(Self::new) + } + + /// Generate a globally unique, monotonically increasing identifier. + pub fn next_id(&self) -> i64 { + self.inner.lock().next_id(self.node_id, self.id_offset) as i64 + } +} + +#[cfg(test)] +mod test { + use std::{collections::HashSet, env::set_var}; + + use super::*; + + #[test] + fn test_unique_ids() { + unsafe { + set_var("NODE_ID", "pgdog-1"); + } + let num_ids = 10_000; + + let mut ids = HashSet::new(); + + for _ in 0..num_ids { + ids.insert(UniqueId::generator().unwrap().next_id()); + } + + assert_eq!(ids.len(), num_ids); + } + + #[test] + fn test_ids_monotonically_increasing() { + let mut state = State::default(); + let node_id = 1u64; + + let mut prev_id = 0u64; + for _ in 0..10_000 { + let id = state.next_id(node_id, 0); + assert!(id > prev_id, "ID {id} not greater than previous {prev_id}"); + prev_id = id; + } + } + + #[test] + fn test_ids_always_positive() { + let mut state = State::default(); + let node_id = MAX_NODE_ID; // Use max node to maximize bits used + + for _ in 0..10_000 { + let id = state.next_id(node_id, 0); + let signed = id as i64; + assert!(signed > 0, "ID should be positive, got {signed}"); + } + } + + #[test] + fn test_bit_layout() { + // Verify the bit allocation: 41 timestamp + 10 node + 12 sequence = 63 bits + assert_eq!(TIMESTAMP_BITS + NODE_BITS + SEQUENCE_BITS, 63); + assert_eq!(TIMESTAMP_SHIFT, 22); + assert_eq!(NODE_SHIFT, 12); + } + + #[test] + fn test_max_values_fit() { + // Construct an ID with max values and verify it stays positive + let max_elapsed = MAX_TIMESTAMP; + let max_node = MAX_NODE_ID; + let max_seq = MAX_SEQUENCE; + + let id = (max_elapsed << TIMESTAMP_SHIFT) | (max_node << NODE_SHIFT) | max_seq; + let signed = id as i64; + + assert!(signed > 0, "Max ID should be positive, got {signed}"); + assert_eq!(id >> 63, 0, "Bit 63 should be clear"); + } + + #[test] + fn test_extract_components() { + let node: u64 = 42; + let mut state = State::default(); + + let id = state.next_id(node, 0); + + // Extract components back + let extracted_seq = id & MAX_SEQUENCE; + let extracted_node = (id >> NODE_SHIFT) & MAX_NODE_ID; + let extracted_elapsed = id >> TIMESTAMP_SHIFT; + + assert_eq!(extracted_node, node); + assert_eq!(extracted_seq, 0); // First ID has sequence 0 + assert!(extracted_elapsed > 0); // Elapsed time since epoch + + // Generate another ID and verify sequence increments + let id2 = state.next_id(node, 0); + let extracted_seq2 = id2 & MAX_SEQUENCE; + let extracted_node2 = (id2 >> NODE_SHIFT) & MAX_NODE_ID; + + assert_eq!(extracted_node2, node); + assert!(matches!(extracted_seq2, 1 | 0)); // Sequence incremented (or time advanced and reset to 0) + } + + #[test] + fn test_id_offset() { + let offset: u64 = 1_000_000_000; + let node: u64 = 5; + let mut state = State::default(); + + for _ in 0..1000 { + let id = state.next_id(node, offset); + assert!( + id > offset, + "ID {id} should be greater than offset {offset}" + ); + } + } + + #[test] + fn test_id_offset_monotonic() { + let offset: u64 = 1_000_000_000; + let node: u64 = 5; + let mut state = State::default(); + + let mut prev_id = 0u64; + for _ in 0..1000 { + let id = state.next_id(node, offset); + assert!(id > prev_id, "ID {id} not greater than previous {prev_id}"); + prev_id = id; + } + } + + #[test] + fn test_max_offset() { + // Verify MAX_OFFSET calculation is correct + let max_base_id = + (MAX_TIMESTAMP << TIMESTAMP_SHIFT) | (MAX_NODE_ID << NODE_SHIFT) | MAX_SEQUENCE; + let result = max_base_id + MAX_OFFSET; + assert!(result <= i64::MAX as u64, "MAX_OFFSET would overflow i64"); + } +} diff --git a/pgdog/src/util.rs b/pgdog/src/util.rs index 29702e4b3..61407976e 100644 --- a/pgdog/src/util.rs +++ b/pgdog/src/util.rs @@ -3,8 +3,8 @@ use chrono::{DateTime, Local, Utc}; use once_cell::sync::Lazy; use pgdog_plugin::comp; -use rand::{distributions::Alphanumeric, Rng}; -use std::{ops::Deref, time::Duration}; +use rand::{distr::Alphanumeric, Rng}; +use std::{env, num::ParseIntError, ops::Deref, time::Duration}; use crate::net::Parameters; // 0.8 @@ -32,7 +32,7 @@ pub fn human_duration(duration: Duration) -> String { let ms = duration.as_millis(); let ms_fmt = |ms: u128, unit: u128, name: &str| -> String { - if ms % unit > 0 { + if !ms.is_multiple_of(unit) { format!("{}ms", ms) } else { format!("{}{}", ms / unit, name) @@ -67,7 +67,7 @@ pub fn postgres_now() -> i64 { /// Generate a random string of length n. pub fn random_string(n: usize) -> String { - rand::thread_rng() + rand::rng() .sample_iter(&Alphanumeric) .take(n) .map(char::from) @@ -76,13 +76,17 @@ pub fn random_string(n: usize) -> String { // Generate a unique 8-character hex instance ID on first access static INSTANCE_ID: Lazy = Lazy::new(|| { - let mut rng = rand::thread_rng(); - (0..8) - .map(|_| { - let n: u8 = rng.gen_range(0..16); - format!("{:x}", n) - }) - .collect() + if let Ok(node_id) = env::var("NODE_ID") { + node_id + } else { + let mut rng = rand::rng(); + (0..8) + .map(|_| { + let n: u8 = rng.random_range(0..16); + format!("{:x}", n) + }) + .collect() + } }); /// Get the instance ID for this pgdog instance. @@ -91,6 +95,13 @@ pub fn instance_id() -> &'static str { &INSTANCE_ID } +/// Get an externally assigned, unique, node identifier +/// for this instance of PgDog. +pub fn node_id() -> Result { + // split always returns at least one element. + instance_id().split("-").last().unwrap().parse() +} + /// Escape PostgreSQL identifiers by doubling any embedded quotes. pub fn escape_identifier(s: &str) -> String { s.replace("\"", "\"\"") @@ -117,6 +128,8 @@ pub fn user_database_from_params(params: &Parameters) -> (&str, &str) { #[cfg(test)] mod test { + use std::env::{remove_var, set_var}; + use super::*; #[test] @@ -153,6 +166,9 @@ mod test { #[test] fn test_instance_id_format() { + unsafe { + remove_var("NODE_ID"); + } let id = instance_id(); assert_eq!(id.len(), 8); // All characters should be valid hex digits (0-9, a-f) @@ -170,4 +186,21 @@ mod test { let id2 = instance_id(); assert_eq!(id1, id2); // Should be the same for lifetime of process } + + #[test] + fn test_node_id_error() { + unsafe { + remove_var("NODE_ID"); + } + assert!(node_id().is_err()); + } + + // These should run in separate processes (if using nextest). + #[test] + fn test_node_id_set() { + unsafe { + set_var("NODE_ID", "pgdog-1"); + } + assert_eq!(node_id(), Ok(1)); + } } diff --git a/pgdog/tests/flame/base/pgdog.toml b/pgdog/tests/flame/base/pgdog.toml new file mode 100644 index 000000000..48a87f5d3 --- /dev/null +++ b/pgdog/tests/flame/base/pgdog.toml @@ -0,0 +1,6 @@ +[[databases]] +name = "pgdog" +host = "127.0.0.1" + +[admin] +password = "pgdog" diff --git a/pgdog/tests/flame/base/users.toml b/pgdog/tests/flame/base/users.toml new file mode 100644 index 000000000..539bb1832 --- /dev/null +++ b/pgdog/tests/flame/base/users.toml @@ -0,0 +1,4 @@ +[[users]] +name = "pgdog" +password = "pgdog" +database = "pgdog" diff --git a/pgdog/tests/flame/profile.sh b/pgdog/tests/flame/profile.sh new file mode 100644 index 000000000..9f09ba483 --- /dev/null +++ b/pgdog/tests/flame/profile.sh @@ -0,0 +1,69 @@ +#!/bin/bash +# +# Utils for performance testing. +# +set -e +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd ) +export PGDOG_BIN=${SCRIPT_DIR}/../../../target/release/pgdog +export PGPASSWORD=pgdog +export PGUSER=pgdog +export PGDATABASE=pgdog +export PGHOST=127.0.0.1 +export PGPORT=6432 + +if [ ! -f ${PGDOG_BIN} ]; then + echo "PgDog is not compiled in release mode (target/release/pgdog is missing)" + echo "Please compile PgDog with:" + echo + printf "\tcargo build --release" + echo + echo + exit 1 +fi + +if ! which samply > /dev/null; then + echo "Samply profiler is not installed on this system" + echo "Please install Samply with:" + echo + printf "\tcargo install samply" + echo + echo + exit 1 +fi + +paranoid=$(cat /proc/sys/kernel/perf_event_paranoid) + +if [ ! "$paranoid" -eq "-1" ]; then + echo "\"/proc/sys/kernel/perf_event_paranoid\" need to be set to -1 for sampling profiler to work" + echo "Please set it manually by running:" + echo + printf "\techo '-1' | sudo tee /proc/sys/kernel/perf_event_paranoid" + echo + echo + exit 1 +fi + +function profile_pgdog() { + pushd $1 + samply record ${PGDOG_BIN} & + pid=$! + + while ! pg_isready > /dev/null; do + sleep 1 + done + + pgbench -i + pgbench -c 10 -t 1000000 -S -P 1 --protocol extended + + kill -TERM $pid + popd +} + +if [ ! -d "$1" ]; then + echo "Benchmark $1 doesn't exist." + echo "Available benchmarks:" + echo "$(ls -d */)" + exit 1 +fi + +profile_pgdog $1 diff --git a/pgdog/tests/flame/read_write/pgdog.toml b/pgdog/tests/flame/read_write/pgdog.toml new file mode 100644 index 000000000..b859628b0 --- /dev/null +++ b/pgdog/tests/flame/read_write/pgdog.toml @@ -0,0 +1,11 @@ +[[databases]] +name = "pgdog" +host = "127.0.0.1" + +[[databases]] +name = "pgdog" +host = "127.0.0.1" +role = "replica" + +[admin] +password = "pgdog" diff --git a/pgdog/tests/flame/read_write/users.toml b/pgdog/tests/flame/read_write/users.toml new file mode 100644 index 000000000..539bb1832 --- /dev/null +++ b/pgdog/tests/flame/read_write/users.toml @@ -0,0 +1,4 @@ +[[users]] +name = "pgdog" +password = "pgdog" +database = "pgdog" diff --git a/pgdog/tests/flame/sharded/pgdog.toml b/pgdog/tests/flame/sharded/pgdog.toml new file mode 100644 index 000000000..e69de29bb diff --git a/pgdog/tests/flame/sharded/users.toml b/pgdog/tests/flame/sharded/users.toml new file mode 100644 index 000000000..e69de29bb diff --git a/pgdog/tests/pgbouncer/pgdog.toml b/pgdog/tests/pgbouncer/pgdog.toml index 572de4642..04ef52419 100644 --- a/pgdog/tests/pgbouncer/pgdog.toml +++ b/pgdog/tests/pgbouncer/pgdog.toml @@ -7,4 +7,4 @@ name = "pgdog" host = "127.0.0.1" [admin] -password = "admin" +password = "pgdog" diff --git a/pgdog/tests/pgbouncer/run.sh b/pgdog/tests/pgbouncer/run.sh deleted file mode 100644 index 5bd90577d..000000000 --- a/pgdog/tests/pgbouncer/run.sh +++ /dev/null @@ -1,11 +0,0 @@ -#!/bin/bash -SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd ) -pushd ${SCRIPT_DIR}/../../../ - -if [[ "$1" == "sample" ]]; then - samply record target/debug/pgdog --config ${SCRIPT_DIR}/pgdog.toml --users ${SCRIPT_DIR}/users.toml -elif [[ "$1" == "dev" ]]; then - cargo run -- --config ${SCRIPT_DIR}/pgdog.toml --users ${SCRIPT_DIR}/users.toml -else - cargo run --release -- --config ${SCRIPT_DIR}/pgdog.toml --users ${SCRIPT_DIR}/users.toml -fi diff --git a/pgdog/tests/psql.sh b/pgdog/tests/psql.sh index d50bdcd0d..9652a6ce1 100644 --- a/pgdog/tests/psql.sh +++ b/pgdog/tests/psql.sh @@ -1,6 +1,6 @@ #!/bin/bash if [[ ! -z "$1" ]]; then - PGPASSWORD=pgdog psql -h 127.0.0.1 -p 6432 -U pgdog pgdog_sharded + PGPASSWORD=pgdog psql -h 127.0.0.1 -p 6432 -U pgdog $1 else PGPASSWORD=pgdog psql -h 127.0.0.1 -p 6432 -U pgdog pgdog fi diff --git a/pgdog/tests/tls_cache.rs b/pgdog/tests/tls_cache.rs new file mode 100644 index 000000000..a886b6fd1 --- /dev/null +++ b/pgdog/tests/tls_cache.rs @@ -0,0 +1,45 @@ +use std::{ + path::PathBuf, + sync::Arc, + time::{Duration, Instant}, +}; + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn tls_acceptor_reuse_is_fast() { + pgdog::logger(); + + pgdog::net::tls::test_reset_acceptor(); + + let base = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tests/tls"); + let cert = base.join("cert.pem"); + let key = base.join("key.pem"); + + let mut cfg = pgdog::config::ConfigAndUsers::default(); + cfg.config.general.tls_certificate = Some(cert.clone()); + cfg.config.general.tls_private_key = Some(key.clone()); + pgdog::config::set(cfg).unwrap(); + + let build_start = Instant::now(); + pgdog::net::tls::reload().unwrap(); + let _build_time = build_start.elapsed(); + + let snapshot = pgdog::net::tls::acceptor().expect("acceptor initialized"); + + let reuse_start = Instant::now(); + for _ in 0..10 { + let candidate = pgdog::net::tls::acceptor().expect("acceptor initialized"); + assert!(Arc::ptr_eq(&snapshot, &candidate)); + } + let reuse_time = reuse_start.elapsed(); + + assert!( + reuse_time < Duration::from_millis(5), + "cached acceptor clones should be fast: {:?}", + reuse_time + ); + + assert_eq!(pgdog::net::tls::test_acceptor_build_count(), 1); + + pgdog::net::tls::test_reset_acceptor(); + pgdog::config::set(pgdog::config::ConfigAndUsers::default()).unwrap(); +} diff --git a/plugins/pgdog-example-plugin/src/plugin.rs b/plugins/pgdog-example-plugin/src/plugin.rs index c452e609e..d6e26adf1 100644 --- a/plugins/pgdog-example-plugin/src/plugin.rs +++ b/plugins/pgdog-example-plugin/src/plugin.rs @@ -55,10 +55,11 @@ pub(crate) fn route_query(context: Context) -> Result { if let NodeEnum::RangeVar(RangeVar { relname, .. }) = table_name { // Got info on last write. - if let Some(last_write) = { WRITE_TIMES.lock().get(relname).cloned() } - && last_write.elapsed() > Duration::from_secs(5) - && context.has_replicas() - { + if let Some(last_write) = { WRITE_TIMES.lock().get(relname).cloned() } { + if last_write.elapsed() > Duration::from_secs(5) && context.has_replicas() { + return Ok(Route::new(Shard::Unknown, ReadWrite::Read)); + } + } else if context.has_replicas() { return Ok(Route::new(Shard::Unknown, ReadWrite::Read)); } }