diff --git a/.github/workflows/ci-new-parser.yml b/.github/workflows/ci-new-parser.yml deleted file mode 100644 index c2252d50c..000000000 --- a/.github/workflows/ci-new-parser.yml +++ /dev/null @@ -1,114 +0,0 @@ -name: ci-new-parser -on: - push: - branches: - - main - pull_request: - types: [opened, synchronize, reopened] - -jobs: - build: - runs-on: blacksmith-4vcpu-ubuntu-2404 - steps: - - uses: actions/checkout@v6 - - name: Install CI deps - run: bash integration/ci/install-deps.sh - # Compute the binary key before restoring rust-cache: rust-cache restores - # target/ build outputs that would otherwise change the key. Use the - # split cache actions so pgdog binaries are saved immediately after the - # build, before rust-cache's post-job cleanup prunes workspace binaries - # from target/{debug,release}. - - name: Compute cache key - id: cache-key - run: echo "key=pgdog-bin-new-parser-${{ runner.os }}-$(bash integration/ci/cache-key.sh)" >> "$GITHUB_OUTPUT" - - name: Restore pgdog binaries - id: pgdog-bin - uses: actions/cache/restore@v5 - with: - path: target/debug/pgdog - key: ${{ steps.cache-key.outputs.key }} - - uses: Swatinem/rust-cache@v2 - if: steps.pgdog-bin.outputs.cache-hit != 'true' - with: - prefix-key: build-new-parser-v1 - - name: Build (debug) - if: steps.pgdog-bin.outputs.cache-hit != 'true' - run: cargo build --no-default-features --features new_parser --bin pgdog - - name: Save pgdog binaries - if: steps.pgdog-bin.outputs.cache-hit != 'true' - uses: actions/cache/save@v5 - with: - path: target/debug/pgdog - key: ${{ steps.cache-key.outputs.key }} - - ci: - runs-on: blacksmith-4vcpu-ubuntu-2404 - needs: [build] - timeout-minutes: 30 - continue-on-error: ${{ matrix.continue_on_error == true }} - strategy: - fail-fast: false - matrix: - include: - - { name: pgbench, script: integration/pgbench/run.sh } - - { name: schema-sync, script: integration/schema_sync/run.sh } - - { name: go, script: integration/go/run.sh } - - { name: js, script: integration/js/pg_tests/run.sh } - - { name: ruby, script: integration/ruby/run.sh } - - { name: java, script: integration/java/run.sh } - - { name: elixir, script: integration/elixir/run.sh, needs_beam: true } - - { name: mirror, script: integration/mirror/run.sh } - - { name: sql, script: integration/sql/run.sh } - - { name: toxi, script: integration/toxi/run.sh } - - { name: rust, script: integration/rust/run.sh } - - { name: python, script: integration/python/run.sh } - - { name: two-pc, script: integration/two_pc/run.sh } - - { name: complex, script: integration/complex/run.sh } - - { name: dry-run, script: integration/dry_run/run.sh } - - { name: copy-data, script: integration/copy_data/run.sh } - - { name: load-balancer, script: integration/load_balancer/run.sh } - - { name: tls, script: integration/tls/run.sh } - - { name: resharding, script: integration/resharding/run.sh } - # plugins/run.sh builds 4 plugin crates with cargo, so it needs - # the workspace target/ cache; the other entries just run the - # cached pgdog binary. - - { name: plugins, script: integration/plugins/run.sh, needs_rust_cache: true, continue_on_error: true } - env: - PGDOG_BIN: ${{ github.workspace }}/target/debug/pgdog - PGDOG_PLUGIN_FEATURES: new_parser - steps: - - uses: actions/checkout@v6 - - name: Install CI deps - run: bash integration/ci/install-deps.sh - - name: Compute cache key - id: cache-key - run: echo "key=pgdog-bin-new-parser-${{ runner.os }}-$(bash integration/ci/cache-key.sh)" >> "$GITHUB_OUTPUT" - # rust-cache must run before the binary restore: it lays down a - # stale target/ that can otherwise wipe target binaries when cargo - # reconciles fingerprints during plugin builds. - - name: Restore Rust cache for plugin builds - if: matrix.needs_rust_cache - uses: Swatinem/rust-cache@v2 - with: - prefix-key: build-new-parser-v1 - - name: Restore pgdog binaries - uses: actions/cache/restore@v5 - with: - path: target/debug/pgdog - key: ${{ steps.cache-key.outputs.key }} - fail-on-cache-miss: true - # The runner image has no BEAM toolchain, and Ubuntu's `elixir` package - # is several years behind, so pull a pinned precompiled pair instead. - - name: Install Erlang/Elixir - if: matrix.needs_beam - uses: erlef/setup-beam@v1 - with: - otp-version: "28" - elixir-version: "1.20.3" - - name: Setup dependencies - run: bash integration/ci/setup.sh --with-toxi - - name: Run ${{ matrix.name }} - run: bash ${{ matrix.script }} - - name: Ensure PgDog stopped - if: always() - run: bash integration/ci/ensure-pgdog-stopped.sh diff --git a/.github/workflows/fmt.yml b/.github/workflows/fmt.yml index 5d9621679..18d8f7fe4 100644 --- a/.github/workflows/fmt.yml +++ b/.github/workflows/fmt.yml @@ -17,5 +17,3 @@ jobs: run: cargo fmt --all -- --check - name: Clippy run: cargo clippy --all-targets -- -D warnings - - name: Clippy (new parser) - run: cargo clippy --all-targets --no-default-features --features new_parser -- -D warnings diff --git a/.github/workflows/package-new-parser.yml b/.github/workflows/package-new-parser.yml deleted file mode 100644 index 41025c3ae..000000000 --- a/.github/workflows/package-new-parser.yml +++ /dev/null @@ -1,168 +0,0 @@ -name: package-new-parser -on: - push: - branches: ['main'] - release: - types: [published] - workflow_dispatch: - -env: - REGISTRY: ghcr.io - IMAGE_NAME: ${{ github.repository }} - BUILDER_BASE: ghcr.io/${{ github.repository }}-base-builder:latest - RUNTIME_BASE: ghcr.io/${{ github.repository }}-base-runtime:latest -jobs: - build: - runs-on: ${{ matrix.runner }} - strategy: - fail-fast: false - matrix: - include: - - platform: linux/amd64 - runner: blacksmith-4vcpu-ubuntu-2404 - - platform: linux/arm64 - runner: blacksmith-4vcpu-ubuntu-2404-arm - permissions: - contents: read - packages: write - attestations: write - id-token: write - steps: - - name: Checkout repository - uses: actions/checkout@v6 - with: - fetch-tags: true - fetch-depth: 0 - - - name: Prepare - run: | - platform='${{ matrix.platform }}' - echo "PLATFORM_PAIR=${platform//\//-}" >> $GITHUB_ENV - - - name: Log in to the Container registry - uses: docker/login-action@v3 - with: - registry: ${{ env.REGISTRY }} - username: ${{ github.actor }} - password: ${{ secrets.GITHUB_TOKEN }} - - - name: Extract metadata (tags, labels) for Docker - id: meta - uses: docker/metadata-action@v5 - with: - images: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }} - - - name: Set up Docker Buildx - uses: docker/setup-buildx-action@v3 - - - name: Build and push Docker image - id: build - uses: docker/build-push-action@v6 - with: - context: . - # Only tag by registry + image name to to push by digest - tags: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME}} - labels: ${{ steps.meta.outputs.labels }} - platforms: ${{ matrix.platform }} - build-args: | - BUILDER_BASE=${{ env.BUILDER_BASE }} - RUNTIME_BASE=${{ env.RUNTIME_BASE }} - FEATURES=new_parser - outputs: type=image,push-by-digest=true,name-canonical=true,push=true - - - name: Export digest - run: | - mkdir -p ${{ runner.temp }}/digests - digest="${{ steps.build.outputs.digest }}" - touch "${{ runner.temp }}/digests/${digest#sha256:}" - - - name: Upload digest - uses: actions/upload-artifact@v4 - with: - name: digests-${{ env.PLATFORM_PAIR }} - path: ${{ runner.temp }}/digests/* - if-no-files-found: error - retention-days: 1 - - - name: Generate artifact attestation - uses: actions/attest-build-provenance@v2 - with: - subject-name: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME}} - subject-digest: ${{ steps.build.outputs.digest }} - push-to-registry: true - - merge: - runs-on: blacksmith-4vcpu-ubuntu-2404 - needs: - - build - permissions: - contents: read - packages: write - attestations: write - id-token: write - steps: - - name: Checkout repository - uses: actions/checkout@v6 - with: - fetch-tags: true - fetch-depth: 0 - - - name: Get short commit SHA - id: commit - run: | - COMMIT_SHA=$(git rev-parse --short HEAD) - echo "sha=$COMMIT_SHA" >> $GITHUB_OUTPUT - echo "Short commit SHA: $COMMIT_SHA" - - - name: Download digests - uses: actions/download-artifact@v4 - with: - path: ${{ runner.temp }}/digests - pattern: digests-* - merge-multiple: true - - - name: Login to Docker Hub - uses: docker/login-action@v3 - with: - registry: ${{ env.REGISTRY }} - username: ${{ github.actor }} - password: ${{ secrets.GITHUB_TOKEN }} - - - name: Set up Docker Buildx - uses: docker/setup-buildx-action@v3 - - - name: Check for git tag - id: tag - run: | - TAG=$(git describe --exact-match --tags HEAD 2>/dev/null || echo "") - echo "tag=$TAG" >> $GITHUB_OUTPUT - if [ -n "$TAG" ]; then - echo "Git tag found: $TAG" - else - echo "No git tag found for current commit" - fi - - - name: Docker meta - id: meta - uses: docker/metadata-action@v5 - with: - images: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }} - flavor: | - suffix=-new-parser - tags: | - type=ref,event=branch - type=ref,event=pr - type=semver,pattern={{version}} - type=semver,pattern={{major}}.{{minor}} - type=raw,value=${{ steps.tag.outputs.tag }},enable=${{ steps.tag.outputs.tag != '' }} - type=raw,value=${{ steps.commit.outputs.sha }} - - - name: Create manifest list and push - working-directory: ${{ runner.temp }}/digests - run: | - docker buildx imagetools create $(jq -cr '.tags | map("-t " + .) | join(" ")' <<< "$DOCKER_METADATA_OUTPUT_JSON") \ - $(printf '${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}@sha256:%s ' *) - - - name: Inspect image - run: | - docker buildx imagetools inspect ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:${{ steps.meta.outputs.version }} diff --git a/.github/workflows/tests-new-parser.yml b/.github/workflows/tests-new-parser.yml deleted file mode 100644 index d180b828a..000000000 --- a/.github/workflows/tests-new-parser.yml +++ /dev/null @@ -1,36 +0,0 @@ -name: tests-new-parser -on: - push: - branches: - - main - pull_request: - types: [opened, synchronize, reopened] - -jobs: - tests: - runs-on: blacksmith-4vcpu-ubuntu-2404 - steps: - - uses: actions/checkout@v6 - - name: Install CI deps - run: bash integration/ci/install-deps.sh - - uses: Swatinem/rust-cache@v2 - with: - prefix-key: "new-parser-v1" # Change this when updating tooling - - name: Setup PostgreSQL - run: bash integration/ci/setup.sh --with-toxi - - name: Run tests with coverage - env: - RUSTFLAGS: "--cfg tokio_unstable -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 --package pgdog --package pgdog-config --package pgdog-vector --package pgdog-stats --package pgdog-postgres-types --no-default-features --features new_parser --filter-expr "package(pgdog) | package(pgdog-config) | package(pgdog-vector) | package(pgdog-stats) | package(pgdog-postgres-types)" - - name: Run documentation tests - run: cargo test --doc --no-default-features --features new_parser - # Requires CODECOV_TOKEN secret for upload - - uses: codecov/codecov-action@v4 - env: - CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }} - with: - files: lcov.info - flags: unit - fail_ci_if_error: true diff --git a/Cargo.lock b/Cargo.lock index 7717242b0..81d68db84 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1805,12 +1805,6 @@ version = "0.4.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0ce7134b9999ecaf8bcd65542e436736ef32ddca1b3e06094cb6ec5755203b80" -[[package]] -name = "fixedbitset" -version = "0.5.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1d674e81391d1e1ab681a28d99df07927c6d4aa5b027d7da16ba32d1d21ecd99" - [[package]] name = "flate2" version = "1.1.9" @@ -2585,15 +2579,6 @@ version = "1.70.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a6cb138bb79a146c1bd460005623e142ef0181e3d0219cb493e02f7d08a35695" -[[package]] -name = "itertools" -version = "0.10.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b0fd2260e829bddf4cb6ea802289de2f86d6a7a690192fbe91b3f46e0f2c8473" -dependencies = [ - "either", -] - [[package]] name = "itertools" version = "0.13.0" @@ -2948,12 +2933,6 @@ dependencies = [ "uuid", ] -[[package]] -name = "multimap" -version = "0.10.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1d87ecb2933e8aeadb3e3a02b828fed80a7528047e68b4f424523a0981a3a084" - [[package]] name = "native-tls" version = "0.2.18" @@ -3289,33 +3268,6 @@ dependencies = [ "sha2 0.10.9", ] -[[package]] -name = "petgraph" -version = "0.7.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3672b37090dbd86368a4145bc067582552b29c27377cad4e0a306c97f9bd7772" -dependencies = [ - "fixedbitset 0.5.7", - "indexmap", -] - -[[package]] -name = "pg_query" -version = "6.1.1" -source = "git+https://github.com/pgdogdev/pg_query.rs.git?rev=97019d0c13ad0b888fe91ee5bed5448b5f409cdd#97019d0c13ad0b888fe91ee5bed5448b5f409cdd" -dependencies = [ - "bindgen 0.72.1", - "cc", - "fs_extra", - "glob", - "itertools 0.10.5", - "prost", - "prost-build", - "serde", - "serde_json", - "thiserror 1.0.69", -] - [[package]] name = "pg_raw_parse" version = "0.1.0" @@ -3376,7 +3328,6 @@ dependencies = [ "moka", "once_cell", "parking_lot", - "pg_query", "pg_raw_parse", "pgdog-config", "pgdog-macros", @@ -3478,7 +3429,6 @@ version = "0.4.0" dependencies = [ "bindgen 0.71.1", "libloading", - "pg_query", "pg_raw_parse", "pgdog-postgres-types", "tracing", @@ -3760,58 +3710,6 @@ dependencies = [ "unicode-ident", ] -[[package]] -name = "prost" -version = "0.13.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2796faa41db3ec313a31f7624d9286acf277b52de526150b7e69f3debf891ee5" -dependencies = [ - "bytes", - "prost-derive", -] - -[[package]] -name = "prost-build" -version = "0.13.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "be769465445e8c1474e9c5dac2018218498557af32d9ed057325ec9a41ae81bf" -dependencies = [ - "heck", - "itertools 0.14.0", - "log", - "multimap", - "once_cell", - "petgraph", - "prettyplease", - "prost", - "prost-types", - "regex", - "syn 2.0.118", - "tempfile", -] - -[[package]] -name = "prost-derive" -version = "0.13.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8a56d757972c98b346a9b766e3f02746cde6dd1cd1d1d563472929fdd74bec4d" -dependencies = [ - "anyhow", - "itertools 0.14.0", - "proc-macro2", - "quote", - "syn 2.0.118", -] - -[[package]] -name = "prost-types" -version = "0.13.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "52c2c1bf36ddb1a1c396b3601a3cec27c2462e45f07c386894ec3ccf5332bd16" -dependencies = [ - "prost", -] - [[package]] name = "ptr_meta" version = "0.1.4" @@ -5214,7 +5112,7 @@ dependencies = [ "fancy-regex", "filedescriptor", "finl_unicode", - "fixedbitset 0.4.2", + "fixedbitset", "hex", "lazy_static", "libc", diff --git a/Cargo.toml b/Cargo.toml index 3db828a23..22208e3e9 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -36,9 +36,6 @@ futures-util = "0.3.32" # [patch.crates-io] # tokio = { path = "../tokio/tokio" } -# [patch."https://github.com/pgdogdev/pg_query.rs.git"] -# pg_query = { path = "../pg_query.rs" } - # [patch."https://github.com/pgdogdev/pg_raw_parse.git"] # pg_raw_parse = { path = "../pg_raw_parse" } diff --git a/README.md b/README.md index de0370d1f..3ff1e342f 100644 --- a/README.md +++ b/README.md @@ -149,7 +149,7 @@ Health checks maximize database availability and protect against bad network con 📘 **[Single endpoint](https://docs.pgdog.dev/features/load-balancer/#single-endpoint)** -PgDog uses [`pg_query`](https://github.com/pganalyze/pg_query.rs), which includes the PostgreSQL native parser. By parsing queries, PgDog can detect writes (e.g. `INSERT`, `UPDATE`, `CREATE TABLE`, etc.) and send them to the primary, leaving the replicas to serve reads (`SELECT`). This allows applications to connect to the same PgDog deployment for both reads and writes. +PgDog uses [`pg_raw_parse`](https://github.com/pgdogdev/pg_raw_parse), which includes the PostgreSQL native parser. By parsing queries, PgDog can detect writes (e.g. `INSERT`, `UPDATE`, `CREATE TABLE`, etc.) and send them to the primary, leaving the replicas to serve reads (`SELECT`). This allows applications to connect to the same PgDog deployment for both reads and writes. ##### Transactions diff --git a/docs/PLUGIN_SYSTEM.md b/docs/PLUGIN_SYSTEM.md index 2e695685f..7fcde9dd2 100644 --- a/docs/PLUGIN_SYSTEM.md +++ b/docs/PLUGIN_SYSTEM.md @@ -64,13 +64,12 @@ Client sends query ↓ PostgreSQL Frontend Parser ├─ Tokenize and parse query - └─ Generate pg_query AST + └─ Generate AST ↓ Query Router: QueryParser::parse() ├─ Create RouterContext (shards, replicas, etc.) ├─ Generate PdRouterContext for plugins │ └─ context.plugin_context(ast, bind_params) - │ ├─ Extract AST from pg_query ParseResult │ ├─ Pack bind parameters into PdParameters │ └─ Include cluster metadata (shards, replicas, transaction state) │ @@ -141,7 +140,7 @@ See the [plugins/pgdog-example-plugin/](../plugins/pgdog-example-plugin/) for a ## Safety & Compatibility - **Rust version:** Plugins must be built with the exact same Rust compiler version as PgDog. Mismatches are skipped at load time. -- **pg_query version:** Plugins must use the same `pg_query` version as PgDog. Use the re-exports from `pgdog-plugin`. +- **pg_raw_parse version:** Plugins must use the same `pg_raw_parse` version as PgDog. (The compiler will automatically enforce this) - **FFI safety:** All FFI types are `#[repr(C)]` and memory is managed to avoid UB. See [pgdog-plugin/src/bindings.rs](../pgdog-plugin/src/bindings.rs) and [pgdog-plugin/src/parameters.rs](../pgdog-plugin/src/parameters.rs). ## FFI & ABI Notes @@ -155,7 +154,7 @@ See the [plugins/pgdog-example-plugin/](../plugins/pgdog-example-plugin/) for a - **Opaque pointers hide implementation details:** Pointer fields (e.g., `void*`) obscure the true layout and ownership, so important rules are only in documentation, not enforced by the type system. - **Fragile container reinterpretation:** Using `Vec::from_raw_parts` and similar tricks relies on *identical* Rust versions, crate versions, and feature flags. Any mismatch can cause undefined behavior or memory corruption. -- **Transitive dependency coupling:** Types like the `pg_query` AST or `bytes::Bytes` add hidden constraints on plugin dependencies, making upgrades and changes risky. +- **Transitive dependency coupling:** Types like the `pg_raw_parse` AST or `bytes::Bytes` add hidden constraints on plugin dependencies, making upgrades and changes risky. - **Hard to evolve:** Internal representation changes (e.g., switching `Bytes` → `Vec`, changing struct layouts) are breaking and require all plugins to be rebuilt in lockstep. - **Debugging cost:** ABI or UB problems are subtle, often nondeterministic, and hard to diagnose or reproduce. - **No ABI versioning:** There is no formal ABI version negotiation; any change in the host or plugin can silently break compatibility. @@ -255,7 +254,6 @@ The following areas lack test coverage: #### Plugin Loading & Lifecycle - ❌ Rust compiler version mismatch scenarios -- ❌ pg_query version verification (currently not implemented) - ❌ Plugin symbol resolution failures - ❌ Plugin with missing required functions - ❌ Plugin init/fini execution order diff --git a/integration/plugins/test-plugins/test-plugin-compatible/Cargo.toml b/integration/plugins/test-plugins/test-plugin-compatible/Cargo.toml index eeeeed422..8b291c532 100644 --- a/integration/plugins/test-plugins/test-plugin-compatible/Cargo.toml +++ b/integration/plugins/test-plugins/test-plugin-compatible/Cargo.toml @@ -9,8 +9,4 @@ edition = "2024" crate-type = ["cdylib"] [dependencies] -pgdog-plugin = { path = "../../../../pgdog-plugin", default-features = false } - -[features] -default = ["pgdog-plugin/pg_query"] -new_parser = ["pgdog-plugin/new_parser"] +pgdog-plugin = { path = "../../../../pgdog-plugin" } diff --git a/integration/plugins/test-plugins/test-plugin-compatible/src/lib.rs b/integration/plugins/test-plugins/test-plugin-compatible/src/lib.rs index 52f2d1d6e..62bfe2673 100644 --- a/integration/plugins/test-plugins/test-plugin-compatible/src/lib.rs +++ b/integration/plugins/test-plugins/test-plugin-compatible/src/lib.rs @@ -31,9 +31,6 @@ impl Plugin for TestPlugin { // query should be accessible let query = context.query; - #[cfg(not(feature = "new_parser"))] - assert!(query.nodes().len() >= 1); - #[cfg(feature = "new_parser")] assert!(query.stmts().next().is_some()); // Write to output file on first call only diff --git a/integration/rust/Cargo.toml b/integration/rust/Cargo.toml index 427b729b1..3cbc0ecb7 100644 --- a/integration/rust/Cargo.toml +++ b/integration/rust/Cargo.toml @@ -6,9 +6,6 @@ edition = "2024" [lib] test = true -[features] -new_parser = [] - [dependencies] tokio-postgres = {version = "0.7.13", features = ["with-uuid-1"]} sqlx = { version = "0.8.6", features = ["postgres", "runtime-tokio", "tls-native-tls", "bigdecimal", "chrono", "json", "rust_decimal"]} diff --git a/integration/rust/tests/integration/cross_shard_oid_drift.rs b/integration/rust/tests/integration/cross_shard_oid_drift.rs index 0c3aeb76a..eef84028c 100644 --- a/integration/rust/tests/integration/cross_shard_oid_drift.rs +++ b/integration/rust/tests/integration/cross_shard_oid_drift.rs @@ -1,4 +1,3 @@ -#![cfg(feature = "new_parser")] use crate::setup::{admin_sqlx, connections_sqlx}; use sqlx::postgres::types::Oid; use sqlx::{Column, Executor, Row}; diff --git a/pgdog-plugin/Cargo.toml b/pgdog-plugin/Cargo.toml index 7b2ef3a16..d8975e68d 100644 --- a/pgdog-plugin/Cargo.toml +++ b/pgdog-plugin/Cargo.toml @@ -15,8 +15,7 @@ crate-type = ["rlib", "cdylib"] [dependencies] libloading = "0.8" -pg_query = { git = "https://github.com/pgdogdev/pg_query.rs.git", rev = "97019d0c13ad0b888fe91ee5bed5448b5f409cdd", optional = true } -pg_raw_parse = { workspace = true, optional = true } +pg_raw_parse.workspace = true pgdog-postgres-types.workspace = true tracing = "0.1" @@ -26,5 +25,4 @@ tracing-subscriber = { version = "0.3", features = ["env-filter", "json"] } bindgen = "0.71.0" [features] -default = ["pg_query"] -new_parser = ["pg_raw_parse"] +default = [] diff --git a/pgdog-plugin/src/context.rs b/pgdog-plugin/src/context.rs index c2960826f..0564f988d 100644 --- a/pgdog-plugin/src/context.rs +++ b/pgdog-plugin/src/context.rs @@ -51,14 +51,6 @@ pub struct Context<'a> { pub in_transaction: bool, /// PgDog strongly believes this statement should go to a primary. pub write_override: bool, - /// pg_query generated Abstract Syntax Tree of the statement. - // Note: This is not at all FFI safe, and is UB across an FFI boundary. - // We are relying on an implemenation detail of the compiler that could - // change at any time to make this work. There is no safe way to pass - // this type, but this won't be an issue with the new parser - #[cfg(not(feature = "new_parser"))] - pub query: &'a pg_query::protobuf::ParseResult, - #[cfg(feature = "new_parser")] /// The parsed Abstract Syntax Tree of the statement(s). pub query: &'a pg_raw_parse::StmtList, /// Bound parameters. @@ -201,22 +193,12 @@ impl Context<'_> { impl Context<'_> { #[doc(hidden)] pub fn doc_test() -> Self { - #[cfg(not(feature = "new_parser"))] - use pg_query::protobuf::ParseResult; - #[cfg(not(feature = "new_parser"))] - static EMPTY_PARSE_RESULT: ParseResult = ParseResult { - version: 0, - stmts: Vec::new(), - }; Context { shards: 1, has_replicas: true, has_primary: true, in_transaction: false, write_override: false, - #[cfg(not(feature = "new_parser"))] - query: &EMPTY_PARSE_RESULT, - #[cfg(feature = "new_parser")] query: pg_raw_parse::list::empty_list(), params: Parameters::default(), } diff --git a/pgdog-plugin/src/lib.rs b/pgdog-plugin/src/lib.rs index 79dd2081c..9bbd6ffd2 100644 --- a/pgdog-plugin/src/lib.rs +++ b/pgdog-plugin/src/lib.rs @@ -19,7 +19,7 @@ //! //! ## Dependencies //! -//! PgDog is using [`pg_query`] to parse SQL. It produces an Abstract Syntax Tree (AST) which plugins can use to inspect queries +//! PgDog is using [`pg_raw_parse`] to parse SQL. It produces an Abstract Syntax Tree (AST) which plugins can use to inspect queries //! and make statement routing decisions. //! //! The AST is computed by PgDog at runtime. It then passes it down to plugins, using a FFI interface. To make this safe, plugins must follow the @@ -68,10 +68,7 @@ //! //! ```no_run //! use pgdog_plugin::prelude::*; -//! #[cfg(feature = "new_parser")] //! use pg_raw_parse::Node; -//! #[cfg(not(feature = "new_parser"))] -//! use pg_query::{protobuf::{Node, RawStmt}, NodeEnum}; //! //! pgdog_plugin::plugin!(MyPlugin); //! @@ -83,20 +80,10 @@ //! } //! //! fn route(context: Context<'_>) -> Route { -//! #[cfg(feature = "new_parser")] //! if let Some(Node::SelectStmt(_)) = context.query.stmts().next() { //! return Route::new(Shard::Unknown, ReadWrite::Read); //! } //! -//! #[cfg(not(feature = "new_parser"))] -//! if let Some(root) = context.query.stmts.first() -//! && let Some(ref stmt) = root.stmt -//! && let Some(ref node) = stmt.node -//! && let NodeEnum::SelectStmt(_) = node -//! { -//! return Route::new(Shard::Unknown, ReadWrite::Read); -//! } -//! //! Route::new(Shard::Unknown, ReadWrite::Write) //! } //! } @@ -188,13 +175,6 @@ //! ``` //! -#[cfg(all(feature = "pg_query", feature = "pg_raw_parse"))] -compile_error!("Cannot build with both the old and new parser"); -#[cfg(not(any(feature = "pg_query", feature = "pg_raw_parse")))] -compile_error!( - r#"pg-plugin must be built with either default features, or features = "new_parser""# -); - mod config; pub mod context; pub mod logging; @@ -213,8 +193,5 @@ pub use string::PdStr; pub use libloading; -#[cfg(feature = "pg_query")] -pub use pg_query; - pub const RUSTC_VERSION: &str = env!("RUSTC_VERSION"); pub const VERSION: &str = env!("CARGO_PKG_VERSION"); diff --git a/pgdog-plugin/src/plugin.rs b/pgdog-plugin/src/plugin.rs index 0432073ff..571b17788 100644 --- a/pgdog-plugin/src/plugin.rs +++ b/pgdog-plugin/src/plugin.rs @@ -37,7 +37,6 @@ pub struct PluginVtable { /// Configure plugin. config: extern "C-unwind" fn(Config<'_>) -> bool, /// Route query. - #[cfg(feature = "new_parser")] route: extern "C-unwind" fn( u64, bool, @@ -47,16 +46,6 @@ pub struct PluginVtable { &pg_raw_parse::StmtList, RawParameters<'_>, ) -> Route, - #[cfg(not(feature = "new_parser"))] - route: extern "C-unwind" fn( - u64, - bool, - bool, - bool, - bool, - &pg_query::protobuf::ParseResult, - RawParameters<'_>, - ) -> Route, /// Logging initialization. logging_init: extern "C-unwind" fn(Config<'_>), } @@ -93,8 +82,7 @@ pub trait Plugin { has_primary: bool, in_transaction: bool, write_override: bool, - #[cfg(not(feature = "new_parser"))] query: &pg_query::protobuf::ParseResult, - #[cfg(feature = "new_parser")] query: &pg_raw_parse::StmtList, + query: &pg_raw_parse::StmtList, params: RawParameters<'_>, ) -> Route { let context = Context { diff --git a/pgdog-plugin/src/prelude.rs b/pgdog-plugin/src/prelude.rs index ecfb4282b..a74c21cd5 100644 --- a/pgdog-plugin/src/prelude.rs +++ b/pgdog-plugin/src/prelude.rs @@ -1,7 +1,5 @@ //! Commonly used structs and re-exports. -#[cfg(feature = "pg_query")] -pub use crate::pg_query; pub use crate::{ Context, ParameterFormat, PdStr, Plugin, ReadWrite, Route, Shard, parameters::{Parameter, ParameterValue, Parameters}, diff --git a/pgdog/Cargo.toml b/pgdog/Cargo.toml index 6987a3456..b7d4cc95e 100644 --- a/pgdog/Cargo.toml +++ b/pgdog/Cargo.toml @@ -11,9 +11,7 @@ readme = "README.md" default-run = "pgdog" [features] -default = ["pg_query", "pgdog-plugin/pg_query"] tui = ["ratatui"] -new_parser = ["pg_raw_parse", "pgdog-plugin/new_parser"] [dependencies] bon.workspace = true @@ -48,7 +46,6 @@ base64 = "0.22" md5 = "0.7" futures = "0.3" csv-core = "0.1" -pg_query = { git = "https://github.com/pgdogdev/pg_query.rs.git", rev = "97019d0c13ad0b888fe91ee5bed5448b5f409cdd", optional = true } regex = "1" memchr = "2" semver = "1" @@ -85,7 +82,7 @@ smallvec = "1" reqwest.workspace = true hex = "0.4" x509-parser = "0.18" -pg_raw_parse = { workspace = true, optional = true } +pg_raw_parse.workspace = true itertools = "0.15.0" dyn-clone = "1.0.20" diff --git a/pgdog/benches/comment_parser.rs b/pgdog/benches/comment_parser.rs index 22afa9ed7..e4cedc232 100644 --- a/pgdog/benches/comment_parser.rs +++ b/pgdog/benches/comment_parser.rs @@ -1,6 +1,4 @@ use brunch::{Bench, benches}; -#[cfg(not(feature = "new_parser"))] -use pg_query::scan_raw; use pgdog::frontend::router::parser::comment::parse_edge_comment; const QUERY_WITH_LEADING: &str = @@ -9,7 +7,6 @@ const QUERY_WITH_TRAILING: &str = "SELECT * FROM users WHERE id = $1 AND name = $2 /* pgdog_role: primary */"; const QUERY_NO_COMMENT: &str = "SELECT * FROM users WHERE id = $1 AND name = $2"; -#[cfg(feature = "new_parser")] benches!( Bench::new("parse_edge_comment(leading)") .run(|| parse_edge_comment(QUERY_WITH_LEADING, &Default::default())), @@ -18,15 +15,3 @@ benches!( Bench::new("parse_edge_comment(no comment)") .run(|| parse_edge_comment(QUERY_NO_COMMENT, &Default::default())), ); -#[cfg(not(feature = "new_parser"))] -benches!( - Bench::new("parse_edge_comment(leading)") - .run(|| parse_edge_comment(QUERY_WITH_LEADING, &Default::default())), - Bench::new("parse_edge_comment(trailing)") - .run(|| parse_edge_comment(QUERY_WITH_TRAILING, &Default::default())), - Bench::new("parse_edge_comment(no comment)") - .run(|| parse_edge_comment(QUERY_NO_COMMENT, &Default::default())), - Bench::new("scan_raw(leading)").run(|| scan_raw(QUERY_WITH_LEADING)), - Bench::new("scan_raw(trailing)").run(|| scan_raw(QUERY_WITH_TRAILING)), - Bench::new("scan_raw(no comment)").run(|| scan_raw(QUERY_NO_COMMENT)), -); diff --git a/pgdog/src/admin/set.rs b/pgdog/src/admin/set.rs index 3b0268426..a06c1514b 100644 --- a/pgdog/src/admin/set.rs +++ b/pgdog/src/admin/set.rs @@ -5,9 +5,6 @@ use crate::{ }; use super::prelude::*; -#[cfg(not(feature = "new_parser"))] -use pg_query::{NodeEnum, parse, protobuf::a_const}; -#[cfg(feature = "new_parser")] use pg_raw_parse::Node; use serde::de::DeserializeOwned; @@ -22,7 +19,6 @@ impl Command for Set { "SET".into() } - #[cfg(feature = "new_parser")] fn parse(sql: &str) -> Result { let stmt = pg_raw_parse::parse(sql).map_err(|_| Error::Syntax)?; let root = stmt.stmts().next().ok_or(Error::Syntax)?; @@ -45,44 +41,6 @@ impl Command for Set { } } - cfg_select! { - not(feature = "new_parser") => { - fn parse(sql: &str) -> Result { - let stmt = parse(sql).map_err(|_| Error::Syntax)?; - let root = stmt.protobuf.stmts.first().cloned().ok_or(Error::Syntax)?; - let stmt = root.stmt.ok_or(Error::Syntax)?; - match stmt.node.ok_or(Error::Syntax)? { - NodeEnum::VariableSetStmt(stmt) => { - let name = stmt.name; - - let setting = stmt.args.first().ok_or(Error::Syntax)?; - let node = setting.node.clone().ok_or(Error::Syntax)?; - match node { - NodeEnum::AConst(a_const) => match a_const.val { - Some(a_const::Val::Ival(val)) => Ok(Self { - name, - value: val.ival.to_string(), - }), - - Some(a_const::Val::Sval(sval)) => Ok(Self { - name, - value: sval.sval.to_string(), - }), - - _ => Err(Error::Syntax), - }, - - _ => Err(Error::Syntax), - } - } - - _ => Err(Error::Syntax), - } - } - } - _ => {} - } - async fn execute(&self) -> Result, Error> { let _lock = databases::lock(); let mut config = (*config()).clone(); diff --git a/pgdog/src/backend/pool/cluster.rs b/pgdog/src/backend/pool/cluster.rs index d10dfb85d..238eb3e9f 100644 --- a/pgdog/src/backend/pool/cluster.rs +++ b/pgdog/src/backend/pool/cluster.rs @@ -776,7 +776,6 @@ impl Cluster { .await } - #[cfg(feature = "new_parser")] pub(crate) fn is_canonicalizing_oids(&self) -> bool { self.canonical_oids.is_some() } diff --git a/pgdog/src/backend/pool/connection/aggregate.rs b/pgdog/src/backend/pool/connection/aggregate.rs index 2e311a732..a15a3bbdb 100644 --- a/pgdog/src/backend/pool/connection/aggregate.rs +++ b/pgdog/src/backend/pool/connection/aggregate.rs @@ -401,9 +401,6 @@ mod test { messages::{Field, Format, RowDescription}, }; use bytes::Bytes; - #[cfg(not(feature = "new_parser"))] - use pg_query::NodeEnum; - #[cfg(feature = "new_parser")] use pg_raw_parse::Node; use pgdog_postgres_types::Double; use std::assert_matches; @@ -445,7 +442,6 @@ mod test { } } - #[cfg(feature = "new_parser")] fn parse(stmt: &str) -> Aggregate { let ast = pg_raw_parse::parse(stmt).unwrap(); let Node::SelectStmt(stmt) = ast.stmts().next().unwrap() else { @@ -454,22 +450,6 @@ mod test { Aggregate::parse(stmt, &Default::default()) } - #[cfg(not(feature = "new_parser"))] - fn parse(stmt: &str) -> Aggregate { - let stmt = pg_query::parse(stmt) - .unwrap() - .protobuf - .stmts - .remove(0) - .stmt - .unwrap(); - let stmt = match stmt.node.unwrap() { - NodeEnum::SelectStmt(stmt) => *stmt, - _ => panic!("not a select"), - }; - Aggregate::parse(&stmt, &Default::default()) - } - #[test] fn aggregate_count_with_int_typecast() { // Regression test for https://github.com/pgdogdev/pgdog/issues/861 diff --git a/pgdog/src/backend/replication/logical/error.rs b/pgdog/src/backend/replication/logical/error.rs index 49589dada..bb840a53f 100644 --- a/pgdog/src/backend/replication/logical/error.rs +++ b/pgdog/src/backend/replication/logical/error.rs @@ -104,10 +104,6 @@ pub enum Error { #[error("missing data")] MissingData, - #[error("pg_query: {0}")] - #[cfg(not(feature = "new_parser"))] - PgQuery(#[from] pg_query::Error), - #[error("copy error")] Copy, diff --git a/pgdog/src/backend/replication/logical/publisher/table.rs b/pgdog/src/backend/replication/logical/publisher/table.rs index 436a27d73..77de017ce 100644 --- a/pgdog/src/backend/replication/logical/publisher/table.rs +++ b/pgdog/src/backend/replication/logical/publisher/table.rs @@ -458,13 +458,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(), - source_cluster, - dest, - #[cfg(not(feature = "new_parser"))] - self.query_parser_engine, - )?; + let mut copy_sub = CopySubscriber::new(copy.statement(), source_cluster, dest)?; copy_sub.connect().await?; // Create sync slot. @@ -530,9 +524,6 @@ mod test { replication::logical::publisher::queries::{PublicationTableColumn, ReplicaIdentity}, server::test::test_server, }; - #[cfg(not(feature = "new_parser"))] - use pg_query::parse; - #[cfg(feature = "new_parser")] use pg_raw_parse::parse; use crate::config::config; diff --git a/pgdog/src/backend/replication/logical/subscriber/copy.rs b/pgdog/src/backend/replication/logical/subscriber/copy.rs index 99ebd4070..def4ec06b 100644 --- a/pgdog/src/backend/replication/logical/subscriber/copy.rs +++ b/pgdog/src/backend/replication/logical/subscriber/copy.rs @@ -2,12 +2,7 @@ //! between N shards. use futures::future::join_all; -#[cfg(not(feature = "new_parser"))] -use pg_query::{NodeEnum, parse_raw}; -#[cfg(feature = "new_parser")] use pg_raw_parse::Node; -#[cfg(not(feature = "new_parser"))] -use pgdog_config::QueryParserEngine; use tracing::debug; use crate::frontend::client::query_engine::TwoPcPhase; @@ -15,7 +10,6 @@ use crate::frontend::client::query_engine::two_pc::{ Manager, TwoPcTransaction, statement::phase_control, }; -#[cfg(feature = "new_parser")] use crate::frontend::router::parser::Error as ParseError; use crate::{ backend::{Cluster, ConnectReason, replication::subscriber::ParallelConnection}, @@ -51,7 +45,6 @@ impl CopySubscriber { /// 1. What kind of encoding we use. /// 2. Which column is used for sharding. /// - #[cfg(feature = "new_parser")] pub fn new( copy_stmt: &CopyStatement, source: &Cluster, @@ -78,53 +71,6 @@ impl CopySubscriber { }) } - cfg_select! { - not(feature = "new_parser") => { - pub fn new( - copy_stmt: &CopyStatement, - source: &Cluster, - 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 - .first() - .ok_or(Error::MissingData)? - .stmt - .as_ref() - .ok_or(Error::MissingData)? - .node - .as_ref() - .ok_or(Error::MissingData)?; - let mut copy = if let NodeEnum::CopyStmt(stmt) = stmt { - CopyParser::new(stmt, cluster).map_err(|_| Error::MissingData)? - } else { - return Err(Error::MissingData); - }; - // The destination's copy of the lookup table may still be - // syncing; the source has the complete mapping. - copy.set_lookup_cluster(source); - - Ok(Self { - copy, - cluster: cluster.clone(), - buffer: vec![], - connections: vec![], - stmt: copy_stmt.clone(), - bytes_sharded: 0, - }) - } - } - _ => {} - } - /// Connect to all shards. One connection per primary. pub async fn connect(&mut self) -> Result<(), Error> { let mut servers = vec![]; @@ -441,14 +387,7 @@ mod test { .await .unwrap(); - let mut subscriber = CopySubscriber::new( - ©, - &cluster, - &cluster, - #[cfg(not(feature = "new_parser"))] - config().config.general.query_parser_engine, - ) - .unwrap(); + let mut subscriber = CopySubscriber::new(©, &cluster, &cluster).unwrap(); subscriber.start_copy().await.unwrap(); let header = CopyData::new(&Header::new().to_bytes()); diff --git a/pgdog/src/backend/schema/sync/error.rs b/pgdog/src/backend/schema/sync/error.rs index a33d80f56..1eda6d9ec 100644 --- a/pgdog/src/backend/schema/sync/error.rs +++ b/pgdog/src/backend/schema/sync/error.rs @@ -25,11 +25,6 @@ pub enum Error { PgDump(String), #[error("{0}")] - #[cfg(not(feature = "new_parser"))] - Syntax(#[from] pg_query::Error), - - #[error("{0}")] - #[cfg(feature = "new_parser")] Syntax(#[from] pg_raw_parse::Error), #[error("parse error, stmt out of bounds")] diff --git a/pgdog/src/backend/schema/sync/pg_dump.rs b/pgdog/src/backend/schema/sync/pg_dump.rs index f1a44a97e..35814ffc7 100644 --- a/pgdog/src/backend/schema/sync/pg_dump.rs +++ b/pgdog/src/backend/schema/sync/pg_dump.rs @@ -9,15 +9,7 @@ use std::{ use lazy_static::lazy_static; use parking_lot::Mutex; -#[cfg(not(feature = "new_parser"))] -use pg_query::{ - Node, NodeEnum, - protobuf::{AlterTableType, ConstrType, ObjectType, ParseResult, RangeVar, String as PgString}, -}; -#[cfg(feature = "new_parser")] use pg_raw_parse::{Node, NodeMut, Owned, StmtList, make, nodes}; -#[cfg(not(feature = "new_parser"))] -use pgdog_config::QueryParserEngine; use regex::Regex; use tracing::{info, trace, warn}; @@ -40,40 +32,10 @@ struct ColumnTypeKey<'a> { column: &'a str, } -#[cfg(not(feature = "new_parser"))] -fn deparse_node(node: NodeEnum) -> Result { - match config().config.general.query_parser_engine { - QueryParserEngine::PgQueryProtobuf => node.deparse(), - QueryParserEngine::PgQueryRaw => node.deparse_raw(), - } -} - -#[cfg(not(feature = "new_parser"))] -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), - } -} - -#[cfg(feature = "new_parser")] fn schema_name(relation: &nodes::RangeVar) -> &str { relation.schemaname().unwrap_or("public") } -cfg_select! { - not(feature = "new_parser") => { - fn schema_name(relation: &RangeVar) -> &str { - if relation.schemaname.is_empty() { - "public" - } else { - relation.schemaname.as_str() - } - } - } - _ => {} -} - fn is_integer_type(type_name: &str) -> bool { matches!( type_name, @@ -86,7 +48,6 @@ fn is_integer_type(type_name: &str) -> bool { /// A column should be converted if: /// 1. It's in the columns_to_convert set (PK/FK that references integer PK), OR /// 2. It's a child partition column where the parent has bigint and child has integer -#[cfg(feature = "new_parser")] fn should_convert_to_bigint<'a>( col: &Column<'a>, col_type_name: Option<&nodes::TypeName>, @@ -123,52 +84,6 @@ fn should_convert_to_bigint<'a>( is_integer_type(type_name) } -cfg_select! { - not(feature = "new_parser") => { - fn should_convert_to_bigint<'a>( - col: &Column<'a>, - col_type_name: Option<&pg_query::protobuf::TypeName>, - columns_to_convert: &HashSet>, - parent_table: Option<&Table<'a>>, - parent_column_types: &HashMap<(Table<'a>, &str), &'static str>, - ) -> bool { - // Check if this column is directly marked for conversion (PK/FK) - if columns_to_convert.contains(col) { - return true; - } - - // Check if this is a child partition where parent has bigint - let Some(parent) = parent_table else { - return false; - }; - - let Some(&parent_type) = parent_column_types.get(&(*parent, col.name)) else { - return false; - }; - - if parent_type != "int8" { - return false; - } - - // Parent has bigint, check if child has integer type - let Some(type_name) = col_type_name else { - return false; - }; - - let Some(last) = type_name.names.last() else { - return false; - }; - - let Some(NodeEnum::String(PgString { sval })) = &last.node else { - return false; - }; - - is_integer_type(sval.as_str()) - } - } - _ => {} -} - use tokio::{process::Command, task::JoinSet}; #[derive(Debug, Clone)] @@ -290,9 +205,6 @@ impl PgDump { let cleaned = Self::clean(&original); trace!("[pg_dump (clean)] {}", cleaned); - #[cfg(not(feature = "new_parser"))] - let stmts = parse(&cleaned)?.protobuf; - #[cfg(feature = "new_parser")] let stmts = pg_raw_parse::parse(&cleaned)?.into_inner(); Ok(PgDumpOutput { @@ -303,16 +215,11 @@ impl PgDump { } #[derive(Debug)] -#[cfg_attr(not(feature = "new_parser"), derive(Clone))] pub struct PgDumpOutput { - #[cfg(not(feature = "new_parser"))] - stmts: ParseResult, - #[cfg(feature = "new_parser")] stmts: Owned, original: String, } -#[cfg(feature = "new_parser")] impl Clone for PgDumpOutput { fn clone(&self) -> Self { Self { @@ -383,7 +290,6 @@ impl<'a> From for Statement<'a> { impl PgDumpOutput { /// Get integer primary key columns (columns that are part of PRIMARY KEY /// constraints and have integer types like int4, int2, serial, etc.). - #[cfg(feature = "new_parser")] pub(crate) fn integer_primary_key_columns(&self) -> HashSet> { let column_types = self.column_types(); let mut result = HashSet::new(); @@ -447,84 +353,7 @@ impl PgDumpOutput { result } - cfg_select! { - not(feature = "new_parser") => { - pub(crate) fn integer_primary_key_columns(&self) -> HashSet> { - let column_types = self.column_types(); - let mut result = HashSet::new(); - - for stmt in &self.stmts.stmts { - let Some(ref node) = stmt.stmt else { - continue; - }; - let Some(NodeEnum::AlterTableStmt(ref alter_stmt)) = node.node else { - continue; - }; - - let Some(ref relation) = alter_stmt.relation else { - continue; - }; - - for cmd in &alter_stmt.cmds { - let Some(NodeEnum::AlterTableCmd(ref cmd)) = cmd.node else { - continue; - }; - - if cmd.subtype() != AlterTableType::AtAddConstraint { - continue; - } - - let Some(ref def) = cmd.def else { - continue; - }; - - let Some(NodeEnum::Constraint(ref cons)) = def.node else { - continue; - }; - - if cons.contype() != ConstrType::ConstrPrimary { - continue; - } - - let schema = schema_name(relation); - let table_name = relation.relname.as_str(); - - for key in &cons.keys { - let Some(NodeEnum::String(PgString { sval })) = &key.node else { - continue; - }; - - let col_name = sval.as_str(); - let type_key = ColumnTypeKey { - schema, - table: table_name, - column: col_name, - }; - - let is_integer = column_types - .get(&type_key) - .map(|t| is_integer_type(t)) - .unwrap_or(false); - - if is_integer { - result.insert(Column { - name: col_name, - table: Some(table_name), - schema: Some(schema), - }); - } - } - } - } - - result - } - } - _ => {} - } - /// Get integer foreign key columns (FK columns that reference integer PKs). - #[cfg(feature = "new_parser")] pub(crate) fn integer_foreign_key_columns(&self) -> HashSet> { let integer_pks = self.integer_primary_key_columns(); let mut result = HashSet::new(); @@ -589,88 +418,7 @@ impl PgDumpOutput { result } - cfg_select! { - not(feature = "new_parser") => { - pub(crate) fn integer_foreign_key_columns(&self) -> HashSet> { - let integer_pks = self.integer_primary_key_columns(); - let mut result = HashSet::new(); - - for stmt in &self.stmts.stmts { - let Some(ref node) = stmt.stmt else { - continue; - }; - let Some(NodeEnum::AlterTableStmt(ref alter_stmt)) = node.node else { - continue; - }; - - let Some(ref fk_table) = alter_stmt.relation else { - continue; - }; - - for cmd in &alter_stmt.cmds { - let Some(NodeEnum::AlterTableCmd(ref cmd)) = cmd.node else { - continue; - }; - - if cmd.subtype() != AlterTableType::AtAddConstraint { - continue; - } - - let Some(ref def) = cmd.def else { - continue; - }; - - let Some(NodeEnum::Constraint(ref cons)) = def.node else { - continue; - }; - - if cons.contype() != ConstrType::ConstrForeign { - continue; - } - - let Some(ref pk_table) = cons.pktable else { - continue; - }; - - let pk_schema = schema_name(pk_table); - let pk_table_name = pk_table.relname.as_str(); - let fk_schema = schema_name(fk_table); - let fk_table_name = fk_table.relname.as_str(); - - for (pk_attr, fk_attr) in cons.pk_attrs.iter().zip(cons.fk_attrs.iter()) { - let ( - Some(NodeEnum::String(PgString { sval: pk_col })), - Some(NodeEnum::String(PgString { sval: fk_col })), - ) = (&pk_attr.node, &fk_attr.node) - else { - continue; - }; - - let pk_column = Column { - name: pk_col.as_str(), - table: Some(pk_table_name), - schema: Some(pk_schema), - }; - - if integer_pks.contains(&pk_column) { - result.insert(Column { - name: fk_col.as_str(), - table: Some(fk_table_name), - schema: Some(fk_schema), - }); - } - } - } - } - - result - } - } - _ => {} - } - /// Get partitioned parent tables (tables with PARTITION BY). - #[cfg(feature = "new_parser")] fn partitioned_tables(&self) -> HashSet> { let mut result = HashSet::new(); @@ -690,36 +438,8 @@ impl PgDumpOutput { result } - cfg_select! { - not(feature = "new_parser") => { - fn partitioned_tables(&self) -> HashSet> { - let mut result = HashSet::new(); - - for stmt in &self.stmts.stmts { - let Some(ref node) = stmt.stmt else { - continue; - }; - let Some(NodeEnum::CreateStmt(ref create_stmt)) = node.node else { - continue; - }; - - // Tables with partspec are partitioned parent tables - if create_stmt.partspec.is_some() - && let Some(ref relation) = create_stmt.relation - { - result.insert(Table::from(relation)); - } - } - - result - } - } - _ => {} - } - /// Get parent-child relationships from ATTACH PARTITION statements. /// Returns a map from child table to parent table. - #[cfg(feature = "new_parser")] fn partition_parents(&self) -> HashMap, Table<'_>> { let mut result = HashMap::new(); @@ -758,59 +478,8 @@ impl PgDumpOutput { result } - cfg_select! { - not(feature = "new_parser") => { - fn partition_parents(&self) -> HashMap, Table<'_>> { - let mut result = HashMap::new(); - - for stmt in &self.stmts.stmts { - let Some(ref node) = stmt.stmt else { - continue; - }; - let Some(NodeEnum::AlterTableStmt(ref alter_stmt)) = node.node else { - continue; - }; - - let Some(ref parent_relation) = alter_stmt.relation else { - continue; - }; - - for cmd in &alter_stmt.cmds { - let Some(NodeEnum::AlterTableCmd(ref cmd)) = cmd.node else { - continue; - }; - - if cmd.subtype() != AlterTableType::AtAttachPartition { - continue; - } - - let Some(ref def) = cmd.def else { - continue; - }; - - let Some(NodeEnum::PartitionCmd(ref partition_cmd)) = def.node else { - continue; - }; - - let Some(ref child_relation) = partition_cmd.name else { - continue; - }; - - let parent = Table::from(parent_relation); - let child = Table::from(child_relation); - result.insert(child, parent); - } - } - - result - } - } - _ => {} - } - /// Get column types for partitioned parent tables after bigint conversion. /// Returns a map from (parent_table, column_name) to the converted type. - #[cfg(feature = "new_parser")] fn partitioned_parent_column_types( &self, columns_to_convert: &HashSet>, @@ -866,70 +535,7 @@ impl PgDumpOutput { result } - cfg_select! { - not(feature = "new_parser") => { - fn partitioned_parent_column_types( - &self, - columns_to_convert: &HashSet>, - ) -> HashMap<(Table<'_>, &str), &'static str> { - let mut result = HashMap::new(); - let partitioned = self.partitioned_tables(); - - for stmt in &self.stmts.stmts { - let Some(ref node) = stmt.stmt else { - continue; - }; - let Some(NodeEnum::CreateStmt(ref create_stmt)) = node.node else { - continue; - }; - - let Some(ref relation) = create_stmt.relation else { - continue; - }; - - let table = Table::from(relation); - if !partitioned.contains(&table) { - continue; - } - - let schema = table.schema().map(|s| s.name).unwrap_or("public"); - let table_name = table.name; - - for elt in &create_stmt.table_elts { - if let Some(NodeEnum::ColumnDef(ref col_def)) = elt.node { - let col = Column { - name: col_def.colname.as_str(), - table: Some(table_name), - schema: Some(schema), - }; - - // Check if this column needs conversion - if columns_to_convert.contains(&col) { - result.insert((table, col_def.colname.as_str()), "int8"); - } else if let Some(ref type_name) = col_def.type_name - && let Some(last_name) = type_name.names.last() - && let Some(NodeEnum::String(PgString { sval })) = &last_name.node - { - // Store original type for non-converted columns - let type_str: &'static str = match sval.as_str() { - "int4" => "int4", - "int8" => "int8", - _ => continue, // Only track integer types - }; - result.insert((table, col_def.colname.as_str()), type_str); - } - } - } - } - - result - } - } - _ => {} - } - /// Get all column types from CREATE TABLE statements. - #[cfg(feature = "new_parser")] fn column_types(&self) -> HashMap, &str> { let mut result = HashMap::new(); @@ -968,53 +574,8 @@ impl PgDumpOutput { result } - cfg_select! { - not(feature = "new_parser") => { - fn column_types(&self) -> HashMap, &str> { - let mut result = HashMap::new(); - - for stmt in &self.stmts.stmts { - let Some(ref node) = stmt.stmt else { - continue; - }; - let Some(NodeEnum::CreateStmt(ref create_stmt)) = node.node else { - continue; - }; - - let Some(ref relation) = create_stmt.relation else { - continue; - }; - - let schema = schema_name(relation); - let table_name = relation.relname.as_str(); - - for elt in &create_stmt.table_elts { - if let Some(NodeEnum::ColumnDef(col_def)) = &elt.node - && let Some(ref type_name) = col_def.type_name - && let Some(last_name) = type_name.names.last() - && let Some(NodeEnum::String(PgString { sval })) = &last_name.node - { - result.insert( - ColumnTypeKey { - schema, - table: table_name, - column: col_def.colname.as_str(), - }, - sval.as_str(), - ); - } - } - } - - result - } - } - _ => {} - } - /// Get schema statements to execute before data sync, /// e.g., CREATE TABLE, primary key. - #[cfg(feature = "new_parser")] pub(crate) fn statements(&self, state: SyncState) -> Result>, Error> { let mut result = vec![]; @@ -1394,378 +955,6 @@ impl PgDumpOutput { Ok(result) } - cfg_select! { - not(feature = "new_parser") => { - pub(crate) fn statements(&self, state: SyncState) -> Result>, Error> { - let mut result = vec![]; - - // Get integer PK and FK columns that need bigint conversion - let columns_to_convert: HashSet> = self - .integer_primary_key_columns() - .union(&self.integer_foreign_key_columns()) - .copied() - .collect(); - - // Get partitioned parent column types and parent-child relationships - let parent_column_types = self.partitioned_parent_column_types(&columns_to_convert); - let partition_parents = self.partition_parents(); - - for stmt in &self.stmts.stmts { - let (_, original_start) = self - .original - .split_at_checked(stmt.stmt_location as usize) - .ok_or(Error::StmtOutOfBounds)?; - let (original, _) = original_start - .split_at_checked(stmt.stmt_len as usize) - .ok_or(Error::StmtOutOfBounds)?; - - if let Some(ref node) = stmt.stmt - && let Some(ref node) = node.node - { - match node { - NodeEnum::CreateStmt(create_stmt) => { - let mut stmt = create_stmt.clone(); - stmt.if_not_exists = true; - - // Get table info - let table = create_stmt - .relation - .as_ref() - .map(Table::from) - .unwrap_or_default(); - let schema = table.schema().map(|s| s.name).unwrap_or("public"); - let table_name = table.name; - - // Check if this table is a child partition - let parent_table = partition_parents.get(&table); - - // Convert integer PK/FK columns to bigint - for elt in &mut stmt.table_elts { - if let Some(NodeEnum::ColumnDef(ref mut col_def)) = elt.node { - let col = Column { - name: col_def.colname.as_str(), - table: Some(table_name), - schema: Some(schema), - }; - - if should_convert_to_bigint( - &col, - col_def.type_name.as_ref(), - &columns_to_convert, - parent_table, - &parent_column_types, - ) && let Some(ref mut type_name) = col_def.type_name - { - type_name.names = vec![ - Node { - node: Some(NodeEnum::String(PgString { - sval: "pg_catalog".to_owned(), - })), - }, - Node { - node: Some(NodeEnum::String(PgString { - sval: "int8".to_owned(), - })), - }, - ]; - } - } - } - - if state == SyncState::PreData { - let sql = deparse_node(NodeEnum::CreateStmt(stmt))?; - result.push(Statement::Table { table, sql }); - } - } - - 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(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()); - } - } - - NodeEnum::AlterTableStmt(stmt) => { - for cmd in &stmt.cmds { - if let Some(NodeEnum::AlterTableCmd(ref cmd)) = cmd.node { - match cmd.subtype() { - AlterTableType::AtAddConstraint => { - if let Some(ref def) = cmd.def - && let Some(NodeEnum::Constraint(ref cons)) = def.node - { - // Only allow primary key constraints. - if matches!( - cons.contype(), - ConstrType::ConstrPrimary - | ConstrType::ConstrNotnull - | ConstrType::ConstrNull - ) { - // Integer PKs are already tracked and converted - // to bigint in CreateStmt handler - if state == SyncState::PreData { - result.push(Statement::Other { - sql: original.to_string(), - idempotent: false, - }); - } - } else if cons.contype() == ConstrType::ConstrForeign { - // FK columns referencing integer PKs are - // computed from fk_columns at the end - if state == SyncState::PostData { - result.push(Statement::Other { - sql: original.to_string(), - idempotent: false, - }); - } - } else if state == SyncState::PostData { - result.push(Statement::Other { - sql: original.to_string(), - idempotent: false, - }); - } - } - } - AlterTableType::AtAttachPartition => { - match stmt.objtype() { - // Index partitions need to be attached to indexes, - // which we create in the post-data step. - ObjectType::ObjectIndex => { - if state == SyncState::PostData { - result.push(Statement::Other { - sql: original.to_string(), - idempotent: false, - }); - } - } - - // Table partitions are attached in pre-data - // after the partition tables are created. - ObjectType::ObjectTable => { - if state == SyncState::PreData { - 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::AtAddIdentity => (), - // AlterTableType::AtChangeOwner => { - // continue; // Don't change owners, for now. - // } - _ => { - if state == SyncState::PreData { - result.push(original.into()); - } - } - } - } - } - } - - NodeEnum::CreateTrigStmt(stmt) => { - let mut stmt = stmt.clone(); - stmt.replace = true; - - if state == SyncState::PreData { - result.push(deparse_node(NodeEnum::CreateTrigStmt(stmt))?.into()); - } - } - - NodeEnum::CreatePublicationStmt(stmt) => { - if state == SyncState::PreData { - // DROP first for idempotency - result.push(Statement::Other { - sql: format!( - "DROP PUBLICATION IF EXISTS \"{}\"", - crate::util::escape_identifier(&stmt.pubname) - ), - idempotent: true, - }); - result.push(Statement::Other { - sql: original.to_string(), - idempotent: false, - }); - } - } - - NodeEnum::AlterPublicationStmt(_) => { - if state == SyncState::PreData { - result.push(Statement::Other { - sql: original.to_string(), - idempotent: false, - }); - } - } - - // Skip these. - NodeEnum::CreateSubscriptionStmt(_) | NodeEnum::AlterSubscriptionStmt(_) => (), - - NodeEnum::AlterSeqStmt(stmt) => { - if matches!(state, SyncState::PreData | SyncState::Cutover) { - let sequence = stmt - .sequence - .as_ref() - .map(Table::from) - .ok_or(Error::MissingEntity)?; - let sequence = Sequence::from(sequence); - let column = stmt.options.first().ok_or(Error::MissingEntity)?; - let column = - Column::try_from(column).map_err(|_| Error::MissingEntity)?; - - if state == SyncState::PreData { - result.push(Statement::SequenceOwner { sql: original }); - } else if state == SyncState::Cutover { - let sql = sequence - .setval_from_column(&column) - .map_err(|_| Error::MissingEntity)?; - result.push(Statement::SequenceSetMax { sql }) - } - } - } - - 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(); - - let index_schema = - stmt.relation.as_ref().map(schema_name).unwrap_or("public"); - result.push(Statement::Other { - sql: format!( - "DROP INDEX IF EXISTS \"{}\".\"{}\"", - index_schema, stmt.idxname - ), - idempotent: true, - }); - - result.push(Statement::Index { table, 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::PreData { - result.push(original.into()); - } - } - } - } - } - - Ok(result) - } - } - _ => {} - } - /// Create objects in destination cluster. pub async fn restore( &self, @@ -1989,9 +1178,6 @@ ALTER TABLE ONLY public.users \unrestrict nu6jB5ogH2xGMn2dB3dMyMbSZ2PsVDqB2IaWK6zZVjngeba0UrnmxMy6s63SwzR "#; - #[cfg(not(feature = "new_parser"))] - let _parse = pg_query::parse(&PgDump::clean(dump)).unwrap(); - #[cfg(feature = "new_parser")] let _parse = pg_raw_parse::parse(&PgDump::clean(dump)).unwrap(); } @@ -2158,16 +1344,10 @@ ALTER TABLE test ADD CONSTRAINT id_pkey PRIMARY KEY (id);"#, statements[0].deref(), "CREATE TABLE IF NOT EXISTS test (id bigint, value text)" ); - #[cfg(feature = "new_parser")] assert_eq!( statements[1].deref(), "ALTER TABLE test ADD CONSTRAINT id_pkey PRIMARY KEY (id)" ); - #[cfg(not(feature = "new_parser"))] - assert_eq!( - statements[1].deref(), - "\nALTER TABLE test ADD CONSTRAINT id_pkey PRIMARY KEY (id)" - ); } #[test] @@ -2193,16 +1373,10 @@ ALTER TABLE child ADD CONSTRAINT child_parent_fk FOREIGN KEY (parent_id) REFEREN statements[1].deref(), "CREATE TABLE IF NOT EXISTS child (id int, parent_id bigint)" ); - #[cfg(feature = "new_parser")] assert_eq!( statements[2].deref(), "ALTER TABLE parent ADD CONSTRAINT parent_pkey PRIMARY KEY (id)" ); - #[cfg(not(feature = "new_parser"))] - assert_eq!( - statements[2].deref(), - "\nALTER TABLE parent ADD CONSTRAINT parent_pkey PRIMARY KEY (id)" - ); } #[test] @@ -2319,23 +1493,10 @@ ALTER TABLE ONLY orders ATTACH PARTITION orders_2024 FOR VALUES FROM ('2024-01-0 ); } - #[cfg(feature = "new_parser")] fn parse(query: &str) -> PgDumpOutput { PgDumpOutput { stmts: pg_raw_parse::parse(query).unwrap().into_inner(), original: query.to_owned(), } } - - cfg_select! { - not(feature = "new_parser") => { - fn parse(query: &str) -> PgDumpOutput { - PgDumpOutput { - stmts: super::parse(query).unwrap().protobuf, - original: query.to_owned(), - } - } - } - _ => {} - } } diff --git a/pgdog/src/config/mod.rs b/pgdog/src/config/mod.rs index a7a298694..5ff23d433 100644 --- a/pgdog/src/config/mod.rs +++ b/pgdog/src/config/mod.rs @@ -128,7 +128,6 @@ fn validate_lookup_queries(config: &ConfigAndUsers) -> Result<(), Error> { Ok(()) } -#[cfg(feature = "new_parser")] fn validate_lookup_query(query: &str) -> Result<(), String> { use itertools::Itertools; use pg_raw_parse::{ @@ -162,38 +161,6 @@ fn validate_lookup_query(query: &str) -> Result<(), String> { Ok(()) } -#[cfg(not(feature = "new_parser"))] -fn validate_lookup_query(query: &str) -> Result<(), String> { - use std::collections::HashSet; - - let ast = - pg_query::parse(query).map_err(|err| format!("\"lookup_query\" is invalid: {}", err))?; - - if ast.protobuf.stmts.len() != 1 { - return Err("\"lookup_query\" must be a single statement".into()); - } - - // Best effort: `nodes()` covers the node types a lookup query - // realistically uses, but isn't guaranteed to visit every subtree. - // A parameter it misses fails loudly at runtime instead, when the - // lookup query runs with a single bound value. - let params = ast - .protobuf - .nodes() - .into_iter() - .filter_map(|(node, ..)| match node { - pg_query::NodeRef::ParamRef(param) => Some(param.number), - _ => None, - }) - .collect::>(); - - if params != HashSet::from([1]) { - return Err("\"lookup_query\" must reference exactly one parameter, \"$1\"".into()); - } - - Ok(()) -} - /// Load configuration from a list of database URLs. pub fn from_urls(urls: &[String]) -> Result { let _lock = LOCK.lock(); diff --git a/pgdog/src/frontend/client/query_engine/test/rewrite_offset.rs b/pgdog/src/frontend/client/query_engine/test/rewrite_offset.rs index 29629110a..3514e3a6c 100644 --- a/pgdog/src/frontend/client/query_engine/test/rewrite_offset.rs +++ b/pgdog/src/frontend/client/query_engine/test/rewrite_offset.rs @@ -171,16 +171,10 @@ async fn test_offset_with_unique_id_simple() { final_sql.contains("LIMIT 15"), "LIMIT should be 10+5=15: {final_sql}" ); - #[cfg(feature = "new_parser")] assert!( !final_sql.contains("OFFSET"), "SQL should not contain OFFSET: {final_sql}", ); - #[cfg(not(feature = "new_parser"))] - assert!( - final_sql.contains("OFFSET 0"), - "OFFSET should be 0: {final_sql}" - ); } #[tokio::test] diff --git a/pgdog/src/frontend/router/parser/aggregate.rs b/pgdog/src/frontend/router/parser/aggregate.rs index 6708b33b3..e985936f4 100644 --- a/pgdog/src/frontend/router/parser/aggregate.rs +++ b/pgdog/src/frontend/router/parser/aggregate.rs @@ -1,9 +1,3 @@ -#[cfg(not(feature = "new_parser"))] -use pg_query::{ - NodeEnum, - protobuf::{Integer, Node, SelectStmt, String as PgQueryString, a_const::Val}, -}; -#[cfg(feature = "new_parser")] use pg_raw_parse::{Node, nodes}; use std::fmt; @@ -68,7 +62,6 @@ pub struct Aggregate { group_by: Vec, } -#[cfg(feature = "new_parser")] fn index_of_column(stmt: &nodes::SelectStmt, qualified_column_name: &[&str]) -> Option { stmt.target_list().iter().position(|node| { let Node::ColumnRef(c) = node.val() else { @@ -83,45 +76,6 @@ fn index_of_column(stmt: &nodes::SelectStmt, qualified_column_name: &[&str]) -> }) } -#[cfg(not(feature = "new_parser"))] -fn target_list_to_index(stmt: &SelectStmt, column_names: &[&str]) -> 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() - && let Some(NodeEnum::ColumnRef(column_ref)) = node_box.node.as_ref() - { - let select_names: Vec<_> = 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.as_str()), - _ => 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: &[&str], select_names: &[&str]) -> bool { if group_by_names == select_names { return true; @@ -140,7 +94,6 @@ fn columns_match(group_by_names: &[&str], select_names: &[&str]) -> bool { impl Aggregate { /// Figure out what aggregates are present and which ones PgDog supports. - #[cfg(feature = "new_parser")] pub(crate) fn parse(stmt: &nodes::SelectStmt, schema: &Schema) -> Self { let group_by = stmt .group_clause() @@ -204,81 +157,6 @@ impl Aggregate { Self { group_by, targets } } - #[cfg(not(feature = "new_parser"))] - pub fn parse(stmt: &SelectStmt, schema: &Schema) -> Self { - let mut targets = vec![]; - let group_by = stmt - .group_clause - .iter() - .filter_map(|node| { - node.node.as_ref().map(|node| match node { - NodeEnum::AConst(aconst) => aconst.val.as_ref().map(|val| match val { - 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<_> = column_ref - .fields - .iter() - .filter_map(|node| match node { - Node { - node: - Some(NodeEnum::String(PgQueryString { sval: column_name })), - } => Some(column_name.as_str()), - _ => None, - }) - .collect(); - Some(target_list_to_index(stmt, &column_names)) - } - _ => None, - }) - }) - .flatten() - .flatten() - .collect::>(); - - for (idx, node) in stmt.target_list.iter().enumerate() { - if let Some(NodeEnum::ResTarget(res)) = &node.node - && let Some(node) = &res.val - && let Ok(func) = Function::try_from(node.as_ref()) - { - let function = match func.name { - "count" => Some(AggregateFunction::Count), - "max" => Some(AggregateFunction::Max), - "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), - fname => { - if schema.aggregate_functions.contains(fname) { - Some(AggregateFunction::Unrecognized(fname.to_owned())) - } else { - None - } - } - }; - - if let Some(function) = function { - let distinct = match node.node.as_ref() { - Some(NodeEnum::FuncCall(func)) => func.agg_distinct, - _ => false, - }; - - targets.push(AggregateTarget { - column: idx, - function, - distinct, - }); - } - } - } - - Self { targets, group_by } - } - pub fn targets(&self) -> &[AggregateTarget] { &self.targets } @@ -321,10 +199,8 @@ impl Aggregate { #[cfg(test)] mod test { use super::*; - #[cfg(feature = "new_parser")] use pg_raw_parse::{Owned, make}; - #[cfg(feature = "new_parser")] fn select(stmt: &str) -> Owned { match pg_raw_parse::parse(stmt).unwrap().stmts().next().unwrap() { Node::SelectStmt(stmt) => make::owned(|mem| mem.make_unique(stmt)), @@ -332,21 +208,6 @@ mod test { } } - #[cfg(not(feature = "new_parser"))] - fn select(stmt: &str) -> SelectStmt { - let stmt = pg_query::parse(stmt) - .unwrap() - .protobuf - .stmts - .remove(0) - .stmt - .unwrap(); - match stmt.node.unwrap() { - NodeEnum::SelectStmt(stmt) => *stmt, - _ => panic!("not a select"), - } - } - fn parse(stmt: &str) -> Aggregate { Aggregate::parse(&select(stmt), &Default::default()) } diff --git a/pgdog/src/frontend/router/parser/cache/ast.rs b/pgdog/src/frontend/router/parser/cache/ast.rs index 52b75faa9..e6308c0a1 100644 --- a/pgdog/src/frontend/router/parser/cache/ast.rs +++ b/pgdog/src/frontend/router/parser/cache/ast.rs @@ -1,6 +1,3 @@ -#[cfg(not(feature = "new_parser"))] -use pg_query::{NodeEnum, ParseResult, parse, parse_raw}; -#[cfg(feature = "new_parser")] use pg_raw_parse::{Node, Owned, StmtList, make}; use pgdog_config::QueryParserEngine; use std::fmt::Debug; @@ -40,11 +37,7 @@ pub struct Ast { #[derive(Debug)] pub struct AstInner { /// Cached AST. - // FIXME(sage): Rename to ast when parser port is done - #[cfg(feature = "new_parser")] pub(crate) ast: Owned, - #[cfg(not(feature = "new_parser"))] - pub(crate) ast: ParseResult, /// AST stats. pub stats: Mutex, /// Rewrite plan. @@ -55,7 +48,6 @@ pub struct AstInner { impl AstInner { /// Create new AST record, with no rewrite or comment routing. - #[cfg(feature = "new_parser")] pub(crate) fn new(ast: Owned) -> Self { Self { ast, @@ -64,16 +56,6 @@ impl AstInner { query_without_comment: "".into(), } } - - #[cfg(not(feature = "new_parser"))] - pub(crate) fn old(ast: ParseResult) -> Self { - Self { - ast, - stats: Mutex::new(Stats::new()), - rewrite_plan: RewritePlan::default(), - query_without_comment: "".into(), - } - } } impl Deref for Ast { @@ -95,13 +77,6 @@ impl Ast { search_path: Option<&ParameterValue>, ) -> Result { let now = Instant::now(); - #[cfg(not(feature = "new_parser"))] - let mut ast = match schema.query_parser_engine { - QueryParserEngine::PgQueryProtobuf => parse(query.query_without_comment), - QueryParserEngine::PgQueryRaw => parse_raw(query.query_without_comment), - } - .map_err(Error::PgQuery)?; - #[cfg(feature = "new_parser")] let ast = pg_raw_parse::parse(query.query_without_comment).map_err(Error::Parse)?; // Run the rewrite unconditionally. Even when a shard comment will @@ -109,8 +84,6 @@ impl Ast { // same query body (without the comment) would require a rewrite, so // `Cache::query` can decide whether this entry is safe to cache. let mut rewriter = StatementRewrite::new(StatementRewriteContext { - #[cfg(not(feature = "new_parser"))] - stmt: &mut ast.protobuf, extended: query.original_query.extended(), prepared: query.original_query.prepared(), prepared_statements, @@ -119,9 +92,7 @@ impl Ast { user, search_path, }); - #[cfg(feature = "new_parser")] let mut rewrite_plan = Default::default(); - #[cfg(feature = "new_parser")] let ast = make::try_owned(|mem| { // FIXME(sage): We should have a parse function on mem so we don't // need to parse and then copy the parsed tree just to throw the @@ -133,9 +104,6 @@ impl Ast { Ok::<_, Error>(copy) })?; - #[cfg(not(feature = "new_parser"))] - let rewrite_plan = rewriter.maybe_rewrite()?; - let elapsed = now.elapsed(); let mut stats = Stats::new(); stats.parse_time += elapsed; @@ -157,9 +125,6 @@ impl Ast { query_parser_engine: schema.query_parser_engine, inner: Arc::new(AstInner { stats: Mutex::new(stats), - #[cfg(feature = "new_parser")] - ast, - #[cfg(not(feature = "new_parser"))] ast, rewrite_plan, query_without_comment: query.query_without_comment.into(), @@ -184,7 +149,6 @@ impl Ast { } /// Record new AST entry, without rewriting or comment-routing. - #[cfg(feature = "new_parser")] pub(crate) fn new_record( query: &str, query_parser_engine: QueryParserEngine, @@ -200,29 +164,7 @@ impl Ast { }) } - cfg_select! { - not(feature = "new_parser") => { - 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, - comment_role: None, - comment_shard: None, - query_parser_engine, - inner: Arc::new(AstInner::old(ast)), - }) - } - } - _ => {} - } - /// Create new AST from a parse result. - #[cfg(feature = "new_parser")] pub fn from_raw_stmts(stmts: Owned) -> Self { Self { cached: true, @@ -233,24 +175,6 @@ impl Ast { } } - /// Create new AST from a parse result. - #[cfg(not(feature = "new_parser"))] - pub(crate) fn from_parse_result(parse_result: ParseResult) -> Self { - Self { - cached: true, - comment_role: None, - comment_shard: None, - query_parser_engine: QueryParserEngine::default(), - inner: Arc::new(AstInner::old(parse_result)), - } - } - - /// Get the reference to the AST. - #[cfg(not(feature = "new_parser"))] - pub(crate) fn parse_result(&self) -> &ParseResult { - &self.ast - } - /// Update stats for this statement, given the route /// calculated by the query parser. pub fn update_stats(&self, route: &Route) { @@ -264,7 +188,6 @@ impl Ast { } /// Get statement type. - #[cfg(feature = "new_parser")] pub(crate) fn statement_type(&self) -> StatementType { let root = self.ast.stmts().next(); @@ -288,41 +211,6 @@ impl Ast { _ => StatementType::Ddl, } } - - cfg_select! { - not(feature = "new_parser") => { - pub(crate) fn statement_type(&self) -> StatementType { - let root = self - .ast - .protobuf - .stmts - .first() - .and_then(|s| s.stmt.as_ref()) - .and_then(|s| s.node.as_ref()); - - match root { - Some(NodeEnum::SelectStmt(_)) - | Some(NodeEnum::InsertStmt(_)) - | Some(NodeEnum::UpdateStmt(_)) - | Some(NodeEnum::DeleteStmt(_)) - | Some(NodeEnum::CopyStmt(_)) - | Some(NodeEnum::ExplainStmt(_)) - | Some(NodeEnum::TransactionStmt(_)) => StatementType::Dml, - - Some(NodeEnum::VariableSetStmt(_)) - | Some(NodeEnum::VariableShowStmt(_)) - | Some(NodeEnum::DeallocateStmt(_)) - | Some(NodeEnum::ListenStmt(_)) - | Some(NodeEnum::NotifyStmt(_)) - | Some(NodeEnum::UnlistenStmt(_)) - | Some(NodeEnum::DiscardStmt(_)) => StatementType::Session, - - _ => StatementType::Ddl, - } - } - } - _ => {} - } } #[derive(Debug, Clone, Copy, PartialEq)] diff --git a/pgdog/src/frontend/router/parser/cache/cache_impl.rs b/pgdog/src/frontend/router/parser/cache/cache_impl.rs index 107a395e0..832a233c0 100644 --- a/pgdog/src/frontend/router/parser/cache/cache_impl.rs +++ b/pgdog/src/frontend/router/parser/cache/cache_impl.rs @@ -1,8 +1,5 @@ use lru::LruCache; use once_cell::sync::Lazy; -#[cfg(not(feature = "new_parser"))] -use pg_query::normalize; -#[cfg(feature = "new_parser")] use pg_raw_parse::normalize::normalize; use pgdog_config::QueryParserEngine; use std::collections::HashMap; @@ -100,7 +97,7 @@ impl Cache { } /// Parse a statement by either getting it from cache - /// or using pg_query parser. + /// or parsing it. /// /// 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 diff --git a/pgdog/src/frontend/router/parser/cache/test.rs b/pgdog/src/frontend/router/parser/cache/test.rs index 7025b8b9e..fda1cb071 100644 --- a/pgdog/src/frontend/router/parser/cache/test.rs +++ b/pgdog/src/frontend/router/parser/cache/test.rs @@ -1,8 +1,5 @@ #![allow(clippy::print_stdout)] -#[cfg(not(feature = "new_parser"))] -use pg_query::parse; -#[cfg(feature = "new_parser")] use pg_raw_parse::parse; use tokio::spawn; diff --git a/pgdog/src/frontend/router/parser/column.rs b/pgdog/src/frontend/router/parser/column.rs index f8785cf93..277c2403e 100644 --- a/pgdog/src/frontend/router/parser/column.rs +++ b/pgdog/src/frontend/router/parser/column.rs @@ -1,8 +1,5 @@ //! Column name reference. -#[cfg(not(feature = "new_parser"))] -use pg_query::{Node, NodeEnum, protobuf::String as PgQueryString}; -#[cfg(feature = "new_parser")] use pg_raw_parse::{list, nodes}; use std::fmt::{Display, Formatter, Result as FmtResult}; @@ -29,18 +26,6 @@ impl<'a> Column<'a> { }) } - #[cfg(not(feature = "new_parser"))] - pub(crate) fn from_string(string: &'a Node) -> Result { - match &string.node { - Some(NodeEnum::String(PgQueryString { sval })) => Ok(Self { - name: sval.as_str(), - ..Default::default() - }), - - _ => Err(()), - } - } - /// Fully-qualify this column with a table. pub(crate) fn qualify(&mut self, table: Table<'a>) { if self.table.is_none() { @@ -77,7 +62,6 @@ impl<'a> Display for Column<'a> { } } -#[cfg(feature = "new_parser")] impl<'a> TryFrom> for Column<'a> { type Error = Error; @@ -93,7 +77,6 @@ impl<'a> TryFrom> for Column<'a> { } } -#[cfg(feature = "new_parser")] impl<'a> TryFrom<&'a nodes::ColumnRef> for Column<'a> { type Error = Error; @@ -102,7 +85,6 @@ impl<'a> TryFrom<&'a nodes::ColumnRef> for Column<'a> { } } -#[cfg(feature = "new_parser")] impl<'a> TryFrom<&'a nodes::ResTarget> for Column<'a> { type Error = Error; @@ -114,7 +96,6 @@ impl<'a> TryFrom<&'a nodes::ResTarget> for Column<'a> { } } -#[cfg(feature = "new_parser")] impl<'a> TryFrom<&'a list::NodeList> for Column<'a> { type Error = Error; @@ -138,115 +119,6 @@ impl<'a> TryFrom<&'a list::NodeList> for Column<'a> { } } -#[cfg(not(feature = "new_parser"))] -impl<'a> TryFrom<&'a Node> for Column<'a> { - type Error = Error; - - fn try_from(value: &'a Node) -> Result { - Column::try_from(&value.node) - } -} - -#[cfg(not(feature = "new_parser"))] -impl<'a> TryFrom<&'a Option> for Column<'a> { - type Error = Error; - - fn try_from(value: &'a Option) -> Result { - fn from_node(node: &Node) -> Option<&str> { - if let Some(NodeEnum::String(PgQueryString { sval })) = &node.node { - Some(sval.as_str()) - } else { - None - } - } - - 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(Error::ColumnDecode)?; - - Ok(Column { - schema, - table, - name, - }) - } - - 2 => { - let table = nodes.first().and_then(from_node); - let name = nodes - .get(1) - .and_then(from_node) - .ok_or(Error::ColumnDecode)?; - - Ok(Column { - schema: None, - table, - name, - }) - } - - 1 => { - let name = nodes - .first() - .and_then(from_node) - .ok_or(Error::ColumnDecode)?; - - Ok(Column { - name, - ..Default::default() - }) - } - - _ => Err(Error::ColumnDecode), - } - } - - match value { - Some(NodeEnum::ResTarget(res_target)) => Ok(Self { - name: res_target.name.as_str(), - ..Default::default() - }), - - Some(NodeEnum::List(list)) => from_slice(&list.items), - - Some(NodeEnum::ColumnRef(column_ref)) => from_slice(&column_ref.fields), - - Some(NodeEnum::DefElem(list)) => { - if list.defname == "owned_by" { - if let Some(ref node) = list.arg { - Ok(Column::try_from(&node.node)?) - } else { - Err(Error::ColumnDecode) - } - } else { - Err(Error::ColumnDecode) - } - } - - _ => Err(Error::ColumnDecode), - } - } -} - -#[cfg(not(feature = "new_parser"))] -impl<'a> TryFrom<&Option<&'a Node>> for Column<'a> { - type Error = Error; - - fn try_from(value: &Option<&'a Node>) -> Result { - if let Some(value) = value { - (*value).try_into() - } else { - Err(Error::ColumnDecode) - } - } -} - impl<'a> From<&'a str> for Column<'a> { fn from(value: &'a str) -> Self { Column { @@ -259,17 +131,12 @@ impl<'a> From<&'a str> for Column<'a> { #[cfg(test)] mod test { - #[cfg(feature = "new_parser")] use itertools::*; - #[cfg(not(feature = "new_parser"))] - use pg_query::{NodeEnum, parse}; - #[cfg(feature = "new_parser")] use pg_raw_parse::*; use super::Column; #[test] - #[cfg(feature = "new_parser")] fn test_column() { let result = parse("INSERT INTO my_table (id, email) VALUES (1, 'test@test.com')").unwrap(); let stmt = result.stmts().exactly_one().ok().unwrap(); @@ -290,39 +157,6 @@ mod test { } #[test] - #[cfg(not(feature = "new_parser"))] - fn test_column() { - let query = parse("INSERT INTO my_table (id, email) VALUES (1, 'test@test.com')").unwrap(); - let select = query.protobuf.stmts.first().unwrap().stmt.as_ref().unwrap(); - match select.node { - Some(NodeEnum::InsertStmt(ref insert)) => { - let columns = insert - .cols - .iter() - .map(Column::try_from) - .collect::, _>>() - .unwrap(); - assert_eq!( - columns, - vec![ - Column { - name: "id", - ..Default::default() - }, - Column { - name: "email", - ..Default::default() - } - ] - ); - } - - _ => panic!("not a select"), - } - } - - #[test] - #[cfg(feature = "new_parser")] fn test_column_sequence() { let result = parse("ALTER SEQUENCE public.user_profiles_id_seq OWNED BY public.user_profiles.id") @@ -343,26 +177,4 @@ mod test { } ); } - - #[test] - #[cfg(not(feature = "new_parser"))] - fn test_column_sequence() { - let query = - parse("ALTER SEQUENCE public.user_profiles_id_seq OWNED BY public.user_profiles.id") - .unwrap(); - let alter = query.protobuf.stmts.first().unwrap().stmt.as_ref().unwrap(); - match alter.node { - Some(NodeEnum::AlterSeqStmt(ref stmt)) => { - if let Some(node) = stmt.options.first() { - let column = Column::try_from(node).unwrap(); - assert_eq!(column.name, "id"); - assert_eq!(column.schema, Some("public")); - assert_eq!(column.table, Some("user_profiles")); - } else { - panic!("no owned by clause"); - } - } - _ => panic!("not an alter sequence"), - } - } } diff --git a/pgdog/src/frontend/router/parser/context.rs b/pgdog/src/frontend/router/parser/context.rs index 3d08a5d2c..e1c0a7193 100644 --- a/pgdog/src/frontend/router/parser/context.rs +++ b/pgdog/src/frontend/router/parser/context.rs @@ -129,7 +129,6 @@ impl<'a> QueryParserContext<'a> { self.router_context.cluster.pooler_mode() == crate::config::PoolerMode::Session } - #[cfg(feature = "new_parser")] pub(super) fn is_canonicalizing_oids(&self) -> bool { self.router_context.cluster.is_canonicalizing_oids() } diff --git a/pgdog/src/frontend/router/parser/copy.rs b/pgdog/src/frontend/router/parser/copy.rs index 627032b6c..c8fdefc80 100644 --- a/pgdog/src/frontend/router/parser/copy.rs +++ b/pgdog/src/frontend/router/parser/copy.rs @@ -1,8 +1,5 @@ //! Parse COPY statement. -#[cfg(not(feature = "new_parser"))] -use pg_query::{NodeEnum, protobuf::CopyStmt}; -#[cfg(feature = "new_parser")] use pg_raw_parse::nodes; use pgdog_config::LookupResult; @@ -105,7 +102,6 @@ impl Default for CopyParser { impl CopyParser { /// Create new copy parser from a COPY statement. - #[cfg(feature = "new_parser")] pub fn new(stmt: &nodes::CopyStmt, cluster: &Cluster) -> Result { let mut parser = Self { is_from: stmt.is_from, @@ -194,111 +190,6 @@ impl CopyParser { Ok(parser) } - cfg_select! { - not(feature = "new_parser") => { - pub fn new(stmt: &CopyStmt, cluster: &Cluster) -> Result { - let mut parser = Self { - is_from: stmt.is_from, - ..Default::default() - }; - - let mut format = CopyFormat::Text; - let mut null_string = "\\N".to_owned(); - - if let Some(ref rel) = stmt.relation { - let mut columns = vec![]; - - for column in &stmt.attlist { - if let Ok(column) = Column::from_string(column) { - columns.push(column); - } - } - - 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; - } - - parser.columns = columns.len(); - - for option in &stmt.options { - if let Some(NodeEnum::DefElem(ref elem)) = option.node { - match elem.defname.to_lowercase().as_str() { - "format" => { - if let Some(ref arg) = elem.arg - && let Some(NodeEnum::String(ref string)) = arg.node - { - match string.sval.to_lowercase().as_str() { - "binary" => { - parser.headers = true; - format = CopyFormat::Binary; - } - "csv" => { - if parser.delimiter.is_none() { - parser.delimiter = Some(','); - } - format = CopyFormat::Csv; - } - _ => (), - } - } - } - - "delimiter" => { - if let Some(ref arg) = elem.arg - && let Some(NodeEnum::String(ref string)) = arg.node - { - parser.delimiter = Some(string.sval.chars().next().unwrap_or(',')); - } - } - - "header" => { - parser.headers = true; - } - - "null" => { - if let Some(ref arg) = elem.arg - && let Some(NodeEnum::String(ref string)) = arg.node - { - null_string = string.sval.clone(); - } - } - - _ => (), - } - } - } - } - - parser.stream = if format == CopyFormat::Binary { - CopyStream::Binary(BinaryStream::default()) - } else { - CopyStream::Text(Box::new(CsvStream::new( - parser.delimiter(), - parser.headers, - format, - &null_string, - ))) - }; - parser.sharding_schema = cluster.sharding_schema(); - parser.lookup_cluster = Some(cluster.clone()); - parser.null_string = null_string; - - Ok(parser) - } - } - _ => {} - } - #[inline] fn delimiter(&self) -> char { self.delimiter.unwrap_or('\t') @@ -497,7 +388,6 @@ impl CopyParser { #[cfg(test)] mod test { use crate::config::config; - #[cfg(feature = "new_parser")] use pg_raw_parse::{Node, Owned, make}; use super::*; @@ -839,7 +729,6 @@ mod test { assert_eq!(rows[2].shard(), &Shard::All); } - #[cfg(feature = "new_parser")] pub(super) fn parse(sql: &str) -> Owned { let stmt = pg_raw_parse::parse(sql).unwrap(); match stmt.stmts().next() { @@ -847,20 +736,6 @@ mod test { _ => panic!("not a copy"), } } - - cfg_select! { - not(feature = "new_parser") => { - pub(super) fn parse(sql: &str) -> Box { - let stmt = pg_query::parse(sql).unwrap(); - let stmt = stmt.protobuf.stmts.first().unwrap(); - match stmt.stmt.clone().unwrap().node.unwrap() { - NodeEnum::CopyStmt(copy) => copy, - _ => panic!("not a copy"), - } - } - } - _ => {} - } } #[cfg(test)] diff --git a/pgdog/src/frontend/router/parser/distinct.rs b/pgdog/src/frontend/router/parser/distinct.rs index ab0a433d3..5373e489f 100644 --- a/pgdog/src/frontend/router/parser/distinct.rs +++ b/pgdog/src/frontend/router/parser/distinct.rs @@ -1,13 +1,4 @@ -#[cfg(not(feature = "new_parser"))] -use super::Error; -#[cfg(feature = "new_parser")] use itertools::*; -#[cfg(not(feature = "new_parser"))] -use pg_query::{ - Node, NodeEnum, - protobuf::{self, AConst, ColumnRef, Integer, SelectStmt, a_const::Val}, -}; -#[cfg(feature = "new_parser")] use pg_raw_parse::{Node, nodes}; #[derive(Debug, PartialEq, Clone)] @@ -24,28 +15,14 @@ pub(crate) enum DistinctBy { #[derive(Debug, Clone)] pub(crate) struct Distinct<'a> { - #[cfg(not(feature = "new_parser"))] - stmt: &'a SelectStmt, - #[cfg(feature = "new_parser")] stmt: &'a nodes::SelectStmt, } impl<'a> Distinct<'a> { - #[cfg(feature = "new_parser")] pub(crate) fn new(stmt: &'a nodes::SelectStmt) -> Self { Self { stmt } } - cfg_select! { - not(feature = "new_parser") => { - pub(crate) fn new(stmt: &'a SelectStmt) -> Self { - Self { stmt } - } - } - _ => {} - } - - #[cfg(feature = "new_parser")] pub(crate) fn distinct(&self) -> Option { match self.stmt.distinct_clause().first() { Some(Node::None) => return Some(DistinctBy::Row), @@ -79,42 +56,4 @@ impl<'a> Distinct<'a> { Some(DistinctBy::Columns(columns)) } - - cfg_select! { - not(feature = "new_parser") => { - pub fn distinct(&self) -> Result, Error> { - match self.stmt.distinct_clause.first() { - Some(Node { node: None }) => return Ok(Some(DistinctBy::Row)), - None => return Ok(None), - _ => (), - } - - let mut columns = vec![]; - - for node in &self.stmt.distinct_clause { - if let Node { node: Some(node) } = node { - match node { - NodeEnum::AConst(AConst { - val: Some(Val::Ival(Integer { ival })), - .. - }) => columns.push(DistinctColumn::Index(*ival as usize - 1)), - NodeEnum::ColumnRef(ColumnRef { fields, .. }) => { - if let Some(Node { - node: Some(NodeEnum::String(protobuf::String { sval })), - }) = fields.first() - { - columns.push(DistinctColumn::Name(sval.to_string())); - } - } - - _ => (), - } - } - } - - Ok(Some(DistinctBy::Columns(columns))) - } - } - _ => {} - } } diff --git a/pgdog/src/frontend/router/parser/error.rs b/pgdog/src/frontend/router/parser/error.rs index c7e0db1e7..276a1237e 100644 --- a/pgdog/src/frontend/router/parser/error.rs +++ b/pgdog/src/frontend/router/parser/error.rs @@ -7,11 +7,6 @@ use crate::frontend::router::sharding; #[derive(Debug, Error)] pub enum Error { - #[error("{0}")] - #[cfg(not(feature = "new_parser"))] - PgQuery(#[from] pg_query::Error), - - #[cfg(feature = "new_parser")] #[error("Error parsing query: {0}")] Parse(#[from] pg_raw_parse::Error), diff --git a/pgdog/src/frontend/router/parser/from_clause.rs b/pgdog/src/frontend/router/parser/from_clause.rs index ace950174..6bb70c29c 100644 --- a/pgdog/src/frontend/router/parser/from_clause.rs +++ b/pgdog/src/frontend/router/parser/from_clause.rs @@ -1,6 +1,3 @@ -#[cfg(not(feature = "new_parser"))] -use pg_query::{Node, NodeEnum}; -#[cfg(feature = "new_parser")] use pg_raw_parse::{Node, list::NodeList}; use super::*; @@ -8,53 +5,22 @@ use super::*; /// Handle FROM clause. #[derive(Copy, Clone, Debug)] pub(crate) struct FromClause<'a> { - #[cfg(feature = "new_parser")] nodes: &'a NodeList, - #[cfg(not(feature = "new_parser"))] - nodes: &'a [Node], } impl<'a> FromClause<'a> { /// Create new FROM clause parser. - #[cfg(all(feature = "new_parser", test))] + #[cfg(test)] pub(crate) fn new(nodes: &'a NodeList) -> Self { Self { nodes } } - cfg_select! { - not(feature = "new_parser") => { - pub(crate) fn new(nodes: &'a [Node]) -> Self { - Self { nodes } - } - } - _ => {} - } - /// Get actual table name from an alias specified in the FROM clause. /// If no alias is specified, the table name is returned as-is. - #[cfg(feature = "new_parser")] pub(crate) fn resolve_alias(&self, name: &str) -> Option<&'a str> { self.nodes.iter().find_map(|node| Self::resolve(name, node)) } - cfg_select! { - not(feature = "new_parser") => { - pub(crate) fn resolve_alias(&self, name: &str) -> Option<&'a str> { - for node in self.nodes { - if let Some(ref node) = node.node - && let Some(name) = Self::resolve(name, node) - { - return Some(name); - } - } - - None - } - } - _ => {} - } - - #[cfg(feature = "new_parser")] fn resolve(name: &str, node: Node<'a>) -> Option<&'a str> { match node { Node::JoinExpr(join) => { @@ -70,58 +36,11 @@ impl<'a> FromClause<'a> { } } - cfg_select! { - not(feature = "new_parser") => { - fn resolve(name: &str, node: &'a NodeEnum) -> Option<&'a str> { - match node { - NodeEnum::JoinExpr(join) => { - for arg in [&join.larg, &join.rarg].into_iter().flatten() { - if let Some(ref node) = arg.node - && let Some(name) = Self::resolve(name, node) - { - return Some(name); - } - } - } - - NodeEnum::RangeVar(range_var) => { - let table = Table::from(range_var); - if table.name_match(name) { - return Some(table.name); - } - } - - _ => (), - } - - None - } - } - _ => {} - } - /// Get table name if the FROM clause contains only one table. - #[cfg(feature = "new_parser")] pub(crate) fn table_name(&self) -> Option<&'a str> { self.nodes.first().and_then(|node| match node { Node::RangeVar(r) => Some(Table::from(r).name), _ => None, }) } - - cfg_select! { - not(feature = "new_parser") => { - pub(crate) fn table_name(&self) -> Option<&'a str> { - if let Some(node) = self.nodes.first() - && let Some(NodeEnum::RangeVar(ref range_var)) = node.node - { - let table = Table::from(range_var); - return Some(table.name); - } - - None - } - } - _ => {} - } } diff --git a/pgdog/src/frontend/router/parser/function.rs b/pgdog/src/frontend/router/parser/function.rs index c4651a95a..a48e6feb5 100644 --- a/pgdog/src/frontend/router/parser/function.rs +++ b/pgdog/src/frontend/router/parser/function.rs @@ -1,6 +1,3 @@ -#[cfg(not(feature = "new_parser"))] -use pg_query::{Node, NodeEnum, protobuf}; -#[cfg(feature = "new_parser")] use pg_raw_parse::{Node, nodes}; const WRITE_ONLY: &[&str] = &["nextval", "setval"]; @@ -39,7 +36,6 @@ impl<'a> Function<'a> { } } - #[cfg(feature = "new_parser")] pub(crate) fn extract_func_call(node: Node<'a>) -> Option<&'a nodes::FuncCall> { match node { Node::FuncCall(func) => Some(func), @@ -50,7 +46,6 @@ impl<'a> Function<'a> { } } -#[cfg(feature = "new_parser")] impl<'a> TryFrom> for Function<'a> { type Error = (); @@ -61,47 +56,13 @@ impl<'a> TryFrom> for Function<'a> { } } -#[cfg(not(feature = "new_parser"))] -impl<'a> TryFrom<&'a Node> for Function<'a> { - type Error = (); - fn try_from(value: &'a Node) -> Result { - match &value.node { - Some(NodeEnum::FuncCall(func)) => { - let strings = func.funcname.iter().filter_map(|s| match &s.node { - Some(NodeEnum::String(protobuf::String { sval })) => Some(sval.as_str()), - _ => None, - }); - Self::from_strings(strings).ok_or(()) - } - - Some(NodeEnum::TypeCast(cast)) if let Some(node) = cast.arg.as_ref() => { - Self::try_from(node.as_ref()) - } - - Some(NodeEnum::ResTarget(res)) if let Some(val) = &res.val => { - Self::try_from(val.as_ref()) - } - - Some(NodeEnum::NullTest(test)) if let Some(node) = test.arg.as_ref() => { - Self::try_from(node.as_ref()) - } - - _ => Err(()), - } - } -} - #[cfg(test)] mod test { - #[cfg(not(feature = "new_parser"))] - use pg_query::parse; - #[cfg(feature = "new_parser")] use pg_raw_parse::parse; use super::*; #[test] - #[cfg(feature = "new_parser")] fn test_function() { let query = "SELECT pg_advisory_lock(234234), pg_try_advisory_lock(23234)::bool"; funcs(query, |func| { @@ -111,28 +72,6 @@ mod test { }); } - #[test] - #[cfg(not(feature = "new_parser"))] - fn test_function() { - let ast = - parse("SELECT pg_advisory_lock(234234), pg_try_advisory_lock(23234)::bool").unwrap(); - let root = ast.protobuf.stmts.first().unwrap().stmt.as_ref().unwrap(); - - match root.node.as_ref() { - Some(NodeEnum::SelectStmt(stmt)) => { - for node in &stmt.target_list { - let func = Function::try_from(node).unwrap(); - assert!(func.name.contains("advisory_lock")); - assert!(func.schema.is_none()); - assert!(!func.behavior().cross_shard); - } - } - - _ => panic!("not a select"), - } - } - - #[cfg(feature = "new_parser")] fn funcs(query: &str, mut check: impl FnMut(Function<'_>)) { let ast = parse(query).unwrap(); let Node::SelectStmt(stmt) = ast.stmts().next().unwrap() else { @@ -145,7 +84,6 @@ mod test { } } - #[cfg(feature = "new_parser")] fn first_func(query: &str, check: impl FnOnce(Function<'_>)) { let mut check = Some(check); funcs(query, |func| { @@ -155,19 +93,6 @@ mod test { }); } - #[cfg(not(feature = "new_parser"))] - fn first_func(query: &str, check: impl FnOnce(Function<'_>) -> R) -> R { - let ast = parse(query).unwrap(); - let root = ast.protobuf.stmts.first().unwrap().stmt.as_ref().unwrap(); - match root.node.as_ref() { - Some(NodeEnum::SelectStmt(stmt)) => { - let target = stmt.target_list.first().unwrap(); - check(Function::try_from(target).unwrap()) - } - _ => panic!("not a select"), - } - } - #[test] fn test_cross_shard_function() { first_func( diff --git a/pgdog/src/frontend/router/parser/limit.rs b/pgdog/src/frontend/router/parser/limit.rs index 8b5dea7c2..e8c43ff62 100644 --- a/pgdog/src/frontend/router/parser/limit.rs +++ b/pgdog/src/frontend/router/parser/limit.rs @@ -1,10 +1,7 @@ -#[cfg(not(feature = "new_parser"))] -use pg_query::{ - Node, NodeEnum, - protobuf::{AConst, Integer, ParamRef, a_const::Val}, +use pg_raw_parse::{ + Node, + nodes::{self, SelectStmt}, }; -#[cfg(feature = "new_parser")] -use pg_raw_parse::{Node, nodes}; use super::Error; use crate::net::Bind; @@ -15,11 +12,6 @@ pub(crate) struct Limit { pub(crate) offset: Option, } -#[cfg(feature = "new_parser")] -type SelectStmt = nodes::SelectStmt; -#[cfg(not(feature = "new_parser"))] -type SelectStmt = pg_query::protobuf::SelectStmt; - #[derive(Debug, Clone)] pub(crate) struct LimitClause<'a> { stmt: &'a SelectStmt, @@ -31,7 +23,6 @@ impl<'a> LimitClause<'a> { Self { stmt, bind } } - #[cfg(feature = "new_parser")] pub(crate) fn limit_offset(&self) -> Result { Ok(Limit { limit: self.decode(self.stmt.limit_count())?, @@ -39,21 +30,6 @@ impl<'a> LimitClause<'a> { }) } - #[cfg(not(feature = "new_parser"))] - pub(crate) fn limit_offset(&self) -> Result { - let mut limit = Limit::default(); - if let Some(ref limit_count) = self.stmt.limit_count { - limit.limit = self.decode(limit_count)?; - } - - if let Some(ref limit_offset) = self.stmt.limit_offset { - limit.offset = self.decode(limit_offset)?; - } - - Ok(limit) - } - - #[cfg(feature = "new_parser")] fn decode(&self, node: Node<'_>) -> Result, Error> { use pg_raw_parse::ConstValue; @@ -86,38 +62,4 @@ impl<'a> LimitClause<'a> { _ => Ok(None), // FIXME: We should error and not silently treat as NULL } } - - #[cfg(not(feature = "new_parser"))] - fn decode(&self, node: &Node) -> Result, Error> { - match &node.node { - Some(NodeEnum::AConst(AConst { - val: Some(Val::Ival(Integer { ival })), - .. - })) => Ok(Some(*ival as usize)), - - Some(NodeEnum::ParamRef(ParamRef { number, .. })) => { - if let Some(bind) = &self.bind { - let param = bind - .parameter(*number as usize - 1)? - .ok_or(Error::MissingParameter(*number as usize))?; - - if param.is_null() { - Ok(None) - } else { - match param.bigint() { - Some(param) => Ok(Some(param as usize)), - None => Err(Error::ParameterNotInteger( - *number as usize, - param.text_debug(), - )), - } - } - } else { - Ok(None) - } - } - - _ => Ok(None), - } - } } diff --git a/pgdog/src/frontend/router/parser/mod.rs b/pgdog/src/frontend/router/parser/mod.rs index e37165447..e890795dc 100644 --- a/pgdog/src/frontend/router/parser/mod.rs +++ b/pgdog/src/frontend/router/parser/mod.rs @@ -26,7 +26,6 @@ pub mod schema; mod sequence; pub mod statement; mod table; -pub(crate) mod util; pub mod value; mod where_clause; diff --git a/pgdog/src/frontend/router/parser/multi_tenant.rs b/pgdog/src/frontend/router/parser/multi_tenant.rs index 7cac86983..40973cea4 100644 --- a/pgdog/src/frontend/router/parser/multi_tenant.rs +++ b/pgdog/src/frontend/router/parser/multi_tenant.rs @@ -1,6 +1,3 @@ -#[cfg(not(feature = "new_parser"))] -use pg_query::{NodeEnum, ParseResult}; -#[cfg(feature = "new_parser")] use pg_raw_parse::Node; use super::Error; @@ -18,10 +15,7 @@ pub struct MultiTenantCheck<'a> { user: &'a str, config: &'a MultiTenant, schema: Schema, - #[cfg(feature = "new_parser")] ast: Node<'a>, - #[cfg(not(feature = "new_parser"))] - ast: &'a ParseResult, search_path: Option<&'a ParameterValue>, } @@ -30,8 +24,7 @@ impl<'a> MultiTenantCheck<'a> { user: &'a str, config: &'a MultiTenant, schema: Schema, - #[cfg(feature = "new_parser")] ast: Node<'a>, - #[cfg(not(feature = "new_parser"))] ast: &'a ParseResult, + ast: Node<'a>, search_path: Option<&'a ParameterValue>, ) -> Self { Self { @@ -43,7 +36,6 @@ impl<'a> MultiTenantCheck<'a> { } } - #[cfg(feature = "new_parser")] pub fn run(&self) -> Result<(), Error> { match self.ast { Node::UpdateStmt(stmt) => { @@ -79,53 +71,6 @@ impl<'a> MultiTenantCheck<'a> { Ok(()) } - cfg_select! { - not(feature = "new_parser") => { - pub fn run(&self) -> Result<(), Error> { - let stmt = self - .ast - .protobuf - .stmts - .first() - .and_then(|s| s.stmt.as_ref()); - - match stmt.and_then(|n| n.node.as_ref()) { - Some(NodeEnum::UpdateStmt(stmt)) => { - 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); - self.check(table, where_clause)?; - } - } - Some(NodeEnum::SelectStmt(stmt)) => { - let table = Table::try_from(&stmt.from_clause).ok(); - - if let Some(table) = table { - let source = TablesSource::from(table); - let where_clause = WhereClause::new(&source, &stmt.where_clause); - self.check(table, where_clause)?; - } - } - Some(NodeEnum::DeleteStmt(stmt)) => { - 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); - self.check(table, where_clause)?; - } - } - - _ => (), - } - Ok(()) - } - } - _ => {} - } - fn check(&self, table: Table, where_clause: Option) -> Result<(), Error> { let search_path = SearchPath::new(self.user, self.search_path, &self.schema); let schemas = search_path.resolve(); @@ -187,7 +132,6 @@ mod tests { } #[test] - #[cfg(feature = "new_parser")] fn multi_tenant_check_passes_with_matching_filter() { let schema = schema_with_tenant_column("tenant_id"); let ast = pg_raw_parse::parse("SELECT * FROM accounts WHERE tenant_id = 1").unwrap(); @@ -200,26 +144,7 @@ mod tests { assert!(check.run().is_ok()); } - cfg_select! { - not(feature = "new_parser") => { - #[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] - #[cfg(feature = "new_parser")] fn multi_tenant_check_requires_tenant_column_in_filter() { let schema = schema_with_tenant_column("tenant_id"); let ast = pg_raw_parse::parse("SELECT * FROM accounts WHERE other_id = 1").unwrap(); @@ -236,27 +161,4 @@ mod tests { .then_some(()) .expect("should return multi-tenant id error"); } - - cfg_select! { - not(feature = "new_parser") => { - #[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/query/ddl.rs b/pgdog/src/frontend/router/parser/query/ddl.rs index 139c8876f..7515a65a2 100644 --- a/pgdog/src/frontend/router/parser/query/ddl.rs +++ b/pgdog/src/frontend/router/parser/query/ddl.rs @@ -4,8 +4,7 @@ impl QueryParser { /// Handle DDL, e.g. CREATE, DROP, ALTER, etc. pub(super) fn ddl( &mut self, - #[cfg(feature = "new_parser")] node: Node<'_>, - #[cfg(not(feature = "new_parser"))] node: &Option, + node: Node<'_>, context: &mut QueryParserContext<'_>, ) -> Result { let command = Self::shard_ddl( @@ -17,7 +16,6 @@ impl QueryParser { Ok(command) } - #[cfg(feature = "new_parser")] pub(super) fn shard_ddl( node: Node<'_>, schema: &ShardingSchema, @@ -209,214 +207,6 @@ impl QueryParser { )) } - cfg_select! { - not(feature = "new_parser") => { - pub(super) fn shard_ddl( - node: &Option, - schema: &ShardingSchema, - calculator: &mut ShardsWithPriority, - ) -> Result { - let mut shard = Shard::All; - let mut schema_changed = false; - - match node { - Some(NodeEnum::CreateStmt(stmt)) => { - schema_changed = true; - 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 - && let Some(schema) = schema.schemas.get(table.schema()) - { - shard = schema.shard().into(); - } - schema_changed = true; - } - - ObjectType::ObjectSchema => { - if let Some(PgNode { - node: Some(NodeEnum::String(string)), - }) = stmt.objects.first() - && 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)) => { - schema_changed = true; - shard = Self::shard_ddl_table(&stmt.view, schema)?.unwrap_or(Shard::All); - } - - Some(NodeEnum::CreateTableAsStmt(stmt)) => { - schema_changed = true; - 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)) => { - schema_changed = true; - 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() - && 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() - && let Some(NodeEnum::DefElem(ref elem)) = inner.node - && let Some(ref arg) = elem.arg - && 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) = pg_query::parse(stmt) - && 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 - && !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()).with_schema_changed(schema_changed), - )) - } - } - _ => {} - } - - #[cfg(feature = "new_parser")] pub(super) fn shard_ddl_table( range_var: Option<&nodes::RangeVar>, schema: &ShardingSchema, @@ -430,25 +220,6 @@ impl QueryParser { Ok(None) } - - cfg_select! { - not(feature = "new_parser") => { - 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 - && let Some(sharded_schema) = schema.schemas.get(table.schema()) - { - return Ok(Some(sharded_schema.shard().into())); - } - - Ok(None) - } - } - _ => {} - } } #[cfg(test)] @@ -476,7 +247,6 @@ mod test { } } - #[cfg(feature = "new_parser")] fn parse_stmt(query: &str) -> Command { let ast = pg_raw_parse::parse(query).unwrap(); let root = ast.stmts().next().unwrap(); @@ -484,26 +254,6 @@ mod test { QueryParser::shard_ddl(root, &test_schema(), &mut calculator).unwrap() } - cfg_select! { - not(feature = "new_parser") => { - fn parse_stmt(query: &str) -> Command { - let root = pg_query::parse(query) - .unwrap() - .protobuf - .stmts - .first() - .unwrap() - .clone() - .stmt - .unwrap() - .node; - let mut calculator = ShardsWithPriority::default(); - QueryParser::shard_ddl(&root, &test_schema(), &mut calculator).unwrap() - } - } - _ => {} - } - #[test] fn test_create_table_sharded_schema() { let command = parse_stmt("CREATE TABLE shard_0.test (id BIGINT)"); diff --git a/pgdog/src/frontend/router/parser/query/delete.rs b/pgdog/src/frontend/router/parser/query/delete.rs index a04594fe8..3dd31e12c 100644 --- a/pgdog/src/frontend/router/parser/query/delete.rs +++ b/pgdog/src/frontend/router/parser/query/delete.rs @@ -3,14 +3,10 @@ use super::*; impl QueryParser { pub(super) fn delete( &mut self, - #[cfg(not(feature = "new_parser"))] stmt: &DeleteStmt, - #[cfg(feature = "new_parser")] stmt: pg_raw_parse::Node<'_>, + stmt: pg_raw_parse::Node<'_>, context: &mut QueryParserContext, ) -> Result { - let mut parser = StatementParser::from_delete( - #[cfg(not(feature = "new_parser"))] - stmt, - #[cfg(feature = "new_parser")] + let mut parser = StatementParser::new( stmt, context.router_context.bind, &context.sharding_schema, diff --git a/pgdog/src/frontend/router/parser/query/explain.rs b/pgdog/src/frontend/router/parser/query/explain.rs index b220c628b..4ffdcbff3 100644 --- a/pgdog/src/frontend/router/parser/query/explain.rs +++ b/pgdog/src/frontend/router/parser/query/explain.rs @@ -1,9 +1,7 @@ use super::*; -#[cfg(feature = "new_parser")] use pg_raw_parse::nodes; impl QueryParser { - #[cfg(feature = "new_parser")] pub(super) fn explain( &mut self, cached_ast: &Ast, @@ -45,67 +43,6 @@ impl QueryParser { } } } - - cfg_select! { - not(feature = "new_parser") => { - 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)?; - - 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(stmt) => self.select( - cached_ast, - stmt, - context, - ), - NodeEnum::InsertStmt(stmt) => self.insert( - stmt, - context, - ), - NodeEnum::UpdateStmt(stmt) => self.update( - stmt, - context, - ), - NodeEnum::DeleteStmt(stmt) => self.delete( - stmt, - context, - ), - - _ => { - // For other statement types, route to all shards - 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) - } - } - } - } - _ => {} - } } #[cfg(test)] diff --git a/pgdog/src/frontend/router/parser/query/mod.rs b/pgdog/src/frontend/router/parser/query/mod.rs index 61adc01e7..5844c8b0c 100644 --- a/pgdog/src/frontend/router/parser/query/mod.rs +++ b/pgdog/src/frontend/router/parser/query/mod.rs @@ -1,8 +1,6 @@ //! Route queries to correct shards. use std::{collections::HashSet, ops::Deref}; -#[cfg(not(feature = "new_parser"))] -use crate::frontend::router::parser::util::{PgStr, pg_str}; use crate::{ backend::ShardingSchema, config::Role, @@ -35,16 +33,9 @@ mod show; mod transaction; mod update; -#[cfg(feature = "new_parser")] use itertools::*; use multi_tenant::MultiTenantCheck; -#[cfg(feature = "new_parser")] use pg_raw_parse::{Node, nodes}; -#[cfg(not(feature = "new_parser"))] -use pgdog_plugin::pg_query::{ - Node as PgNode, NodeEnum, - protobuf::{a_const::Val, *}, -}; use plugins::PluginOutput; use tracing::{debug, trace}; @@ -75,7 +66,6 @@ impl QueryParser { self.explain_recorder.as_mut() } - #[cfg(feature = "new_parser")] fn ensure_explain_recorder(&mut self, node: Node<'_>, context: &QueryParserContext) { if self.explain_recorder.is_some() || !context.expanded_explain() { return; @@ -86,28 +76,6 @@ impl QueryParser { } } - cfg_select! { - not(feature = "new_parser") => { - 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() - && let Some(node) = root.stmt.as_ref().and_then(|stmt| stmt.node.as_ref()) - && matches!(node, NodeEnum::ExplainStmt(_)) - { - self.explain_recorder = Some(ExplainRecorder::new()); - } - } - } - _ => {} - } - fn attach_explain(&mut self, command: &mut Command) { if let (Some(recorder), Command::Query(route)) = (self.explain_recorder.take(), command) { let summary = ExplainSummary { @@ -235,7 +203,6 @@ impl QueryParser { /// /// Returns a `Command` if successful, error otherwise. /// - #[cfg(feature = "new_parser")] fn query(&mut self, context: &mut QueryParserContext) -> Result { let parser_enabled = context.router_context.ast.is_some(); @@ -537,317 +504,7 @@ impl QueryParser { } } - cfg_select! { - not(feature = "new_parser") => { - fn query(&mut self, context: &mut QueryParserContext) -> Result { - let parser_enabled = context.router_context.ast.is_some(); - - debug!( - "parser is {}", - if parser_enabled { - "enabled" - } else { - "disabled" - } - ); - - if !parser_enabled { - // 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 statement = context - .router_context - .ast - .clone() - .ok_or(Error::EmptyQuery)?; - - self.ensure_explain_recorder(statement.parse_result(), context); - - // Parse hardcoded shard from a query comment. - if context.router_needed || context.dry_run { - let mut comment_shard_set = false; - match &statement.comment_shard { - Some(ShardOrLookup::Shard(comment_shard)) => { - context - .shards_calculator - .push(ShardWithPriority::new_comment(comment_shard.clone())); - comment_shard_set = true; - } - // The sharding key in the comment missed the lookup - // cache when the comment was parsed, which happens - // before routing. Check again: on the second routing - // pass the translation has been resolved. The pending - // lookup is cloned only when it actually has to run. - Some(ShardOrLookup::Lookup(pending)) => { - match sharding::lookup::shard_for_pending( - pending, - &context.sharding_schema, - &context.router_context.resolved_lookups, - )? { - Some(shard) => { - context - .shards_calculator - .push(ShardWithPriority::new_comment(shard)); - comment_shard_set = true; - } - None => context.bare_key_lookups.push(pending.clone()), - } - } - None => {} - } - - let role_override = statement.comment_role; - if let Some(role) = role_override { - self.write_override = role == Role::Primary; - } - - if comment_shard_set || 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); - - if let Some(multi_tenant) = context.multi_tenant() { - debug!("running multi-tenant check"); - - MultiTenantCheck::new( - context.router_context.cluster.user(), - multi_tenant, - context.router_context.cluster.schema(), - statement.parse_result(), - context.router_context.parameter_hints.search_path, - ) - .run()?; - } - - let stmts = &statement.parse_result().protobuf.stmts; - - // Handle multi-statement SET commands (e.g. "SET x TO 1; SET y TO 2"). - if stmts.len() > 1 - && let Some(command) = self.try_multi_set(stmts, context)? - { - return Ok(command); - } - - // - // Get the root AST node. - // - // We don't expect clients to send multiple queries. If they do - // only the first one is used for routing. - // - let root = 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( - context.shards_calculator.shard(), - ))); - }; - - let mut command = match root.node { - // SET statements -> return immediately. - Some(NodeEnum::VariableSetStmt(ref stmt)) => { - return self.set(stmt, context); - } - - // SELECT set_config(...) -> treat as SET and return - Some(NodeEnum::SelectStmt(ref stmt)) - if let Some(set_config) = extract_set_config(stmt) => - { - return Ok(self.set_config(set_config, context)); - } - - // SHOW statements -> return immediately. - Some(NodeEnum::VariableShowStmt(ref stmt)) => return self.show(stmt, context), - // DEALLOCATE statements -> return immediately. - Some(NodeEnum::DeallocateStmt(_)) => { - return Ok(Command::Deallocate); - } - // SELECT statements. - 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, - ), - // UPDATE statements. - Some(NodeEnum::UpdateStmt(ref stmt)) => self.update( - stmt, - context, - ), - // DELETE statements. - 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)? { - Command::Query(query) => Ok(Command::Query(query)), - command => return Ok(command), - }, - - // LISTEN ; - Some(NodeEnum::ListenStmt(ref stmt)) => { - let shard = ContextBuilder::from_string(&stmt.conditionname)? - .shards(context.shards) - .build()? - .apply()?; - - return Ok(Command::Listen { - shard, - channel: stmt.conditionname.clone(), - }); - } - - Some(NodeEnum::NotifyStmt(ref stmt)) => { - let shard = ContextBuilder::from_string(&stmt.conditionname)? - .shards(context.shards) - .build()? - .apply()?; - - return Ok(Command::Notify { - shard, - channel: stmt.conditionname.clone(), - payload: stmt.payload.clone(), - }); - } - - Some(NodeEnum::UnlistenStmt(ref stmt)) => { - return Ok(Command::Unlisten(stmt.conditionname.clone())); - } - - Some(NodeEnum::ExplainStmt(ref stmt)) => self.explain(&statement, stmt, context), - - Some(NodeEnum::DiscardStmt { .. }) => { - return Ok(Command::Discard { - extended: !context.query()?.simple(), - }); - } - - _ => self.ddl(&root.node, context), - }?; - - // e.g. Parse, Describe, Flush-style flow. - if !context.router_context.executable - && let Command::Query(ref query) = command - && 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()), - )); - } - - // Run plugins, if any. - self.plugins( - context, - &statement, - match &command { - Command::Query(query) => query.is_read(), - _ => false, - }, - )?; - - // 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(shard); - } - } - - // Set plugin-specified route, if available. - // 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(read); - } - - if let Some(ref shard) = self.plugin_output.shard { - context - .shards_calculator - .push(ShardWithPriority::new_plugin(shard.clone())); - route.set_shard(context.shards_calculator.shard()); - } - } - - // If we only have one shard, set it. - // - // If the query parser couldn't figure it out, - // there is no point of doing a multi-shard query with only one shard - // in the set. - // - if context.shards == 1 - && !context.dry_run - && let Command::Query(ref mut route) = command - { - context - .shards_calculator - .push(ShardWithPriority::new_override_only_one_shard( - Shard::Direct(0), - )); - route.set_shard(context.shards_calculator.shard()); - } - - statement.update_stats(command.route()); - - if context.dry_run { - // Record statement in cache with normalized parameters. - if !statement.cached { - let query = context.query()?.query(); - Cache::get().record_normalized( - query, - command.route(), - context.sharding_schema.query_parser_engine, - )?; - } - Ok(command.dry_run()) - } else { - Ok(command) - } - } - } - _ => {} - } - /// Handle COPY command. - #[cfg(feature = "new_parser")] fn copy(stmt: &nodes::CopyStmt, context: &mut QueryParserContext) -> Result { // Schema-based routing. // @@ -892,66 +549,16 @@ impl QueryParser { } } - cfg_select! { - not(feature = "new_parser") => { - 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 - && 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 !stmt.is_from { - context - .shards_calculator - .push(ShardWithPriority::new_table(Shard::All)); - Ok(Command::Query(Route::read( - context.shards_calculator.shard(), - ))) - } else { - Ok(Command::Copy(Box::new(parser))) - } - } - } - _ => {} - } - /// Handle INSERT statement. /// /// # Arguments /// - /// * `stmt`: INSERT statement from pg_query. + /// * `stmt`: INSERT statement. /// * `context`: Query parser context. /// fn insert( &mut self, - #[cfg(not(feature = "new_parser"))] stmt: &InsertStmt, - #[cfg(feature = "new_parser")] stmt: pg_raw_parse::Node<'_>, + stmt: pg_raw_parse::Node<'_>, context: &mut QueryParserContext, ) -> Result { let schema_lookup = SchemaLookupContext { @@ -959,10 +566,7 @@ impl QueryParser { user: context.router_context.cluster.user(), search_path: context.router_context.parameter_hints.search_path, }; - let mut parser = StatementParser::from_insert( - #[cfg(not(feature = "new_parser"))] - stmt, - #[cfg(feature = "new_parser")] + let mut parser = StatementParser::new( stmt, context.router_context.bind, &context.sharding_schema, @@ -1008,7 +612,6 @@ impl QueryParser { } } -#[cfg(feature = "new_parser")] fn extract_set_config(stmt: &nodes::SelectStmt) -> Option<&nodes::FuncCall> { static SET_CONFIG: &[&[&str]] = &[&["pg_catalog", "set_config"], &["set_config"]]; @@ -1031,35 +634,6 @@ fn extract_set_config(stmt: &nodes::SelectStmt) -> Option<&nodes::FuncCall> { }) } -cfg_select! { - not(feature = "new_parser") => { - fn extract_set_config(stmt: &SelectStmt) -> Option<&FuncCall> { - static SET_CONFIG: &[&[PgStr<'static>]] = &[ - &[pg_str("pg_catalog"), pg_str("set_config")], - &[pg_str("set_config")], - ]; - // FIXME(sage): Dear god we need some pattern macros for this - if let [ - PgNode { - node: Some(NodeEnum::ResTarget(r)), - }, - ] = &*stmt.target_list - && let ResTarget { val: Some(n), .. } = &**r - && let PgNode { - node: Some(NodeEnum::FuncCall(f)), - } = &**n - && SET_CONFIG.iter().any(|&n| n == f.funcname) - { - Some(f) - } else { - None - } - } - } - _ => {} -} - -#[cfg(feature = "new_parser")] fn references_pg_type(stmt: &nodes::SelectStmt) -> bool { use pg_raw_parse::walk::{self, Recurse}; use std::ops::ControlFlow; diff --git a/pgdog/src/frontend/router/parser/query/plugins.rs b/pgdog/src/frontend/router/parser/query/plugins.rs index b7446a2b2..4823aa030 100644 --- a/pgdog/src/frontend/router/parser/query/plugins.rs +++ b/pgdog/src/frontend/router/parser/query/plugins.rs @@ -76,9 +76,6 @@ impl QueryParser { has_primary: !context.write_only, in_transaction: context.router_context.in_transaction(), write_override: self.write_override || !read, // This is set inside `QueryParser::plugins`. - #[cfg(not(feature = "new_parser"))] - query: &statement.parse_result().protobuf, - #[cfg(feature = "new_parser")] query: &statement.ast, params, }; diff --git a/pgdog/src/frontend/router/parser/query/select.rs b/pgdog/src/frontend/router/parser/query/select.rs index b9571502c..4d54d744e 100644 --- a/pgdog/src/frontend/router/parser/query/select.rs +++ b/pgdog/src/frontend/router/parser/query/select.rs @@ -1,15 +1,7 @@ use crate::frontend::router::parser::cache::Ast; -#[cfg(not(feature = "new_parser"))] -use crate::frontend::router::parser::{FromClause, TablesSource}; use super::*; -#[cfg(not(feature = "new_parser"))] -use function::FunctionBehavior; -#[cfg(not(feature = "new_parser"))] -use pg_query::Node as PgNode; -#[cfg(feature = "new_parser")] use pg_raw_parse::walk; -#[cfg(feature = "new_parser")] use pg_raw_parse::{Node, nodes}; use pgdog_config::system_catalogs; use shared::ConvergeAlgorithm; @@ -19,10 +11,9 @@ impl QueryParser { /// /// # Arguments /// - /// * `stmt`: SELECT statement from pg_query. + /// * `stmt`: SELECT statement. /// * `context`: Query parser context. /// - #[cfg(feature = "new_parser")] pub(super) fn select( &mut self, cached_ast: &Ast, @@ -56,7 +47,7 @@ impl QueryParser { } let (advisory_locks, mut omnisharded) = { - let mut parser = StatementParser::from_select( + let mut parser = StatementParser::new( stmt.into(), context.router_context.bind, &context.sharding_schema, @@ -81,7 +72,7 @@ impl QueryParser { let mut shards = HashSet::new(); let (shard, is_sharded, tables, pending_lookups) = { - let mut statement_parser = StatementParser::from_select( + let mut statement_parser = StatementParser::new( stmt.into(), context.router_context.bind, &context.sharding_schema, @@ -262,248 +253,13 @@ impl QueryParser { )) } - cfg_select! { - not(feature = "new_parser") => { - pub(super) fn select( - &mut self, - cached_ast: &Ast, - stmt_old: &SelectStmt, - context: &mut QueryParserContext, - ) -> Result { - let cte_writes = Self::cte_writes(stmt_old); - let has_locking = Self::has_locking_clause(stmt_old); - let FunctionBehavior { - writes, - cross_shard, - } = Self::functions(stmt_old); - - // Write overwrite because of conservative read/write split. - let writes = writes || self.write_override || cte_writes || has_locking; - - if cross_shard { - context - .shards_calculator - .push(ShardWithPriority::new_override_cross_shard_function()); - } - - let (advisory_locks, mut omnisharded) = { - let mut parser = StatementParser::from_select( - stmt_old, - context.router_context.bind, - &context.sharding_schema, - None, - ); - - (parser.extract_advisory_locks(), parser.is_all_omnisharded()) - }; - - let writes = writes || !advisory_locks.is_empty(); - - // Early return for any direct-to-shard queries. - if context.shards_calculator.shard().is_direct() { - return Ok(Command::Query( - Route::read(context.shards_calculator.shard().clone()) - .with_read(!writes) - .with_omnisharded(omnisharded) - .with_advisory_locks(advisory_locks), - )); - } - - let mut shards = HashSet::new(); - - let (shard, is_sharded, tables, pending_lookups) = { - let mut statement_parser = StatementParser::from_select( - stmt_old, - context.router_context.bind, - &context.sharding_schema, - self.recorder_mut(), - ); - statement_parser - .set_resolved_lookups(&context.router_context.resolved_lookups); - - let shard = statement_parser.shard()?; - let pending_lookups = statement_parser.take_pending_lookups(); - - if shard.is_some() { - (shard, true, vec![], pending_lookups) - } else { - ( - None, - statement_parser.is_sharded( - &context.router_context.schema, - context.router_context.cluster.user(), - context.router_context.parameter_hints.search_path, - ), - statement_parser.extract_tables(), - pending_lookups, - ) - } - }; - - context.pending_lookups.extend(pending_lookups); - - if let Some(shard) = shard { - shards.insert(shard); - } - - // SELECT NOW(), SELECT 1 - if shards.is_empty() && stmt_old.from_clause.is_empty() { - let shard = Shard::Direct(round_robin::next() % context.shards); - - if let Some(recorder) = self.recorder_mut() { - recorder.record_entry(Some(shard.clone()), "SELECT omnishard no table".to_string()); - } - - context - .shards_calculator - .push(ShardWithPriority::new_rr_no_table(shard)); - - return Ok(Command::Query( - Route::read(context.shards_calculator.shard().clone()) - .with_read(!writes) - .with_omnisharded(omnisharded) - .with_advisory_locks(advisory_locks), - )); - } - - let order_by = Self::select_sort(&stmt_old.sort_clause, context.router_context.bind); - let from_clause = TablesSource::from(FromClause::new(&stmt_old.from_clause)); - - // Shard by vector in ORDER BY clause. - for order in &order_by { - if let Some((vector, column_name)) = order.vector() { - for table in context.sharding_schema.tables.tables() { - if &table.column == column_name - && (table.name.is_none() - || table.name.as_deref() == from_clause.table_name()) - { - let centroids = Centroids::from(&table.centroids); - 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 stmt = stmt_old; - let shard = Self::converge(&shards, ConvergeAlgorithm::default()); - let aggregates = Aggregate::parse(stmt, &context.router_context.schema); - let limit = LimitClause::new(stmt, context.router_context.bind).limit_offset()?; - let distinct = Distinct::new(stmt_old).distinct()?; - - if let Some(shard) = shard { - debug!("direct-to-shard {}", shard); - - context - .shards_calculator - .push(ShardWithPriority::new_table(shard)); - } else if is_sharded { - debug!("table is sharded, but no sharding key detected"); - - context - .shards_calculator - .push(ShardWithPriority::new_table(Shard::All)); - } else { - let system_catalog_sharded = - if context.sharding_schema.tables().is_system_catalog_sharded() { - { - tables - .iter() - .any(|table| system_catalogs().contains(&table.name)) - } - } else { - Default::default() - }; - - if system_catalog_sharded { - debug!("system catalog sharded"); - - context - .shards_calculator - .push(ShardWithPriority::new_table(Shard::All)); - } else { - debug!( - "table is not sharded, defaulting to omnisharded (schema loaded: {})", - context.router_context.schema.is_loaded() - ); - - // Omnisharded by default. - let sticky = tables.iter().any(|table| { - context - .sharding_schema - .tables() - .is_omnisharded_sticky(table.name) - == Some(true) - }); - - let (rr_index, explain) = if sticky - || context - .sharding_schema - .tables() - .is_omnisharded_sticky_default() - { - (context.router_context.sticky.omni_index, "sticky") - } else { - (round_robin::next(), "round robin") - }; - - let shard = Shard::Direct(rr_index % context.shards); - - // Routed to a single shard via the omnisharded-by-default path - // (non-sharded tables, including system catalogs). - omnisharded = true; - - if let Some(recorder) = self.recorder_mut() { - recorder - .record_entry(Some(shard.clone()), format!("SELECT omnishard {}", explain)); - } - - context - .shards_calculator - .push(ShardWithPriority::new_rr_omni(shard)); - } - } - - let mut query = Route::select( - context.shards_calculator.shard().clone(), - order_by, - aggregates, - limit, - distinct, - ); - - // Only rewrite if query is cross-shard. - if query.is_cross_shard() && context.shards > 1 { - query.set_rewrite_plan(cached_ast.rewrite_plan.aggregates.clone()); - } - - Ok(Command::Query( - query - .with_read(!writes) - .with_omnisharded(omnisharded) - .with_advisory_locks(advisory_locks), - )) - } - } - _ => {} - } - /// Handle the `ORDER BY` clause of a `SELECT` statement. /// /// # Arguments /// - /// * `nodes`: List of pg_query-generated nodes from the ORDER BY clause. + /// * `nodes`: List of parser-generated nodes from the ORDER BY clause. /// * `params`: Bind parameters, if any. /// - #[cfg(feature = "new_parser")] fn select_sort(stmt: &nodes::SelectStmt, params: Option<&Bind>) -> Vec { stmt.sort_clause() .into_iter() @@ -570,181 +326,4 @@ impl QueryParser { }) .collect() } - - #[cfg(not(feature = "new_parser"))] - fn select_sort(nodes: &[PgNode], params: Option<&Bind>) -> Vec { - let mut order_by = vec![]; - for clause in nodes { - if let Some(NodeEnum::SortBy(ref sort_by)) = clause.node { - let asc = matches!(sort_by.sortby_dir, 0..=2); - let Some(ref node) = sort_by.node else { - continue; - }; - let Some(ref node) = node.node else { - continue; - }; - - match node { - NodeEnum::AConst(aconst) => { - if let Some(Val::Ival(ref integer)) = aconst.val { - order_by.push(if asc { - OrderBy::Asc(integer.ival as usize) - } else { - OrderBy::Desc(integer.ival as usize) - }); - } - } - - NodeEnum::ColumnRef(column_ref) => { - // TODO: save the entire column and disambiguate - // when reading data with RowDescription as context. - let Some(field) = column_ref.fields.last() else { - continue; - }; - if let Some(NodeEnum::String(ref string)) = field.node { - order_by.push(if asc { - OrderBy::AscColumn(string.sval.clone()) - } else { - OrderBy::DescColumn(string.sval.clone()) - }); - } - } - - NodeEnum::AExpr(expr) => { - if expr.kind() == AExprKind::AexprOp - && let Some(node) = expr.name.first() - && let Some(NodeEnum::String(String { sval })) = &node.node - { - match sval.as_str() { - "<->" => { - let mut vector: Option = None; - let mut column: Option = None; - - for e in [&expr.lexpr, &expr.rexpr].iter().copied().flatten() { - if let Ok(vec) = Value::try_from(&e.node) { - match vec { - Value::Placeholder(p) => { - if let Some(bind) = params - && let Ok(Some(param)) = - bind.parameter((p - 1) as usize) - { - vector = param.vector(); - } - } - Value::Vector(vec) => vector = Some(vec), - _ => (), - } - } - - if let Ok(col) = Column::try_from(&e.node) { - column = Some(col.name.to_owned()); - } - } - - if let Some(vector) = vector - && let Some(column) = column - { - order_by.push(OrderBy::AscVectorL2Column(column, vector)); - } - } - _ => continue, - } - } - } - - _ => continue, - } - } - } - - order_by - } - - /// Handle Postgres functions that could trigger the SELECT to go to a primary. - /// - /// # Arguments - /// - /// * `stmt`: SELECT statement from pg_query. - /// - #[cfg(not(feature = "new_parser"))] - fn functions(stmt: &SelectStmt) -> FunctionBehavior { - for target in &stmt.target_list { - if let Ok(func) = Function::try_from(target) { - return func.behavior(); - } - } - - // Recurse into CTEs so a write-only function - // nested inside a WITH clause still routes to the primary. - if let Some(ref with_clause) = stmt.with_clause { - for cte in &with_clause.ctes { - if let Some(NodeEnum::CommonTableExpr(ref expr)) = cte.node - && let Some(ref query) = expr.ctequery - && let Some(NodeEnum::SelectStmt(ref inner)) = query.node - { - let behavior = Self::functions(inner); - if behavior.writes { - return behavior; - } - } - } - } - - FunctionBehavior::default() - } - - /// Recursively check for a locking clause (FOR UPDATE, FOR SHARE, etc.) - /// on this statement or any CTE nested within it. - #[cfg(not(feature = "new_parser"))] - fn has_locking_clause(stmt: &SelectStmt) -> bool { - if !stmt.locking_clause.is_empty() { - return true; - } - - if let Some(ref with_clause) = stmt.with_clause { - for cte in &with_clause.ctes { - if let Some(NodeEnum::CommonTableExpr(ref expr)) = cte.node - && let Some(ref query) = expr.ctequery - && let Some(NodeEnum::SelectStmt(ref inner)) = query.node - && Self::has_locking_clause(inner) - { - return true; - } - } - } - - false - } - - /// Check for CTEs that could trigger this query to go to a primary. - /// - /// # Arguments - /// - /// * `stmt`: SELECT statement from pg_query. - /// - #[cfg(not(feature = "new_parser"))] - fn cte_writes(stmt: &SelectStmt) -> bool { - if let Some(ref with_clause) = stmt.with_clause { - for cte in &with_clause.ctes { - if let Some(NodeEnum::CommonTableExpr(ref expr)) = cte.node - && let Some(ref query) = expr.ctequery - && let Some(ref node) = query.node - { - match node { - NodeEnum::SelectStmt(stmt) => { - if Self::cte_writes(stmt) { - return true; - } - } - - _ => { - return true; - } - } - } - } - } - - false - } } diff --git a/pgdog/src/frontend/router/parser/query/set.rs b/pgdog/src/frontend/router/parser/query/set.rs index 95ed30bea..08b652409 100644 --- a/pgdog/src/frontend/router/parser/query/set.rs +++ b/pgdog/src/frontend/router/parser/query/set.rs @@ -1,5 +1,4 @@ use super::*; -#[cfg(feature = "new_parser")] use pg_raw_parse::{Node, nodes, nodes::VariableSetKind::*}; impl QueryParser { @@ -10,7 +9,6 @@ impl QueryParser { /// /// All other SETs change the params on the client and are eventually sent to the server /// when the client is connected to the server. - #[cfg(feature = "new_parser")] pub(super) fn set( &mut self, stmt: &nodes::VariableSetStmt, @@ -34,35 +32,7 @@ impl QueryParser { } } - cfg_select! { - not(feature = "new_parser") => { - pub(super) fn set( - &mut self, - stmt: &VariableSetStmt, - context: &QueryParserContext, - ) -> Result { - if stmt.kind() == VariableSetKind::VarResetAll { - return Ok(Command::ResetAll); - } - - if let Some(param) = Self::parse_set_param(stmt)? { - return Ok(Command::Set { - params: vec![param], - route: Route::write(context.shards_calculator.shard()), - behave_like_select: false, - }); - } - - Ok(Command::Query( - Route::write(context.shards_calculator.shard().clone()).with_read(context.read_only), - )) - } - } - _ => {} - } - /// Parse a single SET statement into a SetParam - #[cfg(feature = "new_parser")] fn parse_set_param(stmt: &nodes::VariableSetStmt) -> Result { let value = if stmt.kind == VAR_SET_VALUE { Some(Self::parse_set_values(stmt)?) @@ -86,35 +56,6 @@ impl QueryParser { } } - cfg_select! { - not(feature = "new_parser") => { - fn parse_set_param(stmt: &VariableSetStmt) -> Result, Error> { - let transaction_state = stmt.name.starts_with("TRANSACTION"); - if transaction_state { - return Ok(None); - } - - let is_reset = stmt.kind() == VariableSetKind::VarReset; - let value = Self::parse_set_value(stmt)?; - - match value { - value @ Some(_) => Ok(Some(SetParam { - name: stmt.name.to_string(), - value, - local: stmt.is_local, - })), - None if is_reset => Ok(Some(SetParam { - name: stmt.name.to_string(), - value: None, - local: false, - })), - None => Ok(None), - } - } - } - _ => {} - } - /// Try to handle multi-statement queries containing SET commands. /// /// - All SETs → returns `Ok(Some(Command::Set { .. }))` @@ -123,7 +64,6 @@ impl QueryParser { /// /// In session mode, returns `Ok(Some(Command::Query(..)))` immediately so that /// all multi-statement queries are forwarded to the server verbatim. - #[cfg(feature = "new_parser")] pub(super) fn try_multi_set<'a>( &mut self, stmts: impl IntoIterator, @@ -164,65 +104,6 @@ impl QueryParser { } } - cfg_select! { - not(feature = "new_parser") => { - pub(super) fn try_multi_set( - &mut self, - stmts: &[RawStmt], - context: &QueryParserContext, - ) -> Result, Error> { - // In session mode, pass through without validation — the server - // owns the session and can handle mixed SET + other statements. - if context.is_session_mode() { - return Ok(Some(Command::Query(Route::write( - context.shards_calculator.shard(), - )))); - } - let mut params = Vec::with_capacity(stmts.len()); - let mut has_set = false; - let mut has_other = false; - - for raw_stmt in stmts { - let node = raw_stmt - .stmt - .as_ref() - .and_then(|s| s.node.as_ref()) - .ok_or(Error::EmptyQuery)?; - - match node { - NodeEnum::VariableSetStmt(stmt) => { - if stmt.name.starts_with("TRANSACTION") { - has_other = true; - } else { - has_set = true; - if let Some(param) = Self::parse_set_param(stmt)? { - params.push(param); - } - } - } - _ => has_other = true, - } - - if has_set && has_other { - return Err(Error::MultiStatementMixedSet); - } - } - - if params.is_empty() { - return Ok(None); - } - - Ok(Some(Command::Set { - params, - route: Route::write(context.shards_calculator.shard()), - behave_like_select: false, - })) - } - } - _ => {} - } - - #[cfg(feature = "new_parser")] fn parse_set_values(stmt: &nodes::VariableSetStmt) -> Result { let mut value = stmt .args() @@ -249,45 +130,4 @@ impl QueryParser { Ok(value) } - - cfg_select! { - not(feature = "new_parser") => { - fn parse_set_value(stmt: &VariableSetStmt) -> Result, Error> { - let mut value = vec![]; - - for node in &stmt.args { - match &node.node { - Some(NodeEnum::AConst(AConst { val: Some(val), .. })) => match val { - Val::Sval(String { sval }) => value.push(sval.to_string()), - Val::Ival(Integer { ival }) => value.push(ival.to_string()), - Val::Fval(Float { fval }) => value.push(fval.to_string()), - Val::Boolval(Boolean { boolval }) => value.push(boolval.to_string()), - _ => (), - }, - // e.g. SET TIME ZONE INTERVAL '+00:00' HOUR TO MINUTE - Some(NodeEnum::TypeCast(tc)) => { - if let Some(ref arg) = tc.arg - && let Some(NodeEnum::AConst(AConst { - val: Some(Val::Sval(String { ref sval })), - .. - })) = arg.node - { - value.push(sval.to_string()); - } - } - _ => (), - } - } - - let value = match value.len() { - 0 => None, - 1 => Some(ParameterValue::String(value.pop().unwrap())), - _ => Some(ParameterValue::Tuple(value)), - }; - - Ok(value) - } - } - _ => {} - } } diff --git a/pgdog/src/frontend/router/parser/query/set_config.rs b/pgdog/src/frontend/router/parser/query/set_config.rs index 3c14a433a..7f3540f4e 100644 --- a/pgdog/src/frontend/router/parser/query/set_config.rs +++ b/pgdog/src/frontend/router/parser/query/set_config.rs @@ -1,15 +1,10 @@ -#[cfg(not(feature = "new_parser"))] -use super::String as PgString; use super::*; -#[cfg(not(feature = "new_parser"))] -use std::string::String; impl QueryParser { /// Handle SELECT set_config('key', 'value', is_local) /// /// If the function arguments are a form we cannot handle, we warn and /// pass through - #[cfg(feature = "new_parser")] pub(super) fn set_config( &mut self, fcall: &nodes::FuncCall, @@ -27,29 +22,9 @@ impl QueryParser { ) } } - - cfg_select! { - not(feature = "new_parser") => { - pub(super) fn set_config(&mut self, fcall: &FuncCall, context: &QueryParserContext) -> Command { - if let Some(param) = parse_args(fcall) { - Command::Set { - params: vec![param], - route: Route::write(context.shards_calculator.shard()), - behave_like_select: true, - } - } else { - Command::Query( - Route::write(context.shards_calculator.shard()).with_read(context.read_only), - ) - } - } - } - _ => {} - } } /// Returns None if the arguments could not be parsed -#[cfg(feature = "new_parser")] fn parse_args(fcall: &nodes::FuncCall) -> Option { let name = parse_config_name(fcall.args().first()?)?; let value = parse_config_value(fcall.args().get(1)?)?; @@ -57,20 +32,7 @@ fn parse_args(fcall: &nodes::FuncCall) -> Option { Some(SetParam { name, value, local }) } -cfg_select! { - not(feature = "new_parser") => { - fn parse_args(fcall: &FuncCall) -> Option { - let name = parse_config_name(fcall.args.first()?)?; - let value = parse_config_value(fcall.args.get(1)?)?; - let local = parse_is_local(fcall.args.get(2)?)?; - Some(SetParam { name, value, local }) - } - } - _ => {} -} - /// Returns None if the name could not be parsed -#[cfg(feature = "new_parser")] fn parse_config_name(arg: Node<'_>) -> Option { match arg { Node::A_Const(c) => c.val()?.string_value().map(ToOwned::to_owned), @@ -79,22 +41,8 @@ fn parse_config_name(arg: Node<'_>) -> Option { } } -/// Returns None if the name could not be parsed -#[cfg(not(feature = "new_parser"))] -fn parse_config_name(arg: &PgNode) -> Option { - match &arg.node { - Some(NodeEnum::AConst(AConst { - val: Some(Val::Sval(PgString { sval })), - .. - })) => Some(sval.clone()), - // Only constant strings can be handled for now - _ => None, - } -} - /// Returns None if the value could not be parsed, Some(None) if the value /// is NULL, and Some if the value was successfully parsed -#[cfg(feature = "new_parser")] fn parse_config_value(arg: Node<'_>) -> Option> { match arg { Node::A_Const(c) => match c.val() { @@ -107,45 +55,10 @@ fn parse_config_value(arg: Node<'_>) -> Option> { } } -cfg_select! { - not(feature = "new_parser") => { - fn parse_config_value(arg: &PgNode) -> Option> { - match &arg.node { - Some(NodeEnum::AConst(AConst { - val: Some(Val::Sval(PgString { sval })), - .. - })) => Some(Some(ParameterValue::String(sval.clone()))), - Some(NodeEnum::AConst(AConst { isnull: true, .. })) => Some(None), - // FIXME(sage): The function only takes text. Do we need to deal with - // other literals? - _ => None, - } - } - } - _ => {} -} - /// Returns None if the node was not a constant boolean -#[cfg(feature = "new_parser")] fn parse_is_local(arg: Node<'_>) -> Option { match arg { Node::A_Const(c) => c.val()?.bool_value(), _ => None, } } - -cfg_select! { - not(feature = "new_parser") => { - fn parse_is_local(arg: &PgNode) -> Option { - match &arg.node { - Some(NodeEnum::AConst(AConst { - val: Some(Val::Boolval(Boolean { boolval })), - .. - })) => Some(*boolval), - // Only constant strings can be handled for now - _ => None, - } - } - } - _ => {} -} diff --git a/pgdog/src/frontend/router/parser/query/show.rs b/pgdog/src/frontend/router/parser/query/show.rs index ac9286360..731cfb9a6 100644 --- a/pgdog/src/frontend/router/parser/query/show.rs +++ b/pgdog/src/frontend/router/parser/query/show.rs @@ -3,7 +3,6 @@ use crate::frontend::router::{parser::Shard, round_robin}; impl QueryParser { /// Handle SHOW command. - #[cfg(feature = "new_parser")] pub(super) fn show( &mut self, stmt: &nodes::VariableShowStmt, @@ -27,35 +26,6 @@ impl QueryParser { } } } - - cfg_select! { - not(feature = "new_parser") => { - pub(super) fn show( - &mut self, - stmt: &VariableShowStmt, - context: &mut QueryParserContext, - ) -> Result { - match stmt.name.as_str() { - "pgdog.shards" => Ok(Command::InternalField { - name: "shards".into(), - value: context.shards.to_string(), - }), - "pgdog.unique_id" => Ok(Command::UniqueId), - _ => { - 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)) - } - } - } - } - _ => {} - } } #[cfg(test)] diff --git a/pgdog/src/frontend/router/parser/query/test/mod.rs b/pgdog/src/frontend/router/parser/query/test/mod.rs index ab6cdb05a..e9e744c85 100644 --- a/pgdog/src/frontend/router/parser/query/test/mod.rs +++ b/pgdog/src/frontend/router/parser/query/test/mod.rs @@ -304,7 +304,6 @@ fn test_select_for_update_in_cte() { } #[test] -#[cfg_attr(not(feature = "new_parser"), should_panic)] fn test_select_for_update_in_subselect() { // FOR UPDATE buried inside a subselect should still route to the primary. let route = diff --git a/pgdog/src/frontend/router/parser/query/test/test_insert.rs b/pgdog/src/frontend/router/parser/query/test/test_insert.rs index e8674b4eb..0a6730c1d 100644 --- a/pgdog/src/frontend/router/parser/query/test/test_insert.rs +++ b/pgdog/src/frontend/router/parser/query/test/test_insert.rs @@ -66,7 +66,6 @@ fn test_insert_multi_row() { } #[test] -#[cfg_attr(not(feature = "new_parser"), should_panic)] // Fixed in port fn test_insert_select() { let mut test = QueryParserTest::new(); @@ -82,7 +81,6 @@ fn test_insert_select() { } #[test] -#[cfg_attr(not(feature = "new_parser"), should_panic)] // Fixed in port fn test_insert_default_values() { let mut test = QueryParserTest::new(); diff --git a/pgdog/src/frontend/router/parser/query/test/test_select.rs b/pgdog/src/frontend/router/parser/query/test/test_select.rs index d06c1c1ea..757356cc3 100644 --- a/pgdog/src/frontend/router/parser/query/test/test_select.rs +++ b/pgdog/src/frontend/router/parser/query/test/test_select.rs @@ -268,29 +268,26 @@ fn test_system_catalog_sharded() { ); assert!(!command.route().is_omnisharded()); - #[cfg(feature = "new_parser")] - { - let command = test.execute(vec![Query::new("SELECT * FROM pg_type").into()]); - assert_eq!( - command.route().shard(), - &Shard::Direct(0), - "pg_type queries should go to shard 0", - ); - - let command = test.execute(vec![Query::new("SELECT $1::regtype").into()]); - assert_eq!( - command.route().shard(), - &Shard::Direct(0), - "regtype casts should go to shard 0", - ); - - let command = test.execute(vec![Query::new("SELECT to_regtype($1)").into()]); - assert_eq!( - command.route().shard(), - &Shard::Direct(0), - "to_regtype should go to shard 0", - ); - } + let command = test.execute(vec![Query::new("SELECT * FROM pg_type").into()]); + assert_eq!( + command.route().shard(), + &Shard::Direct(0), + "pg_type queries should go to shard 0", + ); + + let command = test.execute(vec![Query::new("SELECT $1::regtype").into()]); + assert_eq!( + command.route().shard(), + &Shard::Direct(0), + "regtype casts should go to shard 0", + ); + + let command = test.execute(vec![Query::new("SELECT to_regtype($1)").into()]); + assert_eq!( + command.route().shard(), + &Shard::Direct(0), + "to_regtype should go to shard 0", + ); } #[test] diff --git a/pgdog/src/frontend/router/parser/query/transaction.rs b/pgdog/src/frontend/router/parser/query/transaction.rs index 6280895a4..910e26df2 100644 --- a/pgdog/src/frontend/router/parser/query/transaction.rs +++ b/pgdog/src/frontend/router/parser/query/transaction.rs @@ -1,5 +1,4 @@ use crate::frontend::client::TransactionType; -#[cfg(feature = "new_parser")] use pg_raw_parse::nodes::TransactionStmtKind::*; use super::*; @@ -9,10 +8,9 @@ impl QueryParser { /// /// # Arguments /// - /// * `stmt`: Transaction statement from pg_query. + /// * `stmt`: Transaction statement. /// * `context`: Query parser context. /// - #[cfg(feature = "new_parser")] pub(super) fn transaction( &mut self, stmt: &nodes::TransactionStmt, @@ -61,62 +59,6 @@ impl QueryParser { )) } - cfg_select! { - not(feature = "new_parser") => { - pub(super) fn transaction( - &mut self, - stmt: &TransactionStmt, - 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; - } - - match stmt.kind() { - TransactionStmtKind::TransStmtCommit => { - return Ok(Command::CommitTransaction { extended }); - } - TransactionStmtKind::TransStmtRollback => { - return Ok(Command::RollbackTransaction { extended }); - } - TransactionStmtKind::TransStmtBegin | TransactionStmtKind::TransStmtStart => { - let transaction_type = Self::transaction_type(&stmt.options).unwrap_or_default(); - return Ok(Command::StartTransaction { - query: context.query()?.clone(), - transaction_type, - extended, - route: Route::write(context.shards_calculator.shard()) - .with_read(transaction_type == TransactionType::ReadOnly), - }); - } - TransactionStmtKind::TransStmtRollbackTo => rollback_savepoint = true, - TransactionStmtKind::TransStmtPrepare - | TransactionStmtKind::TransStmtCommitPrepared - | TransactionStmtKind::TransStmtRollbackPrepared - if context.router_context.two_pc => - { - return Err(Error::NoTwoPc); - } - _ => (), - } - - context - .shards_calculator - .push(ShardWithPriority::new_table(Shard::All)); - - Ok(Command::Query( - Route::write(context.shards_calculator.shard()) - .with_rollback_savepoint(rollback_savepoint), - )) - } - } - _ => {} - } - - #[cfg(feature = "new_parser")] fn transaction_type<'a>( options: impl IntoIterator>, ) -> Option { @@ -133,32 +75,6 @@ impl QueryParser { Some(TransactionType::ReadWrite) } - - cfg_select! { - not(feature = "new_parser") => { - fn transaction_type(options: &[PgNode]) -> Option { - for option_node in options { - let node_enum = option_node.node.as_ref()?; - if let NodeEnum::DefElem(def_elem) = node_enum - && def_elem.defname == "transaction_read_only" - { - let arg_node = def_elem.arg.as_ref()?.node.as_ref()?; - if let NodeEnum::AConst(ac) = arg_node { - // 1 => read-only, 0 => read-write - if let Some(a_const::Val::Ival(i)) = ac.val.as_ref() - && i.ival != 0 - { - return Some(TransactionType::ReadOnly); - } - } - } - } - - Some(TransactionType::ReadWrite) - } - } - _ => {} - } } #[cfg(test)] @@ -166,7 +82,6 @@ mod test { use super::*; #[test] - #[cfg(feature = "new_parser")] fn test_detect_transaction_type() { let read_write_queries = [ "BEGIN", @@ -214,80 +129,4 @@ mod test { assert_eq!(t, Some(TransactionType::ReadOnly)); } } - - cfg_select! { - not(feature = "new_parser") => { - #[test] - fn test_detect_transaction_type() { - let read_write_queries = vec![ - "BEGIN", - "BEGIN;", - "begin", - "bEgIn", - "BEGIN WORK", - "BEGIN TRANSACTION", - "BEGIN READ WRITE", - "BEGIN WORK READ WRITE", - "BEGIN TRANSACTION READ WRITE", - "START TRANSACTION", - "START TRANSACTION;", - "start transaction", - "START TRANSACTION READ WRITE", - "BEGIN ISOLATION LEVEL REPEATABLE READ READ WRITE DEFERRABLE", - ]; - - let read_only_queries = vec![ - "BEGIN READ ONLY", - "BEGIN WORK READ ONLY", - "BEGIN TRANSACTION READ ONLY", - "START TRANSACTION READ ONLY", - "BEGIN ISOLATION LEVEL SERIALIZABLE READ ONLY", - "START TRANSACTION ISOLATION LEVEL READ COMMITTED READ ONLY NOT DEFERRABLE", - ]; - - for q in read_write_queries { - let binding = pg_query::parse(q).unwrap(); - let stmt = binding - .protobuf - .stmts - .first() - .as_ref() - .unwrap() - .stmt - .as_ref() - .unwrap(); - - match stmt.node { - Some(NodeEnum::TransactionStmt(ref stmt)) => { - let t = QueryParser::transaction_type(&stmt.options); - assert_eq!(t, Some(TransactionType::ReadWrite)); - } - _ => panic!("not a transaction"), - } - } - - for q in read_only_queries { - let binding = pg_query::parse(q).unwrap(); - let stmt = binding - .protobuf - .stmts - .first() - .as_ref() - .unwrap() - .stmt - .as_ref() - .unwrap(); - - match stmt.node { - Some(NodeEnum::TransactionStmt(ref stmt)) => { - let t = QueryParser::transaction_type(&stmt.options); - assert_eq!(t, Some(TransactionType::ReadOnly)); - } - _ => panic!("not a transaction"), - } - } - } - } - _ => {} - } } diff --git a/pgdog/src/frontend/router/parser/query/update.rs b/pgdog/src/frontend/router/parser/query/update.rs index 2a2a28531..b78a3555b 100644 --- a/pgdog/src/frontend/router/parser/query/update.rs +++ b/pgdog/src/frontend/router/parser/query/update.rs @@ -3,14 +3,10 @@ use super::*; impl QueryParser { pub(super) fn update( &mut self, - #[cfg(not(feature = "new_parser"))] stmt: &UpdateStmt, - #[cfg(feature = "new_parser")] stmt: pg_raw_parse::Node<'_>, + stmt: pg_raw_parse::Node<'_>, context: &mut QueryParserContext, ) -> Result { - let mut parser = StatementParser::from_update( - #[cfg(not(feature = "new_parser"))] - stmt, - #[cfg(feature = "new_parser")] + let mut parser = StatementParser::new( stmt, context.router_context.bind, &context.sharding_schema, @@ -62,11 +58,8 @@ impl QueryParser { #[cfg(test)] mod tests { use super::*; - #[cfg(not(feature = "new_parser"))] - use pg_query::NodeEnum; #[test] - #[cfg(feature = "new_parser")] fn update_preserves_decimal_values() { let parsed = pg_raw_parse::parse( "UPDATE transactions SET amount = 50.00, status = 'completed' WHERE id = 1", @@ -99,59 +92,7 @@ mod tests { assert!(found_string, "Should have found string value"); } - cfg_select! { - not(feature = "new_parser") => { - #[test] - fn update_preserves_decimal_values() { - let parsed = 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 - && 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] - #[cfg(feature = "new_parser")] fn update_with_quoted_decimal() { let parsed = pg_raw_parse::parse("UPDATE transactions SET amount = '50.00' WHERE id = 1").unwrap(); @@ -171,43 +112,4 @@ mod tests { } assert!(found_string, "Should have found string value"); } - - cfg_select! { - not(feature = "new_parser") => { - #[test] - fn update_with_quoted_decimal() { - let parsed = pg_query::parse("UPDATE transactions SET amount = '50.00' 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"), - }; - - // 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 - && let Some(val) = &res.val - { - let value = Value::try_from(&val.node).unwrap(); - if let Value::String(s) = value { - 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/statement/aggregate/engine.rs b/pgdog/src/frontend/router/parser/rewrite/statement/aggregate/engine.rs index fb16cd2db..2a2ef2fb8 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/aggregate/engine.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/aggregate/engine.rs @@ -1,17 +1,6 @@ -#[cfg(feature = "new_parser")] use crate::frontend::router::parser::Function; use crate::frontend::router::parser::aggregate::{Aggregate, AggregateFunction}; -#[cfg(not(feature = "new_parser"))] -use crate::frontend::router::parser::util::pg_string; -#[cfg(feature = "new_parser")] use itertools::*; -#[cfg(not(feature = "new_parser"))] -use pg_query::NodeEnum; -#[cfg(not(feature = "new_parser"))] -use pg_query::protobuf::{ - AConst, FuncCall, Integer, Node, ResTarget, SelectStmt, TypeCast, TypeName, a_const::Val, -}; -#[cfg(feature = "new_parser")] use pg_raw_parse::{Node, make, nodes}; use super::{AggregateRewritePlan, HelperKind, HelperMapping, RewriteOutput}; @@ -24,7 +13,6 @@ pub(crate) struct AggregatesRewrite; impl AggregatesRewrite { /// Rewrite a SELECT query in-place, adding helper aggregates when necessary. - #[cfg(feature = "new_parser")] pub(crate) fn rewrite_select<'a>( select: &mut nodes::SelectStmtMut<'a, '_>, mem: make::MemoryToken<'a>, @@ -76,154 +64,6 @@ impl AggregatesRewrite { } } - #[cfg(not(feature = "new_parser"))] - pub(crate) fn rewrite_select( - &self, - ast: &mut SelectStmt, - aggregate: &Aggregate, - ) -> RewriteOutput { - self.rewrite_parsed(ast, aggregate) - } - - #[cfg(not(feature = "new_parser"))] - fn rewrite_parsed(&self, select: &mut SelectStmt, aggregate: &Aggregate) -> RewriteOutput { - 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() { - 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_{}_col{}", kind.alias_suffix(), target.column()); - - 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_helper(HelperMapping { - target_column: target.column(), - helper_column, - distinct: target.is_distinct(), - kind, - alias: helper_alias, - }); - } - } - - if helper_nodes.is_empty() { - return RewriteOutput::default(); - } - - select.target_list.extend(helper_nodes); - - RewriteOutput::new(plan) - } - - #[cfg(not(feature = "new_parser"))] - 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, - } - } - - #[cfg(not(feature = "new_parser"))] - fn build_count_func(original: &FuncCall, distinct: bool) -> FuncCall { - FuncCall { - funcname: vec![pg_string("count")], - 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(not(feature = "new_parser"))] - fn build_sum_func(original: &FuncCall, distinct: bool) -> FuncCall { - FuncCall { - funcname: vec![pg_string("sum")], - 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(feature = "new_parser")] fn build_sum_of_squares_func<'a>( original: &nodes::FuncCall, mem: make::MemoryToken<'a>, @@ -259,65 +99,6 @@ impl AggregatesRewrite { sumsq } - #[cfg(not(feature = "new_parser"))] - fn build_sum_of_squares_func(original: &FuncCall, distinct: bool) -> FuncCall { - let arg = original.args.first().cloned(); - // POWER will return double even when dealing with integer inputs. - // For any non float type, functions using this helper return numeric. - // We can go from numeric to f64 losslessly, but not the other way, - // so we cast here - let arg = Node { - node: Some(NodeEnum::TypeCast(Box::new(TypeCast { - arg: arg.map(Box::new), - type_name: Some(TypeName { - names: vec![pg_string("pg_catalog"), pg_string("numeric")], - type_oid: 1700, - ..Default::default() - }), - location: original.location, - }))), - }; - - 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![pg_string("power")], - 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, - }; - - FuncCall { - funcname: vec![pg_string("sum")], - 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, - } - } - - #[cfg(feature = "new_parser")] fn helper_specs<'a>( func_call: &nodes::FuncCall, function: &AggregateFunction, @@ -359,7 +140,6 @@ impl AggregatesRewrite { } } - #[cfg(feature = "new_parser")] fn copy_and_rename_function<'a>( func_call: &nodes::FuncCall, name: &str, @@ -370,50 +150,8 @@ impl AggregatesRewrite { .set_funcname(mem.make_list(&[mem.make_string(Some(name)).uncast()])); func } - - #[cfg(not(feature = "new_parser"))] - 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 => { - 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, - }, - HelperSpec { - func: Self::build_sum_of_squares_func(func_call, distinct), - kind: HelperKind::SumSquares, - }, - ] - } - _ => vec![], - } - } } -#[derive(Debug, Clone)] -#[cfg(not(feature = "new_parser"))] -struct HelperSpec { - func: FuncCall, - kind: HelperKind, -} - -#[cfg(feature = "new_parser")] struct HelperSpec<'a> { func: make::Unique<'a, &'a nodes::FuncCall>, kind: HelperKind, @@ -421,28 +159,10 @@ struct HelperSpec<'a> { #[cfg(test)] mod tests { - #![cfg_attr(feature = "new_parser", allow(unused_mut))] use super::*; use crate::frontend::router::parser::aggregate::Aggregate; - #[cfg(not(feature = "new_parser"))] - use pg_query::protobuf::ParseResult; - #[cfg(feature = "new_parser")] use pg_raw_parse::{Node, Owned, make, nodes}; - #[cfg(not(feature = "new_parser"))] - fn select(ast: &mut ParseResult) -> &mut pg_query::protobuf::SelectStmt { - match ast - .stmts - .first_mut() - .and_then(|stmt| stmt.stmt.as_mut()) - .and_then(|stmt| stmt.node.as_mut()) - { - Some(NodeEnum::SelectStmt(select)) => &mut *select, - _ => panic!("not a select"), - } - } - - #[cfg(feature = "new_parser")] fn rewrite(sql: &str) -> (Owned, RewriteOutput) { let ast = pg_raw_parse::parse(sql).unwrap(); @@ -464,30 +184,16 @@ mod tests { (select, output.unwrap()) } - #[cfg(not(feature = "new_parser"))] - fn rewrite(sql: &str) -> (ParseResult, RewriteOutput) { - let mut parsed = pg_query::parse(sql).unwrap().protobuf; - - let stmt_mut = select(&mut parsed); - let aggregate = Aggregate::parse(stmt_mut, &Default::default()); - - let output = AggregatesRewrite.rewrite_select(stmt_mut, &aggregate); - (parsed, output) - } - #[test] fn rewrite_engine_noop() { - let (mut ast, output) = rewrite("SELECT COUNT(price) FROM menu"); + let (ast, output) = rewrite("SELECT COUNT(price) FROM menu"); assert!(output.plan.is_noop()); - #[cfg(feature = "new_parser")] assert_eq!(ast.target_list().len(), 1); - #[cfg(not(feature = "new_parser"))] - assert_eq!(select(&mut ast).target_list.len(), 1); } #[test] fn rewrite_engine_adds_helper() { - let (mut ast, output) = rewrite("SELECT AVG(price) FROM menu"); + let (ast, output) = rewrite("SELECT AVG(price) FROM menu"); assert!(!output.plan.is_noop()); assert_eq!(output.plan.drop_columns().collect::>(), &[1]); assert_eq!(output.plan.helpers().len(), 1); @@ -497,10 +203,7 @@ mod tests { assert!(!helper.distinct); assert!(matches!(helper.kind, HelperKind::Count)); - #[cfg(feature = "new_parser")] let aggregate = Aggregate::parse(&ast, &Default::default()); - #[cfg(not(feature = "new_parser"))] - let aggregate = Aggregate::parse(select(&mut ast), &Default::default()); assert_eq!(aggregate.targets().len(), 2); assert!( aggregate @@ -512,7 +215,7 @@ mod tests { #[test] fn rewrite_engine_handles_mismatched_pair() { - let (mut ast, output) = rewrite("SELECT COUNT(price::numeric), AVG(price) FROM menu"); + let (ast, output) = rewrite("SELECT COUNT(price::numeric), AVG(price) FROM menu"); assert_eq!(output.plan.drop_columns().collect::>(), &[2]); assert_eq!(output.plan.helpers().len(), 1); let helper = &output.plan.helpers()[0]; @@ -521,10 +224,7 @@ mod tests { assert!(!helper.distinct); assert!(matches!(helper.kind, HelperKind::Count)); - #[cfg(feature = "new_parser")] let aggregate = Aggregate::parse(&ast, &Default::default()); - #[cfg(not(feature = "new_parser"))] - let aggregate = Aggregate::parse(select(&mut ast), &Default::default()); assert_eq!(aggregate.targets().len(), 3); assert!( aggregate @@ -538,7 +238,7 @@ mod tests { #[test] fn rewrite_engine_multiple_avg_helpers() { - let (mut ast, output) = rewrite("SELECT AVG(price), AVG(discount) FROM menu"); + let (ast, output) = rewrite("SELECT AVG(price), AVG(discount) FROM menu"); assert_eq!(output.plan.drop_columns().collect::>(), &[2, 3]); assert_eq!(output.plan.helpers().len(), 2); @@ -552,10 +252,7 @@ mod tests { assert_eq!(helper_discount.helper_column, 3); assert!(matches!(helper_discount.kind, HelperKind::Count)); - #[cfg(feature = "new_parser")] let aggregate = Aggregate::parse(&ast, &Default::default()); - #[cfg(not(feature = "new_parser"))] - let aggregate = Aggregate::parse(select(&mut ast), &Default::default()); assert_eq!(aggregate.targets().len(), 4); assert_eq!( aggregate @@ -569,7 +266,7 @@ mod tests { #[test] fn rewrite_engine_stddev_helpers() { - let (mut ast, output) = rewrite("SELECT STDDEV(price) FROM menu"); + let (ast, output) = rewrite("SELECT STDDEV(price) FROM menu"); assert!(!output.plan.is_noop()); assert_eq!(output.plan.drop_columns().collect::>(), &[1, 2, 3]); assert_eq!(output.plan.helpers().len(), 3); @@ -589,9 +286,6 @@ mod tests { assert!(kinds.contains(&HelperKind::SumSquares)); // Expect original STDDEV plus three helpers. - #[cfg(feature = "new_parser")] assert_eq!(ast.target_list().len(), 4); - #[cfg(not(feature = "new_parser"))] - assert_eq!(select(&mut 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 index ea8dd91e0..0ae8e04a5 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/aggregate/mod.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/aggregate/mod.rs @@ -4,9 +4,6 @@ mod plan; use super::{Error, RewritePlan, StatementRewrite}; use crate::backend::schema::Schema; use crate::frontend::router::parser::aggregate::Aggregate; -#[cfg(not(feature = "new_parser"))] -use pg_query::NodeEnum; -#[cfg(feature = "new_parser")] use pg_raw_parse::{make::MemoryToken, nodes::SelectStmtMut}; pub(crate) use engine::AggregatesRewrite; @@ -14,7 +11,6 @@ pub(crate) use plan::{AggregateRewritePlan, HelperKind, HelperMapping, RewriteOu impl StatementRewrite<'_> { /// Add missing COUNT(*) and other helps when using aggregates. - #[cfg(feature = "new_parser")] pub(super) fn rewrite_aggregates<'a>( &mut self, select: &mut SelectStmtMut<'a, '_>, @@ -40,41 +36,4 @@ impl StatementRewrite<'_> { self.rewritten = true; Ok(()) } - - #[cfg(not(feature = "new_parser"))] - pub(super) fn rewrite_aggregates( - &mut self, - plan: &mut RewritePlan, - schema: &Schema, - ) -> Result<(), Error> { - if self.schema.shards == 1 { - return Ok(()); - } - - let Some(raw_stmt) = self.stmt.stmts.first_mut() else { - return Ok(()); - }; - - let Some(stmt) = raw_stmt.stmt.as_mut() else { - return Ok(()); - }; - - let Some(NodeEnum::SelectStmt(select)) = stmt.node.as_mut() else { - return Ok(()); - }; - - let aggregate = Aggregate::parse(select, schema); - if aggregate.is_empty() { - return Ok(()); - } - - let output = AggregatesRewrite.rewrite_select(select, &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/statement/auto_id.rs b/pgdog/src/frontend/router/parser/rewrite/statement/auto_id.rs index c342e4613..c56a074bf 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/auto_id.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/auto_id.rs @@ -1,13 +1,7 @@ //! Auto-inject pgdog.unique_id() for missing BIGINT primary keys in INSERT statements. use indexmap::IndexSet; -#[cfg(feature = "new_parser")] use itertools::*; -#[cfg(not(feature = "new_parser"))] -use pg_query::protobuf::{FuncCall, ResTarget, String as PgString}; -#[cfg(not(feature = "new_parser"))] -use pg_query::{Node as PgNode, NodeEnum}; -#[cfg(feature = "new_parser")] use pg_raw_parse::{Node, NodeMut, make, nodes}; use pgdog_config::RewriteMode; @@ -25,7 +19,6 @@ impl StatementRewrite<'_> { /// /// This runs before unique_id replacement so injected function calls /// will be processed by the unique_id rewriter. - #[cfg(feature = "new_parser")] pub(super) fn inject_auto_id<'a>( &mut self, mut node: nodes::InsertStmtMut<'a, '_>, @@ -100,127 +93,19 @@ impl StatementRewrite<'_> { Ok(()) } - cfg_select! { - not(feature = "new_parser") => { - pub(super) fn inject_auto_id( - &mut self, - plan: &mut RewritePlan, - ) -> Result<(), Error> { - let mode = self.schema.rewrite.primary_key; - - if mode == RewriteMode::Ignore || self.schema.shards == 1 { - return Ok(()); - } - - let Some((table, is_sharded)) = self.get_insert_table() else { - return Ok(()); - }; - - let Some(relation) = self.db_schema.table(table, self.user, self.search_path) else { - return Ok(()); - }; - - // Get the columns specified in the INSERT (preserving order) - let insert_columns: IndexSet<&str> = self.get_insert_column_names_ordered(); - - // Find BIGINT primary key columns - let bigint_pk_columns: Vec<&str> = relation - .columns() - .values() - .filter(|col| col.is_primary_key && is_bigint_type(&col.data_type)) - .map(|col| col.column_name.as_str()) - .collect(); - - if bigint_pk_columns.is_empty() { - return Ok(()); - } - - // Find positions of present PK columns (for DEFAULT replacement) - let present_pk_positions: Vec = bigint_pk_columns - .iter() - .filter_map(|pk_col| insert_columns.get_index_of(pk_col)) - .collect(); - - // Find which PK columns are missing - let missing_columns: Vec<&str> = bigint_pk_columns - .iter() - .filter(|pk_col| !insert_columns.contains(*pk_col)) - .copied() - .collect(); - - let rewrite = - mode == RewriteMode::Rewrite || mode == RewriteMode::RewriteOmni && !is_sharded; - - // Replace DEFAULT values with unique_id() for present columns (only in rewrite mode) - if rewrite { - let replaced = self.replace_set_to_default_at_positions(&present_pk_positions); - if replaced > 0 { - plan.auto_id_injected += replaced as u16; - self.rewritten = true; - } - } - - if missing_columns.is_empty() { - return Ok(()); - } - - if mode == RewriteMode::Error { - return Err(Error::MissingPrimaryKey); - } - - if rewrite { - for column in missing_columns { - self.inject_column_with_unique_id(column)?; - plan.auto_id_injected += 1; - } - self.rewritten = true; - } - - Ok(()) - } - } - _ => {} - } - /// Get the table from an INSERT statement. - #[cfg(feature = "new_parser")] fn get_insert_table<'a>(&self, insert: &'a nodes::InsertStmt) -> (Table<'a>, bool) { let relation = insert.relation().expect("INSERT always has table"); - let is_sharded = StatementParser::from_insert(insert.into(), None, self.schema, None) - .is_sharded(self.db_schema, self.user, self.search_path); + let is_sharded = StatementParser::new(insert.into(), None, self.schema, None).is_sharded( + self.db_schema, + self.user, + self.search_path, + ); (Table::from(relation), is_sharded) } - cfg_select! { - not(feature = "new_parser") => { - fn get_insert_table( - &self, - ) -> Option<(Table<'_>, bool)> { - let stmt = self.stmt.stmts.first()?; - let pg_node = stmt.stmt.as_ref()?; - - if let NodeEnum::InsertStmt(insert) = pg_node.node.as_ref()? { - let relation = insert.relation.as_ref()?; - let is_sharded = StatementParser::from_insert( - insert, - None, - self.schema, - None, - ) - .is_sharded(self.db_schema, self.user, self.search_path); - - return Some((Table::from(relation), is_sharded)); - } - - None - } - } - _ => {} - } - /// Get the column names specified in the INSERT statement, preserving order. - #[cfg(feature = "new_parser")] fn get_insert_column_names_ordered<'a>( &self, insert: &'a nodes::InsertStmt, @@ -235,38 +120,7 @@ impl StatementRewrite<'_> { .collect() } - cfg_select! { - not(feature = "new_parser") => { - fn get_insert_column_names_ordered(&self) -> IndexSet<&str> { - let Some(stmt) = self.stmt.stmts.first() else { - return IndexSet::new(); - }; - let Some(node) = stmt.stmt.as_ref() else { - return IndexSet::new(); - }; - let Some(NodeEnum::InsertStmt(insert)) = node.node.as_ref() else { - return IndexSet::new(); - }; - - insert - .cols - .iter() - .filter_map(|col| { - if let Some(NodeEnum::ResTarget(res)) = &col.node - && !res.name.is_empty() - { - return Some(res.name.as_str()); - } - None - }) - .collect() - } - } - _ => {} - } - /// Replace SetToDefault nodes at the specified column positions with pgdog.unique_id(). - #[cfg(feature = "new_parser")] fn replace_set_to_default_at_positions<'a, 'b>( &mut self, insert: &mut nodes::InsertStmtMut<'a, 'b>, @@ -291,49 +145,7 @@ impl StatementRewrite<'_> { replaced } - cfg_select! { - not(feature = "new_parser") => { - fn replace_set_to_default_at_positions(&mut self, positions: &[usize]) -> usize { - let Some(stmt) = self.stmt.stmts.first_mut() else { - return 0; - }; - let Some(node) = stmt.stmt.as_mut() else { - return 0; - }; - let Some(NodeEnum::InsertStmt(insert)) = node.node.as_mut() else { - return 0; - }; - let Some(select) = insert.select_stmt.as_mut() else { - return 0; - }; - let Some(NodeEnum::SelectStmt(select_stmt)) = select.node.as_mut() else { - return 0; - }; - - let mut replaced = 0; - let unique_id_call = Self::unique_id_func_call(); - - for values_node in &mut select_stmt.values_lists { - if let Some(NodeEnum::List(list)) = &mut values_node.node { - for &pos in positions { - if pos < list.items.len() - && let Some(NodeEnum::SetToDefault(_)) = &list.items[pos].node - { - list.items[pos] = unique_id_call.clone(); - replaced += 1; - } - } - } - } - - replaced - } - } - _ => {} - } - /// Inject a column with pgdog.unique_id() as the value. - #[cfg(feature = "new_parser")] fn inject_column_with_unique_id<'a>( &mut self, insert: &mut nodes::InsertStmtMut<'a, '_>, @@ -356,52 +168,7 @@ impl StatementRewrite<'_> { } } - cfg_select! { - not(feature = "new_parser") => { - fn inject_column_with_unique_id(&mut self, column_name: &str) -> Result<(), Error> { - let Some(stmt) = self.stmt.stmts.first_mut() else { - return Ok(()); - }; - let Some(node) = stmt.stmt.as_mut() else { - return Ok(()); - }; - let Some(NodeEnum::InsertStmt(insert)) = node.node.as_mut() else { - return Ok(()); - }; - - // Add the column to the column list - let col_node = PgNode { - node: Some(NodeEnum::ResTarget(Box::new(ResTarget { - name: column_name.to_string(), - ..Default::default() - }))), - }; - insert.cols.push(col_node); - - // Add pgdog.unique_id() to each values list - let Some(select) = insert.select_stmt.as_mut() else { - return Ok(()); - }; - let Some(NodeEnum::SelectStmt(select_stmt)) = select.node.as_mut() else { - return Ok(()); - }; - - let unique_id_call = Self::unique_id_func_call(); - - for values_node in &mut select_stmt.values_lists { - if let Some(NodeEnum::List(list)) = &mut values_node.node { - list.items.push(unique_id_call.clone()); - } - } - - Ok(()) - } - } - _ => {} - } - /// Create a function call node for pgdog.unique_id(). - #[cfg(feature = "new_parser")] fn unique_id_func_call(mem: make::MemoryToken<'_>) -> make::Unique<'_, &nodes::FuncCall> { mem.make_func_call( mem.make_list(&[ @@ -412,33 +179,6 @@ impl StatementRewrite<'_> { Default::default(), ) } - - cfg_select! { - not(feature = "new_parser") => { - fn unique_id_func_call() -> PgNode { - PgNode { - node: Some(NodeEnum::FuncCall(Box::new(FuncCall { - funcname: vec![ - PgNode { - node: Some(NodeEnum::String(PgString { - sval: "pgdog".to_string(), - })), - }, - PgNode { - node: Some(NodeEnum::String(PgString { - sval: "unique_id".to_string(), - })), - }, - ], - args: vec![], - func_variadic: false, - ..Default::default() - }))), - } - } - } - _ => {} - } } /// Check if a data type is a BIGINT variant. @@ -744,7 +484,6 @@ mod tests { } } - #[cfg(feature = "new_parser")] fn rewrite_sql_with_sharding_schema( sql: &str, db_schema: &Schema, @@ -772,38 +511,6 @@ mod tests { Ok((sql, plan)) } - cfg_select! { - not(feature = "new_parser") => { - fn rewrite_sql_with_sharding_schema( - sql: &str, - db_schema: &Schema, - schema: &ShardingSchema, - ) -> Result<(String, RewritePlan), Error> { - let _guard = set_env_var("NODE_ID", "pgdog-1"); - let mut ast = pg_query::parse(sql).unwrap().protobuf; - let mut prepared = PreparedStatements::default(); - let mut rewriter = StatementRewrite::new(StatementRewriteContext { - stmt: &mut ast, - extended: false, - prepared: false, - prepared_statements: &mut prepared, - schema, - db_schema, - user: "", - search_path: None, - }); - let plan = rewriter.maybe_rewrite()?; - let result = if plan.stmt.is_some() { - plan.stmt.clone().unwrap() - } else { - ast.deparse().unwrap() - }; - Ok((result, plan)) - } - } - _ => {} - } - #[test] fn test_rewrite_omni_skips_sharded_table() { let db_schema = make_schema_with_bigint_pk(); diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/error.rs b/pgdog/src/frontend/router/parser/rewrite/statement/error.rs index 45a760014..18e81a7b5 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/error.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/error.rs @@ -5,12 +5,7 @@ pub enum Error { #[error("unique_id generation failed: {0}")] UniqueId(#[from] crate::unique_id::Error), - #[error("pg_query: {0}")] - #[cfg(not(feature = "new_parser"))] - PgQuery(#[from] pg_query::Error), - #[error("parser: {0}")] - #[cfg(feature = "new_parser")] Parser(#[from] pg_raw_parse::Error), #[error("cache: {0}")] diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/insert.rs b/pgdog/src/frontend/router/parser/rewrite/statement/insert.rs index 4a5d8b4eb..3f6208752 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/insert.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/insert.rs @@ -1,18 +1,10 @@ -#[cfg(feature = "new_parser")] use indexmap::IndexSet; -#[cfg(not(feature = "new_parser"))] -use pg_query::{Node, NodeEnum}; -#[cfg(feature = "new_parser")] use pg_raw_parse::{Node, NodeMut, deparse, make, nodes, walk}; -#[cfg(not(feature = "new_parser"))] -use pgdog_config::QueryParserEngine; use pgdog_config::RewriteMode; use crate::frontend::router::Ast; use crate::frontend::router::parser::Cache; use crate::frontend::{BufferedQuery, ClientRequest}; -#[cfg(any(test, not(feature = "new_parser")))] -use crate::net::messages::bind::{Format, Parameter}; use crate::net::{Bind, Parse, ProtocolMessage, Query}; use super::{Error, RewritePlan, StatementRewrite}; @@ -22,10 +14,7 @@ pub struct InsertSplit { /// Parameter positions in the original Bind message /// that should be used to build the Bind message specific to this /// insert statement. - #[cfg(feature = "new_parser")] params: IndexSet, - #[cfg(not(feature = "new_parser"))] - params: Vec, /// The split up INSERT statement with parameters and/or values. stmt: String, @@ -107,7 +96,6 @@ impl InsertSplit { } /// Extract specific parameters from a Bind message based on this split's param indices. - #[cfg(feature = "new_parser")] fn extract_bind_params(&self, bind: &Bind) -> Result { let mut new = Bind::new_statement(self.statement_name().unwrap_or_default()); for param in &self.params { @@ -119,42 +107,6 @@ impl InsertSplit { Ok(new) } - - cfg_select! { - not(feature = "new_parser") => { - fn extract_bind_params(&self, bind: &Bind) -> Result { - 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()); - - Ok(Bind::new_params_codes(statement_name, ¶ms, &codes)) - } - } - _ => {} - } } /// Build separate ClientRequests for each insert split. @@ -185,7 +137,6 @@ impl StatementRewrite<'_> { /// INSERT INTO my_table (id, value) VALUES ($1, $2) -- These are copied from params $3 and $4 /// ``` /// - #[cfg(feature = "new_parser")] pub(super) fn split_insert( &mut self, insert: &nodes::InsertStmt, @@ -255,84 +206,8 @@ impl StatementRewrite<'_> { Ok(()) } - cfg_select! { - not(feature = "new_parser") => { - 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(); - let ctx = self.ast_context(); - 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, &ctx, 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. - #[cfg(not(feature = "new_parser"))] - 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()? - && !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. - #[cfg(feature = "new_parser")] fn build_single_tuple_select<'mem>( &self, mem: make::MemoryToken<'mem>, @@ -352,89 +227,6 @@ impl StatementRewrite<'_> { select.as_mut().set_values_lists(mem.make_list(&[tuple])); (params, select) } - - cfg_select! { - not(feature = "new_parser") => { - 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() - && let Some(node) = stmt.stmt.as_mut() - && let Some(NodeEnum::InsertStmt(insert)) = node.node.as_mut() - && let Some(select) = insert.select_stmt.as_mut() - && 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. - #[cfg(not(feature = "new_parser"))] - 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. - #[cfg(not(feature = "new_parser"))] - 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)] @@ -446,6 +238,7 @@ mod tests { use crate::backend::schema::Schema; use crate::frontend::PreparedStatements; use crate::frontend::router::parser::StatementRewriteContext; + use crate::net::messages::bind::{Format, Parameter}; fn default_db_schema() -> Schema { Schema::default() @@ -464,11 +257,7 @@ mod tests { } fn parse_and_split(sql: &str) -> Vec { - #[cfg(not(feature = "new_parser"))] - let mut ast = pg_query::parse(sql).unwrap().protobuf; - #[cfg(feature = "new_parser")] let root = pg_raw_parse::parse(sql).unwrap(); - #[cfg(feature = "new_parser")] let insert = match root.stmts().next() { Some(Node::InsertStmt(insert)) => insert, _ => unreachable!(), @@ -477,8 +266,6 @@ mod tests { let schema = default_schema(); let db_schema = default_db_schema(); let mut rewriter = StatementRewrite::new(StatementRewriteContext { - #[cfg(not(feature = "new_parser"))] - stmt: &mut ast, extended: false, prepared: false, prepared_statements: &mut prepared, @@ -488,10 +275,7 @@ mod tests { search_path: None, }); let mut plan = RewritePlan::default(); - #[cfg(feature = "new_parser")] rewriter.split_insert(insert, &mut plan).unwrap(); - #[cfg(not(feature = "new_parser"))] - rewriter.split_insert(&mut plan).unwrap(); plan.insert_split } @@ -502,20 +286,14 @@ mod tests { assert_eq!(splits.len(), 2); // First tuple uses params 0 and 1 (original $1, $2) - #[cfg(feature = "new_parser")] assert_eq!(splits[0].params.as_slice(), &[1, 2]); - #[cfg(not(feature = "new_parser"))] - 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 - #[cfg(feature = "new_parser")] assert_eq!(splits[1].params.as_slice(), &[3, 4]); - #[cfg(not(feature = "new_parser"))] - assert_eq!(splits[1].params, &[2, 3]); assert_eq!( splits[1].stmt(), "INSERT INTO my_table (id, value) VALUES ($1, $2)" @@ -557,19 +335,13 @@ mod tests { assert_eq!(splits.len(), 2); - #[cfg(feature = "new_parser")] assert_eq!(splits[0].params.as_slice(), &[1]); - #[cfg(not(feature = "new_parser"))] - assert_eq!(splits[0].params, &[0]); assert_eq!( splits[0].stmt(), "INSERT INTO my_table (id, value) VALUES ($1, 'a')" ); - #[cfg(feature = "new_parser")] assert_eq!(splits[1].params.as_slice(), &[2]); - #[cfg(not(feature = "new_parser"))] - assert_eq!(splits[1].params, &[1]); assert_eq!( splits[1].stmt(), "INSERT INTO my_table (id, value) VALUES ($1, 'b')" @@ -700,7 +472,6 @@ mod tests { } #[test] - #[cfg(feature = "new_parser")] fn test_extract_bind_params_incorrect_count() { let splits = parse_and_split("INSERT INTO t (a, b) VALUES ($1, $2), ($3, $4)"); let bind = Bind::new_params( diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs b/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs index 3cfd4cc1e..8ef1105d2 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs @@ -1,13 +1,6 @@ //! Statement rewriter. -#[cfg(not(feature = "new_parser"))] -use pg_query::Node as PgNode; -#[cfg(not(feature = "new_parser"))] -use pg_query::protobuf::ParseResult; -#[cfg(feature = "new_parser")] use pg_raw_parse::{Node, NodeMut, make, nodes, transform, walk}; -#[cfg(not(feature = "new_parser"))] -use pgdog_config::QueryParserEngine; use crate::backend::ShardingSchema; use crate::backend::schema::Schema; @@ -24,7 +17,6 @@ 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; @@ -35,9 +27,6 @@ pub(crate) use update::*; /// Statement rewrite engine context. #[derive(Debug)] pub struct StatementRewriteContext<'a> { - /// The AST of the statement we are rewriting. - #[cfg(not(feature = "new_parser"))] - 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 @@ -57,9 +46,6 @@ pub struct StatementRewriteContext<'a> { #[derive(Debug)] pub struct StatementRewrite<'a> { - /// SQL statement. - #[cfg(not(feature = "new_parser"))] - stmt: &'a mut ParseResult, /// The statement was rewritten. rewritten: bool, /// Statement is using the extended protocol, so @@ -88,8 +74,6 @@ impl<'a> StatementRewrite<'a> { /// pub fn new(ctx: StatementRewriteContext<'a>) -> Self { Self { - #[cfg(not(feature = "new_parser"))] - stmt: ctx.stmt, rewritten: false, extended: ctx.extended, prepared: ctx.prepared, @@ -113,7 +97,6 @@ impl<'a> StatementRewrite<'a> { /// Maybe rewrite the statement and produce a rewrite plan /// we can apply to Bind messages. - #[cfg(feature = "new_parser")] pub fn maybe_rewrite<'mem>( &mut self, mut stmt: nodes::RawStmtMut<'mem, '_>, @@ -193,55 +176,4 @@ impl<'a> StatementRewrite<'a> { Ok(plan) } - - #[cfg(not(feature = "new_parser"))] - 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; - } - - // Inject pgdog.unique_id() for missing BIGINT primary keys. - // This must run BEFORE the unique_id rewriter so the injected - // function calls get processed. - self.inject_auto_id(&mut plan)?; - - // Track the next parameter number to use - let mut next_param = plan.params as i32 + 1; - - 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, self.db_schema)?; - self.limit_offset(&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/offset.rs b/pgdog/src/frontend/router/parser/rewrite/statement/offset.rs index 770ea0783..2f9440043 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/offset.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/offset.rs @@ -1,9 +1,3 @@ -#[cfg(not(feature = "new_parser"))] -use pg_query::{ - NodeEnum, - protobuf::{AConst, Integer, ParamRef, ParseResult, a_const::Val}, -}; -#[cfg(feature = "new_parser")] use pg_raw_parse::{ConstValue, Node, Owned, StmtList, nodes}; use crate::frontend::ClientRequest; @@ -91,22 +85,7 @@ impl OffsetPlan { if self.limit.limit.is_some() || self.limit.offset.is_some() { let new_limit = (limit_val.unwrap_or(0) + offset_val.unwrap_or(0)) as i32; let ast = request.ast.as_ref().ok_or(Error::MissingAst)?; - #[cfg(not(feature = "new_parser"))] - let mut protobuf = ast.ast.protobuf.clone(); - #[cfg(not(feature = "new_parser"))] - if rewrite_ast_limit_offset(&mut protobuf, new_limit) { - let result = pg_query::ParseResult::new(protobuf, "".into()); - let new_sql = result.deparse()?; - for message in request.messages.iter_mut() { - match message { - ProtocolMessage::Query(q) => q.set_query(&new_sql), - ProtocolMessage::Parse(p) => p.set_query(&new_sql), - _ => {} - } - } - } - #[cfg(feature = "new_parser")] if let Some(rewritten) = rewrite_ast_limit_offset(&ast.ast, new_limit) { let result = pg_raw_parse::deparse(&*rewritten)?; let new_sql = result.as_str(); @@ -150,7 +129,6 @@ impl LimitValueInfo { } } -#[cfg(feature = "new_parser")] fn extract_limit_value(node: Node<'_>) -> Option { match node { Node::A_Const(c) if let Some(i) = c.val().and_then(|c| c.numeric_value::()) => { @@ -163,21 +141,6 @@ fn extract_limit_value(node: Node<'_>) -> Option { } } -#[cfg(not(feature = "new_parser"))] -fn extract_limit_value(node: &Option) -> Option { - match node { - Some(NodeEnum::AConst(AConst { - val: Some(Val::Ival(Integer { ival })), - .. - })) => Some(LimitValueInfo::Literal(*ival as usize)), - Some(NodeEnum::ParamRef(ParamRef { number, .. })) => { - Some(LimitValueInfo::Param(*number as usize)) - } - _ => None, - } -} - -#[cfg(feature = "new_parser")] fn rewrite_ast_limit_offset(ast: &StmtList, new_limit: i32) -> Option> { let Some(Node::SelectStmt(select)) = ast.stmts().next() else { return None; @@ -193,46 +156,7 @@ fn rewrite_ast_limit_offset(ast: &StmtList, new_limit: i32) -> Option { - fn rewrite_ast_limit_offset(ast: &mut ParseResult, new_limit: i32) -> bool { - let raw_stmt = match ast.stmts.first_mut() { - Some(s) => s, - None => return false, - }; - let stmt = match raw_stmt.stmt.as_mut() { - Some(s) => s, - None => return false, - }; - let select = match &mut stmt.node { - Some(NodeEnum::SelectStmt(s)) => s, - _ => return false, - }; - - select.limit_count = Some(Box::new(pg_query::Node { - node: Some(NodeEnum::AConst(AConst { - val: Some(Val::Ival(Integer { ival: new_limit })), - isnull: false, - location: -1i32, - })), - })); - - select.limit_offset = Some(Box::new(pg_query::Node { - node: Some(NodeEnum::AConst(AConst { - val: Some(Val::Ival(Integer { ival: 0 })), - isnull: false, - location: -1i32, - })), - })); - - true - } - } - _ => {} -} - impl StatementRewrite<'_> { - #[cfg(feature = "new_parser")] pub(super) fn limit_offset(&self, select: &nodes::SelectStmt, plan: &mut RewritePlan) { if self.schema.shards <= 1 { return; @@ -254,54 +178,6 @@ impl StatementRewrite<'_> { offset_param: offset_info.param_index(), }); } - - #[cfg(not(feature = "new_parser"))] - pub(super) fn limit_offset(&mut self, plan: &mut RewritePlan) -> Result<(), Error> { - if self.schema.shards <= 1 { - return Ok(()); - } - - let raw_stmt = match self.stmt.stmts.first() { - Some(s) => s, - None => return Ok(()), - }; - let stmt = match raw_stmt.stmt.as_ref() { - Some(s) => s, - None => return Ok(()), - }; - let select = match &stmt.node { - Some(NodeEnum::SelectStmt(s)) => s, - _ => return Ok(()), - }; - - let offset_node = match &select.limit_offset { - Some(node) => node, - None => return Ok(()), - }; - let limit_node = match &select.limit_count { - Some(node) => node, - None => return Ok(()), - }; - - let limit_info = extract_limit_value(&limit_node.node); - let offset_info = extract_limit_value(&offset_node.node); - - let (limit_info, offset_info) = match (limit_info, offset_info) { - (Some(l), Some(o)) => (l, o), - _ => return Ok(()), - }; - - plan.offset = Some(OffsetPlan { - limit: Limit { - limit: limit_info.literal(), - offset: offset_info.literal(), - }, - limit_param: limit_info.param_index(), - offset_param: offset_info.param_index(), - }); - - Ok(()) - } } #[cfg(test)] @@ -361,16 +237,10 @@ mod tests { } fn run_limit_offset(sql: &str, schema: &ShardingSchema) -> RewritePlan { - #[cfg(not(feature = "new_parser"))] - let mut ast = pg_query::parse(sql).unwrap(); - #[cfg(feature = "new_parser")] let stmt = pg_raw_parse::parse(sql).unwrap(); let db_schema = Schema::default(); let mut ps = PreparedStatements::default(); - #[cfg_attr(feature = "new_parser", allow(unused_mut))] - let mut rewrite = StatementRewrite::new(StatementRewriteContext { - #[cfg(not(feature = "new_parser"))] - stmt: &mut ast.protobuf, + let rewrite = StatementRewrite::new(StatementRewriteContext { extended: false, prepared: false, prepared_statements: &mut ps, @@ -380,9 +250,6 @@ mod tests { search_path: None, }); let mut plan = RewritePlan::default(); - #[cfg(not(feature = "new_parser"))] - rewrite.limit_offset(&mut plan).unwrap(); - #[cfg(feature = "new_parser")] rewrite.limit_offset( if let Node::SelectStmt(stmt) = stmt.stmts().next().unwrap() { stmt @@ -470,10 +337,7 @@ mod tests { ProtocolMessage::Query(q) => q.query().to_owned(), _ => panic!("expected Query"), }; - #[cfg(feature = "new_parser")] assert_eq!(query, "SELECT * FROM t LIMIT 15"); - #[cfg(not(feature = "new_parser"))] - assert_eq!(query, "SELECT * FROM t LIMIT 15 OFFSET 0"); let route = request.route.unwrap(); assert_eq!(route.limit().limit, Some(10)); @@ -563,10 +427,7 @@ mod tests { ProtocolMessage::Parse(p) => p.query().to_owned(), _ => panic!("expected Parse"), }; - #[cfg(feature = "new_parser")] assert_eq!(sql, "SELECT * FROM t LIMIT 15"); - #[cfg(not(feature = "new_parser"))] - assert_eq!(sql, "SELECT * FROM t LIMIT 15 OFFSET 0"); let route = request.route.unwrap(); assert_eq!(route.limit().limit, Some(10)); diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/simple_prepared.rs b/pgdog/src/frontend/router/parser/rewrite/statement/simple_prepared.rs index af88d4243..30a73288a 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/simple_prepared.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/simple_prepared.rs @@ -1,12 +1,5 @@ -#[cfg(not(feature = "new_parser"))] -use pg_query::{Error as PgQueryError, NodeEnum}; -#[cfg(feature = "new_parser")] use pg_raw_parse::{NodeMut, make::MemoryToken}; -#[cfg(not(feature = "new_parser"))] -use pgdog_config::QueryParserEngine; -#[cfg(not(feature = "new_parser"))] -use crate::backend::ShardingSchema; use crate::frontend::PreparedStatements; use crate::net::Parse; @@ -45,7 +38,6 @@ impl StatementRewrite<'_> { /// should prepend `ProtocolMessage::Prepare` to the client request using the returned /// name and statement. /// - #[cfg(feature = "new_parser")] pub(super) fn rewrite_simple_prepared<'a>( &mut self, node: NodeMut<'a, '_>, @@ -70,42 +62,9 @@ impl StatementRewrite<'_> { Ok(result) } - - cfg_select! { - not(feature = "new_parser") => { - 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 - && 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. -#[cfg(feature = "new_parser")] fn rewrite_single_prepared<'a>( node: NodeMut<'a, '_>, mem: MemoryToken<'a>, @@ -146,60 +105,6 @@ fn rewrite_single_prepared<'a>( } } -cfg_select! { - not(feature = "new_parser") => { - 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}; @@ -207,8 +112,6 @@ mod tests { use crate::backend::ShardingSchema; use crate::backend::schema::Schema; use crate::config::PreparedStatements as PreparedStatementsLevel; - #[cfg(not(feature = "new_parser"))] - use pg_query::parse; use pgdog_config::Rewrite; struct TestContext { @@ -235,7 +138,6 @@ mod tests { } } - #[cfg(feature = "new_parser")] fn rewrite(&mut self, sql: &str) -> Result<(String, RewritePlan), Error> { let stmt = pg_raw_parse::parse(sql)?; let mut rewrite = StatementRewrite::new(StatementRewriteContext { @@ -256,27 +158,6 @@ mod tests { let sql = pg_raw_parse::deparse_stmts(&*ast)?; Ok((sql, plan)) } - - cfg_select! { - not(feature = "new_parser") => { - fn rewrite(&mut self, sql: &str) -> Result<(String, 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, - db_schema: &self.db_schema, - user: "", - search_path: None, - }); - let plan = rewrite.maybe_rewrite()?; - Ok((ast.deparse().unwrap(), plan)) - } - } - _ => {} - } } #[test] diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/unique_id.rs b/pgdog/src/frontend/router/parser/rewrite/statement/unique_id.rs index b3d0f4afe..7e6db6b23 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/unique_id.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/unique_id.rs @@ -1,9 +1,4 @@ use super::StatementRewrite; -#[cfg(not(feature = "new_parser"))] -use pg_query::protobuf::{AConst, ParamRef, String as PgString, TypeCast, TypeName, a_const::Val}; -#[cfg(not(feature = "new_parser"))] -use pg_query::{Node, NodeEnum}; -#[cfg(feature = "new_parser")] use pg_raw_parse::{Node, make}; impl StatementRewrite<'_> { @@ -11,7 +6,6 @@ impl StatementRewrite<'_> { /// /// Returns `Ok(Some(replacement_node))` if the node is a unique_id call, /// `Ok(None)` otherwise. Increments `next_param` when in extended mode. - #[cfg(feature = "new_parser")] pub(super) fn rewrite_unique_id<'mem>( node: Node<'_>, mem: make::MemoryToken<'mem>, @@ -46,95 +40,7 @@ impl StatementRewrite<'_> { )) } - cfg_select! { - not(feature = "new_parser") => { - 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 - #[cfg(not(feature = "new_parser"))] - 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 - #[cfg(not(feature = "new_parser"))] - 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). - #[cfg(not(feature = "new_parser"))] - 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(). - #[cfg(feature = "new_parser")] fn is_unique_id(node: Node<'_>) -> bool { let Node::FuncCall(func) = node else { return false; @@ -145,34 +51,6 @@ impl StatementRewrite<'_> { .filter_map(Node::as_str) .eq(["pgdog", "unique_id"]) } - - cfg_select! { - not(feature = "new_parser") => { - 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)] @@ -186,7 +64,6 @@ mod tests { use crate::frontend::router::parser::StatementRewriteContext; use crate::frontend::router::parser::rewrite::statement::RewritePlan; use crate::test_utils::set_env_var; - #[cfg(feature = "new_parser")] use pg_raw_parse::{Owned, nodes}; fn default_schema() -> ShardingSchema { @@ -204,7 +81,6 @@ mod tests { Schema::default() } - #[cfg(feature = "new_parser")] fn parse_first_target(sql: &str) -> Owned { let ast = pg_raw_parse::parse(sql).unwrap(); match ast.stmts().next().unwrap() { @@ -215,69 +91,34 @@ mod tests { } } - cfg_select! { - not(feature = "new_parser") => { - 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()"); - #[cfg(feature = "new_parser")] assert!(StatementRewrite::is_unique_id(node.val())); - #[cfg(not(feature = "new_parser"))] - assert!(StatementRewrite::is_unique_id(&node)); } #[test] fn test_is_unique_id_unqualified() { let node = parse_first_target("SELECT unique_id()"); - #[cfg(feature = "new_parser")] assert!(!StatementRewrite::is_unique_id(node.val())); - #[cfg(not(feature = "new_parser"))] - assert!(!StatementRewrite::is_unique_id(&node)); } #[test] fn test_is_unique_id_wrong_schema() { let node = parse_first_target("SELECT other.unique_id()"); - #[cfg(feature = "new_parser")] assert!(!StatementRewrite::is_unique_id(node.val())); - #[cfg(not(feature = "new_parser"))] - assert!(!StatementRewrite::is_unique_id(&node)); } #[test] fn test_is_unique_id_wrong_function() { let node = parse_first_target("SELECT pgdog.other_func()"); - #[cfg(feature = "new_parser")] assert!(!StatementRewrite::is_unique_id(node.val())); - #[cfg(not(feature = "new_parser"))] - assert!(!StatementRewrite::is_unique_id(&node)); } #[test] fn test_is_unique_id_not_function() { let node = parse_first_target("SELECT 1"); - #[cfg(feature = "new_parser")] assert!(!StatementRewrite::is_unique_id(node.val())); - #[cfg(not(feature = "new_parser"))] - assert!(!StatementRewrite::is_unique_id(&node)); } #[test] @@ -312,12 +153,6 @@ mod tests { let _guard = set_env_var("NODE_ID", "pgdog-1"); let (sql, plan) = run_test("SELECT pgdog.unique_id()", false); - // Value should be a bigint literal cast - #[cfg(not(feature = "new_parser"))] - assert!( - sql.contains("::bigint"), - "Expected ::bigint cast, got: {sql}" - ); assert!( !sql.contains("pgdog.unique_id"), "Function should be replaced: {sql}" @@ -440,7 +275,6 @@ mod tests { assert_eq!(plan.unique_ids, 1); } - #[cfg(feature = "new_parser")] fn run_test(sql: &str, extended: bool) -> (String, RewritePlan) { let stmt = pg_raw_parse::parse(sql).unwrap(); let mut ps = PreparedStatements::default(); @@ -465,30 +299,4 @@ mod tests { let sql = pg_raw_parse::deparse_stmts(&*ast).unwrap(); (sql, plan) } - - cfg_select! { - not(feature = "new_parser") => { - fn run_test(sql: &str, extended: bool) -> (String, RewritePlan) { - let mut ast = pg_query::parse(sql).unwrap().protobuf; - let mut ps = PreparedStatements::default(); - let schema = default_schema(); - let db_schema = default_db_schema(); - let mut rewrite = StatementRewrite::new(StatementRewriteContext { - stmt: &mut ast, - extended, - prepared: false, - prepared_statements: &mut ps, - schema: &schema, - db_schema: &db_schema, - user: "", - search_path: None, - }); - let plan = rewrite - .maybe_rewrite() - .unwrap(); - (ast.deparse().unwrap(), plan) - } - } - _ => {} - } } diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/update.rs b/pgdog/src/frontend/router/parser/rewrite/statement/update.rs index 4581e421f..3a1f91416 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/update.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/update.rs @@ -1,30 +1,10 @@ -#[cfg(feature = "new_parser")] use indexmap::IndexSet; -#[cfg(not(feature = "new_parser"))] -use std::collections::HashMap; use std::{ops::Deref, sync::Arc}; -#[cfg(not(feature = "new_parser"))] -use pg_query::{ - Node as PgNode, NodeEnum, - protobuf::{ - AExpr, AExprKind, AStar, ColumnRef, DeleteStmt, InsertStmt, LimitOption, List, - OverridingKind, ParamRef, ParseResult, RangeVar, RawStmt, ResTarget, SelectStmt, - SetOperation, String as PgString, UpdateStmt, - }, -}; -#[cfg(feature = "new_parser")] use pg_raw_parse::make::{owned, try_owned}; -#[cfg(feature = "new_parser")] use pg_raw_parse::{DeparseResult, Node, NodeMut, Owned, deparse, nodes, walk}; -#[cfg(not(feature = "new_parser"))] -use pgdog_config::QueryParserEngine; use pgdog_config::RewriteMode; -#[cfg(not(feature = "new_parser"))] -use crate::frontend::router::parser::rewrite::statement::visitor::visit_and_mutate_nodes; -#[cfg(not(feature = "new_parser"))] -use crate::net::FromDataType; use crate::{ frontend::{ BufferedQuery, ClientRequest, @@ -46,10 +26,7 @@ use super::*; pub(crate) struct Statement { pub(crate) ast: Ast, pub(crate) stmt: String, - #[cfg(feature = "new_parser")] pub(crate) params: IndexSet, - #[cfg(not(feature = "new_parser"))] - pub(crate) params: Vec, } impl Statement { @@ -136,17 +113,10 @@ impl ShardingKeyUpdate { return false; } - #[cfg(feature = "new_parser")] - { - self.from_update - .target_list() - .iter() - .any(|rt| rt.name() == Some(&*sharded.column)) - } - #[cfg(not(feature = "new_parser"))] - { - self.insert.mapping.contains_key(&sharded.column) - } + self.from_update + .target_list() + .iter() + .any(|rt| rt.name() == Some(&*sharded.column)) }) } } @@ -159,11 +129,7 @@ pub(crate) struct Inner { pub(crate) check: Statement, /// Delete old row from shard. pub(crate) delete: Statement, - /// Partial insert statement. - #[cfg(not(feature = "new_parser"))] - pub(crate) insert: Insert, /// Update this is being constructed from - #[cfg(feature = "new_parser")] // FIXME(sage): There's no reason we need to own this, but this struct is // ultimately a child of AstInner, where the statement is borrowed from, // so we can't add a lifetime here. We should see if we can pass the update @@ -172,7 +138,6 @@ pub(crate) struct Inner { } impl Inner { - #[cfg(feature = "new_parser")] pub(crate) fn target_table(&self) -> Table<'_> { Table::from( self.from_update @@ -181,14 +146,8 @@ impl Inner { ) } - #[cfg(not(feature = "new_parser"))] - pub(crate) fn target_table(&self) -> Table<'_> { - Table::from(&self.insert.table) - } - /// Build an INSERT statement built from an existing /// UPDATE statement and a row returned by a SELECT statement. - #[cfg(feature = "new_parser")] pub(crate) fn build_insert_request( &self, request: &ClientRequest, @@ -268,204 +227,10 @@ impl Inner { Ok(req) } - #[cfg(not(feature = "new_parser"))] - /// Build an INSERT statement built from an existing - /// UPDATE statement and a row returned by a SELECT statement. - pub(crate) fn build_insert_request( - &self, - request: &ClientRequest, - row_description: &RowDescription, - data_row: &DataRow, - ) -> Result { - self.insert - .build_request(request, row_description, data_row) - } - - #[cfg(feature = "new_parser")] /// Do we have to return the rows to the client? pub(crate) fn is_returning(&self) -> bool { self.from_update.returning_clause().is_some() } - - #[cfg(not(feature = "new_parser"))] - /// Do we have to return the rows to the client? - pub(crate) fn is_returning(&self) -> bool { - !self.insert.returning_list.is_empty() && self.insert.returnin_list_deparsed.is_some() - } -} - -/// Partially built INSERT statement. -#[derive(Debug)] -#[cfg(not(feature = "new_parser"))] -pub(crate) struct Insert { - pub(super) table: RangeVar, - /// 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, -} - -#[cfg(not(feature = "new_parser"))] -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.as_ref()).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(itoa::Buffer::new().format(int).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(ryu::Buffer::new().format(f).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(PgNode { - node: Some(NodeEnum::ResTarget(Box::new(ResTarget { - name: field.name.clone(), - ..Default::default() - }))), - }); - - values.push(PgNode { - node: Some(NodeEnum::ParamRef(ParamRef { - number: bind_idx + 1, - ..Default::default() - })), - }); - - bind_idx += 1; - } - - let insert = InsertStmt { - relation: Some(self.table.clone()), - cols: columns, - select_stmt: Some(Box::new(PgNode { - 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![PgNode { - node: Some(NodeEnum::List(List { items: values })), - }], - ..Default::default() - }))), - })), - returning_list: self.returning_list.clone(), - r#override: OverridingKind::OverridingNotSet.into(), - ..Default::default() - }; - - let table = Table::from(&self.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) - } } impl<'a> StatementRewrite<'a> { @@ -473,51 +238,26 @@ impl<'a> StatementRewrite<'a> { /// in the query. pub(super) fn sharding_key_update( &mut self, - #[cfg(feature = "new_parser")] stmt: &nodes::UpdateStmt, + stmt: &nodes::UpdateStmt, plan: &mut RewritePlan, ) -> Result<(), Error> { if self.schema.shards == 1 || self.schema.rewrite.shard_key == RewriteMode::Ignore { return Ok(()); } - #[cfg(not(feature = "new_parser"))] - let Some(NodeEnum::UpdateStmt(stmt)) = self - .stmt - .stmts - .first() - .and_then(|stmt| stmt.stmt.as_ref().map(|stmt| stmt.node.as_ref())) - .flatten() - 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. - #[cfg(feature = "new_parser")] if let Node::None = stmt.where_clause() { return Err(Error::WhereClauseMissing); } - #[cfg(not(feature = "new_parser"))] - if stmt.where_clause.is_none() { - return Err(Error::WhereClauseMissing); - } - plan.sharding_key_update = Some(create_stmts( - stmt, - value, - #[cfg(not(feature = "new_parser"))] - self.schema.query_parser_engine, - )?); + plan.sharding_key_update = Some(create_stmts(stmt, value)?); } Ok(()) } /// Check if the sharding key could be updated. - #[cfg(feature = "new_parser")] fn sharding_key_update_check( &'a self, stmt: &'a nodes::UpdateStmt, @@ -555,72 +295,10 @@ impl<'a> StatementRewrite<'a> { ))) } } - - #[cfg(not(feature = "new_parser"))] - 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| match Column::try_from(&column.node) { - Ok(mut column) => { - column.qualify(table); - self.schema.tables().get_table(column).is_some() - } - _ => 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.as_ref())) - } 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_old(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. -#[cfg(feature = "new_parser")] fn rewrite_params(node: NodeMut<'_, '_>) -> IndexSet { let mut params = IndexSet::new(); walk::walk_mut(node, |node| { @@ -632,143 +310,7 @@ fn rewrite_params(node: NodeMut<'_, '_>) -> IndexSet { params } -#[cfg(not(feature = "new_parser"))] -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_key(|a| a.1); - - Ok(params - .into_iter() - .map(|(original, _)| original as u16) - .collect()) -} - -#[derive(Debug, Clone)] -#[cfg(not(feature = "new_parser"))] -pub(super) enum UpdateValue { - Value(Box), - 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. -/// -#[cfg(not(feature = "new_parser"))] -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(NodeEnum::ResTarget(target)) = target.node.as_ref() { - 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_old( - 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. -#[cfg(not(feature = "new_parser"))] -fn res_target_to_a_expr(res_target: &ResTarget) -> AExpr { - let column_ref = ColumnRef { - fields: vec![PgNode { - node: Some(NodeEnum::String(PgString { - sval: res_target.name.clone(), - })), - }], - location: res_target.location, - }; - - AExpr { - kind: AExprKind::AexprOp.into(), - name: vec![PgNode { - node: Some(NodeEnum::String(PgString { sval: "=".into() })), - }], - lexpr: Some(Box::new(PgNode { - node: Some(NodeEnum::ColumnRef(column_ref)), - })), - rexpr: res_target.val.clone(), - ..Default::default() - } -} - -#[cfg(not(feature = "new_parser"))] -fn select_star() -> Vec { - vec![PgNode { - node: Some(NodeEnum::ResTarget(Box::new(ResTarget { - name: "".into(), - val: Some(Box::new(PgNode { - node: Some(NodeEnum::ColumnRef(ColumnRef { - fields: vec![PgNode { - node: Some(NodeEnum::AStar(AStar {})), - }], - ..Default::default() - })), - })), - ..Default::default() - }))), - }] -} - -#[cfg(not(feature = "new_parser"))] -fn parse_result(node: NodeEnum) -> ParseResult { - ParseResult { - version: pg_query::PG_VERSION_NUM as i32, - stmts: vec![RawStmt { - stmt: Some(Box::new(PgNode { node: Some(node) })), - stmt_location: 0, - stmt_len: 0, - }], - } -} - /// Deparse an expression node by wrapping it in a SELECT statement. -#[cfg(feature = "new_parser")] fn deparse_expr<'a>(nodes: impl IntoIterator>) -> Result { let node = owned(|mem| { let mut select = mem.make_node::(); @@ -785,52 +327,6 @@ fn deparse_expr<'a>(nodes: impl IntoIterator>) -> Result Result { - Ok(deparse_list( - &[PgNode { - 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. -#[cfg(not(feature = "new_parser"))] -fn deparse_list( - list: &[PgNode], - 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)) -} - -#[cfg(feature = "new_parser")] fn create_stmts<'a>( stmt: &'a nodes::UpdateStmt, new_value: &'a nodes::ResTarget, @@ -921,106 +417,10 @@ fn create_stmts<'a>( }) } -#[cfg(not(feature = "new_parser"))] -fn create_stmts( - stmt: &UpdateStmt, - new_value: &ResTarget, - query_parser_engine: QueryParserEngine, -) -> Result { - let select = SelectStmt { - target_list: select_star(), - from_clause: vec![PgNode { - 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![PgNode { - 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(PgNode { - 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().expect("UPDATE always has table"), - 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 crate::frontend::router::sharding::ShardedTable; - #[cfg(feature = "new_parser")] use indexmap::indexset; - #[cfg(not(feature = "new_parser"))] - use pg_query::parse; use pgdog_config::Rewrite; use crate::backend::schema::Schema; @@ -1029,13 +429,6 @@ mod test { use super::*; - #[cfg(not(feature = "new_parser"))] - macro_rules! indexset { - ($($t:tt)*) => { - vec![$($t)*] - }; - } - fn default_db_schema() -> Schema { Schema::default() } @@ -1065,17 +458,12 @@ mod test { } fn run_test(query: &str) -> Result, Error> { - #[cfg(not(feature = "new_parser"))] - let mut stmt_old = parse(query)?; - #[cfg(feature = "new_parser")] let stmt = pg_raw_parse::parse(query)?; let schema = default_schema(); let db_schema = default_db_schema(); let mut stmts = PreparedStatements::new(); let ctx = StatementRewriteContext { - #[cfg(not(feature = "new_parser"))] - stmt: &mut stmt_old.protobuf, schema: &schema, db_schema: &db_schema, extended: true, @@ -1086,7 +474,6 @@ mod test { }; let mut plan = RewritePlan::default(); StatementRewrite::new(ctx).sharding_key_update( - #[cfg(feature = "new_parser")] match stmt.stmts().next().unwrap() { Node::UpdateStmt(stmt) => stmt, _ => panic!("Not an update"), @@ -1485,76 +872,6 @@ mod test { ); } - #[test] - #[cfg(not(feature = "new_parser"))] - 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] - #[cfg(not(feature = "new_parser"))] - 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(); - std::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] - #[cfg(not(feature = "new_parser"))] - 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] - #[cfg(not(feature = "new_parser"))] - 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. @@ -1569,7 +886,6 @@ mod test { Field::bigint("id"), Field::text("email"), Field::text("other_col"), - #[cfg(feature = "new_parser")] Field::text("other_other_col"), ]); @@ -1578,7 +894,6 @@ mod test { 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 - #[cfg(feature = "new_parser")] data_row.add("other_other_value"); // other_other_col - from existing row // Create a simple query request (not prepared statement) diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/visitor.rs b/pgdog/src/frontend/router/parser/rewrite/statement/visitor.rs deleted file mode 100644 index c17de9687..000000000 --- a/pgdog/src/frontend/router/parser/rewrite/statement/visitor.rs +++ /dev/null @@ -1,313 +0,0 @@ -#![cfg(not(feature = "new_parser"))] -//! 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(crate) 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)?; - } - } - - NodeEnum::ExplainStmt(stmt) => { - if let Some(query) = &mut stmt.query { - visit_and_mutate_node(query, 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/route.rs b/pgdog/src/frontend/router/parser/route.rs index 310253468..f65f9a31c 100644 --- a/pgdog/src/frontend/router/parser/route.rs +++ b/pgdog/src/frontend/router/parser/route.rs @@ -523,7 +523,6 @@ impl ShardWithPriority { } } - #[cfg(feature = "new_parser")] pub(crate) fn new_override_canonical_schema_info(shard: Shard) -> Self { Self { shard, diff --git a/pgdog/src/frontend/router/parser/statement.rs b/pgdog/src/frontend/router/parser/statement.rs index a71ac6189..957a660e8 100644 --- a/pgdog/src/frontend/router/parser/statement.rs +++ b/pgdog/src/frontend/router/parser/statement.rs @@ -1,34 +1,13 @@ -#[cfg(feature = "new_parser")] use std::borrow::Cow; use std::collections::{HashMap, HashSet}; use std::sync::Arc; -#[cfg(feature = "new_parser")] use crate::util::ResultControlFlowExt; -#[cfg(feature = "new_parser")] use itertools::*; -#[cfg(not(feature = "new_parser"))] -use pg_query::Node as PgNode; -#[cfg(not(feature = "new_parser"))] -use pg_query::Node; -#[cfg(all(test, not(feature = "new_parser")))] -use pg_query::protobuf::RawStmt; -#[cfg(not(feature = "new_parser"))] -use pg_query::{ - NodeEnum, - protobuf::{ - self, AConst, AExprKind, BoolExprType, DeleteStmt, FuncCall, InsertStmt, Integer, RangeVar, - SelectStmt, UpdateStmt, a_const::Val, - }, -}; -#[cfg(feature = "new_parser")] use pg_raw_parse::walk::Recurse; -#[cfg(feature = "new_parser")] use pg_raw_parse::{Node, list, nodes, walk}; -#[cfg(feature = "new_parser")] use std::ops::ControlFlow; -#[cfg(feature = "new_parser")] fn advisory_locks_from_func_call( func: &nodes::FuncCall, bind: Option<&Bind>, @@ -113,119 +92,15 @@ fn advisory_locks_from_func_call( }] } -#[cfg(not(feature = "new_parser"))] -fn advisory_locks_from_func_call( - func: &FuncCall, - bind: Option<&Bind>, - values_columns: Option<&ValuesColumns<'_>>, -) -> Vec { - // Only unqualified calls (no schema) map to the real advisory lock builtins. - let (schema, name) = match func.funcname.as_slice() { - [only] => (None, name_of_string_node(only)), - [.., s, n] => (name_of_string_node(s), name_of_string_node(n)), - _ => return Vec::new(), - }; - if schema.is_some() { - return Vec::new(); - } - let Some(name) = name else { - return Vec::new(); - }; - - let (unlock, scope) = match name { - "pg_advisory_lock" - | "pg_advisory_lock_shared" - | "pg_try_advisory_lock" - | "pg_try_advisory_lock_shared" => (false, LockScope::Session), - "pg_advisory_xact_lock" - | "pg_advisory_xact_lock_shared" - | "pg_try_advisory_xact_lock" - | "pg_try_advisory_xact_lock_shared" => (false, LockScope::Transaction), - // Session-scoped unlocks. xact locks can't be released by name; - // Postgres drops them automatically at COMMIT/ROLLBACK. - "pg_advisory_unlock" => (true, LockScope::Session), - "pg_advisory_unlock_all" => { - return vec![AdvisoryLock { - id: None, - unlock: true, - scope: LockScope::Session, - }]; - } - _ => return Vec::new(), - }; - - let Some(arg) = func.args.first() else { - return vec![AdvisoryLock { - id: None, - unlock, - scope, - }]; - }; - - // Fast path: the key is a literal / param / cast we can resolve directly. - if let Some(id) = integer_arg(arg, bind) { - return vec![AdvisoryLock { - id: Some(id), - unlock, - scope, - }]; - } - - // If the argument is a parameter placeholder ($1) and we have no Bind message, - // this is just a prepared statement being parsed — the lock isn't actually - // being taken yet. Return empty so we don't route as if a lock is held. - if bind.is_none() && is_param_ref(arg) { - return Vec::new(); - } - - // Slow path: `SELECT pg_advisory_lock(value) FROM (VALUES (1),(2)) AS t(value)`. - // The function is called once per row, so we emit one lock per resolved value. - if let Some(NodeEnum::ColumnRef(cref)) = arg.node.as_ref() - && let Some(col) = last_column_name(&cref.fields) - && let Some(rows) = values_columns.and_then(|m| m.get(col)) - { - return rows - .iter() - // Skip unresolvable param refs when there is no Bind. - .filter(|v| bind.is_some() || !is_param_ref(v)) - .map(|v| AdvisoryLock { - id: integer_arg(v, bind), - unlock, - scope, - }) - .collect(); - } - - vec![AdvisoryLock { - id: None, - unlock, - scope, - }] -} - -#[cfg(feature = "new_parser")] fn last_column_name<'a>(fields: impl IntoIterator>) -> Option<&'a str> { fields.into_iter().last().and_then(Node::as_str) } -#[cfg(not(feature = "new_parser"))] -fn last_column_name(fields: &[Node]) -> Option<&str> { - match fields.last()?.node.as_ref()? { - NodeEnum::String(protobuf::String { sval }) => Some(sval.as_str()), - _ => None, - } -} - -#[cfg(feature = "new_parser")] -type ValuesColumns<'a> = std::collections::HashMap, Vec>>; - /// Map from unqualified VALUES column alias to the list of value nodes — one /// per row — introduced by a `FROM (VALUES (...), ...) AS t(col, ...)` in the /// current SELECT's FROM clause. -#[cfg(not(feature = "new_parser"))] -type ValuesColumns<'a> = std::collections::HashMap<&'a str, Vec<&'a PgNode>>; +type ValuesColumns<'a> = std::collections::HashMap, Vec>>; -#[cfg(feature = "new_parser")] fn collect_values_columns(stmt: &nodes::SelectStmt) -> Option> { let Node::RangeSubselect(rs) = stmt.from_clause().into_iter().exactly_one().ok()? else { return None; @@ -261,56 +136,6 @@ fn collect_values_columns(stmt: &nodes::SelectStmt) -> Option> Some(values) } -#[cfg(not(feature = "new_parser"))] -fn collect_values_columns(stmt: &SelectStmt) -> ValuesColumns<'_> { - let mut out: ValuesColumns<'_> = ValuesColumns::default(); - for node in &stmt.from_clause { - let Some(NodeEnum::RangeSubselect(rs)) = node.node.as_ref() else { - continue; - }; - let Some(alias) = rs.alias.as_ref() else { - continue; - }; - let Some(subquery) = rs.subquery.as_deref() else { - continue; - }; - let Some(NodeEnum::SelectStmt(inner)) = subquery.node.as_ref() else { - continue; - }; - if inner.values_lists.is_empty() { - continue; - } - let colnames: Vec<&str> = alias - .colnames - .iter() - .filter_map(|n| match n.node.as_ref()? { - NodeEnum::String(protobuf::String { sval }) => Some(sval.as_str()), - _ => None, - }) - .collect(); - for row in &inner.values_lists { - let Some(NodeEnum::List(list)) = row.node.as_ref() else { - continue; - }; - for (idx, item) in list.items.iter().enumerate() { - if let Some(col) = colnames.get(idx) { - out.entry(*col).or_default().push(item); - } - } - } - } - out -} - -#[cfg(not(feature = "new_parser"))] -fn name_of_string_node(node: &Node) -> Option<&str> { - match node.node.as_ref()? { - NodeEnum::String(protobuf::String { sval }) => Some(sval.as_str()), - _ => None, - } -} - -#[cfg(feature = "new_parser")] fn integer_arg(node: Node<'_>, bind: Option<&Bind>) -> Option { match node { Node::A_Const(a) => a.val()?.numeric_value(), @@ -324,30 +149,7 @@ fn integer_arg(node: Node<'_>, bind: Option<&Bind>) -> Option { } } -#[cfg(not(feature = "new_parser"))] -fn integer_arg(node: &Node, bind: Option<&Bind>) -> Option { - match node.node.as_ref()? { - NodeEnum::AConst(AConst { val: Some(val), .. }) => match val { - Val::Ival(Integer { ival }) => Some(*ival as i64), - // pg_query stores integers wider than i32 (e.g. bigint keys) as Float - // with a numeric string payload. - Val::Fval(f) => f.fval.parse().ok(), - _ => None, - }, - NodeEnum::TypeCast(cast) => integer_arg(cast.arg.as_deref()?, bind), - // Resolve $N via the Bind message. pg_query numbers parameters from 1. - NodeEnum::ParamRef(param_ref) => { - let bind = bind?; - let index = (param_ref.number as usize).checked_sub(1)?; - let param = bind.parameter(index).ok().flatten()?; - param.decode::() - } - _ => None, - } -} - /// Check whether a node is (or wraps) a parameter placeholder (`$N`). -#[cfg(feature = "new_parser")] fn is_param_ref(node: Node<'_>) -> bool { match node { Node::ParamRef(_) => true, @@ -356,16 +158,6 @@ fn is_param_ref(node: Node<'_>) -> bool { } } -/// Check whether a node is (or wraps) a parameter placeholder (`$N`). -#[cfg(not(feature = "new_parser"))] -fn is_param_ref(node: &Node) -> bool { - match node.node.as_ref() { - Some(NodeEnum::ParamRef(_)) => true, - Some(NodeEnum::TypeCast(cast)) => cast.arg.as_deref().is_some_and(is_param_ref), - _ => false, - } -} - use super::{ super::sharding::Value as ShardingValue, Column, Error, Table, Value, explain_trace::ExplainRecorder, @@ -455,7 +247,6 @@ struct SearchContext<'a> { impl<'a> SearchContext<'a> { /// Build context from a FROM clause, extracting table aliases. - #[cfg(feature = "new_parser")] fn from_from_clause(nodes: &'a list::NodeList) -> Self { let mut aliases = HashMap::new(); @@ -472,29 +263,6 @@ impl<'a> SearchContext<'a> { Self { aliases, table } } - #[cfg(not(feature = "new_parser"))] - fn from_from_clause_old(nodes: &'a [PgNode]) -> Self { - let mut ctx = Self::default(); - ctx.extract_aliases(nodes); - - // Try to get the primary table for simple queries - if nodes.len() == 1 - && let Some(table) = nodes.first().and_then(|n| Table::try_from(n).ok()) - { - ctx.table = Some(table); - } - - ctx - } - - #[cfg(not(feature = "new_parser"))] - fn extract_aliases(&mut self, nodes: &'a [PgNode]) { - for node in nodes { - self.extract_alias_from_node_old(node); - } - } - - #[cfg(feature = "new_parser")] fn extract_alias_from_node(aliases: &mut HashMap<&'a str, Table<'a>>, node: Node<'a>) { match node { Node::RangeVar(rv) if let Some(alias) = rv.alias() => { @@ -525,41 +293,6 @@ impl<'a> SearchContext<'a> { } } - #[cfg(not(feature = "new_parser"))] - fn extract_alias_from_node_old(&mut self, node: &'a PgNode) { - 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_old(larg); - } - if let Some(ref rarg) = join.rarg { - self.extract_alias_from_node_old(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() @@ -571,12 +304,6 @@ enum SearchResult<'a> { Column(Column<'a>), Value(Value<'a>), Values(Vec>), - #[cfg(not(feature = "new_parser"))] - Match(Shard), - #[cfg(not(feature = "new_parser"))] - Matches(Vec), - #[cfg(not(feature = "new_parser"))] - None, } struct ValueIterator<'a, 'b> { @@ -607,30 +334,6 @@ impl<'a, 'b> Iterator for ValueIterator<'a, 'b> { } impl<'a> SearchResult<'a> { - #[cfg(not(feature = "new_parser"))] - fn is_none(&self) -> bool { - matches!(self, Self::None) - } - - #[cfg(not(feature = "new_parser"))] - fn is_match(&self) -> bool { - matches!(self, Self::Match(_) | Self::Matches(_)) - } - - #[cfg(not(feature = "new_parser"))] - 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, @@ -639,14 +342,6 @@ impl<'a> SearchResult<'a> { } } -#[cfg(not(feature = "new_parser"))] -enum Statement<'a> { - Select(&'a SelectStmt), - Update(&'a UpdateStmt), - Delete(&'a DeleteStmt), - Insert(&'a InsertStmt), -} - /// Context for looking up table columns from the database schema. /// Used for INSERT statements without explicit column lists. pub struct SchemaLookupContext<'a> { @@ -659,10 +354,7 @@ pub struct SchemaLookupContext<'a> { } pub struct StatementParser<'a, 'b, 'c> { - #[cfg(not(feature = "new_parser"))] - stmt: Statement<'a>, - #[cfg(feature = "new_parser")] - new_stmt: pg_raw_parse::Node<'a>, + stmt: pg_raw_parse::Node<'a>, bind: Option<&'b Bind>, schema: &'b ShardingSchema, recorder: Option<&'c mut ExplainRecorder>, @@ -679,18 +371,14 @@ pub struct StatementParser<'a, 'b, 'c> { } impl<'a, 'b: 'a, 'c> StatementParser<'a, 'b, 'c> { - fn new( - #[cfg(not(feature = "new_parser"))] stmt: Statement<'a>, - #[cfg(feature = "new_parser")] new_stmt: Node<'a>, + pub(crate) fn new( + stmt: Node<'a>, bind: Option<&'b Bind>, schema: &'b ShardingSchema, recorder: Option<&'c mut ExplainRecorder>, ) -> Self { Self { - #[cfg(not(feature = "new_parser"))] stmt, - #[cfg(feature = "new_parser")] - new_stmt, bind, schema, recorder, @@ -748,78 +436,6 @@ impl<'a, 'b: 'a, 'c> StatementParser<'a, 'b, 'c> { self } - pub(crate) fn from_select( - #[cfg(not(feature = "new_parser"))] stmt: &'a SelectStmt, - #[cfg(feature = "new_parser")] new_stmt: Node<'a>, - bind: Option<&'b Bind>, - schema: &'b ShardingSchema, - recorder: Option<&'c mut ExplainRecorder>, - ) -> Self { - Self::new( - #[cfg(not(feature = "new_parser"))] - Statement::Select(stmt), - #[cfg(feature = "new_parser")] - new_stmt, - bind, - schema, - recorder, - ) - } - - pub(crate) fn from_update( - #[cfg(not(feature = "new_parser"))] stmt: &'a UpdateStmt, - #[cfg(feature = "new_parser")] stmt: Node<'a>, - bind: Option<&'b Bind>, - schema: &'b ShardingSchema, - recorder: Option<&'c mut ExplainRecorder>, - ) -> Self { - Self::new( - #[cfg(not(feature = "new_parser"))] - Statement::Update(stmt), - #[cfg(feature = "new_parser")] - stmt, - bind, - schema, - recorder, - ) - } - - pub(crate) fn from_delete( - #[cfg(not(feature = "new_parser"))] stmt: &'a DeleteStmt, - #[cfg(feature = "new_parser")] stmt: Node<'a>, - bind: Option<&'b Bind>, - schema: &'b ShardingSchema, - recorder: Option<&'c mut ExplainRecorder>, - ) -> Self { - Self::new( - #[cfg(not(feature = "new_parser"))] - Statement::Delete(stmt), - #[cfg(feature = "new_parser")] - stmt, - bind, - schema, - recorder, - ) - } - - pub(crate) fn from_insert( - #[cfg(not(feature = "new_parser"))] stmt: &'a InsertStmt, - #[cfg(feature = "new_parser")] stmt: Node<'a>, - bind: Option<&'b Bind>, - schema: &'b ShardingSchema, - recorder: Option<&'c mut ExplainRecorder>, - ) -> Self { - Self::new( - #[cfg(not(feature = "new_parser"))] - Statement::Insert(stmt), - #[cfg(feature = "new_parser")] - stmt, - bind, - schema, - recorder, - ) - } - /// Record a sharding key match. fn record_sharding_key(&mut self, shard: &Shard, column: Column<'_>, value: &Value<'_>) { self.hooks @@ -841,27 +457,6 @@ impl<'a, 'b: 'a, 'c> StatementParser<'a, 'b, 'c> { } } - #[cfg(test)] - fn from_raw( - #[cfg(not(feature = "new_parser"))] raw: &'a RawStmt, - #[cfg(feature = "new_parser")] stmt: Node<'a>, - bind: Option<&'b Bind>, - schema: &'b ShardingSchema, - recorder: Option<&'c mut ExplainRecorder>, - ) -> Result { - #[cfg(not(feature = "new_parser"))] - return 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), - }; - - #[cfg(feature = "new_parser")] - Ok(Self::new(stmt, bind, schema, recorder)) - } - pub fn shard(&mut self) -> Result, Error> { // Omnisharded config overrides sharded: if all tables are omnisharded, // don't try to find a sharding key - let omnisharded routing handle it @@ -869,16 +464,7 @@ impl<'a, 'b: 'a, 'c> StatementParser<'a, 'b, 'c> { return Ok(None); } - #[cfg(feature = "new_parser")] - let result = self.shard_stmt(self.new_stmt)?; - - #[cfg(not(feature = "new_parser"))] - 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), - }?; + let result = self.shard_stmt(self.stmt)?; // Key-based sharding succeeded if result.is_some() { @@ -974,14 +560,12 @@ impl<'a, 'b: 'a, 'c> StatementParser<'a, 'b, 'c> { } // Are we running? Or walking? MAKE UP YOUR MIND DAMMIT - #[cfg(feature = "new_parser")] fn run_walk(&self) -> Walk<'a> { let mut walk = Walk::default(); - self.walk_stmt(self.new_stmt, &mut walk); + self.walk_stmt(self.stmt, &mut walk); walk } - #[cfg(feature = "new_parser")] fn walk_stmt(&self, stmt: Node<'a>, walk: &mut Walk<'a>) { let values_columns = match stmt { Node::SelectStmt(s) => collect_values_columns(s), @@ -995,292 +579,30 @@ impl<'a, 'b: 'a, 'c> StatementParser<'a, 'b, 'c> { Recurse::no() } - // Extract any advisory locks that this function call may - // represent, using the values clause of the current statement - Node::FuncCall(func) => { - walk.advisory_locks.extend(advisory_locks_from_func_call( - func, - self.bind, - values_columns.as_ref(), - )); - Recurse::no() - } - - Node::RangeVar(r) => { - walk.tables.push(Table::from(r)); - Recurse::yes() - } - - _ => Recurse::yes(), - }); - } - - #[cfg(not(feature = "new_parser"))] - fn run_walk(&self) -> Walk<'a> { - let mut walk = Walk::default(); - match self.stmt { - Statement::Select(stmt) => self.walk_select(stmt, &mut walk), - Statement::Update(stmt) => self.walk_update(stmt, &mut walk), - Statement::Delete(stmt) => self.walk_delete(stmt, &mut walk), - Statement::Insert(stmt) => self.walk_insert(stmt, &mut walk), - } - walk - } - - #[cfg(not(feature = "new_parser"))] - fn walk_select(&self, stmt: &'a SelectStmt, walk: &mut Walk<'a>) { - // Build a VALUES-column lookup for the current SELECT scope so a call - // like `pg_advisory_lock(value) FROM (VALUES (1),(2)) AS t(value)` - // can be expanded to one lock per row. - let values_columns = collect_values_columns(stmt); - let values = if values_columns.is_empty() { - None - } else { - Some(&values_columns) - }; - - // Handle UNION/INTERSECT/EXCEPT - if let Some(ref larg) = stmt.larg { - self.walk_select(larg, walk); - } - if let Some(ref rarg) = stmt.rarg { - self.walk_select(rarg, walk); - } - - // Target list — advisory lock function calls usually live here. - for node in &stmt.target_list { - self.walk_node(node, walk, values); - } - - // FROM clause - for node in &stmt.from_clause { - self.walk_node(node, walk, values); - } - - // 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 - && let Some(ref ctequery) = cte_expr.ctequery - && let Some(NodeEnum::SelectStmt(ref inner_select)) = ctequery.node - { - self.walk_select(inner_select, walk); - } - } - } - - // WHERE clause subqueries - if let Some(ref where_clause) = stmt.where_clause { - self.walk_node(where_clause, walk, values); - } - } - - #[cfg(not(feature = "new_parser"))] - fn walk_update(&self, stmt: &'a UpdateStmt, walk: &mut Walk<'a>) { - if let Some(ref relation) = stmt.relation { - walk.tables.push(Table::from(relation)); - } - - for node in &stmt.from_clause { - self.walk_node(node, walk, None); - } - - 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 - && let Some(ref ctequery) = cte_expr.ctequery - && let Some(NodeEnum::SelectStmt(ref inner_select)) = ctequery.node - { - self.walk_select(inner_select, walk); - } - } - } - - if let Some(ref where_clause) = stmt.where_clause { - self.walk_node(where_clause, walk, None); - } - } - - #[cfg(not(feature = "new_parser"))] - fn walk_delete(&self, stmt: &'a DeleteStmt, walk: &mut Walk<'a>) { - if let Some(ref relation) = stmt.relation { - walk.tables.push(Table::from(relation)); - } - - for node in &stmt.using_clause { - self.walk_node(node, walk, None); - } - - 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 - && let Some(ref ctequery) = cte_expr.ctequery - && let Some(NodeEnum::SelectStmt(ref inner_select)) = ctequery.node - { - self.walk_select(inner_select, walk); - } - } - } - - if let Some(ref where_clause) = stmt.where_clause { - self.walk_node(where_clause, walk, None); - } - } - - #[cfg(not(feature = "new_parser"))] - fn walk_insert(&self, stmt: &'a InsertStmt, walk: &mut Walk<'a>) { - if let Some(ref relation) = stmt.relation { - walk.tables.push(Table::from(relation)); - } - - 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 - && let Some(ref ctequery) = cte_expr.ctequery - && let Some(NodeEnum::SelectStmt(ref inner_select)) = ctequery.node - { - self.walk_select(inner_select, walk); - } - } - } - - if let Some(ref select_stmt) = stmt.select_stmt - && let Some(NodeEnum::SelectStmt(ref inner_select)) = select_stmt.node - { - self.walk_select(inner_select, walk); - } - } - - #[cfg(not(feature = "new_parser"))] - fn walk_node(&self, node: &'a Node, walk: &mut Walk<'a>, values: Option<&ValuesColumns<'a>>) { - match &node.node { - Some(NodeEnum::RangeVar(range_var)) => { - walk.tables.push(Table::from(range_var)); - } - Some(NodeEnum::JoinExpr(join)) => { - if let Some(ref larg) = join.larg { - self.walk_node(larg, walk, values); - } - if let Some(ref rarg) = join.rarg { - self.walk_node(rarg, walk, values); - } - } - Some(NodeEnum::RangeSubselect(subselect)) => { - if let Some(ref subquery) = subselect.subquery - && let Some(NodeEnum::SelectStmt(ref inner_select)) = subquery.node - { - self.walk_select(inner_select, walk); - } - } - Some(NodeEnum::SubLink(sublink)) => { - if let Some(ref subselect) = sublink.subselect - && let Some(NodeEnum::SelectStmt(ref inner_select)) = subselect.node - { - self.walk_select(inner_select, walk); - } - } - Some(NodeEnum::SelectStmt(inner_select)) => { - self.walk_select(inner_select, walk); - } - Some(NodeEnum::BoolExpr(bool_expr)) => { - for arg in &bool_expr.args { - self.walk_node(arg, walk, values); - } - } - Some(NodeEnum::AExpr(a_expr)) => { - if let Some(ref lexpr) = a_expr.lexpr { - self.walk_node(lexpr, walk, values); - } - if let Some(ref rexpr) = a_expr.rexpr { - self.walk_node(rexpr, walk, values); - } - } - Some(NodeEnum::ResTarget(res)) => { - if let Some(ref val) = res.val { - self.walk_node(val, walk, values); - } - } - Some(NodeEnum::TypeCast(cast)) => { - if let Some(ref arg) = cast.arg { - self.walk_node(arg, walk, values); - } - } - Some(NodeEnum::NullTest(test)) => { - if let Some(ref arg) = test.arg { - self.walk_node(arg, walk, values); - } - } - Some(NodeEnum::FuncCall(func)) => { - for lock in advisory_locks_from_func_call(func, self.bind, values) { - walk.advisory_locks.insert(lock); - } - for arg in &func.args { - self.walk_node(arg, walk, values); - } - } - Some(NodeEnum::List(list)) => { - for item in &list.items { - self.walk_node(item, walk, values); - } - } - _ => {} - } - } - - #[cfg(feature = "new_parser")] - fn shard_stmt(&mut self, stmt: Node<'a>) -> Result, Error> { - self.search_stmt(stmt).break_value().transpose() - } - - #[cfg(not(feature = "new_parser"))] - fn shard_select(&mut self, stmt: &'a SelectStmt) -> Result, Error> { - let ctx = SearchContext::from_from_clause_old(&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), - } - } - - #[cfg(not(feature = "new_parser"))] - fn shard_update(&mut self, stmt: &'a UpdateStmt) -> Result, Error> { - let ctx = self.context_from_relation_old(&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), - } - } - - #[cfg(not(feature = "new_parser"))] - fn shard_delete(&mut self, stmt: &'a DeleteStmt) -> Result, Error> { - let ctx = self.context_from_relation_old(&stmt.relation); - let result = self.search_delete_stmt(stmt, &ctx)?; + // Extract any advisory locks that this function call may + // represent, using the values clause of the current statement + Node::FuncCall(func) => { + walk.advisory_locks.extend(advisory_locks_from_func_call( + func, + self.bind, + values_columns.as_ref(), + )); + Recurse::no() + } - match result { - SearchResult::Match(shard) => Ok(Some(shard)), - SearchResult::Matches(shards) => Ok(Self::converge(&shards)), - _ => Ok(None), - } - } + Node::RangeVar(r) => { + walk.tables.push(Table::from(r)); + Recurse::yes() + } - #[cfg(not(feature = "new_parser"))] - fn shard_insert(&mut self, stmt: &'a InsertStmt) -> Result, Error> { - let ctx = self.context_from_relation_old(&stmt.relation); - let result = self.search_insert_stmt(stmt, &ctx)?; + _ => Recurse::yes(), + }); + } - match result { - SearchResult::Match(shard) => Ok(Some(shard)), - SearchResult::Matches(shards) => Ok(Self::converge(&shards)), - _ => Ok(None), - } + fn shard_stmt(&mut self, stmt: Node<'a>) -> Result, Error> { + self.search_stmt(stmt).break_value().transpose() } - #[cfg(feature = "new_parser")] fn context_from_relation(&self, relation: Option<&'a nodes::RangeVar>) -> SearchContext<'a> { let mut ctx = SearchContext::default(); if let Some(range_var) = relation { @@ -1294,19 +616,6 @@ impl<'a, 'b: 'a, 'c> StatementParser<'a, 'b, 'c> { ctx } - #[cfg(not(feature = "new_parser"))] - fn context_from_relation_old(&self, relation: &'a Option) -> SearchContext<'a> { - let mut ctx = SearchContext::default(); - if let Some(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() { @@ -1506,216 +815,6 @@ impl<'a, 'b: 'a, 'c> StatementParser<'a, 'b, 'c> { std::mem::take(&mut self.pending_lookups) } - #[cfg(not(feature = "new_parser"))] - fn select_search( - &mut self, - node: &'a pg_query::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_old(&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() - && let Some(ref table) = ctx.table - { - column.qualify(*table); - } - - Ok(SearchResult::Column(column)) - } - - Some(NodeEnum::AExpr(ref expr)) => { - let kind = expr.kind(); - let supported = match kind { - // Kind carries the full semantic; no operator name to check. - AExprKind::AexprNotDistinct => true, - // Operator-based kinds: accept equality only. - AExprKind::AexprOp | AExprKind::AexprIn | AExprKind::AexprOpAny => { - expr.name - .first() - .map(|node| match node.node { - Some(NodeEnum::String(ref string)) => string.sval.as_str(), - _ => "", - }) - .unwrap_or_default() - == "=" - } - _ => false, - }; - - 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.get_sharded_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), - } - } - - #[cfg(feature = "new_parser")] fn search_stmt(&mut self, stmt: Node<'a>) -> ControlFlow> { use nodes::{A_Expr_Kind, BoolExprType}; @@ -1804,7 +903,6 @@ impl<'a, 'b: 'a, 'c> StatementParser<'a, 'b, 'c> { } } - #[cfg(feature = "new_parser")] fn search_expr( &mut self, node: Node<'a>, @@ -1861,57 +959,6 @@ impl<'a, 'b: 'a, 'c> StatementParser<'a, 'b, 'c> { } } - /// Search a SELECT statement with its own context. - #[cfg(not(feature = "new_parser"))] - 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_old(&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_old(&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, @@ -1941,80 +988,7 @@ impl<'a, 'b: 'a, 'c> StatementParser<'a, 'b, 'c> { Ok(shard) } - /// Search an UPDATE statement for sharding keys. - #[cfg(not(feature = "new_parser"))] - 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. - #[cfg(not(feature = "new_parser"))] - 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) - } - /// Get column names from the INSERT statement, or look them up from schema if not specified. - #[cfg(feature = "new_parser")] fn get_insert_columns( &self, stmt: &'a nodes::InsertStmt, @@ -2048,36 +1022,6 @@ impl<'a, 'b: 'a, 'c> StatementParser<'a, 'b, 'c> { } } - #[cfg(not(feature = "new_parser"))] - fn get_insert_columns(&self, stmt: &InsertStmt, ctx: &SearchContext<'_>) -> Vec { - // First try to get columns from the INSERT statement itself - let cols: Vec = stmt - .cols - .iter() - .filter_map(|node| match &node.node { - Some(NodeEnum::ResTarget(target)) => Some(target.name.clone()), - _ => None, - }) - .collect(); - - if !cols.is_empty() { - return cols; - } - - // No columns specified in INSERT, try to look them up from schema - if let (Some(table), Some(schema_lookup)) = (ctx.table, &self.schema_lookup) - && let Some(relation) = - schema_lookup - .db_schema - .table(table, schema_lookup.user, schema_lookup.search_path) - { - return relation.column_names().map(String::from).collect(); - } - - vec![] - } - - #[cfg(feature = "new_parser")] fn search_insert_stmt(&mut self, stmt: &'a nodes::InsertStmt) -> Result, Error> { let ctx = self.context_from_relation(stmt.relation()); @@ -2152,134 +1096,6 @@ impl<'a, 'b: 'a, 'c> StatementParser<'a, 'b, 'c> { Ok(None) } } - - /// Search an INSERT statement for sharding keys. - #[cfg(not(feature = "new_parser"))] - fn search_insert_stmt( - &mut self, - stmt: &'a InsertStmt, - ctx: &SearchContext<'a>, - ) -> Result, Error> { - // Schema-based routing takes priority for INSERTs - if let Some(table) = ctx.table - && let Some(schema) = self.schema.schemas.get(table.schema()) - { - return Ok(SearchResult::Match(schema.shard().into())); - } - - // Get the column names from INSERT INTO table (col1, col2, ...) or from schema - let columns = self.get_insert_columns(stmt, ctx); - - // Handle different INSERT forms - if let Some(ref select_node) = stmt.select_stmt { - if let Some(NodeEnum::SelectStmt(ref select_stmt)) = select_node.node { - // Multi-row VALUES broadcasts to all shards - if select_stmt.values_lists.len() > 1 { - return Ok(SearchResult::Match(Shard::All)); - } - - // INSERT...SELECT (no VALUES): try to extract sharding key from target list - if select_stmt.values_lists.is_empty() { - // Try to extract constants from SELECT target list - if !select_stmt.target_list.is_empty() { - for (pos, target_node) in select_stmt.target_list.iter().enumerate() { - if let Some(NodeEnum::ResTarget(ref target)) = target_node.node - && let Some(column_name) = columns.get(pos) - { - let table_name = ctx.table.map(|t| t.name); - let table_schema = ctx.table.and_then(|t| t.schema); - let sharded_table = self.get_sharded_table_by_name( - column_name.as_str(), - table_name, - table_schema, - ); - - if sharded_table.is_some() - && let Some(ref val) = target.val - && let Ok(value) = Value::try_from(val.as_ref()) - && let Some(shard) = - self.compute_shard_for_table(sharded_table, value)? - { - return Ok(SearchResult::Match(shard)); - } - } - } - } - - // INSERT...SELECT without extractable key broadcasts - return Ok(SearchResult::Match(Shard::All)); - } - } - } else { - // No select_stmt (DEFAULT VALUES) broadcasts to all shards - return Ok(SearchResult::Match(Shard::All)); - } - - // 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); - } - } - } - - // The select_stmt field contains either VALUES or a SELECT subquery - if let Some(ref select_node) = stmt.select_stmt - && 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 table_name = ctx.table.map(|t| t.name); - let table_schema = ctx.table.and_then(|t| t.schema); - let sharded_table = self.get_sharded_table_by_name( - column_name.as_str(), - table_name, - table_schema, - ); - - if sharded_table.is_some() { - // Try to extract the value directly - if let Ok(value) = Value::try_from(value_node) - && let Some(shard) = - self.compute_shard_for_table(sharded_table, value)? - { - 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); - } - } - } - } - } - } - - // Round-robin fallback: if table is sharded but no sharding key found, - // pick a shard at random - if let Some(table) = ctx.table { - let tables = Tables::new(self.schema); - if tables.sharded(table).is_some() { - return Ok(SearchResult::Match(Shard::Direct( - round_robin::next() % self.schema.shards, - ))); - } - } - - Ok(SearchResult::None) - } } #[cfg(test)] @@ -2350,27 +1166,9 @@ mod test { fn run_test(stmt: &str, bind: Option<&Bind>) -> Result, Error> { let schema = test_schema(); - #[cfg(not(feature = "new_parser"))] - let raw = pg_query::parse(stmt) - .unwrap() - .protobuf - .stmts - .first() - .cloned() - .unwrap(); - #[cfg(feature = "new_parser")] let raw = pg_raw_parse::parse(stmt).unwrap(); - #[cfg(feature = "new_parser")] let stmt = raw.stmts().next().unwrap(); - let mut parser = StatementParser::from_raw( - #[cfg(not(feature = "new_parser"))] - &raw, - #[cfg(feature = "new_parser")] - stmt, - bind, - &schema, - None, - )?; + let mut parser = StatementParser::new(stmt, bind, &schema, None); parser.shard() } @@ -2381,28 +1179,9 @@ mod test { bind: Option<&Bind>, schema: &ShardingSchema, ) -> (Option, Vec) { - #[cfg(not(feature = "new_parser"))] - let raw = pg_query::parse(stmt) - .unwrap() - .protobuf - .stmts - .first() - .cloned() - .unwrap(); - #[cfg(feature = "new_parser")] let raw = pg_raw_parse::parse(stmt).unwrap(); - #[cfg(feature = "new_parser")] let stmt = raw.stmts().next().unwrap(); - let mut parser = StatementParser::from_raw( - #[cfg(not(feature = "new_parser"))] - &raw, - #[cfg(feature = "new_parser")] - stmt, - bind, - schema, - None, - ) - .unwrap(); + let mut parser = StatementParser::new(stmt, bind, schema, None); let shard = parser.shard().unwrap(); (shard, parser.take_pending_lookups()) } @@ -3367,27 +2146,9 @@ mod test { ]), ..Default::default() }; - #[cfg(not(feature = "new_parser"))] - let raw = pg_query::parse(stmt) - .unwrap() - .protobuf - .stmts - .first() - .cloned() - .unwrap(); - #[cfg(feature = "new_parser")] let raw = pg_raw_parse::parse(stmt).unwrap(); - #[cfg(feature = "new_parser")] let stmt = raw.stmts().next().unwrap(); - let mut parser = StatementParser::from_raw( - #[cfg(not(feature = "new_parser"))] - &raw, - #[cfg(feature = "new_parser")] - stmt, - bind, - &schema, - None, - )?; + let mut parser = StatementParser::new(stmt, bind, &schema, None); parser.shard() } @@ -3497,27 +2258,9 @@ mod test { ), ..Default::default() }; - #[cfg(not(feature = "new_parser"))] - let raw = pg_query::parse(stmt) - .unwrap() - .protobuf - .stmts - .first() - .cloned() - .unwrap(); - #[cfg(feature = "new_parser")] let raw = pg_raw_parse::parse(stmt).unwrap(); - #[cfg(feature = "new_parser")] let stmt = raw.stmts().next().unwrap(); - let mut parser = StatementParser::from_raw( - #[cfg(not(feature = "new_parser"))] - &raw, - #[cfg(feature = "new_parser")] - stmt, - bind, - &schema, - None, - )?; + let mut parser = StatementParser::new(stmt, bind, &schema, None); parser.shard() } @@ -3671,28 +2414,10 @@ mod test { user: "test", search_path: None, }; - #[cfg(not(feature = "new_parser"))] - let raw = pg_query::parse(stmt) - .unwrap() - .protobuf - .stmts - .first() - .cloned() - .unwrap(); - #[cfg(feature = "new_parser")] let raw = pg_raw_parse::parse(stmt).unwrap(); - #[cfg(feature = "new_parser")] let stmt = raw.stmts().next().unwrap(); - let mut parser = StatementParser::from_raw( - #[cfg(not(feature = "new_parser"))] - &raw, - #[cfg(feature = "new_parser")] - stmt, - bind, - &sharding_schema, - None, - )? - .with_schema_lookup(schema_lookup); + let mut parser = StatementParser::new(stmt, bind, &sharding_schema, None) + .with_schema_lookup(schema_lookup); parser.shard() } @@ -3809,28 +2534,9 @@ mod test { fn run_is_sharded_test(stmt: &str) -> bool { let schema = make_omnisharded_sharding_schema(); let db_schema = make_omnisharded_db_schema(); - #[cfg(not(feature = "new_parser"))] - let raw = pg_query::parse(stmt) - .unwrap() - .protobuf - .stmts - .first() - .cloned() - .unwrap(); - #[cfg(feature = "new_parser")] let raw = pg_raw_parse::parse(stmt).unwrap(); - #[cfg(feature = "new_parser")] let stmt = raw.stmts().next().unwrap(); - let mut parser = StatementParser::from_raw( - #[cfg(not(feature = "new_parser"))] - &raw, - #[cfg(feature = "new_parser")] - stmt, - None, - &schema, - None, - ) - .unwrap(); + let mut parser = StatementParser::new(stmt, None, &schema, None); parser.is_sharded(&db_schema, "test", None) } @@ -3904,33 +2610,16 @@ mod test { mod advisory_locks { use super::*; - #[cfg(not(feature = "new_parser"))] - use pg_query::parse; fn locks(query: &str) -> Vec { locks_with_bind(query, None) } fn locks_with_bind(query: &str, bind: Option<&Bind>) -> Vec { - #[cfg(not(feature = "new_parser"))] - let ast = parse(query).unwrap().protobuf; let schema = ShardingSchema::default(); - #[cfg(not(feature = "new_parser"))] - let raw = ast.stmts.first().unwrap(); - #[cfg(feature = "new_parser")] let raw = pg_raw_parse::parse(query).unwrap(); - #[cfg(feature = "new_parser")] let stmt = raw.stmts().next().unwrap(); - let mut parser = StatementParser::from_raw( - #[cfg(not(feature = "new_parser"))] - raw, - #[cfg(feature = "new_parser")] - stmt, - bind, - &schema, - None, - ) - .unwrap(); + let mut parser = StatementParser::new(stmt, bind, &schema, None); let mut v: Vec<_> = parser.extract_advisory_locks().iter().copied().collect(); v.sort_by_key(|l| (l.id, l.unlock)); v @@ -3966,7 +2655,7 @@ mod test { #[test] fn bigint_argument() { - // Values larger than i32 are encoded as Float in pg_query. + // Values larger than i32 are encoded as Float in PG internally. assert_eq!( locks("SELECT pg_advisory_lock(9000000000)"), vec![session(Some(9_000_000_000), false)], @@ -4147,7 +2836,6 @@ mod test { } #[test] - #[cfg(feature = "new_parser")] fn advisory_lock_from_values_without_explicit_column_name() { assert_eq!( locks("SELECT pg_advisory_lock(column1) FROM (VALUES (10), (20), (30))",), @@ -4160,7 +2848,6 @@ mod test { } #[test] - #[cfg(feature = "new_parser")] fn advisory_lock_when_client_is_sadistic() { assert_eq!( locks( diff --git a/pgdog/src/frontend/router/parser/table.rs b/pgdog/src/frontend/router/parser/table.rs index 10fc0f150..91965d29c 100644 --- a/pgdog/src/frontend/router/parser/table.rs +++ b/pgdog/src/frontend/router/parser/table.rs @@ -1,13 +1,6 @@ use std::fmt::Display; -#[cfg(not(feature = "new_parser"))] -use pg_query::{ - Node as PgNode, NodeEnum, - protobuf::{List, RangeVar}, -}; -#[cfg(feature = "new_parser")] use pg_raw_parse::list::{CastNodeList, NodeList}; -#[cfg(feature = "new_parser")] use pg_raw_parse::{Node, nodes}; use super::{Error, Schema}; @@ -49,7 +42,6 @@ impl<'a> Table<'a> { } } -#[cfg(feature = "new_parser")] impl<'a> TryFrom> for Table<'a> { type Error = (); @@ -61,75 +53,6 @@ impl<'a> TryFrom> for Table<'a> { } } -#[cfg(not(feature = "new_parser"))] -impl<'a> TryFrom<&'a PgNode> for Table<'a> { - type Error = (); - - fn try_from(value: &'a PgNode) -> Result { - if let Some(NodeEnum::RangeVar(range_var)) = &value.node { - return Ok(range_var.into()); - } - - Err(()) - } -} - -#[cfg(not(feature = "new_parser"))] -impl<'a> TryFrom<&'a Vec> for Table<'a> { - type Error = Error; - - fn try_from(value: &'a Vec) -> Result { - 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 - && let Some(table) = table - { - return Ok(Table { - name: table, - schema: Some(schema), - alias: None, - }); - } - } - - _ => (), - } - - Err(Error::TableDecode) - } -} - -#[cfg(feature = "new_parser")] impl<'a> From<&'a pg_raw_parse::nodes::RangeVar> for Table<'a> { fn from(range_var: &'a pg_raw_parse::nodes::RangeVar) -> Self { let name = range_var.relname().unwrap_or_default(); @@ -143,26 +66,6 @@ impl<'a> From<&'a pg_raw_parse::nodes::RangeVar> for Table<'a> { } } -#[cfg(not(feature = "new_parser"))] -impl<'a> From<&'a RangeVar> for Table<'a> { - fn from(range_var: &'a RangeVar) -> Self { - let (name, alias) = if let Some(ref alias) = range_var.alias { - (range_var.relname.as_str(), Some(alias.aliasname.as_str())) - } else { - (range_var.relname.as_str(), None) - }; - Self { - name, - schema: if !range_var.schemaname.is_empty() { - Some(range_var.schemaname.as_str()) - } else { - None - }, - alias, - } - } -} -#[cfg(feature = "new_parser")] impl<'a> TryFrom<&'a CastNodeList> for Table<'a> { type Error = Error; @@ -188,7 +91,6 @@ impl<'a> TryFrom<&'a CastNodeList> for Table<'a> { } } -#[cfg(feature = "new_parser")] impl<'a> TryFrom<&'a NodeList> for Table<'a> { type Error = Error; @@ -218,44 +120,6 @@ impl<'a> TryFrom<&'a NodeList> for Table<'a> { } } -#[cfg(not(feature = "new_parser"))] -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 { diff --git a/pgdog/src/frontend/router/parser/util.rs b/pgdog/src/frontend/router/parser/util.rs deleted file mode 100644 index 5c9b0e035..000000000 --- a/pgdog/src/frontend/router/parser/util.rs +++ /dev/null @@ -1,29 +0,0 @@ -#![cfg(not(feature = "new_parser"))] -use pg_query::NodeEnum; -use pg_query::protobuf::{Node, String as PgString}; -use std::cmp::PartialEq; - -pub(crate) fn pg_string(s: impl Into) -> Node { - node(NodeEnum::String(PgString { sval: s.into() })) -} - -pub(crate) fn node(n: NodeEnum) -> Node { - Node { node: Some(n) } -} - -/// A const type that can be compared with Node for equality -pub(crate) const fn pg_str(s: &str) -> PgStr<'_> { - PgStr(s) -} - -#[derive(Debug)] -pub(crate) struct PgStr<'a>(&'a str); - -impl PartialEq for PgStr<'_> { - fn eq(&self, rhs: &Node) -> bool { - match &rhs.node { - Some(NodeEnum::String(PgString { sval })) => sval == self.0, - _ => false, - } - } -} diff --git a/pgdog/src/frontend/router/parser/value.rs b/pgdog/src/frontend/router/parser/value.rs index ebbbb572a..1cf00c84e 100644 --- a/pgdog/src/frontend/router/parser/value.rs +++ b/pgdog/src/frontend/router/parser/value.rs @@ -2,14 +2,7 @@ use std::fmt::Display; -#[cfg(feature = "new_parser")] use itertools::*; -#[cfg(not(feature = "new_parser"))] -use pg_query::{ - NodeEnum, - protobuf::{Node as PgNode, a_const::Val, *}, -}; -#[cfg(feature = "new_parser")] use pg_raw_parse::{Node, nodes}; use crate::net::{messages::Vector, vector::str_to_vector}; @@ -59,7 +52,6 @@ impl Value<'_> { } } -#[cfg(feature = "new_parser")] impl<'a> From<&'a nodes::A_Const> for Value<'a> { fn from(node: &'a nodes::A_Const) -> Self { use pg_raw_parse::const_val::ConstValue; @@ -88,59 +80,6 @@ impl<'a> From<&'a nodes::A_Const> for Value<'a> { } } -#[cfg(not(feature = "new_parser"))] -impl<'a> From<&'a AConst> for Value<'a> { - fn from(value: &'a AConst) -> Self { - if value.isnull { - return Value::Null; - } - - match value.val.as_ref() { - Some(Val::Sval(s)) => { - if s.sval.starts_with('[') && s.sval.ends_with(']') { - match str_to_vector(s.sval.as_str()) { - Ok(vector) => Value::Vector(vector), - _ => Value::String(s.sval.as_str()), - } - } else { - match s.sval.parse::() { - Ok(i) => Value::Integer(i), - Err(_) => Value::String(s.sval.as_str()), - } - } - } - Some(Val::Boolval(b)) => Value::Boolean(b.boolval), - Some(Val::Ival(i)) => Value::Integer(i.ival as i64), - Some(Val::Fval(Float { fval })) => { - if fval.contains(".") { - if let Ok(float) = fval.parse() { - Value::Float(float) - } else { - Value::String(fval.as_str()) - } - } else { - 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()), - } - } - } - Some(Val::Bsval(bsval)) => Value::String(bsval.bsval.as_str()), - None => Value::Null, - } - } -} - -#[cfg(not(feature = "new_parser"))] -impl<'a> TryFrom<&'a PgNode> for Value<'a> { - type Error = (); - - fn try_from(value: &'a PgNode) -> Result { - Value::try_from(&value.node) - } -} - -#[cfg(feature = "new_parser")] impl<'a> TryFrom> for Value<'a> { type Error = (); @@ -163,50 +102,11 @@ impl<'a> TryFrom> for Value<'a> { } } -#[cfg(not(feature = "new_parser"))] -impl<'a> TryFrom<&'a Option> for Value<'a> { - type Error = (); - - fn try_from(value: &'a Option) -> Result { - 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)? - } else { - Value::Null - } - } - - Some(NodeEnum::AExpr(expr)) => { - if expr.kind() == AExprKind::AexprOp - && let Some(PgNode { - node: Some(NodeEnum::String(pg_query::protobuf::String { sval })), - }) = expr.name.first() - && sval == "-" - && 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(()), - }) - } -} - #[cfg(test)] mod test { use super::*; #[test] - #[cfg(feature = "new_parser")] fn test_vector_value() { use pgdog_vector::Float; @@ -222,24 +122,6 @@ mod test { } #[test] - #[cfg(not(feature = "new_parser"))] - fn test_vector_value() { - let a_cosnt = AConst { - val: Some(Val::Sval(String { - sval: "[1,2,3]".into(), - })), - isnull: false, - location: 0, - }; - let node = PgNode { - node: Some(NodeEnum::AConst(a_cosnt)), - }; - let vector = Value::try_from(&node).unwrap(); - assert_eq!(vector.vector().unwrap()[0], 1.0.into()); - } - - #[test] - #[cfg(feature = "new_parser")] fn test_negative_numeric_with_cast() { // This will be parsed as a unary negation on a cast node, not a cast on // a negative numeric constant @@ -250,38 +132,6 @@ mod test { assert_eq!(value, Value::Float(-987_654_321.123_456_8)); } - #[test] - #[cfg(not(feature = "new_parser"))] - 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(-987_654_321.123_456_8)); - } - - #[cfg(feature = "new_parser")] fn selected_expr(result: &pg_raw_parse::ParseResult) -> Node<'_> { let stmt = result.stmts().exactly_one().ok().unwrap(); match stmt { diff --git a/pgdog/src/frontend/router/parser/where_clause.rs b/pgdog/src/frontend/router/parser/where_clause.rs index be70cb372..09648a209 100644 --- a/pgdog/src/frontend/router/parser/where_clause.rs +++ b/pgdog/src/frontend/router/parser/where_clause.rs @@ -1,11 +1,5 @@ //! WHERE clause of a UPDATE/SELECT/DELETE query. -#[cfg(not(feature = "new_parser"))] -use pg_query::{ - NodeEnum, - protobuf::{a_const::Val, *}, -}; -#[cfg(feature = "new_parser")] use pg_raw_parse::{ConstValue, Node, nodes}; use std::string::String; @@ -87,7 +81,6 @@ pub(crate) struct WhereClause<'a> { impl<'a> WhereClause<'a> { /// Parse the `WHERE` clause of a statement and extract /// all possible sharding keys. - #[cfg(feature = "new_parser")] pub(crate) fn new(source: &TablesSource<'a>, where_clause: Node<'a>) -> Option { if let Node::None = where_clause { return None; @@ -98,24 +91,6 @@ impl<'a> WhereClause<'a> { Some(Self { output }) } - cfg_select! { - not(feature = "new_parser") => { - pub(crate) fn new( - source: &TablesSource<'a>, - where_clause: &'a Option>, - ) -> Option> { - let Some(where_clause) = where_clause else { - return None; - }; - - let output = Self::parse(source, where_clause, false); - - Some(Self { output }) - } - } - _ => {} - } - pub(crate) fn keys(&self, table_name: Option<&str>, column_name: &str) -> Vec { let mut keys = vec![]; for output in &self.output { @@ -203,18 +178,6 @@ impl<'a> WhereClause<'a> { keys } - #[cfg(not(feature = "new_parser"))] - fn string(node: Option<&Node>) -> Option<&str> { - if let Some(node) = node - && let Some(NodeEnum::String(ref string)) = node.node - { - return Some(string.sval.as_str()); - } - - None - } - - #[cfg(feature = "new_parser")] fn parse(source: &TablesSource<'a>, node: Node<'a>, array: bool) -> Vec> { match node { // Only check for IS NULL, IS NOT NULL definitely doesn't help. @@ -298,129 +261,10 @@ impl<'a> WhereClause<'a> { _ => Vec::new(), } } - - cfg_select! { - not(feature = "new_parser") => { - fn parse(source: &TablesSource<'a>, node: &'a Node, array: bool) -> Vec> { - let mut keys = vec![]; - - match node.node { - Some(NodeEnum::NullTest(ref null_test)) - // Only check for IS NULL, IS NOT NULL definitely doesn't help. - if NullTestType::try_from(null_test.nulltesttype) == Ok(NullTestType::IsNull) => { - let left = null_test - .arg - .as_ref() - .and_then(|node| Self::parse(source, node, array).pop()); - - if let Some(Output::Column(c)) = left { - keys.push(Output::NullCheck(c)); - } - } - - Some(NodeEnum::BoolExpr(ref expr)) => { - // Only AND expressions can really be asserted. - // OR needs both sides to be evaluated and either one - // can direct to a shard. Most cases, this will end up on all shards. - if expr.boolop() != BoolExprType::AndExpr { - return keys; - } - - for arg in &expr.args { - keys.extend(Self::parse(source, arg, array)); - } - } - - Some(NodeEnum::AExpr(ref expr)) => { - let kind = expr.kind(); - if matches!( - kind, - AExprKind::AexprOp | AExprKind::AexprIn | AExprKind::AexprOpAny - ) { - let op = Self::string(expr.name.first()); - if let Some(op) = op - && op != "=" { - return keys; - } - } - let array = matches!(kind, AExprKind::AexprOpAny); - if let Some(ref left) = expr.lexpr - && let Some(ref right) = expr.rexpr { - let left = Self::parse(source, left, array); - let right = Self::parse(source, right, array); - - keys.push(Output::Filter(left, right)); - } - } - - Some(NodeEnum::AConst(ref value)) => { - if let Some(ref val) = value.val { - match val { - Val::Ival(int) => keys.push(Output::Int { - value: int.ival, - array, - }), - Val::Sval(sval) => keys.push(Output::Value { - value: sval.sval.clone(), - array, - }), - Val::Fval(fval) => keys.push(Output::Value { - value: fval.fval.clone(), - array, - }), - _ => (), - } - } - } - - Some(NodeEnum::ColumnRef(ref column)) => { - let name = Self::string(column.fields.last()); - let table = Self::string(column.fields.iter().rev().nth(1)); - let table = if let Some(table) = table { - Some(source.resolve_alias(table)) - } else { - source.table_name() - }; - - if let Some(name) = name { - return vec![Output::Column(Column { name, table })]; - } - } - - Some(NodeEnum::ParamRef(ref param)) => { - keys.push(Output::Parameter { - pos: param.number, - array, - }); - } - - Some(NodeEnum::List(ref list)) => { - for node in &list.items { - keys.extend(Self::parse(source, node, array)); - } - } - - Some(NodeEnum::TypeCast(ref cast)) => { - if let Some(ref arg) = cast.arg { - keys.extend(Self::parse(source, arg, array)); - } - } - - _ => (), - }; - - keys - } - } - _ => {} - } } #[cfg(test)] mod test { - #[cfg(not(feature = "new_parser"))] - use pg_query::{ParseResult, parse}; - #[cfg(feature = "new_parser")] use pg_raw_parse::{ParseResult, parse}; use super::*; @@ -545,7 +389,6 @@ mod test { ); } - #[cfg(feature = "new_parser")] fn where_clause(ast: &ParseResult) -> WhereClause<'_> { let Some(Node::SelectStmt(stmt)) = ast.stmts().next() else { panic!("Not a select"); @@ -554,19 +397,4 @@ mod test { let source = TablesSource::from(from_clause); WhereClause::new(&source, stmt.where_clause()).unwrap() } - - cfg_select! { - not(feature = "new_parser") => { - fn where_clause(ast: &ParseResult) -> WhereClause<'_> { - let stmt = ast.protobuf.stmts.first().as_ref().unwrap().stmt.as_ref().unwrap(); - let Some(NodeEnum::SelectStmt(stmt)) = &stmt.node else { - panic!("Not a select"); - }; - let from_clause = FromClause::new(&stmt.from_clause); - let source = TablesSource::from(from_clause); - WhereClause::new(&source, &stmt.where_clause).unwrap() - } - } - _ => {} - } } diff --git a/pgdog/src/util.rs b/pgdog/src/util.rs index 02b2f05c2..2e3a53f55 100644 --- a/pgdog/src/util.rs +++ b/pgdog/src/util.rs @@ -3,7 +3,6 @@ use chrono::{DateTime, Local, Utc}; use once_cell::sync::Lazy; use rand::{Rng, distr::Alphanumeric}; -#[cfg(feature = "new_parser")] use std::ops::ControlFlow; use std::panic::Location; use std::{env, future::Future, future::pending, num::ParseIntError, time::Duration}; @@ -376,12 +375,10 @@ impl SafeInterval { } } -#[cfg(feature = "new_parser")] pub(crate) trait ResultControlFlowExt { fn break_err(self) -> ControlFlow, T>; } -#[cfg(feature = "new_parser")] impl ResultControlFlowExt for Result { fn break_err(self) -> ControlFlow, T> { match self { diff --git a/plugins/pgdog-example-plugin/Cargo.toml b/plugins/pgdog-example-plugin/Cargo.toml index 4e6f895a5..c7102733a 100644 --- a/plugins/pgdog-example-plugin/Cargo.toml +++ b/plugins/pgdog-example-plugin/Cargo.toml @@ -11,8 +11,4 @@ pgdog-plugin.workspace = true once_cell = "1" parking_lot = "0.12" thiserror = "2" -pg_raw_parse = { workspace = true, optional = true } - -[features] -default = ["pgdog-plugin/pg_query"] -new_parser = ["pg_raw_parse", "pgdog-plugin/new_parser"] +pg_raw_parse.workspace = true diff --git a/plugins/pgdog-example-plugin/src/lib.rs b/plugins/pgdog-example-plugin/src/lib.rs index 106b61164..03f02738d 100644 --- a/plugins/pgdog-example-plugin/src/lib.rs +++ b/plugins/pgdog-example-plugin/src/lib.rs @@ -23,7 +23,7 @@ impl Plugin for ExamplePlugin { /// If defined, this function is called on every query going through PgDog. /// - /// It's provided with the AST generated by pg_query and context on how many databases + /// It's provided with the AST generated by the parser and context on how many databases /// PgDog is proxying. fn route(context: Context<'_>) -> Route { crate::plugin::route_query(context).unwrap_or(Route::unknown()) diff --git a/plugins/pgdog-example-plugin/src/plugin.rs b/plugins/pgdog-example-plugin/src/plugin.rs index b846defab..67934a585 100644 --- a/plugins/pgdog-example-plugin/src/plugin.rs +++ b/plugins/pgdog-example-plugin/src/plugin.rs @@ -5,9 +5,6 @@ use std::{ use once_cell::sync::Lazy; use parking_lot::Mutex; -#[cfg(not(feature = "new_parser"))] -use pg_query::{NodeEnum, protobuf::RangeVar}; -#[cfg(feature = "new_parser")] use pg_raw_parse::Node; use pgdog_plugin::prelude::*; use thiserror::Error; @@ -15,11 +12,6 @@ use thiserror::Error; #[derive(Error, Debug)] pub enum PluginError { #[error("{0}")] - #[cfg(not(feature = "new_parser"))] - PgQuery(#[from] pg_query::Error), - - #[error("{0}")] - #[cfg(feature = "pg_raw_parse")] Parser(#[from] pg_raw_parse::Error), #[error("empty query")] @@ -31,7 +23,6 @@ static WRITE_TIMES: Lazy>> = /// Route query to a replica or a primary, depending on when was the last time /// we wrote to the table. -#[cfg(feature = "new_parser")] pub(crate) fn route_query(context: Context<'_>) -> Result { // PgDog really thinks this should be a write. // This could be because there is an INSERT statement in a CTE, @@ -111,94 +102,6 @@ pub(crate) fn route_query(context: Context<'_>) -> Result { Ok(Route::unknown()) } -cfg_select! { - not(feature = "new_parser") => { - pub(crate) fn route_query(context: Context<'_>) -> Result { - // PgDog really thinks this should be a write. - // This could be because there is an INSERT statement in a CTE, - // or something else. You could override its decision here, but make - // sure you checked the AST first. - let write_override = context.write_override(); - - let root = context - .query - .stmts - .first() - .ok_or(PluginError::EmptyQuery)? - .stmt - .as_ref() - .ok_or(PluginError::EmptyQuery)?; - - match root.node.as_ref() { - Some(NodeEnum::SelectStmt(stmt)) => { - if write_override { - return Ok(Route::unknown()); - } - - let table_name = stmt - .from_clause - .first() - .ok_or(PluginError::EmptyQuery)? - .node - .as_ref() - .ok_or(PluginError::EmptyQuery)?; - - if let NodeEnum::RangeVar(RangeVar { relname, .. }) = table_name { - // Got info on last write. - 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)); - } - } - } - Some(NodeEnum::InsertStmt(stmt)) => { - if let Some(ref relation) = stmt.relation { - WRITE_TIMES - .lock() - .insert(relation.relname.clone(), Instant::now()); - } - } - Some(NodeEnum::UpdateStmt(stmt)) => { - if let Some(ref relation) = stmt.relation { - WRITE_TIMES - .lock() - .insert(relation.relname.clone(), Instant::now()); - } - } - Some(NodeEnum::DeleteStmt(stmt)) => { - if let Some(ref relation) = stmt.relation { - WRITE_TIMES - .lock() - .insert(relation.relname.clone(), Instant::now()); - } - } - _ => {} - } - - // Get prepared statement parameters. - let params = context.parameters(); - if params.parameters.is_empty() { - // No params bound. - } else { - let param = params - .parameters - .first() - .and_then(|p| p.decode(params.parameter_format(0))); - if let Some(param) = param { - println!("Decoded parameter 0 ($1): {:?}", param); - } - } - - // Let PgDog decide. - Ok(Route::unknown()) - } - } - _ => {} -} - #[cfg(test)] mod test { use pgdog_plugin::parameters::Parameters; @@ -208,9 +111,6 @@ mod test { #[test] fn test_routing_plugin() { // Keep protobuf in memory. - #[cfg(not(feature = "new_parser"))] - let query = pg_query::parse("SELECT * FROM users").unwrap().protobuf; - #[cfg(feature = "new_parser")] let query = pg_raw_parse::parse("SELECT * FROM users").unwrap(); let context = pgdog_plugin::Context { shards: 1, diff --git a/plugins/pgdog-primary-only-tables/Cargo.toml b/plugins/pgdog-primary-only-tables/Cargo.toml index c6534b96e..f16b50d7f 100644 --- a/plugins/pgdog-primary-only-tables/Cargo.toml +++ b/plugins/pgdog-primary-only-tables/Cargo.toml @@ -15,8 +15,4 @@ toml = "0.8" serde = { version = "1", features = ["derive"]} arc-swap = "1" tracing = "0.1" -pg_raw_parse = { workspace = true, optional = true } - -[features] -default = ["pgdog-plugin/pg_query"] -new_parser = ["pg_raw_parse", "pgdog-plugin/new_parser"] +pg_raw_parse.workspace = true diff --git a/plugins/pgdog-primary-only-tables/src/lib.rs b/plugins/pgdog-primary-only-tables/src/lib.rs index a7f926159..f9c75ae53 100644 --- a/plugins/pgdog-primary-only-tables/src/lib.rs +++ b/plugins/pgdog-primary-only-tables/src/lib.rs @@ -7,15 +7,10 @@ use std::{ use arc_swap::ArcSwap; use once_cell::sync::Lazy; -#[cfg(feature = "new_parser")] use pg_raw_parse::Node; -#[cfg(feature = "new_parser")] use pg_raw_parse::walk::{self, Recurse}; -#[cfg(not(feature = "new_parser"))] -use pgdog_plugin::pg_query::{NodeEnum, NodeRef}; use pgdog_plugin::{Config as PluginConfig, Context, PdStr, Plugin, ReadWrite, Route, Shard}; use serde::{Deserialize, Serialize}; -#[cfg(feature = "new_parser")] use std::ops::ControlFlow; use tracing::{error, info}; @@ -80,7 +75,6 @@ fn read_config(path: &Path) -> Result<(), Box> { Ok(()) } -#[cfg(feature = "new_parser")] fn route_query(context: Context<'_>) -> Route { let ast = &context.query; @@ -112,52 +106,3 @@ fn route_query(context: Context<'_>) -> Route { route.unwrap_or_default() } - -cfg_select! { - not(feature = "new_parser") => { - fn route_query(context: Context<'_>) -> Route { - let ast = &context.query; - - let root_node = ast - .stmts - .first() - .and_then(|s| s.stmt.as_ref()) - .and_then(|s| s.node.as_ref()); - - let is_select = root_node.is_some_and(|node| match node { - NodeEnum::SelectStmt(_) => true, - NodeEnum::ExplainStmt(stmt) => stmt - .query - .as_ref() - .and_then(|q| q.node.as_ref()) - .is_some_and(|n| matches!(n, NodeEnum::SelectStmt(_))), - _ => false, - }); - - if !is_select { - return Route::default(); - } - - let config = CONFIG.load().clone(); - - for node in ast.nodes() { - if let NodeRef::RangeVar(range_var) = node.0 { - for table in &config.tables { - let name_matches = table.name == range_var.relname; - let schema_matches = match &table.schema { - Some(schema) => schema == &range_var.schemaname, - None => true, - }; - - if name_matches && schema_matches { - return Route::new(Shard::Unknown, ReadWrite::Write); - } - } - } - } - - Route::default() - } - } - _ => {} -}