diff --git a/.github/scripts/rustfmt.sh b/.github/scripts/rustfmt.sh new file mode 100755 index 00000000..605b46db --- /dev/null +++ b/.github/scripts/rustfmt.sh @@ -0,0 +1,6 @@ +#!/usr/bin/env bash +set -euo pipefail + +readonly RUSTFMT_TOOLCHAIN="nightly-2025-08-07" + +exec cargo +"${RUSTFMT_TOOLCHAIN}" fmt "$@" diff --git a/.github/workflows/ci-dotnet.yml b/.github/workflows/ci-dotnet.yml index 85ee7922..33e0ad39 100644 --- a/.github/workflows/ci-dotnet.yml +++ b/.github/workflows/ci-dotnet.yml @@ -30,8 +30,9 @@ jobs: dotnet tool install -g dotnet-format echo "$HOME/.dotnet/tools" >> $GITHUB_PATH - - uses: dtolnay/rust-toolchain@29eef336d9b2848a0b548edc03f92a220660cdb8 # stable + - uses: dtolnay/rust-toolchain@29eef336d9b2848a0b548edc03f92a220660cdb8 # nightly with: + toolchain: nightly-2025-08-07 components: rustfmt - name: Format all files diff --git a/.github/workflows/ci-go.yml b/.github/workflows/ci-go.yml index fdbe2b82..313fa7af 100644 --- a/.github/workflows/ci-go.yml +++ b/.github/workflows/ci-go.yml @@ -40,8 +40,9 @@ jobs: with: go-version: '1.25' - - uses: dtolnay/rust-toolchain@29eef336d9b2848a0b548edc03f92a220660cdb8 # stable + - uses: dtolnay/rust-toolchain@29eef336d9b2848a0b548edc03f92a220660cdb8 # nightly with: + toolchain: nightly-2025-08-07 components: rustfmt - name: Format all files @@ -222,4 +223,4 @@ jobs: - name: Test Proto batch example builds working-directory: go/examples/proto/batch - run: go build -v \ No newline at end of file + run: go build -v diff --git a/.github/workflows/ci-python.yml b/.github/workflows/ci-python.yml index 556b402a..d027b234 100644 --- a/.github/workflows/ci-python.yml +++ b/.github/workflows/ci-python.yml @@ -101,8 +101,9 @@ jobs: with: python-version: '3.12' - - uses: dtolnay/rust-toolchain@29eef336d9b2848a0b548edc03f92a220660cdb8 # stable + - uses: dtolnay/rust-toolchain@29eef336d9b2848a0b548edc03f92a220660cdb8 # nightly with: + toolchain: nightly-2025-08-07 components: rustfmt - name: Copy root LICENSE for maturin @@ -140,4 +141,4 @@ jobs: run: cp ../LICENSE LICENSE - name: Lint code - run: make dev lint \ No newline at end of file + run: make dev lint diff --git a/.github/workflows/ci-rust.yml b/.github/workflows/ci-rust.yml index dd513751..5ae09c46 100644 --- a/.github/workflows/ci-rust.yml +++ b/.github/workflows/ci-rust.yml @@ -19,8 +19,9 @@ jobs: - name: Configure Cargo registry shell: bash run: bash "$GITHUB_WORKSPACE/.github/scripts/configure-cargo.sh" - - uses: dtolnay/rust-toolchain@29eef336d9b2848a0b548edc03f92a220660cdb8 # stable + - uses: dtolnay/rust-toolchain@29eef336d9b2848a0b548edc03f92a220660cdb8 # nightly with: + toolchain: nightly-2025-08-07 components: rustfmt - name: Format all files run: make fmt @@ -80,4 +81,4 @@ jobs: repo-token: ${{ secrets.GITHUB_TOKEN }} - uses: dtolnay/rust-toolchain@29eef336d9b2848a0b548edc03f92a220660cdb8 # stable - name: Run tests - run: cargo test --workspace \ No newline at end of file + run: cargo test --workspace diff --git a/.rustfmt.toml b/.rustfmt.toml new file mode 100644 index 00000000..d9584eec --- /dev/null +++ b/.rustfmt.toml @@ -0,0 +1,13 @@ +# cargo fmt reads each crate's edition from Cargo.toml; this fallback keeps +# direct rustfmt usable for the repository's Rust 2021 crates. +edition = "2021" +style_edition = "2021" +newline_style = "Unix" +group_imports = "StdExternalCrate" +imports_layout = "HorizontalVertical" +imports_granularity = "Module" + +max_width = 100 +fn_call_width = 100 +attr_fn_like_width = 100 +array_width = 100 diff --git a/dotnet/Makefile b/dotnet/Makefile index ee9fbd92..6980fc49 100644 --- a/dotnet/Makefile +++ b/dotnet/Makefile @@ -50,7 +50,7 @@ fmt-dotnet: fmt-rust: @echo "Formatting Rust code..." - cd ../rust/ffi && cargo fmt --all + cd ../rust/ffi && ../../.github/scripts/rustfmt.sh --all lint: lint-dotnet lint-rust diff --git a/go/Makefile b/go/Makefile index 0558b863..46133e52 100644 --- a/go/Makefile +++ b/go/Makefile @@ -85,7 +85,7 @@ fmt-go: fmt-rust: @echo "Formatting Rust code..." - cd ../rust/ffi && cargo fmt --all + cd ../rust/ffi && ../../.github/scripts/rustfmt.sh --all lint: lint-go lint-rust diff --git a/python/Makefile b/python/Makefile index 12128c80..3dd40667 100644 --- a/python/Makefile +++ b/python/Makefile @@ -72,7 +72,7 @@ fmt: $(VENV) -m black zerobus examples tests $(VENV) -m autoflake -ri --exclude '*_pb2*.py' zerobus examples tests $(VENV) -m isort zerobus examples tests - cd rust && cargo fmt --all + cd rust && ../../.github/scripts/rustfmt.sh --all lint: $(VENV) -m pycodestyle --exclude='*_pb2*.py' --max-line-length=120 --ignore=E203,W503 zerobus diff --git a/python/rust/src/arrow.rs b/python/rust/src/arrow.rs index 75a358f6..4b414730 100644 --- a/python/rust/src/arrow.rs +++ b/python/rust/src/arrow.rs @@ -8,14 +8,15 @@ use std::sync::Arc; -use pyo3::prelude::*; -use pyo3::types::PyBytes; -use tokio::sync::RwLock; - use databricks_zerobus_ingest_sdk::{ - StreamBuilder, ZerobusArrowStream as RustZerobusArrowStream, ZerobusError as RustError, + StreamBuilder, + ZerobusArrowStream as RustZerobusArrowStream, + ZerobusError as RustError, ZerobusSdk as RustSdk, }; +use pyo3::prelude::*; +use pyo3::types::PyBytes; +use tokio::sync::RwLock; use crate::auth::HeadersProviderWrapper; use crate::common::map_error; @@ -249,10 +250,7 @@ impl ArrowStreamConfigurationOptions { ("recovery_timeout_ms", self.recovery_timeout_ms), ("recovery_backoff_ms", self.recovery_backoff_ms), ("recovery_retries", self.recovery_retries as i64), - ( - "server_lack_of_ack_timeout_ms", - self.server_lack_of_ack_timeout_ms, - ), + ("server_lack_of_ack_timeout_ms", self.server_lack_of_ack_timeout_ms), ("flush_timeout_ms", self.flush_timeout_ms), ("connection_timeout_ms", self.connection_timeout_ms), ] { diff --git a/python/rust/src/async_wrapper.rs b/python/rust/src/async_wrapper.rs index ea6faed9..c31ba029 100644 --- a/python/rust/src/async_wrapper.rs +++ b/python/rust/src/async_wrapper.rs @@ -2,20 +2,27 @@ use std::future::Future; use std::pin::Pin; use std::sync::{Arc, Mutex}; +use databricks_zerobus_ingest_sdk::{ + StreamBuilder, + ZerobusSdk as RustSdk, + ZerobusStream as RustStream, +}; use pyo3::prelude::*; use pyo3_asyncio::tokio::future_into_py; use tokio::sync::RwLock; -use databricks_zerobus_ingest_sdk::{ - StreamBuilder, ZerobusSdk as RustSdk, ZerobusStream as RustStream, -}; - use crate::arrow; use crate::arrow::{ArrowStreamConfigurationOptions, AsyncZerobusArrowStream}; use crate::auth::HeadersProviderWrapper; use crate::common::{ - apply_grpc_options, encoded_record_to_pybytes, extract_record_payload, extract_record_payloads, - map_error, StreamConfigurationOptions, TableProperties, SDK_IDENTIFIER_PREFIX, + apply_grpc_options, + encoded_record_to_pybytes, + extract_record_payload, + extract_record_payloads, + map_error, + StreamConfigurationOptions, + TableProperties, + SDK_IDENTIFIER_PREFIX, }; // ============================================================================= diff --git a/python/rust/src/auth.rs b/python/rust/src/auth.rs index 0cc08ff2..c292750f 100644 --- a/python/rust/src/auth.rs +++ b/python/rust/src/auth.rs @@ -1,11 +1,13 @@ -use async_trait::async_trait; -use pyo3::exceptions::PyNotImplementedError; -use pyo3::prelude::*; use std::collections::HashMap; +use async_trait::async_trait; use databricks_zerobus_ingest_sdk::{ - HeadersProvider as RustHeadersProvider, ZerobusError as RustError, ZerobusResult as RustResult, + HeadersProvider as RustHeadersProvider, + ZerobusError as RustError, + ZerobusResult as RustResult, }; +use pyo3::exceptions::PyNotImplementedError; +use pyo3::prelude::*; use crate::common::intern_header_name; @@ -41,9 +43,7 @@ impl HeadersProvider { /// Returns: /// List of (header_name, header_value) tuples fn get_headers(&self, _py: Python) -> PyResult { - Err(PyNotImplementedError::new_err( - "Subclasses must implement get_headers()", - )) + Err(PyNotImplementedError::new_err("Subclasses must implement get_headers()")) } } diff --git a/python/rust/src/common.rs b/python/rust/src/common.rs index e79c2aee..c424d9f3 100644 --- a/python/rust/src/common.rs +++ b/python/rust/src/common.rs @@ -1,15 +1,17 @@ use std::collections::HashMap; use std::sync::{Arc, Mutex, OnceLock}; +use databricks_zerobus_ingest_sdk::{ + AckCallback as RustAckCallback, + EncodedRecord, + OffsetId, + StreamBuilder, +}; use prost::Message; use pyo3::exceptions::{PyException, PyValueError}; use pyo3::prelude::*; use pyo3::types::{PyBytes, PyDict, PyList, PyTuple}; -use databricks_zerobus_ingest_sdk::{ - AckCallback as RustAckCallback, EncodedRecord, OffsetId, StreamBuilder, -}; - /// User-agent prefix emitted by this wrapper SDK. Combined with the wrapper /// crate version via `env!("CARGO_PKG_VERSION")` at the call site. pub(crate) const SDK_IDENTIFIER_PREFIX: &str = "zerobus-sdk-py"; @@ -273,24 +275,16 @@ impl StreamConfigurationOptions { /// Validate that all numeric fields are non-negative before casting to unsigned types. pub fn validate(&self) -> PyResult<()> { if self.max_inflight_records < 0 { - return Err(PyValueError::new_err( - "max_inflight_records must be non-negative", - )); + return Err(PyValueError::new_err("max_inflight_records must be non-negative")); } if self.recovery_timeout_ms < 0 { - return Err(PyValueError::new_err( - "recovery_timeout_ms must be non-negative", - )); + return Err(PyValueError::new_err("recovery_timeout_ms must be non-negative")); } if self.recovery_backoff_ms < 0 { - return Err(PyValueError::new_err( - "recovery_backoff_ms must be non-negative", - )); + return Err(PyValueError::new_err("recovery_backoff_ms must be non-negative")); } if self.recovery_retries < 0 { - return Err(PyValueError::new_err( - "recovery_retries must be non-negative", - )); + return Err(PyValueError::new_err("recovery_retries must be non-negative")); } if self.server_lack_of_ack_timeout_ms < 0 { return Err(PyValueError::new_err( @@ -298,9 +292,7 @@ impl StreamConfigurationOptions { )); } if self.flush_timeout_ms < 0 { - return Err(PyValueError::new_err( - "flush_timeout_ms must be non-negative", - )); + return Err(PyValueError::new_err("flush_timeout_ms must be non-negative")); } if let Some(v) = self.stream_paused_max_wait_time_ms { if v < 0 { @@ -454,9 +446,7 @@ pub(crate) fn extract_record_payloads(payloads: &PyAny) -> PyResult( - "Payloads must be a list", - )); + return Err(PyErr::new::("Payloads must be a list")); } Ok(out) diff --git a/python/rust/src/lib.rs b/python/rust/src/lib.rs index c203c6b8..73fbb7f1 100644 --- a/python/rust/src/lib.rs +++ b/python/rust/src/lib.rs @@ -32,14 +32,8 @@ fn _zerobus_core(py: Python, m: &PyModule) -> PyResult<()> { m.add_class::()?; // Add exception types - m.add( - "ZerobusException", - py.get_type::(), - )?; - m.add( - "NonRetriableException", - py.get_type::(), - )?; + m.add("ZerobusException", py.get_type::())?; + m.add("NonRetriableException", py.get_type::())?; // Add authentication classes m.add_class::()?; diff --git a/python/rust/src/sync_wrapper.rs b/python/rust/src/sync_wrapper.rs index c4fcffa7..65eb332b 100644 --- a/python/rust/src/sync_wrapper.rs +++ b/python/rust/src/sync_wrapper.rs @@ -1,20 +1,27 @@ use std::sync::Arc; +use databricks_zerobus_ingest_sdk::{ + StreamBuilder, + ZerobusSdk as RustSdk, + ZerobusStream as RustStream, +}; use pyo3::prelude::*; use pyo3::types::PyBytes; use tokio::runtime::Runtime; use tokio::sync::RwLock; -use databricks_zerobus_ingest_sdk::{ - StreamBuilder, ZerobusSdk as RustSdk, ZerobusStream as RustStream, -}; - use crate::arrow; use crate::arrow::{ArrowStreamConfigurationOptions, ZerobusArrowStream}; use crate::auth::HeadersProviderWrapper; use crate::common::{ - apply_grpc_options, encoded_record_to_pybytes, extract_record_payload, extract_record_payloads, - map_error, StreamConfigurationOptions, TableProperties, SDK_IDENTIFIER_PREFIX, + apply_grpc_options, + encoded_record_to_pybytes, + extract_record_payload, + extract_record_payloads, + map_error, + StreamConfigurationOptions, + TableProperties, + SDK_IDENTIFIER_PREFIX, }; // ============================================================================= diff --git a/rust/CONTRIBUTING.md b/rust/CONTRIBUTING.md index a3d65f58..b522c731 100644 --- a/rust/CONTRIBUTING.md +++ b/rust/CONTRIBUTING.md @@ -10,6 +10,7 @@ This document covers Rust-specific development setup and workflow. - Git - Rust 1.70+ (stable toolchain) +- Pinned formatter: `rustup toolchain install nightly-2025-08-07 --profile minimal --component rustfmt` - Protocol Buffers compiler (`protoc`) ### Setting Up Your Development Environment @@ -37,7 +38,7 @@ Code style is enforced by a formatter check in your pull request. We use `rustfm make fmt ``` -This runs `cargo fmt --all` to format all crates in the workspace. +This runs `cargo fmt --all` with the same pinned rustfmt used by CI. Compilation, linting, and tests continue to use stable Rust. ### Running Linters diff --git a/rust/Makefile b/rust/Makefile index 2d20dd40..b04f2e1c 100644 --- a/rust/Makefile +++ b/rust/Makefile @@ -33,7 +33,7 @@ clean: cargo clean fmt: - cargo fmt --all + ../.github/scripts/rustfmt.sh --all lint: cargo clippy --all -- -D warnings diff --git a/rust/examples/arrow/src/main.rs b/rust/examples/arrow/src/main.rs index 010c51d4..a99b6fce 100644 --- a/rust/examples/arrow/src/main.rs +++ b/rust/examples/arrow/src/main.rs @@ -2,7 +2,11 @@ use std::error::Error; use std::sync::Arc; use arrow_array::{ - Float64Array, Int32Array, LargeStringArray, RecordBatch, TimestampMicrosecondArray, + Float64Array, + Int32Array, + LargeStringArray, + RecordBatch, + TimestampMicrosecondArray, }; use arrow_ipc::CompressionType; use databricks_zerobus_ingest_sdk::{ArrowSchema, DataType, Field, TimeUnit, ZerobusSdk}; @@ -128,11 +132,7 @@ async fn main() -> Result<(), Box> { if (i + 1) % WAIT_EVERY == 0 { stream.wait_for_offset(offset_id).await?; - println!( - "Acknowledged through batch {} (offset ID {})", - i + 1, - offset_id - ); + println!("Acknowledged through batch {} (offset ID {})", i + 1, offset_id); } } diff --git a/rust/examples/json/batch.rs b/rust/examples/json/batch.rs index c4e7a4f4..516e6c2a 100644 --- a/rust/examples/json/batch.rs +++ b/rust/examples/json/batch.rs @@ -97,15 +97,9 @@ async fn ingest_with_offset_api(stream: &mut ZerobusStream) -> Result<(), Box Result<(), Box Result<(), Box Result<(), Box Result<(), Box Result<(), Box Result<(), Box Result<(), Box - no wrapper needed, works the same as ProtoBytes. @@ -198,15 +185,9 @@ async fn ingest_with_offset_api(stream: &mut ZerobusStream) -> Result<(), Box Result<(), Box Result<(), Box using the original length. This is safe because // the pointers were produced by Box::into_raw(vec.into_boxed_slice()), // which guarantees capacity == len. - let ptrs = Box::from_raw(std::ptr::slice_from_raw_parts_mut( - array.batches, - array.count, - )); - let lens = Box::from_raw(std::ptr::slice_from_raw_parts_mut( - array.lengths, - array.count, - )); + let ptrs = + Box::from_raw(std::ptr::slice_from_raw_parts_mut(array.batches, array.count)); + let lens = + Box::from_raw(std::ptr::slice_from_raw_parts_mut(array.lengths, array.count)); for (&ptr, &len) in ptrs.iter().zip(lens.iter()) { if !ptr.is_null() && len > 0 { // Each batch slice was produced by Box::into_raw(bytes.into_boxed_slice()), @@ -620,14 +625,10 @@ pub extern "C" fn zerobus_arrow_free_batch_array(array: CArrowBatchArray) { pub extern "C" fn zerobus_arrow_stream_is_closed(stream: *mut CArrowStream) -> bool { // No CResult out-param; on a caught panic return `true` (treat as closed), // matching the answer for an invalid handle. - ffi_guard( - ptr::null_mut(), - true, - move || match validate_arrow_stream_ptr(stream) { - Ok(s) => s.is_closed(), - Err(_) => true, - }, - ) + ffi_guard(ptr::null_mut(), true, move || match validate_arrow_stream_ptr(stream) { + Ok(s) => s.is_closed(), + Err(_) => true, + }) } /// Returns the default Arrow stream configuration options. diff --git a/rust/ffi/src/builder.rs b/rust/ffi/src/builder.rs index 9c40621b..f9791667 100644 --- a/rust/ffi/src/builder.rs +++ b/rust/ffi/src/builder.rs @@ -1,11 +1,13 @@ //! `ZerobusSdkBuilder` FFI surface. -use crate::common::*; -use databricks_zerobus_ingest_sdk::{NoTlsConfig, ZerobusSdk, ZerobusSdkBuilder}; use std::os::raw::c_char; use std::ptr; use std::sync::Arc; +use databricks_zerobus_ingest_sdk::{NoTlsConfig, ZerobusSdk, ZerobusSdkBuilder}; + +use crate::common::*; + // ============================================================================ // ZerobusSdkBuilder FFI // ============================================================================ diff --git a/rust/ffi/src/common.rs b/rust/ffi/src/common.rs index 0883d2b5..62e1a6c6 100644 --- a/rust/ffi/src/common.rs +++ b/rust/ffi/src/common.rs @@ -1,11 +1,5 @@ //! Shared FFI types and helpers used across the FFI surface modules. -use async_trait::async_trait; -use databricks_zerobus_ingest_sdk::databricks::zerobus::RecordType; -use databricks_zerobus_ingest_sdk::{ - AckCallback, HeadersProvider, OffsetId, ZerobusError, ZerobusResult, ZerobusSdk, ZerobusStream, -}; -use once_cell::sync::Lazy; use std::any::Any; use std::collections::{HashMap, HashSet}; use std::ffi::{CStr, CString}; @@ -14,6 +8,19 @@ use std::panic::{catch_unwind, AssertUnwindSafe}; use std::ptr; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Mutex; + +use async_trait::async_trait; +use databricks_zerobus_ingest_sdk::databricks::zerobus::RecordType; +use databricks_zerobus_ingest_sdk::{ + AckCallback, + HeadersProvider, + OffsetId, + ZerobusError, + ZerobusResult, + ZerobusSdk, + ZerobusStream, +}; +use once_cell::sync::Lazy; use tokio::runtime::Runtime; use tracing_subscriber::{fmt, EnvFilter}; @@ -475,10 +482,7 @@ pub(crate) static ACK_DROP_SENTINEL_CREATE_FAIL_TESTS: u8 = 0; #[cfg(test)] impl Drop for CallbackAckCallback { fn drop(&mut self) { - if std::ptr::eq( - self.user_data as *const u8, - &ACK_DROP_SENTINEL_CREATE_FAIL_TESTS, - ) { + if std::ptr::eq(self.user_data as *const u8, &ACK_DROP_SENTINEL_CREATE_FAIL_TESTS) { ACK_CALLBACK_DROP_COUNT.fetch_add(1, Ordering::SeqCst); } } @@ -587,10 +591,7 @@ mod common_tests { }) .join(); assert!(poisoned.is_err(), "the spawned thread should have panicked"); - assert!( - HEADER_KEY_CACHE.lock().is_err(), - "the lock should now be poisoned" - ); + assert!(HEADER_KEY_CACHE.lock().is_err(), "the lock should now be poisoned"); // Interning must still work despite the poison — no panic, correct value, // and still interned (same pointer on a second call). diff --git a/rust/ffi/src/lib.rs b/rust/ffi/src/lib.rs index d35e556d..31b78c8e 100644 --- a/rust/ffi/src/lib.rs +++ b/rust/ffi/src/lib.rs @@ -36,12 +36,11 @@ mod tests; pub use arrow::*; pub use builder::*; pub use common::*; -pub use proto_schema::*; -pub use sdk::*; -pub use stream::*; - // Re-exported SDK types referenced via `crate::` paths by the test module. #[cfg(test)] pub(crate) use databricks_zerobus_ingest_sdk::databricks::zerobus::RecordType; #[cfg(test)] pub(crate) use databricks_zerobus_ingest_sdk::ZerobusError; +pub use proto_schema::*; +pub use sdk::*; +pub use stream::*; diff --git a/rust/ffi/src/proto_schema.rs b/rust/ffi/src/proto_schema.rs index c2235ce3..689cdcc5 100644 --- a/rust/ffi/src/proto_schema.rs +++ b/rust/ffi/src/proto_schema.rs @@ -1,15 +1,24 @@ //! Dynamic protobuf schema FFI surface. -use crate::common::*; +use std::fmt::Write; +use std::os::raw::c_char; +use std::ptr; + use databricks_zerobus_ingest_sdk::schema::{descriptor_from_uc_schema, UcTableSchema}; use prost::Message; use prost_reflect::{ - Cardinality, DescriptorPool, DeserializeOptions, DynamicMessage, Kind, MapKey, - MessageDescriptor, ReflectMessage, Value, + Cardinality, + DescriptorPool, + DeserializeOptions, + DynamicMessage, + Kind, + MapKey, + MessageDescriptor, + ReflectMessage, + Value, }; -use std::fmt::Write; -use std::os::raw::c_char; -use std::ptr; + +use crate::common::*; // ============================================================================ // Dynamic Protobuf FFI diff --git a/rust/ffi/src/sdk.rs b/rust/ffi/src/sdk.rs index ba71665b..fad8ee6f 100644 --- a/rust/ffi/src/sdk.rs +++ b/rust/ffi/src/sdk.rs @@ -1,10 +1,12 @@ //! Top-level SDK lifecycle FFI surface. -use crate::common::*; -use databricks_zerobus_ingest_sdk::ZerobusSdk; use std::os::raw::c_char; use std::ptr; +use databricks_zerobus_ingest_sdk::ZerobusSdk; + +use crate::common::*; + /// Creates a new ZerobusSdk with default user-agent and TLS settings. /// /// Retained for ABI back-compat with v1.2.x; new code should use the diff --git a/rust/ffi/src/stream.rs b/rust/ffi/src/stream.rs index c8b7a0b5..a582b5e7 100644 --- a/rust/ffi/src/stream.rs +++ b/rust/ffi/src/stream.rs @@ -1,17 +1,23 @@ //! Stream creation and record ingestion FFI surface. -use crate::common::*; -use databricks_zerobus_ingest_sdk::databricks::zerobus::RecordType; -use databricks_zerobus_ingest_sdk::{ - EncodedRecord, HeadersProvider, StreamBuilder, ZerobusError, ZerobusStream, -}; -use prost::Message; use std::ffi::CString; use std::mem::ManuallyDrop; use std::os::raw::c_char; use std::ptr; use std::sync::Arc; +use databricks_zerobus_ingest_sdk::databricks::zerobus::RecordType; +use databricks_zerobus_ingest_sdk::{ + EncodedRecord, + HeadersProvider, + StreamBuilder, + ZerobusError, + ZerobusStream, +}; +use prost::Message; + +use crate::common::*; + // Builder option application helpers fn apply_c_stream_options<'a>( @@ -101,9 +107,8 @@ async fn build_stream_from_parts( headers_callback, user_data, } => { - let headers_provider: Arc = Arc::new( - CallbackHeadersProvider::new(headers_callback, user_data.get()), - ); + let headers_provider: Arc = + Arc::new(CallbackHeadersProvider::new(headers_callback, user_data.get())); sdk_ref .stream_builder() .table(table_name) @@ -122,9 +127,7 @@ async fn build_stream_from_parts( } RecordType::Json => base.json(), RecordType::Unspecified => { - return Err(ZerobusError::InvalidArgument( - "Record type is not specified".to_string(), - )) + return Err(ZerobusError::InvalidArgument("Record type is not specified".to_string())) } }; @@ -174,10 +177,7 @@ fn invoke_offset_async_callback( })) .is_err() { - tracing::error!( - offset, - "async offset callback panicked; contained at FFI boundary" - ); + tracing::error!(offset, "async offset callback panicked; contained at FFI boundary"); } if !callback_result.error_message.is_null() { @@ -201,10 +201,7 @@ fn invoke_bool_async_callback( })) .is_err() { - tracing::error!( - value, - "async bool callback panicked; contained at FFI boundary" - ); + tracing::error!(value, "async bool callback panicked; contained at FFI boundary"); } if !callback_result.error_message.is_null() { @@ -225,11 +222,7 @@ fn invoke_record_array_async_callback( let callback_records = ManuallyDrop::new(records); let panicked = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| unsafe { - callback( - ptr::read(&*callback_records), - callback_result_ptr, - user_data, - ) + callback(ptr::read(&*callback_records), callback_result_ptr, user_data) })) .is_err(); if panicked { @@ -1807,11 +1800,9 @@ pub extern "C" fn zerobus_stream_close_async( pub extern "C" fn zerobus_stream_is_closed(stream: *mut CZerobusStream) -> bool { // No CResult out-param; on a caught panic return `true` (treat as closed), // matching the answer for an invalid handle. - ffi_guard(ptr::null_mut(), true, move || { - match validate_stream_ptr(stream) { - Ok(s) => s.is_closed(), - Err(_) => true, - } + ffi_guard(ptr::null_mut(), true, move || match validate_stream_ptr(stream) { + Ok(s) => s.is_closed(), + Err(_) => true, }) } diff --git a/rust/ffi/src/tests.rs b/rust/ffi/src/tests.rs index c9382ab5..6c627bc1 100644 --- a/rust/ffi/src/tests.rs +++ b/rust/ffi/src/tests.rs @@ -1,26 +1,51 @@ #[cfg(test)] mod tests { - use crate::{ - c_record_type, ffi_guard, intern_header_key, validate_sdk_ptr, validate_stream_ptr, - write_error_result, write_success_result, zerobus_free_error_message, - zerobus_get_default_config, zerobus_sdk_builder_application_name, - zerobus_sdk_builder_build, zerobus_sdk_builder_disable_tls, zerobus_sdk_builder_endpoint, - zerobus_sdk_builder_free, zerobus_sdk_builder_new, zerobus_sdk_builder_sdk_identifier, - zerobus_sdk_builder_unity_catalog_url, zerobus_sdk_create_stream, - zerobus_sdk_create_stream_async, zerobus_sdk_create_stream_with_headers_provider_async, - zerobus_sdk_free, zerobus_sdk_recreate_stream_async, zerobus_stream_close_async, - zerobus_stream_flush_async, zerobus_stream_get_unacked_records_async, - zerobus_stream_ingest_json_record_async, zerobus_stream_ingest_json_records_async, - zerobus_stream_ingest_proto_record_async, zerobus_stream_ingest_proto_records_async, - zerobus_stream_wait_for_offset_async, CHeaders, CRecordArray, CResult, - CallbackHeadersProvider, RecordType, ZerobusError, - }; - use databricks_zerobus_ingest_sdk::HeadersProvider; use std::ffi::{CStr, CString}; use std::ptr; use std::sync::mpsc; use std::time::Duration; + use databricks_zerobus_ingest_sdk::HeadersProvider; + + use crate::{ + c_record_type, + ffi_guard, + intern_header_key, + validate_sdk_ptr, + validate_stream_ptr, + write_error_result, + write_success_result, + zerobus_free_error_message, + zerobus_get_default_config, + zerobus_sdk_builder_application_name, + zerobus_sdk_builder_build, + zerobus_sdk_builder_disable_tls, + zerobus_sdk_builder_endpoint, + zerobus_sdk_builder_free, + zerobus_sdk_builder_new, + zerobus_sdk_builder_sdk_identifier, + zerobus_sdk_builder_unity_catalog_url, + zerobus_sdk_create_stream, + zerobus_sdk_create_stream_async, + zerobus_sdk_create_stream_with_headers_provider_async, + zerobus_sdk_free, + zerobus_sdk_recreate_stream_async, + zerobus_stream_close_async, + zerobus_stream_flush_async, + zerobus_stream_get_unacked_records_async, + zerobus_stream_ingest_json_record_async, + zerobus_stream_ingest_json_records_async, + zerobus_stream_ingest_proto_record_async, + zerobus_stream_ingest_proto_records_async, + zerobus_stream_wait_for_offset_async, + CHeaders, + CRecordArray, + CResult, + CallbackHeadersProvider, + RecordType, + ZerobusError, + }; + // Helper for c_str_to_string since it's private unsafe fn test_c_str_to_string( c_str: *const std::os::raw::c_char, @@ -275,12 +300,8 @@ mod tests { #[test] fn test_builder_empty_strings_are_noops() { // Empty identifier/application_name must not produce a trailing space. - let (sdk, result) = build_via_c_builder( - "https://workspace.zerobus.databricks.com", - "", - Some(""), - Some(""), - ); + let (sdk, result) = + build_via_c_builder("https://workspace.zerobus.databricks.com", "", Some(""), Some("")); assert!(result.success); assert!(!sdk.is_null()); zerobus_sdk_free(sdk); @@ -394,10 +415,7 @@ mod tests { assert!(stream.is_null(), "create_stream should fail"); assert!(!create_result.success, "result should indicate failure"); - assert!( - create_result.is_retryable, - "retryable create failures must set is_retryable=true" - ); + assert!(create_result.is_retryable, "retryable create failures must set is_retryable=true"); zerobus_free_error_message(create_result.error_message); zerobus_sdk_free(sdk); @@ -479,10 +497,7 @@ mod tests { let (stream_is_null, callback_success, callback_retryable, callback_message) = receiver .recv_timeout(Duration::from_secs(2)) .expect("callback should be invoked"); - assert!( - stream_is_null, - "callback should receive a null stream on failure" - ); + assert!(stream_is_null, "callback should receive a null stream on failure"); assert!(!callback_success, "callback result should indicate failure"); assert!( !callback_retryable, @@ -573,19 +588,13 @@ mod tests { &mut create_result as *mut CResult, ); - assert!( - started, - "create_stream_with_headers_provider_async should schedule the task" - ); + assert!(started, "create_stream_with_headers_provider_async should schedule the task"); assert!(create_result.success, "scheduling result should succeed"); let (stream_is_null, callback_success, callback_retryable, callback_message) = receiver .recv_timeout(Duration::from_secs(2)) .expect("callback should be invoked"); - assert!( - stream_is_null, - "callback should receive a null stream on failure" - ); + assert!(stream_is_null, "callback should receive a null stream on failure"); assert!(!callback_success, "callback result should indicate failure"); assert!( !callback_retryable, @@ -926,11 +935,13 @@ mod tests { // Ack callback bridge tests // ======================================================================== - use crate::CallbackAckCallback; - use databricks_zerobus_ingest_sdk::AckCallback as _AckCallbackTrait; use std::os::raw::c_char; use std::sync::atomic::{AtomicI64, Ordering as AtomicOrdering}; + use databricks_zerobus_ingest_sdk::AckCallback as _AckCallbackTrait; + + use crate::CallbackAckCallback; + // extern "C" callbacks can't capture, so they record into these statics. // Tests reset the slots before use. static LAST_ACK_OFFSET: AtomicI64 = AtomicI64::new(-1); @@ -1013,10 +1024,11 @@ mod tests { // code, so the callback runs to completion — `user_data` must outlive the // callback, not merely `teardown()`. - use databricks_zerobus_ingest_sdk::CallbackHandlerHarness; use std::sync::atomic::{AtomicBool, AtomicU64}; use std::sync::Arc; + use databricks_zerobus_ingest_sdk::CallbackHandlerHarness; + // Heap user_data each callback dereferences; a post-free call is a UAF ASan // catches, and `magic` detects a freed/garbage box. `block_ms` lets a callback // stay synchronously in-flight (to outrun a bounded teardown budget); `started` @@ -1044,10 +1056,7 @@ mod tests { // its `user_data` access below) outlives a bounded teardown budget. std::thread::sleep(std::time::Duration::from_millis(data.block_ms)); } - assert_eq!( - data.magic, ACK_MAGIC, - "user_data was freed or corrupted mid-callback" - ); + assert_eq!(data.magic, ACK_MAGIC, "user_data was freed or corrupted mid-callback"); data.calls.fetch_add(1, AtomicOrdering::SeqCst); data.finished.store(true, AtomicOrdering::SeqCst); } @@ -1120,10 +1129,7 @@ mod tests { assert!(harness.send_ack(offset), "enqueue must succeed while live"); } for &(offset, msg) in errors { - assert!( - harness.send_error(offset, msg), - "enqueue must succeed while live" - ); + assert!(harness.send_error(offset, msg), "enqueue must succeed while live"); } let expected = (acks.len() + errors.len()) as u64; wait_for_calls(user_data, expected).await; @@ -1131,14 +1137,8 @@ mod tests { harness.teardown(callback_max_wait_time_ms).await; // Task gone: its receiver is dropped, so no further dispatch is possible. - assert!( - harness.is_task_gone(), - "handler task must be gone after teardown" - ); - assert!( - !harness.send_ack(999), - "enqueue must be rejected once task is gone" - ); + assert!(harness.is_task_gone(), "handler task must be gone after teardown"); + assert!(!harness.send_ack(999), "enqueue must be rejected once task is gone"); assert!(!harness.send_error(1000, "late")); assert_eq!( unsafe { &*user_data }.calls.load(AtomicOrdering::SeqCst), @@ -1236,10 +1236,7 @@ mod tests { harness.teardown(callback_max_wait_time_ms).await; - assert!( - harness.is_task_gone(), - "handler task must be gone after teardown" - ); + assert!(harness.is_task_gone(), "handler task must be gone after teardown"); assert_eq!( unsafe { &*user_data }.calls.load(AtomicOrdering::SeqCst), 0, @@ -1397,10 +1394,7 @@ mod tests { assert!(stream.is_null()); let (success, _retryable, msg) = drain_result(&mut result); - assert!( - !success, - "expected create_stream_with_headers_provider to fail on empty table" - ); + assert!(!success, "expected create_stream_with_headers_provider to fail on empty table"); assert!(!msg.is_empty(), "expected a non-empty error message"); assert_eq!( @@ -1417,13 +1411,16 @@ mod tests { // Dynamic protobuf schema tests // ======================================================================== + use prost::Message; + use prost_reflect::{DescriptorPool, DynamicMessage, MessageDescriptor}; + use crate::{ - zerobus_free_proto_bytes, zerobus_proto_schema_descriptor_bytes, - zerobus_proto_schema_encode_json, zerobus_proto_schema_free, + zerobus_free_proto_bytes, + zerobus_proto_schema_descriptor_bytes, + zerobus_proto_schema_encode_json, + zerobus_proto_schema_free, zerobus_proto_schema_from_uc_json, }; - use prost::Message; - use prost_reflect::{DescriptorPool, DynamicMessage, MessageDescriptor}; // Minimal Unity Catalog table-metadata JSON, shaped like the body of // GET /api/2.1/unity-catalog/tables/{name}. @@ -1522,14 +1519,8 @@ mod tests { let msg_desc = message_descriptor_from_bytes(desc_bytes); let decoded = DynamicMessage::decode(msg_desc, encoded).unwrap(); assert_eq!(decoded.get_field_by_name("id").unwrap().as_i64(), Some(7)); - assert_eq!( - decoded.get_field_by_name("payload").unwrap().as_str(), - Some("hello") - ); - assert_eq!( - decoded.get_field_by_name("ts").unwrap().as_i64(), - Some(1700000000000000) - ); + assert_eq!(decoded.get_field_by_name("payload").unwrap().as_str(), Some("hello")); + assert_eq!(decoded.get_field_by_name("ts").unwrap().as_i64(), Some(1700000000000000)); zerobus_free_proto_bytes(out_data, out_len); zerobus_proto_schema_free(schema); @@ -1782,10 +1773,7 @@ mod tests { // string to preserve precision and scale. let table = uc_table_json_with_column("price", "DECIMAL"); let decoded = encode_and_decode(&table, r#"{"k": 1, "price": "123.45"}"#); - assert_eq!( - decoded.get_field_by_name("price").unwrap().as_str(), - Some("123.45") - ); + assert_eq!(decoded.get_field_by_name("price").unwrap().as_str(), Some("123.45")); } #[test] @@ -1794,10 +1782,7 @@ mod tests { // string round-trips exactly. let table = uc_table_json_with_column("big", "BIGINT"); let decoded = encode_and_decode(&table, r#"{"k": 1, "big": "9223372036854775807"}"#); - assert_eq!( - decoded.get_field_by_name("big").unwrap().as_i64(), - Some(9223372036854775807) - ); + assert_eq!(decoded.get_field_by_name("big").unwrap().as_i64(), Some(9223372036854775807)); } #[test] @@ -1806,10 +1791,7 @@ mod tests { // (a string whose contents are the variant's JSON). let table = uc_table_json_with_column("v", "VARIANT"); let decoded = encode_and_decode(&table, r#"{"k": 1, "v": "{\"a\":1,\"b\":[2,3]}"}"#); - assert_eq!( - decoded.get_field_by_name("v").unwrap().as_str(), - Some(r#"{"a":1,"b":[2,3]}"#) - ); + assert_eq!(decoded.get_field_by_name("v").unwrap().as_str(), Some(r#"{"a":1,"b":[2,3]}"#)); } #[test] @@ -1882,10 +1864,7 @@ mod tests { encode_and_decode(&table, r#"{"k": 1, "addr": {"city": "NYC", "zip": 10001}}"#); let field = decoded.get_field_by_name("addr").unwrap(); let addr = field.as_message().unwrap(); - assert_eq!( - addr.get_field_by_name("city").unwrap().as_str(), - Some("NYC") - ); + assert_eq!(addr.get_field_by_name("city").unwrap().as_str(), Some("NYC")); assert_eq!(addr.get_field_by_name("zip").unwrap().as_i32(), Some(10001)); } @@ -1928,10 +1907,7 @@ mod tests { // integer (not an ISO-8601 string). 19000 days ≈ 2022-01-08. let table = uc_table_json_with_column("d", "DATE"); let decoded = encode_and_decode(&table, r#"{"k": 1, "d": 19000}"#); - assert_eq!( - decoded.get_field_by_name("d").unwrap().as_i32(), - Some(19000) - ); + assert_eq!(decoded.get_field_by_name("d").unwrap().as_i32(), Some(19000)); } #[test] @@ -1940,10 +1916,7 @@ mod tests { // value is microseconds since the epoch, an integer. let table = uc_table_json_with_column("tsn", "TIMESTAMP_NTZ"); let decoded = encode_and_decode(&table, r#"{"k": 1, "tsn": 1700000000000000}"#); - assert_eq!( - decoded.get_field_by_name("tsn").unwrap().as_i64(), - Some(1700000000000000) - ); + assert_eq!(decoded.get_field_by_name("tsn").unwrap().as_i64(), Some(1700000000000000)); } #[test] @@ -1976,14 +1949,8 @@ mod tests { &mut enc_result as *mut CResult, ); assert!(ok, "encode failed"); - assert_eq!( - out_len, 0, - "record with no fields set should encode to zero bytes" - ); - assert!( - !out_data.is_null(), - "buffer pointer should be non-null even when empty" - ); + assert_eq!(out_len, 0, "record with no fields set should encode to zero bytes"); + assert!(!out_data.is_null(), "buffer pointer should be non-null even when empty"); // The assertion is the absence of a leak/crash on free. zerobus_free_proto_bytes(out_data, out_len); @@ -2010,12 +1977,9 @@ mod tests { workers.push(thread::spawn(move || { let handle = handle_addr as *const crate::CZerobusProtoSchema; for i in 0..200 { - let record = CString::new(format!( - r#"{{"id": {}, "payload": "p{}"}}"#, - t * 1000 + i, - i - )) - .unwrap(); + let record = + CString::new(format!(r#"{{"id": {}, "payload": "p{}"}}"#, t * 1000 + i, i)) + .unwrap(); let mut out_data: *mut u8 = ptr::null_mut(); let mut out_len: usize = 0; let mut enc = unwritten_result(); @@ -2272,10 +2236,7 @@ mod tests { ); assert!(!ok, "expected encode to fail"); assert!(!enc.success); - assert!( - !enc.is_retryable, - "a missing required field is a caller error" - ); + assert!(!enc.is_retryable, "a missing required field is a caller error"); assert!(out_data.is_null(), "no buffer should be allocated on error"); assert_eq!(out_len, 0, "length must be cleared on error"); assert!(!enc.error_message.is_null()); @@ -2303,10 +2264,7 @@ mod tests { ) .unwrap(); let msg = encode_expecting_error(&table, r#"{"k": 1, "addr": {"city": "boston"}}"#); - assert!( - msg.contains("addr.zip"), - "error should name the nested path, got: {msg}" - ); + assert!(msg.contains("addr.zip"), "error should name the nested path, got: {msg}"); } #[test] @@ -2326,10 +2284,7 @@ mod tests { .unwrap(); let msg = encode_expecting_error(&table, r#"{"k": 1, "items": [{"id": 5}, {"label": "x"}]}"#); - assert!( - msg.contains("items[1].id"), - "error should name the element path, got: {msg}" - ); + assert!(msg.contains("items[1].id"), "error should name the element path, got: {msg}"); } #[test] @@ -2347,14 +2302,9 @@ mod tests { }"#, ) .unwrap(); - let msg = encode_expecting_error( - &table, - r#"{"k": 1, "lookup": {"home": {"v": 2}, "work": {}}}"#, - ); - assert!( - msg.contains("lookup[work].v"), - "error should name the map-value path, got: {msg}" - ); + let msg = + encode_expecting_error(&table, r#"{"k": 1, "lookup": {"home": {"v": 2}, "work": {}}}"#); + assert!(msg.contains("lookup[work].v"), "error should name the map-value path, got: {msg}"); } #[test] @@ -2410,10 +2360,7 @@ mod tests { .unwrap(); let msg = encode_expecting_error(&table, r#"{"addr": {"geo": {}}, "items": [{"inner": {}}]}"#); - assert!( - msg.contains("addr.geo.lat"), - "should report the 3-level-deep path, got: {msg}" - ); + assert!(msg.contains("addr.geo.lat"), "should report the 3-level-deep path, got: {msg}"); assert!( msg.contains("items[0].inner.id"), "should report the path through an array element's nested struct, got: {msg}" diff --git a/rust/jni/src/arrow_stream.rs b/rust/jni/src/arrow_stream.rs index b656879b..eb166250 100644 --- a/rust/jni/src/arrow_stream.rs +++ b/rust/jni/src/arrow_stream.rs @@ -3,16 +3,18 @@ //! This module provides JNI functions for Arrow Flight stream operations including //! batch ingestion, acknowledgment waiting, flushing, and closing. -use crate::class_cache::{as_jclass, get_class_cache}; -use crate::errors::{throw_from_zerobus_error, throw_zerobus_exception}; -use crate::runtime::block_on; +use std::sync::Arc; + use databricks_zerobus_ingest_sdk::ZerobusArrowStream; use jni::objects::{JByteArray, JClass, JObject, JValue}; use jni::sys::{jboolean, jlong, JNI_FALSE, JNI_TRUE}; use jni::JNIEnv; -use std::sync::Arc; use tokio::sync::Mutex; +use crate::class_cache::{as_jclass, get_class_cache}; +use crate::errors::{throw_from_zerobus_error, throw_zerobus_exception}; +use crate::runtime::block_on; + /// Native Arrow stream handle stored in Java. pub struct NativeArrowStreamHandle { pub stream: Arc>>, @@ -125,11 +127,9 @@ pub extern "system" fn Java_com_databricks_zerobus_ZerobusArrowStream_nativeInge } if batches.is_empty() { - return Err( - databricks_zerobus_ingest_sdk::ZerobusError::InvalidArgument( - "No batches found in Arrow IPC data".to_string(), - ), - ); + return Err(databricks_zerobus_ingest_sdk::ZerobusError::InvalidArgument( + "No batches found in Arrow IPC data".to_string(), + )); } // Ingest the first batch diff --git a/rust/jni/src/async_bridge.rs b/rust/jni/src/async_bridge.rs index e3ff0abc..9f4458d8 100644 --- a/rust/jni/src/async_bridge.rs +++ b/rust/jni/src/async_bridge.rs @@ -3,13 +3,15 @@ //! This module provides utilities for bridging Rust async operations to Java //! CompletableFutures, allowing async results to be propagated back to Java. -use crate::class_cache::{as_jclass, get_class_cache}; -use crate::errors::{create_exception_from_error, create_zerobus_exception}; -use crate::runtime::{get_jvm, spawn}; +use std::future::Future; + use databricks_zerobus_ingest_sdk::ZerobusError; use jni::objects::{GlobalRef, JObject, JValue}; use jni::JNIEnv; -use std::future::Future; + +use crate::class_cache::{as_jclass, get_class_cache}; +use crate::errors::{create_exception_from_error, create_zerobus_exception}; +use crate::runtime::{get_jvm, spawn}; /// Complete a Java CompletableFuture with a successful result. /// @@ -20,12 +22,7 @@ pub fn complete_future<'local>( future: &JObject<'local>, value: JObject<'local>, ) -> Result<(), jni::errors::Error> { - env.call_method( - future, - "complete", - "(Ljava/lang/Object;)Z", - &[JValue::Object(&value)], - )?; + env.call_method(future, "complete", "(Ljava/lang/Object;)Z", &[JValue::Object(&value)])?; Ok(()) } diff --git a/rust/jni/src/callbacks.rs b/rust/jni/src/callbacks.rs index bbf0e155..c9f56495 100644 --- a/rust/jni/src/callbacks.rs +++ b/rust/jni/src/callbacks.rs @@ -3,10 +3,12 @@ //! This module provides a Rust implementation of the SDK's AckCallback trait //! that delegates to a Java AckCallback object. -use crate::runtime::get_jvm; +use std::sync::Arc; + use databricks_zerobus_ingest_sdk::{AckCallback, OffsetId}; use jni::objects::{GlobalRef, JValue}; -use std::sync::Arc; + +use crate::runtime::get_jvm; /// A JNI bridge that implements the Rust AckCallback trait by delegating /// to a Java AckCallback object. diff --git a/rust/jni/src/class_cache.rs b/rust/jni/src/class_cache.rs index 425f1f6f..e87694f9 100644 --- a/rust/jni/src/class_cache.rs +++ b/rust/jni/src/class_cache.rs @@ -9,9 +9,10 @@ //! (which runs on a Java thread with the correct classloader), then reuses them //! from async/daemon threads. +use std::sync::OnceLock; + use jni::objects::{GlobalRef, JClass, JObject}; use jni::JNIEnv; -use std::sync::OnceLock; /// Cached JNI class references, populated during `JNI_OnLoad`. pub struct CachedClasses { diff --git a/rust/jni/src/errors.rs b/rust/jni/src/errors.rs index 7e89884d..78f19a56 100644 --- a/rust/jni/src/errors.rs +++ b/rust/jni/src/errors.rs @@ -3,11 +3,12 @@ //! This module provides utilities for converting Zerobus SDK errors to //! appropriate Java exceptions. -use crate::class_cache::{as_jclass, get_class_cache}; use databricks_zerobus_ingest_sdk::ZerobusError; use jni::objects::{GlobalRef, JString, JThrowable, JValue}; use jni::JNIEnv; +use crate::class_cache::{as_jclass, get_class_cache}; + /// Throw a ZerobusException in Java. /// /// This function creates and throws a Java ZerobusException with the given message. @@ -93,11 +94,7 @@ pub fn create_exception<'local>( }; // Create the exception instance - match env.new_object( - class, - "(Ljava/lang/String;)V", - &[JValue::Object(&j_message.into())], - ) { + match env.new_object(class, "(Ljava/lang/String;)V", &[JValue::Object(&j_message.into())]) { Ok(obj) => Some(JThrowable::from(obj)), Err(e) => { tracing::error!("Failed to create exception instance: {}", e); diff --git a/rust/jni/src/lib.rs b/rust/jni/src/lib.rs index 3162a5aa..855a019e 100644 --- a/rust/jni/src/lib.rs +++ b/rust/jni/src/lib.rs @@ -15,8 +15,6 @@ mod stream; #[cfg(feature = "test-helpers")] mod test_helper; -use jni::sys::jint; -use jni::JavaVM; use std::ffi::c_void; pub use arrow_stream::*; @@ -24,6 +22,8 @@ pub use async_bridge::*; pub use callbacks::*; pub use class_cache::*; pub use errors::*; +use jni::sys::jint; +use jni::JavaVM; pub use options::*; pub use runtime::*; pub use sdk::*; diff --git a/rust/jni/src/options.rs b/rust/jni/src/options.rs index 30be3930..93d853ad 100644 --- a/rust/jni/src/options.rs +++ b/rust/jni/src/options.rs @@ -3,12 +3,14 @@ //! Values are extracted on the JNI thread into JNI-private structs (so they're `Send`), //! then applied to a `StreamBuilder` inside the async task that builds the stream. -use crate::callbacks::JavaAckCallback; +use std::sync::Arc; + use arrow_ipc::CompressionType; use databricks_zerobus_ingest_sdk::{AckCallback, StreamBuilder}; use jni::objects::JObject; use jni::JNIEnv; -use std::sync::Arc; + +use crate::callbacks::JavaAckCallback; /// Extracted gRPC stream options, ready to apply to a `StreamBuilder`. pub struct ExtractedStreamOptions { diff --git a/rust/jni/src/runtime.rs b/rust/jni/src/runtime.rs index a4e96f73..cfa7c244 100644 --- a/rust/jni/src/runtime.rs +++ b/rust/jni/src/runtime.rs @@ -4,8 +4,9 @@ //! in the Zerobus SDK. The runtime is initialized when the JNI library is loaded //! and persists for the lifetime of the JVM. -use jni::JavaVM; use std::sync::OnceLock; + +use jni::JavaVM; use tokio::runtime::Runtime; /// Global Tokio runtime for async operations. diff --git a/rust/jni/src/sdk.rs b/rust/jni/src/sdk.rs index f1a0a1e0..b310e9df 100644 --- a/rust/jni/src/sdk.rs +++ b/rust/jni/src/sdk.rs @@ -2,21 +2,25 @@ //! //! This module provides JNI functions for creating and managing ZerobusSdk instances. +use std::sync::Arc; + +use databricks_zerobus_ingest_sdk::databricks::zerobus::RecordType; +use databricks_zerobus_ingest_sdk::ZerobusSdk; +use jni::objects::{JByteArray, JClass, JObject, JString}; +use jni::sys::{jboolean, jlong, JNI_FALSE}; +use jni::JNIEnv; + use crate::arrow_stream::NativeArrowStreamHandle; use crate::async_bridge::{create_completable_future, spawn_and_complete}; use crate::errors::{throw_from_zerobus_error, throw_zerobus_exception}; use crate::options::{ - apply_arrow_stream_options, apply_stream_options, extract_arrow_stream_options, + apply_arrow_stream_options, + apply_stream_options, + extract_arrow_stream_options, extract_stream_options, }; use crate::runtime::block_on; use crate::stream::NativeStreamHandle; -use databricks_zerobus_ingest_sdk::databricks::zerobus::RecordType; -use databricks_zerobus_ingest_sdk::ZerobusSdk; -use jni::objects::{JByteArray, JClass, JObject, JString}; -use jni::sys::{jboolean, jlong, JNI_FALSE}; -use jni::JNIEnv; -use std::sync::Arc; /// Native SDK handle stored in Java. pub struct NativeSdkHandle { @@ -268,11 +272,9 @@ pub extern "system" fn Java_com_databricks_zerobus_ZerobusSdk_nativeCreateStream } RecordType::Json => base.json(), RecordType::Unspecified => { - return Err( - databricks_zerobus_ingest_sdk::ZerobusError::InvalidArgument( - "Record type is not specified".to_string(), - ), - ); + return Err(databricks_zerobus_ingest_sdk::ZerobusError::InvalidArgument( + "Record type is not specified".to_string(), + )); } }; if let Some(opts) = extracted_options { @@ -281,11 +283,7 @@ pub extern "system" fn Java_com_databricks_zerobus_ZerobusSdk_nativeCreateStream let stream = builder.build().await?; - Ok(NativeStreamHandle::new( - stream, - credentials.0, - credentials.1, - )) + Ok(NativeStreamHandle::new(stream, credentials.0, credentials.1)) }, |handle| handle.into_raw(), ); @@ -351,11 +349,7 @@ pub extern "system" fn Java_com_databricks_zerobus_ZerobusSdk_nativeRecreateStre future_ref, async move { let new_stream = sdk_arc.recreate_stream(&old_stream).await?; - Ok(NativeStreamHandle::new( - new_stream, - client_id, - client_secret, - )) + Ok(NativeStreamHandle::new(new_stream, client_id, client_secret)) }, |handle| handle.into_raw(), ); @@ -483,11 +477,7 @@ pub extern "system" fn Java_com_databricks_zerobus_ZerobusSdk_nativeCreateArrowS let stream = builder.build_arrow().await?; - Ok(NativeArrowStreamHandle::new( - stream, - credentials.0, - credentials.1, - )) + Ok(NativeArrowStreamHandle::new(stream, credentials.0, credentials.1)) }, |handle| handle.into_raw(), ); @@ -553,11 +543,7 @@ pub extern "system" fn Java_com_databricks_zerobus_ZerobusSdk_nativeRecreateArro future_ref, async move { let new_stream = sdk_arc.recreate_arrow_stream(&old_stream).await?; - Ok(NativeArrowStreamHandle::new( - new_stream, - client_id, - client_secret, - )) + Ok(NativeArrowStreamHandle::new(new_stream, client_id, client_secret)) }, |handle| handle.into_raw(), ); diff --git a/rust/jni/src/stream.rs b/rust/jni/src/stream.rs index 6e67e70b..2444f36a 100644 --- a/rust/jni/src/stream.rs +++ b/rust/jni/src/stream.rs @@ -3,18 +3,19 @@ //! This module provides JNI functions for stream operations including //! record ingestion, acknowledgment waiting, flushing, and closing. -use crate::async_bridge::{create_completable_future, spawn_and_complete_void}; -use crate::class_cache::{as_jclass, get_class_cache}; -use crate::errors::{throw_from_zerobus_error, throw_zerobus_exception}; -use crate::runtime::block_on; -use databricks_zerobus_ingest_sdk::ZerobusStream; -use databricks_zerobus_ingest_sdk::{EncodedBatch, EncodedRecord}; +use std::sync::Arc; + +use databricks_zerobus_ingest_sdk::{EncodedBatch, EncodedRecord, ZerobusStream}; use jni::objects::{JByteArray, JClass, JList, JObject, JValue}; use jni::sys::{jboolean, jlong, JNI_FALSE, JNI_TRUE}; use jni::JNIEnv; -use std::sync::Arc; use tokio::sync::Mutex; +use crate::async_bridge::{create_completable_future, spawn_and_complete_void}; +use crate::class_cache::{as_jclass, get_class_cache}; +use crate::errors::{throw_from_zerobus_error, throw_zerobus_exception}; +use crate::runtime::block_on; + /// Native stream handle stored in Java. pub struct NativeStreamHandle { pub stream: Arc>>, @@ -621,12 +622,9 @@ pub extern "system" fn Java_com_databricks_zerobus_BaseZerobusStream_nativeGetUn } }; - if let Err(e) = env.call_method( - &list, - "add", - "(Ljava/lang/Object;)Z", - &[JValue::Object(&batch_obj)], - ) { + if let Err(e) = + env.call_method(&list, "add", "(Ljava/lang/Object;)Z", &[JValue::Object(&batch_obj)]) + { throw_zerobus_exception(&mut env, &format!("Failed to add batch to list: {}", e)); return JObject::null(); } diff --git a/rust/jni/src/test_helper.rs b/rust/jni/src/test_helper.rs index 759f2617..da5fea43 100644 --- a/rust/jni/src/test_helper.rs +++ b/rust/jni/src/test_helper.rs @@ -4,11 +4,12 @@ //! class references work correctly from daemon threads, even when the system //! classloader cannot see the SDK classes. -use crate::class_cache::{as_jclass, get_class_cache}; -use crate::runtime::{get_jvm, get_runtime}; use jni::objects::{JClass, JObject, JString}; use jni::JNIEnv; +use crate::class_cache::{as_jclass, get_class_cache}; +use crate::runtime::{get_jvm, get_runtime}; + /// Test finding a class from a Tokio daemon thread using direct `find_class`. /// /// Returns `"OK"` if the class is found, or an error message string. diff --git a/rust/sdk/src/arrow_configuration.rs b/rust/sdk/src/arrow_configuration.rs index 0d02840d..efec1663 100644 --- a/rust/sdk/src/arrow_configuration.rs +++ b/rust/sdk/src/arrow_configuration.rs @@ -3,9 +3,10 @@ //! **Beta**: Arrow Flight ingestion is in Beta. The API is stabilising but may //! still change before reaching GA. -use crate::stream_options::defaults; use arrow_ipc::CompressionType; +use crate::stream_options::defaults; + /// Configuration options for Arrow Flight stream creation and operation. /// /// These options control the behavior of Arrow Flight ingestion streams, including diff --git a/rust/sdk/src/arrow_stream.rs b/rust/sdk/src/arrow_stream.rs index c6358aaa..03a0b473 100644 --- a/rust/sdk/src/arrow_stream.rs +++ b/rust/sdk/src/arrow_stream.rs @@ -13,10 +13,13 @@ use std::pin::Pin; use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, Ordering}; use std::sync::Arc; +// Re-export arrow types for public API +pub use arrow_array::RecordBatch; use arrow_flight::encode::FlightDataEncoderBuilder; use arrow_flight::error::FlightError; use arrow_flight::{FlightClient, PutResult}; use arrow_ipc::writer::IpcWriteOptions; +pub use arrow_schema::{DataType, Field, Schema as ArrowSchema, TimeUnit}; use bytes::Bytes; use futures::{Stream, StreamExt}; #[cfg(feature = "test-hooks")] @@ -29,10 +32,6 @@ use tonic::metadata::MetadataValue; use tonic::transport::Channel; use tracing::{debug, error, info, instrument, warn}; -// Re-export arrow types for public API -pub use arrow_array::RecordBatch; -pub use arrow_schema::{DataType, Field, Schema as ArrowSchema, TimeUnit}; - use crate::arrow_configuration::ArrowStreamConfigurationOptions; use crate::arrow_metadata::{FlightAckMetadata, FlightBatchMetadata}; use crate::errors::ZerobusError; @@ -152,9 +151,7 @@ fn materialize_ipc(bytes: &Bytes) -> ZerobusResult { })?; let batch = match reader.next() { None => { - return Err(ZerobusError::InvalidArgument( - "IPC stream contains no RecordBatch".into(), - )); + return Err(ZerobusError::InvalidArgument("IPC stream contains no RecordBatch".into())); } Some(Err(e)) => { return Err(ZerobusError::InvalidArgument(format!( @@ -168,9 +165,9 @@ fn materialize_ipc(bytes: &Bytes) -> ZerobusResult { Some(Ok(_)) => Err(ZerobusError::InvalidArgument( "IPC stream must contain exactly one RecordBatch (found extra batch)".into(), )), - Some(Err(e)) => Err(ZerobusError::InvalidArgument(format!( - "IPC: trailing message read failed: {e}" - ))), + Some(Err(e)) => { + Err(ZerobusError::InvalidArgument(format!("IPC: trailing message read failed: {e}"))) + } } } @@ -1700,12 +1697,10 @@ impl ZerobusArrowStream { if let Some(server_error) = error_rx.borrow().clone() { return Err(server_error); } - return Err(ZerobusError::StreamClosedError(tonic::Status::internal( - format!( - "Stream closing or closed during {}", - operation_name.to_lowercase() - ), - ))); + return Err(ZerobusError::StreamClosedError(tonic::Status::internal(format!( + "Stream closing or closed during {}", + operation_name.to_lowercase() + )))); } // Neither arm returns directly. After either watch changes, loop so the @@ -1906,10 +1901,7 @@ impl ZerobusArrowStream { result }; if let Err(e) = &flush_result { - warn!( - "Flush failed during close: {}. Draining pending batches to the failed set.", - e - ); + warn!("Flush failed during close: {}. Draining pending batches to the failed set.", e); } // Reap the supervisor (abort + await) BEFORE clearing the sender, so an in-flight @@ -2100,11 +2092,11 @@ impl Drop for ZerobusArrowStream { #[cfg(test)] mod tests { - use super::*; + use std::collections::HashMap; + use arrow_array::Int32Array; use arrow_schema::{DataType, Field}; use async_trait::async_trait; - use std::collections::HashMap; struct PassthroughTlsConfig; @@ -2129,6 +2121,8 @@ mod tests { } } + use super::*; + #[test] fn test_arrow_table_properties() { let schema = Arc::new(ArrowSchema::new(vec![ @@ -2146,11 +2140,7 @@ mod tests { } fn one_col_schema() -> Arc { - Arc::new(ArrowSchema::new(vec![Field::new( - "id", - DataType::Int32, - false, - )])) + Arc::new(ArrowSchema::new(vec![Field::new("id", DataType::Int32, false)])) } fn batch_with_rows(schema: &Arc, n: i32) -> RecordBatch { @@ -2184,11 +2174,7 @@ mod tests { pending_batch(&sem, batch_with_rows(&schema, 3), 0, 0, 3), pending_batch(&sem, batch_with_rows(&schema, 2), 1, 3, 5), ])); - assert_eq!( - sem.available_permits(), - 2, - "two permits held by pending batches" - ); + assert_eq!(sem.available_permits(), 2, "two permits held by pending batches"); // Stale values that must be overwritten by the atomic install. let cumulative = Arc::new(AtomicU64::new(999)); @@ -2204,11 +2190,7 @@ mod tests { assert!(res.is_err(), "replay must surface the send failure"); let guard = pending.lock().await; - assert_eq!( - guard.len(), - 2, - "pending must retain all batches on replay failure" - ); + assert_eq!(guard.len(), 2, "pending must retain all batches on replay failure"); assert_eq!((guard[0].start_record, guard[0].end_record), (0, 3)); assert_eq!((guard[1].start_record, guard[1].end_record), (3, 5)); drop(guard); @@ -2223,11 +2205,7 @@ mod tests { 0, "watermark must be rebased to 0 atomically with the ranges" ); - assert_eq!( - sem.available_permits(), - 2, - "permits must not be released on replay failure" - ); + assert_eq!(sem.available_permits(), 2, "permits must not be released on replay failure"); } /// With an open receiver, both batches remain pending, replay in order, and reset the @@ -2320,14 +2298,8 @@ mod tests { futures::poll!(fut.as_mut()).is_pending(), "pause_and_detach_sender must block while an ingest holds ingest_mutex" ); - assert!( - !is_paused.load(Ordering::Relaxed), - "is_paused flipped mid-ingest" - ); - assert!( - batch_tx.lock().await.is_some(), - "sender detached mid-ingest" - ); + assert!(!is_paused.load(Ordering::Relaxed), "is_paused flipped mid-ingest"); + assert!(batch_tx.lock().await.is_some(), "sender detached mid-ingest"); // Once the ingest leaves its critical section, the transition completes. drop(guard); @@ -2370,10 +2342,7 @@ mod tests { futures::poll!(fut.as_mut()).is_pending(), "finalize_closed must wait for the in-flight ingest" ); - assert!( - !is_closed.load(Ordering::Relaxed), - "is_closed must not be published mid-ingest" - ); + assert!(!is_closed.load(Ordering::Relaxed), "is_closed must not be published mid-ingest"); // The ingest appends its batch, then releases the mutex. let schema = one_col_schema(); diff --git a/rust/sdk/src/builder/sdk_builder.rs b/rust/sdk/src/builder/sdk_builder.rs index b8b264e9..a3d6b9ab 100644 --- a/rust/sdk/src/builder/sdk_builder.rs +++ b/rust/sdk/src/builder/sdk_builder.rs @@ -8,7 +8,12 @@ use crate::token_cache::DEFAULT_REFRESH_BUFFER; #[cfg(feature = "testing")] use crate::NoTlsConfig; use crate::{ - SecureTlsConfig, TlsConfig, ZerobusError, ZerobusResult, ZerobusSdk, DEFAULT_SDK_IDENTIFIER, + SecureTlsConfig, + TlsConfig, + ZerobusError, + ZerobusResult, + ZerobusSdk, + DEFAULT_SDK_IDENTIFIER, }; /// Builder for creating a [`ZerobusSdk`] instance with fluent configuration. @@ -299,10 +304,7 @@ mod tests { sdk.zerobus_endpoint, "https://my-workspace.zerobus.us-east-1.cloud.databricks.com" ); - assert_eq!( - sdk.unity_catalog_url, - "https://my-workspace.cloud.databricks.com" - ); + assert_eq!(sdk.unity_catalog_url, "https://my-workspace.cloud.databricks.com"); } #[test] @@ -351,10 +353,7 @@ mod tests { .expect("should build successfully with schemeless endpoint"); assert_eq!(sdk.workspace_id, "my-workspace"); - assert_eq!( - sdk.zerobus_endpoint, - "https://my-workspace.zerobus.databricks.com" - ); + assert_eq!(sdk.zerobus_endpoint, "https://my-workspace.zerobus.databricks.com"); } #[test] @@ -427,12 +426,7 @@ mod tests { #[test] fn test_application_name_with_invalid_header_bytes_is_rejected() { // Control bytes that tonic's `user-agent` header rejects. - for bad in [ - "my-app\n1.0", - "my-app\r1.0", - "my-app\u{0}1.0", - "my-app\u{7f}1.0", - ] { + for bad in ["my-app\n1.0", "my-app\r1.0", "my-app\u{0}1.0", "my-app\u{7f}1.0"] { let result = ZerobusSdkBuilder::new() .endpoint("https://workspace.zerobus.databricks.com") .application_name(bad) diff --git a/rust/sdk/src/builder/stream_builder.rs b/rust/sdk/src/builder/stream_builder.rs index abc8e668..9f4a32f9 100644 --- a/rust/sdk/src/builder/stream_builder.rs +++ b/rust/sdk/src/builder/stream_builder.rs @@ -21,6 +21,10 @@ use std::fmt; use std::sync::Arc; +#[cfg(feature = "arrow-flight")] +use crate::arrow_configuration::ArrowStreamConfigurationOptions; +#[cfg(feature = "arrow-flight")] +use crate::arrow_stream::{ArrowSchema, ArrowTableProperties, ZerobusArrowStream}; use crate::callbacks::AckCallback; use crate::databricks::zerobus::RecordType; #[cfg(feature = "testing")] @@ -29,11 +33,6 @@ use crate::headers_provider::{HeadersProvider, OAuthHeadersProvider}; use crate::stream_configuration::StreamConfigurationOptions; use crate::{TableProperties, ZerobusError, ZerobusResult, ZerobusSdk, ZerobusStream}; -#[cfg(feature = "arrow-flight")] -use crate::arrow_configuration::ArrowStreamConfigurationOptions; -#[cfg(feature = "arrow-flight")] -use crate::arrow_stream::{ArrowSchema, ArrowTableProperties, ZerobusArrowStream}; - /// Internal representation of the authentication configuration. enum AuthConfig { OAuth { @@ -490,9 +489,10 @@ impl<'a> StreamBuilder<'a> { #[cfg(test)] mod tests { - use super::*; use std::collections::HashMap; + use super::*; + fn test_sdk() -> ZerobusSdk { ZerobusSdk::new_with_config( "http://localhost:1234".to_string(), @@ -600,10 +600,7 @@ mod tests { .oauth("a", "b") .json() .max_ingest_payload_bytes(5 * 1024 * 1024); - assert_eq!( - builder.grpc_config.max_ingest_payload_bytes, - 5 * 1024 * 1024 - ); + assert_eq!(builder.grpc_config.max_ingest_payload_bytes, 5 * 1024 * 1024); } #[tokio::test] @@ -733,11 +730,7 @@ mod tests { use arrow_schema::{DataType, Field, Schema as ArrowSchema}; let sdk = test_sdk(); - let schema = Arc::new(ArrowSchema::new(vec![Field::new( - "id", - DataType::Int32, - false, - )])); + let schema = Arc::new(ArrowSchema::new(vec![Field::new("id", DataType::Int32, false)])); let _builder = sdk .stream_builder() .table("t") @@ -753,11 +746,7 @@ mod tests { use arrow_schema::{DataType, Field, Schema as ArrowSchema}; let sdk = test_sdk(); - let schema = Arc::new(ArrowSchema::new(vec![Field::new( - "id", - DataType::Int32, - false, - )])); + let schema = Arc::new(ArrowSchema::new(vec![Field::new("id", DataType::Int32, false)])); let builder = sdk .stream_builder() .table("t") @@ -776,9 +765,6 @@ mod tests { assert_eq!(builder.arrow_config.recovery_retries, 2); assert_eq!(builder.arrow_config.server_lack_of_ack_timeout_ms, 10_000); assert_eq!(builder.arrow_config.flush_timeout_ms, 20_000); - assert_eq!( - builder.arrow_config.stream_paused_max_wait_time_ms, - Some(5_000) - ); + assert_eq!(builder.arrow_config.stream_paused_max_wait_time_ms, Some(5_000)); } } diff --git a/rust/sdk/src/callbacks.rs b/rust/sdk/src/callbacks.rs index b98980b8..cb9fb045 100644 --- a/rust/sdk/src/callbacks.rs +++ b/rust/sdk/src/callbacks.rs @@ -86,9 +86,10 @@ pub trait AckCallback: Send + Sync { #[allow(dead_code)] #[cfg(test)] mod tests { - use super::*; use std::sync::atomic::{AtomicBool, AtomicI64, Ordering}; + use super::*; + struct TestCallback { last_ack: AtomicI64, error_called: AtomicBool, diff --git a/rust/sdk/src/client_warnings.rs b/rust/sdk/src/client_warnings.rs index 3b3cf4d2..ce5917d8 100644 --- a/rust/sdk/src/client_warnings.rs +++ b/rust/sdk/src/client_warnings.rs @@ -303,9 +303,6 @@ mod tests { #[test] fn churn_unknown_table_returns_zero() { - assert_eq!( - open_count_in_window_for_testing("cat.sch.churn_unknown_xyz"), - 0 - ); + assert_eq!(open_count_in_window_for_testing("cat.sch.churn_unknown_xyz"), 0); } } diff --git a/rust/sdk/src/default_token_factory.rs b/rust/sdk/src/default_token_factory.rs index 0b178cd7..77a29c82 100644 --- a/rust/sdk/src/default_token_factory.rs +++ b/rust/sdk/src/default_token_factory.rs @@ -181,11 +181,8 @@ impl DefaultTokenFactory { ("scope", "all-apis".to_string()), ( "resource", - format!( - "api://databricks/workspaces/{}/zerobusDirectWriteApi", - workspace_id - ) - .to_string(), + format!("api://databricks/workspaces/{}/zerobusDirectWriteApi", workspace_id) + .to_string(), ), ("authorization_details", authorization_details.to_string()), ]; @@ -320,19 +317,13 @@ impl DefaultTokenFactory { let table = parts[2]; if catalog.is_empty() { - return Err(ZerobusError::InvalidTableName( - "Catalog name cannot be empty".to_string(), - )); + return Err(ZerobusError::InvalidTableName("Catalog name cannot be empty".to_string())); } if schema.is_empty() { - return Err(ZerobusError::InvalidTableName( - "Schema name cannot be empty".to_string(), - )); + return Err(ZerobusError::InvalidTableName("Schema name cannot be empty".to_string())); } if table.is_empty() { - return Err(ZerobusError::InvalidTableName( - "Table name cannot be empty".to_string(), - )); + return Err(ZerobusError::InvalidTableName("Table name cannot be empty".to_string())); } Ok((catalog.to_string(), schema.to_string(), table.to_string())) diff --git a/rust/sdk/src/errors.rs b/rust/sdk/src/errors.rs index 5686fd7e..c9346fbf 100644 --- a/rust/sdk/src/errors.rs +++ b/rust/sdk/src/errors.rs @@ -235,18 +235,12 @@ mod tests { #[test] fn auth_rejection_classification() { - assert!( - ZerobusError::CreateStreamError(tonic::Status::unauthenticated("x")) - .is_auth_rejection() - ); - assert!( - ZerobusError::CreateStreamError(tonic::Status::permission_denied("x")) - .is_auth_rejection() - ); - assert!( - ZerobusError::StreamClosedError(tonic::Status::unauthenticated("x")) - .is_auth_rejection() - ); + assert!(ZerobusError::CreateStreamError(tonic::Status::unauthenticated("x")) + .is_auth_rejection()); + assert!(ZerobusError::CreateStreamError(tonic::Status::permission_denied("x")) + .is_auth_rejection()); + assert!(ZerobusError::StreamClosedError(tonic::Status::unauthenticated("x")) + .is_auth_rejection()); // Non-auth gRPC codes are not rejections. assert!(!ZerobusError::CreateStreamError(tonic::Status::internal("x")).is_auth_rejection()); assert!( @@ -310,6 +304,7 @@ mod tests { #[cfg(feature = "arrow-flight")] fn schema_validation_status(causes: &str) -> tonic::Status { use std::collections::HashMap; + use tonic_types::ErrorDetails; let mut metadata = HashMap::new(); @@ -392,6 +387,7 @@ mod tests { #[test] fn setup_status_with_wrong_domain_falls_back_to_create_stream_error() { use std::collections::HashMap; + use tonic_types::ErrorDetails; let mut metadata = HashMap::new(); diff --git a/rust/sdk/src/headers_provider.rs b/rust/sdk/src/headers_provider.rs index 9c361f04..691ea1d8 100644 --- a/rust/sdk/src/headers_provider.rs +++ b/rust/sdk/src/headers_provider.rs @@ -1,9 +1,11 @@ +use std::collections::HashMap; +use std::sync::Arc; + +use async_trait::async_trait; + use crate::default_token_factory::DefaultTokenFactory; use crate::token_cache::{TokenCache, DEFAULT_REFRESH_BUFFER}; use crate::ZerobusResult; -use async_trait::async_trait; -use std::collections::HashMap; -use std::sync::Arc; /// A trait for providing custom headers for gRPC requests. /// @@ -123,21 +125,16 @@ impl HeadersProvider for OAuthHeadersProvider { async fn get_headers(&self) -> ZerobusResult> { let token = self .token_cache - .get_or_fetch( - &self.client_id, - &self.client_secret, - &self.table_name, - |reason| { - DefaultTokenFactory::fetch_token( - &self.unity_catalog_url, - &self.table_name, - &self.client_id, - &self.client_secret, - &self.workspace_id, - reason, - ) - }, - ) + .get_or_fetch(&self.client_id, &self.client_secret, &self.table_name, |reason| { + DefaultTokenFactory::fetch_token( + &self.unity_catalog_url, + &self.table_name, + &self.client_id, + &self.client_secret, + &self.workspace_id, + reason, + ) + }) .await?; let mut headers = HashMap::new(); headers.insert("authorization", format!("Bearer {}", token)); diff --git a/rust/sdk/src/landing_zone.rs b/rust/sdk/src/landing_zone.rs index f5e61daa..5802f49e 100644 --- a/rust/sdk/src/landing_zone.rs +++ b/rust/sdk/src/landing_zone.rs @@ -276,10 +276,7 @@ mod tests { let lz = Arc::new(LandingZone::::new(10)); let result = lz.remove_observed(); - assert!(matches!( - result, - Err(LandingZoneError::RemovingNonObservedElement) - )); + assert!(matches!(result, Err(LandingZoneError::RemovingNonObservedElement))); } #[tokio::test] diff --git a/rust/sdk/src/lib.rs b/rust/sdk/src/lib.rs index beba3ac3..416ce86b 100644 --- a/rust/sdk/src/lib.rs +++ b/rust/sdk/src/lib.rs @@ -79,8 +79,15 @@ pub use multiplexed_stream::{MessageId, MultiplexedStream}; pub use offset_generator::{OffsetId, OffsetIdGenerator}; pub use proxy::{ConnectorFactory, ProxyConnector}; pub use record_types::{ - EncodedBatch, EncodedBatchIter, EncodedRecord, JsonEncodedRecord, JsonString, JsonValue, - ProtoBytes, ProtoEncodedRecord, ProtoMessage, + EncodedBatch, + EncodedBatchIter, + EncodedRecord, + JsonEncodedRecord, + JsonString, + JsonValue, + ProtoBytes, + ProtoEncodedRecord, + ProtoMessage, }; pub use sdk::{ZerobusSdk, DEFAULT_SDK_IDENTIFIER}; #[cfg(feature = "testing")] diff --git a/rust/sdk/src/multiplexed_stream.rs b/rust/sdk/src/multiplexed_stream.rs index 4abbbb0d..537f0d41 100644 --- a/rust/sdk/src/multiplexed_stream.rs +++ b/rust/sdk/src/multiplexed_stream.rs @@ -18,13 +18,20 @@ //! [`get_unacked_records`](MultiplexedStream::get_unacked_records) or //! [`get_unacked_batches`](MultiplexedStream::get_unacked_batches). -use futures::future::join_all; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::sync::Arc; + +use futures::future::join_all; use tracing::{error, info, warn}; use crate::{ - AckCallback, EncodedBatch, EncodedRecord, OffsetId, ZerobusError, ZerobusResult, ZerobusStream, + AckCallback, + EncodedBatch, + EncodedRecord, + OffsetId, + ZerobusError, + ZerobusResult, + ZerobusStream, }; /// Number of bits reserved for the stream index. @@ -42,12 +49,7 @@ pub struct MessageId(i64); impl std::fmt::Display for MessageId { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!( - f, - "MessageId(stream={}, offset={})", - self.stream_index(), - self.sub_offset() - ) + write!(f, "MessageId(stream={}, offset={})", self.stream_index(), self.sub_offset()) } } @@ -137,10 +139,7 @@ impl MultiplexedStream { /// /// Panics if `streams` is empty or holds more than 64 sub-streams. pub fn new(streams: Vec) -> Self { - assert!( - !streams.is_empty(), - "MultiplexedStream requires at least one sub-stream" - ); + assert!(!streams.is_empty(), "MultiplexedStream requires at least one sub-stream"); assert!( streams.len() <= (1 << STREAM_BITS), "MultiplexedStream supports at most {} sub-streams", @@ -156,9 +155,7 @@ impl MultiplexedStream { #[allow(clippy::result_large_err)] fn check_closed(&self) -> ZerobusResult<()> { if self.is_closed_fast() { - return Err(ZerobusError::InvalidStateError( - "MultiplexedStream is closed".to_string(), - )); + return Err(ZerobusError::InvalidStateError("MultiplexedStream is closed".to_string())); } Ok(()) } @@ -476,9 +473,10 @@ impl Drop for MultiplexedStream { #[cfg(test)] mod tests { - use super::*; use std::sync::Mutex; + use super::*; + #[derive(Default)] struct RecordingMultiplexedCallback { acks: Mutex>, diff --git a/rust/sdk/src/proxy.rs b/rust/sdk/src/proxy.rs index 8b39983e..c5a4111f 100644 --- a/rust/sdk/src/proxy.rs +++ b/rust/sdk/src/proxy.rs @@ -65,14 +65,8 @@ fn build_connector(proxy_uri: &str) -> Result { pub type ConnectorFactory = Arc Option + Send + Sync>; /// Env var names checked for proxy URL, in gRPC core precedence order. -const PROXY_ENV_VARS: &[&str] = &[ - "grpc_proxy", - "GRPC_PROXY", - "https_proxy", - "HTTPS_PROXY", - "http_proxy", - "HTTP_PROXY", -]; +const PROXY_ENV_VARS: &[&str] = + &["grpc_proxy", "GRPC_PROXY", "https_proxy", "HTTPS_PROXY", "http_proxy", "HTTP_PROXY"]; /// Env var names checked for no-proxy list, in gRPC core precedence order. const NO_PROXY_ENV_VARS: &[&str] = &["no_grpc_proxy", "NO_GRPC_PROXY", "no_proxy", "NO_PROXY"]; @@ -161,10 +155,7 @@ mod tests { #[test] fn no_proxy_suffix_match() { - assert!(host_matches_no_proxy( - "workspace.cloud.databricks.com", - "databricks.com" - )); + assert!(host_matches_no_proxy("workspace.cloud.databricks.com", "databricks.com")); assert!(host_matches_no_proxy("foo.example.com", "example.com")); // Must be a subdomain, not just a string suffix assert!(!host_matches_no_proxy("notexample.com", "example.com")); @@ -188,9 +179,6 @@ mod tests { #[test] fn no_proxy_whitespace_handling() { assert!(host_matches_no_proxy("example.com", " example.com ")); - assert!(host_matches_no_proxy( - "example.com", - "other.com , example.com , more.com" - )); + assert!(host_matches_no_proxy("example.com", "other.com , example.com , more.com")); } } diff --git a/rust/sdk/src/record_types.rs b/rust/sdk/src/record_types.rs index e805af37..1ac3ce22 100644 --- a/rust/sdk/src/record_types.rs +++ b/rust/sdk/src/record_types.rs @@ -12,11 +12,15 @@ use prost::Message; use smallvec::{smallvec, SmallVec}; +use crate::databricks::zerobus::ephemeral_stream_request::Payload as RequestPayload; +use crate::databricks::zerobus::ingest_record_batch_request::Batch as IngestRequestBatch; +use crate::databricks::zerobus::ingest_record_request::Record as IngestRequestRecord; use crate::databricks::zerobus::{ - ephemeral_stream_request::Payload as RequestPayload, - ingest_record_batch_request::Batch as IngestRequestBatch, - ingest_record_request::Record as IngestRequestRecord, IngestRecordBatchRequest, - IngestRecordRequest, JsonRecordBatch, ProtoEncodedRecordBatch, RecordType, + IngestRecordBatchRequest, + IngestRecordRequest, + JsonRecordBatch, + ProtoEncodedRecordBatch, + RecordType, }; use crate::OffsetId; @@ -191,28 +195,26 @@ impl EncodedBatch { match record_type { RecordType::Json => batch_iter - .try_fold( - SmallVec::with_capacity(size_hint), - |mut vec, record| match record.into() { + .try_fold(SmallVec::with_capacity(size_hint), |mut vec, record| { + match record.into() { EncodedRecord::Json(value) => { vec.push(value); Some(vec) } _ => None, - }, - ) + } + }) .map(EncodedBatch::Json), RecordType::Proto => batch_iter - .try_fold( - SmallVec::with_capacity(size_hint), - |mut vec, record| match record.into() { + .try_fold(SmallVec::with_capacity(size_hint), |mut vec, record| { + match record.into() { EncodedRecord::Proto(value) => { vec.push(value); Some(vec) } _ => None, - }, - ) + } + }) .map(EncodedBatch::Proto), _ => None, } @@ -230,11 +232,9 @@ impl EncodedBatch { } EncodedBatch::Proto(records) => { RequestPayload::IngestRecordBatch(IngestRecordBatchRequest { - batch: Some(IngestRequestBatch::ProtoEncodedBatch( - ProtoEncodedRecordBatch { - records: records.into_vec(), - }, - )), + batch: Some(IngestRequestBatch::ProtoEncodedBatch(ProtoEncodedRecordBatch { + records: records.into_vec(), + })), offset_id: Some(offset_id), }) } @@ -315,11 +315,12 @@ impl Iterator for EncodedBatchIter { #[cfg(test)] mod tests { - use super::*; use prost::Message as ProstMessage; use serde::Serialize; use smallvec::smallvec; + use super::*; + #[derive(Clone, PartialEq, ProstMessage)] struct TestMessage { #[prost(string, tag = "1")] diff --git a/rust/sdk/src/schema.rs b/rust/sdk/src/schema.rs index 12922edf..af70486d 100644 --- a/rust/sdk/src/schema.rs +++ b/rust/sdk/src/schema.rs @@ -367,10 +367,7 @@ fn parse_type_json(type_json: &str) -> Result { fn type_ref_to_complex(tref: &TypeRef, level: usize) -> Result { if level > MAX_NESTING_DEPTH { - return Err(format!( - "nesting level exceeds maximum depth of {}", - MAX_NESTING_DEPTH - )); + return Err(format!("nesting level exceeds maximum depth of {}", MAX_NESTING_DEPTH)); } match tref { TypeRef::Primitive(s) => parse_primitive_type(s).map(ComplexType::Primitive), @@ -385,9 +382,9 @@ fn type_ref_to_complex(tref: &TypeRef, level: usize) -> Result Ok(ComplexType::Array( - Box::new(type_ref_to_complex(element_type, level + 1)?), - )), + TypeRef::Complex(ComplexTypeJson::Array { element_type }) => { + Ok(ComplexType::Array(Box::new(type_ref_to_complex(element_type, level + 1)?))) + } TypeRef::Complex(ComplexTypeJson::Map { key_type, value_type, @@ -436,10 +433,7 @@ const fn map_primitive_to_protobuf(p: PrimitiveType) -> ProtoType { } const fn is_valid_map_key(p: PrimitiveType) -> bool { - !matches!( - p, - PrimitiveType::Double | PrimitiveType::Float | PrimitiveType::Binary - ) + !matches!(p, PrimitiveType::Double | PrimitiveType::Float | PrimitiveType::Binary) } fn validate_map_key(key: &ComplexType, path: &str) -> Result { @@ -562,10 +556,7 @@ fn generate_struct_message( let path = format!("{}_{}", message_name, f.name); let (field_type, type_name) = map_complex_type_to_protobuf(&f.field_type, &path, &mut local)?; - let is_repeated = matches!( - f.field_type, - ComplexType::Array(_) | ComplexType::Map { .. } - ); + let is_repeated = matches!(f.field_type, ComplexType::Array(_) | ComplexType::Map { .. }); fields.push(field_descriptor( &f.name, (index + 1) as i32, @@ -733,11 +724,7 @@ fn uc_column_to_arrow_field(column: &UcColumn) -> Result Result { - use arrow_schema::{DataType, Field, Fields}; use std::sync::Arc; + use arrow_schema::{DataType, Field, Fields}; + match ct { ComplexType::Primitive(p) => Ok(Field::new(name, map_primitive_to_arrow(*p), nullable)), ComplexType::Struct(st) => { let mut child_fields = Vec::with_capacity(st.fields.len()); for f in &st.fields { validate_field_name(&f.name)?; - child_fields.push(complex_type_to_arrow_field( - &f.name, - &f.field_type, - f.nullable, - )?); + child_fields.push(complex_type_to_arrow_field(&f.name, &f.field_type, f.nullable)?); } - Ok(Field::new( - name, - DataType::Struct(Fields::from(child_fields)), - nullable, - )) + Ok(Field::new(name, DataType::Struct(Fields::from(child_fields)), nullable)) } ComplexType::Array(element) => { // UC's `containsNull` is not surfaced in our AST; default to @@ -826,11 +806,7 @@ fn complex_type_to_arrow_field( ComplexType::Array(_) => return Err(shape_unsupported("nested arrays", name)), ComplexType::Map { .. } => return Err(shape_unsupported("arrays of maps", name)), }; - Ok(Field::new( - name, - DataType::List(Arc::new(item_field)), - nullable, - )) + Ok(Field::new(name, DataType::List(Arc::new(item_field)), nullable)) } ComplexType::Map { key, value } => { let key_primitive = validate_arrow_map_key(key, name)?; @@ -1100,20 +1076,14 @@ mod tests { // should bail out with an InvalidTypeJson rather than overflowing the stack. let mut type_json = String::from("\"integer\""); for _ in 0..MAX_NESTING_DEPTH + 2 { - type_json = format!( - r#"{{"type":"array","elementType":{},"containsNull":true}}"#, - type_json - ); + type_json = + format!(r#"{{"type":"array","elementType":{},"containsNull":true}}"#, type_json); } let cols = vec![complex_col("deep", "ARRAY", &type_json, 0)]; let err = descriptor_from_uc_columns(&cols, "m").unwrap_err(); match err { SchemaError::InvalidTypeJson { reason, .. } => { - assert!( - reason.contains("maximum depth"), - "unexpected reason: {}", - reason - ); + assert!(reason.contains("maximum depth"), "unexpected reason: {}", reason); } other => panic!("expected InvalidTypeJson, got {:?}", other), } @@ -1251,9 +1221,10 @@ mod tests { #[cfg(feature = "arrow-flight")] mod arrow { - use super::*; use arrow_schema::{DataType, TimeUnit}; + use super::*; + fn arrow_field<'a>( schema: &'a arrow_schema::Schema, name: &str, diff --git a/rust/sdk/src/sdk.rs b/rust/sdk/src/sdk.rs index 067035c9..760d3068 100644 --- a/rust/sdk/src/sdk.rs +++ b/rust/sdk/src/sdk.rs @@ -10,14 +10,13 @@ use std::time::Duration; use tonic::transport::{Channel, Endpoint}; use tracing::{error, info, instrument}; +#[cfg(feature = "arrow-flight")] +use crate::arrow_stream::ZerobusArrowStream; use crate::databricks::zerobus::zerobus_client::ZerobusClient; use crate::proxy::{self, ConnectorFactory, ProxyConnector}; use crate::stream::ZerobusStream; use crate::{StreamBuilder, TlsConfig, ZerobusError, ZerobusResult, ZerobusSdkBuilder}; -#[cfg(feature = "arrow-flight")] -use crate::arrow_stream::ZerobusArrowStream; - /// Default identifier the SDK sends as the HTTP `user-agent` header on every /// request. Use [`ZerobusSdkBuilder::application_name`] to append an /// application suffix. diff --git a/rust/sdk/src/stream/grpc/acks.rs b/rust/sdk/src/stream/grpc/acks.rs index 1a4e8839..c835c140 100644 --- a/rust/sdk/src/stream/grpc/acks.rs +++ b/rust/sdk/src/stream/grpc/acks.rs @@ -65,9 +65,10 @@ impl ZerobusStream { if let Some(server_error) = error_rx.borrow().clone() { return Err(server_error); } - return Err(ZerobusError::StreamClosedError(tonic::Status::internal( - format!("Stream closed during {}", operation_name.to_lowercase()), - ))); + return Err(ZerobusError::StreamClosedError(tonic::Status::internal(format!( + "Stream closed during {}", + operation_name.to_lowercase() + )))); } // Race between offset updates and server errors. tokio::select! { @@ -102,9 +103,10 @@ impl ZerobusStream { } } - Err(ZerobusError::StreamClosedError(tonic::Status::internal( - format!("Stream closed during {}", operation_name.to_lowercase()), - ))) + Err(ZerobusError::StreamClosedError(tonic::Status::internal(format!( + "Stream closed during {}", + operation_name.to_lowercase() + )))) }; match tokio::time::timeout( @@ -121,9 +123,10 @@ impl ZerobusStream { } else { error!(table_name = %self.table_properties.table_name, "{} timed out", operation_name); } - Err(ZerobusError::StreamClosedError( - tonic::Status::deadline_exceeded(format!("{} timed out", operation_name)), - )) + Err(ZerobusError::StreamClosedError(tonic::Status::deadline_exceeded(format!( + "{} timed out", + operation_name + )))) } } } diff --git a/rust/sdk/src/stream/grpc/connection.rs b/rust/sdk/src/stream/grpc/connection.rs index 685c748b..de186045 100644 --- a/rust/sdk/src/stream/grpc/connection.rs +++ b/rust/sdk/src/stream/grpc/connection.rs @@ -17,7 +17,10 @@ use crate::databricks::zerobus::ephemeral_stream_request::Payload as RequestPayl use crate::databricks::zerobus::ephemeral_stream_response::Payload as ResponsePayload; use crate::databricks::zerobus::zerobus_client::ZerobusClient; use crate::databricks::zerobus::{ - CreateIngestStreamRequest, EphemeralStreamRequest, EphemeralStreamResponse, RecordType, + CreateIngestStreamRequest, + EphemeralStreamRequest, + EphemeralStreamResponse, + RecordType, }; use crate::{HeadersProvider, TableProperties, ZerobusError, ZerobusResult}; diff --git a/rust/sdk/src/stream/grpc/ingest.rs b/rust/sdk/src/stream/grpc/ingest.rs index ca4c9e89..3d5abcd3 100644 --- a/rust/sdk/src/stream/grpc/ingest.rs +++ b/rust/sdk/src/stream/grpc/ingest.rs @@ -7,6 +7,7 @@ use std::future::Future; use std::sync::atomic::Ordering; + use tracing::{debug, error}; use super::types::IngestRequest; @@ -130,9 +131,7 @@ impl ZerobusStream { ) -> ZerobusResult>> { if self.is_closed.load(Ordering::Relaxed) { error!(table_name = %self.table_properties.table_name, "Stream closed"); - return Err(ZerobusError::StreamClosedError(tonic::Status::internal( - "Stream closed", - ))); + return Err(ZerobusError::StreamClosedError(tonic::Status::internal("Stream closed"))); } let _guard = self.sync_mutex.lock().await; @@ -167,9 +166,7 @@ impl ZerobusStream { }) } else { error!("Stream ID is None"); - Err(ZerobusError::StreamClosedError(tonic::Status::internal( - "Stream ID is None", - ))) + Err(ZerobusError::StreamClosedError(tonic::Status::internal("Stream ID is None"))) } } @@ -188,9 +185,7 @@ impl ZerobusStream { if self.is_closed.load(Ordering::Relaxed) { error!(table_name = %self.table_properties.table_name, "Stream closed"); - return Err(ZerobusError::StreamClosedError(tonic::Status::internal( - "Stream closed", - ))); + return Err(ZerobusError::StreamClosedError(tonic::Status::internal("Stream closed"))); } let _guard = self.sync_mutex.lock().await; diff --git a/rust/sdk/src/stream/grpc/mod.rs b/rust/sdk/src/stream/grpc/mod.rs index 69a84457..1557c256 100644 --- a/rust/sdk/src/stream/grpc/mod.rs +++ b/rust/sdk/src/stream/grpc/mod.rs @@ -31,8 +31,14 @@ use tracing::instrument; use crate::databricks::zerobus::zerobus_client::ZerobusClient; use crate::landing_zone::LandingZone; use crate::{ - HeadersProvider, OffsetId, OffsetIdGenerator, StreamConfigurationOptions, StreamType, - TableProperties, ZerobusError, ZerobusResult, + HeadersProvider, + OffsetId, + OffsetIdGenerator, + StreamConfigurationOptions, + StreamType, + TableProperties, + ZerobusError, + ZerobusResult, }; mod acks; @@ -45,10 +51,9 @@ mod sender; mod supervisor; mod types; -use types::{IngestRequest, OneshotMap, RecordLandingZone}; - #[cfg(feature = "testing")] pub use callback_handler::CallbackHandlerHarness; +use types::{IngestRequest, OneshotMap, RecordLandingZone}; /// Maximum time to wait for the receiver/sender tasks to finish during stream /// teardown. @@ -144,9 +149,8 @@ impl ZerobusStream { let (logical_last_received_offset_id_tx, _logical_last_received_offset_id_rx) = tokio::sync::watch::channel(None); - let landing_zone = Arc::new(LandingZone::>::new( - options.max_inflight_requests, - )); + let landing_zone = + Arc::new(LandingZone::>::new(options.max_inflight_requests)); let oneshot_map = Arc::new(tokio::sync::Mutex::new(HashMap::new())); let is_closed = Arc::new(AtomicBool::new(false)); diff --git a/rust/sdk/src/stream/grpc/receiver.rs b/rust/sdk/src/stream/grpc/receiver.rs index 206f99cc..1dc0eb57 100644 --- a/rust/sdk/src/stream/grpc/receiver.rs +++ b/rust/sdk/src/stream/grpc/receiver.rs @@ -16,7 +16,9 @@ use super::types::{CallbackMessage, OneshotMap, RecordLandingZone}; use super::{ZerobusStream, STREAM_TEARDOWN_DRAIN_TIMEOUT_MS}; use crate::databricks::zerobus::ephemeral_stream_response::Payload as ResponsePayload; use crate::databricks::zerobus::{ - CloseStreamSignal, EphemeralStreamResponse, IngestRecordResponse, + CloseStreamSignal, + EphemeralStreamResponse, + IngestRecordResponse, }; use crate::{OffsetId, StreamConfigurationOptions, ZerobusError, ZerobusResult}; diff --git a/rust/sdk/src/stream/grpc/supervisor.rs b/rust/sdk/src/stream/grpc/supervisor.rs index b08817fa..dc431e28 100644 --- a/rust/sdk/src/stream/grpc/supervisor.rs +++ b/rust/sdk/src/stream/grpc/supervisor.rs @@ -22,8 +22,13 @@ use super::types::{CallbackMessage, OneshotMap, RecordLandingZone}; use super::{ZerobusStream, STREAM_TEARDOWN_DRAIN_TIMEOUT_MS}; use crate::databricks::zerobus::zerobus_client::ZerobusClient; use crate::{ - EncodedBatch, HeadersProvider, OffsetId, StreamConfigurationOptions, TableProperties, - ZerobusError, ZerobusResult, + EncodedBatch, + HeadersProvider, + OffsetId, + StreamConfigurationOptions, + TableProperties, + ZerobusError, + ZerobusResult, }; impl ZerobusStream { @@ -309,10 +314,7 @@ impl ZerobusStream { let _ = sender.send(Err(error.clone())); } if let Some(tx) = callback_tx { - let _ = tx.send(CallbackMessage::Error( - record.offset_id, - error_message.clone(), - )); + let _ = tx.send(CallbackMessage::Error(record.offset_id, error_message.clone())); } } } diff --git a/rust/sdk/src/stream/mod.rs b/rust/sdk/src/stream/mod.rs index 1a6d4f7c..e174754f 100644 --- a/rust/sdk/src/stream/mod.rs +++ b/rust/sdk/src/stream/mod.rs @@ -9,7 +9,6 @@ mod grpc; -pub use grpc::ZerobusStream; - #[cfg(feature = "testing")] pub use grpc::CallbackHandlerHarness; +pub use grpc::ZerobusStream; diff --git a/rust/sdk/src/tls_config.rs b/rust/sdk/src/tls_config.rs index a65f39d6..be7a0212 100644 --- a/rust/sdk/src/tls_config.rs +++ b/rust/sdk/src/tls_config.rs @@ -3,9 +3,10 @@ //! This module provides a strategy pattern for TLS configuration, //! allowing different TLS setups (secure, custom CA, or no TLS for testing). +use tonic::transport::{ClientTlsConfig, Endpoint}; + use crate::errors::ZerobusError; use crate::ZerobusResult; -use tonic::transport::{ClientTlsConfig, Endpoint}; /// Trait for TLS configuration strategies. /// diff --git a/rust/sdk/src/token_cache.rs b/rust/sdk/src/token_cache.rs index 8b99ce67..6410c605 100644 --- a/rust/sdk/src/token_cache.rs +++ b/rust/sdk/src/token_cache.rs @@ -221,9 +221,10 @@ impl TokenCache { #[cfg(test)] mod tests { - use super::*; use std::sync::atomic::{AtomicUsize, Ordering}; + use super::*; + fn fetched(token: &str, ttl_secs: Option) -> FetchedToken { FetchedToken { token: token.to_string(), @@ -252,11 +253,7 @@ mod tests { assert_eq!(a, "tok"); assert_eq!(b, "tok"); - assert_eq!( - calls.load(Ordering::SeqCst), - 1, - "second call should hit cache" - ); + assert_eq!(calls.load(Ordering::SeqCst), 1, "second call should hit cache"); } #[tokio::test] @@ -461,15 +458,10 @@ mod tests { // must surface rather than being masked by the still-valid cached token. let result = cache .get_or_fetch("id", "secret", "c.s.t", |_reason| async { - Err(crate::ZerobusError::InvalidUCTokenError( - "revoked".to_string(), - )) + Err(crate::ZerobusError::InvalidUCTokenError("revoked".to_string())) }) .await; - assert!(matches!( - result, - Err(crate::ZerobusError::InvalidUCTokenError(_)) - )); + assert!(matches!(result, Err(crate::ZerobusError::InvalidUCTokenError(_)))); } #[tokio::test] @@ -487,9 +479,7 @@ mod tests { // A refresh returns a token with no TTL: the caller gets the fresh token, // but the cached valid token must not be discarded. let fresh = cache - .get_or_fetch("id", "secret", "c.s.t", |_reason| async { - Ok(fetched("nottl", None)) - }) + .get_or_fetch("id", "secret", "c.s.t", |_reason| async { Ok(fetched("nottl", None)) }) .await .unwrap(); assert_eq!(fresh, "nottl"); diff --git a/rust/sdk/src/zeroparser/benches/bench_plot.rs b/rust/sdk/src/zeroparser/benches/bench_plot.rs index 4116a160..6dd5adb8 100644 --- a/rust/sdk/src/zeroparser/benches/bench_plot.rs +++ b/rust/sdk/src/zeroparser/benches/bench_plot.rs @@ -5,16 +5,19 @@ use std::path::PathBuf; use std::time::{Duration, Instant}; use common::{ - bench_prost_reflect_decode, bench_prost_typed_decode, bench_zeroparser_decode, - create_sized_message, format_bytes, load_bench_sample, BenchmarkConfig, + bench_prost_reflect_decode, + bench_prost_typed_decode, + bench_zeroparser_decode, + create_sized_message, + format_bytes, + load_bench_sample, + BenchmarkConfig, }; use plotters::prelude::*; use plotters::style::text_anchor::{HPos, Pos, VPos}; -const OUTPUT_PATH: &str = concat!( - env!("CARGO_MANIFEST_DIR"), - "/src/zeroparser/benches/bench_plot.svg" -); +const OUTPUT_PATH: &str = + concat!(env!("CARGO_MANIFEST_DIR"), "/src/zeroparser/benches/bench_plot.svg"); const MIN_MEASURE_SECS: f64 = 1.0; const MAX_ITERATIONS: u64 = 1_000; const TRIALS_PER_MEASUREMENT: usize = 3; @@ -180,13 +183,11 @@ fn measure_scenario( let messages = vec![single; scenario.count]; - let reflect_mbps = measure_mbps(&messages, |m| { - bench_prost_reflect_decode(&scenario.config.msg_desc, m) - }); + let reflect_mbps = + measure_mbps(&messages, |m| bench_prost_reflect_decode(&scenario.config.msg_desc, m)); - let prost_typed_mbps = measure_mbps(&messages, |m| { - bench_prost_typed_decode(scenario.config.prost_typed, m) - }); + let prost_typed_mbps = + measure_mbps(&messages, |m| bench_prost_typed_decode(scenario.config.prost_typed, m)); let zeroparser_mbps = measure_mbps(&messages, |m| { bench_zeroparser_decode(&scenario.config.registry, &scenario.config.fields, m) @@ -263,9 +264,7 @@ fn draw_plot(results: &[Measurement]) -> Result<(), Box> let bars_per_group: &[Bar] = &[ ("prost-reflect", RGBColor(220, 80, 80), |m| m.reflect_mbps), ("prost", RGBColor(235, 145, 70), |m| m.prost_typed_mbps), - ("C++ reflect", RGBColor(120, 100, 180), |m| { - m.cpp_reflect_mbps - }), + ("C++ reflect", RGBColor(120, 100, 180), |m| m.cpp_reflect_mbps), ("C++ typed", RGBColor(95, 165, 110), |m| m.cpp_typed_mbps), ("Zeroparser", RGBColor(60, 130, 200), |m| m.zeroparser_mbps), ]; @@ -343,10 +342,7 @@ fn draw_plot(results: &[Measurement]) -> Result<(), Box> for (name, color, _) in bars_per_group { let c = *color; chart - .draw_series(std::iter::once(Rectangle::new( - [(0.0, 0.0), (0.0, 0.0)], - c.filled(), - )))? + .draw_series(std::iter::once(Rectangle::new([(0.0, 0.0), (0.0, 0.0)], c.filled())))? .label(*name) .legend(move |(x, y)| Rectangle::new([(x, y - 6), (x + 18, y + 6)], c.filled())); } diff --git a/rust/sdk/src/zeroparser/benches/common/mod.rs b/rust/sdk/src/zeroparser/benches/common/mod.rs index ea4f15bf..b26e7f1d 100644 --- a/rust/sdk/src/zeroparser/benches/common/mod.rs +++ b/rust/sdk/src/zeroparser/benches/common/mod.rs @@ -15,22 +15,13 @@ pub const SAMPLE_DATA_JSON: &str = include_str!("../bench_sample_data.json"); pub mod proto { pub mod air_quality { - include!(concat!( - env!("OUT_DIR"), - "/zeroparser.benches.air_quality.rs" - )); + include!(concat!(env!("OUT_DIR"), "/zeroparser.benches.air_quality.rs")); } pub mod wide_schema { - include!(concat!( - env!("OUT_DIR"), - "/zeroparser.benches.wide_schema.rs" - )); + include!(concat!(env!("OUT_DIR"), "/zeroparser.benches.wide_schema.rs")); } pub mod supported_nullable_types { - include!(concat!( - env!("OUT_DIR"), - "/zeroparser.benches.supported_nullable_types.rs" - )); + include!(concat!(env!("OUT_DIR"), "/zeroparser.benches.supported_nullable_types.rs")); } } @@ -159,11 +150,7 @@ impl BenchmarkConfig { fn find_message_and_file<'a>( file_desc_set: &'a FileDescriptorSet, message_name: &str, -) -> ( - &'a DescriptorProto, - &'a prost_types::FileDescriptorProto, - &'a str, -) { +) -> (&'a DescriptorProto, &'a prost_types::FileDescriptorProto, &'a str) { for file in &file_desc_set.file { for msg_desc in &file.message_type { if msg_desc.name.as_deref() == Some(message_name) { diff --git a/rust/sdk/src/zeroparser/benches/parser_bench.rs b/rust/sdk/src/zeroparser/benches/parser_bench.rs index ef9e23bb..540db1c9 100644 --- a/rust/sdk/src/zeroparser/benches/parser_bench.rs +++ b/rust/sdk/src/zeroparser/benches/parser_bench.rs @@ -1,8 +1,13 @@ mod common; use common::{ - bench_prost_reflect_decode, bench_prost_typed_decode, bench_zeroparser_decode, - create_encoded_messages, format_bytes, load_bench_sample, BenchmarkConfig, + bench_prost_reflect_decode, + bench_prost_typed_decode, + bench_zeroparser_decode, + create_encoded_messages, + format_bytes, + load_bench_sample, + BenchmarkConfig, }; use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput}; @@ -62,11 +67,8 @@ fn run_decode_benchmark(c: &mut Criterion, scenario: BenchScenario) { ); let actual_message_size = encoded_messages.first().map(|m| m.len()).unwrap_or(0); let total_bytes = actual_message_size * count; - let description = format_args!( - "{} size x {} messages", - format_bytes(actual_message_size), - count - ); + let description = + format_args!("{} size x {} messages", format_bytes(actual_message_size), count); group.throughput(Throughput::Bytes(total_bytes as u64)); group.bench_with_input( diff --git a/rust/sdk/src/zeroparser/errors.rs b/rust/sdk/src/zeroparser/errors.rs index e19fd166..2341d9be 100644 --- a/rust/sdk/src/zeroparser/errors.rs +++ b/rust/sdk/src/zeroparser/errors.rs @@ -95,11 +95,7 @@ impl fmt::Display for ParseError { } ParseError::InvalidWireType(wt) => write!(f, "Invalid wire type: {}", wt), ParseError::MaxNestingDepthExceeded { max } => { - write!( - f, - "Message nesting depth exceeds maximum allowed limit of {} levels", - max - ) + write!(f, "Message nesting depth exceeds maximum allowed limit of {} levels", max) } ParseError::TruncatedVarint => write!(f, "Truncated varint"), ParseError::TypeMismatch { diff --git a/rust/sdk/src/zeroparser/owned.rs b/rust/sdk/src/zeroparser/owned.rs index 4cc0b0ea..83b41143 100644 --- a/rust/sdk/src/zeroparser/owned.rs +++ b/rust/sdk/src/zeroparser/owned.rs @@ -143,10 +143,7 @@ mod tests { // and triggers ParseError::InvalidWireType regardless of descriptor. let invalid = vec![0x0Eu8]; let (err, recovered) = OwnedParsedMessage::parse(invalid.clone(), registry).unwrap_err(); - assert!( - matches!(err, ParseError::InvalidWireType(6)), - "unexpected error: {err:?}" - ); + assert!(matches!(err, ParseError::InvalidWireType(6)), "unexpected error: {err:?}"); assert_eq!(recovered, invalid); } } diff --git a/rust/sdk/src/zeroparser/parser.rs b/rust/sdk/src/zeroparser/parser.rs index c8582cdd..91c08776 100644 --- a/rust/sdk/src/zeroparser/parser.rs +++ b/rust/sdk/src/zeroparser/parser.rs @@ -7,8 +7,15 @@ use prost_types::field_descriptor_proto::Type; use super::errors::{ParseError, ParseResult}; use super::registry::{DescriptorWithFieldCache, FieldInfo, MessageRegistry}; use super::types::{ - convert_scalar_value, default_value_for_type, ComplexType, FieldValueRef, MapKeyRef, - PackedField, ParsedMapValue, MAP_ENTRY_KEY_FIELD_NUM, MAP_ENTRY_VALUE_FIELD_NUM, + convert_scalar_value, + default_value_for_type, + ComplexType, + FieldValueRef, + MapKeyRef, + PackedField, + ParsedMapValue, + MAP_ENTRY_KEY_FIELD_NUM, + MAP_ENTRY_VALUE_FIELD_NUM, MAX_NESTING_DEPTH, }; use super::wire::{try_parse_field, WireValue}; @@ -330,10 +337,7 @@ impl<'a> ParsedMessage<'a> { /// Parse a protobuf message recursively in a single O(N) pass. #[inline(always)] pub fn parse(bytes: &'a [u8], registry: &'a MessageRegistry) -> ParseResult> { - Self::parse_internal( - bytes, None, /* type_name */ - registry, 0, /* depth */ - ) + Self::parse_internal(bytes, None /* type_name */, registry, 0 /* depth */) } #[inline(always)] @@ -592,7 +596,10 @@ impl<'a, 'b> std::ops::Deref for ParsedFieldValue<'a, 'b> { pub mod tests { use prost_types::field_descriptor_proto::Type; use prost_types::{ - DescriptorProto, FieldDescriptorProto, MessageOptions, OneofDescriptorProto, + DescriptorProto, + FieldDescriptorProto, + MessageOptions, + OneofDescriptorProto, }; use super::*; @@ -678,34 +685,10 @@ pub mod tests { #[test] fn parse_scalar_fields() { let cases: Vec<(i32, &str, Type, &[u8], FieldValueRef)> = vec![ - ( - 1, - "id", - Type::Int32, - &[8, 0x96, 0x01], - FieldValueRef::Int32(150), - ), - ( - 1, - "big", - Type::Int64, - &[8, 0xAC, 0x02], - FieldValueRef::Int64(300), - ), - ( - 1, - "count", - Type::Uint32, - &[8, 42], - FieldValueRef::UInt32(42), - ), - ( - 1, - "ts", - Type::Uint64, - &[8, 0xE8, 0x07], - FieldValueRef::UInt64(1000), - ), + (1, "id", Type::Int32, &[8, 0x96, 0x01], FieldValueRef::Int32(150)), + (1, "big", Type::Int64, &[8, 0xAC, 0x02], FieldValueRef::Int64(300)), + (1, "count", Type::Uint32, &[8, 42], FieldValueRef::UInt32(42)), + (1, "ts", Type::Uint64, &[8, 0xE8, 0x07], FieldValueRef::UInt64(1000)), (1, "delta", Type::Sint32, &[8, 1], FieldValueRef::Int32(-1)), (1, "offset", Type::Sint64, &[8, 3], FieldValueRef::Int64(-2)), (1, "flag", Type::Bool, &[8, 1], FieldValueRef::Bool(true)), @@ -724,13 +707,7 @@ pub mod tests { &[9, 0x08, 0x07, 0x06, 0x05, 0x04, 0x03, 0x02, 0x01], FieldValueRef::UInt64(0x0102030405060708), ), - ( - 1, - "sf32", - Type::Sfixed32, - &[13, 0xFF, 0xFF, 0xFF, 0xFF], - FieldValueRef::Int32(-1), - ), + (1, "sf32", Type::Sfixed32, &[13, 0xFF, 0xFF, 0xFF, 0xFF], FieldValueRef::Int32(-1)), ( 1, "sf64", @@ -779,10 +756,7 @@ pub mod tests { let double_bits = std::f64::consts::E.to_bits().to_le_bytes(); let double_bytes = [&[9u8][..], &double_bits[..]].concat(); - let desc = make_descriptor( - "Test", - vec![make_field(1, "val", Type::Double, false, None)], - ); + let desc = make_descriptor("Test", vec![make_field(1, "val", Type::Double, false, None)]); let registry = MessageRegistry::from_descriptor(&desc); let parsed = ParsedMessage::parse(&double_bytes, ®istry).unwrap(); match parsed.get_scalar(1) { @@ -844,11 +818,7 @@ pub mod tests { "int32_packed", Type::Int32, &[10, 3, 1, 2, 127], - vec![ - FieldValueRef::Int32(1), - FieldValueRef::Int32(2), - FieldValueRef::Int32(127), - ], + vec![FieldValueRef::Int32(1), FieldValueRef::Int32(2), FieldValueRef::Int32(127)], ), ( "fixed32_packed", @@ -937,12 +907,7 @@ pub mod tests { .unwrap_or_else(|_| panic!("{}: parse failed", name)); let values = parsed.get_repeated_scalars(1); - assert_eq!( - values.len(), - expected_values.len(), - "{}: length mismatch", - name - ); + assert_eq!(values.len(), expected_values.len(), "{}: length mismatch", name); for (i, expected_val) in expected_values.iter().enumerate() { assert_eq!(values[i], *expected_val, "{}: value {} mismatch", name, i); } @@ -961,13 +926,7 @@ pub mod tests { let mut outer = make_descriptor( "Outer", - vec![make_field( - 1, - "inner", - Type::Message, - false, - Some(".Outer.Inner"), - )], + vec![make_field(1, "inner", Type::Message, false, Some(".Outer.Inner"))], ); outer.nested_type.push(inner); @@ -980,10 +939,7 @@ pub mod tests { let parsed = ParsedMessage::parse(&wire, ®istry).unwrap(); let inner_parsed = parsed.get_message(1).expect("inner should be parsed"); assert_eq!(inner_parsed.get_scalar(1), Some(&FieldValueRef::Int32(42))); - assert_eq!( - inner_parsed.get_scalar(2), - Some(&FieldValueRef::String("hello")) - ); + assert_eq!(inner_parsed.get_scalar(2), Some(&FieldValueRef::String("hello"))); } #[test] @@ -991,13 +947,7 @@ pub mod tests { let item = make_descriptor("Item", vec![make_field(1, "id", Type::Int32, false, None)]); let mut container = make_descriptor( "Container", - vec![make_field( - 1, - "items", - Type::Message, - true, - Some(".Container.Item"), - )], + vec![make_field(1, "items", Type::Message, true, Some(".Container.Item"))], ); container.nested_type.push(item); @@ -1017,13 +967,7 @@ pub mod tests { let map_entry = make_map_entry_descriptor("MapEntry", Type::String, Type::Int32); let mut outer = make_descriptor( "Outer", - vec![make_field( - 1, - "items", - Type::Message, - true, - Some(".Outer.MapEntry"), - )], + vec![make_field(1, "items", Type::Message, true, Some(".Outer.MapEntry"))], ); outer.nested_type.push(map_entry); @@ -1043,13 +987,7 @@ pub mod tests { let int_key_map_entry = make_map_entry_descriptor("IntKeyMap", Type::Int32, Type::String); let mut outer2 = make_descriptor( "Outer2", - vec![make_field( - 1, - "items", - Type::Message, - true, - Some(".Outer2.IntKeyMap"), - )], + vec![make_field(1, "items", Type::Message, true, Some(".Outer2.IntKeyMap"))], ); outer2.nested_type.push(int_key_map_entry); @@ -1069,10 +1007,8 @@ pub mod tests { #[test] fn parse_map_message_values() { - let value_msg = make_descriptor( - "ValueMsg", - vec![make_field(1, "x", Type::Int32, false, None)], - ); + let value_msg = + make_descriptor("ValueMsg", vec![make_field(1, "x", Type::Int32, false, None)]); let mut map_entry = make_descriptor( "MapEntry", vec![ @@ -1087,13 +1023,7 @@ pub mod tests { let mut outer = make_descriptor( "Outer", - vec![make_field( - 1, - "items", - Type::Message, - true, - Some(".Outer.MapEntry"), - )], + vec![make_field(1, "items", Type::Message, true, Some(".Outer.MapEntry"))], ); outer.nested_type.push(map_entry); outer.nested_type.push(value_msg); @@ -1147,31 +1077,15 @@ pub mod tests { MapKeyRef::String("key"), FieldValueRef::Int32(5), ), - ( - "no_value", - &[10, 1, b'x'], - MapKeyRef::String("x"), - FieldValueRef::Int32(0), - ), - ( - "no_key", - &[16, 99], - MapKeyRef::String(""), - FieldValueRef::Int32(99), - ), + ("no_value", &[10, 1, b'x'], MapKeyRef::String("x"), FieldValueRef::Int32(0)), + ("no_key", &[16, 99], MapKeyRef::String(""), FieldValueRef::Int32(99)), ("empty", &[], MapKeyRef::String(""), FieldValueRef::Int32(0)), ]; let map_entry = make_map_entry_descriptor("MapEntry", Type::String, Type::Int32); let mut outer = make_descriptor( "Outer", - vec![make_field( - 1, - "m", - Type::Message, - true, - Some(".Outer.MapEntry"), - )], + vec![make_field(1, "m", Type::Message, true, Some(".Outer.MapEntry"))], ); outer.nested_type.push(map_entry); let registry = MessageRegistry::from_descriptor(&outer); @@ -1288,13 +1202,7 @@ pub mod tests { let level2 = make_descriptor("L2", vec![make_field(1, "val", Type::Int32, false, None)]); let mut level1 = make_descriptor( "L1", - vec![make_field( - 1, - "l2", - Type::Message, - false, - Some(".Root.L1.L2"), - )], + vec![make_field(1, "l2", Type::Message, false, Some(".Root.L1.L2"))], ); level1.nested_type.push(level2); let mut root = make_descriptor( @@ -1347,11 +1255,7 @@ pub mod tests { // All fields are present even though they have default values. for field_num in 1..=9 { - assert!( - parsed.has_field(field_num), - "Field {} should be present", - field_num - ); + assert!(parsed.has_field(field_num), "Field {} should be present", field_num); } // Case 2: Empty message - no fields present. @@ -1568,13 +1472,7 @@ pub mod tests { let map_entry = make_map_entry_descriptor("MapEntry", Type::String, Type::Int32); let mut outer = make_descriptor( "Outer", - vec![make_field( - 1, - "items", - Type::Message, - true, - Some(".Outer.MapEntry"), - )], + vec![make_field(1, "items", Type::Message, true, Some(".Outer.MapEntry"))], ); outer.nested_type.push(map_entry); @@ -1619,10 +1517,7 @@ pub mod tests { assert_eq!(parsed_dup.get_map_entries_count(1), 1); let dup_entries: Vec<_> = parsed_dup.get_map_entries(1).collect(); assert_eq!(*dup_entries[0].0, MapKeyRef::String("x")); - assert!(matches!( - dup_entries[0].1, - ParsedMapValue::Scalar(FieldValueRef::Int32(20)) - )); + assert!(matches!(dup_entries[0].1, ParsedMapValue::Scalar(FieldValueRef::Int32(20)))); } #[test] @@ -1650,13 +1545,7 @@ pub mod tests { let mut outer = make_descriptor( "Outer", - vec![make_field( - 1, - "items", - Type::Message, - true, - Some(".Outer.MapEntry"), - )], + vec![make_field(1, "items", Type::Message, true, Some(".Outer.MapEntry"))], ); outer.nested_type.push(map_entry); outer.nested_type.push(value_msg); @@ -1753,13 +1642,7 @@ pub mod tests { ); let mut outer = make_descriptor( "Outer", - vec![make_field( - 1, - "inner", - Type::Message, - false, - Some(".Outer.Inner"), - )], + vec![make_field(1, "inner", Type::Message, false, Some(".Outer.Inner"))], ); outer.nested_type.push(inner); let registry = MessageRegistry::from_descriptor(&outer); @@ -1786,15 +1669,11 @@ pub mod tests { make_field(2, "y", Type::Int32, false, None), ], ); - let mut b = make_descriptor( - "B", - vec![make_field(1, "c", Type::Message, false, Some(".A.B.C"))], - ); + let mut b = + make_descriptor("B", vec![make_field(1, "c", Type::Message, false, Some(".A.B.C"))]); b.nested_type.push(c); - let mut a = make_descriptor( - "A", - vec![make_field(1, "b", Type::Message, false, Some(".A.B"))], - ); + let mut a = + make_descriptor("A", vec![make_field(1, "b", Type::Message, false, Some(".A.B"))]); a.nested_type.push(b); let registry = MessageRegistry::from_descriptor(&a); @@ -1833,13 +1712,7 @@ pub mod tests { wrapper.nested_type.push(inner); let mut outer = make_descriptor( "Outer", - vec![make_field( - 1, - "w", - Type::Message, - false, - Some(".Outer.Wrapper"), - )], + vec![make_field(1, "w", Type::Message, false, Some(".Outer.Wrapper"))], ); outer.nested_type.push(wrapper); let registry = MessageRegistry::from_descriptor(&outer); @@ -1850,10 +1723,7 @@ pub mod tests { let p = ParsedMessage::parse(&wire, ®istry).unwrap(); let w = p.get_message(1).unwrap(); assert!(!w.has_field(1) && !w.has_field(2)); - assert_eq!( - w.get_message(3).unwrap().get_scalar(1), - Some(&FieldValueRef::Int32(1)) - ); + assert_eq!(w.get_message(3).unwrap().get_scalar(1), Some(&FieldValueRef::Int32(1))); // B: same message member twice — inner messages merge (Inner has only x, // so scalar last-wins semantics surface inside the merged inner). @@ -1898,13 +1768,7 @@ pub mod tests { wrapper.nested_type.push(map_entry); let mut outer = make_descriptor( "Outer", - vec![make_field( - 1, - "w", - Type::Message, - false, - Some(".Outer.Wrapper"), - )], + vec![make_field(1, "w", Type::Message, false, Some(".Outer.Wrapper"))], ); outer.nested_type.push(wrapper); let registry = MessageRegistry::from_descriptor(&outer); @@ -1926,21 +1790,14 @@ pub mod tests { assert_eq!( w.get_repeated_scalars(1), - &[ - FieldValueRef::Int32(1), - FieldValueRef::Int32(2), - FieldValueRef::Int32(3), - ] + &[FieldValueRef::Int32(1), FieldValueRef::Int32(2), FieldValueRef::Int32(3),] ); let item_vs: Vec<_> = w .get_repeated_messages(2) .iter() .map(|m| m.get_scalar(1).copied()) .collect(); - assert_eq!( - item_vs, - vec![Some(FieldValueRef::Int32(1)), Some(FieldValueRef::Int32(2))] - ); + assert_eq!(item_vs, vec![Some(FieldValueRef::Int32(1)), Some(FieldValueRef::Int32(2))]); let counts: std::collections::HashMap<_, i32> = w .get_map_entries(3) @@ -1970,13 +1827,7 @@ pub mod tests { ); let mut outer = make_descriptor( "Outer", - vec![make_field( - 1, - "w", - Type::Message, - false, - Some(".Outer.Wrapper"), - )], + vec![make_field(1, "w", Type::Message, false, Some(".Outer.Wrapper"))], ); outer.nested_type.push(std::mem::take(&mut wrapper)); let registry = MessageRegistry::from_descriptor(&outer); @@ -2002,13 +1853,7 @@ pub mod tests { "Wrapper", vec![ make_field(1, "id", Type::Int32, false, None), - make_field( - 2, - "inner", - Type::Message, - false, - Some(".Outer.Wrapper.Inner"), - ), + make_field(2, "inner", Type::Message, false, Some(".Outer.Wrapper.Inner")), make_field(3, "counts", Type::Message, true, Some(".Outer.Wrapper.CE")), ], ); @@ -2016,13 +1861,7 @@ pub mod tests { wrapper.nested_type.push(map_entry); let mut outer = make_descriptor( "Outer", - vec![make_field( - 1, - "w", - Type::Message, - false, - Some(".Outer.Wrapper"), - )], + vec![make_field(1, "w", Type::Message, false, Some(".Outer.Wrapper"))], ); outer.nested_type.push(wrapper); let registry = MessageRegistry::from_descriptor(&outer); @@ -2036,10 +1875,7 @@ pub mod tests { let parsed = ParsedMessage::parse(&wire, ®istry).unwrap(); let w = parsed.get_message(1).unwrap(); assert_eq!(w.get_scalar(1), Some(&FieldValueRef::Int32(1))); // preserved - assert_eq!( - w.get_message(2).unwrap().get_scalar(1), - Some(&FieldValueRef::Int32(5)) - ); // installed + assert_eq!(w.get_message(2).unwrap().get_scalar(1), Some(&FieldValueRef::Int32(5))); // installed let counts: Vec<_> = w.get_map_entries(3).collect(); assert_eq!(counts.len(), 1); assert_eq!(*counts[0].0, MapKeyRef::String("a")); @@ -2050,19 +1886,11 @@ pub mod tests { // Second occurrence parses to a completely empty ParsedMessage; // merge_from must leave self untouched. Targets the "all None / Empty // in other" path of every loop in merge_from. - let wrapper = make_descriptor( - "Wrapper", - vec![make_field(1, "id", Type::Int32, false, None)], - ); + let wrapper = + make_descriptor("Wrapper", vec![make_field(1, "id", Type::Int32, false, None)]); let mut outer = make_descriptor( "Outer", - vec![make_field( - 1, - "w", - Type::Message, - false, - Some(".Outer.Wrapper"), - )], + vec![make_field(1, "w", Type::Message, false, Some(".Outer.Wrapper"))], ); outer.nested_type.push(wrapper); let registry = MessageRegistry::from_descriptor(&outer); @@ -2182,13 +2010,7 @@ pub mod tests { let mut outer = make_descriptor( "Outer", - vec![make_field( - 1, - "items", - Type::Message, - true, - Some(".Outer.MapEntry"), - )], + vec![make_field(1, "items", Type::Message, true, Some(".Outer.MapEntry"))], ); outer.nested_type.push(map_entry); @@ -2222,23 +2044,14 @@ pub mod tests { // Note: We only test float/double/bytes here because: // - message (structs) → comes through as Bytes in wire format, so already covered // - repeated/map → invalid protobuf syntax, can't be declared as key types - let invalid_key_types = vec![ - ("float", Type::Float), - ("double", Type::Double), - ("bytes", Type::Bytes), - ]; + let invalid_key_types = + vec![("float", Type::Float), ("double", Type::Double), ("bytes", Type::Bytes)]; for (type_name, key_type) in invalid_key_types { let map_entry = make_map_entry_descriptor("MapEntry", key_type, Type::Int32); let mut outer = make_descriptor( "Outer", - vec![make_field( - 1, - "items", - Type::Message, - true, - Some(".Outer.MapEntry"), - )], + vec![make_field(1, "items", Type::Message, true, Some(".Outer.MapEntry"))], ); outer.nested_type.push(map_entry); @@ -2267,11 +2080,7 @@ pub mod tests { wire.extend_from_slice(&entry_wire); let result = ParsedMessage::parse(&wire, ®istry); - assert!( - result.is_err(), - "Expected error for invalid map key type: {}", - type_name - ); + assert!(result.is_err(), "Expected error for invalid map key type: {}", type_name); assert!( matches!(result, Err(ParseError::InvalidMapKeyType { .. })), "Expected InvalidMapKeyType error for {}, got {:?}", @@ -2369,13 +2178,7 @@ pub mod tests { let map_entry = make_map_entry_descriptor("MapEntry", key_type, Type::Int32); let mut outer = make_descriptor( "Outer", - vec![make_field( - 1, - "items", - Type::Message, - true, - Some(".Outer.MapEntry"), - )], + vec![make_field(1, "items", Type::Message, true, Some(".Outer.MapEntry"))], ); outer.nested_type.push(map_entry); @@ -2393,11 +2196,7 @@ pub mod tests { assert_eq!(parsed.get_map_entries_count(1), 1, "case: {}", case_name); let entries: Vec<_> = parsed.get_map_entries(1).collect(); - assert_eq!( - *entries[0].0, expected_key, - "case: {} key mismatch", - case_name - ); + assert_eq!(*entries[0].0, expected_key, "case: {} key mismatch", case_name); match entries[0].1 { ParsedMapValue::Scalar(v) => { assert_eq!(v, &FieldValueRef::Int32(99), "case: {} value", case_name) @@ -2507,14 +2306,7 @@ pub mod tests { vec![ make_field(1, "tag", Type::Int32, false, None), make_oneof_field(2, "str_val", Type::String, false, None, Some(0)), - make_oneof_field( - 3, - "msg_val", - Type::Message, - false, - Some(".Outer.Inner"), - Some(0), - ), + make_oneof_field(3, "msg_val", Type::Message, false, Some(".Outer.Inner"), Some(0)), ], vec!["payload"], ); @@ -2558,14 +2350,7 @@ pub mod tests { "Outer", vec![ make_oneof_field(2, "str_val", Type::String, false, None, Some(0)), - make_oneof_field( - 3, - "msg_val", - Type::Message, - false, - Some(".Outer.Inner"), - Some(0), - ), + make_oneof_field(3, "msg_val", Type::Message, false, Some(".Outer.Inner"), Some(0)), ], vec!["payload"], ); @@ -2629,10 +2414,7 @@ pub mod tests { let desc = make_descriptor_with_oneofs( "Test", - vec![ - make_oneof_field(1, "real_a", Type::Int32, false, None, Some(0)), - field2, - ], + vec![make_oneof_field(1, "real_a", Type::Int32, false, None, Some(0)), field2], vec!["_opt_field"], ); let registry = MessageRegistry::from_descriptor(&desc); @@ -2664,13 +2446,7 @@ pub mod tests { let inner = make_descriptor("Inner", vec![make_field(1, "a", Type::Int32, false, None)]); let mut outer = make_descriptor( "Outer", - vec![make_field( - 1, - "inner", - Type::Message, - false, - Some(".Outer.Inner"), - )], + vec![make_field(1, "inner", Type::Message, false, Some(".Outer.Inner"))], ); outer.nested_type.push(inner); diff --git a/rust/sdk/src/zeroparser/registry.rs b/rust/sdk/src/zeroparser/registry.rs index 19e9deb7..c762d594 100644 --- a/rust/sdk/src/zeroparser/registry.rs +++ b/rust/sdk/src/zeroparser/registry.rs @@ -203,10 +203,7 @@ impl MessageRegistry { format!("{current_prefix}.{name}") }; - acc.insert( - full_name.clone(), - DescriptorWithFieldCache::from_descriptor(desc), - ); + acc.insert(full_name.clone(), DescriptorWithFieldCache::from_descriptor(desc)); // Push nested types onto the stack (in reverse order to maintain processing order). for nested in desc.nested_type.iter().rev() { @@ -245,13 +242,7 @@ pub mod tests { let fields = vec![ make_field(1, "id", field_descriptor_proto::Type::Int32, false, None), make_field(2, "name", field_descriptor_proto::Type::String, false, None), - make_field( - 200, - "large", - field_descriptor_proto::Type::String, - false, - None, - ), + make_field(200, "large", field_descriptor_proto::Type::String, false, None), make_field(3, "items", field_descriptor_proto::Type::Int32, true, None), ]; let desc = make_descriptor("TestMessage", fields); @@ -297,13 +288,7 @@ pub mod tests { #[test] fn message_registry_lookup() { - let fields = vec![make_field( - 1, - "id", - field_descriptor_proto::Type::Int32, - false, - None, - )]; + let fields = vec![make_field(1, "id", field_descriptor_proto::Type::Int32, false, None)]; let desc = make_descriptor("RootMessage", fields); let registry = MessageRegistry::from_descriptor(&desc); @@ -316,13 +301,7 @@ pub mod tests { fn message_registry_nested() { let level3 = make_descriptor( "Level3", - vec![make_field( - 1, - "field3", - field_descriptor_proto::Type::Bool, - false, - None, - )], + vec![make_field(1, "field3", field_descriptor_proto::Type::Bool, false, None)], ); let mut level2 = make_descriptor("Level2", vec![]); level2.nested_type.push(level3); diff --git a/rust/sdk/src/zeroparser/tests/common/mod.rs b/rust/sdk/src/zeroparser/tests/common/mod.rs index 97b8c954..c21b207a 100644 --- a/rust/sdk/src/zeroparser/tests/common/mod.rs +++ b/rust/sdk/src/zeroparser/tests/common/mod.rs @@ -63,10 +63,7 @@ pub fn create_registry_for_version(version: ProtoVersion, message_name: &str) -> let file_desc_set = load_descriptor_set(version); let (msg_desc, file) = find_message_and_file(&file_desc_set, version.package(), message_name); let mut descriptor = msg_desc.clone(); - descriptor.name = Some(format!( - "{}.{message_name}", - file.package.as_deref().unwrap_or("") - )); + descriptor.name = Some(format!("{}.{message_name}", file.package.as_deref().unwrap_or(""))); MessageRegistry::from_descriptor(&descriptor) } diff --git a/rust/sdk/src/zeroparser/tests/e2e.rs b/rust/sdk/src/zeroparser/tests/e2e.rs index 14ea2713..d4388b06 100644 --- a/rust/sdk/src/zeroparser/tests/e2e.rs +++ b/rust/sdk/src/zeroparser/tests/e2e.rs @@ -3,9 +3,15 @@ mod common; // Proto2 and proto3 have usually the same field ordinals, so we can use the same constants. use all_types_fields::{proto2 as all_types_ordinals, proto3 as all_types_ordinals_proto3}; use common::{ - all_types_fields, complex_nested_fields, create_registry_for_version, - deeply_nested_message_fields, encode_message_for_version, field_num, nested_message_fields, - supported_types_fields, ProtoVersion, + all_types_fields, + complex_nested_fields, + create_registry_for_version, + deeply_nested_message_fields, + encode_message_for_version, + field_num, + nested_message_fields, + supported_types_fields, + ProtoVersion, }; use databricks_zerobus_ingest_sdk::zeroparser::parser::ParsedMessage; use databricks_zerobus_ingest_sdk::zeroparser::types::{FieldValueRef, MapKeyRef}; @@ -58,10 +64,7 @@ fn test_all_scalar_types(#[case] version: ProtoVersion) { let parsed = ParsedMessage::parse(&buf, ®istry).unwrap(); - assert_eq!( - parsed.get_scalar(all_types_ordinals::F_INT32), - Some(&FieldValueRef::Int32(-42)) - ); + assert_eq!(parsed.get_scalar(all_types_ordinals::F_INT32), Some(&FieldValueRef::Int32(-42))); assert_eq!( parsed.get_scalar(all_types_ordinals::F_INT64), Some(&FieldValueRef::Int64(-9223372036854775807)) @@ -98,18 +101,12 @@ fn test_all_scalar_types(#[case] version: ProtoVersion) { parsed.get_scalar(all_types_ordinals::F_SFIXED64), Some(&FieldValueRef::Int64(-123456789012345)) ); - assert_eq!( - parsed.get_scalar(all_types_ordinals::F_FLOAT), - Some(&FieldValueRef::Float(3.25)) - ); + assert_eq!(parsed.get_scalar(all_types_ordinals::F_FLOAT), Some(&FieldValueRef::Float(3.25))); assert_eq!( parsed.get_scalar(all_types_ordinals::F_DOUBLE), Some(&FieldValueRef::Double(2.125)) ); - assert_eq!( - parsed.get_scalar(all_types_ordinals::F_BOOL), - Some(&FieldValueRef::Bool(true)) - ); + assert_eq!(parsed.get_scalar(all_types_ordinals::F_BOOL), Some(&FieldValueRef::Bool(true))); assert_eq!( parsed.get_scalar(all_types_ordinals::F_STRING), Some(&FieldValueRef::String("test_string")) @@ -118,10 +115,7 @@ fn test_all_scalar_types(#[case] version: ProtoVersion) { parsed.get_scalar(all_types_ordinals::F_BYTES), Some(&FieldValueRef::Bytes(&[0xDE, 0xAD, 0xBE, 0xEF])) ); - assert_eq!( - parsed.get_scalar(all_types_ordinals::F_ENUM), - Some(&FieldValueRef::Int32(2)) - ); + assert_eq!(parsed.get_scalar(all_types_ordinals::F_ENUM), Some(&FieldValueRef::Int32(2))); } #[rstest] @@ -502,14 +496,8 @@ fn test_map_fields(#[case] version: ProtoVersion) { let parsed_empty = ParsedMessage::parse(&buf_empty, ®istry_empty).unwrap(); assert_eq!(parsed_empty.get_map_entries_count(map_int_string_field), 0); - assert_eq!( - parsed_empty.get_map_entries_count(map_string_string_field), - 0 - ); - assert_eq!( - parsed_empty.get_map_entries_count(map_string_message_field), - 0 - ); + assert_eq!(parsed_empty.get_map_entries_count(map_string_string_field), 0); + assert_eq!(parsed_empty.get_map_entries_count(map_string_message_field), 0); } #[rstest] @@ -600,9 +588,7 @@ fn test_string_and_binary_data(#[case] version: ProtoVersion) { // Verify binary data assert_eq!( parsed.get_scalar(all_types_ordinals::F_BYTES), - Some(&FieldValueRef::Bytes(&[ - 0x00, 0xFF, 0xAB, 0xCD, 0xEF, 0x12, 0x34, 0x56 - ])) + Some(&FieldValueRef::Bytes(&[0x00, 0xFF, 0xAB, 0xCD, 0xEF, 0x12, 0x34, 0x56])) ); // Verify repeated strings with international characters @@ -641,31 +627,19 @@ fn test_unknown_fields(#[case] version: ProtoVersion) { parsed.get_scalar(supported_types_fields::APPROVED), Some(&FieldValueRef::Bool(true)) ); - assert_eq!( - parsed.get_scalar(supported_types_fields::DAY_NUM), - Some(&FieldValueRef::Int32(10)) - ); - assert_eq!( - parsed.get_scalar(supported_types_fields::COST), - Some(&FieldValueRef::Int64(5000)) - ); + assert_eq!(parsed.get_scalar(supported_types_fields::DAY_NUM), Some(&FieldValueRef::Int32(10))); + assert_eq!(parsed.get_scalar(supported_types_fields::COST), Some(&FieldValueRef::Int64(5000))); assert_eq!( parsed.get_scalar(supported_types_fields::DESCRIPTION), Some(&FieldValueRef::String("test")) ); - assert!( - !parsed.has_field(supported_types_fields::DISCOUNT), - "discount should be unknown" - ); + assert!(!parsed.has_field(supported_types_fields::DISCOUNT), "discount should be unknown"); assert!( !parsed.has_field(supported_types_fields::COST_WITH_DISCOUNT), "cost_with_discount should be unknown" ); - assert!( - !parsed.has_field(supported_types_fields::PHOTO), - "photo should be unknown" - ); + assert!(!parsed.has_field(supported_types_fields::PHOTO), "photo should be unknown"); assert_eq!( parsed .get_repeated_scalars(supported_types_fields::TAGS) @@ -742,10 +716,7 @@ fn test_negative_numbers(#[case] version: ProtoVersion) { let parsed = ParsedMessage::parse(&buf, ®istry).unwrap(); - assert_eq!( - parsed.get_scalar(all_types_ordinals::F_INT32), - Some(&FieldValueRef::Int32(-12345)) - ); + assert_eq!(parsed.get_scalar(all_types_ordinals::F_INT32), Some(&FieldValueRef::Int32(-12345))); assert_eq!( parsed.get_scalar(all_types_ordinals::F_INT64), Some(&FieldValueRef::Int64(-9876543210)) @@ -766,10 +737,7 @@ fn test_negative_numbers(#[case] version: ProtoVersion) { parsed.get_scalar(all_types_ordinals::F_SFIXED64), Some(&FieldValueRef::Int64(-9999999999)) ); - assert_eq!( - parsed.get_scalar(all_types_ordinals::F_FLOAT), - Some(&FieldValueRef::Float(-3.25)) - ); + assert_eq!(parsed.get_scalar(all_types_ordinals::F_FLOAT), Some(&FieldValueRef::Float(-3.25))); assert_eq!( parsed.get_scalar(all_types_ordinals::F_DOUBLE), Some(&FieldValueRef::Double(-2.125)) @@ -913,18 +881,9 @@ fn test_oneof_fields(#[case] version: ProtoVersion) { let parsed_int = ParsedMessage::parse(&buf_int, ®istry_int).unwrap(); - assert!( - parsed_int.has_field(oneof_int_field), - "oneof_int should be present" - ); - assert_eq!( - parsed_int.get_scalar(oneof_int_field), - Some(&FieldValueRef::Int32(42)) - ); - assert!( - !parsed_int.has_field(oneof_string_field), - "oneof_string should be absent" - ); + assert!(parsed_int.has_field(oneof_int_field), "oneof_int should be present"); + assert_eq!(parsed_int.get_scalar(oneof_int_field), Some(&FieldValueRef::Int32(42))); + assert!(!parsed_int.has_field(oneof_string_field), "oneof_string should be absent"); assert!( parsed_int.get_message(oneof_message_field).is_none(), "oneof_message should be absent" @@ -954,14 +913,8 @@ fn test_oneof_fields(#[case] version: ProtoVersion) { let parsed_string = ParsedMessage::parse(&buf_string, ®istry_string).unwrap(); - assert!( - !parsed_string.has_field(oneof_int_field), - "oneof_int should be absent" - ); - assert!( - parsed_string.has_field(oneof_string_field), - "oneof_string should be present" - ); + assert!(!parsed_string.has_field(oneof_int_field), "oneof_int should be absent"); + assert!(parsed_string.has_field(oneof_string_field), "oneof_string should be present"); assert_eq!( parsed_string.get_scalar(oneof_string_field), Some(&FieldValueRef::String("oneof_test")) @@ -1001,14 +954,8 @@ fn test_oneof_fields(#[case] version: ProtoVersion) { let parsed_message = ParsedMessage::parse(&buf_message, ®istry_message).unwrap(); - assert!( - !parsed_message.has_field(oneof_int_field), - "oneof_int should be absent" - ); - assert!( - !parsed_message.has_field(oneof_string_field), - "oneof_string should be absent" - ); + assert!(!parsed_message.has_field(oneof_int_field), "oneof_int should be absent"); + assert!(!parsed_message.has_field(oneof_string_field), "oneof_string should be absent"); let nested = parsed_message .get_message(oneof_message_field) .expect("oneof_message should be present"); @@ -1040,9 +987,7 @@ fn test_oneof_last_writer_wins_on_wire(#[case] version: ProtoVersion) { }; let msg_string = proto2::AllTypesMessage { f_required: 1, - f_oneof: Some(proto2::all_types_message::FOneof::OneofString( - "winner".to_string(), - )), + f_oneof: Some(proto2::all_types_message::FOneof::OneofString("winner".to_string())), ..Default::default() }; let msg_message = proto2::AllTypesMessage { @@ -1076,10 +1021,7 @@ fn test_oneof_last_writer_wins_on_wire(#[case] version: ProtoVersion) { !parsed.has_field(oneof_int_field), "oneof_int should be cleared by later oneof_string" ); - assert_eq!( - parsed.get_scalar(oneof_string_field), - Some(&FieldValueRef::String("winner")) - ); + assert_eq!(parsed.get_scalar(oneof_string_field), Some(&FieldValueRef::String("winner"))); assert!(parsed.get_message(oneof_message_field).is_none()); // Case 2: scalar → message. Message wins; scalar is cleared. @@ -1108,10 +1050,7 @@ fn test_oneof_last_writer_wins_on_wire(#[case] version: ProtoVersion) { "oneof_message should be cleared by later oneof_int" ); assert!(!parsed.has_field(oneof_string_field)); - assert_eq!( - parsed.get_scalar(oneof_int_field), - Some(&FieldValueRef::Int32(42)) - ); + assert_eq!(parsed.get_scalar(oneof_int_field), Some(&FieldValueRef::Int32(42))); } #[rstest] @@ -1127,9 +1066,7 @@ fn test_proto3_optional_coexists_with_oneof(#[case] version: ProtoVersion) { let msg_a = proto3::AllTypesMessage { f_optional_int32: Some(1), f_optional_string: Some("keep".to_string()), - f_oneof: Some(proto3::all_types_message::FOneof::OneofString( - "first".to_string(), - )), + f_oneof: Some(proto3::all_types_message::FOneof::OneofString("first".to_string())), ..Default::default() }; let msg_b = proto3::AllTypesMessage { @@ -1362,10 +1299,7 @@ fn test_complex_nested_all_field_types(#[case] version: ProtoVersion) { let parsed = ParsedMessage::parse(&buf, ®istry).unwrap(); - assert_eq!( - parsed.get_scalar(complex_nested_fields::ID), - Some(&FieldValueRef::Int32(100)) - ); + assert_eq!(parsed.get_scalar(complex_nested_fields::ID), Some(&FieldValueRef::Int32(100))); assert_eq!( parsed.get_scalar(complex_nested_fields::NAME), Some(&FieldValueRef::String("complex")) @@ -1411,20 +1345,11 @@ fn test_complex_nested_all_field_types(#[case] version: ProtoVersion) { Some(&FieldValueRef::Int64(222)) ); - assert_eq!( - parsed.get_map_entries_count(complex_nested_fields::STRING_TO_INT), - 2 - ); + assert_eq!(parsed.get_map_entries_count(complex_nested_fields::STRING_TO_INT), 2); - assert_eq!( - parsed.get_map_entries_count(complex_nested_fields::INT_TO_STRING), - 2 - ); + assert_eq!(parsed.get_map_entries_count(complex_nested_fields::INT_TO_STRING), 2); - assert_eq!( - parsed.get_map_entries_count(complex_nested_fields::STRING_TO_MESSAGE), - 1 - ); + assert_eq!(parsed.get_map_entries_count(complex_nested_fields::STRING_TO_MESSAGE), 1); let data_with_maps = parsed .get_message(complex_nested_fields::DATA_WITH_MAPS) @@ -1456,10 +1381,7 @@ fn test_complex_nested_all_field_types(#[case] version: ProtoVersion) { .len(), 2 ); - assert_eq!( - items[0].get_map_entries_count(complex_nested_fields::complex_item::ATTRIBUTES), - 1 - ); + assert_eq!(items[0].get_map_entries_count(complex_nested_fields::complex_item::ATTRIBUTES), 1); // Verify tree structure let tree_root = parsed @@ -1557,42 +1479,15 @@ fn test_all_types_message_with_null_values(#[case] version: ProtoVersion) { let parsed = ParsedMessage::parse(&buf, ®istry).unwrap(); // All optional scalar fields should be absent - assert!( - !parsed.has_field(all_types_ordinals::F_INT32), - "f_int32 should be absent" - ); - assert!( - !parsed.has_field(all_types_ordinals::F_INT64), - "f_int64 should be absent" - ); - assert!( - !parsed.has_field(all_types_ordinals::F_FLOAT), - "f_float should be absent" - ); - assert!( - !parsed.has_field(all_types_ordinals::F_DOUBLE), - "f_double should be absent" - ); - assert!( - !parsed.has_field(all_types_ordinals::F_BOOL), - "f_bool should be absent" - ); - assert!( - !parsed.has_field(all_types_ordinals::F_STRING), - "f_string should be absent" - ); - assert!( - !parsed.has_field(all_types_ordinals::F_BYTES), - "f_bytes should be absent" - ); - assert!( - !parsed.has_field(all_types_ordinals::F_ENUM), - "f_enum should be absent" - ); - assert!( - !parsed.has_field(all_types_ordinals::F_DEFAULT_INT), - "f_default_int should be absent" - ); + assert!(!parsed.has_field(all_types_ordinals::F_INT32), "f_int32 should be absent"); + assert!(!parsed.has_field(all_types_ordinals::F_INT64), "f_int64 should be absent"); + assert!(!parsed.has_field(all_types_ordinals::F_FLOAT), "f_float should be absent"); + assert!(!parsed.has_field(all_types_ordinals::F_DOUBLE), "f_double should be absent"); + assert!(!parsed.has_field(all_types_ordinals::F_BOOL), "f_bool should be absent"); + assert!(!parsed.has_field(all_types_ordinals::F_STRING), "f_string should be absent"); + assert!(!parsed.has_field(all_types_ordinals::F_BYTES), "f_bytes should be absent"); + assert!(!parsed.has_field(all_types_ordinals::F_ENUM), "f_enum should be absent"); + assert!(!parsed.has_field(all_types_ordinals::F_DEFAULT_INT), "f_default_int should be absent"); assert!( !parsed.has_field(all_types_ordinals::F_DEFAULT_STRING), "f_default_string should be absent" @@ -1601,23 +1496,14 @@ fn test_all_types_message_with_null_values(#[case] version: ProtoVersion) { !parsed.has_field(all_types_ordinals::F_DEFAULT_BOOL), "f_default_bool should be absent" ); - assert!( - !parsed.has_field(all_types_ordinals::F_NESTED), - "f_nested should be absent" - ); + assert!(!parsed.has_field(all_types_ordinals::F_NESTED), "f_nested should be absent"); assert!( !parsed.has_field(all_types_ordinals::F_DEEPLY_NESTED), "f_deeply_nested should be absent" ); // Required field should be present - assert!( - parsed.has_field(all_types_ordinals::F_REQUIRED), - "f_required should be present" - ); - assert_eq!( - parsed.get_scalar(all_types_ordinals::F_REQUIRED), - Some(&FieldValueRef::Int32(1)) - ); + assert!(parsed.has_field(all_types_ordinals::F_REQUIRED), "f_required should be present"); + assert_eq!(parsed.get_scalar(all_types_ordinals::F_REQUIRED), Some(&FieldValueRef::Int32(1))); // Empty repeated fields assert_eq!( parsed @@ -1650,23 +1536,11 @@ fn test_all_types_message_with_null_values(#[case] version: ProtoVersion) { 0 ); // Empty maps - assert_eq!( - parsed.get_map_entries_count(all_types_ordinals::F_MAP_INT_STRING), - 0 - ); - assert_eq!( - parsed.get_map_entries_count(all_types_ordinals::F_MAP_STRING_STRING), - 0 - ); - assert_eq!( - parsed.get_map_entries_count(all_types_ordinals::F_MAP_STRING_MESSAGE), - 0 - ); + assert_eq!(parsed.get_map_entries_count(all_types_ordinals::F_MAP_INT_STRING), 0); + assert_eq!(parsed.get_map_entries_count(all_types_ordinals::F_MAP_STRING_STRING), 0); + assert_eq!(parsed.get_map_entries_count(all_types_ordinals::F_MAP_STRING_MESSAGE), 0); // Oneof should have no value set - assert!( - !parsed.has_field(all_types_ordinals::F_ONEOF_INT), - "oneof_int should not be set" - ); + assert!(!parsed.has_field(all_types_ordinals::F_ONEOF_INT), "oneof_int should not be set"); assert!( !parsed.has_field(all_types_ordinals::F_ONEOF_STRING), "oneof_string should not be set" @@ -1698,10 +1572,7 @@ fn test_all_types_message_with_null_values(#[case] version: ProtoVersion) { // Required field should be present in proto2 if version == ProtoVersion::Proto2 { - assert!( - parsed.has_field(all_types_ordinals::F_REQUIRED), - "f_required should be present" - ); + assert!(parsed.has_field(all_types_ordinals::F_REQUIRED), "f_required should be present"); assert_eq!( parsed.get_scalar(all_types_ordinals::F_REQUIRED), Some(&FieldValueRef::Int32(1)) @@ -1775,10 +1646,7 @@ fn test_complex_nested_with_null_values(#[case] version: ProtoVersion) { match version { ProtoVersion::Proto2 => { // Top level null field - assert!( - !parsed.has_field(complex_nested_fields::ID), - "id should be absent" - ); + assert!(!parsed.has_field(complex_nested_fields::ID), "id should be absent"); assert_eq!( parsed.get_scalar(complex_nested_fields::NAME), Some(&FieldValueRef::String("partial")) @@ -1910,18 +1778,9 @@ fn test_complex_nested_with_null_values(#[case] version: ProtoVersion) { ); // Empty maps - assert_eq!( - parsed.get_map_entries_count(complex_nested_fields::STRING_TO_INT), - 0 - ); - assert_eq!( - parsed.get_map_entries_count(complex_nested_fields::INT_TO_STRING), - 0 - ); - assert_eq!( - parsed.get_map_entries_count(complex_nested_fields::STRING_TO_MESSAGE), - 0 - ); + assert_eq!(parsed.get_map_entries_count(complex_nested_fields::STRING_TO_INT), 0); + assert_eq!(parsed.get_map_entries_count(complex_nested_fields::INT_TO_STRING), 0); + assert_eq!(parsed.get_map_entries_count(complex_nested_fields::STRING_TO_MESSAGE), 0); // Null nested message assert!( @@ -2056,11 +1915,7 @@ fn test_map_duplicate_keys_last_value_wins(#[case] version: ProtoVersion) { let map_int_string_field = field_num(version, "f_map_int_string"); - assert_eq!( - parsed.get_map_entries_count(map_int_string_field), - 3, - "Should have 3 unique keys" - ); + assert_eq!(parsed.get_map_entries_count(map_int_string_field), 3, "Should have 3 unique keys"); let mut found_overridden = false; let mut found_second = false; @@ -2081,10 +1936,7 @@ fn test_map_duplicate_keys_last_value_wins(#[case] version: ProtoVersion) { } } - assert!( - found_overridden, - "Key 1 should have 'overridden' (last value wins)" - ); + assert!(found_overridden, "Key 1 should have 'overridden' (last value wins)"); assert!(found_second, "Key 2 should have 'second'"); assert!(found_third, "Key 3 should have 'third'"); } @@ -2124,10 +1976,7 @@ fn test_map_with_complex_nested_values(#[case] version: ProtoVersion) { let (buf, registry) = encode_message_for_version(version, &msg, "ComplexNested"); let parsed = ParsedMessage::parse(&buf, ®istry).unwrap(); - assert_eq!( - parsed.get_map_entries_count(complex_nested_fields::STRING_TO_MESSAGE), - 1 - ); + assert_eq!(parsed.get_map_entries_count(complex_nested_fields::STRING_TO_MESSAGE), 1); let string_to_message: Vec<_> = parsed .get_map_entries(complex_nested_fields::STRING_TO_MESSAGE) @@ -2201,11 +2050,7 @@ fn test_interleaved_repeated_fields(#[case] version: ProtoVersion) { let repeated_string_field = field_num(version, "f_repeated_string"); let repeated_int32 = parsed.get_repeated_scalars(repeated_int32_field); - assert_eq!( - repeated_int32.len(), - 5, - "Should accumulate all repeated int32 values" - ); + assert_eq!(repeated_int32.len(), 5, "Should accumulate all repeated int32 values"); assert_eq!(repeated_int32[0], FieldValueRef::Int32(1)); assert_eq!(repeated_int32[1], FieldValueRef::Int32(2)); assert_eq!(repeated_int32[2], FieldValueRef::Int32(3)); @@ -2213,11 +2058,7 @@ fn test_interleaved_repeated_fields(#[case] version: ProtoVersion) { assert_eq!(repeated_int32[4], FieldValueRef::Int32(5)); let repeated_string = parsed.get_repeated_scalars(repeated_string_field); - assert_eq!( - repeated_string.len(), - 3, - "Should accumulate all repeated string values" - ); + assert_eq!(repeated_string.len(), 3, "Should accumulate all repeated string values"); assert_eq!(repeated_string[0], FieldValueRef::String("first")); assert_eq!(repeated_string[1], FieldValueRef::String("second")); assert_eq!(repeated_string[2], FieldValueRef::String("third")); @@ -2237,11 +2078,7 @@ fn test_mixed_packed_unpacked_encoding(#[case] version: ProtoVersion) { encode_key(repeated_int32_field as u32, WireType::Varint, &mut buf); encode_varint(100, &mut buf); - encode_key( - repeated_int32_field as u32, - WireType::LengthDelimited, - &mut buf, - ); + encode_key(repeated_int32_field as u32, WireType::LengthDelimited, &mut buf); let mut packed_data = Vec::new(); encode_varint(200, &mut packed_data); encode_varint(300, &mut packed_data); @@ -2254,11 +2091,7 @@ fn test_mixed_packed_unpacked_encoding(#[case] version: ProtoVersion) { let parsed = ParsedMessage::parse(&buf, ®istry).unwrap(); let repeated = parsed.get_repeated_scalars(repeated_int32_field); - assert_eq!( - repeated.len(), - 4, - "Should handle mixed packed and unpacked encoding" - ); + assert_eq!(repeated.len(), 4, "Should handle mixed packed and unpacked encoding"); assert_eq!(repeated[0], FieldValueRef::Int32(100)); assert_eq!(repeated[1], FieldValueRef::Int32(200)); assert_eq!(repeated[2], FieldValueRef::Int32(300)); @@ -2326,14 +2159,8 @@ fn test_length_delimited_field_with_zero_length(#[case] version: ProtoVersion) { let nested = parsed .get_message(all_types_ordinals::F_NESTED) .expect("Empty nested message should be present in proto2"); - assert!( - !nested.has_field(1), - "Empty nested message should have no fields" - ); - assert!( - !nested.has_field(2), - "Empty nested message should have no fields" - ); + assert!(!nested.has_field(1), "Empty nested message should have no fields"); + assert!(!nested.has_field(2), "Empty nested message should have no fields"); } ProtoVersion::Proto3 => { assert!( @@ -2347,14 +2174,8 @@ fn test_length_delimited_field_with_zero_length(#[case] version: ProtoVersion) { let nested = parsed.get_message(all_types_ordinals::F_NESTED); if let Some(nested_msg) = nested { - assert!( - !nested_msg.has_field(1), - "Empty nested message should have no fields" - ); - assert!( - !nested_msg.has_field(2), - "Empty nested message should have no fields" - ); + assert!(!nested_msg.has_field(1), "Empty nested message should have no fields"); + assert!(!nested_msg.has_field(2), "Empty nested message should have no fields"); } } } @@ -2838,10 +2659,7 @@ fn test_invalid_field_numbers(#[case] version: ProtoVersion) { let test_cases = vec![ (vec![0x00, 0x42], "field number 0"), - ( - vec![0x80, 0x80, 0x80, 0x80, 0x10, 0x42], - "field number > max", - ), + (vec![0x80, 0x80, 0x80, 0x80, 0x10, 0x42], "field number > max"), ]; for (malformed, description) in test_cases { @@ -2897,18 +2715,9 @@ fn test_invalid_utf8_various_errors(#[case] version: ProtoVersion) { (vec![0x62, 0x02, 0xC0, 0x81], "overlong encoding"), (vec![0x62, 0x02, 0xC2, 0x00], "invalid continuation byte"), (vec![0x62, 0x02, 0xE0, 0xA0], "truncated multibyte sequence"), - ( - vec![0x62, 0x03, 0xED, 0xA0, 0x80], - "UTF-16 surrogate halves", - ), - ( - vec![0x62, 0x03, 0x80, 0x80, 0x80], - "continuation without start byte", - ), - ( - vec![0x62, 0x04, 0xF5, 0x80, 0x80, 0x80], - "invalid 4-byte sequence (out of range)", - ), + (vec![0x62, 0x03, 0xED, 0xA0, 0x80], "UTF-16 surrogate halves"), + (vec![0x62, 0x03, 0x80, 0x80, 0x80], "continuation without start byte"), + (vec![0x62, 0x04, 0xF5, 0x80, 0x80, 0x80], "invalid 4-byte sequence (out of range)"), ]; for (malformed, description) in test_cases { @@ -3026,22 +2835,13 @@ fn test_proto2_required_and_default_values(#[case] version: ProtoVersion) { encode_message_for_version(version, &msg_unset, "AllTypesMessage"); let parsed_unset = ParsedMessage::parse(&buf_unset, ®istry_unset).unwrap(); - assert!( - parsed_unset.has_field(all_types_ordinals::F_REQUIRED), - "f_required should be present" - ); + assert!(parsed_unset.has_field(all_types_ordinals::F_REQUIRED), "f_required should be present"); assert_eq!( parsed_unset.get_scalar(all_types_ordinals::F_REQUIRED), Some(&FieldValueRef::Int32(42)) ); - assert!( - !parsed_unset.has_field(all_types_ordinals::F_INT32), - "f_int32 should be absent" - ); - assert!( - !parsed_unset.has_field(all_types_ordinals::F_STRING), - "f_string should be absent" - ); + assert!(!parsed_unset.has_field(all_types_ordinals::F_INT32), "f_int32 should be absent"); + assert!(!parsed_unset.has_field(all_types_ordinals::F_STRING), "f_string should be absent"); assert!( !parsed_unset.has_field(all_types_ordinals::F_DEFAULT_INT), "f_default_int should be absent when not set" @@ -3139,38 +2939,14 @@ fn test_zero_values_proto2(#[case] version: ProtoVersion) { let parsed = ParsedMessage::parse(&buf, ®istry).unwrap(); - assert!( - parsed.has_field(all_types_ordinals::F_INT32), - "f_int32 should be present" - ); - assert_eq!( - parsed.get_scalar(all_types_ordinals::F_INT32), - Some(&FieldValueRef::Int32(0)) - ); - assert!( - parsed.has_field(all_types_ordinals::F_INT64), - "f_int64 should be present" - ); - assert_eq!( - parsed.get_scalar(all_types_ordinals::F_INT64), - Some(&FieldValueRef::Int64(0)) - ); - assert!( - parsed.has_field(all_types_ordinals::F_BOOL), - "f_bool should be present" - ); - assert_eq!( - parsed.get_scalar(all_types_ordinals::F_BOOL), - Some(&FieldValueRef::Bool(false)) - ); - assert!( - parsed.has_field(all_types_ordinals::F_STRING), - "f_string should be present" - ); - assert_eq!( - parsed.get_scalar(all_types_ordinals::F_STRING), - Some(&FieldValueRef::String("")) - ); + assert!(parsed.has_field(all_types_ordinals::F_INT32), "f_int32 should be present"); + assert_eq!(parsed.get_scalar(all_types_ordinals::F_INT32), Some(&FieldValueRef::Int32(0))); + assert!(parsed.has_field(all_types_ordinals::F_INT64), "f_int64 should be present"); + assert_eq!(parsed.get_scalar(all_types_ordinals::F_INT64), Some(&FieldValueRef::Int64(0))); + assert!(parsed.has_field(all_types_ordinals::F_BOOL), "f_bool should be present"); + assert_eq!(parsed.get_scalar(all_types_ordinals::F_BOOL), Some(&FieldValueRef::Bool(false))); + assert!(parsed.has_field(all_types_ordinals::F_STRING), "f_string should be present"); + assert_eq!(parsed.get_scalar(all_types_ordinals::F_STRING), Some(&FieldValueRef::String(""))); } // ============================================================================ @@ -3271,10 +3047,7 @@ fn test_zero_values_proto3(#[case] version: ProtoVersion) { !parsed.has_field(all_types_ordinals::F_INT64), "f_int64 zero value should not be present" ); - assert!( - !parsed.has_field(all_types_ordinals::F_BOOL), - "f_bool false should not be present" - ); + assert!(!parsed.has_field(all_types_ordinals::F_BOOL), "f_bool false should not be present"); assert!( !parsed.has_field(all_types_ordinals::F_STRING), "f_string empty should not be present" @@ -3290,11 +3063,7 @@ fn test_proto3_packed_vs_unpacked_encoding(#[case] version: ProtoVersion) { let repeated_int32_field = field_num(version, "f_repeated_int32"); let mut buf_packed = Vec::new(); - encode_key( - repeated_int32_field as u32, - WireType::LengthDelimited, - &mut buf_packed, - ); + encode_key(repeated_int32_field as u32, WireType::LengthDelimited, &mut buf_packed); let values = vec![10, 20, 30]; let mut packed_data = Vec::new(); for val in &values { @@ -3312,11 +3081,7 @@ fn test_proto3_packed_vs_unpacked_encoding(#[case] version: ProtoVersion) { let mut buf_unpacked = Vec::new(); for val in &values { - encode_key( - repeated_int32_field as u32, - WireType::Varint, - &mut buf_unpacked, - ); + encode_key(repeated_int32_field as u32, WireType::Varint, &mut buf_unpacked); encode_varint(*val as u64, &mut buf_unpacked); } @@ -3342,11 +3107,7 @@ fn test_proto2_packed_unpacked_cross_acceptance(#[case] version: ProtoVersion) { let packed_field = field_num(version, "f_repeated_packed"); let unpacked_field = field_num(version, "f_repeated_unpacked"); let values = [10i32, 20, 30]; - let expected = [ - FieldValueRef::Int32(10), - FieldValueRef::Int32(20), - FieldValueRef::Int32(30), - ]; + let expected = [FieldValueRef::Int32(10), FieldValueRef::Int32(20), FieldValueRef::Int32(30)]; // (A) `[packed = true]` field encoded with UNPACKED wire bytes. let mut buf = Vec::new(); diff --git a/rust/sdk/src/zeroparser/types.rs b/rust/sdk/src/zeroparser/types.rs index 142010d5..6e023f87 100644 --- a/rust/sdk/src/zeroparser/types.rs +++ b/rust/sdk/src/zeroparser/types.rs @@ -44,13 +44,13 @@ pub(crate) fn convert_scalar_value<'a>( match field_type { Type::String => Ok(FieldValueRef::String(wire_value.try_as_str(field_num)?)), Type::Int32 => Ok(FieldValueRef::Int32(wire_value.try_as_i32(field_num)?)), - Type::Sint32 => Ok(FieldValueRef::Int32(decode_zigzag32( - wire_value.try_as_u32(field_num)?, - ))), + Type::Sint32 => { + Ok(FieldValueRef::Int32(decode_zigzag32(wire_value.try_as_u32(field_num)?))) + } Type::Int64 => Ok(FieldValueRef::Int64(wire_value.try_as_i64(field_num)?)), - Type::Sint64 => Ok(FieldValueRef::Int64(decode_zigzag64( - wire_value.try_as_u64(field_num)?, - ))), + Type::Sint64 => { + Ok(FieldValueRef::Int64(decode_zigzag64(wire_value.try_as_u64(field_num)?))) + } Type::Uint32 | Type::Fixed32 => { Ok(FieldValueRef::UInt32(wire_value.try_as_u32(field_num)?)) } @@ -355,10 +355,7 @@ mod tests { fn field_value_ref_traits() { assert_eq!(FieldValueRef::Int32(42), FieldValueRef::Int32(42)); assert_ne!(FieldValueRef::Int32(42), FieldValueRef::Int64(42)); - assert_eq!( - FieldValueRef::String("hello"), - FieldValueRef::String("hello") - ); + assert_eq!(FieldValueRef::String("hello"), FieldValueRef::String("hello")); let val = FieldValueRef::UInt64(12345); let copy = val; @@ -515,69 +512,24 @@ mod tests { #[test] fn default_value_for_all_types() { - assert_eq!( - default_value_for_type(Type::String), - FieldValueRef::String("") - ); + assert_eq!(default_value_for_type(Type::String), FieldValueRef::String("")); assert_eq!(default_value_for_type(Type::Int32), FieldValueRef::Int32(0)); assert_eq!(default_value_for_type(Type::Int64), FieldValueRef::Int64(0)); - assert_eq!( - default_value_for_type(Type::Uint32), - FieldValueRef::UInt32(0) - ); - assert_eq!( - default_value_for_type(Type::Uint64), - FieldValueRef::UInt64(0) - ); - assert_eq!( - default_value_for_type(Type::Sint32), - FieldValueRef::Int32(0) - ); - assert_eq!( - default_value_for_type(Type::Sint64), - FieldValueRef::Int64(0) - ); - assert_eq!( - default_value_for_type(Type::Fixed32), - FieldValueRef::UInt32(0) - ); - assert_eq!( - default_value_for_type(Type::Fixed64), - FieldValueRef::UInt64(0) - ); - assert_eq!( - default_value_for_type(Type::Sfixed32), - FieldValueRef::Int32(0) - ); - assert_eq!( - default_value_for_type(Type::Sfixed64), - FieldValueRef::Int64(0) - ); - assert_eq!( - default_value_for_type(Type::Bool), - FieldValueRef::Bool(false) - ); - assert_eq!( - default_value_for_type(Type::Float), - FieldValueRef::Float(0.0) - ); - assert_eq!( - default_value_for_type(Type::Double), - FieldValueRef::Double(0.0) - ); - assert_eq!( - default_value_for_type(Type::Bytes), - FieldValueRef::Bytes(&[]) - ); - assert_eq!( - default_value_for_type(Type::Message), - FieldValueRef::Bytes(&[]) - ); + assert_eq!(default_value_for_type(Type::Uint32), FieldValueRef::UInt32(0)); + assert_eq!(default_value_for_type(Type::Uint64), FieldValueRef::UInt64(0)); + assert_eq!(default_value_for_type(Type::Sint32), FieldValueRef::Int32(0)); + assert_eq!(default_value_for_type(Type::Sint64), FieldValueRef::Int64(0)); + assert_eq!(default_value_for_type(Type::Fixed32), FieldValueRef::UInt32(0)); + assert_eq!(default_value_for_type(Type::Fixed64), FieldValueRef::UInt64(0)); + assert_eq!(default_value_for_type(Type::Sfixed32), FieldValueRef::Int32(0)); + assert_eq!(default_value_for_type(Type::Sfixed64), FieldValueRef::Int64(0)); + assert_eq!(default_value_for_type(Type::Bool), FieldValueRef::Bool(false)); + assert_eq!(default_value_for_type(Type::Float), FieldValueRef::Float(0.0)); + assert_eq!(default_value_for_type(Type::Double), FieldValueRef::Double(0.0)); + assert_eq!(default_value_for_type(Type::Bytes), FieldValueRef::Bytes(&[])); + assert_eq!(default_value_for_type(Type::Message), FieldValueRef::Bytes(&[])); assert_eq!(default_value_for_type(Type::Enum), FieldValueRef::Int32(0)); - assert_eq!( - default_value_for_type(Type::Group), - FieldValueRef::Bytes(&[]) - ); + assert_eq!(default_value_for_type(Type::Group), FieldValueRef::Bytes(&[])); } #[test] @@ -714,10 +666,7 @@ mod tests { 1, /* field_num */ ); assert!(result.is_err()); - assert!(matches!( - result, - Err(ParseError::InvalidPackedFieldType { field_num: 1 }) - )); + assert!(matches!(result, Err(ParseError::InvalidPackedFieldType { field_num: 1 }))); // Fixed32 with non-fixed32 type. let mut dest = Vec::new(); @@ -787,10 +736,7 @@ mod tests { // Int64. let key = MapKeyRef::Int64(-1_000_000_000_000); - assert_eq!( - key.to_field_value(), - FieldValueRef::Int64(-1_000_000_000_000) - ); + assert_eq!(key.to_field_value(), FieldValueRef::Int64(-1_000_000_000_000)); // UInt32. let key = MapKeyRef::UInt32(42); @@ -798,10 +744,7 @@ mod tests { // UInt64. let key = MapKeyRef::UInt64(1_000_000_000_000); - assert_eq!( - key.to_field_value(), - FieldValueRef::UInt64(1_000_000_000_000) - ); + assert_eq!(key.to_field_value(), FieldValueRef::UInt64(1_000_000_000_000)); // Bool true. let key = MapKeyRef::Bool(true); diff --git a/rust/sdk/src/zeroparser/wire.rs b/rust/sdk/src/zeroparser/wire.rs index 0a48eb56..85a690fa 100644 --- a/rust/sdk/src/zeroparser/wire.rs +++ b/rust/sdk/src/zeroparser/wire.rs @@ -226,10 +226,9 @@ pub fn try_read_varint(data: &[u8]) -> ParseResult<(u64, &[u8])> { // Only 2 bytes but continuation bit set. [_, _] => Err(ParseError::TruncatedVarint), // Fast path: 3-byte varint (values 16384-2097151). - [b0, b1, b2, ref rest @ ..] if b2 < 0x80 => Ok(( - ((b0 & 0x7f) as u64) | (((b1 & 0x7f) as u64) << 7) | ((b2 as u64) << 14), - rest, - )), + [b0, b1, b2, ref rest @ ..] if b2 < 0x80 => { + Ok((((b0 & 0x7f) as u64) | (((b1 & 0x7f) as u64) << 7) | ((b2 as u64) << 14), rest)) + } // Only 3 bytes but continuation bit set. [_, _, _] => Err(ParseError::TruncatedVarint), // Fast path: 4-byte varint (values 2097152-268435455). @@ -400,27 +399,16 @@ mod tests { fn varint_errors() { // Truncated varints. assert_eq!(try_read_varint(&[0x80]), Err(ParseError::TruncatedVarint)); - assert_eq!( - try_read_varint(&[0x80, 0x80]), - Err(ParseError::TruncatedVarint) - ); - assert_eq!( - try_read_varint(&[0x80, 0x80, 0x80]), - Err(ParseError::TruncatedVarint) - ); - assert_eq!( - try_read_varint(&[0x80, 0x80, 0x80, 0x80]), - Err(ParseError::TruncatedVarint) - ); + assert_eq!(try_read_varint(&[0x80, 0x80]), Err(ParseError::TruncatedVarint)); + assert_eq!(try_read_varint(&[0x80, 0x80, 0x80]), Err(ParseError::TruncatedVarint)); + assert_eq!(try_read_varint(&[0x80, 0x80, 0x80, 0x80]), Err(ParseError::TruncatedVarint)); assert_eq!( try_read_varint(&[0x80, 0x80, 0x80, 0x80, 0x80]), Err(ParseError::TruncatedVarint) ); // Varint too long (11 bytes). - let too_long = &[ - 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x01, - ]; + let too_long = &[0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x01]; assert_eq!(try_read_varint(too_long), Err(ParseError::VarintTooLong)); } @@ -449,37 +437,20 @@ mod tests { // Invalid: 10th byte = 0x02 (bit 1 set, would overflow u64). let overflow_bit1 = &[0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x02]; - assert_eq!( - try_read_varint(overflow_bit1), - Err(ParseError::VarintTooLong) - ); + assert_eq!(try_read_varint(overflow_bit1), Err(ParseError::VarintTooLong)); // Invalid: 10th byte = 0x7F (bits 1-6 all set, would overflow u64). let overflow_bits = &[0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x7F]; - assert_eq!( - try_read_varint(overflow_bits), - Err(ParseError::VarintTooLong) - ); + assert_eq!(try_read_varint(overflow_bits), Err(ParseError::VarintTooLong)); // Invalid: 10th byte = 0x80 (continuation bit set, would need 11+ bytes). let continuation = &[0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80]; - assert_eq!( - try_read_varint(continuation), - Err(ParseError::VarintTooLong) - ); + assert_eq!(try_read_varint(continuation), Err(ParseError::VarintTooLong)); } #[test] fn zigzag_decoding() { - let cases32 = [ - (0, 0), - (1, -1), - (2, 1), - (3, -2), - (4, 2), - (99, -50), - (100, 50), - ]; + let cases32 = [(0, 0), (1, -1), (2, 1), (3, -2), (4, 2), (99, -50), (100, 50)]; for (encoded, expected) in cases32 { assert_eq!(decode_zigzag32(encoded), expected, "zigzag32({})", encoded); } @@ -504,14 +475,8 @@ mod tests { assert_eq!(WireType::try_from(val), Ok(expected)); } - assert_eq!( - WireType::try_from(6u64), - Err(ParseError::InvalidWireType(6)) - ); - assert_eq!( - WireType::try_from(7u64), - Err(ParseError::InvalidWireType(7)) - ); + assert_eq!(WireType::try_from(6u64), Err(ParseError::InvalidWireType(6))); + assert_eq!(WireType::try_from(7u64), Err(ParseError::InvalidWireType(7))); } #[test] @@ -521,11 +486,7 @@ mod tests { // Varint: field 1, value 150. Tag = 8, 150 = 0x96 0x01. (&[8, 0x96, 0x01][..], 1, WireValue::Varint(150)), // I32: field 1, tag = 13, value 0x01020304 little-endian. - ( - &[13, 0x04, 0x03, 0x02, 0x01][..], - 1, - WireValue::I32(0x01020304), - ), + (&[13, 0x04, 0x03, 0x02, 0x01][..], 1, WireValue::I32(0x01020304)), // I64: field 1, tag = 9. ( &[9, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08][..], @@ -533,11 +494,7 @@ mod tests { WireValue::I64(0x0807060504030201), ), // Len: field 1, tag = 10, length 3. - ( - &[10, 3, 0xAA, 0xBB, 0xCC][..], - 1, - WireValue::Len(&[0xAA, 0xBB, 0xCC]), - ), + (&[10, 3, 0xAA, 0xBB, 0xCC][..], 1, WireValue::Len(&[0xAA, 0xBB, 0xCC])), ]; for (data, expected_num, expected_val) in cases { let (field, rest) = try_parse_field(data).unwrap(); @@ -553,10 +510,7 @@ mod tests { assert_eq!(try_parse_field(&[14]), Err(ParseError::InvalidWireType(6))); // Group wire type: tag = (1 << 3) | 3 = 11. - assert_eq!( - try_parse_field(&[11]), - Err(ParseError::UnsupportedGroupWireType) - ); + assert_eq!(try_parse_field(&[11]), Err(ParseError::UnsupportedGroupWireType)); // Buffer too short for I32: tag 13, only 2 bytes. assert_eq!( @@ -589,10 +543,7 @@ mod tests { ); // Invalid field number 0: tag = (0 << 3) | 0 = 0. - assert_eq!( - try_parse_field(&[0]), - Err(ParseError::InvalidFieldNumber { field_num: 0 }) - ); + assert_eq!(try_parse_field(&[0]), Err(ParseError::InvalidFieldNumber { field_num: 0 })); // Invalid field number 536_870_912 (2^29) - exceeds max valid field number 536_870_911. // Tag = (536_870_912 << 3) | 0 = 4_294_967_296 -> varint [0x80, 0x80, 0x80, 0x80, 0x10]. @@ -611,10 +562,7 @@ mod tests { #[test] fn wire_value_accessors() { - assert_eq!( - WireValue::Len(b"hello").try_as_str(1 /* field_num */), - Ok("hello") - ); + assert_eq!(WireValue::Len(b"hello").try_as_str(1 /* field_num */), Ok("hello")); assert_eq!( WireValue::Len(&[0xFF, 0xFE]).try_as_str(1 /* field_num */), Err(ParseError::InvalidUtf8 { field_num: 1 }) @@ -624,32 +572,17 @@ mod tests { Err(ParseError::TypeMismatch { .. }) )); - assert_eq!( - WireValue::Len(&[1, 2, 3]).try_as_bytes(1 /* field_num */), - Ok(&[1, 2, 3][..]) - ); + assert_eq!(WireValue::Len(&[1, 2, 3]).try_as_bytes(1 /* field_num */), Ok(&[1, 2, 3][..])); assert_eq!(WireValue::Varint(42).try_as_i32(1 /* field_num */), Ok(42)); assert_eq!(WireValue::I32(100).try_as_i32(1 /* field_num */), Ok(100)); - assert_eq!( - WireValue::Varint(1000).try_as_u64(1 /* field_num */), - Ok(1000) - ); + assert_eq!(WireValue::Varint(1000).try_as_u64(1 /* field_num */), Ok(1000)); assert_eq!(WireValue::I64(2000).try_as_u64(1 /* field_num */), Ok(2000)); - assert_eq!( - WireValue::Varint(0).try_as_bool(1 /* field_num */), - Ok(false) - ); - assert_eq!( - WireValue::Varint(1).try_as_bool(1 /* field_num */), - Ok(true) - ); - assert_eq!( - WireValue::Varint(42).try_as_bool(1 /* field_num */), - Ok(true) - ); + assert_eq!(WireValue::Varint(0).try_as_bool(1 /* field_num */), Ok(false)); + assert_eq!(WireValue::Varint(1).try_as_bool(1 /* field_num */), Ok(true)); + assert_eq!(WireValue::Varint(42).try_as_bool(1 /* field_num */), Ok(true)); assert!( (WireValue::I32(std::f32::consts::PI.to_bits()) diff --git a/rust/tests/src/arrow_tests.rs b/rust/tests/src/arrow_tests.rs index 71a3eab3..75b645ec 100644 --- a/rust/tests/src/arrow_tests.rs +++ b/rust/tests/src/arrow_tests.rs @@ -9,9 +9,15 @@ mod arrow_flight_tests { use crate::mock_arrow_flight::{start_mock_flight_server, MockFlightResponse}; use crate::utils::{ - create_test_arrow_schema, create_test_dict_record_batch, create_test_dict_schema, - create_test_record_batch, record_batch_to_ipc_bytes, setup_tracing, - CountingHeadersProvider, HangingInvalidationHeadersProvider, TestHeadersProvider, + create_test_arrow_schema, + create_test_dict_record_batch, + create_test_dict_schema, + create_test_record_batch, + record_batch_to_ipc_bytes, + setup_tracing, + CountingHeadersProvider, + HangingInvalidationHeadersProvider, + TestHeadersProvider, }; const TABLE_NAME: &str = "test_catalog.test_schema.test_table"; @@ -56,11 +62,7 @@ mod arrow_flight_tests { .build_arrow() .await; - assert!( - result.is_ok(), - "Failed to create Arrow Flight stream: {:?}", - result.err() - ); + assert!(result.is_ok(), "Failed to create Arrow Flight stream: {:?}", result.err()); let stream = result.unwrap(); assert_eq!(stream.table_name(), TABLE_NAME); @@ -109,11 +111,7 @@ mod arrow_flight_tests { "expected an empty-batch InvalidArgument, got: {}", err ); - assert!( - !err.is_retryable(), - "empty-batch rejection is non-retryable, got: {}", - err - ); + assert!(!err.is_retryable(), "empty-batch rejection is non-retryable, got: {}", err); Ok(()) } @@ -577,15 +575,9 @@ mod arrow_flight_tests { .expect("close must reach the finalization barrier"); drop(close_future); - assert!( - !stream.is_closed(), - "cancelled teardown is not finalized yet" - ); + assert!(!stream.is_closed(), "cancelled teardown is not finalized yet"); let batch = create_test_record_batch(schema, vec![2], vec![Some("late")]); - assert!( - stream.ingest_batch(batch).await.is_err(), - "Closing must reject new ingests" - ); + assert!(stream.ingest_batch(batch).await.is_err(), "Closing must reject new ingests"); assert!( stream.get_unacked_batches().await.is_err(), "unacked retrieval is allowed only after Closed" @@ -736,10 +728,7 @@ mod arrow_flight_tests { let schema = create_test_arrow_schema(); mock_server - .inject_responses( - TABLE_NAME, - vec![MockFlightResponse::CloseStream { delay_ms: 0 }], - ) + .inject_responses(TABLE_NAME, vec![MockFlightResponse::CloseStream { delay_ms: 0 }]) .await; let sdk = ZerobusSdk::builder() @@ -976,12 +965,9 @@ mod arrow_flight_tests { // The auth rejection must not terminate the stream: recovery re-mints and the // batch is eventually acknowledged. - tokio::time::timeout( - std::time::Duration::from_secs(5), - stream.wait_for_offset(offset), - ) - .await - .expect("recovery after auth re-mint should complete")?; + tokio::time::timeout(std::time::Duration::from_secs(5), stream.wait_for_offset(offset)) + .await + .expect("recovery after auth re-mint should complete")?; // Exactly one invalidation: the single auth rejection on conn2 (conn1's // retriable error and conn3's successful setup must not invalidate). @@ -1062,10 +1048,7 @@ mod arrow_flight_tests { "must surface the original auth rejection, got: {}", err ); - assert!( - stream.is_closed(), - "stream must close after invalidation timeout" - ); + assert!(stream.is_closed(), "stream must close after invalidation timeout"); assert_eq!( invalidations.load(std::sync::atomic::Ordering::SeqCst), 1, @@ -1150,10 +1133,7 @@ mod arrow_flight_tests { 1, "terminal auth cleanup should be attempted once" ); - assert!( - stream.is_closed(), - "terminal auth failure must close the stream" - ); + assert!(stream.is_closed(), "terminal auth failure must close the stream"); Ok(()) } @@ -1389,21 +1369,15 @@ mod arrow_flight_tests { use arrow_array::Int32Array; use arrow_schema::{DataType, Field, Schema}; - let wrong_schema = Arc::new(Schema::new(vec![Field::new( - "different_field", - DataType::Int32, - false, - )])); + let wrong_schema = + Arc::new(Schema::new(vec![Field::new("different_field", DataType::Int32, false)])); let wrong_batch = arrow_array::RecordBatch::try_new( wrong_schema, vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], )?; let result = stream.ingest_batch(wrong_batch).await; - assert!( - result.is_err(), - "Expected schema mismatch error, but got Ok" - ); + assert!(result.is_err(), "Expected schema mismatch error, but got Ok"); Ok(()) } @@ -1689,10 +1663,7 @@ mod arrow_flight_tests { let _offset = stream.ingest_batch(batch).await?; let close_result = stream.close().await; - assert!( - close_result.is_err(), - "close() must propagate the error when its flush fails" - ); + assert!(close_result.is_err(), "close() must propagate the error when its flush fails"); // The stream is still torn down, and the unacked batch is recoverable. assert!(stream.is_closed()); @@ -1900,11 +1871,7 @@ mod arrow_flight_tests { let joined = tokio::time::timeout(std::time::Duration::from_secs(3), handle) .await .expect("2nd ingest_batch should unblock after the ack frees a permit")?; - assert!( - joined.is_ok(), - "2nd ingest_batch failed: {:?}", - joined.err() - ); + assert!(joined.is_ok(), "2nd ingest_batch failed: {:?}", joined.err()); Ok(()) } @@ -2003,10 +1970,7 @@ mod arrow_flight_tests { stream.close().await?; let unacked = stream.get_unacked_batches().await?; - assert!( - unacked.is_empty(), - "All batches were acked, should be empty" - ); + assert!(unacked.is_empty(), "All batches were acked, should be empty"); Ok(()) } @@ -2139,11 +2103,7 @@ mod arrow_flight_tests { vec![2, 3], "batch A must be sliced to its un-acked suffix (ids [2, 3])" ); - assert_eq!( - batch_ids(&unacked[1]), - vec![4], - "batch B must be fully retained" - ); + assert_eq!(batch_ids(&unacked[1]), vec![4], "batch B must be fully retained"); Ok(()) } @@ -2220,10 +2180,7 @@ mod arrow_flight_tests { vec![vec![2, 3], vec![4]], "first snapshot should be sliced A suffix (ids [2,3]) + full B (id [4])" ); - assert_eq!( - first, second, - "repeated get_unacked_batches must be idempotent" - ); + assert_eq!(first, second, "repeated get_unacked_batches must be idempotent"); Ok(()) } @@ -2392,10 +2349,7 @@ mod arrow_flight_tests { // But every auto-ack must be connection-relative: exactly the 3 rows // replayed on the recovered connection, never the cumulative 6. let acks = mock_server.get_auto_ack_records().await; - assert!( - !acks.is_empty(), - "expected an auto-ack on the recovered connection" - ); + assert!(!acks.is_empty(), "expected an auto-ack on the recovered connection"); assert!( acks.iter().all(|&r| r == 3), "auto-ack must exclude rows from the first connection (expected all == 3), got {:?}", @@ -2570,18 +2524,11 @@ mod arrow_flight_tests { tokio::time::timeout(std::time::Duration::from_secs(5), stream.close()) .await .expect("close() must not hang while a reconnect is parked"); - assert!( - close_result.is_err(), - "close() should surface the flush timeout" - ); + assert!(close_result.is_err(), "close() should surface the flush timeout"); // The un-acked batch was moved to the failed set and is retrievable. let unacked = stream.get_unacked_batches().await?; - assert_eq!( - unacked.len(), - 1, - "pending batch must be moved to the failed set" - ); + assert_eq!(unacked.len(), 1, "pending batch must be moved to the failed set"); Ok(()) } @@ -2839,10 +2786,7 @@ mod arrow_flight_tests { tokio::time::sleep(Duration::from_millis(20)).await; waited += Duration::from_millis(20); } - assert!( - stream.is_closed(), - "stream should be closed by the schema reconnect failure" - ); + assert!(stream.is_closed(), "stream should be closed by the schema reconnect failure"); // A flush() starting after close must still surface the typed error. match stream.flush().await { @@ -3077,11 +3021,7 @@ mod arrow_flight_tests { let offset2 = stream.ingest_batch(batch2).await?; let result = stream.wait_for_offset(offset2).await; - assert!( - result.is_ok(), - "Expected partial batch recovery to succeed: {:?}", - result - ); + assert!(result.is_ok(), "Expected partial batch recovery to succeed: {:?}", result); stream.close().await?; @@ -3557,10 +3497,7 @@ mod arrow_flight_tests { .ingest_ipc_batch(bytes::Bytes::from_static(b"not valid arrow ipc")) .await; - assert!( - result.is_err(), - "Expected InvalidArgument for garbage bytes" - ); + assert!(result.is_err(), "Expected InvalidArgument for garbage bytes"); Ok(()) } @@ -3662,11 +3599,7 @@ mod arrow_flight_tests { .await?; let result = stream.wait_for_offset(offset2).await; - assert!( - result.is_ok(), - "Expected IPC batch recovery to succeed: {:?}", - result - ); + assert!(result.is_ok(), "Expected IPC batch recovery to succeed: {:?}", result); Ok(()) } @@ -3757,11 +3690,8 @@ mod arrow_flight_tests { // Create IPC bytes with a different schema. use arrow_array::Int32Array; use arrow_schema::{DataType, Field, Schema}; - let wrong_schema = Arc::new(Schema::new(vec![Field::new( - "different_field", - DataType::Int32, - false, - )])); + let wrong_schema = + Arc::new(Schema::new(vec![Field::new("different_field", DataType::Int32, false)])); let wrong_batch = arrow_array::RecordBatch::try_new( wrong_schema, vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], @@ -4216,9 +4146,10 @@ mod arrow_flight_tests { } mod graceful_close_tests { - use super::*; use std::time::Instant; + use super::*; + #[tokio::test] async fn test_default_graceful_close_waits_for_full_server_duration( ) -> Result<(), Box> { diff --git a/rust/tests/src/mock_arrow_flight.rs b/rust/tests/src/mock_arrow_flight.rs index 07acf1f2..9068bfa6 100644 --- a/rust/tests/src/mock_arrow_flight.rs +++ b/rust/tests/src/mock_arrow_flight.rs @@ -7,8 +7,18 @@ use std::time::Duration; use arrow_flight::flight_service_server::{FlightService, FlightServiceServer}; use arrow_flight::{ - Action, ActionType, Criteria, Empty, FlightData, FlightDescriptor, FlightInfo, - HandshakeRequest, HandshakeResponse, PutResult, SchemaResult, Ticket, + Action, + ActionType, + Criteria, + Empty, + FlightData, + FlightDescriptor, + FlightInfo, + HandshakeRequest, + HandshakeResponse, + PutResult, + SchemaResult, + Ticket, }; use futures::Stream; use serde::{Deserialize, Serialize}; @@ -491,10 +501,7 @@ impl FlightService for MockFlightServer { }; let ack_bytes = serde_json::to_vec(&ack_metadata).unwrap(); - debug!( - "Auto-acking offset: {}, records: {}", - metadata.offset_id, records - ); + debug!("Auto-acking offset: {}, records: {}", metadata.offset_id, records); let put_result = PutResult { app_metadata: ack_bytes.into(), }; diff --git a/rust/tests/src/mock_grpc.rs b/rust/tests/src/mock_grpc.rs index 54003d6d..d99ad1c2 100644 --- a/rust/tests/src/mock_grpc.rs +++ b/rust/tests/src/mock_grpc.rs @@ -8,11 +8,14 @@ pub mod databricks { tonic::include_proto!("databricks.zerobus"); } } +use databricks::zerobus::ephemeral_stream_request::Payload as RequestPayload; +use databricks::zerobus::ephemeral_stream_response::Payload as ResponsePayload; +use databricks::zerobus::zerobus_server::{Zerobus, ZerobusServer}; use databricks::zerobus::{ - ephemeral_stream_request::Payload as RequestPayload, - ephemeral_stream_response::Payload as ResponsePayload, - zerobus_server::{Zerobus, ZerobusServer}, - CloseStreamSignal, CreateIngestStreamResponse, EphemeralStreamRequest, EphemeralStreamResponse, + CloseStreamSignal, + CreateIngestStreamResponse, + EphemeralStreamRequest, + EphemeralStreamResponse, IngestRecordResponse, }; use prost_types::Duration as ProtobufDuration; @@ -426,11 +429,9 @@ async fn handle_mock_response( request_type, ack_up_to_offset ); let response = EphemeralStreamResponse { - payload: Some(ResponsePayload::IngestRecordResponse( - IngestRecordResponse { - durability_ack_up_to_offset: Some(*ack_up_to_offset), - }, - )), + payload: Some(ResponsePayload::IngestRecordResponse(IngestRecordResponse { + durability_ack_up_to_offset: Some(*ack_up_to_offset), + })), }; if tx.send(Ok(response)).await.is_err() { return (false, current_index); @@ -447,10 +448,7 @@ async fn handle_mock_response( if *delay_ms > 0 { sleep(Duration::from_millis(*delay_ms)).await; } - info!( - "Sending CloseStreamSignal with duration: {}s", - duration_seconds - ); + info!("Sending CloseStreamSignal with duration: {}s", duration_seconds); let response = EphemeralStreamResponse { payload: Some(ResponsePayload::CloseStreamSignal(CloseStreamSignal { duration: Some(ProtobufDuration { diff --git a/rust/tests/src/multiplexed_stream_tests.rs b/rust/tests/src/multiplexed_stream_tests.rs index c8cc8538..19c2bd3f 100644 --- a/rust/tests/src/multiplexed_stream_tests.rs +++ b/rust/tests/src/multiplexed_stream_tests.rs @@ -4,7 +4,12 @@ mod utils; use std::sync::Arc; use databricks_zerobus_ingest_sdk::{ - MessageId, MultiplexedStream, NoTlsConfig, ZerobusError, ZerobusSdk, ZerobusStream, + MessageId, + MultiplexedStream, + NoTlsConfig, + ZerobusError, + ZerobusSdk, + ZerobusStream, }; use mock_grpc::{start_mock_server, MockResponse}; use tracing::info; @@ -227,9 +232,10 @@ mod single_stream_tests { } mod multi_stream_tests { - use super::*; use std::time::Duration; + use super::*; + /// Use separate table names per stream so each gRPC connection gets its own response sequence. const TABLE_A: &str = "multi.schema.table_a"; const TABLE_B: &str = "multi.schema.table_b"; @@ -675,10 +681,7 @@ mod failure_tests { // The sub-stream closed (non-retryable error), so the mux should be poisoned // and further ingest should fail with InvalidStateError. - assert!( - mux.is_closed(), - "Expected mux to be poisoned after sub-stream close" - ); + assert!(mux.is_closed(), "Expected mux to be poisoned after sub-stream close"); let ingest_after = mux.ingest_record(b"record3".to_vec()).await; assert!( matches!(ingest_after, Err(ZerobusError::InvalidStateError(_))), @@ -721,10 +724,7 @@ mod failure_tests { // Non-retryable error → sub-stream closes → flush errors → mux poisoned. let flush_result = mux.flush().await; assert!(flush_result.is_err(), "Expected flush to fail"); - assert!( - mux.is_closed(), - "Expected mux poisoned after sub-stream close" - ); + assert!(mux.is_closed(), "Expected mux poisoned after sub-stream close"); let ingest_after = mux.ingest_record(b"data".to_vec()).await; assert!( @@ -816,10 +816,7 @@ mod failure_tests { let unacked: Vec<_> = mux.get_unacked_records().await?.collect(); assert!(unacked.is_empty(), "All records were acked"); - assert!( - mux.is_closed(), - "get_unacked_records should have closed the mux" - ); + assert!(mux.is_closed(), "get_unacked_records should have closed the mux"); Ok(()) } diff --git a/rust/tests/src/proxy_tests.rs b/rust/tests/src/proxy_tests.rs index 7b5bd589..f5e1f652 100644 --- a/rust/tests/src/proxy_tests.rs +++ b/rust/tests/src/proxy_tests.rs @@ -130,10 +130,7 @@ async fn test_proxy_and_no_proxy() -> Result<(), Box> { .await; let (proxy_url, connect_count) = start_mock_proxy().await; - info!( - "Mock proxy at: {}, mock server at: {}", - proxy_url, server_url - ); + info!("Mock proxy at: {}, mock server at: {}", proxy_url, server_url); std::env::set_var("grpc_proxy", &proxy_url); @@ -143,10 +140,7 @@ async fn test_proxy_and_no_proxy() -> Result<(), Box> { let connects = connect_count.load(Ordering::SeqCst); info!("Proxy received {} CONNECT requests", connects); - assert!( - connects > 0, - "Expected proxy to receive CONNECT requests, got 0" - ); + assert!(connects > 0, "Expected proxy to receive CONNECT requests, got 0"); } // === Part 2: Verify no_proxy bypasses the proxy === diff --git a/rust/tests/src/rust_tests.rs b/rust/tests/src/rust_tests.rs index ce3761fd..da2dbfe6 100644 --- a/rust/tests/src/rust_tests.rs +++ b/rust/tests/src/rust_tests.rs @@ -45,11 +45,7 @@ mod stream_initialization_and_basic_lifecycle_tests { .recovery(false) .build() .await; - assert!( - result.is_ok(), - "Failed to create a stream: {:?}", - result.err() - ); + assert!(result.is_ok(), "Failed to create a stream: {:?}", result.err()); let stream = result.unwrap(); assert_eq!(stream.stream_type, StreamType::Ephemeral); @@ -378,10 +374,7 @@ mod stream_initialization_and_basic_lifecycle_tests { let ingest_result = stream .ingest_record_offset(b"test record data".to_vec()) .await; - assert!(matches!( - ingest_result, - Err(ZerobusError::StreamClosedError(_)) - )); + assert!(matches!(ingest_result, Err(ZerobusError::StreamClosedError(_)))); Ok(()) } @@ -424,10 +417,7 @@ mod stream_initialization_and_basic_lifecycle_tests { let batch = vec![b"record 1".to_vec(), b"record 2".to_vec()]; let ingest_result = stream.ingest_records_offset(batch).await; - assert!(matches!( - ingest_result, - Err(ZerobusError::StreamClosedError(_)) - )); + assert!(matches!(ingest_result, Err(ZerobusError::StreamClosedError(_)))); Ok(()) } @@ -525,11 +515,7 @@ mod stream_initialization_and_basic_lifecycle_tests { let duration = start.elapsed(); // Drop should be nearly instantaneous. - assert!( - duration.as_millis() < 100, - "Drop should be immediate, took {:?}", - duration - ); + assert!(duration.as_millis() < 100, "Drop should be immediate, took {:?}", duration); Ok(()) } @@ -585,10 +571,7 @@ mod stream_initialization_and_basic_lifecycle_tests { // Second close should also work. let second_close = stream.close().await; assert!(second_close.is_ok(), "Second close should be idempotent"); - assert!( - stream.is_closed(), - "Stream should still be closed after second close()" - ); + assert!(stream.is_closed(), "Stream should still be closed after second close()"); Ok(()) } @@ -1040,10 +1023,7 @@ mod schema_tests { let exact = vec![0u8; limit]; let result = stream.ingest_record_offset(exact).await; - assert!( - result.is_ok(), - "payload of exactly the limit should be accepted, got {result:?}" - ); + assert!(result.is_ok(), "payload of exactly the limit should be accepted, got {result:?}"); Ok(()) } @@ -1144,10 +1124,7 @@ mod schema_tests { .build() .await; - assert!( - result.is_ok(), - "JSON stream should be created even with descriptor" - ); + assert!(result.is_ok(), "JSON stream should be created even with descriptor"); Ok(()) } @@ -1306,11 +1283,7 @@ mod standard_operation_and_state_management_tests { .build() .await?; - let batch = vec![ - b"record 1".to_vec(), - b"record 2".to_vec(), - b"record 3".to_vec(), - ]; + let batch = vec![b"record 1".to_vec(), b"record 2".to_vec(), b"record 3".to_vec()]; let ack_future = stream.ingest_records_offset(batch).await?; if let Some(off) = ack_future { @@ -1744,10 +1717,7 @@ mod standard_operation_and_state_management_tests { if let Err(ZerobusError::StreamClosedError(_)) = flush_result { // Expected timeout error } else { - panic!( - "Expected StreamClosedError with timeout, got: {:?}", - flush_result - ); + panic!("Expected StreamClosedError with timeout, got: {:?}", flush_result); } Ok(()) @@ -1882,10 +1852,7 @@ mod concurrency_and_race_condition_tests { "Offsets should be a complete sequence from 0 to NUM_RECORDS - 1" ); assert_eq!(mock_server.get_write_count().await, NUM_RECORDS as u64); - assert_eq!( - mock_server.get_max_offset_sent().await, - (NUM_RECORDS - 1) as i64 - ); + assert_eq!(mock_server.get_max_offset_sent().await, (NUM_RECORDS - 1) as i64); Ok(()) } @@ -1976,10 +1943,7 @@ mod concurrency_and_race_condition_tests { } assert_eq!(mock_server.get_write_count().await, TOTAL_REQUESTS as u64); - assert_eq!( - mock_server.get_max_offset_sent().await, - (TOTAL_REQUESTS - 1) as i64 - ); + assert_eq!(mock_server.get_max_offset_sent().await, (TOTAL_REQUESTS - 1) as i64); Ok(()) } @@ -2071,14 +2035,8 @@ mod concurrency_and_race_condition_tests { } // Total writes = TOTAL_BATCHES * RECORDS_PER_BATCH - assert_eq!( - mock_server.get_write_count().await, - (TOTAL_BATCHES * RECORDS_PER_BATCH) as u64 - ); - assert_eq!( - mock_server.get_max_offset_sent().await, - (TOTAL_BATCHES - 1) as i64 - ); + assert_eq!(mock_server.get_write_count().await, (TOTAL_BATCHES * RECORDS_PER_BATCH) as u64); + assert_eq!(mock_server.get_max_offset_sent().await, (TOTAL_BATCHES - 1) as i64); Ok(()) } @@ -2767,10 +2725,7 @@ mod failure_scenarios_tests { .await; let duration = start_time.elapsed(); - assert!( - result.is_err(), - "Expected stream creation to fail after exhausting retries" - ); + assert!(result.is_err(), "Expected stream creation to fail after exhausting retries"); let error = result.err().unwrap(); assert!( @@ -2859,11 +2814,7 @@ mod failure_scenarios_tests { write_count ); - assert_eq!( - max_offset, 4, - "Expected max offset of 4 (records 0-4), got {}", - max_offset - ); + assert_eq!(max_offset, 4, "Expected max offset of 4 (records 0-4), got {}", max_offset); Ok(()) } @@ -2927,11 +2878,7 @@ mod failure_scenarios_tests { let max_offset = mock_server.get_max_offset_sent().await; assert_eq!(write_count, 5, "Expected 5 writes, got {}", write_count); - assert_eq!( - max_offset, 4, - "Expected max offset of 4 (records 0-4), got {}", - max_offset - ); + assert_eq!(max_offset, 4, "Expected max offset of 4 (records 0-4), got {}", max_offset); Ok(()) } @@ -2996,10 +2943,7 @@ mod failure_scenarios_tests { let ingest_after_failed_close = stream.ingest_record_offset(b"more data".to_vec()).await; assert!( - matches!( - ingest_after_failed_close, - Err(ZerobusError::StreamClosedError(_)) - ), + matches!(ingest_after_failed_close, Err(ZerobusError::StreamClosedError(_))), "Expected StreamClosedError after failed close" ); @@ -3082,11 +3026,7 @@ mod failure_scenarios_tests { write_count ); - assert_eq!( - max_offset, 2, - "Expected max physical offset of 2, got {}", - max_offset - ); + assert_eq!(max_offset, 2, "Expected max physical offset of 2, got {}", max_offset); Ok(()) } @@ -3283,11 +3223,7 @@ mod failure_scenarios_tests { // Original stream: 2 + 1 + 2 = 5 records written // Recreated stream: 1 + 2 = 3 records re-written (at least some should complete) // Total: at least 6 (may be up to 8 depending on timing) - assert!( - write_count >= 6, - "Expected at least 6 writes, got {}", - write_count - ); + assert!(write_count >= 6, "Expected at least 6 writes, got {}", write_count); // Verify the stream was successfully recreated assert_eq!(new_stream.stream_type, StreamType::Ephemeral); @@ -4071,11 +4007,7 @@ mod api_offset_tests { .build() .await?; - let batch = vec![ - b"record 1".to_vec(), - b"record 2".to_vec(), - b"record 3".to_vec(), - ]; + let batch = vec![b"record 1".to_vec(), b"record 2".to_vec(), b"record 3".to_vec()]; // ingest_records_offset returns the offset directly without nested future let offset = stream.ingest_records_offset(batch).await?; @@ -4142,10 +4074,7 @@ mod api_offset_tests { let elapsed = start.elapsed(); - assert!( - wait_result.is_err(), - "Expected wait_for_offset to fail with server error" - ); + assert!(wait_result.is_err(), "Expected wait_for_offset to fail with server error"); if let Err(e) = wait_result { let error_msg = e.to_string(); @@ -4419,10 +4348,7 @@ mod api_offset_tests { let ingest_result = stream .ingest_record_offset(b"test record data".to_vec()) .await; - assert!(matches!( - ingest_result, - Err(ZerobusError::StreamClosedError(_)) - )); + assert!(matches!(ingest_result, Err(ZerobusError::StreamClosedError(_)))); Ok(()) } @@ -4466,10 +4392,7 @@ mod api_offset_tests { let batch = vec![b"record 1".to_vec(), b"record 2".to_vec()]; let ingest_result = stream.ingest_records_offset(batch).await; - assert!(matches!( - ingest_result, - Err(ZerobusError::StreamClosedError(_)) - )); + assert!(matches!(ingest_result, Err(ZerobusError::StreamClosedError(_)))); Ok(()) } @@ -4563,10 +4486,7 @@ mod api_offset_tests { stream.flush().await?; assert_eq!(mock_server.get_write_count().await, TOTAL_REQUESTS as u64); - assert_eq!( - mock_server.get_max_offset_sent().await, - (TOTAL_REQUESTS - 1) as i64 - ); + assert_eq!(mock_server.get_max_offset_sent().await, (TOTAL_REQUESTS - 1) as i64); Ok(()) } @@ -4704,11 +4624,7 @@ mod callback_tests { // Other two should have error callbacks let errors = callback.get_errors(); - assert!( - errors.len() >= 2, - "Expected at least 2 error callbacks, got {}", - errors.len() - ); + assert!(errors.len() >= 2, "Expected at least 2 error callbacks, got {}", errors.len()); Ok(()) } diff --git a/rust/third_party/arrow-flight/src/client.rs b/rust/third_party/arrow-flight/src/client.rs index b2059a81..907cd459 100644 --- a/rust/third_party/arrow-flight/src/client.rs +++ b/rust/third_party/arrow-flight/src/client.rs @@ -15,27 +15,36 @@ // specific language governing permissions and limitations // under the License. -use crate::{ - Action, ActionType, Criteria, Empty, FlightData, FlightDescriptor, FlightEndpoint, FlightInfo, - HandshakeRequest, PollInfo, PutResult, Ticket, - decode::FlightRecordBatchStream, - flight_service_client::FlightServiceClient, - r#gen::{CancelFlightInfoRequest, CancelFlightInfoResult, RenewFlightEndpointRequest}, - trailers::extract_lazy_trailers, -}; use arrow_schema::Schema; use bytes::Bytes; -use futures::{ - Stream, StreamExt, TryStreamExt, - future::ready, - stream::{self, BoxStream}, -}; +use futures::future::ready; +use futures::stream::{self, BoxStream}; +use futures::{Stream, StreamExt, TryStreamExt}; use prost::Message; use tonic::codegen::{Body, StdError}; -use tonic::{metadata::MetadataMap, transport::Channel}; +use tonic::metadata::MetadataMap; +use tonic::transport::Channel; +use crate::decode::FlightRecordBatchStream; use crate::error::{FlightError, Result}; +use crate::flight_service_client::FlightServiceClient; +use crate::r#gen::{CancelFlightInfoRequest, CancelFlightInfoResult, RenewFlightEndpointRequest}; use crate::streams::{FallibleRequestStream, FallibleTonicResponseStream}; +use crate::trailers::extract_lazy_trailers; +use crate::{ + Action, + ActionType, + Criteria, + Empty, + FlightData, + FlightDescriptor, + FlightEndpoint, + FlightInfo, + HandshakeRequest, + PollInfo, + PutResult, + Ticket, +}; /// A "Mid level" [Apache Arrow Flight](https://arrow.apache.org/docs/format/Flight.html) client. /// @@ -163,9 +172,7 @@ where if let Some(response) = response_stream.next().await.transpose()? { // check if there is another response if response_stream.next().await.is_some() { - return Err(FlightError::protocol( - "Got unexpected second response from handshake", - )); + return Err(FlightError::protocol("Got unexpected second response from handshake")); } Ok(response.payload) @@ -625,9 +632,8 @@ where ) -> Result { let action = Action::new("CancelFlightInfo", request.encode_to_vec()); let response = self.do_action(action).await?.try_next().await?; - let response = response.ok_or(FlightError::protocol( - "Received no response for cancel_flight_info call", - ))?; + let response = response + .ok_or(FlightError::protocol("Received no response for cancel_flight_info call"))?; CancelFlightInfoResult::decode(response) .map_err(|e| FlightError::DecodeError(e.to_string())) } @@ -664,9 +670,8 @@ where ) -> Result { let action = Action::new("RenewFlightEndpoint", request.encode_to_vec()); let response = self.do_action(action).await?.try_next().await?; - let response = response.ok_or(FlightError::protocol( - "Received no response for renew_flight_endpoint call", - ))?; + let response = response + .ok_or(FlightError::protocol("Received no response for renew_flight_endpoint call"))?; FlightEndpoint::decode(response).map_err(|e| FlightError::DecodeError(e.to_string())) } @@ -681,19 +686,14 @@ where #[cfg(test)] mod tests { - use super::FlightClient; - use crate::encode::FlightDataEncoderBuilder; - use crate::flight_service_server::{FlightService, FlightServiceServer}; - use crate::{ - Action, ActionType, Criteria, Empty, FlightData, FlightDescriptor, FlightInfo, - HandshakeRequest, HandshakeResponse, PollInfo, PutResult, SchemaResult, Ticket, - }; - use arrow_array::{RecordBatch, UInt64Array}; - use bytes::Bytes; - use futures::{StreamExt, TryStreamExt, stream::BoxStream}; use std::net::SocketAddr; use std::sync::{Arc, Mutex}; use std::time::Duration; + + use arrow_array::{RecordBatch, UInt64Array}; + use bytes::Bytes; + use futures::stream::BoxStream; + use futures::{StreamExt, TryStreamExt}; use tokio::net::TcpListener; use tokio::task::JoinHandle; use tonic::metadata::MetadataMap; @@ -702,6 +702,25 @@ mod tests { use tonic::{Request, Response, Status, Streaming}; use uuid::Uuid; + use super::FlightClient; + use crate::encode::FlightDataEncoderBuilder; + use crate::flight_service_server::{FlightService, FlightServiceServer}; + use crate::{ + Action, + ActionType, + Criteria, + Empty, + FlightData, + FlightDescriptor, + FlightInfo, + HandshakeRequest, + HandshakeResponse, + PollInfo, + PutResult, + SchemaResult, + Ticket, + }; + /// Minimal `FlightService` that records request metadata and serves a /// configured `do_get` response. Other RPCs return `Unimplemented`. #[derive(Debug, Clone, Default)] diff --git a/rust/third_party/arrow-flight/src/decode.rs b/rust/third_party/arrow-flight/src/decode.rs index 1f649249..d25e7168 100644 --- a/rust/third_party/arrow-flight/src/decode.rs +++ b/rust/third_party/arrow-flight/src/decode.rs @@ -15,17 +15,24 @@ // specific language governing permissions and limitations // under the License. -use crate::{FlightData, trailers::LazyTrailers}; +use std::collections::HashMap; +use std::fmt::Debug; +use std::pin::Pin; +use std::sync::Arc; +use std::task::Poll; + use arrow_array::{ArrayRef, RecordBatch}; use arrow_buffer::Buffer; use arrow_data::UnsafeFlag; use arrow_schema::{Schema, SchemaRef}; use bytes::Bytes; -use futures::{Stream, StreamExt, ready, stream::BoxStream}; -use std::{collections::HashMap, fmt::Debug, pin::Pin, sync::Arc, task::Poll}; +use futures::stream::BoxStream; +use futures::{ready, Stream, StreamExt}; use tonic::metadata::MetadataMap; use crate::error::{FlightError, Result}; +use crate::trailers::LazyTrailers; +use crate::FlightData; /// Decodes a [Stream] of [`FlightData`] back into /// [`RecordBatch`]es. This can be used to decode the response from an @@ -296,9 +303,7 @@ impl FlightDataDecoder { let state = if let Some(state) = self.state.as_mut() { state } else { - return Err(FlightError::protocol( - "Received DictionaryBatch prior to Schema", - )); + return Err(FlightError::protocol("Received DictionaryBatch prior to Schema")); }; let buffer = Buffer::from(data.data_body); @@ -326,9 +331,7 @@ impl FlightDataDecoder { let state = if let Some(state) = self.state.as_ref() { state } else { - return Err(FlightError::protocol( - "Received RecordBatch prior to Schema", - )); + return Err(FlightError::protocol("Received RecordBatch prior to Schema")); }; let record_batch = message.header_as_record_batch().ok_or_else(|| { diff --git a/rust/third_party/arrow-flight/src/encode.rs b/rust/third_party/arrow-flight/src/encode.rs index 82749ec5..517a7e23 100644 --- a/rust/third_party/arrow-flight/src/encode.rs +++ b/rust/third_party/arrow-flight/src/encode.rs @@ -15,16 +15,21 @@ // specific language governing permissions and limitations // under the License. -use std::{collections::VecDeque, fmt::Debug, pin::Pin, sync::Arc, task::Poll}; - -use crate::{FlightData, FlightDescriptor, SchemaAsIpc, error::Result}; +use std::collections::VecDeque; +use std::fmt::Debug; +use std::pin::Pin; +use std::sync::Arc; +use std::task::Poll; use arrow_array::{Array, ArrayRef, RecordBatch, RecordBatchOptions, UnionArray}; use arrow_ipc::writer::{DictionaryTracker, IpcDataGenerator, IpcWriteContext, IpcWriteOptions}; - use arrow_schema::{DataType, Field, FieldRef, Fields, Schema, SchemaRef, UnionMode}; use bytes::Bytes; -use futures::{Stream, StreamExt, ready, stream::BoxStream}; +use futures::stream::BoxStream; +use futures::{ready, Stream, StreamExt}; + +use crate::error::Result; +use crate::{FlightData, FlightDescriptor, SchemaAsIpc}; /// Creates a [`Stream`] of [`FlightData`]s from a /// `Stream` of [`Result`]<[`RecordBatch`], [`FlightError`]>. @@ -512,10 +517,7 @@ fn prepare_field_for_flight( let (type_ids, new_fields): (Vec, Vec) = fields .iter() .map(|(type_id, f)| { - ( - type_id, - prepare_field_for_flight(f, dictionary_tracker, send_dictionaries), - ) + (type_id, prepare_field_for_flight(f, dictionary_tracker, send_dictionaries)) }) .unzip(); @@ -524,11 +526,8 @@ fn prepare_field_for_flight( DataType::Dictionary(_, value_type) => { if !send_dictionaries { // Recurse into value type to handle nested dicts being stripped - let value_field = Field::new( - field.name(), - value_type.as_ref().clone(), - field.is_nullable(), - ); + let value_field = + Field::new(field.name(), value_type.as_ref().clone(), field.is_nullable()); prepare_field_for_flight( &Arc::new(value_field), dictionary_tracker, @@ -572,11 +571,7 @@ fn prepare_field_for_flight( DataType::FixedSizeList(inner, size) => Field::new( field.name(), DataType::FixedSizeList( - Arc::new(prepare_field_for_flight( - inner, - dictionary_tracker, - send_dictionaries, - )), + Arc::new(prepare_field_for_flight(inner, dictionary_tracker, send_dictionaries)), *size, ), field.is_nullable(), @@ -586,11 +581,7 @@ fn prepare_field_for_flight( field.name(), DataType::RunEndEncoded( run_ends.clone(), - Arc::new(prepare_field_for_flight( - values, - dictionary_tracker, - send_dictionaries, - )), + Arc::new(prepare_field_for_flight(values, dictionary_tracker, send_dictionaries)), ), field.is_nullable(), ) @@ -761,9 +752,7 @@ fn hydrate_dictionaries(batch: &RecordBatch, schema: SchemaRef) -> Result Result>>()?, )?) @@ -794,21 +780,27 @@ fn hydrate_dictionary(array: &ArrayRef, data_type: &DataType) -> Result::new(), - ))], + vec![Box::new(builder::ListBuilder::new(StringDictionaryBuilder::::new()))], ); struct_builder @@ -1111,11 +1099,7 @@ mod tests { let arr2 = struct_builder.finish(); - let schema = Arc::new(Schema::new(vec![Field::new_struct( - "struct", - struct_fields, - true, - )])); + let schema = Arc::new(Schema::new(vec![Field::new_struct("struct", struct_fields, true)])); let batch1 = RecordBatch::try_new(schema.clone(), vec![Arc::new(arr1)]).unwrap(); let batch2 = RecordBatch::try_new(schema, vec![Arc::new(arr2)]).unwrap(); @@ -1127,11 +1111,7 @@ mod tests { let mut decoder = FlightDataDecoder::new(encoder); let expected_schema = Schema::new(vec![Field::new_struct( "struct", - vec![Field::new_list( - "dict_list", - Field::new_list_field(DataType::Utf8, true), - true, - )], + vec![Field::new_list("dict_list", Field::new_list_field(DataType::Utf8, true), true)], true, )]); @@ -1173,9 +1153,7 @@ mod tests { let mut struct_builder = StructBuilder::new( struct_fields.clone(), - vec![Box::new(builder::ListBuilder::new( - StringDictionaryBuilder::::new(), - ))], + vec![Box::new(builder::ListBuilder::new(StringDictionaryBuilder::::new()))], ); struct_builder.field_builder::>>>(0) @@ -1192,11 +1170,7 @@ mod tests { let arr2 = struct_builder.finish(); - let schema = Arc::new(Schema::new(vec![Field::new_struct( - "struct", - struct_fields, - true, - )])); + let schema = Arc::new(Schema::new(vec![Field::new_struct("struct", struct_fields, true)])); let batch1 = RecordBatch::try_new(schema.clone(), vec![Arc::new(arr1)]).unwrap(); let batch2 = RecordBatch::try_new(schema, vec![Arc::new(arr2)]).unwrap(); @@ -1221,10 +1195,7 @@ mod tests { true, )), ), - ( - 1, - Arc::new(Field::new_struct("struct", struct_fields.clone(), true)), - ), + (1, Arc::new(Field::new_struct("struct", struct_fields.clone(), true))), (2, Arc::new(Field::new("string", DataType::Utf8, true))), ] .into_iter() @@ -1307,18 +1278,11 @@ mod tests { let mut decoder = FlightDataDecoder::new(encoder); - let hydrated_struct_fields = vec![Field::new_list( - "dict_list", - Field::new_list_field(DataType::Utf8, true), - true, - )]; + let hydrated_struct_fields = + vec![Field::new_list("dict_list", Field::new_list_field(DataType::Utf8, true), true)]; let hydrated_union_fields = vec![ - Field::new_list( - "dict_list", - Field::new_list_field(DataType::Utf8, true), - true, - ), + Field::new_list("dict_list", Field::new_list_field(DataType::Utf8, true), true), Field::new_struct("struct", hydrated_struct_fields.clone(), true), Field::new("string", DataType::Utf8, true), ]; @@ -1390,10 +1354,7 @@ mod tests { true, )), ), - ( - 1, - Arc::new(Field::new_struct("struct", struct_fields.clone(), true)), - ), + (1, Arc::new(Field::new_struct("struct", struct_fields.clone(), true))), (2, Arc::new(Field::new("string", DataType::Utf8, true))), ] .into_iter() @@ -1421,11 +1382,7 @@ mod tests { union_fields.clone(), type_id_buffer, None, - vec![ - Arc::new(arr1), - new_null_array(struct_ty, 1), - new_null_array(string_ty, 1), - ], + vec![Arc::new(arr1), new_null_array(struct_ty, 1), new_null_array(string_ty, 1)], ) .unwrap(); @@ -1439,11 +1396,7 @@ mod tests { union_fields.clone(), type_id_buffer, None, - vec![ - new_null_array(dict_list_ty, 1), - Arc::new(arr2), - new_null_array(string_ty, 1), - ], + vec![new_null_array(dict_list_ty, 1), Arc::new(arr2), new_null_array(string_ty, 1)], ) .unwrap(); @@ -1538,20 +1491,12 @@ mod tests { // array without dictionary fields let arr1 = MapArray::from_vec_of_maps::( - vec![Some(vec![ - ("k1", Some("a")), - ("k2", None), - ("k3", Some("b")), - ])], + vec![Some(vec![("k1", Some("a")), ("k2", None), ("k3", Some("b"))])], false, ); let arr2 = MapArray::from_vec_of_maps::( - vec![Some(vec![ - ("k1", Some("c")), - ("k2", None), - ("k3", Some("d")), - ])], + vec![Some(vec![("k1", Some("c")), ("k2", None), ("k3", Some("d"))])], false, ); @@ -1633,11 +1578,7 @@ mod tests { let run_ends2 = Int32Array::from(vec![1, 2]); let arr2 = RunArray::try_new(&run_ends2, &dict_values2).unwrap(); - let schema = Arc::new(Schema::new(vec![Field::new( - "ree", - arr1.data_type().clone(), - true, - )])); + let schema = Arc::new(Schema::new(vec![Field::new("ree", arr1.data_type().clone(), true)])); let batch1 = RecordBatch::try_new(schema.clone(), vec![Arc::new(arr1)]).unwrap(); let batch2 = RecordBatch::try_new(schema, vec![Arc::new(arr2)]).unwrap(); @@ -1662,14 +1603,8 @@ mod tests { let int_array = Int32Array::from(vec![10, 20, 30, 40, 50]); let struct_array = StructArray::from(vec![ - ( - Arc::new(struct_fields[0].clone()), - Arc::new(inner_dict) as ArrayRef, - ), - ( - Arc::new(struct_fields[1].clone()), - Arc::new(int_array) as ArrayRef, - ), + (Arc::new(struct_fields[0].clone()), Arc::new(inner_dict) as ArrayRef), + (Arc::new(struct_fields[1].clone()), Arc::new(int_array) as ArrayRef), ]); let outer_keys = Int8Array::from_iter_values([0, 0, 1, 2]); @@ -1681,24 +1616,15 @@ mod tests { let int_array2 = Int32Array::from(vec![100, 200, 300]); let struct_array2 = StructArray::from(vec![ - ( - Arc::new(struct_fields[0].clone()), - Arc::new(inner_dict2) as ArrayRef, - ), - ( - Arc::new(struct_fields[1].clone()), - Arc::new(int_array2) as ArrayRef, - ), + (Arc::new(struct_fields[0].clone()), Arc::new(inner_dict2) as ArrayRef), + (Arc::new(struct_fields[1].clone()), Arc::new(int_array2) as ArrayRef), ]); let outer_keys2 = Int8Array::from_iter_values([0, 1]); let arr2 = DictionaryArray::new(outer_keys2, Arc::new(struct_array2)); - let schema = Arc::new(Schema::new(vec![Field::new( - "dict_struct", - arr1.data_type().clone(), - false, - )])); + let schema = + Arc::new(Schema::new(vec![Field::new("dict_struct", arr1.data_type().clone(), false)])); let batch1 = RecordBatch::try_new(schema.clone(), vec![Arc::new(arr1)]).unwrap(); let batch2 = RecordBatch::try_new(schema, vec![Arc::new(arr2)]).unwrap(); @@ -1716,12 +1642,7 @@ mod tests { builder.append_value(vec![Some("c"), None, Some("d")]); let arr2 = builder.finish(); - let inner = Arc::new(Field::new_dictionary( - "item", - DataType::UInt16, - DataType::Utf8, - true, - )); + let inner = Arc::new(Field::new_dictionary("item", DataType::UInt16, DataType::Utf8, true)); let dt = if O::IS_LARGE { DataType::LargeListView(inner) } else { @@ -1808,9 +1729,8 @@ mod tests { #[test] fn test_schema_metadata_encoded() { - let schema = Schema::new(vec![Field::new("data", DataType::Int32, false)]).with_metadata( - HashMap::from([("some_key".to_owned(), "some_value".to_owned())]), - ); + let schema = Schema::new(vec![Field::new("data", DataType::Int32, false)]) + .with_metadata(HashMap::from([("some_key".to_owned(), "some_value".to_owned())])); let mut dictionary_tracker = DictionaryTracker::new(false); @@ -1846,12 +1766,7 @@ mod tests { let mut ipc_write_context = IpcWriteContext::default(); let (encoded_dictionaries, encoded_batch) = data_gen - .encode( - batch, - &mut dictionary_tracker, - options, - &mut ipc_write_context, - ) + .encode(batch, &mut dictionary_tracker, options, &mut ipc_write_context) .expect("DictionaryTracker configured above to not error on replacement"); let flight_dictionaries = encoded_dictionaries.into_iter().map(Into::into).collect(); @@ -1882,10 +1797,7 @@ mod tests { let split: Vec<_> = split_batch_for_grpc_response(batch.clone(), max_flight_data_size).collect(); assert_eq!(split.len(), 3); - assert_eq!( - split.iter().map(|batch| batch.num_rows()).sum::(), - n_rows - ); + assert_eq!(split.iter().map(|batch| batch.num_rows()).sum::(), n_rows); let a = pretty_format_batches(&split).unwrap().to_string(); let b = pretty_format_batches(&[batch]).unwrap().to_string(); assert_eq!(a, b); @@ -2141,9 +2053,6 @@ mod tests { // ensure that the specified overage is exactly the maxmium so // that when the splitting logic improves, the tests must be // updated to reflect the better logic - assert_eq!( - allowed_overage, max_overage_seen, - "Specified overage was too high" - ); + assert_eq!(allowed_overage, max_overage_seen, "Specified overage was too high"); } } diff --git a/rust/third_party/arrow-flight/src/lib.rs b/rust/third_party/arrow-flight/src/lib.rs index 94b44942..eef7b6ba 100644 --- a/rust/third_party/arrow-flight/src/lib.rs +++ b/rust/third_party/arrow-flight/src/lib.rs @@ -47,15 +47,17 @@ // The unused_crate_dependencies lint does not work well for crates defining additional examples/bin targets #![allow(unused_crate_dependencies)] -use arrow_ipc::{convert, writer, writer::EncodedData, writer::IpcWriteOptions}; -use arrow_schema::{ArrowError, Schema}; +use std::fmt; +use std::ops::Deref; use arrow_ipc::convert::try_schema_from_ipc_buffer; -use base64::Engine; +use arrow_ipc::writer::{EncodedData, IpcWriteOptions}; +use arrow_ipc::{convert, writer}; +use arrow_schema::{ArrowError, Schema}; use base64::prelude::BASE64_STANDARD; +use base64::Engine; use bytes::Bytes; use prost_types::Timestamp; -use std::{fmt, ops::Deref}; type ArrowResult = std::result::Result; @@ -68,22 +70,24 @@ mod r#gen { /// Defines a `Flight` for generation or retrieval. pub mod flight_descriptor { - use super::r#gen; pub use r#gen::flight_descriptor::DescriptorType; + + use super::r#gen; } /// Low Level [tonic] [`FlightServiceClient`](gen::flight_service_client::FlightServiceClient). pub mod flight_service_client { - use super::r#gen; pub use r#gen::flight_service_client::FlightServiceClient; + + use super::r#gen; } /// Low Level [tonic] [`FlightServiceServer`](gen::flight_service_server::FlightServiceServer) /// and [`FlightService`](gen::flight_service_server::FlightService). pub mod flight_service_server { + pub use r#gen::flight_service_server::{FlightService, FlightServiceServer}; + use super::r#gen; - pub use r#gen::flight_service_server::FlightService; - pub use r#gen::flight_service_server::FlightServiceServer; } /// Mid Level [`FlightClient`] @@ -101,27 +105,29 @@ pub mod encode; /// Common error types pub mod error; -pub use r#gen::Action; -pub use r#gen::ActionType; -pub use r#gen::BasicAuth; -pub use r#gen::CancelFlightInfoRequest; -pub use r#gen::CancelFlightInfoResult; -pub use r#gen::CancelStatus; -pub use r#gen::Criteria; -pub use r#gen::Empty; -pub use r#gen::FlightData; -pub use r#gen::FlightDescriptor; -pub use r#gen::FlightEndpoint; -pub use r#gen::FlightInfo; -pub use r#gen::HandshakeRequest; -pub use r#gen::HandshakeResponse; -pub use r#gen::Location; -pub use r#gen::PollInfo; -pub use r#gen::PutResult; -pub use r#gen::RenewFlightEndpointRequest; -pub use r#gen::Result; -pub use r#gen::SchemaResult; -pub use r#gen::Ticket; +pub use r#gen::{ + Action, + ActionType, + BasicAuth, + CancelFlightInfoRequest, + CancelFlightInfoResult, + CancelStatus, + Criteria, + Empty, + FlightData, + FlightDescriptor, + FlightEndpoint, + FlightInfo, + HandshakeRequest, + HandshakeResponse, + Location, + PollInfo, + PutResult, + RenewFlightEndpointRequest, + Result, + SchemaResult, + Ticket, +}; /// Helper to extract HTTP/gRPC trailers from a tonic stream. mod trailers; @@ -418,9 +424,7 @@ impl TryFrom<&FlightData> for Schema { type Error = ArrowError; fn try_from(data: &FlightData) -> ArrowResult { convert::try_schema_from_flatbuffer_bytes(&data.data_header[..]).map_err(|err| { - ArrowError::ParseError(format!( - "Unable to convert flight data to Arrow schema: {err}" - )) + ArrowError::ParseError(format!("Unable to convert flight data to Arrow schema: {err}")) }) } } @@ -829,10 +833,11 @@ impl FlightEndpoint { #[cfg(test)] mod tests { - use super::*; use arrow_ipc::MetadataVersion; use arrow_schema::{DataType, Field, TimeUnit}; + use super::*; + struct TestVector(Vec, usize); impl fmt::Display for TestVector { diff --git a/rust/third_party/arrow-flight/src/sql/client.rs b/rust/third_party/arrow-flight/src/sql/client.rs index 5476d4ed..5b9d96dc 100644 --- a/rust/third_party/arrow-flight/src/sql/client.rs +++ b/rust/third_party/arrow-flight/src/sql/client.rs @@ -17,50 +17,76 @@ //! A FlightSQL Client [`FlightSqlServiceClient`] +use std::collections::HashMap; +use std::str::FromStr; + +use arrow_array::RecordBatch; use arrow_buffer::Buffer; -use arrow_ipc::MessageHeader; use arrow_ipc::convert::fb_to_schema; use arrow_ipc::reader::read_record_batch; -use arrow_ipc::root_as_message; -use arrow_schema::SchemaRef; -use base64::Engine; +use arrow_ipc::{root_as_message, MessageHeader}; +use arrow_schema::{ArrowError, Schema, SchemaRef}; use base64::prelude::BASE64_STANDARD; +use base64::Engine; use bytes::Bytes; -use std::collections::HashMap; -use std::str::FromStr; +use futures::{stream, Stream, TryStreamExt}; +use prost::Message; +use tonic::codegen::{Body, StdError}; use tonic::metadata::AsciiMetadataKey; +use tonic::{IntoRequest, IntoStreamingRequest, Streaming}; use crate::decode::FlightRecordBatchStream; use crate::encode::FlightDataEncoderBuilder; -use crate::error::FlightError; -use crate::error::Result; +use crate::error::{FlightError, Result}; use crate::flight_service_client::FlightServiceClient; use crate::sql::r#gen::action_end_transaction_request::EndTransaction; use crate::sql::server::{ - BEGIN_TRANSACTION, CLOSE_PREPARED_STATEMENT, CREATE_PREPARED_STATEMENT, END_TRANSACTION, + BEGIN_TRANSACTION, + CLOSE_PREPARED_STATEMENT, + CREATE_PREPARED_STATEMENT, + END_TRANSACTION, }; use crate::sql::{ - ActionBeginTransactionRequest, ActionBeginTransactionResult, - ActionClosePreparedStatementRequest, ActionCreatePreparedStatementRequest, - ActionCreatePreparedStatementResult, ActionEndTransactionRequest, Any, CommandGetCatalogs, - CommandGetCrossReference, CommandGetDbSchemas, CommandGetExportedKeys, CommandGetImportedKeys, - CommandGetPrimaryKeys, CommandGetSqlInfo, CommandGetTableTypes, CommandGetTables, - CommandGetXdbcTypeInfo, CommandPreparedStatementQuery, CommandPreparedStatementUpdate, - CommandStatementIngest, CommandStatementQuery, CommandStatementUpdate, - DoPutPreparedStatementResult, DoPutUpdateResult, ProstMessageExt, SqlInfo, + ActionBeginTransactionRequest, + ActionBeginTransactionResult, + ActionClosePreparedStatementRequest, + ActionCreatePreparedStatementRequest, + ActionCreatePreparedStatementResult, + ActionEndTransactionRequest, + Any, + CommandGetCatalogs, + CommandGetCrossReference, + CommandGetDbSchemas, + CommandGetExportedKeys, + CommandGetImportedKeys, + CommandGetPrimaryKeys, + CommandGetSqlInfo, + CommandGetTableTypes, + CommandGetTables, + CommandGetXdbcTypeInfo, + CommandPreparedStatementQuery, + CommandPreparedStatementUpdate, + CommandStatementIngest, + CommandStatementQuery, + CommandStatementUpdate, + DoPutPreparedStatementResult, + DoPutUpdateResult, + ProstMessageExt, + SqlInfo, }; use crate::streams::FallibleRequestStream; use crate::trailers::extract_lazy_trailers; use crate::{ - Action, FlightData, FlightDescriptor, FlightInfo, HandshakeRequest, HandshakeResponse, - IpcMessage, PutResult, Ticket, + Action, + FlightData, + FlightDescriptor, + FlightInfo, + HandshakeRequest, + HandshakeResponse, + IpcMessage, + PutResult, + Ticket, }; -use arrow_array::RecordBatch; -use arrow_schema::{ArrowError, Schema}; -use futures::{Stream, TryStreamExt, stream}; -use prost::Message; -use tonic::codegen::{Body, StdError}; -use tonic::{IntoRequest, IntoStreamingRequest, Streaming}; /// A FlightSQLServiceClient is an endpoint for retrieving or storing Arrow data /// by FlightSQL protocol. @@ -196,9 +222,7 @@ where let resp = match responses.as_slice() { [resp] => resp.payload.clone(), [] => Bytes::new(), - _ => Err(ArrowError::ParseError( - "Multiple handshake responses".to_string(), - ))?, + _ => Err(ArrowError::ParseError("Multiple handshake responses".to_string()))?, }; Ok(resp) } @@ -572,9 +596,7 @@ where .with_flight_descriptor(Some(descriptor)) .with_schema(params_batch.schema()); let flight_data = flight_stream_builder - .build(futures::stream::iter( - self.parameter_binding.clone().map(Ok), - )) + .build(futures::stream::iter(self.parameter_binding.clone().map(Ok))) .try_collect::>() .await?; diff --git a/rust/third_party/arrow-flight/src/sql/metadata/catalogs.rs b/rust/third_party/arrow-flight/src/sql/metadata/catalogs.rs index e27c63c3..39d86b78 100644 --- a/rust/third_party/arrow-flight/src/sql/metadata/catalogs.rs +++ b/rust/third_party/arrow-flight/src/sql/metadata/catalogs.rs @@ -92,13 +92,8 @@ fn get_catalogs_schema() -> SchemaRef { } /// The schema for GetCatalogs -static GET_CATALOG_SCHEMA: Lazy = Lazy::new(|| { - Arc::new(Schema::new(vec![Field::new( - "catalog_name", - DataType::Utf8, - false, - )])) -}); +static GET_CATALOG_SCHEMA: Lazy = + Lazy::new(|| Arc::new(Schema::new(vec![Field::new("catalog_name", DataType::Utf8, false)]))); #[cfg(test)] mod tests { diff --git a/rust/third_party/arrow-flight/src/sql/metadata/db_schemas.rs b/rust/third_party/arrow-flight/src/sql/metadata/db_schemas.rs index c182140e..f844a69f 100644 --- a/rust/third_party/arrow-flight/src/sql/metadata/db_schemas.rs +++ b/rust/third_party/arrow-flight/src/sql/metadata/db_schemas.rs @@ -22,10 +22,12 @@ use std::sync::Arc; use arrow_arith::boolean::and; -use arrow_array::{ArrayRef, RecordBatch, StringArray, builder::StringBuilder}; +use arrow_array::builder::StringBuilder; +use arrow_array::{ArrayRef, RecordBatch, StringArray}; use arrow_ord::cmp::eq; use arrow_schema::{DataType, Field, Schema, SchemaRef}; -use arrow_select::{filter::filter_record_batch, take::take}; +use arrow_select::filter::filter_record_batch; +use arrow_select::take::take; use arrow_string::like::like; use once_cell::sync::Lazy; @@ -139,10 +141,7 @@ impl GetDbSchemasBuilder { let batch = RecordBatch::try_new( schema, - vec![ - Arc::new(catalog_name) as ArrayRef, - Arc::new(db_schema_name) as ArrayRef, - ], + vec![Arc::new(catalog_name) as ArrayRef, Arc::new(db_schema_name) as ArrayRef], )?; // Apply the filters if needed @@ -184,9 +183,10 @@ static GET_DB_SCHEMAS_SCHEMA: Lazy = Lazy::new(|| { #[cfg(test)] mod tests { - use super::*; use arrow_array::{StringArray, UInt32Array}; + use super::*; + fn get_ref_batch() -> RecordBatch { RecordBatch::try_new( get_db_schemas_schema(), @@ -197,9 +197,8 @@ mod tests { "b_catalog", "b_catalog", ])) as ArrayRef, - Arc::new(StringArray::from(vec![ - "a_schema", "b_schema", "a_schema", "b_schema", - ])) as ArrayRef, + Arc::new(StringArray::from(vec!["a_schema", "b_schema", "a_schema", "b_schema"])) + as ArrayRef, ], ) .unwrap() diff --git a/rust/third_party/arrow-flight/src/sql/metadata/mod.rs b/rust/third_party/arrow-flight/src/sql/metadata/mod.rs index b914cd38..54c4e3a2 100644 --- a/rust/third_party/arrow-flight/src/sql/metadata/mod.rs +++ b/rust/third_party/arrow-flight/src/sql/metadata/mod.rs @@ -37,17 +37,14 @@ mod table_types; mod tables; mod xdbc_info; +use arrow_array::{ArrayRef, UInt32Array}; +use arrow_row::{RowConverter, SortField}; pub use catalogs::GetCatalogsBuilder; pub use db_schemas::GetDbSchemasBuilder; pub use sql_info::{SqlInfoData, SqlInfoDataBuilder}; pub use tables::GetTablesBuilder; pub use xdbc_info::{XdbcTypeInfo, XdbcTypeInfoData, XdbcTypeInfoDataBuilder}; -use arrow_array::ArrayRef; -use arrow_array::UInt32Array; -use arrow_row::RowConverter; -use arrow_row::SortField; - /// Helper function to sort all the columns in an array fn lexsort_to_indices(arrays: &[ArrayRef]) -> UInt32Array { let fields = arrays diff --git a/rust/third_party/arrow-flight/src/sql/metadata/sql_info.rs b/rust/third_party/arrow-flight/src/sql/metadata/sql_info.rs index 155946ea..d41db29b 100644 --- a/rust/third_party/arrow-flight/src/sql/metadata/sql_info.rs +++ b/rust/third_party/arrow-flight/src/sql/metadata/sql_info.rs @@ -30,8 +30,15 @@ use std::sync::Arc; use arrow_arith::boolean::or; use arrow_array::array::{Array, UInt32Array, UnionArray}; use arrow_array::builder::{ - ArrayBuilder, BooleanBuilder, Int8Builder, Int32Builder, Int64Builder, ListBuilder, MapBuilder, - StringBuilder, UInt32Builder, + ArrayBuilder, + BooleanBuilder, + Int32Builder, + Int64Builder, + Int8Builder, + ListBuilder, + MapBuilder, + StringBuilder, + UInt32Builder, }; use arrow_array::{RecordBatch, Scalar}; use arrow_data::ArrayData; @@ -495,10 +502,7 @@ mod tests { // bool builder.append(SqlInfo::SqlDdlCatalog, false); // i32 - builder.append( - SqlInfo::SqlNullOrdering, - SqlNullOrdering::SqlNullsSortedHigh as i32, - ); + builder.append(SqlInfo::SqlNullOrdering, SqlNullOrdering::SqlNullsSortedHigh as i32); // i64 builder.append(SqlInfo::SqlMaxBinaryLiteralLength, i32::MAX as i64); // [str] diff --git a/rust/third_party/arrow-flight/src/sql/metadata/table_types.rs b/rust/third_party/arrow-flight/src/sql/metadata/table_types.rs index 7f525da0..a9282b9b 100644 --- a/rust/third_party/arrow-flight/src/sql/metadata/table_types.rs +++ b/rust/third_party/arrow-flight/src/sql/metadata/table_types.rs @@ -21,16 +21,16 @@ use std::sync::Arc; -use arrow_array::{ArrayRef, RecordBatch, builder::StringBuilder}; +use arrow_array::builder::StringBuilder; +use arrow_array::{ArrayRef, RecordBatch}; use arrow_schema::{DataType, Field, Schema, SchemaRef}; use arrow_select::take::take; use once_cell::sync::Lazy; +use super::lexsort_to_indices; use crate::error::*; use crate::sql::CommandGetTableTypes; -use super::lexsort_to_indices; - /// A builder for a [`CommandGetTableTypes`] response. /// /// Builds rows like this: @@ -101,19 +101,15 @@ fn get_table_types_schema() -> SchemaRef { } /// The schema for [`CommandGetTableTypes`]. -static GET_TABLE_TYPES_SCHEMA: Lazy = Lazy::new(|| { - Arc::new(Schema::new(vec![Field::new( - "table_type", - DataType::Utf8, - false, - )])) -}); +static GET_TABLE_TYPES_SCHEMA: Lazy = + Lazy::new(|| Arc::new(Schema::new(vec![Field::new("table_type", DataType::Utf8, false)]))); #[cfg(test)] mod tests { - use super::*; use arrow_array::StringArray; + use super::*; + fn get_ref_batch() -> RecordBatch { RecordBatch::try_new( get_table_types_schema(), diff --git a/rust/third_party/arrow-flight/src/sql/metadata/tables.rs b/rust/third_party/arrow-flight/src/sql/metadata/tables.rs index 2cd16fdc..e224bd54 100644 --- a/rust/third_party/arrow-flight/src/sql/metadata/tables.rs +++ b/rust/third_party/arrow-flight/src/sql/metadata/tables.rs @@ -26,7 +26,8 @@ use arrow_array::builder::{BinaryBuilder, StringBuilder}; use arrow_array::{ArrayRef, RecordBatch, StringArray}; use arrow_ord::cmp::eq; use arrow_schema::{DataType, Field, Schema, SchemaRef}; -use arrow_select::{filter::filter_record_batch, take::take}; +use arrow_select::filter::filter_record_batch; +use arrow_select::take::take; use arrow_string::like::like; use once_cell::sync::Lazy; @@ -311,9 +312,10 @@ static GET_TABLES_SCHEMA_WITH_TABLE_SCHEMA: Lazy = Lazy::new(|| { #[cfg(test)] mod tests { - use super::*; use arrow_array::{StringArray, UInt32Array}; + use super::*; + fn get_ref_batch() -> RecordBatch { RecordBatch::try_new( get_tables_schema(false), @@ -371,13 +373,7 @@ mod tests { ); for (catalog_name, schema_name, table_name, table_type) in tables { builder - .append( - catalog_name, - schema_name, - table_name, - table_type, - &dummy_schema, - ) + .append(catalog_name, schema_name, table_name, table_type, &dummy_schema) .unwrap(); } builder @@ -461,13 +457,7 @@ mod tests { ); for (catalog_name, schema_name, table_name, table_type) in tables { builder - .append( - catalog_name, - schema_name, - table_name, - table_type, - &dummy_schema, - ) + .append(catalog_name, schema_name, table_name, table_type, &dummy_schema) .unwrap(); } let table_batch = builder.build().unwrap(); diff --git a/rust/third_party/arrow-flight/src/sql/mod.rs b/rust/third_party/arrow-flight/src/sql/mod.rs index e076f7aa..7c58e137 100644 --- a/rust/third_party/arrow-flight/src/sql/mod.rs +++ b/rust/third_party/arrow-flight/src/sql/mod.rs @@ -50,63 +50,66 @@ mod r#gen { include!("arrow.flight.protocol.sql.rs"); } -pub use r#gen::ActionBeginSavepointRequest; -pub use r#gen::ActionBeginSavepointResult; -pub use r#gen::ActionBeginTransactionRequest; -pub use r#gen::ActionBeginTransactionResult; -pub use r#gen::ActionCancelQueryRequest; -pub use r#gen::ActionCancelQueryResult; -pub use r#gen::ActionClosePreparedStatementRequest; -pub use r#gen::ActionCreatePreparedStatementRequest; -pub use r#gen::ActionCreatePreparedStatementResult; -pub use r#gen::ActionCreatePreparedSubstraitPlanRequest; -pub use r#gen::ActionEndSavepointRequest; -pub use r#gen::ActionEndTransactionRequest; -pub use r#gen::CommandGetCatalogs; -pub use r#gen::CommandGetCrossReference; -pub use r#gen::CommandGetDbSchemas; -pub use r#gen::CommandGetExportedKeys; -pub use r#gen::CommandGetImportedKeys; -pub use r#gen::CommandGetPrimaryKeys; -pub use r#gen::CommandGetSqlInfo; -pub use r#gen::CommandGetTableTypes; -pub use r#gen::CommandGetTables; -pub use r#gen::CommandGetXdbcTypeInfo; -pub use r#gen::CommandPreparedStatementQuery; -pub use r#gen::CommandPreparedStatementUpdate; -pub use r#gen::CommandStatementIngest; -pub use r#gen::CommandStatementQuery; -pub use r#gen::CommandStatementSubstraitPlan; -pub use r#gen::CommandStatementUpdate; -pub use r#gen::DoPutPreparedStatementResult; -pub use r#gen::DoPutUpdateResult; -pub use r#gen::Nullable; -pub use r#gen::Searchable; -pub use r#gen::SqlInfo; -pub use r#gen::SqlNullOrdering; -pub use r#gen::SqlOuterJoinsSupportLevel; -pub use r#gen::SqlSupportedCaseSensitivity; -pub use r#gen::SqlSupportedElementActions; -pub use r#gen::SqlSupportedGroupBy; -pub use r#gen::SqlSupportedPositionedCommands; -pub use r#gen::SqlSupportedResultSetConcurrency; -pub use r#gen::SqlSupportedResultSetType; -pub use r#gen::SqlSupportedSubqueries; -pub use r#gen::SqlSupportedTransaction; -pub use r#gen::SqlSupportedTransactions; -pub use r#gen::SqlSupportedUnions; -pub use r#gen::SqlSupportsConvert; -pub use r#gen::SqlTransactionIsolationLevel; -pub use r#gen::SubstraitPlan; -pub use r#gen::SupportedSqlGrammar; -pub use r#gen::TicketStatementQuery; -pub use r#gen::UpdateDeleteRules; -pub use r#gen::XdbcDataType; -pub use r#gen::XdbcDatetimeSubcode; pub use r#gen::action_end_transaction_request::EndTransaction; -pub use r#gen::command_statement_ingest::TableDefinitionOptions; pub use r#gen::command_statement_ingest::table_definition_options::{ - TableExistsOption, TableNotExistOption, + TableExistsOption, + TableNotExistOption, +}; +pub use r#gen::command_statement_ingest::TableDefinitionOptions; +pub use r#gen::{ + ActionBeginSavepointRequest, + ActionBeginSavepointResult, + ActionBeginTransactionRequest, + ActionBeginTransactionResult, + ActionCancelQueryRequest, + ActionCancelQueryResult, + ActionClosePreparedStatementRequest, + ActionCreatePreparedStatementRequest, + ActionCreatePreparedStatementResult, + ActionCreatePreparedSubstraitPlanRequest, + ActionEndSavepointRequest, + ActionEndTransactionRequest, + CommandGetCatalogs, + CommandGetCrossReference, + CommandGetDbSchemas, + CommandGetExportedKeys, + CommandGetImportedKeys, + CommandGetPrimaryKeys, + CommandGetSqlInfo, + CommandGetTableTypes, + CommandGetTables, + CommandGetXdbcTypeInfo, + CommandPreparedStatementQuery, + CommandPreparedStatementUpdate, + CommandStatementIngest, + CommandStatementQuery, + CommandStatementSubstraitPlan, + CommandStatementUpdate, + DoPutPreparedStatementResult, + DoPutUpdateResult, + Nullable, + Searchable, + SqlInfo, + SqlNullOrdering, + SqlOuterJoinsSupportLevel, + SqlSupportedCaseSensitivity, + SqlSupportedElementActions, + SqlSupportedGroupBy, + SqlSupportedPositionedCommands, + SqlSupportedResultSetConcurrency, + SqlSupportedResultSetType, + SqlSupportedSubqueries, + SqlSupportedTransaction, + SqlSupportedTransactions, + SqlSupportedUnions, + SqlSupportsConvert, + SqlTransactionIsolationLevel, + SubstraitPlan, + SupportedSqlGrammar, + TicketStatementQuery, + UpdateDeleteRules, + XdbcDataType, + XdbcDatetimeSubcode, }; pub mod client; diff --git a/rust/third_party/arrow-flight/src/sql/server.rs b/rust/third_party/arrow-flight/src/sql/server.rs index 871a67b7..1b9e0642 100644 --- a/rust/third_party/arrow-flight/src/sql/server.rs +++ b/rust/third_party/arrow-flight/src/sql/server.rs @@ -20,27 +20,64 @@ use std::fmt::{Display, Formatter}; use std::pin::Pin; +use futures::stream::Peekable; +use futures::{Stream, StreamExt}; +use prost::Message; +use tonic::{Request, Response, Status, Streaming}; + use super::{ - ActionBeginSavepointRequest, ActionBeginSavepointResult, ActionBeginTransactionRequest, - ActionBeginTransactionResult, ActionCancelQueryRequest, ActionCancelQueryResult, - ActionClosePreparedStatementRequest, ActionCreatePreparedStatementRequest, - ActionCreatePreparedStatementResult, ActionCreatePreparedSubstraitPlanRequest, - ActionEndSavepointRequest, ActionEndTransactionRequest, Any, Command, CommandGetCatalogs, - CommandGetCrossReference, CommandGetDbSchemas, CommandGetExportedKeys, CommandGetImportedKeys, - CommandGetPrimaryKeys, CommandGetSqlInfo, CommandGetTableTypes, CommandGetTables, - CommandGetXdbcTypeInfo, CommandPreparedStatementQuery, CommandPreparedStatementUpdate, - CommandStatementIngest, CommandStatementQuery, CommandStatementSubstraitPlan, - CommandStatementUpdate, DoPutPreparedStatementResult, DoPutUpdateResult, ProstMessageExt, - SqlInfo, TicketStatementQuery, + ActionBeginSavepointRequest, + ActionBeginSavepointResult, + ActionBeginTransactionRequest, + ActionBeginTransactionResult, + ActionCancelQueryRequest, + ActionCancelQueryResult, + ActionClosePreparedStatementRequest, + ActionCreatePreparedStatementRequest, + ActionCreatePreparedStatementResult, + ActionCreatePreparedSubstraitPlanRequest, + ActionEndSavepointRequest, + ActionEndTransactionRequest, + Any, + Command, + CommandGetCatalogs, + CommandGetCrossReference, + CommandGetDbSchemas, + CommandGetExportedKeys, + CommandGetImportedKeys, + CommandGetPrimaryKeys, + CommandGetSqlInfo, + CommandGetTableTypes, + CommandGetTables, + CommandGetXdbcTypeInfo, + CommandPreparedStatementQuery, + CommandPreparedStatementUpdate, + CommandStatementIngest, + CommandStatementQuery, + CommandStatementSubstraitPlan, + CommandStatementUpdate, + DoPutPreparedStatementResult, + DoPutUpdateResult, + ProstMessageExt, + SqlInfo, + TicketStatementQuery, }; +use crate::flight_service_server::FlightService; +use crate::r#gen::PollInfo; use crate::{ - Action, ActionType, Criteria, Empty, FlightData, FlightDescriptor, FlightInfo, - HandshakeRequest, HandshakeResponse, PutResult, SchemaResult, Ticket, - flight_service_server::FlightService, r#gen::PollInfo, + Action, + ActionType, + Criteria, + Empty, + FlightData, + FlightDescriptor, + FlightInfo, + HandshakeRequest, + HandshakeResponse, + PutResult, + SchemaResult, + Ticket, }; -use futures::{Stream, StreamExt, stream::Peekable}; -use prost::Message; -use tonic::{Request, Response, Status, Streaming}; pub(crate) static CREATE_PREPARED_STATEMENT: &str = "CreatePreparedStatement"; pub(crate) static CLOSE_PREPARED_STATEMENT: &str = "ClosePreparedStatement"; @@ -66,9 +103,7 @@ pub trait FlightSqlService: Sync + Send + Sized + 'static { Response> + Send>>>, Status, > { - Err(Status::unimplemented( - "Handshake has no default implementation", - )) + Err(Status::unimplemented("Handshake has no default implementation")) } /// Implementors may override to handle additional calls to do_get() @@ -89,9 +124,7 @@ pub trait FlightSqlService: Sync + Send + Sized + 'static { _query: CommandStatementQuery, _request: Request, ) -> Result, Status> { - Err(Status::unimplemented( - "get_flight_info_statement has no default implementation", - )) + Err(Status::unimplemented("get_flight_info_statement has no default implementation")) } /// Get a FlightInfo for executing a substrait plan. @@ -100,9 +133,7 @@ pub trait FlightSqlService: Sync + Send + Sized + 'static { _query: CommandStatementSubstraitPlan, _request: Request, ) -> Result, Status> { - Err(Status::unimplemented( - "get_flight_info_substrait_plan has no default implementation", - )) + Err(Status::unimplemented("get_flight_info_substrait_plan has no default implementation")) } /// Get a FlightInfo for executing an already created prepared statement. @@ -122,9 +153,7 @@ pub trait FlightSqlService: Sync + Send + Sized + 'static { _query: CommandGetCatalogs, _request: Request, ) -> Result, Status> { - Err(Status::unimplemented( - "get_flight_info_catalogs has no default implementation", - )) + Err(Status::unimplemented("get_flight_info_catalogs has no default implementation")) } /// Get a FlightInfo for listing schemas. @@ -133,9 +162,7 @@ pub trait FlightSqlService: Sync + Send + Sized + 'static { _query: CommandGetDbSchemas, _request: Request, ) -> Result, Status> { - Err(Status::unimplemented( - "get_flight_info_schemas has no default implementation", - )) + Err(Status::unimplemented("get_flight_info_schemas has no default implementation")) } /// Get a FlightInfo for listing tables. @@ -144,9 +171,7 @@ pub trait FlightSqlService: Sync + Send + Sized + 'static { _query: CommandGetTables, _request: Request, ) -> Result, Status> { - Err(Status::unimplemented( - "get_flight_info_tables has no default implementation", - )) + Err(Status::unimplemented("get_flight_info_tables has no default implementation")) } /// Get a FlightInfo to extract information about the table types. @@ -155,9 +180,7 @@ pub trait FlightSqlService: Sync + Send + Sized + 'static { _query: CommandGetTableTypes, _request: Request, ) -> Result, Status> { - Err(Status::unimplemented( - "get_flight_info_table_types has no default implementation", - )) + Err(Status::unimplemented("get_flight_info_table_types has no default implementation")) } /// Get a FlightInfo for retrieving other information (See SqlInfo). @@ -166,9 +189,7 @@ pub trait FlightSqlService: Sync + Send + Sized + 'static { _query: CommandGetSqlInfo, _request: Request, ) -> Result, Status> { - Err(Status::unimplemented( - "get_flight_info_sql_info has no default implementation", - )) + Err(Status::unimplemented("get_flight_info_sql_info has no default implementation")) } /// Get a FlightInfo to extract information about primary and foreign keys. @@ -177,9 +198,7 @@ pub trait FlightSqlService: Sync + Send + Sized + 'static { _query: CommandGetPrimaryKeys, _request: Request, ) -> Result, Status> { - Err(Status::unimplemented( - "get_flight_info_primary_keys has no default implementation", - )) + Err(Status::unimplemented("get_flight_info_primary_keys has no default implementation")) } /// Get a FlightInfo to extract information about exported keys. @@ -188,9 +207,7 @@ pub trait FlightSqlService: Sync + Send + Sized + 'static { _query: CommandGetExportedKeys, _request: Request, ) -> Result, Status> { - Err(Status::unimplemented( - "get_flight_info_exported_keys has no default implementation", - )) + Err(Status::unimplemented("get_flight_info_exported_keys has no default implementation")) } /// Get a FlightInfo to extract information about imported keys. @@ -199,9 +216,7 @@ pub trait FlightSqlService: Sync + Send + Sized + 'static { _query: CommandGetImportedKeys, _request: Request, ) -> Result, Status> { - Err(Status::unimplemented( - "get_flight_info_imported_keys has no default implementation", - )) + Err(Status::unimplemented("get_flight_info_imported_keys has no default implementation")) } /// Get a FlightInfo to extract information about cross reference. @@ -210,9 +225,7 @@ pub trait FlightSqlService: Sync + Send + Sized + 'static { _query: CommandGetCrossReference, _request: Request, ) -> Result, Status> { - Err(Status::unimplemented( - "get_flight_info_cross_reference has no default implementation", - )) + Err(Status::unimplemented("get_flight_info_cross_reference has no default implementation")) } /// Get a FlightInfo to extract information about the supported XDBC types. @@ -221,9 +234,7 @@ pub trait FlightSqlService: Sync + Send + Sized + 'static { _query: CommandGetXdbcTypeInfo, _request: Request, ) -> Result, Status> { - Err(Status::unimplemented( - "get_flight_info_xdbc_type_info has no default implementation", - )) + Err(Status::unimplemented("get_flight_info_xdbc_type_info has no default implementation")) } /// Implementors may override to handle additional calls to get_flight_info() @@ -246,9 +257,7 @@ pub trait FlightSqlService: Sync + Send + Sized + 'static { _ticket: TicketStatementQuery, _request: Request, ) -> Result::DoGetStream>, Status> { - Err(Status::unimplemented( - "do_get_statement has no default implementation", - )) + Err(Status::unimplemented("do_get_statement has no default implementation")) } /// Get a FlightDataStream containing the prepared statement query results. @@ -257,9 +266,7 @@ pub trait FlightSqlService: Sync + Send + Sized + 'static { _query: CommandPreparedStatementQuery, _request: Request, ) -> Result::DoGetStream>, Status> { - Err(Status::unimplemented( - "do_get_prepared_statement has no default implementation", - )) + Err(Status::unimplemented("do_get_prepared_statement has no default implementation")) } /// Get a FlightDataStream containing the list of catalogs. @@ -268,9 +275,7 @@ pub trait FlightSqlService: Sync + Send + Sized + 'static { _query: CommandGetCatalogs, _request: Request, ) -> Result::DoGetStream>, Status> { - Err(Status::unimplemented( - "do_get_catalogs has no default implementation", - )) + Err(Status::unimplemented("do_get_catalogs has no default implementation")) } /// Get a FlightDataStream containing the list of schemas. @@ -279,9 +284,7 @@ pub trait FlightSqlService: Sync + Send + Sized + 'static { _query: CommandGetDbSchemas, _request: Request, ) -> Result::DoGetStream>, Status> { - Err(Status::unimplemented( - "do_get_schemas has no default implementation", - )) + Err(Status::unimplemented("do_get_schemas has no default implementation")) } /// Get a FlightDataStream containing the list of tables. @@ -290,9 +293,7 @@ pub trait FlightSqlService: Sync + Send + Sized + 'static { _query: CommandGetTables, _request: Request, ) -> Result::DoGetStream>, Status> { - Err(Status::unimplemented( - "do_get_tables has no default implementation", - )) + Err(Status::unimplemented("do_get_tables has no default implementation")) } /// Get a FlightDataStream containing the data related to the table types. @@ -301,9 +302,7 @@ pub trait FlightSqlService: Sync + Send + Sized + 'static { _query: CommandGetTableTypes, _request: Request, ) -> Result::DoGetStream>, Status> { - Err(Status::unimplemented( - "do_get_table_types has no default implementation", - )) + Err(Status::unimplemented("do_get_table_types has no default implementation")) } /// Get a FlightDataStream containing the list of SqlInfo results. @@ -312,9 +311,7 @@ pub trait FlightSqlService: Sync + Send + Sized + 'static { _query: CommandGetSqlInfo, _request: Request, ) -> Result::DoGetStream>, Status> { - Err(Status::unimplemented( - "do_get_sql_info has no default implementation", - )) + Err(Status::unimplemented("do_get_sql_info has no default implementation")) } /// Get a FlightDataStream containing the data related to the primary and foreign keys. @@ -323,9 +320,7 @@ pub trait FlightSqlService: Sync + Send + Sized + 'static { _query: CommandGetPrimaryKeys, _request: Request, ) -> Result::DoGetStream>, Status> { - Err(Status::unimplemented( - "do_get_primary_keys has no default implementation", - )) + Err(Status::unimplemented("do_get_primary_keys has no default implementation")) } /// Get a FlightDataStream containing the data related to the exported keys. @@ -334,9 +329,7 @@ pub trait FlightSqlService: Sync + Send + Sized + 'static { _query: CommandGetExportedKeys, _request: Request, ) -> Result::DoGetStream>, Status> { - Err(Status::unimplemented( - "do_get_exported_keys has no default implementation", - )) + Err(Status::unimplemented("do_get_exported_keys has no default implementation")) } /// Get a FlightDataStream containing the data related to the imported keys. @@ -345,9 +338,7 @@ pub trait FlightSqlService: Sync + Send + Sized + 'static { _query: CommandGetImportedKeys, _request: Request, ) -> Result::DoGetStream>, Status> { - Err(Status::unimplemented( - "do_get_imported_keys has no default implementation", - )) + Err(Status::unimplemented("do_get_imported_keys has no default implementation")) } /// Get a FlightDataStream containing the data related to the cross reference. @@ -356,9 +347,7 @@ pub trait FlightSqlService: Sync + Send + Sized + 'static { _query: CommandGetCrossReference, _request: Request, ) -> Result::DoGetStream>, Status> { - Err(Status::unimplemented( - "do_get_cross_reference has no default implementation", - )) + Err(Status::unimplemented("do_get_cross_reference has no default implementation")) } /// Get a FlightDataStream containing the data related to the supported XDBC types. @@ -367,9 +356,7 @@ pub trait FlightSqlService: Sync + Send + Sized + 'static { _query: CommandGetXdbcTypeInfo, _request: Request, ) -> Result::DoGetStream>, Status> { - Err(Status::unimplemented( - "do_get_xdbc_type_info has no default implementation", - )) + Err(Status::unimplemented("do_get_xdbc_type_info has no default implementation")) } // do_put @@ -401,9 +388,7 @@ pub trait FlightSqlService: Sync + Send + Sized + 'static { _ticket: CommandStatementUpdate, _request: Request, ) -> Result { - Err(Status::unimplemented( - "do_put_statement_update has no default implementation", - )) + Err(Status::unimplemented("do_put_statement_update has no default implementation")) } /// Execute a bulk ingestion. @@ -412,9 +397,7 @@ pub trait FlightSqlService: Sync + Send + Sized + 'static { _ticket: CommandStatementIngest, _request: Request, ) -> Result { - Err(Status::unimplemented( - "do_put_statement_ingest has no default implementation", - )) + Err(Status::unimplemented("do_put_statement_ingest has no default implementation")) } /// Bind parameters to given prepared statement. @@ -427,9 +410,7 @@ pub trait FlightSqlService: Sync + Send + Sized + 'static { _query: CommandPreparedStatementQuery, _request: Request, ) -> Result { - Err(Status::unimplemented( - "do_put_prepared_statement_query has no default implementation", - )) + Err(Status::unimplemented("do_put_prepared_statement_query has no default implementation")) } /// Execute an update SQL prepared statement. @@ -438,9 +419,7 @@ pub trait FlightSqlService: Sync + Send + Sized + 'static { _query: CommandPreparedStatementUpdate, _request: Request, ) -> Result { - Err(Status::unimplemented( - "do_put_prepared_statement_update has no default implementation", - )) + Err(Status::unimplemented("do_put_prepared_statement_update has no default implementation")) } /// Execute a substrait plan @@ -449,9 +428,7 @@ pub trait FlightSqlService: Sync + Send + Sized + 'static { _query: CommandStatementSubstraitPlan, _request: Request, ) -> Result { - Err(Status::unimplemented( - "do_put_substrait_plan has no default implementation", - )) + Err(Status::unimplemented("do_put_substrait_plan has no default implementation")) } // do_action @@ -511,9 +488,7 @@ pub trait FlightSqlService: Sync + Send + Sized + 'static { _query: ActionBeginTransactionRequest, _request: Request, ) -> Result { - Err(Status::unimplemented( - "do_action_begin_transaction has no default implementation", - )) + Err(Status::unimplemented("do_action_begin_transaction has no default implementation")) } /// End a transaction @@ -522,9 +497,7 @@ pub trait FlightSqlService: Sync + Send + Sized + 'static { _query: ActionEndTransactionRequest, _request: Request, ) -> Result<(), Status> { - Err(Status::unimplemented( - "do_action_end_transaction has no default implementation", - )) + Err(Status::unimplemented("do_action_end_transaction has no default implementation")) } /// Begin a savepoint @@ -533,9 +506,7 @@ pub trait FlightSqlService: Sync + Send + Sized + 'static { _query: ActionBeginSavepointRequest, _request: Request, ) -> Result { - Err(Status::unimplemented( - "do_action_begin_savepoint has no default implementation", - )) + Err(Status::unimplemented("do_action_begin_savepoint has no default implementation")) } /// End a savepoint @@ -544,9 +515,7 @@ pub trait FlightSqlService: Sync + Send + Sized + 'static { _query: ActionEndSavepointRequest, _request: Request, ) -> Result<(), Status> { - Err(Status::unimplemented( - "do_action_end_savepoint has no default implementation", - )) + Err(Status::unimplemented("do_action_end_savepoint has no default implementation")) } /// Cancel a query @@ -555,9 +524,7 @@ pub trait FlightSqlService: Sync + Send + Sized + 'static { _query: ActionCancelQueryRequest, _request: Request, ) -> Result { - Err(Status::unimplemented( - "do_action_cancel_query has no default implementation", - )) + Err(Status::unimplemented("do_action_cancel_query has no default implementation")) } /// do_exchange diff --git a/rust/third_party/arrow-flight/src/streams.rs b/rust/third_party/arrow-flight/src/streams.rs index 8a9d5ab3..2901e0dd 100644 --- a/rust/third_party/arrow-flight/src/streams.rs +++ b/rust/third_party/arrow-flight/src/streams.rs @@ -17,13 +17,13 @@ //! [`FallibleRequestStream`] and [`FallibleTonicResponseStream`] adapters -use crate::error::FlightError; -use futures::{ - FutureExt, Stream, StreamExt, - channel::oneshot::{Receiver, Sender}, -}; use std::pin::Pin; -use std::task::{Poll, ready}; +use std::task::{ready, Poll}; + +use futures::channel::oneshot::{Receiver, Sender}; +use futures::{FutureExt, Stream, StreamExt}; + +use crate::error::FlightError; /// Wrapper around a fallible stream (one that returns errors) that makes it infallible. /// diff --git a/rust/third_party/arrow-flight/src/trailers.rs b/rust/third_party/arrow-flight/src/trailers.rs index 7929b53a..f236c474 100644 --- a/rust/third_party/arrow-flight/src/trailers.rs +++ b/rust/third_party/arrow-flight/src/trailers.rs @@ -15,14 +15,13 @@ // specific language governing permissions and limitations // under the License. -use std::{ - pin::Pin, - sync::{Arc, Mutex}, - task::{Context, Poll}, -}; +use std::pin::Pin; +use std::sync::{Arc, Mutex}; +use std::task::{Context, Poll}; -use futures::{FutureExt, Stream, StreamExt, ready}; -use tonic::{Status, Streaming, metadata::MetadataMap}; +use futures::{ready, FutureExt, Stream, StreamExt}; +use tonic::metadata::MetadataMap; +use tonic::{Status, Streaming}; /// Extract [`LazyTrailers`] from [`Streaming`] [tonic] response. /// diff --git a/rust/third_party/arrow-flight/src/utils.rs b/rust/third_party/arrow-flight/src/utils.rs index 07a18019..87c0a763 100644 --- a/rust/third_party/arrow-flight/src/utils.rs +++ b/rust/third_party/arrow-flight/src/utils.rs @@ -17,17 +17,18 @@ //! Utilities to assist with reading and writing Arrow data as Flight messages -use crate::{FlightData, SchemaAsIpc}; use std::collections::HashMap; use std::sync::Arc; use arrow_array::{ArrayRef, RecordBatch}; use arrow_buffer::Buffer; use arrow_ipc::convert::fb_to_schema; -use arrow_ipc::writer::IpcWriteContext; -use arrow_ipc::{reader, root_as_message, writer, writer::IpcWriteOptions}; +use arrow_ipc::writer::{IpcWriteContext, IpcWriteOptions}; +use arrow_ipc::{reader, root_as_message, writer}; use arrow_schema::{ArrowError, Schema, SchemaRef}; +use crate::{FlightData, SchemaAsIpc}; + /// Convert a slice of wire protocol `FlightData`s into a vector of `RecordBatch`es pub fn flight_data_to_batches(flight_data: &[FlightData]) -> Result, ArrowError> { let schema = flight_data.first().ok_or_else(|| { @@ -95,12 +96,8 @@ pub fn batches_to_flight_data( let mut ipc_write_context = IpcWriteContext::default(); for batch in batches.iter() { - let (encoded_dictionaries, encoded_batch) = data_gen.encode( - batch, - &mut dictionary_tracker, - &options, - &mut ipc_write_context, - )?; + let (encoded_dictionaries, encoded_batch) = + data_gen.encode(batch, &mut dictionary_tracker, &options, &mut ipc_write_context)?; dictionaries.extend(encoded_dictionaries.into_iter().map(Into::into)); flight_data.push(encoded_batch.into()); diff --git a/rust/tools/generate_files/src/generate.rs b/rust/tools/generate_files/src/generate.rs index 521df1e7..bfeb6ec8 100644 --- a/rust/tools/generate_files/src/generate.rs +++ b/rust/tools/generate_files/src/generate.rs @@ -4,11 +4,11 @@ use std::fs::File; use std::io::Write; use std::path::PathBuf; -use anyhow::{Context, Result, anyhow}; -use databricks_zerobus_ingest_sdk::schema::{UcColumn, descriptor_from_uc_columns}; +use anyhow::{anyhow, Context, Result}; +use databricks_zerobus_ingest_sdk::schema::{descriptor_from_uc_columns, UcColumn}; use prost_types::field_descriptor_proto::{Label, Type}; use prost_types::{DescriptorProto, FieldDescriptorProto}; -use reqwest::header::{AUTHORIZATION, CONTENT_TYPE, HeaderMap, HeaderValue}; +use reqwest::header::{HeaderMap, HeaderValue, AUTHORIZATION, CONTENT_TYPE}; use serde::Deserialize; use urlencoding::encode; @@ -50,10 +50,7 @@ pub async fn fetch_table_info(endpoint: &str, token: &str, table: &str) -> Resul let url = format!("{base}/api/2.1/unity-catalog/tables/{encoded_table}"); let mut headers = HeaderMap::new(); - headers.insert( - AUTHORIZATION, - HeaderValue::from_str(&format!("Bearer {}", token))?, - ); + headers.insert(AUTHORIZATION, HeaderValue::from_str(&format!("Bearer {}", token))?); headers.insert(CONTENT_TYPE, HeaderValue::from_static("application/json")); let client = reqwest::Client::builder() @@ -153,14 +150,7 @@ fn write_field( Label::Repeated => "repeated", }; let type_str = type_to_str(f.r#type(), f.type_name.as_deref()); - out.push_str(&format!( - "{}{} {} {} = {};\n", - indent, - label, - type_str, - f.name(), - f.number() - )); + out.push_str(&format!("{}{} {} {} = {};\n", indent, label, type_str, f.name(), f.number())); } struct MapEntry { @@ -248,10 +238,7 @@ pub fn generate_rust_and_descriptor( tonic_prost_build::configure() .out_dir(output_dir) .file_descriptor_set_path(&desc_file) - .compile_protos( - &[proto_file.to_str().unwrap()], - &[proto_dir.to_str().unwrap()], - ) + .compile_protos(&[proto_file.to_str().unwrap()], &[proto_dir.to_str().unwrap()]) .context("protoc compilation failed")?; Ok(()) @@ -385,12 +372,10 @@ mod tests { let result = generate_proto_file("TestMessage", &columns, &proto_path, &output_dir); assert!(result.is_err()); - assert!( - result - .unwrap_err() - .to_string() - .contains("map keys must be primitive types") - ); + assert!(result + .unwrap_err() + .to_string() + .contains("map keys must be primitive types")); } #[test] @@ -409,12 +394,10 @@ mod tests { let result = generate_proto_file("TestMessage", &columns, &proto_path, &output_dir); assert!(result.is_err()); - assert!( - result - .unwrap_err() - .to_string() - .contains("maps with complex value types not supported") - ); + assert!(result + .unwrap_err() + .to_string() + .contains("maps with complex value types not supported")); } #[test] @@ -432,12 +415,10 @@ mod tests { let result = generate_proto_file("TestMessage", &columns, &proto_path, &output_dir); assert!(result.is_err()); - assert!( - result - .unwrap_err() - .to_string() - .contains("nested arrays not supported") - ); + assert!(result + .unwrap_err() + .to_string() + .contains("nested arrays not supported")); } #[test] @@ -470,12 +451,10 @@ mod tests { let result = generate_proto_file("TestMessage", &columns, &proto_path, &output_dir); assert!(result.is_err()); - assert!( - result - .unwrap_err() - .to_string() - .contains("invalid field name 'invalid-name'") - ); + assert!(result + .unwrap_err() + .to_string() + .contains("invalid field name 'invalid-name'")); } #[test] @@ -488,12 +467,10 @@ mod tests { let result = generate_proto_file("TestMessage", &columns, &proto_path, &output_dir); assert!(result.is_err()); - assert!( - result - .unwrap_err() - .to_string() - .contains("cannot start with a digit") - ); + assert!(result + .unwrap_err() + .to_string() + .contains("cannot start with a digit")); } #[test] diff --git a/rust/tools/generate_files/src/main.rs b/rust/tools/generate_files/src/main.rs index 0670c9e5..51f2d171 100644 --- a/rust/tools/generate_files/src/main.rs +++ b/rust/tools/generate_files/src/main.rs @@ -12,7 +12,10 @@ use clap::Parser; mod generate; use generate::{ - clean_filename, fetch_table_info, generate_proto_file, generate_rust_and_descriptor, + clean_filename, + fetch_table_info, + generate_proto_file, + generate_rust_and_descriptor, }; mod token_factory; @@ -82,19 +85,10 @@ async fn main() -> Result<()> { let proto_output_path = args.output_dir.join(&output); - generate_proto_file( - &msg_name, - &table_info.columns, - &proto_output_path, - &args.output_dir, - ) - .context("failed to write proto file")?; + generate_proto_file(&msg_name, &table_info.columns, &proto_output_path, &args.output_dir) + .context("failed to write proto file")?; - generate_rust_and_descriptor( - proto_output_path.to_str().unwrap(), - &msg_name, - &args.output_dir, - )?; + generate_rust_and_descriptor(proto_output_path.to_str().unwrap(), &msg_name, &args.output_dir)?; Ok(()) } diff --git a/rust/tools/generate_files/src/token_factory.rs b/rust/tools/generate_files/src/token_factory.rs index 5464584e..a1fcc87f 100644 --- a/rust/tools/generate_files/src/token_factory.rs +++ b/rust/tools/generate_files/src/token_factory.rs @@ -1,4 +1,4 @@ -use anyhow::{Result, anyhow}; +use anyhow::{anyhow, Result}; pub async fn get_token( uc_endpoint: &str, @@ -84,9 +84,5 @@ fn parse_table_name(table_name: &str) -> Result<(String, String, String)> { table_name )); } - Ok(( - parts[0].to_string(), - parts[1].to_string(), - parts[2].to_string(), - )) + Ok((parts[0].to_string(), parts[1].to_string(), parts[2].to_string())) } diff --git a/typescript/CLAUDE.md b/typescript/CLAUDE.md index 21bb65f0..e0fdc2d8 100644 --- a/typescript/CLAUDE.md +++ b/typescript/CLAUDE.md @@ -42,7 +42,7 @@ Run from `typescript/`: - `npm run build:arrow` — Build with Arrow Flight support - `npm test` — Run all tests - `npm run test:unit` / `npm run test:integration` — Targeted test runs -- `cargo fmt --all` and `cargo clippy --all-targets --all-features` — Lint/format Rust code +- `../.github/scripts/rustfmt.sh --all` and `cargo clippy --all-targets --all-features` — Format/lint Rust code ## FFI boundary: NAPI-RS diff --git a/typescript/CONTRIBUTING.md b/typescript/CONTRIBUTING.md index 31fb7548..687ddc7a 100644 --- a/typescript/CONTRIBUTING.md +++ b/typescript/CONTRIBUTING.md @@ -85,7 +85,7 @@ npm run test:integration Follow standard Rust formatting: ```bash -cargo fmt --all +../.github/scripts/rustfmt.sh --all cargo clippy --all-targets --all-features ``` diff --git a/typescript/src/lib.rs b/typescript/src/lib.rs index c0390e58..809efd69 100644 --- a/typescript/src/lib.rs +++ b/typescript/src/lib.rs @@ -12,22 +12,25 @@ #![deny(clippy::all)] -use napi::bindgen_prelude::*; -use napi::threadsafe_function::{ErrorStrategy, ThreadsafeFunction}; -use napi::{Env, JsFunction, JsGlobal, JsObject, JsString, JsUnknown, ValueType}; -use napi_derive::napi; +use std::collections::HashMap; +use std::sync::Arc; use async_trait::async_trait; use databricks_zerobus_ingest_sdk::databricks::zerobus::RecordType as RustRecordType; use databricks_zerobus_ingest_sdk::{ - DefaultTokenFactory, EncodedRecord as RustRecordPayload, - HeadersProvider as RustHeadersProvider, ZerobusError as RustZerobusError, - ZerobusResult as RustZerobusResult, ZerobusSdk as RustZerobusSdk, + DefaultTokenFactory, + EncodedRecord as RustRecordPayload, + HeadersProvider as RustHeadersProvider, + ZerobusError as RustZerobusError, + ZerobusResult as RustZerobusResult, + ZerobusSdk as RustZerobusSdk, ZerobusStream as RustZerobusStream, }; +use napi::bindgen_prelude::*; +use napi::threadsafe_function::{ErrorStrategy, ThreadsafeFunction}; +use napi::{Env, JsFunction, JsGlobal, JsObject, JsString, JsUnknown, ValueType}; +use napi_derive::napi; use prost_types; -use std::collections::HashMap; -use std::sync::Arc; use tokio::sync::Mutex; /// User-Agent header value for TypeScript SDK requests. @@ -1164,7 +1167,8 @@ impl ZerobusSdk { /// Helper function to decode base64 strings. fn base64_decode(input: &str) -> std::result::Result, String> { - use base64::{engine::general_purpose::STANDARD, Engine}; + use base64::engine::general_purpose::STANDARD; + use base64::Engine; STANDARD .decode(input) .map_err(|e| format!("Base64 decode error: {}", e)) @@ -1181,8 +1185,11 @@ use arrow_ipc::writer::StreamWriter; use bytes::Bytes; #[cfg(feature = "arrow-flight")] use databricks_zerobus_ingest_sdk::{ - ArrowSchema as RustArrowSchema, DataType as RustDataType, Field as RustField, - RecordBatch as RustRecordBatch, ZerobusArrowStream as RustZerobusArrowStream, + ArrowSchema as RustArrowSchema, + DataType as RustDataType, + Field as RustField, + RecordBatch as RustRecordBatch, + ZerobusArrowStream as RustZerobusArrowStream, }; /// IPC compression type for Arrow Flight streams.