From c9a7209b1a1c69041d6259f95ae08b1944664f3b Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Sun, 23 Aug 2026 15:51:23 +0800 Subject: [PATCH] feat: add snapshot-consistent query service --- Cargo.lock | 142 ++ Cargo.toml | 2 + DEPENDENCIES.rust.tsv | 11 + crates/query-service-server/Cargo.toml | 53 + .../DEPENDENCIES.rust.tsv | 415 ++++ crates/query-service-server/README.md | 224 ++ crates/query-service-server/src/auth.rs | 380 +++ crates/query-service-server/src/lib.rs | 2179 +++++++++++++++++ crates/query-service-server/src/main.rs | 27 + crates/query-service/Cargo.toml | 42 + crates/query-service/DEPENDENCIES.rust.tsv | 392 +++ crates/query-service/src/error.rs | 105 + crates/query-service/src/key.rs | 1011 ++++++++ crates/query-service/src/lib.rs | 34 + crates/query-service/src/lookup.rs | 863 +++++++ crates/query-service/src/model.rs | 114 + crates/query-service/src/policy.rs | 194 ++ crates/query-service/tests/lookup.rs | 656 +++++ 18 files changed, 6844 insertions(+) create mode 100644 crates/query-service-server/Cargo.toml create mode 100644 crates/query-service-server/DEPENDENCIES.rust.tsv create mode 100644 crates/query-service-server/README.md create mode 100644 crates/query-service-server/src/auth.rs create mode 100644 crates/query-service-server/src/lib.rs create mode 100644 crates/query-service-server/src/main.rs create mode 100644 crates/query-service/Cargo.toml create mode 100644 crates/query-service/DEPENDENCIES.rust.tsv create mode 100644 crates/query-service/src/error.rs create mode 100644 crates/query-service/src/key.rs create mode 100644 crates/query-service/src/lib.rs create mode 100644 crates/query-service/src/lookup.rs create mode 100644 crates/query-service/src/model.rs create mode 100644 crates/query-service/src/policy.rs create mode 100644 crates/query-service/tests/lookup.rs diff --git a/Cargo.lock b/Cargo.lock index 627ac1d4..289af618 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3849,6 +3849,15 @@ version = "0.1.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "58093314a45e00c77d5c508f76e77c3396afbbc0d01506e7fae47b018bac2b1d" +[[package]] +name = "matchers" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d1525a2a28c7f4fa0fc98bb91ae755d1e2d1505079e05539e35bc876b5d65ae9" +dependencies = [ + "regex-automata", +] + [[package]] name = "matchit" version = "0.7.3" @@ -4067,6 +4076,15 @@ dependencies = [ "syn 1.0.109", ] +[[package]] +name = "nu-ansi-term" +version = "0.50.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" +dependencies = [ + "windows-sys 0.61.2", +] + [[package]] name = "num" version = "0.4.3" @@ -4716,6 +4734,44 @@ dependencies = [ "zstd", ] +[[package]] +name = "paimon-query-service" +version = "0.4.0" +dependencies = [ + "arrow-array", + "arrow-schema", + "base64", + "chrono", + "futures", + "moka", + "paimon", + "serde", + "serde_json", + "tempfile", + "tokio", +] + +[[package]] +name = "paimon-query-service-server" +version = "0.4.0" +dependencies = [ + "axum", + "futures", + "hyper-util", + "paimon", + "paimon-query-service", + "serde", + "serde_json", + "sha2 0.10.9", + "subtle", + "tempfile", + "tokio", + "tower", + "tracing", + "tracing-appender", + "tracing-subscriber", +] + [[package]] name = "paimon-rest-server" version = "0.4.0" @@ -6305,6 +6361,15 @@ version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5f179d4e11094a893b82fff208f74d448a7512f99f5a0acbd5c679b705f83ed9" +[[package]] +name = "sharded-slab" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f40ca3c46823713e0d4209592e8d6e826aa57e928f09752619fc696c499637f6" +dependencies = [ + "lazy_static", +] + [[package]] name = "shlex" version = "1.3.0" @@ -6607,6 +6672,12 @@ version = "2.6.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292" +[[package]] +name = "symlink" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7973cce6668464ea31f176d85b13c7ab3bba2cb3b77a2ed26abd7801688010a" + [[package]] name = "syn" version = "1.0.109" @@ -6928,6 +6999,15 @@ dependencies = [ "syn 3.0.2", ] +[[package]] +name = "thread_local" +version = "1.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ad99c4c6d32803332c548b1af0540b357b3f5fc0be8f6c6bfe8b2e6ae784070" +dependencies = [ + "cfg-if", +] + [[package]] name = "thrift" version = "0.17.0" @@ -7164,6 +7244,19 @@ dependencies = [ "tracing-core", ] +[[package]] +name = "tracing-appender" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "050686193eb999b4bb3bc2acfa891a13da00f79734704c4b8b4ef1a10b368a3c" +dependencies = [ + "crossbeam-channel", + "symlink", + "thiserror 2.0.19", + "time", + "tracing-subscriber", +] + [[package]] name = "tracing-attributes" version = "0.1.31" @@ -7182,6 +7275,49 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "db97caf9d906fbde555dd62fa95ddba9eecfd14cb388e4f491a66d74cd5fb79a" dependencies = [ "once_cell", + "valuable", +] + +[[package]] +name = "tracing-log" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ee855f1f400bd0e5c02d150ae5de3840039a3f54b025156404e34c23c03f47c3" +dependencies = [ + "log", + "once_cell", + "tracing-core", +] + +[[package]] +name = "tracing-serde" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "704b1aeb7be0d0a84fc9828cae51dab5970fee5088f83d1dd7ee6f6246fc6ff1" +dependencies = [ + "serde", + "tracing-core", +] + +[[package]] +name = "tracing-subscriber" +version = "0.3.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb7f578e5945fb242538965c2d0b04418d38ec25c79d160cd279bf0731c8d319" +dependencies = [ + "matchers", + "nu-ansi-term", + "once_cell", + "regex-automata", + "serde", + "serde_json", + "sharded-slab", + "smallvec", + "thread_local", + "tracing", + "tracing-core", + "tracing-log", + "tracing-serde", ] [[package]] @@ -7324,6 +7460,12 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "valuable" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba73ea9cf16a25df0c8caa16c51acb937d5712a8429db78a3ee29d5dcacd3a65" + [[package]] name = "vcpkg" version = "0.2.15" diff --git a/Cargo.toml b/Cargo.toml index 08fe20d3..dd534f6c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -19,6 +19,8 @@ resolver = "2" members = [ "crates/paimon", + "crates/query-service", + "crates/query-service-server", "crates/paimon-rest-server", "crates/integration_tests", "bindings/c", diff --git a/DEPENDENCIES.rust.tsv b/DEPENDENCIES.rust.tsv index 5b708be8..9f5e3c0e 100644 --- a/DEPENDENCIES.rust.tsv +++ b/DEPENDENCIES.rust.tsv @@ -330,6 +330,7 @@ lz4_flex@0.13.1 X lzokay-native@0.1.0 X macro_rules_attribute@0.1.3 X macro_rules_attribute-proc_macro@0.1.3 X +matchers@0.2.0 X matchit@0.7.3 X X matrixmultiply@0.3.11 X X md-5@0.10.6 X X @@ -352,6 +353,7 @@ no_std_io2@0.9.4 X X nom@7.1.3 X nougat@0.2.4 X X X nougat-proc_macros@0.2.4 X X X +nu-ansi-term@0.50.3 X num@0.4.3 X X num-bigint@0.4.8 X X num-bigint-dig@0.8.6 X X @@ -396,6 +398,8 @@ paimon-datafusion@0.4.0 X paimon-ftindex-core@0.1.0 X paimon-integration-tests@0.4.0 X paimon-mosaic-core@0.2.0 X +paimon-query-service@0.4.0 X +paimon-query-service-server@0.4.0 X paimon-rest-server@0.4.0 X paimon-tpcds-bench@0.4.0 X paimon-vindex-core@0.3.0 X @@ -528,6 +532,7 @@ sha1@0.11.0 X X sha2@0.10.9 X X sha2@0.11.0 X X sha2-const-stable@0.1.0 X X +sharded-slab@0.1.7 X shlex@1.3.0 X X shlex@2.0.1 X X signal-hook-registry@1.4.8 X X @@ -562,6 +567,7 @@ strsim@0.11.1 X strum@0.27.2 X strum_macros@0.27.2 X subtle@2.6.1 X +symlink@0.1.0 X X syn@1.0.109 X X syn@2.0.119 X X syn@3.0.2 X X @@ -588,6 +594,7 @@ thiserror@1.0.69 X X thiserror@2.0.19 X X thiserror-impl@1.0.69 X X thiserror-impl@2.0.19 X X +thread_local@1.1.10 X X thrift@0.17.0 X time@0.3.54 X X time-core@0.1.9 X X @@ -607,8 +614,12 @@ tower-http@0.6.11 X tower-layer@0.3.3 X tower-service@0.3.3 X tracing@0.1.44 X +tracing-appender@0.2.5 X tracing-attributes@0.1.31 X tracing-core@0.1.36 X +tracing-log@0.2.0 X +tracing-serde@0.2.0 X +tracing-subscriber@0.3.23 X try-lock@0.2.5 X twox-hash@2.1.3 X typed-builder@0.19.1 X X diff --git a/crates/query-service-server/Cargo.toml b/crates/query-service-server/Cargo.toml new file mode 100644 index 00000000..7c31b7b1 --- /dev/null +++ b/crates/query-service-server/Cargo.toml @@ -0,0 +1,53 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +[package] +name = "paimon-query-service-server" +version.workspace = true +edition.workspace = true +homepage.workspace = true +repository.workspace = true +license.workspace = true +rust-version.workspace = true +description = "HTTP service for snapshot-consistent Paimon point queries" +publish = false + +[lib] +name = "paimon_query_service_server" +path = "src/lib.rs" + +[[bin]] +name = "paimon-query-service-server" +path = "src/main.rs" + +[dependencies] +paimon = { workspace = true } +paimon-query-service = { path = "../query-service" } +axum = { version = "0.7", features = ["macros", "tokio", "http1", "http2"] } +hyper-util = { version = "0.1", features = ["server-auto", "server-graceful", "service", "tokio"] } +serde = { version = "1", features = ["derive"] } +serde_json = "1.0.120" +sha2 = "0.10" +subtle = "2.6" +tokio = { version = "1.39.2", features = ["io-util", "macros", "net", "rt-multi-thread", "signal", "sync", "time"] } +tracing = "0.1" +tracing-appender = "0.2" +tracing-subscriber = { version = "0.3", features = ["env-filter", "json"] } + +[dev-dependencies] +futures = "0.3" +tempfile = "3" +tower = { version = "0.5", features = ["util"] } diff --git a/crates/query-service-server/DEPENDENCIES.rust.tsv b/crates/query-service-server/DEPENDENCIES.rust.tsv new file mode 100644 index 00000000..fb48bc42 --- /dev/null +++ b/crates/query-service-server/DEPENDENCIES.rust.tsv @@ -0,0 +1,415 @@ +crate 0BSD Apache-2.0 Apache-2.0 WITH LLVM-exception BSD-2-Clause BSD-3-Clause BSL-1.0 CC0-1.0 CDLA-Permissive-2.0 ISC LGPL-2.1-or-later MIT MIT-0 Unicode-3.0 Unlicense Zlib +adler2@2.0.1 X X X +ahash@0.8.12 X X +aho-corasick@1.1.4 X X +alloc-no-stdlib@2.0.4 X +alloc-stdlib@0.2.4 X +allocator-api2@0.2.21 X X +android_system_properties@0.1.5 X X +anyhow@1.0.104 X X +apache-avro@0.21.0 X +approx@0.5.1 X +arrow@58.3.0 X +arrow-arith@58.3.0 X +arrow-array@58.3.0 X X +arrow-buffer@58.3.0 X +arrow-cast@58.3.0 X +arrow-csv@58.3.0 X +arrow-data@58.3.0 X +arrow-ipc@58.3.0 X +arrow-json@58.3.0 X +arrow-ord@58.3.0 X +arrow-row@58.3.0 X +arrow-schema@58.3.0 X +arrow-select@58.3.0 X +arrow-string@58.3.0 X +async-lock@3.4.2 X X +async-stream@0.3.6 X +async-stream-impl@0.3.6 X +async-trait@0.1.91 X X +atoi@2.0.0 X +atomic-waker@1.1.2 X X +autocfg@1.5.1 X X +aws-lc-rs@1.17.3 X X +aws-lc-sys@0.43.0 X X X X X +axum@0.7.9 X +axum-core@0.4.5 X +axum-macros@0.4.2 X +backon@1.6.0 X +base64@0.22.1 X X +bigdecimal@0.4.10 X X +bitflags@2.13.1 X X +block-buffer@0.10.4 X X +block-buffer@0.12.1 X X +bon@3.9.3 X X +bon-macros@3.9.3 X X +brotli@8.0.4 X X +brotli-decompressor@5.0.3 X X +bumpalo@3.20.3 X X +bytemuck@1.25.2 X X X +byteorder@1.5.0 X X +bytes@1.12.1 X +cc@1.3.0 X X +cfg-if@1.0.4 X X +chrono@0.4.45 X X +chrono-tz@0.10.4 X X +cmake@0.1.58 X X +cmov@0.5.4 X X +combine@4.6.7 X +comfy-table@7.2.2 X +concurrent-queue@2.5.0 X X +const-oid@0.10.2 X X +const-random@0.1.18 X X +const-random-macro@0.1.16 X X +core-foundation@0.10.1 X X +core-foundation@0.9.4 X X +core-foundation-sys@0.8.7 X X +cpufeatures@0.2.17 X X +cpufeatures@0.3.0 X X +crc32fast@1.5.0 X X +crossbeam-channel@0.5.16 X X +crossbeam-deque@0.8.7 X X +crossbeam-epoch@0.9.20 X X +crossbeam-utils@0.8.22 X X +crunchy@0.2.4 X +crypto-common@0.1.7 X X +crypto-common@0.2.2 X X +csv@1.4.0 X X +csv-core@0.1.13 X X +ctutils@0.4.2 X X +darling@0.23.0 X +darling_core@0.23.0 X +darling_macro@0.23.0 X +deranged@0.5.8 X X +diff@0.1.13 X X +digest@0.10.7 X X +digest@0.11.3 X X +displaydoc@0.2.6 X X +dlv-list@0.5.2 X X +dunce@1.0.5 X X X +either@1.16.0 X X +encoding_rs@0.8.35 X X X +equivalent@1.0.2 X X +errno@0.3.14 X X +event-listener@5.4.1 X X +event-listener-strategy@0.5.4 X X +fallible-streaming-iterator@0.1.9 X X +fastrand@2.5.0 X X +find-msvc-tools@0.1.9 X X +flatbuffers@25.12.19 X +flate2@1.1.9 X X +fnv@1.0.7 X X +foldhash@0.2.0 X +foreign-types@0.3.2 X X +foreign-types-shared@0.1.1 X X +form_urlencoded@1.2.2 X X +fs_extra@1.3.0 X +futures@0.3.33 X X +futures-channel@0.3.33 X X +futures-core@0.3.33 X X +futures-executor@0.3.33 X X +futures-io@0.3.33 X X +futures-macro@0.3.33 X X +futures-sink@0.3.33 X X +futures-task@0.3.33 X X +futures-util@0.3.33 X X +generic-array@0.14.7 X +getrandom@0.2.17 X X +getrandom@0.3.4 X X +getrandom@0.4.3 X X +gloo-timers@0.3.0 X X +h2@0.4.16 X +half@2.7.1 X X +hashbrown@0.14.5 X X +hashbrown@0.17.1 X X +heck@0.5.0 X X +hex@0.4.3 X X +hmac@0.12.1 X X +hmac@0.13.0 X X +http@1.4.2 X X +http-body@1.1.0 X +http-body-util@0.1.4 X +httparse@1.10.1 X X +httpdate@1.0.3 X X +hybrid-array@0.4.13 X X +hyper@1.10.1 X +hyper-rustls@0.27.9 X X X +hyper-tls@0.6.0 X X +hyper-util@0.1.20 X +iana-time-zone@0.1.65 X X +iana-time-zone-haiku@0.1.2 X X +icu_collections@2.2.0 X +icu_locale_core@2.2.0 X +icu_normalizer@2.2.0 X +icu_normalizer_data@2.2.0 X +icu_properties@2.2.0 X +icu_properties_data@2.2.0 X +icu_provider@2.2.0 X +ident_case@1.0.1 X X +idna@1.1.0 X X +idna_adapter@1.2.2 X X +indexmap@2.14.0 X X +integer-encoding@3.0.4 X +ipnet@2.12.0 X X +itertools@0.14.0 X X +itoa@1.0.18 X X +jiff@0.2.34 X X +jiff-core@0.1.0 X X +jiff-tzdb@0.1.8 X X +jiff-tzdb-platform@0.1.3 X X +jni@0.22.4 X X +jni-macros@0.22.4 X X +jni-sys@0.4.1 X X +jni-sys-macros@0.4.1 X X +jobserver@0.1.35 X X +js-sys@0.3.103 X X +lazy_static@1.5.0 X X +lexical-core@1.0.6 X X +lexical-parse-float@1.0.6 X X +lexical-parse-integer@1.0.6 X X +lexical-util@1.0.7 X X +lexical-write-float@1.0.6 X X +lexical-write-integer@1.0.6 X X +libc@0.2.186 X X +libloading@0.9.0 X +libm@0.2.16 X +linux-raw-sys@0.12.1 X X X +litemap@0.8.2 X +lock_api@0.4.14 X X +log@0.4.33 X X +lru@0.18.2 X +lz4_flex@0.11.6 X +lz4_flex@0.13.1 X +lzokay-native@0.1.0 X +matchers@0.2.0 X +matchit@0.7.3 X X +matrixmultiply@0.3.11 X X +md-5@0.10.6 X X +md-5@0.11.0 X X +mea@0.6.4 X +memchr@2.8.3 X X +mime@0.3.17 X X +miniz_oxide@0.8.9 X X X +mio@1.2.2 X +moka@0.12.15 X X +nalgebra@0.33.3 X +nalgebra-macros@0.2.2 X +native-tls@0.2.18 X X +nu-ansi-term@0.50.3 X +num@0.4.3 X X +num-bigint@0.4.8 X X +num-complex@0.4.6 X X +num-conv@0.2.2 X X +num-integer@0.1.46 X X +num-iter@0.1.46 X X +num-rational@0.4.2 X X +num-traits@0.2.19 X X +once_cell@1.21.4 X X +opendal-core@0.58.0 X +opendal-http-transport-reqwest@0.58.0 X +opendal-layer-retry@0.58.0 X +opendal-service-fs@0.58.0 X +opendal-service-oss@0.58.0 X +openssl@0.10.81 X +openssl-macros@0.1.1 X X +openssl-probe@0.2.1 X X +openssl-sys@0.9.117 X +orc-rust@0.8.0 X +ordered-float@2.10.1 X +ordered-multimap@0.7.3 X +paimon@0.4.0 X +paimon-mosaic-core@0.2.0 X +paimon-query-service@0.4.0 X +paimon-query-service-server@0.4.0 X +paimon-vindex-core@0.3.0 X +parking@2.2.1 X X +parking_lot@0.12.5 X X +parking_lot_core@0.9.12 X X +parquet@58.3.0 X +paste@1.0.15 X X +percent-encoding@2.3.2 X X +phf@0.12.1 X +phf_shared@0.12.1 X +pin-project-lite@0.2.17 X X +pkg-config@0.3.33 X X +portable-atomic@1.14.0 X X +portable-atomic-util@0.2.7 X X +potential_utf@0.1.5 X +powerfmt@0.2.0 X X +ppv-lite86@0.2.21 X X +pretty_assertions@1.4.1 X X +prettyplease@0.2.37 X X +proc-macro2@1.0.107 X X +prost@0.13.5 X +prost-derive@0.13.5 X +quad-rand@0.2.3 X +quick-xml@0.41.0 X +quote@1.0.47 X X +r-efi@5.3.0 X X X +r-efi@6.0.0 X X X +rand@0.8.7 X X +rand@0.9.5 X X +rand_chacha@0.3.1 X X +rand_chacha@0.9.0 X X +rand_core@0.6.4 X X +rand_core@0.9.5 X X +rawpointer@0.2.1 X X +rayon@1.12.0 X X +rayon-core@1.13.0 X X +redox_syscall@0.5.18 X +regex@1.13.1 X X +regex-automata@0.4.16 X X +regex-lite@0.1.9 X X +regex-syntax@0.8.11 X X +reqsign-aliyun-oss@3.1.1 X +reqsign-core@3.1.0 X +reqsign-file-read-tokio@3.0.2 X +reqwest@0.12.28 X X +reqwest@0.13.4 X X +roaring@0.11.4 X X +rust-ini@0.21.3 X +rustc_version@0.4.1 X X +rustix@1.1.4 X X X +rustls@0.23.42 X X X +rustls-native-certs@0.8.4 X X X +rustls-pki-types@1.15.0 X X +rustls-platform-verifier@0.7.0 X X +rustls-platform-verifier-android@0.1.1 X X +rustls-webpki@0.103.13 X +rustversion@1.0.23 X X +ryu@1.0.23 X X +safe_arch@0.7.4 X X X +same-file@1.0.6 X X +schannel@0.1.29 X +scopeguard@1.2.0 X X +security-framework@3.7.0 X X +security-framework-sys@2.17.0 X X +semver@1.0.28 X X +seq-macro@0.3.6 X X +serde@1.0.229 X X +serde_bytes@0.11.19 X X +serde_core@1.0.229 X X +serde_derive@1.0.229 X X +serde_json@1.0.151 X X +serde_path_to_error@0.1.20 X X +serde_repr@0.1.21 X X +serde_urlencoded@0.7.1 X X +serde_with@3.21.0 X X +serde_with_macros@3.21.0 X X +sha1@0.10.7 X X +sha1@0.11.0 X X +sha2@0.10.9 X X +sha2@0.11.0 X X +sharded-slab@0.1.7 X +shlex@2.0.1 X X +signal-hook-registry@1.4.8 X X +simba@0.9.1 X +simd-adler32@0.3.10 X +simd_cesu8@1.2.0 X X +simdutf8@0.1.5 X X +siphasher@1.0.3 X X +slab@0.4.12 X +smallvec@1.15.2 X X +snafu@0.8.9 X X +snafu@0.9.1 X X +snafu-derive@0.8.9 X X +snafu-derive@0.9.1 X X +snap@1.1.2 X +socket2@0.6.5 X X +stable_deref_trait@1.2.1 X X +strsim@0.11.1 X +strum@0.27.2 X +strum_macros@0.27.2 X +subtle@2.6.1 X +symlink@0.1.0 X X +syn@2.0.119 X X +syn@3.0.2 X X +sync_wrapper@1.0.2 X +synstructure@0.13.2 X +system-configuration@0.7.0 X X +system-configuration-sys@0.6.0 X X +tagptr@0.2.0 X X +tempfile@3.27.0 X X +thiserror@1.0.69 X X +thiserror@2.0.19 X X +thiserror-impl@1.0.69 X X +thiserror-impl@2.0.19 X X +thread_local@1.1.10 X X +thrift@0.17.0 X +time@0.3.54 X X +time-core@0.1.9 X X +tiny-keccak@2.0.2 X +tinystr@0.8.3 X +tokio@1.53.0 X +tokio-macros@2.7.1 X +tokio-native-tls@0.3.1 X +tokio-rustls@0.26.4 X X +tokio-util@0.7.18 X +tower@0.5.3 X +tower-http@0.6.11 X +tower-layer@0.3.3 X +tower-service@0.3.3 X +tracing@0.1.44 X +tracing-appender@0.2.5 X +tracing-attributes@0.1.31 X +tracing-core@0.1.36 X +tracing-log@0.2.0 X +tracing-serde@0.2.0 X +tracing-subscriber@0.3.23 X +try-lock@0.2.5 X +twox-hash@2.1.3 X +typed-builder@0.19.1 X X +typed-builder-macro@0.19.1 X X +typenum@1.20.1 X X +unicode-ident@1.0.24 X X X +unicode-segmentation@1.13.2 X X +unicode-width@0.2.2 X X +untrusted@0.9.0 X +url@2.5.8 X X +urlencoding@2.1.3 X +utf8_iter@1.0.4 X X +uuid@1.24.0 X X +vcpkg@0.2.15 X X +version_check@0.9.5 X X +walkdir@2.5.0 X X +want@0.3.1 X +wasi@0.11.1+wasi-snapshot-preview1 X X X +wasip2@1.0.4+wasi-0.2.12 X X X +wasm-bindgen@0.2.126 X X +wasm-bindgen-futures@0.4.76 X X +wasm-bindgen-macro@0.2.126 X X +wasm-bindgen-macro-support@0.2.126 X X +wasm-bindgen-shared@0.2.126 X X +wasm-streams@0.5.0 X X +web-sys@0.3.103 X X +web-time@1.1.0 X X +webpki-root-certs@1.0.9 X +wide@0.7.33 X X X +winapi-util@0.1.11 X X +windows-core@0.62.2 X X +windows-implement@0.60.2 X X +windows-interface@0.59.3 X X +windows-link@0.2.1 X X +windows-registry@0.6.1 X X +windows-result@0.4.1 X X +windows-strings@0.5.1 X X +windows-sys@0.61.2 X X +wit-bindgen@0.57.1 X X X +writeable@0.6.3 X +xattr@1.6.1 X X +yansi@1.0.1 X X +yoke@0.8.3 X +yoke-derive@0.8.2 X +zerocopy@0.8.54 X X X +zerocopy-derive@0.8.54 X X X +zerofrom@0.1.8 X +zerofrom-derive@0.1.7 X +zeroize@1.9.0 X X +zerotrie@0.2.4 X +zerovec@0.11.6 X +zerovec-derive@0.11.3 X +zlib-rs@0.6.6 X +zmij@1.0.23 X +zstd@0.13.3 X +zstd-safe@7.2.4 X X +zstd-sys@2.0.16+zstd.1.5.7 X X diff --git a/crates/query-service-server/README.md b/crates/query-service-server/README.md new file mode 100644 index 00000000..10116644 --- /dev/null +++ b/crates/query-service-server/README.md @@ -0,0 +1,224 @@ + + +# Paimon Query Service + +This service hosts snapshot-consistent, budgeted point-query capabilities over +a Paimon catalog. Its first capability is `BlobDescriptor` batch lookup; the +service boundary is intentionally generic so other projected fields can be +added without creating another process. + +Create a JSON configuration file: + +```json +{ + "listen": "127.0.0.1:8081", + "queryTimeoutMs": 5000, + "httpRequestTimeoutMs": 10000, + "maxConcurrentQueries": 64, + "maxConcurrentIndexReads": 64, + "maxConnections": 1024, + "httpHeaderReadTimeoutMs": 5000, + "httpConnectionIdleTimeoutMs": 60000, + "httpConnectionMaxAgeMs": 300000, + "http2MaxConcurrentStreams": 64, + "gracefulShutdownTimeoutMs": 10000, + "queueTimeoutMs": 100, + "maxRequestBodyBytes": 1048576, + "maxResponseBodyBytes": 4194304, + "readinessCacheTtlMs": 1000, + "tableMetadataCacheTtlMs": 30000, + "descriptorCacheTtlMs": 60000, + "descriptorCacheMaxBytes": 67108864, + "catalog": { + "metastore": "filesystem", + "warehouse": "/tmp/paimon-warehouse" + }, + "policies": [ + { + "table": {"database": "default", "table": "assets"}, + "keyFields": ["asset_id"], + "blobFields": ["picture", "thumbnail"], + "strategy": "GLOBAL_BTREE", + "budget": { + "maxBatchKeys": 200, + "maxPlannedFiles": 16, + "maxPlannedBytes": 134217728 + } + } + ], + "principals": [ + { + "name": "asset-reader", + "bearerTokenEnv": "QUERY_SERVICE_ASSET_READER_TOKEN", + "grants": [ + { + "table": {"database": "default", "table": "assets"}, + "blobFields": ["picture"] + } + ] + } + ] +} +``` + +Start the server: + +```bash +QUERY_SERVICE_ASSET_READER_TOKEN=replace-with-a-long-random-token \ +QUERY_SERVICE_CONFIG=/path/to/query-service.json \ + cargo run -p paimon-query-service-server +``` + +`BLOB_QUERY_CONFIG` remains accepted as a compatibility fallback when +`QUERY_SERVICE_CONFIG` is unset. + +Query descriptors: + +```bash +curl -X POST \ + -H 'Authorization: Bearer replace-with-a-long-random-token' \ + -H 'Content-Type: application/json' \ + http://127.0.0.1:8081/api/blob/v1/databases/default/tables/assets/descriptors:batchGet \ + -d '{ + "keys": [{"asset_id": "9007199254740993"}], + "blobFields": ["picture"], + "descriptorFormat": "PAIMON_BASE64" + }' +``` + +Lookup key values use schema-driven JSON encoding: + +| Paimon key type | JSON encoding | +| --- | --- | +| `BOOLEAN` | JSON boolean | +| `TINYINT` / `SMALLINT` / `INT` | JSON integer or decimal integer string | +| `BIGINT` | Decimal integer string is recommended to avoid JSON number precision loss | +| `FLOAT` / `DOUBLE` | Finite JSON number or numeric string | +| `CHAR` / `VARCHAR` | JSON string | +| `BINARY` / `VARBINARY` | `{"base64":"..."}` | +| `DECIMAL(p,s)` | Exact decimal string without exponent notation, with at most `s` fractional digits | +| `DATE` | ISO string `YYYY-MM-DD` | +| `TIME(p)` | ISO string `HH:MM:SS[.fraction]`; lookup values currently have millisecond resolution | +| `TIMESTAMP(p)` | ISO string `YYYY-MM-DDTHH:MM:SS[.fraction]` without a time zone | +| `TIMESTAMP(p) WITH LOCAL TIME ZONE` | RFC 3339 string with `Z` or an explicit UTC offset | + +Fractional seconds may not exceed the field's declared precision. The server +rejects non-finite floats and values that would require rounding or truncation. +Collection, variant, vector, and BLOB values are not supported as lookup keys. +`keys` and `blobFields` must both be non-empty. + +`GLOBAL_BTREE` policies additionally require every key field to use a scalar +type supported by Paimon's sorted global indexes. In particular, binary key +types can be normalized for primary-key lookup but cannot be configured for +`GLOBAL_BTREE`. Index files are snapshot data and may have temporarily partial +coverage after writes; readiness validates table capabilities and key-type +compatibility, while per-policy planned-file and planned-byte budgets protect +the unindexed fallback path. Build or refresh the indexes with Paimon's +`create_global_index` procedure as part of the table maintenance workflow. + +All HTTP errors use a JSON body: + +```json +{"code":"INVALID_REQUEST","message":"keys must not be empty"} +``` + +Stable transport-level codes include `INVALID_JSON`, `INVALID_REQUEST`, +`UNSUPPORTED_MEDIA_TYPE`, `REQUEST_TOO_LARGE`, `ROUTE_NOT_FOUND`, and +`METHOD_NOT_ALLOWED`. `REQUEST_TIMEOUT` covers the complete HTTP lifecycle, +including slow request uploads; `RESPONSE_TOO_LARGE` rejects a successful +lookup whose serialized JSON exceeds the configured response envelope. The +`x-request-id` response header is present on these errors as well as successful +responses. + +Each principal receives explicit table and optional BLOB-field grants. Tokens +must contain at least 16 bytes and are retained by the running service only as +SHA-256 hashes; `bearerTokenEnv` avoids storing the secret in the JSON file. +Using `bearerToken` inside a principal is also supported. The legacy top-level +`bearerToken`/`bearerTokenEnv` grants access to all configured policies, while +anonymous access is disabled by default. To run without authentication, set +`"allowAnonymous": true` explicitly; it cannot be combined with bearer-token +authentication. Unknown configuration fields, including nested policy and grant +fields, fail startup so security-sensitive spelling mistakes cannot broaden +access. + +Missing or invalid credentials receive `401` with `WWW-Authenticate: Bearer`; +authenticated requests outside their grants receive `403`. Continue to place +the service behind TLS and the platform's normal identity boundary. + +Every response includes an `x-request-id` header. A valid caller-provided ID is +preserved; otherwise the server generates one. Access logs are emitted as JSON +through a bounded non-blocking queue to stderr and contain the request ID, +authenticated principal, method, path, status, and elapsed time. When stderr +cannot keep up, log events are dropped rather than blocking Tokio workers; the +dropped count is exported in Prometheus metrics. + +`maxConcurrentQueries` limits admitted descriptor requests per process and is +acquired before the JSON body is read. A request that cannot enter within +`queueTimeoutMs` receives `429 SERVER_BUSY`; a lookup that +exceeds `queryTimeoutMs` receives `504 QUERY_TIMEOUT`. The JSON body limit is +controlled by `maxRequestBodyBytes`, and successful JSON responses are bounded +by `maxResponseBodyBytes`. `httpRequestTimeoutMs` is an outer deadline covering +request upload, admission, lookup, and response construction. Keep it larger +than `queueTimeoutMs + queryTimeoutMs`. The values above are also the defaults. +`maxConcurrentIndexReads` bounds the product of admitted queries and each +query's global-index scanner width; it must be at least `maxConcurrentQueries`. + +At the transport layer, `maxConnections` bounds parsed and partially parsed TCP +connections, `httpHeaderReadTimeoutMs` closes slow HTTP/1 header uploads, and +also bounds the initial HTTP/1 versus HTTP/2 protocol detection. +`httpConnectionIdleTimeoutMs` closes read-idle HTTP/1 and HTTP/2 connections and +must exceed `httpRequestTimeoutMs`. At `httpConnectionMaxAgeMs`, the server +starts a graceful connection shutdown (HTTP/2 GOAWAY) and gives active work up +to `httpRequestTimeoutMs` to drain. This bounds HTTP/2 clients which continuously +drip incomplete frame or header bytes without resetting valid in-flight work; +the max age must exceed `httpConnectionIdleTimeoutMs`. +`http2MaxConcurrentStreams` bounds multiplexing per HTTP/2 connection. +`gracefulShutdownTimeoutMs` bounds connection drain after SIGTERM or Ctrl-C. + +Operational endpoints do not require the bearer token: + +- `GET /healthz` is a process liveness probe. +- `GET /readyz` reloads and validates every configured table, including catalog + connectivity and lookup-policy compatibility. Results are reused for + `readinessCacheTtlMs`, and concurrent refreshes collapse into one catalog + traversal; set the TTL to `0` to force every probe to refresh. +- `GET /metrics` exposes Prometheus text metrics for HTTP status classes, + latency histograms, authentication failures, admission, timeouts, in-flight work, + lookup outcomes, planned files/bytes, readiness refresh/cache activity, + total HTTP timeouts, oversized responses, and table/descriptor-cache activity. + +Keep these operational endpoints on an internal network or restrict them at the +ingress layer; `/readyz` deliberately performs catalog metadata I/O. + +Table metadata is cached per process for `tableMetadataCacheTtlMs`; set it to +`0` to disable caching. Snapshot IDs are still resolved and pinned per request, +and a readiness probe refreshes all cached table entries. + +Successful descriptor results are cached per process for +`descriptorCacheTtlMs`, bounded by `descriptorCacheMaxBytes`. Set either value +to `0` to disable this cache. The cache key contains the resolved snapshot ID, +schema ID, table location, immutable snapshot metadata fingerprint, requested +keys, BLOB fields, and descriptor format, +so a request for the latest data resolves and pins its snapshot before cache +lookup. Concurrent identical misses are collapsed into one scan. Responses +include `cacheHit`; `scan` describes the original plan, while Prometheus +planned-file/byte counters count only scans actually executed by this process. +A single cached entry is also bounded by `maxResponseBodyBytes`, so an oversized +result that cannot be returned is never retained in the descriptor cache. +Miss coalescing is independent of long-term cache admission, so concurrent +identical oversized results and failures also execute only one scan. diff --git a/crates/query-service-server/src/auth.rs b/crates/query-service-server/src/auth.rs new file mode 100644 index 00000000..6be4243d --- /dev/null +++ b/crates/query-service-server/src/auth.rs @@ -0,0 +1,380 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::{BTreeMap, BTreeSet, HashSet}; + +use paimon_query_service::{TableLookupPolicy, TableRef}; +use serde::Deserialize; +use sha2::{Digest, Sha256}; +use subtle::ConstantTimeEq; + +const MIN_TOKEN_BYTES: usize = 16; + +#[derive(Debug, Clone, Deserialize)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct PrincipalConfig { + pub name: String, + #[serde(default)] + pub bearer_token: Option, + #[serde(default)] + pub bearer_token_env: Option, + pub grants: Vec, +} + +#[derive(Debug, Clone, Deserialize)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct TableGrant { + pub table: TableRef, + #[serde(default)] + pub blob_fields: Option>, +} + +#[derive(Debug, Clone)] +pub(crate) struct Principal { + name: String, + grants: Option>>, +} + +impl Principal { + pub(crate) fn name(&self) -> &str { + &self.name + } + + pub(crate) fn allows(&self, table: &TableRef, blob_fields: &[String]) -> bool { + let Some(grants) = &self.grants else { + return true; + }; + grants + .get(table) + .is_some_and(|allowed| blob_fields.iter().all(|field| allowed.contains(field))) + } +} + +#[derive(Clone)] +pub(crate) struct Credential { + token_hash: [u8; 32], + principal: Principal, +} + +#[derive(Clone)] +pub(crate) enum AuthPolicy { + Anonymous(Principal), + Credentials(Vec), +} + +impl AuthPolicy { + pub(crate) fn build( + allow_anonymous: bool, + legacy_token: Option<&str>, + legacy_token_env: Option<&str>, + principals: &[PrincipalConfig], + policies: &[TableLookupPolicy], + ) -> Result { + if allow_anonymous + && (!principals.is_empty() || legacy_token.is_some() || legacy_token_env.is_some()) + { + return Err(AuthConfigError( + "allowAnonymous cannot be combined with bearer-token authentication".to_string(), + )); + } + if !principals.is_empty() && (legacy_token.is_some() || legacy_token_env.is_some()) { + return Err(AuthConfigError( + "principals cannot be combined with bearerToken or bearerTokenEnv".to_string(), + )); + } + + if principals.is_empty() { + let Some(token) = resolve_token(legacy_token, legacy_token_env, "legacy bearer token")? + else { + if !allow_anonymous { + return Err(AuthConfigError( + "authentication is required unless allowAnonymous is explicitly true" + .to_string(), + )); + } + return Ok(Self::Anonymous(Principal { + name: "anonymous".to_string(), + grants: None, + })); + }; + validate_token(&token, "legacy bearer token")?; + return Ok(Self::Credentials(vec![Credential { + token_hash: hash_token(&token), + principal: Principal { + name: "legacy".to_string(), + grants: None, + }, + }])); + } + + let policy_by_table = policies + .iter() + .map(|policy| (&policy.table, policy)) + .collect::>(); + let mut names = HashSet::new(); + let mut token_hashes = HashSet::new(); + let mut credentials = Vec::with_capacity(principals.len()); + for config in principals { + if config.name.trim().is_empty() || !names.insert(config.name.clone()) { + return Err(AuthConfigError(format!( + "principal name must be non-empty and unique: {:?}", + config.name + ))); + } + let token = resolve_token( + config.bearer_token.as_deref(), + config.bearer_token_env.as_deref(), + &format!("principal '{}'", config.name), + )? + .ok_or_else(|| { + AuthConfigError(format!( + "principal '{}' must configure bearerToken or bearerTokenEnv", + config.name + )) + })?; + validate_token(&token, &format!("principal '{}'", config.name))?; + let token_hash = hash_token(&token); + if !token_hashes.insert(token_hash) { + return Err(AuthConfigError( + "bearer tokens must be unique across principals".to_string(), + )); + } + if config.grants.is_empty() { + return Err(AuthConfigError(format!( + "principal '{}' has no table grants", + config.name + ))); + } + + let mut grants = BTreeMap::new(); + for grant in &config.grants { + let policy = policy_by_table.get(&grant.table).ok_or_else(|| { + AuthConfigError(format!( + "principal '{}' grants unknown table {}.{}", + config.name, grant.table.database, grant.table.table + )) + })?; + let allowed = match &grant.blob_fields { + None => policy.blob_fields.clone(), + Some(fields) if fields.is_empty() => { + return Err(AuthConfigError(format!( + "principal '{}' has an empty BLOB field grant for {}.{}", + config.name, grant.table.database, grant.table.table + ))) + } + Some(fields) if fields.is_subset(&policy.blob_fields) => fields.clone(), + Some(fields) => { + let invalid = fields + .difference(&policy.blob_fields) + .cloned() + .collect::>(); + return Err(AuthConfigError(format!( + "principal '{}' grants disallowed BLOB fields {:?} for {}.{}", + config.name, invalid, grant.table.database, grant.table.table + ))); + } + }; + if grants.insert(grant.table.clone(), allowed).is_some() { + return Err(AuthConfigError(format!( + "principal '{}' has duplicate grants for {}.{}", + config.name, grant.table.database, grant.table.table + ))); + } + } + credentials.push(Credential { + token_hash, + principal: Principal { + name: config.name.clone(), + grants: Some(grants), + }, + }); + } + Ok(Self::Credentials(credentials)) + } + + pub(crate) fn authenticate(&self, token: Option<&str>) -> Option { + match self { + Self::Anonymous(principal) => Some(principal.clone()), + Self::Credentials(credentials) => { + let candidate = hash_token(token?); + let mut matched = None; + for credential in credentials { + if bool::from(candidate.ct_eq(&credential.token_hash)) { + matched = Some(credential.principal.clone()); + } + } + matched + } + } + } +} + +fn resolve_token( + inline: Option<&str>, + environment: Option<&str>, + owner: &str, +) -> Result, AuthConfigError> { + match (inline, environment) { + (Some(_), Some(_)) => Err(AuthConfigError(format!( + "{owner} must configure only one of bearerToken and bearerTokenEnv" + ))), + (Some(token), None) => Ok(Some(token.to_string())), + (None, Some(variable)) if variable.trim().is_empty() => Err(AuthConfigError(format!( + "{owner} bearerTokenEnv must not be empty" + ))), + (None, Some(variable)) => std::env::var(variable).map(Some).map_err(|error| { + AuthConfigError(format!( + "failed to read bearer token for {owner} from environment variable '{variable}': {error}" + )) + }), + (None, None) => Ok(None), + } +} + +fn validate_token(token: &str, owner: &str) -> Result<(), AuthConfigError> { + if token.len() < MIN_TOKEN_BYTES { + return Err(AuthConfigError(format!( + "{owner} bearer token must contain at least {MIN_TOKEN_BYTES} bytes" + ))); + } + Ok(()) +} + +fn hash_token(token: &str) -> [u8; 32] { + Sha256::digest(token.as_bytes()).into() +} + +#[derive(Debug)] +pub(crate) struct AuthConfigError(String); + +impl std::fmt::Display for AuthConfigError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "invalid authentication configuration: {}", self.0) + } +} + +impl std::error::Error for AuthConfigError {} + +#[cfg(test)] +mod tests { + use std::collections::BTreeSet; + + use paimon_query_service::{LookupStrategy, QueryBudget}; + + use super::*; + + fn policy() -> TableLookupPolicy { + TableLookupPolicy { + table: TableRef::new("db", "assets"), + key_fields: vec!["id".to_string()], + blob_fields: BTreeSet::from(["picture".to_string(), "thumbnail".to_string()]), + strategy: LookupStrategy::GlobalBtree, + budget: QueryBudget::default(), + } + } + + #[test] + fn principal_is_limited_to_granted_blob_fields() { + let auth = AuthPolicy::build( + false, + None, + None, + &[PrincipalConfig { + name: "image-reader".to_string(), + bearer_token: Some("long-test-token-1234".to_string()), + bearer_token_env: None, + grants: vec![TableGrant { + table: TableRef::new("db", "assets"), + blob_fields: Some(BTreeSet::from(["picture".to_string()])), + }], + }], + &[policy()], + ) + .unwrap(); + + let principal = auth.authenticate(Some("long-test-token-1234")).unwrap(); + assert_eq!(principal.name(), "image-reader"); + assert!(principal.allows(&TableRef::new("db", "assets"), &["picture".to_string()])); + assert!(!principal.allows(&TableRef::new("db", "assets"), &["thumbnail".to_string()])); + assert!(auth.authenticate(Some("invalid-test-token")).is_none()); + } + + #[test] + fn rejects_unknown_tables_and_duplicate_tokens() { + let unknown = match AuthPolicy::build( + false, + None, + None, + &[PrincipalConfig { + name: "reader".to_string(), + bearer_token: Some("long-test-token-1234".to_string()), + bearer_token_env: None, + grants: vec![TableGrant { + table: TableRef::new("db", "missing"), + blob_fields: None, + }], + }], + &[policy()], + ) { + Ok(_) => panic!("unknown grant table must be rejected"), + Err(error) => error, + }; + assert!(unknown.to_string().contains("unknown table")); + + let duplicate = match AuthPolicy::build( + false, + None, + None, + &[ + PrincipalConfig { + name: "one".to_string(), + bearer_token: Some("long-test-token-1234".to_string()), + bearer_token_env: None, + grants: vec![TableGrant { + table: TableRef::new("db", "assets"), + blob_fields: None, + }], + }, + PrincipalConfig { + name: "two".to_string(), + bearer_token: Some("long-test-token-1234".to_string()), + bearer_token_env: None, + grants: vec![TableGrant { + table: TableRef::new("db", "assets"), + blob_fields: None, + }], + }, + ], + &[policy()], + ) { + Ok(_) => panic!("duplicate bearer tokens must be rejected"), + Err(error) => error, + }; + assert!(duplicate.to_string().contains("unique")); + } + + #[test] + fn anonymous_access_requires_explicit_opt_in() { + let error = match AuthPolicy::build(false, None, None, &[], &[policy()]) { + Ok(_) => panic!("implicit anonymous access must be rejected"), + Err(error) => error, + }; + assert!(error.to_string().contains("allowAnonymous")); + + let auth = AuthPolicy::build(true, None, None, &[], &[policy()]).unwrap(); + assert_eq!(auth.authenticate(None).unwrap().name(), "anonymous"); + } +} diff --git a/crates/query-service-server/src/lib.rs b/crates/query-service-server/src/lib.rs new file mode 100644 index 00000000..b496f880 --- /dev/null +++ b/crates/query-service-server/src/lib.rs @@ -0,0 +1,2179 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +mod auth; + +use std::collections::HashMap; +use std::future::Future; +use std::net::SocketAddr; +use std::path::Path; +use std::pin::Pin; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{Arc, Mutex, OnceLock}; +use std::task::{Context, Poll}; +use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; + +use axum::extract::rejection::JsonRejection; +use axum::extract::{DefaultBodyLimit, FromRequestParts, Path as AxumPath, Request, State}; +use axum::http::request::Parts; +use axum::http::{header, HeaderValue, StatusCode}; +use axum::middleware::{self, Next}; +use axum::response::{IntoResponse, Response}; +use axum::routing::{get, post}; +use axum::{Json, Router}; +use hyper_util::rt::{TokioExecutor, TokioIo, TokioTimer}; +use hyper_util::server::conn::auto; +use hyper_util::service::TowerToHyperService; +use paimon::{CatalogFactory, Options}; +use paimon_query_service::{ + BatchGetRequest, BatchGetResponse, BlobLookupOptions, BlobLookupService, DescriptorCacheStats, + DescriptorFormat, LookupError, LookupKey, LookupStatus, TableCacheStats, TableLookupPolicy, + TableRef, +}; +use serde::{Deserialize, Serialize}; +use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, ReadBuf}; +use tokio::sync::{watch, Mutex as AsyncMutex, OwnedSemaphorePermit, Semaphore}; +use tokio::task::JoinSet; +use tracing::Instrument; +use tracing_appender::non_blocking::{ErrorCounter, NonBlockingBuilder, WorkerGuard}; +use tracing_subscriber::EnvFilter; + +use crate::auth::{AuthPolicy, Principal}; + +pub use crate::auth::{PrincipalConfig, TableGrant}; + +pub type BoxError = Box; + +const REQUEST_ID_HEADER: &str = "x-request-id"; +const HTTP_DURATION_BUCKETS: [(&str, u64); 10] = [ + ("0.001", 1_000), + ("0.005", 5_000), + ("0.01", 10_000), + ("0.025", 25_000), + ("0.05", 50_000), + ("0.1", 100_000), + ("0.25", 250_000), + ("0.5", 500_000), + ("1", 1_000_000), + ("5", 5_000_000), +]; +static REQUEST_SEQUENCE: AtomicU64 = AtomicU64::new(1); +static LOG_ERROR_COUNTER: OnceLock = OnceLock::new(); + +pub struct LoggingGuard { + _guard: WorkerGuard, +} + +pub fn init_logging() -> Result { + let (writer, guard) = NonBlockingBuilder::default() + .buffered_lines_limit(8_192) + .lossy(true) + .finish(std::io::stderr()); + let error_counter = writer.error_counter(); + let filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info")); + let subscriber = tracing_subscriber::fmt() + .with_env_filter(filter) + .json() + .with_writer(writer) + .finish(); + tracing::subscriber::set_global_default(subscriber)?; + let _ = LOG_ERROR_COUNTER.set(error_counter); + Ok(LoggingGuard { _guard: guard }) +} + +#[derive(Debug, Clone, Deserialize)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct ServerConfig { + #[serde(default = "default_listen")] + pub listen: String, + pub catalog: HashMap, + pub policies: Vec, + #[serde(default)] + pub bearer_token: Option, + #[serde(default)] + pub bearer_token_env: Option, + #[serde(default)] + pub principals: Vec, + #[serde(default)] + pub allow_anonymous: bool, + #[serde(default = "default_query_timeout_ms")] + pub query_timeout_ms: u64, + #[serde(default = "default_http_request_timeout_ms")] + pub http_request_timeout_ms: u64, + #[serde(default = "default_max_concurrent_queries")] + pub max_concurrent_queries: usize, + #[serde(default = "default_max_concurrent_index_reads")] + pub max_concurrent_index_reads: usize, + #[serde(default = "default_max_connections")] + pub max_connections: usize, + #[serde(default = "default_http_header_read_timeout_ms")] + pub http_header_read_timeout_ms: u64, + #[serde(default = "default_http_connection_idle_timeout_ms")] + pub http_connection_idle_timeout_ms: u64, + #[serde(default = "default_http_connection_max_age_ms")] + pub http_connection_max_age_ms: u64, + #[serde(default = "default_http2_max_concurrent_streams")] + pub http2_max_concurrent_streams: u32, + #[serde(default = "default_graceful_shutdown_timeout_ms")] + pub graceful_shutdown_timeout_ms: u64, + #[serde(default = "default_queue_timeout_ms")] + pub queue_timeout_ms: u64, + #[serde(default = "default_max_request_body_bytes")] + pub max_request_body_bytes: usize, + #[serde(default = "default_max_response_body_bytes")] + pub max_response_body_bytes: usize, + #[serde(default = "default_readiness_cache_ttl_ms")] + pub readiness_cache_ttl_ms: u64, + #[serde(default = "default_table_metadata_cache_ttl_ms")] + pub table_metadata_cache_ttl_ms: u64, + #[serde(default = "default_descriptor_cache_ttl_ms")] + pub descriptor_cache_ttl_ms: u64, + #[serde(default = "default_descriptor_cache_max_bytes")] + pub descriptor_cache_max_bytes: u64, +} + +fn default_listen() -> String { + "127.0.0.1:8081".to_string() +} + +fn default_query_timeout_ms() -> u64 { + 5_000 +} + +fn default_http_request_timeout_ms() -> u64 { + 10_000 +} + +fn default_max_concurrent_queries() -> usize { + 64 +} + +fn default_max_concurrent_index_reads() -> usize { + 64 +} + +fn default_max_connections() -> usize { + 1_024 +} + +fn default_http_header_read_timeout_ms() -> u64 { + 5_000 +} + +fn default_http_connection_idle_timeout_ms() -> u64 { + 60_000 +} + +fn default_http_connection_max_age_ms() -> u64 { + 300_000 +} + +fn default_http2_max_concurrent_streams() -> u32 { + 64 +} + +fn default_graceful_shutdown_timeout_ms() -> u64 { + 10_000 +} + +fn default_queue_timeout_ms() -> u64 { + 100 +} + +fn default_max_request_body_bytes() -> usize { + 1024 * 1024 +} + +fn default_max_response_body_bytes() -> usize { + 4 * 1024 * 1024 +} + +fn default_readiness_cache_ttl_ms() -> u64 { + 1_000 +} + +fn default_table_metadata_cache_ttl_ms() -> u64 { + 30_000 +} + +fn default_descriptor_cache_ttl_ms() -> u64 { + 60_000 +} + +fn default_descriptor_cache_max_bytes() -> u64 { + 64 * 1024 * 1024 +} + +impl ServerConfig { + fn validate(&self) -> Result<(), BoxError> { + if self.query_timeout_ms == 0 { + return Err(invalid_config("queryTimeoutMs must be positive")); + } + if self.http_request_timeout_ms == 0 { + return Err(invalid_config("httpRequestTimeoutMs must be positive")); + } + if self.max_concurrent_queries == 0 || self.max_concurrent_queries > Semaphore::MAX_PERMITS + { + return Err(invalid_config(format!( + "maxConcurrentQueries must be between 1 and {}", + Semaphore::MAX_PERMITS + ))); + } + if self.max_concurrent_index_reads < self.max_concurrent_queries { + return Err(invalid_config( + "maxConcurrentIndexReads must be at least maxConcurrentQueries", + )); + } + if self.max_connections == 0 || self.max_connections > Semaphore::MAX_PERMITS { + return Err(invalid_config(format!( + "maxConnections must be between 1 and {}", + Semaphore::MAX_PERMITS + ))); + } + if self.http_header_read_timeout_ms == 0 { + return Err(invalid_config("httpHeaderReadTimeoutMs must be positive")); + } + if self.http_connection_idle_timeout_ms <= self.http_request_timeout_ms { + return Err(invalid_config( + "httpConnectionIdleTimeoutMs must be greater than httpRequestTimeoutMs", + )); + } + if self.http_connection_max_age_ms <= self.http_connection_idle_timeout_ms { + return Err(invalid_config( + "httpConnectionMaxAgeMs must be greater than httpConnectionIdleTimeoutMs", + )); + } + if self.http2_max_concurrent_streams == 0 { + return Err(invalid_config("http2MaxConcurrentStreams must be positive")); + } + if self.graceful_shutdown_timeout_ms == 0 { + return Err(invalid_config("gracefulShutdownTimeoutMs must be positive")); + } + if self.queue_timeout_ms == 0 { + return Err(invalid_config("queueTimeoutMs must be positive")); + } + if self.http_request_timeout_ms + <= self.query_timeout_ms.saturating_add(self.queue_timeout_ms) + { + return Err(invalid_config( + "httpRequestTimeoutMs must be greater than queryTimeoutMs + queueTimeoutMs", + )); + } + if self.max_request_body_bytes == 0 { + return Err(invalid_config("maxRequestBodyBytes must be positive")); + } + if self.max_response_body_bytes == 0 { + return Err(invalid_config("maxResponseBodyBytes must be positive")); + } + Ok(()) + } +} + +fn invalid_config(message: impl Into) -> BoxError { + std::io::Error::new(std::io::ErrorKind::InvalidInput, message.into()).into() +} + +pub fn load_config(path: impl AsRef) -> Result { + let bytes = std::fs::read(path)?; + Ok(serde_json::from_slice(&bytes)?) +} + +pub async fn build_app(config: &ServerConfig) -> Result { + config.validate()?; + let descriptor_cache_max_entry_bytes = u64::try_from(config.max_response_body_bytes) + .map_err(|_| invalid_config("maxResponseBodyBytes does not fit the cache size type"))?; + let auth = AuthPolicy::build( + config.allow_anonymous, + config.bearer_token.as_deref(), + config.bearer_token_env.as_deref(), + &config.principals, + &config.policies, + )?; + let catalog = CatalogFactory::create(Options::from_map(config.catalog.clone())).await?; + let lookup = BlobLookupService::new_with_options( + catalog, + config.policies.clone(), + BlobLookupOptions { + table_cache_ttl: Duration::from_millis(config.table_metadata_cache_ttl_ms), + descriptor_cache_ttl: Duration::from_millis(config.descriptor_cache_ttl_ms), + descriptor_cache_max_bytes: config.descriptor_cache_max_bytes, + descriptor_cache_max_entry_bytes, + global_index_thread_num: config.max_concurrent_index_reads + / config.max_concurrent_queries, + }, + )?; + let metrics = Arc::new(ServerMetrics::default()); + Ok(router( + AppState { + lookup, + auth, + query_permits: Arc::new(Semaphore::new(config.max_concurrent_queries)), + query_timeout: Duration::from_millis(config.query_timeout_ms), + queue_timeout: Duration::from_millis(config.queue_timeout_ms), + max_response_body_bytes: config.max_response_body_bytes, + readiness: ReadinessCache::new(Duration::from_millis(config.readiness_cache_ttl_ms)), + metrics, + }, + config.max_request_body_bytes, + Duration::from_millis(config.http_request_timeout_ms), + )) +} + +pub async fn serve(config: ServerConfig) -> Result<(), BoxError> { + let address: SocketAddr = config.listen.parse()?; + let app = build_app(&config).await?; + let transport = TransportConfig { + max_connections: config.max_connections, + header_read_timeout: Duration::from_millis(config.http_header_read_timeout_ms), + connection_idle_timeout: Duration::from_millis(config.http_connection_idle_timeout_ms), + connection_max_age: Duration::from_millis(config.http_connection_max_age_ms), + request_drain_timeout: Duration::from_millis(config.http_request_timeout_ms), + http2_max_concurrent_streams: config.http2_max_concurrent_streams, + graceful_shutdown_timeout: Duration::from_millis(config.graceful_shutdown_timeout_ms), + }; + let listener = tokio::net::TcpListener::bind(address).await?; + let bound_address = listener.local_addr()?; + tracing::info!( + event = "server_listening", + address = %bound_address, + max_connections = transport.max_connections, + "query service is listening" + ); + serve_listener(listener, app, transport).await?; + Ok(()) +} + +#[derive(Debug, Clone, Copy)] +struct TransportConfig { + max_connections: usize, + header_read_timeout: Duration, + connection_idle_timeout: Duration, + connection_max_age: Duration, + request_drain_timeout: Duration, + http2_max_concurrent_streams: u32, + graceful_shutdown_timeout: Duration, +} + +async fn serve_listener( + listener: tokio::net::TcpListener, + app: Router, + transport: TransportConfig, +) -> Result<(), BoxError> { + let connection_permits = Arc::new(Semaphore::new(transport.max_connections)); + let (connection_shutdown, shutdown_receiver) = watch::channel(false); + let mut connections = JoinSet::new(); + let shutdown = shutdown_signal(); + tokio::pin!(shutdown); + + loop { + let accepted = tokio::select! { + accepted = listener.accept() => Some(accepted), + _ = &mut shutdown => None, + completed = connections.join_next(), if !connections.is_empty() => { + log_connection_task_result(completed); + continue; + } + }; + let Some(accepted) = accepted else { + break; + }; + let (stream, peer_address) = match accepted { + Ok(connection) => connection, + Err(error) => { + tracing::warn!(%error, "failed to accept TCP connection"); + tokio::time::sleep(Duration::from_millis(100)).await; + continue; + } + }; + let permit = match connection_permits.clone().try_acquire_owned() { + Ok(permit) => permit, + Err(_) => { + tracing::warn!(%peer_address, "connection limit reached"); + drop(stream); + continue; + } + }; + let shutdown_receiver = shutdown_receiver.clone(); + let service = TowerToHyperService::new(app.clone()); + connections.spawn(async move { + let _permit = permit; + let stream = match read_protocol_prefix(stream, transport.header_read_timeout).await { + Ok(stream) => stream, + Err(error) => { + tracing::debug!(%peer_address, %error, "HTTP protocol detection failed"); + return; + } + }; + let io = TokioIo::new(ReadIdleTimeout::new( + stream, + transport.connection_idle_timeout, + )); + let mut server = auto::Builder::new(TokioExecutor::new()); + server + .http1() + .timer(TokioTimer::new()) + .header_read_timeout(transport.header_read_timeout) + .max_headers(64); + server + .http2() + .timer(TokioTimer::new()) + .max_concurrent_streams(transport.http2_max_concurrent_streams); + let connection = server.serve_connection_with_upgrades(io, service); + let result = drive_connection( + connection, + transport.connection_max_age, + transport.request_drain_timeout, + transport.graceful_shutdown_timeout, + shutdown_receiver, + |connection| connection.graceful_shutdown(), + ) + .await; + match result { + ConnectionEnd::Completed(Ok(())) => {} + ConnectionEnd::Drained { + reason, + output: Ok(()), + } => { + tracing::debug!(%peer_address, ?reason, "HTTP connection drained"); + } + ConnectionEnd::Completed(Err(error)) => { + tracing::debug!(%peer_address, %error, "HTTP connection closed with an error"); + } + ConnectionEnd::Drained { + reason, + output: Err(error), + } => { + tracing::debug!(%peer_address, ?reason, %error, "HTTP connection drain ended with an error"); + } + ConnectionEnd::DrainTimedOut(ConnectionShutdownReason::MaxAge) => { + tracing::debug!(%peer_address, "HTTP connection max-age drain timed out"); + } + ConnectionEnd::DrainTimedOut(ConnectionShutdownReason::ServerShutdown) => { + tracing::debug!(%peer_address, "HTTP connection shutdown drain timed out"); + } + } + }); + } + + let _ = connection_shutdown.send(true); + let drain_connections = async { + while let Some(completed) = connections.join_next().await { + log_connection_task_result(Some(completed)); + } + }; + if tokio::time::timeout(transport.graceful_shutdown_timeout, drain_connections) + .await + .is_err() + { + tracing::warn!("graceful shutdown deadline expired"); + connections.abort_all(); + while let Some(completed) = connections.join_next().await { + log_connection_task_result(Some(completed)); + } + } + Ok(()) +} + +fn log_connection_task_result(result: Option>) { + if let Some(Err(error)) = result { + if !error.is_cancelled() { + tracing::warn!(%error, "HTTP connection task failed"); + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum ConnectionShutdownReason { + MaxAge, + ServerShutdown, +} + +enum ConnectionEnd { + Completed(T), + Drained { + reason: ConnectionShutdownReason, + output: T, + }, + DrainTimedOut(ConnectionShutdownReason), +} + +async fn drive_connection( + connection: F, + max_age: Duration, + max_age_drain_timeout: Duration, + server_drain_timeout: Duration, + mut shutdown: watch::Receiver, + graceful_shutdown: G, +) -> ConnectionEnd +where + F: Future, + G: FnOnce(Pin<&mut F>), +{ + let mut connection = Box::pin(connection); + let max_age = tokio::time::sleep(max_age); + tokio::pin!(max_age); + let reason = tokio::select! { + output = connection.as_mut() => return ConnectionEnd::Completed(output), + _ = &mut max_age => ConnectionShutdownReason::MaxAge, + _ = shutdown.changed() => ConnectionShutdownReason::ServerShutdown, + }; + graceful_shutdown(connection.as_mut()); + let drain_timeout = match reason { + ConnectionShutdownReason::MaxAge => max_age_drain_timeout, + ConnectionShutdownReason::ServerShutdown => server_drain_timeout, + }; + match tokio::time::timeout(drain_timeout, connection.as_mut()).await { + Ok(output) => ConnectionEnd::Drained { reason, output }, + Err(_) => ConnectionEnd::DrainTimedOut(reason), + } +} + +const HTTP2_PREFACE: &[u8] = b"PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n"; + +async fn read_protocol_prefix(mut stream: T, timeout: Duration) -> std::io::Result> +where + T: AsyncRead + Unpin, +{ + let mut prefix = Vec::with_capacity(HTTP2_PREFACE.len()); + tokio::time::timeout(timeout, async { + loop { + let mut byte = [0_u8; 1]; + if stream.read(&mut byte).await? == 0 { + return Err(std::io::Error::new( + std::io::ErrorKind::UnexpectedEof, + "connection closed before HTTP protocol detection", + )); + } + prefix.push(byte[0]); + if !HTTP2_PREFACE.starts_with(&prefix) || prefix.len() == HTTP2_PREFACE.len() { + return Ok(()); + } + } + }) + .await + .map_err(|_| { + std::io::Error::new( + std::io::ErrorKind::TimedOut, + "HTTP protocol preface exceeded its deadline", + ) + })??; + Ok(PrefixedIo { + inner: stream, + prefix, + prefix_offset: 0, + }) +} + +struct PrefixedIo { + inner: T, + prefix: Vec, + prefix_offset: usize, +} + +impl AsyncRead for PrefixedIo { + fn poll_read( + mut self: Pin<&mut Self>, + context: &mut Context<'_>, + buffer: &mut ReadBuf<'_>, + ) -> Poll> { + if self.prefix_offset < self.prefix.len() && buffer.remaining() > 0 { + let remaining = &self.prefix[self.prefix_offset..]; + let length = remaining.len().min(buffer.remaining()); + buffer.put_slice(&remaining[..length]); + self.prefix_offset += length; + return Poll::Ready(Ok(())); + } + Pin::new(&mut self.inner).poll_read(context, buffer) + } +} + +impl AsyncWrite for PrefixedIo { + fn poll_write( + mut self: Pin<&mut Self>, + context: &mut Context<'_>, + buffer: &[u8], + ) -> Poll> { + Pin::new(&mut self.inner).poll_write(context, buffer) + } + + fn poll_flush( + mut self: Pin<&mut Self>, + context: &mut Context<'_>, + ) -> Poll> { + Pin::new(&mut self.inner).poll_flush(context) + } + + fn poll_shutdown( + mut self: Pin<&mut Self>, + context: &mut Context<'_>, + ) -> Poll> { + Pin::new(&mut self.inner).poll_shutdown(context) + } +} + +struct ReadIdleTimeout { + inner: T, + timeout: Duration, + sleep: Pin>, +} + +impl ReadIdleTimeout { + fn new(inner: T, timeout: Duration) -> Self { + Self { + inner, + timeout, + sleep: Box::pin(tokio::time::sleep(timeout)), + } + } +} + +impl AsyncRead for ReadIdleTimeout { + fn poll_read( + mut self: Pin<&mut Self>, + context: &mut Context<'_>, + buffer: &mut ReadBuf<'_>, + ) -> Poll> { + let before = buffer.filled().len(); + match Pin::new(&mut self.inner).poll_read(context, buffer) { + Poll::Ready(result) => { + if result.is_ok() && buffer.filled().len() > before { + let deadline = tokio::time::Instant::now() + self.timeout; + self.sleep.as_mut().reset(deadline); + } + Poll::Ready(result) + } + Poll::Pending => match self.sleep.as_mut().poll(context) { + Poll::Ready(()) => Poll::Ready(Err(std::io::Error::new( + std::io::ErrorKind::TimedOut, + "HTTP connection exceeded its read-idle timeout", + ))), + Poll::Pending => Poll::Pending, + }, + } + } +} + +impl AsyncWrite for ReadIdleTimeout { + fn poll_write( + mut self: Pin<&mut Self>, + context: &mut Context<'_>, + buffer: &[u8], + ) -> Poll> { + Pin::new(&mut self.inner).poll_write(context, buffer) + } + + fn poll_flush( + mut self: Pin<&mut Self>, + context: &mut Context<'_>, + ) -> Poll> { + Pin::new(&mut self.inner).poll_flush(context) + } + + fn poll_shutdown( + mut self: Pin<&mut Self>, + context: &mut Context<'_>, + ) -> Poll> { + Pin::new(&mut self.inner).poll_shutdown(context) + } +} + +#[derive(Clone)] +struct AppState { + lookup: BlobLookupService, + auth: AuthPolicy, + query_permits: Arc, + query_timeout: Duration, + queue_timeout: Duration, + max_response_body_bytes: usize, + readiness: ReadinessCache, + metrics: Arc, +} + +#[derive(Clone)] +struct ReadinessCache { + ttl: Duration, + cached: Arc>>, +} + +#[derive(Clone, Copy)] +struct CachedReadiness { + checked_at: Instant, + ready: bool, +} + +impl ReadinessCache { + fn new(ttl: Duration) -> Self { + Self { + ttl, + cached: Arc::new(AsyncMutex::new(None)), + } + } + + async fn check(&self, lookup: &BlobLookupService, metrics: &ServerMetrics) -> bool { + // Keep the guard during refresh so concurrent probes collapse into a + // single catalog traversal. Probe traffic must not amplify metadata I/O. + let mut cached = self.cached.lock().await; + if cached + .as_ref() + .is_some_and(|value| value.checked_at.elapsed() < self.ttl) + { + metrics.readiness_cache_hits.fetch_add(1, Ordering::Relaxed); + return cached.as_ref().is_some_and(|value| value.ready); + } + + metrics.readiness_checks.fetch_add(1, Ordering::Relaxed); + let ready = lookup.check_ready().await.is_ok(); + *cached = Some(CachedReadiness { + checked_at: Instant::now(), + ready, + }); + ready + } +} + +#[derive(Clone)] +struct HttpTimeoutState { + timeout: Duration, + metrics: Arc, +} + +fn router( + state: AppState, + max_request_body_bytes: usize, + http_request_timeout: Duration, +) -> Router { + let metrics = state.metrics.clone(); + let timeout = HttpTimeoutState { + timeout: http_request_timeout, + metrics: metrics.clone(), + }; + let query_admission = state.clone(); + Router::new() + .route("/healthz", get(health)) + .route("/readyz", get(ready)) + .route("/metrics", get(metrics_endpoint)) + .route( + "/api/blob/v1/databases/:database/tables/:table/descriptors:batchGet", + post(batch_get) + .route_layer(middleware::from_fn_with_state(query_admission, admit_query)), + ) + .fallback(not_found) + .method_not_allowed_fallback(method_not_allowed) + .with_state(state) + .layer(DefaultBodyLimit::max(max_request_body_bytes)) + .layer(middleware::from_fn_with_state( + timeout, + enforce_http_timeout, + )) + .layer(middleware::from_fn_with_state(metrics, request_context)) +} + +async fn health() -> StatusCode { + StatusCode::NO_CONTENT +} + +async fn ready(State(state): State) -> Response { + match tokio::time::timeout( + state.query_timeout, + state.readiness.check(&state.lookup, &state.metrics), + ) + .await + { + Ok(true) => StatusCode::NO_CONTENT.into_response(), + Ok(false) => api_error( + StatusCode::SERVICE_UNAVAILABLE, + "NOT_READY", + "one or more configured tables are unavailable or incompatible", + ), + Err(_) => api_error( + StatusCode::SERVICE_UNAVAILABLE, + "READINESS_TIMEOUT", + "table readiness verification exceeded its deadline", + ), + } +} + +async fn metrics_endpoint(State(state): State) -> Response { + let body = state.metrics.render( + state.lookup.table_cache_stats(), + state.lookup.descriptor_cache_stats(), + ); + ( + [( + header::CONTENT_TYPE, + "text/plain; version=0.0.4; charset=utf-8", + )], + body, + ) + .into_response() +} + +async fn not_found() -> Response { + api_error(StatusCode::NOT_FOUND, "ROUTE_NOT_FOUND", "route not found") +} + +async fn method_not_allowed() -> Response { + api_error( + StatusCode::METHOD_NOT_ALLOWED, + "METHOD_NOT_ALLOWED", + "HTTP method is not allowed for this route", + ) +} + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +struct ApiBatchGetRequest { + keys: Vec, + blob_fields: Vec, + #[serde(default)] + snapshot_id: Option, + #[serde(default)] + descriptor_format: DescriptorFormat, +} + +#[derive(Clone)] +struct Authenticated(Principal); + +#[axum::async_trait] +impl FromRequestParts for Authenticated { + type Rejection = Response; + + async fn from_request_parts( + parts: &mut Parts, + state: &AppState, + ) -> Result { + let token = parts + .headers + .get(header::AUTHORIZATION) + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.strip_prefix("Bearer ")); + let principal = state.auth.authenticate(token).ok_or_else(|| { + state + .metrics + .auth_unauthorized + .fetch_add(1, Ordering::Relaxed); + unauthorized_response() + })?; + if let Some(context) = parts.extensions.get::() { + context.set_principal(principal.name()); + } + Ok(Self(principal)) + } +} + +async fn batch_get( + Authenticated(principal): Authenticated, + AxumPath((database, table)): AxumPath<(String, String)>, + State(state): State, + request: Result, JsonRejection>, +) -> Response { + let Json(request) = match request { + Ok(request) => request, + Err(rejection) => return json_rejection_response(rejection), + }; + let table = TableRef::new(database, table); + if !principal.allows(&table, &request.blob_fields) { + state.metrics.auth_forbidden.fetch_add(1, Ordering::Relaxed); + return api_error( + StatusCode::FORBIDDEN, + "FORBIDDEN", + "principal is not authorized for the requested table or BLOB fields", + ); + } + + let request = BatchGetRequest { + table, + keys: request.keys, + blob_fields: request.blob_fields, + snapshot_id: request.snapshot_id, + descriptor_format: request.descriptor_format, + }; + let response = + match tokio::time::timeout(state.query_timeout, state.lookup.batch_get(request)).await { + Ok(Ok(response)) => { + state.metrics.record_lookup(&response); + success_response(response, state.max_response_body_bytes, &state.metrics) + } + Ok(Err(error)) => lookup_error_response(error), + Err(_) => { + state.metrics.query_timeouts.fetch_add(1, Ordering::Relaxed); + api_error( + StatusCode::GATEWAY_TIMEOUT, + "QUERY_TIMEOUT", + "blob descriptor query exceeded its deadline", + ) + } + }; + response +} + +async fn admit_query(State(state): State, request: Request, next: Next) -> Response { + let _permit = match acquire_query_permit(state.query_permits.clone(), state.queue_timeout).await + { + Ok(permit) => permit, + Err(AdmissionError::TimedOut) => { + state.metrics.query_rejected.fetch_add(1, Ordering::Relaxed); + let mut response = api_error( + StatusCode::TOO_MANY_REQUESTS, + "SERVER_BUSY", + "query concurrency limit reached", + ); + response + .headers_mut() + .insert(header::RETRY_AFTER, HeaderValue::from_static("1")); + return response; + } + Err(AdmissionError::Closed) => { + state.metrics.query_rejected.fetch_add(1, Ordering::Relaxed); + return api_error( + StatusCode::SERVICE_UNAVAILABLE, + "SERVER_SHUTTING_DOWN", + "query admission is closed", + ); + } + }; + state.metrics.query_started.fetch_add(1, Ordering::Relaxed); + let _inflight = InflightQueryGuard::new(state.metrics.clone()); + next.run(request).await +} + +fn success_response( + response: BatchGetResponse, + max_response_body_bytes: usize, + metrics: &ServerMetrics, +) -> Response { + let body = match serde_json::to_vec(&response) { + Ok(body) => body, + Err(_) => { + return api_error( + StatusCode::INTERNAL_SERVER_ERROR, + "RESPONSE_SERIALIZATION_FAILED", + "failed to serialize blob descriptor response", + ) + } + }; + if body.len() > max_response_body_bytes { + metrics.responses_too_large.fetch_add(1, Ordering::Relaxed); + return api_error( + StatusCode::INSUFFICIENT_STORAGE, + "RESPONSE_TOO_LARGE", + "blob descriptor response exceeds the configured limit", + ); + } + ( + StatusCode::OK, + [(header::CONTENT_TYPE, "application/json")], + body, + ) + .into_response() +} + +fn json_rejection_response(rejection: JsonRejection) -> Response { + match rejection { + JsonRejection::JsonSyntaxError(_) => api_error( + StatusCode::BAD_REQUEST, + "INVALID_JSON", + "request body is not valid JSON", + ), + JsonRejection::JsonDataError(error) => api_error( + StatusCode::BAD_REQUEST, + "INVALID_REQUEST", + &error.body_text(), + ), + JsonRejection::MissingJsonContentType(_) => api_error( + StatusCode::UNSUPPORTED_MEDIA_TYPE, + "UNSUPPORTED_MEDIA_TYPE", + "Content-Type must be application/json", + ), + JsonRejection::BytesRejection(error) if error.status() == StatusCode::PAYLOAD_TOO_LARGE => { + api_error( + StatusCode::PAYLOAD_TOO_LARGE, + "REQUEST_TOO_LARGE", + "JSON request body exceeds the configured limit", + ) + } + other => api_error(other.status(), "INVALID_REQUEST", &other.body_text()), + } +} + +fn unauthorized_response() -> Response { + let mut response = api_error( + StatusCode::UNAUTHORIZED, + "UNAUTHORIZED", + "missing or invalid bearer token", + ); + response + .headers_mut() + .insert(header::WWW_AUTHENTICATE, HeaderValue::from_static("Bearer")); + response +} + +#[derive(Debug, Default)] +struct ServerMetrics { + http_requests: AtomicU64, + http_status_2xx: AtomicU64, + http_status_3xx: AtomicU64, + http_status_4xx: AtomicU64, + http_status_5xx: AtomicU64, + http_duration_micros: AtomicU64, + http_duration_buckets: [AtomicU64; HTTP_DURATION_BUCKETS.len()], + http_timeouts: AtomicU64, + query_started: AtomicU64, + query_rejected: AtomicU64, + query_timeouts: AtomicU64, + query_inflight: AtomicU64, + auth_unauthorized: AtomicU64, + auth_forbidden: AtomicU64, + lookup_keys_found: AtomicU64, + lookup_keys_not_found: AtomicU64, + lookup_keys_non_unique: AtomicU64, + planned_files: AtomicU64, + planned_bytes: AtomicU64, + responses_too_large: AtomicU64, + readiness_checks: AtomicU64, + readiness_cache_hits: AtomicU64, +} + +impl ServerMetrics { + fn record_http(&self, status: StatusCode, elapsed: Duration) { + self.http_requests.fetch_add(1, Ordering::Relaxed); + self.http_duration_micros.fetch_add( + u64::try_from(elapsed.as_micros()).unwrap_or(u64::MAX), + Ordering::Relaxed, + ); + let elapsed_micros = u64::try_from(elapsed.as_micros()).unwrap_or(u64::MAX); + for (index, (_, upper_bound)) in HTTP_DURATION_BUCKETS.iter().enumerate() { + if elapsed_micros <= *upper_bound { + self.http_duration_buckets[index].fetch_add(1, Ordering::Relaxed); + } + } + match status.as_u16() / 100 { + 2 => &self.http_status_2xx, + 3 => &self.http_status_3xx, + 4 => &self.http_status_4xx, + 5 => &self.http_status_5xx, + _ => return, + } + .fetch_add(1, Ordering::Relaxed); + } + + fn record_lookup(&self, response: &BatchGetResponse) { + for result in &response.results { + match result.status { + LookupStatus::Found => &self.lookup_keys_found, + LookupStatus::NotFound => &self.lookup_keys_not_found, + LookupStatus::NonUnique => &self.lookup_keys_non_unique, + } + .fetch_add(1, Ordering::Relaxed); + } + if !response.cache_hit { + self.planned_files.fetch_add( + u64::try_from(response.scan.planned_files).unwrap_or(u64::MAX), + Ordering::Relaxed, + ); + self.planned_bytes + .fetch_add(response.scan.planned_bytes, Ordering::Relaxed); + } + } + + fn render( + &self, + table_cache: TableCacheStats, + descriptor_cache: DescriptorCacheStats, + ) -> String { + let requests = self.http_requests.load(Ordering::Relaxed); + let duration_seconds = + self.http_duration_micros.load(Ordering::Relaxed) as f64 / 1_000_000.0; + let duration_buckets = HTTP_DURATION_BUCKETS + .iter() + .enumerate() + .map(|(index, (label, _))| { + format!( + "paimon_query_service_http_request_duration_seconds_bucket{{le=\"{label}\"}} {}\n", + self.http_duration_buckets[index].load(Ordering::Relaxed) + ) + }) + .collect::(); + format!( + concat!( + "# HELP paimon_query_service_http_requests_total HTTP requests handled by this process.\n", + "# TYPE paimon_query_service_http_requests_total counter\n", + "paimon_query_service_http_requests_total {requests}\n", + "# HELP paimon_query_service_http_responses_total HTTP responses by status class.\n", + "# TYPE paimon_query_service_http_responses_total counter\n", + "paimon_query_service_http_responses_total{{status_class=\"2xx\"}} {status_2xx}\n", + "paimon_query_service_http_responses_total{{status_class=\"3xx\"}} {status_3xx}\n", + "paimon_query_service_http_responses_total{{status_class=\"4xx\"}} {status_4xx}\n", + "paimon_query_service_http_responses_total{{status_class=\"5xx\"}} {status_5xx}\n", + "# HELP paimon_query_service_http_request_duration_seconds Total HTTP request duration.\n", + "# TYPE paimon_query_service_http_request_duration_seconds histogram\n", + "{duration_buckets}", + "paimon_query_service_http_request_duration_seconds_bucket{{le=\"+Inf\"}} {requests}\n", + "paimon_query_service_http_request_duration_seconds_sum {duration_seconds:.6}\n", + "paimon_query_service_http_request_duration_seconds_count {requests}\n", + "# HELP paimon_query_service_http_request_timeouts_total Requests terminated by the total HTTP deadline.\n", + "# TYPE paimon_query_service_http_request_timeouts_total counter\n", + "paimon_query_service_http_request_timeouts_total {http_timeouts}\n", + "# HELP paimon_query_service_queries_started_total Admitted descriptor queries.\n", + "# TYPE paimon_query_service_queries_started_total counter\n", + "paimon_query_service_queries_started_total {query_started}\n", + "# HELP paimon_query_service_queries_rejected_total Queries rejected by admission control.\n", + "# TYPE paimon_query_service_queries_rejected_total counter\n", + "paimon_query_service_queries_rejected_total {query_rejected}\n", + "# HELP paimon_query_service_query_timeouts_total Descriptor query timeouts.\n", + "# TYPE paimon_query_service_query_timeouts_total counter\n", + "paimon_query_service_query_timeouts_total {query_timeouts}\n", + "# HELP paimon_query_service_queries_inflight Currently executing descriptor queries.\n", + "# TYPE paimon_query_service_queries_inflight gauge\n", + "paimon_query_service_queries_inflight {query_inflight}\n", + "# HELP paimon_query_service_auth_failures_total Authentication and authorization failures.\n", + "# TYPE paimon_query_service_auth_failures_total counter\n", + "paimon_query_service_auth_failures_total{{reason=\"unauthorized\"}} {auth_unauthorized}\n", + "paimon_query_service_auth_failures_total{{reason=\"forbidden\"}} {auth_forbidden}\n", + "# HELP paimon_query_service_lookup_keys_total Requested keys by lookup result.\n", + "# TYPE paimon_query_service_lookup_keys_total counter\n", + "paimon_query_service_lookup_keys_total{{status=\"found\"}} {keys_found}\n", + "paimon_query_service_lookup_keys_total{{status=\"not_found\"}} {keys_not_found}\n", + "paimon_query_service_lookup_keys_total{{status=\"non_unique\"}} {keys_non_unique}\n", + "# HELP paimon_query_service_planned_files_total Data files selected by lookup plans.\n", + "# TYPE paimon_query_service_planned_files_total counter\n", + "paimon_query_service_planned_files_total {planned_files}\n", + "# HELP paimon_query_service_planned_bytes_total Known data-file bytes selected by lookup plans.\n", + "# TYPE paimon_query_service_planned_bytes_total counter\n", + "paimon_query_service_planned_bytes_total {planned_bytes}\n", + "# HELP paimon_query_service_responses_rejected_total Successful query results rejected by the response resource envelope.\n", + "# TYPE paimon_query_service_responses_rejected_total counter\n", + "paimon_query_service_responses_rejected_total{{reason=\"too_large\"}} {responses_too_large}\n", + "# HELP paimon_query_service_readiness_checks_total Catalog readiness traversals executed.\n", + "# TYPE paimon_query_service_readiness_checks_total counter\n", + "paimon_query_service_readiness_checks_total {readiness_checks}\n", + "# HELP paimon_query_service_readiness_cache_hits_total Readiness probes served from the short-lived cache.\n", + "# TYPE paimon_query_service_readiness_cache_hits_total counter\n", + "paimon_query_service_readiness_cache_hits_total {readiness_cache_hits}\n", + "# HELP paimon_query_service_log_events_dropped_total Log events dropped because the non-blocking queue was full.\n", + "# TYPE paimon_query_service_log_events_dropped_total counter\n", + "paimon_query_service_log_events_dropped_total {log_events_dropped}\n", + "# HELP paimon_query_service_table_cache_hits_total Table metadata cache hits.\n", + "# TYPE paimon_query_service_table_cache_hits_total counter\n", + "paimon_query_service_table_cache_hits_total {cache_hits}\n", + "# HELP paimon_query_service_table_cache_misses_total Table metadata cache misses.\n", + "# TYPE paimon_query_service_table_cache_misses_total counter\n", + "paimon_query_service_table_cache_misses_total {cache_misses}\n", + "# HELP paimon_query_service_table_cache_entries Table metadata cache entries.\n", + "# TYPE paimon_query_service_table_cache_entries gauge\n", + "paimon_query_service_table_cache_entries {table_cache_entries}\n", + "# HELP paimon_query_service_descriptor_cache_hits_total Snapshot-pinned descriptor cache hits.\n", + "# TYPE paimon_query_service_descriptor_cache_hits_total counter\n", + "paimon_query_service_descriptor_cache_hits_total {descriptor_cache_hits}\n", + "# HELP paimon_query_service_descriptor_cache_misses_total Snapshot-pinned descriptor cache misses.\n", + "# TYPE paimon_query_service_descriptor_cache_misses_total counter\n", + "paimon_query_service_descriptor_cache_misses_total {descriptor_cache_misses}\n", + "# HELP paimon_query_service_descriptor_cache_entries Snapshot-pinned descriptor cache entries.\n", + "# TYPE paimon_query_service_descriptor_cache_entries gauge\n", + "paimon_query_service_descriptor_cache_entries {descriptor_cache_entries}\n", + "# HELP paimon_query_service_descriptor_cache_bytes Approximate weighted descriptor cache bytes.\n", + "# TYPE paimon_query_service_descriptor_cache_bytes gauge\n", + "paimon_query_service_descriptor_cache_bytes {descriptor_cache_bytes}\n", + ), + requests = requests, + duration_seconds = duration_seconds, + duration_buckets = duration_buckets, + status_2xx = self.http_status_2xx.load(Ordering::Relaxed), + status_3xx = self.http_status_3xx.load(Ordering::Relaxed), + status_4xx = self.http_status_4xx.load(Ordering::Relaxed), + status_5xx = self.http_status_5xx.load(Ordering::Relaxed), + http_timeouts = self.http_timeouts.load(Ordering::Relaxed), + query_started = self.query_started.load(Ordering::Relaxed), + query_rejected = self.query_rejected.load(Ordering::Relaxed), + query_timeouts = self.query_timeouts.load(Ordering::Relaxed), + query_inflight = self.query_inflight.load(Ordering::Relaxed), + auth_unauthorized = self.auth_unauthorized.load(Ordering::Relaxed), + auth_forbidden = self.auth_forbidden.load(Ordering::Relaxed), + keys_found = self.lookup_keys_found.load(Ordering::Relaxed), + keys_not_found = self.lookup_keys_not_found.load(Ordering::Relaxed), + keys_non_unique = self.lookup_keys_non_unique.load(Ordering::Relaxed), + planned_files = self.planned_files.load(Ordering::Relaxed), + planned_bytes = self.planned_bytes.load(Ordering::Relaxed), + responses_too_large = self.responses_too_large.load(Ordering::Relaxed), + readiness_checks = self.readiness_checks.load(Ordering::Relaxed), + readiness_cache_hits = self.readiness_cache_hits.load(Ordering::Relaxed), + log_events_dropped = LOG_ERROR_COUNTER + .get() + .map(ErrorCounter::dropped_lines) + .unwrap_or_default(), + cache_hits = table_cache.hits, + cache_misses = table_cache.misses, + table_cache_entries = table_cache.entries, + descriptor_cache_hits = descriptor_cache.hits, + descriptor_cache_misses = descriptor_cache.misses, + descriptor_cache_entries = descriptor_cache.entries, + descriptor_cache_bytes = descriptor_cache.weighted_bytes, + ) + } +} + +struct InflightQueryGuard { + metrics: Arc, +} + +impl InflightQueryGuard { + fn new(metrics: Arc) -> Self { + metrics.query_inflight.fetch_add(1, Ordering::Relaxed); + Self { metrics } + } +} + +impl Drop for InflightQueryGuard { + fn drop(&mut self) { + self.metrics.query_inflight.fetch_sub(1, Ordering::Relaxed); + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum AdmissionError { + TimedOut, + Closed, +} + +async fn acquire_query_permit( + permits: Arc, + queue_timeout: Duration, +) -> Result { + match tokio::time::timeout(queue_timeout, permits.acquire_owned()).await { + Ok(Ok(permit)) => Ok(permit), + Ok(Err(_)) => Err(AdmissionError::Closed), + Err(_) => Err(AdmissionError::TimedOut), + } +} + +async fn enforce_http_timeout( + State(state): State, + request: Request, + next: Next, +) -> Response { + match tokio::time::timeout(state.timeout, next.run(request)).await { + Ok(response) => response, + Err(_) => { + state.metrics.http_timeouts.fetch_add(1, Ordering::Relaxed); + api_error( + StatusCode::REQUEST_TIMEOUT, + "REQUEST_TIMEOUT", + "HTTP request exceeded its total deadline", + ) + } + } +} + +async fn request_context( + State(metrics): State>, + mut request: Request, + next: Next, +) -> Response { + let request_id = request + .headers() + .get(REQUEST_ID_HEADER) + .and_then(|value| value.to_str().ok()) + .filter(|value| valid_request_id(value)) + .map(str::to_string) + .unwrap_or_else(new_request_id); + let method = request.method().to_string(); + let path = request.uri().path().to_string(); + let access_context = AccessContext::default(); + request.extensions_mut().insert(access_context.clone()); + let started = Instant::now(); + let span = tracing::info_span!( + "http_request", + request_id = %request_id, + method = %method, + path = %path, + ); + let mut response = next.run(request).instrument(span.clone()).await; + let elapsed = started.elapsed(); + let status = response.status().as_u16(); + let principal = access_context.principal(); + response.headers_mut().insert( + REQUEST_ID_HEADER, + HeaderValue::from_str(&request_id).expect("generated request ID must be a header value"), + ); + metrics.record_http(response.status(), elapsed); + tracing::info!( + parent: &span, + event = "http_request_complete", + principal = principal.as_deref().unwrap_or(""), + status, + elapsed_ms = u64::try_from(elapsed.as_millis()).unwrap_or(u64::MAX), + "HTTP request completed" + ); + response +} + +#[derive(Clone, Default)] +struct AccessContext(Arc>>); + +impl AccessContext { + fn set_principal(&self, principal: &str) { + if let Ok(mut current) = self.0.lock() { + *current = Some(principal.to_string()); + } + } + + fn principal(&self) -> Option { + self.0.lock().ok().and_then(|current| current.clone()) + } +} + +fn valid_request_id(value: &str) -> bool { + !value.is_empty() + && value.len() <= 128 + && value + .bytes() + .all(|byte| byte.is_ascii_alphanumeric() || b"-_.:".contains(&byte)) +} + +fn new_request_id() -> String { + let epoch_millis = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_millis(); + let sequence = REQUEST_SEQUENCE.fetch_add(1, Ordering::Relaxed); + format!("{epoch_millis:x}-{:x}-{sequence:x}", std::process::id()) +} + +fn lookup_error_response(error: LookupError) -> Response { + let mut root = &error; + while let LookupError::Shared(inner) = root { + root = inner; + } + let (status, code, message) = match root { + LookupError::InvalidRequest(_) | LookupError::InvalidKeyValue { .. } => ( + StatusCode::BAD_REQUEST, + "INVALID_REQUEST", + error.to_string(), + ), + LookupError::UnsupportedKeyType { .. } => ( + StatusCode::UNPROCESSABLE_ENTITY, + "UNSUPPORTED_KEY_TYPE", + error.to_string(), + ), + LookupError::QueryBudgetExceeded { .. } => ( + StatusCode::TOO_MANY_REQUESTS, + "QUERY_BUDGET_EXCEEDED", + error.to_string(), + ), + LookupError::Paimon(paimon::Error::TableNotExist { .. }) => ( + StatusCode::NOT_FOUND, + "TABLE_NOT_FOUND", + "requested table does not exist".to_string(), + ), + LookupError::Paimon(paimon::Error::DatabaseNotExist { .. }) => ( + StatusCode::NOT_FOUND, + "DATABASE_NOT_FOUND", + "requested database does not exist".to_string(), + ), + LookupError::Paimon(paimon::Error::SnapshotNotExist { .. }) => ( + StatusCode::NOT_FOUND, + "SNAPSHOT_NOT_FOUND", + "requested snapshot does not exist".to_string(), + ), + LookupError::Paimon(_) | LookupError::PaimonUnavailable => ( + StatusCode::SERVICE_UNAVAILABLE, + "PAIMON_ERROR", + "table metadata or data is temporarily unavailable".to_string(), + ), + LookupError::LoadCancelled => ( + StatusCode::SERVICE_UNAVAILABLE, + "LOOKUP_CANCELLED", + "the shared lookup attempt was cancelled; retry the request".to_string(), + ), + LookupError::InvalidPolicy(_) => ( + StatusCode::INTERNAL_SERVER_ERROR, + "INVALID_POLICY", + "the configured lookup policy is incompatible with the table".to_string(), + ), + LookupError::SnapshotMismatch { .. } + | LookupError::InvalidDescriptor { .. } + | LookupError::UnexpectedResult(_) => ( + StatusCode::INTERNAL_SERVER_ERROR, + "LOOKUP_INVARIANT_VIOLATION", + "the lookup result failed an internal consistency check".to_string(), + ), + LookupError::Shared(_) => unreachable!("shared lookup errors are unwrapped above"), + }; + if status.is_server_error() { + tracing::error!(%error, status = status.as_u16(), code, "blob descriptor lookup failed"); + } + api_error(status, code, &message) +} + +#[derive(Serialize)] +#[serde(rename_all = "camelCase")] +struct ApiError<'a> { + code: &'a str, + message: &'a str, +} + +fn api_error(status: StatusCode, code: &str, message: &str) -> Response { + (status, Json(ApiError { code, message })).into_response() +} + +async fn shutdown_signal() { + let ctrl_c = async { + tokio::signal::ctrl_c() + .await + .expect("failed to install Ctrl-C handler"); + }; + + #[cfg(unix)] + let terminate = async { + tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) + .expect("failed to install SIGTERM handler") + .recv() + .await; + }; + + #[cfg(not(unix))] + let terminate = std::future::pending::<()>(); + + tokio::select! { + _ = ctrl_c => {}, + _ = terminate => {}, + } +} + +#[cfg(test)] +mod tests { + use std::collections::BTreeMap; + + use axum::body::{to_bytes, Body}; + use axum::http::Request; + use tower::ServiceExt; + + use super::*; + + #[test] + fn parses_minimal_config_with_defaults() { + let config: ServerConfig = serde_json::from_value(serde_json::json!({ + "catalog": {"warehouse": "/tmp/warehouse"}, + "policies": [{ + "table": {"database": "db", "table": "assets"}, + "keyFields": ["id"], + "blobFields": ["picture"], + "strategy": "GLOBAL_BTREE" + }] + })) + .unwrap(); + assert_eq!(config.listen, "127.0.0.1:8081"); + assert_eq!(config.policies[0].budget, Default::default()); + assert_eq!(config.query_timeout_ms, 5_000); + assert_eq!(config.http_request_timeout_ms, 10_000); + assert_eq!(config.max_concurrent_queries, 64); + assert_eq!(config.max_concurrent_index_reads, 64); + assert_eq!(config.max_connections, 1_024); + assert_eq!(config.http_header_read_timeout_ms, 5_000); + assert_eq!(config.http_connection_idle_timeout_ms, 60_000); + assert_eq!(config.http_connection_max_age_ms, 300_000); + assert_eq!(config.http2_max_concurrent_streams, 64); + assert_eq!(config.graceful_shutdown_timeout_ms, 10_000); + assert_eq!(config.queue_timeout_ms, 100); + assert_eq!(config.max_request_body_bytes, 1024 * 1024); + assert_eq!(config.max_response_body_bytes, 4 * 1024 * 1024); + assert_eq!(config.readiness_cache_ttl_ms, 1_000); + assert_eq!(config.table_metadata_cache_ttl_ms, 30_000); + assert_eq!(config.descriptor_cache_ttl_ms, 60_000); + assert_eq!(config.descriptor_cache_max_bytes, 64 * 1024 * 1024); + assert!(config.principals.is_empty()); + assert!(!config.allow_anonymous); + } + + #[test] + fn rejects_unknown_security_and_policy_config_fields() { + let root_typo = serde_json::from_value::(serde_json::json!({ + "catalog": {}, + "policies": [], + "principlas": [] + })) + .unwrap_err(); + assert!(root_typo.to_string().contains("principlas")); + + let grant_typo = serde_json::from_value::(serde_json::json!({ + "catalog": {}, + "policies": [{ + "table": {"database": "db", "table": "assets"}, + "keyFields": ["id"], + "blobFields": ["picture"], + "strategy": "GLOBAL_BTREE" + }], + "principals": [{ + "name": "reader", + "bearerToken": "long-test-token-1234", + "grants": [{ + "table": {"database": "db", "table": "assets"}, + "blobFeilds": ["picture"] + }] + }] + })) + .unwrap_err(); + assert!(grant_typo.to_string().contains("blobFeilds")); + } + + #[tokio::test] + async fn rejects_implicit_anonymous_access_before_catalog_creation() { + let config: ServerConfig = serde_json::from_value(serde_json::json!({ + "catalog": {}, + "policies": [] + })) + .unwrap(); + let error = match build_app(&config).await { + Ok(_) => panic!("anonymous access must require an explicit opt-in"), + Err(error) => error, + }; + assert!(error.to_string().contains("allowAnonymous")); + } + + #[tokio::test] + async fn rejects_zero_resource_limits_before_catalog_creation() { + let config: ServerConfig = serde_json::from_value(serde_json::json!({ + "catalog": {}, + "policies": [], + "maxConcurrentQueries": 0 + })) + .unwrap(); + let error = match build_app(&config).await { + Ok(_) => panic!("zero concurrency must be rejected"), + Err(error) => error, + }; + assert!(error.to_string().contains("maxConcurrentQueries")); + + let config: ServerConfig = serde_json::from_value(serde_json::json!({ + "catalog": {}, + "policies": [], + "maxConcurrentQueries": 2, + "maxConcurrentIndexReads": 1 + })) + .unwrap(); + let error = match build_app(&config).await { + Ok(_) => panic!("an incoherent index-read envelope must be rejected"), + Err(error) => error, + }; + assert!(error.to_string().contains("maxConcurrentIndexReads")); + } + + #[tokio::test] + async fn rejects_total_deadline_that_cannot_contain_queue_and_query() { + let config: ServerConfig = serde_json::from_value(serde_json::json!({ + "catalog": {}, + "policies": [], + "queryTimeoutMs": 5000, + "queueTimeoutMs": 100, + "httpRequestTimeoutMs": 5100 + })) + .unwrap(); + let error = match build_app(&config).await { + Ok(_) => panic!("incoherent HTTP deadline must be rejected"), + Err(error) => error, + }; + assert!(error.to_string().contains("httpRequestTimeoutMs")); + } + + #[tokio::test] + async fn rejects_connection_max_age_shorter_than_idle_timeout() { + let config: ServerConfig = serde_json::from_value(serde_json::json!({ + "catalog": {}, + "policies": [], + "httpConnectionIdleTimeoutMs": 60_000, + "httpConnectionMaxAgeMs": 60_000 + })) + .unwrap(); + let error = match build_app(&config).await { + Ok(_) => panic!("connection max age must be an absolute outer bound"), + Err(error) => error, + }; + assert!(error.to_string().contains("httpConnectionMaxAgeMs")); + } + + #[tokio::test] + async fn http_route_enforces_auth_and_maps_missing_table() { + let warehouse = tempfile::TempDir::new().unwrap(); + let config: ServerConfig = serde_json::from_value(serde_json::json!({ + "catalog": { + "warehouse": warehouse.path().to_str().unwrap() + }, + "bearerToken": "legacy-secret-token-1234", + "policies": [{ + "table": {"database": "db", "table": "assets"}, + "keyFields": ["id"], + "blobFields": ["picture"], + "strategy": "GLOBAL_BTREE" + }] + })) + .unwrap(); + let app = build_app(&config).await.unwrap(); + let uri = "/api/blob/v1/databases/db/tables/assets/descriptors:batchGet"; + let body = r#"{"keys":[{"id":1}],"blobFields":["picture"]}"#; + + let unauthorized = app + .clone() + .oneshot( + Request::post(uri) + .header(header::CONTENT_TYPE, "application/json") + .header(REQUEST_ID_HEADER, "caller-123") + .body(Body::from("{not-json")) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(unauthorized.status(), StatusCode::UNAUTHORIZED); + assert_eq!(unauthorized.headers()[REQUEST_ID_HEADER], "caller-123"); + assert_eq!(unauthorized.headers()[header::WWW_AUTHENTICATE], "Bearer"); + + let missing = app + .clone() + .oneshot( + Request::post(uri) + .header(header::CONTENT_TYPE, "application/json") + .header(header::AUTHORIZATION, "Bearer legacy-secret-token-1234") + .body(Body::from(body)) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(missing.status(), StatusCode::NOT_FOUND); + assert!(valid_request_id( + missing.headers()[REQUEST_ID_HEADER].to_str().unwrap() + )); + + let not_ready = app + .oneshot(Request::get("/readyz").body(Body::empty()).unwrap()) + .await + .unwrap(); + assert_eq!(not_ready.status(), StatusCode::SERVICE_UNAVAILABLE); + } + + #[tokio::test] + async fn principal_grants_enforce_table_and_blob_field_access() { + let warehouse = tempfile::TempDir::new().unwrap(); + let config: ServerConfig = serde_json::from_value(serde_json::json!({ + "catalog": { + "warehouse": warehouse.path().to_str().unwrap() + }, + "policies": [{ + "table": {"database": "db", "table": "assets"}, + "keyFields": ["id"], + "blobFields": ["picture", "thumbnail"], + "strategy": "GLOBAL_BTREE" + }], + "principals": [{ + "name": "image-reader", + "bearerToken": "image-reader-token-1234", + "grants": [{ + "table": {"database": "db", "table": "assets"}, + "blobFields": ["picture"] + }] + }] + })) + .unwrap(); + let app = build_app(&config).await.unwrap(); + let uri = "/api/blob/v1/databases/db/tables/assets/descriptors:batchGet"; + + let forbidden = app + .clone() + .oneshot( + Request::post(uri) + .header(header::CONTENT_TYPE, "application/json") + .header(header::AUTHORIZATION, "Bearer image-reader-token-1234") + .body(Body::from( + r#"{"keys":[{"id":1}],"blobFields":["thumbnail"]}"#, + )) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(forbidden.status(), StatusCode::FORBIDDEN); + + let allowed = app + .clone() + .oneshot( + Request::post(uri) + .header(header::CONTENT_TYPE, "application/json") + .header(header::AUTHORIZATION, "Bearer image-reader-token-1234") + .body(Body::from( + r#"{"keys":[{"id":1}],"blobFields":["picture"]}"#, + )) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(allowed.status(), StatusCode::NOT_FOUND); + + let metrics = app + .oneshot(Request::get("/metrics").body(Body::empty()).unwrap()) + .await + .unwrap(); + let body = to_bytes(metrics.into_body(), 64 * 1024).await.unwrap(); + let body = std::str::from_utf8(&body).unwrap(); + assert!(body.contains("paimon_query_service_auth_failures_total{reason=\"forbidden\"} 1")); + } + + #[tokio::test] + async fn rejects_oversized_json_body_and_adds_request_id() { + let warehouse = tempfile::TempDir::new().unwrap(); + let config: ServerConfig = serde_json::from_value(serde_json::json!({ + "catalog": { + "warehouse": warehouse.path().to_str().unwrap() + }, + "allowAnonymous": true, + "policies": [], + "maxRequestBodyBytes": 16 + })) + .unwrap(); + let response = build_app(&config) + .await + .unwrap() + .oneshot( + Request::post("/api/blob/v1/databases/db/tables/assets/descriptors:batchGet") + .header(header::CONTENT_TYPE, "application/json") + .body(Body::from( + r#"{"keys":[{"id":1}],"blobFields":["picture"]}"#, + )) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::PAYLOAD_TOO_LARGE); + assert!(valid_request_id( + response.headers()[REQUEST_ID_HEADER].to_str().unwrap() + )); + assert_api_error(response, StatusCode::PAYLOAD_TOO_LARGE, "REQUEST_TOO_LARGE").await; + } + + #[tokio::test] + async fn returns_json_errors_for_invalid_json_content_type_route_and_method() { + let warehouse = tempfile::TempDir::new().unwrap(); + let config: ServerConfig = serde_json::from_value(serde_json::json!({ + "catalog": { + "warehouse": warehouse.path().to_str().unwrap() + }, + "allowAnonymous": true, + "policies": [] + })) + .unwrap(); + let app = build_app(&config).await.unwrap(); + let uri = "/api/blob/v1/databases/db/tables/assets/descriptors:batchGet"; + + let invalid_json = app + .clone() + .oneshot( + Request::post(uri) + .header(header::CONTENT_TYPE, "application/json") + .body(Body::from("{not-json")) + .unwrap(), + ) + .await + .unwrap(); + assert_api_error(invalid_json, StatusCode::BAD_REQUEST, "INVALID_JSON").await; + + let missing_content_type = app + .clone() + .oneshot( + Request::post(uri) + .body(Body::from(r#"{"keys":[],"blobFields":[]}"#)) + .unwrap(), + ) + .await + .unwrap(); + assert_api_error( + missing_content_type, + StatusCode::UNSUPPORTED_MEDIA_TYPE, + "UNSUPPORTED_MEDIA_TYPE", + ) + .await; + + let unknown_request_field = app + .clone() + .oneshot( + Request::post(uri) + .header(header::CONTENT_TYPE, "application/json") + .body(Body::from( + r#"{"keys":[{"id":1}],"blobFields":["picture"],"snapshotID":1}"#, + )) + .unwrap(), + ) + .await + .unwrap(); + assert_api_error( + unknown_request_field, + StatusCode::BAD_REQUEST, + "INVALID_REQUEST", + ) + .await; + + let unknown_route = app + .clone() + .oneshot(Request::get("/unknown").body(Body::empty()).unwrap()) + .await + .unwrap(); + assert_api_error(unknown_route, StatusCode::NOT_FOUND, "ROUTE_NOT_FOUND").await; + + let wrong_method = app + .oneshot(Request::post("/healthz").body(Body::empty()).unwrap()) + .await + .unwrap(); + assert_api_error( + wrong_method, + StatusCode::METHOD_NOT_ALLOWED, + "METHOD_NOT_ALLOWED", + ) + .await; + } + + #[tokio::test] + async fn admission_times_out_when_all_query_slots_are_held() { + let permits = Arc::new(Semaphore::new(1)); + let _held = permits.clone().acquire_owned().await.unwrap(); + let error = match acquire_query_permit(permits, Duration::from_millis(1)).await { + Ok(_) => panic!("admission should time out"), + Err(error) => error, + }; + assert_eq!(error, AdmissionError::TimedOut); + } + + #[tokio::test] + async fn protocol_detection_times_out_before_hyper_protocol_sniffing() { + let (_client, server) = tokio::io::duplex(64); + let error = match read_protocol_prefix(server, Duration::from_millis(1)).await { + Ok(_) => panic!("an idle connection must not occupy a permit indefinitely"), + Err(error) => error, + }; + assert_eq!(error.kind(), std::io::ErrorKind::TimedOut); + } + + #[tokio::test] + async fn connection_max_age_gracefully_drains_active_work() { + let (_shutdown_sender, shutdown_receiver) = watch::channel(false); + let (graceful_sender, graceful_receiver) = tokio::sync::oneshot::channel(); + let connection = async move { + graceful_receiver.await.unwrap(); + tokio::time::sleep(Duration::from_millis(5)).await; + 42 + }; + let result = drive_connection( + connection, + Duration::from_millis(1), + Duration::from_millis(100), + Duration::from_millis(100), + shutdown_receiver, + move |_| graceful_sender.send(()).unwrap(), + ) + .await; + match result { + ConnectionEnd::Drained { reason, output } => { + assert_eq!(reason, ConnectionShutdownReason::MaxAge); + assert_eq!(output, 42); + } + _ => panic!("max age must signal graceful shutdown and await active work"), + } + } + + #[tokio::test] + async fn connection_max_age_still_has_a_hard_drain_bound() { + let (_shutdown_sender, shutdown_receiver) = watch::channel(false); + let result = drive_connection( + std::future::pending::<()>(), + Duration::from_millis(1), + Duration::from_millis(1), + Duration::from_millis(100), + shutdown_receiver, + |_| {}, + ) + .await; + assert!(matches!( + result, + ConnectionEnd::DrainTimedOut(ConnectionShutdownReason::MaxAge) + )); + } + + #[tokio::test] + async fn admission_limit_is_acquired_before_json_body_parsing() { + let warehouse = tempfile::TempDir::new().unwrap(); + let config: ServerConfig = serde_json::from_value(serde_json::json!({ + "catalog": {"warehouse": warehouse.path().to_str().unwrap()}, + "policies": [], + "allowAnonymous": true, + "maxConcurrentQueries": 1, + "maxConcurrentIndexReads": 1, + "queueTimeoutMs": 1, + "queryTimeoutMs": 100, + "httpRequestTimeoutMs": 1000 + })) + .unwrap(); + let app = build_app(&config).await.unwrap(); + let uri = "/api/blob/v1/databases/db/tables/assets/descriptors:batchGet"; + let pending_body = Body::from_stream(futures::stream::pending::< + Result, std::convert::Infallible>, + >()); + let first = tokio::spawn( + app.clone().oneshot( + Request::post(uri) + .header(header::CONTENT_TYPE, "application/json") + .body(pending_body) + .unwrap(), + ), + ); + tokio::time::sleep(Duration::from_millis(10)).await; + + let rejected = app + .oneshot( + Request::post(uri) + .header(header::CONTENT_TYPE, "application/json") + .body(Body::from(r#"{"keys":[],"blobFields":[]}"#)) + .unwrap(), + ) + .await + .unwrap(); + assert_api_error(rejected, StatusCode::TOO_MANY_REQUESTS, "SERVER_BUSY").await; + first.abort(); + } + + #[tokio::test] + async fn total_http_deadline_returns_json_error_and_request_id() { + let metrics = Arc::new(ServerMetrics::default()); + let timeout = HttpTimeoutState { + timeout: Duration::from_millis(1), + metrics: metrics.clone(), + }; + let app = Router::new() + .route( + "/slow", + get(|| async { + tokio::time::sleep(Duration::from_millis(50)).await; + StatusCode::NO_CONTENT + }), + ) + .layer(middleware::from_fn_with_state( + timeout, + enforce_http_timeout, + )) + .layer(middleware::from_fn_with_state( + metrics.clone(), + request_context, + )); + + let response = app + .oneshot(Request::get("/slow").body(Body::empty()).unwrap()) + .await + .unwrap(); + assert_api_error(response, StatusCode::REQUEST_TIMEOUT, "REQUEST_TIMEOUT").await; + assert_eq!(metrics.http_timeouts.load(Ordering::Relaxed), 1); + assert_eq!(metrics.http_status_4xx.load(Ordering::Relaxed), 1); + } + + #[tokio::test] + async fn rejects_success_response_over_configured_limit() { + let metrics = ServerMetrics::default(); + let result = BatchGetResponse { + table: TableRef::new("db", "assets"), + snapshot_id: Some(1), + schema_id: 1, + cache_hit: false, + scan: Default::default(), + results: vec![], + }; + + let accepted = success_response(result.clone(), 4096, &metrics); + assert_eq!(accepted.status(), StatusCode::OK); + assert_eq!(accepted.headers()[header::CONTENT_TYPE], "application/json"); + let body = to_bytes(accepted.into_body(), 64 * 1024).await.unwrap(); + let decoded: BatchGetResponse = serde_json::from_slice(&body).unwrap(); + assert_eq!(decoded, result); + + let response = success_response(result, 1, &metrics); + + assert_eq!(response.status(), StatusCode::INSUFFICIENT_STORAGE); + let body = to_bytes(response.into_body(), 64 * 1024).await.unwrap(); + let body: serde_json::Value = serde_json::from_slice(&body).unwrap(); + assert_eq!(body["code"], "RESPONSE_TOO_LARGE"); + assert_eq!(metrics.responses_too_large.load(Ordering::Relaxed), 1); + } + + #[tokio::test] + async fn hides_internal_paimon_errors_and_maps_missing_snapshots() { + let internal = lookup_error_response(LookupError::Paimon(paimon::Error::DataInvalid { + message: "secret path /warehouse/db/table".to_string(), + source: None, + })); + assert_eq!(internal.status(), StatusCode::SERVICE_UNAVAILABLE); + let body = to_bytes(internal.into_body(), 64 * 1024).await.unwrap(); + let body: serde_json::Value = serde_json::from_slice(&body).unwrap(); + assert_eq!(body["code"], "PAIMON_ERROR"); + assert!(!body["message"].as_str().unwrap().contains("/warehouse")); + + let missing = lookup_error_response(LookupError::Paimon(paimon::Error::SnapshotNotExist { + snapshot_id: 17, + })); + assert_eq!(missing.status(), StatusCode::NOT_FOUND); + let body = to_bytes(missing.into_body(), 64 * 1024).await.unwrap(); + let body: serde_json::Value = serde_json::from_slice(&body).unwrap(); + assert_eq!(body["code"], "SNAPSHOT_NOT_FOUND"); + + let cancelled = lookup_error_response(LookupError::LoadCancelled); + assert_eq!(cancelled.status(), StatusCode::SERVICE_UNAVAILABLE); + let body = to_bytes(cancelled.into_body(), 64 * 1024).await.unwrap(); + let body: serde_json::Value = serde_json::from_slice(&body).unwrap(); + assert_eq!(body["code"], "LOOKUP_CANCELLED"); + } + + #[test] + fn validates_caller_request_ids() { + assert!(valid_request_id("trace-123:attempt_2")); + assert!(!valid_request_id("")); + assert!(!valid_request_id("has spaces")); + assert!(!valid_request_id(&"x".repeat(129))); + } + + #[tokio::test] + async fn exposes_liveness_readiness_and_prometheus_metrics() { + let warehouse = tempfile::TempDir::new().unwrap(); + let config: ServerConfig = serde_json::from_value(serde_json::json!({ + "catalog": { + "warehouse": warehouse.path().to_str().unwrap() + }, + "allowAnonymous": true, + "policies": [] + })) + .unwrap(); + let app = build_app(&config).await.unwrap(); + + let health = app + .clone() + .oneshot(Request::get("/healthz").body(Body::empty()).unwrap()) + .await + .unwrap(); + assert_eq!(health.status(), StatusCode::NO_CONTENT); + + let ready = app + .clone() + .oneshot(Request::get("/readyz").body(Body::empty()).unwrap()) + .await + .unwrap(); + assert_eq!(ready.status(), StatusCode::NO_CONTENT); + + let cached_ready = app + .clone() + .oneshot(Request::get("/readyz").body(Body::empty()).unwrap()) + .await + .unwrap(); + assert_eq!(cached_ready.status(), StatusCode::NO_CONTENT); + + let metrics = app + .oneshot(Request::get("/metrics").body(Body::empty()).unwrap()) + .await + .unwrap(); + assert_eq!(metrics.status(), StatusCode::OK); + assert_eq!( + metrics.headers()[header::CONTENT_TYPE], + "text/plain; version=0.0.4; charset=utf-8" + ); + let body = to_bytes(metrics.into_body(), 64 * 1024).await.unwrap(); + let body = std::str::from_utf8(&body).unwrap(); + assert!(body.contains("paimon_query_service_http_requests_total 3")); + assert!(body.contains("paimon_query_service_http_responses_total{status_class=\"2xx\"} 3")); + assert!(body.contains("paimon_query_service_readiness_checks_total 1")); + assert!(body.contains("paimon_query_service_readiness_cache_hits_total 1")); + assert!(body.contains("paimon_query_service_table_cache_entries 0")); + } + + #[test] + fn inflight_metric_is_released_by_guard_drop() { + let metrics = Arc::new(ServerMetrics::default()); + { + let _guard = InflightQueryGuard::new(metrics.clone()); + assert_eq!(metrics.query_inflight.load(Ordering::Relaxed), 1); + } + assert_eq!(metrics.query_inflight.load(Ordering::Relaxed), 0); + } + + #[test] + fn records_lookup_outcomes_and_planned_work() { + use paimon_query_service::{LookupResult, LookupScanStats}; + + let metrics = ServerMetrics::default(); + let response = BatchGetResponse { + table: TableRef::new("db", "assets"), + snapshot_id: Some(1), + schema_id: 1, + cache_hit: false, + scan: LookupScanStats { + planned_files: 3, + planned_bytes: 1024, + }, + results: vec![ + LookupResult { + key: BTreeMap::new(), + status: LookupStatus::Found, + blobs: BTreeMap::new(), + }, + LookupResult { + key: BTreeMap::new(), + status: LookupStatus::NotFound, + blobs: BTreeMap::new(), + }, + LookupResult { + key: BTreeMap::new(), + status: LookupStatus::NonUnique, + blobs: BTreeMap::new(), + }, + ], + }; + metrics.record_lookup(&response); + metrics.record_lookup(&BatchGetResponse { + cache_hit: true, + ..response + }); + + let rendered = metrics.render( + TableCacheStats::default(), + DescriptorCacheStats { + hits: 7, + misses: 2, + entries: 3, + weighted_bytes: 4096, + }, + ); + assert!(rendered.contains("lookup_keys_total{status=\"found\"} 2")); + assert!(rendered.contains("lookup_keys_total{status=\"not_found\"} 2")); + assert!(rendered.contains("lookup_keys_total{status=\"non_unique\"} 2")); + assert!(rendered.contains("planned_files_total 3")); + assert!(rendered.contains("planned_bytes_total 1024")); + assert!(rendered.contains("descriptor_cache_hits_total 7")); + assert!(rendered.contains("descriptor_cache_misses_total 2")); + assert!(rendered.contains("descriptor_cache_entries 3")); + assert!(rendered.contains("descriptor_cache_bytes 4096")); + } + + async fn assert_api_error(response: Response, status: StatusCode, code: &str) { + assert_eq!(response.status(), status); + assert_eq!(response.headers()[header::CONTENT_TYPE], "application/json"); + assert!(valid_request_id( + response.headers()[REQUEST_ID_HEADER].to_str().unwrap() + )); + let body = to_bytes(response.into_body(), 64 * 1024).await.unwrap(); + let body: serde_json::Value = serde_json::from_slice(&body).unwrap(); + assert_eq!(body["code"], code); + assert!(body["message"] + .as_str() + .is_some_and(|message| !message.is_empty())); + } +} diff --git a/crates/query-service-server/src/main.rs b/crates/query-service-server/src/main.rs new file mode 100644 index 00000000..ca06454b --- /dev/null +++ b/crates/query-service-server/src/main.rs @@ -0,0 +1,27 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use paimon_query_service_server::{init_logging, load_config, serve, BoxError}; + +#[tokio::main] +async fn main() -> Result<(), BoxError> { + let _logging_guard = init_logging()?; + let path = std::env::var("QUERY_SERVICE_CONFIG") + .or_else(|_| std::env::var("BLOB_QUERY_CONFIG")) + .unwrap_or_else(|_| "query-service.json".to_string()); + let config = load_config(&path)?; + serve(config).await +} diff --git a/crates/query-service/Cargo.toml b/crates/query-service/Cargo.toml new file mode 100644 index 00000000..e39b3a7b --- /dev/null +++ b/crates/query-service/Cargo.toml @@ -0,0 +1,42 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +[package] +name = "paimon-query-service" +version.workspace = true +edition.workspace = true +homepage.workspace = true +repository.workspace = true +license.workspace = true +rust-version.workspace = true +description = "Snapshot-consistent point-query capabilities for Paimon tables" +publish = false + +[dependencies] +paimon = { workspace = true } +arrow-array = { workspace = true } +base64 = "0.22" +chrono = "0.4.38" +futures = "0.3" +moka = { version = "0.12", features = ["future"] } +serde = { version = "1", features = ["derive"] } +serde_json = "1.0.120" +tokio = { version = "1.39.2", features = ["sync"] } + +[dev-dependencies] +arrow-schema = { workspace = true } +tempfile = "3" +tokio = { version = "1.39.2", features = ["macros", "rt-multi-thread"] } diff --git a/crates/query-service/DEPENDENCIES.rust.tsv b/crates/query-service/DEPENDENCIES.rust.tsv new file mode 100644 index 00000000..3d6cb1ae --- /dev/null +++ b/crates/query-service/DEPENDENCIES.rust.tsv @@ -0,0 +1,392 @@ +crate 0BSD Apache-2.0 Apache-2.0 WITH LLVM-exception BSD-2-Clause BSD-3-Clause BSL-1.0 CC0-1.0 CDLA-Permissive-2.0 ISC LGPL-2.1-or-later MIT MIT-0 Unicode-3.0 Unlicense Zlib +adler2@2.0.1 X X X +ahash@0.8.12 X X +aho-corasick@1.1.4 X X +alloc-no-stdlib@2.0.4 X +alloc-stdlib@0.2.4 X +allocator-api2@0.2.21 X X +android_system_properties@0.1.5 X X +anyhow@1.0.104 X X +apache-avro@0.21.0 X +approx@0.5.1 X +arrow@58.3.0 X +arrow-arith@58.3.0 X +arrow-array@58.3.0 X X +arrow-buffer@58.3.0 X +arrow-cast@58.3.0 X +arrow-csv@58.3.0 X +arrow-data@58.3.0 X +arrow-ipc@58.3.0 X +arrow-json@58.3.0 X +arrow-ord@58.3.0 X +arrow-row@58.3.0 X +arrow-schema@58.3.0 X +arrow-select@58.3.0 X +arrow-string@58.3.0 X +async-lock@3.4.2 X X +async-stream@0.3.6 X +async-stream-impl@0.3.6 X +async-trait@0.1.91 X X +atoi@2.0.0 X +atomic-waker@1.1.2 X X +autocfg@1.5.1 X X +aws-lc-rs@1.17.3 X X +aws-lc-sys@0.43.0 X X X X X +backon@1.6.0 X +base64@0.22.1 X X +bigdecimal@0.4.10 X X +bitflags@2.13.1 X X +block-buffer@0.10.4 X X +block-buffer@0.12.1 X X +bon@3.9.3 X X +bon-macros@3.9.3 X X +brotli@8.0.4 X X +brotli-decompressor@5.0.3 X X +bumpalo@3.20.3 X X +bytemuck@1.25.2 X X X +byteorder@1.5.0 X X +bytes@1.12.1 X +cc@1.3.0 X X +cfg-if@1.0.4 X X +chrono@0.4.45 X X +chrono-tz@0.10.4 X X +cmake@0.1.58 X X +cmov@0.5.4 X X +combine@4.6.7 X +comfy-table@7.2.2 X +concurrent-queue@2.5.0 X X +const-oid@0.10.2 X X +const-random@0.1.18 X X +const-random-macro@0.1.16 X X +core-foundation@0.10.1 X X +core-foundation@0.9.4 X X +core-foundation-sys@0.8.7 X X +cpufeatures@0.2.17 X X +cpufeatures@0.3.0 X X +crc32fast@1.5.0 X X +crossbeam-channel@0.5.16 X X +crossbeam-deque@0.8.7 X X +crossbeam-epoch@0.9.20 X X +crossbeam-utils@0.8.22 X X +crunchy@0.2.4 X +crypto-common@0.1.7 X X +crypto-common@0.2.2 X X +csv@1.4.0 X X +csv-core@0.1.13 X X +ctutils@0.4.2 X X +darling@0.23.0 X +darling_core@0.23.0 X +darling_macro@0.23.0 X +diff@0.1.13 X X +digest@0.10.7 X X +digest@0.11.3 X X +displaydoc@0.2.6 X X +dlv-list@0.5.2 X X +dunce@1.0.5 X X X +either@1.16.0 X X +encoding_rs@0.8.35 X X X +equivalent@1.0.2 X X +errno@0.3.14 X X +event-listener@5.4.1 X X +event-listener-strategy@0.5.4 X X +fallible-streaming-iterator@0.1.9 X X +fastrand@2.5.0 X X +find-msvc-tools@0.1.9 X X +flatbuffers@25.12.19 X +flate2@1.1.9 X X +fnv@1.0.7 X X +foldhash@0.2.0 X +foreign-types@0.3.2 X X +foreign-types-shared@0.1.1 X X +form_urlencoded@1.2.2 X X +fs_extra@1.3.0 X +futures@0.3.33 X X +futures-channel@0.3.33 X X +futures-core@0.3.33 X X +futures-executor@0.3.33 X X +futures-io@0.3.33 X X +futures-macro@0.3.33 X X +futures-sink@0.3.33 X X +futures-task@0.3.33 X X +futures-util@0.3.33 X X +generic-array@0.14.7 X +getrandom@0.2.17 X X +getrandom@0.3.4 X X +getrandom@0.4.3 X X +gloo-timers@0.3.0 X X +h2@0.4.16 X +half@2.7.1 X X +hashbrown@0.14.5 X X +hashbrown@0.17.1 X X +heck@0.5.0 X X +hex@0.4.3 X X +hmac@0.12.1 X X +hmac@0.13.0 X X +http@1.4.2 X X +http-body@1.1.0 X +http-body-util@0.1.4 X +httparse@1.10.1 X X +httpdate@1.0.3 X X +hybrid-array@0.4.13 X X +hyper@1.10.1 X +hyper-rustls@0.27.9 X X X +hyper-tls@0.6.0 X X +hyper-util@0.1.20 X +iana-time-zone@0.1.65 X X +iana-time-zone-haiku@0.1.2 X X +icu_collections@2.2.0 X +icu_locale_core@2.2.0 X +icu_normalizer@2.2.0 X +icu_normalizer_data@2.2.0 X +icu_properties@2.2.0 X +icu_properties_data@2.2.0 X +icu_provider@2.2.0 X +ident_case@1.0.1 X X +idna@1.1.0 X X +idna_adapter@1.2.2 X X +indexmap@2.14.0 X X +integer-encoding@3.0.4 X +ipnet@2.12.0 X X +itertools@0.14.0 X X +itoa@1.0.18 X X +jiff@0.2.34 X X +jiff-core@0.1.0 X X +jiff-tzdb@0.1.8 X X +jiff-tzdb-platform@0.1.3 X X +jni@0.22.4 X X +jni-macros@0.22.4 X X +jni-sys@0.4.1 X X +jni-sys-macros@0.4.1 X X +jobserver@0.1.35 X X +js-sys@0.3.103 X X +lexical-core@1.0.6 X X +lexical-parse-float@1.0.6 X X +lexical-parse-integer@1.0.6 X X +lexical-util@1.0.7 X X +lexical-write-float@1.0.6 X X +lexical-write-integer@1.0.6 X X +libc@0.2.186 X X +libloading@0.9.0 X +libm@0.2.16 X +linux-raw-sys@0.12.1 X X X +litemap@0.8.2 X +lock_api@0.4.14 X X +log@0.4.33 X X +lru@0.18.2 X +lz4_flex@0.11.6 X +lz4_flex@0.13.1 X +lzokay-native@0.1.0 X +matrixmultiply@0.3.11 X X +md-5@0.10.6 X X +md-5@0.11.0 X X +mea@0.6.4 X +memchr@2.8.3 X X +mime@0.3.17 X X +miniz_oxide@0.8.9 X X X +mio@1.2.2 X +moka@0.12.15 X X +nalgebra@0.33.3 X +nalgebra-macros@0.2.2 X +native-tls@0.2.18 X X +num@0.4.3 X X +num-bigint@0.4.8 X X +num-complex@0.4.6 X X +num-integer@0.1.46 X X +num-iter@0.1.46 X X +num-rational@0.4.2 X X +num-traits@0.2.19 X X +once_cell@1.21.4 X X +opendal-core@0.58.0 X +opendal-http-transport-reqwest@0.58.0 X +opendal-layer-retry@0.58.0 X +opendal-service-fs@0.58.0 X +opendal-service-oss@0.58.0 X +openssl@0.10.81 X +openssl-macros@0.1.1 X X +openssl-probe@0.2.1 X X +openssl-sys@0.9.117 X +orc-rust@0.8.0 X +ordered-float@2.10.1 X +ordered-multimap@0.7.3 X +paimon@0.4.0 X +paimon-mosaic-core@0.2.0 X +paimon-query-service@0.4.0 X +paimon-vindex-core@0.3.0 X +parking@2.2.1 X X +parking_lot@0.12.5 X X +parking_lot_core@0.9.12 X X +parquet@58.3.0 X +paste@1.0.15 X X +percent-encoding@2.3.2 X X +phf@0.12.1 X +phf_shared@0.12.1 X +pin-project-lite@0.2.17 X X +pkg-config@0.3.33 X X +portable-atomic@1.14.0 X X +portable-atomic-util@0.2.7 X X +potential_utf@0.1.5 X +ppv-lite86@0.2.21 X X +pretty_assertions@1.4.1 X X +prettyplease@0.2.37 X X +proc-macro2@1.0.107 X X +prost@0.13.5 X +prost-derive@0.13.5 X +quad-rand@0.2.3 X +quick-xml@0.41.0 X +quote@1.0.47 X X +r-efi@5.3.0 X X X +r-efi@6.0.0 X X X +rand@0.8.7 X X +rand@0.9.5 X X +rand_chacha@0.3.1 X X +rand_chacha@0.9.0 X X +rand_core@0.6.4 X X +rand_core@0.9.5 X X +rawpointer@0.2.1 X X +rayon@1.12.0 X X +rayon-core@1.13.0 X X +redox_syscall@0.5.18 X +regex@1.13.1 X X +regex-automata@0.4.16 X X +regex-lite@0.1.9 X X +regex-syntax@0.8.11 X X +reqsign-aliyun-oss@3.1.1 X +reqsign-core@3.1.0 X +reqsign-file-read-tokio@3.0.2 X +reqwest@0.12.28 X X +reqwest@0.13.4 X X +roaring@0.11.4 X X +rust-ini@0.21.3 X +rustc_version@0.4.1 X X +rustix@1.1.4 X X X +rustls@0.23.42 X X X +rustls-native-certs@0.8.4 X X X +rustls-pki-types@1.15.0 X X +rustls-platform-verifier@0.7.0 X X +rustls-platform-verifier-android@0.1.1 X X +rustls-webpki@0.103.13 X +rustversion@1.0.23 X X +ryu@1.0.23 X X +safe_arch@0.7.4 X X X +same-file@1.0.6 X X +schannel@0.1.29 X +scopeguard@1.2.0 X X +security-framework@3.7.0 X X +security-framework-sys@2.17.0 X X +semver@1.0.28 X X +seq-macro@0.3.6 X X +serde@1.0.229 X X +serde_bytes@0.11.19 X X +serde_core@1.0.229 X X +serde_derive@1.0.229 X X +serde_json@1.0.151 X X +serde_repr@0.1.21 X X +serde_urlencoded@0.7.1 X X +serde_with@3.21.0 X X +serde_with_macros@3.21.0 X X +sha1@0.10.7 X X +sha1@0.11.0 X X +sha2@0.10.9 X X +sha2@0.11.0 X X +shlex@2.0.1 X X +simba@0.9.1 X +simd-adler32@0.3.10 X +simd_cesu8@1.2.0 X X +simdutf8@0.1.5 X X +siphasher@1.0.3 X X +slab@0.4.12 X +smallvec@1.15.2 X X +snafu@0.8.9 X X +snafu@0.9.1 X X +snafu-derive@0.8.9 X X +snafu-derive@0.9.1 X X +snap@1.1.2 X +socket2@0.6.5 X X +stable_deref_trait@1.2.1 X X +strsim@0.11.1 X +strum@0.27.2 X +strum_macros@0.27.2 X +subtle@2.6.1 X +syn@2.0.119 X X +syn@3.0.2 X X +sync_wrapper@1.0.2 X +synstructure@0.13.2 X +system-configuration@0.7.0 X X +system-configuration-sys@0.6.0 X X +tagptr@0.2.0 X X +tempfile@3.27.0 X X +thiserror@1.0.69 X X +thiserror@2.0.19 X X +thiserror-impl@1.0.69 X X +thiserror-impl@2.0.19 X X +thrift@0.17.0 X +tiny-keccak@2.0.2 X +tinystr@0.8.3 X +tokio@1.53.0 X +tokio-macros@2.7.1 X +tokio-native-tls@0.3.1 X +tokio-rustls@0.26.4 X X +tokio-util@0.7.18 X +tower@0.5.3 X +tower-http@0.6.11 X +tower-layer@0.3.3 X +tower-service@0.3.3 X +tracing@0.1.44 X +tracing-core@0.1.36 X +try-lock@0.2.5 X +twox-hash@2.1.3 X +typed-builder@0.19.1 X X +typed-builder-macro@0.19.1 X X +typenum@1.20.1 X X +unicode-ident@1.0.24 X X X +unicode-segmentation@1.13.2 X X +unicode-width@0.2.2 X X +untrusted@0.9.0 X +url@2.5.8 X X +urlencoding@2.1.3 X +utf8_iter@1.0.4 X X +uuid@1.24.0 X X +vcpkg@0.2.15 X X +version_check@0.9.5 X X +walkdir@2.5.0 X X +want@0.3.1 X +wasi@0.11.1+wasi-snapshot-preview1 X X X +wasip2@1.0.4+wasi-0.2.12 X X X +wasm-bindgen@0.2.126 X X +wasm-bindgen-futures@0.4.76 X X +wasm-bindgen-macro@0.2.126 X X +wasm-bindgen-macro-support@0.2.126 X X +wasm-bindgen-shared@0.2.126 X X +wasm-streams@0.5.0 X X +web-sys@0.3.103 X X +web-time@1.1.0 X X +webpki-root-certs@1.0.9 X +wide@0.7.33 X X X +winapi-util@0.1.11 X X +windows-core@0.62.2 X X +windows-implement@0.60.2 X X +windows-interface@0.59.3 X X +windows-link@0.2.1 X X +windows-registry@0.6.1 X X +windows-result@0.4.1 X X +windows-strings@0.5.1 X X +windows-sys@0.61.2 X X +wit-bindgen@0.57.1 X X X +writeable@0.6.3 X +xattr@1.6.1 X X +yansi@1.0.1 X X +yoke@0.8.3 X +yoke-derive@0.8.2 X +zerocopy@0.8.54 X X X +zerocopy-derive@0.8.54 X X X +zerofrom@0.1.8 X +zerofrom-derive@0.1.7 X +zeroize@1.9.0 X X +zerotrie@0.2.4 X +zerovec@0.11.6 X +zerovec-derive@0.11.3 X +zlib-rs@0.6.6 X +zmij@1.0.23 X +zstd@0.13.3 X +zstd-safe@7.2.4 X X +zstd-sys@2.0.16+zstd.1.5.7 X X diff --git a/crates/query-service/src/error.rs b/crates/query-service/src/error.rs new file mode 100644 index 00000000..f894ae9b --- /dev/null +++ b/crates/query-service/src/error.rs @@ -0,0 +1,105 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::fmt::{Display, Formatter}; +use std::sync::Arc; + +pub type Result = std::result::Result; + +#[derive(Debug)] +pub enum LookupError { + Paimon(paimon::Error), + PaimonUnavailable, + LoadCancelled, + InvalidRequest(String), + InvalidPolicy(String), + UnsupportedKeyType { + field: String, + data_type: String, + }, + InvalidKeyValue { + field: String, + message: String, + }, + QueryBudgetExceeded { + files: usize, + bytes: u64, + max_files: usize, + max_bytes: u64, + }, + SnapshotMismatch { + expected: i64, + actual: Option, + }, + InvalidDescriptor { + field: String, + message: String, + }, + UnexpectedResult(String), + Shared(Arc), +} + +impl Display for LookupError { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + match self { + Self::Paimon(error) => Display::fmt(error, f), + Self::PaimonUnavailable => write!(f, "Paimon metadata or data is unavailable"), + Self::LoadCancelled => write!(f, "in-flight descriptor lookup was cancelled"), + Self::InvalidRequest(message) => write!(f, "invalid lookup request: {message}"), + Self::InvalidPolicy(message) => write!(f, "invalid lookup policy: {message}"), + Self::UnsupportedKeyType { field, data_type } => { + write!(f, "unsupported lookup key type for '{field}': {data_type}") + } + Self::InvalidKeyValue { field, message } => { + write!(f, "invalid value for lookup key '{field}': {message}") + } + Self::QueryBudgetExceeded { + files, + bytes, + max_files, + max_bytes, + } => write!( + f, + "lookup plan exceeds budget: files={files}/{max_files}, bytes={bytes}/{max_bytes}" + ), + Self::SnapshotMismatch { expected, actual } => write!( + f, + "lookup planned an unexpected snapshot: expected={expected}, actual={actual:?}" + ), + Self::InvalidDescriptor { field, message } => { + write!(f, "invalid BlobDescriptor in field '{field}': {message}") + } + Self::UnexpectedResult(message) => write!(f, "unexpected lookup result: {message}"), + Self::Shared(error) => Display::fmt(error, f), + } + } +} + +impl std::error::Error for LookupError { + fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { + match self { + Self::Paimon(error) => Some(error), + Self::Shared(error) => Some(error.as_ref()), + _ => None, + } + } +} + +impl From for LookupError { + fn from(value: paimon::Error) -> Self { + Self::Paimon(value) + } +} diff --git a/crates/query-service/src/key.rs b/crates/query-service/src/key.rs new file mode 100644 index 00000000..30836310 --- /dev/null +++ b/crates/query-service/src/key.rs @@ -0,0 +1,1011 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::BTreeSet; + +use arrow_array::{ + Array, BinaryArray, BinaryViewArray, BooleanArray, Date32Array, Decimal128Array, Float32Array, + Float64Array, Int16Array, Int32Array, Int64Array, Int8Array, LargeBinaryArray, + LargeStringArray, StringArray, StringViewArray, Time32MillisecondArray, + TimestampMicrosecondArray, TimestampMillisecondArray, TimestampNanosecondArray, +}; +use base64::Engine; +use chrono::{DateTime, Datelike, NaiveDate, NaiveDateTime, NaiveTime, Timelike, Utc}; +use paimon::spec::{DataField, DataType, Datum, Predicate, PredicateBuilder}; +use serde_json::Value; + +use crate::error::{LookupError, Result}; +use crate::model::LookupKey; + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub(crate) enum KeyComponent { + Bool(bool), + TinyInt(i8), + SmallInt(i16), + Int(i32), + Long(i64), + Float(u32), + Double(u64), + String(String), + Bytes(Vec), + Date(i32), + Time(i32), + Timestamp { + millis: i64, + nanos: i32, + }, + LocalZonedTimestamp { + millis: i64, + nanos: i32, + }, + Decimal { + unscaled: i128, + precision: u32, + scale: u32, + }, +} + +impl KeyComponent { + fn datum(&self) -> Datum { + match self { + Self::Bool(value) => Datum::Bool(*value), + Self::TinyInt(value) => Datum::TinyInt(*value), + Self::SmallInt(value) => Datum::SmallInt(*value), + Self::Int(value) => Datum::Int(*value), + Self::Long(value) => Datum::Long(*value), + Self::Float(bits) => Datum::Float(f32::from_bits(*bits)), + Self::Double(bits) => Datum::Double(f64::from_bits(*bits)), + Self::String(value) => Datum::String(value.clone()), + Self::Bytes(value) => Datum::Bytes(value.clone()), + Self::Date(value) => Datum::Date(*value), + Self::Time(value) => Datum::Time(*value), + Self::Timestamp { millis, nanos } => Datum::Timestamp { + millis: *millis, + nanos: *nanos, + }, + Self::LocalZonedTimestamp { millis, nanos } => Datum::LocalZonedTimestamp { + millis: *millis, + nanos: *nanos, + }, + Self::Decimal { + unscaled, + precision, + scale, + } => Datum::Decimal { + unscaled: *unscaled, + precision: *precision, + scale: *scale, + }, + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub(crate) struct NormalizedKey(pub(crate) Vec); + +#[derive(Debug, Clone)] +pub(crate) struct PreparedKey { + pub(crate) original: LookupKey, + pub(crate) normalized: NormalizedKey, +} + +pub(crate) fn validate_key_type(field: &str, data_type: &DataType) -> Result<()> { + if matches!( + data_type, + DataType::Boolean(_) + | DataType::TinyInt(_) + | DataType::SmallInt(_) + | DataType::Int(_) + | DataType::BigInt(_) + | DataType::Float(_) + | DataType::Double(_) + | DataType::Char(_) + | DataType::VarChar(_) + | DataType::Binary(_) + | DataType::VarBinary(_) + | DataType::Date(_) + | DataType::Time(_) + | DataType::Timestamp(_) + | DataType::LocalZonedTimestamp(_) + | DataType::Decimal(_) + ) { + Ok(()) + } else { + Err(LookupError::UnsupportedKeyType { + field: field.to_string(), + data_type: format!("{data_type:?}"), + }) + } +} + +pub(crate) fn supports_global_btree(data_type: &DataType) -> bool { + matches!( + data_type, + DataType::Boolean(_) + | DataType::TinyInt(_) + | DataType::SmallInt(_) + | DataType::Int(_) + | DataType::BigInt(_) + | DataType::Float(_) + | DataType::Double(_) + | DataType::Char(_) + | DataType::VarChar(_) + | DataType::Date(_) + | DataType::Time(_) + | DataType::Timestamp(_) + | DataType::LocalZonedTimestamp(_) + | DataType::Decimal(_) + ) +} + +pub(crate) fn prepare_keys( + keys: Vec, + key_fields: &[String], + schema_fields: &[DataField], +) -> Result> { + let expected = key_fields.iter().cloned().collect::>(); + let field_types = key_fields + .iter() + .map(|name| { + schema_fields + .iter() + .find(|field| field.name() == name) + .map(|field| field.data_type()) + .ok_or_else(|| { + LookupError::InvalidPolicy(format!( + "lookup key field '{name}' is missing from the table schema" + )) + }) + }) + .collect::>>()?; + + keys.into_iter() + .map(|original| { + let actual = original.keys().cloned().collect::>(); + if actual != expected { + return Err(LookupError::InvalidRequest(format!( + "lookup key fields must be exactly {:?}, got {:?}", + key_fields, + original.keys().collect::>() + ))); + } + let components = key_fields + .iter() + .zip(&field_types) + .map(|(name, data_type)| json_to_component(name, &original[name], data_type)) + .collect::>>()?; + Ok(PreparedKey { + original, + normalized: NormalizedKey(components), + }) + }) + .collect() +} + +pub(crate) fn build_batch_predicate( + prepared: &[PreparedKey], + key_fields: &[String], + schema_fields: &[DataField], +) -> Result { + if prepared.is_empty() { + return Ok(Predicate::AlwaysFalse); + } + + let builder = PredicateBuilder::new(schema_fields); + if key_fields.len() == 1 { + let mut seen = std::collections::HashSet::new(); + let values = prepared + .iter() + .filter_map(|key| { + let component = key.normalized.0[0].clone(); + seen.insert(component.clone()).then(|| component.datum()) + }) + .collect(); + return builder + .is_in(&key_fields[0], values) + .map_err(LookupError::from); + } + + let mut seen = std::collections::HashSet::new(); + let disjunction = prepared + .iter() + .filter(|key| seen.insert(key.normalized.clone())) + .map(|key| { + key_fields + .iter() + .zip(&key.normalized.0) + .map(|(field, value)| builder.equal(field, value.datum())) + .collect::>>() + .map(Predicate::and) + }) + .collect::>>()?; + Ok(Predicate::or(disjunction)) +} + +pub(crate) fn normalized_key_from_batch( + batch: &arrow_array::RecordBatch, + row: usize, + key_fields: &[String], + schema_fields: &[DataField], +) -> Result { + let mut components = Vec::with_capacity(key_fields.len()); + for (column_index, name) in key_fields.iter().enumerate() { + let field = schema_fields + .iter() + .find(|field| field.name() == name) + .ok_or_else(|| { + LookupError::InvalidPolicy(format!( + "lookup key field '{name}' is missing from the table schema" + )) + })?; + components.push(array_to_component( + name, + batch.column(column_index).as_ref(), + row, + field.data_type(), + )?); + } + Ok(NormalizedKey(components)) +} + +fn json_to_component(field: &str, value: &Value, data_type: &DataType) -> Result { + match data_type { + DataType::Boolean(_) => value + .as_bool() + .map(KeyComponent::Bool) + .ok_or_else(|| invalid_value(field, "expected a JSON boolean")), + DataType::TinyInt(_) => parse_signed(field, value).and_then(|value| { + i8::try_from(value) + .map(KeyComponent::TinyInt) + .map_err(|_| invalid_value(field, "value is outside TINYINT range")) + }), + DataType::SmallInt(_) => parse_signed(field, value).and_then(|value| { + i16::try_from(value) + .map(KeyComponent::SmallInt) + .map_err(|_| invalid_value(field, "value is outside SMALLINT range")) + }), + DataType::Int(_) => parse_signed(field, value).and_then(|value| { + i32::try_from(value) + .map(KeyComponent::Int) + .map_err(|_| invalid_value(field, "value is outside INT range")) + }), + DataType::BigInt(_) => parse_signed(field, value).map(KeyComponent::Long), + DataType::Float(_) => parse_float(field, value), + DataType::Double(_) => parse_double(field, value), + DataType::Char(_) | DataType::VarChar(_) => value + .as_str() + .map(|value| KeyComponent::String(value.to_string())) + .ok_or_else(|| invalid_value(field, "expected a JSON string")), + DataType::Binary(_) | DataType::VarBinary(_) => { + let encoded = value + .as_object() + .and_then(|object| object.get("base64")) + .and_then(Value::as_str) + .ok_or_else(|| invalid_value(field, "expected {\"base64\":\"...\"}"))?; + base64::engine::general_purpose::STANDARD + .decode(encoded) + .map(KeyComponent::Bytes) + .map_err(|error| invalid_value(field, format!("invalid base64: {error}"))) + } + DataType::Date(_) => parse_date(field, value), + DataType::Time(data_type) => parse_time(field, value, data_type.precision()), + DataType::Timestamp(data_type) => { + parse_timestamp(field, value, data_type.precision(), false) + } + DataType::LocalZonedTimestamp(data_type) => { + parse_timestamp(field, value, data_type.precision(), true) + } + DataType::Decimal(data_type) => { + parse_decimal(field, value, data_type.precision(), data_type.scale()) + } + other => Err(LookupError::UnsupportedKeyType { + field: field.to_string(), + data_type: format!("{other:?}"), + }), + } +} + +fn array_to_component( + field: &str, + array: &dyn Array, + row: usize, + data_type: &DataType, +) -> Result { + if array.is_null(row) { + return Err(LookupError::UnexpectedResult(format!( + "lookup key field '{field}' was NULL" + ))); + } + + macro_rules! primitive { + ($array:ty, $variant:ident) => { + array + .as_any() + .downcast_ref::<$array>() + .map(|array| KeyComponent::$variant(array.value(row))) + .ok_or_else(|| arrow_type_mismatch(field, array, data_type)) + }; + } + + match data_type { + DataType::Boolean(_) => primitive!(BooleanArray, Bool), + DataType::TinyInt(_) => primitive!(Int8Array, TinyInt), + DataType::SmallInt(_) => primitive!(Int16Array, SmallInt), + DataType::Int(_) => primitive!(Int32Array, Int), + DataType::BigInt(_) => primitive!(Int64Array, Long), + DataType::Float(_) => array + .as_any() + .downcast_ref::() + .map(|array| KeyComponent::Float(canonical_f32_bits(array.value(row)))) + .ok_or_else(|| arrow_type_mismatch(field, array, data_type)), + DataType::Double(_) => array + .as_any() + .downcast_ref::() + .map(|array| KeyComponent::Double(canonical_f64_bits(array.value(row)))) + .ok_or_else(|| arrow_type_mismatch(field, array, data_type)), + DataType::Char(_) | DataType::VarChar(_) => { + string_component(array, row).ok_or_else(|| arrow_type_mismatch(field, array, data_type)) + } + DataType::Binary(_) | DataType::VarBinary(_) => { + binary_component(array, row).ok_or_else(|| arrow_type_mismatch(field, array, data_type)) + } + DataType::Date(_) => primitive!(Date32Array, Date), + DataType::Time(_) => primitive!(Time32MillisecondArray, Time), + DataType::Timestamp(timestamp_type) => timestamp_component( + field, + array, + row, + timestamp_type.precision(), + false, + data_type, + ), + DataType::LocalZonedTimestamp(timestamp_type) => timestamp_component( + field, + array, + row, + timestamp_type.precision(), + true, + data_type, + ), + DataType::Decimal(decimal_type) => array + .as_any() + .downcast_ref::() + .map(|array| KeyComponent::Decimal { + unscaled: array.value(row), + precision: decimal_type.precision(), + scale: decimal_type.scale(), + }) + .ok_or_else(|| arrow_type_mismatch(field, array, data_type)), + other => Err(LookupError::UnsupportedKeyType { + field: field.to_string(), + data_type: format!("{other:?}"), + }), + } +} + +fn parse_float(field: &str, value: &Value) -> Result { + let value = match value { + Value::Number(value) => value + .as_f64() + .map(|value| value as f32) + .ok_or_else(|| invalid_value(field, "expected a finite JSON number"))?, + Value::String(value) => value + .parse::() + .map_err(|error| invalid_value(field, format!("invalid FLOAT: {error}")))?, + _ => return Err(invalid_value(field, "expected a number or numeric string")), + }; + if !value.is_finite() { + return Err(invalid_value(field, "FLOAT must be finite")); + } + Ok(KeyComponent::Float(canonical_f32_bits(value))) +} + +fn parse_double(field: &str, value: &Value) -> Result { + let value = match value { + Value::Number(value) => value + .as_f64() + .ok_or_else(|| invalid_value(field, "expected a finite JSON number"))?, + Value::String(value) => value + .parse::() + .map_err(|error| invalid_value(field, format!("invalid DOUBLE: {error}")))?, + _ => return Err(invalid_value(field, "expected a number or numeric string")), + }; + if !value.is_finite() { + return Err(invalid_value(field, "DOUBLE must be finite")); + } + Ok(KeyComponent::Double(canonical_f64_bits(value))) +} + +fn canonical_f32_bits(value: f32) -> u32 { + if value == 0.0 { + 0.0_f32.to_bits() + } else { + value.to_bits() + } +} + +fn canonical_f64_bits(value: f64) -> u64 { + if value == 0.0 { + 0.0_f64.to_bits() + } else { + value.to_bits() + } +} + +fn parse_decimal(field: &str, value: &Value, precision: u32, scale: u32) -> Result { + let text = value.as_str().ok_or_else(|| { + invalid_value(field, "expected a decimal string without exponent notation") + })?; + let (negative, unsigned) = match text.as_bytes().first() { + Some(b'-') => (true, &text[1..]), + Some(b'+') => (false, &text[1..]), + _ => (false, text), + }; + let mut parts = unsigned.split('.'); + let integer = parts.next().unwrap_or_default(); + let fraction = parts.next().unwrap_or_default(); + if integer.is_empty() + || !integer.bytes().all(|byte| byte.is_ascii_digit()) + || !fraction.bytes().all(|byte| byte.is_ascii_digit()) + || parts.next().is_some() + || (unsigned.contains('.') && fraction.is_empty()) + { + return Err(invalid_value( + field, + "expected a decimal string such as \"123.45\"", + )); + } + if fraction.len() > scale as usize { + return Err(invalid_value( + field, + format!("fractional digits exceed DECIMAL scale {scale}"), + )); + } + + let mut digits = String::with_capacity(integer.len() + scale as usize); + digits.push_str(integer); + digits.push_str(fraction); + digits.extend(std::iter::repeat_n('0', scale as usize - fraction.len())); + let significant = digits.trim_start_matches('0'); + if significant.len() > precision as usize { + return Err(invalid_value( + field, + format!("value exceeds DECIMAL({precision}, {scale}) precision"), + )); + } + let magnitude = if significant.is_empty() { + 0 + } else { + significant.parse::().map_err(|error| { + invalid_value(field, format!("decimal cannot be represented: {error}")) + })? + }; + let unscaled = if negative { -magnitude } else { magnitude }; + Ok(KeyComponent::Decimal { + unscaled, + precision, + scale, + }) +} + +fn parse_date(field: &str, value: &Value) -> Result { + let text = value + .as_str() + .ok_or_else(|| invalid_value(field, "expected an ISO date string YYYY-MM-DD"))?; + let date = NaiveDate::parse_from_str(text, "%Y-%m-%d") + .map_err(|error| invalid_value(field, format!("invalid ISO date: {error}")))?; + validate_paimon_year(field, date.year())?; + let epoch = NaiveDate::from_ymd_opt(1970, 1, 1).expect("Unix epoch must be valid"); + let days = i32::try_from(date.signed_duration_since(epoch).num_days()) + .map_err(|_| invalid_value(field, "date is outside DATE range"))?; + Ok(KeyComponent::Date(days)) +} + +fn parse_time(field: &str, value: &Value, precision: u32) -> Result { + let text = value + .as_str() + .ok_or_else(|| invalid_value(field, "expected an ISO time string HH:MM:SS[.fraction]"))?; + validate_fraction_precision(field, text, precision)?; + let time = NaiveTime::parse_from_str(text, "%H:%M:%S%.f") + .map_err(|error| invalid_value(field, format!("invalid ISO time: {error}")))?; + if time.nanosecond() >= 1_000_000_000 { + return Err(invalid_value(field, "leap seconds are not supported")); + } + if time.nanosecond() % 1_000_000 != 0 { + return Err(invalid_value( + field, + "TIME lookup values are limited to millisecond precision", + )); + } + let millis = i32::try_from( + u64::from(time.num_seconds_from_midnight()) * 1_000 + + u64::from(time.nanosecond() / 1_000_000), + ) + .expect("millis of day always fit i32"); + Ok(KeyComponent::Time(millis)) +} + +fn parse_timestamp( + field: &str, + value: &Value, + precision: u32, + local_zoned: bool, +) -> Result { + let text = value.as_str().ok_or_else(|| { + invalid_value( + field, + if local_zoned { + "expected an RFC 3339 timestamp with Z or an explicit offset" + } else { + "expected an ISO timestamp YYYY-MM-DDTHH:MM:SS[.fraction] without a time zone" + }, + ) + })?; + validate_fraction_precision(field, text, precision)?; + + let (seconds, subsecond_nanos) = if local_zoned { + let timestamp = DateTime::parse_from_rfc3339(text).map_err(|error| { + invalid_value(field, format!("invalid RFC 3339 timestamp: {error}")) + })?; + validate_paimon_year(field, timestamp.year())?; + let timestamp = timestamp.with_timezone(&Utc); + (timestamp.timestamp(), timestamp.timestamp_subsec_nanos()) + } else { + let timestamp = NaiveDateTime::parse_from_str(text, "%Y-%m-%dT%H:%M:%S%.f") + .map_err(|error| invalid_value(field, format!("invalid ISO timestamp: {error}")))?; + validate_paimon_year(field, timestamp.year())?; + let timestamp = timestamp.and_utc(); + (timestamp.timestamp(), timestamp.timestamp_subsec_nanos()) + }; + if subsecond_nanos >= 1_000_000_000 { + return Err(invalid_value(field, "leap seconds are not supported")); + } + let (millis, nanos) = checked_timestamp_parts(field, seconds, subsecond_nanos, precision)?; + Ok(if local_zoned { + KeyComponent::LocalZonedTimestamp { millis, nanos } + } else { + KeyComponent::Timestamp { millis, nanos } + }) +} + +fn validate_paimon_year(field: &str, year: i32) -> Result<()> { + if (0..=9_999).contains(&year) { + Ok(()) + } else { + Err(invalid_value( + field, + "year is outside the Paimon range 0000 through 9999", + )) + } +} + +fn validate_fraction_precision(field: &str, value: &str, precision: u32) -> Result<()> { + let fraction_digits = value + .split_once('.') + .map(|(_, suffix)| { + suffix + .bytes() + .take_while(|byte| byte.is_ascii_digit()) + .count() + }) + .unwrap_or(0); + if fraction_digits > precision as usize { + return Err(invalid_value( + field, + format!("fractional digits exceed declared precision {precision}"), + )); + } + Ok(()) +} + +fn checked_timestamp_parts( + field: &str, + seconds: i64, + subsecond_nanos: u32, + precision: u32, +) -> Result<(i64, i32)> { + let value_fits_arrow = match precision { + 0..=3 => seconds + .checked_mul(1_000) + .and_then(|value| value.checked_add(i64::from(subsecond_nanos / 1_000_000))) + .is_some(), + 4..=6 => seconds + .checked_mul(1_000_000) + .and_then(|value| value.checked_add(i64::from(subsecond_nanos / 1_000))) + .is_some(), + 7..=9 => seconds + .checked_mul(1_000_000_000) + .and_then(|value| value.checked_add(i64::from(subsecond_nanos))) + .is_some(), + _ => false, + }; + if !value_fits_arrow { + return Err(invalid_value( + field, + "timestamp is outside the Arrow range for the field precision", + )); + } + let millis = seconds + .checked_mul(1_000) + .and_then(|value| value.checked_add(i64::from(subsecond_nanos / 1_000_000))) + .ok_or_else(|| invalid_value(field, "timestamp milliseconds overflow i64"))?; + Ok((millis, (subsecond_nanos % 1_000_000) as i32)) +} + +fn timestamp_component( + field: &str, + array: &dyn Array, + row: usize, + precision: u32, + local_zoned: bool, + data_type: &DataType, +) -> Result { + let (millis, nanos) = match precision { + 0..=3 => array + .as_any() + .downcast_ref::() + .map(|array| (array.value(row), 0)), + 4..=6 => array + .as_any() + .downcast_ref::() + .map(|array| { + let micros = array.value(row); + ( + micros.div_euclid(1_000), + (micros.rem_euclid(1_000) * 1_000) as i32, + ) + }), + 7..=9 => array + .as_any() + .downcast_ref::() + .map(|array| { + let nanos = array.value(row); + ( + nanos.div_euclid(1_000_000), + nanos.rem_euclid(1_000_000) as i32, + ) + }), + _ => None, + } + .ok_or_else(|| arrow_type_mismatch(field, array, data_type))?; + Ok(if local_zoned { + KeyComponent::LocalZonedTimestamp { millis, nanos } + } else { + KeyComponent::Timestamp { millis, nanos } + }) +} + +fn string_component(array: &dyn Array, row: usize) -> Option { + if let Some(array) = array.as_any().downcast_ref::() { + Some(KeyComponent::String(array.value(row).to_string())) + } else if let Some(array) = array.as_any().downcast_ref::() { + Some(KeyComponent::String(array.value(row).to_string())) + } else { + array + .as_any() + .downcast_ref::() + .map(|array| KeyComponent::String(array.value(row).to_string())) + } +} + +fn binary_component(array: &dyn Array, row: usize) -> Option { + if let Some(array) = array.as_any().downcast_ref::() { + Some(KeyComponent::Bytes(array.value(row).to_vec())) + } else if let Some(array) = array.as_any().downcast_ref::() { + Some(KeyComponent::Bytes(array.value(row).to_vec())) + } else { + array + .as_any() + .downcast_ref::() + .map(|array| KeyComponent::Bytes(array.value(row).to_vec())) + } +} + +fn parse_signed(field: &str, value: &Value) -> Result { + if let Some(value) = value.as_i64() { + return Ok(value); + } + value + .as_str() + .ok_or_else(|| invalid_value(field, "expected an integer or decimal integer string"))? + .parse::() + .map_err(|error| invalid_value(field, format!("invalid integer: {error}"))) +} + +fn invalid_value(field: &str, message: impl Into) -> LookupError { + LookupError::InvalidKeyValue { + field: field.to_string(), + message: message.into(), + } +} + +fn arrow_type_mismatch(field: &str, array: &dyn Array, data_type: &DataType) -> LookupError { + LookupError::UnexpectedResult(format!( + "lookup key field '{field}' expected {data_type:?}, got Arrow {:?}", + array.data_type() + )) +} + +#[cfg(test)] +mod tests { + use std::collections::BTreeMap; + use std::sync::Arc; + + use arrow_array::{ + ArrayRef, BinaryViewArray, Date32Array, Decimal128Array, Float32Array, Float64Array, + RecordBatch, StringViewArray, Time32MillisecondArray, TimestampMicrosecondArray, + TimestampNanosecondArray, + }; + use arrow_schema::{Field, Schema}; + use paimon::spec::{ + BigIntType, DataField, DataType, DateType, DecimalType, DoubleType, FloatType, IntType, + LocalZonedTimestampType, TimeType, TimestampType, VarBinaryType, VarCharType, + }; + use serde_json::json; + + use super::*; + + fn fields() -> Vec { + vec![ + DataField::new( + 0, + "tenant".into(), + DataType::VarChar(VarCharType::string_type()), + ), + DataField::new(1, "id".into(), DataType::BigInt(BigIntType::new())), + DataField::new(2, "shard".into(), DataType::Int(IntType::new())), + ] + } + + #[test] + fn prepares_bigint_strings_in_policy_order() { + let original = BTreeMap::from([ + ("id".to_string(), json!("9007199254740993")), + ("tenant".to_string(), json!("t1")), + ]); + let prepared = + prepare_keys(vec![original], &["tenant".into(), "id".into()], &fields()).unwrap(); + assert_eq!( + prepared[0].normalized, + NormalizedKey(vec![ + KeyComponent::String("t1".into()), + KeyComponent::Long(9_007_199_254_740_993), + ]) + ); + } + + #[test] + fn rejects_missing_and_extra_fields() { + let error = prepare_keys( + vec![BTreeMap::from([ + ("tenant".to_string(), json!("t1")), + ("extra".to_string(), json!(1)), + ])], + &["tenant".into(), "id".into()], + &fields(), + ) + .unwrap_err(); + assert!(matches!(error, LookupError::InvalidRequest(_))); + } + + #[test] + fn single_key_batch_uses_in_predicate() { + let prepared = prepare_keys( + vec![ + BTreeMap::from([("shard".to_string(), json!(1))]), + BTreeMap::from([("shard".to_string(), json!(2))]), + ], + &["shard".into()], + &fields(), + ) + .unwrap(); + let predicate = build_batch_predicate(&prepared, &["shard".into()], &fields()).unwrap(); + assert_eq!(predicate.to_string(), "shard IN (1, 2)"); + } + + #[test] + fn normalizes_extended_json_types_like_arrow_rows() { + let fields = extended_fields(); + let key_fields = fields + .iter() + .map(|field| field.name().to_string()) + .collect::>(); + let key = BTreeMap::from([ + ("float".to_string(), json!(-0.0)), + ("double".to_string(), json!("1.25")), + ("decimal".to_string(), json!("123.45")), + ("date".to_string(), json!("2024-02-29")), + ("time".to_string(), json!("12:34:56.789")), + ("timestamp".to_string(), json!("2024-02-29T12:34:56.123456")), + ( + "local_timestamp".to_string(), + json!("2024-02-29T20:34:56.123456789+08:00"), + ), + ]); + let mut prepared = prepare_keys(vec![key], &key_fields, &fields).unwrap(); + build_batch_predicate(&prepared, &key_fields, &fields).unwrap(); + let expected = prepared.remove(0).normalized; + + let instant = DateTime::parse_from_rfc3339("2024-02-29T12:34:56.123456789Z").unwrap(); + let timestamp_micros = instant.timestamp_micros(); + let timestamp_nanos = instant.timestamp_nanos_opt().unwrap(); + let arrays: Vec = vec![ + Arc::new(Float32Array::from(vec![0.0])), + Arc::new(Float64Array::from(vec![1.25])), + Arc::new( + Decimal128Array::from(vec![12345_i128]) + .with_precision_and_scale(10, 2) + .unwrap(), + ), + Arc::new(Date32Array::from(vec![19_782])), + Arc::new(Time32MillisecondArray::from(vec![45_296_789])), + Arc::new(TimestampMicrosecondArray::from(vec![timestamp_micros])), + Arc::new(TimestampNanosecondArray::from(vec![timestamp_nanos]).with_timezone("UTC")), + ]; + let arrow_fields = arrays + .iter() + .zip(&key_fields) + .map(|(array, name)| Field::new(name, array.data_type().clone(), false)) + .collect::>(); + let batch = RecordBatch::try_new(Arc::new(Schema::new(arrow_fields)), arrays).unwrap(); + + assert_eq!( + normalized_key_from_batch(&batch, 0, &key_fields, &fields).unwrap(), + expected + ); + } + + #[test] + fn normalizes_negative_subsecond_timestamp() { + let data_type = DataType::Timestamp(TimestampType::new(6).unwrap()); + let json = json_to_component( + "timestamp", + &json!("1969-12-31T23:59:59.999999"), + &data_type, + ) + .unwrap(); + let arrow = TimestampMicrosecondArray::from(vec![-1]); + assert_eq!( + array_to_component("timestamp", &arrow, 0, &data_type).unwrap(), + json + ); + assert_eq!( + json, + KeyComponent::Timestamp { + millis: -1, + nanos: 999_000, + } + ); + } + + #[test] + fn rejects_lossy_or_ambiguous_extended_values() { + let decimal = DataType::Decimal(DecimalType::new(5, 2).unwrap()); + assert_invalid_key(json_to_component("amount", &json!(12.34), &decimal)); + assert_invalid_key(json_to_component("amount", &json!("1e2"), &decimal)); + assert_invalid_key(json_to_component("amount", &json!("1.234"), &decimal)); + assert_invalid_key(json_to_component("amount", &json!("1234.00"), &decimal)); + + let time = DataType::Time(TimeType::new(6).unwrap()); + assert_invalid_key(json_to_component("time", &json!("12:00:00.000001"), &time)); + assert_invalid_key(json_to_component("time", &json!("23:59:60"), &time)); + + let timestamp = DataType::Timestamp(TimestampType::new(6).unwrap()); + assert_invalid_key(json_to_component( + "timestamp", + &json!("2024-01-01T00:00:00.1234567"), + ×tamp, + )); + assert_invalid_key(json_to_component( + "timestamp", + &json!("2024-01-01T00:00:00Z"), + ×tamp, + )); + + let local_timestamp = + DataType::LocalZonedTimestamp(LocalZonedTimestampType::new(9).unwrap()); + assert_invalid_key(json_to_component( + "local_timestamp", + &json!("2024-01-01T00:00:00"), + &local_timestamp, + )); + + assert_invalid_key(json_to_component( + "float", + &json!("NaN"), + &DataType::Float(FloatType::new()), + )); + assert_invalid_key(json_to_component( + "double", + &json!("inf"), + &DataType::Double(DoubleType::new()), + )); + } + + #[test] + fn accepts_arrow_view_arrays_for_text_and_binary_keys() { + let text = StringViewArray::from(vec!["tenant-a"]); + let binary = BinaryViewArray::from_iter_values([b"asset-1".as_slice()]); + assert_eq!( + array_to_component( + "tenant", + &text, + 0, + &DataType::VarChar(VarCharType::string_type()), + ) + .unwrap(), + KeyComponent::String("tenant-a".into()) + ); + assert_eq!( + array_to_component( + "asset", + &binary, + 0, + &DataType::VarBinary(VarBinaryType::new(32).unwrap()), + ) + .unwrap(), + KeyComponent::Bytes(b"asset-1".to_vec()) + ); + } + + #[test] + fn validates_supported_key_type_set() { + for field in extended_fields() { + validate_key_type(field.name(), field.data_type()).unwrap(); + } + let error = + validate_key_type("blob", &DataType::Blob(paimon::spec::BlobType::new())).unwrap_err(); + assert!(matches!(error, LookupError::UnsupportedKeyType { .. })); + } + + #[test] + fn global_btree_rejects_binary_keys_but_primary_key_normalization_supports_them() { + let binary = DataType::VarBinary(VarBinaryType::new(32).unwrap()); + validate_key_type("asset", &binary).unwrap(); + assert!(!supports_global_btree(&binary)); + assert!(supports_global_btree(&DataType::Decimal( + DecimalType::new(10, 2).unwrap() + ))); + } + + fn extended_fields() -> Vec { + vec![ + DataField::new(0, "float".into(), DataType::Float(FloatType::new())), + DataField::new(1, "double".into(), DataType::Double(DoubleType::new())), + DataField::new( + 2, + "decimal".into(), + DataType::Decimal(DecimalType::new(10, 2).unwrap()), + ), + DataField::new(3, "date".into(), DataType::Date(DateType::new())), + DataField::new(4, "time".into(), DataType::Time(TimeType::new(3).unwrap())), + DataField::new( + 5, + "timestamp".into(), + DataType::Timestamp(TimestampType::new(6).unwrap()), + ), + DataField::new( + 6, + "local_timestamp".into(), + DataType::LocalZonedTimestamp(LocalZonedTimestampType::new(9).unwrap()), + ), + ] + } + + fn assert_invalid_key(result: Result) { + assert!(matches!(result, Err(LookupError::InvalidKeyValue { .. }))); + } +} diff --git a/crates/query-service/src/lib.rs b/crates/query-service/src/lib.rs new file mode 100644 index 00000000..6937cf6a --- /dev/null +++ b/crates/query-service/src/lib.rs @@ -0,0 +1,34 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Snapshot-consistent, budgeted point queries over Paimon tables. +//! +//! The initial query capability returns BLOB descriptors. Additional projected +//! field types can be added behind this crate's service boundary. + +mod error; +mod key; +mod lookup; +mod model; +mod policy; + +pub use error::{LookupError, Result}; +pub use lookup::{BlobLookupOptions, BlobLookupService, DescriptorCacheStats, TableCacheStats}; +pub use model::{ + BatchGetRequest, BatchGetResponse, BlobDescriptorDto, DescriptorFormat, LookupKey, + LookupResult, LookupScanStats, LookupStatus, TableRef, +}; +pub use policy::{LookupStrategy, QueryBudget, TableLookupPolicy}; diff --git a/crates/query-service/src/lookup.rs b/crates/query-service/src/lookup.rs new file mode 100644 index 00000000..d33b0673 --- /dev/null +++ b/crates/query-service/src/lookup.rs @@ -0,0 +1,863 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::{BTreeMap, HashMap}; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{Arc, Mutex, RwLock, Weak}; +use std::time::{Duration, Instant}; + +use arrow_array::{Array, BinaryArray, RecordBatch}; +use base64::Engine; +use futures::lock::Mutex as AsyncMutex; +use futures::{StreamExt, TryStreamExt}; +use moka::future::Cache; +use paimon::spec::BlobDescriptor; +use paimon::{Catalog, Table}; +use tokio::sync::Notify; + +use crate::error::{LookupError, Result}; +use crate::key::{ + build_batch_predicate, normalized_key_from_batch, prepare_keys, NormalizedKey, PreparedKey, +}; +use crate::model::{ + BatchGetRequest, BatchGetResponse, BlobDescriptorDto, DescriptorFormat, LookupResult, + LookupScanStats, LookupStatus, TableRef, +}; +use crate::policy::TableLookupPolicy; + +#[derive(Debug, Clone)] +pub struct BlobLookupOptions { + pub table_cache_ttl: Duration, + pub descriptor_cache_ttl: Duration, + pub descriptor_cache_max_bytes: u64, + pub descriptor_cache_max_entry_bytes: u64, + pub global_index_thread_num: usize, +} + +impl Default for BlobLookupOptions { + fn default() -> Self { + Self { + table_cache_ttl: Duration::from_secs(30), + descriptor_cache_ttl: Duration::from_secs(60), + descriptor_cache_max_bytes: 64 * 1024 * 1024, + descriptor_cache_max_entry_bytes: 4 * 1024 * 1024, + global_index_thread_num: 1, + } + } +} + +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct TableCacheStats { + pub hits: u64, + pub misses: u64, + pub entries: usize, +} + +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct DescriptorCacheStats { + pub hits: u64, + pub misses: u64, + pub entries: u64, + pub weighted_bytes: u64, +} + +#[derive(Debug, Clone)] +struct CachedTable { + table: Table, + loaded_at: Instant, +} + +#[derive(Debug, Default)] +struct TableCacheCounters { + hits: AtomicU64, + misses: AtomicU64, +} + +#[derive(Debug, Default)] +struct DescriptorCacheCounters { + hits: AtomicU64, + misses: AtomicU64, +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +struct DescriptorCacheKey { + table: TableRef, + table_location: String, + snapshot_id: i64, + schema_id: i64, + snapshot_fingerprint: Vec, + request_fingerprint: Vec, +} + +#[derive(Debug)] +struct CachedDescriptorResponse { + response: BatchGetResponse, + weight: u32, +} + +type SharedLookupResult = Arc>>; + +#[derive(Debug, Default)] +struct DescriptorLoad { + result: Mutex>, + notify: Notify, +} + +impl DescriptorLoad { + fn complete(&self, result: SharedLookupResult) -> Result<()> { + *self.result.lock().map_err(|_| cache_lock_error())? = Some(result); + self.notify.notify_waiters(); + Ok(()) + } + + async fn wait(&self) -> Result { + loop { + let notified = self.notify.notified(); + if let Some(result) = self.result.lock().map_err(|_| cache_lock_error())?.clone() { + return Ok(result); + } + notified.await; + } + } +} + +struct DescriptorLoadGuard { + load: Arc, + completed: bool, +} + +impl DescriptorLoadGuard { + fn new(load: Arc) -> Self { + Self { + load, + completed: false, + } + } + + fn complete(mut self, result: SharedLookupResult) -> Result<()> { + self.load.complete(result)?; + self.completed = true; + Ok(()) + } +} + +impl Drop for DescriptorLoadGuard { + fn drop(&mut self) { + if !self.completed { + let _ = self + .load + .complete(Arc::new(Err(Arc::new(LookupError::LoadCancelled)))); + } + } +} + +#[derive(Clone)] +pub struct BlobLookupService { + catalog: Arc, + policies: Arc>, + options: BlobLookupOptions, + table_cache: Arc>>, + table_load_gates: Arc>>>>, + table_cache_counters: Arc, + descriptor_cache: Option>>, + descriptor_loads: Arc>>>, + descriptor_cache_counters: Arc, +} + +impl std::fmt::Debug for BlobLookupService { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("BlobLookupService") + .field("policies", &self.policies.keys().collect::>()) + .field("options", &self.options) + .field("table_cache_stats", &self.table_cache_stats()) + .field("descriptor_cache_stats", &self.descriptor_cache_stats()) + .finish_non_exhaustive() + } +} + +impl BlobLookupService { + pub fn new( + catalog: Arc, + policies: impl IntoIterator, + ) -> Result { + Self::new_with_options(catalog, policies, BlobLookupOptions::default()) + } + + pub fn new_with_options( + catalog: Arc, + policies: impl IntoIterator, + options: BlobLookupOptions, + ) -> Result { + if options.global_index_thread_num == 0 { + return Err(LookupError::InvalidPolicy( + "global index thread count must be positive".to_string(), + )); + } + let mut by_table = BTreeMap::new(); + for policy in policies { + policy.validate_definition()?; + let table = policy.table.clone(); + if by_table.insert(table.clone(), policy).is_some() { + return Err(LookupError::InvalidPolicy(format!( + "duplicate policy for {}", + table.full_name() + ))); + } + } + let descriptor_cache = if options.descriptor_cache_ttl.is_zero() + || options.descriptor_cache_max_bytes == 0 + || options.descriptor_cache_max_entry_bytes == 0 + { + None + } else { + Some( + Cache::builder() + .max_capacity(options.descriptor_cache_max_bytes) + .time_to_live(options.descriptor_cache_ttl) + .weigher( + |_key: &DescriptorCacheKey, value: &Arc| { + value.weight + }, + ) + .build(), + ) + }; + Ok(Self { + catalog, + policies: Arc::new(by_table), + options, + table_cache: Arc::new(RwLock::new(HashMap::new())), + table_load_gates: Arc::new(Mutex::new(HashMap::new())), + table_cache_counters: Arc::new(TableCacheCounters::default()), + descriptor_cache, + descriptor_loads: Arc::new(Mutex::new(HashMap::new())), + descriptor_cache_counters: Arc::new(DescriptorCacheCounters::default()), + }) + } + + pub fn table_cache_stats(&self) -> TableCacheStats { + TableCacheStats { + hits: self.table_cache_counters.hits.load(Ordering::Relaxed), + misses: self.table_cache_counters.misses.load(Ordering::Relaxed), + entries: self + .table_cache + .read() + .map(|cache| { + cache + .values() + .filter(|cached| cached.loaded_at.elapsed() < self.options.table_cache_ttl) + .count() + }) + .unwrap_or_default(), + } + } + + pub fn descriptor_cache_stats(&self) -> DescriptorCacheStats { + match &self.descriptor_cache { + Some(cache) => DescriptorCacheStats { + hits: self.descriptor_cache_counters.hits.load(Ordering::Relaxed), + misses: self + .descriptor_cache_counters + .misses + .load(Ordering::Relaxed), + entries: cache.entry_count(), + weighted_bytes: cache.weighted_size(), + }, + None => DescriptorCacheStats::default(), + } + } + + /// Reload and validate every configured table. This intentionally bypasses + /// the TTL so readiness checks also verify catalog connectivity. + pub async fn check_ready(&self) -> Result<()> { + const MAX_CONCURRENT_READINESS_CHECKS: usize = 8; + futures::stream::iter(self.policies.values().cloned()) + .map(|policy| async move { + let table = self.catalog.get_table(&policy.table.identifier()).await?; + policy.validate_table(&table)?; + self.cache_table(policy.table.clone(), table)?; + Ok::<(), LookupError>(()) + }) + .buffer_unordered(MAX_CONCURRENT_READINESS_CHECKS) + .try_collect::>() + .await?; + Ok(()) + } + + pub async fn batch_get(&self, request: BatchGetRequest) -> Result { + let policy = self.policies.get(&request.table).ok_or_else(|| { + LookupError::InvalidRequest(format!( + "no lookup policy is configured for {}", + request.table.full_name() + )) + })?; + self.batch_get_with_policy(policy, request).await + } + + async fn batch_get_with_policy( + &self, + policy: &TableLookupPolicy, + request: BatchGetRequest, + ) -> Result { + validate_request(policy, &request)?; + let base = self.get_table(&request.table).await?; + + let snapshot_id = match request.snapshot_id { + Some(snapshot_id) if snapshot_id > 0 => Some(snapshot_id), + Some(snapshot_id) => { + return Err(LookupError::InvalidRequest(format!( + "snapshotId must be positive, got {snapshot_id}" + ))) + } + None => base.snapshot_manager().get_latest_snapshot_id().await?, + }; + + let Some(snapshot_id) = snapshot_id else { + policy.validate_table(&base)?; + let prepared = prepare_keys(request.keys, &policy.key_fields, base.schema().fields())?; + return Ok(not_found_response( + request.table, + None, + base.schema().id(), + prepared, + )); + }; + + let table = base + .copy_with_time_travel_strict(HashMap::from([ + ("scan.snapshot-id".to_string(), snapshot_id.to_string()), + ("blob-as-descriptor".to_string(), "true".to_string()), + ("global-index.search-mode".to_string(), "full".to_string()), + ( + "global-index.thread-num".to_string(), + self.options.global_index_thread_num.to_string(), + ), + ])) + .await?; + policy.validate_table_for_request(&table, &request.blob_fields)?; + + let cache = self.descriptor_cache.clone(); + let cache_key = descriptor_cache_key(&request, snapshot_id, &table)?; + if let Some(cache) = &cache { + if let Some(response) = cache.get(&cache_key).await { + self.descriptor_cache_counters + .hits + .fetch_add(1, Ordering::Relaxed); + return Ok(as_cache_hit(response.response.clone())); + } + } + + let (load, leader) = self.descriptor_load(&cache_key)?; + if !leader { + if cache.is_some() { + self.descriptor_cache_counters + .hits + .fetch_add(1, Ordering::Relaxed); + } + return shared_lookup_result(load.wait().await?, true); + } + let guard = DescriptorLoadGuard::new(load); + if cache.is_some() { + self.descriptor_cache_counters + .misses + .fetch_add(1, Ordering::Relaxed); + } + match execute_snapshot_lookup(policy, request, snapshot_id, table).await { + Ok(response) => { + if let Some(cache) = cache { + let weight = descriptor_cache_weight(&cache_key, &response); + if u64::from(weight) <= self.options.descriptor_cache_max_entry_bytes { + cache + .insert( + cache_key, + Arc::new(CachedDescriptorResponse { + response: response.clone(), + weight, + }), + ) + .await; + } + } + guard.complete(Arc::new(Ok(response.clone())))?; + Ok(response) + } + Err(error) => { + let shared_error = Arc::new(clone_lookup_error_for_follower(&error)); + guard.complete(Arc::new(Err(shared_error)))?; + Err(error) + } + } + } + + fn descriptor_load( + &self, + cache_key: &DescriptorCacheKey, + ) -> Result<(Arc, bool)> { + let mut loads = self + .descriptor_loads + .lock() + .map_err(|_| cache_lock_error())?; + loads.retain(|_, load| load.strong_count() > 0); + if let Some(load) = loads.get(cache_key).and_then(Weak::upgrade) { + return Ok((load, false)); + } + let load = Arc::new(DescriptorLoad::default()); + loads.insert(cache_key.clone(), Arc::downgrade(&load)); + Ok((load, true)) + } + + async fn get_table(&self, table_ref: &TableRef) -> Result { + if self.options.table_cache_ttl.is_zero() { + self.table_cache_counters + .misses + .fetch_add(1, Ordering::Relaxed); + return Ok(self.catalog.get_table(&table_ref.identifier()).await?); + } + if let Some(table) = self.cached_table(table_ref)? { + self.table_cache_counters + .hits + .fetch_add(1, Ordering::Relaxed); + return Ok(table); + } + + // Collapse concurrent cold loads for the same table into one catalog + // request. Different tables use different gates and can load in parallel. + let gate = self.table_load_gate(table_ref)?; + let _guard = gate.lock().await; + if let Some(table) = self.cached_table(table_ref)? { + self.table_cache_counters + .hits + .fetch_add(1, Ordering::Relaxed); + return Ok(table); + } + + self.table_cache_counters + .misses + .fetch_add(1, Ordering::Relaxed); + let table = self.catalog.get_table(&table_ref.identifier()).await?; + self.cache_table(table_ref.clone(), table.clone())?; + Ok(table) + } + + fn cached_table(&self, table_ref: &TableRef) -> Result> { + Ok(self + .table_cache + .read() + .map_err(|_| cache_lock_error())? + .get(table_ref) + .filter(|cached| cached.loaded_at.elapsed() < self.options.table_cache_ttl) + .map(|cached| cached.table.clone())) + } + + fn table_load_gate(&self, table_ref: &TableRef) -> Result>> { + Ok(self + .table_load_gates + .lock() + .map_err(|_| cache_lock_error())? + .entry(table_ref.clone()) + .or_default() + .clone()) + } + + fn cache_table(&self, table_ref: TableRef, table: Table) -> Result<()> { + if self.options.table_cache_ttl.is_zero() { + return Ok(()); + } + self.table_cache + .write() + .map_err(|_| cache_lock_error())? + .insert( + table_ref, + CachedTable { + table, + loaded_at: Instant::now(), + }, + ); + Ok(()) + } +} + +async fn execute_snapshot_lookup( + policy: &TableLookupPolicy, + request: BatchGetRequest, + snapshot_id: i64, + table: Table, +) -> Result { + let prepared = prepare_keys(request.keys, &policy.key_fields, table.schema().fields())?; + let predicate = build_batch_predicate(&prepared, &policy.key_fields, table.schema().fields())?; + let mut projection = policy.key_fields.clone(); + projection.extend(request.blob_fields.iter().cloned()); + let projection_refs = projection.iter().map(String::as_str).collect::>(); + + let mut builder = table.new_read_builder(); + builder + .with_projection(&projection_refs)? + .with_filter(predicate); + let (plan, trace) = builder.new_scan().plan_with_trace().await?; + if trace.snapshot_id != Some(snapshot_id) { + return Err(LookupError::SnapshotMismatch { + expected: snapshot_id, + actual: trace.snapshot_id, + }); + } + enforce_budget(policy, trace.final_files, trace.planned_data_file_bytes)?; + + let read = builder.new_read()?; + let mut stream = read.to_arrow(plan.splits())?; + let mut matches = HashMap::::new(); + while let Some(batch) = stream.try_next().await? { + collect_batch( + &batch, + &prepared, + &policy.key_fields, + &request.blob_fields, + table.schema().fields(), + request.descriptor_format, + &mut matches, + )?; + } + + Ok(BatchGetResponse { + table: request.table, + snapshot_id: Some(snapshot_id), + schema_id: table.schema().id(), + cache_hit: false, + scan: LookupScanStats { + planned_files: trace.final_files, + planned_bytes: trace.planned_data_file_bytes, + }, + results: build_results(prepared, matches), + }) +} + +fn descriptor_cache_key( + request: &BatchGetRequest, + snapshot_id: i64, + table: &Table, +) -> Result { + let snapshot = table.travel_snapshot().ok_or_else(|| { + LookupError::UnexpectedResult(format!( + "snapshot {snapshot_id} was not resolved for descriptor lookup" + )) + })?; + if snapshot.id() != snapshot_id { + return Err(LookupError::SnapshotMismatch { + expected: snapshot_id, + actual: Some(snapshot.id()), + }); + } + let snapshot_fingerprint = serde_json::to_vec(&( + snapshot.id(), + snapshot.schema_id(), + snapshot.base_manifest_list(), + snapshot.delta_manifest_list(), + snapshot.changelog_manifest_list(), + snapshot.index_manifest(), + snapshot.commit_user(), + snapshot.commit_identifier(), + snapshot.time_millis(), + snapshot.statistics(), + snapshot.next_row_id(), + )) + .map_err(|error| { + LookupError::UnexpectedResult(format!("failed to encode snapshot cache identity: {error}")) + })?; + let request_fingerprint = serde_json::to_vec(&( + &request.keys, + &request.blob_fields, + request.descriptor_format, + )) + .map_err(|error| { + LookupError::UnexpectedResult(format!("failed to encode descriptor cache key: {error}")) + })?; + Ok(DescriptorCacheKey { + table: request.table.clone(), + table_location: table.location().to_string(), + snapshot_id, + schema_id: table.schema().id(), + snapshot_fingerprint, + request_fingerprint, + }) +} + +fn as_cache_hit(mut response: BatchGetResponse) -> BatchGetResponse { + response.cache_hit = true; + response +} + +fn shared_lookup_result(result: SharedLookupResult, cache_hit: bool) -> Result { + match result.as_ref() { + Ok(response) if cache_hit => Ok(as_cache_hit(response.clone())), + Ok(response) => Ok(response.clone()), + Err(error) => Err(LookupError::Shared(error.clone())), + } +} + +fn clone_lookup_error_for_follower(error: &LookupError) -> LookupError { + match error { + LookupError::InvalidRequest(message) => LookupError::InvalidRequest(message.clone()), + LookupError::InvalidPolicy(message) => LookupError::InvalidPolicy(message.clone()), + LookupError::UnsupportedKeyType { field, data_type } => LookupError::UnsupportedKeyType { + field: field.clone(), + data_type: data_type.clone(), + }, + LookupError::InvalidKeyValue { field, message } => LookupError::InvalidKeyValue { + field: field.clone(), + message: message.clone(), + }, + LookupError::QueryBudgetExceeded { + files, + bytes, + max_files, + max_bytes, + } => LookupError::QueryBudgetExceeded { + files: *files, + bytes: *bytes, + max_files: *max_files, + max_bytes: *max_bytes, + }, + LookupError::SnapshotMismatch { expected, actual } => LookupError::SnapshotMismatch { + expected: *expected, + actual: *actual, + }, + LookupError::InvalidDescriptor { field, message } => LookupError::InvalidDescriptor { + field: field.clone(), + message: message.clone(), + }, + LookupError::UnexpectedResult(message) => LookupError::UnexpectedResult(message.clone()), + LookupError::LoadCancelled => LookupError::LoadCancelled, + LookupError::Shared(error) => LookupError::Shared(error.clone()), + LookupError::Paimon(paimon::Error::TableNotExist { full_name }) => { + LookupError::Paimon(paimon::Error::TableNotExist { + full_name: full_name.clone(), + }) + } + LookupError::Paimon(paimon::Error::DatabaseNotExist { database }) => { + LookupError::Paimon(paimon::Error::DatabaseNotExist { + database: database.clone(), + }) + } + LookupError::Paimon(paimon::Error::SnapshotNotExist { snapshot_id }) => { + LookupError::Paimon(paimon::Error::SnapshotNotExist { + snapshot_id: *snapshot_id, + }) + } + LookupError::Paimon(_) | LookupError::PaimonUnavailable => LookupError::PaimonUnavailable, + } +} + +fn descriptor_cache_weight(key: &DescriptorCacheKey, response: &BatchGetResponse) -> u32 { + let response_bytes = serde_json::to_vec(response) + .map(|value| value.len()) + .unwrap_or(usize::MAX); + let total = response_bytes + .saturating_add(key.request_fingerprint.len()) + .saturating_add(key.table.database.len()) + .saturating_add(key.table.table.len()) + .saturating_add(key.table_location.len()) + .saturating_add(key.snapshot_fingerprint.len()); + u32::try_from(total).unwrap_or(u32::MAX).max(1) +} + +fn cache_lock_error() -> LookupError { + LookupError::UnexpectedResult("lookup cache lock was poisoned".to_string()) +} + +#[derive(Debug)] +struct MatchedRow { + count: usize, + blobs: BTreeMap>, +} + +fn validate_request(policy: &TableLookupPolicy, request: &BatchGetRequest) -> Result<()> { + if request.keys.is_empty() { + return Err(LookupError::InvalidRequest( + "keys must not be empty".to_string(), + )); + } + if request.keys.len() > policy.budget.max_batch_keys { + return Err(LookupError::InvalidRequest(format!( + "batch has {} keys, maximum is {}", + request.keys.len(), + policy.budget.max_batch_keys + ))); + } + if request.blob_fields.is_empty() { + return Err(LookupError::InvalidRequest( + "blobFields must not be empty".to_string(), + )); + } + let mut seen = std::collections::HashSet::new(); + for field in &request.blob_fields { + if !seen.insert(field) { + return Err(LookupError::InvalidRequest(format!( + "blobFields contains duplicate field '{field}'" + ))); + } + if !policy.blob_fields.contains(field) { + return Err(LookupError::InvalidRequest(format!( + "BLOB field '{field}' is not allowed for {}", + request.table.full_name() + ))); + } + } + Ok(()) +} + +fn enforce_budget(policy: &TableLookupPolicy, files: usize, bytes: u64) -> Result<()> { + if files > policy.budget.max_planned_files || bytes > policy.budget.max_planned_bytes { + return Err(LookupError::QueryBudgetExceeded { + files, + bytes, + max_files: policy.budget.max_planned_files, + max_bytes: policy.budget.max_planned_bytes, + }); + } + Ok(()) +} + +fn collect_batch( + batch: &RecordBatch, + prepared: &[PreparedKey], + key_fields: &[String], + blob_fields: &[String], + schema_fields: &[paimon::spec::DataField], + format: DescriptorFormat, + matches: &mut HashMap, +) -> Result<()> { + let requested = prepared + .iter() + .map(|key| &key.normalized) + .collect::>(); + for row in 0..batch.num_rows() { + let key = normalized_key_from_batch(batch, row, key_fields, schema_fields)?; + if !requested.contains(&key) { + return Err(LookupError::UnexpectedResult( + "reader returned a row outside the requested key set".to_string(), + )); + } + let blobs = decode_descriptors(batch, row, key_fields.len(), blob_fields, format)?; + matches + .entry(key) + .and_modify(|matched| { + matched.count += 1; + matched.blobs.clear(); + }) + .or_insert(MatchedRow { count: 1, blobs }); + } + Ok(()) +} + +fn decode_descriptors( + batch: &RecordBatch, + row: usize, + blob_start: usize, + blob_fields: &[String], + format: DescriptorFormat, +) -> Result>> { + blob_fields + .iter() + .enumerate() + .map(|(index, field)| { + let array = batch + .column(blob_start + index) + .as_any() + .downcast_ref::() + .ok_or_else(|| { + LookupError::UnexpectedResult(format!( + "BLOB field '{field}' did not produce an Arrow BinaryArray" + )) + })?; + if array.is_null(row) { + return Ok((field.clone(), None)); + } + let raw = array.value(row); + let descriptor = BlobDescriptor::deserialize(raw).map_err(|error| { + LookupError::InvalidDescriptor { + field: field.clone(), + message: error.to_string(), + } + })?; + descriptor + .validate() + .map_err(|error| LookupError::InvalidDescriptor { + field: field.clone(), + message: error.to_string(), + })?; + let encoded = (format == DescriptorFormat::PaimonBase64) + .then(|| base64::engine::general_purpose::STANDARD.encode(raw)); + Ok(( + field.clone(), + Some(BlobDescriptorDto { + version: descriptor.version(), + uri: descriptor.uri().to_string(), + offset: descriptor.offset(), + length: descriptor.length(), + encoded, + }), + )) + }) + .collect() +} + +fn build_results( + prepared: Vec, + matches: HashMap, +) -> Vec { + prepared + .into_iter() + .map(|key| match matches.get(&key.normalized) { + None => LookupResult { + key: key.original, + status: LookupStatus::NotFound, + blobs: BTreeMap::new(), + }, + Some(matched) if matched.count == 1 => LookupResult { + key: key.original, + status: LookupStatus::Found, + blobs: matched.blobs.clone(), + }, + Some(_) => LookupResult { + key: key.original, + status: LookupStatus::NonUnique, + blobs: BTreeMap::new(), + }, + }) + .collect() +} + +fn not_found_response( + table: TableRef, + snapshot_id: Option, + schema_id: i64, + prepared: Vec, +) -> BatchGetResponse { + BatchGetResponse { + table, + snapshot_id, + schema_id, + cache_hit: false, + scan: LookupScanStats::default(), + results: prepared + .into_iter() + .map(|key| LookupResult { + key: key.original, + status: LookupStatus::NotFound, + blobs: BTreeMap::new(), + }) + .collect(), + } +} diff --git a/crates/query-service/src/model.rs b/crates/query-service/src/model.rs new file mode 100644 index 00000000..624eca14 --- /dev/null +++ b/crates/query-service/src/model.rs @@ -0,0 +1,114 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::BTreeMap; + +use paimon::catalog::Identifier; +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct TableRef { + pub database: String, + pub table: String, +} + +impl TableRef { + pub fn new(database: impl Into, table: impl Into) -> Self { + Self { + database: database.into(), + table: table.into(), + } + } + + pub(crate) fn identifier(&self) -> Identifier { + Identifier::new(&self.database, &self.table) + } + + pub(crate) fn full_name(&self) -> String { + format!("{}.{}", self.database, self.table) + } +} + +pub type LookupKey = BTreeMap; + +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "SCREAMING_SNAKE_CASE")] +pub enum DescriptorFormat { + #[default] + Json, + PaimonBase64, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct BatchGetRequest { + pub table: TableRef, + pub keys: Vec, + pub blob_fields: Vec, + #[serde(default)] + pub snapshot_id: Option, + #[serde(default)] + pub descriptor_format: DescriptorFormat, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "SCREAMING_SNAKE_CASE")] +pub enum LookupStatus { + Found, + NotFound, + NonUnique, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct BlobDescriptorDto { + pub version: u8, + pub uri: String, + pub offset: i64, + pub length: i64, + #[serde(skip_serializing_if = "Option::is_none")] + pub encoded: Option, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct LookupResult { + pub key: LookupKey, + pub status: LookupStatus, + pub blobs: BTreeMap>, +} + +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct LookupScanStats { + pub planned_files: usize, + pub planned_bytes: u64, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct BatchGetResponse { + pub table: TableRef, + pub snapshot_id: Option, + pub schema_id: i64, + #[serde(default)] + pub cache_hit: bool, + #[serde(default)] + pub scan: LookupScanStats, + pub results: Vec, +} diff --git a/crates/query-service/src/policy.rs b/crates/query-service/src/policy.rs new file mode 100644 index 00000000..365caf6f --- /dev/null +++ b/crates/query-service/src/policy.rs @@ -0,0 +1,194 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::BTreeSet; + +use paimon::table::Table; +use serde::{Deserialize, Serialize}; + +use crate::error::{LookupError, Result}; +use crate::key::{supports_global_btree, validate_key_type}; +use crate::model::TableRef; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "SCREAMING_SNAKE_CASE")] +pub enum LookupStrategy { + PrimaryKey, + GlobalBtree, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct QueryBudget { + pub max_batch_keys: usize, + pub max_planned_files: usize, + pub max_planned_bytes: u64, +} + +impl Default for QueryBudget { + fn default() -> Self { + Self { + max_batch_keys: 200, + max_planned_files: 16, + max_planned_bytes: 128 * 1024 * 1024, + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct TableLookupPolicy { + pub table: TableRef, + pub key_fields: Vec, + pub blob_fields: BTreeSet, + pub strategy: LookupStrategy, + #[serde(default)] + pub budget: QueryBudget, +} + +impl TableLookupPolicy { + pub fn validate_definition(&self) -> Result<()> { + if self.key_fields.is_empty() { + return Err(LookupError::InvalidPolicy(format!( + "{} has no key fields", + self.table.full_name() + ))); + } + if self.key_fields.iter().collect::>().len() != self.key_fields.len() { + return Err(LookupError::InvalidPolicy(format!( + "{} has duplicate key fields", + self.table.full_name() + ))); + } + if self.blob_fields.is_empty() { + return Err(LookupError::InvalidPolicy(format!( + "{} has no allowed BLOB fields", + self.table.full_name() + ))); + } + if self.budget.max_batch_keys == 0 + || self.budget.max_planned_files == 0 + || self.budget.max_planned_bytes == 0 + { + return Err(LookupError::InvalidPolicy(format!( + "{} has a zero query budget", + self.table.full_name() + ))); + } + Ok(()) + } + + pub(crate) fn validate_table(&self, table: &Table) -> Result<()> { + self.validate_table_for_blob_fields( + table, + self.blob_fields.iter().map(String::as_str), + false, + ) + } + + pub(crate) fn validate_table_for_request( + &self, + table: &Table, + blob_fields: &[String], + ) -> Result<()> { + self.validate_table_for_blob_fields(table, blob_fields.iter().map(String::as_str), true) + } + + fn validate_table_for_blob_fields<'a>( + &self, + table: &Table, + blob_fields: impl IntoIterator, + request_context: bool, + ) -> Result<()> { + self.validate_definition()?; + let schema = table.schema(); + for key in &self.key_fields { + let field = schema + .fields() + .iter() + .find(|field| field.name() == key) + .ok_or_else(|| { + LookupError::InvalidPolicy(format!( + "lookup key field '{key}' does not exist in {}", + self.table.full_name() + )) + })?; + validate_key_type(key, field.data_type())?; + if self.strategy == LookupStrategy::GlobalBtree + && !supports_global_btree(field.data_type()) + { + return Err(LookupError::InvalidPolicy(format!( + "GLOBAL_BTREE lookup key field '{key}' in {} has a type that cannot be indexed by a sorted global index: {:?}", + self.table.full_name(), + field.data_type() + ))); + } + } + for blob in blob_fields { + let Some(field) = schema.fields().iter().find(|field| field.name() == blob) else { + let message = format!( + "BLOB field '{blob}' does not exist in the selected snapshot of {}", + self.table.full_name() + ); + return Err(if request_context { + LookupError::InvalidRequest(message) + } else { + LookupError::InvalidPolicy(message) + }); + }; + if !field.data_type().is_blob_type() { + let message = format!( + "field '{blob}' in the selected snapshot of {} is not a scalar BLOB", + self.table.full_name() + ); + return Err(if request_context { + LookupError::InvalidRequest(message) + } else { + LookupError::InvalidPolicy(message) + }); + } + } + + match self.strategy { + LookupStrategy::PrimaryKey => { + let configured = self.key_fields.iter().collect::>(); + let primary = schema.primary_keys().iter().collect::>(); + if primary.is_empty() || configured != primary { + return Err(LookupError::InvalidPolicy(format!( + "{} PRIMARY_KEY lookup must configure the complete primary key {:?}", + self.table.full_name(), + schema.primary_keys() + ))); + } + } + LookupStrategy::GlobalBtree => { + let options = schema.core_options(); + if !options.data_evolution_enabled() + || !options.row_tracking_enabled() + || !options.global_index_enabled() + || options.deletion_vectors_enabled() + || !schema.primary_keys().is_empty() + { + return Err(LookupError::InvalidPolicy(format!( + "{} GLOBAL_BTREE lookup requires a row-tracking data-evolution table with global indexes enabled, no primary key, and no deletion vectors", + self.table.full_name() + ))); + } + } + } + Ok(()) + } +} diff --git a/crates/query-service/tests/lookup.rs b/crates/query-service/tests/lookup.rs new file mode 100644 index 00000000..8a3f7582 --- /dev/null +++ b/crates/query-service/tests/lookup.rs @@ -0,0 +1,656 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::{BTreeMap, BTreeSet}; +use std::sync::Arc; +use std::time::Duration; + +use arrow_array::{BinaryArray, Int32Array, RecordBatch}; +use arrow_schema::{DataType as ArrowDataType, Field as ArrowField, Schema as ArrowSchema}; +use base64::Engine; +use futures::future::join_all; +use paimon::catalog::{Catalog, Identifier}; +use paimon::spec::{ + BlobDescriptor, BlobType, DataType, IntType, Schema, SchemaChange, VarBinaryType, +}; +use paimon::{CatalogOptions, FileSystemCatalog, Options}; +use paimon_query_service::{ + BatchGetRequest, BlobLookupOptions, BlobLookupService, DescriptorFormat, LookupError, + LookupStatus, LookupStrategy, QueryBudget, TableLookupPolicy, TableRef, +}; +use serde_json::json; +use tempfile::TempDir; + +struct Fixture { + _warehouse: TempDir, + catalog: Arc, + table: TableRef, +} + +impl Fixture { + async fn new(descriptor_field: bool) -> Self { + let warehouse = TempDir::new().unwrap(); + let mut options = Options::new(); + options.set( + CatalogOptions::WAREHOUSE, + warehouse.path().to_str().unwrap(), + ); + let catalog = Arc::new(FileSystemCatalog::new(options).unwrap()); + catalog + .create_database("db", false, Default::default()) + .await + .unwrap(); + + let mut schema = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column("picture", DataType::Blob(BlobType::new())) + .option("data-evolution.enabled", "true") + .option("row-tracking.enabled", "true") + .option("global-index.enabled", "true"); + if descriptor_field { + schema = schema.option("blob-descriptor-field", "picture"); + } + catalog + .create_table( + &Identifier::new("db", "assets"), + schema.build().unwrap(), + false, + ) + .await + .unwrap(); + + Self { + _warehouse: warehouse, + catalog, + table: TableRef::new("db", "assets"), + } + } + + fn service(&self) -> BlobLookupService { + self.service_with_budget(QueryBudget { + max_batch_keys: 10, + max_planned_files: 10, + max_planned_bytes: 16 * 1024 * 1024, + }) + } + + fn service_with_budget(&self, budget: QueryBudget) -> BlobLookupService { + self.service_with_options(budget, BlobLookupOptions::default()) + } + + fn service_with_options( + &self, + budget: QueryBudget, + options: BlobLookupOptions, + ) -> BlobLookupService { + BlobLookupService::new_with_options( + self.catalog.clone(), + [TableLookupPolicy { + table: self.table.clone(), + key_fields: vec!["id".to_string()], + blob_fields: BTreeSet::from(["picture".to_string()]), + strategy: LookupStrategy::GlobalBtree, + budget, + }], + options, + ) + .unwrap() + } + + async fn append(&self, ids: Vec, pictures: Vec>>) -> i64 { + let picture_refs = pictures + .iter() + .map(|value| value.as_deref()) + .collect::>(); + let batch = RecordBatch::try_new( + Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", ArrowDataType::Int32, true), + ArrowField::new("picture", ArrowDataType::Binary, true), + ])), + vec![ + Arc::new(Int32Array::from(ids)), + Arc::new(BinaryArray::from(picture_refs)), + ], + ) + .unwrap(); + + let table = self + .catalog + .get_table(&Identifier::new("db", "assets")) + .await + .unwrap(); + let builder = table.new_write_builder(); + let mut writer = builder.new_write().unwrap(); + writer.write_arrow_batch(&batch).await.unwrap(); + let messages = writer.prepare_commit().await.unwrap(); + builder.new_commit().commit(messages).await.unwrap(); + table + .snapshot_manager() + .get_latest_snapshot_id() + .await + .unwrap() + .unwrap() + } + + fn request(&self, ids: &[i32]) -> BatchGetRequest { + BatchGetRequest { + table: self.table.clone(), + keys: ids + .iter() + .map(|id| BTreeMap::from([("id".to_string(), json!(id))])) + .collect(), + blob_fields: vec!["picture".to_string()], + snapshot_id: None, + descriptor_format: DescriptorFormat::PaimonBase64, + } + } +} + +#[tokio::test] +async fn batch_get_returns_descriptors_and_preserves_request_order() { + let fixture = Fixture::new(false).await; + fixture + .append( + vec![1, 2, 3], + vec![Some(b"first".to_vec()), Some(b"second".to_vec()), None], + ) + .await; + + let response = fixture + .service() + .batch_get(fixture.request(&[2, 99, 1, 3])) + .await + .unwrap(); + + assert_eq!(response.results.len(), 4); + assert!(response.scan.planned_files >= 1); + assert!(response.scan.planned_bytes > 0); + assert_eq!(response.results[0].status, LookupStatus::Found); + assert_eq!(response.results[1].status, LookupStatus::NotFound); + assert_eq!(response.results[2].status, LookupStatus::Found); + assert_eq!(response.results[3].status, LookupStatus::Found); + assert_eq!(response.results[0].key["id"], json!(2)); + + let descriptor = response.results[0].blobs["picture"].as_ref().unwrap(); + assert_eq!(descriptor.length, 6); + assert!(descriptor.uri.ends_with(".blob")); + let encoded = descriptor.encoded.as_ref().unwrap(); + let raw = base64::engine::general_purpose::STANDARD + .decode(encoded) + .unwrap(); + let decoded = BlobDescriptor::deserialize(&raw).unwrap(); + assert_eq!(decoded.length(), 6); + + assert_eq!(response.results[3].blobs["picture"], None); +} + +#[tokio::test] +async fn explicit_snapshot_isolated_from_later_duplicate_key() { + let fixture = Fixture::new(false).await; + let first_snapshot = fixture.append(vec![1], vec![Some(b"old".to_vec())]).await; + fixture.append(vec![1], vec![Some(b"new".to_vec())]).await; + + let service = fixture.service(); + let latest = service.batch_get(fixture.request(&[1])).await.unwrap(); + assert!(!latest.cache_hit); + assert_eq!(latest.results[0].status, LookupStatus::NonUnique); + + let mut historical_request = fixture.request(&[1]); + historical_request.snapshot_id = Some(first_snapshot); + let historical = service.batch_get(historical_request.clone()).await.unwrap(); + assert!(!historical.cache_hit); + assert_eq!(historical.snapshot_id, Some(first_snapshot)); + assert_eq!(historical.results[0].status, LookupStatus::Found); + assert_eq!( + historical.results[0].blobs["picture"] + .as_ref() + .unwrap() + .length, + 3 + ); + + let cached_historical = service.batch_get(historical_request).await.unwrap(); + assert!(cached_historical.cache_hit); + assert_eq!(cached_historical.snapshot_id, Some(first_snapshot)); + assert_eq!(service.descriptor_cache_stats().misses, 2); + assert_eq!(service.descriptor_cache_stats().hits, 1); +} + +#[tokio::test] +async fn cached_explicit_snapshot_is_rejected_after_snapshot_deletion() { + let fixture = Fixture::new(false).await; + let snapshot_id = fixture.append(vec![1], vec![Some(b"old".to_vec())]).await; + let service = fixture.service(); + let mut request = fixture.request(&[1]); + request.snapshot_id = Some(snapshot_id); + + service.batch_get(request.clone()).await.unwrap(); + let table = fixture + .catalog + .get_table(&Identifier::new("db", "assets")) + .await + .unwrap(); + table + .snapshot_manager() + .delete_snapshot(snapshot_id) + .await + .unwrap(); + + let error = service.batch_get(request).await.unwrap_err(); + assert!(matches!( + error, + LookupError::Paimon(paimon::Error::SnapshotNotExist { + snapshot_id: missing + }) if missing == snapshot_id + )); +} + +#[tokio::test] +async fn historical_lookup_validates_only_requested_blob_fields() { + let fixture = Fixture::new(false).await; + let snapshot_id = fixture.append(vec![1], vec![Some(b"old".to_vec())]).await; + fixture + .catalog + .alter_table( + &Identifier::new("db", "assets"), + vec![SchemaChange::add_column( + "thumbnail".to_string(), + DataType::Blob(BlobType::new()), + )], + false, + ) + .await + .unwrap(); + let service = BlobLookupService::new( + fixture.catalog.clone(), + [TableLookupPolicy { + table: fixture.table.clone(), + key_fields: vec!["id".to_string()], + blob_fields: BTreeSet::from(["picture".to_string(), "thumbnail".to_string()]), + strategy: LookupStrategy::GlobalBtree, + budget: QueryBudget { + max_batch_keys: 10, + max_planned_files: 10, + max_planned_bytes: 16 * 1024 * 1024, + }, + }], + ) + .unwrap(); + let mut request = fixture.request(&[1]); + request.snapshot_id = Some(snapshot_id); + + let response = service.batch_get(request.clone()).await.unwrap(); + assert_eq!(response.snapshot_id, Some(snapshot_id)); + assert_eq!(response.results[0].status, LookupStatus::Found); + + request.blob_fields = vec!["thumbnail".to_string()]; + let error = service.batch_get(request).await.unwrap_err(); + assert!(matches!( + error, + LookupError::InvalidRequest(message) + if message.contains("selected snapshot") && message.contains("thumbnail") + )); +} + +#[tokio::test] +async fn descriptor_cache_does_not_cross_drop_and_recreate() { + let fixture = Fixture::new(false).await; + fixture.append(vec![1], vec![Some(b"old".to_vec())]).await; + let service = fixture.service_with_options( + QueryBudget { + max_batch_keys: 10, + max_planned_files: 10, + max_planned_bytes: 16 * 1024 * 1024, + }, + BlobLookupOptions { + table_cache_ttl: Duration::ZERO, + ..BlobLookupOptions::default() + }, + ); + let first = service.batch_get(fixture.request(&[1])).await.unwrap(); + assert_eq!( + first.results[0].blobs["picture"].as_ref().unwrap().length, + 3 + ); + + let identifier = Identifier::new("db", "assets"); + fixture + .catalog + .drop_table(&identifier, false) + .await + .unwrap(); + fixture + .catalog + .create_table( + &identifier, + Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column("picture", DataType::Blob(BlobType::new())) + .option("data-evolution.enabled", "true") + .option("row-tracking.enabled", "true") + .option("global-index.enabled", "true") + .build() + .unwrap(), + false, + ) + .await + .unwrap(); + fixture + .append(vec![1], vec![Some(b"recreated".to_vec())]) + .await; + + let recreated = service.batch_get(fixture.request(&[1])).await.unwrap(); + assert!(!recreated.cache_hit); + assert_eq!( + recreated.results[0].blobs["picture"] + .as_ref() + .unwrap() + .length, + 9 + ); +} + +#[tokio::test] +async fn inline_descriptor_field_round_trips_without_blob_payload_read() { + let fixture = Fixture::new(true).await; + let expected = BlobDescriptor::new("s3://example/object".to_string(), 7, -1); + fixture + .append(vec![7], vec![Some(expected.serialize())]) + .await; + + let response = fixture + .service() + .batch_get(fixture.request(&[7])) + .await + .unwrap(); + let actual = response.results[0].blobs["picture"].as_ref().unwrap(); + assert_eq!(actual.uri, "s3://example/object"); + assert_eq!(actual.offset, 7); + assert_eq!(actual.length, -1); +} + +#[tokio::test] +async fn rejects_unindexed_fallback_that_exceeds_the_data_file_budget() { + let fixture = Fixture::new(false).await; + fixture + .append(vec![1], vec![Some(b"payload".to_vec())]) + .await; + let service = fixture.service_with_budget(QueryBudget { + max_batch_keys: 10, + max_planned_files: 10, + max_planned_bytes: 1, + }); + + let error = service.batch_get(fixture.request(&[1])).await.unwrap_err(); + assert!(matches!(error, LookupError::QueryBudgetExceeded { .. })); +} + +#[tokio::test] +async fn caches_table_metadata_and_can_refresh_it_for_readiness() { + let fixture = Fixture::new(false).await; + fixture + .append(vec![1], vec![Some(b"payload".to_vec())]) + .await; + let service = fixture.service(); + + let first = service.batch_get(fixture.request(&[1])).await.unwrap(); + let second = service.batch_get(fixture.request(&[1])).await.unwrap(); + assert!(!first.cache_hit); + assert!(second.cache_hit); + assert_eq!(service.table_cache_stats().hits, 1); + assert_eq!(service.table_cache_stats().misses, 1); + assert_eq!(service.table_cache_stats().entries, 1); + assert_eq!(service.descriptor_cache_stats().hits, 1); + assert_eq!(service.descriptor_cache_stats().misses, 1); + + service.check_ready().await.unwrap(); + assert_eq!(service.table_cache_stats().entries, 1); +} + +#[tokio::test] +async fn zero_ttl_disables_table_metadata_cache() { + let fixture = Fixture::new(false).await; + fixture + .append(vec![1], vec![Some(b"payload".to_vec())]) + .await; + let service = fixture.service_with_options( + QueryBudget { + max_batch_keys: 10, + max_planned_files: 10, + max_planned_bytes: 16 * 1024 * 1024, + }, + BlobLookupOptions { + table_cache_ttl: Duration::ZERO, + ..BlobLookupOptions::default() + }, + ); + + service.batch_get(fixture.request(&[1])).await.unwrap(); + service.batch_get(fixture.request(&[1])).await.unwrap(); + assert_eq!(service.table_cache_stats().hits, 0); + assert_eq!(service.table_cache_stats().misses, 2); + assert_eq!(service.table_cache_stats().entries, 0); +} + +#[tokio::test] +async fn zero_limits_disable_descriptor_cache() { + let fixture = Fixture::new(false).await; + fixture + .append(vec![1], vec![Some(b"payload".to_vec())]) + .await; + let service = fixture.service_with_options( + QueryBudget { + max_batch_keys: 10, + max_planned_files: 10, + max_planned_bytes: 16 * 1024 * 1024, + }, + BlobLookupOptions { + descriptor_cache_ttl: Duration::ZERO, + ..BlobLookupOptions::default() + }, + ); + + let first = service.batch_get(fixture.request(&[1])).await.unwrap(); + let second = service.batch_get(fixture.request(&[1])).await.unwrap(); + assert!(!first.cache_hit); + assert!(!second.cache_hit); + assert_eq!(service.descriptor_cache_stats(), Default::default()); +} + +#[tokio::test] +async fn disabled_long_term_cache_still_coalesces_concurrent_scans() { + let fixture = Fixture::new(false).await; + fixture + .append(vec![1], vec![Some(b"payload".to_vec())]) + .await; + let service = fixture.service_with_options( + QueryBudget { + max_batch_keys: 10, + max_planned_files: 10, + max_planned_bytes: 16 * 1024 * 1024, + }, + BlobLookupOptions { + descriptor_cache_ttl: Duration::ZERO, + ..BlobLookupOptions::default() + }, + ); + + let queries = (0..8).map(|_| { + let service = service.clone(); + let request = fixture.request(&[1]); + async move { service.batch_get(request).await.unwrap() } + }); + let responses = join_all(queries).await; + + assert_eq!(responses.len(), 8); + assert_eq!( + responses + .iter() + .filter(|response| response.cache_hit) + .count(), + 7 + ); + assert_eq!(service.descriptor_cache_stats(), Default::default()); +} + +#[tokio::test] +async fn oversized_descriptor_response_is_not_cached() { + let fixture = Fixture::new(false).await; + fixture + .append(vec![1], vec![Some(b"payload".to_vec())]) + .await; + let service = fixture.service_with_options( + QueryBudget { + max_batch_keys: 10, + max_planned_files: 10, + max_planned_bytes: 16 * 1024 * 1024, + }, + BlobLookupOptions { + descriptor_cache_max_entry_bytes: 1, + ..BlobLookupOptions::default() + }, + ); + + let first = service.batch_get(fixture.request(&[1])).await.unwrap(); + let second = service.batch_get(fixture.request(&[1])).await.unwrap(); + assert!(!first.cache_hit); + assert!(!second.cache_hit); + assert_eq!(service.descriptor_cache_stats().hits, 0); + assert_eq!(service.descriptor_cache_stats().misses, 2); +} + +#[tokio::test] +async fn concurrent_oversized_cache_misses_share_one_scan() { + let fixture = Fixture::new(false).await; + fixture + .append(vec![1], vec![Some(b"payload".to_vec())]) + .await; + let service = fixture.service_with_options( + QueryBudget { + max_batch_keys: 10, + max_planned_files: 10, + max_planned_bytes: 16 * 1024 * 1024, + }, + BlobLookupOptions { + descriptor_cache_max_entry_bytes: 1, + ..BlobLookupOptions::default() + }, + ); + + let queries = (0..8).map(|_| { + let service = service.clone(); + let request = fixture.request(&[1]); + async move { service.batch_get(request).await.unwrap() } + }); + let responses = join_all(queries).await; + + assert_eq!(responses.len(), 8); + assert_eq!(service.descriptor_cache_stats().misses, 1); + assert_eq!(service.descriptor_cache_stats().hits, 7); + assert_eq!(service.descriptor_cache_stats().entries, 0); +} + +#[tokio::test] +async fn collapses_concurrent_cold_table_metadata_loads() { + let fixture = Fixture::new(false).await; + fixture + .append(vec![1], vec![Some(b"payload".to_vec())]) + .await; + let service = fixture.service(); + + let queries = (0..8).map(|_| { + let service = service.clone(); + let request = fixture.request(&[1]); + async move { service.batch_get(request).await.unwrap() } + }); + let responses = join_all(queries).await; + assert_eq!(responses.len(), 8); + assert_eq!( + responses + .iter() + .filter(|response| response.cache_hit) + .count(), + 7 + ); + assert_eq!(service.table_cache_stats().misses, 1); + assert_eq!(service.table_cache_stats().hits, 7); + assert_eq!(service.descriptor_cache_stats().misses, 1); + assert_eq!(service.descriptor_cache_stats().hits, 7); +} + +#[tokio::test] +async fn rejects_empty_key_batches_before_scanning() { + let fixture = Fixture::new(false).await; + let error = fixture + .service() + .batch_get(fixture.request(&[])) + .await + .unwrap_err(); + assert!(matches!( + error, + LookupError::InvalidRequest(message) if message.contains("keys must not be empty") + )); +} + +#[tokio::test] +async fn readiness_rejects_global_btree_binary_key_fields() { + let warehouse = TempDir::new().unwrap(); + let mut options = Options::new(); + options.set( + CatalogOptions::WAREHOUSE, + warehouse.path().to_str().unwrap(), + ); + let catalog = Arc::new(FileSystemCatalog::new(options).unwrap()); + catalog + .create_database("db", false, Default::default()) + .await + .unwrap(); + catalog + .create_table( + &Identifier::new("db", "binary_assets"), + Schema::builder() + .column( + "asset_key", + DataType::VarBinary(VarBinaryType::new(32).unwrap()), + ) + .column("picture", DataType::Blob(BlobType::new())) + .option("data-evolution.enabled", "true") + .option("row-tracking.enabled", "true") + .option("global-index.enabled", "true") + .build() + .unwrap(), + false, + ) + .await + .unwrap(); + let service = BlobLookupService::new( + catalog, + [TableLookupPolicy { + table: TableRef::new("db", "binary_assets"), + key_fields: vec!["asset_key".to_string()], + blob_fields: BTreeSet::from(["picture".to_string()]), + strategy: LookupStrategy::GlobalBtree, + budget: QueryBudget::default(), + }], + ) + .unwrap(); + + let error = service.check_ready().await.unwrap_err(); + assert!(matches!( + error, + LookupError::InvalidPolicy(message) if message.contains("cannot be indexed") + )); +}