Conciseness cleanups - #30
Open
tonyalaribe wants to merge 316 commits into
Open
Conversation
…ize would have 128mb targets
DataFusion's CommonSubexprEliminate optimizer can wrap the UPDATE assignment Projection in a second Projection that defines synthetic `__common_expr_*` columns. extract_dml_info was overwriting the real assignments with the inner Projection's contents (which contain pass- throughs plus the CSE definition), so mem_buffer.update saw an assignment for a column that doesn't exist in the table schema and failed with 'Column __common_expr_1 not found'. Now only the topmost Projection is treated as the source of assignments; deeper Projections have their aliases inlined into the existing assignment value exprs.
- mem_buffer::tests::test_delete_with_qualified_predicate
- mem_buffer::tests::test_update_with_qualified_predicate_and_assignment
Verified to fail with 'FieldNotFound { field: Column { relation: Some(...) } }'
before strip_column_qualifiers was added.
- test_dml_operations::test_update_with_common_subexpression
Verified to fail (Bob's duration stays at 200 instead of 300) before
inline_projection_aliases was added — DataFusion's CSE optimizer hoists
'duration + 100' into a __common_expr_* synthetic column, and
extract_dml_info used to overwrite the real assignments with that inner
Projection's contents.
WAL UPDATE/DELETE entries serialize predicate + assignment exprs as SQL
strings. On recovery, mem_buffer.{update,delete}_by_sql re-parses them
via SqlToRel. The old EmptyContextProvider returned None from
get_function_meta — so any function reference (CAST is a built-in, but
coalesce/to_char/upper/variant_get/etc. are UDFs) failed planning with
'Internal error: No functions registered with this context' and the
entry was silently quarantined under a misleading 'WAL CORRUPTION' log.
In practice this dropped in-flight UPDATEs on every restart whenever the
SET/WHERE used a UDF — rows stayed at their pre-UPDATE values forever.
Fix:
- BufferedWriteLayer takes an Arc<dyn FunctionRegistry> via
with_function_registry(); main.rs builds it from a SessionContext
pre-populated with register_custom_functions before recover_from_wal.
- mem_buffer's parser also takes the table's DFSchema so column refs
nested inside function args resolve (empty schema rejects them).
- Renamed misleading 'WAL CORRUPTION' log lines for replay-side failures
to 'WAL REPLAY FAILED'. True deserialization corruption stays labeled.
Regression tests in mem_buffer::tests:
- parse_sql_predicate_without_registry_rejects_udf
- parse_sql_predicate_with_registry_handles_udf
- update_by_sql_with_udf_replays_when_registry_present
All three verified to fail when the registry plumbing is reverted.
Also adds .github/workflows/autoformat.yml: on push to master, runs
cargo fmt + cargo clippy --fix and pushes the result back, so PRs don't
need to ping-pong over format/clippy nits.
- Add timefusion::functions::function_registry() helper; collapse three call sites (main.rs + two tests) into one-line uses. - Add type alias mem_buffer::FnRegistry for the dyn FunctionRegistry + Send + Sync bound. Drops six 60-char trait-object signatures. - Inline parse_sql_expr — it was a one-line forwarder to parse_sql_predicate; remove the indirection. - Trim narrative comments above field, setter, regression tests, and main.rs registry construction. Keep one-line WHY; drop WHAT/incident references that git blame already carries.
… reuse runtime SessionContext Make BufferedWriteLayer's function_registry non-Option, passed positionally to with_config(). Drops the with_function_registry setter, the Option field, the runtime warn!, and the unenforceable 'MUST be called before recover_from_wal' doc-contract — now a compile-time guarantee. main.rs no longer builds a throwaway bootstrap SessionContext. Instead: 1. Create the real SessionContext early. 2. Database::setup_session_udfs() registers UDFs only (no buffered_layer dep). 3. Hand its FunctionRegistry to the layer. 4. recover_from_wal. 5. Attach buffered_layer to db. 6. Database::setup_session_tables() registers routing + stats + pg_settings. Database::setup_session_context preserved for existing call sites; it just calls the two new pieces in order. register_custom_functions runs once. Updated all internal/external/bench test callers to pass a registry via the timefusion::functions::function_registry() helper.
- Move FnRegistry alias from mem_buffer.rs (data-structure module) to functions.rs, alongside function_registry() — keeps the registry types in one place. - function_registry() is now a process-wide OnceLock singleton. Test harnesses that build many BufferedWriteLayers no longer re-register ~20 UDFs per layer. Production builds it once at startup either way. - Add timefusion::test_utils::test_helpers::test_layer(cfg) helper. Collapses 10 sites (lib tests, integration tests, benches) of BufferedWriteLayer::with_config(cfg, function_registry()?) into one call. Drops 3 now-unused imports. - Trim main.rs's three-line narration above setup_session_tables (the method name carries it) and the 'Used by main.rs' clause from Database::setup_session_udfs' doc — couples API to a caller.
Monoscope's Log Explorer emits to_jsonb(...) (Postgres syntax); DataFusion planning failed with 'Invalid function to_jsonb'. Both return JSON text over PGWire here, so a name alias is sufficient.
`do_query` gated `Response::Execution` (CommandComplete) vs. `Response::Query` (TuplesOk) on `matches!(statement, Statement::Insert(_))`. Anything that didn't decode to exactly that AST variant — postgres-synonym rewrites, plan-rewriter hooks, cached plans — fell through to the Query branch, so pgwire shipped RowDescription + DataRow for what was really a write. Clients expecting CommandComplete for INSERTs (e.g. row-count-only decoders) saw `TuplesOk` and treated it as a hard error. Replace both query paths (simple + extended) with `dml_completion`, which gates on `LogicalPlan::Dml`/`LogicalPlan::Copy` from the executed DataFrame and emits the right tag per WriteOp (INSERT/UPDATE/DELETE/ COPY/TRUNCATE/SELECT-via-CTAS). Driving off the plan keeps both query paths consistent with what DataFusion actually runs. Regression tests cover both paths: - `dml_returns_command_complete` — simple-query end-to-end via do_query - `dml_completion_survives_logical_optimisation` — extended-query gating after `state.optimize(...)` to confirm Dml survives the optimiser
Foundation for the Delta-derived cursor work in docs/plans/zero-replay-shutdown-steps-5-6.md. Step 5 needs walrus to expose its read-cursor position so we can snapshot it at bucket-seal time, record it in Delta commit metadata, and fast-forward the cursor on startup when Delta is ahead of locally-fsynced walrus state. Vendors walrus-rust 0.2.0 under vendor/walrus-rust (Cargo.toml points at the path) and adds two public APIs: - Walrus::current_position(topic) -> WalPosition — snapshots the tail position for a topic without consuming entries. - Walrus::set_persisted_read_position(topic, pos) -> io::Result<()> — fast-forwards the persisted-read cursor directly (atomic fsync via WalIndex::set), bypassing the read_next walk that advance_by_counts uses today. Fixes two pre-existing walrus bugs uncovered while writing tests for the new APIs: 1. The chain-fold block in read_next unconditionally cleared persisted_tail even when the sealed chain was empty (no fold happened). The tail-path else-branch then reset the cursor to offset 0 on every read. Now persisted_tail is only cleared after a successful fold. 2. The tail-path else-branch reset offset to 0 whenever no persisted_tail was present, even when in-memory state from a prior read on the same instance had advanced past 0. Subsequent reads would persist (block_id, 0) on top of a live cursor. Now the else-branch consults info.tail_block_id/tail_offset (also mirrored from the index at hydration time) before falling back to the offset-0 init. Plus 5 walrus-side tests in tests/position.rs covering origin, monotonicity, set-to-tail, set-to-origin, and snapshot-then-set recovery. Existing walrus test suite (16 configuration tests + others) still passes. Wires the new APIs through src/wal.rs as WalManager::current_position and WalManager::set_persisted_positions (per-shard variants over the existing shards_per_topic infrastructure). Docs include the original Step 4 plan and the new Step 5 + 6 plan (steps-5-6).
Step 1 of the zero-replay-shutdown work: plumb a CancellationToken through pgwire so its accept loop exits on shutdown signal, then await in-flight connections to drain before the BufferedWriteLayer flush fires. Without this, fresh INSERTs land in MemBuffer+WAL throughout the drain — "flush everything" becomes a moving target and the next container starts in a multi-GiB WAL replay. Step 2: per-phase shutdown timeout default raised 5s → 180s (each of pgwire drain, gRPC drain, buffered flush gets its own budget), and the memory_mb/100 heuristic dropped from compute_shutdown_timeout (never calibrated against real flush throughput, capped below realistic flush time for 5 GiB+ buffers). Vendor: datafusion-postgres serve_with_handlers takes an `impl Future<Output = ()> + Send + 'static` shutdown parameter. Pass std::future::pending() from anywhere that doesn't care. Note: Step 4 (per-shard count snapshot, advance_by_counts) lives in the same files (mem_buffer, wal, buffered_write_layer, main) as the Step 5 work and is in the follow-up commit due to file-level overlap.
…data Closes the crash-mid-flush window where Delta committed but `advance_by_counts` didn't finish — without this, restart replays entries already in Delta and the next flush double-writes them. Step 4 (already in the prior commit's territory): each MemBuffer TimeBucket accumulates per-shard WAL entry counts via `record_wal_append`; at seal time these counts drive `advance_by_counts` so the walrus cursor moves only past entries belonging to the flushed bucket, never into the open follow-on bucket. Step 5 (this commit): - TimeBucket also accumulates per-shard *positions* alongside counts. On every WAL append, `record_wal_append` queries `WalManager::current_position` and stores the post-append tail position on the bucket. At seal time both vectors snapshot into `FlushableBucket` (counts → advance, positions → watermark). - `DeltaWriteCallback` signature gains a `DeltaWatermark` parameter carrying the per-shard positions for the bucket being flushed. - `insert_records_batch` takes an `Option<&DeltaWatermark>` and, when present, attaches a `CommitProperties::with_metadata` map to the Delta write. `build_watermark_commit_properties` serializes the watermark under the `timefusion.wal_watermark` key in `commitInfo.info`. delta-rs inserts `Action::CommitInfo` at the head of the same actions vector as the Add file actions, so the metadata lands atomically in the same `_delta_log/N.json` write. - On startup, `Database::derive_wal_cursors_from_delta` scans the most recent 16 commits per known WAL topic and takes the per-shard MAX watermark. For each shard where the Delta watermark exceeds walrus's locally-fsynced cursor, walrus is fast-forwarded via `set_persisted_read_position`. Runs before `recover_from_wal` so replay starts past any entries already durable in Delta. Pure best-effort — missing/older metadata falls back to the local walrus state (today's at-least-once behaviour), so this can't make recovery worse than before. Walrus additions: `persisted_read_position` getter so the cursor comparison can take `max(local, delta)` without blindly stepping backward. Plus: `flush_callback_receives_per_shard_watermark` regression test proving the seal → callback pipeline carries the right per-shard positions. Existing `insert_records_batch` callsites updated to pass `None` (no-op for any path that doesn't go through the flush callback).
Factors the watermark serialization out of build_watermark_commit_properties and the parsing out of derive_wal_cursor_for_table into pure helpers (serialize_watermark_to_json, parse_watermark_from_json, max_watermark_across_commits) with a single WAL_WATERMARK_KEY constant so writer and reader can't drift. Four unit tests pin the format and aggregation: - serialize → parse roundtrip preserves per-shard positions, including the absent-shard distinction needed for MAX aggregation. - all-None watermark serializes to empty map (no metadata written), matching pre-feature commits on the recovery side. - per-shard MAX across commits picks the furthest position per shard; commits missing the key (replay-derived) contribute nothing and cannot reset the MAX backward. - out-of-range and malformed shard keys are dropped silently — a future writer with more shards than this reader configures won't crash recovery. Live-Delta E2E tests (real commit metadata roundtrip, walrus cursor fast-forward after simulated crash) need MinIO infrastructure and are deferred to a follow-up.
- src/wal.rs: extract `for_each_shard` + `check_shard_len` so `current_position`/`persisted_read_positions`/`set_persisted_positions`/ `advance_by_counts` stop duplicating the same per-shard iterator and length-check pattern. Add `current_position_for_shard` (no-allocation single-shard variant) and use it in `BufferedWriteLayer::insert` — the hot ingest path no longer builds an N-element Vec per append. Add `list_topic_pairs` so callers iterating topics don't reparse the `:`-joined convention. - src/mem_buffer.rs: merge `TimeBucket.wal_shard_counts` and `wal_positions` into a single `Mutex<WalShardState>` so one append takes one lock instead of two, and a snapshot can't see counts ahead of positions. `record_wal_append` keeps both in sync under one lock; `snapshot_wal_shard_state` returns both vectors. Use `Ord::max` for the monotonicity guard. - src/database.rs: parallelize `derive_wal_cursors_from_delta` (cap 8 concurrent S3 history reads — startup-blocking otherwise scales linearly with table count). Hoist `build_watermark_commit_properties` out of the retry loop. Flatten `derive_wal_cursor_for_table`'s shard-walk into a single `map_or` and drop the redundant all-None early-return. Unify the per-shard `max` via `Ord::max`. - src/config.rs: drop the dead `_current_memory_mb` parameter on `compute_shutdown_timeout` (formula was removed; arg was kept as underscore-prefixed dead weight). Tests: all 7 buffered_write_layer tests + all 4 watermark unit tests pass. MinIO-dependent tests in src/database.rs were failing before this commit too (no MinIO running locally).
Schema YAML grows a `dedup_keys` list (e.g. [id, timestamp]); the buffered write layer collapses dupes inside a bucket via Arrow RowConverter before the Delta commit and tantivy index build, so client retries land as one row. WAL still records every original, so crash replay is deterministic. Cross-bucket dupes are left for a follow-up read-side row_number rewrite. Opted in for otel_logs_and_spans on (id, timestamp).
#28) * gitignore: ignore vendored crate target/ dirs * pgwire: reply NoData for DML Describe so strict clients don't poison rows datafusion-postgres synthesises a `count: UInt64` output schema on every LogicalPlan::Dml / LogicalPlan::Copy. The default ExtendedQueryHandler fed that schema straight back to clients as a RowDescription at Describe-Statement time, so prepared INSERT/UPDATE/DELETE responses were tagged TuplesOk instead of NoData. Lenient drivers (tokio-postgres' `execute`, asyncpg) silently discard the extra DataRow; strict prepare-validating clients (Hasql, pgjdbc, Npgsql, psycopg3, sqlx) reject the protocol mismatch and drop the write. monoscope's Hasql background-ingest job was logging POISON_ROW_DROPPED with `UnexpectedResultStatementError "TuplesOk. Expecting [CommandOk]"` for every prepared INSERT. The prior attempt (7d053b2 "tag DML responses off LogicalPlan") only patched the Execute response via dml_completion and never touched Describe, which is where Hasql actually fails — before Bind/Execute is ever sent. That's why the bug survived two rounds of integration tests using tokio-postgres `execute`. Fix: in Parser::get_result_schema, recognise the synthetic DML count schema (Dml/Copy plan with one UInt64 `count` field) and return an empty result schema so pgwire emits NoData. Guard on the exact field shape so a future RETURNING-capable plan (wider schema) falls through to the normal path. Wire-level regression tests in tests/pgwire_dml_tag_test.rs cover INSERT/UPDATE/DELETE/Variant-INSERT Describe via tokio-postgres and sqlx (the latter independently exercises the same wire fact a strict Rust client would observe), plus an end-to-end Execute+SELECT to keep the Execute path honest, plus the simple-query path. Removed the prior false-confidence vendor unit tests that hit the wrong handler. * vendor: in-handlers unit test for DML get_result_schema → empty Adds get_result_schema_returns_no_data_for_dml: parses INSERT/UPDATE/ DELETE/SELECT and asserts the DML cases return an empty FieldInfo list (SELECT returns columns as the positive control). Lives next to the code it covers, no test infra dependencies, suitable as the regression test in an upstream PR. * address PR review: source-pointer, over-match guard, test-server cleanups - Comment on the get_result_schema guard now names the upstream source (datafusion/expr/src/logical_plan/dml.rs `make_count_schema`) so a future rename/widen of the synthetic DML schema points the maintainer straight at the right symbol to update. - Unit test gains a `SELECT COUNT(*) AS count FROM t` case as the over-match guard: ensures a SELECT producing a single UInt64 `count` column is NOT suppressed — only Dml/Copy of that shape is. If the `LogicalPlan::Dml | Copy` matcher is ever dropped, this case fails. - COPY can't be reached via `state.statement_to_plan` (rejected as unsupported), so the Copy arm stays defensive only — noted in test. - tests: bind to 127.0.0.1 (was 0.0.0.0), widen port range to ~2000 (was 100), and explain why simple_query uses interpolated SQL (no parameters in that wire path). Drop redundant `#[cfg(test)]` from the integration-test module. No RETURNING wire test added: DataFusion still rejects RETURNING at planning today, and the inverse (`SELECT count(*)`) covers the same "don't over-suppress" invariant from the other side. * fmt: nightly cargo fmt on tests/pgwire_dml_tag_test.rs * unbreak master CI: missing dedup_keys, DeltaWriteCallback arity, clippy Pre-existing on master, surfaced by this PR's CI run: - src/wal.rs:562 — needless-range-loop in advance_by_counts. Rewrite per clippy's suggestion (counts.iter().enumerate()). Behavior unchanged; counts.len() == shards_per_topic is asserted by check_shard_len above. - tests/tantivy_{search,index,storage}_test.rs, benches/tantivy_benchmarks.rs — TableSchema literals missing `dedup_keys: vec![]` (field added in 8d2d18d but these call sites weren't updated). - benches/{tantivy,core}_benchmarks.rs — insert_records_batch grew a 5th arg (wal watermark Option) and DeltaWriteCallback grew a 4th arg (the per-shard DeltaWatermark) when the zero-replay-shutdown work landed. Pass None / `_watermark` placeholder respectively; benches don't drive the watermark path. cargo check --all-targets --all-features and cargo clippy --all-targets --all-features -- -D warnings now both pass locally. * address review round 2: OS-assigned port, explicit table create, timeout context - TestServer: bind 127.0.0.1:0 to pull an OS-assigned port instead of random-guessing in a range. Eliminates collision with other test binaries / CI jobs sharing the host. The micro race window between drop and re-bind is harmless in practice. - TestServer: also pre-create variant_bench so Variant INSERT failures in the Describe test surface deterministically (table init errors) rather than as confusing lazy-create errors during prepare(). - connect(): explicit 10s deadline + propagate the last connect error with context. Silent 10s spins are gone. - Unit test: explicit "COPY: not tested — unreachable via prepare path" entry where the case would live, so a future reader doesn't add a COPY test and get a NotImplemented error.
- metrics.rs: replace the 13 hand-written counter registrations with a counter_registry! macro (single source of truth for field + id + desc), collapse 5 observable gauges into a layer_gauge! macro, and generate the no-arg record_* helpers via a simple_recorders! macro. - Remove a duplicated #[instrument] attribute on head_cached. - Extract config::is_insecure_auth_allowed(), shared by the pgwire and gRPC auth paths instead of inlining the env-var check twice. - Extract record_query_span() in pgwire_handlers, used by both the simple and extended query handlers.
Adds derive_more (debug + display features) and uses it for:
- StorageConfig: #[debug("[redacted]")] on the two credential fields keeps
them out of {:?} output without a hand-rolled fmt impl.
- DmlQueryPlanner / DmlExec: #[debug(skip)] on the non-printable plan/session
fields, preserving the previous Debug output.
- FoyerObjectStoreCache: container-level #[display]/#[debug] over inner.
Behavior-preserving; removes ~45 lines of boilerplate fmt code.
- Route the two error-swallowing payload-collection sites (proactive and
background _last_checkpoint refresh) through the existing collect_payload
helper instead of re-inlining the Stream/File match. The two sites that
propagate errors or slice ranges keep their bespoke handling.
- Extract table_path_from_uri(), shared by both invalidate_checkpoint_cache
impls instead of duplicating the scheme-strip + trim logic.
- Extract is_parquet_file() and use it at the 7 inline .ends_with(".parquet")
call sites.
ToCharUDF and AtTimeZoneUDF both inlined the same ~20-line block to pull a constant UTF-8 string out of a scalar-or-length-1-array argument (StringView or String). Factor it into extract_scalar_string(arg, label); the label parameterizes the error messages. Preserves the && short-circuit so is_null(0) is never evaluated on an empty array.
The as_any/name/signature trio was spelled out identically in every UDF
ScalarUDFImpl block. A scalar_udf_boilerplate!("name") macro emits all three;
return_type and invoke_with_args stay per-impl. Covers JsonToPgText, ToChar,
AtTimeZone, JsonBuildArray, ToJson, ExtractEpoch, ApproxPercentile and
JsonbPathExists.
- functions.rs: push_json_primitive! macro collapses the identical downcast + null-aware json!() loop for the Int64/Float64/Boolean arms. - manifest.rs: extract mutate() for the shared load->apply->save skeleton behind upsert() and remove_many().
query() and query_partitioned() both matched on the optional compiled predicate to filter a bucket snapshot. Factor that into filter_snapshot(); the text_match query variant keeps its bespoke id-set + predicate flow. Verified by the mem_buffer unit tests (27 passed).
The DynamoDB-locking storage-option block was written twice: once in AwsConfig::build_storage_options and again, hand-inlined, in the per-project custom-table path in database.rs. Extract a single add_dynamodb_locking_options method and call it from both, so the locking config has one source of truth.
Code Review: Conciseness CleanupsThis is a solid refactoring PR that achieves genuine reduction of repetition (-180 net lines across 12 files) without behavioral changes. The extractions are well-scoped and improve long-term maintainability. A few minor items worth noting. Summary of Changes
Security: Improvement Worth HighlightingReplacing the manual Minor Issues1. Subtle error message regression in The shared helper normalizes messages. For the non-scalar array path, the original 2. The original impl called 3. This was the same behaviour before (the read was inlined), but extraction makes it more visible. It is only called at startup/connection time so there is no performance concern. The real risk is in tests that mutate env vars between calls — both auth paths could see different values. Caching in 4. PR description is not filled in The "Closes #" line and checklist items are left at template defaults. Even for a cleanup PR it helps to note that no tests are required (pure refactor) and which issue, if any, prompted it. Nits
VerdictApproved with minor suggestions. Each extraction is small enough to verify at a glance and improves the code without over-abstracting. The credential-redaction change is a net security improvement. The only thing worth confirming before merge is that no tests assert on the exact error strings that were normalized. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Closes #
How to test
Checklist