diff --git a/.github/runner-scale-sets/query-regression/README.md b/.github/runner-scale-sets/query-regression/README.md index a1d58de7166..c0e34a4c40d 100644 --- a/.github/runner-scale-sets/query-regression/README.md +++ b/.github/runner-scale-sets/query-regression/README.md @@ -1,321 +1,184 @@ # Query regression self-hosted runners -The `Query Regression` workflow uses the dedicated ARC runner scale set -`perf-regression-8-cores`. ARC runner Pods run in the target Kubernetes cluster -and connect outbound to GitHub. The live scale set is currently **paused**: -`minRunners=0`, `maxRunners=0`, and no runner Pods. Do not resume it without -explicit approval. +The `Query Regression` workflow runs on **Aliyun ECS ephemeral runners** +(`aliyun-ecs`, the default and only automated path): a `provision` job on +`ubuntu-latest` creates one pay-as-you-go ECS instance per run, and the +instance registers itself as an ephemeral GitHub runner with a per-run +label. A `teardown` job (`if: always()`) deletes the instance; +`query-regression-janitor.yml` sweeps tagged leftovers older than 4 hours +daily. -## Prerequisites and trust admission +Build caches live on the instance's system disk, so **every run compiles +cold**; only the within-run reuse (base build warms the candidate build +through the shared target dir and sccache) applies. There is no retained +data disk. -Install the ARC scale set controller if it is not already installed: +Dispatching with any other `runner` value uses it as a literal self-hosted +runner label, which is how a manually prepared host (see +`ecs-image/bootstrap-runner-host.sh`) runs the workflow. For PR labels, set +the repository variable `QUERY_REGRESSION_PR_RUNNER` to such a label to +redirect PR runs away from ECS. + +The office ARC scale set `perf-regression-8-cores` that previously ran this +workflow is retired; see git history for its values files and pause/deploy +procedures. Tearing down the office cluster (helm release, runner namespace, +and the `query-regression-build-cache` PVC) is a manual operator action +outside this repository. This directory keeps its historical +`runner-scale-sets` name for path stability. + +## Nightly vs previous nightly + +`query-regression-nightly.yml` runs after a successful `GreptimeDB Nightly +Build` (`workflow_run`). It resolves that run's `head_sha` as the candidate +and the previous successful nightly (same branch, typically Friday when +Monday's nightly fires) as the base, then calls `query-regression.yml` with +those immutable SHAs. Both binaries are still compiled on the ECS runner; +this is not an artifact-download path. `workflow_dispatch` can pass explicit +`base_ref` / `candidate_ref` or a nightly run id. If there is no previous +nightly, or both nightlies built the same commit, the comparison is skipped. + +## Aliyun ECS path + +Configuration lives in repository variables/secrets: + +| Kind | Name | Purpose | +| --- | --- | --- | +| secret | `ALICLOUD_ECS_ACCESS_KEY_ID` / `ALICLOUD_ECS_ACCESS_KEY_SECRET` | RAM user scoped to ECS RunInstances/DeleteInstances/Describe*/CreateImage/RunCommand. Used only by provision/teardown jobs on `ubuntu-latest`; never reaches the ECS instance. | +| secret | `GH_PERSONAL_ACCESS_TOKEN` | Creates the short-lived runner registration token (shared with the jsonbench EC2 path). | +| vars | `ALIYUN_ECS_REGION_ID` / `ALIYUN_ECS_VSWITCH_ID` / `ALIYUN_ECS_SECURITY_GROUP_ID` | Network placement. The security group should allow egress only; no inbound rules are needed. The vSwitch pins the zone. | +| vars | `ALIYUN_ECS_INSTANCE_TYPE` | Dedicated (non-burstable, non-shared) instance family. Prefer 32 GiB (e.g. `ecs.g8i.2xlarge`); `ecs.c9i.2xlarge` is 8c16g and nightly thin-LTO of greptime peaks above that. Both base and candidate clusters run on the same machine, so noisy neighbors break thresholds. | +| vars | `QUERY_REGRESSION_ECS_IMAGE_ID` | Custom image built by `ecs-image/build-ecs-image.py`. | + +The system disk is 50 GiB, which covers the image, a 16 GiB swapfile, the +checkout, and cold build caches (target dir, cargo registry, sccache). ENOSPC +stops the runner itself from writing logs, which GitHub reports as `The +operation was canceled` with no telemetry, indistinguishable from a +platform-side cancellation. +cloud-init masks `systemd-oomd`, disables `unattended-upgrades` / +`apt-daily-upgrade`, creates `/swapfile`, and sets `OOMPolicy=continue` +on the runner unit. Ubuntu 24.04 defaults to `DefaultOOMPolicy=stop`, +which SIGTERM-s the whole unit when rustc is OOM-killed and GitHub +reports `The operation was canceled` with no telemetry. Unattended +upgrades can do the same via `systemctl restart` of the runner after a +library update; GitHub then records `UserCancelled` even though the job +was still valid. Swap is a safety net for 16 GiB types, not a substitute +for 32 GiB; linking on swap is slow. + +When a run dies with an unexplained `The operation was canceled` (no +"Canceled by" banner, healthy machine), re-dispatch with `keep_instance` +checked: the teardown job is skipped and the instance survives for +post-mortem inspection. The security group is egress-only, so inspect via +Cloud Assistant (`RunCommand`) or VNC: the runner's `_diag` logs under the +runner home record reconnects and worker crashes, `journalctl -u +ephemeral-github-runner.service` mirrors the console stream, `dmesg -T` +shows kernel OOM kills, and `machine-telemetry.log` in the job workspace +has the 30 s sampler history. The janitor sweep still deletes the instance +after its TTL, so finish the inspection within that window. + +Trust model on the ECS path: the instance receives only the one-hour runner +registration token via user data and holds no cloud credentials; the Aliyun +AK/SK exist only in the control-plane jobs. Instance tags +(`managed-by=query-regression-ci`, `query-regression-run-id`) feed the janitor +sweep. + +### Building and updating the ECS image + +The runner `Dockerfile` in the parent directory stays the single source of the +tool contract. Build a new ECS image from it: ```bash -helm upgrade --install arc \ - oci://ghcr.io/actions/actions-runner-controller-charts/gha-runner-scale-set-controller \ - --namespace arc-systems \ - --create-namespace \ - --version 0.14.2 +ALIBABA_CLOUD_ACCESS_KEY_ID=... ALIBABA_CLOUD_ACCESS_KEY_SECRET=... \ +uv run .github/runner-scale-sets/query-regression/ecs-image/build-ecs-image.py \ + --region-id --vswitch-id --security-group-id \ + --base-image-id ``` -Create the GitHub App secret in the runner namespace. Prefer an App limited to -`GreptimeTeam/greptimedb`: +The script boots a temporary builder instance, `docker build`s the runner +image, materializes `/opt/rustup`, `/opt/cargo`, `/usr/local/bin` tools, and +`/home/runner` (actions-runner) onto the host, installs the ephemeral-runner +systemd unit from `ecs-image/`, and snapshots a custom image. It prints the +image id; set it as `QUERY_REGRESSION_ECS_IMAGE_ID`, and bump +`RUNNER_IMAGE_EPOCH` in `query-regression.yml` at the same time so the target +cache invalidates. Builder sentinel polling requires the Cloud Assistant +agent, which Aliyun public Ubuntu images include. -```bash -kubectl -n arc-runners create secret generic greptimedb-arc-github-app \ - --from-literal=github_app_id= \ - --from-literal=github_app_installation_id= \ - --from-file=github_app_private_key= -``` - -The values files here reference that secret by name. +## Trust admission for PR runs A maintainer applying the `query-regression` or `heavy-regression` label is -**trust admission for that exact PR revision**. `query-regression` runs the six -routine default cases; `heavy-regression` runs only the high-cardinality -`prom_remote_write_7913` remote-write case. The admitted job may use this scale -set's dedicated, writable persistent cache. `pull_request: labeled` is the only -PR trigger: the label event snapshots its merge, head, and base SHAs. A queued -job fetches that immutable event merge SHA directly, verifies it is a two-parent -merge whose parents include the snapshotted head exactly once, and uses its -other parent as the actual base build revision. The snapshotted event base is -retained for audit only, so a difference from the merge's non-head parent is -not a failure. The job never follows a newer mutable PR merge ref. An -unavailable event merge, or one that does not contain exactly one snapshotted -head parent, fails closed. A later PR head change does not retarget an already -queued run: it may execute only its previously trusted event revision if that -revision remains fetchable. To run the new revision, the maintainer must review -it, remove the label, and re-add the desired regression label; cancel the old -run if it is no longer wanted. An existing label does not automatically rerun -the benchmark. +**trust admission for that exact PR revision**. `query-regression` runs the +six routine default cases; `heavy-regression` runs only the high-cardinality +`prom_remote_write_7913` remote-write case. `pull_request: labeled` is the only PR +trigger: the label event snapshots its merge, head, and base SHAs. A queued +job fetches that immutable event merge SHA directly, verifies it is a +two-parent merge whose parents include the snapshotted head exactly once, and +uses its other parent as the actual base build revision. The snapshotted +event base is retained for audit only, so a difference from the merge's +non-head parent is not a failure. The job never follows a newer mutable PR +merge ref. An unavailable event merge, or one that does not contain exactly +one snapshotted head parent, fails closed. A later PR head change does not +retarget an already queued run: it may execute only its previously trusted +event revision if that revision remains fetchable. To run the new revision, +the maintainer must review it, remove the label, and re-add the desired +regression label; cancel the old run if it is no longer wanted. An existing +label does not automatically rerun the benchmark. -Admission does not relax runner hardening or GitHub permissions. Keep -service-account token mounting disabled; do not mount host paths, the Docker -socket, kubeconfig, or long-lived credentials. The runner and cache initializer -use UID/GID 1001, disallow privilege escalation, drop all capabilities, and use -the RuntimeDefault seccomp profile. Keep GitHub tokens least-privilege and -review workflow changes before admission. Where the CNI supports it, restrict -egress to required GitHub Actions, artifact/cache, Rust/crate/toolchain, DNS, -and image-registry endpoints; block unrelated cluster services, private ranges, -and metadata endpoints unless a case requires them. - -### Network routing prerequisite - -Required split routing is an **external environment-specific prerequisite**. The -responsible network operator must route GitHub Actions, GitHub content, -artifact/cache, crates.io, Rust toolchain, and image-registry traffic through -the approved path rather than the VPN where applicable. Neither this repository -nor Kubernetes configures that route. Verify it with the responsible network -operator before any canary. +Admission does not relax runner hardening or GitHub permissions. Keep the ECS +instance free of cloud credentials and long-lived tokens, keep the security +group egress-only, keep GitHub tokens least-privilege, and review workflow +changes before admission. ## Runner image and workflow tools -Build and push the derived runner image; it preserves the official -`/home/runner/run.sh` entrypoint and supplies CI tools needed at runtime. The -image builds `otelgen` from +The runner `Dockerfile` builds `otelgen` from [`WenyXu/otelgen`](https://github.com/WenyXu/otelgen) commit [`863a3f395d062c7322cc1de08a38774b7fdaa6c8`](https://github.com/WenyXu/otelgen/commit/863a3f395d062c7322cc1de08a38774b7fdaa6c8) -so trace cases do not download or compile tools during a benchmark run: +so trace cases do not download or compile tools during a benchmark run. It is +built only as a toolchain factory: `build-ecs-image.py` and +`bootstrap-runner-host.sh` both materialize its contents onto a host, and the +benchmark itself runs host-native. -```bash -docker build \ - --platform linux/amd64 \ - -f .github/runner-scale-sets/query-regression/Dockerfile \ - -t greptime-registry.cn-hangzhou.cr.aliyuncs.com/greptime/greptimedb-query-regression-runner:latest \ - .github/runner-scale-sets/query-regression - -docker push greptime-registry.cn-hangzhou.cr.aliyuncs.com/greptime/greptimedb-query-regression-runner:latest -``` - -Deploy by digest, not mutable tag, by updating both image references in -`values-8-cores.yaml` after a rebuild. Update `RUNNER_IMAGE_DIGEST` and bump -`RUNNER_IMAGE_EPOCH` in `query-regression.yml` at the same time. If the registry -is private, use a dedicated read-only pull secret only as `imagePullSecrets`; -never expose registry credentials to runner containers. -Both digest-pinned init and runner containers use `IfNotPresent`: the immutable -digest makes a cached image safe and avoids adding a registry dependency to every -runner startup. - -The runner optionally imports only the non-sensitive `HTTP_PROXY`, -`HTTPS_PROXY`, and `NO_PROXY` variables from the -`query-regression-runner-local-env` ConfigMap. Manage that ConfigMap locally in -the target namespace; private endpoint configuration must not be committed, and -credentials or secrets must never be placed in a ConfigMap. - -Before builds, the workflow asserts UID/GID 1001 and exact image tool versions: -`libprotoc 3.21.12`, `uv 0.11.26`, `mold 2.30.0`, `Python 3.12.3`, `sccache -0.16.0`, `otelgen` commit `863a3f395d062c7322cc1de08a38774b7fdaa6c8`, -root-owned `rustup 1.29.0`, and the image-baked -`nightly-2026-03-21` Rust toolchain. Rustup, Cargo, and Rustc must resolve from -`/opt/cargo/bin`; the runner cannot write `/opt/rustup` or `/opt/cargo/bin`. -Protobuf well-known includes, including `google/protobuf/any.proto` and -`google/protobuf/empty.proto`, are an image contract and must compile with -`protoc`. +Before builds, the workflow asserts the runner UID/GID (1001 in the ECS image; +overridable via `QUERY_REGRESSION_RUNNER_UID`/`QUERY_REGRESSION_RUNNER_GID`) +and exact tool versions: `libprotoc 3.21.12`, `uv 0.11.26`, `mold 2.40.4`, +`Python 3.14.4`, `sccache 0.16.0`, `otelgen` commit +`863a3f395d062c7322cc1de08a38774b7fdaa6c8`, root-owned `rustup 1.29.0`, and +the image-baked `nightly-2026-03-21` Rust toolchain. `mold` and `python3` +come from apt at image-build time (not Ubuntu 24.04's default 3.12); bump +the Verify pins together with `QUERY_REGRESSION_ECS_IMAGE_ID` when the +image is rebuilt. Rustup, Cargo, and Rustc +must resolve from `/opt/cargo/bin`; the runner cannot write `/opt/rustup` or +`/opt/cargo/bin`. Protobuf well-known includes, including +`google/protobuf/any.proto` and `google/protobuf/empty.proto`, are an image +contract and must compile with `protoc`. `actions-rust-lang/setup-rust-toolchain@v1` is intentionally removed. The -workflow sets its warning-denying mold `RUSTFLAGS` directly, disables automatic -Rustup installation, and performs no runtime toolchain downloads. +workflow sets its warning-denying mold `RUSTFLAGS` directly, disables +automatic Rustup installation, and performs no runtime toolchain downloads. -The workflow no longer uses GitHub `rust-cache`, `setup-protoc`, `setup-uv`, or -runtime Rust setup: the image establishes immutable executable state and the PVC -supplies only reusable Cargo data. Do not reintroduce those actions unless the -corresponding cache or image contract changes. +The workflow no longer uses GitHub `rust-cache`, `setup-protoc`, `setup-uv`, +or runtime Rust setup: the image establishes immutable executable state and +each instance's system disk holds only that run's Cargo data. Do not +reintroduce those actions unless the image contract changes. -## Capacity and persistent cache +## Capacity -`values-8-cores.yaml` is normal operation: `minRunners=0`, `maxRunners=1`. -`values-paused.yaml` is the mandatory pause overlay: `minRunners=0`, -`maxRunners=0`. The job uses group `query-regression-persistent-cache-v1`, -`queue: max`, and `cancel-in-progress: false`; admitted jobs queue rather than -replacing older pending jobs. During maintenance, cancel admitted queued runs as -well as pausing ARC. Runner Pods have `activeDeadlineSeconds=12600`. -The runner requests 6 CPU and limits at 8 CPU to preserve `minipc-3` -allocatable-capacity scheduling headroom; do not reset the request to 8 CPU -without revalidating scheduling capacity. +Each run provisions its own ECS instance, so overlapping dispatches proceed +in parallel. There is no workflow `concurrency` group. All Cargo state +(`CARGO_HOME` including registry/git, `CARGO_TARGET_DIR`, sccache, cache +metadata) lives on the 50 GiB system disk and is discarded with the VM. +`RUSTUP_HOME=/opt/rustup` and `/opt/cargo/bin` are image-owned. The runner +sets `RUSTC_WRAPPER=/usr/local/bin/sccache`, +`SCCACHE_DIR=/home/runner/.cache/sccache`, `SCCACHE_CACHE_SIZE=10G`, and +`CARGO_INCREMENTAL=0`. sccache uses its local disk backend and self-evicts +at 10G. Base and candidate builds share the target on that disk; Cargo +fingerprints invalidate source and dependency changes. -The cache claim `query-regression-build-cache` is a nominal 600Gi `local-path` -PVC in `arc-runners`. It is `ReadWriteOnce`; `local-path` uses -WaitForFirstConsumer binding and Delete reclaim behavior, produces a -node-affine local PV, is non-expandable, and the 600Gi request is not a hard -storage quota. The runner's `minipc-3` selector is its only consumer candidate. - -The initializer mounts the PVC root at `/cache`, creates and write-tests these -versioned subpaths as non-root UID/GID 1001, and the runner mounts them as: - -| Persistent state | PVC subpath | Runner mount | -| --- | --- | --- | -| Ephemeral Cargo home | `emptyDir` | `/home/runner/.cargo` | -| Cargo registry data | `cargo-registry-v1` | `/home/runner/.cargo/registry` | -| Cargo Git data | `cargo-git-v1` | `/home/runner/.cargo/git` | -| Cargo target | `query-regression-target-v1` | `/home/runner/query-regression-target` | -| Cache metadata | `meta-v1` | `/home/runner/query-regression-cache-meta` | -| sccache local disk cache | `sccache-v1` | `/home/runner/.cache/sccache` | -| Immutable Rust toolchain | image-owned | `/opt/rustup`, `/opt/cargo/bin` | - -The Pod security context uses UID/GID and `fsGroup` 1001 with -`fsGroupChangePolicy: OnRootMismatch`; no privileged `chown` or raw `hostPath` -is used. `CARGO_HOME` is a per-Pod `emptyDir`; only its nested `registry` and -`git` mounts are persistent. `RUSTUP_HOME=/opt/rustup` and `/opt/cargo/bin` are -image-owned immutable paths, while `CARGO_TARGET_DIR`, cache metadata, and -`SCCACHE_DIR` are persistent absolute paths. The runner sets -`RUSTC_WRAPPER=/usr/local/bin/sccache`, -`SCCACHE_DIR=/home/runner/.cache/sccache`, `SCCACHE_CACHE_SIZE=40G`, and -`CARGO_INCREMENTAL=0`. sccache uses its local PVC disk backend and self-evicts -at 40G; do not add runtime downloads, object storage, or a shared backend. - -The repository's `.cargo/config.toml` remains a trusted per-revision build input. -In contrast, `$CARGO_HOME/config*`, credentials, installed bins, and Cargo -metadata outside the persistent `registry` and `git` data mounts are ephemeral -and cannot survive to another Pod. - -The local disk backend has a one-server constraint. `maxRunners=1` and the -unchanged `query-regression-persistent-cache-v1` workflow concurrency group -serialize runs; do not increase runner capacity or relax that serialization -while this backend is in use. Base and candidate builds share the target; Cargo -fingerprints invalidate source and dependency changes. The workflow records the -sccache version and relevant environment in the target ABI marker, starts and -zeros sccache after cache and toolchain checks, shows initial/base/candidate -statistics, and resets statistics between base and candidate builds. - -### Disk preflight and cleanup contract - -Before applying or unpausing, verify the backing filesystem on `minipc-3` has -at least 900GiB free. The current local-path provisioner source is -`/opt/local-path-provisioner`; measure the filesystem containing it: - -```bash -df -PB1G /opt/local-path-provisioner -``` - -The workflow reports `du`, `df -P`, human-readable free space, and inode -availability before builds and in an always-run report. Its cleanup is narrow -and non-destructive: - -- warn at target size 400GiB; at 450GiB clear only the complete target root; -- warn at Cargo registry-plus-Git data size 60GiB; at 80GiB remove only - `registry/src` and `git/checkouts`, then abort if that persistent data remains - at least 80GiB; -- below 300GiB backing free space, clear the complete target root first, - remeasure, then remove only those Cargo extracted trees and checkouts; abort - if free space is still below 300GiB; -- never automatically remove Cargo registry cache/index, Git database, the - image-owned Cargo bin or Rustup toolchain, cache metadata, the self-evicting - sccache directory, or the PVC. - -The target clear uses fixed absolute roots and removes all entries, including -dotfiles. After migration validation, remove obsolete `cargo-home-v1` and -`rustup-home-v1` only in explicit maintenance while ARC is 0/0 and no runner Pod -exists; they are not mounted by the current configuration. - -## Deploy and pause safely - -First verify the configured external network route and the disk preflight. Apply the PVC; while -the scale set is paused, it is expected to remain `Pending` because -WaitForFirstConsumer has no scheduled runner: - -```bash -kubectl apply --dry-run=server \ - -f .github/runner-scale-sets/query-regression/cache-pvc.yaml -kubectl apply -f .github/runner-scale-sets/query-regression/cache-pvc.yaml -``` - -Render normal and paused configurations. Normal values are always first; the -pause overlay is always last: - -```bash -helm template perf-regression-8-cores \ - oci://ghcr.io/actions/actions-runner-controller-charts/gha-runner-scale-set \ - --namespace arc-runners --version 0.14.2 \ - --set controllerServiceAccount.name=arc-gha-rs-controller \ - --set controllerServiceAccount.namespace=arc-systems \ - -f .github/runner-scale-sets/query-regression/values-8-cores.yaml - -helm template perf-regression-8-cores \ - oci://ghcr.io/actions/actions-runner-controller-charts/gha-runner-scale-set \ - --namespace arc-runners --version 0.14.2 \ - --set controllerServiceAccount.name=arc-gha-rs-controller \ - --set controllerServiceAccount.namespace=arc-systems \ - -f .github/runner-scale-sets/query-regression/values-8-cores.yaml \ - -f .github/runner-scale-sets/query-regression/values-paused.yaml -``` - -The **first post-merge Helm deployment must reconcile the release in paused -mode**. Keep the pause overlay last: - -```bash -# First post-merge deployment and every return to paused mode: 0/0. -helm upgrade --install perf-regression-8-cores \ - oci://ghcr.io/actions/actions-runner-controller-charts/gha-runner-scale-set \ - --namespace arc-runners --create-namespace --version 0.14.2 \ - --reset-values --wait \ - -f .github/runner-scale-sets/query-regression/values-8-cores.yaml \ - -f .github/runner-scale-sets/query-regression/values-paused.yaml - -# Expect 0/0 and no runner resources before considering normal mode. -kubectl -n arc-runners get autoscalingrunnerset perf-regression-8-cores \ - -o jsonpath='{.spec.minRunners}{"/"}{.spec.maxRunners}{"\n"}' -kubectl -n arc-runners get ephemeralrunners,pods \ - -l actions.github.com/scale-set-name=perf-regression-8-cores -``` - -Only after that verification and separate explicit approval, apply normal 0/1 -operation without the pause overlay: - -```bash -helm upgrade --install perf-regression-8-cores \ - oci://ghcr.io/actions/actions-runner-controller-charts/gha-runner-scale-set \ - --namespace arc-runners --create-namespace --version 0.14.2 \ - --reset-values --wait \ - -f .github/runner-scale-sets/query-regression/values-8-cores.yaml -``` - -Do not use bare `helm rollback`, `--atomic`, or `--reuse-values`: a stored -revision can restore nonzero runner capacity. Inspect rendered manifests for -capacity, the `minipc-3` selector, cache claim and mounts, initializer, -security context, and resources. After approved normal mode receives its first -canary, the PVC binds to `minipc-3`. - -For local-PV node loss, cache recovery is intentionally disposable: return to -0/0, recreate the PVC on a healthy node, cold-fill it, and run a new canary. - -## Canary and rollback - -With explicit approval, run two identical `workflow_dispatch` canaries on -`perf-regression-8-cores`, using immutable full base and candidate commit SHAs -and `cargo_profile=nightly`. The first is the cold fill; the second verifies warm -reuse. Record the workflow's base/candidate build elapsed logs and cache -ABI-marker output, initial/base/candidate sccache statistics, and cache report. -Confirm the image tool contract (root-owned Rustup/Cargo paths, baked nightly, -and non-writable `/opt` roots) and that the ephemeral Cargo home contains only -the mounted registry/Git data before Cargo creates per-Pod state. -Obtain dependency and tool network byte counters from the configured -environment counter source, filtered to `minipc-3` and the dependency/tool -destinations. - -Accept the canary only when all of the following hold: - -- exactly one runner Pod runs on `minipc-3`, and both jobs use the same bound PV; -- UID/GID 1001 cache mounts are writable; the warm run does not invalidate the - target or bulk-redownload crates or toolchains; sccache reports separate base - and candidate build statistics without server or cache-path errors; immutable - Rustup/Cargo roots remain non-writable and only registry/Git data persists; -- warm base build time is at most 50% of cold base build time; -- warm dependency/tool network bytes are at most 10% of cold fill bytes; -- cache sizes remain below soft watermarks, node free space remains at least - 300GiB, and the benchmark is correct without TLS EOFs or timeouts; -- the configured environment counter source confirms this traffic is outside VPN - accounting. - -Immediately return to 0/0 after either canary unless ongoing normal operation -has been explicitly approved; return immediately on any traffic, cache, disk, -TLS, or correctness failure. To roll back, use the paused Helm upgrade above, -or another explicit `helm upgrade --install` with known-good values followed by -`values-paused.yaml`, `--reset-values`, and `--wait`. Do not delete the PVC -automatically; preserve it for diagnosis unless intentionally discarding cache. +The workflow reports `du` / `df` in telemetry. It does not try to reclaim +space across runs: a cold 50 GiB disk that fills up fails the build. ## Future optional phases -The current phase uses a digest-pinned image with sccache 0.16.0 and no shared -cache service. Optional follow-ups are an image additionally seeded with the -exact Rust toolchain and `cargo fetch --locked`; or an internal read/write -sccache backend or Cargo/Git mirror. Evaluate them only if persistent PVC reuse -is insufficient. +The current phase uses a materialized runner toolchain with sccache 0.16.0 and +no shared cache service. Optional follow-ups are an image additionally seeded +with `cargo fetch --locked` results, or an internal read/write sccache +backend or Cargo/Git mirror. Evaluate them only if cold compile time on the +system disk becomes the bottleneck. diff --git a/.github/runner-scale-sets/query-regression/cache-pvc.yaml b/.github/runner-scale-sets/query-regression/cache-pvc.yaml deleted file mode 100644 index 692f3a2d6e2..00000000000 --- a/.github/runner-scale-sets/query-regression/cache-pvc.yaml +++ /dev/null @@ -1,15 +0,0 @@ -apiVersion: v1 -kind: PersistentVolumeClaim -metadata: - name: query-regression-build-cache - namespace: arc-runners - labels: - app.kubernetes.io/part-of: perf-regression-8-cores -spec: - volumeMode: Filesystem - accessModes: - - ReadWriteOnce - storageClassName: local-path - resources: - requests: - storage: 600Gi diff --git a/.github/runner-scale-sets/query-regression/ecs-image/bootstrap-runner-host.sh b/.github/runner-scale-sets/query-regression/ecs-image/bootstrap-runner-host.sh new file mode 100644 index 00000000000..883bbede092 --- /dev/null +++ b/.github/runner-scale-sets/query-regression/ecs-image/bootstrap-runner-host.sh @@ -0,0 +1,195 @@ +#!/usr/bin/env bash +# Copyright 2023 Greptime Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Bootstrap an Ubuntu 24.04 host as a query-regression self-hosted runner. +# +# This is the host-native equivalent of the ECS image contract: it +# installs docker-ce from Docker's official repository, builds the runner +# image from the sibling Dockerfile (the single source of the tool contract), +# materializes the toolchain onto the host, and optionally registers the +# runner and installs its systemd service. +# +# Large state lives under --data-root (a data disk): docker's data-root, the +# runner home, and the rustup/cargo toolchain roots, symlinked back to the +# hard-coded contract paths (/home/runner, /opt/rustup, /opt/cargo). +# +# The script is idempotent: re-running refreshes the toolchain in place, and +# the runner registration (.runner/.credentials) and _work survive. A live +# runner service is stopped first and restarted at the end. +# +# Usage: +# sudo bash bootstrap-runner-host.sh --repo-dir /path/to/greptimedb +# sudo RUNNER_TOKEN= bash bootstrap-runner-host.sh --repo-dir . \ +# --register --repo-url https://github.com// +# +# Options: +# --repo-dir PATH greptimedb checkout containing the runner Dockerfile (required) +# --data-root PATH data disk mount point (default: /data) +# --uid / --gid N runner user id (default: 3141; the ECS image contract is 1001) +# --register also register the runner and install the systemd service +# --repo-url URL repository URL for registration (required with --register) +# --runner-name N runner name (default: qreg-host) +# --labels L runner labels (default: perf-regression-8-cores) +# +# With --register, provide a fresh registration token via RUNNER_TOKEN. + +set -euo pipefail + +REPO_DIR="" +DATA_ROOT="/data" +RUNNER_UID="3141" +RUNNER_GID="3141" +REGISTER="false" +REPO_URL="" +RUNNER_NAME="qreg-host" +RUNNER_LABELS="perf-regression-8-cores" +IMAGE_TAG="qreg-runner:manual" + +while [[ $# -gt 0 ]]; do + case "$1" in + --repo-dir) REPO_DIR="$2"; shift 2 ;; + --data-root) DATA_ROOT="$2"; shift 2 ;; + --uid) RUNNER_UID="$2"; shift 2 ;; + --gid) RUNNER_GID="$2"; shift 2 ;; + --register) REGISTER="true"; shift ;; + --repo-url) REPO_URL="$2"; shift 2 ;; + --runner-name) RUNNER_NAME="$2"; shift 2 ;; + --labels) RUNNER_LABELS="$2"; shift 2 ;; + -h | --help) sed -n '17,40p' "$0"; exit 0 ;; + *) echo "Unknown argument: $1" >&2; exit 2 ;; + esac +done + +[[ $EUID -eq 0 ]] || { echo "Run as root, e.g. sudo bash $0 ..." >&2; exit 1; } +[[ -n "${REPO_DIR}" ]] || { echo "--repo-dir is required" >&2; exit 2; } +DOCKERFILE="${REPO_DIR}/.github/runner-scale-sets/query-regression/Dockerfile" +[[ -f "${DOCKERFILE}" ]] || { echo "Dockerfile not found at ${DOCKERFILE}" >&2; exit 1; } +if [[ "${REGISTER}" == "true" ]]; then + [[ -n "${REPO_URL}" ]] || { echo "--repo-url is required with --register" >&2; exit 2; } + [[ -n "${RUNNER_TOKEN:-}" ]] || { echo "RUNNER_TOKEN env is required with --register" >&2; exit 2; } +fi + +step() { printf '\n==> %s\n' "$*"; } + +# A re-run refreshes files under the (possibly live) runner service, so stop +# it first and restart it at the end (unless --register re-creates it). +RUNNER_SERVICE="$(systemctl list-units --type=service --all --no-legend 'actions.runner.*' 2>/dev/null | awk '{print $1}' | head -n 1 || true)" +RUNNER_WAS_ACTIVE="false" +if [[ -n "${RUNNER_SERVICE}" ]] && systemctl is-active --quiet "${RUNNER_SERVICE}"; then + step "Stop running runner service ${RUNNER_SERVICE}" + systemctl stop "${RUNNER_SERVICE}" + RUNNER_WAS_ACTIVE="true" +fi + +step "Install docker-ce from Docker's official repository" +for pkg in docker.io docker-doc docker-compose docker-compose-v2 podman-docker containerd runc; do + apt-get remove -y "${pkg}" 2>/dev/null || true +done +apt-get update +apt-get install -y ca-certificates curl +install -m 0755 -d /etc/apt/keyrings +curl -fsSL https://download.docker.com/linux/ubuntu/gpg -o /etc/apt/keyrings/docker.asc +chmod a+r /etc/apt/keyrings/docker.asc +# shellcheck disable=SC1091 +. /etc/os-release +tee /etc/apt/sources.list.d/docker.sources > /dev/null <&2 + exit 1 +fi +echo "{ \"data-root\": \"${DATA_ROOT}/docker\" }" > /etc/docker/daemon.json +# containerd keeps its own content store (where build-time image data lands); +# move it off the system disk too. +mkdir -p /etc/containerd +containerd config default | sed "s|root = \"/var/lib/containerd\"|root = \"${DATA_ROOT}/containerd\"|" > /etc/containerd/config.toml +systemctl start containerd docker +docker info | grep "Docker Root Dir" + +step "Build the runner image" +docker build --platform linux/amd64 -f "${DOCKERFILE}" -t "${IMAGE_TAG}" "${REPO_DIR}" + +step "Materialize the toolchain onto the host" +mkdir -p "${DATA_ROOT}/opt" "${DATA_ROOT}/runner" /opt /home +container="$(docker create "${IMAGE_TAG}")" +trap 'docker rm -f "${container}" >/dev/null 2>&1 || true' EXIT +# The `/.` suffix copies directory *contents*, so re-running over an existing +# materialization refreshes it in place instead of nesting runner/runner. +# Files absent from the image (runner registration, _work) are preserved. +docker cp "${container}:/home/runner/." "${DATA_ROOT}/runner" +docker cp "${container}:/opt/rustup/." "${DATA_ROOT}/opt/rustup" +docker cp "${container}:/opt/cargo/." "${DATA_ROOT}/opt/cargo" +for tool in uv uvx otelgen sccache; do + docker cp "${container}:/usr/local/bin/${tool}" "/usr/local/bin/${tool}" +done +docker rm "${container}" > /dev/null +trap - EXIT +ln -sfn "${DATA_ROOT}/runner" /home/runner +ln -sfn "${DATA_ROOT}/opt/rustup" /opt/rustup +ln -sfn "${DATA_ROOT}/opt/cargo" /opt/cargo +docker image rm "${IMAGE_TAG}" > /dev/null + +step "Install system packages" +apt-get install -y --no-install-recommends \ + build-essential clang cmake git gzip jq libprotobuf-dev libssl-dev mold \ + openssh-client pkg-config protobuf-compiler python3 sudo tar unzip wget xz-utils zip zstd + +step "Create runner user (${RUNNER_UID}:${RUNNER_GID}) and environment" +getent group "${RUNNER_GID}" > /dev/null || groupadd -g "${RUNNER_GID}" runner +id -u runner > /dev/null 2>&1 || useradd -u "${RUNNER_UID}" -g "${RUNNER_GID}" -d /home/runner -s /bin/bash runner +chown -R "${RUNNER_UID}:${RUNNER_GID}" "${DATA_ROOT}/runner" +echo 'PATH=/opt/cargo/bin:/usr/local/bin:/usr/bin:/bin' > "${DATA_ROOT}/runner/.env" +chown "${RUNNER_UID}:${RUNNER_GID}" "${DATA_ROOT}/runner/.env" + +if [[ "${REGISTER}" == "true" ]]; then + step "Register runner ${RUNNER_NAME} (labels: ${RUNNER_LABELS})" + if [[ -n "${RUNNER_SERVICE}" ]]; then + (cd /home/runner && ./svc.sh uninstall) || true + fi + runuser -u runner -- bash -c "cd /home/runner && HOME=/home/runner ./config.sh \ + --url '${REPO_URL}' --token '${RUNNER_TOKEN}' --name '${RUNNER_NAME}' \ + --labels '${RUNNER_LABELS}' --unattended --replace --disableupdate" + + step "Install and start the runner service" + cd /home/runner + ./svc.sh install runner + ./svc.sh start + ./svc.sh status +else + step "Skipping registration (pass --register --repo-url ... with RUNNER_TOKEN to enable)" + if [[ "${RUNNER_WAS_ACTIVE}" == "true" ]]; then + step "Restart runner service ${RUNNER_SERVICE}" + systemctl start "${RUNNER_SERVICE}" + systemctl --no-pager status "${RUNNER_SERVICE}" || true + fi +fi + +step "Done" +echo "Runner home: /home/runner -> ${DATA_ROOT}/runner" +echo "Before each workflow run on this persistent host, clean transient cargo state:" +echo " sudo rm -f /home/runner/.cargo/.package-cache" +echo " sudo rm -rf /home/runner/.cargo/.global-cache" diff --git a/.github/runner-scale-sets/query-regression/ecs-image/build-ecs-image.py b/.github/runner-scale-sets/query-regression/ecs-image/build-ecs-image.py new file mode 100644 index 00000000000..7031917f910 --- /dev/null +++ b/.github/runner-scale-sets/query-regression/ecs-image/build-ecs-image.py @@ -0,0 +1,392 @@ +#!/usr/bin/env python3 +# Copyright 2023 Greptime Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# PEP 723 inline metadata (see .github/scripts/aliyun-ecs-runner-provision.py +# for the convention). +# /// script +# requires-python = ">=3.10" +# dependencies = [ +# "alibabacloud_ecs20140526>=4.1.0,<6", +# "alibabacloud_tea_openapi>=0.3.12,<1", +# ] +# /// + +"""Build the query-regression ECS custom image (manual ops tool). + +Boots a temporary pay-as-you-go ECS instance from a public Ubuntu 24.04 image, +builds the existing runner container image (the Dockerfile in the parent +directory remains the single source of the tool contract), materializes the +tool directories onto the host filesystem so the workflow's "Verify runner +image tools" step holds unchanged, installs the ephemeral-runner systemd unit, +and snapshots the result as a custom image. The temporary instance is deleted +afterwards. + +Sentinel polling reads the instance's serial console output +(GetInstanceConsoleOutput) and looks for marker lines the user data writes to +/dev/console. This has no in-guest agent dependency. + +Usage: + uv run .github/runner-scale-sets/query-regression/ecs-image/build-ecs-image.py \ + --region-id cn-hangzhou --vswitch-id vsw-... --security-group-id sg-... \ + --base-image-id ubuntu_24_04_x64_20G_alibase_*.vhd +""" + +from __future__ import annotations + +import argparse +import base64 +import os +import time +from pathlib import Path + +ASSETS_DIR = Path(__file__).resolve().parent +DONE_MARKER = "QREG_IMAGE_BUILD_DONE" +FAILED_MARKER = "QREG_IMAGE_BUILD_FAILED" +POLL_INTERVAL_SECONDS = 15 +CONSOLE_POLL_INTERVAL_SECONDS = 30 +BUILD_TIMEOUT_SECONDS = 60 * 60 + +# Same apt package contract as the runner Dockerfile; the base +# actions-runner image is Ubuntu 24.04, so an Ubuntu 24.04 host resolves the +# same tool versions (protoc 3.21.12, mold 2.40.4, Python 3.14.4). +# Docker itself comes from Docker's official repository (docker-ce), not the +# distribution-packaged docker.io. +APT_PACKAGES = [ + "build-essential", + "ca-certificates", + "clang", + "cmake", + "curl", + "git", + "gpg", + "gzip", + "jq", + "libprotobuf-dev", + "libssl-dev", + "mold", + "openssh-client", + "pkg-config", + "protobuf-compiler", + "python3", + "sudo", + "tar", + "unzip", + "wget", + "xz-utils", + "zip", + "zstd", +] + +DOCKER_CE_PACKAGES = "docker-ce docker-ce-cli containerd.io docker-buildx-plugin" + + +def render_user_data(dockerfile: str, start_runner: str, unit: str) -> str: + dockerfile_b64 = base64.b64encode(dockerfile.encode()).decode() + start_runner_b64 = base64.b64encode(start_runner.encode()).decode() + unit_b64 = base64.b64encode(unit.encode()).decode() + packages = " ".join(APT_PACKAGES) + return f"""#!/bin/bash +set -euo pipefail +trap 'echo "{FAILED_MARKER} at line $LINENO" > /dev/console' ERR +# Stream the full build log to the serial console so the poller (and anyone +# watching GetInstanceConsoleOutput) sees real progress, not a silent login +# prompt for the whole build. +exec > >(tee -a /dev/console) 2>&1 + +apt-get update +DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends {packages} + +# Docker from the official repository, not the distribution-packaged docker.io. +install -m 0755 -d /etc/apt/keyrings +curl -fsSL https://download.docker.com/linux/ubuntu/gpg | gpg --dearmor -o /etc/apt/keyrings/docker.gpg +chmod a+r /etc/apt/keyrings/docker.gpg +. /etc/os-release +echo "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.gpg] https://download.docker.com/linux/ubuntu ${{VERSION_CODENAME}} stable" \ + > /etc/apt/sources.list.d/docker.list +apt-get update +DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends {DOCKER_CE_PACKAGES} + +base64 -d > /tmp/Dockerfile <<'EOF' +{dockerfile_b64} +EOF +mkdir -p /tmp/image-context +docker build --platform linux/amd64 -f /tmp/Dockerfile -t qreg-runner:local /tmp/image-context + +# Materialize the tool contract onto the host filesystem. +container="$(docker create qreg-runner:local)" +trap 'docker rm -f "${{container}}" >/dev/null 2>&1 || true' EXIT +docker cp "${{container}}:/opt/rustup" /opt/rustup +docker cp "${{container}}:/opt/cargo" /opt/cargo +for tool in uv uvx otelgen sccache; do + docker cp "${{container}}:/usr/local/bin/${{tool}}" "/usr/local/bin/${{tool}}" +done +docker cp "${{container}}:/home/runner" /home/runner + +# Runner identity mirrors the container image (UID/GID 1001). +groupadd --gid 1001 runner 2>/dev/null || true +useradd --uid 1001 --gid 1001 --home-dir /home/runner --shell /bin/bash runner 2>/dev/null || true +chown -R 1001:1001 /home/runner + +mkdir -p /opt/ephemeral-github-runner +base64 -d > /opt/ephemeral-github-runner/start-runner.sh <<'EOF' +{start_runner_b64} +EOF +chmod 0755 /opt/ephemeral-github-runner/start-runner.sh +base64 -d > /etc/systemd/system/ephemeral-github-runner.service <<'EOF' +{unit_b64} +EOF +systemctl daemon-reload +systemctl enable ephemeral-github-runner.service + +# Keep the image free of the build-time docker state (also shrinks the +# snapshot: buildkit cache is several GB). +docker rm -f "${{container}}" >/dev/null +docker system prune -af >/dev/null +trap - EXIT + +echo "{DONE_MARKER}" > /dev/console +""" + + +def make_ecs_client(region_id: str): + from alibabacloud_ecs20140526.client import Client as EcsClient + from alibabacloud_tea_openapi.models import Config as OpenApiConfig + + return EcsClient( + OpenApiConfig( + access_key_id=os.environ["ALIBABA_CLOUD_ACCESS_KEY_ID"], + access_key_secret=os.environ["ALIBABA_CLOUD_ACCESS_KEY_SECRET"], + region_id=region_id, + endpoint=f"ecs.{region_id}.aliyuncs.com", + ) + ) + + +def wait_for_instance_status(client, region_id: str, instance_id: str, wanted: str, deadline: float) -> None: + from alibabacloud_ecs20140526 import models as ecs_models + + import json + + while time.monotonic() < deadline: + response = call_api_with_retry( + lambda: client.describe_instances( + ecs_models.DescribeInstancesRequest( + region_id=region_id, instance_ids=json.dumps([instance_id]) + ) + ), + "DescribeInstances", + ) + instances = response.body.instances.instance + if instances and instances[0].status == wanted: + return + time.sleep(POLL_INTERVAL_SECONDS) + raise TimeoutError(f"Instance {instance_id} did not reach status {wanted} in time") + + +# Aliyun error codes worth retrying: throttling and server-side faults. Client +# errors (4xx: permissions, bad parameters) are configuration problems and must +# fail fast instead of being retried. +TRANSIENT_ERROR_CODES = {"Throttling", "Throttling.User", "InternalError", "ServiceUnavailable"} + + +def is_transient(error: Exception) -> bool: + # UnretryableException comes from the darabonba network layer: connection + # reset, read timeout, DNS blip. + if type(error).__name__ == "UnretryableException": + return True + code = getattr(error, "code", "") or "" + status = getattr(error, "statusCode", None) + return code in TRANSIENT_ERROR_CODES or (isinstance(status, int) and status >= 500) + + +def call_api_with_retry(fn, description: str, attempts: int = 5): + """Retry transient API/network failures; the build is too long to die on a blip. + + Only call this with idempotent or safely-repeatable requests (reads, + StopInstance, CreateImage). Never with RunInstances: if the request + succeeded but the response was lost, a retry double-creates instances. + """ + for attempt in range(1, attempts + 1): + try: + return fn() + except Exception as error: # noqa: BLE001 + if attempt == attempts or not is_transient(error): + raise + print(f"{description} failed (attempt {attempt}/{attempts}): {error}", flush=True) + time.sleep(POLL_INTERVAL_SECONDS) + raise RuntimeError("unreachable: retry loop exited without returning") + + +def read_console_output(client, region_id: str, instance_id: str) -> str: + """Fetch the instance's serial console output; no in-guest agent needed.""" + from alibabacloud_ecs20140526 import models as ecs_models + + response = call_api_with_retry( + lambda: client.get_instance_console_output( + ecs_models.GetInstanceConsoleOutputRequest(region_id=region_id, instance_id=instance_id) + ), + "GetInstanceConsoleOutput", + ) + return base64.b64decode(response.body.console_output or "").decode("utf-8", "replace") + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--region-id", default=os.environ.get("ALIYUN_ECS_REGION_ID")) + parser.add_argument("--vswitch-id", default=os.environ.get("ALIYUN_ECS_VSWITCH_ID")) + parser.add_argument("--security-group-id", default=os.environ.get("ALIYUN_ECS_SECURITY_GROUP_ID")) + parser.add_argument("--base-image-id", default=os.environ.get("ALIYUN_ECS_BASE_IMAGE_ID")) + parser.add_argument("--resource-group-id", default=os.environ.get("ALIYUN_ECS_RESOURCE_GROUP_ID")) + parser.add_argument("--instance-type", default="ecs.g7.xlarge") + parser.add_argument("--image-name", default=None, help="Defaults to a timestamped name.") + args = parser.parse_args() + + for name in ("region_id", "vswitch_id", "security_group_id", "base_image_id"): + if not getattr(args, name): + raise SystemExit(f"Missing required configuration: --{name.replace('_', '-')}") + + from alibabacloud_ecs20140526 import models as ecs_models + + client = make_ecs_client(args.region_id) + image_name = args.image_name or time.strftime( + "greptimedb-query-regression-runner-%Y%m%d%H%M%S", time.gmtime() + ) + + user_data = base64.b64encode( + render_user_data( + (ASSETS_DIR.parent / "Dockerfile").read_text(), + (ASSETS_DIR / "start-runner.sh").read_text(), + (ASSETS_DIR / "ephemeral-github-runner.service").read_text(), + ).encode() + ).decode() + + instance_id = None + try: + response = client.run_instances( + ecs_models.RunInstancesRequest( + region_id=args.region_id, + image_id=args.base_image_id, + resource_group_id=args.resource_group_id, + instance_type=args.instance_type, + v_switch_id=args.vswitch_id, + security_group_id=args.security_group_id, + instance_name=f"build-{image_name}", + description="Temporary builder for the query-regression ECS image", + amount=1, + instance_charge_type="PostPaid", + spot_strategy="NoSpot", + internet_charge_type="PayByTraffic", + internet_max_bandwidth_out=100, + # Peak usage is ~18G (OS+apt, docker image, and the materialized + # toolchain coexist briefly): deliberately tight, and a smaller + # disk makes the image snapshot faster. Instances created from + # the image get a larger system disk from the provision side. + system_disk=ecs_models.RunInstancesRequestSystemDisk( + category="cloud_essd", size="20" + ), + user_data=user_data, + tag=[ + ecs_models.RunInstancesRequestTag(key="managed-by", value="query-regression-ci"), + ecs_models.RunInstancesRequestTag(key="role", value="image-builder"), + ], + ) + ) + instance_id = response.body.instance_id_sets.instance_id_set[0] + print(f"Builder instance: {instance_id}", flush=True) + wait_for_instance_status(client, args.region_id, instance_id, "Running", time.monotonic() + 10 * 60) + + deadline = time.monotonic() + BUILD_TIMEOUT_SECONDS + while time.monotonic() < deadline: + console = read_console_output(client, args.region_id, instance_id) + if DONE_MARKER in console: + break + if FAILED_MARKER in console: + tail = "\n".join(console.splitlines()[-20:]) + raise RuntimeError(f"Image build failed on the builder; console tail:\n{tail}") + lines = console.splitlines() + print(f"Build in progress; last console line: {lines[-1] if lines else '(none yet)'}", flush=True) + time.sleep(CONSOLE_POLL_INTERVAL_SECONDS) + else: + raise TimeoutError("Image build did not finish in time") + + print("Stopping builder before image creation", flush=True) + try: + call_api_with_retry( + lambda: client.stop_instance( + ecs_models.StopInstanceRequest( + instance_id=instance_id, + # The instance is deleted right after the snapshot, but + # there is no reason to keep billing vCPU during the stop. + stopped_mode="StopCharging", + ) + ), + "StopInstance", + ) + except Exception as error: # noqa: BLE001 + # Instance families with local disks do not support StopCharging. + print(f"StopCharging unavailable ({error}); stopping with default mode", flush=True) + call_api_with_retry( + lambda: client.stop_instance(ecs_models.StopInstanceRequest(instance_id=instance_id)), + "StopInstance", + ) + wait_for_instance_status(client, args.region_id, instance_id, "Stopped", time.monotonic() + 10 * 60) + + image = call_api_with_retry( + lambda: client.create_image( + ecs_models.CreateImageRequest( + region_id=args.region_id, instance_id=instance_id, image_name=image_name + ) + ), + "CreateImage", + ) + image_id = image.body.image_id + deadline = time.monotonic() + 60 * 60 + while time.monotonic() < deadline: + description = call_api_with_retry( + lambda: client.describe_images( + ecs_models.DescribeImagesRequest(region_id=args.region_id, image_id=image_id) + ), + "DescribeImages", + ) + images = description.body.images.image + if images: + status = images[0].status + print(f"Image {image_id} status: {status} (progress {images[0].progress})", flush=True) + if status == "Available": + break + time.sleep(POLL_INTERVAL_SECONDS) + else: + raise TimeoutError( + f"Image {image_id} did not become Available in time. The snapshot " + f"continues server-side: check its status in the console and reuse it " + f"once Available instead of rebuilding." + ) + + print(f"Custom image ready: {image_id} ({image_name})", flush=True) + print(f"Set the repo variable QUERY_REGRESSION_ECS_IMAGE_ID={image_id}", flush=True) + return 0 + finally: + if instance_id: + try: + client.delete_instance( + ecs_models.DeleteInstanceRequest(instance_id=instance_id, force=True) + ) + print(f"Deleted builder instance {instance_id}", flush=True) + except Exception as error: # noqa: BLE001 + print(f"Failed to delete builder instance {instance_id}: {error}", flush=True) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/.github/runner-scale-sets/query-regression/ecs-image/ephemeral-github-runner.service b/.github/runner-scale-sets/query-regression/ecs-image/ephemeral-github-runner.service new file mode 100644 index 00000000000..a1cbe6a4503 --- /dev/null +++ b/.github/runner-scale-sets/query-regression/ecs-image/ephemeral-github-runner.service @@ -0,0 +1,45 @@ +# Copyright 2023 Greptime Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +[Unit] +Description=Ephemeral GitHub Actions runner +After=network-online.target +Wants=network-online.target +# cloud-final runs the provision user data, which writes +# /etc/ephemeral-github-runner.env and the drop-ins below. Ordering after it +# makes the runner come up configured on first boot, no restart needed. +# (cloud-init's `systemctl restart` in user data remains as the bridge for +# images built before this ordering existed.) +After=cloud-final.service + +[Service] +Type=simple +# The toolchain lives in the image at /opt/cargo/bin (the Dockerfile's ENV +# PATH does not survive materialization onto the host). Publish it here so +# every runner job inherits it. /etc/ephemeral-github-runner.env currently +# carries the same PATH line as a bridge for images built before this +# directive existed; EnvironmentFile is applied after Environment=, and the +# two values are kept identical, so there is no conflict. +Environment=PATH=/opt/cargo/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin +EnvironmentFile=/etc/ephemeral-github-runner.env +ExecStart=/opt/ephemeral-github-runner/start-runner.sh +Restart=no +# If the kernel OOM killer kills a job child (rustc/cargo), keep the runner +# service alive so the job is reported as a normal failure with logs instead +# of the runner vanishing and GitHub reporting a cancellation. cloud-init's +# oom.conf drop-in carries the same directive for images built before this. +OOMPolicy=continue + +[Install] +WantedBy=multi-user.target diff --git a/.github/runner-scale-sets/query-regression/ecs-image/start-runner.sh b/.github/runner-scale-sets/query-regression/ecs-image/start-runner.sh new file mode 100644 index 00000000000..a6aff87ca22 --- /dev/null +++ b/.github/runner-scale-sets/query-regression/ecs-image/start-runner.sh @@ -0,0 +1,37 @@ +#!/bin/bash +# Copyright 2023 Greptime Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Configure and start the ephemeral GitHub Actions runner on an ECS instance. +# Baked into the custom image at /opt/ephemeral-github-runner/ by +# build-ecs-image.py; the per-run environment is written by cloud-init to +# /etc/ephemeral-github-runner.env. +set -euo pipefail + +# shellcheck disable=SC1091 +source /etc/ephemeral-github-runner.env + +export HOME=/home/runner +cd /home/runner + +if [[ ! -f /home/runner/.runner ]]; then + runuser -u runner -- ./config.sh \ + --url "${REPO_URL}" \ + --token "${RUNNER_TOKEN}" \ + --name "${RUNNER_NAME}" \ + --labels "${RUNNER_LABELS}" \ + --ephemeral --unattended --replace --disableupdate +fi + +exec runuser -u runner -- ./run.sh diff --git a/.github/runner-scale-sets/query-regression/values-8-cores.yaml b/.github/runner-scale-sets/query-regression/values-8-cores.yaml deleted file mode 100644 index cff1322680b..00000000000 --- a/.github/runner-scale-sets/query-regression/values-8-cores.yaml +++ /dev/null @@ -1,146 +0,0 @@ -githubConfigUrl: "https://github.com/GreptimeTeam/greptimedb" -githubConfigSecret: greptimedb-arc-github-app - -runnerScaleSetName: "perf-regression-8-cores" -minRunners: 0 -maxRunners: 1 - -template: - spec: - automountServiceAccountToken: false - activeDeadlineSeconds: 12600 - nodeSelector: - kubernetes.io/hostname: minipc-3 - securityContext: - runAsNonRoot: true - runAsUser: 1001 - runAsGroup: 1001 - fsGroup: 1001 - fsGroupChangePolicy: OnRootMismatch - seccompProfile: - type: RuntimeDefault - volumes: - - name: cargo-home - emptyDir: {} - - name: build-cache - persistentVolumeClaim: - claimName: query-regression-build-cache - initContainers: - - name: initialize-build-cache - image: greptime-registry.cn-hangzhou.cr.aliyuncs.com/greptime/greptimedb-query-regression-runner@sha256:e713b294e23b7e15184e558866c90025e59930033e72c97650dbc7f1ca022d11 - imagePullPolicy: IfNotPresent - command: - - /bin/sh - - -ec - - | - umask 0002 - for directory in cargo-registry-v1 cargo-git-v1 query-regression-target-v1 meta-v1 sccache-v1; do - cache_directory="/cache/${directory}" - mkdir -p "${cache_directory}" - test -w "${cache_directory}" - probe_file="${cache_directory}/.write-probe" - : > "${probe_file}" - test -f "${probe_file}" - rm "${probe_file}" - done - securityContext: - runAsNonRoot: true - runAsUser: 1001 - runAsGroup: 1001 - allowPrivilegeEscalation: false - capabilities: - drop: - - ALL - resources: - requests: - cpu: 100m - memory: 128Mi - ephemeral-storage: 1Gi - limits: - cpu: 500m - memory: 512Mi - ephemeral-storage: 1Gi - volumeMounts: - - name: build-cache - mountPath: /cache - containers: - - name: runner - image: greptime-registry.cn-hangzhou.cr.aliyuncs.com/greptime/greptimedb-query-regression-runner@sha256:e713b294e23b7e15184e558866c90025e59930033e72c97650dbc7f1ca022d11 - imagePullPolicy: IfNotPresent - command: ["/home/runner/run.sh"] - env: - - name: CARGO_HOME - value: /home/runner/.cargo - - name: HTTP_PROXY - valueFrom: - configMapKeyRef: - name: query-regression-runner-local-env - key: HTTP_PROXY - optional: true - - name: HTTPS_PROXY - valueFrom: - configMapKeyRef: - name: query-regression-runner-local-env - key: HTTPS_PROXY - optional: true - - name: NO_PROXY - valueFrom: - configMapKeyRef: - name: query-regression-runner-local-env - key: NO_PROXY - optional: true - - name: UV_CACHE_DIR - value: /home/runner/.cargo/uv-cache - - name: RUSTUP_HOME - value: /opt/rustup - - name: RUSTUP_TOOLCHAIN - value: nightly-2026-03-21 - - name: RUSTUP_AUTO_INSTALL - value: "0" - - name: CARGO_TARGET_DIR - value: /home/runner/query-regression-target - - name: QUERY_REGRESSION_CACHE_META - value: /home/runner/query-regression-cache-meta - - name: RUSTC_WRAPPER - value: /usr/local/bin/sccache - - name: SCCACHE_DIR - value: /home/runner/.cache/sccache - - name: SCCACHE_CACHE_SIZE - value: 40G - - name: CARGO_INCREMENTAL - value: "0" - securityContext: - runAsNonRoot: true - runAsUser: 1001 - runAsGroup: 1001 - allowPrivilegeEscalation: false - capabilities: - drop: - - ALL - resources: - requests: - cpu: "6" - memory: 24Gi - ephemeral-storage: 80Gi - limits: - cpu: "8" - memory: 24Gi - ephemeral-storage: 80Gi - volumeMounts: - - name: cargo-home - mountPath: /home/runner/.cargo - - name: build-cache - mountPath: /home/runner/.cargo/registry - subPath: cargo-registry-v1 - - name: build-cache - mountPath: /home/runner/.cargo/git - subPath: cargo-git-v1 - - name: build-cache - mountPath: /home/runner/query-regression-target - subPath: query-regression-target-v1 - - name: build-cache - mountPath: /home/runner/query-regression-cache-meta - subPath: meta-v1 - - name: build-cache - mountPath: /home/runner/.cache/sccache - subPath: sccache-v1 diff --git a/.github/runner-scale-sets/query-regression/values-paused.yaml b/.github/runner-scale-sets/query-regression/values-paused.yaml deleted file mode 100644 index fe734ef7615..00000000000 --- a/.github/runner-scale-sets/query-regression/values-paused.yaml +++ /dev/null @@ -1,2 +0,0 @@ -minRunners: 0 -maxRunners: 0 diff --git a/.github/scripts/aliyun-ecs-runner-provision.py b/.github/scripts/aliyun-ecs-runner-provision.py new file mode 100644 index 00000000000..307ee336ebf --- /dev/null +++ b/.github/scripts/aliyun-ecs-runner-provision.py @@ -0,0 +1,488 @@ +#!/usr/bin/env python3 +# Copyright 2023 Greptime Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# This is the first `.github/scripts/` user of PEP 723 inline script metadata: +# `uv run` installs the pinned Aliyun SDK before executing. The SDK import stays +# lazy (inside `make_ecs_client`) so unit tests can import this module with a +# plain stdlib interpreter. +# /// script +# requires-python = ">=3.10" +# dependencies = [ +# # Pin exactly after the first live run against the Aliyun API. +# "alibabacloud_ecs20140526>=4.1.0,<6", +# "alibabacloud_tea_openapi>=0.3.12,<1", +# ] +# /// + +"""Provision an ephemeral Aliyun ECS query-regression runner. + +Creates one pay-as-you-go ECS instance from the query-regression custom image +and waits until the instance registers itself as an ephemeral GitHub Actions +runner with a per-run label. Build caches live on the instance system disk +and disappear with the VM. The instance id is written to job outputs as soon +as the VM exists so the teardown job is the single DeleteInstance caller (a +timeout here used to delete and then teardown deleted again). On runner +online timeout the console is dumped and the job fails; teardown releases +the instance. + +The Aliyun credentials are used only in this control-plane job (ubuntu-latest) +and are never passed to the ECS instance; the instance receives only a +short-lived GitHub runner registration token via user data. +""" + +from __future__ import annotations + +import argparse +import base64 +import json +import os +import time +import urllib.error +import urllib.request +from dataclasses import dataclass + +RUNNER_NAME_PREFIX = "qreg-ecs" +RUNNER_LABEL_PREFIX = "query-regression-ecs" +MANAGED_BY_TAG_KEY = "managed-by" +MANAGED_BY_TAG_VALUE = "query-regression-ci" +RUN_TAG_KEY = "query-regression-run-id" + +# Runner cache paths on the instance system disk. They are created empty +# every provision and deleted with the VM. +CACHE_PATHS = ( + "/home/runner/.cargo/registry", + "/home/runner/.cargo/git", + "/home/runner/query-regression-target", + "/home/runner/query-regression-cache-meta", + "/home/runner/.cache/sccache", +) + +RUNNER_ONLINE_TIMEOUT_SECONDS = 10 * 60 +POLL_INTERVAL_SECONDS = 5 + +# Nightly thin-LTO of greptime peaks above ecs.c9i.2xlarge's 16 GiB. A swap +# file plus masking systemd-oomd lets the kernel reclaim rustc pages instead +# of SIGTERM-ing the runner service cgroup (GitHub then reports the step as +# "The operation was canceled"). Sized to match that 16 GiB instance; a +# 32 GiB type still benefits as a safety net. Lives on the system disk next +# to the cold build caches. +SWAP_FILE = "/swapfile" +SWAP_SIZE_GIB = 16 +# Image (~12G), 16G swap, checkout, base+candidate data homes, and cold +# build caches. Deleted with the instance. +SYSTEM_DISK_GIB = 50 + + +@dataclass(frozen=True) +class ProvisionConfig: + region_id: str + vswitch_id: str + security_group_id: str + image_id: str + instance_type: str + repo: str + run_id: str + github_token: str + # Optional resource group; required when the RAM grant is scoped to one. + resource_group_id: str | None = None + # Runner identity inside the image; the workflow asserts the same values. + runner_uid: str = "1001" + runner_gid: str = "1001" + + +def runner_name_for_run(run_id: str) -> str: + return f"{RUNNER_NAME_PREFIX}-{run_id}" + + +def runner_label_for_run(run_id: str) -> str: + return f"{RUNNER_LABEL_PREFIX}-{run_id}" + + +def render_user_data( + runner_name: str, + runner_label: str, + runner_token: str, + repo: str, + runner_uid: str = "1001", + runner_gid: str = "1001", +) -> str: + """Render the cloud-init shell script for the runner instance.""" + destinations = " ".join(f'"{dst}"' for dst in CACHE_PATHS) + cache_setup = f"""# Caches live on the system disk, are deleted with the instance, and every +# run compiles cold. Within-run reuse (base warming candidate via the shared +# target dir and sccache) still applies. +mkdir -p {destinations} +chown -R {runner_uid}:{runner_gid} /home/runner""" + swap_setup = f"""# Mask systemd-oomd before enabling swap: Ubuntu 24.04 kills the whole +# service cgroup on PSI pressure, and swap thrashing looks like pressure. +systemctl disable --now systemd-oomd.socket systemd-oomd.service || true +systemctl mask systemd-oomd.socket systemd-oomd.service || true +# apt-daily-upgrade / unattended-upgrades can `systemctl restart` services +# whose libraries were updated. That SIGTERM-s the runner mid-job; GitHub +# records UserCancelled (job still valid, child exit 143). These VMs live +# ~1h; security updates belong in the image, not at job time. +systemctl disable --now unattended-upgrades.service apt-daily.timer apt-daily-upgrade.timer apt-daily.service apt-daily-upgrade.service || true +systemctl mask unattended-upgrades.service apt-daily.timer apt-daily-upgrade.timer apt-daily.service apt-daily-upgrade.service || true +systemctl stop unattended-upgrades.service || true +killall -9 unattended-upgr || true +cat > /etc/apt/apt.conf.d/99disable-auto-updates <<'APTEOF' +APT::Periodic::Update-Package-Lists "0"; +APT::Periodic::Unattended-Upgrade "0"; +APT::Periodic::Download-Upgradeable-Packages "0"; +APTEOF +if [[ ! -f "{SWAP_FILE}" ]]; then + fallocate --length {SWAP_SIZE_GIB}G "{SWAP_FILE}" + chmod 600 "{SWAP_FILE}" + mkswap "{SWAP_FILE}" +fi +swapon "{SWAP_FILE}" +sysctl --write vm.swappiness=10 +swapon --show +free --human""" + return f"""#!/bin/bash +set -euo pipefail + +{cache_setup} + +{swap_setup} + +cat > /etc/ephemeral-github-runner.env <<'ENVEOF' +RUNNER_NAME={runner_name} +RUNNER_LABELS={runner_label} +RUNNER_TOKEN={runner_token} +REPO_URL=https://github.com/{repo} +# The container image baked /opt/cargo/bin into PATH via Dockerfile ENV; on +# the host nothing inherits that, so publish it through the unit's +# EnvironmentFile: runner job processes inherit the runner's environment. +PATH=/opt/cargo/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin +ENVEOF + +# Stream the runner service's own logs (config output, job lifecycle, crash +# messages) to the serial console, next to the kernel's OOM-killer records. +# Teardown dumps the console tail before deleting the instance, so the +# witness of what killed the runner survives the machine. +mkdir -p /etc/systemd/system/ephemeral-github-runner.service.d +cat > /etc/systemd/system/ephemeral-github-runner.service.d/console.conf <<'CONFEOF' +[Service] +StandardOutput=journal+console +StandardError=journal+console +CONFEOF +# systemd's DefaultOOMPolicy=stop SIGTERM-s the whole unit when *any* +# process in the cgroup is OOM-killed (typically rustc during thin LTO). +# That is the "Session terminated, killing shell..." / job Canceled path: +# runuser dies, GitHub deletes the ephemeral registration, always() steps +# never run. continue keeps the runner up so cargo can report signal 9. +cat > /etc/systemd/system/ephemeral-github-runner.service.d/oom.conf <<'CONFEOF' +[Service] +OOMPolicy=continue +CONFEOF +systemctl daemon-reload +# restart (not start): the image enables this unit, so it may already +# have come up without the drop-ins if it raced cloud-init. restart +# applies OOMPolicy=continue to the running cgroup. --no-block matters: +# newer image units carry After=cloud-final.service, and this script IS +# cloud-final — a blocking restart would deadlock until TimeoutStartSec. +systemctl restart --no-block ephemeral-github-runner.service +""" + + +def encode_user_data(script: str) -> str: + return base64.b64encode(script.encode("utf-8")).decode("ascii") + + +def github_api(token: str, method: str, path: str, body: dict | None = None) -> dict: + data = json.dumps(body).encode("utf-8") if body is not None else None + request = urllib.request.Request( + f"https://api.github.com{path}", + data=data, + method=method, + headers={ + "Accept": "application/vnd.github+json", + "Authorization": f"Bearer {token}", + "X-GitHub-Api-Version": "2022-11-28", + }, + ) + try: + with urllib.request.urlopen(request, timeout=30) as response: + return json.loads(response.read().decode("utf-8")) + except urllib.error.HTTPError as error: + # GitHub's error body says exactly why (e.g. "Must have admin rights to + # Repository" for a PAT without the required scope); surface it instead + # of a bare "HTTP Error 403". + body = error.read().decode("utf-8", "replace") + raise SystemExit( + f"GitHub API {method} {path} failed: HTTP {error.code}: {body}\n" + "The token comes from the GH_PERSONAL_ACCESS_TOKEN secret; it needs " + "'repo' scope (classic PAT) or 'Administration: write' on the " + "repository (fine-grained PAT)." + ) from error + + +def create_registration_token(github_token: str, repo: str) -> str: + response = github_api( + github_token, + "POST", + f"/repos/{repo}/actions/runners/registration-token", + body={}, + ) + return response["token"] + + +def find_runner_by_name(github_token: str, repo: str, name: str) -> dict | None: + page = 1 + while True: + response = github_api( + github_token, + "GET", + f"/repos/{repo}/actions/runners?per_page=100&page={page}", + ) + runners = response.get("runners", []) + for runner in runners: + if runner.get("name") == name: + return runner + if len(runners) < 100: + return None + page += 1 + + +def make_ecs_client(config: ProvisionConfig): + from alibabacloud_ecs20140526.client import Client as EcsClient + from alibabacloud_tea_openapi.models import Config as OpenApiConfig + + access_key_id = os.environ.get("ALIBABA_CLOUD_ACCESS_KEY_ID", "") + access_key_secret = os.environ.get("ALIBABA_CLOUD_ACCESS_KEY_SECRET", "") + if not access_key_id or not access_key_secret: + raise SystemExit( + "ALIBABA_CLOUD_ACCESS_KEY_ID/SECRET are empty; in CI they come from the " + "ALICLOUD_ECS_ACCESS_KEY_ID/SECRET repository secrets (a missing or " + "misnamed secret expands to an empty string)." + ) + return EcsClient( + OpenApiConfig( + access_key_id=access_key_id, + access_key_secret=access_key_secret, + region_id=config.region_id, + endpoint=f"ecs.{config.region_id}.aliyuncs.com", + ) + ) + + +def wait_for_instance_status(client, region_id: str, instance_id: str, wanted: str, deadline: float) -> None: + from alibabacloud_ecs20140526 import models as ecs_models + + while time.monotonic() < deadline: + response = client.describe_instances( + ecs_models.DescribeInstancesRequest( + region_id=region_id, instance_ids=json.dumps([instance_id]) + ) + ) + instances = response.body.instances.instance + if instances and instances[0].status == wanted: + return + time.sleep(POLL_INTERVAL_SECONDS) + raise TimeoutError(f"Instance {instance_id} did not reach status {wanted} in time") + + +def fetch_console_output(client, region_id: str, instance_id: str) -> str: + from alibabacloud_ecs20140526 import models as ecs_models + + response = client.get_instance_console_output( + ecs_models.GetInstanceConsoleOutputRequest(region_id=region_id, instance_id=instance_id) + ) + return base64.b64decode(response.body.console_output or "").decode("utf-8", "replace") + + +class ConsoleTailer: + """Incrementally prints new serial-console lines so cloud-init progress is + visible in the CI log while waiting, not only after a failure dump.""" + + def __init__(self, client, region_id: str, instance_id: str) -> None: + self.client = client + self.region_id = region_id + self.instance_id = instance_id + self.printed_lines = 0 + + def poll(self) -> None: + try: + lines = fetch_console_output(self.client, self.region_id, self.instance_id).splitlines() + except Exception as error: # noqa: BLE001 + print(f"[console] unable to fetch console output: {error}", flush=True) + return + # The API returns a bounded tail; if it ever shrinks, restart from the + # beginning of the new buffer rather than skipping lines. + if len(lines) < self.printed_lines: + self.printed_lines = 0 + for line in lines[self.printed_lines :]: + print(f"[console] {line}", flush=True) + self.printed_lines = len(lines) + + +def dump_console_output(client, region_id: str, instance_id: str) -> None: + try: + output = fetch_console_output(client, region_id, instance_id) + except Exception as error: # noqa: BLE001 + print(f"Unable to fetch console output for {instance_id}: {error}", flush=True) + return + tail = "\n".join(output.splitlines()[-80:]) + print(f"::group::Console output tail for {instance_id}\n{tail}\n::endgroup::", flush=True) + summary_path = os.environ.get("GITHUB_STEP_SUMMARY") + if summary_path: + with open(summary_path, "a", encoding="utf-8") as summary: + summary.write(f"
ECS console output ({instance_id})\n\n```\n") + summary.write(tail) + summary.write("\n```\n
\n") + + +def append_github_output(name: str, value: str) -> None: + output_path = os.environ.get("GITHUB_OUTPUT") + if output_path: + with open(output_path, "a", encoding="utf-8") as output: + output.write(f"{name}={value}\n") + + +def run_instance(client, config: ProvisionConfig, user_data: str) -> str: + from alibabacloud_ecs20140526 import models as ecs_models + + request = ecs_models.RunInstancesRequest( + region_id=config.region_id, + image_id=config.image_id, + resource_group_id=config.resource_group_id, + instance_type=config.instance_type, + v_switch_id=config.vswitch_id, + security_group_id=config.security_group_id, + instance_name=runner_name_for_run(config.run_id), + host_name=runner_name_for_run(config.run_id), + description=f"Ephemeral query-regression runner for run {config.run_id}", + amount=1, + instance_charge_type="PostPaid", + spot_strategy="NoSpot", + internet_charge_type="PayByTraffic", + internet_max_bandwidth_out=100, + system_disk=ecs_models.RunInstancesRequestSystemDisk( + category="cloud_essd", + size=str(SYSTEM_DISK_GIB), + ), + user_data=user_data, + tag=[ + ecs_models.RunInstancesRequestTag(key=MANAGED_BY_TAG_KEY, value=MANAGED_BY_TAG_VALUE), + ecs_models.RunInstancesRequestTag(key=RUN_TAG_KEY, value=config.run_id), + ], + ) + response = client.run_instances(request) + instance_id = response.body.instance_id_sets.instance_id_set[0] + return instance_id + + +def provision(config: ProvisionConfig) -> int: + client = make_ecs_client(config) + runner_name = runner_name_for_run(config.run_id) + runner_label = runner_label_for_run(config.run_id) + + registration_token = create_registration_token(config.github_token, config.repo) + user_data = encode_user_data( + render_user_data( + runner_name, + runner_label, + registration_token, + config.repo, + config.runner_uid, + config.runner_gid, + ) + ) + + print(f"::group::Run instance {runner_name}", flush=True) + instance_id = run_instance(client, config, user_data) + print(f"Instance id: {instance_id}", flush=True) + # Teardown is the only DeleteInstance caller. Write outputs immediately + # so a later status/online failure still releases this VM. + append_github_output("label", runner_label) + append_github_output("instance_id", instance_id) + append_github_output("runner_name", runner_name) + wait_for_instance_status(client, config.region_id, instance_id, "Running", time.monotonic() + 5 * 60) + print("::endgroup::", flush=True) + + print("::group::Wait for runner online", flush=True) + deadline = time.monotonic() + RUNNER_ONLINE_TIMEOUT_SECONDS + tailer = ConsoleTailer(client, config.region_id, instance_id) + last_console_poll = 0.0 + while time.monotonic() < deadline: + runner = find_runner_by_name(config.github_token, config.repo, runner_name) + if runner is not None: + status = runner.get("status") + print(f"Runner {runner_name} status: {status}", flush=True) + if status == "online": + print("::endgroup::", flush=True) + print(f"Runner {runner_name} is online with label {runner_label}", flush=True) + return 0 + # The serial console lags the guest by a minute or so; 30s polling is + # enough to follow cloud-init without tripping API throttling. + if time.monotonic() - last_console_poll >= 30: + tailer.poll() + last_console_poll = time.monotonic() + time.sleep(POLL_INTERVAL_SECONDS) + + print("::endgroup::", flush=True) + print( + f"Runner {runner_name} did not come online in time; " + "leaving instance for the teardown job to delete", + flush=True, + ) + dump_console_output(client, config.region_id, instance_id) + return 1 + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--region-id", default=os.environ.get("ALIYUN_ECS_REGION_ID")) + parser.add_argument("--vswitch-id", default=os.environ.get("ALIYUN_ECS_VSWITCH_ID")) + parser.add_argument("--security-group-id", default=os.environ.get("ALIYUN_ECS_SECURITY_GROUP_ID")) + parser.add_argument("--image-id", default=os.environ.get("QUERY_REGRESSION_ECS_IMAGE_ID")) + parser.add_argument("--instance-type", default=os.environ.get("ALIYUN_ECS_INSTANCE_TYPE")) + parser.add_argument("--repo", default=os.environ.get("GITHUB_REPOSITORY")) + parser.add_argument("--run-id", default=os.environ.get("GITHUB_RUN_ID")) + parser.add_argument("--github-token", default=os.environ.get("GH_PERSONAL_ACCESS_TOKEN")) + parser.add_argument("--runner-uid", default=os.environ.get("QUERY_REGRESSION_RUNNER_UID", "1001")) + parser.add_argument("--runner-gid", default=os.environ.get("QUERY_REGRESSION_RUNNER_GID", "1001")) + parser.add_argument("--resource-group-id", default=os.environ.get("ALIYUN_ECS_RESOURCE_GROUP_ID")) + args = parser.parse_args() + + missing = [ + name + for name, value in vars(args).items() + if name not in ("resource_group_id",) + and (value is None or (isinstance(value, str) and not value)) + ] + if missing: + flags = ", ".join(f"--{name.replace('_', '-')}" for name in missing) + raise SystemExit(f"Missing required configuration: {flags} (or their env defaults)") + + config = ProvisionConfig( + region_id=args.region_id, + vswitch_id=args.vswitch_id, + security_group_id=args.security_group_id, + image_id=args.image_id, + instance_type=args.instance_type, + repo=args.repo, + run_id=args.run_id, + github_token=args.github_token, + resource_group_id=args.resource_group_id or None, + runner_uid=args.runner_uid, + runner_gid=args.runner_gid, + ) + return provision(config) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/.github/scripts/aliyun-ecs-runner-teardown.py b/.github/scripts/aliyun-ecs-runner-teardown.py new file mode 100644 index 00000000000..36c80253b95 --- /dev/null +++ b/.github/scripts/aliyun-ecs-runner-teardown.py @@ -0,0 +1,237 @@ +#!/usr/bin/env python3 +# Copyright 2023 Greptime Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# PEP 723 inline metadata (see aliyun-ecs-runner-provision.py for the +# convention). The SDK import stays lazy so unit tests run on a plain +# stdlib interpreter. +# /// script +# requires-python = ">=3.10" +# dependencies = [ +# "alibabacloud_ecs20140526>=4.1.0,<6", +# "alibabacloud_tea_openapi>=0.3.12,<1", +# ] +# /// + +"""Tear down ephemeral Aliyun ECS query-regression runners. + +Two modes: + +- Targeted (per workflow run): delete one instance by id and deregister its + runner by name. Both steps are idempotent and best-effort; a missing + instance or runner is not an error. +- Sweep (scheduled janitor): delete every instance tagged as managed by + query-regression CI whose creation time is older than the given TTL, and + deregister the matching runners. This is the safety net for runs whose + teardown job never executed. +""" + +from __future__ import annotations + +import argparse +import importlib.util +import os +import sys +from datetime import datetime, timedelta, timezone +from pathlib import Path + +# Reuse the GitHub API client, runner lookup, and tag constants from the +# provision script. +_PROVISION_SPEC = importlib.util.spec_from_file_location( + "aliyun_ecs_runner_provision", + Path(__file__).resolve().parent / "aliyun-ecs-runner-provision.py", +) +assert _PROVISION_SPEC is not None and _PROVISION_SPEC.loader is not None +provision = importlib.util.module_from_spec(_PROVISION_SPEC) +sys.modules[_PROVISION_SPEC.name] = provision +_PROVISION_SPEC.loader.exec_module(provision) + +SWEEP_TTL = timedelta(hours=4) +# ECS DescribeInstances creation_time is ISO8601 UTC, e.g. "2026-08-17T02:13Z". +CREATION_TIME_FORMATS = ("%Y-%m-%dT%H:%M:%SZ", "%Y-%m-%dT%H:%MZ") + + +def parse_creation_time(value: str) -> datetime: + for fmt in CREATION_TIME_FORMATS: + try: + return datetime.strptime(value, fmt).replace(tzinfo=timezone.utc) + except ValueError: + continue + raise ValueError(f"Unparsable ECS creation time: {value}") + + +def expired_instance_names( + instances: list[tuple[str, str, str]], now: datetime, ttl: timedelta +) -> list[tuple[str, str]]: + """Pick (instance_id, instance_name) pairs whose creation is older than ttl. + + `instances` items are (instance_id, instance_name, creation_time). + """ + expired = [] + for instance_id, instance_name, creation_time in instances: + age = now - parse_creation_time(creation_time) + if age >= ttl: + expired.append((instance_id, instance_name)) + return expired + + +def make_ecs_client(region_id: str): + from alibabacloud_ecs20140526.client import Client as EcsClient + from alibabacloud_tea_openapi.models import Config as OpenApiConfig + + access_key_id = os.environ.get("ALIBABA_CLOUD_ACCESS_KEY_ID", "") + access_key_secret = os.environ.get("ALIBABA_CLOUD_ACCESS_KEY_SECRET", "") + if not access_key_id or not access_key_secret: + raise SystemExit( + "ALIBABA_CLOUD_ACCESS_KEY_ID/SECRET are empty; in CI they come from the " + "ALICLOUD_ECS_ACCESS_KEY_ID/SECRET repository secrets (a missing or " + "misnamed secret expands to an empty string)." + ) + return EcsClient( + OpenApiConfig( + access_key_id=access_key_id, + access_key_secret=access_key_secret, + region_id=region_id, + endpoint=f"ecs.{region_id}.aliyuncs.com", + ) + ) + + +def delete_instance(client, instance_id: str, region_id: str | None = None) -> bool: + import time + + from alibabacloud_ecs20140526 import models as ecs_models + + # Export the serial console before destroying the evidence: a cancelled + # run still runs this teardown job, and the console tail (cloud-init, + # runner service logs, kernel OOM records) is the only witness of what + # happened on the machine. + if region_id is not None: + provision.dump_console_output(client, region_id, instance_id) + + # DeleteInstance rejects instances that are still Initializing (or in + # another transitional status). Teardown runs right after a cancelled + # run, when the instance may be only seconds old, so retry the transient + # status errors for a few minutes instead of leaking the instance. + deadline = time.monotonic() + 5 * 60 + while True: + try: + client.delete_instance(ecs_models.DeleteInstanceRequest(instance_id=instance_id, force=True)) + print(f"Deleted instance {instance_id}", flush=True) + return True + except Exception as error: # noqa: BLE001 + message = str(error) + if "InvalidInstanceId.NotFound" in message: + print(f"Instance {instance_id} already gone", flush=True) + return True + if "IncorrectInstanceStatus" in message and time.monotonic() < deadline: + print(f"Instance {instance_id} is in a transitional status; retrying delete", flush=True) + time.sleep(15) + continue + print(f"Failed to delete instance {instance_id}: {error}", flush=True) + return False + + +def deregister_runner(token: str, repo: str, runner_name: str) -> bool: + runner = provision.find_runner_by_name(token, repo, runner_name) + if runner is None: + print(f"Runner {runner_name} is not registered", flush=True) + return True + try: + provision.github_api(token, "DELETE", f"/repos/{repo}/actions/runners/{runner['id']}") + print(f"Deregistered runner {runner_name} (id {runner['id']})", flush=True) + return True + except Exception as error: # noqa: BLE001 + print(f"Failed to deregister runner {runner_name}: {error}", flush=True) + return False + + +def list_managed_instances(client, region_id: str) -> list[tuple[str, str, str]]: + from alibabacloud_ecs20140526 import models as ecs_models + + result: list[tuple[str, str, str]] = [] + next_token = None + while True: + request = ecs_models.DescribeInstancesRequest( + region_id=region_id, + tag=[ + ecs_models.DescribeInstancesRequestTag( + key=provision.MANAGED_BY_TAG_KEY, value=provision.MANAGED_BY_TAG_VALUE + ) + ], + max_results=100, + next_token=next_token, + ) + response = client.describe_instances(request) + for instance in response.body.instances.instance: + result.append((instance.instance_id, instance.instance_name, instance.creation_time)) + next_token = response.body.next_token + if not next_token: + return result + + +def sweep(client, region_id: str, repo: str, github_token: str, ttl: timedelta) -> int: + instances = list_managed_instances(client, region_id) + print(f"Found {len(instances)} managed instance(s) in {region_id}", flush=True) + expired = expired_instance_names(instances, datetime.now(timezone.utc), ttl) + ok = True + for instance_id, instance_name in expired: + print(f"Instance {instance_id} ({instance_name}) exceeds TTL {ttl}; deleting", flush=True) + # Runner names mirror instance names by construction in the provision + # script (both are qreg-ecs-). + ok &= delete_instance(client, instance_id, region_id) + ok &= deregister_runner(github_token, repo, instance_name) + return 0 if ok else 1 + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--region-id", default=os.environ.get("ALIYUN_ECS_REGION_ID")) + parser.add_argument("--repo", default=os.environ.get("GITHUB_REPOSITORY")) + parser.add_argument("--github-token", default=os.environ.get("GH_PERSONAL_ACCESS_TOKEN")) + parser.add_argument("--instance-id", default=os.environ.get("QUERY_REGRESSION_ECS_INSTANCE_ID")) + parser.add_argument("--runner-name", default=os.environ.get("QUERY_REGRESSION_ECS_RUNNER_NAME")) + parser.add_argument( + "--sweep", + action="store_true", + help="Janitor mode: delete all managed instances older than --sweep-ttl-hours.", + ) + parser.add_argument("--sweep-ttl-hours", type=float, default=SWEEP_TTL.total_seconds() / 3600) + args = parser.parse_args() + + for name in ("region_id", "repo", "github_token"): + if not getattr(args, name): + raise SystemExit(f"Missing required configuration: --{name.replace('_', '-')}") + + client = make_ecs_client(args.region_id) + + if args.sweep: + return sweep( + client, args.region_id, args.repo, args.github_token, timedelta(hours=args.sweep_ttl_hours) + ) + + ok = True + if args.instance_id: + ok &= delete_instance(client, args.instance_id, args.region_id) + else: + print("No instance id given; skipping instance deletion", flush=True) + if args.runner_name: + ok &= deregister_runner(args.github_token, args.repo, args.runner_name) + else: + print("No runner name given; skipping runner deregistration", flush=True) + return 0 if ok else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/.github/scripts/query-regression-nightly-refs.py b/.github/scripts/query-regression-nightly-refs.py new file mode 100644 index 00000000000..6e9f54682d4 --- /dev/null +++ b/.github/scripts/query-regression-nightly-refs.py @@ -0,0 +1,270 @@ +#!/usr/bin/env python3 +# Copyright 2023 Greptime Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Resolve consecutive successful Nightly Build SHAs for query-regression. + +Query regression builds both binaries from git; this script only picks the +head SHAs of two Nightly Build workflow runs (today vs the previous success). +""" + +from __future__ import annotations + +import argparse +import json +import os +import sys +import urllib.error +import urllib.parse +import urllib.request +from dataclasses import dataclass +from typing import Any + + +NIGHTLY_WORKFLOW = "nightly-build.yml" +NIGHTLY_EVENTS = frozenset({"schedule", "workflow_dispatch"}) + + +@dataclass(frozen=True) +class WorkflowRun: + id: int + head_sha: str + head_branch: str + html_url: str + created_at: str + conclusion: str + event: str + + +@dataclass(frozen=True) +class RefPair: + base: WorkflowRun | None + candidate: WorkflowRun | None + skip: bool + reason: str + + +def parse_run(payload: dict[str, Any]) -> WorkflowRun: + return WorkflowRun( + id=int(payload["id"]), + head_sha=str(payload["head_sha"]), + head_branch=str(payload.get("head_branch") or ""), + html_url=str(payload.get("html_url") or ""), + created_at=str(payload.get("created_at") or ""), + conclusion=str(payload.get("conclusion") or ""), + event=str(payload.get("event") or ""), + ) + + +def github_get(token: str, path: str, query: dict[str, str] | None = None) -> dict[str, Any]: + encoded = f"?{urllib.parse.urlencode(query)}" if query else "" + request = urllib.request.Request( + f"https://api.github.com{path}{encoded}", + method="GET", + headers={ + "Accept": "application/vnd.github+json", + "Authorization": f"Bearer {token}", + "X-GitHub-Api-Version": "2022-11-28", + }, + ) + try: + with urllib.request.urlopen(request, timeout=30) as response: + return json.loads(response.read().decode("utf-8")) + except urllib.error.HTTPError as error: + body = error.read().decode("utf-8", "replace") + raise SystemExit( + f"GitHub API GET {path} failed: HTTP {error.code}: {body}" + ) from error + + +def list_successful_nightly_runs( + token: str, + repo: str, + *, + branch: str | None = None, + per_page: int = 30, +) -> list[WorkflowRun]: + query = { + "status": "success", + "per_page": str(per_page), + } + if branch: + query["branch"] = branch + payload = github_get( + token, + f"/repos/{repo}/actions/workflows/{NIGHTLY_WORKFLOW}/runs", + query, + ) + runs = [parse_run(item) for item in payload.get("workflow_runs") or []] + return [ + run + for run in runs + if run.conclusion == "success" and run.event in NIGHTLY_EVENTS and run.head_sha + ] + + +def fetch_run(token: str, repo: str, run_id: int) -> WorkflowRun: + payload = github_get(token, f"/repos/{repo}/actions/runs/{run_id}") + return parse_run(payload) + + +def select_base_and_candidate( + runs: list[WorkflowRun], + *, + candidate_run_id: int | None = None, + candidate: WorkflowRun | None = None, +) -> RefPair: + """Pick candidate (newest / requested) and the previous successful run. + + `runs` must be newest-first. Same-SHA consecutive nightlies are skipped + (no commits that day). + """ + if candidate is None: + if candidate_run_id is not None: + candidate = next((run for run in runs if run.id == candidate_run_id), None) + if candidate is None: + return RefPair( + None, + None, + True, + f"candidate nightly run {candidate_run_id} was not a successful " + f"{NIGHTLY_WORKFLOW} run", + ) + elif runs: + candidate = runs[0] + else: + return RefPair(None, None, True, "no successful Nightly Build runs") + + branch = candidate.head_branch + older = [ + run + for run in runs + if run.id != candidate.id + and (not branch or not run.head_branch or run.head_branch == branch) + and (not candidate.created_at or run.created_at <= candidate.created_at) + ] + if not older: + return RefPair( + None, + candidate, + True, + "no previous successful Nightly Build to compare against", + ) + base = older[0] + if base.head_sha.lower() == candidate.head_sha.lower(): + return RefPair( + base, + candidate, + True, + f"previous nightly SHA {base.head_sha} matches candidate; nothing to compare", + ) + return RefPair(base, candidate, False, "") + + +def _override_run(sha: str) -> WorkflowRun: + return WorkflowRun( + id=0, + head_sha=sha, + head_branch="", + html_url="", + created_at="", + conclusion="success", + event="workflow_dispatch", + ) + + +def override_pair(base_ref: str, candidate_ref: str) -> RefPair: + return RefPair(_override_run(base_ref), _override_run(candidate_ref), False, "explicit refs") + + +def write_outputs(pair: RefPair) -> None: + values = { + "skip": "true" if pair.skip else "false", + "reason": pair.reason, + "base_sha": pair.base.head_sha if pair.base else "", + "candidate_sha": pair.candidate.head_sha if pair.candidate else "", + "base_run_url": pair.base.html_url if pair.base else "", + "candidate_run_url": pair.candidate.html_url if pair.candidate else "", + "base_run_id": str(pair.base.id) if pair.base and pair.base.id else "", + "candidate_run_id": str(pair.candidate.id) if pair.candidate and pair.candidate.id else "", + } + output_path = os.environ.get("GITHUB_OUTPUT") + if output_path: + with open(output_path, "a", encoding="utf-8") as handle: + for key, value in values.items(): + handle.write(f"{key}={value}\n") + for key, value in values.items(): + print(f"{key}={value}") + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--repo", default=os.environ.get("GITHUB_REPOSITORY")) + parser.add_argument( + "--token", + default=os.environ.get("GITHUB_TOKEN") or os.environ.get("GH_TOKEN"), + ) + parser.add_argument("--branch", default=os.environ.get("NIGHTLY_BRANCH") or "") + parser.add_argument( + "--candidate-run-id", + default=os.environ.get("CANDIDATE_RUN_ID") or "", + ) + parser.add_argument("--base-ref", default=os.environ.get("BASE_REF") or "") + parser.add_argument("--candidate-ref", default=os.environ.get("CANDIDATE_REF") or "") + return parser.parse_args(argv) + + +def main(argv: list[str] | None = None) -> int: + args = parse_args(argv) + base_ref = args.base_ref.strip() + candidate_ref = args.candidate_ref.strip() + if base_ref and candidate_ref: + write_outputs(override_pair(base_ref, candidate_ref)) + return 0 + if bool(base_ref) != bool(candidate_ref): + print("Provide both --base-ref and --candidate-ref, or neither.", file=sys.stderr) + return 2 + if not args.repo or not args.token: + print("--repo and --token (or GITHUB_REPOSITORY / GITHUB_TOKEN) are required.", file=sys.stderr) + return 2 + + candidate_run_id = int(args.candidate_run_id) if str(args.candidate_run_id).strip() else None + candidate: WorkflowRun | None = None + branch = args.branch.strip() or None + if candidate_run_id is not None: + candidate = fetch_run(args.token, args.repo, candidate_run_id) + if candidate.conclusion != "success": + write_outputs( + RefPair( + None, + candidate, + True, + f"nightly run {candidate_run_id} conclusion is {candidate.conclusion}", + ) + ) + return 0 + branch = branch or candidate.head_branch or None + + runs = list_successful_nightly_runs(args.token, args.repo, branch=branch) + pair = select_base_and_candidate( + runs, + candidate_run_id=candidate_run_id if candidate is None else None, + candidate=candidate, + ) + write_outputs(pair) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/.github/workflows/checks.yml b/.github/workflows/checks.yml index 7f064f52eb5..8d7a37080b3 100644 --- a/.github/workflows/checks.yml +++ b/.github/workflows/checks.yml @@ -80,6 +80,14 @@ jobs: run: python3 scripts/check-enterprise-license-test.py - name: Run fuzz orchestration script tests run: .github/scripts/run-fuzz-targets-test.sh + - name: Run query-regression Python tooling tests + run: | + python3 tests/perf/test_query_regression_runner_compaction_toctou.py + python3 tests/perf/test_query_regression_runner_otlp_trace_load.py + python3 tests/perf/test_query_regression_summary_otlp.py + python3 tests/perf/test_query_regression_case_selection.py + python3 tests/perf/test_query_regression_nightly_refs.py + python3 tests/perf/test_aliyun_ecs_runner_scripts.py check: if: ${{ github.repository == 'GreptimeTeam/greptimedb' }} diff --git a/.github/workflows/query-regression-janitor.yml b/.github/workflows/query-regression-janitor.yml new file mode 100644 index 00000000000..ce277ce2f02 --- /dev/null +++ b/.github/workflows/query-regression-janitor.yml @@ -0,0 +1,36 @@ +name: Query Regression ECS Janitor + +# Safety net for the aliyun-ecs query-regression path: deletes managed ECS +# instances (and their runner registrations) that outlived their workflow run, +# so a failed teardown never leaks pay-as-you-go spend. + +on: + schedule: + - cron: "23 3 * * *" + workflow_dispatch: + +permissions: + contents: read + +jobs: + sweep: + runs-on: ubuntu-latest + timeout-minutes: 15 + steps: + - name: Checkout teardown script + uses: actions/checkout@v4 + with: + persist-credentials: false + + - name: Install uv + uses: astral-sh/setup-uv@v6 + + - name: Sweep expired ECS runners + env: + ALIBABA_CLOUD_ACCESS_KEY_ID: ${{ secrets.ALICLOUD_ECS_ACCESS_KEY_ID }} + ALIBABA_CLOUD_ACCESS_KEY_SECRET: ${{ secrets.ALICLOUD_ECS_ACCESS_KEY_SECRET }} + GH_PERSONAL_ACCESS_TOKEN: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }} + ALIYUN_ECS_REGION_ID: ${{ vars.ALIYUN_ECS_REGION_ID }} + run: >- + uv run .github/scripts/aliyun-ecs-runner-teardown.py + --sweep --sweep-ttl-hours 4 diff --git a/.github/workflows/query-regression-nightly.yml b/.github/workflows/query-regression-nightly.yml new file mode 100644 index 00000000000..027fcd16987 --- /dev/null +++ b/.github/workflows/query-regression-nightly.yml @@ -0,0 +1,104 @@ +name: Query Regression Nightly + +# After a successful GreptimeDB Nightly Build, compare that commit against +# the previous successful nightly. The reusable Query Regression workflow +# still compiles both SHAs; this wrapper only resolves which two SHAs. + +on: + workflow_run: + workflows: + - GreptimeDB Nightly Build + types: + - completed + workflow_dispatch: + inputs: + base_ref: + description: Base ref/SHA (empty = previous successful nightly) + required: false + type: string + default: "" + candidate_ref: + description: Candidate ref/SHA (empty = latest successful nightly) + required: false + type: string + default: "" + candidate_run_id: + description: Nightly Build run id to treat as candidate (empty = latest) + required: false + type: string + default: "" + case: + description: Query perf case path(s), all, or heavy + required: false + type: string + default: all + +permissions: + contents: read + actions: read + +jobs: + resolve-refs: + name: Resolve previous vs current nightly SHAs + if: ${{ github.event_name == 'workflow_dispatch' || github.event.workflow_run.conclusion == 'success' }} + runs-on: ubuntu-latest + timeout-minutes: 5 + outputs: + skip: ${{ steps.refs.outputs.skip }} + reason: ${{ steps.refs.outputs.reason }} + base_sha: ${{ steps.refs.outputs.base_sha }} + candidate_sha: ${{ steps.refs.outputs.candidate_sha }} + steps: + - name: Checkout ref resolver + uses: actions/checkout@v4 + with: + persist-credentials: false + + - name: Resolve nightly SHAs + id: refs + env: + GITHUB_TOKEN: ${{ github.token }} + GITHUB_REPOSITORY: ${{ github.repository }} + CANDIDATE_RUN_ID: ${{ github.event_name == 'workflow_run' && github.event.workflow_run.id || inputs.candidate_run_id || '' }} + BASE_REF: ${{ github.event_name == 'workflow_dispatch' && inputs.base_ref || '' }} + CANDIDATE_REF: ${{ github.event_name == 'workflow_dispatch' && inputs.candidate_ref || '' }} + run: python3 .github/scripts/query-regression-nightly-refs.py + + - name: Summarize comparison + if: always() + env: + SKIP: ${{ steps.refs.outputs.skip }} + REASON: ${{ steps.refs.outputs.reason }} + BASE_SHA: ${{ steps.refs.outputs.base_sha }} + CANDIDATE_SHA: ${{ steps.refs.outputs.candidate_sha }} + BASE_RUN_URL: ${{ steps.refs.outputs.base_run_url }} + CANDIDATE_RUN_URL: ${{ steps.refs.outputs.candidate_run_url }} + run: | + set -euo pipefail + { + if [[ "${SKIP}" == "true" ]]; then + printf 'Skipping query-regression nightly: %s\n' "${REASON}" + else + printf 'Comparing previous nightly `%s` -> current nightly `%s`\n' \ + "${BASE_SHA}" "${CANDIDATE_SHA}" + if [[ -n "${BASE_RUN_URL}" ]]; then + printf -- '- Previous Nightly Build: %s\n' "${BASE_RUN_URL}" + fi + if [[ -n "${CANDIDATE_RUN_URL}" ]]; then + printf -- '- Current Nightly Build: %s\n' "${CANDIDATE_RUN_URL}" + fi + fi + } | tee -a "${GITHUB_STEP_SUMMARY}" + + query-regression: + name: Query regression nightly + needs: [resolve-refs] + if: ${{ needs.resolve-refs.outputs.skip != 'true' && needs.resolve-refs.outputs.base_sha != '' && needs.resolve-refs.outputs.candidate_sha != '' }} + uses: ./.github/workflows/query-regression.yml + secrets: inherit + with: + case: ${{ github.event_name == 'workflow_dispatch' && inputs.case || 'all' }} + base_ref: ${{ needs.resolve-refs.outputs.base_sha }} + candidate_ref: ${{ needs.resolve-refs.outputs.candidate_sha }} + cargo_profile: nightly + runner: aliyun-ecs diff --git a/.github/workflows/query-regression.yml b/.github/workflows/query-regression.yml index 0806d4d3e38..209c8d76fff 100644 --- a/.github/workflows/query-regression.yml +++ b/.github/workflows/query-regression.yml @@ -32,10 +32,10 @@ on: type: string default: nightly runner: - description: Self-hosted runner label or ARC runner scale set for this query regression run + description: Self-hosted runner label; aliyun-ecs provisions a fresh ECS instance per run required: false type: string - default: perf-regression-8-cores + default: aliyun-ecs workflow_dispatch: inputs: case: @@ -68,12 +68,20 @@ on: - release - dev runner: - description: Self-hosted runner label or ARC runner scale set for this query regression run + description: >- + Self-hosted runner label; aliyun-ecs provisions a fresh ECS instance + per run, any other value is used as a literal runner label required: true - type: choice - default: perf-regression-8-cores - options: - - perf-regression-8-cores + type: string + default: aliyun-ecs + keep_instance: + description: >- + Debug: keep the ECS instance after the run (skip teardown) so its + runner _diag logs, journal, and telemetry can be inspected; the + janitor still sweeps it after the TTL + required: false + type: boolean + default: false pull_request: types: [labeled] @@ -81,19 +89,99 @@ permissions: contents: read jobs: - query-regression: + test-tooling: + # Stdlib unittests for the workflow Python (case selection, report + # helpers, nightly SHA picking, and rendered ECS user-data). They do + # not talk to Aliyun or the Actions runner process; running them on + # ubuntu-latest fails fast before any ECS spend. Not gated on the + # regression labels: those labels boot a VM, these tests should not. + # Ordinary PRs also run the same tests from checks.yml, because this + # workflow only starts on `labeled` (or dispatch / workflow_call). + runs-on: ubuntu-latest + timeout-minutes: 10 + steps: + - name: Checkout + uses: actions/checkout@v4 + with: + persist-credentials: false + + - name: Test query regression tooling + run: | + python3 tests/perf/test_query_regression_runner_compaction_toctou.py + python3 tests/perf/test_query_regression_runner_otlp_trace_load.py + python3 tests/perf/test_query_regression_summary_otlp.py + python3 tests/perf/test_query_regression_case_selection.py + python3 tests/perf/test_query_regression_nightly_refs.py + python3 tests/perf/test_aliyun_ecs_runner_scripts.py + + provision: + # Runs when the aliyun-ecs path is selected: explicitly via the runner + # input, or for PR labels by default (a QUERY_REGRESSION_PR_RUNNER + # repository variable set to another value redirects PRs to that literal + # runner label instead). Uses trusted scripts from the PR base (or the + # dispatched ref), never from candidate code. Waits for test-tooling so + # a broken user-data template does not still create a VM. + needs: [test-tooling] if: >- - ${{ github.event_name != 'pull_request' || - (github.event_name == 'pull_request' && - !github.event.pull_request.draft && - (github.event.label.name == 'query-regression' || - github.event.label.name == 'heavy-regression')) }} - runs-on: ${{ github.event_name != 'pull_request' && inputs.runner || 'perf-regression-8-cores' }} + ${{ !failure() && !cancelled() && + ((github.event_name != 'pull_request' && inputs.runner == 'aliyun-ecs') || + (github.event_name == 'pull_request' && + !github.event.pull_request.draft && + (github.event.label.name == 'query-regression' || + github.event.label.name == 'heavy-regression') && + (vars.QUERY_REGRESSION_PR_RUNNER || 'aliyun-ecs') == 'aliyun-ecs')) }} + runs-on: ubuntu-latest + timeout-minutes: 45 + outputs: + label: ${{ steps.provision.outputs.label }} + instance_id: ${{ steps.provision.outputs.instance_id }} + runner_name: ${{ steps.provision.outputs.runner_name }} + steps: + - name: Checkout trusted provisioning scripts + uses: actions/checkout@v4 + with: + repository: ${{ github.repository }} + ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.base.sha || github.sha }} + persist-credentials: false + + - name: Install uv + uses: astral-sh/setup-uv@v6 + + - name: Provision ECS runner + id: provision + env: + ALIBABA_CLOUD_ACCESS_KEY_ID: ${{ secrets.ALICLOUD_ECS_ACCESS_KEY_ID }} + ALIBABA_CLOUD_ACCESS_KEY_SECRET: ${{ secrets.ALICLOUD_ECS_ACCESS_KEY_SECRET }} + GH_PERSONAL_ACCESS_TOKEN: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }} + ALIYUN_ECS_REGION_ID: ${{ vars.ALIYUN_ECS_REGION_ID }} + ALIYUN_ECS_VSWITCH_ID: ${{ vars.ALIYUN_ECS_VSWITCH_ID }} + ALIYUN_ECS_SECURITY_GROUP_ID: ${{ vars.ALIYUN_ECS_SECURITY_GROUP_ID }} + ALIYUN_ECS_INSTANCE_TYPE: ${{ vars.ALIYUN_ECS_INSTANCE_TYPE }} + ALIYUN_ECS_RESOURCE_GROUP_ID: ${{ vars.ALIYUN_ECS_RESOURCE_GROUP_ID }} + QUERY_REGRESSION_ECS_IMAGE_ID: ${{ vars.QUERY_REGRESSION_ECS_IMAGE_ID }} + QUERY_REGRESSION_RUNNER_UID: ${{ vars.QUERY_REGRESSION_RUNNER_UID || '1001' }} + QUERY_REGRESSION_RUNNER_GID: ${{ vars.QUERY_REGRESSION_RUNNER_GID || '1001' }} + run: >- + uv run .github/scripts/aliyun-ecs-runner-provision.py + + query-regression: + needs: [provision, test-tooling] + # `!failure() && !cancelled()` both suppresses the implicit success() and + # lets the job run when provision was intentionally skipped because a + # literal runner label was selected; a failed or cancelled provision still + # blocks the run because no ECS runner would be waiting. + if: >- + ${{ !failure() && !cancelled() && + (github.event_name != 'pull_request' || + (github.event_name == 'pull_request' && + !github.event.pull_request.draft && + (github.event.label.name == 'query-regression' || + github.event.label.name == 'heavy-regression'))) }} + runs-on: >- + ${{ needs.provision.outputs.label || + (github.event_name != 'pull_request' && inputs.runner || + (vars.QUERY_REGRESSION_PR_RUNNER || 'aliyun-ecs')) }} timeout-minutes: 180 - concurrency: - group: query-regression-persistent-cache-v1 - queue: max - cancel-in-progress: false env: CARGO_PROFILE: ${{ github.event_name == 'pull_request' && 'nightly' || inputs.cargo_profile }} CARGO_HOME: /home/runner/.cargo @@ -105,7 +193,8 @@ jobs: QUERY_REGRESSION_CACHE_META: /home/runner/query-regression-cache-meta RUSTC_WRAPPER: /usr/local/bin/sccache SCCACHE_DIR: /home/runner/.cache/sccache - SCCACHE_CACHE_SIZE: 40G + # Caps the local sccache on the system disk next to target and cargo. + SCCACHE_CACHE_SIZE: 10G CARGO_INCREMENTAL: "0" RUSTFLAGS: -D warnings -C link-arg=-fuse-ld=mold QUERY_REGRESSION_CACHE_EPOCH: "1" @@ -114,7 +203,55 @@ jobs: EVENT_MERGE_SHA: ${{ github.event_name == 'pull_request' && github.sha || '' }} EVENT_HEAD_SHA: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || '' }} EVENT_BASE_SHA: ${{ github.event_name == 'pull_request' && github.event.pull_request.base.sha || '' }} + # Runner identity contract. The ECS image uses 1001; a manually prepared + # host may override via repo variables when 1001 is already taken. + EXPECTED_RUNNER_UID: ${{ vars.QUERY_REGRESSION_RUNNER_UID || '1001' }} + EXPECTED_RUNNER_GID: ${{ vars.QUERY_REGRESSION_RUNNER_GID || '1001' }} steps: + - name: Report provisioned ECS runner + if: ${{ needs.provision.outputs.instance_id != '' }} + shell: bash + run: | + set -euo pipefail + { + printf -- '- ECS instance: `%s` (type `%s`, image `%s`, region `%s`)\n' \ + "${{ needs.provision.outputs.instance_id }}" \ + "${{ vars.ALIYUN_ECS_INSTANCE_TYPE }}" \ + "${{ vars.QUERY_REGRESSION_ECS_IMAGE_ID }}" \ + "${{ vars.ALIYUN_ECS_REGION_ID }}" + } >> "${GITHUB_STEP_SUMMARY}" + + - name: Start machine telemetry sampler + shell: bash + # The ECS instance is deleted when the run ends, so machine state must + # be captured while the job runs. A background sampler appends load, + # memory, disk, and top-memory processes to a log every 30s; the + # "Dump machine telemetry" step (if: always()) prints and uploads it. + # Note: when the job is *cancelled*, even always() steps are killed — + # telemetry covers failures, not cancellations. + run: | + set -euo pipefail + log="${GITHUB_WORKSPACE}/machine-telemetry.log" + { + echo "== baseline $(date -u +%Y-%m-%dT%H:%M:%SZ) ==" + nproc + free -h + df -h / /home/runner + } >> "${log}" + nohup bash -c " + while true; do + { + date -u '+== %Y-%m-%dT%H:%M:%SZ ==' + uptime + free -m + df -h / /home/runner + ps -eo pid,comm,%mem,%cpu --sort=-%mem | head -6 + } >> '${log}' 2>&1 + sleep 30 + done + " >/dev/null 2>&1 & + echo "Telemetry sampler started (pid $!), logging to ${log}" + - name: Checkout base source uses: actions/checkout@v4 with: @@ -263,14 +400,64 @@ jobs: shell: bash run: | set -euo pipefail - [[ "$(id -u)" == "1001" ]] - [[ "$(id -g)" == "1001" ]] - [[ "${UV_CACHE_DIR}" == "/home/runner/.cargo/uv-cache" ]] - [[ "$(protoc --version)" == "libprotoc 3.21.12" ]] + errors=() + record() { + printf 'FAIL %s\n' "$*" + errors+=("$*") + } + require_eq() { + local name="$1" actual="$2" expected="$3" + printf 'check %s: %s\n' "${name}" "${actual}" + [[ "${actual}" == "${expected}" ]] || record "${name}: expected '${expected}', got '${actual}'" + } + require_match() { + local name="$1" actual="$2" pattern="$3" + printf 'check %s: %s\n' "${name}" "${actual}" + [[ "${actual}" =~ $pattern ]] || record "${name}: expected to match ${pattern}, got '${actual}'" + } + require() { + local name="$1" + shift + printf 'check %s\n' "${name}" + "$@" || record "${name}" + } + capture() { + local out + if out="$("$@" 2>&1)"; then + printf '%s' "${out}" + else + printf '' "$*" + fi + } + + printf 'uid=%s gid=%s PATH=%s CARGO_HOME=%s\n' "$(id -u)" "$(id -g)" "${PATH}" "${CARGO_HOME}" + ls -la "${CARGO_HOME}" 2>&1 || printf '(CARGO_HOME missing)\n' + + # Image/host hygiene: CARGO_HOME must be empty of config and + # install-state before cargo/rustup run. Those tools create + # .package-cache / bin on first use, so this cannot come after. + cargo_home_dirty=() + for entry in config config.toml credentials credentials.toml bin .crates.toml .crates2.json .global-cache .package-cache; do + if [[ -e "${CARGO_HOME}/${entry}" ]]; then + cargo_home_dirty+=("${entry}") + record "CARGO_HOME must not contain ${entry}" + fi + done + if (( ${#cargo_home_dirty[@]} > 0 )); then + printf 'CARGO_HOME contents:\n' >&2 + ls -la "${CARGO_HOME}" >&2 || true + fi + + require_eq uid "$(id -u)" "${EXPECTED_RUNNER_UID}" + require_eq gid "$(id -g)" "${EXPECTED_RUNNER_GID}" + require_eq UV_CACHE_DIR "${UV_CACHE_DIR}" "/home/runner/.cargo/uv-cache" + require_eq protoc "$(capture protoc --version)" "libprotoc 3.21.12" temporary_proto_dir="$(mktemp --directory)" trap 'rm -rf "${temporary_proto_dir}"' EXIT - test -r /usr/include/google/protobuf/any.proto - test -r /usr/include/google/protobuf/empty.proto + require "readable /usr/include/google/protobuf/any.proto" \ + test -r /usr/include/google/protobuf/any.proto + require "readable /usr/include/google/protobuf/empty.proto" \ + test -r /usr/include/google/protobuf/empty.proto printf '%s\n' \ 'syntax = "proto3";' \ 'package smoke;' \ @@ -278,47 +465,51 @@ jobs: 'import "google/protobuf/empty.proto";' \ 'message Smoke { google.protobuf.Any any = 1; google.protobuf.Empty empty = 2; }' \ > "${temporary_proto_dir}/smoke.proto" - protoc --proto_path="${temporary_proto_dir}" --proto_path=/usr/include \ + if protoc --proto_path="${temporary_proto_dir}" --proto_path=/usr/include \ --descriptor_set_out="${temporary_proto_dir}/smoke.pb" \ - "${temporary_proto_dir}/smoke.proto" - test -s "${temporary_proto_dir}/smoke.pb" - [[ "$(uv --version)" =~ ^uv[[:space:]]0\.11\.26([[:space:]]|$) ]] - mold_version="$(mold --version)" - [[ "${mold_version}" =~ ^mold[[:space:]]2\.30\.0([[:space:]]|$) ]] - [[ "$(python3 --version)" == "Python 3.12.3" ]] + "${temporary_proto_dir}/smoke.proto"; then + require "protoc smoke descriptor is non-empty" test -s "${temporary_proto_dir}/smoke.pb" + else + record "protoc smoke compile failed" + fi + require_match uv "$(capture uv --version)" '^uv[[:space:]]0\.11\.26([[:space:]]|$)' + require_match mold "$(capture mold --version)" '^mold[[:space:]]2\.40\.4([[:space:]]|$)' + require_eq python3 "$(capture python3 --version)" "Python 3.14.4" otelgen_path="$(command -v otelgen || true)" otelgen_version="" if [[ -n "${otelgen_path}" ]]; then - otelgen_version="$("${otelgen_path}" --version 2>&1 || true)" + otelgen_version="$(capture "${otelgen_path}" --version)" fi - printf 'otelgen path: %s\n' "${otelgen_path:-}" - printf 'otelgen version: %s\n' "${otelgen_version}" - [[ "${otelgen_path}" == "/usr/local/bin/otelgen" ]] - [[ "${otelgen_version}" == *"863a3f395d062c7322cc1de08a38774b7fdaa6c8"* ]] - sccache_version="$(sccache --version)" - [[ "${sccache_version}" =~ ^sccache[[:space:]]0\.16\.0([[:space:]]|$) ]] - [[ "$(command -v rustup)" == "/opt/cargo/bin/rustup" ]] - [[ "$(command -v cargo)" == "/opt/cargo/bin/cargo" ]] - [[ "$(command -v rustc)" == "/opt/cargo/bin/rustc" ]] - [[ "$(rustup --version)" =~ ^rustup[[:space:]]1\.29\.0([[:space:]]|$) ]] - cargo_version="$(cargo --version)" - [[ "${cargo_version}" =~ ^cargo[[:space:]]1\.96\.0-nightly[[:space:]]\(cbb9bb8bd[[:space:]][0-9]{4}-[0-9]{2}-[0-9]{2}\)$ ]] - rustc_version="$(rustc --version)" - [[ "${rustc_version}" =~ ^rustc[[:space:]]1\.96\.0-nightly[[:space:]]\(ac7f9ec7d[[:space:]][0-9]{4}-[0-9]{2}-[0-9]{2}\)$ ]] - active_toolchain="$(rustup show active-toolchain)" - [[ "${active_toolchain}" =~ ^nightly-2026-03-21-x86_64-unknown-linux-gnu([[:space:]]|$) ]] - [[ "${RUSTUP_HOME}" == "/opt/rustup" ]] - [[ "${RUSTUP_TOOLCHAIN}" == "nightly-2026-03-21" ]] - [[ "${RUSTUP_AUTO_INSTALL}" == "0" ]] - test -r /opt/rustup && test -x /opt/rustup - test ! -w /opt/rustup - test ! -w /opt/cargo/bin - for entry in config config.toml credentials credentials.toml bin .crates.toml .crates2.json .global-cache .package-cache; do - test ! -e "${CARGO_HOME}/${entry}" - done + require_eq otelgen_path "${otelgen_path}" "/usr/local/bin/otelgen" + printf 'check otelgen_version: %s\n' "${otelgen_version}" + [[ "${otelgen_version}" == *"863a3f395d062c7322cc1de08a38774b7fdaa6c8"* ]] \ + || record "otelgen_version: expected commit 863a3f395d062c7322cc1de08a38774b7fdaa6c8, got '${otelgen_version}'" + require_match sccache "$(capture sccache --version)" '^sccache[[:space:]]0\.16\.0([[:space:]]|$)' + require_eq rustup_path "$(command -v rustup || true)" "/opt/cargo/bin/rustup" + require_eq cargo_path "$(command -v cargo || true)" "/opt/cargo/bin/cargo" + require_eq rustc_path "$(command -v rustc || true)" "/opt/cargo/bin/rustc" + require_match rustup "$(capture rustup --version)" '^rustup[[:space:]]1\.29\.0([[:space:]]|$)' + require_match cargo "$(capture cargo --version)" \ + '^cargo[[:space:]]1\.96\.0-nightly[[:space:]]\(cbb9bb8bd[[:space:]][0-9]{4}-[0-9]{2}-[0-9]{2}\)$' + require_match rustc "$(capture rustc --version)" \ + '^rustc[[:space:]]1\.96\.0-nightly[[:space:]]\(ac7f9ec7d[[:space:]][0-9]{4}-[0-9]{2}-[0-9]{2}\)$' + require_match active_toolchain "$(capture rustup show active-toolchain)" \ + '^nightly-2026-03-21-x86_64-unknown-linux-gnu([[:space:]]|$)' + require_eq RUSTUP_HOME "${RUSTUP_HOME}" "/opt/rustup" + require_eq RUSTUP_TOOLCHAIN "${RUSTUP_TOOLCHAIN}" "nightly-2026-03-21" + require_eq RUSTUP_AUTO_INSTALL "${RUSTUP_AUTO_INSTALL}" "0" + require "readable /opt/rustup" test -r /opt/rustup + require "executable /opt/rustup" test -x /opt/rustup + require "runner cannot write /opt/rustup" test ! -w /opt/rustup + require "runner cannot write /opt/cargo/bin" test ! -w /opt/cargo/bin mkdir -p "${CARGO_HOME}/registry" "${CARGO_HOME}/git" + if (( ${#errors[@]} > 0 )); then + printf '\nVerify runner image tools failed (%d checks):\n' "${#errors[@]}" >&2 + printf ' - %s\n' "${errors[@]}" >&2 + exit 1 + fi - - name: Prepare persistent query regression cache + - name: Prepare query regression cache shell: bash working-directory: src run: | @@ -362,24 +553,6 @@ jobs: find "${root}" -mindepth 1 -maxdepth 1 -exec rm -rf -- {} + } - clear_cargo_extracted_trees() { - require_expected_root CARGO_REGISTRY "${CARGO_HOME}/registry" "${EXPECTED_CARGO_REGISTRY}" - require_expected_root CARGO_GIT "${CARGO_HOME}/git" "${EXPECTED_CARGO_GIT}" - rm -rf -- "${CARGO_HOME}/registry/src" "${CARGO_HOME}/git/checkouts" - } - - cargo_size_kib() { - du -sk -- "${CARGO_HOME}/registry" "${CARGO_HOME}/git" | awk '{ total += $1 } END { print total }' - } - - target_size_kib() { - du -sk -- "${CARGO_TARGET_DIR}" | cut -f1 - } - - free_kib() { - df -Pk "${CARGO_TARGET_DIR}" | awk 'NR == 2 { print $4 }' - } - report_cache_usage() { du -sh -- "${CARGO_HOME}" "${CARGO_HOME}/registry" "${CARGO_HOME}/git" \ "${RUSTUP_HOME}" "${CARGO_TARGET_DIR}" "${QUERY_REGRESSION_CACHE_META}" "${SCCACHE_DIR}" @@ -400,7 +573,7 @@ jobs: printf 'Refusing unexpected RUSTC_WRAPPER: %s\n' "${RUSTC_WRAPPER}" >&2 exit 1 } - [[ "${SCCACHE_CACHE_SIZE}" == "40G" ]] || { + [[ "${SCCACHE_CACHE_SIZE}" == "10G" ]] || { printf 'Refusing unexpected SCCACHE_CACHE_SIZE: %s\n' "${SCCACHE_CACHE_SIZE}" >&2 exit 1 } @@ -487,46 +660,10 @@ jobs: report_cache_usage - target_size="$(target_size_kib)" - if (( target_size >= 400 * 1024 * 1024 )); then - printf 'Warning: target cache is at least 400 GiB (%s KiB)\n' "${target_size}" >&2 - fi - if (( target_size >= 450 * 1024 * 1024 )); then - printf 'Target cache is at least 450 GiB; clearing complete target root\n' >&2 - clear_directory "${CARGO_TARGET_DIR}" - fi - - cargo_size="$(cargo_size_kib)" - if (( cargo_size >= 60 * 1024 * 1024 )); then - printf 'Warning: Cargo cache is at least 60 GiB (%s KiB)\n' "${cargo_size}" >&2 - fi - if (( cargo_size >= 80 * 1024 * 1024 )); then - printf 'Cargo cache is at least 80 GiB; removing extracted sources and checkouts\n' >&2 - clear_cargo_extracted_trees - cargo_size="$(cargo_size_kib)" - if (( cargo_size >= 80 * 1024 * 1024 )); then - printf 'Cargo cache remains at least 80 GiB after cleanup (%s KiB)\n' "${cargo_size}" >&2 - exit 1 - fi - fi - - free_space="$(free_kib)" - if (( free_space < 300 * 1024 * 1024 )); then - printf 'Free space is below 300 GiB; clearing complete target root\n' >&2 - clear_directory "${CARGO_TARGET_DIR}" - free_space="$(free_kib)" - if (( free_space < 300 * 1024 * 1024 )); then - printf 'Free space remains below 300 GiB; removing Cargo extracted sources and checkouts\n' >&2 - clear_cargo_extracted_trees - free_space="$(free_kib)" - if (( free_space < 300 * 1024 * 1024 )); then - printf 'Free space remains below 300 GiB after cleanup (%s KiB)\n' "${free_space}" >&2 - exit 1 - fi - fi - fi - - report_cache_usage + # Every run starts from a fresh system disk, so size/free-space + # watermarks from the retired retained-disk era are not restored. + # A run that overflows the disk fails the build outright, which the + # telemetry step makes diagnosable. sccache --start-server sccache --zero-stats @@ -567,14 +704,6 @@ jobs: git reset --hard "${VERIFIED_CANDIDATE_SHA}" git clean -ffdx - - name: Test query regression tooling - working-directory: src - run: | - uv run --no-project python tests/perf/test_query_regression_runner_compaction_toctou.py - uv run --no-project python tests/perf/test_query_regression_runner_otlp_trace_load.py - uv run --no-project python tests/perf/test_query_regression_summary_otlp.py - uv run --no-project python tests/perf/test_query_regression_case_selection.py - - name: Build candidate greptime and query regression helpers working-directory: src run: | @@ -643,6 +772,7 @@ jobs: query-regression-work/**/logs/** query-regression-work/**/otelgen/** query-regression-summary.md + machine-telemetry.log if-no-files-found: warn retention-days: 7 @@ -657,7 +787,7 @@ jobs: if-no-files-found: warn retention-days: 7 - - name: Report persistent cache usage + - name: Report cache usage if: ${{ always() }} shell: bash run: | @@ -697,6 +827,55 @@ jobs: printf 'sccache is unavailable (report only)\n' >&2 fi + - name: Dump machine telemetry + if: ${{ always() }} + shell: bash + run: | + log="${GITHUB_WORKSPACE}/machine-telemetry.log" + if [[ -f "${log}" ]]; then + echo "::group::Machine telemetry (last 200 lines)" + tail -n 200 "${log}" + echo "::endgroup::" + else + echo "No telemetry log found (sampler never started?)" + fi + echo "::group::dmesg tail (OOM killer records)" + sudo dmesg -T 2>/dev/null | tail -n 50 || dmesg -T 2>/dev/null | tail -n 50 || \ + echo "dmesg unavailable without root" + echo "::endgroup::" + - name: Fail on regression failure if: ${{ steps.run.outputs.status != '0' }} run: exit 1 + + teardown: + # Releases the dynamically provisioned ECS runner. Runs even when the + # benchmark job fails or is cancelled; skipped when a literal runner label + # was selected because the provision outputs are empty, or when the + # dispatch set keep_instance to preserve the machine for post-mortem + # debugging (the janitor sweep still reclaims it after the TTL). + if: ${{ always() && needs.provision.outputs.instance_id != '' && !inputs.keep_instance }} + needs: [provision, query-regression] + runs-on: ubuntu-latest + timeout-minutes: 15 + steps: + - name: Checkout trusted teardown scripts + uses: actions/checkout@v4 + with: + repository: ${{ github.repository }} + ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.base.sha || github.sha }} + persist-credentials: false + + - name: Install uv + uses: astral-sh/setup-uv@v6 + + - name: Teardown ECS runner + env: + ALIBABA_CLOUD_ACCESS_KEY_ID: ${{ secrets.ALICLOUD_ECS_ACCESS_KEY_ID }} + ALIBABA_CLOUD_ACCESS_KEY_SECRET: ${{ secrets.ALICLOUD_ECS_ACCESS_KEY_SECRET }} + GH_PERSONAL_ACCESS_TOKEN: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }} + ALIYUN_ECS_REGION_ID: ${{ vars.ALIYUN_ECS_REGION_ID }} + QUERY_REGRESSION_ECS_INSTANCE_ID: ${{ needs.provision.outputs.instance_id }} + QUERY_REGRESSION_ECS_RUNNER_NAME: ${{ needs.provision.outputs.runner_name }} + run: >- + uv run .github/scripts/aliyun-ecs-runner-teardown.py diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 52d62299a07..0a9c39c9034 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -510,7 +510,7 @@ jobs: base_ref: ${{ needs.prepare-release-validation.outputs.previous-release-tag }} candidate_ref: ${{ needs.prepare-release-validation.outputs.candidate-ref }} cargo_profile: nightly - runner: perf-regression-8-cores + runner: aliyun-ecs release-images-to-dockerhub: name: Build and push images to DockerHub diff --git a/Cargo.lock b/Cargo.lock index ec477e5c9a3..8fe6d74dd2f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2167,6 +2167,7 @@ dependencies = [ "common-time", "common-version", "common-wal", + "daemonize", "datafusion", "datafusion-common", "datafusion-physical-plan", @@ -3570,6 +3571,15 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "daemonize" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ab8bfdaacb3c887a54d41bdf48d3af8873b3f5566469f8ba21b92057509f116e" +dependencies = [ + "libc", +] + [[package]] name = "darling" version = "0.14.4" @@ -4628,10 +4638,23 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e7c1832837b905bbfb5101e07cc24c8deddf52f93225eee6ead5f4d63d53ddcb" dependencies = [ "const-oid 0.9.6", + "der_derive", + "flagset", "pem-rfc7468", "zeroize", ] +[[package]] +name = "der_derive" +version = "0.7.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8034092389675178f570469e6c3b0465d3d30b4505c294a6550db47f3c17ad18" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.117", +] + [[package]] name = "deranged" version = "0.5.8" @@ -5334,6 +5357,12 @@ version = "0.5.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1d674e81391d1e1ab681a28d99df07927c6d4aa5b027d7da16ba32d1d21ecd99" +[[package]] +name = "flagset" +version = "0.4.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b7ac824320a75a52197e8f2d787f6a38b6718bb6897a35142d749af3c0e8f4fe" + [[package]] name = "flatbuffers" version = "25.2.10" @@ -5973,7 +6002,7 @@ dependencies = [ "cfg-if", "js-sys", "libc", - "wasi", + "wasi 0.11.1+wasi-snapshot-preview1", "wasm-bindgen", ] @@ -7866,13 +7895,14 @@ checksum = "f9fbbcab51052fe104eb5e5d351cf728d30a5be1fe14d9be8a3b097481fb97de" [[package]] name = "libredox" -version = "0.1.4" +version = "0.1.20" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1580801010e535496706ba011c15f8532df6b42297d2e471fec38ceadd8c0638" +checksum = "28d0a00925a9f930d679b6789b721e3a7f9ed110f41b86d2497caa780c3a070a" dependencies = [ "bitflags 2.12.1", "libc", - "redox_syscall 0.5.13", + "plain", + "redox_syscall 0.9.3", ] [[package]] @@ -8543,7 +8573,7 @@ checksum = "02bd0af71c67b473010cbbc60715ee815645a4dc942899111f494b4b737d6fda" dependencies = [ "libc", "log", - "wasi", + "wasi 0.11.1+wasi-snapshot-preview1", "windows-sys 0.61.2", ] @@ -9358,6 +9388,15 @@ dependencies = [ "objc2-encode", ] +[[package]] +name = "objc2-core-foundation" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a180dd8642fa45cdb7dd721cd4c11b1cadd4929ce112ebd8b9f5803cc79d536" +dependencies = [ + "bitflags 2.12.1", +] + [[package]] name = "objc2-encode" version = "4.1.0" @@ -9374,6 +9413,15 @@ dependencies = [ "objc2", ] +[[package]] +name = "objc2-system-configuration" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7216bd11cbda54ccabcab84d523dc93b858ec75ecfb3a7d89513fa22464da396" +dependencies = [ + "objc2-core-foundation", +] + [[package]] name = "object" version = "0.36.7" @@ -10552,16 +10600,7 @@ dependencies = [ "tokio", "tokio-rustls", "tokio-util", - "x509-certificate 0.25.0", -] - -[[package]] -name = "phf" -version = "0.11.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1fd6780a80ae0c52cc120a26a1a42c1ae51b247a253e4e06113d23d2c2edd078" -dependencies = [ - "phf_shared 0.11.3", + "x509-certificate", ] [[package]] @@ -10814,6 +10853,12 @@ version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7edddbd0b52d732b21ad9a5fab5c704c14cd949e5e9a1ec5929a24fded1b904c" +[[package]] +name = "plain" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4596b6d070b27117e987119b4dac604f3c58cfb0b191112e24771b2faeac1a6" + [[package]] name = "plist" version = "1.7.2" @@ -10924,16 +10969,16 @@ dependencies = [ [[package]] name = "postgres-types" -version = "0.2.9" +version = "0.2.14" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "613283563cd90e1dfc3518d548caee47e0e725455ed619881f5cf21f36de4b48" +checksum = "851ca9db4932932d69f3ea811b1abe63087a0f740a47692619dd40d4899b68be" dependencies = [ "array-init", "bytes", "chrono", "fallible-iterator", "postgres-protocol", - "serde", + "serde_core", "serde_json", ] @@ -11299,7 +11344,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "22505a5c94da8e3b7c2996394d1c933236c4d743e81a410bcca4e6989fc066a4" dependencies = [ "bytes", - "heck 0.4.1", + "heck 0.5.0", "itertools 0.12.1", "log", "multimap", @@ -11319,7 +11364,7 @@ version = "0.14.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ac6c3320f9abac597dcbc668774ef006702672474aad53c6d596b62e487b40b1" dependencies = [ - "heck 0.4.1", + "heck 0.5.0", "itertools 0.14.0", "log", "multimap", @@ -12087,6 +12132,15 @@ dependencies = [ "bitflags 2.12.1", ] +[[package]] +name = "redox_syscall" +version = "0.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d678d17679829e73d371e96880897e98fee2ded7acc0a50bdf8af2affa4b2fe5" +dependencies = [ + "bitflags 2.12.1", +] + [[package]] name = "redox_users" version = "0.4.6" @@ -13785,7 +13839,7 @@ version = "0.8.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1961e2ef424c1424204d3a5d6975f934f56b6d50ff5732382d84ebf460e147f7" dependencies = [ - "heck 0.4.1", + "heck 0.5.0", "proc-macro2", "quote", "syn 2.0.117", @@ -14117,7 +14171,7 @@ dependencies = [ "stringprep", "thiserror 2.0.17", "tracing", - "whoami", + "whoami 1.6.0", ] [[package]] @@ -14156,7 +14210,7 @@ dependencies = [ "stringprep", "thiserror 2.0.17", "tracing", - "whoami", + "whoami 1.6.0", ] [[package]] @@ -15175,6 +15229,27 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" +[[package]] +name = "tls_codec" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0de2e01245e2bb89d6f05801c564fa27624dbd7b1846859876c7dad82e90bf6b" +dependencies = [ + "tls_codec_derive", + "zeroize", +] + +[[package]] +name = "tls_codec_derive" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d2e76690929402faae40aebdda620a2c0e25dd6d3b9afe48867dfd95991f4bd" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.117", +] + [[package]] name = "tokio" version = "1.52.3" @@ -15241,9 +15316,9 @@ dependencies = [ [[package]] name = "tokio-postgres" -version = "0.7.13" +version = "0.7.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6c95d533c83082bb6490e0189acaa0bbeef9084e60471b696ca6988cd0541fb0" +checksum = "a528f7d280f6d5b9cd149635c8705b0dd049754bc67d81d31fa25169a93809d3" dependencies = [ "async-trait", "byteorder", @@ -15254,29 +15329,29 @@ dependencies = [ "log", "parking_lot 0.12.4", "percent-encoding", - "phf 0.11.3", + "phf 0.13.1", "pin-project-lite", "postgres-protocol", "postgres-types", - "rand 0.9.4", - "socket2 0.5.10", + "rand 0.10.1", + "socket2 0.6.4", "tokio", "tokio-util", - "whoami", + "whoami 2.1.3", ] [[package]] name = "tokio-postgres-rustls" -version = "0.12.0" +version = "0.14.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "04fb792ccd6bbcd4bba408eb8a292f70fc4a3589e5d793626f45190e6454b6ab" +checksum = "4c2ad44aa0ae96db89c4742212ed41645b2f597311ff6e1945542a4d9fadc2fb" dependencies = [ - "ring", "rustls", + "sha2 0.11.0", "tokio", "tokio-postgres", "tokio-rustls", - "x509-certificate 0.23.1", + "x509-cert", ] [[package]] @@ -16284,6 +16359,15 @@ version = "0.11.1+wasi-snapshot-preview1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" +[[package]] +name = "wasi" +version = "0.14.7+wasi-0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "883478de20367e224c0090af9cf5f9fa85bed63a95c1abf3afc5c083ebc06e8c" +dependencies = [ + "wasip2", +] + [[package]] name = "wasip2" version = "1.0.2+wasi-0.2.9" @@ -16308,6 +16392,15 @@ version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b8dad83b4f25e74f184f64c43b150b91efe7647395b42289f38e50566d82855b" +[[package]] +name = "wasite" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "66fe902b4a6b8028a753d5424909b764ccf79b7a209eac9bf97e59cda9f71a42" +dependencies = [ + "wasi 0.14.7+wasi-0.2.4", +] + [[package]] name = "wasm-bindgen" version = "0.2.122" @@ -16530,7 +16623,19 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6994d13118ab492c3c80c1f81928718159254c53c472bf9ce36f8dae4add02a7" dependencies = [ "redox_syscall 0.5.13", - "wasite", + "wasite 0.1.0", +] + +[[package]] +name = "whoami" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "626c4bac6755d76ffc12cb01b2eac751db1996b9e0041de9aa02c8c211ddc82c" +dependencies = [ + "libc", + "libredox", + "objc2-system-configuration", + "wasite 1.0.2", "web-sys", ] @@ -17139,22 +17244,15 @@ dependencies = [ ] [[package]] -name = "x509-certificate" -version = "0.23.1" +name = "x509-cert" +version = "0.2.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "66534846dec7a11d7c50a74b7cdb208b9a581cad890b7866430d438455847c85" +checksum = "1301e935010a701ae5f8655edc0ad17c44bad3ac5ce8c39185f75453b720ae94" dependencies = [ - "bcder", - "bytes", - "chrono", + "const-oid 0.9.6", "der", - "hex", - "pem", - "ring", - "signature", "spki", - "thiserror 1.0.69", - "zeroize", + "tls_codec", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 7c3fccfbbf8..fab7b2aa736 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -255,7 +255,7 @@ strum = { version = "0.27", features = ["derive"] } sysinfo = "0.33" tempfile = "3" tokio = { version = "1.47", features = ["full"] } -tokio-postgres = "0.7" +tokio-postgres = "0.7.18" tokio-rustls = { version = "0.26.2", default-features = false } tokio-stream = "0.1" tokio-util = { version = "0.7", features = ["io-util", "compat"] } diff --git a/config/config.md b/config/config.md index 91cb4fe7669..e8a89e94e1f 100644 --- a/config/config.md +++ b/config/config.md @@ -76,6 +76,7 @@ | `otlp.enable` | Bool | `true` | Whether to enable OpenTelemetry protocol in HTTP API. | | `otlp.experimental_enable_exponential_histogram` | Bool | `false` | Experimental: enable cumulative OTLP exponential histogram ingestion. | | `otlp.trace_ingest_chunk_size` | Integer | `512` | Maximum spans per trace ingest chunk. Set to 0 to disable splitting. | +| `otlp.experimental_enable_resource_info` | Bool | `false` | Whether to synthesize the `greptime_otel_resource_info` table from OTLP metric
resource attributes, so metrics-only services reach the semantic graph. | | `prom_store` | -- | -- | Prometheus remote storage options | | `prom_store.enable` | Bool | `true` | Whether to enable Prometheus remote write and read in HTTP API. | | `prom_store.with_metric_engine` | Bool | `true` | Whether to store the data from Prometheus remote write in metric engine. | @@ -160,7 +161,7 @@ | `region_engine.mito.max_background_purges` | Integer | Auto | Max number of running background purge jobs (default: number of cpu cores). | | `region_engine.mito.experimental_compaction_memory_limit` | String | 0 | Memory budget for compaction tasks.
Supports absolute size (e.g., "2GiB", "512MB") or percentage of system memory (e.g., "50%").
Setting it to 0 or "unlimited" disables the limit. | | `region_engine.mito.experimental_compaction_on_exhausted` | String | wait | Behavior when compaction cannot acquire memory from the budget.
Options: "wait" (default, 10s), "wait()", "fail" | -| `region_engine.mito.auto_flush_interval` | String | `1h` | Interval to auto flush a region if it has not flushed yet. | +| `region_engine.mito.auto_flush_interval` | String | `10m` | Interval to auto flush a region if it has not flushed yet. | | `region_engine.mito.global_write_buffer_size` | String | Auto | Global write buffer size for all regions. If not set, it's default to 1/8 of OS memory with a max limitation of 1GB. | | `region_engine.mito.global_write_buffer_reject_size` | String | Auto | Global write buffer size threshold to reject write requests. If not set, it's default to 2 times of `global_write_buffer_size`. | | `region_engine.mito.default_region_write_buffer_size` | String | `0` | Default write buffer size for each region. Regions stall at this size and reject writes at twice this size. Setting it to 0 disables both limits unless the table specifies `write_buffer_size`. | @@ -236,7 +237,7 @@ | `tracing.tokio_console_addr` | String | Unset | The tokio console address. | | `event_recorder` | -- | -- | Configuration options for the event recorder. | | `event_recorder.ttl` | String | `90d` | TTL for the events table that will be used to store the events. Default is `90d`. | -| `event_recorder.event_types` | Array | -- | Event types to record. Current available event types: `create_database`,
`alter_database`, `drop_database`, `create_flow`, `drop_flow`,
`create_table`, `create_logical_tables`, `alter_table`, `alter_logical_tables`,
`drop_table`, `undrop_table`, `purge_dropped_table`, `truncate_table`,
`create_view`, `drop_view`, `admin_function`.
When omitted, all current and future event types are recorded.
Set to an empty array to disable event recording. | +| `event_recorder.event_types` | Array | -- | Event types to record. Current available event types: `create_database`,
`alter_database`, `drop_database`, `create_flow`, `drop_flow`,
`create_table`, `create_logical_tables`, `alter_table`, `alter_logical_tables`,
`drop_table`, `undrop_table`, `purge_dropped_table`, `truncate_table`,
`create_view`, `drop_view`, `admin_function`, `reconcile_table`.
When omitted, all current and future event types are recorded.
Set to an empty array to disable event recording. | | `memory` | -- | -- | The memory options. | | `memory.enable_heap_profiling` | Bool | `true` | Whether to enable heap profiling activation during startup.
When enabled, heap profiling will be activated if the `MALLOC_CONF` environment variable
is set to "prof:true,prof_active:false". The official image adds this env variable.
Default is true. | @@ -319,6 +320,7 @@ | `otlp.enable` | Bool | `true` | Whether to enable OpenTelemetry protocol in HTTP API. | | `otlp.experimental_enable_exponential_histogram` | Bool | `false` | Experimental: enable cumulative OTLP exponential histogram ingestion. | | `otlp.trace_ingest_chunk_size` | Integer | `512` | Maximum spans per trace ingest chunk. Set to 0 to disable splitting. | +| `otlp.experimental_enable_resource_info` | Bool | `false` | Whether to synthesize the `greptime_otel_resource_info` table from OTLP metric
resource attributes, so metrics-only services reach the semantic graph. | | `prom_store` | -- | -- | Prometheus remote storage options | | `prom_store.enable` | Bool | `true` | Whether to enable Prometheus remote write and read in HTTP API. | | `prom_store.with_metric_engine` | Bool | `true` | Whether to store the data from Prometheus remote write in metric engine. | @@ -454,7 +456,7 @@ | `wal.create_topic_timeout` | String | `30s` | The timeout for creating a Kafka topic.
**It's only used when the provider is `kafka`**. | | `event_recorder` | -- | -- | Configuration options for the event recorder. | | `event_recorder.ttl` | String | `90d` | TTL for the events table that will be used to store the events. Default is `90d`. | -| `event_recorder.event_types` | Array | -- | Event types to record. Current available event types: `region_migration`,
`create_database`, `alter_database`, `drop_database`, `create_flow`,
`drop_flow`, `create_table`, `create_logical_tables`, `alter_table`,
`alter_logical_tables`, `drop_table`, `undrop_table`, `purge_dropped_table`,
`truncate_table`, `create_view`, `drop_view`, `repartition`,
`repartition_group`, `wal_prune`, `batch_gc`.
When omitted, all current and future event types are recorded.
Set to an empty array to disable event recording. | +| `event_recorder.event_types` | Array | -- | Event types to record. Current available event types: `region_migration`,
`create_database`, `alter_database`, `drop_database`, `create_flow`,
`drop_flow`, `create_table`, `create_logical_tables`, `alter_table`,
`alter_logical_tables`, `drop_table`, `undrop_table`, `purge_dropped_table`,
`truncate_table`, `create_view`, `drop_view`, `repartition`,
`repartition_group`, `wal_prune`, `batch_gc`, `reconcile_table`.
When omitted, all current and future event types are recorded.
Set to an empty array to disable event recording. | | `stats_persistence` | -- | -- | Configuration options for the stats persistence. | | `stats_persistence.ttl` | String | `0s` | TTL for the stats table that will be used to store the stats.
Set to `0s` to disable stats persistence.
Default is `0s`.
If you want to enable stats persistence, set the TTL to a value greater than 0.
It is recommended to set a small value, e.g., `3h`. | | `stats_persistence.interval` | String | `10m` | The interval to persist the stats. Default is `10m`.
The minimum value is `10m`, if the value is less than `10m`, it will be overridden to `10m`. | @@ -583,7 +585,7 @@ | `region_engine.mito.max_background_purges` | Integer | Auto | Max number of running background purge jobs (default: number of cpu cores). | | `region_engine.mito.experimental_compaction_memory_limit` | String | 0 | Memory budget for compaction tasks.
Supports absolute size (e.g., "2GiB", "512MB") or percentage of system memory (e.g., "50%").
Setting it to 0 or "unlimited" disables the limit. | | `region_engine.mito.experimental_compaction_on_exhausted` | String | wait | Behavior when compaction cannot acquire memory from the budget.
Options: "wait" (default, 10s), "wait()", "fail" | -| `region_engine.mito.auto_flush_interval` | String | `1h` | Interval to auto flush a region if it has not flushed yet. | +| `region_engine.mito.auto_flush_interval` | String | `10m` | Interval to auto flush a region if it has not flushed yet. | | `region_engine.mito.global_write_buffer_size` | String | Auto | Global write buffer size for all regions. If not set, it's default to 1/8 of OS memory with a max limitation of 1GB. | | `region_engine.mito.global_write_buffer_reject_size` | String | Auto | Global write buffer size threshold to reject write requests. If not set, it's default to 2 times of `global_write_buffer_size` | | `region_engine.mito.default_region_write_buffer_size` | String | `0` | Default write buffer size for each region. Regions stall at this size and reject writes at twice this size. Setting it to 0 disables both limits unless the table specifies `write_buffer_size`. | diff --git a/config/datanode.example.toml b/config/datanode.example.toml index 5966405452a..20c84fd9ee4 100644 --- a/config/datanode.example.toml +++ b/config/datanode.example.toml @@ -462,7 +462,7 @@ compress_manifest = false #+ experimental_compaction_on_exhausted = "wait" ## Interval to auto flush a region if it has not flushed yet. -auto_flush_interval = "1h" +auto_flush_interval = "10m" ## Global write buffer size for all regions. If not set, it's default to 1/8 of OS memory with a max limitation of 1GB. ## @toml2docs:none-default="Auto" diff --git a/config/frontend.example.toml b/config/frontend.example.toml index 3635a63f063..80534855721 100644 --- a/config/frontend.example.toml +++ b/config/frontend.example.toml @@ -241,6 +241,9 @@ enable = true experimental_enable_exponential_histogram = false ## Maximum spans per trace ingest chunk. Set to 0 to disable splitting. trace_ingest_chunk_size = 512 +## Whether to synthesize the `greptime_otel_resource_info` table from OTLP metric +## resource attributes, so metrics-only services reach the semantic graph. +experimental_enable_resource_info = false ## Prometheus remote storage options [prom_store] diff --git a/config/metasrv.example.toml b/config/metasrv.example.toml index 6017bfcc49a..ae373d07482 100644 --- a/config/metasrv.example.toml +++ b/config/metasrv.example.toml @@ -320,7 +320,7 @@ ttl = "90d" ## `drop_flow`, `create_table`, `create_logical_tables`, `alter_table`, ## `alter_logical_tables`, `drop_table`, `undrop_table`, `purge_dropped_table`, ## `truncate_table`, `create_view`, `drop_view`, `repartition`, -## `repartition_group`, `wal_prune`, `batch_gc`. +## `repartition_group`, `wal_prune`, `batch_gc`, `reconcile_table`. ## When omitted, all current and future event types are recorded. ## Set to an empty array to disable event recording. #+ event_types = ["region_migration"] diff --git a/config/standalone.example.toml b/config/standalone.example.toml index 72ba992cce0..1fd8fe5a522 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -208,6 +208,9 @@ enable = true experimental_enable_exponential_histogram = false ## Maximum spans per trace ingest chunk. Set to 0 to disable splitting. trace_ingest_chunk_size = 512 +## Whether to synthesize the `greptime_otel_resource_info` table from OTLP metric +## resource attributes, so metrics-only services reach the semantic graph. +experimental_enable_resource_info = false ## Prometheus remote storage options [prom_store] @@ -622,7 +625,7 @@ compress_manifest = false #+ experimental_compaction_on_exhausted = "wait" ## Interval to auto flush a region if it has not flushed yet. -auto_flush_interval = "1h" +auto_flush_interval = "10m" ## Global write buffer size for all regions. If not set, it's default to 1/8 of OS memory with a max limitation of 1GB. ## @toml2docs:none-default="Auto" @@ -927,7 +930,7 @@ ttl = "90d" ## `alter_database`, `drop_database`, `create_flow`, `drop_flow`, ## `create_table`, `create_logical_tables`, `alter_table`, `alter_logical_tables`, ## `drop_table`, `undrop_table`, `purge_dropped_table`, `truncate_table`, -## `create_view`, `drop_view`, `admin_function`. +## `create_view`, `drop_view`, `admin_function`, `reconcile_table`. ## When omitted, all current and future event types are recorded. ## Set to an empty array to disable event recording. #+ event_types = ["create_database"] diff --git a/docs/rfcs/2026-06-25-entity-relationships-and-graph-query.md b/docs/rfcs/2026-06-25-entity-relationships-and-graph-query.md index 995cf6cab7e..1c78e91f252 100644 --- a/docs/rfcs/2026-06-25-entity-relationships-and-graph-query.md +++ b/docs/rfcs/2026-06-25-entity-relationships-and-graph-query.md @@ -149,7 +149,7 @@ The decisions behind the shape: - **`provenance` is part of edge identity**, so a declared edge and a derived edge between the same pair coexist, and a user-declared edge survives when derivation is off. - **`confidence` expresses derivation certainty, not statistical completeness**: `1.0` for a successfully paired or declared edge, `< 1.0` for virtual-node or agent-inferred edges. It does not correct for trace sampling. - **RED metrics describe the observed span-pair population** — the `service_graph` connector semantics. Under sampling, counts understate true traffic, and ratios (error rate, mean duration) are representative only when sampling is unbiased with respect to status and latency — tail sampling that keeps errors and slow traces skews both. -- **Endpoint encoding (v1).** A single-attribute identity is the value verbatim; a composite identity is the attributes sorted by key, rendered `k1=v1,k2=v2` — a human-readable traversal key, not a collision-free canonical encoding. Nodes carry the structured `entity_id_attrs`; edges carry only the string key. Known v1 limitation: composite components containing unescaped `,` or `=` are not guaranteed to be distinguished. +- **Endpoint encoding.** An identity renders as its values in declared order (broad to narrow), joined by `,` with `,` and `\` escaped in each value. The identifying attribute *keys* are deliberately absent: the same identity reaches the graph under different column names per signal — a trace table's `service_name` against a metric table's `job` — and encoding the keys would split one entity into one entity per signal. Where a source carries a qualifier separately (`service.namespace` beside `service.name`) the convention folds it in as `/`, reproducing the identity a source that pre-composed it (Prometheus `job`) already emits. Nodes carry the structured `entity_id_attrs`; edges carry only the string key. - **Hand-declared edges** (`provenance = 'declared'`) live in a physical table of the same shape (plus a business validity window) in `greptime_private` and are unioned into the computed `semantic_relationships`. The relationship vocabulary is small, typed, and inverse-paired; stored direction is `src -> dst`, the inverse is a query concern: @@ -197,7 +197,7 @@ GROUP BY 1, src_id, dst_id; Edge endpoints are built from each trace table's `service` entity declaration (the same id expression the registry uses), so edges land on exactly the entity ids the registry emits, self-calls are judged on the full endpoint id, and a table whose declaration is unusable contributes no edges rather than silently switching identity. The join is bounded by a time-proximity window and the results of **all** trace tables combine into one edge per window with merged RED metrics. One assumption to note: a deployment can route traces to multiple tables (`x-greptime-trace-table-name`), so one distributed trace is not guaranteed to live in a single table — pairing must account for that, and how is an implementation concern. Uninstrumented peers become **virtual nodes**: a client span with no matching server span yields an edge to a synthetic node named from `peer.service`/`db.name`/`server.address`. -**The entity registry.** Per declaring table: filter to the window, project (time bin, casts, JSON assembly), `DISTINCT`. The per-table results are unioned without any cross-source merge: no descriptive-conflict or JSON-merging policy lives in the derivation. Rows whose identity columns are NULL identify nothing and are filtered out; single-column ids are used verbatim, composite ids render the sorted `k=v` form plus the `entity_id_attrs` JSON. Consumers deduplicate with `SELECT DISTINCT entity_type, entity_id`; a unique node set is the snapshot relation's job (see *Querying the graph*). +**The entity registry.** Per declaring table: filter to the window, project (time bin, casts, JSON assembly), `DISTINCT`. The per-table results are unioned without any cross-source merge: no descriptive-conflict or JSON-merging policy lives in the derivation. Rows whose identity columns are NULL identify nothing and are filtered out; ids render as their escaped values joined in declared order, with the `entity_id_attrs` JSON alongside once an id is assembled from more than one column. Consumers deduplicate with `SELECT DISTINCT entity_type, entity_id`; a unique node set is the snapshot relation's job (see *Querying the graph*). **Attribute edges.** Shared attributes provide join keys; co-declaration or a relationship template provides the relationship semantics. A shared value alone determines neither direction nor edge type, so the derivation is rule-based, not join-everything. A table declaring both a `service.instance` and a `host` identity on the same rows derives `runs_on` between them: the row itself witnesses the relationship, and the built-in vocabulary fixes the direction. Cross-table edges require a declared relationship template (endpoint mappings, `rel_type`, direction): a derivation *rule* declared on schema, not per-instance edge data. @@ -284,7 +284,7 @@ Drawbacks: - Read-time derivation costs a self-join over trace tables on every scan of the computed tables. The default one-hour window bounds it, but very large deployments will want materialisation (Future Work). - `semantic_entities` exposes per-observation rows; deduplication falls on the consumer (`DISTINCT`) until the snapshot relation lands in M2. -- The v1 `k=v` id encoding can collide on hostile values (unescaped `,`/`=`); hardening is an open question. +- The id encoding is readable rather than canonical: escaping keeps components distinguishable, but a value containing the qualifier separator (`/`) is still ambiguous against a qualified name. - The graph is only as connected as the identity columns tables actually share: the `job`/`service_name` misalignment above is the canonical example, and enrichment discipline gates connectivity. # Implementation Plan @@ -306,7 +306,7 @@ Drawbacks: # Open Questions -1. **Endpoint encoding hardening.** The v1 `k=v` rendering (with `entity_id_attrs` JSON as the source of truth) is readable but collides when values contain `,`/`=`. Open: a fixed-width `entity_id_hash` join key (New-Relic-style); typed value encoding; the criteria for introducing the `semantic_key=column` mapping form; `schema_url`-gated entity merge. Related and unresolved: whether the identifying attribute *key* should participate in entity equality — a single-attribute id keeps only the value, so `job=api` and `service_name=api` currently name the same entity even when the keys differ semantically. +1. **Endpoint encoding hardening.** The rendering (with `entity_id_attrs` JSON as the source of truth) is readable rather than canonical. Open: a fixed-width `entity_id_hash` join key (New-Relic-style); typed value encoding; the criteria for introducing the `semantic_key=column` mapping form; `schema_url`-gated entity merge. Settled: identifying attribute *keys* do not participate in entity equality, so `job=api` and `service_name=api` name the same entity — the property that lets one service reach the graph from several signals. 2. **Identity reconciliation across signals.** Trace edges key on the trace table's service declaration; host edges key on `host.id`/`k8s.pod.uid`. Linking them relies on tables carrying both columns; how aggressively to enrich (a `k8sattributes`-equivalent) gates how connected the graph is. 3. **`GRAPH_TABLE` surface scope.** How much of SQL/PGQ to implement vs. a minimal single-hop subset. Start minimal and grow with demand. 4. **Visibility beyond table grants.** The derivation contract filters by table-read permission; whether topology needs finer policies (row-level, tenant-scoped visibility of cross-schema edges) is open. diff --git a/src/client/src/database.rs b/src/client/src/database.rs index 4fff3636b63..7c33794f477 100644 --- a/src/client/src/database.rs +++ b/src/client/src/database.rs @@ -12,20 +12,25 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::collections::HashMap; use std::pin::Pin; use std::str::FromStr; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, RwLock}; use std::task::{Context, Poll}; +use std::time::Duration; use api::v1::auth_header::AuthScheme; +#[cfg(feature = "testing")] use api::v1::ddl_request::Expr as DdlExpr; use api::v1::greptime_database_client::GreptimeDatabaseClient; use api::v1::greptime_request::Request; use api::v1::query_request::Query; +#[cfg(feature = "testing")] +use api::v1::{AlterTableExpr, CreateTableExpr, DdlRequest}; use api::v1::{ - AlterTableExpr, AuthHeader, Basic, CreateTableExpr, DdlRequest, GreptimeRequest, - InsertRequests, QueryRequest, RequestHeader, RowInsertRequests, + AuthHeader, Basic, GreptimeRequest, InsertRequests, QueryRequest, RequestHeader, + RowInsertRequests, }; use arc_swap::ArcSwapOption; use arrow_flight::{FlightData, Ticket}; @@ -34,7 +39,7 @@ use base64::Engine; use base64::prelude::BASE64_STANDARD; use common_catalog::build_db_string; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; -use common_error::ext::BoxedError; +use common_error::ext::{BoxedError, ErrorExt}; use common_grpc::flight::do_put::DoPutResponse; use common_grpc::flight::{ FLOW_EXTENSIONS_METADATA_KEY, FlightDecoder, FlightMessage, SNAPSHOT_SEQS_METADATA_KEY, @@ -49,21 +54,23 @@ use common_telemetry::{error, warn}; use futures::future; use futures_util::{Stream, StreamExt, TryStreamExt}; use prost::Message; -use snafu::ResultExt; +use snafu::{IntoError, ResultExt}; use tonic::metadata::{AsciiMetadataKey, AsciiMetadataValue, MetadataMap, MetadataValue}; use tonic::transport::Channel; use crate::error::{ - ConvertFlightDataSnafu, Error, FlightGetSnafu, IllegalFlightMessagesSnafu, + ConvertFlightDataSnafu, Error, FlightGetSnafu, FlightStreamSnafu, IllegalFlightMessagesSnafu, InvalidTonicMetadataValueSnafu, }; -use crate::flight::decode_flight_data; +use crate::flight::{FlightMessageReader, decode_flight_data}; use crate::{Client, Result, error, from_grpc_response}; type FlightDataStream = Pin + Send>>; type DoPutResponseStream = Pin>>>; +const HINTS_METADATA_KEY: &str = "x-greptime-hints"; + /// Terminal metrics associated with a query output. /// /// For streaming outputs, metrics are only final after the stream is fully @@ -257,19 +264,17 @@ fn attach_terminal_metrics(output: Output, terminal_metrics: &OutputMetrics) -> } async fn output_from_flight_message_stream( - mut flight_message_stream: S, + remote_addr: String, + flight_message_stream: S, ) -> Result where S: Stream> + Send + Unpin + 'static, { - let Some(first_flight_message) = flight_message_stream.next().await else { - return IllegalFlightMessagesSnafu { - reason: "Expect the response not to be empty", - } - .fail(); - }; - - let first_flight_message = first_flight_message?; + let mut reader = FlightMessageReader::new(remote_addr, flight_message_stream); + let first_flight_message = reader + .read_first() + .await + .map_err(|error| flight_stream_error(reader.remote_addr(), error))?; match first_flight_message { FlightMessage::AffectedRows { rows, metrics } => { @@ -277,7 +282,10 @@ where if let Some(metrics) = metrics { terminal_metrics.update(Some(parse_terminal_metrics(&metrics)?)); } - let next_message = flight_message_stream.next().await.transpose()?; + let next_message = reader + .read_next() + .await + .map_err(|error| flight_stream_error(reader.remote_addr(), error))?; match next_message { None => terminal_metrics.mark_ready(), Some(FlightMessage::Metrics(s)) if terminal_metrics.get().is_none() => { @@ -316,15 +324,19 @@ where ); let schema_cloned = schema.clone(); let stream = Box::pin(stream!({ - while let Some(flight_message_item) = flight_message_stream.next().await { - let flight_message = match flight_message_item { - Ok(message) => message, - Err(e) => { - yield Err(BoxedError::new(e)).context(ExternalSnafu); + loop { + let flight_message = match reader.read_next().await { + Ok(Some(message)) => message, + Ok(None) => break, + Err(error) => { + yield Err(BoxedError::new(flight_stream_error( + reader.remote_addr(), + error, + ))) + .context(ExternalSnafu); break; } }; - match flight_message { FlightMessage::RecordBatch(arrow_batch) => { yield Ok(RecordBatch::from_df_record_batch( @@ -371,6 +383,25 @@ where } } +fn flight_stream_error(addr: &str, error: Error) -> Error { + let tonic_code = error.tonic_code().unwrap_or(tonic::Code::Unknown); + let message = error.to_string(); + if error.status_code().should_log_error() { + error!( + error; "Failed to receive Flight data, addr: {}, code: {}", + addr, + tonic_code + ); + } + + FlightStreamSnafu { + addr: addr.to_string(), + tonic_code, + message, + } + .into_error(BoxedError::new(error)) +} + #[derive(Clone, Debug, Default)] pub struct Database { // The "catalog" and "schema" to be used in processing the requests at the server side. @@ -389,6 +420,42 @@ pub struct Database { ctx: FlightContext, } +#[derive(Default)] +struct FlightRequestOptions { + hints: Option, + flow_extensions: Option, + snapshot_seqs: Option, + timeout: Option, +} + +impl FlightRequestOptions { + fn apply_to(self, request: &mut tonic::Request) -> Result<()> { + let metadata = request.metadata_mut(); + if let Some(hints) = self.hints { + Database::put_metadata_value(metadata, HINTS_METADATA_KEY, hints)?; + } + if let Some(flow_extensions) = self.flow_extensions { + Database::put_metadata_value(metadata, FLOW_EXTENSIONS_METADATA_KEY, flow_extensions)?; + } + if let Some(snapshot_seqs) = self.snapshot_seqs { + Database::put_metadata_value(metadata, SNAPSHOT_SEQS_METADATA_KEY, snapshot_seqs)?; + } + if let Some(timeout) = self.timeout { + request.set_timeout(timeout); + } + Ok(()) + } +} + +/// A single Flight DoGet request to a [`Database`]. +/// +/// The builder carries request-scoped metadata and options. It does not modify +/// the underlying [`Database`], so its configuration cannot affect later RPCs. +pub struct DatabaseFlightRequest<'a> { + database: &'a Database, + options: FlightRequestOptions, +} + pub struct DatabaseClient { pub addr: String, pub inner: GreptimeDatabaseClient, @@ -483,6 +550,14 @@ impl Database { }); } + /// Creates a builder for a single Flight DoGet request. + pub fn flight_request(&self) -> DatabaseFlightRequest<'_> { + DatabaseFlightRequest { + database: self, + options: FlightRequestOptions::default(), + } + } + /// Make an InsertRequests request to the database. pub async fn insert(&self, requests: InsertRequests) -> Result { self.handle(Request::Inserts(requests)).await @@ -538,44 +613,31 @@ impl Database { } fn put_hints(metadata: &mut MetadataMap, hints: &[(&str, &str)]) -> Result<()> { - let Some(value) = hints - .iter() - .map(|(k, v)| format!("{}={}", k, v)) - .reduce(|a, b| format!("{},{}", a, b)) - else { + let Some(value) = Self::encode_hints(hints) else { return Ok(()); }; - let key = AsciiMetadataKey::from_static("x-greptime-hints"); - let value = AsciiMetadataValue::from_str(&value).context(InvalidTonicMetadataValueSnafu)?; - metadata.insert(key, value); - Ok(()) + Self::put_metadata_value(metadata, HINTS_METADATA_KEY, value) } - fn put_flow_extensions( - metadata: &mut MetadataMap, - flow_extensions: &[(&str, &str)], - ) -> Result<()> { - if flow_extensions.is_empty() { - return Ok(()); - } - - let value = serde_json::to_string(&flow_extensions.to_vec()) - .expect("flow extension pairs should serialize"); - Self::put_metadata_value(metadata, FLOW_EXTENSIONS_METADATA_KEY, value) + fn encode_hints(hints: &[(&str, &str)]) -> Option { + hints + .iter() + .map(|(k, v)| format!("{}={}", k, v)) + .reduce(|a, b| format!("{},{}", a, b)) } - fn put_snapshot_seqs( - metadata: &mut MetadataMap, - snapshot_seqs: &std::collections::HashMap, - ) -> Result<()> { - if snapshot_seqs.is_empty() { - return Ok(()); - } + fn encode_flow_extensions(flow_extensions: &[(&str, &str)]) -> Option { + (!flow_extensions.is_empty()).then(|| { + serde_json::to_string(&flow_extensions.to_vec()) + .expect("flow extension pairs should serialize") + }) + } - let value = - serde_json::to_string(snapshot_seqs).expect("snapshot sequence map should serialize"); - Self::put_metadata_value(metadata, SNAPSHOT_SEQS_METADATA_KEY, value) + fn encode_snapshot_seqs(snapshot_seqs: &HashMap) -> Option { + (!snapshot_seqs.is_empty()).then(|| { + serde_json::to_string(snapshot_seqs).expect("snapshot sequence map should serialize") + }) } fn put_metadata_value( @@ -673,7 +735,7 @@ impl Database { where S: AsRef, { - self.sql_with_hint(sql, &[]).await + self.flight_request().sql(sql).await } /// Executes a SQL query with optional hints for query optimization. @@ -681,12 +743,7 @@ impl Database { where S: AsRef, { - let request = Request::Query(QueryRequest { - query: Some(Query::Sql(sql.as_ref().to_string())), - }); - self.do_get(request, hints, &[], &Default::default()) - .await - .map(OutputWithMetrics::into_output) + self.flight_request().with_hints(hints).sql(sql).await } /// Executes a SQL query and returns the output with terminal metrics. @@ -701,77 +758,33 @@ impl Database { where S: AsRef, { - self.query_with_terminal_metrics_and_flow_extensions( - QueryRequest { - query: Some(Query::Sql(sql.as_ref().to_string())), - }, - hints, - &[], - &Default::default(), - ) - .await + self.flight_request() + .with_hints(hints) + .sql_with_terminal_metrics(sql) + .await } /// Executes a logical plan directly without SQL parsing. pub async fn logical_plan(&self, logical_plan: Vec) -> Result { - self.query_with_terminal_metrics_and_flow_extensions( - QueryRequest { - query: Some(Query::LogicalPlan(logical_plan)), - }, - &[], - &[], - &Default::default(), - ) - .await - .map(OutputWithMetrics::into_output) - } - - /// Executes a query and carries flow extensions through Flight metadata. - /// - /// This is the lower-level terminal-metrics API for Flow callers that need - /// to pass JSON-bearing flow extensions without going through hint metadata. - pub async fn query_with_terminal_metrics_and_flow_extensions( - &self, - request: QueryRequest, - hints: &[(&str, &str)], - flow_extensions: &[(&str, &str)], - snapshot_seqs: &std::collections::HashMap, - ) -> Result { - self.do_get( - Request::Query(request), - hints, - flow_extensions, - snapshot_seqs, - ) - .await + self.flight_request().logical_plan(logical_plan).await } /// Creates a new table using the provided table expression. + #[cfg(feature = "testing")] pub async fn create(&self, expr: CreateTableExpr) -> Result { - let request = Request::Ddl(DdlRequest { - expr: Some(DdlExpr::CreateTable(expr)), - }); - self.do_get(request, &[], &[], &Default::default()) - .await - .map(OutputWithMetrics::into_output) + self.flight_request().create(expr).await } /// Alters an existing table using the provided alter expression. + #[cfg(feature = "testing")] pub async fn alter(&self, expr: AlterTableExpr) -> Result { - let request = Request::Ddl(DdlRequest { - expr: Some(DdlExpr::AlterTable(expr)), - }); - self.do_get(request, &[], &[], &Default::default()) - .await - .map(OutputWithMetrics::into_output) + self.flight_request().alter(expr).await } async fn do_get( &self, request: Request, - hints: &[(&str, &str)], - flow_extensions: &[(&str, &str)], - snapshot_seqs: &std::collections::HashMap, + options: FlightRequestOptions, ) -> Result { let request = self.to_rpc_request(request); let request = Ticket { @@ -779,12 +792,10 @@ impl Database { }; let mut request = tonic::Request::new(request); - let metadata = request.metadata_mut(); - Self::put_hints(metadata, hints)?; - Self::put_flow_extensions(metadata, flow_extensions)?; - Self::put_snapshot_seqs(metadata, snapshot_seqs)?; + options.apply_to(&mut request)?; let mut client = self.client.make_flight_client(false, false)?; + let remote_addr = client.addr().to_string(); let response = client.mut_inner().do_get(request).await.or_else(|e| { let tonic_code = e.code(); @@ -796,7 +807,7 @@ impl Database { e ); Err(BoxedError::new(e)).with_context(|_| FlightGetSnafu { - addr: client.addr().to_string(), + addr: remote_addr.clone(), tonic_code, }) })?; @@ -808,7 +819,7 @@ impl Database { future::ready(decode_flight_data(&mut decoder, flight_data)) }); - output_from_flight_message_stream(flight_message_stream).await + output_from_flight_message_stream(remote_addr, flight_message_stream).await } /// Ingest a stream of [RecordBatch]es that belong to a table, using Arrow Flight's "`DoPut`" @@ -847,6 +858,98 @@ impl Database { } } +impl<'a> DatabaseFlightRequest<'a> { + /// Adds query optimization hints to this Flight request. + pub fn with_hints(mut self, hints: &[(&str, &str)]) -> Self { + self.options.hints = Database::encode_hints(hints); + self + } + + /// Adds Flow extensions to this Flight request. + pub fn with_flow_extensions(mut self, flow_extensions: &[(&str, &str)]) -> Self { + self.options.flow_extensions = Database::encode_flow_extensions(flow_extensions); + self + } + + /// Adds snapshot sequence fences to this Flight request. + pub fn with_snapshot_seqs(mut self, snapshot_seqs: &HashMap) -> Self { + self.options.snapshot_seqs = Database::encode_snapshot_seqs(snapshot_seqs); + self + } + + /// Sets a timeout for this Flight request only. + pub fn with_timeout(mut self, timeout: Duration) -> Self { + self.options.timeout = Some(timeout); + self + } + + /// Executes a SQL query. + pub async fn sql(self, sql: S) -> Result + where + S: AsRef, + { + let request = Request::Query(QueryRequest { + query: Some(Query::Sql(sql.as_ref().to_string())), + }); + self.do_get(request) + .await + .map(OutputWithMetrics::into_output) + } + + /// Executes a SQL query and returns terminal metrics. + pub async fn sql_with_terminal_metrics(self, sql: S) -> Result + where + S: AsRef, + { + self.query_with_terminal_metrics(QueryRequest { + query: Some(Query::Sql(sql.as_ref().to_string())), + }) + .await + } + + /// Executes a logical plan directly without SQL parsing. + pub async fn logical_plan(self, logical_plan: Vec) -> Result { + self.query_with_terminal_metrics(QueryRequest { + query: Some(Query::LogicalPlan(logical_plan)), + }) + .await + .map(OutputWithMetrics::into_output) + } + + /// Executes a query and returns terminal metrics. + pub async fn query_with_terminal_metrics( + self, + request: QueryRequest, + ) -> Result { + self.do_get(Request::Query(request)).await + } + + /// Creates a new table using the provided table expression. + #[cfg(feature = "testing")] + pub async fn create(self, expr: CreateTableExpr) -> Result { + self.do_get(Request::Ddl(DdlRequest { + expr: Some(DdlExpr::CreateTable(expr)), + })) + .await + .map(OutputWithMetrics::into_output) + } + + /// Alters an existing table using the provided alter expression. + #[cfg(feature = "testing")] + pub async fn alter(self, expr: AlterTableExpr) -> Result { + self.do_get(Request::Ddl(DdlRequest { + expr: Some(DdlExpr::AlterTable(expr)), + })) + .await + .map(OutputWithMetrics::into_output) + } + + async fn do_get(self, request: Request) -> Result { + let Self { database, options } = self; + database.do_get(request, options).await + } +} + /// by grpc standard, only `Unavailable` is retryable, see: https://github.com/grpc/grpc/blob/master/doc/statuscodes.md#status-codes-and-their-use-in-grpc pub fn is_grpc_retryable(err: &tonic::Status) -> bool { matches!(err.code(), tonic::Code::Unavailable) @@ -934,12 +1037,14 @@ mod tests { #[test] fn test_put_flow_extensions_preserves_comma_bearing_values() { let mut metadata = MetadataMap::new(); - Database::put_flow_extensions( + Database::put_metadata_value( &mut metadata, - &[ + FLOW_EXTENSIONS_METADATA_KEY, + Database::encode_flow_extensions(&[ ("flow.return_region_seq", "true"), ("flow.incremental_after_seqs", r#"{"1":10,"2":20}"#), - ], + ]) + .unwrap(), ) .unwrap(); @@ -969,7 +1074,12 @@ mod tests { (9_007_199_254_740_993_u64, 9_007_199_254_740_995_u64), ]); - Database::put_snapshot_seqs(&mut metadata, &snapshot_seqs).unwrap(); + Database::put_metadata_value( + &mut metadata, + SNAPSHOT_SEQS_METADATA_KEY, + Database::encode_snapshot_seqs(&snapshot_seqs).unwrap(), + ) + .unwrap(); let value = metadata .get(SNAPSHOT_SEQS_METADATA_KEY) @@ -980,6 +1090,50 @@ mod tests { assert_eq!(decoded, snapshot_seqs); } + #[test] + fn test_flight_request_builder_applies_request_options() { + let database = Database::new("greptime", "public", Client::default()); + let snapshot_seqs = HashMap::from([(42, 99)]); + let request = database + .flight_request() + .with_hints(&[("query_parallelism", "1")]) + .with_flow_extensions(&[("flow.return_region_seq", "true")]) + .with_snapshot_seqs(&snapshot_seqs) + .with_timeout(Duration::from_millis(50)); + let mut tonic_request = tonic::Request::new(()); + + request.options.apply_to(&mut tonic_request).unwrap(); + + let metadata = tonic_request.metadata(); + assert_eq!( + metadata.get(HINTS_METADATA_KEY).unwrap(), + "query_parallelism=1" + ); + assert_eq!( + serde_json::from_str::>( + metadata + .get(FLOW_EXTENSIONS_METADATA_KEY) + .unwrap() + .to_str() + .unwrap(), + ) + .unwrap(), + vec![("flow.return_region_seq".to_string(), "true".to_string())] + ); + assert_eq!( + serde_json::from_str::>( + metadata + .get(SNAPSHOT_SEQS_METADATA_KEY) + .unwrap() + .to_str() + .unwrap(), + ) + .unwrap(), + snapshot_seqs + ); + assert!(metadata.get("grpc-timeout").is_some()); + } + #[test] fn test_flight_ctx() { let mut ctx = FlightContext::default(); @@ -1020,6 +1174,28 @@ mod tests { assert_eq!(expected.should_retry(), actual.should_retry()); } + #[test] + fn test_flight_stream_error_preserves_addr_and_message() { + let error = flight_stream_error( + "127.0.0.1:4001", + Status::out_of_range("message length too large").into(), + ); + + assert!(matches!( + &error, + Error::FlightStream { + addr, + tonic_code: Code::OutOfRange, + message, + .. + } if addr == "127.0.0.1:4001" && message == "message length too large" + )); + assert_eq!( + "Failed to receive Flight data from 127.0.0.1:4001, code: Operation was attempted past the valid range: message length too large", + error.to_string(), + ); + } + #[test] fn test_from_tonic_status_with_retry_hint() { let mut headers = HeaderMap::new(); @@ -1117,13 +1293,13 @@ mod tests { #[tokio::test] async fn test_affected_rows_inline_metrics_are_parsed() { - let output = output_from_flight_message_stream(futures_util::stream::iter(vec![Ok( - FlightMessage::AffectedRows { + let output = output_from_flight_message_stream( + "test-peer".to_string(), + futures_util::stream::iter(vec![Ok(FlightMessage::AffectedRows { rows: 3, metrics: Some(terminal_metrics_json()), - }, - )] - as Vec>)) + })] as Vec>), + ) .await .unwrap(); @@ -1138,14 +1314,16 @@ mod tests { #[tokio::test] async fn test_affected_rows_inline_metrics_rejects_trailing_metrics() { let metrics_json = terminal_metrics_json(); - let err = output_from_flight_message_stream(futures_util::stream::iter(vec![ - Ok(FlightMessage::AffectedRows { - rows: 3, - metrics: Some(metrics_json.clone()), - }), - Ok(FlightMessage::Metrics(metrics_json)), - ] - as Vec>)) + let err = output_from_flight_message_stream( + "test-peer".to_string(), + futures_util::stream::iter(vec![ + Ok(FlightMessage::AffectedRows { + rows: 3, + metrics: Some(metrics_json.clone()), + }), + Ok(FlightMessage::Metrics(metrics_json)), + ] as Vec>), + ) .await .unwrap_err(); @@ -1167,12 +1345,14 @@ mod tests { vec![Arc::new(Int32Vector::from_slice([1])) as VectorRef], ) .unwrap(); - let output = output_from_flight_message_stream(futures_util::stream::iter(vec![ - Ok(FlightMessage::Schema(schema.arrow_schema().clone())), - Ok(FlightMessage::RecordBatch(batch.into_df_record_batch())), - Ok(FlightMessage::Metrics("{not-json}".to_string())), - ] - as Vec>)) + let output = output_from_flight_message_stream( + "test-peer".to_string(), + futures_util::stream::iter(vec![ + Ok(FlightMessage::Schema(schema.arrow_schema().clone())), + Ok(FlightMessage::RecordBatch(batch.into_df_record_batch())), + Ok(FlightMessage::Metrics("{not-json}".to_string())), + ] as Vec>), + ) .await .unwrap(); let terminal_metrics = output.metrics.clone(); @@ -1211,18 +1391,20 @@ mod tests { vec![Arc::new(Int32Vector::from_slice([2])) as VectorRef], ) .unwrap(); - let output = output_from_flight_message_stream(futures_util::stream::iter(vec![ - Ok(FlightMessage::Schema(schema.arrow_schema().clone())), - Ok(FlightMessage::RecordBatch( - first_batch.into_df_record_batch(), - )), - Ok(FlightMessage::Metrics(terminal_metrics_json_with_seq(1))), - Ok(FlightMessage::RecordBatch( - second_batch.into_df_record_batch(), - )), - Ok(FlightMessage::Metrics(terminal_metrics_json_with_seq(2))), - ] - as Vec>)) + let output = output_from_flight_message_stream( + "test-peer".to_string(), + futures_util::stream::iter(vec![ + Ok(FlightMessage::Schema(schema.arrow_schema().clone())), + Ok(FlightMessage::RecordBatch( + first_batch.into_df_record_batch(), + )), + Ok(FlightMessage::Metrics(terminal_metrics_json_with_seq(1))), + Ok(FlightMessage::RecordBatch( + second_batch.into_df_record_batch(), + )), + Ok(FlightMessage::Metrics(terminal_metrics_json_with_seq(2))), + ] as Vec>), + ) .await .unwrap(); let terminal_metrics = output.metrics.clone(); diff --git a/src/client/src/error.rs b/src/client/src/error.rs index 1db0858f087..168ba422715 100644 --- a/src/client/src/error.rs +++ b/src/client/src/error.rs @@ -40,6 +40,21 @@ pub enum Error { source: BoxedError, }, + #[snafu(display( + "Failed to receive Flight data from {}, code: {}: {}", + addr, + tonic_code, + message + ))] + FlightStream { + addr: String, + tonic_code: Code, + message: String, + source: BoxedError, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Failed to convert FlightData"))] ConvertFlightData { #[snafu(implicit)] @@ -154,6 +169,7 @@ impl ErrorExt for Error { Error::Server { code, .. } | Error::Tonic { code, .. } => *code, Error::FlightGet { source, .. } + | Error::FlightStream { source, .. } | Error::RegionServer { source, .. } | Error::FlowServer { source, .. } => source.status_code(), Error::CreateChannel { source, .. } @@ -174,6 +190,7 @@ impl ErrorExt for Error { match self { Error::Tonic { retry_hint, .. } => *retry_hint, Error::FlightGet { source, .. } + | Error::FlightStream { source, .. } | Error::RegionServer { source, .. } | Error::FlowServer { source, .. } | Error::External { source, .. } => source.retry_hint(), @@ -193,6 +210,7 @@ impl Error { pub fn tonic_code(&self) -> Option { match self { Self::FlightGet { tonic_code, .. } + | Self::FlightStream { tonic_code, .. } | Self::RegionServer { code: tonic_code, .. } diff --git a/src/client/src/flight.rs b/src/client/src/flight.rs index 4d00cba962e..1753fdd5246 100644 --- a/src/client/src/flight.rs +++ b/src/client/src/flight.rs @@ -12,12 +12,82 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::pin::Pin; + use arrow_flight::FlightData; use common_grpc::flight::{FlightDecoder, FlightMessage}; -use snafu::ResultExt; +use futures_util::stream::Peekable; +use futures_util::{Stream, StreamExt}; +use snafu::{OptionExt, ResultExt}; use crate::Result; -use crate::error::{ConvertFlightDataSnafu, Error}; +use crate::error::{ConvertFlightDataSnafu, Error, IllegalFlightMessagesSnafu}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum FlightMessageKind { + Schema, + RecordBatch, + AffectedRows, + Metrics, +} + +impl From<&FlightMessage> for FlightMessageKind { + fn from(message: &FlightMessage) -> Self { + match message { + FlightMessage::Schema(_) => Self::Schema, + FlightMessage::RecordBatch(_) => Self::RecordBatch, + FlightMessage::AffectedRows { .. } => Self::AffectedRows, + FlightMessage::Metrics(_) => Self::Metrics, + } + } +} + +pub(crate) struct FlightMessageReader { + /// Remote Flight peer associated with this response stream. + remote_addr: String, + messages: Peekable, +} + +impl FlightMessageReader +where + S: Stream> + Unpin, +{ + pub(crate) fn new(remote_addr: impl Into, messages: S) -> Self { + Self { + remote_addr: remote_addr.into(), + messages: messages.peekable(), + } + } + + pub(crate) fn remote_addr(&self) -> &str { + &self.remote_addr + } + + pub(crate) async fn read_first(&mut self) -> Result { + self.read_next().await?.context(IllegalFlightMessagesSnafu { + reason: "Expect the response not to be empty", + }) + } + + pub(crate) async fn read_next(&mut self) -> Result> { + self.messages.next().await.transpose() + } + + pub(crate) async fn peek_next_message_kind(&mut self) -> Result> { + match Pin::new(&mut self.messages).peek().await { + Some(Ok(message)) => Ok(Some(message.into())), + None => Ok(None), + Some(Err(_)) => match self.read_next().await { + // `peek` only borrows the error; consume it to preserve the source error. + Err(error) => Err(error), + Ok(_) => IllegalFlightMessagesSnafu { + reason: "Flight stream changed after peek".to_string(), + } + .fail(), + }, + } + } +} pub(crate) fn decode_flight_data( decoder: &mut FlightDecoder, diff --git a/src/client/src/region.rs b/src/client/src/region.rs index 2bb5f7dd0a6..cd155810ea9 100644 --- a/src/client/src/region.rs +++ b/src/client/src/region.rs @@ -47,7 +47,7 @@ use crate::error::{ self, FlightGetSnafu, IllegalDatabaseResponseSnafu, IllegalFlightMessagesSnafu, MissingFieldSnafu, Result, ServerSnafu, }; -use crate::flight::decode_flight_data; +use crate::flight::{FlightMessageKind, FlightMessageReader, decode_flight_data}; use crate::{Client, metrics}; const FLIGHT_DO_GET_TIMEOUT: Duration = Duration::from_secs(10); @@ -203,26 +203,22 @@ impl RegionRequester { async fn recordbatches_from_flight_message_stream( addr: String, - mut flight_message_stream: S, + flight_message_stream: S, ) -> Result where S: Stream> + Send + Unpin + 'static, { - let Some(first_flight_message) = flight_message_stream.next().await else { - return IllegalFlightMessagesSnafu { - reason: "Expect the response not to be empty", - } - .fail() - .map_err(|error| flight_stream_error(&addr, error)); - }; - let FlightMessage::Schema(schema) = - first_flight_message.map_err(|e| flight_stream_error(&addr, e))? + let mut reader = FlightMessageReader::new(addr.clone(), flight_message_stream); + let FlightMessage::Schema(schema) = reader + .read_first() + .await + .map_err(|error| flight_stream_error(reader.remote_addr(), error))? else { return IllegalFlightMessagesSnafu { reason: "Expect schema to be the first flight message", } .fail() - .map_err(|error| flight_stream_error(&addr, error)); + .map_err(|error| flight_stream_error(reader.remote_addr(), error)); }; let metrics = Arc::new(ArcSwapOption::from(None)); @@ -233,31 +229,23 @@ where let schema = Arc::new(datatypes::schema::Schema::try_from(schema).context(error::ConvertSchemaSnafu)?); let schema_cloned = schema.clone(); - let stream_addr = addr.clone(); + let stream_addr = addr; let stream = Box::pin(stream!({ let _span = tracing_context.attach(common_telemetry::tracing::info_span!( "poll_flight_data_stream" )); - let mut buffered_message: Option = None; let mut stream_ended = false; while !stream_ended { - // get the next message from the buffered message or read from the flight message stream - let flight_message_item = if let Some(msg) = buffered_message.take() { - Some(Ok(msg)) - } else { - flight_message_stream.next().await - }; - - let flight_message = match flight_message_item { - Some(Ok(message)) => message, - Some(Err(e)) => { - yield Err(BoxedError::new(flight_stream_error(&stream_addr, e))) + let flight_message = match reader.read_next().await { + Ok(Some(message)) => message, + Ok(None) => break, + Err(error) => { + yield Err(BoxedError::new(flight_stream_error(&stream_addr, error))) .context(ExternalSnafu); break; } - None => break, }; match flight_message { @@ -265,42 +253,54 @@ where let result_to_yield = RecordBatch::from_df_record_batch(schema_cloned.clone(), record_batch); - // get the next message from the stream. normally it should be a metrics message. - if let Some(next_flight_message_result) = flight_message_stream.next().await { - match next_flight_message_result { - Ok(FlightMessage::Metrics(s)) => { - let m = serde_json::from_str(&s).ok().map(Arc::new); - metrics_ref.swap(m); - } - Ok(FlightMessage::RecordBatch(rb)) => { - // for some reason it's not a metrics message, so we need to buffer this record batch - // and yield it in the next iteration. - buffered_message = Some(FlightMessage::RecordBatch(rb)); - } - Ok(_) => { - yield IllegalFlightMessagesSnafu { - reason: "A RecordBatch message can only be succeeded by a Metrics message or another RecordBatch message" - } - .fail() - .map_err(BoxedError::new) - .context(ExternalSnafu); - break; - } - Err(e) => { - yield Err(BoxedError::new(flight_stream_error(&stream_addr, e))) + // Metrics follow a batch so MergeScan can observe them before yielding it. + match reader.peek_next_message_kind().await { + Ok(Some(FlightMessageKind::Metrics)) => { + let metrics_message = match reader.read_next().await { + Ok(Some(FlightMessage::Metrics(metrics))) => metrics, + Ok(Some(_) | None) => { + yield IllegalFlightMessagesSnafu { + reason: "Flight stream changed after peek", + } + .fail() + .map_err(BoxedError::new) .context(ExternalSnafu); - break; + break; + } + Err(error) => { + yield Err(BoxedError::new(flight_stream_error( + &stream_addr, + error, + ))) + .context(ExternalSnafu); + break; + } + }; + let metrics = serde_json::from_str(&metrics_message).ok().map(Arc::new); + metrics_ref.swap(metrics); + } + Ok(Some(FlightMessageKind::RecordBatch)) => {} + Ok(Some(FlightMessageKind::Schema | FlightMessageKind::AffectedRows)) => { + yield IllegalFlightMessagesSnafu { + reason: "A RecordBatch message can only be succeeded by a Metrics message or another RecordBatch message" } + .fail() + .map_err(BoxedError::new) + .context(ExternalSnafu); + break; + } + Ok(None) => stream_ended = true, + Err(error) => { + yield Err(BoxedError::new(flight_stream_error(&stream_addr, error))) + .context(ExternalSnafu); + break; } - } else { - // the stream has ended - stream_ended = true; } yield Ok(result_to_yield); } FlightMessage::Metrics(s) => { - // just a branch in case of some metrics message comes after other things. + // Metrics may arrive before the next RecordBatch. let m = serde_json::from_str(&s).ok().map(Arc::new); metrics_ref.swap(m); continue; @@ -412,6 +412,7 @@ mod test { RemoteDynFilterUnregister, RemoteDynFilterUpdate, region_request, remote_dyn_filter_request, }; use common_recordbatch::adapter::RecordBatchMetrics; + use datatypes::arrow::array::Int32Array; use datatypes::prelude::{ConcreteDataType, VectorRef}; use datatypes::schema::{ColumnSchema, Schema}; use datatypes::vectors::Int32Vector; @@ -598,6 +599,86 @@ mod test { assert_eq!(metrics.elapsed_compute, 7); } + #[tokio::test] + async fn test_record_batch_stream_updates_following_metrics_before_yielding_batch() { + let schema = test_schema(); + let batch = RecordBatch::new( + schema.clone(), + vec![Arc::new(Int32Vector::from_slice([1])) as VectorRef], + ) + .unwrap(); + + let mut recordbatches = recordbatches_from_flight_message_stream( + "test-peer".to_string(), + stream::iter(vec![ + Ok(FlightMessage::Schema(schema.arrow_schema().clone())), + Ok(FlightMessage::RecordBatch(batch.into_df_record_batch())), + Ok(FlightMessage::Metrics(test_metrics_json())), + ]), + ) + .await + .unwrap(); + + let batch = recordbatches.next().await.unwrap().unwrap(); + assert_eq!(batch.num_rows(), 1); + + let metrics = recordbatches.metrics().unwrap(); + assert_eq!(metrics.elapsed_compute, 7); + assert!(recordbatches.next().await.is_none()); + } + + #[tokio::test] + async fn test_record_batch_stream_preserves_peeked_record_batch() { + let schema = test_schema(); + let first_batch = RecordBatch::new( + schema.clone(), + vec![Arc::new(Int32Vector::from_slice([1])) as VectorRef], + ) + .unwrap(); + let second_batch = RecordBatch::new( + schema.clone(), + vec![Arc::new(Int32Vector::from_slice([2])) as VectorRef], + ) + .unwrap(); + + let mut recordbatches = recordbatches_from_flight_message_stream( + "test-peer".to_string(), + stream::iter(vec![ + Ok(FlightMessage::Schema(schema.arrow_schema().clone())), + Ok(FlightMessage::RecordBatch( + first_batch.into_df_record_batch(), + )), + Ok(FlightMessage::RecordBatch( + second_batch.into_df_record_batch(), + )), + ]), + ) + .await + .unwrap(); + + let first_batch = recordbatches.next().await.unwrap().unwrap(); + let second_batch = recordbatches.next().await.unwrap().unwrap(); + assert_eq!( + first_batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + 1 + ); + assert_eq!( + second_batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + 2 + ); + assert!(recordbatches.next().await.is_none()); + } + #[tokio::test] async fn test_record_batch_stream_exposes_error_after_pre_batch_metrics() { let schema = test_schema(); diff --git a/src/cmd/Cargo.toml b/src/cmd/Cargo.toml index 579e5b1454c..8cb8cff4300 100644 --- a/src/cmd/Cargo.toml +++ b/src/cmd/Cargo.toml @@ -118,6 +118,7 @@ tonic.workspace = true tracing-appender.workspace = true [target.'cfg(unix)'.dependencies] +daemonize = "0.5" pprof = { version = "0.14", features = [ "flamegraph", ] } diff --git a/src/cmd/src/bin/greptime.rs b/src/cmd/src/bin/greptime.rs index f5f4a590487..2fa75ff4171 100644 --- a/src/cmd/src/bin/greptime.rs +++ b/src/cmd/src/bin/greptime.rs @@ -17,6 +17,8 @@ use clap::{Parser, Subcommand}; use cmd::datanode::builder::InstanceBuilder; +#[cfg(unix)] +use cmd::error::DaemonizeSnafu; use cmd::error::{InitTlsProviderSnafu, Result}; use cmd::options::GlobalOptions; use cmd::{App, cli, datanode, flownode, frontend, metasrv, standalone, user}; @@ -73,20 +75,24 @@ static ALLOC: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc; #[cfg(debug_assertions)] fn main() -> Result<()> { use snafu::ResultExt; + + let cli = Command::parse(); + maybe_daemonize(&cli)?; + // Set the stack size to 8MB for the thread so it wouldn't overflow on large stack usage in debug mode // see https://github.com/GreptimeTeam/greptimedb/pull/4317 // and https://github.com/rust-lang/rust/issues/34283 std::thread::Builder::new() .name("main_spawn".to_string()) .stack_size(8 * 1024 * 1024) - .spawn(|| { + .spawn(move || { { tokio::runtime::Builder::new_multi_thread() .thread_stack_size(8 * 1024 * 1024) .enable_all() .build() .expect("Failed building the Runtime") - .block_on(main_body()) + .block_on(main_body(cli)) } }) .context(cmd::error::SpawnThreadSnafu)? @@ -95,15 +101,53 @@ fn main() -> Result<()> { } #[cfg(not(debug_assertions))] -#[tokio::main] -async fn main() -> Result<()> { - main_body().await +fn main() -> Result<()> { + let cli = Command::parse(); + maybe_daemonize(&cli)?; + + tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build() + .expect("Failed building the Runtime") + .block_on(main_body(cli)) } -async fn main_body() -> Result<()> { +async fn main_body(cli: Command) -> Result<()> { setup_human_panic(); install_default_crypto_provider().map_err(|msg| InitTlsProviderSnafu { msg }.build())?; - start(Command::parse()).await + start(cli).await +} + +/// If `--daemon` was passed for `standalone start`, fork into the background +/// before any Tokio runtime is built. On non-Unix platforms the `--daemon` +/// flag does not exist, so this is a no-op. +fn maybe_daemonize(cli: &Command) -> Result<()> { + let daemon = matches!(&cli.subcmd, SubCommand::Standalone(cmd) if cmd.is_daemon()); + if !daemon { + return Ok(()); + } + + #[cfg(unix)] + { + let cwd = std::env::current_dir().map_err(|e| { + DaemonizeSnafu { + msg: format!("failed to read current working directory: {e}"), + } + .build() + })?; + + daemonize::Daemonize::new() + .working_directory(cwd) + .start() + .map_err(|e| { + DaemonizeSnafu { + msg: format!("failed to daemonize: {e}"), + } + .build() + })?; + } + + Ok(()) } async fn start(cli: Command) -> Result<()> { diff --git a/src/cmd/src/error.rs b/src/cmd/src/error.rs index b80828d13ac..82273c0b67d 100644 --- a/src/cmd/src/error.rs +++ b/src/cmd/src/error.rs @@ -162,6 +162,13 @@ pub enum Error { location: Location, }, + #[snafu(display("Failed to daemonize: {msg}"))] + Daemonize { + msg: String, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Unsupported selector type: {}", selector_type))] UnsupportedSelectorType { selector_type: String, @@ -361,7 +368,8 @@ impl ErrorExt for Error { Error::SerdeJson { .. } | Error::FileIo { .. } | Error::SpawnThread { .. } - | Error::InitTlsProvider { .. } => StatusCode::Unexpected, + | Error::InitTlsProvider { .. } + | Error::Daemonize { .. } => StatusCode::Unexpected, Error::Other { source, .. } => source.status_code(), diff --git a/src/cmd/src/frontend.rs b/src/cmd/src/frontend.rs index e3737088e8d..6cae195f3e6 100644 --- a/src/cmd/src/frontend.rs +++ b/src/cmd/src/frontend.rs @@ -13,7 +13,9 @@ // limitations under the License. use std::fmt::Debug; +use std::future::Future; use std::path::Path; +use std::pin::Pin; use std::sync::Arc; use std::time::Duration; @@ -61,6 +63,7 @@ use crate::options::{GlobalOptions, GreptimeOptions}; use crate::{App, create_resource_limit_metrics, log_versions, maybe_activate_heap_profile}; type FrontendOptions = GreptimeOptions; +type HeartbeatExtensionSetupFuture<'a> = Pin> + Send + 'a>>; pub struct Instance { frontend: Frontend, @@ -117,7 +120,35 @@ pub struct Command { impl Command { pub async fn build(&self, opts: FrontendOptions) -> Result { - self.subcmd.build(opts).await + self.build_with_heartbeat_extensions(opts, FrontendHeartbeatExtensions::default()) + .await + } + + /// Builds a frontend with pre-registered heartbeat extensions. + /// + /// Register extensions before calling this method. The supplied registry is installed before + /// normal plugin setup. After plugin heartbeat setup completes, the registry is frozen before + /// heartbeat handlers and the heartbeat task are built, so later registrations are rejected + /// and every heartbeat consumer observes the same extension membership. + pub async fn build_with_heartbeat_extensions( + &self, + opts: FrontendOptions, + heartbeat_extensions: FrontendHeartbeatExtensions, + ) -> Result { + let plugins = plugins_with_heartbeat_extensions(heartbeat_extensions); + self.forward_build_with_plugins(opts, plugins, |command, opts, plugins| { + command.build_with_plugins(opts, plugins) + }) + .await + } + + fn forward_build_with_plugins<'a, T>( + &'a self, + opts: FrontendOptions, + plugins: Plugins, + build: impl FnOnce(&'a StartCommand, FrontendOptions, Plugins) -> T, + ) -> T { + self.subcmd.forward_build_with_plugins(opts, plugins, build) } pub fn load_options(&self, global_options: &GlobalOptions) -> Result { @@ -131,9 +162,14 @@ pub enum SubCommand { } impl SubCommand { - async fn build(&self, opts: FrontendOptions) -> Result { + fn forward_build_with_plugins<'a, T>( + &'a self, + opts: FrontendOptions, + plugins: Plugins, + build: impl FnOnce(&'a StartCommand, FrontendOptions, Plugins) -> T, + ) -> T { match self { - SubCommand::Start(cmd) => cmd.build(opts).await, + SubCommand::Start(cmd) => build(cmd, opts, plugins), } } @@ -329,7 +365,11 @@ impl StartCommand { Ok(()) } - async fn build(&self, opts: FrontendOptions) -> Result { + async fn build_with_plugins( + &self, + opts: FrontendOptions, + plugins: Plugins, + ) -> Result { let guard = common_telemetry::init_global_logging( APP_NAME, &opts.component.logging, @@ -380,15 +420,8 @@ impl StartCommand { .await .context(error::MetaClientInitSnafu)?; - let mut plugins = Plugins::new(); - plugins::setup_frontend_plugins_pre_build( - &mut plugins, - &plugin_opts, - &opts, - Some(&meta_config), - ) - .await - .context(error::StartFrontendSnafu)?; + let mut plugins = + prepare_frontend_plugins(plugins, &plugin_opts, &opts, Some(&meta_config)).await?; // now initialize the meta_client with plugins let meta_client = meta_client::create_meta_client( @@ -434,6 +467,8 @@ impl StartCommand { // Some queries are expected to take long time. let mut channel_config = opts.datanode.client.channel_config(); channel_config.timeout = None; + // Source Flight streams and sink unary responses share pooled connections. + channel_config.http2_adaptive_window = Some(true); if opts.grpc.flight_compression.transport_compression() { channel_config.accept_compression = true; channel_config.send_compression = true; @@ -497,12 +532,20 @@ impl StartCommand { let instance = Arc::new(instance); - plugins::setup_frontend_heartbeat_extensions(&mut plugins, &plugin_opts, &instance) - .await - .context(error::StartFrontendSnafu)?; - let heartbeat_extensions = plugins - .get::() - .unwrap_or_default(); + let heartbeat_instance = instance.clone(); + let heartbeat_extensions = + setup_and_freeze_frontend_heartbeat_extensions(&mut plugins, move |plugins| { + Box::pin(async move { + plugins::setup_frontend_heartbeat_extensions( + plugins, + &plugin_opts, + &heartbeat_instance, + ) + .await + .context(error::StartFrontendSnafu) + }) + }) + .await?; let heartbeat_task = Some(create_heartbeat_task_with_extensions( &opts, meta_client, @@ -524,6 +567,36 @@ impl StartCommand { } } +async fn prepare_frontend_plugins( + mut plugins: Plugins, + plugin_opts: &[PluginOptions], + opts: &frontend::frontend::FrontendOptions, + meta_config: Option<&[PluginOptions]>, +) -> Result { + plugins::setup_frontend_plugins_pre_build(&mut plugins, plugin_opts, opts, meta_config) + .await + .context(error::StartFrontendSnafu)?; + Ok(plugins) +} + +fn plugins_with_heartbeat_extensions(heartbeat_extensions: FrontendHeartbeatExtensions) -> Plugins { + let plugins = Plugins::new(); + plugins.insert(heartbeat_extensions); + plugins +} + +async fn setup_and_freeze_frontend_heartbeat_extensions( + plugins: &mut Plugins, + setup: impl for<'a> FnOnce(&'a mut Plugins) -> HeartbeatExtensionSetupFuture<'a>, +) -> Result { + setup(plugins).await?; + let extensions = plugins + .get::() + .unwrap_or_default(); + extensions.freeze(); + Ok(extensions) +} + pub fn create_heartbeat_task( options: &frontend::frontend::FrontendOptions, meta_client: MetaClientRef, @@ -563,6 +636,7 @@ fn create_heartbeat_task_with_extensions( #[cfg(test)] mod tests { use std::io::Write; + use std::sync::Arc; use std::time::Duration; use auth::{Identity, Password, UserProviderRef}; @@ -576,6 +650,109 @@ mod tests { use super::*; use crate::options::GlobalOptions; + #[derive(Debug)] + struct TestExtension(&'static str); + + #[async_trait] + impl frontend::heartbeat::FrontendHeartbeatExtension for TestExtension { + fn name(&self) -> &'static str { + self.0 + } + } + + #[test] + fn test_command_forwards_heartbeat_extensions_to_start_build() { + let command = Command { + subcmd: SubCommand::Start(StartCommand::default()), + }; + let extensions = FrontendHeartbeatExtensions::default(); + let shared_extensions = extensions.clone(); + + let forwarded_plugins = command.forward_build_with_plugins( + FrontendOptions::default(), + plugins_with_heartbeat_extensions(extensions), + |_, _, plugins| plugins, + ); + let forwarded_extensions = forwarded_plugins + .get::() + .unwrap(); + + assert_registry_identity(&shared_extensions, &forwarded_extensions); + } + + #[tokio::test] + async fn test_prefilled_heartbeat_extensions_survive_pre_build_setup() { + let extensions = FrontendHeartbeatExtensions::default(); + let shared_extensions = extensions.clone(); + + let options = frontend::frontend::FrontendOptions { + meta_client: Some(MetaClientOptions::default()), + ..Default::default() + }; + let plugins = prepare_frontend_plugins( + plugins_with_heartbeat_extensions(extensions), + &[], + &options, + None, + ) + .await + .unwrap(); + + let setup_extensions = plugins.get_or_insert(FrontendHeartbeatExtensions::default); + assert_registry_identity(&shared_extensions, &setup_extensions); + } + + #[tokio::test] + async fn test_heartbeat_extension_setup_completes_before_freeze() { + let extensions = FrontendHeartbeatExtensions::default(); + assert_eq!( + extensions.try_register(Arc::new(TestExtension("caller"))), + Ok(()) + ); + let mut plugins = plugins_with_heartbeat_extensions(extensions.clone()); + + let finalized = setup_and_freeze_frontend_heartbeat_extensions(&mut plugins, |plugins| { + Box::pin(async move { + let setup_extensions = plugins.get_or_insert(FrontendHeartbeatExtensions::default); + assert_eq!( + setup_extensions.try_register(Arc::new(TestExtension("plugin"))), + Ok(()) + ); + Ok(()) + }) + }) + .await + .unwrap(); + + assert_eq!( + finalized + .extensions() + .iter() + .map(|extension| extension.name()) + .collect::>(), + ["caller", "plugin"] + ); + assert_eq!( + extensions.try_register(Arc::new(TestExtension("late"))), + Err(frontend::heartbeat::RegistrationError::Frozen) + ); + assert_eq!(finalized.len(), 2); + } + + fn assert_registry_identity( + expected: &FrontendHeartbeatExtensions, + actual: &FrontendHeartbeatExtensions, + ) { + let extension = Arc::new(TestExtension("cmd-test-extension")); + assert!(expected.register(extension.clone())); + let registered = actual.extensions(); + assert_eq!(registered.len(), 1); + assert!(Arc::ptr_eq( + &(extension as Arc), + ®istered[0] + )); + } + #[test] fn test_try_from_start_command() { let command = StartCommand { diff --git a/src/cmd/src/standalone.rs b/src/cmd/src/standalone.rs index 2b5bd5e8cec..d628461ba75 100644 --- a/src/cmd/src/standalone.rs +++ b/src/cmd/src/standalone.rs @@ -152,6 +152,11 @@ impl Command { ) -> Result> { self.subcmd.load_options(global_options) } + + /// Whether the `standalone start` command requested daemonization. + pub fn is_daemon(&self) -> bool { + self.subcmd.is_daemon() + } } #[derive(Parser)] @@ -174,6 +179,12 @@ impl SubCommand { SubCommand::Start(cmd) => cmd.load_options(global_options), } } + + fn is_daemon(&self) -> bool { + match self { + SubCommand::Start(cmd) => cmd.is_daemon(), + } + } } pub struct Instance { @@ -291,9 +302,24 @@ pub struct StartCommand { /// The working home directory of this standalone instance. #[clap(long)] data_home: Option, + /// Run in the background as a daemon. + #[cfg(unix)] + #[clap(short, long)] + daemon: bool, } impl StartCommand { + /// Whether the `standalone start` command requested daemonization. + #[cfg(unix)] + pub(crate) fn is_daemon(&self) -> bool { + self.daemon + } + + #[cfg(not(unix))] + pub(crate) fn is_daemon(&self) -> bool { + false + } + /// Load the GreptimeDB options from various sources (command line, config file or env). pub fn load_options( &self, @@ -1356,6 +1382,19 @@ mod tests { assert_eq!(command.grpc_bind_addr.as_deref(), Some("127.0.0.1:34001")); } + #[cfg(unix)] + #[test] + fn test_parse_daemon_flag() { + let command = StartCommand::try_parse_from(["standalone", "--daemon"]).unwrap(); + assert!(command.is_daemon()); + + let command = StartCommand::try_parse_from(["standalone", "-d"]).unwrap(); + assert!(command.is_daemon()); + + let command = StartCommand::try_parse_from(["standalone"]).unwrap(); + assert!(!command.is_daemon()); + } + #[test] fn test_help_uses_grpc_option_names() { let mut cmd = StartCommand::command(); diff --git a/src/cmd/tests/load_config_test.rs b/src/cmd/tests/load_config_test.rs index 70594d0653f..8af216256f9 100644 --- a/src/cmd/tests/load_config_test.rs +++ b/src/cmd/tests/load_config_test.rs @@ -113,7 +113,7 @@ fn test_load_datanode_example_config() { }, region_engine: vec![ RegionEngineConfig::Mito(MitoConfig { - auto_flush_interval: Duration::from_secs(3600), + auto_flush_interval: Duration::from_secs(10 * 60), default_region_write_buffer_size: ReadableSize::mb(0), write_cache_ttl: Some(Duration::from_secs(60 * 60 * 8)), scan_memory_limit: MemoryLimit::Unlimited, @@ -354,7 +354,7 @@ fn test_load_standalone_example_config() { }), region_engine: vec![ RegionEngineConfig::Mito(MitoConfig { - auto_flush_interval: Duration::from_secs(3600), + auto_flush_interval: Duration::from_secs(10 * 60), default_region_write_buffer_size: ReadableSize::mb(0), write_cache_ttl: Some(Duration::from_secs(60 * 60 * 8)), scan_memory_limit: MemoryLimit::Unlimited, diff --git a/src/cmd/tests/standalone_daemon.rs b/src/cmd/tests/standalone_daemon.rs new file mode 100644 index 00000000000..5b7c7427d5f --- /dev/null +++ b/src/cmd/tests/standalone_daemon.rs @@ -0,0 +1,109 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#![cfg(unix)] + +use std::net::TcpStream; +use std::process::{Command, Stdio}; +use std::time::{Duration, Instant}; + +/// Stops the orphaned background daemon when the test finishes (including on +/// panic). The daemon's command line contains the unique `--data-home` path, so +/// we match on it to locate the process. +struct DaemonGuard(String); + +impl Drop for DaemonGuard { + fn drop(&mut self) { + let _ = Command::new("pkill") + .args(["-f", &self.0]) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .status(); + } +} + +#[test] +#[cfg(unix)] +fn test_standalone_daemon_mode_detaches_from_shell() { + let bin = env!("CARGO_BIN_EXE_greptime"); + let data_home = tempfile::tempdir().unwrap(); + let data_home_path = data_home.path().to_str().unwrap().to_string(); + + let http_addr = format!("127.0.0.1:{}", common_test_util::ports::get_port()); + let grpc_addr = format!("127.0.0.1:{}", common_test_util::ports::get_port()); + let mysql_addr = format!("127.0.0.1:{}", common_test_util::ports::get_port()); + let postgres_addr = format!("127.0.0.1:{}", common_test_util::ports::get_port()); + + let mut child = Command::new(bin) + .args([ + "standalone", + "start", + "--daemon", + "--data-home", + &data_home_path, + "--http-addr", + &http_addr, + "--grpc-bind-addr", + &grpc_addr, + "--mysql-addr", + &mysql_addr, + "--postgres-addr", + &postgres_addr, + ]) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .spawn() + .unwrap(); + + let _guard = DaemonGuard(data_home_path); + + // When daemonization works, the parent forks and exits right away. When it + // does not, the process stays in the foreground and never returns. + let detach_timeout = Duration::from_secs(30); + let started = Instant::now(); + loop { + match child.try_wait().unwrap() { + Some(status) => { + assert!( + status.success(), + "daemon start should exit successfully, got {status:?}" + ); + break; + } + None => { + assert!( + started.elapsed() < detach_timeout, + "daemon mode should detach and return quickly, but it blocked for over {detach_timeout:?}" + ); + std::thread::sleep(Duration::from_millis(200)); + } + } + } + + // A successful detach is not enough: a broken implementation could fork, + // have the parent exit, then crash the child. Poll the HTTP address to + // confirm the daemon actually came up and is listening. + let ready_timeout = Duration::from_secs(30); + let started = Instant::now(); + loop { + if TcpStream::connect(&http_addr).is_ok() { + break; + } + assert!( + started.elapsed() < ready_timeout, + "daemon should come up and listen on {http_addr}, but it was unreachable for {ready_timeout:?}" + ); + std::thread::sleep(Duration::from_millis(200)); + } +} diff --git a/src/common/catalog/src/consts.rs b/src/common/catalog/src/consts.rs index b976cbc256e..63aa103b83d 100644 --- a/src/common/catalog/src/consts.rs +++ b/src/common/catalog/src/consts.rs @@ -220,6 +220,11 @@ pub const SEMANTIC_RELATIONSHIPS_TABLE_NAME: &str = "semantic_relationships"; /// `semantic_relationships`; see [`is_ddl_reserved_table`] for its lifecycle. pub const SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME: &str = "semantic_relationships_declared"; +/// Width of the window the graph derivation bins observations into. Sources +/// synthesizing observations must land a row in every window they describe, +/// so this is shared rather than restated per crate. +pub const SEMANTIC_GRAPH_WINDOW_NANOS: i64 = 60 * 1_000_000_000; + // Column names of the graph tables: `catalog` exposes these schemas and the // read-time plans in `operator` must project exactly them. pub const OBSERVED_AT_COLUMN: &str = "observed_at"; diff --git a/src/common/meta/Cargo.toml b/src/common/meta/Cargo.toml index 8e0cbbdf354..8deb3d00814 100644 --- a/src/common/meta/Cargo.toml +++ b/src/common/meta/Cargo.toml @@ -89,7 +89,7 @@ strum.workspace = true table = { workspace = true, features = ["testing"] } tokio.workspace = true tokio-postgres = { workspace = true, optional = true } -tokio-postgres-rustls = { version = "0.12", optional = true } +tokio-postgres-rustls = { version = "0.14", optional = true } tonic.workspace = true tracing.workspace = true typetag.workspace = true diff --git a/src/common/meta/src/datanode.rs b/src/common/meta/src/datanode.rs index 8af86462129..d50c37c381b 100644 --- a/src/common/meta/src/datanode.rs +++ b/src/common/meta/src/datanode.rs @@ -31,6 +31,8 @@ use crate::heartbeat::utils::get_datanode_workloads; const DATANODE_STAT_PREFIX: &str = "__meta_datanode_stat"; pub const REGION_STATISTIC_KEY: &str = "__region_statistic"; +/// The private table name for persisted region statistics. +pub const REGION_STATS_HISTORY_TABLE_NAME: &str = "region_statistics_history"; lazy_static! { pub(crate) static ref DATANODE_LEASE_KEY_PATTERN: Regex = diff --git a/src/common/meta/src/reconciliation.rs b/src/common/meta/src/reconciliation.rs index 351a415566a..06ff9e704af 100644 --- a/src/common/meta/src/reconciliation.rs +++ b/src/common/meta/src/reconciliation.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +pub(crate) mod event; pub mod manager; pub(crate) mod reconcile_catalog; pub(crate) mod reconcile_database; diff --git a/src/common/meta/src/reconciliation/event.rs b/src/common/meta/src/reconciliation/event.rs new file mode 100644 index 00000000000..9430fa575f0 --- /dev/null +++ b/src/common/meta/src/reconciliation/event.rs @@ -0,0 +1,351 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::any::Any; + +use api::v1::value::ValueData; +use api::v1::{ColumnSchema, Row}; +use common_event_recorder::Event; +use common_event_recorder::error::{Result, SerializeEventSnafu}; +use common_event_recorder::event_table::{ + CATALOG_NAME_COLUMN, PHYSICAL_TABLE_ID_COLUMN, SCHEMA_NAME_COLUMN, TABLE_ID_COLUMN, + TABLE_NAME_COLUMN, column_schemas, nullable_string, nullable_value, +}; +use serde::Serialize; +use snafu::ResultExt; +use store_api::storage::TableId; + +use crate::reconciliation::ResolveStrategy; + +/// Stable event type stored for physical table reconciliation procedures. +pub(crate) const RECONCILE_TABLE_EVENT_TYPE: &str = "reconcile_table"; +const PAYLOAD_VERSION: u8 = 1; + +#[derive(Debug, Clone, Copy)] +enum ReconciliationEventType { + Table, +} + +impl ReconciliationEventType { + const fn as_str(self) -> &'static str { + match self { + Self::Table => RECONCILE_TABLE_EVENT_TYPE, + } + } +} + +/// Nullable object locators shared by all reconciliation event types. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub(crate) struct ReconciliationLocator { + catalog_name: Option, + schema_name: Option, + table_name: Option, + table_id: Option, + physical_table_id: Option, +} + +impl ReconciliationLocator { + /// Creates a locator for a physical table with its fully qualified name and ID. + pub(crate) fn physical_table( + catalog_name: &str, + schema_name: &str, + table_name: &str, + table_id: TableId, + ) -> Self { + Self { + catalog_name: Some(catalog_name.to_string()), + schema_name: Some(schema_name.to_string()), + table_name: Some(table_name.to_string()), + table_id: Some(table_id), + ..Default::default() + } + } +} + +#[derive(Debug, Serialize)] +#[serde(untagged)] +enum ReconciliationPayload { + TableSubmitted(TableSubmittedPayload), + TableResult(TableResultPayload), +} + +#[derive(Debug, Serialize)] +struct TableSubmittedPayload { + version: u8, + resolve_strategy: &'static str, + is_subprocedure: bool, +} + +#[derive(Debug, Serialize)] +struct TableResultPayload { + version: u8, + complete: bool, + metadata_state: Option<&'static str>, + resolution_strategy_applied: Option<&'static str>, + resolved_column_count: Option, + scanned_region_count: usize, + updated_region_count: usize, + table_info_updated: bool, + last_completed_phase: Option<&'static str>, +} + +/// Event representation shared by reconciliation procedures. +#[derive(Debug)] +pub(crate) struct ReconciliationEvent { + event_type: ReconciliationEventType, + locator: ReconciliationLocator, + payload: Option, +} + +impl ReconciliationEvent { + /// Builds the bounded intent event emitted when table reconciliation is submitted. + pub(crate) fn table_submitted( + locator: ReconciliationLocator, + resolve_strategy: ResolveStrategy, + is_subprocedure: bool, + ) -> Self { + Self { + event_type: ReconciliationEventType::Table, + locator, + payload: Some(ReconciliationPayload::TableSubmitted( + TableSubmittedPayload { + version: PAYLOAD_VERSION, + resolve_strategy: resolve_strategy_name(resolve_strategy), + is_subprocedure, + }, + )), + } + } + + /// Builds a terminal event from the bounded reconciliation result summary. + #[allow(clippy::too_many_arguments)] + pub(crate) fn table_result( + locator: ReconciliationLocator, + complete: bool, + metadata_state: Option<&'static str>, + resolution_strategy_applied: Option, + resolved_column_count: Option, + scanned_region_count: usize, + updated_region_count: usize, + table_info_updated: bool, + last_completed_phase: Option<&'static str>, + ) -> Self { + Self { + event_type: ReconciliationEventType::Table, + locator, + payload: Some(ReconciliationPayload::TableResult(TableResultPayload { + version: PAYLOAD_VERSION, + complete, + metadata_state, + resolution_strategy_applied: resolution_strategy_applied.map(resolve_strategy_name), + resolved_column_count, + scanned_region_count, + updated_region_count, + table_info_updated, + last_completed_phase, + })), + } + } + + /// Builds a table lifecycle event whose reconciliation payload is null. + pub(crate) fn table_lifecycle(locator: ReconciliationLocator) -> Self { + Self::lifecycle(ReconciliationEventType::Table, locator) + } + + fn lifecycle(event_type: ReconciliationEventType, locator: ReconciliationLocator) -> Self { + Self { + event_type, + locator, + payload: None, + } + } + + fn schema() -> Vec { + column_schemas([ + &CATALOG_NAME_COLUMN, + &SCHEMA_NAME_COLUMN, + &TABLE_NAME_COLUMN, + &TABLE_ID_COLUMN, + &PHYSICAL_TABLE_ID_COLUMN, + ]) + } +} + +impl Event for ReconciliationEvent { + fn event_type(&self) -> &str { + self.event_type.as_str() + } + + fn json_payload(&self) -> Result { + match &self.payload { + Some(payload) => serde_json::to_value(payload).context(SerializeEventSnafu), + None => Ok(serde_json::Value::Null), + } + } + + fn extra_schema(&self) -> Vec { + Self::schema() + } + + fn extra_rows(&self) -> Result> { + Ok(vec![Row { + values: vec![ + nullable_string(self.locator.catalog_name.as_deref()), + nullable_string(self.locator.schema_name.as_deref()), + nullable_string(self.locator.table_name.as_deref()), + nullable_table_id(self.locator.table_id), + nullable_table_id(self.locator.physical_table_id), + ], + }]) + } + + fn as_any(&self) -> &dyn Any { + self + } +} + +fn resolve_strategy_name(strategy: ResolveStrategy) -> &'static str { + match strategy { + ResolveStrategy::UseLatest => "use_latest", + ResolveStrategy::UseMetasrv => "use_metasrv", + ResolveStrategy::AbortOnConflict => "abort_on_conflict", + } +} + +fn nullable_table_id(value: Option) -> api::v1::Value { + nullable_value(value.map(ValueData::U32Value)) +} + +#[cfg(test)] +mod tests { + use api::v1::value::ValueData; + use api::v1::{ColumnDataType, Row, SemanticType, Value}; + use common_event_recorder::Event; + use serde_json::json; + + use super::*; + + #[test] + fn reconciliation_events_use_the_shared_locator_contract() { + let table = ReconciliationEvent::table_lifecycle(ReconciliationLocator::physical_table( + "greptime", "public", "metrics", 42, + )); + assert_eq!(table.event_type(), RECONCILE_TABLE_EVENT_TYPE); + assert_eq!( + table + .extra_schema() + .into_iter() + .map(|column| { + ( + column.column_name, + ColumnDataType::try_from(column.datatype).unwrap(), + SemanticType::try_from(column.semantic_type).unwrap(), + ) + }) + .collect::>(), + vec![ + ( + "catalog_name".to_string(), + ColumnDataType::String, + SemanticType::Field, + ), + ( + "schema_name".to_string(), + ColumnDataType::String, + SemanticType::Field, + ), + ( + "table_name".to_string(), + ColumnDataType::String, + SemanticType::Field, + ), + ( + "table_id".to_string(), + ColumnDataType::Uint32, + SemanticType::Field, + ), + ( + "physical_table_id".to_string(), + ColumnDataType::Uint32, + SemanticType::Field, + ), + ] + ); + assert_eq!( + table.extra_rows().unwrap(), + vec![Row { + values: vec![ + ValueData::StringValue("greptime".to_string()).into(), + ValueData::StringValue("public".to_string()).into(), + ValueData::StringValue("metrics".to_string()).into(), + ValueData::U32Value(42).into(), + Value::default(), + ], + }] + ); + assert_eq!(table.json_payload().unwrap(), serde_json::Value::Null); + } + + #[test] + fn submitted_payloads_are_versioned_and_use_stable_strategy_names() { + for (strategy, expected) in [ + (ResolveStrategy::UseLatest, "use_latest"), + (ResolveStrategy::UseMetasrv, "use_metasrv"), + (ResolveStrategy::AbortOnConflict, "abort_on_conflict"), + ] { + let table = ReconciliationEvent::table_submitted( + ReconciliationLocator::physical_table("greptime", "public", "metrics", 42), + strategy, + true, + ); + assert_eq!( + table.json_payload().unwrap(), + json!({ + "version": 1, + "resolve_strategy": expected, + "is_subprocedure": true, + }) + ); + } + } + + #[test] + fn terminal_payloads_distinguish_complete_and_partial_results() { + let table = ReconciliationEvent::table_result( + ReconciliationLocator::physical_table("greptime", "public", "metrics", 42), + false, + Some("inconsistent"), + Some(ResolveStrategy::UseMetasrv), + Some(4), + 3, + 2, + true, + Some("update_table_info"), + ); + assert_eq!( + table.json_payload().unwrap(), + json!({ + "version": 1, + "complete": false, + "metadata_state": "inconsistent", + "resolution_strategy_applied": "use_metasrv", + "resolved_column_count": 4, + "scanned_region_count": 3, + "updated_region_count": 2, + "table_info_updated": true, + "last_completed_phase": "update_table_info", + }) + ); + } +} diff --git a/src/common/meta/src/reconciliation/reconcile_table.rs b/src/common/meta/src/reconciliation/reconcile_table.rs index 004b93d38c5..562a2061b0a 100644 --- a/src/common/meta/src/reconciliation/reconcile_table.rs +++ b/src/common/meta/src/reconciliation/reconcile_table.rs @@ -18,13 +18,12 @@ pub(crate) mod reconciliation_start; pub(crate) mod resolve_column_metadata; pub(crate) mod update_table_info; -use std::any::Any; use std::fmt::Debug; use common_procedure::error::{FromJsonSnafu, ToJsonSnafu}; use common_procedure::{ - Context as ProcedureContext, Error as ProcedureError, LockKey, Procedure, - Result as ProcedureResult, Status, + Context as ProcedureContext, Error as ProcedureError, EventContext, EventTrigger, LockKey, + Procedure, Result as ProcedureResult, Status, }; use serde::{Deserialize, Serialize}; use snafu::ResultExt; @@ -42,6 +41,9 @@ use crate::key::{DeserializedValueWithBytes, TableMetadataManagerRef}; use crate::lock_key::{CatalogLock, SchemaLock, TableNameLock}; use crate::metrics; use crate::node_manager::NodeManagerRef; +use crate::reconciliation::event::{ + RECONCILE_TABLE_EVENT_TYPE, ReconciliationEvent, ReconciliationLocator, +}; use crate::reconciliation::reconcile_table::reconciliation_start::ReconciliationStart; use crate::reconciliation::reconcile_table::resolve_column_metadata::ResolveStrategy; use crate::reconciliation::utils::{ @@ -56,6 +58,105 @@ pub struct ReconcileTableContext { pub volatile_ctx: VolatileContext, } +#[derive(Debug, Clone, Copy)] +enum TablePhase { + Start, + ResolveColumnMetadata, + ReconcileRegions, + UpdateTableInfo, +} + +impl TablePhase { + const fn as_event_value(self) -> &'static str { + match self { + Self::Start => "start", + Self::ResolveColumnMetadata => "resolve_column_metadata", + Self::ReconcileRegions => "reconcile_regions", + Self::UpdateTableInfo => "update_table_info", + } + } +} + +#[derive(Debug, Clone, Copy)] +enum TableMetadataState { + Consistent, + Inconsistent, +} + +impl TableMetadataState { + const fn as_event_value(self) -> &'static str { + match self { + Self::Consistent => "consistent", + Self::Inconsistent => "inconsistent", + } + } +} + +#[derive(Debug, Default)] +struct ReconcileTableResultSummary { + metadata_state: Option, + resolution_strategy_applied: Option, + resolved_column_count: Option, + scanned_region_count: usize, + updated_region_count: usize, + table_info_updated: bool, + last_completed_phase: Option, +} + +impl ReconcileTableResultSummary { + fn record_scanned_regions(&mut self, scanned_region_count: usize) { + self.scanned_region_count = scanned_region_count; + } + + fn mark_start_completed(&mut self) { + self.last_completed_phase = Some(TablePhase::Start); + } + + fn record_metadata_state(&mut self, metadata_state: TableMetadataState) { + self.metadata_state = Some(metadata_state); + } + + fn record_resolution_strategy(&mut self, resolve_strategy: ResolveStrategy) { + self.resolution_strategy_applied = Some(resolve_strategy); + } + + fn record_resolved_columns( + &mut self, + metadata_state: TableMetadataState, + resolution_strategy_applied: Option, + resolved_column_count: Option, + ) { + self.metadata_state = Some(metadata_state); + self.resolution_strategy_applied = resolution_strategy_applied; + self.resolved_column_count = resolved_column_count; + self.last_completed_phase = Some(TablePhase::ResolveColumnMetadata); + } + + fn record_updated_regions(&mut self, updated_region_count: usize) { + self.updated_region_count = updated_region_count; + } + + fn mark_region_phase_completed(&mut self) { + self.last_completed_phase = Some(TablePhase::ReconcileRegions); + } + + fn mark_table_info_updated(&mut self) { + self.table_info_updated = true; + } + + fn mark_table_info_phase_completed(&mut self) { + self.last_completed_phase = Some(TablePhase::UpdateTableInfo); + } + + fn metadata_state(&self) -> Option<&'static str> { + self.metadata_state.map(TableMetadataState::as_event_value) + } + + fn last_completed_phase(&self) -> Option<&'static str> { + self.last_completed_phase.map(TablePhase::as_event_value) + } +} + impl ReconcileTableContext { /// Creates a new [`ReconcileTableContext`] with the given [`Context`] and [`PersistentContext`]. pub fn new(ctx: Context, persistent_ctx: PersistentContext) -> Self { @@ -147,6 +248,7 @@ impl PersistentContext { pub(crate) struct VolatileContext { pub(crate) table_meta: Option, pub(crate) metrics: ReconcileTableMetrics, + result_summary: ReconcileTableResultSummary, } pub struct ReconcileTableProcedure { @@ -256,6 +358,56 @@ impl Procedure for ReconcileTableProcedure { TableNameLock::new(table_ref.catalog, table_ref.schema, table_ref.table).into(), ]) } + + fn event(&self, ctx: &EventContext<'_>) -> Option> { + if !ctx.event_type_filter.allows(RECONCILE_TABLE_EVENT_TYPE) { + return None; + } + + let persistent_ctx = &self.context.persistent_ctx; + let table_name = &persistent_ctx.table_name; + let locator = ReconciliationLocator::physical_table( + &table_name.catalog_name, + &table_name.schema_name, + &table_name.table_name, + persistent_ctx.table_id, + ); + let event = match ctx.trigger { + EventTrigger::Submitted => ReconciliationEvent::table_submitted( + locator, + persistent_ctx.resolve_strategy, + persistent_ctx.is_subprocedure, + ), + EventTrigger::Succeeded => { + Self::result_event(locator, &self.context.volatile_ctx.result_summary, true) + } + EventTrigger::Failed | EventTrigger::Poisoned => { + Self::result_event(locator, &self.context.volatile_ctx.result_summary, false) + } + _ => ReconciliationEvent::table_lifecycle(locator), + }; + Some(Box::new(event)) + } +} + +impl ReconcileTableProcedure { + fn result_event( + locator: ReconciliationLocator, + summary: &ReconcileTableResultSummary, + complete: bool, + ) -> ReconciliationEvent { + ReconciliationEvent::table_result( + locator, + complete, + summary.metadata_state(), + summary.resolution_strategy_applied, + summary.resolved_column_count, + summary.scanned_region_count, + summary.updated_region_count, + summary.table_info_updated, + summary.last_completed_phase(), + ) + } } #[async_trait::async_trait] @@ -272,6 +424,390 @@ pub(crate) trait State: Sync + Send + Debug { ctx: &mut ReconcileTableContext, procedure_ctx: &ProcedureContext, ) -> Result<(Box, Status)>; +} - fn as_any(&self) -> &dyn Any; +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use common_event_recorder::{EventTypeFilter, EventTypeFilterRef}; + use common_procedure::{ + ChildSubmissionOutcome, EventContext, EventTrigger, Procedure, ProcedureId, ProcedureState, + RetryPhase, + }; + use common_procedure_test::MockContextProvider; + use serde_json::{Value, json}; + use store_api::storage::RegionId; + + use super::*; + use crate::ddl::test_util::datanode_handler::PartialSuccessDatanodeHandler; + use crate::key::DeserializedValueWithBytes; + use crate::key::table_info::TableInfoValue; + use crate::key::table_route::PhysicalTableRouteValue; + use crate::key::test_utils::new_test_table_info_with_name; + use crate::peer::Peer; + use crate::reconciliation::reconcile_table::reconcile_regions::ReconcileRegions; + use crate::reconciliation::utils::build_column_metadata_from_table_info; + use crate::rpc::router::{Region, RegionRoute}; + use crate::test_util::{MockDatanodeManager, new_ddl_context}; + + struct TableEventHarness { + procedure_id: ProcedureId, + lifecycle_state: ProcedureState, + event_type_filter: EventTypeFilterRef, + } + + impl TableEventHarness { + fn all() -> Self { + Self { + procedure_id: ProcedureId::random(), + lifecycle_state: ProcedureState::Running, + event_type_filter: Arc::new(EventTypeFilter::All), + } + } + + fn selected(event_types: impl IntoIterator) -> Self { + Self { + event_type_filter: Arc::new(EventTypeFilter::Only( + event_types.into_iter().map(str::to_string).collect(), + )), + ..Self::all() + } + } + + fn event( + &self, + procedure: &dyn Procedure, + trigger: EventTrigger, + ) -> Option> { + procedure.event(&EventContext { + procedure_id: self.procedure_id, + lifecycle_state: &self.lifecycle_state, + trigger, + event_type_filter: self.event_type_filter.clone(), + event_context: None, + }) + } + } + + #[test] + fn table_submitted_events_cover_root_and_child_intent() { + let events = TableEventHarness::all(); + let root = test_procedure(false); + let child = test_procedure(true); + for (procedure, is_subprocedure) in [(&root, false), (&child, true)] { + let submitted = events.event(procedure, EventTrigger::Submitted).unwrap(); + assert_eq!(submitted.event_type(), RECONCILE_TABLE_EVENT_TYPE); + assert_eq!( + submitted.json_payload().unwrap(), + json!({ + "version": 1, + "resolve_strategy": "use_latest", + "is_subprocedure": is_subprocedure, + }) + ); + } + } + + #[test] + fn table_non_terminal_lifecycle_events_have_null_payloads() { + let events = TableEventHarness::all(); + let mut procedure = test_procedure(true); + procedure.context.volatile_ctx.result_summary = populated_summary(); + + for trigger in [ + EventTrigger::Recovered, + EventTrigger::ChildSubmitted { + procedure_id: ProcedureId::random(), + outcome: ChildSubmissionOutcome::Accepted, + }, + EventTrigger::Retrying { + phase: RetryPhase::Execute, + attempt: 2, + }, + EventTrigger::RollingBack, + ] { + assert_eq!( + events + .event(&procedure, trigger) + .unwrap() + .json_payload() + .unwrap(), + Value::Null + ); + } + } + + #[test] + fn table_terminal_events_report_bounded_results() { + let events = TableEventHarness::all(); + let mut procedure = test_procedure(true); + procedure.context.volatile_ctx.result_summary = populated_summary(); + + for (trigger, complete) in [ + (EventTrigger::Succeeded, true), + (EventTrigger::Failed, false), + (EventTrigger::Poisoned, false), + ] { + assert_eq!( + events + .event(&procedure, trigger) + .unwrap() + .json_payload() + .unwrap(), + json!({ + "version": 1, + "complete": complete, + "metadata_state": "inconsistent", + "resolution_strategy_applied": "use_latest", + "resolved_column_count": 3, + "scanned_region_count": 2, + "updated_region_count": 1, + "table_info_updated": true, + "last_completed_phase": "update_table_info", + }) + ); + } + } + + #[test] + fn table_event_filtering_uses_the_reconciliation_event_type() { + let procedure = test_procedure(false); + assert!( + TableEventHarness::selected([RECONCILE_TABLE_EVENT_TYPE]) + .event(&procedure, EventTrigger::Submitted) + .is_some() + ); + assert!( + TableEventHarness::selected(["create_table"]) + .event(&procedure, EventTrigger::Submitted) + .is_none() + ); + assert!( + TableEventHarness::selected([]) + .event(&procedure, EventTrigger::Submitted) + .is_none() + ); + } + + #[test] + fn table_result_summary_is_not_persisted() { + let events = TableEventHarness::all(); + let mut procedure = test_procedure(false); + procedure.context.volatile_ctx.result_summary = populated_summary(); + let value: Value = serde_json::from_str(&procedure.dump().unwrap()).unwrap(); + assert!(value["persistent_ctx"].get("result_summary").is_none()); + + let loaded = ReconcileTableProcedure::from_json( + test_context(), + &serde_json::to_string(&value).unwrap(), + ) + .unwrap(); + assert_eq!( + events + .event(&loaded, EventTrigger::Succeeded) + .unwrap() + .json_payload() + .unwrap(), + json!({ + "version": 1, + "complete": true, + "metadata_state": null, + "resolution_strategy_applied": null, + "resolved_column_count": null, + "scanned_region_count": 0, + "updated_region_count": 0, + "table_info_updated": false, + "last_completed_phase": null, + }) + ); + } + + #[test] + fn table_partial_summary_preserves_completed_work_without_advancing_phases() { + let events = TableEventHarness::all(); + let mut procedure = test_procedure(false); + let summary = &mut procedure.context.volatile_ctx.result_summary; + summary.record_scanned_regions(2); + + assert_eq!( + events + .event(&procedure, EventTrigger::Failed) + .unwrap() + .json_payload() + .unwrap(), + json!({ + "version": 1, + "complete": false, + "metadata_state": null, + "resolution_strategy_applied": null, + "resolved_column_count": null, + "scanned_region_count": 2, + "updated_region_count": 0, + "table_info_updated": false, + "last_completed_phase": null, + }) + ); + + let summary = &mut procedure.context.volatile_ctx.result_summary; + summary.mark_start_completed(); + summary.record_metadata_state(TableMetadataState::Inconsistent); + assert_eq!( + events + .event(&procedure, EventTrigger::Failed) + .unwrap() + .json_payload() + .unwrap(), + json!({ + "version": 1, + "complete": false, + "metadata_state": "inconsistent", + "resolution_strategy_applied": null, + "resolved_column_count": null, + "scanned_region_count": 2, + "updated_region_count": 0, + "table_info_updated": false, + "last_completed_phase": "start", + }) + ); + + let summary = &mut procedure.context.volatile_ctx.result_summary; + summary.record_resolved_columns( + TableMetadataState::Inconsistent, + Some(ResolveStrategy::UseLatest), + Some(3), + ); + summary.record_updated_regions(1); + assert_eq!( + events + .event(&procedure, EventTrigger::Failed) + .unwrap() + .json_payload() + .unwrap(), + json!({ + "version": 1, + "complete": false, + "metadata_state": "inconsistent", + "resolution_strategy_applied": "use_latest", + "resolved_column_count": 3, + "scanned_region_count": 2, + "updated_region_count": 1, + "table_info_updated": false, + "last_completed_phase": "resolve_column_metadata", + }) + ); + } + + #[tokio::test] + async fn table_mixed_region_outcome_preserves_partial_result() { + let ddl_context = new_ddl_context(Arc::new(MockDatanodeManager::new( + PartialSuccessDatanodeHandler { retryable: false }, + ))); + let context = Context { + node_manager: ddl_context.node_manager, + table_metadata_manager: ddl_context.table_metadata_manager, + cache_invalidator: ddl_context.cache_invalidator, + }; + let mut procedure = ReconcileTableProcedure::new( + context, + 42, + TableName::new("greptime", "public", "metrics"), + ResolveStrategy::UseLatest, + false, + ); + + let mut table_info = new_test_table_info_with_name(42, "metrics"); + table_info.meta.column_ids = vec![0, 1, 2]; + let name_to_ids = table_info.name_to_ids().unwrap(); + let column_metadatas = build_column_metadata_from_table_info( + table_info.meta.schema.column_schemas(), + &table_info.meta.primary_key_indices, + &name_to_ids, + ) + .unwrap(); + let region_ids = [RegionId::new(42, 1), RegionId::new(42, 2)]; + procedure.context.persistent_ctx.table_info_value = Some( + DeserializedValueWithBytes::from_inner(TableInfoValue::new(table_info)), + ); + procedure.context.persistent_ctx.physical_table_route = + Some(PhysicalTableRouteValue::new(vec![ + RegionRoute { + region: Region::new_test(region_ids[0]), + leader_peer: Some(Peer::empty(1)), + ..Default::default() + }, + RegionRoute { + region: Region::new_test(region_ids[1]), + leader_peer: Some(Peer::empty(2)), + ..Default::default() + }, + ])); + + let summary = &mut procedure.context.volatile_ctx.result_summary; + summary.record_scanned_regions(region_ids.len()); + summary.mark_start_completed(); + summary.record_resolved_columns( + TableMetadataState::Inconsistent, + Some(ResolveStrategy::UseLatest), + Some(column_metadatas.len()), + ); + procedure.state = Box::new(ReconcileRegions::new(column_metadatas, region_ids.to_vec())); + + let procedure_ctx = ProcedureContext { + procedure_id: ProcedureId::random(), + provider: Arc::new(MockContextProvider::default()), + event_context: None, + }; + assert!(procedure.execute(&procedure_ctx).await.is_err()); + + assert_eq!( + TableEventHarness::all() + .event(&procedure, EventTrigger::Failed) + .unwrap() + .json_payload() + .unwrap(), + json!({ + "version": 1, + "complete": false, + "metadata_state": "inconsistent", + "resolution_strategy_applied": "use_latest", + "resolved_column_count": 3, + "scanned_region_count": 2, + "updated_region_count": 1, + "table_info_updated": false, + "last_completed_phase": "resolve_column_metadata", + }) + ); + } + + fn populated_summary() -> ReconcileTableResultSummary { + ReconcileTableResultSummary { + metadata_state: Some(TableMetadataState::Inconsistent), + resolution_strategy_applied: Some(ResolveStrategy::UseLatest), + resolved_column_count: Some(3), + scanned_region_count: 2, + updated_region_count: 1, + table_info_updated: true, + last_completed_phase: Some(TablePhase::UpdateTableInfo), + } + } + + fn test_procedure(is_subprocedure: bool) -> ReconcileTableProcedure { + ReconcileTableProcedure::new( + test_context(), + 42, + TableName::new("greptime", "public", "metrics"), + ResolveStrategy::UseLatest, + is_subprocedure, + ) + } + + fn test_context() -> Context { + let ddl_context = new_ddl_context(Arc::new(MockDatanodeManager::new(()))); + Context { + node_manager: ddl_context.node_manager, + table_metadata_manager: ddl_context.table_metadata_manager, + cache_invalidator: ddl_context.cache_invalidator, + } + } } diff --git a/src/common/meta/src/reconciliation/reconcile_table/reconcile_regions.rs b/src/common/meta/src/reconciliation/reconcile_table/reconcile_regions.rs index 6ee3df70f37..fd5d55fb0cb 100644 --- a/src/common/meta/src/reconciliation/reconcile_table/reconcile_regions.rs +++ b/src/common/meta/src/reconciliation/reconcile_table/reconcile_regions.rs @@ -12,7 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::any::Any; use std::collections::{HashMap, HashSet}; use api::v1::column_def::try_as_column_def; @@ -65,7 +64,6 @@ impl State for ReconcileRegions { let table_meta = ctx.build_table_meta(&self.column_metadatas)?; ctx.volatile_ctx.table_meta = Some(table_meta); let table_id = ctx.table_id(); - let table_name = ctx.table_name(); let primary_keys = self .column_metadatas @@ -126,10 +124,12 @@ impl State for ReconcileRegions { } } - let mut results = future::join_all(sync_column_tsks) - .await - .into_iter() - .collect::>>()?; + let results = future::join_all(sync_column_tsks).await; + let updated_region_count = results.iter().filter(|result| result.is_ok()).count(); + ctx.volatile_ctx + .result_summary + .record_updated_regions(updated_region_count); + let mut results = results.into_iter().collect::>>()?; // Ensures all the column metadatas are the same. let column_metadatas = @@ -137,17 +137,22 @@ impl State for ReconcileRegions { UnexpectedSnafu { err_msg: format!( "The table column metadata schemas from datanodes are not the same, table: {}, table_id: {}", - table_name, + ctx.table_name(), table_id ), }, )?; + ctx.volatile_ctx + .result_summary + .mark_region_phase_completed(); + // Checks all column metadatas are consistent, and updates the table info if needed. if column_metadatas != self.column_metadatas { info!( "Datanode column metadatas are not consistent with metasrv, updating metasrv's column metadatas, table: {}, table_id: {}", - table_name, table_id + ctx.table_name(), + table_id ); // Safety: fetched in the above. let table_info_value = ctx.persistent_ctx.table_info_value.clone().unwrap(); @@ -159,10 +164,6 @@ impl State for ReconcileRegions { Ok((Box::new(ReconciliationEnd), Status::executing(false))) } - - fn as_any(&self) -> &dyn Any { - self - } } /// Makes an alter region request to sync columns. diff --git a/src/common/meta/src/reconciliation/reconcile_table/reconciliation_end.rs b/src/common/meta/src/reconciliation/reconcile_table/reconciliation_end.rs index 9f28d7f4be5..b6e695dca46 100644 --- a/src/common/meta/src/reconciliation/reconcile_table/reconciliation_end.rs +++ b/src/common/meta/src/reconciliation/reconcile_table/reconciliation_end.rs @@ -12,8 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::any::Any; - use common_procedure::{Context as ProcedureContext, Status}; use common_telemetry::info; use serde::{Deserialize, Serialize}; @@ -46,8 +44,4 @@ impl State for ReconciliationEnd { Ok((Box::new(ReconciliationEnd), Status::done())) } - - fn as_any(&self) -> &dyn Any { - self - } } diff --git a/src/common/meta/src/reconciliation/reconcile_table/reconciliation_start.rs b/src/common/meta/src/reconciliation/reconcile_table/reconciliation_start.rs index 2fcc8ded0c0..48a7fff8b96 100644 --- a/src/common/meta/src/reconciliation/reconcile_table/reconciliation_start.rs +++ b/src/common/meta/src/reconciliation/reconcile_table/reconciliation_start.rs @@ -12,8 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::any::Any; - use common_procedure::{Context as ProcedureContext, Status}; use common_telemetry::info; use serde::{Deserialize, Serialize}; @@ -46,7 +44,6 @@ impl State for ReconciliationStart { procedure_ctx: &ProcedureContext, ) -> Result<(Box, Status)> { let table_id = ctx.table_id(); - let table_name = ctx.table_name(); let (physical_table_id, physical_table_route) = ctx .table_metadata_manager @@ -58,14 +55,17 @@ impl State for ReconciliationStart { error::UnexpectedSnafu { err_msg: format!( "Reconcile table only works for physical table, but got logical table: {}, table_id: {}", - table_name, table_id + ctx.table_name(), + table_id ), } ); info!( "Reconciling table: {}, table_id: {}, procedure_id: {}", - table_name, table_id, procedure_ctx.procedure_id + ctx.table_name(), + table_id, + procedure_ctx.procedure_id ); // TODO(weny): Repairs the table route if needed. let region_metadata_lister = RegionMetadataLister::new(ctx.node_manager.clone()); @@ -79,6 +79,9 @@ impl State for ReconciliationStart { .list(physical_table_id, &physical_table_route.region_routes) .await? }; + ctx.volatile_ctx + .result_summary + .record_scanned_regions(region_metadatas.len()); ensure!(!region_metadatas.is_empty(), { metrics::METRIC_META_RECONCILIATION_STATS @@ -92,7 +95,8 @@ impl State for ReconciliationStart { error::UnexpectedSnafu { err_msg: format!( "No region metadata found for table: {}, table_id: {}", - table_name, table_id + ctx.table_name(), + table_id ), } }); @@ -109,11 +113,14 @@ impl State for ReconciliationStart { error::UnexpectedSnafu { err_msg: format!( "Some regions are not opened, table: {}, table_id: {}", - table_name, table_id + ctx.table_name(), + table_id ), } }); + ctx.volatile_ctx.result_summary.mark_start_completed(); + // Persist the physical table route. // TODO(weny): refetch the physical table route if repair is needed. ctx.persistent_ctx.physical_table_route = Some(physical_table_route); @@ -127,8 +134,4 @@ impl State for ReconciliationStart { Status::executing(false), )) } - - fn as_any(&self) -> &dyn Any { - self - } } diff --git a/src/common/meta/src/reconciliation/reconcile_table/resolve_column_metadata.rs b/src/common/meta/src/reconciliation/reconcile_table/resolve_column_metadata.rs index d45cfe464ea..27950f0fd5d 100644 --- a/src/common/meta/src/reconciliation/reconcile_table/resolve_column_metadata.rs +++ b/src/common/meta/src/reconciliation/reconcile_table/resolve_column_metadata.rs @@ -12,8 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::any::Any; - use async_trait::async_trait; use common_procedure::{Context as ProcedureContext, Status}; use common_telemetry::info; @@ -25,7 +23,7 @@ use strum::AsRefStr; use crate::error::{self, MissingColumnIdsSnafu, Result}; use crate::reconciliation::reconcile_table::reconcile_regions::ReconcileRegions; use crate::reconciliation::reconcile_table::update_table_info::UpdateTableInfo; -use crate::reconciliation::reconcile_table::{ReconcileTableContext, State}; +use crate::reconciliation::reconcile_table::{ReconcileTableContext, State, TableMetadataState}; use crate::reconciliation::utils::{ ResolveColumnMetadataResult, build_column_metadata_from_table_info, check_column_metadatas_consistent, resolve_column_metadatas_with_latest, @@ -101,6 +99,12 @@ impl State for ResolveColumnMetadata { table_name, table_id ); + ctx.volatile_ctx.result_summary.record_resolved_columns( + TableMetadataState::Consistent, + None, + Some(column_metadatas.len()), + ); + // Update metrics. ctx.mut_metrics().resolve_column_metadata_result = Some(ResolveColumnMetadataResult::Consistent); @@ -110,6 +114,10 @@ impl State for ResolveColumnMetadata { )); }; + ctx.volatile_ctx + .result_summary + .record_metadata_state(TableMetadataState::Inconsistent); + match self.strategy { ResolveStrategy::UseMetasrv => { let table_info_value = ctx.persistent_ctx.table_info_value.as_ref().unwrap(); @@ -126,6 +134,12 @@ impl State for ResolveColumnMetadata { let region_ids = resolve_column_metadatas_with_metasrv(&column_metadata, &self.region_metadata)?; + ctx.volatile_ctx.result_summary.record_resolved_columns( + TableMetadataState::Inconsistent, + Some(self.strategy), + Some(column_metadata.len()), + ); + // Update metrics. let metrics = ctx.mut_metrics(); metrics.resolve_column_metadata_result = @@ -139,6 +153,12 @@ impl State for ResolveColumnMetadata { let (column_metadatas, region_ids) = resolve_column_metadatas_with_latest(&self.region_metadata)?; + ctx.volatile_ctx.result_summary.record_resolved_columns( + TableMetadataState::Inconsistent, + Some(self.strategy), + Some(column_metadatas.len()), + ); + // Update metrics. let metrics = ctx.mut_metrics(); metrics.resolve_column_metadata_result = @@ -151,6 +171,10 @@ impl State for ResolveColumnMetadata { ResolveStrategy::AbortOnConflict => { let table_name = table_name.to_string(); + ctx.volatile_ctx + .result_summary + .record_resolution_strategy(self.strategy); + // Update metrics. let metrics = ctx.mut_metrics(); metrics.resolve_column_metadata_result = @@ -163,8 +187,4 @@ impl State for ResolveColumnMetadata { } } } - - fn as_any(&self) -> &dyn Any { - self - } } diff --git a/src/common/meta/src/reconciliation/reconcile_table/update_table_info.rs b/src/common/meta/src/reconciliation/reconcile_table/update_table_info.rs index 790184c07ab..7ef62265de3 100644 --- a/src/common/meta/src/reconciliation/reconcile_table/update_table_info.rs +++ b/src/common/meta/src/reconciliation/reconcile_table/update_table_info.rs @@ -12,8 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::any::Any; - use common_procedure::{Context as ProcedureContext, Status}; use common_telemetry::info; use serde::{Deserialize, Serialize}; @@ -81,6 +79,9 @@ impl State for UpdateTableInfo { ctx.table_name(), ctx.table_id() ); + ctx.volatile_ctx + .result_summary + .mark_table_info_phase_completed(); return Ok((Box::new(ReconciliationEnd), Status::executing(true))); } @@ -98,6 +99,7 @@ impl State for UpdateTableInfo { new_table_info, ) .await?; + ctx.volatile_ctx.result_summary.mark_table_info_updated(); let table_ref = ctx.table_name().table_ref(); let table_id = ctx.table_id(); @@ -119,11 +121,10 @@ impl State for UpdateTableInfo { // Update metrics. let metrics = ctx.mut_metrics(); metrics.update_table_info = true; + ctx.volatile_ctx + .result_summary + .mark_table_info_phase_completed(); Ok((Box::new(ReconciliationEnd), Status::executing(true))) } - - fn as_any(&self) -> &dyn Any { - self - } } diff --git a/src/common/recordbatch/src/cursor.rs b/src/common/recordbatch/src/cursor.rs index a741953ccc2..125a68f4f13 100644 --- a/src/common/recordbatch/src/cursor.rs +++ b/src/common/recordbatch/src/cursor.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +use datatypes::schema::SchemaRef; use futures::StreamExt; use tokio::sync::Mutex; @@ -28,12 +29,15 @@ struct Inner { /// A cursor on RecordBatchStream that fetches data batch by batch pub struct RecordBatchStreamCursor { + schema: SchemaRef, inner: Mutex, } impl RecordBatchStreamCursor { pub fn new(stream: SendableRecordBatchStream) -> RecordBatchStreamCursor { + let schema = stream.schema(); Self { + schema, inner: Mutex::new(Inner { stream, current_row_index: 0, @@ -43,6 +47,10 @@ impl RecordBatchStreamCursor { } } + pub fn schema(&self) -> SchemaRef { + self.schema.clone() + } + /// Take `size` of row from the `RecordBatchStream` and create a new /// `RecordBatch` for these rows. pub async fn take(&self, size: usize) -> Result { diff --git a/src/common/time/src/timestamp.rs b/src/common/time/src/timestamp.rs index 79f97d412f5..1acc9514cd8 100644 --- a/src/common/time/src/timestamp.rs +++ b/src/common/time/src/timestamp.rs @@ -483,6 +483,31 @@ impl Timestamp { ParseTimestampSnafu { raw: s }.fail() } + /// Interprets a timezone-less [`NaiveDateTime`] in the given timezone and + /// returns the corresponding timestamp. + /// + /// Datetimes that fall into a DST gap (a local time that does not exist) + /// are rejected, while ambiguous datetimes (from a repeated local time) + /// are resolved to the earlier instant. This policy is shared with + /// [`Timestamp::from_str`] so that the text and binary protocols interpret + /// datetimes consistently. + pub fn from_naive_datetime( + datetime: NaiveDateTime, + timezone: &Timezone, + ) -> crate::error::Result { + match datetime_to_utc(&datetime, timezone) { + LocalResult::Single(utc) | LocalResult::Ambiguous(utc, _) => { + Timestamp::from_chrono_datetime(utc).context(ParseTimestampSnafu { + raw: format!("{datetime} (timezone {timezone})"), + }) + } + LocalResult::None => ParseTimestampSnafu { + raw: format!("{datetime} (timezone {timezone})"), + } + .fail(), + } + } + pub fn negative(mut self) -> Self { self.value = -self.value; self @@ -531,12 +556,8 @@ fn naive_datetime_to_timestamp( .context(ParseTimestampSnafu { raw: s }); }; - match datetime_to_utc(&datetime, timezone) { - LocalResult::None => ParseTimestampSnafu { raw: s }.fail(), - LocalResult::Single(utc) | LocalResult::Ambiguous(utc, _) => { - Timestamp::from_chrono_datetime(utc).context(ParseTimestampSnafu { raw: s }) - } - } + Timestamp::from_naive_datetime(datetime, timezone) + .map_err(|_| ParseTimestampSnafu { raw: s }.build()) } impl From for Timestamp { @@ -923,6 +944,48 @@ mod tests { ); } + #[test] + fn test_from_naive_datetime() { + let datetime = NaiveDate::from_ymd_opt(2026, 8, 13) + .unwrap() + .and_hms_opt(8, 0, 0) + .unwrap(); + + // A fixed-offset timezone shifts the datetime by a constant amount. + let shanghai = Timezone::from_tz_string("Asia/Shanghai").unwrap(); + assert_eq!( + "2026-08-13 00:00:00", + Timestamp::from_naive_datetime(datetime, &shanghai) + .unwrap() + .to_chrono_datetime() + .unwrap() + .to_string() + ); + + // 2026-03-08 02:30 does not exist in America/New_York (DST gap). + let new_york = Timezone::from_tz_string("America/New_York").unwrap(); + let gap = NaiveDate::from_ymd_opt(2026, 3, 8) + .unwrap() + .and_hms_opt(2, 30, 0) + .unwrap(); + assert!(Timestamp::from_naive_datetime(gap, &new_york).is_err()); + + // 2026-11-01 01:30 is ambiguous in America/New_York; picks the first + // instant (EDT, UTC-4). + let ambiguous = NaiveDate::from_ymd_opt(2026, 11, 1) + .unwrap() + .and_hms_opt(1, 30, 0) + .unwrap(); + assert_eq!( + "2026-11-01 05:30:00", + Timestamp::from_naive_datetime(ambiguous, &new_york) + .unwrap() + .to_chrono_datetime() + .unwrap() + .to_string() + ); + } + #[test] fn test_to_iso8601_string() { set_default_timezone(Some("Asia/Shanghai")).unwrap(); diff --git a/src/common/time/src/timezone.rs b/src/common/time/src/timezone.rs index 41cc1f78428..dbd8dffcd49 100644 --- a/src/common/time/src/timezone.rs +++ b/src/common/time/src/timezone.rs @@ -108,6 +108,15 @@ impl Timezone { } } + /// A named zone merely sitting at +00:00 today is not UTC: it may have been + /// elsewhere at the timestamp being converted. + pub fn is_utc(&self) -> bool { + match self { + Self::Offset(offset) => offset.local_minus_utc() == 0, + Self::Named(tz) => matches!(tz, Tz::UTC), + } + } + /// Returns the number of seconds to add to convert from UTC to the local time. pub fn local_minus_utc(&self) -> i64 { match self { diff --git a/src/datanode/src/region_server.rs b/src/datanode/src/region_server.rs index 0d9398b0889..79d0dd19152 100644 --- a/src/datanode/src/region_server.rs +++ b/src/datanode/src/region_server.rs @@ -67,7 +67,8 @@ use servers::error::{ }; use servers::grpc::FlightCompression; use servers::grpc::flight::{ - FlightCraft, FlightRecordBatchSource, FlightRecordBatchStream, TonicStream, + FlightCraft, FlightRecordBatchSource, FlightRecordBatchStream, FlightRecordBatchStreamInput, + TonicStream, }; use servers::grpc::region_server::RegionServerHandler; use session::context::{ @@ -990,7 +991,11 @@ impl FlightCraft for RegionServer { }; let stream = Box::pin(FlightRecordBatchStream::new( - FlightRecordBatchSource::initializer(initializer), + FlightRecordBatchStreamInput::initializer(async move { + initializer + .await + .map(FlightRecordBatchSource::RecordBatches) + }), tracing_context, self.flight_compression, query_ctx, diff --git a/src/flow/src/batching_mode.rs b/src/flow/src/batching_mode.rs index bb946e75740..3981446721b 100644 --- a/src/flow/src/batching_mode.rs +++ b/src/flow/src/batching_mode.rs @@ -44,6 +44,8 @@ pub struct BatchingModeOptions { /// The gRPC connection timeout #[serde(with = "humantime_serde")] pub grpc_conn_timeout: Duration, + #[serde(with = "humantime_serde")] + pub experimental_flight_do_get_timeout: Duration, /// The gRPC max retry number pub experimental_grpc_max_retries: u32, /// Flow wait for available frontend timeout, @@ -72,6 +74,7 @@ impl Default for BatchingModeOptions { slow_query_threshold: Duration::from_secs(60), experimental_min_refresh_duration: Duration::new(5, 0), grpc_conn_timeout: Duration::from_secs(5), + experimental_flight_do_get_timeout: Duration::from_secs(10), experimental_grpc_max_retries: 3, experimental_frontend_scan_timeout: Duration::from_secs(30), experimental_max_filter_num_per_query: 20, diff --git a/src/flow/src/batching_mode/frontend_client.rs b/src/flow/src/batching_mode/frontend_client.rs index aa94b91b2ae..48314dfa21b 100644 --- a/src/flow/src/batching_mode/frontend_client.rs +++ b/src/flow/src/batching_mode/frontend_client.rs @@ -413,12 +413,12 @@ impl FrontendClient { peer: db.peer.clone(), }); db.database - .query_with_terminal_metrics_and_flow_extensions( - request, - &hints, - extensions, - snapshot_seqs, - ) + .flight_request() + .with_hints(&hints) + .with_flow_extensions(extensions) + .with_snapshot_seqs(snapshot_seqs) + .with_timeout(batch_opts.experimental_flight_do_get_timeout) + .query_with_terminal_metrics(request) .await .map_err(BoxedError::new) .context(ExternalSnafu) @@ -623,6 +623,7 @@ mod tests { use std::task::{Context, Poll}; use std::time::Duration; + use api::v1::query_request::Query; use arrow_flight::flight_service_server::FlightServiceServer; use arrow_flight::{FlightData, Ticket}; use common_query::{Output, OutputData}; diff --git a/src/frontend/src/heartbeat.rs b/src/frontend/src/heartbeat.rs index 0d9cde74fbe..b7e3a7aff48 100644 --- a/src/frontend/src/heartbeat.rs +++ b/src/frontend/src/heartbeat.rs @@ -16,8 +16,9 @@ mod tests; use std::collections::{HashMap, HashSet}; +use std::fmt::{Display, Formatter}; use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; -use std::sync::{Arc, Mutex as StdMutex}; +use std::sync::{Arc, Mutex as StdMutex, MutexGuard, PoisonError}; use api::v1::meta::heartbeat_request::NodeWorkloads; use api::v1::meta::{FrontendWorkloads, HeartbeatRequest, HeartbeatResponse, NodeInfo, Peer}; @@ -93,34 +94,88 @@ pub trait FrontendHeartbeatExtension: Send + Sync { } /// A shareable, ordered registry of frontend heartbeat extensions. +/// +/// Registration is available until [`FrontendHeartbeatExtensions::freeze`] is called. After that, +/// membership is immutable so every heartbeat consumer observes the same extensions. #[derive(Clone, Default)] pub struct FrontendHeartbeatExtensions { inner: Arc>, } +/// The reason a frontend heartbeat extension was not registered. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum RegistrationError { + /// An extension with the same name is already registered. + Duplicate, + /// The registry is frozen and no longer accepts registrations. + Frozen, +} + +impl Display for RegistrationError { + fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result { + match self { + Self::Duplicate => formatter.write_str("heartbeat extension is already registered"), + Self::Frozen => formatter.write_str("heartbeat extension registry is frozen"), + } + } +} + +impl std::error::Error for RegistrationError {} + #[derive(Default)] struct FrontendHeartbeatExtensionsInner { names: HashSet, extensions: Vec>, + frozen: bool, } impl FrontendHeartbeatExtensions { + fn lock(&self) -> MutexGuard<'_, FrontendHeartbeatExtensionsInner> { + self.inner.lock().unwrap_or_else(PoisonError::into_inner) + } + /// Registers an extension without replacing an existing extension of the same name. /// - /// Returns `true` for a new registration and `false` for an idempotent duplicate. + /// Returns `true` for a new registration and `false` for an idempotent duplicate or when the + /// registry is frozen. Use [`FrontendHeartbeatExtensions::try_register`] when the rejection + /// reason must not be discarded. pub fn register(&self, extension: Arc) -> bool { - let mut inner = self.inner.lock().unwrap(); + self.try_register(extension).is_ok() + } + + /// Registers an extension and returns the reason when registration is rejected. + /// + /// If the registry is already frozen, this method returns [`RegistrationError::Frozen`] + /// without calling [`FrontendHeartbeatExtension::name`]. If freezing races with name + /// evaluation, the registry state is checked again before insertion. + pub fn try_register( + &self, + extension: Arc, + ) -> std::result::Result<(), RegistrationError> { + if self.lock().frozen { + return Err(RegistrationError::Frozen); + } + let name = extension.name().to_string(); + let mut inner = self.lock(); + if inner.frozen { + return Err(RegistrationError::Frozen); + } if !inner.names.insert(name) { - return false; + return Err(RegistrationError::Duplicate); } inner.extensions.push(extension); - true + Ok(()) + } + + /// Prevents further registrations while preserving the current registration order. + pub fn freeze(&self) { + self.lock().frozen = true; } /// Returns the registered extensions in registration order. pub fn extensions(&self) -> Vec> { - self.inner.lock().unwrap().extensions.clone() + self.lock().extensions.clone() } /// Returns response handlers in extension registration order. @@ -133,7 +188,7 @@ impl FrontendHeartbeatExtensions { /// Returns the number of registered extensions. pub fn len(&self) -> usize { - self.inner.lock().unwrap().extensions.len() + self.lock().extensions.len() } /// Returns whether no extension is registered. diff --git a/src/frontend/src/heartbeat/tests.rs b/src/frontend/src/heartbeat/tests.rs index d3c651de326..8f5bfa614b2 100644 --- a/src/frontend/src/heartbeat/tests.rs +++ b/src/frontend/src/heartbeat/tests.rs @@ -14,8 +14,9 @@ use std::collections::{HashMap, VecDeque}; use std::future::pending; +use std::panic::{AssertUnwindSafe, catch_unwind}; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; -use std::sync::{Arc, Mutex}; +use std::sync::{Arc, Barrier, Mutex}; use std::time::Duration; use api::v1::meta::heartbeat_request::NodeWorkloads; @@ -494,6 +495,170 @@ impl FrontendHeartbeatExtension for TestExtension { } } +struct PanickingNameExtension; + +#[async_trait] +impl FrontendHeartbeatExtension for PanickingNameExtension { + fn name(&self) -> &str { + panic!("extension name panic") + } +} + +#[test] +fn test_panicking_extension_name_does_not_corrupt_registry() { + let extensions = FrontendHeartbeatExtensions::default(); + + let result = catch_unwind(AssertUnwindSafe(|| { + extensions.register(Arc::new(PanickingNameExtension)); + })); + assert!(result.is_err()); + + let valid = TestExtension::new("valid"); + assert!(extensions.register(valid.clone())); + let registered = extensions.extensions(); + assert_eq!(registered.len(), 1); + assert!(Arc::ptr_eq( + &(valid as Arc), + ®istered[0] + )); + assert_eq!(extensions.len(), 1); +} + +#[test] +fn test_registry_recovers_from_poisoned_lock() { + let extensions = FrontendHeartbeatExtensions::default(); + let poisoned = extensions.clone(); + + let result = catch_unwind(AssertUnwindSafe(move || { + let _guard = poisoned.lock(); + panic!("poison registry lock") + })); + assert!(result.is_err()); + + let valid = TestExtension::new("valid"); + assert!(extensions.register(valid.clone())); + let registered = extensions.extensions(); + assert_eq!(registered.len(), 1); + assert!(Arc::ptr_eq( + &(valid as Arc), + ®istered[0] + )); + assert_eq!(extensions.len(), 1); +} + +#[test] +fn test_registry_freeze_rejects_late_registration_without_changing_membership() { + let extensions = FrontendHeartbeatExtensions::default(); + let calls = Arc::new(AtomicUsize::new(0)); + let handler: HeartbeatResponseHandlerRef = Arc::new(ShortCircuitHandler { + result: ShortCircuitResult::Done, + calls, + }); + let registered = TestExtension::with_handler("registered", handler.clone()); + + assert!(extensions.register(registered.clone())); + assert!(!extensions.register(TestExtension::new("registered"))); + extensions.freeze(); + + assert!(!extensions.register(TestExtension::new("late"))); + let frozen_extensions = extensions.extensions(); + assert_eq!(frozen_extensions.len(), 1); + assert!(Arc::ptr_eq( + &(registered as Arc), + &frozen_extensions[0] + )); + let frozen_handlers = extensions.response_handlers(); + assert_eq!(frozen_handlers.len(), 1); + assert!(Arc::ptr_eq(&handler, &frozen_handlers[0])); + assert_eq!(extensions.len(), 1); +} + +#[test] +fn test_try_register_distinguishes_duplicate_and_frozen() { + let extensions = FrontendHeartbeatExtensions::default(); + + assert_eq!( + extensions.try_register(TestExtension::new("registered")), + Ok(()) + ); + assert_eq!( + extensions.try_register(TestExtension::new("registered")), + Err(RegistrationError::Duplicate) + ); + + extensions.freeze(); + assert_eq!( + extensions.try_register(TestExtension::new("late")), + Err(RegistrationError::Frozen) + ); +} + +#[test] +fn test_frozen_registration_does_not_call_extension_name() { + let extensions = FrontendHeartbeatExtensions::default(); + extensions.freeze(); + + let compatible_result = catch_unwind(AssertUnwindSafe(|| { + extensions.register(Arc::new(PanickingNameExtension)) + })); + assert!(!compatible_result.unwrap()); + + let typed_result = catch_unwind(AssertUnwindSafe(|| { + extensions.try_register(Arc::new(PanickingNameExtension)) + })); + + assert_eq!(typed_result.unwrap(), Err(RegistrationError::Frozen)); + assert!(extensions.is_empty()); +} + +struct BlockingNameExtension { + entered_name: Arc, + release_name: Arc, +} + +#[async_trait] +impl FrontendHeartbeatExtension for BlockingNameExtension { + fn name(&self) -> &str { + self.entered_name.wait(); + self.release_name.wait(); + "blocking" + } +} + +#[test] +fn test_freeze_wins_while_registration_name_is_running() { + let extensions = FrontendHeartbeatExtensions::default(); + let entered_name = Arc::new(Barrier::new(2)); + let release_name = Arc::new(Barrier::new(2)); + let registering = extensions.clone(); + let extension = Arc::new(BlockingNameExtension { + entered_name: entered_name.clone(), + release_name: release_name.clone(), + }); + + let registration = std::thread::spawn(move || registering.try_register(extension)); + entered_name.wait(); + let freezing = extensions.clone(); + let (freeze_done_tx, freeze_done_rx) = std::sync::mpsc::sync_channel(1); + let freeze = std::thread::spawn(move || { + freezing.freeze(); + let _ = freeze_done_tx.send(()); + }); + if let Err(error) = freeze_done_rx.recv_timeout(Duration::from_secs(2)) { + release_name.wait(); + let registration = registration.join(); + let freeze = freeze.join(); + panic!( + "freeze did not complete while extension name was blocked: {error:?}; registration: {registration:?}; freeze: {freeze:?}" + ); + } + release_name.wait(); + + freeze.join().unwrap(); + assert_eq!(registration.join().unwrap(), Err(RegistrationError::Frozen)); + assert!(extensions.is_empty()); +} + fn test_task( connector: Arc, extensions: FrontendHeartbeatExtensions, diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs index 78b701a3442..90650014e6a 100644 --- a/src/frontend/src/instance.rs +++ b/src/frontend/src/instance.rs @@ -59,6 +59,7 @@ use common_recordbatch::error::StreamTimeoutSnafu; use common_telemetry::logging::SlowQueryOptions; use common_telemetry::{debug, error, tracing}; use dashmap::DashMap; +use datafusion::dataframe::DataFrame; use datafusion::physical_plan::ExecutionPlan; use datafusion_expr::LogicalPlan; use futures::{Stream, StreamExt, future}; @@ -134,6 +135,7 @@ pub struct Instance { slow_query_options: SlowQueryOptions, influxdb_default_merge_mode: InfluxdbMergeMode, trace_ingest_chunk_size: usize, + otlp_resource_info: bool, suspend: Arc, // cache for otlp metrics @@ -846,6 +848,22 @@ impl Instance { query_interceptor.pre_execute(stmt.as_ref(), Some(&plan), query_ctx.clone())?; + // TQL EXPLAIN/ANALYZE formats are consumed from the query context at + // execution time (see `optimize_physical_plan`); re-apply the side + // effect of `plan_tql` that was lost when the plan was built during + // Describe. `explain_format` is per-query state, so this never + // overwrites anything. + if let Some(Statement::Tql(tql)) = &stmt { + let format = match tql { + Tql::Explain(explain) => explain.format.as_ref(), + Tql::Analyze(analyze) => analyze.format.as_ref(), + Tql::Eval(_) => None, + }; + if let Some(format) = format { + query_ctx.set_explain_format(format.to_string()); + } + } + let query = stmt .as_ref() .map(|s| s.to_string()) @@ -927,6 +945,59 @@ impl Instance { vec![result] } + /// Builds the [`DataFrame`] for an information-schema-backed `SHOW` + /// statement; `None` for other statements. The future is boxed to keep + /// `do_describe_inner`'s state machine small. + fn show_statement_dataframe<'a>( + &'a self, + stmt: &'a Statement, + query_ctx: &'a QueryContextRef, + ) -> Pin>> + Send + 'a>> { + Box::pin(async move { + let engine = &self.query_engine; + let catalog_manager = self.catalog_manager(); + let ctx = query_ctx.clone(); + let dataframe = match stmt { + Statement::ShowDatabases(show) => { + query::sql::show_databases_dataframe(show, engine, catalog_manager, ctx).await + } + Statement::ShowTables(show) => { + query::sql::show_tables_dataframe(show, engine, catalog_manager, ctx).await + } + Statement::ShowViews(show) => { + query::sql::show_views_dataframe(show, engine, catalog_manager, ctx).await + } + Statement::ShowFlows(show) => { + query::sql::show_flows_dataframe(show, engine, catalog_manager, ctx).await + } + Statement::ShowColumns(show) => { + query::sql::show_columns_dataframe(show, engine, catalog_manager, ctx).await + } + Statement::ShowTableStatus(show) => { + query::sql::show_table_status_dataframe(show, engine, catalog_manager, ctx) + .await + } + Statement::ShowCharset(kind) => { + query::sql::show_charsets_dataframe(kind, engine, catalog_manager, ctx).await + } + Statement::ShowCollation(kind) => { + query::sql::show_collations_dataframe(kind, engine, catalog_manager, ctx).await + } + Statement::ShowIndex(show) => { + query::sql::show_index_dataframe(show, engine, catalog_manager, ctx).await + } + Statement::ShowRegion(show) => { + query::sql::show_region_dataframe(show, engine, catalog_manager, ctx).await + } + Statement::ShowProcesslist(show) => { + query::sql::show_processlist_dataframe(show, engine, catalog_manager, ctx).await + } + _ => return None, + }; + Some(dataframe) + }) + } + async fn do_describe_inner( &self, stmt: Statement, @@ -946,6 +1017,37 @@ impl Instance { let plannable = is_inner_plannable(&stmt) || matches!(&stmt, Statement::Explain(explain) if is_inner_plannable(explain.statement.as_ref())); + if let Statement::Tql(tql) = stmt { + // TQL produces a logical plan; describe it from the plan so the + // extended-protocol RowDescription matches the executed DataRows. + self.check_sql_permission(&Statement::Tql(tql.clone()), &query_ctx) + .await?; + let plan = self.statement_executor.plan_tql(tql, &query_ctx).await?; + return self + .query_engine + .describe(plan, query_ctx) + .await + .map(Some) + .context(error::DescribeStatementSnafu); + } + + // Describe SHOW statements from the same projection the executor builds. + if let Some(dataframe) = self + .show_statement_dataframe(&stmt, &query_ctx) + .await + .transpose() + .context(PlanStatementSnafu)? + { + self.check_sql_permission(&stmt, &query_ctx).await?; + let plan = dataframe.into_unoptimized_plan(); + return self + .query_engine + .describe(plan, query_ctx) + .await + .map(Some) + .context(error::DescribeStatementSnafu); + } + if plannable { self.check_sql_permission(&stmt, &query_ctx).await?; diff --git a/src/frontend/src/instance/builder.rs b/src/frontend/src/instance/builder.rs index c49a8394beb..0f87dc424f0 100644 --- a/src/frontend/src/instance/builder.rs +++ b/src/frontend/src/instance/builder.rs @@ -367,6 +367,7 @@ impl FrontendBuilder { slow_query_options: self.options.slow_query.clone(), influxdb_default_merge_mode: self.options.influxdb.default_merge_mode, trace_ingest_chunk_size: self.options.otlp.trace_ingest_chunk_size, + otlp_resource_info: self.options.otlp.experimental_enable_resource_info, suspend: Arc::new(AtomicBool::new(false)), }) } diff --git a/src/frontend/src/instance/entity_graph.rs b/src/frontend/src/instance/entity_graph.rs index 14f18ed1b79..a0817cf2806 100644 --- a/src/frontend/src/instance/entity_graph.rs +++ b/src/frontend/src/instance/entity_graph.rs @@ -23,7 +23,7 @@ //! the query engine. Injected into the catalog manager after the engine is built, //! breaking the `catalog -> query` cycle. -use std::collections::HashMap; +use std::collections::{BTreeMap, HashMap}; use std::sync::Weak; use async_trait::async_trait; @@ -60,8 +60,9 @@ use table::TableRef; use table::metadata::TableInfo; use table::predicate::{TimeRangeExtraction, extract_time_range_strict}; use table::requests::{ - EntityRole, SEMANTIC_SIGNAL_TYPE, SEMANTIC_SOURCE, SIGNAL_TYPE_METRIC, SOURCE_PROMETHEUS, - is_trace_v1_table, parse_entity_columns, parse_entity_option_key, + EntityRole, SEMANTIC_METRIC_TYPE, SEMANTIC_SIGNAL_TYPE, SEMANTIC_SOURCE, SIGNAL_TYPE_METRIC, + SOURCE_OPENTELEMETRY, SOURCE_PROMETHEUS, is_trace_v1_table, parse_entity_columns, + parse_entity_option_key, }; use crate::error; @@ -177,6 +178,7 @@ impl EntityGraphProviderImpl { time_index: time_index.clone(), entity_type, id_columns, + id_qualifier: None, descriptive_columns, scope_columns, }) @@ -188,10 +190,11 @@ impl EntityGraphProviderImpl { /// All entity declarations of one table: the explicit options plus the /// zero-configuration conventions (`otlp_trace_entities` for trace-v1 /// tables — including the `service` identity of tables created before the - /// ingest-side auto-stamp — and the Prometheus descriptor whitelist). An - /// explicit declaration of a type always suppresses the implicit one, even - /// when the explicit declaration is invalid and skipped: silently falling - /// back would change entity identity behind the user's back. + /// ingest-side auto-stamp — and the Prometheus/OTel descriptor + /// whitelists). An explicit declaration of a type always suppresses the + /// implicit one, even when the explicit declaration is invalid and + /// skipped: silently falling back would change entity identity behind the + /// user's back. fn declarations_for( table_info: &TableInfo, conventions: &Conventions, @@ -204,7 +207,20 @@ impl EntityGraphProviderImpl { &mut declarations, ); } - Self::extend_with_prometheus_conventions(table_info, conventions, &mut declarations); + Self::extend_with_info_metric_conventions( + table_info, + &conventions.prometheus_info_metrics, + SOURCE_PROMETHEUS, + None, + &mut declarations, + ); + Self::extend_with_info_metric_conventions( + table_info, + &conventions.otel_info_metrics, + SOURCE_OPENTELEMETRY, + Some(servers::semantic::METRIC_TYPE_INFO), + &mut declarations, + ); declarations } @@ -216,29 +232,34 @@ impl EntityGraphProviderImpl { }) } - /// Implicit declarations of the well-known Prometheus entity-descriptor - /// metrics (the `prometheus_info_metrics` whitelist of `conventions.yaml`), - /// gated on the ingest-stamped `signal_type=metric` + `source=prometheus` - /// options and keyed by table name. The metric engine's physical table + /// Implicit declarations of the well-known entity-descriptor metrics + /// (the `prometheus_info_metrics` / `otel_info_metrics` whitelists of + /// `conventions.yaml`), gated on the ingest-stamped `signal_type=metric` + /// option plus the whitelist's expected `source`; OTel descriptors also + /// require `metric.type=info`. The metric engine's physical table /// aggregates every logical table's columns and must not contribute a /// duplicate source. - fn extend_with_prometheus_conventions( + fn extend_with_info_metric_conventions( table_info: &TableInfo, - conventions: &Conventions, + whitelist: &BTreeMap>, + expected_source: &str, + expected_metric_type: Option<&str>, declarations: &mut Vec, ) { - let Some(implicit_entities) = conventions.prometheus_info_metrics.get(&table_info.name) - else { + let Some(implicit_entities) = whitelist.get(&table_info.name) else { return; }; let options = &table_info.meta.options.extra_options; if options.get(SEMANTIC_SIGNAL_TYPE).map(String::as_str) != Some(SIGNAL_TYPE_METRIC) - || options.get(SEMANTIC_SOURCE).map(String::as_str) != Some(SOURCE_PROMETHEUS) + || options.get(SEMANTIC_SOURCE).map(String::as_str) != Some(expected_source) + || expected_metric_type.is_some_and(|expected| { + options.get(SEMANTIC_METRIC_TYPE).map(String::as_str) != Some(expected) + }) || table_info.is_physical_table() { debug!( - "Table `{}` matches the info-metric whitelist but is not a prometheus logical \ - metric table; skipping its implicit declarations", + "Table `{}` matches the info-metric whitelist but is not an eligible \ + `{expected_source}` info-metric source; skipping its implicit declarations", table_info.name ); return; @@ -296,12 +317,19 @@ impl EntityGraphProviderImpl { .cloned() .collect() }; + // A table predating the qualifier column keeps the unqualified + // identity rather than losing the declaration. + let id_qualifier = implicit + .qualified_by + .clone() + .filter(|c| schema.column_schema_by_name(c).is_some()); declarations.push(EntityDeclaration { schema: table_info.schema_name.clone(), table: table_info.name.clone(), time_index: time_index.clone(), entity_type: implicit.entity.clone(), id_columns: implicit.id.clone(), + id_qualifier, descriptive_columns, scope_columns: vec![], }); @@ -832,6 +860,24 @@ mod tests { declarations[2].id_columns, vec!["service_name", "resource_attributes.service.instance.id"] ); + assert_eq!(declarations[1].id_qualifier, None); + + let namespaced = table_info( + &[ + "service_name", + "resource_attributes.service.namespace", + "resource_attributes.service.instance.id", + ], + &[(TABLE_DATA_MODEL, TABLE_DATA_MODEL_TRACE_V1)], + ); + for declaration in sorted_declarations(&namespaced) { + assert_eq!( + declaration.id_qualifier.as_deref(), + Some("resource_attributes.service.namespace"), + "{} must qualify its identity like the metric side's job", + declaration.entity_type + ); + } // Missing uid column: no pod entity synthesized, no name-based guess. let no_uid = table_info( @@ -856,6 +902,47 @@ mod tests { assert!(sorted_declarations(&invalid_explicit).is_empty()); } + #[test] + fn trace_table_host_and_container_require_stable_ids() { + let with_ids = table_info( + &[ + "service_name", + "resource_attributes.host.id", + "resource_attributes.host.name", + "resource_attributes.container.id", + ], + &[(TABLE_DATA_MODEL, TABLE_DATA_MODEL_TRACE_V1)], + ); + let declarations = sorted_declarations(&with_ids); + let types: Vec<&str> = declarations + .iter() + .map(|d| d.entity_type.as_str()) + .collect(); + assert_eq!(types, vec!["container", "host", "service"]); + assert_eq!( + declarations[1].id_columns, + vec!["resource_attributes.host.id"] + ); + assert_eq!( + declarations[1].descriptive_columns, + vec!["resource_attributes.host.name"] + ); + + let names_only = table_info( + &[ + "service_name", + "resource_attributes.host.name", + "resource_attributes.container.name", + ], + &[(TABLE_DATA_MODEL, TABLE_DATA_MODEL_TRACE_V1)], + ); + let types: Vec = sorted_declarations(&names_only) + .into_iter() + .map(|d| d.entity_type) + .collect(); + assert_eq!(types, vec!["service"]); + } + #[test] fn prometheus_implicit_declarations_are_gated() { let labels: &[&str] = &["namespace", "pod", "node"]; @@ -900,6 +987,111 @@ mod tests { assert_eq!(declarations[1].id_columns, vec!["pod"]); } + const OTEL_STAMPS: &[(&str, &str)] = &[ + (SEMANTIC_SIGNAL_TYPE, SIGNAL_TYPE_METRIC), + (SEMANTIC_SOURCE, SOURCE_OPENTELEMETRY), + (SEMANTIC_METRIC_TYPE, servers::semantic::METRIC_TYPE_INFO), + ]; + + #[test] + fn greptime_otel_resource_info_gets_implicit_declarations() { + let info = prom_table_info( + "greptime_otel_resource_info", + &[ + "job", + "instance", + "service.name", + "service.namespace", + "host.id", + "host.name", + "container.id", + "container.name", + "k8s.pod.uid", + "k8s.pod.name", + "k8s.namespace.name", + ], + OTEL_STAMPS, + ); + let declarations = sorted_declarations(&info); + let types: Vec<&str> = declarations + .iter() + .map(|d| d.entity_type.as_str()) + .collect(); + assert_eq!( + types, + vec![ + "container", + "host", + "k8s.pod", + "service", + "service.instance" + ] + ); + assert_eq!(declarations[1].id_columns, vec!["host.id"]); + assert_eq!(declarations[1].descriptive_columns, vec!["host.name"]); + assert_eq!(declarations[3].id_columns, vec!["job"]); + assert_eq!( + declarations[3].descriptive_columns, + vec!["service.name", "service.namespace"] + ); + assert_eq!(declarations[4].id_columns, vec!["job", "instance"]); + assert!(declarations[4].descriptive_columns.is_empty()); + + let partial = prom_table_info( + "greptime_otel_resource_info", + &["job", "service.name", "host.id"], + OTEL_STAMPS, + ); + let types: Vec = sorted_declarations(&partial) + .into_iter() + .map(|d| d.entity_type) + .collect(); + assert_eq!(types, vec!["host", "service"]); + } + + #[test] + fn otel_implicit_declarations_are_gated() { + let labels: &[&str] = &["job", "instance", "host.id"]; + assert!( + sorted_declarations(&prom_table_info( + "greptime_otel_resource_info", + labels, + PROM_STAMPS + )) + .is_empty() + ); + let mut stamps = OTEL_STAMPS.to_vec(); + stamps.push((PHYSICAL_TABLE_METADATA_KEY, "true")); + assert!( + sorted_declarations(&prom_table_info( + "greptime_otel_resource_info", + labels, + &stamps + )) + .is_empty() + ); + assert!( + conventions() + .unwrap() + .otel_info_metrics + .contains_key(servers::otlp::metrics::OTEL_RESOURCE_INFO_TABLE_NAME) + ); + + let wrong_type = [ + (SEMANTIC_SIGNAL_TYPE, SIGNAL_TYPE_METRIC), + (SEMANTIC_SOURCE, SOURCE_OPENTELEMETRY), + (SEMANTIC_METRIC_TYPE, servers::semantic::METRIC_TYPE_GAUGE), + ]; + assert!( + sorted_declarations(&prom_table_info( + "greptime_otel_resource_info", + labels, + &wrong_type + )) + .is_empty() + ); + } + #[test] fn target_info_descriptive_rest_covers_remaining_tags() { let info = prom_table_info( diff --git a/src/frontend/src/instance/otlp.rs b/src/frontend/src/instance/otlp.rs index 08086198536..ccc09c9978c 100644 --- a/src/frontend/src/instance/otlp.rs +++ b/src/frontend/src/instance/otlp.rs @@ -50,7 +50,7 @@ use table::requests::{ use self::trace_ingest::trace_conventions; use crate::instance::Instance; -use crate::metrics::{OTLP_LOGS_ROWS, OTLP_METRICS_ROWS}; +use crate::metrics::{OTLP_LOGS_ROWS, OTLP_METRICS_ROWS, OTLP_RESOURCE_INFO_WRITE_ERRORS}; fn trace_permission_targets( table_name: &str, @@ -119,9 +119,15 @@ impl OpenTelemetryProtocolHandler for Instance { .cloned() .unwrap_or_default(); metric_ctx.is_legacy = is_legacy; + metric_ctx.resource_info = self.otlp_resource_info; - let (requests, rows, semantic_index, mut outcome) = - otlp::metrics::to_grpc_insert_requests(request, &mut metric_ctx)?; + let otlp::metrics::MetricsConversion { + requests, + rows, + semantic_index, + resource_info, + mut outcome, + } = otlp::metrics::to_grpc_insert_requests(request, &mut metric_ctx)?; if outcome.rejected_data_points > 0 { warn!( "Rejected {} OTLP exponential histogram data points: {}", @@ -156,7 +162,7 @@ impl OpenTelemetryProtocolHandler for Instance { // OTLP tables have one sample field in both the legacy and physical paths. let output = if metric_ctx.is_legacy || !metric_ctx.with_metric_engine { - self.handle_row_inserts(requests, ctx, false, true) + self.handle_row_inserts(requests, ctx.clone(), false, true) .await .map_err(BoxedError::new) .context(error::ExecuteGrpcQuerySnafu) @@ -165,12 +171,43 @@ impl OpenTelemetryProtocolHandler for Instance { .extension(PHYSICAL_TABLE_PARAM) .unwrap_or(GREPTIME_PHYSICAL_TABLE) .to_string(); - self.handle_metric_row_inserts(requests, ctx, physical_table) + self.handle_metric_row_inserts(requests, ctx.clone(), physical_table) .await .map_err(BoxedError::new) .context(error::ExecuteGrpcQuerySnafu) }?; outcome.write_cost = output.meta.cost; + + // Derived enrichment, written after the metric data is committed: + // failing here would make the client retry data the server already + // accepted, so every failure degrades to a warning instead. + if let Some(resource_info) = resource_info { + let written = match self.check_row_insert_permission( + &resource_info, + &ctx, + PermissionReq::Action(OTLP_WRITE), + ) { + Ok(_) => self + .handle_row_inserts(resource_info, ctx, false, false) + .await + .map_err(BoxedError::new) + .map_err(|e| e.to_string()), + Err(e) => Err(e.to_string()), + }; + match written { + Ok(descriptor_output) => outcome.write_cost += descriptor_output.meta.cost, + Err(e) => { + OTLP_RESOURCE_INFO_WRITE_ERRORS.inc(); + warn!("Failed to write the OTLP resource descriptor table: {e}"); + outcome.error_message.get_or_insert(format!( + "metric data was accepted, but writing the resource \ + descriptor table `{}` failed: {e}", + otlp::metrics::OTEL_RESOURCE_INFO_TABLE_NAME + )); + } + } + } + Ok(outcome) } diff --git a/src/frontend/src/lib.rs b/src/frontend/src/lib.rs index c1702360739..1468fe7c0b2 100644 --- a/src/frontend/src/lib.rs +++ b/src/frontend/src/lib.rs @@ -12,6 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. +#![recursion_limit = "256"] + pub mod error; pub mod events; pub mod frontend; diff --git a/src/frontend/src/metrics.rs b/src/frontend/src/metrics.rs index aba33637cfc..f649db0b46f 100644 --- a/src/frontend/src/metrics.rs +++ b/src/frontend/src/metrics.rs @@ -45,6 +45,14 @@ lazy_static! { ) .unwrap(); + /// Failed writes of the synthesized OTLP resource descriptor table; these + /// surface as OTLP partial-success warnings, not request failures. + pub static ref OTLP_RESOURCE_INFO_WRITE_ERRORS: IntCounter = register_int_counter!( + "greptime_frontend_otlp_resource_info_write_errors", + "frontend otlp resource descriptor write errors" + ) + .unwrap(); + /// The number of OpenTelemetry traces send by frontend node. pub static ref OTLP_TRACES_ROWS: IntCounter = register_int_counter!( "greptime_frontend_otlp_traces_rows", diff --git a/src/frontend/src/service_config/otlp.rs b/src/frontend/src/service_config/otlp.rs index 5dcab2c8ce2..d6a36f702e4 100644 --- a/src/frontend/src/service_config/otlp.rs +++ b/src/frontend/src/service_config/otlp.rs @@ -23,6 +23,11 @@ pub struct OtlpOptions { pub experimental_enable_exponential_histogram: bool, /// Maximum spans per trace ingest chunk. Set to 0 to disable splitting. pub trace_ingest_chunk_size: usize, + /// Whether to synthesize the `greptime_otel_resource_info` descriptor + /// table from the resource attributes of OTLP metrics, so metrics-only + /// services reach the semantic graph. Off by default: it creates and + /// writes a table the user did not send. + pub experimental_enable_resource_info: bool, } impl Default for OtlpOptions { @@ -31,6 +36,7 @@ impl Default for OtlpOptions { enable: true, experimental_enable_exponential_histogram: false, trace_ingest_chunk_size: DEFAULT_TRACE_INGEST_CHUNK_SIZE, + experimental_enable_resource_info: false, } } } @@ -45,6 +51,7 @@ mod tests { assert!(default.enable); assert!(!default.experimental_enable_exponential_histogram); assert_eq!(default.trace_ingest_chunk_size, 512); + assert!(!default.experimental_enable_resource_info); let options: OtlpOptions = toml::from_str("enable = false").unwrap(); assert!(!options.enable); diff --git a/src/meta-srv/src/handler/persist_stats_handler.rs b/src/meta-srv/src/handler/persist_stats_handler.rs index b1378c39e6b..fcdb9fa04bc 100644 --- a/src/meta-srv/src/handler/persist_stats_handler.rs +++ b/src/meta-srv/src/handler/persist_stats_handler.rs @@ -22,7 +22,7 @@ use client::inserter::{Context as InserterContext, Inserter}; use common_catalog::consts::DEFAULT_PRIVATE_SCHEMA_NAME; use common_macro::{Schema, ToRow}; use common_meta::DatanodeId; -use common_meta::datanode::RegionStat; +use common_meta::datanode::{REGION_STATS_HISTORY_TABLE_NAME, RegionStat}; use common_telemetry::warn; use dashmap::DashMap; use store_api::region_engine::RegionRole; @@ -40,8 +40,6 @@ pub struct PersistStatsHandler { persist_interval: Duration, } -/// The name of the table to persist region stats. -const META_REGION_STATS_HISTORY_TABLE_NAME: &str = "region_statistics_history"; /// The default context to persist region stats. const DEFAULT_CONTEXT: InserterContext = InserterContext { catalog: DEFAULT_CATALOG_NAME, @@ -207,7 +205,7 @@ impl PersistStatsHandler { &DEFAULT_CONTEXT, RowInsertRequests { inserts: vec![RowInsertRequest { - table_name: META_REGION_STATS_HISTORY_TABLE_NAME.to_string(), + table_name: REGION_STATS_HISTORY_TABLE_NAME.to_string(), rows: Some(Rows { schema: PersistRegionStat::schema(), rows, @@ -537,7 +535,7 @@ mod tests { }; assert_eq!( request.table_name, - META_REGION_STATS_HISTORY_TABLE_NAME.to_string() + REGION_STATS_HISTORY_TABLE_NAME.to_string() ); assert_eq!(request.rows.unwrap().rows, vec![expected_row]); diff --git a/src/metric-engine/src/engine/flush.rs b/src/metric-engine/src/engine/flush.rs index 62b2266249c..21e642de0b2 100644 --- a/src/metric-engine/src/engine/flush.rs +++ b/src/metric-engine/src/engine/flush.rs @@ -121,6 +121,7 @@ mod tests { .map(|path| path.replace(&e.file_id, "")); e.file_id = "".to_string(); e.index_version = 0; + e.max_row_group_uncompressed_size = 0; // Round down sizes to nearest 1000 to avoid exact size // comparisons that break when the SST format changes. e.file_size = e.file_size / 1000 * 1000; @@ -133,12 +134,12 @@ mod tests { assert_eq!( debug_format, r#" -ManifestSstEntry { table_dir: "test_metric_region/", region_id: 47244640257(11, 1), table_id: 11, region_number: 1, region_group: 0, region_sequence: 1, file_id: "", index_version: 0, level: 0, file_path: "test_metric_region/11_0000000001/data/.parquet", file_size: 3000, index_file_path: Some("test_metric_region/11_0000000001/data/index/.puffin"), index_file_size: Some(0), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9::Millisecond, sequence: Some(20), origin_region_id: 47244640257(11, 1), node_id: None, visible: true, primary_key_min: Some(b"\x80\0\0\x04\x01\0\0\x04W\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x02\x01\x01tag_0\0\0\0\x05"), primary_key_max: Some(b"\x80\0\0\x04\x01\0\0\x04W\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x02\x01\x01tag_0\0\0\0\x05") } -ManifestSstEntry { table_dir: "test_metric_region/", region_id: 47244640258(11, 2), table_id: 11, region_number: 2, region_group: 0, region_sequence: 2, file_id: "", index_version: 0, level: 0, file_path: "test_metric_region/11_0000000002/data/.parquet", file_size: 3000, index_file_path: Some("test_metric_region/11_0000000002/data/index/.puffin"), index_file_size: Some(0), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9::Millisecond, sequence: Some(10), origin_region_id: 47244640258(11, 2), node_id: None, visible: true, primary_key_min: Some(b"\x80\0\0\x04\x01\0\0\x04W\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x02\x01\x01tag_0\0\0\0\x05"), primary_key_max: Some(b"\x80\0\0\x04\x01\0\0\x04W\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x02\x01\x01tag_0\0\0\0\x05") } -ManifestSstEntry { table_dir: "test_metric_region/", region_id: 47261417473(11, 16777217), table_id: 11, region_number: 16777217, region_group: 1, region_sequence: 1, file_id: "", index_version: 0, level: 0, file_path: "test_metric_region/11_0000000001/metadata/.parquet", file_size: 4000, index_file_path: None, index_file_size: None, num_rows: 8, num_row_groups: 1, num_series: Some(8), min_ts: 0::Millisecond, max_ts: 0::Millisecond, sequence: Some(8), origin_region_id: 47261417473(11, 16777217), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01__column\t_4771708\t665867_Z\t3JlcHRpb\tWVfdGltZ\tXN0YW1w\0\x07"), primary_key_max: Some(b"\x01\x01__region\t_4771708\t665868\0\0\x06") } -ManifestSstEntry { table_dir: "test_metric_region/", region_id: 47261417474(11, 16777218), table_id: 11, region_number: 16777218, region_group: 1, region_sequence: 2, file_id: "", index_version: 0, level: 0, file_path: "test_metric_region/11_0000000002/metadata/.parquet", file_size: 4000, index_file_path: None, index_file_size: None, num_rows: 4, num_row_groups: 1, num_series: Some(4), min_ts: 0::Millisecond, max_ts: 0::Millisecond, sequence: Some(4), origin_region_id: 47261417474(11, 16777218), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01__column\t_4771708\t665858_Z\t3JlcHRpb\tWVfdGltZ\tXN0YW1w\0\x07"), primary_key_max: Some(b"\x01\x01__region\t_4771708\t665858\0\0\x06") } -ManifestSstEntry { table_dir: "test_metric_region/", region_id: 94489280554(22, 42), table_id: 22, region_number: 42, region_group: 0, region_sequence: 42, file_id: "", index_version: 0, level: 0, file_path: "test_metric_region/22_0000000042/data/.parquet", file_size: 3000, index_file_path: Some("test_metric_region/22_0000000042/data/index/.puffin"), index_file_size: Some(0), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9::Millisecond, sequence: Some(10), origin_region_id: 94489280554(22, 42), node_id: None, visible: true, primary_key_min: Some(b"\x80\0\0\x04\x01\0\0\x08\xae\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x02\x01\x01tag_0\0\0\0\x05"), primary_key_max: Some(b"\x80\0\0\x04\x01\0\0\x08\xae\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x02\x01\x01tag_0\0\0\0\x05") } -ManifestSstEntry { table_dir: "test_metric_region/", region_id: 94506057770(22, 16777258), table_id: 22, region_number: 16777258, region_group: 1, region_sequence: 42, file_id: "", index_version: 0, level: 0, file_path: "test_metric_region/22_0000000042/metadata/.parquet", file_size: 4000, index_file_path: None, index_file_size: None, num_rows: 4, num_row_groups: 1, num_series: Some(4), min_ts: 0::Millisecond, max_ts: 0::Millisecond, sequence: Some(4), origin_region_id: 94506057770(22, 16777258), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01__column\t_9543417\t331754_Z\t3JlcHRpb\tWVfdGltZ\tXN0YW1w\0\x07"), primary_key_max: Some(b"\x01\x01__region\t_9543417\t331754\0\0\x06") }"#, +ManifestSstEntry { table_dir: "test_metric_region/", region_id: 47244640257(11, 1), table_id: 11, region_number: 1, region_group: 0, region_sequence: 1, file_id: "", index_version: 0, level: 0, file_path: "test_metric_region/11_0000000001/data/.parquet", file_size: 3000, max_row_group_uncompressed_size: 0, index_file_path: Some("test_metric_region/11_0000000001/data/index/.puffin"), index_file_size: Some(0), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9::Millisecond, sequence: Some(20), partition_expr: None, origin_region_id: 47244640257(11, 1), node_id: None, visible: true, primary_key_min: Some(b"\x80\0\0\x04\x01\0\0\x04W\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x02\x01\x01tag_0\0\0\0\x05"), primary_key_max: Some(b"\x80\0\0\x04\x01\0\0\x04W\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x02\x01\x01tag_0\0\0\0\x05") } +ManifestSstEntry { table_dir: "test_metric_region/", region_id: 47244640258(11, 2), table_id: 11, region_number: 2, region_group: 0, region_sequence: 2, file_id: "", index_version: 0, level: 0, file_path: "test_metric_region/11_0000000002/data/.parquet", file_size: 3000, max_row_group_uncompressed_size: 0, index_file_path: Some("test_metric_region/11_0000000002/data/index/.puffin"), index_file_size: Some(0), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9::Millisecond, sequence: Some(10), partition_expr: None, origin_region_id: 47244640258(11, 2), node_id: None, visible: true, primary_key_min: Some(b"\x80\0\0\x04\x01\0\0\x04W\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x02\x01\x01tag_0\0\0\0\x05"), primary_key_max: Some(b"\x80\0\0\x04\x01\0\0\x04W\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x02\x01\x01tag_0\0\0\0\x05") } +ManifestSstEntry { table_dir: "test_metric_region/", region_id: 47261417473(11, 16777217), table_id: 11, region_number: 16777217, region_group: 1, region_sequence: 1, file_id: "", index_version: 0, level: 0, file_path: "test_metric_region/11_0000000001/metadata/.parquet", file_size: 4000, max_row_group_uncompressed_size: 0, index_file_path: None, index_file_size: None, num_rows: 8, num_row_groups: 1, num_series: Some(8), min_ts: 0::Millisecond, max_ts: 0::Millisecond, sequence: Some(8), partition_expr: None, origin_region_id: 47261417473(11, 16777217), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01__column\t_4771708\t665867_Z\t3JlcHRpb\tWVfdGltZ\tXN0YW1w\0\x07"), primary_key_max: Some(b"\x01\x01__region\t_4771708\t665868\0\0\x06") } +ManifestSstEntry { table_dir: "test_metric_region/", region_id: 47261417474(11, 16777218), table_id: 11, region_number: 16777218, region_group: 1, region_sequence: 2, file_id: "", index_version: 0, level: 0, file_path: "test_metric_region/11_0000000002/metadata/.parquet", file_size: 4000, max_row_group_uncompressed_size: 0, index_file_path: None, index_file_size: None, num_rows: 4, num_row_groups: 1, num_series: Some(4), min_ts: 0::Millisecond, max_ts: 0::Millisecond, sequence: Some(4), partition_expr: None, origin_region_id: 47261417474(11, 16777218), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01__column\t_4771708\t665858_Z\t3JlcHRpb\tWVfdGltZ\tXN0YW1w\0\x07"), primary_key_max: Some(b"\x01\x01__region\t_4771708\t665858\0\0\x06") } +ManifestSstEntry { table_dir: "test_metric_region/", region_id: 94489280554(22, 42), table_id: 22, region_number: 42, region_group: 0, region_sequence: 42, file_id: "", index_version: 0, level: 0, file_path: "test_metric_region/22_0000000042/data/.parquet", file_size: 3000, max_row_group_uncompressed_size: 0, index_file_path: Some("test_metric_region/22_0000000042/data/index/.puffin"), index_file_size: Some(0), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9::Millisecond, sequence: Some(10), partition_expr: None, origin_region_id: 94489280554(22, 42), node_id: None, visible: true, primary_key_min: Some(b"\x80\0\0\x04\x01\0\0\x08\xae\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x02\x01\x01tag_0\0\0\0\x05"), primary_key_max: Some(b"\x80\0\0\x04\x01\0\0\x08\xae\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x02\x01\x01tag_0\0\0\0\x05") } +ManifestSstEntry { table_dir: "test_metric_region/", region_id: 94506057770(22, 16777258), table_id: 22, region_number: 16777258, region_group: 1, region_sequence: 42, file_id: "", index_version: 0, level: 0, file_path: "test_metric_region/22_0000000042/metadata/.parquet", file_size: 4000, max_row_group_uncompressed_size: 0, index_file_path: None, index_file_size: None, num_rows: 4, num_row_groups: 1, num_series: Some(4), min_ts: 0::Millisecond, max_ts: 0::Millisecond, sequence: Some(4), partition_expr: None, origin_region_id: 94506057770(22, 16777258), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01__column\t_9543417\t331754_Z\t3JlcHRpb\tWVfdGltZ\tXN0YW1w\0\x07"), primary_key_max: Some(b"\x01\x01__region\t_9543417\t331754\0\0\x06") }"#, ); // list from storage let storage_entries = mito diff --git a/src/mito2/benches/bench_compaction_picker.rs b/src/mito2/benches/bench_compaction_picker.rs index 74142127aee..439b9e6f226 100644 --- a/src/mito2/benches/bench_compaction_picker.rs +++ b/src/mito2/benches/bench_compaction_picker.rs @@ -17,7 +17,6 @@ use std::hint::black_box; use criterion::{Criterion, criterion_group, criterion_main}; use mito2::compaction::run::{ Item, Ranged, SortedRun, find_overlapping_items, find_sorted_runs, find_sorted_runs_original, - merge_seq_files, reduce_runs, }; #[derive(Clone, Debug, Eq, Hash, PartialEq)] @@ -49,20 +48,6 @@ impl Item for MockFile { } } -fn generate_test_files(n: usize) -> Vec { - let mut files = Vec::with_capacity(n); - for _ in 0..n { - // Create slightly overlapping ranges to force multiple sorted runs - files.push(MockFile { - start: 0, - end: 10, - pk: 0, - size: 10, - }); - } - files -} - fn generate_same_timestamp_files(total_files: usize, files_per_timestamp: usize) -> Vec { const TIMESTAMP_INTERVAL_SECS: i64 = 10 * 60; @@ -107,21 +92,6 @@ fn bench_find_sorted_runs(c: &mut Criterion) { group.finish(); } -fn bench_reduce_runs(c: &mut Criterion) { - let mut group = c.benchmark_group("reduce_runs"); - - for size in [10, 100, 1000].iter() { - group.bench_function(format!("size_{}", size), |b| { - let mut files = generate_test_files(*size); - let runs = find_sorted_runs(&mut files); - b.iter(|| { - reduce_runs(black_box(runs.clone())); - }); - }); - } - group.finish(); -} - fn bench_find_overlapping_items(c: &mut Criterion) { let mut group = c.benchmark_group("find_overlapping_items"); @@ -158,45 +128,9 @@ fn bench_find_overlapping_items(c: &mut Criterion) { group.finish(); } -fn bench_merge_seq_files(c: &mut Criterion) { - let mut group = c.benchmark_group("merge_seq_files"); - - for size in [10, 100, 1000].iter() { - group.bench_function(format!("size_{}", size), |b| { - // Create a set of files with varying sizes - let mut files = Vec::with_capacity(*size); - - for i in 0..*size { - // Create files with different sizes to test the scoring algorithm - let file_size = if i % 3 == 0 { - 5 - } else if i % 3 == 1 { - 10 - } else { - 15 - }; - - files.push(MockFile { - start: i as i64, - end: (i + 1) as i64, - pk: 0, - size: file_size, - }); - } - - b.iter(|| { - merge_seq_files(black_box(&files), black_box(Some(50))); - }); - }); - } - group.finish(); -} - criterion_group!( benches, bench_find_sorted_runs, - bench_reduce_runs, bench_find_overlapping_items, - bench_merge_seq_files ); criterion_main!(benches); diff --git a/src/mito2/src/compaction/run.rs b/src/mito2/src/compaction/run.rs index 1ba45690215..5698a6ce667 100644 --- a/src/mito2/src/compaction/run.rs +++ b/src/mito2/src/compaction/run.rs @@ -20,14 +20,9 @@ use std::collections::BinaryHeap; use bytes::{Buf, Bytes}; use common_base::BitVec; -use common_base::readable_size::ReadableSize; use common_time::Timestamp; -use smallvec::{SmallVec, smallvec}; -use crate::sst::file::{FileHandle, RegionFileId}; - -/// Default max compaction output file size when not specified. -const DEFAULT_MAX_OUTPUT_SIZE: u64 = ReadableSize::gb(2).as_bytes(); +use crate::sst::file::FileHandle; /// Trait for any items with specific range (both boundaries are inclusive). pub trait Ranged { @@ -152,71 +147,11 @@ pub trait Item: Ranged + Clone { fn size(&self) -> usize; } -/// A group of files that are created by the same compaction task. -#[derive(Debug, Clone)] -pub struct FileGroup { - files: SmallVec<[FileHandle; 2]>, - size: usize, - num_rows: usize, - min_timestamp: Timestamp, - max_timestamp: Timestamp, - primary_key_range: Option<(Bytes, Bytes)>, -} - -impl FileGroup { - pub(crate) fn new_with_file(file: FileHandle) -> Self { - let size = file.size() as usize; - let (min_timestamp, max_timestamp) = file.time_range(); - let num_rows = file.num_rows(); - let primary_key_range = file.primary_key_range(); - Self { - files: smallvec![file], - size, - num_rows, - min_timestamp, - max_timestamp, - primary_key_range, - } - } - - pub(crate) fn num_rows(&self) -> usize { - self.num_rows - } - - pub(crate) fn add_file(&mut self, file: FileHandle) { - self.size += file.size() as usize; - self.num_rows += file.num_rows(); - let (min_timestamp, max_timestamp) = file.time_range(); - self.min_timestamp = self.min_timestamp.min(min_timestamp); - self.max_timestamp = self.max_timestamp.max(max_timestamp); - self.primary_key_range = - merge_primary_key_ranges(self.primary_key_range.take(), file.primary_key_range()); - self.files.push(file); - } - - pub(crate) fn num_files(&self) -> usize { - self.files.len() - } - - #[cfg(test)] - pub(crate) fn files(&self) -> &[FileHandle] { - &self.files[..] - } - - pub(crate) fn file_ids(&self) -> SmallVec<[RegionFileId; 2]> { - SmallVec::from_iter(self.files.iter().map(|f| f.file_id())) - } - - pub(crate) fn into_files(self) -> impl Iterator { - self.files.into_iter() - } -} - -impl Ranged for FileGroup { +impl Ranged for FileHandle { type BoundType = Timestamp; fn range(&self) -> (Self::BoundType, Self::BoundType) { - (self.min_timestamp, self.max_timestamp) + self.time_range() } fn overlap(&self, other: &Self) -> bool { @@ -226,7 +161,7 @@ impl Ranged for FileGroup { return false; } - match (&self.primary_key_range, &other.primary_key_range) { + match (&self.primary_key_range(), &other.primary_key_range()) { (Some(lhs), Some(rhs)) => primary_key_ranges_overlap(lhs, rhs), _ => true, } @@ -239,16 +174,16 @@ impl Ranged for FileGroup { return false; } - match (&self.primary_key_range, &other.primary_key_range) { + match (&self.primary_key_range(), &other.primary_key_range()) { (Some(lhs), Some(rhs)) => primary_key_ranges_overlap(lhs, rhs), _ => true, } } } -impl Item for FileGroup { +impl Item for FileHandle { fn size(&self) -> usize { - self.size + self.size() as usize } } @@ -553,124 +488,8 @@ where runs.into_iter().map(|x| x.run).collect() } -/// Finds a set of files with minimum penalty to merge that can reduce the total num of runs. -/// The penalty of merging is defined as the size of all overlapping files between two runs. -pub fn reduce_runs(mut runs: Vec>) -> Vec { - assert!(runs.len() > 1); - // sort runs by size - runs.sort_unstable_by_key(|a| a.size); - // limit max probe runs to 100 - let probe_end = runs.len().min(100); - let mut min_penalty = usize::MAX; - let mut files = vec![]; - let mut temp_files = vec![]; - for i in 0..probe_end { - for j in i + 1..probe_end { - let (a, b) = runs.split_at_mut(j); - find_overlapping_items(&mut a[i], &mut b[0], &mut temp_files); - let penalty = temp_files.iter().map(|e| e.size()).sum(); - if penalty < min_penalty { - min_penalty = penalty; - files.clear(); - files.extend_from_slice(&temp_files); - } - } - } - files -} - -/// Finds the optimal set of adjacent files to merge based on a scoring system. -/// -/// This function evaluates all possible contiguous subsets of files to find the best -/// candidates for merging, considering: -/// -/// 1. File reduction - prioritizes merging more files to reduce the total count -/// 2. Write amplification - minimizes the ratio of largest file to total size -/// 3. Size efficiency - prefers merges that utilize available space effectively -/// -/// When multiple merge candidates have the same score, older files (those with lower indices) -/// are preferred. -/// -/// # Arguments -/// * `input_files` - Slice of files to consider for merging -/// * `max_file_size` - Optional maximum size constraint for the merged file. -/// If None, uses 1.5 times the average file size. -/// -/// # Returns -/// A vector containing the best set of adjacent files to merge. -/// Returns an empty vector if input is empty or contains only one file. -pub fn merge_seq_files(input_files: &[T], max_file_size: Option) -> Vec { - if input_files.is_empty() || input_files.len() == 1 { - return vec![]; - } - - // Limit the number of files to process to 100 to control time complexity - let files_to_process = if input_files.len() > 100 { - &input_files[0..100] - } else { - input_files - }; - - // Calculate target size based on max_file_size or average file size - let target_size = match max_file_size { - Some(size) => size as usize, - None => { - // Calculate 1.5*average_file_size if max_file_size is not provided and clamp to 2GB - let total_size: usize = files_to_process.iter().map(|f| f.size()).sum(); - ((((total_size as f64) / (files_to_process.len() as f64)) * 1.5) as usize) - .min(DEFAULT_MAX_OUTPUT_SIZE as usize) - } - }; - - // Find the best group of adjacent files to merge - let mut best_group = Vec::new(); - let mut best_score = f64::NEG_INFINITY; - - // Try different starting positions - iterate from end to start to prefer older files - for start_idx in (0..files_to_process.len()).rev() { - // Try different ending positions - also iterate from end to start - for end_idx in (start_idx + 1..files_to_process.len()).rev() { - let group = &files_to_process[start_idx..=end_idx]; - let total_size: usize = group.iter().map(|f| f.size()).sum(); - - // Skip if total size exceeds target size - if total_size > target_size { - continue; // Use continue instead of break to check smaller ranges - } - - // Calculate amplification factor (largest file size / total size) - let largest_file_size = group.iter().map(|f| f.size()).max().unwrap_or(0); - let amplification_factor = largest_file_size as f64 / total_size as f64; - - // Calculate file reduction (number of files that will be reduced) - let file_reduction = group.len() - 1; - - // Calculate score based on multiple factors: - // 1. File reduction (higher is better) - // 2. Amplification factor (lower is better) - // 3. Size efficiency (how close to target size) - let file_reduction_score = file_reduction as f64 / files_to_process.len() as f64; - let amp_factor_score = (1.0 - amplification_factor) * 1.5; // Lower amplification is better - let size_efficiency = (total_size as f64 / target_size as f64).min(1.0); // Reward using available space - - let score = file_reduction_score + amp_factor_score + size_efficiency; - - // Check if this group is better than our current best - // Use >= instead of > to prefer older files (which we encounter first due to reverse iteration) - if score >= best_score { - best_score = score; - best_group = group.to_vec(); - } - } - } - - best_group -} - #[cfg(test)] mod tests { - use std::collections::HashSet; - use bytes::Bytes; use store_api::storage::FileId; @@ -709,17 +528,6 @@ mod tests { .collect() } - fn build_items_with_size(items: &[(i64, i64, usize)]) -> Vec { - items - .iter() - .map(|(start, end, size)| MockFile { - start: *start, - end: *end, - size: *size, - }) - .collect() - } - fn pk_range(min: &'static [u8], max: &'static [u8]) -> Option<(Bytes, Bytes)> { Some((Bytes::from_static(min), Bytes::from_static(max))) } @@ -844,99 +652,6 @@ mod tests { } } - fn check_reduce_runs( - files: &[(i64, i64)], - expected_runs: &[Vec<(i64, i64)>], - expected: &[(i64, i64)], - ) { - let runs = check_sorted_runs(files, expected_runs); - if runs.len() <= 1 { - assert!(expected.is_empty()); - return; - } - let files_to_merge = reduce_runs(runs); - let file_to_merge_timestamps = files_to_merge - .into_iter() - .map(|f| (f.start, f.end)) - .collect::>(); - - let expected = expected.iter().cloned().collect::>(); - assert_eq!(&expected, &file_to_merge_timestamps); - } - - #[test] - fn test_reduce_runs() { - // [1..3] [5..6] - // [2..4] - check_reduce_runs( - &[(1, 3), (2, 4), (5, 6)], - &[vec![(1, 3), (5, 6)], vec![(2, 4)]], - &[(1, 3), (2, 4)], - ); - - // [1..2][3..5] - // [4..6] - check_reduce_runs( - &[(1, 2), (3, 5), (4, 6)], - &[vec![(1, 2), (3, 5)], vec![(4, 6)]], - &[(3, 5), (4, 6)], - ); - - // [1..2][3..4] [7..8] - // [4..6] - check_reduce_runs( - &[(1, 2), (3, 4), (4, 6), (7, 8)], - &[vec![(1, 2), (3, 4), (4, 6), (7, 8)]], - &[], - ); - - // [1..2][3........6][7..8][8..9] - // [3..4][5..6] - check_reduce_runs( - &[(1, 2), (3, 4), (5, 6), (3, 6), (7, 8), (8, 9)], - &[vec![(1, 2), (3, 6), (7, 8), (8, 9)], vec![(3, 4), (5, 6)]], - &[(5, 6), (3, 4), (3, 6)], // already satisfied - ); - - // [1..2][3........6][7..8][8..9] - // [3..4][5..6] - check_reduce_runs( - &[(1, 2), (3, 4), (5, 6), (3, 6), (7, 8), (8, 9)], - &[vec![(1, 2), (3, 6), (7, 8), (8, 9)], vec![(3, 4), (5, 6)]], - &[(3, 4), (3, 6), (5, 6)], - ); - - // [10..20] [30..40] [50........80][80...100][100..110] - // [50..60] [80..90] - // - check_reduce_runs( - &[ - (10, 20), - (30, 40), - (50, 60), - (50, 80), - (80, 90), - (80, 100), - (100, 110), - ], - &[ - vec![(10, 20), (30, 40), (50, 80), (80, 100), (100, 110)], - vec![(50, 60), (80, 90)], - ], - &[(50, 80), (80, 100), (50, 60), (80, 90)], - ); - - // [0..10] - // [0...11] - // [0....12] - // [0.....13] - check_reduce_runs( - &[(0, 10), (0, 11), (0, 12), (0, 13)], - &[vec![(0, 13)], vec![(0, 12)], vec![(0, 11)], vec![(0, 10)]], - &[(0, 10), (0, 11)], - ); - } - #[test] fn test_find_overlapping_items() { let mut result = Vec::new(); @@ -1038,27 +753,25 @@ mod tests { } #[test] - fn test_file_group_overlap_time_overlap_pk_disjoint() { - let lhs = - FileGroup::new_with_file(new_file_handle_with_size_sequence_and_primary_key_range( - FileId::random(), - 0, - 100, - 0, - 1, - 10, - pk_range(b"a", b"f"), - )); - let rhs = - FileGroup::new_with_file(new_file_handle_with_size_sequence_and_primary_key_range( - FileId::random(), - 50, - 150, - 0, - 2, - 10, - pk_range(b"x", b"z"), - )); + fn test_file_overlap_time_overlap_pk_disjoint() { + let lhs = new_file_handle_with_size_sequence_and_primary_key_range( + FileId::random(), + 0, + 100, + 0, + 1, + 10, + pk_range(b"a", b"f"), + ); + let rhs = new_file_handle_with_size_sequence_and_primary_key_range( + FileId::random(), + 50, + 150, + 0, + 2, + 10, + pk_range(b"x", b"z"), + ); assert!(!lhs.overlap(&rhs)); } @@ -1066,7 +779,7 @@ mod tests { #[test] fn test_find_sorted_runs_collapses_pk_disjoint_files_into_one_run() { let mut files = vec![ - FileGroup::new_with_file(new_file_handle_with_size_sequence_and_primary_key_range( + new_file_handle_with_size_sequence_and_primary_key_range( FileId::random(), 0, 100, @@ -1074,8 +787,8 @@ mod tests { 1, 10, pk_range(b"a", b"f"), - )), - FileGroup::new_with_file(new_file_handle_with_size_sequence_and_primary_key_range( + ), + new_file_handle_with_size_sequence_and_primary_key_range( FileId::random(), 50, 150, @@ -1083,7 +796,7 @@ mod tests { 2, 10, pk_range(b"x", b"z"), - )), + ), ]; let runs = find_sorted_runs(&mut files); @@ -1095,7 +808,7 @@ mod tests { #[test] fn test_find_sorted_runs_handles_2d_transitivity_break() { let mut files = vec![ - FileGroup::new_with_file(new_file_handle_with_size_sequence_and_primary_key_range( + new_file_handle_with_size_sequence_and_primary_key_range( FileId::random(), 0, 100, @@ -1103,8 +816,8 @@ mod tests { 1, 10, pk_range(b"a", b"f"), - )), - FileGroup::new_with_file(new_file_handle_with_size_sequence_and_primary_key_range( + ), + new_file_handle_with_size_sequence_and_primary_key_range( FileId::random(), 50, 150, @@ -1112,8 +825,8 @@ mod tests { 2, 10, pk_range(b"x", b"z"), - )), - FileGroup::new_with_file(new_file_handle_with_size_sequence_and_primary_key_range( + ), + new_file_handle_with_size_sequence_and_primary_key_range( FileId::random(), 50, 150, @@ -1121,7 +834,7 @@ mod tests { 3, 10, pk_range(b"a", b"f"), - )), + ), ]; let runs = find_sorted_runs(&mut files); @@ -1133,7 +846,7 @@ mod tests { #[test] fn test_find_overlapping_items_skips_pk_disjoint_pairs() { - let mut left = SortedRun::from(vec![FileGroup::new_with_file( + let mut left = SortedRun::from(vec![ new_file_handle_with_size_sequence_and_primary_key_range( FileId::random(), 0, @@ -1143,8 +856,8 @@ mod tests { 10, pk_range(b"a", b"f"), ), - )]); - let mut right = SortedRun::from(vec![FileGroup::new_with_file( + ]); + let mut right = SortedRun::from(vec![ new_file_handle_with_size_sequence_and_primary_key_range( FileId::random(), 50, @@ -1154,7 +867,7 @@ mod tests { 10, pk_range(b"x", b"z"), ), - )]); + ]); let mut result = Vec::new(); find_overlapping_items(&mut left, &mut right, &mut result); @@ -1163,102 +876,26 @@ mod tests { } #[test] - fn test_file_group_touching_time_boundary_with_same_pk_is_not_overlap() { - let lhs = - FileGroup::new_with_file(new_file_handle_with_size_sequence_and_primary_key_range( - FileId::random(), - 0, - 100, - 0, - 1, - 10, - pk_range(b"a", b"f"), - )); - let rhs = - FileGroup::new_with_file(new_file_handle_with_size_sequence_and_primary_key_range( - FileId::random(), - 100, - 150, - 0, - 2, - 10, - pk_range(b"a", b"f"), - )); + fn test_file_touching_time_boundary_with_same_pk_is_not_overlap() { + let lhs = new_file_handle_with_size_sequence_and_primary_key_range( + FileId::random(), + 0, + 100, + 0, + 1, + 10, + pk_range(b"a", b"f"), + ); + let rhs = new_file_handle_with_size_sequence_and_primary_key_range( + FileId::random(), + 100, + 150, + 0, + 2, + 10, + pk_range(b"a", b"f"), + ); assert!(!lhs.overlap(&rhs)); } - - #[test] - fn test_merge_seq_files() { - // Test empty input - let files = Vec::::new(); - assert_eq!(merge_seq_files(&files, None), Vec::::new()); - - // Test single file input (should return empty vec as no merge needed) - let files = build_items(&[(1, 5)]); - assert_eq!(merge_seq_files(&files, None), Vec::::new()); - - // Test the example case: [10, 1, 1, 1] - should merge the last three files - let files = build_items_with_size(&[(1, 2, 10), (3, 4, 1), (5, 6, 1), (7, 8, 1)]); - let result = merge_seq_files(&files, None); - assert_eq!(result.len(), 3); - assert_eq!(result[0].size, 1); - assert_eq!(result[1].size, 1); - assert_eq!(result[2].size, 1); - - // Test with files of equal size - should merge as many as possible - let files = build_items_with_size(&[(1, 2, 5), (3, 4, 5), (5, 6, 5), (7, 8, 5)]); - let result = merge_seq_files(&files, Some(20)); - assert_eq!(result.len(), 4); // Should merge all 4 files as total size is 20 - - // Test with max_file_size constraint - let files = build_items_with_size(&[(1, 2, 5), (3, 4, 5), (5, 6, 5), (7, 8, 5)]); - let result = merge_seq_files(&files, Some(10)); - assert_eq!(result.len(), 2); // Should merge only 2 files as max size is 10 - - // Test with uneven file sizes - should prioritize reducing file count - let files = build_items_with_size(&[(1, 2, 2), (3, 4, 3), (5, 6, 4), (7, 8, 10)]); - let result = merge_seq_files(&files, Some(10)); - assert_eq!(result.len(), 3); // Should merge the first 3 files (total size 9) - - // Test amplification factor prioritization - // Two possible merges: [5, 5] (amp factor 0.5) vs [10, 1, 1] (amp factor 0.83) - let files = - build_items_with_size(&[(1, 2, 5), (3, 4, 5), (5, 6, 10), (7, 8, 1), (9, 10, 1)]); - let result = merge_seq_files(&files, Some(12)); - assert_eq!(result.len(), 2); - assert_eq!(result[0].size, 5); - assert_eq!(result[1].size, 5); - - // Test with large file preventing merges - let files = build_items_with_size(&[(1, 2, 100), (3, 4, 1), (5, 6, 1), (7, 8, 1)]); - let result = merge_seq_files(&files, Some(10)); - assert_eq!(result.len(), 3); // Should merge the last 3 small files - assert_eq!(result[0].size, 1); - assert_eq!(result[1].size, 1); - assert_eq!(result[2].size, 1); - - let files = build_items_with_size(&[(1, 2, 100), (3, 4, 20), (5, 6, 20), (7, 8, 20)]); - let result = merge_seq_files(&files, Some(200)); - assert_eq!(result.len(), 4); - - let files = build_items_with_size(&[(1, 2, 160), (3, 4, 20), (5, 6, 20), (7, 8, 20)]); - let result = merge_seq_files(&files, None); - assert_eq!(result.len(), 3); - assert_eq!(result[0].size, 20); - assert_eq!(result[1].size, 20); - assert_eq!(result[2].size, 20); - - let files = build_items_with_size(&[(1, 2, 100), (3, 4, 1)]); - let result = merge_seq_files(&files, Some(200)); - assert_eq!(result.len(), 2); - assert_eq!(result[0].size, 100); - assert_eq!(result[1].size, 1); - - let files = build_items_with_size(&[(1, 2, 20), (3, 4, 20), (5, 6, 20), (7, 8, 20)]); - let result = merge_seq_files(&files, Some(40)); - assert_eq!(result.len(), 2); - assert_eq!(result[0].start, 1); - assert_eq!(result[1].start, 3); - } } diff --git a/src/mito2/src/compaction/scheduler.rs b/src/mito2/src/compaction/scheduler.rs index 69a7acc5a5a..6aebb25eebc 100644 --- a/src/mito2/src/compaction/scheduler.rs +++ b/src/mito2/src/compaction/scheduler.rs @@ -264,9 +264,13 @@ impl CompactionScheduler { /// /// # Effects /// - /// Notifies waiters, schedules a pending manual or explicit automatic - /// follow-up, or removes the region status. The returned transition reports - /// a dispatched automatic follow-up or DDLs that are now safe to execute. + /// Notifies waiters, schedules a pending manual or automatic follow-up, or + /// removes the region status. A follow-up is scheduled whenever the cycle + /// latched an automatic trigger, or when `made_progress` is true (the + /// execution removed more files than it added): successful compaction keeps + /// draining the region until the picker returns no plan. The returned + /// transition reports a dispatched automatic follow-up or DDLs that are now + /// safe to execute. /// /// # Constraints /// @@ -279,12 +283,18 @@ impl CompactionScheduler { execution: &CompactionExecution, manifest_ctx: &ManifestContextRef, schema_metadata_manager: SchemaMetadataManagerRef, + made_progress: bool, ) -> CompactionTransition { if !self.is_current_execution(region_id, execution) { return CompactionTransition::NoAction; } - self.on_compaction_finished(region_id, manifest_ctx, schema_metadata_manager) - .await + self.on_compaction_finished( + region_id, + manifest_ctx, + schema_metadata_manager, + made_progress, + ) + .await } /// Completes a cooperatively canceled execution. @@ -571,6 +581,7 @@ impl CompactionScheduler { region_id: RegionId, manifest_ctx: &ManifestContextRef, schema_metadata_manager: SchemaMetadataManagerRef, + made_progress: bool, ) -> CompactionTransition { if !self.region_status.contains_key(®ion_id) { return CompactionTransition::NoAction; @@ -605,7 +616,11 @@ impl CompactionScheduler { return CompactionTransition::DdlReady(pending_ddl_requests); } - if status.active.reset_automatic_followup() + // Keep draining when the cycle latched an automatic trigger, or when the + // execution reduced the file count. A no-progress rewrite stops here so a + // split-heavy output cannot loop forever; the next flush trigger resumes. + let should_continue = status.active.reset_automatic_followup() || made_progress; + if should_continue && self.schedule_automatic_followup(region_id, manifest_ctx, schema_metadata_manager) { return CompactionTransition::AutomaticFollowupScheduled; diff --git a/src/mito2/src/compaction/scheduler_test.rs b/src/mito2/src/compaction/scheduler_test.rs index d3548ee8a5f..9545d0eea2a 100644 --- a/src/mito2/src/compaction/scheduler_test.rs +++ b/src/mito2/src/compaction/scheduler_test.rs @@ -826,6 +826,7 @@ async fn test_stale_plan_execution_does_not_affect_replacement_status() { &stale_execution, &manifest_ctx, schema_metadata_manager, + true, ) .await; assert!(pending_ddls.is_empty()); @@ -907,7 +908,7 @@ async fn test_schedule_compaction_skips_task_exceeding_memory_limit() { } #[tokio::test] -async fn test_execution_finished_without_followup_removes_status() { +async fn test_execution_finished_drains_until_no_plan() { common_telemetry::init_default_ut_logging(); let job_scheduler = Arc::new(VecScheduler::default()); let env = SchedulerEnv::new().await.scheduler(job_scheduler.clone()); @@ -975,23 +976,59 @@ async fn test_execution_finished_without_followup_removes_status() { &file_metas, purger.clone(), ); - // A completed execution without an explicit follow-up removes its status. - scheduler - .on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager.clone()) + // A completed execution that reduced the file count chains another pick even + // without an explicit trigger, because the layout is still compactable. + let transition = scheduler + .on_compaction_finished( + region_id, + &manifest_ctx, + schema_metadata_manager.clone(), + true, + ) .await; + assert_matches!(transition, CompactionTransition::AutomaticFollowupScheduled); + assert!(scheduler.region_status.contains_key(®ion_id)); + let finished = recv_compaction_pick_finished(&mut rx).await; + scheduler + .handle_compaction_pick_finished(finished, &manifest_ctx, schema_metadata_manager.clone()) + .await; + assert_eq!(2, job_scheduler.num_jobs()); + + // Replace the layout with a single file: the drain must stop once the picker + // finds nothing to do. + let data = version_control.current(); + let file_metas: Vec<_> = data.version.ssts.levels()[0] + .files + .values() + .map(|file| file.meta_ref().clone()) + .collect(); + apply_edit(&version_control, &[(0, end)], &file_metas, purger.clone()); + let transition = scheduler + .on_compaction_finished( + region_id, + &manifest_ctx, + schema_metadata_manager.clone(), + true, + ) + .await; + assert_matches!(transition, CompactionTransition::AutomaticFollowupScheduled); + let finished = recv_compaction_pick_finished(&mut rx).await; + let transition = scheduler + .handle_compaction_pick_finished(finished, &manifest_ctx, schema_metadata_manager.clone()) + .await; + assert_matches!(transition, CompactionTransition::NoAction); assert!(scheduler.region_status.is_empty()); - assert_eq!(1, job_scheduler.num_jobs()); + assert_eq!(2, job_scheduler.num_jobs()); } #[tokio::test] -async fn test_time_range_compaction_when_compaction_in_progress() { +async fn test_execution_finished_without_progress_removes_status() { common_telemetry::init_default_ut_logging(); let job_scheduler = Arc::new(VecScheduler::default()); let env = SchedulerEnv::new().await.scheduler(job_scheduler.clone()); let (tx, mut rx) = mpsc::channel(4); let mut scheduler = env.mock_compaction_scheduler(tx); let mut builder = VersionControlBuilder::new(); - let purger = builder.file_purger(); let region_id = builder.region_id(); let (schema_metadata_manager, kv_backend) = mock_schema_metadata_manager(); @@ -1017,6 +1054,70 @@ async fn test_time_range_compaction_when_compaction_in_progress() { .push_l0_file(90, end) .build(), ); + let manifest_ctx = env + .mock_manifest_context(version_control.current().version.metadata.clone()) + .await; + let scheduled = scheduler + .schedule_automatic_compaction( + compact_request::Options::Regular(Default::default()), + &version_control, + &env.access_layer, + &manifest_ctx, + schema_metadata_manager.clone(), + ) + .unwrap(); + assert!(scheduled); + let finished = recv_compaction_pick_finished(&mut rx).await; + scheduler + .handle_compaction_pick_finished(finished, &manifest_ctx, schema_metadata_manager.clone()) + .await; + assert_eq!(1, job_scheduler.num_jobs()); + + // The execution rewrote files without reducing the file count (e.g. its output + // was split into more files than its input). Chaining would loop without making + // progress, so the lifecycle ends here; the next flush trigger resumes compaction. + let transition = scheduler + .on_compaction_finished( + region_id, + &manifest_ctx, + schema_metadata_manager.clone(), + false, + ) + .await; + assert_matches!(transition, CompactionTransition::NoAction); + assert!(scheduler.region_status.is_empty()); + assert_eq!(1, job_scheduler.num_jobs()); +} + +#[tokio::test] +async fn test_time_range_compaction_when_compaction_in_progress() { + common_telemetry::init_default_ut_logging(); + let job_scheduler = Arc::new(VecScheduler::default()); + let env = SchedulerEnv::new().await.scheduler(job_scheduler.clone()); + let (tx, mut rx) = mpsc::channel(4); + let mut scheduler = env.mock_compaction_scheduler(tx); + let mut builder = VersionControlBuilder::new(); + let purger = builder.file_purger(); + let region_id = builder.region_id(); + + let (schema_metadata_manager, kv_backend) = mock_schema_metadata_manager(); + schema_metadata_manager + .register_region_table_info( + builder.region_id().table_id(), + "test_table", + "test_catalog", + "test_schema", + None, + kv_backend, + ) + .await; + + // 40 files to compact. The first task picks 32, leaving 8 for the pending request. + let end = 1000 * 1000; + for offset in 0..40 { + builder.push_l0_file(offset * 10, end); + } + let version_control = Arc::new(builder.build()); let manifest_ctx = env .mock_manifest_context(version_control.current().version.metadata.clone()) .await; @@ -1027,13 +1128,9 @@ async fn test_time_range_compaction_when_compaction_in_progress() { .map(|file| file.meta_ref().clone()) .collect(); - // 5 files for next compaction and removes old files. - apply_edit( - &version_control, - &[(0, end), (20, end), (40, end), (60, end), (80, end)], - &file_metas, - purger.clone(), - ); + // Replaces the files before scheduling the first compaction. + let next_files = (0..40).map(|offset| (offset * 20, end)).collect::>(); + apply_edit(&version_control, &next_files, &file_metas, purger.clone()); scheduler .schedule_automatic_compaction( @@ -1094,7 +1191,12 @@ async fn test_time_range_compaction_when_compaction_in_progress() { // On compaction finished and schedule next compaction. scheduler - .on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager.clone()) + .on_compaction_finished( + region_id, + &manifest_ctx, + schema_metadata_manager.clone(), + true, + ) .await; assert_eq!(1, scheduler.region_status.len()); assert!( @@ -1285,10 +1387,21 @@ async fn test_ranged_compaction_stops_without_followup() { assert_eq!(1, job_scheduler.num_jobs()); scheduler - .on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager.clone()) + .on_compaction_finished( + region_id, + &manifest_ctx, + schema_metadata_manager.clone(), + true, + ) .await; - assert!(!scheduler.region_status.contains_key(®ion_id)); - assert_eq!(1, job_scheduler.num_jobs()); + // The execution reduced the file count and the remaining window is still + // compactable, so the scheduler drains it with an unrestricted follow-up. + assert!(scheduler.region_status.contains_key(®ion_id)); + let followup = recv_compaction_pick_finished(&mut rx).await; + scheduler + .handle_compaction_pick_finished(followup, &manifest_ctx, schema_metadata_manager.clone()) + .await; + assert_eq!(2, job_scheduler.num_jobs()); } #[tokio::test] @@ -1369,7 +1482,12 @@ async fn test_automatic_trigger_during_execution_clears_continuation_scope() { // Finishing the manual execution must chain an unrestricted follow-up instead of // continuing with the manual request's time range. scheduler - .on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager.clone()) + .on_compaction_finished( + region_id, + &manifest_ctx, + schema_metadata_manager.clone(), + true, + ) .await; let continuation = recv_compaction_pick_finished(&mut rx).await; @@ -1936,7 +2054,7 @@ async fn test_on_compaction_finished_returns_pending_ddl_requests() { }); let pending_ddls = scheduler - .on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager) + .on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager, true) .await; assert_eq!(pending_ddls.len(), 1); @@ -1985,7 +2103,12 @@ async fn test_planning_terminal_prioritizes_pending_ddl_over_automatic_followup( assert!(result.is_ok()); let pending_ddls = scheduler - .on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager.clone()) + .on_compaction_finished( + region_id, + &manifest_ctx, + schema_metadata_manager.clone(), + true, + ) .await; assert!(pending_ddls.is_empty()); @@ -2034,7 +2157,7 @@ async fn test_on_compaction_finished_dispatches_pending_ddl_before_chained_regul }); let pending_ddls = scheduler - .on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager) + .on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager, true) .await; assert_eq!(pending_ddls.len(), 1); @@ -2056,7 +2179,7 @@ async fn test_on_compaction_finished_returns_empty_when_region_absent() { let (schema_metadata_manager, _kv_backend) = mock_schema_metadata_manager(); let pending_ddls = scheduler - .on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager) + .on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager, true) .await; assert!(pending_ddls.is_empty()); @@ -2111,7 +2234,12 @@ async fn test_on_compaction_finished_manual_schedule_error_cleans_status() { }); let pending_ddls = scheduler - .on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager.clone()) + .on_compaction_finished( + region_id, + &manifest_ctx, + schema_metadata_manager.clone(), + true, + ) .await; assert!(pending_ddls.is_empty()); diff --git a/src/mito2/src/compaction/twcs.rs b/src/mito2/src/compaction/twcs.rs index 342426bbefe..cdd613b3424 100644 --- a/src/mito2/src/compaction/twcs.rs +++ b/src/mito2/src/compaction/twcs.rs @@ -15,8 +15,7 @@ use std::collections::hash_map::Entry; use std::collections::{BTreeMap, HashMap, HashSet}; use std::fmt::Debug; -use std::num::NonZeroU64; -use std::sync::Arc; +use std::sync::{Arc, LazyLock}; use common_base::readable_size::ReadableSize; use common_telemetry::{debug, info}; @@ -32,8 +31,8 @@ use crate::compaction::buckets::infer_time_bucket; use crate::compaction::compactor::CompactionRegion; use crate::compaction::picker::{Picker, PickerOutput, get_expired_ssts}; use crate::compaction::run::{ - FileGroup, Item, Ranged, find_sorted_runs, find_sorted_runs_by_time_range, - merge_primary_key_ranges, merge_seq_files, primary_key_ranges_overlap, reduce_runs, + Ranged, SortedRun, find_sorted_runs, find_sorted_runs_by_time_range, merge_primary_key_ranges, + primary_key_ranges_overlap, }; use crate::error::{JoinSnafu, Result}; use crate::sst::file::{FileHandle, Level, overlaps}; @@ -41,8 +40,24 @@ use crate::sst::version::LevelMeta; const LEVEL_COMPACTED: Level = 1; -/// Default value for max compaction input file num. -const DEFAULT_MAX_INPUT_FILE_NUM: usize = 32; +/// Default maximum number of input SST files in one compaction input. +const DEFAULT_MAX_INPUT_FILES: usize = 32; + +const MAX_INPUT_FILES_ENV: &str = "GREPTIME_TWCS_MAX_INPUT_FILES"; + +/// Maximum number of input SST files in one compaction input. +/// Configurable via [`MAX_INPUT_FILES_ENV`]. +static MAX_INPUT_FILES: LazyLock = LazyLock::new(|| { + let env_value = std::env::var(MAX_INPUT_FILES_ENV).ok(); + parse_max_input_files(env_value.as_deref()) +}); + +fn parse_max_input_files(env_value: Option<&str>) -> usize { + env_value + .and_then(|env_value| env_value.parse().ok()) + .filter(|max_input_files| *max_input_files >= 2) + .unwrap_or(DEFAULT_MAX_INPUT_FILES) +} /// `TwcsPicker` picks files of which the max timestamp are in the same time window as compaction /// candidates. @@ -52,7 +67,8 @@ pub struct TwcsPicker { pub trigger_file_num: usize, /// Compaction time window in seconds. pub time_window_seconds: Option, - /// Max allowed compaction output file size. + /// Max allowed compaction output file size. The picker also uses it to predict + /// output splits when scoring candidates. pub max_output_file_size: Option, /// Whether the target region is in append mode. pub append_mode: bool, @@ -113,7 +129,7 @@ impl TwcsPicker { output.push(CompactionOutput { output_level: LEVEL_COMPACTED, // always compact to l1 - inputs: inputs.into_iter().flat_map(|fg| fg.into_files()).collect(), + inputs, filter_deleted, output_time_range: None, // we do not enforce output time range in twcs compactions. }); @@ -138,7 +154,11 @@ impl TwcsPicker { active_window: Option, files: &Window, windows: &BTreeMap, - ) -> (Vec, bool) { + ) -> (Vec, bool) { + if files.files.len() < self.trigger_file_num { + return (vec![], false); + } + let window = &files.time_window; let mut files_to_merge: Vec<_> = files.files().cloned().collect(); @@ -148,7 +168,7 @@ impl TwcsPicker { { let (kept_files, ignored_files) = files_to_merge .into_iter() - .partition(|fg| fg.size() <= max_size as usize); + .partition(|file| file.size() <= max_size); files_to_merge = kept_files; if !ignored_files.is_empty() { info!( @@ -167,59 +187,13 @@ impl TwcsPicker { find_sorted_runs_by_time_range(&mut files_to_merge) }; let found_runs = sorted_runs.len(); - // We only remove deletion markers if we found less than 2 runs and not in append mode. - // because after compaction there will be no overlapping files. - let mut filter_deleted = - found_runs <= 2 && !self.append_mode && !window_has_overlap(files, windows); - if found_runs == 0 { - return (vec![], filter_deleted); - } - - let mut inputs = if found_runs > 1 { - reduce_runs(sorted_runs) - } else { - let run = sorted_runs.last().unwrap(); - if run.items().len() < self.trigger_file_num { - return (vec![], filter_deleted); - } - // no overlapping files, try merge small files - merge_seq_files(run.items(), self.max_output_file_size) - }; - - // Limits the number of input files. - let total_input_files: usize = inputs.iter().map(|fg| fg.num_files()).sum(); - if total_input_files > DEFAULT_MAX_INPUT_FILE_NUM { - // Sorts file groups by size first. - inputs.sort_unstable_by_key(|fg| fg.size()); - let mut num_picked_files = 0; - inputs = inputs - .into_iter() - .take_while(|fg| { - let current_group_file_num = fg.num_files(); - if current_group_file_num + num_picked_files <= DEFAULT_MAX_INPUT_FILE_NUM { - num_picked_files += current_group_file_num; - true - } else { - false - } - }) - .collect::>(); - info!( - "Compaction for region {} enforces max input file num limit: {}, current total: {}, input: {:?}", - region_id, DEFAULT_MAX_INPUT_FILE_NUM, total_input_files, inputs - ); - } - - // The inputs are only a subset of the window: `reduce_runs` and `merge_seq_files` both - // narrow the selection down and the file num limit above narrows it further. A deletion - // marker may only be dropped when every file that can hold a row it masks is compacted - // along with it, so re-check the final selection against what stays in the window. - if filter_deleted && overlaps_files_left_behind(&inputs, &files_to_merge) { - filter_deleted = false; - } + let inputs = pick_count_first(sorted_runs, self.max_output_file_size); + let filter_deleted = !self.append_mode + && !window_has_overlap(files, windows) + && !selected_overlaps_unselected(&inputs, files); if inputs.len() > 1 { - // If we have more than one group to compact. + // If we have more than one file to compact. log_pick_result( region_id, *window, @@ -235,18 +209,225 @@ impl TwcsPicker { } } -/// Returns whether any file group of `window_files` that is not part of `inputs` overlaps -/// `inputs`. Such a group keeps rows a deletion marker in `inputs` masks, so the compaction -/// must not filter deleted rows out. +#[derive(Debug)] +struct OrderedFile<'a> { + file: &'a FileHandle, + run_id: usize, + position_in_run: usize, +} + +/// Metrics of a candidate compaction input, accumulated incrementally as the +/// candidate interval expands one file at a time. +#[derive(Debug, Default)] +struct Candidate { + /// Number of files in the interval. + num_files: usize, + /// Total input bytes of the interval. + total_size: usize, + /// Size of the largest single file in the interval. + largest_file_size: usize, + /// Files in the interval that overlap another interval file from a different + /// sorted run. Resolving these overlaps is what reduces sorted runs. + overlap_participants: usize, +} + +impl Candidate { + /// Absorbs `file` into the candidate. `preceding` holds the interval files added + /// before it and `participations` tracks which of them already participate in an + /// intra-interval overlap; both exist only for overlap accounting. + fn absorb( + &mut self, + file: &OrderedFile, + preceding: &[OrderedFile], + participations: &mut Vec, + ) { + self.num_files += 1; + let file_size = file.file.size() as usize; + self.total_size += file_size; + self.largest_file_size = self.largest_file_size.max(file_size); + + let mut participates = false; + for (offset, other) in preceding.iter().enumerate() { + if file.run_id != other.run_id && file.file.overlap_inclusive(other.file) { + if !participations[offset] { + participations[offset] = true; + self.overlap_participants += 1; + } + participates = true; + } + } + participations.push(participates); + if participates { + self.overlap_participants += 1; + } + } + + /// Predicted number of output files after compacting this candidate. Compaction + /// output holds at most the input bytes, so a split threshold of + /// `max_output_file_size` yields at most `ceil(total_size / threshold)` files. + /// Without a threshold the output is a single file. + fn predicted_output_files(&self, max_output_file_size: Option) -> usize { + match max_output_file_size { + Some(max) if max > 0 => self.total_size.div_ceil(max as usize).max(1), + _ => 1, + } + } + + /// Net file count reduction compacting this candidate would achieve, i.e. how + /// many fewer physical SSTs the window holds afterwards. Zero when the output + /// would be split back into at least as many files as the input. + fn file_reduction(&self, max_output_file_size: Option) -> usize { + self.num_files + .saturating_sub(self.predicted_output_files(max_output_file_size)) + } + + /// A large historical file is only rewritten once its peers contribute at least + /// the same amount of data, so every rewrite moves it into a larger size tier. + fn is_balanced(&self) -> bool { + self.largest_file_size <= self.total_size - self.largest_file_size + } + + /// A candidate is worth compacting only if it makes progress on at least one + /// axis: it reduces the physical file count, or it resolves at least one + /// overlap (merging sorted runs and reducing read amplification). A pure + /// rewrite that achieves neither only burns I/O. + fn makes_progress(&self, max_output_file_size: Option) -> bool { + self.file_reduction(max_output_file_size) > 0 || self.overlap_participants > 0 + } + + fn score(&self, max_output_file_size: Option) -> CandidateScore { + CandidateScore { + file_reduction: self.file_reduction(max_output_file_size), + overlap_participants: self.overlap_participants, + total_size: self.total_size, + } + } +} + +/// Score of a candidate compaction input. Higher is better: first by the predicted +/// net file count reduction, then by the number of files participating in an +/// overlap, then by smaller input bytes. +#[derive(Debug)] +struct CandidateScore { + file_reduction: usize, + overlap_participants: usize, + total_size: usize, +} + +impl CandidateScore { + fn is_better_than(&self, other: &Self) -> bool { + self.file_reduction + .cmp(&other.file_reduction) + .then_with(|| self.overlap_participants.cmp(&other.overlap_participants)) + .then_with(|| other.total_size.cmp(&self.total_size)) + .is_gt() + } +} + +/// Picks a contiguous (in global time order) interval of files to compact. /// -/// Overlapping is inclusive on both boundaries here: two files that only share a boundary -/// timestamp can still hold the same row. -fn overlaps_files_left_behind(inputs: &[FileGroup], window_files: &[FileGroup]) -> bool { - let picked: HashSet<_> = inputs.iter().flat_map(|fg| fg.file_ids()).collect(); - window_files +/// The picker reorders all files from all sorted runs by `(start asc, end desc)` and +/// enumerates every interval of at most [`MAX_INPUT_FILES`] files as a [`Candidate`]. +/// An interval is eligible when it +/// +/// - holds at least 2 files, +/// - is balanced: no single file dominates it (`largest <= sum of the others`), +/// - makes progress on at least one axis: it reduces the physical file count given +/// the output split threshold `max_output_file_size`, or it resolves at least one +/// overlap between sorted runs. +/// +/// The best interval wins by [`CandidateScore`]; ties keep the earliest interval. +fn pick_count_first( + sorted_runs: Vec>, + max_output_file_size: Option, +) -> Vec { + let files = ordered_files(&sorted_runs); + + let mut best = None; + for left in 0..files.len() { + let mut candidate = Candidate::default(); + let right_bound = left.saturating_add(*MAX_INPUT_FILES).min(files.len()); + let mut participations: Vec = Vec::with_capacity(right_bound - left); + for right in left..right_bound { + candidate.absorb(&files[right], &files[left..right], &mut participations); + if candidate.num_files < 2 + || !candidate.is_balanced() + || !candidate.makes_progress(max_output_file_size) + { + continue; + } + + let score = candidate.score(max_output_file_size); + if best + .as_ref() + .is_none_or(|(best_score, _)| score.is_better_than(best_score)) + { + best = Some((score, &files[left..=right])); + } + } + } + + let Some((_, best)) = best else { + return vec![]; + }; + best.iter().map(|file| file.file.clone()).collect() +} + +/// Flattens sorted runs into files ordered by `(start asc, end desc)`, breaking ties +/// by run and position to keep the order deterministic. +fn ordered_files(sorted_runs: &[SortedRun]) -> Vec> { + let mut files = sorted_runs .iter() - .filter(|fg| !fg.file_ids().iter().any(|id| picked.contains(id))) - .any(|fg| inputs.iter().any(|input| input.overlap_inclusive(fg))) + .enumerate() + .flat_map(|(run_id, run)| { + run.items() + .iter() + .enumerate() + .map(move |(position_in_run, file)| OrderedFile { + file, + run_id, + position_in_run, + }) + }) + .collect::>(); + files.sort_unstable_by(|lhs, rhs| { + let (lhs_start, lhs_end) = lhs.file.range(); + let (rhs_start, rhs_end) = rhs.file.range(); + lhs_start + .cmp(&rhs_start) + .then_with(|| rhs_end.cmp(&lhs_end)) + .then_with(|| lhs.run_id.cmp(&rhs.run_id)) + .then_with(|| lhs.position_in_run.cmp(&rhs.position_in_run)) + }); + files +} + +fn selected_overlaps_unselected(selected: &[FileHandle], window: &Window) -> bool { + // The overall time span of the selection: a file outside it cannot overlap any + // selected file (ranges are inclusive), so it needs no precise overlap check. + let Some((span_start, span_end)) = selected + .iter() + .map(Ranged::range) + .reduce(|(start_a, end_a), (start_b, end_b)| (start_a.min(start_b), end_a.max(end_b))) + else { + return false; + }; + let selected_file_ids = selected + .iter() + .map(FileHandle::file_id) + .collect::>(); + window + .files() + .filter(|file| { + let (start, end) = file.range(); + start <= span_end && span_start <= end + }) + .filter(|file| !selected_file_ids.contains(&file.file_id())) + .any(|unselected| { + selected + .iter() + .any(|selected| selected.overlap_inclusive(unselected)) + }) } #[allow(clippy::too_many_arguments)] @@ -258,7 +439,7 @@ fn log_pick_result( file_num: usize, max_output_file_size: Option, filter_deleted: bool, - inputs: &[FileGroup], + inputs: &[FileHandle], ) { let input_file_str: Vec = inputs .iter() @@ -268,11 +449,11 @@ fn log_pick_result( let end = range.1.to_iso8601_string(); let num_rows = f.num_rows(); format!( - "FileGroup{{id: {:?}, range: ({}, {}), size: {}, num rows: {} }}", - f.file_ids(), + "File{{id: {:?}, range: ({}, {}), size: {}, num rows: {} }}", + f.file_id(), start, end, - ReadableSize(f.size() as u64), + ReadableSize(f.size()), num_rows ) }) @@ -369,9 +550,7 @@ impl Picker for TwcsPicker { struct Window { start: Timestamp, end: Timestamp, - // Mapping from file sequence to file groups. Files with the same sequence is considered - // created from the same compaction task. - files: HashMap, FileGroup>, + files: Vec, time_window: i64, primary_key_range: Option<(bytes::Bytes, bytes::Bytes)>, } @@ -381,11 +560,10 @@ impl Window { fn new_with_file(file: FileHandle) -> Self { let (start, end) = file.time_range(); let primary_key_range = file.primary_key_range(); - let files = HashMap::from([(file.meta_ref().sequence, FileGroup::new_with_file(file))]); Self { start, end, - files, + files: vec![file], time_window: 0, primary_key_range, } @@ -403,19 +581,11 @@ impl Window { self.end = self.end.max(end); self.primary_key_range = merge_primary_key_ranges(self.primary_key_range.take(), file.primary_key_range()); - - match self.files.entry(file.meta_ref().sequence) { - Entry::Occupied(mut o) => { - o.get_mut().add_file(file); - } - Entry::Vacant(v) => { - v.insert(FileGroup::new_with_file(file)); - } - } + self.files.push(file); } - fn files(&self) -> impl Iterator { - self.files.values() + fn files(&self) -> impl Iterator { + self.files.iter() } } @@ -547,6 +717,18 @@ mod tests { use crate::test_util::memtable_util::metadata_for_test; use crate::test_util::scheduler_util::SchedulerEnv; + #[test] + fn test_valid_max_input_files_env_overrides_default() { + assert_eq!(64, parse_max_input_files(Some("64"))); + } + + #[test] + fn test_invalid_max_input_files_env_falls_back_to_default() { + for env_value in [None, Some(""), Some("invalid"), Some("0"), Some("1")] { + assert_eq!(32, parse_max_input_files(env_value)); + } + } + async fn compaction_region_with_expired_sst() -> CompactionRegion { let env = SchedulerEnv::new().await; let metadata = metadata_for_test(); @@ -664,8 +846,7 @@ mod tests { 3, ); let fgs = &windows.get(&0).unwrap().files; - assert_eq!(1, fgs.len()); - assert_eq!(fgs.values().map(|f| f.files().len()).sum::(), 5); + assert_eq!(5, fgs.len()); let files = [FileId::random(); 3]; let windows = assign_to_windows( @@ -679,26 +860,41 @@ mod tests { ); assert_eq!( files[0], - windows.get(&0).unwrap().files().next().unwrap().files()[0] + windows + .get(&0) + .unwrap() + .files() + .next() + .unwrap() .file_id() .file_id() ); assert_eq!( files[1], - windows.get(&3).unwrap().files().next().unwrap().files()[0] + windows + .get(&3) + .unwrap() + .files() + .next() + .unwrap() .file_id() .file_id() ); assert_eq!( files[2], - windows.get(&12).unwrap().files().next().unwrap().files()[0] + windows + .get(&12) + .unwrap() + .files() + .next() + .unwrap() .file_id() .file_id() ); } #[test] - fn test_assign_file_groups_to_windows() { + fn test_assign_files_to_windows() { let files = [ FileId::random(), FileId::random(), @@ -716,25 +912,14 @@ mod tests { 3, ); assert_eq!(windows.len(), 1); - let fgs = &windows.get(&0).unwrap().files; - assert_eq!(2, fgs.len()); + let window_files = &windows.get(&0).unwrap().files; + assert_eq!(4, window_files.len()); assert_eq!( - fgs.get(&NonZeroU64::new(1)) - .unwrap() - .files() + window_files .iter() .map(|f| f.file_id().file_id()) .collect::>(), - [files[0], files[1]].into_iter().collect() - ); - assert_eq!( - fgs.get(&NonZeroU64::new(2)) - .unwrap() - .files() - .iter() - .map(|f| f.file_id().file_id()) - .collect::>(), - [files[2], files[3]].into_iter().collect() + files.into_iter().collect() ); } @@ -751,11 +936,10 @@ mod tests { files[2].set_compacting(true); let mut windows = assign_to_windows(files.iter(), 3); let window0 = windows.remove(&0).unwrap(); - assert_eq!(1, window0.files.len()); + assert_eq!(3, window0.files.len()); let candidates = window0 .files - .into_values() - .flat_map(|fg| fg.into_files()) + .iter() .map(|f| f.file_id().file_id()) .collect::>(); assert_eq!(candidates.len(), 3); @@ -801,12 +985,10 @@ mod tests { assert_eq!(*overlapping, actual_overlapping); let mut file_ranges = actual_window .files - .values() - .flat_map(|f| { - f.files().iter().map(|f| { - let (s, e) = f.time_range(); - (s.value(), e.value()) - }) + .iter() + .map(|f| { + let (s, e) = f.time_range(); + (s.value(), e.value()) }) .collect::>(); file_ranges.sort_unstable_by(|l, r| l.0.cmp(&r.0).then(l.1.cmp(&r.1))); @@ -988,7 +1170,7 @@ mod tests { let active_window = find_latest_window_in_seconds(self.input_files.iter(), self.window_size); let output = TwcsPicker { - trigger_file_num: 4, + trigger_file_num: 2, time_window_seconds: None, max_output_file_size: None, append_mode: false, @@ -1075,7 +1257,7 @@ mod tests { .to_vec(), expected_outputs: vec![ ExpectedOutput { - input_files: vec![2, 4], + input_files: vec![2, 3, 4], output_level: 1, }, ExpectedOutput { @@ -1088,8 +1270,9 @@ mod tests { .await; // Case 3: - // A compaction may split output into several files that have overlapping time ranges and same sequence, - // we should treat these files as one FileGroup. + // A compaction may split output into several files that have overlapping time ranges and + // the same sequence. They are ordinary compaction candidates now: the picker merges the + // overlapping ones in the window that reaches the trigger. let file_ids = (0..6).map(|_| FileId::random()).collect::>(); CompactionPickerTestCase { window_size: 3, @@ -1101,10 +1284,16 @@ mod tests { new_file_handle_with_sequence(file_ids[4], 11, 2990, 0, 3), ] .to_vec(), - expected_outputs: vec![ExpectedOutput { - input_files: vec![0, 1, 4], - output_level: 1, - }], + expected_outputs: vec![ + ExpectedOutput { + input_files: vec![2, 3], + output_level: 1, + }, + ExpectedOutput { + input_files: vec![0, 1, 4], + output_level: 1, + }, + ], } .check() .await; @@ -1160,30 +1349,24 @@ mod tests { let small_file_2 = new_file_handle_with_size_and_sequence(file_ids[2], 0, 999, 0, 3, 800); let large_file_2 = new_file_handle_with_size_and_sequence(file_ids[3], 0, 999, 0, 4, 2000); - // Create file groups (each file is in its own group due to different sequences) - let mut files_to_merge = vec![ - FileGroup::new_with_file(small_file_1), - FileGroup::new_with_file(large_file_1), - FileGroup::new_with_file(small_file_2), - FileGroup::new_with_file(large_file_2), - ]; + let mut files_to_merge = vec![small_file_1, large_file_1, small_file_2, large_file_2]; // Test filtering logic directly let original_count = files_to_merge.len(); // Apply append mode filtering - files_to_merge.retain(|fg| fg.size() <= max_output_file_size as usize); + files_to_merge.retain(|file| file.size() <= max_output_file_size); // Should have filtered out 2 large files, leaving 2 small files assert_eq!(files_to_merge.len(), 2); assert_eq!(original_count, 4); // Verify the remaining files are the small ones - for fg in &files_to_merge { + for file in &files_to_merge { assert!( - fg.size() <= max_output_file_size as usize, + file.size() <= max_output_file_size, "File size {} should be <= {}", - fg.size(), + file.size(), max_output_file_size ); } @@ -1191,7 +1374,7 @@ mod tests { #[tokio::test] async fn test_build_output_multiple_windows_with_zero_runs() { - let file_ids = (0..6).map(|_| FileId::random()).collect::>(); + let file_ids = (0..7).map(|_| FileId::random()).collect::>(); let files = [ // Window 0: Contains 3 files but not forming any runs (not enough files in sequence to reach trigger_file_num) @@ -1202,6 +1385,7 @@ mod tests { new_file_handle_with_sequence(file_ids[3], 3000, 3999, 0, 4), new_file_handle_with_sequence(file_ids[4], 3000, 3999, 0, 5), new_file_handle_with_sequence(file_ids[5], 3000, 3999, 0, 6), + new_file_handle_with_sequence(file_ids[6], 3000, 3999, 0, 7), ]; let windows = assign_to_windows(files.iter(), 3); @@ -1431,11 +1615,11 @@ mod tests { .unwrap(); assert_eq!(1, output.len()); - assert_eq!(output[0].inputs.len(), 2); + assert_eq!(output[0].inputs.len(), num_files); } #[tokio::test] - async fn test_limit_max_input_files() { + async fn test_window_trigger_can_exceed_input_limit() { common_telemetry::init_default_ut_logging(); let num_files = 50; @@ -1460,7 +1644,7 @@ mod tests { let windows = assign_to_windows(files.iter(), 3); let picker = TwcsPicker { - trigger_file_num: 4, + trigger_file_num: num_files, time_window_seconds: Some(3), max_output_file_size: None, append_mode: false, @@ -1475,7 +1659,7 @@ mod tests { .unwrap(); assert_eq!(1, output.len()); - assert_eq!(output[0].inputs.len(), 32); + assert_eq!(output[0].inputs.len(), num_files.min(*MAX_INPUT_FILES)); } #[tokio::test] @@ -1671,5 +1855,394 @@ mod tests { ); } + #[tokio::test] + async fn test_count_first_prefers_more_files_over_smaller_overlap() { + let files = [ + new_file_handle_with_size_and_sequence(FileId::random(), 0, 10, 0, 1, 10), + new_file_handle_with_size_and_sequence(FileId::random(), 20, 30, 0, 2, 10), + new_file_handle_with_size_and_sequence(FileId::random(), 40, 50, 0, 3, 10), + new_file_handle_with_size_and_sequence(FileId::random(), 5, 15, 0, 4, 10), + ]; + let windows = assign_to_windows(files.iter(), 100); + let picker = TwcsPicker { + trigger_file_num: 2, + time_window_seconds: Some(100), + max_output_file_size: None, + append_mode: false, + max_background_tasks: None, + time_range: None, + }; + + let output = picker + .build_output_with_time_range(RegionId::from_u64(1), windows, Some(0), None) + .await + .unwrap(); + + assert_eq!(1, output.len()); + assert_eq!(4, output[0].inputs.len()); + } + + #[tokio::test] + async fn test_count_first_trigger_counts_physical_ssts() { + let files = [ + new_file_handle_with_size_and_sequence(FileId::random(), 0, 10, 0, 1, 10), + new_file_handle_with_size_and_sequence(FileId::random(), 0, 10, 0, 1, 10), + new_file_handle_with_size_and_sequence(FileId::random(), 20, 30, 0, 2, 10), + ]; + let windows = assign_to_windows(files.iter(), 100); + let picker = TwcsPicker { + trigger_file_num: 3, + time_window_seconds: Some(100), + max_output_file_size: None, + append_mode: false, + max_background_tasks: None, + time_range: None, + }; + + let output = picker + .build_output_with_time_range(RegionId::from_u64(1), windows, Some(0), None) + .await + .unwrap(); + + assert_eq!(1, output.len()); + assert_eq!(3, output[0].inputs.len()); + } + + #[tokio::test] + async fn test_count_first_does_not_compact_overlap_below_trigger() { + let files = [ + new_file_handle_with_size_and_sequence(FileId::random(), 0, 20, 0, 1, 10), + new_file_handle_with_size_and_sequence(FileId::random(), 10, 30, 0, 2, 10), + ]; + let windows = assign_to_windows(files.iter(), 100); + let picker = TwcsPicker { + trigger_file_num: 3, + time_window_seconds: Some(100), + max_output_file_size: None, + append_mode: false, + max_background_tasks: None, + time_range: None, + }; + + let output = picker + .build_output_with_time_range(RegionId::from_u64(1), windows, Some(0), None) + .await + .unwrap(); + + assert!(output.is_empty()); + } + + #[tokio::test] + async fn test_filter_deleted_is_false_when_selected_files_overlap_unselected_file() { + let mut files = (0..32) + .map(|idx| { + new_file_handle_with_size_and_sequence( + FileId::random(), + idx * 10, + idx * 10 + 9, + 0, + idx as u64 + 1, + 10, + ) + }) + .collect::>(); + files.push(new_file_handle_with_size_and_sequence( + FileId::random(), + 0, + 320, + 0, + 33, + 10, + )); + let windows = assign_to_windows(files.iter(), 1000); + let picker = TwcsPicker { + trigger_file_num: 2, + time_window_seconds: Some(1000), + max_output_file_size: None, + append_mode: false, + max_background_tasks: None, + time_range: None, + }; + + let output = picker + .build_output_with_time_range(RegionId::from_u64(1), windows, Some(0), None) + .await + .unwrap(); + + assert_eq!(1, output.len()); + assert_eq!(DEFAULT_MAX_INPUT_FILES, output[0].inputs.len()); + assert!(!output[0].filter_deleted); + } + + fn new_file(start: i64, end: i64, sequence: u64, file_size: u64) -> FileHandle { + new_file_handle_with_size_and_sequence(FileId::random(), start, end, 0, sequence, file_size) + } + + fn picked_ranges(files: &[FileHandle]) -> Vec<(i64, i64)> { + files + .iter() + .map(|file| { + let (start, end) = file.range(); + (start.value(), end.value()) + }) + .collect() + } + + #[test] + fn test_count_first_rejects_dominant_historical_file() { + let picked = pick_count_first( + vec![SortedRun::from(vec![ + new_file(0, 9, 1, 400), + new_file(20, 29, 2, 100), + ])], + None, + ); + + assert!(picked.is_empty()); + } + + #[test] + fn test_count_first_accepts_balanced_historical_file() { + let picked = pick_count_first( + vec![SortedRun::from(vec![ + new_file(0, 9, 1, 400), + new_file(20, 29, 2, 100), + new_file(40, 49, 3, 100), + new_file(60, 69, 4, 100), + new_file(80, 89, 5, 100), + ])], + None, + ); + + assert_eq!( + vec![(0, 9), (20, 29), (40, 49), (60, 69), (80, 89)], + picked_ranges(&picked) + ); + } + + #[test] + fn test_count_first_finds_smaller_balanced_interval_when_larger_one_is_unbalanced() { + let picked = pick_count_first( + vec![SortedRun::from(vec![ + new_file(0, 9, 1, 1000), + new_file(20, 29, 2, 100), + new_file(40, 49, 3, 100), + new_file(60, 69, 4, 100), + new_file(80, 89, 5, 100), + ])], + None, + ); + + assert_eq!( + vec![(20, 29), (40, 49), (60, 69), (80, 89)], + picked_ranges(&picked) + ); + } + + #[test] + fn test_count_first_prefers_overlap_participants_when_file_counts_match() { + let first_run = (0..DEFAULT_MAX_INPUT_FILES) + .map(|idx| { + let start = idx as i64 * 20; + let end = if idx + 1 == DEFAULT_MAX_INPUT_FILES { + 700 + } else { + start + 9 + }; + new_file(start, end, idx as u64 + 1, 10) + }) + .collect::>(); + let overlapping = new_file(690, 710, 100, 10); + + let picked = pick_count_first( + vec![ + SortedRun::from(first_run), + SortedRun::from(vec![overlapping]), + ], + None, + ); + let ranges = picked_ranges(&picked); + + assert_eq!(DEFAULT_MAX_INPUT_FILES, ranges.len()); + assert!(!ranges.contains(&(0, 9))); + assert!(ranges.contains(&(690, 710))); + } + + #[test] + fn test_count_first_prefers_smaller_bytes_when_file_counts_match() { + let files = (0..=DEFAULT_MAX_INPUT_FILES) + .map(|idx| { + let start = idx as i64 * 20; + let size = if idx == 0 { + 100 + } else if idx == DEFAULT_MAX_INPUT_FILES { + 1 + } else { + 10 + }; + new_file(start, start + 9, idx as u64 + 1, size) + }) + .collect::>(); + + let picked = pick_count_first(vec![SortedRun::from(files)], None); + let ranges = picked_ranges(&picked); + + assert_eq!(DEFAULT_MAX_INPUT_FILES, ranges.len()); + assert_eq!(Some(&(20, 29)), ranges.first()); + assert_eq!(Some(&(640, 649)), ranges.last()); + } + + #[test] + fn test_count_first_skips_pure_rewrite_without_progress() { + // Four non-overlapping files whose combined size splits back into at least + // four outputs: compacting them reduces neither the file count nor any + // overlap, so there is nothing worth picking. + let picked = pick_count_first( + vec![SortedRun::from(vec![ + new_file(0, 9, 1, 600), + new_file(20, 29, 2, 600), + new_file(40, 49, 3, 600), + new_file(60, 69, 4, 600), + ])], + Some(512), + ); + + assert!(picked.is_empty()); + } + + #[test] + fn test_count_first_allows_overlap_resolution_without_file_reduction() { + // Two large overlapping files from different runs: compacting them cannot + // reduce the file count (the output splits back into just as many files), + // but it resolves the overlap and merges two sorted runs into one. + let picked = pick_count_first( + vec![ + SortedRun::from(vec![new_file(0, 19, 1, 600)]), + SortedRun::from(vec![new_file(10, 29, 2, 600)]), + ], + Some(512), + ); + + assert_eq!(vec![(0, 19), (10, 29)], picked_ranges(&picked)); + } + + #[test] + fn test_count_first_prefers_guaranteed_reduction_over_pure_rewrite() { + // A window mixing large and small files: the large triple on its own is a + // pure rewrite, and every interval containing large files rewrites more + // bytes for the same predicted reduction, so the small triple wins. + let picked = pick_count_first( + vec![SortedRun::from(vec![ + new_file(0, 9, 1, 600), + new_file(10, 19, 2, 600), + new_file(20, 29, 3, 600), + new_file(30, 39, 4, 10), + new_file(40, 49, 5, 10), + new_file(50, 59, 6, 10), + ])], + Some(512), + ); + + assert_eq!(vec![(30, 39), (40, 49), (50, 59)], picked_ranges(&picked)); + } + + #[test] + fn test_count_first_prefers_earlier_time_on_exact_tie() { + let files = (0..=DEFAULT_MAX_INPUT_FILES) + .map(|idx| { + let start = idx as i64 * 20; + new_file(start, start + 9, idx as u64 + 1, 10) + }) + .collect::>(); + + let picked = pick_count_first(vec![SortedRun::from(files)], None); + let ranges = picked_ranges(&picked); + + assert_eq!(DEFAULT_MAX_INPUT_FILES, ranges.len()); + assert_eq!(Some(&(0, 9)), ranges.first()); + assert_eq!(Some(&(620, 629)), ranges.last()); + } + + #[test] + fn test_count_first_keeps_each_interleaved_run_contiguous() { + let picked = pick_count_first( + vec![ + SortedRun::from(vec![ + new_file(0, 9, 1, 1), + new_file(20, 29, 2, 1), + new_file(40, 49, 3, 1), + ]), + SortedRun::from(vec![new_file(10, 19, 4, 1), new_file(30, 39, 5, 1)]), + ], + None, + ); + + assert_eq!( + vec![(0, 9), (10, 19), (20, 29), (30, 39), (40, 49)], + picked_ranges(&picked) + ); + } + // TODO(hl): TTL tester that checks if get_expired_ssts function works as expected. + + // Reproduces the leftover-small-files deadlock observed in the compaction bench: + // an "EngineFull" flush emits several SSTs sharing one sequence. They must be ordinary + // candidates so a contiguous candidate can span them instead of treating them as an + // indivisible block that fragments the timeline. + #[test] + fn test_count_first_candidate_spans_same_sequence_files() { + let picked = pick_count_first( + vec![SortedRun::from(vec![ + new_file(0, 9, 1, 10), + // One EngineFull flush output: three SSTs with the same sequence. + new_file(10, 19, 2, 10), + new_file(10, 19, 2, 10), + new_file(10, 19, 2, 10), + new_file(20, 29, 3, 10), + ])], + None, + ); + + assert_eq!( + vec![(0, 9), (10, 19), (10, 19), (10, 19), (20, 29)], + picked_ranges(&picked) + ); + } + + // End-to-end repro of the compaction-bench leftover files: a window where + // EngineFull flush groups (multiple SSTs sharing a sequence) sit between + // singleton flushes. The picker must be able to merge the whole contiguous + // range; otherwise these files can never be compacted once the active window + // moves on. + #[tokio::test] + async fn test_count_first_merges_window_with_interleaved_flush_groups() { + let files = [ + // Singleton flush. + new_file_handle_with_size_and_sequence(FileId::random(), 0, 9, 0, 1, 10), + // EngineFull flush: three SSTs sharing sequence 2. + new_file_handle_with_size_and_sequence(FileId::random(), 10, 19, 0, 2, 10), + new_file_handle_with_size_and_sequence(FileId::random(), 10, 19, 0, 2, 10), + new_file_handle_with_size_and_sequence(FileId::random(), 10, 19, 0, 2, 10), + // Singleton flush. + new_file_handle_with_size_and_sequence(FileId::random(), 20, 29, 0, 3, 10), + ]; + let windows = assign_to_windows(files.iter(), 100); + let picker = TwcsPicker { + trigger_file_num: 3, + time_window_seconds: Some(100), + max_output_file_size: None, + append_mode: false, + max_background_tasks: None, + time_range: None, + }; + + let output = picker + .build_output_with_time_range(RegionId::from_u64(1), windows, Some(0), None) + .await + .unwrap(); + + assert_eq!(1, output.len()); + // All five SSTs of the window are merged in one compaction. + assert_eq!(5, output[0].inputs.len()); + } } diff --git a/src/mito2/src/config.rs b/src/mito2/src/config.rs index 54445549a05..be74fa67c60 100644 --- a/src/mito2/src/config.rs +++ b/src/mito2/src/config.rs @@ -102,7 +102,7 @@ pub struct MitoConfig { pub experimental_compaction_on_exhausted: OnExhaustedPolicy, // Flush configs: - /// Interval to auto flush a region if it has not flushed yet (default 30 min). + /// Interval to auto flush a region if it has not flushed yet (default 10 min). #[serde(with = "humantime_serde")] pub auto_flush_interval: Duration, /// Global write buffer size threshold to trigger flush. @@ -210,7 +210,7 @@ impl Default for MitoConfig { max_background_purges: get_total_cpu_cores(), experimental_compaction_memory_limit: MemoryLimit::Unlimited, experimental_compaction_on_exhausted: OnExhaustedPolicy::default(), - auto_flush_interval: Duration::from_secs(30 * 60), + auto_flush_interval: Duration::from_secs(10 * 60), global_write_buffer_size: ReadableSize::gb(1), global_write_buffer_reject_size: ReadableSize::gb(2), default_region_write_buffer_size: ReadableSize::mb(0), @@ -381,6 +381,14 @@ impl MitoConfig { mod tests { use super::*; + #[test] + fn test_default_auto_flush_interval() { + assert_eq!( + Duration::from_secs(10 * 60), + MitoConfig::default().auto_flush_interval + ); + } + #[test] fn test_adjust_sst_metadata_and_prefilter_cache_caps_independently() { let mut config = MitoConfig::default(); diff --git a/src/mito2/src/engine/apply_staging_manifest_test.rs b/src/mito2/src/engine/apply_staging_manifest_test.rs index e10ed0cbf00..49cc97c7d51 100644 --- a/src/mito2/src/engine/apply_staging_manifest_test.rs +++ b/src/mito2/src/engine/apply_staging_manifest_test.rs @@ -12,20 +12,24 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::collections::HashSet; use std::sync::Arc; use std::{assert_matches, fs}; use api::v1::Rows; +use api::v1::region::{StrictWindow, compact_request}; use common_function::utils::partition_expr_version; use common_recordbatch::RecordBatches; +use datatypes::arrow::array::AsArray; +use datatypes::arrow::datatypes::Float64Type; use datatypes::value::Value; use partition::expr::{PartitionExpr, col}; use store_api::region_engine::{ RegionEngine, RegionRole, RemapManifestsRequest, SettableRegionRoleState, }; use store_api::region_request::{ - ApplyStagingManifestRequest, EnterStagingRequest, RegionFlushRequest, RegionPutRequest, - RegionRequest, StagingPartitionDirective, + ApplyStagingManifestRequest, EnterStagingRequest, RegionCompactRequest, RegionFlushRequest, + RegionPutRequest, RegionRequest, StagingPartitionDirective, }; use store_api::storage::{FileId, RegionId}; @@ -37,7 +41,9 @@ use crate::manifest::action::{ }; use crate::sst::FormatType; use crate::sst::file::FileMeta; -use crate::test_util::{CreateRequestBuilder, TestEnv, build_rows, put_rows, rows_schema}; +use crate::test_util::{ + CreateRequestBuilder, TestEnv, build_rows, build_rows_for_key, put_rows, rows_schema, +}; fn range_expr(col_name: &str, start: i64, end: i64) -> PartitionExpr { col(col_name) @@ -45,6 +51,209 @@ fn range_expr(col_name: &str, start: i64, end: i64) -> PartitionExpr { .and(col(col_name).lt(Value::Int64(end))) } +#[tokio::test] +async fn test_apply_staging_manifest_sequence_domain() { + common_telemetry::init_default_ut_logging(); + test_apply_staging_manifest_sequence_domain_with_format(false).await; + test_apply_staging_manifest_sequence_domain_with_format(true).await; +} + +async fn test_apply_staging_manifest_sequence_domain_with_format(flat_format: bool) { + let mut env = TestEnv::with_prefix("apply-staging-sequence-domain").await; + let engine = env + .create_engine(MitoConfig { + default_flat_format: flat_format, + ..Default::default() + }) + .await; + let source = RegionId::new(1, 1); + let target = RegionId::new(1, 2); + let request = CreateRequestBuilder::new().build(); + let schema = rows_schema(&request); + + engine + .handle_request(source, RegionRequest::Create(request.clone())) + .await + .unwrap(); + for value in 0..3 { + put_rows( + &engine, + source, + Rows { + schema: schema.clone(), + rows: build_rows_for_key("0", 0, 1, value), + }, + ) + .await; + } + engine + .handle_request(source, RegionRequest::Flush(RegionFlushRequest::default())) + .await + .unwrap(); + let source_manifest = engine + .get_region(source) + .unwrap() + .manifest_ctx + .manifest() + .await; + assert_eq!(source_manifest.files.len(), 1); + assert_eq!( + source_manifest.files.values().next().unwrap().sequence, + Some(std::num::NonZeroU64::new(3).unwrap()) + ); + + engine + .set_region_role_state_gracefully(source, SettableRegionRoleState::StagingLeader) + .await + .unwrap(); + let partition_expr = float_range_expr("field_0", 0.1, 100.1) + .as_json_str() + .unwrap(); + let result = engine + .remap_manifests(RemapManifestsRequest { + region_id: source, + input_regions: vec![source], + region_mapping: [(source, vec![target])].into_iter().collect(), + new_partition_exprs: [(target, partition_expr.clone())].into_iter().collect(), + }) + .await + .unwrap(); + engine + .handle_request(target, RegionRequest::Create(request.clone())) + .await + .unwrap(); + engine + .handle_request( + target, + RegionRequest::EnterStaging(EnterStagingRequest { + partition_directive: StagingPartitionDirective::UpdatePartitionExpr( + partition_expr.clone(), + ), + }), + ) + .await + .unwrap(); + engine + .handle_request( + target, + RegionRequest::ApplyStagingManifest(ApplyStagingManifestRequest { + partition_expr, + central_region_id: source, + manifest_path: result.manifest_paths[&target].clone(), + }), + ) + .await + .unwrap(); + + let manifest = engine + .get_region(target) + .unwrap() + .manifest_ctx + .manifest() + .await; + assert_eq!(manifest.files.len(), 1); + assert_eq!(manifest.committed_sequence, Some(1)); + let imported_file = manifest.files.values().next().unwrap(); + assert_eq!(imported_file.region_id, source); + assert_eq!( + imported_file.sequence, + Some(std::num::NonZeroU64::new(1).unwrap()) + ); + + put_rows( + &engine, + target, + Rows { + schema: schema.clone(), + rows: build_rows_for_key("0", 0, 1, 99), + }, + ) + .await; + assert_target_value(&engine, target, 99.0).await; + + engine + .handle_request(target, RegionRequest::Flush(RegionFlushRequest::default())) + .await + .unwrap(); + let input_ids = current_file_ids(&engine, target); + assert_eq!( + input_ids.len(), + 2, + "imported and target-write SSTs must both exist" + ); + + engine + .handle_request( + target, + RegionRequest::Compact(RegionCompactRequest { + options: compact_request::Options::StrictWindow(StrictWindow { + window_seconds: 60, + }), + parallelism: None, + time_range: None, + }), + ) + .await + .unwrap(); + let output_version = engine.get_region(target).unwrap().version(); + let output_files = output_version + .ssts + .levels() + .iter() + .flat_map(|level| level.files.values()) + .collect::>(); + let output_ids = output_files + .iter() + .map(|file| file.meta_ref().file_id) + .collect::>(); + assert_eq!(output_ids.len(), 1); + assert!( + output_files + .iter() + .all(|file| file.meta_ref().region_id == target) + ); + assert!( + input_ids.iter().all(|id| !output_ids.contains(id)), + "real compaction must replace both input SSTs" + ); + assert_target_value(&engine, target, 99.0).await; +} + +fn current_file_ids(engine: &crate::engine::MitoEngine, region_id: RegionId) -> HashSet { + engine + .get_region(region_id) + .unwrap() + .version() + .ssts + .levels() + .iter() + .flat_map(|level| level.files.values()) + .map(|file| file.meta_ref().file_id) + .collect() +} + +async fn assert_target_value( + engine: &crate::engine::MitoEngine, + region_id: RegionId, + expected: f64, +) { + let scan = engine + .scan_to_stream(region_id, ScanRequest::default()) + .await + .unwrap(); + let batches = RecordBatches::try_collect(scan).await.unwrap(); + assert_eq!( + batches.iter().map(|batch| batch.num_rows()).sum::(), + 1 + ); + let batch = batches.iter().next().unwrap(); + let values = batch + .column_by_name("field_0") + .unwrap() + .as_primitive::(); + assert_eq!(values.value(0), expected); +} + fn float_range_expr(col_name: &str, start: f64, end: f64) -> PartitionExpr { col(col_name) .gt_eq(Value::Float64(start.into())) diff --git a/src/mito2/src/engine/basic_test.rs b/src/mito2/src/engine/basic_test.rs index da36c827665..6d90f90b784 100644 --- a/src/mito2/src/engine/basic_test.rs +++ b/src/mito2/src/engine/basic_test.rs @@ -882,9 +882,9 @@ async fn test_cache_null_primary_key_with_format(flat_format: bool) { #[tokio::test] async fn test_list_ssts() { test_list_ssts_with_format(false, r#" -ManifestSstEntry { table_dir: "test/", region_id: 47244640257(11, 1), table_id: 11, region_number: 1, region_group: 0, region_sequence: 1, file_id: "", index_version: 0, level: 0, file_path: "test/11_0000000001/.parquet", file_size: 2951, index_file_path: Some("test/11_0000000001/index/.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), origin_region_id: 47244640257(11, 1), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") } -ManifestSstEntry { table_dir: "test/", region_id: 47244640258(11, 2), table_id: 11, region_number: 2, region_group: 0, region_sequence: 2, file_id: "", index_version: 0, level: 0, file_path: "test/11_0000000002/.parquet", file_size: 2951, index_file_path: Some("test/11_0000000002/index/.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), origin_region_id: 47244640258(11, 2), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") } -ManifestSstEntry { table_dir: "test/", region_id: 94489280554(22, 42), table_id: 22, region_number: 42, region_group: 0, region_sequence: 42, file_id: "", index_version: 0, level: 0, file_path: "test/22_0000000042/.parquet", file_size: 2951, index_file_path: Some("test/22_0000000042/index/.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), origin_region_id: 94489280554(22, 42), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") }"# , +ManifestSstEntry { table_dir: "test/", region_id: 47244640257(11, 1), table_id: 11, region_number: 1, region_group: 0, region_sequence: 1, file_id: "", index_version: 0, level: 0, file_path: "test/11_0000000001/.parquet", file_size: 2951, max_row_group_uncompressed_size: 0, index_file_path: Some("test/11_0000000001/index/.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), partition_expr: None, origin_region_id: 47244640257(11, 1), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") } +ManifestSstEntry { table_dir: "test/", region_id: 47244640258(11, 2), table_id: 11, region_number: 2, region_group: 0, region_sequence: 2, file_id: "", index_version: 0, level: 0, file_path: "test/11_0000000002/.parquet", file_size: 2951, max_row_group_uncompressed_size: 0, index_file_path: Some("test/11_0000000002/index/.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), partition_expr: None, origin_region_id: 47244640258(11, 2), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") } +ManifestSstEntry { table_dir: "test/", region_id: 94489280554(22, 42), table_id: 22, region_number: 42, region_group: 0, region_sequence: 42, file_id: "", index_version: 0, level: 0, file_path: "test/22_0000000042/.parquet", file_size: 2951, max_row_group_uncompressed_size: 0, index_file_path: Some("test/22_0000000042/index/.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), partition_expr: None, origin_region_id: 94489280554(22, 42), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") }"# , r#" StorageSstEntry { file_path: "test/11_0000000001/.parquet", file_size: None, last_modified_ms: None, node_id: None } StorageSstEntry { file_path: "test/11_0000000001/index/.puffin", file_size: None, last_modified_ms: None, node_id: None } @@ -893,9 +893,9 @@ StorageSstEntry { file_path: "test/11_0000000002/index/.puffin", file_s StorageSstEntry { file_path: "test/22_0000000042/.parquet", file_size: None, last_modified_ms: None, node_id: None } StorageSstEntry { file_path: "test/22_0000000042/index/.puffin", file_size: None, last_modified_ms: None, node_id: None }"#).await; test_list_ssts_with_format(true, r#" -ManifestSstEntry { table_dir: "test/", region_id: 47244640257(11, 1), table_id: 11, region_number: 1, region_group: 0, region_sequence: 1, file_id: "", index_version: 0, level: 0, file_path: "test/11_0000000001/.parquet", file_size: 3365, index_file_path: Some("test/11_0000000001/index/.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), origin_region_id: 47244640257(11, 1), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") } -ManifestSstEntry { table_dir: "test/", region_id: 47244640258(11, 2), table_id: 11, region_number: 2, region_group: 0, region_sequence: 2, file_id: "", index_version: 0, level: 0, file_path: "test/11_0000000002/.parquet", file_size: 3365, index_file_path: Some("test/11_0000000002/index/.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), origin_region_id: 47244640258(11, 2), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") } -ManifestSstEntry { table_dir: "test/", region_id: 94489280554(22, 42), table_id: 22, region_number: 42, region_group: 0, region_sequence: 42, file_id: "", index_version: 0, level: 0, file_path: "test/22_0000000042/.parquet", file_size: 3365, index_file_path: Some("test/22_0000000042/index/.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), origin_region_id: 94489280554(22, 42), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") }"#, +ManifestSstEntry { table_dir: "test/", region_id: 47244640257(11, 1), table_id: 11, region_number: 1, region_group: 0, region_sequence: 1, file_id: "", index_version: 0, level: 0, file_path: "test/11_0000000001/.parquet", file_size: 3365, max_row_group_uncompressed_size: 0, index_file_path: Some("test/11_0000000001/index/.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), partition_expr: None, origin_region_id: 47244640257(11, 1), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") } +ManifestSstEntry { table_dir: "test/", region_id: 47244640258(11, 2), table_id: 11, region_number: 2, region_group: 0, region_sequence: 2, file_id: "", index_version: 0, level: 0, file_path: "test/11_0000000002/.parquet", file_size: 3365, max_row_group_uncompressed_size: 0, index_file_path: Some("test/11_0000000002/index/.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), partition_expr: None, origin_region_id: 47244640258(11, 2), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") } +ManifestSstEntry { table_dir: "test/", region_id: 94489280554(22, 42), table_id: 22, region_number: 42, region_group: 0, region_sequence: 42, file_id: "", index_version: 0, level: 0, file_path: "test/22_0000000042/.parquet", file_size: 3365, max_row_group_uncompressed_size: 0, index_file_path: Some("test/22_0000000042/index/.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), partition_expr: None, origin_region_id: 94489280554(22, 42), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") }"#, r#" StorageSstEntry { file_path: "test/11_0000000001/.parquet", file_size: None, last_modified_ms: None, node_id: None } StorageSstEntry { file_path: "test/11_0000000001/index/.puffin", file_size: None, last_modified_ms: None, node_id: None } @@ -967,6 +967,7 @@ async fn test_list_ssts_with_format( .map(|p| p.replace(&e.file_id, "")); e.file_id = "".to_string(); e.index_version = 0; + e.max_row_group_uncompressed_size = 0; format!("\n{:?}", e) }) .sorted() diff --git a/src/mito2/src/engine/compaction_test.rs b/src/mito2/src/engine/compaction_test.rs index 1fc3a8f785d..04c97bdfbba 100644 --- a/src/mito2/src/engine/compaction_test.rs +++ b/src/mito2/src/engine/compaction_test.rs @@ -768,6 +768,7 @@ async fn test_truncate_waits_for_non_cancellable_compaction_commit() { .await; let create = CreateRequestBuilder::new() .insert_option("compaction.type", "twcs") + .insert_option("compaction.twcs.trigger_file_num", "4") .build(); let column_schemas = create .column_metadatas @@ -779,7 +780,9 @@ async fn test_truncate_waits_for_non_cancellable_compaction_commit() { .await .unwrap(); put_and_flush(&engine, region_id, &column_schemas, 0..10).await; - put_and_flush(&engine, region_id, &column_schemas, 5..20).await; + put_and_flush(&engine, region_id, &column_schemas, 10..20).await; + put_and_flush(&engine, region_id, &column_schemas, 20..30).await; + put_and_flush(&engine, region_id, &column_schemas, 30..40).await; let commit_guard = gate.arm_commit(); let compact_engine = engine.clone(); @@ -803,7 +806,7 @@ async fn test_truncate_waits_for_non_cancellable_compaction_commit() { RegionRequest::Truncate(RegionTruncateRequest::ByTimeRanges { time_ranges: vec![( Timestamp::new_millisecond(0), - Timestamp::new_millisecond(19_000), + Timestamp::new_millisecond(39_000), )], }), ) @@ -981,12 +984,10 @@ async fn test_compaction_region_with_format(flat_format: bool) { // [20....29] // -[15.........29]- (delete) // [15.....24] - // Output: - // [0..9] - // [10............29] (contains delete) - // [15....24] + // Count-first compaction consumes the first 4 SSTs as soon as the trigger is reached. + // The compacted output and final flush leave 2 SSTs. assert_eq!( - 3, + 2, scanner.num_files(), "unexpected files: {:?}", scanner.file_ids() @@ -1530,6 +1531,7 @@ async fn test_readonly_during_compaction_with_format(flat_format: bool) { let request = CreateRequestBuilder::new() .insert_option("compaction.type", "twcs") + .insert_option("compaction.twcs.trigger_file_num", "4") .build(); let column_schemas = request @@ -1542,7 +1544,9 @@ async fn test_readonly_during_compaction_with_format(flat_format: bool) { .await .unwrap(); let listener_guard = CompactionListenerGuard::new(listener.clone()); - // Flush 2 SSTs for compaction. + // Flush 4 balanced SSTs for compaction. + put_and_flush(&engine, region_id, &column_schemas, 0..10).await; + put_and_flush(&engine, region_id, &column_schemas, 5..20).await; put_and_flush(&engine, region_id, &column_schemas, 0..10).await; put_and_flush(&engine, region_id, &column_schemas, 5..20).await; @@ -1573,7 +1577,7 @@ async fn test_readonly_during_compaction_with_format(flat_format: bool) { .await .unwrap(); assert_eq!( - 2, + 4, scanner.num_files(), "unexpected files: {:?}", scanner.file_ids() @@ -1614,6 +1618,7 @@ async fn test_local_compaction_cancellation_notifies_before_pending_ddl_dispatch let request = CreateRequestBuilder::new() .insert_option("compaction.type", "twcs") + .insert_option("compaction.twcs.trigger_file_num", "4") .build(); let column_schemas = request .column_metadatas @@ -1628,6 +1633,8 @@ async fn test_local_compaction_cancellation_notifies_before_pending_ddl_dispatch let merge_guard = gate.arm_merge(); put_and_flush(&engine, region_id, &column_schemas, 0..10).await; put_and_flush(&engine, region_id, &column_schemas, 5..20).await; + put_and_flush(&engine, region_id, &column_schemas, 0..10).await; + put_and_flush(&engine, region_id, &column_schemas, 5..20).await; tokio::time::timeout(Duration::from_secs(5), gate.wait_until_merge_entered()) .await diff --git a/src/mito2/src/engine/merge_mode_test.rs b/src/mito2/src/engine/merge_mode_test.rs index 40a87642ae4..afb984e77bd 100644 --- a/src/mito2/src/engine/merge_mode_test.rs +++ b/src/mito2/src/engine/merge_mode_test.rs @@ -127,6 +127,7 @@ async fn test_merge_mode_compaction_with_format(flat_format: bool) { let request = CreateRequestBuilder::new() .field_num(2) .insert_option("compaction.type", "twcs") + .insert_option("compaction.twcs.trigger_file_num", "3") .insert_option("merge_mode", "last_non_null") .build(); let table_dir = request.table_dir.clone(); @@ -208,7 +209,7 @@ async fn test_merge_mode_compaction_with_format(flat_format: bool) { .scanner(region_id, ScanRequest::default()) .await .unwrap(); - assert_eq!(2, scanner.num_files()); + assert_eq!(1, scanner.num_files()); assert_eq!(1, scanner.num_memtables()); scanner.set_target_partitions(2); let stream = scanner.scan().await.unwrap(); diff --git a/src/mito2/src/gc/worker_test.rs b/src/mito2/src/gc/worker_test.rs index cd8b875e9e8..e0b03484996 100644 --- a/src/mito2/src/gc/worker_test.rs +++ b/src/mito2/src/gc/worker_test.rs @@ -399,7 +399,8 @@ async fn test_gc_worker_basic_compact() { let region = engine.get_region(region_id).unwrap(); let manifest = region.manifest_ctx.manifest().await; - assert_eq!(manifest.removed_files.removed_files[0].files.len(), 3); + let removed_file_num = manifest.removed_files.removed_files[0].files.len(); + assert!(removed_file_num > 0); let version = manifest.manifest_version; @@ -413,7 +414,10 @@ async fn test_gc_worker_basic_compact() { let gc_worker = create_gc_worker(&engine, regions, &file_ref_manifest, true).await; let report = gc_worker.run().await.unwrap(); - assert_eq!(report.deleted_files.get(®ion_id).unwrap().len(), 3,); + assert_eq!( + report.deleted_files.get(®ion_id).unwrap().len(), + removed_file_num + ); assert!(report.need_retry_regions.is_empty()); } @@ -476,7 +480,7 @@ async fn test_gc_worker_compact_with_ref() { let region = engine.get_region(region_id).unwrap(); let manifest = region.manifest_ctx.manifest().await; - assert_eq!(manifest.removed_files.removed_files[0].files.len(), 3); + assert!(!manifest.removed_files.removed_files[0].files.is_empty()); let version = manifest.manifest_version; diff --git a/src/mito2/src/region.rs b/src/mito2/src/region.rs index 76204da9cf6..b68b0d1cccd 100644 --- a/src/mito2/src/region.rs +++ b/src/mito2/src/region.rs @@ -839,6 +839,7 @@ impl MitoRegion { level: meta.level, file_path: sst_file_path(table_dir, meta.file_id(), path_type), file_size: meta.file_size, + max_row_group_uncompressed_size: meta.max_row_group_uncompressed_size, index_file_path, index_file_size, num_rows: meta.num_rows, @@ -847,6 +848,7 @@ impl MitoRegion { min_ts: meta.time_range.0, max_ts: meta.time_range.1, sequence: meta.sequence.map(|s| s.get()), + partition_expr: meta.partition_expr.as_ref().map(ToString::to_string), origin_region_id, node_id: None, visible, diff --git a/src/mito2/src/sst/index.rs b/src/mito2/src/sst/index.rs index 581d1edebdf..7954af2b920 100644 --- a/src/mito2/src/sst/index.rs +++ b/src/mito2/src/sst/index.rs @@ -706,6 +706,14 @@ pub struct IndexBuildTask { pub region_id: RegionId, /// The SST file handle to build index for. pub file: FileHandle, + /// The target region metadata used to decode rows from the SST. + /// + /// An SST may originate in another region while being visible in the target + /// manifest. This metadata defines the target schema and sequence domain; + /// applying the staging manifest only makes imported files visible. Index + /// rebuild happens later when a flush, compaction, schema change, or manual + /// index build request schedules it. + pub(crate) target_region_metadata: RegionMetadataRef, /// The manifest state this build is based on. pub(crate) source: IndexBuildSource, pub reason: IndexBuildType, @@ -846,6 +854,7 @@ impl IndexBuildTask { let mut parquet_reader = self .access_layer .read_sst(self.file.clone()) // use the latest file handle instead of creating a new one + .expected_metadata(Some(self.target_region_metadata.clone())) .build() .await?; @@ -1494,8 +1503,12 @@ mod tests { use datatypes::schema::{ ColumnSchema, FulltextOptions, SkippingIndexOptions, SkippingIndexType, }; + use datatypes::value::Value; + use index::inverted_index::format::reader::InvertedIndexReader; use object_store::ObjectStore; use object_store::services::Memory; + use partition::expr::col; + use puffin::puffin_manager::{PuffinManager, PuffinReader}; use puffin_manager::PuffinManagerFactory; use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder}; use tokio::sync::mpsc; @@ -2088,6 +2101,7 @@ mod tests { let task = IndexBuildTask { region_id, file, + target_region_metadata: version_control.current().version.metadata.clone(), source: IndexBuildSource::new( file_meta, version_control.current().version.metadata.schema_version, @@ -2120,16 +2134,34 @@ mod tests { } #[tokio::test] - async fn test_index_build_task_increments_legacy_index_version() { + async fn test_index_build_task_foreign_file_uses_target_metadata() { let env = SchedulerEnv::new().await; let mut scheduler = env.mock_index_build_scheduler(4); - let metadata = Arc::new(sst_region_metadata()); - let manifest_ctx = env.mock_manifest_context(metadata.clone()).await; - let region_id = metadata.region_id; + let source_metadata = Arc::new(sst_region_metadata()); + let mut target_metadata = (*source_metadata).clone(); + target_metadata.region_id = RegionId::new(1, 3); + let mut target_builder = RegionMetadataBuilder::new(target_metadata.region_id); + for mut column_metadata in target_metadata.column_metadatas.clone() { + if column_metadata.column_id == 2 { + column_metadata.column_schema = + column_metadata.column_schema.with_inverted_index(true); + } + target_builder.push_column_metadata(column_metadata); + } + let partition_expr = col("field_0") + .gt_eq(Value::UInt64(100)) + .and(col("field_0").lt(Value::UInt64(200))); + target_builder + .primary_key(target_metadata.primary_key.clone()) + .partition_expr_json(Some(partition_expr.as_json_str().unwrap())) + .bump_version(); + let target_metadata = Arc::new(target_builder.build().unwrap()); + let manifest_ctx = env.mock_manifest_context(target_metadata.clone()).await; + let region_id = target_metadata.region_id; let file_purger = Arc::new(NoopFilePurger {}); - let sst_info = mock_sst_file(metadata.clone(), &env, IndexBuildMode::Async).await; + let sst_info = mock_sst_file(source_metadata.clone(), &env, IndexBuildMode::Async).await; let file_meta = FileMeta { - region_id, + region_id: source_metadata.region_id, file_id: sst_info.file_id, file_size: sst_info.file_size, max_row_group_uncompressed_size: sst_info.max_row_group_uncompressed_size, @@ -2144,8 +2176,8 @@ mod tests { seed_manifest_file(&manifest_ctx, &file_meta).await; let files = HashMap::from([(file_meta.file_id, file_meta.clone())]); let version_control = - mock_version_control(metadata.clone(), file_purger.clone(), files).await; - let indexer_builder = mock_indexer_builder(metadata.clone(), &env).await; + mock_version_control(target_metadata.clone(), file_purger.clone(), files).await; + let indexer_builder = mock_indexer_builder(target_metadata.clone(), &env).await; let file = FileHandle::new(file_meta.clone(), file_purger.clone()); @@ -2155,6 +2187,7 @@ mod tests { let task = IndexBuildTask { region_id, file, + target_region_metadata: version_control.current().version.metadata.clone(), source: IndexBuildSource::new( file_meta.clone(), version_control.current().version.metadata.schema_version, @@ -2199,9 +2232,40 @@ mod tests { assert!(updated_meta.index_file_size > 0); assert_eq!(updated_meta.file_id, file_meta.file_id); assert_eq!(updated_meta.index_version, 1); + let field_0_index = updated_meta + .indexes + .iter() + .find(|index| index.column_id == 2) + .expect("field_0 should have an inverted index"); + assert_eq!( + field_0_index.created_indexes.as_slice(), + [IndexType::InvertedIndex] + ); } _ => panic!("Unexpected worker request: {:?}", worker_req), } + + let puffin_reader = env + .access_layer + .build_puffin_manager() + .reader(&RegionIndexId::new( + RegionFileId::new(source_metadata.region_id, file_meta.file_id), + 1, + )) + .await + .unwrap(); + let blob = puffin_reader + .blob(inverted_index::INDEX_BLOB_TYPE) + .await + .unwrap(); + let blob_reader = blob.reader().await.unwrap(); + let index_metadata = + index::inverted_index::format::reader::InvertedIndexBlobReader::new(blob_reader) + .metadata(None) + .await + .unwrap(); + assert!(index_metadata.metas.contains_key("2")); + assert_eq!(index_metadata.total_row_count, 100); } async fn schedule_index_build_task_with_mode(build_mode: IndexBuildMode) { @@ -2236,6 +2300,7 @@ mod tests { let task = IndexBuildTask { region_id, file, + target_region_metadata: version_control.current().version.metadata.clone(), source: IndexBuildSource::new( file_meta.clone(), version_control.current().version.metadata.schema_version, @@ -2346,6 +2411,7 @@ mod tests { let task = IndexBuildTask { region_id, file, + target_region_metadata: version_control.current().version.metadata.clone(), source: IndexBuildSource::new( file_meta.clone(), version_control.current().version.metadata.schema_version, @@ -2444,6 +2510,7 @@ mod tests { let task = IndexBuildTask { region_id, file, + target_region_metadata: version_control.current().version.metadata.clone(), source: IndexBuildSource::new( file_meta.clone(), version_control.current().version.metadata.schema_version, @@ -2505,7 +2572,7 @@ mod tests { let schema_version = metadata.schema_version; let manifest_ctx = env.mock_manifest_context(metadata.clone()).await; let file_purger = Arc::new(NoopFilePurger {}); - let indexer_builder = mock_indexer_builder(metadata, env).await; + let indexer_builder = mock_indexer_builder(metadata.clone(), env).await; let (tx, _rx) = mpsc::channel(4); let (result_tx, result_rx) = mpsc::channel::>(4); @@ -2521,6 +2588,7 @@ mod tests { let task = IndexBuildTask { region_id, file, + target_region_metadata: metadata, source: IndexBuildSource::new(file_meta, schema_version), reason, access_layer: env.access_layer.clone(), diff --git a/src/mito2/src/sst/parquet.rs b/src/mito2/src/sst/parquet.rs index e67dcf2bee5..8a05c099e98 100644 --- a/src/mito2/src/sst/parquet.rs +++ b/src/mito2/src/sst/parquet.rs @@ -118,8 +118,8 @@ mod tests { use datafusion_expr::{BinaryExpr, Expr, Literal, Operator, col, lit}; use datatypes::arrow; use datatypes::arrow::array::{ - ArrayRef, BinaryDictionaryBuilder, RecordBatch, StringArray, StringDictionaryBuilder, - TimestampMillisecondArray, UInt8Array, UInt64Array, + ArrayRef, AsArray, BinaryDictionaryBuilder, RecordBatch, StringArray, + StringDictionaryBuilder, TimestampMillisecondArray, UInt8Array, UInt64Array, }; use datatypes::arrow::datatypes::{DataType, Field, Schema, UInt32Type}; use datatypes::arrow::util::pretty::pretty_format_batches; @@ -1625,91 +1625,178 @@ mod tests { #[tokio::test] async fn test_read_with_override_sequence() { + test_read_with_override_sequence_with_format(false).await; + test_read_with_override_sequence_with_format(true).await; + } + + async fn test_read_with_override_sequence_with_format(flat_format: bool) { let mut env = TestEnv::new().await; let object_store = env.init_object_store_manager(); - let handle = sst_file_handle(0, 1000); - let file_path = FixedPathProvider { - region_file_id: handle.file_id(), - }; let metadata = Arc::new(sst_region_metadata()); - // Create batches with sequence 0 to trigger override functionality. - let source = new_flat_source_from_record_batches(vec![ - new_record_batch_with_custom_sequence(&["a", "d"], 0, 60, 0), - new_record_batch_with_custom_sequence(&["b", "f"], 0, 40, 0), - ]); + async fn read_sequences(builder: ParquetReaderBuilder) -> Vec { + let mut reader = builder.build().await.unwrap().unwrap(); + let mut sequences = Vec::new(); + while let Some(batch) = reader.next_record_batch().await.unwrap() { + let sequence = batch + .column(batch.num_columns() - 2) + .as_primitive::(); + sequences.extend((0..sequence.len()).map(|idx| sequence.value(idx))); + } + sequences + } - let write_opts = WriteOptions { - row_group_size: 50, - ..Default::default() - }; + async fn write_sst( + object_store: ObjectStore, + metadata: Arc, + handle: FileHandle, + flat_format: bool, + sequence: u64, + ) { + let file_path = FixedPathProvider { + region_file_id: handle.file_id(), + }; + let source = new_flat_source_from_record_batches(vec![ + new_record_batch_with_custom_sequence(&["a", "d"], 0, 60, sequence), + new_record_batch_with_custom_sequence(&["b", "f"], 0, 40, sequence), + ]); + let write_opts = WriteOptions { + row_group_size: 50, + ..Default::default() + }; + let mut metrics = Metrics::new(WriteType::Flush); + let mut writer = ParquetWriter::new_with_object_store( + object_store, + metadata, + IndexConfig::default(), + NoopIndexBuilder, + file_path, + &mut metrics, + ) + .await; + if flat_format { + writer + .write_all_flat(source, None, &write_opts) + .await + .unwrap(); + } else { + writer + .write_all_flat_as_primary_key(source, None, &write_opts) + .await + .unwrap(); + } + } - let mut metrics = Metrics::new(WriteType::Flush); - let mut writer = ParquetWriter::new_with_object_store( + let custom_sequence = 12345u64; + let local_zero_handle = sst_file_handle(0, 1000); + write_sst( object_store.clone(), metadata.clone(), - IndexConfig::default(), - NoopIndexBuilder, - file_path, - &mut metrics, + local_zero_handle.clone(), + flat_format, + 0, ) .await; - writer - .write_all_flat_as_primary_key(source, None, &write_opts) - .await - .unwrap() - .remove(0); + // Local all-zero SSTs retain the compatibility override. + let local_zero_none = read_sequences( + ParquetReaderBuilder::new( + FILE_DIR.to_string(), + PathType::Bare, + local_zero_handle.clone(), + object_store.clone(), + ) + .expected_metadata(Some(metadata.clone())), + ) + .await; + assert!(local_zero_none.iter().all(|sequence| *sequence == 0)); - // Read without override sequence (should read sequence 0) - let builder = ParquetReaderBuilder::new( - FILE_DIR.to_string(), - PathType::Bare, - handle.clone(), - object_store.clone(), - ); - let mut reader = builder.build().await.unwrap().unwrap(); - let mut normal_batches = Vec::new(); - while let Some(batch) = reader.next_record_batch().await.unwrap() { - normal_batches.push(batch); - } - - // Read with override sequence using FileMeta.sequence - let custom_sequence = 12345u64; - let file_meta = handle.meta_ref(); - let mut override_file_meta = file_meta.clone(); - override_file_meta.sequence = Some(std::num::NonZero::new(custom_sequence).unwrap()); - let override_handle = FileHandle::new( - override_file_meta, + let mut local_zero_meta = local_zero_handle.meta_ref().clone(); + local_zero_meta.sequence = Some(std::num::NonZeroU64::new(custom_sequence).unwrap()); + let local_zero_override_handle = FileHandle::new( + local_zero_meta, Arc::new(crate::sst::file_purger::NoopFilePurger), ); - - let builder = ParquetReaderBuilder::new( - FILE_DIR.to_string(), - PathType::Bare, - override_handle, - object_store.clone(), + let local_zero_override = read_sequences( + ParquetReaderBuilder::new( + FILE_DIR.to_string(), + PathType::Bare, + local_zero_override_handle, + object_store.clone(), + ) + .expected_metadata(Some(metadata.clone())), + ) + .await; + assert!( + local_zero_override + .iter() + .all(|sequence| *sequence == custom_sequence) ); - let mut reader = builder.build().await.unwrap().unwrap(); - let mut override_batches = Vec::new(); - while let Some(batch) = reader.next_record_batch().await.unwrap() { - override_batches.push(batch); - } - // Compare the results - assert_eq!(normal_batches.len(), override_batches.len()); - for (normal, override_batch) in normal_batches.into_iter().zip(override_batches.iter()) { - let expected_batch = { - let mut columns = normal.columns().to_vec(); - let num_cols = columns.len(); - columns[num_cols - 2] = - Arc::new(UInt64Array::from_value(custom_sequence, normal.num_rows())); - RecordBatch::try_new(normal.schema(), columns).unwrap() - }; + let local_nonzero_handle = sst_file_handle(0, 1000); + write_sst( + object_store.clone(), + metadata.clone(), + local_nonzero_handle.clone(), + flat_format, + 7, + ) + .await; - // Override batch should match expected batch - assert_eq!(*override_batch, expected_batch); - } + // Local nonzero SSTs retain physical per-row sequences, even with FileMeta.sequence. + let mut local_nonzero_meta = local_nonzero_handle.meta_ref().clone(); + local_nonzero_meta.sequence = Some(std::num::NonZeroU64::new(custom_sequence).unwrap()); + let local_nonzero_override_handle = FileHandle::new( + local_nonzero_meta, + Arc::new(crate::sst::file_purger::NoopFilePurger), + ); + let local_nonzero_override = read_sequences( + ParquetReaderBuilder::new( + FILE_DIR.to_string(), + PathType::Bare, + local_nonzero_override_handle, + object_store.clone(), + ) + .expected_metadata(Some(metadata.clone())), + ) + .await; + assert!(local_nonzero_override.iter().all(|sequence| *sequence == 7)); + + // None never overrides a local nonzero physical sequence. + let local_nonzero_none = read_sequences( + ParquetReaderBuilder::new( + FILE_DIR.to_string(), + PathType::Bare, + local_nonzero_handle.clone(), + object_store.clone(), + ) + .expected_metadata(Some(metadata.clone())), + ) + .await; + assert!(local_nonzero_none.iter().all(|sequence| *sequence == 7)); + + // A source-owned handle is foreign when read against target metadata, so the + // target-local manifest barrier is applied even for nonzero physical sequences. + let mut target_metadata = (*metadata).clone(); + target_metadata.region_id = RegionId::new(0, 1); + let target_metadata = Arc::new(target_metadata); + let mut foreign_meta = local_nonzero_handle.meta_ref().clone(); + foreign_meta.sequence = Some(std::num::NonZeroU64::new(custom_sequence).unwrap()); + let foreign_handle = FileHandle::new( + foreign_meta, + Arc::new(crate::sst::file_purger::NoopFilePurger), + ); + let foreign = read_sequences( + ParquetReaderBuilder::new( + FILE_DIR.to_string(), + PathType::Bare, + foreign_handle, + object_store, + ) + .expected_metadata(Some(target_metadata)), + ) + .await; + assert!(foreign.iter().all(|sequence| *sequence == custom_sequence)); } #[tokio::test] diff --git a/src/mito2/src/sst/parquet/reader.rs b/src/mito2/src/sst/parquet/reader.rs index 80f6dee679c..3640fde4526 100644 --- a/src/mito2/src/sst/parquet/reader.rs +++ b/src/mito2/src/sst/parquet/reader.rs @@ -22,7 +22,7 @@ use std::time::{Duration, Instant}; use api::v1::SemanticType; use common_recordbatch::filter::{SimpleFilterEvaluator, TimestampUnitCast}; -use common_telemetry::{error, tracing, warn}; +use common_telemetry::{debug, error, tracing, warn}; use datafusion::physical_plan::PhysicalExpr; use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion}; use datafusion_expr::utils::expr_to_columns; @@ -471,7 +471,24 @@ impl ParquetReaderBuilder { &file_path, skip_auto_convert, )?; - if need_override_sequence(&parquet_meta) { + // `region_meta` comes from the Parquet/source file and must not be used as the + // target identity. When the caller has no current metadata, the handle is the + // only local identity available and therefore denotes a local read. + let expected_region_id = self + .expected_metadata + .as_ref() + .map(|metadata| metadata.region_id) + .unwrap_or(self.file_handle.region_id()); + let is_foreign = self.file_handle.region_id() != expected_region_id; + if is_foreign { + debug!( + "Reading foreign SST, file_id: {}, source_region_id: {}, expected_region_id: {}", + self.file_handle.file_id().file_id(), + self.file_handle.region_id(), + expected_region_id, + ); + } + if is_foreign || need_override_sequence(&parquet_meta) { read_format .set_override_sequence(self.file_handle.meta_ref().sequence.map(|x| x.get())); } diff --git a/src/mito2/src/worker/handle_compaction.rs b/src/mito2/src/worker/handle_compaction.rs index 15dc884a1e5..1f8a66887cd 100644 --- a/src/mito2/src/worker/handle_compaction.rs +++ b/src/mito2/src/worker/handle_compaction.rs @@ -131,6 +131,11 @@ impl RegionWorkerLoop { return; } let execution = request.execution.clone(); + // Whether this execution reduced the SST file count. The scheduler keeps + // draining while compaction makes progress; a rewrite that did not reduce + // files (e.g. its output was split into more files than its input) must end + // the chain to avoid a no-progress loop. + let made_progress = request.edit.files_to_remove.len() > request.edit.files_to_add.len(); region.version_control.apply_edit( Some(request.edit.clone()), @@ -167,6 +172,7 @@ impl RegionWorkerLoop { &execution, ®ion.manifest_ctx, self.schema_metadata_manager.clone(), + made_progress, ) .await; match transition { diff --git a/src/mito2/src/worker/handle_rebuild_index.rs b/src/mito2/src/worker/handle_rebuild_index.rs index ad0a09e1d32..3eb45b1ab21 100644 --- a/src/mito2/src/worker/handle_rebuild_index.rs +++ b/src/mito2/src/worker/handle_rebuild_index.rs @@ -79,6 +79,7 @@ impl RegionWorkerLoop { IndexBuildTask { region_id: region.region_id, file: file.clone(), + target_region_metadata: version.metadata.clone(), source: IndexBuildSource::new(file_meta, version.metadata.schema_version), reason: build_type, access_layer: access_layer.clone(), diff --git a/src/operator/src/insert.rs b/src/operator/src/insert.rs index ead517ab80f..52a73892dc0 100644 --- a/src/operator/src/insert.rs +++ b/src/operator/src/insert.rs @@ -29,12 +29,15 @@ use api::v1::{ use catalog::CatalogManagerRef; use client::{OutputData, OutputMeta}; use common_catalog::consts::{ - PARENT_SPAN_ID_COLUMN, SERVICE_NAME_COLUMN, TRACE_ID_COLUMN, TRACE_TABLE_NAME, - TRACE_TABLE_NAME_SESSION_KEY, default_engine, is_ddl_reserved_table, + DEFAULT_PRIVATE_SCHEMA_NAME, PARENT_SPAN_ID_COLUMN, SERVICE_NAME_COLUMN, TRACE_ID_COLUMN, + TRACE_TABLE_NAME, TRACE_TABLE_NAME_SESSION_KEY, default_engine, is_ddl_reserved_table, trace_operations_table_name, trace_services_table_name, }; +use common_event_recorder::DEFAULT_EVENTS_TABLE_NAME; +use common_frontend::slow_query_event::SLOW_QUERY_TABLE_NAME; use common_grpc_expr::util::ColumnExpr; use common_meta::cache::TableFlownodeSetCacheRef; +use common_meta::datanode::REGION_STATS_HISTORY_TABLE_NAME; use common_meta::node_manager::{AffectedRows, NodeManagerRef}; use common_meta::peer::Peer; use common_meta::rpc::ddl::TriggerReason; @@ -490,9 +493,9 @@ impl Inserter { Ok(inserts) } - /// Returns `None` if auto table creation is allowed, or `Some(reason)` if - /// disabled by either the global config or the request hint. The reason tells - /// which one, for a clearer error. + /// Returns `Some(reason)` if the config or request hint disables automatic + /// table creation. Exempt private system tables are handled by + /// [`Self::is_auto_create_exempt_private_table`]. fn auto_create_disabled_reason(&self, ctx: &QueryContextRef) -> Result> { let auto_create_table_hint = ctx .extension(AUTO_CREATE_TABLE_KEY) @@ -514,6 +517,16 @@ impl Inserter { }) } + /// Returns whether a private system table may infer and reconcile its schema + /// even when automatic table creation is disabled. + fn is_auto_create_exempt_private_table(schema: &str, table: &str) -> bool { + schema == DEFAULT_PRIVATE_SCHEMA_NAME + && matches!( + table, + DEFAULT_EVENTS_TABLE_NAME | SLOW_QUERY_TABLE_NAME | REGION_STATS_HISTORY_TABLE_NAME + ) + } + /// Ensures a trace table has the request-global schema without requiring a /// padded data row to drive on-demand creation or alteration. When /// `alter_existing` is false, a table created after planning is left for the @@ -571,12 +584,21 @@ impl Inserter { let _timer = crate::metrics::CREATE_ALTER_ON_DEMAND .with_label_values(&[auto_create_table_type.as_str()]) .start_timer(); - let catalog = ctx.current_catalog(); let schema = ctx.current_schema(); + let auto_create_disabled_reason = self.auto_create_disabled_reason(ctx)?; + // Enabled batches permit every table, so only disabled batches need a whitelist scan. + let has_auto_create_exempt_table = auto_create_disabled_reason.is_some() + && requests + .inserts + .iter() + .any(|req| Self::is_auto_create_exempt_private_table(&schema, &req.table_name)); let mut table_infos = HashMap::new(); - if let Some(disabled_reason) = self.auto_create_disabled_reason(ctx)? { + // Without exempt tables, verify existing tables and reject missing ones without inferring schemas. + if let Some(disabled_reason) = auto_create_disabled_reason + && !has_auto_create_exempt_table + { let mut instant_table_ids = HashSet::new(); for req in &requests.inserts { let table = match self.get_table(catalog, &schema, &req.table_name).await? { @@ -618,20 +640,25 @@ impl Inserter { let mut per_table_semantics: Option> = None; for req in &mut requests.inserts { + // Mixed batches need a per-table decision so an exempt table cannot authorize others. + let auto_create_allowed = auto_create_disabled_reason.is_none() + || Self::is_auto_create_exempt_private_table(&schema, &req.table_name); match self.get_table(catalog, &schema, &req.table_name).await? { Some(table) => { let table_info = table.table_info(); if table_info.is_ttl_instant_table() { instant_table_ids.insert(table_info.table_id()); } - if let Some(alter_expr) = self.get_alter_table_expr_on_demand( - req, - &table, - ctx, - accommodate_existing_schema, - is_single_value, - auto_create_table_type.alter_existing(), - )? { + if auto_create_allowed + && let Some(alter_expr) = self.get_alter_table_expr_on_demand( + req, + &table, + ctx, + accommodate_existing_schema, + is_single_value, + auto_create_table_type.alter_existing(), + )? + { alter_tables.push(alter_expr); need_refresh_table_infos.insert(( catalog.to_string(), @@ -655,6 +682,17 @@ impl Inserter { } table_infos.insert(table_info.table_id(), table_info); } + None if !auto_create_allowed + && let Some(disabled_reason) = auto_create_disabled_reason => + { + return InvalidInsertRequestSnafu { + reason: format!( + "Table `{}` does not exist, and {}", + req.table_name, disabled_reason, + ), + } + .fail(); + } None => { let semantic_index = per_table_semantics .get_or_insert_with(|| parse_per_table_semantic_index(ctx)) diff --git a/src/operator/src/statement.rs b/src/operator/src/statement.rs index 478f06e83f7..21c3ed61464 100644 --- a/src/operator/src/statement.rs +++ b/src/operator/src/statement.rs @@ -84,7 +84,7 @@ use table::table_reference::TableReference; pub use self::admin::{ AdminEventRecorderHandle, AdminFunctionLayer, AdminFunctionLayerRef, AdminFunctionRecordingLayer, AdminFunctionRequest, AdminFunctionResponse, AdminFunctionService, - AdminFunctionServiceRef, + AdminFunctionServiceRef, admin_output_schema, }; use self::set::{ set_bytea_output, set_datestyle, set_intervalstyle, set_timezone, validate_client_encoding, diff --git a/src/operator/src/statement/admin.rs b/src/operator/src/statement/admin.rs index 2b4a909e3cd..ff8b26915cc 100644 --- a/src/operator/src/statement/admin.rs +++ b/src/operator/src/statement/admin.rs @@ -19,6 +19,7 @@ use std::sync::Arc; use common_function::function::FunctionContext; use common_function::function_registry::{FUNCTION_REGISTRY, get_admin_function}; +use common_function::state::FunctionState; use common_query::Output; use common_recordbatch::{RecordBatch, RecordBatches}; use common_sql::convert::sql_value_to_value; @@ -77,6 +78,103 @@ struct CoreAdminFunctionService { query_engine: query::QueryEngineRef, } +/// Parts of an `ADMIN` call needed both for execution and schema derivation. +struct ResolvedAdminFunction { + admin_udf: datafusion_expr::ScalarUDF, + fn_name: String, + args: Vec, + arg_types: Vec, + ret_type: ArrowDataType, +} + +/// Resolves the function, parses its literal arguments and derives its +/// return type, without executing it. +fn resolve_admin_function( + stmt: &Admin, + query_ctx: &QueryContextRef, + state: Arc, +) -> Result { + let Admin::Func(func) = stmt; + // the function name should be in lower case. + let func_name = func.name.to_string().to_lowercase(); + let factory = get_admin_function(&func_name) + .or_else(|| FUNCTION_REGISTRY.get_function(&func_name)) + .context(error::AdminFunctionNotFoundSnafu { + name: func_name.clone(), + })?; + + let func_ctx = FunctionContext { + query_ctx: query_ctx.clone(), + state, + }; + + let admin_udf = factory.provide(func_ctx); + admin_udf + .as_async() + .context(error::AdminFunctionNotFoundSnafu { name: func_name })?; + + let fn_name = admin_udf.name().to_string(); + let signature = admin_udf.signature(); + + // Parse function arguments + let FunctionArguments::List(args) = &func.args else { + return error::BuildAdminFunctionArgsSnafu { + msg: format!("unsupported function args {} for {}", func.args, fn_name), + } + .fail(); + }; + let arg_values = args + .args + .iter() + .map(|arg| { + let FunctionArg::Unnamed(FunctionArgExpr::Expr(Expr::Value(value))) = arg else { + return error::BuildAdminFunctionArgsSnafu { + msg: format!("unsupported function arg {arg} for {}", fn_name), + } + .fail(); + }; + Ok(&value.value) + }) + .collect::>>()?; + + let args = args_to_vector(&signature.type_signature, &arg_values, query_ctx)?; + let arg_types = args + .iter() + .map(|arg| arg.data_type().as_arrow_type()) + .collect::>(); + let ret_type = + admin_udf + .return_type(&arg_types) + .map_err(|e| error::Error::BuildAdminFunctionArgs { + msg: format!( + "Failed to get return type of admin function {}: {}", + fn_name, e + ), + })?; + + Ok(ResolvedAdminFunction { + admin_udf, + fn_name, + args, + arg_types, + ret_type, + }) +} + +/// Output schema of an `ADMIN` statement, mirroring what +/// [`CoreAdminFunctionService::execute`] produces. `None` if the function +/// or arguments are unresolvable; execution will then surface the error. +pub fn admin_output_schema(stmt: &Admin, query_ctx: &QueryContextRef) -> Option { + let resolved = + resolve_admin_function(stmt, query_ctx, Arc::new(FunctionState::default())).ok()?; + Some(Schema::new(vec![ColumnSchema::new( + // Use statement as the result column name + stmt.to_string(), + ConcreteDataType::from_arrow_type(&resolved.ret_type), + false, + )])) +} + impl CoreAdminFunctionService { fn new(query_engine: query::QueryEngineRef) -> Self { Self { query_engine } @@ -88,62 +186,23 @@ impl CoreAdminFunctionService { query_ctx, } = request; - let Admin::Func(func) = &stmt; - // the function name should be in lower case. - let func_name = func.name.to_string().to_lowercase(); - let factory = get_admin_function(&func_name) - .or_else(|| FUNCTION_REGISTRY.get_function(&func_name)) - .context(error::AdminFunctionNotFoundSnafu { - name: func_name.clone(), - })?; - - let func_ctx = FunctionContext { - query_ctx: query_ctx.clone(), - state: self.query_engine.engine_state().function_state(), - }; - - let admin_udf = factory.provide(func_ctx); + let resolved = resolve_admin_function( + &stmt, + &query_ctx, + self.query_engine.engine_state().function_state(), + )?; + let ResolvedAdminFunction { + admin_udf, + fn_name, + args, + arg_types, + ret_type, + } = resolved; let admin_async_fn = admin_udf .as_async() - .context(error::AdminFunctionNotFoundSnafu { name: func_name })?; - - let fn_name = admin_udf.name(); - let signature = admin_udf.signature(); - - // Parse function arguments - let FunctionArguments::List(args) = &func.args else { - return error::BuildAdminFunctionArgsSnafu { - msg: format!("unsupported function args {} for {}", func.args, fn_name), - } - .fail(); - }; - let arg_values = args - .args - .iter() - .map(|arg| { - let FunctionArg::Unnamed(FunctionArgExpr::Expr(Expr::Value(value))) = arg else { - return error::BuildAdminFunctionArgsSnafu { - msg: format!("unsupported function arg {arg} for {}", fn_name), - } - .fail(); - }; - Ok(&value.value) - }) - .collect::>>()?; - - let args = args_to_vector(&signature.type_signature, &arg_values, &query_ctx)?; - let arg_types = args - .iter() - .map(|arg| arg.data_type().as_arrow_type()) - .collect::>(); - let ret_type = admin_udf.return_type(&arg_types).map_err(|e| { - error::Error::BuildAdminFunctionArgs { - msg: format!( - "Failed to get return type of admin function {}: {}", - fn_name, e - ), - } - })?; + .context(error::AdminFunctionNotFoundSnafu { + name: fn_name.clone(), + })?; // Convert arguments to DataFusion ColumnarValue format let columnar_args: Vec = args @@ -174,9 +233,7 @@ impl CoreAdminFunctionService { let result_columnar = admin_async_fn .invoke_async_with_args(func_args) .await - .with_context(|_| ExecuteAdminFunctionSnafu { - msg: fn_name.to_string(), - })?; + .with_context(|_| ExecuteAdminFunctionSnafu { msg: fn_name })?; // Convert result back to VectorRef let result_columnar: common_query::prelude::ColumnarValue = diff --git a/src/operator/src/statement/semantic_graph.rs b/src/operator/src/statement/semantic_graph.rs index 7014a94f8ac..86827cf8c7a 100644 --- a/src/operator/src/statement/semantic_graph.rs +++ b/src/operator/src/statement/semantic_graph.rs @@ -41,8 +41,9 @@ use common_catalog::consts::{ ENTITY_DESCRIPTIVE_COLUMN, ENTITY_ID_ATTRS_COLUMN, ENTITY_ID_COLUMN, ENTITY_SCOPE_COLUMN, ENTITY_TYPE_COLUMN, ERROR_COUNT_COLUMN, FRESH_UNTIL_COLUMN, GENERATION_ID_COLUMN, OBSERVED_AT_COLUMN, PROVENANCE_COLUMN, REL_TYPE_COLUMN, REQUEST_COUNT_COLUMN, - SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME, SOURCE_TABLES_COLUMN, SRC_ID_COLUMN, - SRC_TYPE_COLUMN, VALID_FROM_COLUMN, VALID_UNTIL_COLUMN, WINDOW_END_COLUMN, WINDOW_START_COLUMN, + SEMANTIC_GRAPH_WINDOW_NANOS, SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME, SOURCE_TABLES_COLUMN, + SRC_ID_COLUMN, SRC_TYPE_COLUMN, VALID_FROM_COLUMN, VALID_UNTIL_COLUMN, WINDOW_END_COLUMN, + WINDOW_START_COLUMN, }; use common_function::function::FunctionContext; use common_function::function_registry::FUNCTION_REGISTRY; @@ -60,7 +61,7 @@ use datafusion::dataframe::DataFrame; use datafusion::functions::{core as core_fns, datetime as datetime_fns, string as string_fns}; use datafusion::functions_nested::expr_fn::make_array; use datafusion_common::{Column, Result as DfResult, ScalarValue}; -use datafusion_expr::{Expr, LogicalPlan, ScalarUDF, cast, ident, lit}; +use datafusion_expr::{Case, Expr, LogicalPlan, ScalarUDF, cast, ident, lit}; pub use relationships::{ CallsSource, CoDeclaredSource, DeclaredSource, RelationshipSources, build_relationships_plan, }; @@ -122,9 +123,10 @@ pub fn declared_relationships_schema_matches(table_info: &table::metadata::Table }) } -/// Bin width for the temporal window of derived rows: 60s buckets, matching the -/// service-graph convention. -const BIN_NANOS: i64 = 60 * 1_000_000_000; +/// Bin width for the temporal window of derived rows, matching the +/// service-graph convention. Shared so ingestion-synthesized observations +/// land in the same buckets. +const BIN_NANOS: i64 = SEMANTIC_GRAPH_WINDOW_NANOS; /// Default retention for the declared-edge table; expiry slides the topology window. const DEFAULT_DECLARED_RELATIONSHIPS_TTL: &str = "90d"; @@ -284,8 +286,12 @@ pub struct EntityDeclaration { /// The table's time index column, used for the temporal window filter. pub time_index: String, pub entity_type: String, - /// Identifying columns (>= 1). One column → id verbatim; several → composite. + /// Identifying columns (>= 1), ordered broad to narrow. pub id_columns: Vec, + /// Optional column qualifying the first identity component, so a source + /// carrying the parts separately (a trace table's `service.namespace`) + /// yields the same id as one carrying them pre-composed (`job`). + pub id_qualifier: Option, /// Descriptive columns snapshotted into the `descriptive` JSON (may be empty). pub descriptive_columns: Vec, /// Scope columns (namespace/environment). One column → scope verbatim; @@ -417,28 +423,91 @@ fn cast_string_or_empty(column: &str) -> Expr { core_fns::coalesce().call(vec![cast(ident(column), DataType::Utf8), lit("")]) } -/// The canonical entity-id expression for `id_columns`: the value verbatim for -/// a single column, the sorted `k=v,k=v` rendering for a composite. `col` -/// constructs the column reference (unqualified for registry branches, -/// join-side-qualified for the calls derivation). /// A row identifies an entity only when every identity component is present /// and non-empty: kube-state-metrics descriptors emit empty-string labels (an /// unscheduled pod's `node`, an owner-less pod's `owner_*`), and an empty -/// string is never a meaningful entity id. +/// string is never a meaningful entity id. The qualifier is optional by +/// construction and so is not guarded. pub(crate) fn identifies(column: &str) -> Expr { ident(column) .is_not_null() .and(cast(ident(column), DataType::Utf8).not_eq(lit(""))) } -fn entity_id_expr(id_columns: &[String], col: &dyn Fn(&str) -> Expr) -> Expr { - if let [id] = id_columns { - cast(col(id), DataType::Utf8) - } else { - let mut cols = id_columns.to_vec(); - cols.sort(); - sorted_kv_expr_with(&cols, false, col) +const ID_SEPARATOR: &str = ","; +const ID_ESCAPE: &str = "\\"; +/// Left unescaped: `/` is how Prometheus renders `job`. +const ID_QUALIFIER_SEPARATOR: &str = "/"; + +/// Escapes so a composite id decodes back to its components. The escape +/// character goes first, or it would double the escapes the separator pass +/// introduces. +fn escaped_id_value(value: Expr) -> Expr { + let escape_escape = string_fns::replace().call(vec![ + value, + lit(ID_ESCAPE), + lit(format!("{ID_ESCAPE}{ID_ESCAPE}")), + ]); + string_fns::replace().call(vec![ + escape_escape, + lit(ID_SEPARATOR), + lit(format!("{ID_ESCAPE}{ID_SEPARATOR}")), + ]) +} + +/// The identity values in declared order (broad to narrow), escaped and +/// joined. +/// +/// The identifying *column names* are deliberately absent: the same identity +/// reaches us under different names per source — a trace table's +/// `service_name` against a metric table's `job` — and encoding them would +/// split one entity into one per signal. `entity_id_attrs` keeps the +/// structured form. +/// +/// `col` constructs the column reference (unqualified for registry branches, +/// join-side-qualified for the calls derivation). +fn entity_id_expr( + id_columns: &[String], + qualifier: Option<&str>, + col: &dyn Fn(&str) -> Expr, +) -> Expr { + let mut parts = Vec::with_capacity(id_columns.len() * 2); + for (index, column) in id_columns.iter().enumerate() { + if index > 0 { + parts.push(lit(ID_SEPARATOR)); + } + let value = escaped_id_value(cast(col(column), DataType::Utf8)); + parts.push(match qualifier { + Some(qualifier) if index == 0 => qualified_id_expr(qualifier, value, col), + _ => value, + }); } + if let [single] = parts.as_slice() { + single.clone() + } else { + concat_expr(parts) + } +} + +/// `/`, or bare `value` when the qualifier is empty on this +/// row — the spec's rule composing `job` from `service.namespace` and +/// `service.name`. +fn qualified_id_expr(qualifier: &str, value: Expr, col: &dyn Fn(&str) -> Expr) -> Expr { + let qualifier = escaped_id_value( + core_fns::coalesce().call(vec![cast(col(qualifier), DataType::Utf8), lit("")]), + ); + Expr::Case(Case::new( + None, + vec![( + Box::new(qualifier.clone().eq(lit(""))), + Box::new(value.clone()), + )], + Some(Box::new(concat_expr(vec![ + qualifier, + lit(ID_QUALIFIER_SEPARATOR), + value, + ]))), + )) } /// The `parse_json` UDF, shared by all derivation plans. Resolved from the @@ -583,14 +652,22 @@ fn registry_source( let mut rows = Vec::with_capacity(1 + rest.len()); for decl in std::iter::once(first).chain(rest) { // CAST even a single-column id: id columns need not be strings, and - // the computed table declares entity_id STRING. Composite ids - // additionally carry a JSON object of the id columns in - // entity_id_attrs. - let entity_id = entity_id_expr(&decl.id_columns, &|c| ident(c)); - let entity_id_attrs = if decl.id_columns.len() == 1 { + // the computed table declares entity_id STRING. An id assembled from + // more than one column additionally carries its parts as a JSON + // object in entity_id_attrs, the structured form of the id string. + let entity_id = entity_id_expr(&decl.id_columns, decl.id_qualifier.as_deref(), &|c| { + ident(c) + }); + let id_parts = decl + .id_qualifier + .iter() + .chain(&decl.id_columns) + .cloned() + .collect::>(); + let entity_id_attrs = if id_parts.len() == 1 { null_json() } else { - json_object_expr(&decl.id_columns) + json_object_expr(&id_parts) }; let scope = match decl.scope_columns.as_slice() { @@ -811,11 +888,74 @@ mod tests { time_index: "ts".to_string(), entity_type: entity_type.to_string(), id_columns: id_columns.iter().map(|s| s.to_string()).collect(), + id_qualifier: None, descriptive_columns: vec![], scope_columns: vec![], } } + /// The id string must decode back to its components, and a qualifier must + /// reproduce the identity a source that pre-composed it already emits. + #[tokio::test] + async fn registry_id_escaping_and_qualifier() { + let schema = Arc::new(Schema::new(vec![ + Field::new( + "ts", + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + Field::new("service_name", DataType::Utf8, false), + Field::new("instance", DataType::Utf8, false), + Field::new("namespace", DataType::Utf8, true), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(vec![1_000, 2_000, 3_000])) as ArrayRef, + Arc::new(StringArray::from(vec!["cart", "a,b", "we\\ird"])), + Arc::new(StringArray::from(vec!["i-1", "c", "i-3"])), + Arc::new(StringArray::from(vec![Some("shop"), None, Some("")])), + ], + ) + .unwrap(); + let ctx = SessionContext::new(); + ctx.register_table( + "svc", + Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap()), + ) + .unwrap(); + + let mut declaration = decl("service.instance", &["service_name", "instance"]); + declaration.id_qualifier = Some("namespace".to_string()); + declaration.time_index = "ts".to_string(); + let plan = build_registry_plan( + vec![RegistrySource { + declarations: vec![declaration], + scan: ctx.table("svc").await.unwrap(), + }], + &test_window(), + ) + .unwrap() + .unwrap(); + + let mut ids: Vec = collect(&ctx, plan) + .await + .iter() + .flat_map(|batch| strings(batch, 5)) + .collect(); + ids.sort(); + assert_eq!( + ids, + vec![ + // an absent qualifier leaves the identity bare, and a value + // holding the separator stays distinguishable from two values + "a\\,b,c".to_string(), + "shop/cart,i-1".to_string(), + "we\\\\ird,i-3".to_string(), + ] + ); + } + #[tokio::test] async fn registry_single_column_identity() { let ctx = metric_table_ctx(); @@ -900,24 +1040,24 @@ mod tests { .collect(); rows.sort(); - // Composite id -> sorted `k=v,k=v` plus a JSON object of the id columns; - // descriptive JSON keeps `\`, `"` and control characters intact in - // runtime values, NULL -> "". + // Composite id -> values in declared order plus a JSON object of the + // id columns; descriptive JSON keeps `\`, `"` and control characters + // intact in runtime values, NULL -> "". assert_eq!( rows, vec![ ( - "pid=42,service_name=cart".to_string(), + "cart,42".to_string(), Some(r#"{"pid":"42","service_name":"cart"}"#.to_string()), Some(r#"{"host":""}"#.to_string()), ), ( - "pid=42,service_name=cart".to_string(), + "cart,42".to_string(), Some(r#"{"pid":"42","service_name":"cart"}"#.to_string()), Some(r#"{"host":"h2"}"#.to_string()), ), ( - "pid=42,service_name=cart".to_string(), + "cart,42".to_string(), Some(r#"{"pid":"42","service_name":"cart"}"#.to_string()), Some(r#"{"host":"we\"ird\\\nhost"}"#.to_string()), ), diff --git a/src/operator/src/statement/semantic_graph/conventions.rs b/src/operator/src/statement/semantic_graph/conventions.rs index 36d058921b5..9ec6f52e158 100644 --- a/src/operator/src/statement/semantic_graph/conventions.rs +++ b/src/operator/src/statement/semantic_graph/conventions.rs @@ -19,7 +19,7 @@ //! with the binary, not an operator-editable configuration surface; explicit //! `greptime.semantic.entity.*` declarations always override it. -use std::collections::{BTreeMap, HashSet}; +use std::collections::{BTreeMap, HashMap, HashSet}; use std::sync::LazyLock; use serde::Deserialize; @@ -43,15 +43,20 @@ pub struct VirtualDstCandidate { pub connection_type: String, } -/// One implicit entity declaration: of a whitelisted Prometheus info metric, -/// or of a trace-v1 table's flattened resource attributes. +/// One implicit entity declaration: of a whitelisted Prometheus or OTel info +/// metric, or of a trace-v1 table's flattened resource attributes. #[derive(Debug, Deserialize)] #[serde(deny_unknown_fields)] pub struct ImplicitEntity { pub entity: String, - /// Identifying label columns; every one must exist on the table for the - /// declaration to apply. + /// Identifying label columns, ordered broad to narrow; every one must + /// exist on the table for the declaration to apply. pub id: Vec, + /// Column qualifying `id[0]` as `/`, so a source + /// carrying the parts separately matches one carrying them pre-composed. + /// Skipped when the column is absent or empty. + #[serde(default)] + pub qualified_by: Option, /// Descriptive label columns, filtered to those present (kube-state-metrics /// label sets vary across versions). #[serde(default)] @@ -70,7 +75,12 @@ pub struct Conventions { pub trace_co_declared_edges: Vec, pub virtual_dst_candidates: Vec, pub otlp_trace_entities: Vec, + /// Table name -> the entities that table declares, for Prometheus-sourced + /// descriptor metrics (`source = prometheus`). pub prometheus_info_metrics: BTreeMap>, + /// Table name -> the entities that table declares, for OTLP-sourced + /// descriptor tables (`source = opentelemetry`). + pub otel_info_metrics: BTreeMap>, } /// The built-in entity-type vocabulary. User-declared types are open-ended; @@ -78,6 +88,7 @@ pub struct Conventions { pub const ENTITY_TYPE_SERVICE: &str = "service"; pub const ENTITY_TYPE_SERVICE_INSTANCE: &str = "service.instance"; pub const ENTITY_TYPE_HOST: &str = "host"; +pub const ENTITY_TYPE_CONTAINER: &str = "container"; pub const ENTITY_TYPE_PROCESS: &str = "process"; pub const ENTITY_TYPE_K8S_POD: &str = "k8s.pod"; pub const ENTITY_TYPE_K8S_NODE: &str = "k8s.node"; @@ -88,10 +99,11 @@ pub const ENTITY_TYPE_GEN_AI_AGENT: &str = "gen_ai.agent"; pub const ENTITY_TYPE_GEN_AI_MODEL: &str = "gen_ai.model"; pub const ENTITY_TYPE_GEN_AI_TOOL: &str = "gen_ai.tool"; -const ENTITY_TYPES: [&str; 12] = [ +const ENTITY_TYPES: [&str; 13] = [ ENTITY_TYPE_SERVICE, ENTITY_TYPE_SERVICE_INSTANCE, ENTITY_TYPE_HOST, + ENTITY_TYPE_CONTAINER, ENTITY_TYPE_PROCESS, ENTITY_TYPE_K8S_POD, ENTITY_TYPE_K8S_NODE, @@ -200,11 +212,16 @@ fn validate(conventions: &Conventions) -> Result<(), String> { let per_table = conventions .prometheus_info_metrics .iter() + .chain(&conventions.otel_info_metrics) .map(|(table, entities)| (table.as_str(), entities)) .chain(std::iter::once(( "otlp traces", &conventions.otlp_trace_entities, ))); + // An entity type declared with a different number of id columns by two + // sources yields ids that can never match, silently splitting one entity + // in two. + let mut arity = HashMap::new(); for (table, entities) in per_table { let mut seen_types = HashSet::new(); for implicit in entities { @@ -226,6 +243,12 @@ fn validate(conventions: &Conventions) -> Result<(), String> { implicit.entity )); } + if implicit.qualified_by.as_ref().is_some_and(String::is_empty) { + return Err(format!( + "entity `{}` of info metric `{table}` has an empty qualified_by", + implicit.entity + )); + } if implicit.descriptive_rest && !implicit.descriptive.is_empty() { return Err(format!( "entity `{}` of info metric `{table}` sets both descriptive and \ @@ -233,6 +256,17 @@ fn validate(conventions: &Conventions) -> Result<(), String> { implicit.entity )); } + match arity.insert(implicit.entity.as_str(), implicit.id.len()) { + Some(previous) if previous != implicit.id.len() => { + return Err(format!( + "entity `{}` is declared with {previous} id columns elsewhere but \ + {} for `{table}`", + implicit.entity, + implicit.id.len() + )); + } + _ => {} + } } } Ok(()) @@ -249,25 +283,28 @@ mod tests { /// Each case must fail on exactly the rule it names, so the shared /// boilerplate is valid and the mutated part is inside the vocabulary. - fn broken(edges: &str, info_metrics: &str) -> String { + fn broken(edges: &str, info_metrics: &str, otel_metrics: &str) -> String { format!( "co_declared_edges: [{edges}]\ntrace_co_declared_edges: []\n\ virtual_dst_candidates: []\notlp_trace_entities: []\n\ - prometheus_info_metrics: {{{info_metrics}}}" + prometheus_info_metrics: {{{info_metrics}}}\n\ + otel_info_metrics: {{{otel_metrics}}}" ) } #[test] fn validation_rejects_broken_conventions() { - let err = |edges, info| parse(&broken(edges, info)).unwrap_err(); + let err = |edges, info, otel| parse(&broken(edges, info, otel)).unwrap_err(); - assert!(err("{src: host, dst: service, rel: pets}", "").contains("unknown rel_type")); + assert!(err("{src: host, dst: service, rel: pets}", "", "").contains("unknown rel_type")); assert!( - err("{src: k8s.pods, dst: k8s.node, rel: runs_on}", "").contains("unknown entity type") + err("{src: k8s.pods, dst: k8s.node, rel: runs_on}", "", "") + .contains("unknown entity type") ); assert!( err( "{src: host, dst: service, rel: uses}, {src: host, dst: service, rel: uses}", + "", "" ) .contains("duplicate edge rule") @@ -275,14 +312,27 @@ mod tests { assert!( err( "", - "t: [{entity: host, id: [x], descriptive: [y], descriptive_rest: true}]" + "t: [{entity: host, id: [x], descriptive: [y], descriptive_rest: true}]", + "" ) .contains("descriptive_rest") ); + // the otel map runs through the same per-table validation + assert!(err("", "", "t: [{entity: hosts, id: [x]}]").contains("unknown entity type")); + // ids of a different arity for one type can never match each other + assert!( + err( + "", + "a: [{entity: host, id: [x]}]", + "b: [{entity: host, id: [x, y]}]" + ) + .contains("id columns") + ); // Unknown YAML keys are rejected, catching typos in the embedded file. assert!( parse(&broken( "{src: host, dst: service, rel: uses, direction: down}", + "", "" )) .is_err() diff --git a/src/operator/src/statement/semantic_graph/conventions.yaml b/src/operator/src/statement/semantic_graph/conventions.yaml index 9fc37fe2a02..808fe0f2658 100644 --- a/src/operator/src/statement/semantic_graph/conventions.yaml +++ b/src/operator/src/statement/semantic_graph/conventions.yaml @@ -2,6 +2,10 @@ # versioned with the binary (include_str!), not an operator-editable surface; # non-standard deployments override per table with explicit # greptime.semantic.entity.* declarations, which always win. +# +# Every `id` list is ordered broad to narrow, and that order is part of the +# identity: the id renders as the values joined in this order, so all +# declarations of one entity type must agree on it or they name two entities. co_declared_edges: - { src: service.instance, dst: host, rel: runs_on } @@ -11,6 +15,8 @@ co_declared_edges: - { src: service.instance, dst: service, rel: part_of } - { src: k8s.pod, dst: k8s.workload, rel: part_of } - { src: service.instance, dst: k8s.pod, rel: runs_on } + - { src: service.instance, dst: container, rel: runs_on } + - { src: container, dst: host, rel: runs_on } # Applied only to trace sources; the endpoints are Greptime entity types # derived from GenAI semantic-convention attributes. @@ -30,12 +36,25 @@ virtual_dst_candidates: - { column: span_attributes.server.address, connection_type: virtual_node } # Implicit declarations for greptime_trace_v1 tables; the pod UID bridges -# trace-side pods onto the kube-state-metrics entities below. +# trace-side pods onto the kube-state-metrics entities below. host.name and +# container.name vary by SDK and resource detector, so identity is the id. +# qualified_by folds the namespace into the service identity the same way the +# metric sources' `job` already carries it, so both signals name one entity. otlp_trace_entities: - entity: service id: [service_name] + qualified_by: resource_attributes.service.namespace - entity: service.instance id: [service_name, resource_attributes.service.instance.id] + qualified_by: resource_attributes.service.namespace + - entity: host + id: [resource_attributes.host.id] + descriptive: + - resource_attributes.host.name + - entity: container + id: [resource_attributes.container.id] + descriptive: + - resource_attributes.container.name - entity: k8s.pod id: [resource_attributes.k8s.pod.uid] descriptive: @@ -113,3 +132,32 @@ prometheus_info_metrics: - entity: service.instance id: [job, instance] descriptive_rest: true + +# Implicit declarations for the OTLP resource descriptor synthesized at +# metrics ingestion, gated on source=opentelemetry and metric.type=info. +# job/instance are the derived Prometheus-compatible identity, so services +# line up with target_info. Unlike target_info the attribute set is known, so +# descriptive columns are listed instead of swept up by descriptive_rest, +# which would misfile other entities' ids under service.instance. +otel_info_metrics: + greptime_otel_resource_info: + - entity: service + id: [job] + descriptive: + - service.name + - service.namespace + - entity: service.instance + id: [job, instance] + - entity: host + id: [host.id] + descriptive: + - host.name + - entity: container + id: [container.id] + descriptive: + - container.name + - entity: k8s.pod + id: [k8s.pod.uid] + descriptive: + - k8s.pod.name + - k8s.namespace.name diff --git a/src/operator/src/statement/semantic_graph/relationships.rs b/src/operator/src/statement/semantic_graph/relationships.rs index 3b5c6d37b78..257c844a9a8 100644 --- a/src/operator/src/statement/semantic_graph/relationships.rs +++ b/src/operator/src/statement/semantic_graph/relationships.rs @@ -200,9 +200,9 @@ fn co_declared_branch( bin.clone() + bin_interval(), bin.clone() + bin_interval(), lit(src.entity_type.as_str()), - entity_id_expr(&src.id_columns, &|c| ident(c)), + entity_id_expr(&src.id_columns, src.id_qualifier.as_deref(), &|c| ident(c)), lit(dst.entity_type.as_str()), - entity_id_expr(&dst.id_columns, &|c| ident(c)), + entity_id_expr(&dst.id_columns, dst.id_qualifier.as_deref(), &|c| ident(c)), lit(*rel_type), lit(*provenance), lit(1.0_f64), @@ -501,6 +501,13 @@ fn span_predicate(service: &EntityDeclaration, window: &GraphQueryWindow, strict } /// One trace table's client spans, normalized to +/// A trace table's `duration_nano` is UInt64 on tables created before the +/// signed-integer ingest change and Int64 after it. The per-table selects are +/// unioned, and those two have no common integer type. +fn duration_nano_expr() -> Expr { + cast(ident(DURATION_NANO_COLUMN), DataType::Int64).alias(DURATION_NANO_COLUMN) +} + /// `(timestamp, trace_id, span_id, src_id, status_code, duration_nano, /// virtual_dst, virtual_conn)`. `src_id` is built from the table's `service` /// declaration, so edges land on exactly the entity ids the registry emits (a @@ -553,9 +560,12 @@ fn client_spans( ident(SPAN_ID_COLUMN), // The cast inside entity_id_expr also normalizes tag columns, which // come out of the storage engine dictionary-encoded. - entity_id_expr(&service.id_columns, &|c| ident(c)).alias(SRC_ID_COLUMN), + entity_id_expr(&service.id_columns, service.id_qualifier.as_deref(), &|c| { + ident(c) + }) + .alias(SRC_ID_COLUMN), ident(SPAN_STATUS_CODE_COLUMN), - ident(DURATION_NANO_COLUMN), + duration_nano_expr(), virtual_dst.alias("virtual_dst"), virtual_conn.alias("virtual_conn"), ]) @@ -578,9 +588,12 @@ fn server_spans( ident(TRACE_TIMESTAMP_COLUMN), ident(TRACE_ID_COLUMN), ident(PARENT_SPAN_ID_COLUMN), - entity_id_expr(&service.id_columns, &|c| ident(c)).alias(DST_ID_COLUMN), + entity_id_expr(&service.id_columns, service.id_qualifier.as_deref(), &|c| { + ident(c) + }) + .alias(DST_ID_COLUMN), ident(SPAN_STATUS_CODE_COLUMN), - ident(DURATION_NANO_COLUMN), + duration_nano_expr(), ]) } @@ -609,7 +622,10 @@ fn agent_calls_branch( ident(TRACE_TIMESTAMP_COLUMN), ident(TRACE_ID_COLUMN), ident(SPAN_ID_COLUMN), - entity_id_expr(&agent.id_columns, &|c| ident(c)).alias(SRC_ID_COLUMN), + entity_id_expr(&agent.id_columns, agent.id_qualifier.as_deref(), &|c| { + ident(c) + }) + .alias(SRC_ID_COLUMN), ])?; let child = trace .scan @@ -619,9 +635,12 @@ fn agent_calls_branch( ident(TRACE_TIMESTAMP_COLUMN), ident(TRACE_ID_COLUMN), ident(PARENT_SPAN_ID_COLUMN), - entity_id_expr(&agent.id_columns, &|c| ident(c)).alias(DST_ID_COLUMN), + entity_id_expr(&agent.id_columns, agent.id_qualifier.as_deref(), &|c| { + ident(c) + }) + .alias(DST_ID_COLUMN), ident(SPAN_STATUS_CODE_COLUMN), - ident(DURATION_NANO_COLUMN), + duration_nano_expr(), ])?; parents = union_all(parents, parent)?; children = union_all(children, child)?; @@ -732,6 +751,16 @@ mod tests { name: &str, extra: &[(&str, &[Option<&str>])], spans: &[Span<'_>], + ) { + register_typed_trace_table(ctx, name, extra, spans, DataType::UInt64) + } + + fn register_typed_trace_table( + ctx: &SessionContext, + name: &str, + extra: &[(&str, &[Option<&str>])], + spans: &[Span<'_>], + duration_type: DataType, ) { let mut fields = vec![ Field::new( @@ -745,7 +774,7 @@ mod tests { Field::new(SPAN_KIND_COLUMN, DataType::Utf8, false), Field::new(SPAN_STATUS_CODE_COLUMN, DataType::Utf8, false), Field::new(SERVICE_NAME_COLUMN, DataType::Utf8, false), - Field::new(DURATION_NANO_COLUMN, DataType::UInt64, false), + Field::new(DURATION_NANO_COLUMN, duration_type.clone(), false), ]; for (column, _) in extra { fields.push(Field::new(*column, DataType::Utf8, true)); @@ -774,9 +803,14 @@ mod tests { Arc::new(StringArray::from( spans.iter().map(|s| s.6).collect::>(), )), - Arc::new(UInt64Array::from( - spans.iter().map(|s| s.7).collect::>(), - )), + match duration_type { + DataType::Int64 => Arc::new(Int64Array::from( + spans.iter().map(|s| s.7 as i64).collect::>(), + )) as ArrayRef, + _ => Arc::new(UInt64Array::from( + spans.iter().map(|s| s.7).collect::>(), + )), + }, ]; for (_, values) in extra { columns.push(Arc::new(StringArray::from(values.to_vec()))); @@ -843,6 +877,7 @@ mod tests { time_index: TRACE_TIMESTAMP_COLUMN.to_string(), entity_type: "service".to_string(), id_columns: id_columns.iter().map(|s| s.to_string()).collect(), + id_qualifier: None, descriptive_columns: vec![], scope_columns: vec![], } @@ -972,7 +1007,7 @@ mod tests { ), ], ); - register_trace_table( + register_typed_trace_table( &ctx, "trace_b", &[], @@ -989,6 +1024,7 @@ mod tests { 1_500_000_000, ), ], + DataType::Int64, ); let a = ctx.table("trace_a").await.unwrap(); let b = ctx.table("trace_b").await.unwrap(); @@ -1294,6 +1330,7 @@ mod tests { time_index: "ts".to_string(), entity_type: entity_type.to_string(), id_columns: id_columns.iter().map(|s| s.to_string()).collect(), + id_qualifier: None, descriptive_columns: vec![], scope_columns: vec![], } @@ -1544,16 +1581,10 @@ mod tests { let total: usize = batches.iter().map(|b| b.num_rows()).sum(); assert_eq!(total, 1); let batch = &batches[0]; - // Composite service identity renders the same sorted `k=v` form the - // registry emits, so edges land on registry entity ids. - assert_eq!( - strings(batch, 5), - vec!["service_name=frontend,service_namespace=ns1"] - ); - assert_eq!( - strings(batch, 7), - vec!["service_name=cart,service_namespace=ns1"] - ); + // Composite service identity renders exactly as the registry emits it, + // so edges land on registry entity ids. + assert_eq!(strings(batch, 5), vec!["frontend,ns1"]); + assert_eq!(strings(batch, 7), vec!["cart,ns1"]); } fn build_relationships_plan_for_test( diff --git a/src/query/src/analyze.rs b/src/query/src/analyze.rs index ec0abf25ba6..2ab3ec5176e 100644 --- a/src/query/src/analyze.rs +++ b/src/query/src/analyze.rs @@ -46,6 +46,17 @@ const STAGE: &str = "stage"; const NODE: &str = "node"; const PLAN: &str = "plan"; +/// Fixed output schema of [`DistAnalyzeExec`], for `Describe` handlers: +/// execution rewrites the plan in `optimize_physical_plan`, so this schema +/// differs from the logical `Analyze` plan's. +pub fn dist_analyze_output_schema() -> SchemaRef { + SchemaRef::new(Schema::new(vec![ + Field::new(STAGE, DataType::UInt32, true), + Field::new(NODE, DataType::UInt32, true), + Field::new(PLAN, DataType::Utf8, true), + ])) +} + #[derive(Debug)] pub struct DistAnalyzeExec { input: Arc, @@ -58,11 +69,7 @@ pub struct DistAnalyzeExec { impl DistAnalyzeExec { /// Create a new DistAnalyzeExec pub fn new(input: Arc, verbose: bool, format: AnalyzeFormat) -> Self { - let schema = SchemaRef::new(Schema::new(vec![ - Field::new(STAGE, DataType::UInt32, true), - Field::new(NODE, DataType::UInt32, true), - Field::new(PLAN, DataType::Utf8, true), - ])); + let schema = dist_analyze_output_schema(); let properties = Arc::new(Self::compute_properties(&input, schema.clone())); Self { input, diff --git a/src/query/src/datafusion.rs b/src/query/src/datafusion.rs index 754199a78fc..f9d87782a03 100644 --- a/src/query/src/datafusion.rs +++ b/src/query/src/datafusion.rs @@ -30,7 +30,7 @@ use common_function::function::FunctionContext; use common_function::function_factory::ScalarFunctionFactory; use common_query::{Output, OutputData, OutputMeta}; use common_recordbatch::adapter::{RecordBatchStreamAdapter, RegionQueryStatCounters}; -use common_recordbatch::{EmptyRecordBatchStream, SendableRecordBatchStream}; +use common_recordbatch::{EmptyRecordBatchStream, RecordBatch, SendableRecordBatchStream}; use common_telemetry::tracing; use datafusion::catalog::TableFunction; use datafusion::dataframe::DataFrame; @@ -44,6 +44,7 @@ use datafusion_expr::{ use datatypes::prelude::VectorRef; use datatypes::schema::Schema; use futures_util::StreamExt; +use futures_util::future::try_join; use session::context::QueryContextRef; use snafu::{OptionExt, ResultExt, ensure}; use sqlparser::ast::AnalyzeFormat; @@ -80,6 +81,29 @@ pub const QUERY_PARALLELISM_HINT: &str = "query_parallelism"; /// Whether to fallback to the original plan when failed to push down. pub const QUERY_FALLBACK_HINT: &str = "query_fallback"; +// An unbounded queue keeps draining source RPCs while mutation RPCs on a shared +// HTTP/2 connection are pending, trading bounded memory for request liveness. +async fn forward_record_batches( + mut stream: SendableRecordBatchStream, + batch_tx: tokio::sync::mpsc::UnboundedSender>, +) -> Result<()> { + while let Some(batch) = stream.next().await { + match batch.context(CreateRecordBatchSnafu) { + Ok(batch) => { + if batch_tx.send(Ok(batch)).is_err() { + break; + } + tokio::task::yield_now().await; + } + Err(error) => { + let _ = batch_tx.send(Err(error)); + break; + } + } + } + Ok(()) +} + fn query_load_region_id(plan: &Arc) -> Option { let mut region_id = None; let mut stack = vec![plan.clone()]; @@ -205,7 +229,7 @@ impl DatafusionQueryEngine { let Output { data, meta } = self .exec_query_plan((*dml.input).clone(), query_ctx.clone()) .await?; - let mut stream = match data { + let stream = match data { OutputData::RecordBatches(batches) => batches.as_stream(), OutputData::Stream(stream) => stream, _ => unreachable!(), @@ -214,30 +238,35 @@ impl DatafusionQueryEngine { let mut affected_rows = 0; let mut insert_cost = 0; - while let Some(batch) = stream.next().await { - let batch = batch.context(CreateRecordBatchSnafu)?; - let column_vectors = batch - .column_vectors(&table_name.to_string(), table.schema()) - .map_err(BoxedError::new) - .context(QueryExecutionSnafu)?; - - match dml.op { - WriteOp::Insert(_) => { - // We ignore the insert op. - let output = self - .insert(&table_name, column_vectors, query_ctx.clone()) - .await?; - let (rows, cost) = output.extract_rows_and_cost(); - affected_rows += rows; - insert_cost += cost; - } - WriteOp::Delete => { - affected_rows += self - .delete(&table_name, &table, column_vectors, query_ctx.clone()) - .await?; - } - _ => unreachable!("guarded by the 'ensure!' at the beginning"), + match dml.op { + WriteOp::Insert(_) => { + let (batch_tx, batch_rx) = tokio::sync::mpsc::unbounded_channel(); + let producer = forward_record_batches(stream, batch_tx); + let consumer = self.consume_insert_record_batches( + batch_rx, + &table_name, + table.schema(), + query_ctx.clone(), + ); + let ((), (rows, cost)) = try_join(producer, consumer).await?; + affected_rows += rows; + insert_cost += cost; } + WriteOp::Delete => { + // Keep DELETE on the same producer/consumer schedule as INSERT so the source + // stream can continue draining while mutation RPCs are pending. + let (batch_tx, batch_rx) = tokio::sync::mpsc::unbounded_channel(); + let producer = forward_record_batches(stream, batch_tx); + let consumer = self.consume_delete_record_batches( + batch_rx, + &table_name, + &table, + query_ctx.clone(), + ); + let (rows, ()) = try_join(consumer, producer).await?; + affected_rows += rows; + } + _ => unreachable!("guarded by the 'ensure!' at the beginning"), } Ok(Output::new( OutputData::AffectedRows(affected_rows), @@ -245,6 +274,53 @@ impl DatafusionQueryEngine { )) } + async fn consume_insert_record_batches( + &self, + mut batch_rx: tokio::sync::mpsc::UnboundedReceiver>, + table_name: &ResolvedTableReference, + table_schema: Arc, + query_ctx: QueryContextRef, + ) -> Result<(usize, usize)> { + let mut affected_rows = 0; + let mut insert_cost = 0; + while let Some(batch) = batch_rx.recv().await { + let batch = batch?; + let column_vectors = batch + .column_vectors(&table_name.to_string(), table_schema.clone()) + .map_err(BoxedError::new) + .context(QueryExecutionSnafu)?; + // We ignore the insert op. + let output = self + .insert(table_name, column_vectors, query_ctx.clone()) + .await?; + let (rows, cost) = output.extract_rows_and_cost(); + affected_rows += rows; + insert_cost += cost; + } + Ok((affected_rows, insert_cost)) + } + + async fn consume_delete_record_batches( + &self, + mut batch_rx: tokio::sync::mpsc::UnboundedReceiver>, + table_name: &ResolvedTableReference, + table: &TableRef, + query_ctx: QueryContextRef, + ) -> Result { + let mut affected_rows = 0; + while let Some(batch) = batch_rx.recv().await { + let batch = batch?; + let column_vectors = batch + .column_vectors(&table_name.to_string(), table.schema()) + .map_err(BoxedError::new) + .context(QueryExecutionSnafu)?; + affected_rows += self + .delete(table_name, table, column_vectors, query_ctx.clone()) + .await?; + } + Ok(affected_rows) + } + #[tracing::instrument(skip_all)] async fn delete( &self, diff --git a/src/query/src/lib.rs b/src/query/src/lib.rs index 68d8ff3a9eb..4b2a2d5264d 100644 --- a/src/query/src/lib.rs +++ b/src/query/src/lib.rs @@ -45,7 +45,7 @@ pub(crate) mod test_util; #[cfg(test)] mod tests; -pub use crate::analyze::analyze_plan_metrics_to_json_value; +pub use crate::analyze::{analyze_plan_metrics_to_json_value, dist_analyze_output_schema}; pub use crate::datafusion::DfContextProviderAdapter; pub use crate::query_engine::{ QueryEngine, QueryEngineContext, QueryEngineFactory, QueryEngineRef, diff --git a/src/query/src/optimizer.rs b/src/query/src/optimizer.rs index 8827d48ed63..a6b36e6fdae 100644 --- a/src/query/src/optimizer.rs +++ b/src/query/src/optimizer.rs @@ -17,6 +17,7 @@ pub mod constant_term; pub mod count_nest_aggr; pub mod count_wildcard; pub mod global_limit; +pub(crate) mod insert_assignment; pub(crate) mod json_type_concretize; pub mod parallelize_scan; pub mod pass_distribution; diff --git a/src/query/src/optimizer/insert_assignment.rs b/src/query/src/optimizer/insert_assignment.rs new file mode 100644 index 00000000000..ad77736ce0a --- /dev/null +++ b/src/query/src/optimizer/insert_assignment.rs @@ -0,0 +1,307 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; + +use common_time::Timezone; +use datafusion::config::ConfigOptions; +use datafusion_common::{DFSchemaRef, Result, ScalarValue}; +use datafusion_expr::expr::{Alias, Cast}; +use datafusion_expr::{Distinct, Expr, ExprSchemable, LogicalPlan, Projection, Values}; +use datafusion_optimizer::analyzer::AnalyzerRule; +use datafusion_optimizer::analyzer::type_coercion::TypeCoercion; +use datatypes::arrow::datatypes::{DataType, TimeUnit}; +use session::context::QueryContextRef; + +use crate::optimizer::type_conversion::cast_string_to_timestamp; + +/// Interprets strings assigned to timestamp columns at an `INSERT` boundary +/// using the session timezone. `plan` is the assignment projection under a +/// `WriteOp::Insert`. +/// +/// DataFusion plans `INSERT` as a projection casting each source column to its +/// target column type, and that cast reads a naive string as UTC. Arrow does +/// apply a timezone when the cast target carries one, so the assignment is +/// routed through `Timestamp(unit, Some(tz))` and back. Stripping the timezone +/// afterwards is value-preserving — arrow only shifts values in the opposite +/// direction. +/// +/// The source query is left untouched. Reinterpreting a value where it is +/// *produced* would change what the source query means: pushing the conversion +/// below a `UNION`'s `DISTINCT`, for instance, moves the dedup key from the raw +/// strings to parsed instants and silently drops rows. +/// +/// # Why `TypeCoercion` runs here +/// +/// The rewrite reads source types, and those are only settled once a `UNION`'s +/// branch types have been reconciled: before coercion a union carries its loose +/// schema (the first branch's types), so a mixed +/// `SELECT 'string' UNION ALL SELECT CAST(.. AS TIMESTAMP)` still looks like a +/// string. Retargeting that cast would leave `Timestamp(None) -> +/// Timestamp(Some(tz))` behind once coercion retypes the union — the one +/// direction in which arrow shifts the value instead of relabelling it. +/// +/// Coercing here rather than deferring to the analyzer is forced by where an +/// INSERT is still identifiable: `exec_dml_statement` strips the `Dml` node and +/// executes its input, so by the time the analyzer runs, an assignment +/// projection is indistinguishable from any other projection. +/// +/// Explicit casts stay out of this: the SQL layer turns a user's +/// `CAST(x AS TIMESTAMP)` into an `arrow_cast` call, which only becomes an +/// `Expr::Cast` in the optimizer's `SimplifyExpressions`. Assignment casts are +/// therefore the only `Expr::Cast` reaching a timestamp column here. +/// +/// # Reach +/// +/// The emitted cast only carries its timezone where the expression is evaluated +/// on this node. Substrait drops the timezone name when a plan is pushed down — +/// it encodes any zoned timestamp as `PrecisionTimestampTz` and decodes it back +/// as UTC — so a source reading from a table falls back to UTC, the behaviour it +/// had before this rule existed. Sources that never leave this node (literals, +/// `VALUES`, and `UNION`s of them) keep the session timezone, and those are what +/// an INSERT's timestamp assignment is in practice. +pub(crate) fn rewrite_insert_assignments( + plan: LogicalPlan, + query_ctx: &QueryContextRef, + config: &ConfigOptions, +) -> Result { + let Some(timezone) = session_timezone(query_ctx) else { + return Ok(plan); + }; + + let plan = TypeCoercion::new().analyze(plan, config)?; + rewrite_assignment(plan, &timezone) +} + +/// Session timezone, in both forms the rewrite needs. +struct SessionTimezone { + /// Parses literals, matching the plain `INSERT ... VALUES` path. + parsed: Timezone, + /// Names the intermediate arrow cast target. + name: Arc, +} + +fn session_timezone(query_ctx: &QueryContextRef) -> Option { + let parsed = query_ctx.timezone(); + + // A UTC session already gets UTC semantics from the plain assignment cast. + if parsed.is_utc() { + return None; + } + + Some(SessionTimezone { + name: Arc::from(parsed.to_string()), + parsed, + }) +} + +fn rewrite_assignment(plan: LogicalPlan, timezone: &SessionTimezone) -> Result { + let LogicalPlan::Projection(assignment) = plan else { + return Ok(plan); + }; + + let mut exprs = assignment.expr.clone(); + let mut changed = false; + for expr in &mut exprs { + changed |= retarget_assignment_cast( + expr, + assignment.input.schema(), + Some(assignment.input.as_ref()), + timezone, + )?; + } + + // The planner types `VALUES` against the target table, so the assignment + // cast lands inside the `Values` rows instead of on the projection above. + let mut input = assignment.input.clone(); + if let LogicalPlan::Values(values) = assignment.input.as_ref() + && let Some(rewritten) = rewrite_values(values, timezone)? + { + input = Arc::new(LogicalPlan::Values(rewritten)); + changed = true; + } + + if !changed { + return Ok(LogicalPlan::Projection(assignment)); + } + Projection::try_new(exprs, input).map(LogicalPlan::Projection) +} + +fn rewrite_values(values: &Values, timezone: &SessionTimezone) -> Result> { + let mut rewritten = values.clone(); + let mut changed = false; + for row in &mut rewritten.values { + for expr in row.iter_mut() { + changed |= retarget_assignment_cast(expr, &values.schema, None, timezone)?; + } + } + + Ok(changed.then_some(rewritten)) +} + +/// Reinterprets one assignment cast, returning whether it was rewritten. +/// +/// `source_plan` is the projection's input, used to resolve a literal behind a +/// column reference; `Values` rows carry their expression inline and pass `None`. +fn retarget_assignment_cast( + expr: &mut Expr, + schema: &DFSchemaRef, + source_plan: Option<&LogicalPlan>, + timezone: &SessionTimezone, +) -> Result { + let expr = unalias_mut(expr); + let Expr::Cast(Cast { + expr: source, + data_type: DataType::Timestamp(unit, None), + }) = expr + else { + return Ok(false); + }; + let unit = *unit; + + if !matches!( + source.get_type(schema)?, + DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View + ) { + return Ok(false); + } + + // Fold literals with the same parser the plain `INSERT ... VALUES` path + // uses, so a given string means the same thing however it reaches a column. + // The parsers disagree on ambiguous local times: this one resolves them, + // arrow rejects them. + let folded = source_literal(source.as_ref(), source_plan) + .and_then(|literal| convert_literal(&literal, unit, &timezone.parsed)); + if let Some(folded) = folded { + *expr = folded; + return Ok(true); + } + + let source = source.as_ref().clone(); + *expr = Expr::Cast(Cast::new( + Box::new(Expr::Cast(Cast::new( + Box::new(source), + DataType::Timestamp(unit, Some(timezone.name.clone())), + ))), + DataType::Timestamp(unit, None), + )); + Ok(true) +} + +fn source_literal(source: &Expr, source_plan: Option<&LogicalPlan>) -> Option { + match source { + Expr::Literal(value, _) => Some(value.clone()), + Expr::Column(column) => { + let plan = source_plan?; + let index = plan.schema().maybe_index_of_column(column)?; + lineage_literal(plan, index).cloned() + } + _ => None, + } +} + +/// Resolves a literal when every row carries the same value at `output_idx`. +/// +/// Read-only: the literal is folded into the assignment above, so nodes that +/// drop, reorder or deduplicate rows can be traversed — none of them changes +/// the value a surviving row carries, and folding above them leaves their keys +/// on the original strings. +fn lineage_literal(plan: &LogicalPlan, output_idx: usize) -> Option<&ScalarValue> { + if output_idx >= plan.schema().fields().len() { + return None; + } + + match plan { + LogicalPlan::Projection(projection) => match unalias(&projection.expr[output_idx]) { + Expr::Literal(value, _) => Some(value), + Expr::Column(column) => { + let input_idx = projection.input.schema().maybe_index_of_column(column)?; + lineage_literal(projection.input.as_ref(), input_idx) + } + _ => None, + }, + LogicalPlan::Filter(_) + | LogicalPlan::Sort(_) + | LogicalPlan::Limit(_) + | LogicalPlan::SubqueryAlias(_) + | LogicalPlan::Distinct(Distinct::All(_)) => { + let inputs = plan.inputs(); + let [input] = inputs.as_slice() else { + return None; + }; + lineage_literal(input, output_idx) + } + _ => None, + } +} + +fn convert_literal(value: &ScalarValue, unit: TimeUnit, timezone: &Timezone) -> Option { + let ScalarValue::Utf8(Some(value)) = value else { + return None; + }; + cast_string_to_timestamp(value, &DataType::Timestamp(unit, None), Some(timezone)) + .ok() + .filter(|value| !value.is_null()) + .map(|value| Expr::Literal(value, None)) +} + +fn unalias(expr: &Expr) -> &Expr { + match expr { + Expr::Alias(Alias { expr, .. }) => unalias(expr), + expr => expr, + } +} + +fn unalias_mut(expr: &mut Expr) -> &mut Expr { + match expr { + Expr::Alias(Alias { expr, .. }) => unalias_mut(expr), + expr => expr, + } +} + +#[cfg(test)] +mod tests { + use datafusion_common::DFSchema; + use datafusion_expr::expr::Placeholder; + + use super::*; + + fn shanghai() -> SessionTimezone { + let parsed = Timezone::from_tz_string("Asia/Shanghai").unwrap(); + SessionTimezone { + name: Arc::from(parsed.to_string()), + parsed, + } + } + + /// A prepared `INSERT ... VALUES (?)` arrives here as a cast over an untyped + /// placeholder, which must survive for parameter substitution. + #[test] + fn test_untyped_placeholder_assignment_is_left_alone() { + let schema = Arc::new(DFSchema::empty()); + let mut expr = Expr::Cast(Cast::new( + Box::new(Expr::Placeholder(Placeholder::new_with_field( + "$1".to_string(), + None, + ))), + DataType::Timestamp(TimeUnit::Millisecond, None), + )); + let original = expr.clone(); + + let changed = retarget_assignment_cast(&mut expr, &schema, None, &shanghai()).unwrap(); + + assert!(!changed); + assert_eq!(expr, original); + } +} diff --git a/src/query/src/optimizer/type_conversion.rs b/src/query/src/optimizer/type_conversion.rs index 3941fedfe8a..40bb43e791e 100644 --- a/src/query/src/optimizer/type_conversion.rs +++ b/src/query/src/optimizer/type_conversion.rs @@ -12,8 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -mod insert_assignment; - use std::sync::Arc; use common_time::Timezone; @@ -31,7 +29,7 @@ use session::context::QueryContextRef; use crate::QueryEngineContext; use crate::optimizer::ExtensionAnalyzerRule; -use crate::optimizer::type_conversion::insert_assignment::rewrite_insert_assignments; +use crate::optimizer::insert_assignment::rewrite_insert_assignments; use crate::plan::ExtractExpr; /// TypeConversionRule converts some literal values in logical plan to other types according @@ -46,7 +44,7 @@ impl ExtensionAnalyzerRule for TypeConversionRule { &self, plan: LogicalPlan, ctx: &QueryEngineContext, - _config: &ConfigOptions, + config: &ConfigOptions, ) -> Result { plan.transform_up_with_subqueries(|plan| match plan { LogicalPlan::Filter(filter) => { @@ -129,7 +127,8 @@ impl ExtensionAnalyzerRule for TypeConversionRule { LogicalPlan::Dml(mut dml) if matches!(dml.op, WriteOp::Insert(_)) => { dml.input = Arc::new(rewrite_insert_assignments( dml.input.as_ref().clone(), - ctx.query_ctx(), + &ctx.query_ctx(), + config, )?); Ok(Transformed::yes(LogicalPlan::Dml(dml))) } @@ -324,7 +323,7 @@ fn timestamp_to_timestamp_ms_expr(val: i64, unit: TimeUnit) -> Expr { ) } -fn cast_string_to_timestamp( +pub(crate) fn cast_string_to_timestamp( string: &str, target_type: &DataType, timezone: Option<&Timezone>, diff --git a/src/query/src/optimizer/type_conversion/insert_assignment.rs b/src/query/src/optimizer/type_conversion/insert_assignment.rs deleted file mode 100644 index 11ab7e7829e..00000000000 --- a/src/query/src/optimizer/type_conversion/insert_assignment.rs +++ /dev/null @@ -1,336 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -use std::sync::Arc; - -use datafusion_common::{Column, Result, ScalarValue}; -use datafusion_expr::expr::{Alias, Cast}; -use datafusion_expr::{Distinct, Expr, ExprSchemable, LogicalPlan, Projection, Union, Values}; -use datatypes::arrow::datatypes::DataType; -use session::context::QueryContextRef; - -use crate::optimizer::type_conversion::cast_string_to_timestamp; -use crate::plan::ExtractExpr; - -/// Rewrites string literals that feed timestamp columns at an INSERT boundary. -/// Constants are folded at the assignment to avoid changing source types; -/// `VALUES` and `UNION` inputs are rewritten per row or branch. Explicit casts -/// stay on DataFusion's existing path. -pub(super) fn rewrite_insert_assignments( - plan: LogicalPlan, - query_ctx: QueryContextRef, -) -> Result { - let LogicalPlan::Projection(assignment) = plan else { - return Ok(plan); - }; - - let converter = InsertAssignmentConverter { query_ctx }; - let mut exprs = assignment.expr.clone(); - let mut input = assignment.input.as_ref().clone(); - let mut changed = false; - for (output_idx, expr) in assignment.expr.iter().enumerate() { - let target_type = assignment.schema.field(output_idx).data_type(); - if !matches!(target_type, DataType::Timestamp(_, _)) { - continue; - } - - let Some(column) = assignment_input_column(expr) else { - continue; - }; - let Some(input_idx) = input.schema().maybe_index_of_column(column) else { - continue; - }; - - let literal = lineage_literal(&input, input_idx).cloned(); - if let Some(literal) = literal - && let Some(folded) = converter.convert_literal(&literal, target_type) - { - let (qualifier, field) = assignment.schema.qualified_field(output_idx); - exprs[output_idx] = folded.alias_qualified(qualifier.cloned(), field.name()); - changed = true; - continue; - } - - // A hand-built DML plan (e.g. from flow) may share one source column - // between targets; rewriting it in place would retype every reader. - if assignment - .expr - .iter() - .enumerate() - .any(|(idx, other)| idx != output_idx && other.column_refs().contains(column)) - { - continue; - } - if let Some(rewritten) = converter.rewrite_output_column(&input, input_idx, target_type)? { - input = rewritten; - changed = true; - } - } - - if !changed { - return Ok(LogicalPlan::Projection(assignment)); - } - Projection::try_new(exprs, Arc::new(input)).map(LogicalPlan::Projection) -} - -/// Resolves a literal when every row carries the same value at `output_idx`. -fn lineage_literal(plan: &LogicalPlan, output_idx: usize) -> Option<&ScalarValue> { - if output_idx >= plan.schema().fields().len() { - return None; - } - - match plan { - LogicalPlan::Projection(projection) => match unalias(&projection.expr[output_idx]) { - Expr::Literal(value, _) => Some(value), - Expr::Column(column) => { - let input_idx = projection.input.schema().maybe_index_of_column(column)?; - lineage_literal(projection.input.as_ref(), input_idx) - } - _ => None, - }, - // These nodes drop, reorder or deduplicate rows without touching the - // value a surviving row carries, and their schema stays positional. - LogicalPlan::Filter(_) - | LogicalPlan::Sort(_) - | LogicalPlan::Limit(_) - | LogicalPlan::SubqueryAlias(_) - | LogicalPlan::Distinct(Distinct::All(_)) => { - let inputs = plan.inputs(); - let [input] = inputs.as_slice() else { - return None; - }; - lineage_literal(input, output_idx) - } - _ => None, - } -} - -struct InsertAssignmentConverter { - query_ctx: QueryContextRef, -} - -impl InsertAssignmentConverter { - fn rewrite_output_column( - &self, - plan: &LogicalPlan, - output_idx: usize, - target_type: &DataType, - ) -> Result> { - if output_idx >= plan.schema().fields().len() { - return Ok(None); - } - - match plan { - // DataFusion pushes INSERT assignment casts into Values, so inspect - // them even though the Values schema already has the target type. - LogicalPlan::Values(values) => self.rewrite_values(values, output_idx, target_type), - // Once the source query has produced the target type, its casts and - // coercions belong to the source query rather than INSERT assignment. - _ if plan.schema().field(output_idx).data_type() == target_type => Ok(None), - LogicalPlan::Projection(projection) => { - self.rewrite_projection(projection, output_idx, target_type) - } - LogicalPlan::Union(union) => self.rewrite_union(union, output_idx, target_type), - // Filter and Sort are excluded here: their predicates and keys read - // the retyped column, which would change the source query. Constants - // still reach the assignment through `lineage_literal`. Distinct::All - // holds no expressions; allowing it shifts dedup keys from raw - // strings to parsed instants. - LogicalPlan::Limit(_) - | LogicalPlan::SubqueryAlias(_) - | LogicalPlan::Distinct(Distinct::All(_)) => { - self.rewrite_passthrough(plan, output_idx, target_type) - } - _ => Ok(None), - } - } - - fn rewrite_projection( - &self, - projection: &Projection, - output_idx: usize, - target_type: &DataType, - ) -> Result> { - match unalias(&projection.expr[output_idx]) { - Expr::Literal(value, _) => { - let Some(converted) = self.convert_literal(value, target_type) else { - return Ok(None); - }; - let (qualifier, field) = projection.schema.qualified_field(output_idx); - let mut exprs = projection.expr.clone(); - exprs[output_idx] = converted.alias_qualified(qualifier.cloned(), field.name()); - - Projection::try_new(exprs, projection.input.clone()) - .map(LogicalPlan::Projection) - .map(Some) - } - Expr::Column(column) => { - // Retyping the input column affects every output column reading - // it, so only follow lineage with a single consumer. - if projection - .expr - .iter() - .enumerate() - .any(|(idx, other)| idx != output_idx && other.column_refs().contains(column)) - { - return Ok(None); - } - - let input = projection.input.as_ref(); - let Some(input_idx) = input.schema().maybe_index_of_column(column) else { - return Ok(None); - }; - let Some(rewritten) = self.rewrite_output_column(input, input_idx, target_type)? - else { - return Ok(None); - }; - - Projection::try_new(projection.expr.clone(), Arc::new(rewritten)) - .map(LogicalPlan::Projection) - .map(Some) - } - _ => Ok(None), - } - } - - fn rewrite_values( - &self, - values: &Values, - output_idx: usize, - target_type: &DataType, - ) -> Result> { - let mut rewritten = values.clone(); - let mut changed = false; - for row in &mut rewritten.values { - let Some(expr) = row.get_mut(output_idx) else { - return Ok(None); - }; - - if let Expr::Cast(Cast { - expr: inner, - data_type, - }) = expr - && data_type == target_type - && let Expr::Literal(value, _) = inner.as_ref() - && let Some(value) = self.convert_literal(value, target_type) - { - *expr = value; - changed = true; - continue; - } - - if expr.get_type(values.schema.as_ref()).ok().as_ref() != Some(target_type) { - return Ok(None); - } - } - - Ok(changed.then_some(LogicalPlan::Values(rewritten))) - } - - fn rewrite_union( - &self, - union: &Union, - output_idx: usize, - target_type: &DataType, - ) -> Result> { - let mut inputs = Vec::with_capacity(union.inputs.len()); - // A union has one shared schema, so rewrite it only when the same output - // column is an INSERT-assigned literal in every branch. - for input in &union.inputs { - let Some(rewritten) = self.rewrite_output_column(input, output_idx, target_type)? - else { - return Ok(None); - }; - inputs.push(Arc::new(rewritten)); - } - - // Untouched columns may still have mismatched branch types before - // TypeCoercion runs, so rebuild loosely like the SQL planner does. - Union::try_new_with_loose_types(inputs) - .map(LogicalPlan::Union) - .map(Some) - } - - fn rewrite_passthrough( - &self, - plan: &LogicalPlan, - output_idx: usize, - target_type: &DataType, - ) -> Result> { - let inputs = plan.inputs(); - let [input] = inputs.as_slice() else { - return Ok(None); - }; - let Some(rewritten_input) = self.rewrite_output_column(input, output_idx, target_type)? - else { - return Ok(None); - }; - - plan.with_new_exprs(plan.expressions_consider_join(), vec![rewritten_input]) - .map(Some) - } - - fn convert_literal(&self, value: &ScalarValue, target_type: &DataType) -> Option { - let ScalarValue::Utf8(Some(value)) = value else { - return None; - }; - cast_string_to_timestamp(value, target_type, Some(&self.query_ctx.timezone())) - .ok() - .filter(|value| !value.is_null()) - .map(|value| Expr::Literal(value, None)) - } -} - -fn assignment_input_column(expr: &Expr) -> Option<&Column> { - let expr = match unalias(expr) { - Expr::Cast(Cast { expr, .. }) => expr.as_ref(), - expr => expr, - }; - let Expr::Column(column) = expr else { - return None; - }; - Some(column) -} - -fn unalias(expr: &Expr) -> &Expr { - match expr { - Expr::Alias(Alias { expr, .. }) => unalias(expr), - expr => expr, - } -} - -#[cfg(test)] -mod tests { - use datafusion_common::ScalarValue; - use datafusion_common::arrow::datatypes::TimeUnit; - use session::context::QueryContext; - - use super::*; - - #[test] - fn test_convert_literal_falls_back_for_unsupported_literal() { - let converter = InsertAssignmentConverter { - query_ctx: QueryContext::arc(), - }; - let target_type = DataType::Timestamp(TimeUnit::Nanosecond, None); - - for literal in ["1970-01-01", "-8-01-01 00:00:01.5"] { - assert_eq!( - converter - .convert_literal(&ScalarValue::Utf8(Some(literal.to_string())), &target_type,), - None - ); - } - } -} diff --git a/src/query/src/planner.rs b/src/query/src/planner.rs index a7d4a84cb28..c3871a96afa 100644 --- a/src/query/src/planner.rs +++ b/src/query/src/planner.rs @@ -866,6 +866,30 @@ mod tests { engine.planner().plan(&stmt, query_ctx).await.unwrap() } + /// Plans `sql` and runs the DataFusion analyzer, which is where + /// `InsertAssignmentRule` sits. Planning alone stops short of it, so these + /// assertions would not see the assignment rewrite at all. + async fn analyze_insert( + engine: &QueryEngineRef, + sql: &str, + query_ctx: &QueryContextRef, + ) -> String { + let stmt = QueryLanguageParser::parse_sql(sql, query_ctx).unwrap(); + let plan = engine + .planner() + .plan(&stmt, query_ctx.clone()) + .await + .unwrap(); + let context = engine.engine_context(query_ctx.clone()); + let state = context.state(); + state + .analyzer() + .execute_and_check(plan, state.config_options(), |_, _| {}) + .unwrap() + .display_indent_schema() + .to_string() + } + #[tokio::test] async fn test_insert_timestamp_literals_use_query_timezone() { let query_ctx = Arc::new( @@ -898,20 +922,6 @@ mod tests { ) AS source", &[1_785_902_400_001_i64][..], ), - ( - "INSERT INTO timestamps (ts, st) \ - SELECT '2026-08-06 12:00:00.001', now() \ - UNION ALL \ - SELECT '2026-08-07 12:00:00.001', now()", - &[1_785_988_800_001_i64, 1_786_075_200_001_i64][..], - ), - ( - "INSERT INTO timestamps (ts, st) \ - SELECT '2026-08-16 12:00:00.001', now() \ - UNION \ - SELECT '2026-08-17 12:00:00.001', now()", - &[1_786_852_800_001_i64, 1_786_939_200_001_i64][..], - ), ( "INSERT INTO timestamps (ts, st) \ SELECT '2026-08-18 12:00:00.001', max(st) \ @@ -940,14 +950,7 @@ mod tests { &[1_786_766_400_001_i64][..], ), ] { - let stmt = QueryLanguageParser::parse_sql(sql, &query_ctx).unwrap(); - let plan = engine - .planner() - .plan(&stmt, query_ctx.clone()) - .await - .unwrap() - .display_indent() - .to_string(); + let plan = analyze_insert(&engine, sql, &query_ctx).await; for expected_timestamp in expected_timestamps { assert!( @@ -968,19 +971,19 @@ mod tests { let engine = create_timestamp_test_engine().await; let sql = "INSERT INTO timestamps (ts, st) \ VALUES (CAST('2026-08-08 12:00:00.001' AS TIMESTAMP), now())"; - let stmt = QueryLanguageParser::parse_sql(sql, &query_ctx).unwrap(); - let plan = engine - .planner() - .plan(&stmt, query_ctx) - .await - .unwrap() - .display_indent() - .to_string(); + let plan = analyze_insert(&engine, sql, &query_ctx).await; + // An explicit cast reaches the analyzer as an `arrow_cast` call rather + // than an `Expr::Cast`, which is how it stays out of the rewrite. assert!( plan.contains("arrow_cast(Utf8(\"2026-08-08 12:00:00.001\")"), "{plan}" ); + // 12:00:00.001 read as Shanghai local time; the source query keeps UTC. + assert!( + !plan.contains("TimestampMillisecond(1786104000001, None)"), + "{plan}" + ); } #[tokio::test] @@ -1017,14 +1020,7 @@ mod tests { ][..], ), ] { - let stmt = QueryLanguageParser::parse_sql(sql, &query_ctx).unwrap(); - let plan = engine - .planner() - .plan(&stmt, query_ctx.clone()) - .await - .unwrap() - .display_indent_schema() - .to_string(); + let plan = analyze_insert(&engine, sql, &query_ctx).await; for expected in expected { assert!(plan.contains(expected), "{plan}"); @@ -1033,30 +1029,37 @@ mod tests { } #[tokio::test] - async fn test_insert_union_tolerates_uncoerced_untouched_column() { + async fn test_insert_union_converts_via_assignment_cast() { let query_ctx = Arc::new( QueryContextBuilder::default() .timezone(Timezone::from_tz_string("Asia/Shanghai").unwrap()) .build(), ); let engine = create_timestamp_test_engine().await; - // `st` stays Timestamp vs Null across branches until TypeCoercion runs. - let sql = "INSERT INTO timestamps (ts, st) \ - SELECT '2026-08-06 12:00:00.001', now() \ - UNION ALL \ - SELECT '2026-08-07 12:00:00.001', NULL"; - let stmt = QueryLanguageParser::parse_sql(sql, &query_ctx).unwrap(); - let plan = engine - .planner() - .plan(&stmt, query_ctx) - .await - .unwrap() - .display_indent() - .to_string(); + // Branches disagree, so the conversion stays as a cast on the + // assignment instead of folding. One cast covers every branch, which is + // why a NULL branch no longer cancels the conversion for the column and + // why UNION's dedup keys stay on the original strings. + for sql in [ + "INSERT INTO timestamps (ts, st) \ + SELECT '2026-08-06 12:00:00.001', now() \ + UNION ALL \ + SELECT '2026-08-07 12:00:00.001', NULL", + "INSERT INTO timestamps (ts, st) \ + SELECT '2026-08-16 12:00:00.001', now() \ + UNION \ + SELECT '2026-08-17 12:00:00.001', now()", + ] { + let plan = analyze_insert(&engine, sql, &query_ctx).await; - for expected_timestamp in [1_785_988_800_001_i64, 1_786_075_200_001_i64] { assert!( - plan.contains(&format!("TimestampMillisecond({expected_timestamp}, None)")), + plan.contains("AS Timestamp(ms, \"Asia/Shanghai\")"), + "{plan}" + ); + // The branches themselves are untouched. + assert!( + plan.contains("Utf8(\"2026-08-07 12:00:00.001\")") + || plan.contains("Utf8(\"2026-08-17 12:00:00.001\")"), "{plan}" ); } @@ -1074,14 +1077,8 @@ mod tests { SELECT '2026-08-10 12:00:00.001', now() \ UNION ALL \ SELECT CAST('2026-08-11 12:00:00.001' AS TIMESTAMP), now()"; - let stmt = QueryLanguageParser::parse_sql(sql, &query_ctx).unwrap(); - let plan = engine - .planner() - .plan(&stmt, query_ctx) - .await - .unwrap() - .display_indent_schema() - .to_string(); + let plan = analyze_insert(&engine, sql, &query_ctx).await; + assert!( !plan.contains("TimestampMillisecond(1786334400001, None)"), "{plan}" @@ -1090,6 +1087,11 @@ mod tests { plan.contains("arrow_cast(Utf8(\"2026-08-11 12:00:00.001\")"), "{plan}" ); + // TypeCoercion has already settled this union to timestamp, so the + // assignment has nothing left to reinterpret. Retargeting the cast here + // would leave a Timestamp(None) -> Timestamp(Some(tz)) step behind, + // which shifts the value instead of relabelling it. + assert!(!plan.contains("Asia/Shanghai"), "{plan}"); } #[tokio::test] diff --git a/src/query/src/sql.rs b/src/query/src/sql.rs index b2d979c6664..93e74ca33ec 100644 --- a/src/query/src/sql.rs +++ b/src/query/src/sql.rs @@ -41,6 +41,7 @@ use common_recordbatch::RecordBatches; use common_recordbatch::adapter::RecordBatchStreamAdapter; use common_time::Timestamp; use common_time::timezone::get_timezone; +use datafusion::dataframe::DataFrame; use datafusion::prelude::SessionContext; use datafusion_expr::{Expr, SortExpr, col, lit}; use datatypes::prelude::*; @@ -109,7 +110,7 @@ const INDEX_KEY_NAME_COLUMN: &str = "Key_name"; const INDEX_SEQ_IN_INDEX_COLUMN: &str = "Seq_in_index"; const INDEX_COLUMN_NAME_COLUMN: &str = "Column_name"; -static DESCRIBE_TABLE_OUTPUT_SCHEMA: Lazy> = Lazy::new(|| { +pub static DESCRIBE_TABLE_OUTPUT_SCHEMA: Lazy> = Lazy::new(|| { Arc::new(Schema::new(vec![ ColumnSchema::new( COLUMN_NAME_COLUMN, @@ -178,6 +179,18 @@ pub async fn show_databases( catalog_manager: &CatalogManagerRef, query_ctx: QueryContextRef, ) -> Result { + let dataframe = + show_databases_dataframe(&stmt, query_engine, catalog_manager, query_ctx).await?; + dataframe_to_output(dataframe).await +} + +/// Builds the [`DataFrame`] for `SHOW DATABASES` without executing it. +pub async fn show_databases_dataframe( + stmt: &ShowDatabases, + query_engine: &QueryEngineRef, + catalog_manager: &CatalogManagerRef, + query_ctx: QueryContextRef, +) -> Result { let projects = if stmt.full { vec![ (schemata::SCHEMA_NAME, SCHEMAS_COLUMN), @@ -191,7 +204,7 @@ pub async fn show_databases( let like_field = Some(schemata::SCHEMA_NAME); let sort = vec![col(schemata::SCHEMA_NAME).sort(true, true)]; - query_from_information_schema_table( + query_from_information_schema_dataframe( query_engine, catalog_manager, query_ctx, @@ -201,7 +214,7 @@ pub async fn show_databases( filters, like_field, sort, - stmt.kind, + &stmt.kind, ) .await } @@ -249,6 +262,45 @@ async fn query_from_information_schema_table( sort: Vec, kind: ShowKind, ) -> Result { + let dataframe = query_from_information_schema_dataframe( + query_engine, + catalog_manager, + query_ctx, + table_name, + select, + projects, + filters, + like_field, + sort, + &kind, + ) + .await?; + dataframe_to_output(dataframe).await +} + +async fn dataframe_to_output(dataframe: DataFrame) -> Result { + let stream = dataframe.execute_stream().await?; + Ok(Output::new_with_stream(Box::pin( + RecordBatchStreamAdapter::try_new(stream).context(error::CreateRecordBatchSnafu)?, + ))) +} + +/// Builds the [`DataFrame`] for a `SHOW` statement without executing it, +/// so `Describe` handlers can derive the output schema from the same +/// projection the executor uses. +#[allow(clippy::too_many_arguments)] +async fn query_from_information_schema_dataframe( + query_engine: &QueryEngineRef, + catalog_manager: &CatalogManagerRef, + query_ctx: QueryContextRef, + table_name: &str, + select: Vec, + projects: Vec<(&str, &str)>, + filters: Vec, + like_field: Option<&str>, + sort: Vec, + kind: &ShowKind, +) -> Result { let table = catalog_manager .table( query_ctx.current_catalog(), @@ -336,7 +388,7 @@ async fn query_from_information_schema_table( .expect("Must be the datafusion planner"); let filter = planner - .sql_to_expr(filter, dataframe.schema(), false, query_ctx) + .sql_to_expr(filter.clone(), dataframe.schema(), false, query_ctx) .await?; // Apply the `where` clause filters @@ -344,11 +396,7 @@ async fn query_from_information_schema_table( } }; - let stream = dataframe.execute_stream().await?; - - Ok(Output::new_with_stream(Box::pin( - RecordBatchStreamAdapter::try_new(stream).context(error::CreateRecordBatchSnafu)?, - ))) + Ok(dataframe) } /// Execute `SHOW COLUMNS` statement. @@ -358,8 +406,19 @@ pub async fn show_columns( catalog_manager: &CatalogManagerRef, query_ctx: QueryContextRef, ) -> Result { - let schema_name = if let Some(database) = stmt.database { - database + let dataframe = show_columns_dataframe(&stmt, query_engine, catalog_manager, query_ctx).await?; + dataframe_to_output(dataframe).await +} + +/// Builds the [`DataFrame`] for `SHOW COLUMNS` without executing it. +pub async fn show_columns_dataframe( + stmt: &ShowColumns, + query_engine: &QueryEngineRef, + catalog_manager: &CatalogManagerRef, + query_ctx: QueryContextRef, +) -> Result { + let schema_name = if let Some(database) = &stmt.database { + database.clone() } else { query_ctx.current_schema() }; @@ -397,7 +456,7 @@ pub async fn show_columns( let like_field = Some(columns::COLUMN_NAME); let sort = vec![col(columns::COLUMN_NAME).sort(true, true)]; - query_from_information_schema_table( + query_from_information_schema_dataframe( query_engine, catalog_manager, query_ctx, @@ -407,7 +466,7 @@ pub async fn show_columns( filters, like_field, sort, - stmt.kind, + &stmt.kind, ) .await } @@ -419,8 +478,19 @@ pub async fn show_index( catalog_manager: &CatalogManagerRef, query_ctx: QueryContextRef, ) -> Result { - let schema_name = if let Some(database) = stmt.database { - database + let dataframe = show_index_dataframe(&stmt, query_engine, catalog_manager, query_ctx).await?; + dataframe_to_output(dataframe).await +} + +/// Builds the [`DataFrame`] for `SHOW INDEX` without executing it. +pub async fn show_index_dataframe( + stmt: &ShowIndex, + query_engine: &QueryEngineRef, + catalog_manager: &CatalogManagerRef, + query_ctx: QueryContextRef, +) -> Result { + let schema_name = if let Some(database) = &stmt.database { + database.clone() } else { query_ctx.current_schema() }; @@ -471,7 +541,7 @@ pub async fn show_index( col(statistics::SEQ_IN_INDEX).sort(true, true), ]; - query_from_information_schema_table( + query_from_information_schema_dataframe( query_engine, catalog_manager, query_ctx, @@ -481,7 +551,7 @@ pub async fn show_index( filters, like_field, sort, - stmt.kind, + &stmt.kind, ) .await } @@ -493,8 +563,19 @@ pub async fn show_region( catalog_manager: &CatalogManagerRef, query_ctx: QueryContextRef, ) -> Result { - let schema_name = if let Some(database) = stmt.database { - database + let dataframe = show_region_dataframe(&stmt, query_engine, catalog_manager, query_ctx).await?; + dataframe_to_output(dataframe).await +} + +/// Builds the [`DataFrame`] for `SHOW REGION` without executing it. +pub async fn show_region_dataframe( + stmt: &ShowRegion, + query_engine: &QueryEngineRef, + catalog_manager: &CatalogManagerRef, + query_ctx: QueryContextRef, +) -> Result { + let schema_name = if let Some(database) = &stmt.database { + database.clone() } else { query_ctx.current_schema() }; @@ -517,7 +598,7 @@ pub async fn show_region( col(columns::PEER_ID).sort(true, true), ]; - query_from_information_schema_table( + query_from_information_schema_dataframe( query_engine, catalog_manager, query_ctx, @@ -527,7 +608,7 @@ pub async fn show_region( filters, like_field, sort, - stmt.kind, + &stmt.kind, ) .await } @@ -539,8 +620,19 @@ pub async fn show_tables( catalog_manager: &CatalogManagerRef, query_ctx: QueryContextRef, ) -> Result { - let schema_name = if let Some(database) = stmt.database { - database + let dataframe = show_tables_dataframe(&stmt, query_engine, catalog_manager, query_ctx).await?; + dataframe_to_output(dataframe).await +} + +/// Builds the [`DataFrame`] for [`ShowTables`] without executing it. +pub async fn show_tables_dataframe( + stmt: &ShowTables, + query_engine: &QueryEngineRef, + catalog_manager: &CatalogManagerRef, + query_ctx: QueryContextRef, +) -> Result { + let schema_name = if let Some(database) = &stmt.database { + database.clone() } else { query_ctx.current_schema() }; @@ -564,15 +656,16 @@ pub async fn show_tables( // Transform the WHERE clause for backward compatibility: // Replace "Tables" with "Tables_in_{schema}" to support old queries - let kind = match stmt.kind { - ShowKind::Where(mut filter) => { + let rewritten_kind = match &stmt.kind { + ShowKind::Where(filter) => { + let mut filter = filter.clone(); replace_column_in_expr(&mut filter, "Tables", &tables_column); ShowKind::Where(filter) } - other => other, + kind => kind.clone(), }; - query_from_information_schema_table( + query_from_information_schema_dataframe( query_engine, catalog_manager, query_ctx, @@ -582,7 +675,7 @@ pub async fn show_tables( filters, like_field, sort, - kind, + &rewritten_kind, ) .await } @@ -594,8 +687,20 @@ pub async fn show_table_status( catalog_manager: &CatalogManagerRef, query_ctx: QueryContextRef, ) -> Result { - let schema_name = if let Some(database) = stmt.database { - database + let dataframe = + show_table_status_dataframe(&stmt, query_engine, catalog_manager, query_ctx).await?; + dataframe_to_output(dataframe).await +} + +/// Builds the [`DataFrame`] for [`ShowTableStatus`] without executing it. +pub async fn show_table_status_dataframe( + stmt: &ShowTableStatus, + query_engine: &QueryEngineRef, + catalog_manager: &CatalogManagerRef, + query_ctx: QueryContextRef, +) -> Result { + let schema_name = if let Some(database) = &stmt.database { + database.clone() } else { query_ctx.current_schema() }; @@ -629,7 +734,7 @@ pub async fn show_table_status( let like_field = Some(tables::TABLE_NAME); let sort = vec![col(tables::TABLE_NAME).sort(true, true)]; - query_from_information_schema_table( + query_from_information_schema_dataframe( query_engine, catalog_manager, query_ctx, @@ -639,7 +744,7 @@ pub async fn show_table_status( filters, like_field, sort, - stmt.kind, + &stmt.kind, ) .await } @@ -651,6 +756,18 @@ pub async fn show_collations( catalog_manager: &CatalogManagerRef, query_ctx: QueryContextRef, ) -> Result { + let dataframe = + show_collations_dataframe(&kind, query_engine, catalog_manager, query_ctx).await?; + dataframe_to_output(dataframe).await +} + +/// Builds the [`DataFrame`] for `SHOW COLLATION` without executing it. +pub async fn show_collations_dataframe( + kind: &ShowKind, + query_engine: &QueryEngineRef, + catalog_manager: &CatalogManagerRef, + query_ctx: QueryContextRef, +) -> Result { // Refer to https://dev.mysql.com/doc/refman/8.0/en/show-collation.html let projects = vec![ ("collation_name", "Collation"), @@ -665,7 +782,7 @@ pub async fn show_collations( let like_field = Some("collation_name"); let sort = vec![]; - query_from_information_schema_table( + query_from_information_schema_dataframe( query_engine, catalog_manager, query_ctx, @@ -687,6 +804,18 @@ pub async fn show_charsets( catalog_manager: &CatalogManagerRef, query_ctx: QueryContextRef, ) -> Result { + let dataframe = + show_charsets_dataframe(&kind, query_engine, catalog_manager, query_ctx).await?; + dataframe_to_output(dataframe).await +} + +/// Builds the [`DataFrame`] for `SHOW CHARSET` without executing it. +pub async fn show_charsets_dataframe( + kind: &ShowKind, + query_engine: &QueryEngineRef, + catalog_manager: &CatalogManagerRef, + query_ctx: QueryContextRef, +) -> Result { // Refer to https://dev.mysql.com/doc/refman/8.0/en/show-character-set.html let projects = vec![ ("character_set_name", "Charset"), @@ -699,7 +828,7 @@ pub async fn show_charsets( let like_field = Some("character_set_name"); let sort = vec![]; - query_from_information_schema_table( + query_from_information_schema_dataframe( query_engine, catalog_manager, query_ctx, @@ -916,8 +1045,19 @@ pub async fn show_views( catalog_manager: &CatalogManagerRef, query_ctx: QueryContextRef, ) -> Result { - let schema_name = if let Some(database) = stmt.database { - database + let dataframe = show_views_dataframe(&stmt, query_engine, catalog_manager, query_ctx).await?; + dataframe_to_output(dataframe).await +} + +/// Builds the [`DataFrame`] for [`ShowViews`] without executing it. +pub async fn show_views_dataframe( + stmt: &ShowViews, + query_engine: &QueryEngineRef, + catalog_manager: &CatalogManagerRef, + query_ctx: QueryContextRef, +) -> Result { + let schema_name = if let Some(database) = &stmt.database { + database.clone() } else { query_ctx.current_schema() }; @@ -930,7 +1070,7 @@ pub async fn show_views( let like_field = Some(tables::TABLE_NAME); let sort = vec![col(tables::TABLE_NAME).sort(true, true)]; - query_from_information_schema_table( + query_from_information_schema_dataframe( query_engine, catalog_manager, query_ctx, @@ -940,7 +1080,7 @@ pub async fn show_views( filters, like_field, sort, - stmt.kind, + &stmt.kind, ) .await } @@ -952,12 +1092,23 @@ pub async fn show_flows( catalog_manager: &CatalogManagerRef, query_ctx: QueryContextRef, ) -> Result { + let dataframe = show_flows_dataframe(&stmt, query_engine, catalog_manager, query_ctx).await?; + dataframe_to_output(dataframe).await +} + +/// Builds the [`DataFrame`] for [`ShowFlows`] without executing it. +pub async fn show_flows_dataframe( + stmt: &ShowFlows, + query_engine: &QueryEngineRef, + catalog_manager: &CatalogManagerRef, + query_ctx: QueryContextRef, +) -> Result { let projects = vec![(flows::FLOW_NAME, FLOWS_COLUMN)]; let filters = vec![col(flows::TABLE_CATALOG).eq(lit(query_ctx.current_catalog()))]; let like_field = Some(flows::FLOW_NAME); let sort = vec![col(flows::FLOW_NAME).sort(true, true)]; - query_from_information_schema_table( + query_from_information_schema_dataframe( query_engine, catalog_manager, query_ctx, @@ -967,7 +1118,7 @@ pub async fn show_flows( filters, like_field, sort, - stmt.kind, + &stmt.kind, ) .await } @@ -1340,6 +1491,18 @@ pub async fn show_processlist( catalog_manager: &CatalogManagerRef, query_ctx: QueryContextRef, ) -> Result { + let dataframe = + show_processlist_dataframe(&stmt, query_engine, catalog_manager, query_ctx).await?; + dataframe_to_output(dataframe).await +} + +/// Builds the [`DataFrame`] for `SHOW PROCESSLIST` without executing it. +pub async fn show_processlist_dataframe( + stmt: &ShowProcessList, + query_engine: &QueryEngineRef, + catalog_manager: &CatalogManagerRef, + query_ctx: QueryContextRef, +) -> Result { let projects = if stmt.full { vec![ (process_list::ID, "Id"), @@ -1367,17 +1530,17 @@ pub async fn show_processlist( }; let like_field = None; let sort = vec![col("id").sort(true, true)]; - query_from_information_schema_table( + query_from_information_schema_dataframe( query_engine, catalog_manager, query_ctx.clone(), "process_list", vec![], - projects.clone(), + projects, filters, like_field, sort, - ShowKind::All, + &ShowKind::All, ) .await } diff --git a/src/servers/Cargo.toml b/src/servers/Cargo.toml index 7936c96fbda..bb358ea76a6 100644 --- a/src/servers/Cargo.toml +++ b/src/servers/Cargo.toml @@ -167,8 +167,8 @@ servers = { workspace = true, features = ["testing"] } session = { workspace = true, features = ["testing"] } table.workspace = true tempfile = "3.0.0" -tokio-postgres = "0.7" -tokio-postgres-rustls = "0.12" +tokio-postgres.workspace = true +tokio-postgres-rustls = "0.14" [target.'cfg(unix)'.dev-dependencies] pprof = { version = "0.14", features = ["criterion", "flamegraph"] } diff --git a/src/servers/dashboard/VERSION b/src/servers/dashboard/VERSION index b68b26e6f78..0822cdb1e68 100644 --- a/src/servers/dashboard/VERSION +++ b/src/servers/dashboard/VERSION @@ -1 +1 @@ -v0.13.13 +v0.13.14 diff --git a/src/servers/src/grpc/flight.rs b/src/servers/src/grpc/flight.rs index e6f35e17b29..0b7a9dd5704 100644 --- a/src/servers/src/grpc/flight.rs +++ b/src/servers/src/grpc/flight.rs @@ -30,8 +30,7 @@ use bytes::{self, Bytes}; use common_error::ext::ErrorExt; use common_grpc::flight::do_put::{DoPutMetadata, DoPutResponse}; use common_grpc::flight::{ - FLOW_EXTENSIONS_METADATA_KEY, FlightDecoder, FlightEncoder, FlightMessage, - SNAPSHOT_SEQS_METADATA_KEY, + FLOW_EXTENSIONS_METADATA_KEY, FlightDecoder, FlightMessage, SNAPSHOT_SEQS_METADATA_KEY, }; use common_memory_manager::MemoryGuard; use common_query::{Output, OutputData}; @@ -46,14 +45,16 @@ use prost::Message; use query::metrics::terminal_recordbatch_metrics_from_plan_if_requested; use query::options::FlowQueryExtensions; use session::context::{Channel, QueryContextRef}; -use snafu::{IntoError, ResultExt, ensure}; +use snafu::{IntoError, OptionExt, ResultExt, ensure}; use table::table_name::TableName; use tokio::sync::mpsc; use tokio_stream::wrappers::ReceiverStream; use tonic::{Request, Response, Status, Streaming}; -use crate::error::{InvalidParameterSnafu, Result, ToJsonSnafu}; -pub use crate::grpc::flight::stream::{FlightRecordBatchSource, FlightRecordBatchStream}; +use crate::error::{InvalidParameterSnafu, InvalidQuerySnafu, Result, ToJsonSnafu}; +pub use crate::grpc::flight::stream::{ + FlightRecordBatchSource, FlightRecordBatchStream, FlightRecordBatchStreamInput, +}; use crate::grpc::greptime_handler::{ GreptimeRequestHandler, create_query_context, get_request_type, }; @@ -221,9 +222,12 @@ impl FlightCraft for GreptimeRequestHandler { ); let flight_compression = self.flight_compression; async { - let output = self - .handle_request_with_query_ctx(request, query_ctx.clone()) + let query = request.request.context(InvalidQuerySnafu { + reason: "Expecting non-empty GreptimeRequest.", + })?; + self.authenticate_request_with_query_ctx(request.header.as_ref(), &query_ctx) .await?; + let output = self.handle_request_with_query_ctx(query, query_ctx.clone()); let stream = to_flight_data_stream( output, TracingContext::from_current_span(), @@ -573,32 +577,36 @@ fn extract_json_metadata( Ok(Some(parsed)) } -fn to_flight_data_stream( - output: Output, +fn to_flight_data_stream( + output: F, tracing_context: TracingContext, flight_compression: FlightCompression, query_ctx: QueryContextRef, should_emit_terminal_metrics: bool, -) -> TonicStream { +) -> TonicStream +where + F: std::future::Future> + Send + 'static, +{ + let initializer = async move { + let output = output.await.map_err(Status::from)?; + output_to_flight_record_batch_source(output, should_emit_terminal_metrics) + }; + let stream = FlightRecordBatchStream::new( + FlightRecordBatchStreamInput::initializer(initializer), + tracing_context, + flight_compression, + query_ctx, + ); + Box::pin(stream) as _ +} + +fn output_to_flight_record_batch_source( + output: Output, + should_emit_terminal_metrics: bool, +) -> TonicResult { match output.data { - OutputData::Stream(stream) => { - let stream = FlightRecordBatchStream::new( - FlightRecordBatchSource::RecordBatches(stream), - tracing_context, - flight_compression, - query_ctx, - ); - Box::pin(stream) as _ - } - OutputData::RecordBatches(x) => { - let stream = FlightRecordBatchStream::new( - FlightRecordBatchSource::RecordBatches(x.as_stream()), - tracing_context, - flight_compression, - query_ctx, - ); - Box::pin(stream) as _ - } + OutputData::Stream(stream) => Ok(FlightRecordBatchSource::RecordBatches(stream)), + OutputData::RecordBatches(x) => Ok(FlightRecordBatchSource::RecordBatches(x.as_stream())), OutputData::AffectedRows(rows) => { let terminal_metrics = match terminal_recordbatch_metrics_from_plan_if_requested( output.meta.plan, @@ -607,20 +615,17 @@ fn to_flight_data_stream( Some(metrics) => match serde_json::to_string(&metrics) { Ok(metrics) => Some(metrics), Err(e) => { - let stream = tokio_stream::once(Err(Status::internal(format!( + return Err(Status::internal(format!( "Failed to serialize terminal metrics: {e}" - )))); - return Box::pin(stream) as _; + ))); } }, None => None, }; - let affected_rows = FlightEncoder::default().encode(FlightMessage::AffectedRows { + Ok(FlightRecordBatchSource::AffectedRows { rows, metrics: terminal_metrics, - }); - let stream = tokio_stream::iter(affected_rows.into_iter().map(Ok)); - Box::pin(stream) as _ + }) } } } diff --git a/src/servers/src/grpc/flight/stream.rs b/src/servers/src/grpc/flight/stream.rs index 319a797d9ca..b740f12d02e 100644 --- a/src/servers/src/grpc/flight/stream.rs +++ b/src/servers/src/grpc/flight/stream.rs @@ -41,27 +41,34 @@ use crate::error; use crate::grpc::FlightCompression; use crate::grpc::flight::TonicResult; -pub type FlightRecordBatchStreamInitializer = Pin< - Box< - dyn Future> - + Send - + 'static, - >, ->; - pub enum FlightRecordBatchSource { RecordBatches(SendableRecordBatchStream), - Initializer(FlightRecordBatchStreamInitializer), + AffectedRows { + rows: usize, + metrics: Option, + }, } -impl FlightRecordBatchSource { - pub fn initializer(initializer: F) -> Self - where - F: Future> - + Send - + 'static, - { - Self::Initializer(Box::pin(initializer)) +/// Determines whether a Flight result is ready now or initialized asynchronously. +pub enum FlightRecordBatchStreamInput>> +{ + Ready(FlightRecordBatchSource), + Initializer(F), +} + +impl FlightRecordBatchStreamInput { + /// Creates an input from a source that is already available. + pub fn ready(source: FlightRecordBatchSource) -> Self { + Self::Ready(source) + } +} + +impl FlightRecordBatchStreamInput { + /// Creates an input that obtains its source asynchronously. + /// + /// Errors from the initializer are returned through the Flight response stream. + pub fn initializer(initializer: F) -> Self { + Self::Initializer(initializer) } } @@ -160,31 +167,32 @@ impl FlightRecordBatchStream { Self::send_metrics(tx, metrics, metrics_str).await } - pub fn new( - source: FlightRecordBatchSource, + pub fn new( + input: FlightRecordBatchStreamInput, tracing_context: TracingContext, compression: FlightCompression, query_ctx: QueryContextRef, - ) -> Self { - let should_send_partial_metrics = query_ctx.explain_verbose(); - let can_send_metrics_before_batch = query_ctx.explain_verbose() - && query_ctx.live_analyze_metrics_enabled() - && query_ctx - .remote_query_id() - .zip(query_ctx.extension(SUPPORT_FLIGHT_METRICS_BEFORE_BATCH_EXTENSION_KEY)) - .is_some_and(|(remote_query_id, capability)| capability == remote_query_id); + ) -> Self + where + F: Future> + Send + 'static, + { let (mut tx, rx) = mpsc::channel::>(1); - let source_type = match &source { - FlightRecordBatchSource::RecordBatches(_) => "record_batches", - FlightRecordBatchSource::Initializer(_) => "initializer", + let source_type = match &input { + FlightRecordBatchStreamInput::Ready(FlightRecordBatchSource::RecordBatches(_)) => { + "record_batches" + } + FlightRecordBatchStreamInput::Ready(FlightRecordBatchSource::AffectedRows { + .. + }) => "affected_rows", + FlightRecordBatchStreamInput::Initializer(_) => "initializer", }; let initializer_tracing_context = tracing_context.clone(); let join_handle = common_runtime::spawn_global( async move { - let recordbatches = async move { - match source { - FlightRecordBatchSource::RecordBatches(recordbatches) => Ok(recordbatches), - FlightRecordBatchSource::Initializer(initializer) => initializer.await, + let source = async move { + match input { + FlightRecordBatchStreamInput::Ready(source) => Ok(source), + FlightRecordBatchStreamInput::Initializer(initializer) => initializer.await, } } .trace( @@ -193,8 +201,20 @@ impl FlightRecordBatchStream { ) .await; - match recordbatches { - Ok(recordbatches) => { + match source { + Ok(FlightRecordBatchSource::RecordBatches(recordbatches)) => { + let should_send_partial_metrics = query_ctx.explain_verbose(); + let can_send_metrics_before_batch = + query_ctx.explain_verbose() + && query_ctx.live_analyze_metrics_enabled() + && query_ctx + .remote_query_id() + .zip(query_ctx.extension( + SUPPORT_FLIGHT_METRICS_BEFORE_BATCH_EXTENSION_KEY, + )) + .is_some_and(|(remote_query_id, capability)| { + capability == remote_query_id + }); Self::flight_data_stream( recordbatches, tx, @@ -203,6 +223,11 @@ impl FlightRecordBatchStream { ) .await; } + Ok(FlightRecordBatchSource::AffectedRows { rows, metrics }) => { + let _ = tx + .send(Ok(FlightMessage::AffectedRows { rows, metrics })) + .await; + } Err(status) => { let _ = tx.send(Err(status)).await; } @@ -485,7 +510,9 @@ mod test { .unwrap() .as_stream(); let mut stream = FlightRecordBatchStream::new( - FlightRecordBatchSource::RecordBatches(recordbatches), + FlightRecordBatchStreamInput::ready(FlightRecordBatchSource::RecordBatches( + recordbatches, + )), TracingContext::default(), FlightCompression::default(), QueryContext::arc(), @@ -519,10 +546,34 @@ mod test { } } + #[tokio::test] + async fn test_flight_record_batch_stream_encodes_affected_rows() { + let mut stream = FlightRecordBatchStream::new( + FlightRecordBatchStreamInput::ready(FlightRecordBatchSource::AffectedRows { + rows: 42, + metrics: Some(r#"{"region_watermarks":[]}"#.to_string()), + }), + TracingContext::default(), + FlightCompression::default(), + QueryContext::arc(), + ); + + let data = stream.next().await.unwrap().unwrap(); + let message = FlightDecoder::default().try_decode(&data).unwrap().unwrap(); + assert!(matches!( + message, + FlightMessage::AffectedRows { + rows: 42, + metrics: Some(_), + } + )); + assert!(stream.next().await.is_none()); + } + #[tokio::test] async fn test_flight_record_batch_stream_forwards_initializer_error() { let mut stream = FlightRecordBatchStream::new( - FlightRecordBatchSource::initializer(async { + FlightRecordBatchStreamInput::initializer(async { Err(tonic::Status::unavailable( "remote read initialization failed", )) @@ -552,9 +603,12 @@ mod test { metrics, }); let query_ctx = query_context_with_live_metrics_and_matching_capability(); - query_ctx.set_explain_verbose(true); + let initializer_query_ctx = query_ctx.clone(); let mut stream = FlightRecordBatchStream::new( - FlightRecordBatchSource::RecordBatches(recordbatches), + FlightRecordBatchStreamInput::initializer(async move { + initializer_query_ctx.set_explain_verbose(true); + Ok(FlightRecordBatchSource::RecordBatches(recordbatches)) + }), TracingContext::default(), FlightCompression::default(), query_ctx, @@ -609,7 +663,9 @@ mod test { let query_ctx = query_context_with_live_metrics_and_matching_capability(); query_ctx.set_explain_verbose(true); let mut stream = FlightRecordBatchStream::new( - FlightRecordBatchSource::RecordBatches(recordbatches), + FlightRecordBatchStreamInput::ready(FlightRecordBatchSource::RecordBatches( + recordbatches, + )), TracingContext::default(), FlightCompression::default(), query_ctx, @@ -663,7 +719,9 @@ mod test { let query_ctx = query_context_with_matching_capability(); query_ctx.set_explain_verbose(true); let mut stream = FlightRecordBatchStream::new( - FlightRecordBatchSource::RecordBatches(recordbatches), + FlightRecordBatchStreamInput::ready(FlightRecordBatchSource::RecordBatches( + recordbatches, + )), TracingContext::default(), FlightCompression::default(), query_ctx, @@ -706,7 +764,9 @@ mod test { let query_ctx = Arc::new(query_ctx); query_ctx.set_explain_verbose(true); let mut stream = FlightRecordBatchStream::new( - FlightRecordBatchSource::RecordBatches(recordbatches), + FlightRecordBatchStreamInput::ready(FlightRecordBatchSource::RecordBatches( + recordbatches, + )), TracingContext::default(), FlightCompression::default(), query_ctx, @@ -745,7 +805,9 @@ mod test { }); let query_ctx = query_context_with_matching_capability(); let mut stream = FlightRecordBatchStream::new( - FlightRecordBatchSource::RecordBatches(recordbatches), + FlightRecordBatchStreamInput::ready(FlightRecordBatchSource::RecordBatches( + recordbatches, + )), TracingContext::default(), FlightCompression::default(), query_ctx, diff --git a/src/servers/src/grpc/greptime_handler.rs b/src/servers/src/grpc/greptime_handler.rs index ad7779435b8..98786d30adc 100644 --- a/src/servers/src/grpc/greptime_handler.rs +++ b/src/servers/src/grpc/greptime_handler.rs @@ -15,11 +15,13 @@ //! Handler for Greptime Database service. It's implemented by frontend. use std::collections::HashMap; +use std::future::Future; use std::str::FromStr; use std::sync::{Arc, RwLock}; use std::time::Instant; use api::helper::request_type; +use api::v1::greptime_request::Request as QueryRequest; use api::v1::{GreptimeRequest, RequestHeader}; use auth::UserProviderRef; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; @@ -79,30 +81,39 @@ impl GreptimeRequestHandler { ) -> Result { let header = request.header.as_ref(); let query_ctx = create_query_context(Channel::Grpc, header, hints, HashMap::new())?; - self.handle_request_with_query_ctx(request, query_ctx).await - } - - pub(crate) async fn handle_request_with_query_ctx( - &self, - request: GreptimeRequest, - query_ctx: QueryContextRef, - ) -> Result { let query = request.request.context(InvalidQuerySnafu { reason: "Expecting non-empty GreptimeRequest.", })?; + self.authenticate_request_with_query_ctx(request.header.as_ref(), &query_ctx) + .await?; + self.handle_request_with_query_ctx(query, query_ctx).await + } - let header = request.header.as_ref(); - let user_info = context_auth::auth(self.user_provider.clone(), header, &query_ctx).await?; + pub(crate) async fn authenticate_request_with_query_ctx( + &self, + header: Option<&RequestHeader>, + query_ctx: &QueryContextRef, + ) -> Result<()> { + let user_info = context_auth::auth(self.user_provider.clone(), header, query_ctx).await?; query_ctx.set_current_user(user_info); + Ok(()) + } + pub(crate) fn handle_request_with_query_ctx( + &self, + query: QueryRequest, + query_ctx: QueryContextRef, + ) -> impl Future> + Send + 'static { let handler = self.handler.clone(); + let runtime = self.runtime.clone(); let request_type = request_type(&query).to_string(); let db = query_ctx.get_db_string(); let timer = RequestTimer::new(db.clone(), request_type); let tracing_context = TracingContext::from_current_span(); - let result_future = async move { - handler + async move { + let result_future = async move { + handler .do_query(query, query_ctx) .trace(tracing_context.attach(tracing::info_span!( "GreptimeRequestHandler::handle_request_runtime" @@ -118,26 +129,27 @@ impl GreptimeRequestHandler { } e }) - }; + }; - match &self.runtime { - Some(runtime) => { - // Executes requests in another runtime to - // 1. prevent the execution from being cancelled unexpected by Tonic runtime; - // - Refer to our blog for the rational behind it: - // https://www.greptime.com/blogs/2023-01-12-hidden-control-flow.html - // - Obtaining a `JoinHandle` to get the panic message (if there's any). - // From its docs, `JoinHandle` is cancel safe. The task keeps running even it's handle been dropped. - // 2. avoid the handler blocks the gRPC runtime incidentally. - runtime - .spawn(result_future) - .await - .context(JoinTaskSnafu) - .inspect_err(|e| { - timer.record(e.status_code()); - })? + match runtime { + Some(runtime) => { + // Executes requests in another runtime to + // 1. prevent the execution from being cancelled unexpected by Tonic runtime; + // - Refer to our blog for the rational behind it: + // https://www.greptime.com/blogs/2023-01-12-hidden-control-flow.html + // - Obtaining a `JoinHandle` to get the panic message (if there's any). + // From its docs, `JoinHandle` is cancel safe. The task keeps running even it's handle been dropped. + // 2. avoid the handler blocks the gRPC runtime incidentally. + runtime + .spawn(result_future) + .await + .context(JoinTaskSnafu) + .inspect_err(|e| { + timer.record(e.status_code()); + })? + } + None => result_future.await, } - None => result_future.await, } } diff --git a/src/servers/src/http/otlp.rs b/src/servers/src/http/otlp.rs index d04022cdf88..3ecaba74a4e 100644 --- a/src/servers/src/http/otlp.rs +++ b/src/servers/src/http/otlp.rs @@ -118,8 +118,9 @@ pub async fn metrics( promote_scope_attrs: http_opts.promote_scope_attrs, with_metric_engine, experimental_enable_exponential_histogram, - // set is_legacy later + // set by the frontend from its config is_legacy: false, + resource_info: false, metric_type: MetricType::Init, metric_translation_strategy: http_opts.metric_translation_strategy, })); @@ -128,7 +129,7 @@ pub async fn metrics( handler.metrics(request, query_ctx).await.map(|outcome| { if outcome.accepted_data_points == 0 && outcome.rejected_data_points > 0 { OtlpMetricsResponse::Failure(outcome) - } else if outcome.rejected_data_points > 0 { + } else if outcome.rejected_data_points > 0 || outcome.error_message.is_some() { OtlpMetricsResponse::PartialSuccess(outcome) } else { OtlpMetricsResponse::FullSuccess(outcome) diff --git a/src/servers/src/mysql/handler.rs b/src/servers/src/mysql/handler.rs index b0eb33f8fd5..30579d8846a 100644 --- a/src/servers/src/mysql/handler.rs +++ b/src/servers/src/mysql/handler.rs @@ -25,6 +25,7 @@ use common_catalog::parse_optional_catalog_and_schema_from_db_string; use common_error::ext::ErrorExt; use common_query::Output; use common_telemetry::{debug, error, tracing, warn}; +use common_time::Timezone; use datafusion_common::ParamValues; use datafusion_expr::LogicalPlan; use datatypes::prelude::ConcreteDataType; @@ -289,12 +290,13 @@ impl MysqlInstanceShim { .fail(); } + let timezone = query_ctx.timezone(); let replaced_plan = match params { Params::ProtocolParams(params) => { - replace_params_with_values(&plan, param_types, ¶ms) + replace_params_with_values(&plan, param_types, ¶ms, &timezone) } Params::CliParams(params) => { - replace_params_with_exprs(&plan, param_types, ¶ms) + replace_params_with_exprs(&plan, param_types, ¶ms, &timezone) } }?; @@ -807,6 +809,7 @@ fn replace_params_with_values( plan: &LogicalPlan, param_types: HashMap>, params: &[ParamValue], + timezone: &Timezone, ) -> Result { debug_assert_eq!(param_types.len(), params.len()); @@ -824,7 +827,7 @@ fn replace_params_with_values( for (i, param) in params.iter().enumerate() { if let Some(Some(t)) = param_types.get(&format_placeholder(i + 1)) { - let value = helper::convert_value(param, t)?; + let value = helper::convert_value(param, t, timezone)?; values.push(value.into()); } @@ -839,6 +842,7 @@ fn replace_params_with_exprs( plan: &LogicalPlan, param_types: HashMap>, params: &[sql::ast::Expr], + timezone: &Timezone, ) -> Result { debug_assert_eq!(param_types.len(), params.len()); @@ -853,7 +857,7 @@ fn replace_params_with_exprs( for (i, param) in params.iter().enumerate() { if let Some(Some(t)) = param_types.get(&format_placeholder(i + 1)) { - let value = helper::convert_expr_to_scalar_value(param, t)?; + let value = helper::convert_expr_to_scalar_value(param, t, timezone)?; values.push(value.into()); } diff --git a/src/servers/src/mysql/helper.rs b/src/servers/src/mysql/helper.rs index a4f289cdfb1..f9ab5325925 100644 --- a/src/servers/src/mysql/helper.rs +++ b/src/servers/src/mysql/helper.rs @@ -15,10 +15,11 @@ use std::ops::ControlFlow; use std::time::Duration; -use chrono::NaiveDate; +use chrono::{NaiveDate, NaiveDateTime}; use common_query::prelude::ScalarValue; use common_sql::convert::sql_value_to_value; -use common_time::{Date, Timestamp}; +use common_time::timestamp::TimeUnit; +use common_time::{Date, Timestamp, Timezone}; use datatypes::prelude::{ConcreteDataType, DataType}; use datatypes::schema::ColumnSchema; use datatypes::types::TimestampType; @@ -137,11 +138,15 @@ where /// Convert [`ParamValue`] into [`Value`] according to param type. /// It will try it's best to do type conversions if possible -pub fn convert_value(param: &ParamValue, t: &ConcreteDataType) -> Result { +pub fn convert_value( + param: &ParamValue, + t: &ConcreteDataType, + timezone: &Timezone, +) -> Result { if let ConcreteDataType::Dictionary(dictionary) = t { return Ok(ScalarValue::Dictionary( Box::new(dictionary.key_type().as_arrow_type()), - Box::new(convert_value(param, dictionary.value_type())?), + Box::new(convert_value(param, dictionary.value_type(), timezone)?), )); } @@ -221,7 +226,9 @@ pub fn convert_value(param: &ParamValue, t: &ConcreteDataType) -> Result Ok(ScalarValue::Binary(Some(b.to_vec()))), - ConcreteDataType::Timestamp(ts_type) => convert_bytes_to_timestamp(b, ts_type), + ConcreteDataType::Timestamp(ts_type) => { + convert_bytes_to_timestamp(b, ts_type, timezone) + } ConcreteDataType::Date(_) => convert_bytes_to_date(b), _ => error::PreparedStmtTypeMismatchSnafu { expected: t, @@ -233,29 +240,26 @@ pub fn convert_value(param: &ParamValue, t: &ConcreteDataType) -> Result { - let timestamp_millis = to_naive_datetime(param.value) - .map_err(|e| { + ValueInner::Datetime(_) => match t { + ConcreteDataType::Timestamp(_) => { + let datetime = to_naive_datetime(param.value).map_err(|e| { error::MysqlValueConversionSnafu { err_msg: e.to_string(), } .build() - })? - .and_utc() - .timestamp_millis(); - - match t { - ConcreteDataType::Timestamp(_) => Ok(ScalarValue::TimestampMillisecond( + })?; + let timestamp_millis = datetime_to_timestamp_millis(datetime, timezone)?; + Ok(ScalarValue::TimestampMillisecond( Some(timestamp_millis), None, - )), - _ => error::PreparedStmtTypeMismatchSnafu { - expected: t, - actual: param.coltype, - } - .fail(), + )) } - } + _ => error::PreparedStmtTypeMismatchSnafu { + expected: t, + actual: param.coltype, + } + .fail(), + }, ValueInner::Time(_) => Ok(ScalarValue::Time64Nanosecond(Some( Duration::from(param.value).as_millis() as i64, ))), @@ -264,13 +268,18 @@ pub fn convert_value(param: &ParamValue, t: &ConcreteDataType) -> Result Result { +pub fn convert_expr_to_scalar_value( + param: &Expr, + t: &ConcreteDataType, + timezone: &Timezone, +) -> Result { if let ConcreteDataType::Dictionary(dictionary) = t { return Ok(ScalarValue::Dictionary( Box::new(dictionary.key_type().as_arrow_type()), Box::new(convert_expr_to_scalar_value( param, dictionary.value_type(), + timezone, )?), )); } @@ -278,7 +287,7 @@ pub fn convert_expr_to_scalar_value(param: &Expr, t: &ConcreteDataType) -> Resul let column_schema = ColumnSchema::new("", t.clone(), true); match param { Expr::Value(v) => { - let v = sql_value_to_value(&column_schema, &v.value, None, None, true); + let v = sql_value_to_value(&column_schema, &v.value, Some(timezone), None, true); match v { Ok(v) => v .try_to_scalar_value(t) @@ -290,7 +299,7 @@ pub fn convert_expr_to_scalar_value(param: &Expr, t: &ConcreteDataType) -> Resul } } Expr::UnaryOp { op, expr } if let Expr::Value(v) = &**expr => { - let v = sql_value_to_value(&column_schema, &v.value, None, Some(*op), true); + let v = sql_value_to_value(&column_schema, &v.value, Some(timezone), Some(*op), true); match v { Ok(v) => v .try_to_scalar_value(t) @@ -308,8 +317,32 @@ pub fn convert_expr_to_scalar_value(param: &Expr, t: &ConcreteDataType) -> Resul } } -fn convert_bytes_to_timestamp(bytes: &[u8], ts_type: &TimestampType) -> Result { - let ts = Timestamp::from_str_utc(&String::from_utf8_lossy(bytes)) +/// Interprets a timezone-less datetime in the given timezone and returns the +/// corresponding epoch timestamp in milliseconds. +fn datetime_to_timestamp_millis(datetime: NaiveDateTime, timezone: &Timezone) -> Result { + let ts = Timestamp::from_naive_datetime(datetime, timezone) + .map_err(|e| { + error::MysqlValueConversionSnafu { + err_msg: e.to_string(), + } + .build() + })? + .convert_to(TimeUnit::Millisecond) + .ok_or_else(|| { + error::MysqlValueConversionSnafu { + err_msg: "Overflow when converting datetime to milliseconds".to_string(), + } + .build() + })?; + Ok(ts.value()) +} + +fn convert_bytes_to_timestamp( + bytes: &[u8], + ts_type: &TimestampType, + timezone: &Timezone, +) -> Result { + let ts = Timestamp::from_str(&String::from_utf8_lossy(bytes), Some(timezone)) .map_err(|e| { error::MysqlValueConversionSnafu { err_msg: e.to_string(), @@ -446,19 +479,20 @@ mod tests { #[test] fn test_convert_expr_to_scalar_value() { + let utc = Timezone::from_tz_string("UTC").unwrap(); let expr = Expr::Value(ValueExpr::Number("123".to_string(), false).into()); let t = ConcreteDataType::int32_datatype(); - let v = convert_expr_to_scalar_value(&expr, &t).unwrap(); + let v = convert_expr_to_scalar_value(&expr, &t, &utc).unwrap(); assert_eq!(ScalarValue::Int32(Some(123)), v); let expr = Expr::Value(ValueExpr::Number("123.456789".to_string(), false).into()); let t = ConcreteDataType::float64_datatype(); - let v = convert_expr_to_scalar_value(&expr, &t).unwrap(); + let v = convert_expr_to_scalar_value(&expr, &t, &utc).unwrap(); assert_eq!(ScalarValue::Float64(Some(123.456789)), v); let expr = Expr::Value(ValueExpr::SingleQuotedString("2001-01-02".to_string()).into()); let t = ConcreteDataType::date_datatype(); - let v = convert_expr_to_scalar_value(&expr, &t).unwrap(); + let v = convert_expr_to_scalar_value(&expr, &t, &utc).unwrap(); let scalar_v = ScalarValue::Utf8(Some("2001-01-02".to_string())) .cast_to(&arrow_schema::DataType::Date32) .unwrap(); @@ -467,7 +501,7 @@ mod tests { let expr = Expr::Value(ValueExpr::SingleQuotedString("2001-01-02 03:04:05".to_string()).into()); let t = ConcreteDataType::timestamp_microsecond_datatype(); - let v = convert_expr_to_scalar_value(&expr, &t).unwrap(); + let v = convert_expr_to_scalar_value(&expr, &t, &utc).unwrap(); let scalar_v = ScalarValue::Utf8(Some("2001-01-02 03:04:05".to_string())) .cast_to(&arrow_schema::DataType::Timestamp( arrow_schema::TimeUnit::Microsecond, @@ -478,14 +512,14 @@ mod tests { let expr = Expr::Value(ValueExpr::SingleQuotedString("hello".to_string()).into()); let t = ConcreteDataType::string_datatype(); - let v = convert_expr_to_scalar_value(&expr, &t).unwrap(); + let v = convert_expr_to_scalar_value(&expr, &t, &utc).unwrap(); assert_eq!(ScalarValue::Utf8(Some("hello".to_string())), v); let t = ConcreteDataType::dictionary_datatype( ConcreteDataType::uint32_datatype(), ConcreteDataType::string_datatype(), ); - let v = convert_expr_to_scalar_value(&expr, &t).unwrap(); + let v = convert_expr_to_scalar_value(&expr, &t, &utc).unwrap(); assert_eq!( ScalarValue::Dictionary( Box::new(arrow_schema::DataType::UInt32), @@ -496,7 +530,7 @@ mod tests { let expr = Expr::Value(ValueExpr::Null.into()); let t = ConcreteDataType::time_microsecond_datatype(); - let v = convert_expr_to_scalar_value(&expr, &t).unwrap(); + let v = convert_expr_to_scalar_value(&expr, &t, &utc).unwrap(); assert_eq!(ScalarValue::Time64Microsecond(None), v); } @@ -577,12 +611,81 @@ mod tests { ), ]; + let utc = Timezone::from_tz_string("UTC").unwrap(); for (input, ts_type, expected) in test_cases { - let result = convert_bytes_to_timestamp(input.as_bytes(), &ts_type).unwrap(); + let result = convert_bytes_to_timestamp(input.as_bytes(), &ts_type, &utc).unwrap(); assert_eq!(result, expected); } } + fn utc_millis(year: i32, month: u32, day: u32, hour: u32, minute: u32, second: u32) -> i64 { + NaiveDate::from_ymd_opt(year, month, day) + .unwrap() + .and_hms_opt(hour, minute, second) + .unwrap() + .and_utc() + .timestamp_millis() + } + + #[test] + fn test_datetime_to_timestamp_millis() { + let datetime = NaiveDate::from_ymd_opt(2026, 8, 13) + .unwrap() + .and_hms_opt(8, 0, 0) + .unwrap(); + + let utc = Timezone::from_tz_string("UTC").unwrap(); + assert_eq!( + datetime_to_timestamp_millis(datetime, &utc).unwrap(), + utc_millis(2026, 8, 13, 8, 0, 0) + ); + + let shanghai = Timezone::from_tz_string("Asia/Shanghai").unwrap(); + assert_eq!( + datetime_to_timestamp_millis(datetime, &shanghai).unwrap(), + utc_millis(2026, 8, 13, 0, 0, 0) + ); + + let minus_7 = Timezone::from_tz_string("-07:00").unwrap(); + assert_eq!( + datetime_to_timestamp_millis(datetime, &minus_7).unwrap(), + utc_millis(2026, 8, 13, 15, 0, 0) + ); + + // 2026-03-08 02:30 does not exist in America/New_York (DST gap). + let new_york = Timezone::from_tz_string("America/New_York").unwrap(); + let gap = NaiveDate::from_ymd_opt(2026, 3, 8) + .unwrap() + .and_hms_opt(2, 30, 0) + .unwrap(); + assert!(datetime_to_timestamp_millis(gap, &new_york).is_err()); + + // 2026-11-01 01:30 is ambiguous in America/New_York; picks the first (EDT, UTC-4). + let ambiguous = NaiveDate::from_ymd_opt(2026, 11, 1) + .unwrap() + .and_hms_opt(1, 30, 0) + .unwrap(); + assert_eq!( + datetime_to_timestamp_millis(ambiguous, &new_york).unwrap(), + utc_millis(2026, 11, 1, 5, 30, 0) + ); + } + + #[test] + fn test_convert_bytes_to_timestamp_with_timezone() { + let shanghai = Timezone::from_tz_string("Asia/Shanghai").unwrap(); + let result = convert_bytes_to_timestamp( + "2026-08-13 08:00:00".as_bytes(), + &TimestampType::Millisecond(TimestampMillisecondType), + &shanghai, + ) + .unwrap(); + assert_eq!( + result, + ScalarValue::TimestampMillisecond(Some(utc_millis(2026, 8, 13, 0, 0, 0)), None) + ); + } + #[test] fn test_convert_bytes_to_date() { let test_cases = vec![ diff --git a/src/servers/src/otlp/metrics.rs b/src/servers/src/otlp/metrics.rs index bd0bfb2ce77..d65e117e07c 100644 --- a/src/servers/src/otlp/metrics.rs +++ b/src/servers/src/otlp/metrics.rs @@ -25,6 +25,7 @@ use common_query::native_histogram::{ }; use common_query::prelude::{GREPTIME_COUNT, greptime_timestamp, greptime_value}; use common_query::prometheus::PROMETHEUS_STALE_NAN_BITS; +use common_telemetry::warn; use lazy_static::lazy_static; use otel_arrow_rust::proto::opentelemetry::collector::metrics::v1::ExportMetricsServiceRequest; use otel_arrow_rust::proto::opentelemetry::common::v1::{AnyValue, KeyValue, any_value}; @@ -36,7 +37,7 @@ use table::requests::{ }; use crate::error::{self, Result}; -use crate::otlp::trace::{KEY_SERVICE_INSTANCE_ID, KEY_SERVICE_NAME}; +use crate::otlp::trace::{KEY_SERVICE_INSTANCE_ID, KEY_SERVICE_NAME, KEY_SERVICE_NAMESPACE}; use crate::query_handler::MetricsIngestOutcome; use crate::row_writer::{self, MultiTableData, TableData}; pub use crate::semantic::SemanticIndex; @@ -45,8 +46,11 @@ use crate::semantic::{ METRIC_TYPE_UPDOWN_COUNTER, }; +mod resource_info; mod translator; +pub use resource_info::OTEL_RESOURCE_INFO_TABLE_NAME; +use resource_info::ResourceInfoData; pub use translator::legacy_normalize_otlp_name; pub(crate) use translator::ucum_to_openmetrics_unit; use translator::{translate_label_name, translate_metric_name}; @@ -96,28 +100,43 @@ const MIN_EXPONENTIAL_HISTOGRAM_SCALE: i32 = -4; const MAX_EXPONENTIAL_HISTOGRAM_SCALE: i32 = 8; const MAX_REJECTION_MESSAGE_BYTES: usize = 512; +/// Result of converting one OTLP metrics request. +#[derive(Debug)] +pub struct MetricsConversion { + pub requests: RowInsertRequests, + /// Row count of `requests`; the resource descriptor is not counted. + pub rows: usize, + /// Per-table semantic index for the auto-create path to stamp as table + /// options; covers the descriptor table too. + pub semantic_index: SemanticIndex, + /// The synthesized resource descriptor, written separately from the + /// metric data. See [`resource_info`]. + pub resource_info: Option, + pub outcome: MetricsIngestOutcome, +} + /// Convert OpenTelemetry metrics to GreptimeDB insert requests /// /// See /// /// for data structure of OTLP metrics. -/// -/// Returns `InsertRequests`, total number of rows to ingest, and the per-table -/// semantic index for the auto-create path to stamp as table options. pub fn to_grpc_insert_requests( request: ExportMetricsServiceRequest, metric_ctx: &mut OtlpMetricCtx, -) -> Result<( - RowInsertRequests, - usize, - SemanticIndex, - MetricsIngestOutcome, -)> { +) -> Result { let mut table_writer = MultiTableData::default(); let mut semantic_index = SemanticIndex::default(); let mut outcome = MetricsIngestOutcome::default(); + let mut resource_info = ResourceInfoData::default(); for resource in &request.resource_metrics { + if metric_ctx.resource_info + && !metric_ctx.is_legacy + && let Some(r) = resource.resource.as_ref() + { + resource_info.observe(&r.attributes, resource, metric_ctx); + } + let resource_attrs = resource.resource.as_ref().map(|r| { let mut attrs = r.attributes.clone(); process_resource_attrs(&mut attrs, metric_ctx); @@ -149,8 +168,46 @@ pub fn to_grpc_insert_requests( } let (requests, rows) = table_writer.into_row_insert_requests(); + validate_sample_kinds(&requests)?; - Ok((requests, rows, semantic_index, outcome)) + + // The metric and the descriptor would fight over the same table across + // two engines. The user's metric wins. + let resource_info = if !metric_ctx.resource_info { + None + } else if requests + .inserts + .iter() + .any(|r| r.table_name == OTEL_RESOURCE_INFO_TABLE_NAME) + { + warn!( + "Skipping OTLP resource descriptor synthesis: the request writes \ + a metric named `{OTEL_RESOURCE_INFO_TABLE_NAME}`" + ); + None + } else { + resource_info.into_row_insert_requests()? + }; + if resource_info.is_some() { + semantic_index.record_scalar( + OTEL_RESOURCE_INFO_TABLE_NAME, + SEMANTIC_METRIC_TYPE, + crate::semantic::METRIC_TYPE_INFO, + ); + semantic_index.record_scalar( + OTEL_RESOURCE_INFO_TABLE_NAME, + SEMANTIC_METRIC_METADATA_QUALITY, + METADATA_QUALITY_DECLARED, + ); + } + + Ok(MetricsConversion { + requests, + rows, + semantic_index, + outcome, + resource_info, + }) } fn validate_sample_kinds(requests: &RowInsertRequests) -> Result<()> { @@ -287,30 +344,61 @@ fn from_metric_type(data: &metric::Data) -> MetricType { } } +/// Non-scalar values (bool, arrays, maps, bytes) are not representable as tags. +fn scalar_value_string(value: Option<&AnyValue>) -> Option { + match value.and_then(|v| v.value.as_ref())? { + any_value::Value::StringValue(s) => Some(s.clone()), + any_value::Value::IntValue(v) => Some(v.to_string()), + any_value::Value::DoubleValue(v) => Some(v.to_string()), + _ => None, + } +} + +/// Prometheus-style `(job, instance)` identity. Per the OTel Prometheus +/// compatibility spec `job` folds in `service.namespace` when present; a +/// resource without `service.name` gets no job rather than a fabricated one. +/// Both fields are optional strings, so a named type keeps a swapped +/// destructuring from type-checking at the call sites. +pub(crate) struct ServiceIdentity { + pub job: Option, + pub instance: Option, +} + +pub(crate) fn service_identity(attrs: &[KeyValue]) -> ServiceIdentity { + let mut name = None; + let mut namespace = None; + let mut instance = None; + for kv in attrs { + match kv.key.as_str() { + KEY_SERVICE_NAME => name = scalar_value_string(kv.value.as_ref()), + KEY_SERVICE_NAMESPACE => namespace = scalar_value_string(kv.value.as_ref()), + KEY_SERVICE_INSTANCE_ID => instance = scalar_value_string(kv.value.as_ref()), + _ => {} + } + } + let job = name.map(|name| match namespace { + Some(ns) if !ns.is_empty() => format!("{ns}/{name}"), + _ => name, + }); + ServiceIdentity { job, instance } +} + +fn string_key_value(key: &str, value: String) -> KeyValue { + KeyValue { + key: key.to_string(), + value: Some(AnyValue { + value: Some(any_value::Value::StringValue(value)), + }), + } +} + fn process_resource_attrs(attrs: &mut Vec, metric_ctx: &OtlpMetricCtx) { if metric_ctx.is_legacy { return; } - // remap service.name and service.instance.id to job and instance - let mut tmp = Vec::with_capacity(2); - for kv in attrs.iter() { - match &kv.key as &str { - KEY_SERVICE_NAME => { - tmp.push(KeyValue { - key: JOB_KEY.to_string(), - value: kv.value.clone(), - }); - } - KEY_SERVICE_INSTANCE_ID => { - tmp.push(KeyValue { - key: INSTANCE_KEY.to_string(), - value: kv.value.clone(), - }); - } - _ => {} - } - } + // remap the service identity attributes to job and instance + let ServiceIdentity { job, instance } = service_identity(attrs); // if promote all, then exclude the list, else, include the list if metric_ctx.promote_all_resource_attrs { @@ -322,7 +410,12 @@ fn process_resource_attrs(attrs: &mut Vec, metric_ctx: &OtlpMetricCtx) }); } - attrs.extend(tmp); + if let Some(job) = job { + attrs.push(string_key_value(JOB_KEY, job)); + } + if let Some(instance) = instance { + attrs.push(string_key_value(INSTANCE_KEY, instance)); + } } fn process_scope_attrs(scope: &ScopeMetrics, metric_ctx: &OtlpMetricCtx) -> Option> { @@ -536,35 +629,13 @@ fn encode_exponential_histogram( metric_ctx: &OtlpMetricCtx, outcome: &mut MetricsIngestOutcome, ) -> Result { - if !metric_ctx.experimental_enable_exponential_histogram { + if let Err(rejection) = exponential_histogram_gate(histogram, metric_ctx) { reject_data_points(outcome, histogram.data_points.len(), || { - format!( - "metric `{name}` uses OTLP exponential histograms; set otlp.experimental_enable_exponential_histogram = true to enable ingestion" - ) + rejection.message(name) })?; return Ok(false); } - match AggregationTemporality::try_from(histogram.aggregation_temporality) { - Ok(AggregationTemporality::Cumulative) => {} - Ok(AggregationTemporality::Delta) => { - reject_data_points(outcome, histogram.data_points.len(), || { - format!( - "metric `{name}` uses delta OTLP exponential histograms; only cumulative temporality is supported" - ) - })?; - return Ok(false); - } - _ => { - reject_data_points(outcome, histogram.data_points.len(), || { - format!( - "metric `{name}` has unspecified OTLP exponential histogram temporality; cumulative temporality is required" - ) - })?; - return Ok(false); - } - } - let column_schema = native_histogram_column_schema().map_err(|error| { error::InternalSnafu { err_msg: error.to_string(), @@ -611,7 +682,46 @@ fn encode_exponential_histogram( Ok(emitted) } -fn exponential_histogram_value( +pub(crate) enum ExponentialHistogramRejection { + Disabled, + DeltaTemporality, + UnspecifiedTemporality, +} + +impl ExponentialHistogramRejection { + fn message(&self, name: &str) -> String { + match self { + Self::Disabled => format!( + "metric `{name}` uses OTLP exponential histograms; set otlp.experimental_enable_exponential_histogram = true to enable ingestion" + ), + Self::DeltaTemporality => format!( + "metric `{name}` uses delta OTLP exponential histograms; only cumulative temporality is supported" + ), + Self::UnspecifiedTemporality => format!( + "metric `{name}` has unspecified OTLP exponential histogram temporality; cumulative temporality is required" + ), + } + } +} + +/// Whole-metric acceptance, decided once for the encoder and for the resource +/// descriptor, which must not describe a resource whose data was rejected. +/// Individual points can still fail [`exponential_histogram_value`]. +pub(crate) fn exponential_histogram_gate( + histogram: &ExponentialHistogram, + metric_ctx: &OtlpMetricCtx, +) -> std::result::Result<(), ExponentialHistogramRejection> { + if !metric_ctx.experimental_enable_exponential_histogram { + return Err(ExponentialHistogramRejection::Disabled); + } + match AggregationTemporality::try_from(histogram.aggregation_temporality) { + Ok(AggregationTemporality::Cumulative) => Ok(()), + Ok(AggregationTemporality::Delta) => Err(ExponentialHistogramRejection::DeltaTemporality), + _ => Err(ExponentialHistogramRejection::UnspecifiedTemporality), + } +} + +pub(crate) fn exponential_histogram_value( data_point: &ExponentialHistogramDataPoint, ) -> std::result::Result<(ValueData, i64), String> { if data_point.start_time_unix_nano > data_point.time_unix_nano { @@ -823,29 +933,21 @@ fn write_attributes( }; let tags = attrs.iter().filter_map(|attr| { - attr.value - .as_ref() - .and_then(|v| v.value.as_ref()) - .and_then(|val| { - let key = match attribute_type { - AttributeType::Resource | AttributeType::DataPoint => { - translate_label_name(&attr.key, metric_ctx.metric_translation_strategy) - } - AttributeType::Scope => { - format!( - "otel_scope_{}", - translate_label_name(&attr.key, metric_ctx.metric_translation_strategy) - ) - } - AttributeType::Legacy => legacy_normalize_otlp_name(&attr.key), - }; - match val { - any_value::Value::StringValue(s) => Some((key, s.clone())), - any_value::Value::IntValue(v) => Some((key, v.to_string())), - any_value::Value::DoubleValue(v) => Some((key, v.to_string())), - _ => None, // TODO(sunng87): allow different type of values - } - }) + // TODO(sunng87): allow different type of values + let value = scalar_value_string(attr.value.as_ref())?; + let key = match attribute_type { + AttributeType::Resource | AttributeType::DataPoint => { + translate_label_name(&attr.key, metric_ctx.metric_translation_strategy) + } + AttributeType::Scope => { + format!( + "otel_scope_{}", + translate_label_name(&attr.key, metric_ctx.metric_translation_strategy) + ) + } + AttributeType::Legacy => legacy_normalize_otlp_name(&attr.key), + }; + Some((key, value)) }); row_writer::write_tags(writer, tags, row)?; @@ -1295,6 +1397,169 @@ mod tests { } } + fn descriptor_ctx() -> OtlpMetricCtx { + OtlpMetricCtx { + resource_info: true, + ..Default::default() + } + } + + fn attr_value(attrs: &[KeyValue], key: &str) -> Option { + attrs + .iter() + .find(|kv| kv.key == key) + .and_then(|kv| scalar_value_string(kv.value.as_ref())) + } + + fn gauge_request( + resource_attrs: Vec, + metric_name: &str, + ) -> ExportMetricsServiceRequest { + use otel_arrow_rust::proto::opentelemetry::resource::v1::Resource; + ExportMetricsServiceRequest { + resource_metrics: vec![ResourceMetrics { + resource: Some(Resource { + attributes: resource_attrs, + ..Default::default() + }), + scope_metrics: vec![ScopeMetrics { + metrics: vec![Metric { + name: metric_name.to_string(), + data: Some(metric::Data::Gauge(Gauge { + data_points: vec![NumberDataPoint { + time_unix_nano: 1_000_000, + value: Some(Value::AsInt(1)), + ..Default::default() + }], + })), + ..Default::default() + }], + ..Default::default() + }], + ..Default::default() + }], + } + } + + fn column_names(request: &RowInsertRequests, table: &str) -> Vec { + request + .inserts + .iter() + .find(|r| r.table_name == table) + .unwrap_or_else(|| panic!("missing table {table}")) + .rows + .as_ref() + .unwrap() + .schema + .iter() + .map(|c| c.column_name.clone()) + .collect() + } + + #[test] + fn test_conversion_synthesizes_resource_descriptor() { + set_default_prefix(None).unwrap(); + let request = gauge_request( + vec![keyvalue("service.name", "api"), keyvalue("host.id", "h-1")], + "my_gauge", + ); + let conversion = to_grpc_insert_requests(request, &mut descriptor_ctx()).unwrap(); + + // descriptor keeps raw OTel keys while the metric table's labels went + // through the (default underscore-escaping) translation strategy + let resource_info = conversion.resource_info.expect("descriptor synthesized"); + let descriptor_cols = column_names(&resource_info, OTEL_RESOURCE_INFO_TABLE_NAME); + assert!(descriptor_cols.contains(&"host.id".to_string())); + assert!(descriptor_cols.contains(&"service.name".to_string())); + assert!(descriptor_cols.contains(&"job".to_string())); + let metric_cols = column_names(&conversion.requests, "my_gauge"); + assert!(metric_cols.contains(&"service_name".to_string())); + assert!(!metric_cols.contains(&"service.name".to_string())); + // host.id is not in the promote list: only the descriptor keeps it + assert!(!metric_cols.contains(&"host_id".to_string())); + + let decoded = decode(&conversion.semantic_index); + let t = &decoded[OTEL_RESOURCE_INFO_TABLE_NAME]; + assert_eq!( + t.get(SEMANTIC_METRIC_TYPE).map(String::as_str), + Some("info") + ); + assert_eq!( + t.get(SEMANTIC_METRIC_METADATA_QUALITY).map(String::as_str), + Some("declared") + ); + } + + #[test] + fn test_conversion_skips_descriptor_for_legacy_mode() { + set_default_prefix(None).unwrap(); + let request = gauge_request( + vec![keyvalue("service.name", "api"), keyvalue("host.id", "h-1")], + "my_gauge", + ); + let mut ctx = OtlpMetricCtx { + is_legacy: true, + ..descriptor_ctx() + }; + let conversion = to_grpc_insert_requests(request, &mut ctx).unwrap(); + assert!(conversion.resource_info.is_none()); + + // legacy tables predate job/instance and the promote filter; adding + // either would alter the schema of tables already in use + let cols = column_names(&conversion.requests, "my_gauge"); + assert!(!cols.contains(&"job".to_string())); + assert!(cols.contains(&"host_id".to_string())); + } + + #[test] + fn test_conversion_skips_descriptor_on_metric_name_collision() { + set_default_prefix(None).unwrap(); + let request = gauge_request( + vec![keyvalue("service.name", "api")], + OTEL_RESOURCE_INFO_TABLE_NAME, + ); + let conversion = to_grpc_insert_requests(request, &mut descriptor_ctx()).unwrap(); + assert!(conversion.resource_info.is_none()); + // the metric itself still goes through the main path + assert!( + conversion + .requests + .inserts + .iter() + .any(|r| r.table_name == OTEL_RESOURCE_INFO_TABLE_NAME) + ); + } + + #[test] + fn test_job_composition_follows_service_namespace() { + let mut attrs = vec![ + keyvalue("service.name", "api"), + keyvalue("service.namespace", "shop"), + ]; + process_resource_attrs(&mut attrs, &OtlpMetricCtx::default()); + assert_eq!(attr_value(&attrs, "job").as_deref(), Some("shop/api")); + + let mut attrs = vec![keyvalue("service.name", "api")]; + process_resource_attrs(&mut attrs, &OtlpMetricCtx::default()); + assert_eq!(attr_value(&attrs, "job").as_deref(), Some("api")); + + let mut attrs = vec![ + keyvalue("service.name", "api"), + keyvalue("service.namespace", ""), + ]; + process_resource_attrs(&mut attrs, &OtlpMetricCtx::default()); + assert_eq!(attr_value(&attrs, "job").as_deref(), Some("api")); + + // no service.name: no job is fabricated + let mut attrs = vec![ + keyvalue("service.namespace", "shop"), + keyvalue("service.instance.id", "inst-1"), + ]; + process_resource_attrs(&mut attrs, &OtlpMetricCtx::default()); + assert_eq!(attr_value(&attrs, "job"), None); + assert_eq!(attr_value(&attrs, "instance").as_deref(), Some("inst-1")); + } + #[test] fn test_encode_gauge() { let mut tables = MultiTableData::default(); @@ -2104,8 +2369,12 @@ mod tests { AggregationTemporality::Cumulative, ), ]); - let (requests, _, semantic_index, outcome) = - to_grpc_insert_requests(request, &mut OtlpMetricCtx::default()).unwrap(); + let MetricsConversion { + requests, + semantic_index, + outcome, + .. + } = to_grpc_insert_requests(request, &mut OtlpMetricCtx::default()).unwrap(); assert_eq!(outcome.accepted_data_points, 1); assert_eq!(outcome.rejected_data_points, 1); @@ -2127,8 +2396,9 @@ mod tests { vec![], AggregationTemporality::Cumulative, )]); - let (_, _, _, outcome) = - to_grpc_insert_requests(empty, &mut OtlpMetricCtx::default()).unwrap(); + let outcome = to_grpc_insert_requests(empty, &mut OtlpMetricCtx::default()) + .unwrap() + .outcome; assert_eq!(outcome.rejected_data_points, 0); assert_eq!(outcome.error_message, None); } @@ -2208,8 +2478,12 @@ mod tests { .collect(), }; - let (requests, rows, _, outcome) = - to_grpc_insert_requests(request, &mut OtlpMetricCtx::default()).unwrap(); + let MetricsConversion { + requests, + rows, + outcome, + .. + } = to_grpc_insert_requests(request, &mut OtlpMetricCtx::default()).unwrap(); assert_eq!(outcome.accepted_data_points, 2); assert_eq!(rows, 6); @@ -2243,8 +2517,13 @@ mod tests { experimental_enable_exponential_histogram: true, ..Default::default() }; - let (requests, rows, semantic_index, outcome) = - to_grpc_insert_requests(request, &mut ctx).unwrap(); + let MetricsConversion { + requests, + rows, + semantic_index, + outcome, + .. + } = to_grpc_insert_requests(request, &mut ctx).unwrap(); assert_eq!(outcome.accepted_data_points, 0); assert_eq!(outcome.rejected_data_points, 1); @@ -2269,14 +2548,17 @@ mod tests { experimental_enable_exponential_histogram: true, ..Default::default() }; - let (new_requests, _, _, _) = - to_grpc_insert_requests(request.clone(), &mut new_ctx).unwrap(); + let new_requests = to_grpc_insert_requests(request.clone(), &mut new_ctx) + .unwrap() + .requests; let mut legacy_ctx = OtlpMetricCtx { experimental_enable_exponential_histogram: true, is_legacy: true, ..Default::default() }; - let (legacy_requests, _, _, _) = to_grpc_insert_requests(request, &mut legacy_ctx).unwrap(); + let legacy_requests = to_grpc_insert_requests(request, &mut legacy_ctx) + .unwrap() + .requests; let new_insert = &new_requests.inserts[0]; let legacy_insert = &legacy_requests.inserts[0]; @@ -2325,8 +2607,9 @@ mod tests { vec![ExponentialHistogramDataPoint::default()], AggregationTemporality::Cumulative, )]); - let (_, _, _, outcome) = - to_grpc_insert_requests(request, &mut OtlpMetricCtx::default()).unwrap(); + let outcome = to_grpc_insert_requests(request, &mut OtlpMetricCtx::default()) + .unwrap() + .outcome; assert_eq!(outcome.rejected_data_points, 1); assert!(outcome.error_message.unwrap().len() <= MAX_REJECTION_MESSAGE_BYTES); diff --git a/src/servers/src/otlp/metrics/resource_info.rs b/src/servers/src/otlp/metrics/resource_info.rs new file mode 100644 index 00000000000..cd3427f7163 --- /dev/null +++ b/src/servers/src/otlp/metrics/resource_info.rs @@ -0,0 +1,393 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Resource descriptor synthesized from OTLP metrics requests, so the entity +//! graph reads one table instead of scanning every logical metric table for +//! attributes the promote filter may have dropped. +//! +//! Columns are a fixed allowlist under the raw OTel attribute names: the +//! conventions whitelist matches fixed names, so they must not follow the +//! per-request label translation strategy or the promote/ignore headers. + +use std::collections::BTreeMap; + +use api::v1::RowInsertRequests; +use common_catalog::consts::SEMANTIC_GRAPH_WINDOW_NANOS; +use common_grpc::precision::Precision; +use common_query::prelude::{greptime_timestamp, greptime_value}; +use otel_arrow_rust::proto::opentelemetry::common::v1::KeyValue; +use otel_arrow_rust::proto::opentelemetry::metrics::v1::{ResourceMetrics, metric}; +use session::protocol_ctx::OtlpMetricCtx; + +use crate::error::Result; +use crate::otlp::metrics::{ + INSTANCE_KEY, JOB_KEY, ServiceIdentity, exponential_histogram_gate, + exponential_histogram_value, scalar_value_string, service_identity, +}; +use crate::otlp::trace::{ + KEY_CONTAINER_ID, KEY_CONTAINER_NAME, KEY_HOST_ID, KEY_HOST_NAME, KEY_K8S_NAMESPACE_NAME, + KEY_K8S_POD_NAME, KEY_K8S_POD_UID, KEY_SERVICE_NAME, KEY_SERVICE_NAMESPACE, +}; +use crate::row_writer::{self, MultiTableData}; + +/// Prefixed like the other engine-managed tables, so a user metric is +/// unlikely to claim the name. +pub const OTEL_RESOURCE_INFO_TABLE_NAME: &str = "greptime_otel_resource_info"; + +/// Attributes projected under their raw OTel keys. `service.instance.id` is +/// absent on purpose: it lands in `instance`. +/// +/// Matched instead of scanned: this runs for every attribute of every +/// resource, and the compiler turns it into a length-and-prefix dispatch. +fn is_projected_attr(key: &str) -> bool { + matches!( + key, + KEY_SERVICE_NAME + | KEY_SERVICE_NAMESPACE + | KEY_HOST_ID + | KEY_HOST_NAME + | KEY_CONTAINER_ID + | KEY_CONTAINER_NAME + | KEY_K8S_POD_UID + | KEY_K8S_POD_NAME + | KEY_K8S_NAMESPACE_NAME + ) +} + +/// Upper bound of [`is_projected_attr`] plus the derived `job`/`instance`, +/// used to size the per-row buffers. +const MAX_PROJECTED_TAGS: usize = 11; + +/// Projected attributes (sorted `(name, value)` pairs) -> graph window -> +/// the newest data-point time seen in that window, which is what the row for +/// that window is stamped with. +/// +/// Windows are an inner map so the attributes are stored, and moved, once per +/// resource. They are keyed separately because one request may carry data for +/// several of them: a single row per resource would describe only the newest +/// window, leaving the earlier ones with metric rows but no entities. +#[derive(Debug, Default)] +pub struct ResourceInfoData { + rows: BTreeMap, BTreeMap>, +} + +impl ResourceInfoData { + /// Takes the raw attributes, before the promote filter runs on them. + pub fn observe( + &mut self, + raw_attrs: &[KeyValue], + resource: &ResourceMetrics, + metric_ctx: &OtlpMetricCtx, + ) { + let mut tags = Vec::with_capacity(MAX_PROJECTED_TAGS); + let ServiceIdentity { job, instance } = service_identity(raw_attrs); + if let Some(job) = job { + tags.push((JOB_KEY.to_string(), job)); + } + if let Some(instance) = instance { + tags.push((INSTANCE_KEY.to_string(), instance)); + } + for kv in raw_attrs { + if is_projected_attr(&kv.key) + && let Some(value) = scalar_value_string(kv.value.as_ref()) + { + tags.push((kv.key.clone(), value)); + } + } + if tags.is_empty() { + return; + } + // Sorted so equal attribute sets share a key, and so the emitted + // columns keep a stable order. + tags.sort_unstable(); + + let mut observed: BTreeMap = BTreeMap::new(); + for_each_encoded_time(resource, metric_ctx, |ts| { + let window = ts - ts.rem_euclid(SEMANTIC_GRAPH_WINDOW_NANOS); + observed + .entry(window) + .and_modify(|newest| *newest = (*newest).max(ts)) + .or_insert(ts); + }); + if observed.is_empty() { + return; + } + + let windows = self.rows.entry(tags).or_default(); + for (window, newest) in observed { + windows + .entry(window) + .and_modify(|seen| *seen = (*seen).max(newest)) + .or_insert(newest); + } + } + + /// Every projected attribute becomes a tag, so auto-create puts it in the + /// primary key where the conventions expect it. + pub fn into_row_insert_requests(self) -> Result> { + if self.rows.is_empty() { + return Ok(None); + } + + let mut writer = MultiTableData::default(); + let table = writer.get_or_default_table_data( + OTEL_RESOURCE_INFO_TABLE_NAME, + MAX_PROJECTED_TAGS + 2, + self.rows.values().map(BTreeMap::len).sum(), + ); + for (tags, windows) in &self.rows { + for ts_nanos in windows.values().copied() { + let mut row = table.alloc_one_row(); + row_writer::write_tags(table, tags.iter().cloned(), &mut row)?; + row_writer::write_f64(table, greptime_value(), 1.0, &mut row)?; + row_writer::write_ts_to_millis( + table, + greptime_timestamp(), + Some(ts_nanos), + Precision::Nanosecond, + &mut row, + )?; + table.add_row(row); + } + } + + let (requests, _) = writer.into_row_insert_requests(); + Ok(Some(requests)) + } +} + +/// Visits the times of the data points the encoder writes rows for, so a +/// resource is described exactly where it is measured rather than wherever +/// its request happens to reach. +fn for_each_encoded_time( + resource: &ResourceMetrics, + metric_ctx: &OtlpMetricCtx, + mut visit: impl FnMut(i64), +) { + fn visit_all(points: impl Iterator, visit: &mut impl FnMut(i64)) { + for ts in points { + visit(ts as i64); + } + } + for scope in &resource.scope_metrics { + for m in &scope.metrics { + match &m.data { + Some(metric::Data::Gauge(g)) => { + visit_all(g.data_points.iter().map(|p| p.time_unix_nano), &mut visit) + } + Some(metric::Data::Sum(s)) => { + visit_all(s.data_points.iter().map(|p| p.time_unix_nano), &mut visit) + } + Some(metric::Data::Histogram(h)) => { + visit_all(h.data_points.iter().map(|p| p.time_unix_nano), &mut visit) + } + Some(metric::Data::Summary(s)) => { + visit_all(s.data_points.iter().map(|p| p.time_unix_nano), &mut visit) + } + Some(metric::Data::ExponentialHistogram(h)) + if exponential_histogram_gate(h, metric_ctx).is_ok() => + { + for point in &h.data_points { + if let Ok((_, ts)) = exponential_histogram_value(point) { + visit(ts); + } + } + } + Some(metric::Data::ExponentialHistogram(_)) => {} + None => {} + } + } + } +} + +#[cfg(test)] +mod tests { + use api::v1::SemanticType; + use api::v1::value::ValueData; + use common_query::prelude::set_default_prefix; + use otel_arrow_rust::proto::opentelemetry::common::v1::{AnyValue, any_value}; + use otel_arrow_rust::proto::opentelemetry::metrics::v1::{ + AggregationTemporality, ExponentialHistogram, ExponentialHistogramDataPoint, Gauge, Metric, + NumberDataPoint, ScopeMetrics, + }; + + use super::*; + + fn kv(key: &str, value: &str) -> KeyValue { + KeyValue { + key: key.into(), + value: Some(AnyValue { + value: Some(any_value::Value::StringValue(value.into())), + }), + } + } + + fn gauge_at(times: &[i64]) -> ResourceMetrics { + ResourceMetrics { + scope_metrics: vec![ScopeMetrics { + metrics: vec![Metric { + data: Some(metric::Data::Gauge(Gauge { + data_points: times + .iter() + .map(|ts| NumberDataPoint { + time_unix_nano: *ts as u64, + ..Default::default() + }) + .collect(), + })), + ..Default::default() + }], + ..Default::default() + }], + ..Default::default() + } + } + + #[test] + fn observe_projects_allowlist_and_dedups_per_request() { + let mut data = ResourceInfoData::default(); + let attrs = vec![ + kv("service.name", "api"), + kv("service.namespace", "shop"), + kv("service.instance.id", "inst-1"), + kv("host.id", "h-1"), + kv("os.type", "linux"), + ]; + data.observe(&attrs, &gauge_at(&[100, 50]), &OtlpMetricCtx::default()); + assert_eq!(data.rows.len(), 1); + let (tags, windows) = data.rows.iter().next().unwrap(); + assert_eq!(windows.values().copied().collect::>(), vec![100]); + assert!(tags.contains(&("job".to_string(), "shop/api".to_string()))); + assert!(tags.contains(&("instance".to_string(), "inst-1".to_string()))); + assert!(tags.contains(&("service.name".to_string(), "api".to_string()))); + assert!( + tags.iter() + .all(|(k, _)| k != "os.type" && k != "service.instance.id") + ); + + data.observe( + &[kv("host.id", "h-2")], + &gauge_at(&[10]), + &OtlpMetricCtx::default(), + ); + assert_eq!(data.rows.len(), 2); + + let mut empty = ResourceInfoData::default(); + empty.observe( + &[kv("os.type", "linux")], + &gauge_at(&[100]), + &OtlpMetricCtx::default(), + ); + assert!(empty.into_row_insert_requests().unwrap().is_none()); + } + + /// Earlier windows would keep their metric rows but lose their entities. + #[test] + fn observe_keeps_one_row_per_graph_window() { + let window = SEMANTIC_GRAPH_WINDOW_NANOS; + let mut data = ResourceInfoData::default(); + data.observe( + &[kv("service.name", "api")], + &gauge_at(&[window + 1, window + 2, 3 * window + 7]), + &OtlpMetricCtx::default(), + ); + + let windows = data.rows.values().next().unwrap(); + assert_eq!( + windows.iter().collect::>(), + vec![(&window, &(window + 2)), (&(3 * window), &(3 * window + 7))] + ); + } + + /// Describing a resource whose only data the encoder drops invents an + /// entity with no measurements. + #[test] + fn observe_ignores_data_the_encoder_drops() { + let exponential = |temporality: AggregationTemporality| ResourceMetrics { + scope_metrics: vec![ScopeMetrics { + metrics: vec![Metric { + data: Some(metric::Data::ExponentialHistogram(ExponentialHistogram { + data_points: vec![ExponentialHistogramDataPoint { + time_unix_nano: 100, + ..Default::default() + }], + aggregation_temporality: temporality as i32, + })), + ..Default::default() + }], + ..Default::default() + }], + ..Default::default() + }; + let enabled = OtlpMetricCtx { + experimental_enable_exponential_histogram: true, + ..Default::default() + }; + + for (resource, ctx) in [ + ( + exponential(AggregationTemporality::Cumulative), + OtlpMetricCtx::default(), + ), + (exponential(AggregationTemporality::Delta), enabled), + ] { + let mut data = ResourceInfoData::default(); + data.observe(&[kv("service.name", "api")], &resource, &ctx); + assert!(data.into_row_insert_requests().unwrap().is_none()); + } + } + + #[test] + fn rows_carry_raw_key_tags_value_and_millis_timestamp() { + set_default_prefix(None).unwrap(); + let mut data = ResourceInfoData::default(); + data.observe( + &[kv("service.name", "api"), kv("host.id", "h-1")], + &gauge_at(&[1_700_000_000_123_456_789]), + &OtlpMetricCtx::default(), + ); + let requests = data.into_row_insert_requests().unwrap().unwrap(); + assert_eq!(requests.inserts.len(), 1); + let insert = &requests.inserts[0]; + assert_eq!(insert.table_name, OTEL_RESOURCE_INFO_TABLE_NAME); + + let rows = insert.rows.as_ref().unwrap(); + let names = rows + .schema + .iter() + .map(|c| c.column_name.as_str()) + .collect::>(); + assert_eq!( + names, + vec![ + "host.id", + "job", + "service.name", + greptime_value(), + greptime_timestamp() + ] + ); + for column in &rows.schema[..3] { + assert_eq!(column.semantic_type, SemanticType::Tag as i32); + } + + assert_eq!(rows.rows.len(), 1); + let values = &rows.rows[0].values; + assert_eq!(values[3].value_data, Some(ValueData::F64Value(1.0))); + assert_eq!( + values[4].value_data, + Some(ValueData::TimestampMillisecondValue(1_700_000_000_123)) + ); + } +} diff --git a/src/servers/src/otlp/trace.rs b/src/servers/src/otlp/trace.rs index a8439cdb2c4..9ee16507bda 100644 --- a/src/servers/src/otlp/trace.rs +++ b/src/servers/src/otlp/trace.rs @@ -46,7 +46,15 @@ pub const TRACE_STATE_COLUMN: &str = "trace_state"; // const keys pub const KEY_SERVICE_NAME: &str = "service.name"; +pub const KEY_SERVICE_NAMESPACE: &str = "service.namespace"; pub const KEY_SERVICE_INSTANCE_ID: &str = "service.instance.id"; +pub const KEY_HOST_ID: &str = "host.id"; +pub const KEY_HOST_NAME: &str = "host.name"; +pub const KEY_CONTAINER_ID: &str = "container.id"; +pub const KEY_CONTAINER_NAME: &str = "container.name"; +pub const KEY_K8S_POD_UID: &str = "k8s.pod.uid"; +pub const KEY_K8S_POD_NAME: &str = "k8s.pod.name"; +pub const KEY_K8S_NAMESPACE_NAME: &str = "k8s.namespace.name"; pub const KEY_SPAN_KIND: &str = "span.kind"; // jaeger const keys, not sure if they are general diff --git a/src/servers/src/postgres/handler.rs b/src/servers/src/postgres/handler.rs index 484bb6a1f11..aca4362bd34 100644 --- a/src/servers/src/postgres/handler.rs +++ b/src/servers/src/postgres/handler.rs @@ -28,6 +28,7 @@ use datafusion_pg_catalog::sql::PostgresCompatibilityParser; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{Schema, SchemaRef}; use futures::{Sink, SinkExt, Stream, StreamExt, future, stream}; +use operator::statement::admin_output_schema; use pgwire::api::portal::{Format, Portal}; use pgwire::api::query::{ExtendedQueryHandler, SimpleQueryHandler}; use pgwire::api::results::{ @@ -40,8 +41,10 @@ use pgwire::error::{ErrorInfo, PgWireError, PgWireResult}; use pgwire::messages::PgWireBackendMessage; use pgwire::messages::copy::CopyData; use pgwire::messages::data::DataRow; +use query::dist_analyze_output_schema; use query::planner::DfLogicalPlanner; use query::query_engine::DescribeResult; +use query::sql::DESCRIBE_TABLE_OUTPUT_SCHEMA; use session::Session; use session::context::QueryContextRef; use snafu::ResultExt; @@ -554,6 +557,13 @@ fn describe_fields( session: &Arc, ) -> PgWireResult> { match sql_plan { + // Execution swaps in DistAnalyzeExec (stage/node/plan), whose schema + // differs from the logical `Analyze` plan's (plan_type/plan). + SqlPlan::Plan(LogicalPlan::Analyze(_), _) => { + let schema: Schema = + Schema::try_from(dist_analyze_output_schema()).map_err(convert_err)?; + schema_to_pg(&schema, format, None).map_err(convert_err) + } // query SqlPlan::Plan(plan, _) if !matches!(plan, LogicalPlan::Dml(_) | LogicalPlan::Ddl(_)) => { let schema: Schema = plan.schema().clone().try_into().map_err(convert_err)?; @@ -647,17 +657,6 @@ fn describe_fields( ), ]), - // single column show statements - SqlPlan::Statement( - Statement::ShowTables(_) | Statement::ShowFlows(_) | Statement::ShowViews(_), - _, - ) => Ok(vec![FieldInfo::new( - "name".to_string(), - None, - None, - Type::TEXT, - format.format_for(0), - )]), #[cfg(feature = "enterprise")] SqlPlan::Statement(Statement::ShowTriggers(_), _) => Ok(vec![FieldInfo::new( "name".to_string(), @@ -679,6 +678,61 @@ fn describe_fields( Ok(vec![]) } } + // Single column named after the variable (see `query::sql::show_variable`). + SqlPlan::Statement(Statement::ShowVariables(show), _) => Ok(vec![FieldInfo::new( + show.variable.to_string().to_uppercase(), + None, + None, + Type::TEXT, + format.format_for(0), + )]), + // Mirrors `query::sql::show_status` (currently always empty). + SqlPlan::Statement(Statement::ShowStatus(_), _) => Ok(vec![ + FieldInfo::new( + "Variable_name".to_string(), + None, + None, + Type::TEXT, + format.format_for(0), + ), + FieldInfo::new( + "Value".to_string(), + None, + None, + Type::TEXT, + format.format_for(1), + ), + ]), + SqlPlan::Statement(Statement::ShowSearchPath(_), _) => Ok(vec![FieldInfo::new( + "search_path".to_string(), + None, + None, + Type::TEXT, + format.format_for(0), + )]), + // Mirrors `query::sql::describe_table`. + SqlPlan::Statement(Statement::DescribeTable(_), _) => { + schema_to_pg(&DESCRIBE_TABLE_OUTPUT_SCHEMA, format, None).map_err(convert_err) + } + // Single column typed with the function's return type (see + // `operator::statement::admin_output_schema`). + SqlPlan::Statement(Statement::Admin(admin), _) => { + let query_ctx = session.new_query_context(); + match admin_output_schema(admin, &query_ctx) { + Some(schema) => schema_to_pg(&schema, format, None).map_err(convert_err), + // Unresolvable; execution will surface the error. + None => Ok(vec![]), + } + } + // Describe from the declared cursor's schema. + SqlPlan::Statement(Statement::FetchCursor(fetch), _) => { + let cursor_name = fetch.cursor_name.to_string(); + match session.get_cursor(&cursor_name) { + Some(cursor) => schema_to_pg(&cursor.schema(), format, None).map_err(convert_err), + // Cursor not declared yet; execution will error. + None => Ok(vec![]), + } + } _ => { // NoData Ok(vec![]) diff --git a/src/session/src/lib.rs b/src/session/src/lib.rs index 110560b3243..c1414fc5a3b 100644 --- a/src/session/src/lib.rs +++ b/src/session/src/lib.rs @@ -111,6 +111,12 @@ impl Session { .into() } + /// Cursors are shared across query contexts created from this session. + pub fn get_cursor(&self, name: &str) -> Option> { + let guard = self.mutable_inner.read().unwrap(); + guard.cursors.get(name).cloned() + } + pub fn conn_info(&self) -> &ConnInfo { &self.conn_info } diff --git a/src/session/src/protocol_ctx.rs b/src/session/src/protocol_ctx.rs index 537dec5c414..d9a4a386f91 100644 --- a/src/session/src/protocol_ctx.rs +++ b/src/session/src/protocol_ctx.rs @@ -56,6 +56,9 @@ pub struct OtlpMetricCtx { pub with_metric_engine: bool, pub experimental_enable_exponential_histogram: bool, pub is_legacy: bool, + /// Set from the server's `otlp.experimental_enable_resource_info`; off + /// means the resource descriptor is not synthesized at all. + pub resource_info: bool, pub metric_type: MetricType, pub metric_translation_strategy: OtlpMetricTranslationStrategy, } diff --git a/src/store-api/src/sst_entry.rs b/src/store-api/src/sst_entry.rs index f1bb1f9c36f..c70e909ebd6 100644 --- a/src/store-api/src/sst_entry.rs +++ b/src/store-api/src/sst_entry.rs @@ -65,6 +65,8 @@ pub struct ManifestSstEntry { pub file_path: String, /// File size in bytes. pub file_size: u64, + /// Maximum uncompressed row group size in bytes. Zero means unknown. + pub max_row_group_uncompressed_size: u64, /// Full path of the index file in object store. pub index_file_path: Option, /// File size of the index file in object store. @@ -81,6 +83,8 @@ pub struct ManifestSstEntry { pub max_ts: Timestamp, /// The sequence number associated with this file. pub sequence: Option, + /// Human-readable partition expression associated with this file. + pub partition_expr: Option, /// The region id of region that creates the file. pub origin_region_id: RegionId, /// The node id fetched from the manifest. @@ -122,6 +126,12 @@ impl ManifestSstEntry { ColumnSchema::new("visible", Ty::boolean_datatype(), false), ColumnSchema::new("primary_key_min", Ty::binary_datatype(), true), ColumnSchema::new("primary_key_max", Ty::binary_datatype(), true), + ColumnSchema::new( + "max_row_group_uncompressed_size", + Ty::uint64_datatype(), + false, + ), + ColumnSchema::new("partition_expr", Ty::string_datatype(), true), ])) } @@ -139,6 +149,8 @@ impl ManifestSstEntry { let levels = entries.iter().map(|e| e.level); let file_paths = entries.iter().map(|e| e.file_path.as_str()); let file_sizes = entries.iter().map(|e| e.file_size); + let max_row_group_uncompressed_sizes = + entries.iter().map(|e| e.max_row_group_uncompressed_size); let index_file_paths = entries.iter().map(|e| e.index_file_path.as_ref()); let index_file_sizes = entries.iter().map(|e| e.index_file_size); let num_rows = entries.iter().map(|e| e.num_rows); @@ -155,6 +167,7 @@ impl ManifestSstEntry { .map(|ts| ts.value()) }); let sequences = entries.iter().map(|e| e.sequence); + let partition_exprs = entries.iter().map(|e| e.partition_expr.as_ref()); let origin_region_ids = entries.iter().map(|e| e.origin_region_id.as_u64()); let node_ids = entries.iter().map(|e| e.node_id); let visible_flags = entries.iter().map(|e| Some(e.visible)); @@ -186,6 +199,10 @@ impl ManifestSstEntry { Arc::new(BooleanArray::from_iter(visible_flags)), Arc::new(BinaryArray::from_iter(primary_key_min)), Arc::new(BinaryArray::from_iter(primary_key_max)), + Arc::new(UInt64Array::from_iter_values( + max_row_group_uncompressed_sizes, + )), + Arc::new(StringArray::from_iter(partition_exprs)), ]; DfRecordBatch::try_new(schema.arrow_schema().clone(), columns) @@ -458,6 +475,7 @@ mod tests { level: 1, file_path: "/p1".to_string(), file_size: 100, + max_row_group_uncompressed_size: 80, index_file_path: None, index_file_size: None, num_rows: 10, @@ -466,6 +484,7 @@ mod tests { min_ts: Timestamp::new_millisecond(1000), // 1s -> 1_000_000_000ns max_ts: Timestamp::new_second(2), // 2s -> 2_000_000_000ns sequence: None, + partition_expr: Some("a < 10".to_string()), origin_region_id: region_id1, node_id: Some(1), visible: false, @@ -484,6 +503,7 @@ mod tests { level: 3, file_path: "/p2".to_string(), file_size: 200, + max_row_group_uncompressed_size: 160, index_file_path: Some("idx".to_string()), index_file_size: Some(11), num_rows: 20, @@ -492,6 +512,7 @@ mod tests { min_ts: Timestamp::new_nanosecond(5), // 5ns max_ts: Timestamp::new_microsecond(2000), // 2ms -> 2_000_000ns sequence: Some(9), + partition_expr: None, origin_region_id: region_id2, node_id: None, visible: true, @@ -506,6 +527,44 @@ mod tests { // Schema checks assert_eq!(schema.arrow_schema().fields().len(), batch.num_columns()); assert_eq!(2, batch.num_rows()); + let expected_columns = [ + "table_dir", + "region_id", + "table_id", + "region_number", + "region_group", + "region_sequence", + "file_id", + "index_version", + "level", + "file_path", + "file_size", + "index_file_path", + "index_file_size", + "num_rows", + "num_row_groups", + "num_series", + "min_ts", + "max_ts", + "sequence", + "origin_region_id", + "node_id", + "visible", + "primary_key_min", + "primary_key_max", + "max_row_group_uncompressed_size", + "partition_expr", + ]; + assert_eq!( + expected_columns, + schema + .arrow_schema() + .fields() + .iter() + .map(|field| field.name().as_str()) + .collect::>() + .as_slice() + ); for (i, f) in schema.arrow_schema().fields().iter().enumerate() { assert_eq!(f.name(), batch.schema().field(i).name()); assert_eq!(f.is_nullable(), batch.schema().field(i).is_nullable()); @@ -704,6 +763,22 @@ mod tests { .unwrap(); assert_eq!(b"zzz", primary_key_max.value(0)); assert!(primary_key_max.is_null(1)); + + let max_row_group_uncompressed_sizes = batch + .column(24) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(80, max_row_group_uncompressed_sizes.value(0)); + assert_eq!(160, max_row_group_uncompressed_sizes.value(1)); + + let partition_exprs = batch + .column(25) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!("a < 10", partition_exprs.value(0)); + assert!(partition_exprs.is_null(1)); } #[test] diff --git a/src/table/src/requests/semantic.rs b/src/table/src/requests/semantic.rs index 2daa440c11d..471be8b891c 100644 --- a/src/table/src/requests/semantic.rs +++ b/src/table/src/requests/semantic.rs @@ -98,6 +98,10 @@ pub const SEMANTIC_ENTITY_SERVICE_ID: &str = "greptime.semantic.entity.service.i /// ([`has_stable_string_form`]). #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum EntityRole { + /// Identifying columns. Their **order is part of the identity**: the + /// entity id is their values joined in the declared order, so two tables + /// naming the same entity must list them the same way (broad to narrow) + /// or they name two entities. Id, Descriptive, Scope, diff --git a/tests-integration/fixtures/etcd-tls-certs/ca-key.pem b/tests-integration/fixtures/etcd-tls-certs/ca-key.pem index 2066dfb7b47..07a0a86ee39 100644 --- a/tests-integration/fixtures/etcd-tls-certs/ca-key.pem +++ b/tests-integration/fixtures/etcd-tls-certs/ca-key.pem @@ -1,28 +1,28 @@ -----BEGIN PRIVATE KEY----- -MIIEvQIBADANBgkqhkiG9w0BAQEFAASCBKcwggSjAgEAAoIBAQCfIi0iTllUQyNY -ka+d9Zt30ILA0QBzLKHwez9QJj8bQLpfKR26d794CQn/OO6XO7MWsJ6nvqwRNYM0 -WX+UwUtuT9MrttYWpS5yEAVO79dMtF67gCNafXKjQyK8L1vz6Avq4XJN4VqZEifT -Fdk3s2MmS0twv4sgrK//jh7aZzBQJtw+ROe9RN1dqlKoa3TtpB7iwNTDxGiR0EYj -wHj+GdiNHtDMWvrnTNNX/IKeEwLP/Mjirzr7GnaOjR25ruqcLuwWrufu+EMhqFTU -r+LIQpZD0TFuPX6DPlQZOs/4UpjteoC84m5AnQXNlh0vzjbW1efEYWYQoKDesx9e -oE5jcABpAgMBAAECggEAA3Rw/r6CYaTG1tdBilA47LF/XTihu15mh6Y4BNadESDn -IEYav0p2mDW4sgH7lWyg4ecPaBHo124jfUFM1p8Zsvk9sERwbR5vSIpWKyq5hbtM -Fo2yIXakGgJM9ev19vDR7WqxRMVAkR6GnyZpcwmh0e+ENvGZpNR1n7qSAJOJRibW -Ye8qMF/zMLF9VekBRPFJFIZBvKgHxSGCLbXRzsi5LKMbwfBoJmxEZSOC9zODl43o -Z2i+LvlS0R/Rsf2KG+Hwgj1P7HH1znJeA5bVJFDewljjSwu1mk0JpQL7Jgp2xpqs -F/mRZmtoQSH3EsTm2Kjz1sLGXwUb9JaQ274aV4t7pwKBgQDXYF+qFYip7k+b9FTF -sFbfp3CjOqxXiAxnoE5hlfV/h0w4V1Su42AdejnbIVYFX8bvmezkNWPXhB6P5oHU -4nO+oAuhsRT4317mfRM5NWO6mv2eglSBnS04GLsX142A7jMPr9e+Cqh4dYfw/A1L -bSdsdILemTFQrak27pC11mHZOwKBgQC9JhDhQf5uRh6k8GvgKGu9ZEZmkYCVu2sr -7kpD60UhklnyQo6DKYsa6OcChRQMv9ss3Hqu3Y1G0/sTe1Tjz5zwrWWAZoLlbdOt -vipesvwNTxvJoUWPtrjOoFZtmyg0Q0TqxaxfEWAssDU9IfnIy+zx7sAg4zZIPUyy -6kudaY1SqwKBgQCAQ/71Fjn7qddzc4GA8lHqhJeKPokg3/8zP78uUtaQCo2UCD6A -oR0+sOn/3MyUCsQ5MZxpFHrPgPmKjabIl8yCvGHw+7sXtD+aWOa37VnlaiSc39Vg -E7E4dVIHEvJM1I9ISlrb7REEHErHc/Se9PTDnGfMFcPO3n2mH1HDWVeQvQKBgGGn -BHH3a08tXmbTRS5uT+lwmrQbjKJBJ3x/wtG75m4Fq/BaEk9/JDUZZyKy5/4JEzPf -BGvBME4P5QFS3CndJu5O5ydaRVwDzpRVqHRJvb11SShY3Zvrvw/WUai2wRPyYuM+ -eNaAFwIbWvEb2GSle8gP9htEkuLK2w1HzxAOzYqPAoGAKl5l/WsoLfQP5UIUjSQe -JtVTDVmoz8m8Mg2nn56UdUetG246EVhDx5ngj84kZDM0GDXhfVWV6MWnIZr5NQsr -bORIDNP7Czat8FnlXeAQ+R6A/87o5g5p6ydH9jYqrrLxYjaZyXJptE+ja4zW7Lbl -ogWAc76iZvXLXqKRfepG0zQ= +MIIEvAIBADANBgkqhkiG9w0BAQEFAASCBKYwggSiAgEAAoIBAQC05De2Tk8TMZ0J +ME+7Rp+zvvLkEBPPXCCFcXfo91gYVHMMr1TGQJuVtbx7rIqC1tFhc2luAGfsBmW3 +IfTXZjzY6834ZUzmPL7uuBfoRpd0Paj5X4owpn0eHIgcsqx4juHSTYc3fdnrTLpq +lNMulZEemkNULLCE7pyxg3hhYIx9JEe9nv26fkyr1vW5z1cPlEpyOK+SUIeqHgXi +VtWPU9WGv5sdQeXnD4E1JTX2URihb7C8xali3YVbI3SUxy/gAhM6swQgZ4SLgW8k +f87a+PcP9C1N5Gujfw4fT90dJoutyiTSbcDrjZLjF03TgjcDaACvntb0qN85XD96 +dIvhg5jfAgMBAAECggEAAd3vHe6d+KTMDZ1HwPAAKS151ipnpX3W3aRgi/FC6OMm +r0VUpWC4jCJlif3kYvaID66yFeucJveD75CJw6VaF4bSnTa9GmKM5F2BQhsrhKj+ +69jY+MU8cztGfJMdgczNfi79qDc8zSW6TUFpoMhmD6Og8icRCODbkbRWDDnnQi1d +1CVe/lZ4yIapCGgzHfQxdl9I8LEZWyeAfqMCvOTI2eE4Z4ZYto/EmuRoJN3oqOx1 +MAtpS5mw6RTFVitO4b0j5gwAN2u6Cgo8jDc5UTzHGAiGc/dAn7S+mNGv+IXP1/aL +UA/ZsBwQhZbqtGCkSmCo5gGDL+anDTiIy6nRCwIWWQKBgQDtE7V4pNj5J2MciK+h +KAR/kT36NuGoZSnf585HJVB/xAs/d8F4kU5pCofmrtm2Bl5Lx4sCE3wPmv/HMYcI +QH2zt/YxPA3F/9zoBowfTNB8/iCDlzNb/pUry+qkBhKG2HgWaUkkGEtiO16PeCF2 +bidYA+h+XZBnHKexAKL74ee71wKBgQDDVHJ6HRo1uTAT/8Uz3AF5NRj+qj9QC3UX +On6dtzbtDGsJFkcqW5CV+LwUPN6zQ7MW7U6AEeH99q1UAUDkd5Q8/4bqs+0y2Ss7 +0UUWr1LDJVYo6FIs/f2ON0e9LzcJcG9TgZ6T95TL1kGEG1Ng96Je3hAZWolLmr+x +yfjTSYKqOQKBgEwRWtTO799hx+dL5C5tTKQx0hUKrvT9IKZ7FjC1xFJ6cLF4l1c8 +KFCD1H8r8yb4fCEMcYnE/rVzIkajmZQIPU0A5bl+b1zsb9Dy6NrSJsM0NvKB/TSz +RuG6mBrw59jkdAOc3J78PJPUQM7/2JzLU0xmVJ7XHpI3G4crkSAIp/YZAoGATX6m +gF4ldOUI9xZFleKWTxFK3laLEeXJybJyY34582g22v8UsvBq96UccWcI79RPLCxw +NY1ivNBuSeLJbRsoG99BFsLVu5O/fFb1cx+R1Uxt14L8f08xlofGFX+y9TK/aEEH +uirCxPA3RANXXCRDLiIp/vUVfYJixVWdO65xgbkCgYAVQhvJ3jPXZLXPl9fVC3Ro +qweEZHs0knPc4mqL/7bE+BCE1HcccU7O7+ZKbhW2q11EMrPJ7cHUqweHE3GsTgRC +C5FIvRxInZUgtPNx2UIZAVNjxRdCZJHBc/59eu8wVp9f5ETiSFmdz5xlsJECSxdO +mBakpZMqnESl1/6FMtcqeQ== -----END PRIVATE KEY----- diff --git a/tests-integration/fixtures/etcd-tls-certs/ca.crt b/tests-integration/fixtures/etcd-tls-certs/ca.crt index 8723f7ffedd..ecd404c8918 100644 --- a/tests-integration/fixtures/etcd-tls-certs/ca.crt +++ b/tests-integration/fixtures/etcd-tls-certs/ca.crt @@ -1,21 +1,21 @@ -----BEGIN CERTIFICATE----- -MIIDeTCCAmGgAwIBAgIUeRD4bMVd0XEzS3UHo2ptnknsWzEwDQYJKoZIhvcNAQEL +MIIDeTCCAmGgAwIBAgIUM4xko7n+KubCgGJnX4EaUzP7ZicwDQYJKoZIhvcNAQEL BQAwTDELMAkGA1UEBhMCVVMxCzAJBgNVBAgMAkNBMQswCQYDVQQHDAJTRjERMA8G -A1UECgwIR3JlcHRpbWUxEDAOBgNVBAMMB2V0Y2QtY2EwHhcNMjUwODI1MTgwNjU1 -WhcNMjYwODI1MTgwNjU1WjBMMQswCQYDVQQGEwJVUzELMAkGA1UECAwCQ0ExCzAJ +A1UECgwIR3JlcHRpbWUxEDAOBgNVBAMMB2V0Y2QtY2EwHhcNMjYwODI2MDQxMzQy +WhcNMjcwODI2MDQxMzQyWjBMMQswCQYDVQQGEwJVUzELMAkGA1UECAwCQ0ExCzAJ BgNVBAcMAlNGMREwDwYDVQQKDAhHcmVwdGltZTEQMA4GA1UEAwwHZXRjZC1jYTCC -ASIwDQYJKoZIhvcNAQEBBQADggEPADCCAQoCggEBAJ8iLSJOWVRDI1iRr531m3fQ -gsDRAHMsofB7P1AmPxtAul8pHbp3v3gJCf847pc7sxawnqe+rBE1gzRZf5TBS25P -0yu21halLnIQBU7v10y0XruAI1p9cqNDIrwvW/PoC+rhck3hWpkSJ9MV2TezYyZL -S3C/iyCsr/+OHtpnMFAm3D5E571E3V2qUqhrdO2kHuLA1MPEaJHQRiPAeP4Z2I0e -0Mxa+udM01f8gp4TAs/8yOKvOvsado6NHbmu6pwu7Bau5+74QyGoVNSv4shClkPR -MW49foM+VBk6z/hSmO16gLzibkCdBc2WHS/ONtbV58RhZhCgoN6zH16gTmNwAGkC -AwEAAaNTMFEwHQYDVR0OBBYEFOCkI3Uyx7F38LtXtSlg3ORE9ur6MB8GA1UdIwQY -MBaAFOCkI3Uyx7F38LtXtSlg3ORE9ur6MA8GA1UdEwEB/wQFMAMBAf8wDQYJKoZI -hvcNAQELBQADggEBAJXbWJJ9b+5nXjRiFXOg+wQnzn7kzMf6sWKIFk+AuKJWWt78 -O00t6vyAz6zel5Cj3ho9yAaMFNy8vYEnJCYngy5pT/2hOncnz/w7IKTeoEhzqAnf -MCEmCgHbTKDoFfMfrrRwtyoePVfx4xbiGVWMQFTPG+WNlE/ivFMRvFwsgPJ+7SUK -mR2FscH6DVo4sqF6s729lTmr6/U7bOD5l2HYGpKJ2cjCI8+HDv55aAT43tsTwBoJ -BgX+RT4w8ryGhYV+hrRTjPlMHWiHsNbzJeAi5bNA0f6vUP6k6zHn/Ur3MNb1avcf -cSsdYU2dhs7PeB6IpoJ0QCeJw9MIFK3XTnmbeVY= +ASIwDQYJKoZIhvcNAQEBBQADggEPADCCAQoCggEBALTkN7ZOTxMxnQkwT7tGn7O+ +8uQQE89cIIVxd+j3WBhUcwyvVMZAm5W1vHusioLW0WFzaW4AZ+wGZbch9NdmPNjr +zfhlTOY8vu64F+hGl3Q9qPlfijCmfR4ciByyrHiO4dJNhzd92etMumqU0y6VkR6a +Q1QssITunLGDeGFgjH0kR72e/bp+TKvW9bnPVw+USnI4r5JQh6oeBeJW1Y9T1Ya/ +mx1B5ecPgTUlNfZRGKFvsLzFqWLdhVsjdJTHL+ACEzqzBCBnhIuBbyR/ztr49w/0 +LU3ka6N/Dh9P3R0mi63KJNJtwOuNkuMXTdOCNwNoAK+e1vSo3zlcP3p0i+GDmN8C +AwEAAaNTMFEwHQYDVR0OBBYEFDjvuYaWEoViLOKyCse6mAwH66mIMB8GA1UdIwQY +MBaAFDjvuYaWEoViLOKyCse6mAwH66mIMA8GA1UdEwEB/wQFMAMBAf8wDQYJKoZI +hvcNAQELBQADggEBAJuGxtprq/GF/fwUJ5hk0cG80VY24G+iKKi7Jw7PWPEMvnlQ +/ZIKgNDupE8mSSyokicKZMrtwT8QF6wlDLVcz4LTPCtKUmYU3FiuNE4V57+yw7zt +XZn6rwiWGxeyfY0ZktLXiSYA5FydOJ9GQQWgxkpbd5QG1L9aoFW49cbz8ob6efxG +p+YRgr0Fr4hXn7ImJWWqeawf2tg4iLwpV4sem5atpI7y9qamxrGGUp2h/CnADqUz +YWYIYP1RCPp2bXzxO40Ni1NYmFsGUU/SL1vtJw5n2Zmx/Phe7u/CLe05dH1O54Qf +vCbk6X0CthFI0HzSepbEVAuniVHkrpWwDLHcQ6s= -----END CERTIFICATE----- diff --git a/tests-integration/fixtures/etcd-tls-certs/ca.srl b/tests-integration/fixtures/etcd-tls-certs/ca.srl index 69ac026da25..ad32d06e16e 100644 --- a/tests-integration/fixtures/etcd-tls-certs/ca.srl +++ b/tests-integration/fixtures/etcd-tls-certs/ca.srl @@ -1 +1 @@ -4484A9BBD4F25F32994F3C03D80D294105B119E6 +4484A9BBD4F25F32994F3C03D80D294105B119E8 diff --git a/tests-integration/fixtures/etcd-tls-certs/client-key.pem b/tests-integration/fixtures/etcd-tls-certs/client-key.pem index 13a23c53478..1d64fc66a3d 100644 --- a/tests-integration/fixtures/etcd-tls-certs/client-key.pem +++ b/tests-integration/fixtures/etcd-tls-certs/client-key.pem @@ -1,28 +1,28 @@ -----BEGIN PRIVATE KEY----- -MIIEvAIBADANBgkqhkiG9w0BAQEFAASCBKYwggSiAgEAAoIBAQCOcrzRmE3sEGYo -pazDrpkkzbiaBhHClfmBiZAAGV2cvSXY8e9pc21GOgrGBNeyPPieS305rAE02EZl -oIAt97XRFxkZniUy8C5+PT3eicAKPRcrgGhIrYxUs39rdRocqVVpZsiZGzhHhboB -bXNW1l37SDQpuxAQJuhHu2wu9J7c5/TBXWN9ZOxzWTlrIdEQOA46VclE/N1TlaKs -CSWw8qH87YLwkLEXcFLOJfmWNZxOOhdL5MSzKKNcQKbpmkd65B/PHNwBIBZHpI+K -xphK1SybA+ESkN09qZaa1Kj3e0I62UjHNXUCuAod6RhRRmpiFDFgggXhD9T9Ttuf -1Lq38XxDAgMBAAECggEAANNhqQldwHzi9WCxXObLoufBe9lM73+1Df1IQaDHkZR3 -jHliixxF18XQW7qZ5tAJqjsDWLPm4BxSKXozsjsvBdgIUjbqkyx1dsGlO9TjeGu5 -H2/8jlH14boIMWJ9Y0JBWwvU7EnVU5IkUOEiwsxvFlzM5Q0H+841CrSD+Xd7iSZE -ZrpyZgpHNmvc/u7AmseZVccjX21dxwWtcd7cGUmFWC9o2a/hkgG6+baA8zH2ze3J -Smeowa/oOyOVkTM1EHLCBvddE+/W+BeMGPY7A6SztVw0gLqlSOL3cr/g2llUTmFg -T6MOS6FZZ7MD5c2HBL0V3OXIhxNABzD+c0O3CaRbSQKBgQC/+QuhMpDCzxBPaL/K -sqRp4sQgURDCXPA1nAfI+NoHos/CbjwcMGp6LUSiBW0tLTJ+LZoVgZXj6j/uxji9 -DWxQfULvQ7oanm/q/D+40NnupQQTzmMT04z/mV3XjFtzEhbO30YZfqwXy+tqB8vR -g1KajXVqLmCxKDtsFedAlIVxRQKBgQC99TKGcBwsKlYJJThU2eIIy8G0GkjjXleK -ROMO5i0keJtLac1oRvXo4JN5UBLJM6ZpSy0TEa2bxfzy6Uea+41aZmdCpPg0Kcsq -Pb1S3rsPpVdsfxSxsdKvwZmHXOzAdF+fom1+ZRAdo/8ALIPaNQSSaw+KkmmpHPEV -qdzzWJ4b5wKBgFubEtKcF3nuZwENoh+ueUhRvncRV+b3hGSAjTJ4lUn5hhxoj+R/ -sf+VJGAQKNXa8HJHfnRuvsDgYhulmSOViS8rZspXzjGvkwZV0m51stjvA3AUFzE5 -zNmXLLGTt3vEkP+siX3W9XXxh+ezyq2ydbNsdy/w65D9+sUL+qrVdIvlAoGAOOQy -2ajCB0g2tE59bIxE8jV0MiidI9uhhDvVdSTi6EVm3VM2vcBi7fg0suSUe8YIVQi6 -2zc0M688btQHKhek4ipBSuh1ncnWmzQae7NRewIeCNSWshF79D+bZ7sg/RLdgMX4 -3R4PkZEIUlkCtFuknuWJpgrrskaEveQ91HP6BokCgYAX42nKPANIMrTXuJ2vACWu -qKQ9atAC1ngVzYe2PI8rHjA4gCDraGfCaILewrd89rK9ukBeSDIm5Mn448q3ystw -p/QMxmMdHADCD6UwCjlt9N0Xigv+IKbDzDzhartRVIPpZdTViQXW3cSHt9id8zvO -HTu5C6W+N9Vb7j8Ak/mvzw== +MIIEvgIBADANBgkqhkiG9w0BAQEFAASCBKgwggSkAgEAAoIBAQC5F3hVsri2FBzP +dtRe+aJFh9yqhXnJ7zN+xbUmNWsTZVnPwrbVo9H6PtGWzkbZ6YMBzMlQSYP2txaS +ROjE8wwEs38eKbyJrHO2Ccq3iR9NfDxSbHTClngsOjf9fEoYExif84VKWLO6+EKL +KqqOfj3+iuI3hbabr0MqGu4e2Nr6MlQtViBxTs5aSnahc8hmZbujBcWRTsw/GNx5 +Am0a8BrrJRVmVR+BzlYtzbf1ltfLxIhnc8dAGyT9B6GTwcMNST6m4WtTzR/uOr6S +2YEbIqwCHB1VQ+oJUGJvExqaNLoM/x5a8sUNIWx3LBN0/XC4D+MMmh/bbyb2cAOE +gyUYu+hLAgMBAAECggEAEkt+wksVP9toEkecFpEqgss65VimdkPK/WkHDpzxawGL +FDfvq4Py9tRrvLIXRbseObjSCI9jrai2Ne9Sv1jeDhQ8kzDWk8MfItX/8A6Whf8A +KrjRUlneIfz/2Mcxaa9wMSs4YxC2dIYidFqwxhSLWyluJ8Ud4hs7kndb+beUmW3i +DgkLsQ9ytQ0HDQrAERnrbRf19fmVaqQapfxpYT2gg80BMIq+zlzoiXareDz3r/1g +aoponxArWV6yCvQOYYjWcDr4WZPvIA2APEX5isxyvMBvb72Fpd2z+0JX+4LVe6N1 +FqZLNpSeJn6O+Q1vJSScSJcoJYkLWCfZbFLHPttZGQKBgQDdVr5PtgPIVM+67cvv +p3x7NNC10Gp387PgQS/FlG+QtbP87RTBnANT1uFd/FMFqL35or7Y46YKDUdk00OZ +LinpM4ARFRYFqkthM6bcUWatpF6Fg2dfsKuD4H6ARp8YHhlmCnDCQ0swZp1HCmAB +6vltUNTZjQjRdd44GlQ0icX0mQKBgQDWE50TXRszea75S+zUfXeyvxMyptIOS/rY +L50XVuhGVknH4OuBQbGF6UZulIGSpN7gmmbO29Yx1HKmqiJVaWlvB71baK4hWLDd +2grIUy0zuq1B7I1IOLmOpMcL3xau76g7ZWRUprX26hxAuxYyPb9V8gNo1P7HN8si +VLJVE8ZugwKBgQCaxoGmM91JRSVNzeOB3ljJvxEDUo5g+uWZt3u0aivpwWXvQ8nz +6Sjag7RsiHl1x52w5wEVoXsGJGr8Mk9e2k0saXrwdxJDO+YiPoA8KB/o5LvEGTM8 +UspdGarcAIZX0xRnqn1XGr+FRPxOJQ8lyC5LJu7wghLchdOy35Zqdr0aYQKBgQC/ +38yFsonS1VnS8A5RVjOW7lPSrlrPnaIzalmutaJyiJyQnjP3Il5u2+rY6hpIyaVK +QpmrBrcw6m3om80yKMzrS1CZQXXxRYEhF3Fao9J77vGjiNYIyW7nPyF4rneyS/PJ +aNNIXDP0H1k7W3RFi7qW2dfceivxezyChM9iGdtc6QKBgAIfb63+1q98MNlAB7eD +x7K8vuSxHrJ9gX9KcLrG1ScLOxnAC4Op+8xf06ZHt828CajchAOqee2nh52MsFOr +AdPNdhS/8mTq8KosTfgVuawv4yTHpZViMR2KSOcTKhKaBmhufY+k8mJ0YAAOTnzO +JJqjDSJdjWsDG+7UlNgFcvuU -----END PRIVATE KEY----- diff --git a/tests-integration/fixtures/etcd-tls-certs/client.crt b/tests-integration/fixtures/etcd-tls-certs/client.crt index 935351a48b1..0dee5915ffa 100644 --- a/tests-integration/fixtures/etcd-tls-certs/client.crt +++ b/tests-integration/fixtures/etcd-tls-certs/client.crt @@ -1,20 +1,20 @@ -----BEGIN CERTIFICATE----- -MIIDMjCCAhqgAwIBAgIURISpu9TyXzKZTzwD2A0pQQWxGeYwDQYJKoZIhvcNAQEL +MIIDMjCCAhqgAwIBAgIURISpu9TyXzKZTzwD2A0pQQWxGegwDQYJKoZIhvcNAQEL BQAwTDELMAkGA1UEBhMCVVMxCzAJBgNVBAgMAkNBMQswCQYDVQQHDAJTRjERMA8G -A1UECgwIR3JlcHRpbWUxEDAOBgNVBAMMB2V0Y2QtY2EwHhcNMjUwODI1MTgwNjU1 -WhcNMjYwODI1MTgwNjU1WjAWMRQwEgYDVQQDDAtldGNkLWNsaWVudDCCASIwDQYJ -KoZIhvcNAQEBBQADggEPADCCAQoCggEBAI5yvNGYTewQZiilrMOumSTNuJoGEcKV -+YGJkAAZXZy9Jdjx72lzbUY6CsYE17I8+J5LfTmsATTYRmWggC33tdEXGRmeJTLw -Ln49Pd6JwAo9FyuAaEitjFSzf2t1GhypVWlmyJkbOEeFugFtc1bWXftINCm7EBAm -6Ee7bC70ntzn9MFdY31k7HNZOWsh0RA4DjpVyUT83VOVoqwJJbDyofztgvCQsRdw -Us4l+ZY1nE46F0vkxLMoo1xApumaR3rkH88c3AEgFkekj4rGmErVLJsD4RKQ3T2p -lprUqPd7QjrZSMc1dQK4Ch3pGFFGamIUMWCCBeEP1P1O25/UurfxfEMCAwEAAaNC -MEAwHQYDVR0OBBYEFINcPBk3iqCw1B2sZzcDpSF6ZNGnMB8GA1UdIwQYMBaAFOCk -I3Uyx7F38LtXtSlg3ORE9ur6MA0GCSqGSIb3DQEBCwUAA4IBAQBgyF1XSpAyiQNA -sFCqXrBfwPnv6V+R5jhO6Glmn+Nhj0gbzNBlU8EKsK/WZllDrJjmxubPBqb563Zt -J+QcNLqugZdioXPbxmmi9xj3oK25pUCRBj03nlHpGqIwuCJyn/ZxlfXMpNYE1KjM -nprFNUKgnYPk2AGffAngQdm0yrnmSpxXhNTuhQE8avgRhd5nEUEDyPIvHWwo4iU6 -pw2aP2FpWorESv9LVna+JortThYDQrIFgCtPM6BP1N7Aeqn0H846+JsXlq2ch1Iy -aBRqKNQRc1Q2Qb9/QCDEJ3dEixgFjxztIaqrIcr18ZJ797MhTFS3aVEwE8clSRR5 -9ow9jK8B +A1UECgwIR3JlcHRpbWUxEDAOBgNVBAMMB2V0Y2QtY2EwHhcNMjYwODI2MDQxMzQy +WhcNMjcwODI2MDQxMzQyWjAWMRQwEgYDVQQDDAtldGNkLWNsaWVudDCCASIwDQYJ +KoZIhvcNAQEBBQADggEPADCCAQoCggEBALkXeFWyuLYUHM921F75okWH3KqFecnv +M37FtSY1axNlWc/CttWj0fo+0ZbORtnpgwHMyVBJg/a3FpJE6MTzDASzfx4pvIms +c7YJyreJH018PFJsdMKWeCw6N/18ShgTGJ/zhUpYs7r4Qosqqo5+Pf6K4jeFtpuv +Qyoa7h7Y2voyVC1WIHFOzlpKdqFzyGZlu6MFxZFOzD8Y3HkCbRrwGuslFWZVH4HO +Vi3Nt/WW18vEiGdzx0AbJP0HoZPBww1JPqbha1PNH+46vpLZgRsirAIcHVVD6glQ +Ym8TGpo0ugz/HlryxQ0hbHcsE3T9cLgP4wyaH9tvJvZwA4SDJRi76EsCAwEAAaNC +MEAwHQYDVR0OBBYEFD1aTqJO87ovA2usg/v8UckSsnwhMB8GA1UdIwQYMBaAFDjv +uYaWEoViLOKyCse6mAwH66mIMA0GCSqGSIb3DQEBCwUAA4IBAQArP3mtjdj5a9uT +XTys4DUjp+CJaSmsZ/4ICf/KeddaEyijx6uoiqbNsAtfwlD5A5BtyjJrxg3SpjqD +Fm5mTV0HlwvYHxcFPTUlR56//sOyeRhm+kVPp4xdVq2p6IZXFF/hDkeT/AETWOjj +l4+90gAsJsUHrKUZmzlqHBeOXoAhzo2PwSh8KLcGK0705FGql4QnS731PHo9a9qp +e8uxk1KKoc+RHm2OjJvZ7hvkCXuCm3TJeOHQEtjZ5buQg40BRrt11GVnsPePpnQR +8OtV0DvPT3UcOss88kLDQuRuZeJf8UyaevGGvItv3+Yk5cwUBSaJ0mY6sgrmkXn0 +TstUzYIT -----END CERTIFICATE----- diff --git a/tests-integration/fixtures/etcd-tls-certs/server-key.pem b/tests-integration/fixtures/etcd-tls-certs/server-key.pem index fca8a0122f9..a73ac974434 100644 --- a/tests-integration/fixtures/etcd-tls-certs/server-key.pem +++ b/tests-integration/fixtures/etcd-tls-certs/server-key.pem @@ -1,28 +1,28 @@ -----BEGIN PRIVATE KEY----- -MIIEvQIBADANBgkqhkiG9w0BAQEFAASCBKcwggSjAgEAAoIBAQCtYFxH+oet6jTD -0eF55ISO+bwm1NPTluP9zLJrQFNDkqlddyiIasK5q7qN1pPXOw9mJF/EN8Hf1lWj -42an6tr7PgE0NHLcOrqDv4dRjegRi87IWiPkfep8peUE3glI4IOMXfZmAUNp+nRO -JP/ymmR0pzHYJlx+D082OYmeBny7uFQxo+vxgVNo2MVVwy6oTv+jtjLgZ4iOrIvh -iILoXtXcdg/tjB1rS2hSABGPJMNHwKswcrt6xXv+7Llsxroq0cx0E2PrAbc5qB40 -v1Qp7PjFTJ0drgf0g51LkOb7znpB1viPD2BLMxpELTZqa4OruarcDNxShrK8WWMP -1mygLHLxAgMBAAECggEAC4TgtLaQP62VImKGLs1QQliU1+ahiUgX70Ojog0UyyNK -JefWFVQ0ilX+z9AvI+hsYj6t7zE+LBNHPtuLtUHdGT66IUAP1pJ/VGQMB07cmZfW -ngihJFv6UZxLDkr7RnCGRPP0PDw+wKKPiiaaq8F2xapbHS+VSxnUyzdA7bMkI+uj -+gLAFj+kEU0FNbIp/mBzz8sksDzcB29oy3VnZkR1211KkL1BJDvr9oKiUg8NFKaN -NF2WXRYI/m0CQMlX6zlTyv7+h/Ay7y8oNG7cfHavmCrWvYIMfaDggYxbIpNdGxb5 -MAdbrtXJ6daUBa2sezxgQUyJX8+8GGF2zpGvCfhkfQKBgQDiAwA8rlxrBnXclwRK -9s/apILKCxPEGwU6XZi/8na3ugIO14y2WD0zZTjUK/3xwoKN6j6+ObLkUBX2RRYZ -ewwxpqos2cyxi5EBKPhsydqbJQlat68I5+NOG4KOZQkqO92tBtgY3gGFWVntaqtE -UVE/Clh2gFcc/5nR/5C5Mil8LQKBgQDEYXtrl2kl4qcyXuuXKDA/bLSnDuIEMhLO -4oIiqmywXQgjPjGAWy0eztjMq2a353pqTFCqto/cfOif+rjkB7oFfi9gMm0plIKI -+F8QsBB7HpGqWos3zb1yt8X9hbxKuFFwZS/un3Li4lwPlk8EZ1YIvhC44BHAvNkD -VzxzBwoYVQKBgArRe/RroC7bS074x5LTB5X+o+gJ6bNMW860Zjhh4b7fn3OYa7ra -tGs+YB7/0BL/bYJfgQtX9bEqCDMWkX08v5Os1553+m1RMeqtTF7gtp8Qgcce3bj+ -aIn3lSM9wNeNsAm1NyjRj58TbNOJdJM7lTkARMW/VOwla/Z6VjIXLZctAoGBAJKf -xira7eMfi36MaJJ/qyZv36Ir9ozzZh+Z91gyrtwvWfgWY5dWfCXYgv6tqw/8gOYE -/OW5UUhq6rUn2gxHyJh5Up4ciGzXOW9TIoevLV7/v/rVh8SulJimpelYhPG1FPk6 -U8NywbCtGdd5fp3nGdGFN68Rfa/OUKmx5KxtwRfRAoGAWpfzmoI8EmCBZNXYk03k -hnDKmICsKi9seDclcQq2PbyNvE7LDEd8O4pDnwV6RwKVEDTer9Y4IGcjkEBZcJVm -za3l4En9+M27Xlf1I8XeB8o8RYbFjoDQyScFFIgQpvTDsOezXeDR55FGPZN7V6Si -lKmd+8BuzJ7d31KJtRjoS1U= +MIIEvgIBADANBgkqhkiG9w0BAQEFAASCBKgwggSkAgEAAoIBAQDQLtVfdk9l1rdH +/ssuy/BK9LpmNPx7gatp/E3EqpgekBjizMXIxnCHDFNM4Dtq1v3afaxeijXnnW/b +0CVKSeGkOBmb0dO0K21xe86Y5VfNHkGotQyCZ8mPOzep/EQOGJY5+dprErL0YKM+ +2drgiRNGwidmNk/jD3OtzXSMR+8UvDWJSxUoZ1ncvgmqf3zTTiaTiABHrfLdXaGI +23+JftyLgoeU4vEI+x+RQJMAa6Gzjo9vALxX3kTtR/uH34hfm8vb3yo68UHGKVMp +i3yi6YDP4YEao5lOpm7fIyEes10fPkD224++S34KHykAk4q+eTPetWkDYFCYzdSp +AR2Ob3cxAgMBAAECggEAMDgAxO0mw7xBVGQgFJU48WuQtvaj2kl09gXxz1ECDeYr +VXC/iNrpmmYQ7zfqmzrzrkU4hOc3SA/PplamJHhLUpmJ2Oz3P35liYj3F7PbK8/L +vnM81AGNDmdVY8Jh0u//76q+29kHaRHvDbIw/5vQQq3aqVKAG2PrU8DIM2u/5Qmt +0GHJTMmyVADAOLe7j2agN9hrOrjeoWobcELwtwgKXxmU0l+g0XX1OdJpbjA3qDPa +S5a4MzcucGINeB3bohcqbcx2f1HwVUoC19mti3Cz+KvxF2YZV9bS9izW4hSEsAwU +0upNHF3MwTUQg7ftLPAZcfNoXIZMFXwYMS1HW2v7HwKBgQDuHHkFRN2lvor/Ycjr +gJD6LGb9RlQ7NMYBQHB+5g49Y/q8oDZJgyl4CfgYWdVjjOfEX8oMst1WU/Stbva5 +HlHR7yYFyebScQYOPHZaPyVKGIbcl2p7Qy9H4pxJc1ZMNxO8SxoSC68LZEan08FD +/Qsish51fZ2Sqf37b+8Uz/CGdwKBgQDf0sIp88MPGbhtzdYtF9gA4AoIOPo9UxeI +F6wTb05uCqZyK+hcdK37av1xcNpczd/yvlALxRdM9xGedJkcS2cgeExtMRcfhAzL +as0dYPzuhRkpEIfN/y9cqtHj2wigrbu35eDgRRnJDgXgL9dYl6r+Swqn3jBaasP7 +KtQdQb/RlwKBgQC7qe0n3fLi4p4iUStNkPKyebRiAb/5Ocqkyejf2ul2MQo5B/xB +TAKu/QxwBL1NzIwOFYDlKUOQ+nJpDn+dvuu1jcpl3Y7yZOnk5npQ/luhXltMGHpv +06+79DpBGYn2X6JKUNanSlYXoFyfgSFdOF5CZifjabF7Gkd2l+3SdWCYWQKBgQDB +3k4sBGZKaB7ljUscl/CTIXvPD3tBLv3M9aQo2Vp32mW9suZ7Xt1sTonkfrnFdNWr +7shqyXabRc5PD/OnHHDhIRIh6kl7FOf4MjQkZGPxPfxDI3xeI9EkVRmkYY6hjppw +eX9FAtWI3sqcGxROOmD0Do/WQ5BiYOQMZFaCWPcLVQKBgGVT+iUYkYTs+UJeD4RC +QeL6QEJDszvAI/gavUp7sD/LWn7Ig+hJiAl2O+HFYMurh7eda0JnE2KF40UyPiYA +DDl19xRAXrLPOlOj9YKDA9HkROcaVZCQCqyDalKsbw1IqMrAkRmLCIZW3EHaMrUf +UNjg+GaxwvrRj11gbd4IPgai -----END PRIVATE KEY----- diff --git a/tests-integration/fixtures/etcd-tls-certs/server.crt b/tests-integration/fixtures/etcd-tls-certs/server.crt index 6ed8b1054e7..8dcb7b13b6d 100644 --- a/tests-integration/fixtures/etcd-tls-certs/server.crt +++ b/tests-integration/fixtures/etcd-tls-certs/server.crt @@ -1,21 +1,21 @@ -----BEGIN CERTIFICATE----- -MIIDjDCCAnSgAwIBAgIURISpu9TyXzKZTzwD2A0pQQWxGeUwDQYJKoZIhvcNAQEL +MIIDjDCCAnSgAwIBAgIURISpu9TyXzKZTzwD2A0pQQWxGecwDQYJKoZIhvcNAQEL BQAwTDELMAkGA1UEBhMCVVMxCzAJBgNVBAgMAkNBMQswCQYDVQQHDAJTRjERMA8G -A1UECgwIR3JlcHRpbWUxEDAOBgNVBAMMB2V0Y2QtY2EwHhcNMjUwODI1MTgwNjU1 -WhcNMjYwODI1MTgwNjU1WjATMREwDwYDVQQDDAhldGNkLXRsczCCASIwDQYJKoZI -hvcNAQEBBQADggEPADCCAQoCggEBAK1gXEf6h63qNMPR4XnkhI75vCbU09OW4/3M -smtAU0OSqV13KIhqwrmruo3Wk9c7D2YkX8Q3wd/WVaPjZqfq2vs+ATQ0ctw6uoO/ -h1GN6BGLzshaI+R96nyl5QTeCUjgg4xd9mYBQ2n6dE4k//KaZHSnMdgmXH4PTzY5 -iZ4GfLu4VDGj6/GBU2jYxVXDLqhO/6O2MuBniI6si+GIguhe1dx2D+2MHWtLaFIA -EY8kw0fAqzByu3rFe/7suWzGuirRzHQTY+sBtzmoHjS/VCns+MVMnR2uB/SDnUuQ -5vvOekHW+I8PYEszGkQtNmprg6u5qtwM3FKGsrxZYw/WbKAscvECAwEAAaOBnjCB +A1UECgwIR3JlcHRpbWUxEDAOBgNVBAMMB2V0Y2QtY2EwHhcNMjYwODI2MDQxMzQy +WhcNMjcwODI2MDQxMzQyWjATMREwDwYDVQQDDAhldGNkLXRsczCCASIwDQYJKoZI +hvcNAQEBBQADggEPADCCAQoCggEBANAu1V92T2XWt0f+yy7L8Er0umY0/HuBq2n8 +TcSqmB6QGOLMxcjGcIcMU0zgO2rW/dp9rF6KNeedb9vQJUpJ4aQ4GZvR07QrbXF7 +zpjlV80eQai1DIJnyY87N6n8RA4Yljn52msSsvRgoz7Z2uCJE0bCJ2Y2T+MPc63N +dIxH7xS8NYlLFShnWdy+Cap/fNNOJpOIAEet8t1doYjbf4l+3IuCh5Ti8Qj7H5FA +kwBrobOOj28AvFfeRO1H+4ffiF+by9vfKjrxQcYpUymLfKLpgM/hgRqjmU6mbt8j +IR6zXR8+QPbbj75LfgofKQCTir55M961aQNgUJjN1KkBHY5vdzECAwEAAaOBnjCB mzAJBgNVHRMEAjAAMAsGA1UdDwQEAwIEMDBBBgNVHREEOjA4gglsb2NhbGhvc3SC CGV0Y2QtdGxzggkxMjcuMC4wLjGHBH8AAAGHEAAAAAAAAAAAAAAAAAAAAAEwHQYD -VR0OBBYEFK8y9gWv8Si9AusP280BrgDF6y3jMB8GA1UdIwQYMBaAFOCkI3Uyx7F3 -8LtXtSlg3ORE9ur6MA0GCSqGSIb3DQEBCwUAA4IBAQBMBsAPIlk17I2ioN9xwxwl -yYbSjZlTs+18wSZAoCNLfxwWIYAQ08dPoEUdALfMUPEEe1Ol2IAp/qx2JoDrZWje -maA43hppBUpFFCkTKWUMYxsetN7d7BVxWL43GDoMwoD/k36nhoUKBUjlbF0+nkem -dOMr8SA4GZbEg7qk2cL93g0UHv/Z/dZgKf3epZkR9hEN4/R2jSP6OPmY9XIvQCoZ -8D9jgbQKavAAXmUcP5a81alQMRroGaBCzI1f5OlS3EuVE7ZTEBgxKK1idbouSrPt -4UHEbOGz9zXDjmfut6CTm247+lJm9jzYe2Xx+XGw29l0pzd/8tGFN9zIvW8mCv7a +VR0OBBYEFNqGaSuuVHwEIvQvYxt9H0kdWFPSMB8GA1UdIwQYMBaAFDjvuYaWEoVi +LOKyCse6mAwH66mIMA0GCSqGSIb3DQEBCwUAA4IBAQAHttiuyEZieeTsggAKzAIC +UXsjWCQMAmjRxery6mzM8P33CdCKA4cwGyy6tzoFPiZC/wh+gPq+LiUO9G+vkqNA +NQ/DSJS+fsqsOBWUt5DSYekD25o80mZVE/4H/ynFOgENBqgkVl1ikpkXu9BK8wTm +Wij+p9RWyNup8lr9OFYZRpy7wrxPp7opwp3qG2Bfu/BGY1YEr04CIrLMKnk1zRMm +73X1lcrlrM9ycCmYcfXxyU69G0iFKaPcbS7T0fngm1NciA1qxYYLuMyzDtF7/SyC +J2ynvDBBnc9+/X8piu8p9D0I9mbcbEySv6zZfYMT0kqXbsn2GijZTH5mCJ3P8cwc -----END CERTIFICATE----- diff --git a/tests-integration/src/cluster.rs b/tests-integration/src/cluster.rs index d0cd2e584c3..e40c707fe98 100644 --- a/tests-integration/src/cluster.rs +++ b/tests-integration/src/cluster.rs @@ -169,6 +169,7 @@ pub struct GreptimeDbClusterBuilder { metasrv_wal_config: MetasrvWalConfig, datanode_gc_config: GcConfig, metasrv_gc_config: GcSchedulerOptions, + frontend_auto_create_table: bool, shared_home_dir: Option>, meta_selector: Option, } @@ -202,6 +203,7 @@ impl GreptimeDbClusterBuilder { metasrv_wal_config: MetasrvWalConfig::default(), datanode_gc_config: GcConfig::default(), metasrv_gc_config: GcSchedulerOptions::default(), + frontend_auto_create_table: true, shared_home_dir: None, meta_selector: None, } @@ -248,6 +250,13 @@ impl GreptimeDbClusterBuilder { self } + /// Sets whether the test frontend automatically creates tables on write. + #[must_use] + pub fn with_frontend_auto_create_table(mut self, auto_create_table: bool) -> Self { + self.frontend_auto_create_table = auto_create_table; + self + } + #[must_use] pub fn with_shared_home_dir(mut self, shared_home_dir: Arc) -> Self { self.shared_home_dir = Some(shared_home_dir); @@ -529,7 +538,10 @@ impl GreptimeDbClusterBuilder { } fn build_frontend_options(&self) -> FrontendOptions { - let mut fe_opts = FrontendOptions::default(); + let mut fe_opts = FrontendOptions { + auto_create_table: self.frontend_auto_create_table, + ..Default::default() + }; // Choose a random unused port between [14000, 24000] for local test to avoid conflicts. let port_range = 14000..=24000; diff --git a/tests-integration/src/grpc/flight.rs b/tests-integration/src/grpc/flight.rs index cf75ae7fff1..1a2d841600f 100644 --- a/tests-integration/src/grpc/flight.rs +++ b/tests-integration/src/grpc/flight.rs @@ -16,40 +16,45 @@ mod test { use std::collections::HashMap; use std::net::SocketAddr; + use std::pin::Pin; use std::sync::Arc; use std::time::{Duration, Instant}; use api::v1::auth_header::AuthScheme; + use api::v1::greptime_request::Request as GreptimeQueryRequest; use api::v1::query_request::Query; use api::v1::{Basic, ColumnDataType, ColumnDef, CreateTableExpr, QueryRequest, SemanticType}; use arrow_flight::flight_service_server::FlightServiceServer; use arrow_flight::{FlightData, FlightDescriptor, Ticket}; use auth::user_provider_from_option; + use client::region::RegionRequester; use client::{Client, Database}; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; use common_grpc::channel_manager::{ChannelConfig, ChannelManager}; - use common_grpc::flight::do_put::DoPutMetadata; + use common_grpc::flight::do_put::{DoPutMetadata, DoPutResponse}; use common_grpc::flight::{FlightEncoder, FlightMessage}; - use common_query::OutputData; + use common_query::{Output, OutputData}; use common_recordbatch::adapter::RegionWatermarkEntry; use common_recordbatch::{RecordBatch, RecordBatches, SendableRecordBatchStream}; use common_telemetry::tracing_context::TracingContext; use datatypes::prelude::{ConcreteDataType, ScalarVector, VectorRef}; use datatypes::schema::{ColumnSchema, Schema}; use datatypes::vectors::{Int32Vector, StringVector, TimestampMillisecondVector}; - use futures_util::StreamExt; + use futures_util::{Stream, StreamExt}; use hyper_util::rt::TokioIo; use itertools::Itertools; use servers::grpc::builder::GrpcServerBuilder; use servers::grpc::flight::{ FlightCraft, FlightCraftWrapper, FlightRecordBatchSource, FlightRecordBatchStream, - TonicStream, + FlightRecordBatchStreamInput, PutRecordBatchRequestStream, TonicStream, }; use servers::grpc::greptime_handler::GreptimeRequestHandler; use servers::grpc::{FlightCompression, GrpcServerConfig}; + use servers::query_handler::grpc::GrpcQueryHandler; use servers::server::Server; - use tonic::Response; + use session::context::QueryContextRef; use tonic::transport::Server as TonicServer; + use tonic::{Response, Status}; use tower::service_fn; use crate::cluster::GreptimeDbClusterBuilder; @@ -59,6 +64,12 @@ mod test { struct SlowFlightCraft; + struct ErrorFlightCraft; + + struct SlowRemoteQueryHandler { + region_requester: RegionRequester, + } + fn slow_recordbatch_stream() -> SendableRecordBatchStream { let schema = Arc::new(Schema::new(vec![ColumnSchema::new( "value", @@ -83,9 +94,11 @@ mod test { _: tonic::Request, ) -> std::result::Result>, tonic::Status> { let stream = FlightRecordBatchStream::new( - FlightRecordBatchSource::initializer(async { + FlightRecordBatchStreamInput::initializer(async { tokio::time::sleep(Duration::from_secs(2)).await; - Ok(slow_recordbatch_stream()) + Ok(FlightRecordBatchSource::RecordBatches( + slow_recordbatch_stream(), + )) }), TracingContext::default(), FlightCompression::default(), @@ -96,14 +109,72 @@ mod test { } } - #[tokio::test(flavor = "multi_thread")] - async fn test_do_get_timeout_does_not_cancel_slow_flight_stream() { + #[async_trait::async_trait] + impl FlightCraft for ErrorFlightCraft { + async fn do_get( + &self, + _: tonic::Request, + ) -> std::result::Result>, tonic::Status> { + let stream = FlightRecordBatchStream::new( + FlightRecordBatchStreamInput::initializer(async { + Err(Status::internal("deferred initializer detail")) + }), + TracingContext::default(), + FlightCompression::default(), + session::context::QueryContext::arc(), + ); + Ok(Response::new(Box::pin(stream))) + } + } + + #[async_trait::async_trait] + impl GrpcQueryHandler for SlowRemoteQueryHandler { + async fn do_query( + &self, + _: GreptimeQueryRequest, + _: QueryContextRef, + ) -> servers::error::Result { + let stream = self + .region_requester + .do_get_inner(Ticket::default()) + .await + .unwrap(); + Ok(Output::new_with_stream(stream)) + } + + fn handle_put_record_batch_stream( + &self, + _: PutRecordBatchRequestStream, + _: QueryContextRef, + ) -> Pin> + Send>> { + Box::pin(futures::stream::empty()) + } + } + + fn client_for_flight_craft(addr: &'static str, craft: T) -> Client + where + T: FlightCraft, + { + client_for_flight_craft_with_max_encoding(addr, craft, None) + } + + fn client_for_flight_craft_with_max_encoding( + addr: &'static str, + craft: T, + max_encoding_message_size: Option, + ) -> Client + where + T: FlightCraft, + { let (client_io, server_io) = tokio::io::duplex(1024); tokio::spawn(async move { + let flight_service = FlightServiceServer::new(FlightCraftWrapper(craft)); + let flight_service = match max_encoding_message_size { + Some(size) => flight_service.max_encoding_message_size(size), + None => flight_service, + }; TonicServer::builder() - .add_service(FlightServiceServer::new(FlightCraftWrapper( - SlowFlightCraft, - ))) + .add_service(flight_service) .serve_with_incoming(futures::stream::iter(vec![Ok::<_, std::io::Error>( server_io, )])) @@ -115,7 +186,7 @@ mod test { let mut client_io = Some(client_io); channel_manager .reset_with_connector( - "slow-flight", + addr, service_fn(move |_| { let client_io = client_io.take(); @@ -127,7 +198,12 @@ mod test { }), ) .unwrap(); - let client = Client::with_manager_and_urls(channel_manager, ["slow-flight"]); + Client::with_manager_and_urls(channel_manager, [addr]) + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_do_get_timeout_does_not_cancel_slow_flight_stream() { + let client = client_for_flight_craft("slow-flight", SlowFlightCraft); let mut flight_client = client.make_flight_client(false, false).unwrap(); let start = Instant::now(); @@ -143,6 +219,64 @@ mod test { assert!(stream.message().await.unwrap().is_some()); } + + #[tokio::test(flavor = "multi_thread")] + async fn test_deferred_flight_initializer_error_preserves_message() { + let client = client_for_flight_craft("error-flight", ErrorFlightCraft); + let mut flight_client = client.make_flight_client(false, false).unwrap(); + let mut stream = flight_client + .mut_inner() + .do_get(tonic::Request::new(Ticket::default())) + .await + .unwrap() + .into_inner(); + let error = stream.message().await.unwrap_err(); + assert_eq!(tonic::Code::Internal, error.code()); + assert_eq!("deferred initializer detail", error.message()); + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_deferred_flight_encode_error_preserves_message() { + let client = + client_for_flight_craft_with_max_encoding("limited-flight", SlowFlightCraft, Some(1)); + let mut flight_client = client.make_flight_client(false, false).unwrap(); + let mut stream = flight_client + .mut_inner() + .do_get(tonic::Request::new(Ticket::default())) + .await + .unwrap() + .into_inner(); + let error = stream.message().await.unwrap_err(); + assert_eq!(tonic::Code::OutOfRange, error.code()); + assert!(error.message().contains("message length too large")); + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_flight_request_timeout_does_not_cancel_slow_datanode_stream() { + let datanode_client = client_for_flight_craft("slow-datanode", SlowFlightCraft); + let frontend_handler = GreptimeRequestHandler::new( + Arc::new(SlowRemoteQueryHandler { + region_requester: RegionRequester::new(datanode_client, false, false), + }), + None, + None, + FlightCompression::default(), + ); + let frontend_client = client_for_flight_craft("slow-frontend", frontend_handler); + let database = Database::new(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, frontend_client); + + let start = Instant::now(); + let output = database + .flight_request() + .with_timeout(Duration::from_secs(1)) + .sql("select 1") + .await + .unwrap(); + // `sql()` waits for the first Flight message. Succeeding after this delay + // proves the request timeout applied only to the response header. + assert!(start.elapsed() >= Duration::from_secs(1)); + assert!(matches!(output.data, OutputData::Stream(_))); + } #[tokio::test(flavor = "multi_thread")] async fn test_standalone_flight_do_put() { common_telemetry::init_default_ut_logging(); @@ -446,16 +580,14 @@ mod test { ); let result = client - .query_with_terminal_metrics_and_flow_extensions( - QueryRequest { - query: Some(Query::Sql( - "select ts, a, `B` from foo order by ts".to_string(), - )), - }, - &[], - &[("flow.return_region_seq", "true")], - &HashMap::from([(region_id.as_u64(), stale_snapshot_seq)]), - ) + .flight_request() + .with_flow_extensions(&[("flow.return_region_seq", "true")]) + .with_snapshot_seqs(&HashMap::from([(region_id.as_u64(), stale_snapshot_seq)])) + .query_with_terminal_metrics(QueryRequest { + query: Some(Query::Sql( + "select ts, a, `B` from foo order by ts".to_string(), + )), + }) .await .unwrap(); diff --git a/tests-integration/src/instance.rs b/tests-integration/src/instance.rs index 08d2a1aff15..9f48d01a467 100644 --- a/tests-integration/src/instance.rs +++ b/tests-integration/src/instance.rs @@ -691,4 +691,119 @@ mod tests { "denied source leaked into entities:\n{pretty_print}" ); } + + /// Denying the descriptor's table must not fail the request: the metric + /// data is already committed, so the client would retry what the server + /// accepted. + #[tokio::test(flavor = "multi_thread")] + async fn test_otlp_descriptor_permission_denial_degrades() { + use otel_arrow_rust::proto::opentelemetry::collector::metrics::v1::ExportMetricsServiceRequest; + use otel_arrow_rust::proto::opentelemetry::common::v1::{AnyValue, KeyValue, any_value}; + use otel_arrow_rust::proto::opentelemetry::metrics::v1::number_data_point::Value; + use otel_arrow_rust::proto::opentelemetry::metrics::v1::{ + Gauge, Metric, NumberDataPoint, ResourceMetrics, ScopeMetrics, metric, + }; + use otel_arrow_rust::proto::opentelemetry::resource::v1::Resource; + use servers::query_handler::OpenTelemetryProtocolHandler; + + common_telemetry::init_default_ut_logging(); + + /// Denies writes to the descriptor table only. + struct DenyResourceInfo; + + impl auth::PermissionChecker for DenyResourceInfo { + fn check_permission( + &self, + _user_info: auth::UserInfoRef, + _req: auth::PermissionReq, + ) -> auth::error::Result { + Ok(auth::PermissionResp::Allow) + } + + fn check_permission_with_table_targets( + &self, + _user_info: auth::UserInfoRef, + _req: auth::PermissionReq, + targets: auth::PermissionTableTargets, + ) -> auth::error::Result { + if let auth::PermissionTableTargets::Resolved(targets) = targets + && targets + .iter() + .any(|target| target.table == "greptime_otel_resource_info") + { + return Ok(auth::PermissionResp::Reject); + } + Ok(auth::PermissionResp::Allow) + } + } + + let plugins = Plugins::new(); + plugins.insert::(Arc::new(DenyResourceInfo)); + + let standalone = GreptimeDbStandaloneBuilder::new("test_otlp_descriptor_denied") + .with_plugin(plugins) + .build() + .await; + let instance = standalone.fe_instance().clone(); + + let attr = |key: &str, value: &str| KeyValue { + key: key.to_string(), + value: Some(AnyValue { + value: Some(any_value::Value::StringValue(value.to_string())), + }), + }; + let request = ExportMetricsServiceRequest { + resource_metrics: vec![ResourceMetrics { + resource: Some(Resource { + attributes: vec![attr("service.name", "api"), attr("host.id", "h-1")], + ..Default::default() + }), + scope_metrics: vec![ScopeMetrics { + metrics: vec![Metric { + name: "denied_descriptor_gauge".to_string(), + data: Some(metric::Data::Gauge(Gauge { + data_points: vec![NumberDataPoint { + time_unix_nano: 1_700_000_000_000_000_000, + value: Some(Value::AsDouble(1.0)), + ..Default::default() + }], + })), + ..Default::default() + }], + ..Default::default() + }], + ..Default::default() + }], + }; + + let outcome = + OpenTelemetryProtocolHandler::metrics(&*instance, request, QueryContext::arc()) + .await + .unwrap(); + let warning = outcome + .error_message + .expect("denial must surface as a warning"); + assert!( + warning.contains("greptime_otel_resource_info"), + "unexpected warning: {warning}" + ); + + // the metric data is committed, the descriptor table was never created + let sql = "select table_name from information_schema.tables \ + where table_schema = 'public' order by table_name"; + let output = query(&instance, sql).await; + let OutputData::Stream(s) = output.data else { + unreachable!() + }; + let batches = common_recordbatch::util::collect_batches(s).await.unwrap(); + let pretty_print = batches.pretty_print().unwrap(); + assert!( + pretty_print.contains("denied_descriptor_gauge"), + "metric data was not committed:\n{pretty_print}" + ); + assert!( + !pretty_print.contains("greptime_otel_resource_info"), + "denied descriptor table was created:\n{pretty_print}" + ); + } } diff --git a/tests-integration/src/standalone.rs b/tests-integration/src/standalone.rs index 1085c21983e..0704b184e89 100644 --- a/tests-integration/src/standalone.rs +++ b/tests-integration/src/standalone.rs @@ -27,6 +27,7 @@ use common_base::Plugins; use common_catalog::consts::{MIN_USER_FLOW_ID, MIN_USER_TABLE_ID}; use common_config::KvBackendConfig; use common_datasource::object_store::LocalFileAccess; +use common_event_recorder::EventRecorderOptions; use common_meta::cache::LayeredCacheRegistryBuilder; use common_meta::ddl::flow_meta::FlowMetadataAllocator; use common_meta::ddl::table_meta::TableMetadataAllocator; @@ -84,6 +85,7 @@ pub struct GreptimeDbStandaloneBuilder { default_store: Option, plugin: Option, slow_query_options: SlowQueryOptions, + event_recorder_options: EventRecorderOptions, auto_create_table: bool, } @@ -102,6 +104,7 @@ impl GreptimeDbStandaloneBuilder { threshold: Duration::from_secs(1), ..Default::default() }, + event_recorder_options: EventRecorderOptions::default(), auto_create_table: true, } } @@ -112,6 +115,16 @@ impl GreptimeDbStandaloneBuilder { self } + /// Sets the event recorder options for the standalone instance. + #[must_use] + pub fn with_event_recorder_options( + mut self, + event_recorder_options: EventRecorderOptions, + ) -> Self { + self.event_recorder_options = event_recorder_options; + self + } + #[must_use] pub fn with_default_store_type(self, store_type: StorageType) -> Self { Self { @@ -368,7 +381,13 @@ impl GreptimeDbStandaloneBuilder { wal: self.metasrv_wal_config.clone().into(), grpc: GrpcOptions::default().with_server_addr("127.0.0.1:4001"), slow_query: self.slow_query_options.clone(), + event_recorder: self.event_recorder_options.clone(), auto_create_table: self.auto_create_table, + // Tests cover the descriptor, so they run with it enabled. + otlp: frontend::service_config::OtlpOptions { + experimental_enable_resource_info: true, + ..Default::default() + }, ..StandaloneOptions::default() }; diff --git a/tests-integration/tests/database_ddl_event.rs b/tests-integration/tests/database_ddl_event.rs index 4dfda4f32b9..c08b81211de 100644 --- a/tests-integration/tests/database_ddl_event.rs +++ b/tests-integration/tests/database_ddl_event.rs @@ -15,20 +15,163 @@ use std::sync::Arc; use client::{Database, OutputData}; +use common_event_recorder::{EventRecorderOptions, EventTypeFilter}; use common_test_util::temp_dir::create_temp_dir; use tests_integration::cluster::GreptimeDbClusterBuilder; use tests_integration::standalone::GreptimeDbStandaloneBuilder; use tests_integration::test_util::{ - StorageType, get_test_store_config, setup_authenticated_grpc_database, + StorageType, execute_sql, get_test_store_config, setup_authenticated_grpc_database, }; use crate::event_recorder_test_util::{ - assert_procedure_actor, assert_single_event, find_eventually_string, + assert_eventually_eq, assert_procedure_actor, assert_single_event, find_eventually_string, }; const DATABASE_NAME: &str = "database_ddl_events"; const PROCEDURE_ACTOR: &str = "procedure_actor"; const PROCEDURE_ACTOR_PASSWORD: &str = "procedure_actor_pwd"; +const CREATION_DATABASE_NAME: &str = "event_schema_creation"; +const RECONCILIATION_DATABASE_NAME: &str = "event_schema_reconciliation"; + +#[tokio::test(flavor = "multi_thread")] +async fn test_event_table_auto_creation_with_auto_create_disabled() { + common_telemetry::init_default_ut_logging(); + let standalone = GreptimeDbStandaloneBuilder::new("test_event_table_auto_creation") + .with_auto_create_table(false) + .with_event_recorder_options(EventRecorderOptions { + event_types: Arc::new(EventTypeFilter::Only( + [String::from("create_database")].into_iter().collect(), + )), + ..Default::default() + }) + .build() + .await; + let instance = standalone.fe_instance(); + + let (database, _grpc_server) = setup_authenticated_grpc_database( + instance.clone(), + PROCEDURE_ACTOR, + PROCEDURE_ACTOR_PASSWORD, + ) + .await; + database + .sql(format!("CREATE DATABASE {CREATION_DATABASE_NAME}")) + .await + .unwrap(); + + assert_eventually_eq( + instance, + "SELECT count(*) AS events_tables \ + FROM information_schema.tables \ + WHERE table_catalog = 'greptime' \ + AND table_schema = 'greptime_private' \ + AND table_name = 'events'", + "+---------------+\n| events_tables |\n+---------------+\n| 1 |\n+---------------+", + ) + .await; + + let procedure_id = find_eventually_string( + instance, + &format!( + "SELECT procedure_id FROM greptime_private.events \ + WHERE type = 'create_database' \ + AND schema_name = '{CREATION_DATABASE_NAME}' \ + AND json_path_match(procedure_trigger, '$.type == \"Submitted\"') \ + ORDER BY timestamp DESC LIMIT 1" + ), + "procedure_id", + ) + .await; + assert_procedure_actor(instance, &procedure_id, Some(PROCEDURE_ACTOR)).await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn test_event_table_schema_reconciliation_with_auto_create_disabled() { + common_telemetry::init_default_ut_logging(); + let standalone = GreptimeDbStandaloneBuilder::new("test_event_table_schema_reconciliation") + .with_auto_create_table(false) + .with_event_recorder_options(EventRecorderOptions { + event_types: Arc::new(EventTypeFilter::Only( + [String::from("create_database")].into_iter().collect(), + )), + ..Default::default() + }) + .build() + .await; + let instance = standalone.fe_instance(); + + // Matches the pre-actor procedure-event schema. The recorder filter keeps + // this setup DDL from creating an event before the regression is exercised. + execute_sql( + instance, + r#" + CREATE TABLE greptime_private.events ( + "type" STRING, + payload JSON, + "timestamp" TIMESTAMP(9) NOT NULL, + procedure_id STRING, + procedure_state STRING, + procedure_error STRING, + procedure_trigger JSON, + catalog_name STRING, + schema_name STRING, + event_context JSON, + TIME INDEX ("timestamp"), + PRIMARY KEY ("type") + ) WITH (append_mode = 'true') + "#, + ) + .await; + + assert_eventually_eq( + instance, + "SELECT count(*) AS actor_columns \ + FROM information_schema.columns \ + WHERE table_catalog = 'greptime' \ + AND table_schema = 'greptime_private' \ + AND table_name = 'events' \ + AND column_name = 'actor'", + "+---------------+\n| actor_columns |\n+---------------+\n| 0 |\n+---------------+", + ) + .await; + + let (database, _grpc_server) = setup_authenticated_grpc_database( + instance.clone(), + PROCEDURE_ACTOR, + PROCEDURE_ACTOR_PASSWORD, + ) + .await; + database + .sql(format!("CREATE DATABASE {RECONCILIATION_DATABASE_NAME}")) + .await + .unwrap(); + + assert_eventually_eq( + instance, + "SELECT count(*) AS actor_columns \ + FROM information_schema.columns \ + WHERE table_catalog = 'greptime' \ + AND table_schema = 'greptime_private' \ + AND table_name = 'events' \ + AND column_name = 'actor'", + "+---------------+\n| actor_columns |\n+---------------+\n| 1 |\n+---------------+", + ) + .await; + + let procedure_id = find_eventually_string( + instance, + &format!( + "SELECT procedure_id FROM greptime_private.events \ + WHERE type = 'create_database' \ + AND schema_name = '{RECONCILIATION_DATABASE_NAME}' \ + AND json_path_match(procedure_trigger, '$.type == \"Submitted\"') \ + ORDER BY timestamp DESC LIMIT 1" + ), + "procedure_id", + ) + .await; + assert_procedure_actor(instance, &procedure_id, Some(PROCEDURE_ACTOR)).await; +} #[tokio::test(flavor = "multi_thread")] async fn test_database_ddl_events() { @@ -43,6 +186,7 @@ async fn test_database_ddl_events() { let cluster = GreptimeDbClusterBuilder::new("test_database_ddl_events") .await .with_datanodes(1) + .with_frontend_auto_create_table(false) .with_store_config(store_config) .with_shared_home_dir(Arc::new(home_dir)) .build(true) diff --git a/tests-integration/tests/grpc.rs b/tests-integration/tests/grpc.rs index 737b4fb6c1d..f06ffa65224 100644 --- a/tests-integration/tests/grpc.rs +++ b/tests-integration/tests/grpc.rs @@ -26,9 +26,12 @@ use api::v1::{ use auth::user_provider_from_option; use base64::prelude::{BASE64_STANDARD, Engine as _}; use client::{Client, DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, Database, OutputData}; -use common_catalog::consts::MITO_ENGINE; +use common_catalog::consts::{DEFAULT_PRIVATE_SCHEMA_NAME, MITO_ENGINE}; +use common_event_recorder::DEFAULT_EVENTS_TABLE_NAME; +use common_frontend::slow_query_event::SLOW_QUERY_TABLE_NAME; use common_grpc::channel_manager::ClientTlsOption; use common_memory_manager::OnExhaustedPolicy; +use common_meta::datanode::REGION_STATS_HISTORY_TABLE_NAME; use common_query::Output; use common_recordbatch::RecordBatches; use common_runtime::Runtime; @@ -98,6 +101,8 @@ macro_rules! grpc_tests { test_auto_create_table, test_auto_create_table_with_hints, test_auto_create_table_disabled_by_config, + test_private_system_tables_auto_create_table_with_global_disabled, + test_private_system_tables_bypass_auto_create_hint, test_otel_arrow_auth, test_otel_arrow_exponential_histogram, test_insert_and_select, @@ -679,6 +684,154 @@ pub async fn test_auto_create_table_disabled_by_config(store_type: StorageType) let _ = fe_grpc_server.shutdown().await; } +pub async fn test_private_system_tables_auto_create_table_with_global_disabled( + store_type: StorageType, +) { + let (_db, fe_grpc_server) = setup_grpc_server_with_auto_create_table_disabled( + store_type, + "test_private_system_tables_auto_create_table_with_global_disabled", + ) + .await; + let addr = fe_grpc_server.bind_addr().unwrap().to_string(); + + let grpc_client = Client::with_urls(vec![addr]); + let db = Database::new( + DEFAULT_CATALOG_NAME, + DEFAULT_PRIVATE_SCHEMA_NAME, + grpc_client, + ); + let (host, cpu, mem, ts) = expect_data(); + + for table_name in [ + DEFAULT_EVENTS_TABLE_NAME, + SLOW_QUERY_TABLE_NAME, + REGION_STATS_HISTORY_TABLE_NAME, + ] { + let result = db + .insert(InsertRequests { + inserts: vec![InsertRequest { + table_name: table_name.to_string(), + columns: vec![host.clone(), cpu.clone(), mem.clone(), ts.clone()], + row_count: 4, + }], + }) + .await; + assert_eq!(result.unwrap(), 4); + } + + let load = Column { + column_name: "load".to_string(), + values: Some(column::Values { + f64_values: vec![0.4, 0.5, 0.6, 0.7], + ..Default::default() + }), + semantic_type: SemanticType::Field as i32, + datatype: ColumnDataType::Float64 as i32, + ..Default::default() + }; + let result = db + .insert(InsertRequests { + inserts: vec![InsertRequest { + table_name: DEFAULT_EVENTS_TABLE_NAME.to_string(), + columns: vec![host, cpu, mem, ts, load], + row_count: 4, + }], + }) + .await; + assert_eq!(result.unwrap(), 4); + + let output = db + .sql(format!("SHOW CREATE TABLE {DEFAULT_EVENTS_TABLE_NAME}")) + .await + .unwrap(); + let record_batches = match output.data { + OutputData::RecordBatches(record_batches) => record_batches, + OutputData::Stream(stream) => RecordBatches::try_collect(stream).await.unwrap(), + OutputData::AffectedRows(_) => unreachable!(), + }; + assert!(record_batches.pretty_print().unwrap().contains("\"load\"")); + + let _ = fe_grpc_server.shutdown().await; +} + +pub async fn test_private_system_tables_bypass_auto_create_hint(store_type: StorageType) { + let (_db, fe_grpc_server) = setup_grpc_server( + store_type, + "test_private_system_tables_bypass_auto_create_hint", + ) + .await; + let addr = fe_grpc_server.bind_addr().unwrap().to_string(); + + let grpc_client = Client::with_urls(vec![addr]); + let db = Database::new( + DEFAULT_CATALOG_NAME, + DEFAULT_PRIVATE_SCHEMA_NAME, + grpc_client, + ); + let (host, cpu, mem, ts) = expect_data(); + + for table_name in [ + DEFAULT_EVENTS_TABLE_NAME, + SLOW_QUERY_TABLE_NAME, + REGION_STATS_HISTORY_TABLE_NAME, + ] { + let result = db + .insert_with_hints( + InsertRequests { + inserts: vec![InsertRequest { + table_name: table_name.to_string(), + columns: vec![host.clone(), cpu.clone(), mem.clone(), ts.clone()], + row_count: 4, + }], + }, + &[("auto_create_table", "false")], + ) + .await; + assert_eq!(result.unwrap(), 4); + } + + let ordinary_table = "ordinary_private_table"; + let result = db + .insert_with_hints( + InsertRequests { + inserts: vec![ + InsertRequest { + table_name: DEFAULT_EVENTS_TABLE_NAME.to_string(), + columns: vec![host.clone(), cpu.clone(), mem.clone(), ts.clone()], + row_count: 4, + }, + InsertRequest { + table_name: ordinary_table.to_string(), + columns: vec![host, cpu, mem, ts], + row_count: 4, + }, + ], + }, + &[("auto_create_table", "false")], + ) + .await; + let err = result.unwrap_err().to_string(); + assert!( + err.contains(ordinary_table) && err.contains("auto_create_table"), + "unexpected error: {err}" + ); + + let output = db.sql("SHOW TABLES").await.unwrap(); + let record_batches = match output.data { + OutputData::RecordBatches(record_batches) => record_batches, + OutputData::Stream(stream) => RecordBatches::try_collect(stream).await.unwrap(), + OutputData::AffectedRows(_) => unreachable!(), + }; + assert!( + !record_batches + .pretty_print() + .unwrap() + .contains(ordinary_table) + ); + + let _ = fe_grpc_server.shutdown().await; +} + fn expect_data() -> (Column, Column, Column, Column) { // testing data: let expected_host_col = Column { diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index 7a787833376..f4cfda24dc5 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -164,6 +164,7 @@ macro_rules! http_tests { test_otlp_metrics_new, test_otlp_exponential_histogram, test_otlp_metric_translation_strategies, + test_otlp_metrics_resource_info_conflicts, test_otlp_traces_v0, test_otlp_traces_v1, test_otlp_traces_v1_entity_graph, @@ -2208,6 +2209,7 @@ enable = true enable = true experimental_enable_exponential_histogram = false trace_ingest_chunk_size = 512 +experimental_enable_resource_info = true [prom_store] enable = true @@ -2250,6 +2252,7 @@ query_timeout = "10m" slow_query_threshold = "1m" experimental_min_refresh_duration = "5s" grpc_conn_timeout = "5s" +experimental_flight_do_get_timeout = "10s" experimental_grpc_max_retries = 3 experimental_frontend_scan_timeout = "30s" experimental_max_filter_num_per_query = 20 @@ -2275,7 +2278,7 @@ experimental_manifest_keep_removed_file_ttl = "1h" compress_manifest = false experimental_compaction_memory_limit = "unlimited" experimental_compaction_on_exhausted = "wait" -auto_flush_interval = "30m" +auto_flush_interval = "10m" default_region_write_buffer_size = "0KiB" enable_write_cache = false write_cache_path = "" @@ -6500,7 +6503,7 @@ pub async fn test_otlp_metrics_new(store_type: StorageType) { .await; assert_eq!(StatusCode::OK, res.status()); - let expected = "[[\"claude_code_cost_usage_USD_total\"],[\"claude_code_token_usage_tokens_total\"],[\"demo\"],[\"greptime_physical_table\"],[\"numbers\"]]"; + let expected = "[[\"claude_code_cost_usage_USD_total\"],[\"claude_code_token_usage_tokens_total\"],[\"demo\"],[\"greptime_otel_resource_info\"],[\"greptime_physical_table\"],[\"numbers\"]]"; validate_data("otlp_metrics_all_tables", &client, "show tables;", expected).await; // Metric-engine logical table carries the semantic identity. Match substrings @@ -6567,6 +6570,33 @@ pub async fn test_otlp_metrics_new(store_type: StorageType) { ) .await; + // The synthesized resource descriptor: a plain mito info table keyed by + // the raw OTel attribute keys (service.name is the only allowlisted + // resource attribute in this payload), independent of the promote/ignore + // headers. The timestamp is the newest data-point time. + validate_data( + "otlp_metrics_resource_info", + &client, + "select job, \"service.name\", greptime_value, greptime_timestamp from greptime_otel_resource_info;", + "[[\"claude-code\",\"claude-code\",1.0,1753780559836]]", + ) + .await; + // All four ownership markers must land on auto-create: the descriptor + // write path refuses tables that miss any of them, so a missing stamp + // here would silently stop every follow-up descriptor write. + validate_data( + "otlp_metrics_resource_info_options", + &client, + "select count(*) from information_schema.tables where table_name = 'greptime_otel_resource_info' \ + and engine = 'mito' \ + and create_options like '%greptime.semantic.metric.type=info%' \ + and create_options like '%greptime.semantic.metric.metadata_quality=declared%' \ + and create_options like '%greptime.semantic.signal_type=metric%' \ + and create_options like '%greptime.semantic.source=opentelemetry%';", + "[[1]]", + ) + .await; + // drop table let res = client .get("/v1/sql?sql=drop table `claude_code_cost_usage_USD_total`;") @@ -6578,6 +6608,11 @@ pub async fn test_otlp_metrics_new(store_type: StorageType) { .send() .await; assert_eq!(res.status(), StatusCode::OK); + let res = client + .get("/v1/sql?sql=drop table greptime_otel_resource_info;") + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); // write metrics data // with scope attrs @@ -6704,6 +6739,16 @@ pub async fn test_otlp_metrics_new(store_type: StorageType) { ) .await; + // the descriptor was auto-created again after the earlier drop + validate_data( + "otlp_metrics_resource_info_recreated", + &client, + "select count(*) from information_schema.tables \ + where table_name = 'greptime_otel_resource_info' and engine = 'mito';", + "[[1]]", + ) + .await; + // drop table let res = client .get("/v1/sql?sql=drop table `claude_code_cost_usage_USD_total`;") @@ -6715,6 +6760,44 @@ pub async fn test_otlp_metrics_new(store_type: StorageType) { .send() .await; assert_eq!(res.status(), StatusCode::OK); + let res = client + .get("/v1/sql?sql=drop table greptime_otel_resource_info;") + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + + // job composition end-to-end: service.namespace folds into "/" + // on both the descriptor and the metric tables + let content = r#" +{"resourceMetrics":[{"resource":{"attributes":[{"key":"service.name","value":{"stringValue":"api"}},{"key":"service.namespace","value":{"stringValue":"shop"}},{"key":"host.id","value":{"stringValue":"h-1"}}]},"scopeMetrics":[{"scope":{"name":"s"},"metrics":[{"name":"ns_gauge","gauge":{"dataPoints":[{"timeUnixNano":"1753780559836000000","asDouble":1.0}]}}]}]}]} + "#; + let req: ExportMetricsServiceRequest = serde_json::from_str(content).unwrap(); + let res = send_req( + &client, + vec![( + HeaderName::from_static("content-type"), + HeaderValue::from_static("application/x-protobuf"), + )], + "/v1/otlp/v1/metrics", + req.encode_to_vec(), + false, + ) + .await; + assert_eq!(StatusCode::OK, res.status()); + validate_data( + "otlp_metrics_namespace_job_descriptor", + &client, + "select job, \"service.name\", \"service.namespace\", \"host.id\" from greptime_otel_resource_info;", + "[[\"shop/api\",\"api\",\"shop\",\"h-1\"]]", + ) + .await; + validate_data( + "otlp_metrics_namespace_job_metric_table", + &client, + "select job from ns_gauge;", + "[[\"shop/api\"]]", + ) + .await; guard.remove_all().await; } @@ -6912,6 +6995,24 @@ pub async fn test_otlp_metric_translation_strategies(store_type: StorageType) { ) .await; + // the descriptor's columns are the fixed raw keys regardless of the + // translation strategy, and only allowlisted attributes are projected + validate_data( + "otlp_metric_translation_strategy_resource_info", + &client, + "select job, \"service.name\" from greptime_otel_resource_info;", + "[[\"strategy-service\",\"strategy-service\"]]", + ) + .await; + validate_data( + "otlp_metric_translation_strategy_resource_info_allowlist", + &client, + "select count(*) from information_schema.columns \ + where table_name = 'greptime_otel_resource_info' and column_name = 'resource.attr';", + "[[0]]", + ) + .await; + let res = send_req( &client, vec![ @@ -6934,6 +7035,122 @@ pub async fn test_otlp_metric_translation_strategies(store_type: StorageType) { guard.remove_all().await; } +pub async fn test_otlp_metrics_resource_info_conflicts(store_type: StorageType) { + common_telemetry::init_default_ut_logging(); + let (app, mut guard) = + setup_test_http_app_with_frontend(store_type, "test_otlp_metrics_resource_info_conflicts") + .await; + let client = TestClient::new(app).await; + + let content_type = || { + ( + HeaderName::from_static("content-type"), + HeaderValue::from_static("application/x-protobuf"), + ) + }; + + // A metric named like the descriptor wins the table name: it goes through + // the normal metric-engine path and no descriptor is synthesized. + let content = r#" +{"resourceMetrics":[{"resource":{"attributes":[{"key":"service.name","value":{"stringValue":"api"}},{"key":"host.id","value":{"stringValue":"h-1"}}]},"scopeMetrics":[{"scope":{"name":"s"},"metrics":[{"name":"greptime_otel_resource_info","gauge":{"dataPoints":[{"timeUnixNano":"1753780559836000000","asDouble":1.0}]}}]}]}]} + "#; + let req: ExportMetricsServiceRequest = serde_json::from_str(content).unwrap(); + let res = send_req( + &client, + vec![content_type()], + "/v1/otlp/v1/metrics", + req.encode_to_vec(), + false, + ) + .await; + assert_eq!(StatusCode::OK, res.status()); + let body = ExportMetricsServiceResponse::decode(res.bytes().await).unwrap(); + assert!(body.partial_success.is_none()); + validate_data( + "otlp_metrics_name_collision_engine", + &client, + "select count(*) from information_schema.tables \ + where table_name = 'greptime_otel_resource_info' and engine = 'metric';", + "[[1]]", + ) + .await; + let res = execute_sql(&client, "drop table greptime_otel_resource_info;").await; + assert_eq!(res.status(), StatusCode::OK); + + // A pre-existing table with an incompatible schema fails the descriptor + // write, which degrades to an OTLP partial-success warning while the + // metric data itself is accepted. + let res = execute_sql( + &client, + "create table greptime_otel_resource_info (ts timestamp time index, job double, greptime_value double);", + ) + .await; + assert_eq!(res.status(), StatusCode::OK); + + let content = r#" +{"resourceMetrics":[{"resource":{"attributes":[{"key":"service.name","value":{"stringValue":"api"}}]},"scopeMetrics":[{"scope":{"name":"s"},"metrics":[{"name":"conflict_gauge","gauge":{"dataPoints":[{"timeUnixNano":"1753780559836000000","asDouble":7.0}]}}]}]}]} + "#; + let req: ExportMetricsServiceRequest = serde_json::from_str(content).unwrap(); + let res = send_req( + &client, + vec![content_type()], + "/v1/otlp/v1/metrics", + req.encode_to_vec(), + false, + ) + .await; + assert_eq!(StatusCode::OK, res.status()); + let body = ExportMetricsServiceResponse::decode(res.bytes().await).unwrap(); + let partial_success = body.partial_success.as_ref().unwrap(); + assert_eq!(partial_success.rejected_data_points, 0); + assert!( + partial_success + .error_message + .contains("greptime_otel_resource_info"), + "unexpected warning: {}", + partial_success.error_message + ); + validate_data( + "otlp_metrics_conflict_main_data_accepted", + &client, + "select greptime_value from conflict_gauge;", + "[[7.0]]", + ) + .await; + + // The descriptor is auto-created again once the conflicting table is + // gone, and a follow-up write into it does not degrade. + let res = execute_sql(&client, "drop table greptime_otel_resource_info;").await; + assert_eq!(res.status(), StatusCode::OK); + for round in 0..2 { + let req: ExportMetricsServiceRequest = serde_json::from_str(content).unwrap(); + let res = send_req( + &client, + vec![content_type()], + "/v1/otlp/v1/metrics", + req.encode_to_vec(), + false, + ) + .await; + assert_eq!(StatusCode::OK, res.status()); + let body = ExportMetricsServiceResponse::decode(res.bytes().await).unwrap(); + assert!( + body.partial_success.is_none(), + "round {round} unexpectedly degraded: {:?}", + body.partial_success + ); + } + validate_data( + "otlp_metrics_owned_descriptor_keeps_accepting", + &client, + "select count(*) from greptime_otel_resource_info;", + "[[1]]", + ) + .await; + + guard.remove_all().await; +} + pub async fn test_otlp_traces_v0(store_type: StorageType) { // init common_telemetry::init_default_ut_logging(); @@ -7136,7 +7353,7 @@ pub async fn test_otlp_traces_v1_entity_graph(store_type: StorageType) { let res = send_trace_v1_req(&client, "graph_traces", req, false).await; assert_eq!(StatusCode::OK, res.status()); - let expected = r#"[["service","frontend","service","cart","calls"],["service.instance","resource_attributes.service.instance.id=binst-1,service_name=batch","service","batch","part_of"],["service.instance","resource_attributes.service.instance.id=inst-1,service_name=frontend","service","frontend","part_of"],["service.instance","resource_attributes.service.instance.id=inst-1,service_name=frontend","k8s.pod","uid-9","runs_on"]]"#; + let expected = r#"[["service","frontend","service","cart","calls"],["service.instance","batch,binst-1","service","batch","part_of"],["service.instance","frontend,inst-1","service","frontend","part_of"],["service.instance","frontend,inst-1","k8s.pod","uid-9","runs_on"]]"#; validate_data( "otlp_traces_entity_graph", &client, diff --git a/tests-integration/tests/main.rs b/tests-integration/tests/main.rs index 965a3a67462..56108d4d3ff 100644 --- a/tests-integration/tests/main.rs +++ b/tests-integration/tests/main.rs @@ -37,6 +37,7 @@ mod repartition_event; #[macro_use] mod repartition_expr_version; mod mysql; +mod reconciliation_event; mod view_ddl_event; mod wal_prune_event; diff --git a/tests-integration/tests/reconciliation_event.rs b/tests-integration/tests/reconciliation_event.rs new file mode 100644 index 00000000000..c4b4b88ea2d --- /dev/null +++ b/tests-integration/tests/reconciliation_event.rs @@ -0,0 +1,144 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use client::OutputData; +use common_meta::reconciliation::ResolveStrategy; +use common_procedure::watcher; +use frontend::instance::Instance; +use servers::query_handler::sql::SqlQueryHandler; +use session::context::QueryContext; +use table::table_reference::TableReference; +use tests_integration::cluster::GreptimeDbClusterBuilder; + +use crate::event_recorder_test_util::{assert_eventually_eq, find_eventually_string}; + +const EVENTS_TABLE: &str = "greptime_private.events"; +const CATALOG: &str = "greptime"; +const DATABASE: &str = "reconciliation_event_database"; + +#[tokio::test(flavor = "multi_thread")] +async fn test_table_reconciliation_events() { + common_telemetry::init_default_ut_logging(); + + let cluster = GreptimeDbClusterBuilder::new("table_reconciliation_events") + .await + .with_datanodes(1) + .build(true) + .await; + let frontend = cluster.fe_instance().clone(); + run_sql(frontend.as_ref(), &format!("CREATE DATABASE {DATABASE}")).await; + run_sql( + frontend.as_ref(), + &format!("CREATE TABLE {DATABASE}.metrics (ts TIMESTAMP TIME INDEX)"), + ) + .await; + + let table_procedure_id = cluster + .metasrv + .reconciliation_manager() + .reconcile_table( + TableReference { + catalog: CATALOG, + schema: DATABASE, + table: "metrics", + }, + ResolveStrategy::UseLatest, + ) + .await + .unwrap(); + let mut procedure_watcher = cluster + .metasrv + .procedure_manager() + .procedure_watcher(table_procedure_id) + .unwrap(); + watcher::wait(&mut procedure_watcher).await.unwrap(); + let table_procedure_id = table_procedure_id.to_string(); + + assert_eventually_eq( + &frontend, + &format!( + "SELECT count(*) = 1 AS matches FROM {EVENTS_TABLE} \ + WHERE type = 'reconcile_table' AND procedure_id = '{table_procedure_id}' \ + AND json_get_string(procedure_trigger, 'type') = 'Submitted' \ + AND catalog_name = '{CATALOG}' AND schema_name = '{DATABASE}' \ + AND table_name = 'metrics' AND table_id IS NOT NULL \ + AND physical_table_id IS NULL \ + AND json_get_int(payload, 'version') = 1 \ + AND json_get_string(payload, 'resolve_strategy') = 'use_latest' \ + AND json_get_bool(payload, 'is_subprocedure') = false" + ), + "\ ++---------+ +| matches | ++---------+ +| true | ++---------+", + ) + .await; + + let table_result_payload = find_eventually_string( + &frontend, + &format!( + "SELECT json_to_string(payload) AS payload FROM {EVENTS_TABLE} \ + WHERE type = 'reconcile_table' AND procedure_id = '{table_procedure_id}' \ + AND json_get_string(procedure_trigger, 'type') = 'Succeeded' \ + AND catalog_name = '{CATALOG}' AND schema_name = '{DATABASE}' \ + AND table_name = 'metrics' AND table_id IS NOT NULL \ + AND physical_table_id IS NULL LIMIT 1" + ), + "payload", + ) + .await; + assert_eq!( + serde_json::from_str::(&table_result_payload).unwrap(), + serde_json::json!({ + "version": 1, + "complete": true, + "metadata_state": "consistent", + "resolution_strategy_applied": null, + "resolved_column_count": 1, + "scanned_region_count": 1, + "updated_region_count": 0, + "table_info_updated": true, + "last_completed_phase": "update_table_info", + }) + ); + + assert_eventually_eq( + &frontend, + &format!( + "SELECT count(*) AS terminal_event_count FROM {EVENTS_TABLE} \ + WHERE type = 'reconcile_table' AND procedure_id = '{table_procedure_id}' \ + AND json_get_string(procedure_trigger, 'type') = 'Succeeded' \ + AND catalog_name = '{CATALOG}' AND schema_name = '{DATABASE}' \ + AND table_name = 'metrics' AND table_id IS NOT NULL \ + AND physical_table_id IS NULL" + ), + "\ ++----------------------+ +| terminal_event_count | ++----------------------+ +| 1 | ++----------------------+", + ) + .await; +} + +async fn run_sql(instance: &Instance, sql: &str) { + let output = SqlQueryHandler::do_query(instance, sql, QueryContext::arc()) + .await + .remove(0) + .unwrap(); + assert!(matches!(output.data, OutputData::AffectedRows(_)), "{sql}"); +} diff --git a/tests-integration/tests/sql.rs b/tests-integration/tests/sql.rs index e46a8f58d1b..c95adb483ec 100644 --- a/tests-integration/tests/sql.rs +++ b/tests-integration/tests/sql.rs @@ -89,7 +89,10 @@ macro_rules! sql_tests { test_postgres_explain_bind_parameter, test_postgres_array_types, test_mysql_prepare_stmt_insert_timestamp, + test_mysql_prepare_stmt_timezone, test_mysql_federated_prepare_stmt, + test_mysql_prepare_tql_and_show, + test_postgres_extended_query_row_returning_statements, test_declare_fetch_close_cursor, test_alter_update_on, ); @@ -1486,6 +1489,280 @@ fn assert_pg_numeric_range_error(error: tokio_postgres::Error) { assert_eq!("numeric_value_out_of_range", error.message()); } +pub async fn test_postgres_extended_query_row_returning_statements(store_type: StorageType) { + // Regression test for the tokio-postgres >= 0.7.14 DataRow/RowDescription + // mismatch: statements answered with NoData at Describe but emitting + // DataRows at Execute must describe their real output schema. + let (mut guard, fe_pg_server) = setup_pg_server(store_type, "test_pg_extended_row_stmts").await; + let addr = fe_pg_server.bind_addr().unwrap().to_string(); + + let (client, connection) = tokio_postgres::connect(&format!("postgres://{addr}/public"), NoTls) + .await + .unwrap(); + + let (tx, rx) = tokio::sync::oneshot::channel(); + tokio::spawn(async move { + connection.await.unwrap(); + tx.send(()).unwrap(); + }); + + client + .execute( + "CREATE TABLE demo_metrics (ts timestamp time index, val double, host string primary key skipping index)", + &[], + ) + .await + .unwrap(); + client + .execute( + "INSERT INTO demo_metrics (ts, host, val) VALUES (1000, 'host-a', 1.0), (2000, 'host-b', 2.0)", + &[], + ) + .await + .unwrap(); + + // ---- SHOW DATABASES: single `Database` column ---- + let rows = client.query("SHOW DATABASES", &[]).await.unwrap(); + assert!(!rows.is_empty()); + assert_eq!(1, rows[0].columns().len()); + assert_eq!("Database", rows[0].columns()[0].name()); + assert!(rows.iter().any(|r| r.get::<_, String>(0) == "public")); + + // ---- SHOW FULL DATABASES: `Database` + `Options` columns ---- + let rows = client.query("SHOW FULL DATABASES", &[]).await.unwrap(); + assert!(!rows.is_empty()); + assert_eq!(2, rows[0].columns().len()); + assert_eq!("Database", rows[0].columns()[0].name()); + assert_eq!("Options", rows[0].columns()[1].name()); + assert!(rows.iter().any(|r| r.get::<_, String>(0) == "public")); + + // ---- SHOW TABLES: single `Tables_in_` column ---- + let rows = client.query("SHOW TABLES", &[]).await.unwrap(); + assert!(!rows.is_empty()); + assert_eq!(1, rows[0].columns().len()); + assert_eq!("Tables_in_public", rows[0].columns()[0].name()); + assert!(rows.iter().any(|r| r.get::<_, String>(0) == "demo_metrics")); + + // ---- SHOW FULL TABLES: `Tables_in_` + `Table_type` columns ---- + let rows = client.query("SHOW FULL TABLES", &[]).await.unwrap(); + assert!(!rows.is_empty()); + assert_eq!(2, rows[0].columns().len()); + assert_eq!("Tables_in_public", rows[0].columns()[0].name()); + assert_eq!("Table_type", rows[0].columns()[1].name()); + + // ---- SHOW VIEWS / SHOW FLOWS: empty results, still described with one column ---- + let rows = client.query("SHOW VIEWS", &[]).await.unwrap(); + assert!(rows.is_empty()); + let stmt = client.prepare("SHOW VIEWS").await.unwrap(); + assert_eq!(1, stmt.columns().len()); + assert_eq!("Views", stmt.columns()[0].name()); + let rows = client.query("SHOW FLOWS", &[]).await.unwrap(); + assert!(rows.is_empty()); + let stmt = client.prepare("SHOW FLOWS").await.unwrap(); + assert_eq!(1, stmt.columns().len()); + assert_eq!("Flows", stmt.columns()[0].name()); + + // ---- SHOW TABLE STATUS: fixed eighteen-column schema ---- + let rows = client.query("SHOW TABLE STATUS", &[]).await.unwrap(); + assert!(!rows.is_empty()); + let names: Vec<&str> = rows[0].columns().iter().map(|c| c.name()).collect(); + assert_eq!( + vec![ + "Name", + "Engine", + "Version", + "Row_format", + "Rows", + "Avg_row_length", + "Data_length", + "Max_data_length", + "Index_length", + "Data_free", + "Auto_increment", + "Create_time", + "Update_time", + "Check_time", + "Collation", + "Checksum", + "Create_options", + "Comment", + ], + names + ); + + // ---- SHOW COLUMNS / SHOW FULL COLUMNS ---- + let rows = client + .query("SHOW COLUMNS FROM demo_metrics", &[]) + .await + .unwrap(); + assert_eq!(3, rows.len()); + let names: Vec<&str> = rows[0].columns().iter().map(|c| c.name()).collect(); + assert_eq!( + vec![ + "Field", + "Type", + "Null", + "Key", + "Default", + "Extra", + "Greptime_type" + ], + names + ); + let rows = client + .query("SHOW FULL COLUMNS FROM demo_metrics", &[]) + .await + .unwrap(); + assert_eq!(10, rows[0].columns().len()); + + // ---- SHOW CHARSET / SHOW COLLATION ---- + let rows = client.query("SHOW CHARSET", &[]).await.unwrap(); + assert!(!rows.is_empty()); + assert_eq!(4, rows[0].columns().len()); + let rows = client.query("SHOW COLLATION", &[]).await.unwrap(); + assert!(!rows.is_empty()); + assert_eq!(6, rows[0].columns().len()); + + // ---- SHOW INDEX: fixed fifteen-column schema ---- + let rows = client + .query("SHOW INDEX IN demo_metrics", &[]) + .await + .unwrap(); + assert!(!rows.is_empty()); + let names: Vec<&str> = rows[0].columns().iter().map(|c| c.name()).collect(); + assert_eq!( + vec![ + "Table", + "Non_unique", + "Key_name", + "Seq_in_index", + "Column_name", + "Collation", + "Cardinality", + "Sub_part", + "Packed", + "Null", + "Index_type", + "Comment", + "Index_comment", + "Visible", + "Expression", + ], + names + ); + + // ---- SHOW REGION ---- + let rows = client + .query("SHOW REGION IN demo_metrics", &[]) + .await + .unwrap(); + assert!(!rows.is_empty()); + assert_eq!(4, rows[0].columns().len()); + + // ---- SHOW SEARCH_PATH: single string column ---- + let rows = client.query("SHOW SEARCH_PATH", &[]).await.unwrap(); + assert_eq!(1, rows.len()); + assert_eq!(1, rows[0].columns().len()); + assert_eq!("search_path", rows[0].columns()[0].name()); + assert_eq!("public", rows[0].get::<_, String>(0)); + + // ---- SHOW VARIABLES: single column named after the variable ---- + let rows = client.query("SHOW VARIABLES timezone", &[]).await.unwrap(); + assert_eq!(1, rows.len()); + assert_eq!(1, rows[0].columns().len()); + assert_eq!("TIMEZONE", rows[0].columns()[0].name()); + let _ = rows[0].get::<_, String>(0); + + // ---- DESCRIBE TABLE: fixed six-column string schema ---- + let rows = client + .query("DESCRIBE TABLE demo_metrics", &[]) + .await + .unwrap(); + assert_eq!(3, rows.len()); + let names: Vec<&str> = rows[0].columns().iter().map(|c| c.name()).collect(); + assert_eq!( + vec!["Column", "Type", "Key", "Null", "Default", "Semantic Type"], + names + ); + // first column of each row is the column name; ensure values decode as TEXT + let columns: Vec = rows.iter().map(|r| r.get::<_, String>(0)).collect(); + assert!(columns.contains(&"ts".to_string())); + assert!(columns.contains(&"val".to_string())); + assert!(columns.contains(&"host".to_string())); + + // ---- ADMIN function: single column named after the statement ---- + let rows = client + .query("ADMIN flush_table('demo_metrics')", &[]) + .await + .unwrap(); + assert_eq!(1, rows.len()); + assert_eq!(1, rows[0].columns().len()); + assert!(rows[0].columns()[0].name().contains("flush_table")); + + // ---- TQL EVAL / EXPLAIN / ANALYZE: described from the planned query ---- + let rows = client + .query("TQL EVAL (0, 3000, '1s') demo_metrics", &[]) + .await + .unwrap(); + assert!(!rows.is_empty()); + let names: Vec<&str> = rows[0].columns().iter().map(|c| c.name()).collect(); + assert_eq!(vec!["ts", "val", "host"], names); + + let rows = client + .query("TQL EXPLAIN (0, 3000, '1s') demo_metrics", &[]) + .await + .unwrap(); + assert!(!rows.is_empty()); + assert_eq!(2, rows[0].columns().len()); + + let rows = client + .query("TQL ANALYZE (0, 3000, '1s') demo_metrics", &[]) + .await + .unwrap(); + assert!(!rows.is_empty()); + assert_eq!(3, rows[0].columns().len()); + + // FORMAT JSON variants keep working through the plan-based execution path + let rows = client + .query("TQL EXPLAIN FORMAT JSON (0, 3000, '1s') demo_metrics", &[]) + .await + .unwrap(); + assert!(!rows.is_empty()); + let rows = client + .query("TQL ANALYZE FORMAT JSON (0, 3000, '1s') demo_metrics", &[]) + .await + .unwrap(); + assert!(!rows.is_empty()); + + // ---- the same statements also work over the simple query protocol ---- + for sql in [ + "SHOW DATABASES", + "SHOW FULL TABLES", + "SHOW TABLE STATUS", + "SHOW COLUMNS FROM demo_metrics", + "SHOW CHARSET", + "SHOW COLLATION", + "SHOW INDEX IN demo_metrics", + "SHOW REGION IN demo_metrics", + "SHOW SEARCH_PATH", + "DESCRIBE TABLE demo_metrics", + "ADMIN flush_table('demo_metrics')", + "TQL EVAL (0, 3000, '1s') demo_metrics", + ] { + let msgs = client.simple_query(sql).await.unwrap(); + assert!( + msgs.iter().any(|m| matches!(m, SimpleQueryMessage::Row(_))), + "simple query {sql} should return rows" + ); + } + + drop(client); + rx.await.unwrap(); + + let _ = fe_pg_server.shutdown().await; + guard.remove_all().await; +} + pub async fn test_postgres_explain_bind_parameter(store_type: StorageType) { // Regression test for #8029: EXPLAIN / EXPLAIN ANALYZE must accept bind // parameters over the Postgres extended query protocol. @@ -1779,6 +2056,77 @@ pub async fn test_mysql_prepare_stmt_insert_timestamp(store_type: StorageType) { guard.remove_all().await; } +pub async fn test_mysql_prepare_stmt_timezone(store_type: StorageType) { + let (mut guard, server) = + setup_mysql_server(store_type, "test_mysql_prepare_stmt_timezone").await; + let addr = server.bind_addr().unwrap().to_string(); + + let mut conn = MySqlConnection::connect(&format!("mysql://{addr}/public")) + .await + .unwrap(); + + conn.execute("create table demo(i bigint, ts timestamp time index)") + .await + .unwrap(); + conn.execute("SET time_zone = 'Asia/Shanghai'") + .await + .unwrap(); + + // Server-side prepared statement: the binary DATETIME parameter must be + // interpreted in the session timezone. + sqlx::query("insert into demo values(?, ?)") + .bind(1) + .bind( + NaiveDate::from_ymd_opt(2026, 8, 13) + .and_then(|x| x.and_hms_opt(8, 0, 0)) + .unwrap(), + ) + .execute(&mut conn) + .await + .unwrap(); + + // Text protocol with an equivalent timezone-less literal one hour later. + sqlx::query("insert into demo values(2, '2026-08-13 09:00:00')") + .execute(&mut conn) + .await + .unwrap(); + + // Timestamps are read back in the session timezone, and sqlx parses the + // timezone-less wire representation as UTC. + let rows = sqlx::query("select i, ts from demo order by i") + .fetch_all(&mut conn) + .await + .unwrap(); + assert_eq!(rows.len(), 2); + let ts: DateTime = rows[0].get("ts"); + assert_eq!(ts.to_string(), "2026-08-13 08:00:00 UTC"); + let ts: DateTime = rows[1].get("ts"); + assert_eq!(ts.to_string(), "2026-08-13 09:00:00 UTC"); + + // The prepared predicate compares the same instant as the text literal. + let rows = sqlx::query("select i from demo where ts = ? order by i") + .bind( + NaiveDate::from_ymd_opt(2026, 8, 13) + .and_then(|x| x.and_hms_opt(8, 0, 0)) + .unwrap(), + ) + .fetch_all(&mut conn) + .await + .unwrap(); + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].get::("i"), 1); + + let rows = sqlx::query("select i from demo where ts = '2026-08-13 08:00:00'") + .fetch_all(&mut conn) + .await + .unwrap(); + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].get::("i"), 1); + + let _ = server.shutdown().await; + guard.remove_all().await; +} + pub async fn test_mysql_federated_prepare_stmt(store_type: StorageType) { common_telemetry::init_default_ut_logging(); @@ -1818,6 +2166,73 @@ pub async fn test_mysql_federated_prepare_stmt(store_type: StorageType) { guard.remove_all().await; } +pub async fn test_mysql_prepare_tql_and_show(store_type: StorageType) { + // `do_describe` now plans TQL and information-schema-backed SHOW + // statements, so MySQL prepared statements derive their column metadata + // from the planned query and execute the plan directly. This exercises + // that path end-to-end. + common_telemetry::init_default_ut_logging(); + + let (mut guard, fe_mysql_server) = + setup_mysql_server(store_type, "test_mysql_prepare_tql_and_show").await; + let addr = fe_mysql_server.bind_addr().unwrap().to_string(); + + let pool = MySqlPoolOptions::new() + .max_connections(2) + .connect(&format!("mysql://{addr}/public")) + .await + .unwrap(); + + sqlx::query( + "CREATE TABLE demo_metrics (ts timestamp time index, val double, host string primary key)", + ) + .execute(&pool) + .await + .unwrap(); + sqlx::query("INSERT INTO demo_metrics (ts, host, val) VALUES (1000, 'host-a', 1.0), (2000, 'host-b', 2.0)") + .execute(&pool) + .await + .unwrap(); + + // sqlx::query uses the binary prepared statement protocol + // (COM_STMT_PREPARE + COM_STMT_EXECUTE). + let rows = sqlx::query("TQL EVAL (0, 3000, '1s') demo_metrics") + .fetch_all(&pool) + .await + .unwrap(); + assert!(!rows.is_empty()); + assert_eq!(3, rows[0].columns().len()); + + let rows = sqlx::query("TQL ANALYZE (0, 3000, '1s') demo_metrics") + .fetch_all(&pool) + .await + .unwrap(); + assert!(!rows.is_empty()); + + // SHOW statements share the same describe path; prepared SHOW FULL + // TABLES must report both columns. + let rows = sqlx::query("SHOW TABLES").fetch_all(&pool).await.unwrap(); + assert!(!rows.is_empty()); + assert_eq!(1, rows[0].columns().len()); + + let rows = sqlx::query("SHOW FULL TABLES") + .fetch_all(&pool) + .await + .unwrap(); + assert!(!rows.is_empty()); + assert_eq!(2, rows[0].columns().len()); + + let rows = sqlx::query("SHOW DATABASES") + .fetch_all(&pool) + .await + .unwrap(); + assert!(!rows.is_empty()); + assert_eq!(1, rows[0].columns().len()); + + let _ = fe_mysql_server.shutdown().await; + guard.remove_all().await; +} + pub async fn test_postgres_array_types(store_type: StorageType) { let (mut guard, fe_pg_server) = setup_pg_server(store_type, "test_postgres_array_types").await; let addr = fe_pg_server.bind_addr().unwrap().to_string(); diff --git a/tests/cases/standalone/common/information_schema/ssts.result b/tests/cases/standalone/common/information_schema/ssts.result index caf6a239607..9777e3b12a9 100644 --- a/tests/cases/standalone/common/information_schema/ssts.result +++ b/tests/cases/standalone/common/information_schema/ssts.result @@ -1,33 +1,35 @@ DESC TABLE information_schema.ssts_manifest; -+------------------+---------------------+-----+------+---------+---------------+ -| Column | Type | Key | Null | Default | Semantic Type | -+------------------+---------------------+-----+------+---------+---------------+ -| table_dir | String | | NO | | FIELD | -| region_id | UInt64 | | NO | | FIELD | -| table_id | UInt32 | | NO | | FIELD | -| region_number | UInt32 | | NO | | FIELD | -| region_group | UInt8 | | NO | | FIELD | -| region_sequence | UInt32 | | NO | | FIELD | -| file_id | String | | NO | | FIELD | -| index_version | UInt64 | | NO | | FIELD | -| level | UInt8 | | NO | | FIELD | -| file_path | String | | NO | | FIELD | -| file_size | UInt64 | | NO | | FIELD | -| index_file_path | String | | YES | | FIELD | -| index_file_size | UInt64 | | YES | | FIELD | -| num_rows | UInt64 | | NO | | FIELD | -| num_row_groups | UInt64 | | NO | | FIELD | -| num_series | UInt64 | | YES | | FIELD | -| min_ts | TimestampNanosecond | | YES | | FIELD | -| max_ts | TimestampNanosecond | | YES | | FIELD | -| sequence | UInt64 | | YES | | FIELD | -| origin_region_id | UInt64 | | NO | | FIELD | -| node_id | UInt64 | | YES | | FIELD | -| visible | Boolean | | NO | | FIELD | -| primary_key_min | Binary | | YES | | FIELD | -| primary_key_max | Binary | | YES | | FIELD | -+------------------+---------------------+-----+------+---------+---------------+ ++---------------------------------+---------------------+-----+------+---------+---------------+ +| Column | Type | Key | Null | Default | Semantic Type | ++---------------------------------+---------------------+-----+------+---------+---------------+ +| table_dir | String | | NO | | FIELD | +| region_id | UInt64 | | NO | | FIELD | +| table_id | UInt32 | | NO | | FIELD | +| region_number | UInt32 | | NO | | FIELD | +| region_group | UInt8 | | NO | | FIELD | +| region_sequence | UInt32 | | NO | | FIELD | +| file_id | String | | NO | | FIELD | +| index_version | UInt64 | | NO | | FIELD | +| level | UInt8 | | NO | | FIELD | +| file_path | String | | NO | | FIELD | +| file_size | UInt64 | | NO | | FIELD | +| index_file_path | String | | YES | | FIELD | +| index_file_size | UInt64 | | YES | | FIELD | +| num_rows | UInt64 | | NO | | FIELD | +| num_row_groups | UInt64 | | NO | | FIELD | +| num_series | UInt64 | | YES | | FIELD | +| min_ts | TimestampNanosecond | | YES | | FIELD | +| max_ts | TimestampNanosecond | | YES | | FIELD | +| sequence | UInt64 | | YES | | FIELD | +| origin_region_id | UInt64 | | NO | | FIELD | +| node_id | UInt64 | | YES | | FIELD | +| visible | Boolean | | NO | | FIELD | +| primary_key_min | Binary | | YES | | FIELD | +| primary_key_max | Binary | | YES | | FIELD | +| max_row_group_uncompressed_size | UInt64 | | NO | | FIELD | +| partition_expr | String | | YES | | FIELD | ++---------------------------------+---------------------+-----+------+---------+---------------+ DESC TABLE information_schema.ssts_storage; @@ -99,13 +101,13 @@ ADMIN FLUSH_TABLE('sst_case'); -- SQLNESS REPLACE (/public/\d+) /public/ SELECT * FROM information_schema.ssts_manifest order by file_path; -+----------------------------+---------------+----------+---------------+--------------+-----------------+--------------------------------------+---------------+-------+----------------------------------------------------------------------------------------+-----------+---------------------------------------------------------------------------------------------+-----------------+----------+----------------+------------+-------------------------+-------------------------+----------+------------------+---------+---------+-----------------+-----------------+ -| table_dir | region_id | table_id | region_number | region_group | region_sequence | file_id | index_version | level | file_path | file_size | index_file_path | index_file_size | num_rows | num_row_groups | num_series | min_ts | max_ts | sequence | origin_region_id | node_id | visible | primary_key_min | primary_key_max | -+----------------------------+---------------+----------+---------------+--------------+-----------------+--------------------------------------+---------------+-------+----------------------------------------------------------------------------------------+-----------+---------------------------------------------------------------------------------------------+-----------------+----------+----------------+------------+-------------------------+-------------------------+----------+------------------+---------+---------+-----------------+-----------------+ -| data/greptime/public// |||||| ||| data/greptime/public//_/.parquet || data/greptime/public//_/index/.puffin ||||| | |||| true | 01800001f4 | 01800001f4 | -| data/greptime/public// |||||| ||| data/greptime/public//_/.parquet || data/greptime/public//_/index/.puffin ||||| | |||| true | 01800005dc | 01800005dc | -| data/greptime/public// |||||| ||| data/greptime/public//_/.parquet || data/greptime/public//_/index/.puffin ||||| | |||| true | 01800009c4 | 01800009c4 | -+----------------------------+---------------+----------+---------------+--------------+-----------------+--------------------------------------+---------------+-------+----------------------------------------------------------------------------------------+-----------+---------------------------------------------------------------------------------------------+-----------------+----------+----------------+------------+-------------------------+-------------------------+----------+------------------+---------+---------+-----------------+-----------------+ ++----------------------------+---------------+----------+---------------+--------------+-----------------+--------------------------------------+---------------+-------+----------------------------------------------------------------------------------------+-----------+---------------------------------------------------------------------------------------------+-----------------+----------+----------------+------------+-------------------------+-------------------------+----------+------------------+---------+---------+-----------------+-----------------+---------------------------------+------------------------+ +| table_dir | region_id | table_id | region_number | region_group | region_sequence | file_id | index_version | level | file_path | file_size | index_file_path | index_file_size | num_rows | num_row_groups | num_series | min_ts | max_ts | sequence | origin_region_id | node_id | visible | primary_key_min | primary_key_max | max_row_group_uncompressed_size | partition_expr | ++----------------------------+---------------+----------+---------------+--------------+-----------------+--------------------------------------+---------------+-------+----------------------------------------------------------------------------------------+-----------+---------------------------------------------------------------------------------------------+-----------------+----------+----------------+------------+-------------------------+-------------------------+----------+------------------+---------+---------+-----------------+-----------------+---------------------------------+------------------------+ +| data/greptime/public// |||||| ||| data/greptime/public//_/.parquet || data/greptime/public//_/index/.puffin ||||| | |||| true | 01800001f4 | 01800001f4 || a <| +| data/greptime/public// |||||| ||| data/greptime/public//_/.parquet || data/greptime/public//_/index/.puffin ||||| | |||| true | 01800005dc | 01800005dc || a >=AND a <| +| data/greptime/public// |||||| ||| data/greptime/public//_/.parquet || data/greptime/public//_/index/.puffin ||||| | |||| true | 01800009c4 | 01800009c4 || a >=| ++----------------------------+---------------+----------+---------------+--------------+-----------------+--------------------------------------+---------------+-------+----------------------------------------------------------------------------------------+-----------+---------------------------------------------------------------------------------------------+-----------------+----------+----------------+------------+-------------------------+-------------------------+----------+------------------+---------+---------+-----------------+-----------------+---------------------------------+------------------------+ -- SQLNESS REPLACE (\s+\d+\s+) -- SQLNESS REPLACE ([0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}) @@ -168,15 +170,15 @@ ADMIN FLUSH_TABLE('sst_case'); -- SQLNESS REPLACE (/public/\d+) /public/ SELECT * FROM information_schema.ssts_manifest order by region_id, sequence; -+----------------------------+---------------+----------+---------------+--------------+-----------------+--------------------------------------+---------------+-------+----------------------------------------------------------------------------------------+-----------+---------------------------------------------------------------------------------------------+-----------------+----------+----------------+------------+-------------------------+-------------------------+----------+------------------+---------+---------+-----------------+-----------------+ -| table_dir | region_id | table_id | region_number | region_group | region_sequence | file_id | index_version | level | file_path | file_size | index_file_path | index_file_size | num_rows | num_row_groups | num_series | min_ts | max_ts | sequence | origin_region_id | node_id | visible | primary_key_min | primary_key_max | -+----------------------------+---------------+----------+---------------+--------------+-----------------+--------------------------------------+---------------+-------+----------------------------------------------------------------------------------------+-----------+---------------------------------------------------------------------------------------------+-----------------+----------+----------------+------------+-------------------------+-------------------------+----------+------------------+---------+---------+-----------------+-----------------+ -| data/greptime/public// |||||| ||| data/greptime/public//_/.parquet || data/greptime/public//_/index/.puffin ||||| | |||| true | 01800001f4 | 01800001f4 | -| data/greptime/public// |||||| ||| data/greptime/public//_/.parquet || data/greptime/public//_/index/.puffin ||||| | |||| true || 018000007c | -| data/greptime/public// |||||| ||| data/greptime/public//_/.parquet || data/greptime/public//_/index/.puffin ||||| | |||| true | 01800005dc | 01800005dc | -| data/greptime/public// |||||| ||| data/greptime/public//_/.parquet || data/greptime/public//_/index/.puffin ||||| | |||| true ||| -| data/greptime/public// |||||| ||| data/greptime/public//_/.parquet || data/greptime/public//_/index/.puffin ||||| | |||| true | 01800009c4 | 01800009c4 | -+----------------------------+---------------+----------+---------------+--------------+-----------------+--------------------------------------+---------------+-------+----------------------------------------------------------------------------------------+-----------+---------------------------------------------------------------------------------------------+-----------------+----------+----------------+------------+-------------------------+-------------------------+----------+------------------+---------+---------+-----------------+-----------------+ ++----------------------------+---------------+----------+---------------+--------------+-----------------+--------------------------------------+---------------+-------+----------------------------------------------------------------------------------------+-----------+---------------------------------------------------------------------------------------------+-----------------+----------+----------------+------------+-------------------------+-------------------------+----------+------------------+---------+---------+-----------------+-----------------+---------------------------------+------------------------+ +| table_dir | region_id | table_id | region_number | region_group | region_sequence | file_id | index_version | level | file_path | file_size | index_file_path | index_file_size | num_rows | num_row_groups | num_series | min_ts | max_ts | sequence | origin_region_id | node_id | visible | primary_key_min | primary_key_max | max_row_group_uncompressed_size | partition_expr | ++----------------------------+---------------+----------+---------------+--------------+-----------------+--------------------------------------+---------------+-------+----------------------------------------------------------------------------------------+-----------+---------------------------------------------------------------------------------------------+-----------------+----------+----------------+------------+-------------------------+-------------------------+----------+------------------+---------+---------+-----------------+-----------------+---------------------------------+------------------------+ +| data/greptime/public// |||||| ||| data/greptime/public//_/.parquet || data/greptime/public//_/index/.puffin ||||| | |||| true | 01800001f4 | 01800001f4 || a <| +| data/greptime/public// |||||| ||| data/greptime/public//_/.parquet || data/greptime/public//_/index/.puffin ||||| | |||| true || 018000007c || a <| +| data/greptime/public// |||||| ||| data/greptime/public//_/.parquet || data/greptime/public//_/index/.puffin ||||| | |||| true | 01800005dc | 01800005dc || a >=AND a <| +| data/greptime/public// |||||| ||| data/greptime/public//_/.parquet || data/greptime/public//_/index/.puffin ||||| | |||| true |||| a >=AND a <| +| data/greptime/public// |||||| ||| data/greptime/public//_/.parquet || data/greptime/public//_/index/.puffin ||||| | |||| true | 01800009c4 | 01800009c4 || a >=| ++----------------------------+---------------+----------+---------------+--------------+-----------------+--------------------------------------+---------------+-------+----------------------------------------------------------------------------------------+-----------+---------------------------------------------------------------------------------------------+-----------------+----------+----------------+------------+-------------------------+-------------------------+----------+------------------+---------+---------+-----------------+-----------------+---------------------------------+------------------------+ -- SQLNESS REPLACE (\s+\d+\s+) -- SQLNESS REPLACE ([0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}) diff --git a/tests/cases/standalone/common/insert/insert_default_timezone.result b/tests/cases/standalone/common/insert/insert_default_timezone.result index 413eb5d46ae..41d8ca6cba8 100644 --- a/tests/cases/standalone/common/insert/insert_default_timezone.result +++ b/tests/cases/standalone/common/insert/insert_default_timezone.result @@ -114,6 +114,27 @@ INSERT INTO test3 (ts, st, ts_ns) VALUES ( Affected Rows: 1 +-- a NULL branch must not cancel the conversion for the whole column +INSERT INTO test3 (ts, ts_ns) +SELECT '2026-08-13 12:00:00.001' AS a, '2026-08-13 12:00:00.123456789' AS b +UNION ALL +SELECT '2026-08-14 12:00:00.001', NULL; + +Affected Rows: 2 + +-- NULL in the first branch: the union's schema starts out as Null +INSERT INTO test3 (ts, ts_ns) +SELECT '2026-08-15 12:00:00.001' AS a, NULL AS b +UNION ALL +SELECT '2026-08-19 12:00:00.001', '2026-08-19 12:00:00.123456789'; + +Affected Rows: 2 + +-- the assignment cast also lands on non-literal VALUES expressions +INSERT INTO test3 (ts, st) VALUES (concat('2026-08-20 ', '12:00:00.001'), now()); + +Affected Rows: 1 + SELECT ts, ts_ns FROM test3 ORDER BY ts; +-------------------------+-------------------------------+ @@ -130,10 +151,43 @@ SELECT ts, ts_ns FROM test3 ORDER BY ts; | 2026-08-10T12:00:00.001 | | | 2026-08-11T12:00:00.001 | | | 2026-08-12T04:00:00.123 | 2026-08-12T04:00:00.123456789 | +| 2026-08-13T04:00:00.001 | 2026-08-13T04:00:00.123456789 | +| 2026-08-14T04:00:00.001 | | +| 2026-08-15T04:00:00.001 | | | 2026-08-16T04:00:00.001 | | | 2026-08-17T04:00:00.001 | | +| 2026-08-19T04:00:00.001 | 2026-08-19T04:00:00.123456789 | +| 2026-08-20T04:00:00.001 | | +-------------------------+-------------------------------+ +-- UNION dedup keys must stay on the source strings: these two spell the same +-- instant differently, so the source query yields two rows and both are kept. +CREATE TABLE test4 (ts TIMESTAMP TIME INDEX) WITH ('append_mode'='true'); + +Affected Rows: 0 + +INSERT INTO test4 (ts) +SELECT '2026-08-06 04:00:00' UNION SELECT '2026-08-06 04:00:00.000'; + +Affected Rows: 2 + +SELECT count(*) FROM test4; + ++----------+ +| count(*) | ++----------+ +| 2 | ++----------+ + +SELECT ts FROM test4 ORDER BY ts; + ++---------------------+ +| ts | ++---------------------+ +| 2026-08-05T20:00:00 | +| 2026-08-05T20:00:00 | ++---------------------+ + SET time_zone = 'UTC'; Affected Rows: 0 @@ -150,3 +204,7 @@ DROP TABLE test3; Affected Rows: 0 +DROP TABLE test4; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/insert/insert_default_timezone.sql b/tests/cases/standalone/common/insert/insert_default_timezone.sql index 7cbef15d6d9..f7c196699b7 100644 --- a/tests/cases/standalone/common/insert/insert_default_timezone.sql +++ b/tests/cases/standalone/common/insert/insert_default_timezone.sql @@ -56,8 +56,34 @@ INSERT INTO test3 (ts, st, ts_ns) VALUES ( '2026-08-09 12:00:00.123456789' ); +-- a NULL branch must not cancel the conversion for the whole column +INSERT INTO test3 (ts, ts_ns) +SELECT '2026-08-13 12:00:00.001' AS a, '2026-08-13 12:00:00.123456789' AS b +UNION ALL +SELECT '2026-08-14 12:00:00.001', NULL; + +-- NULL in the first branch: the union's schema starts out as Null +INSERT INTO test3 (ts, ts_ns) +SELECT '2026-08-15 12:00:00.001' AS a, NULL AS b +UNION ALL +SELECT '2026-08-19 12:00:00.001', '2026-08-19 12:00:00.123456789'; + +-- the assignment cast also lands on non-literal VALUES expressions +INSERT INTO test3 (ts, st) VALUES (concat('2026-08-20 ', '12:00:00.001'), now()); + SELECT ts, ts_ns FROM test3 ORDER BY ts; +-- UNION dedup keys must stay on the source strings: these two spell the same +-- instant differently, so the source query yields two rows and both are kept. +CREATE TABLE test4 (ts TIMESTAMP TIME INDEX) WITH ('append_mode'='true'); + +INSERT INTO test4 (ts) +SELECT '2026-08-06 04:00:00' UNION SELECT '2026-08-06 04:00:00.000'; + +SELECT count(*) FROM test4; + +SELECT ts FROM test4 ORDER BY ts; + SET time_zone = 'UTC'; DROP TABLE test1; @@ -65,3 +91,5 @@ DROP TABLE test1; DROP TABLE test2; DROP TABLE test3; + +DROP TABLE test4; diff --git a/tests/cases/standalone/common/system/information_schema.result b/tests/cases/standalone/common/system/information_schema.result index e05cd064e1d..17e1f0ef80d 100644 --- a/tests/cases/standalone/common/system/information_schema.result +++ b/tests/cases/standalone/common/system/information_schema.result @@ -436,6 +436,7 @@ order by table_schema, table_name, column_name; | greptime | information_schema | ssts_manifest | index_file_size | 13 | | | 20 | 0 | | | | | | select,insert | | UInt64 | bigint unsigned | FIELD | | YES | bigint unsigned | | | | greptime | information_schema | ssts_manifest | index_version | 8 | | | 20 | 0 | | | | | | select,insert | | UInt64 | bigint unsigned | FIELD | | NO | bigint unsigned | | | | greptime | information_schema | ssts_manifest | level | 9 | | | 3 | 0 | | | | | | select,insert | | UInt8 | tinyint unsigned | FIELD | | NO | tinyint unsigned | | | +| greptime | information_schema | ssts_manifest | max_row_group_uncompressed_size | 25 | | | 20 | 0 | | | | | | select,insert | | UInt64 | bigint unsigned | FIELD | | NO | bigint unsigned | | | | greptime | information_schema | ssts_manifest | max_ts | 18 | | | | | 9 | | | | | select,insert | | TimestampNanosecond | timestamp(9) | FIELD | | YES | timestamp(9) | | | | greptime | information_schema | ssts_manifest | min_ts | 17 | | | | | 9 | | | | | select,insert | | TimestampNanosecond | timestamp(9) | FIELD | | YES | timestamp(9) | | | | greptime | information_schema | ssts_manifest | node_id | 21 | | | 20 | 0 | | | | | | select,insert | | UInt64 | bigint unsigned | FIELD | | YES | bigint unsigned | | | @@ -443,6 +444,7 @@ order by table_schema, table_name, column_name; | greptime | information_schema | ssts_manifest | num_rows | 14 | | | 20 | 0 | | | | | | select,insert | | UInt64 | bigint unsigned | FIELD | | NO | bigint unsigned | | | | greptime | information_schema | ssts_manifest | num_series | 16 | | | 20 | 0 | | | | | | select,insert | | UInt64 | bigint unsigned | FIELD | | YES | bigint unsigned | | | | greptime | information_schema | ssts_manifest | origin_region_id | 20 | | | 20 | 0 | | | | | | select,insert | | UInt64 | bigint unsigned | FIELD | | NO | bigint unsigned | | | +| greptime | information_schema | ssts_manifest | partition_expr | 26 | 2147483647 | 2147483647 | | | | utf8 | utf8_bin | | | select,insert | | String | string | FIELD | | YES | string | | | | greptime | information_schema | ssts_manifest | primary_key_max | 24 | | | | | | | | | | select,insert | | Binary | varbinary | FIELD | | YES | varbinary | | | | greptime | information_schema | ssts_manifest | primary_key_min | 23 | | | | | | | | | | select,insert | | Binary | varbinary | FIELD | | YES | varbinary | | | | greptime | information_schema | ssts_manifest | region_group | 5 | | | 3 | 0 | | | | | | select,insert | | UInt8 | tinyint unsigned | FIELD | | NO | tinyint unsigned | | | diff --git a/tests/cases/standalone/common/system/semantic_graph.result b/tests/cases/standalone/common/system/semantic_graph.result index 6af28cfa2b5..f8894319842 100644 --- a/tests/cases/standalone/common/system/semantic_graph.result +++ b/tests/cases/standalone/common/system/semantic_graph.result @@ -94,26 +94,26 @@ select entity_type, entity_id, entity_id_attrs, scope, descriptive, source_table from greptime_private.semantic_entities order by entity_type, entity_id; -+------------------+---------------------------+-------------------------------------+---------+-------------------+------------------------------+ -| entity_type | entity_id | entity_id_attrs | scope | descriptive | source_tables | -+------------------+---------------------------+-------------------------------------+---------+-------------------+------------------------------+ -| host | h1 | | | | ["public.graph_app_metrics"] | -| process | host=h1,service_name=cart | {"host":"h1","service_name":"cart"} | | {"env":"us-east"} | ["public.graph_app_metrics"] | -| service | cart | | us-east | | ["public.graph_app_metrics"] | -| service.instance | cart-0 | | | | ["public.graph_app_metrics"] | -+------------------+---------------------------+-------------------------------------+---------+-------------------+------------------------------+ ++------------------+-----------+-------------------------------------+---------+-------------------+------------------------------+ +| entity_type | entity_id | entity_id_attrs | scope | descriptive | source_tables | ++------------------+-----------+-------------------------------------+---------+-------------------+------------------------------+ +| host | h1 | | | | ["public.graph_app_metrics"] | +| process | cart,h1 | {"host":"h1","service_name":"cart"} | | {"env":"us-east"} | ["public.graph_app_metrics"] | +| service | cart | | us-east | | ["public.graph_app_metrics"] | +| service.instance | cart-0 | | | | ["public.graph_app_metrics"] | ++------------------+-----------+-------------------------------------+---------+-------------------+------------------------------+ select src_type, src_id, dst_type, dst_id, rel_type, provenance, confidence from greptime_private.semantic_relationships order by rel_type, src_id; -+------------------+---------------------------+----------+--------+----------+------------+------------+ -| src_type | src_id | dst_type | dst_id | rel_type | provenance | confidence | -+------------------+---------------------------+----------+--------+----------+------------+------------+ -| service.instance | cart-0 | service | cart | part_of | attribute | 1.0 | -| service.instance | cart-0 | host | h1 | runs_on | attribute | 1.0 | -| process | host=h1,service_name=cart | host | h1 | runs_on | attribute | 1.0 | -+------------------+---------------------------+----------+--------+----------+------------+------------+ ++------------------+---------+----------+--------+----------+------------+------------+ +| src_type | src_id | dst_type | dst_id | rel_type | provenance | confidence | ++------------------+---------+----------+--------+----------+------------+------------+ +| service.instance | cart-0 | service | cart | part_of | attribute | 1.0 | +| process | cart,h1 | host | h1 | runs_on | attribute | 1.0 | +| service.instance | cart-0 | host | h1 | runs_on | attribute | 1.0 | ++------------------+---------+----------+--------+----------+------------+------------+ drop table graph_app_metrics; @@ -570,38 +570,38 @@ select entity_type, entity_id, source_tables from greptime_private.semantic_entities order by entity_type, entity_id, source_tables; -+------------------+-----------------------------------------------------------+------------------------------------+ -| entity_type | entity_id | source_tables | -+------------------+-----------------------------------------------------------+------------------------------------+ -| k8s.container | container=main,uid=uid-1 | ["public.kube_pod_container_info"] | -| k8s.node | node-a | ["public.kube_pod_info"] | -| k8s.pod | uid-1 | ["public.kube_pod_info"] | -| k8s.pod | uid-1 | ["public.kube_pod_owner"] | -| k8s.pod | uid-1 | ["public.kube_pod_container_info"] | -| k8s.pod | uid-2 | ["public.kube_pod_info"] | -| k8s.pod | uid-2 | ["public.kube_pod_owner"] | -| k8s.pod | uid-3 | ["public.kube_pod_info"] | -| k8s.service | svc-uid-1 | ["public.kube_service_info"] | -| k8s.workload | namespace=default,owner_kind=ReplicaSet,owner_name=api-rs | ["public.kube_pod_owner"] | -| service | shop/api | ["public.target_info"] | -| service.instance | instance=inst-1,job=shop/api | ["public.target_info"] | -+------------------+-----------------------------------------------------------+------------------------------------+ ++------------------+---------------------------+------------------------------------+ +| entity_type | entity_id | source_tables | ++------------------+---------------------------+------------------------------------+ +| k8s.container | uid-1,main | ["public.kube_pod_container_info"] | +| k8s.node | node-a | ["public.kube_pod_info"] | +| k8s.pod | uid-1 | ["public.kube_pod_info"] | +| k8s.pod | uid-1 | ["public.kube_pod_owner"] | +| k8s.pod | uid-1 | ["public.kube_pod_container_info"] | +| k8s.pod | uid-2 | ["public.kube_pod_info"] | +| k8s.pod | uid-2 | ["public.kube_pod_owner"] | +| k8s.pod | uid-3 | ["public.kube_pod_info"] | +| k8s.service | svc-uid-1 | ["public.kube_service_info"] | +| k8s.workload | default,ReplicaSet,api-rs | ["public.kube_pod_owner"] | +| service | shop/api | ["public.target_info"] | +| service.instance | shop/api,inst-1 | ["public.target_info"] | ++------------------+---------------------------+------------------------------------+ -- SQLNESS PROTOCOL MYSQL select src_type, src_id, dst_type, dst_id, rel_type, provenance from greptime_private.semantic_relationships order by rel_type, src_id, dst_id; -+------------------+------------------------------+---------------+-----------------------------------------------------------+----------+------------+ -| src_type | src_id | dst_type | dst_id | rel_type | provenance | -+------------------+------------------------------+---------------+-----------------------------------------------------------+----------+------------+ -| k8s.pod | uid-1 | k8s.container | container=main,uid=uid-1 | contains | attribute | -| service.instance | instance=inst-1,job=shop/api | service | shop/api | part_of | attribute | -| k8s.pod | uid-1 | k8s.workload | namespace=default,owner_kind=ReplicaSet,owner_name=api-rs | part_of | attribute | -| k8s.pod | uid-2 | k8s.workload | namespace=default,owner_kind=ReplicaSet,owner_name=api-rs | part_of | attribute | -| k8s.pod | uid-1 | k8s.node | node-a | runs_on | attribute | -| k8s.pod | uid-2 | k8s.node | node-a | runs_on | attribute | -+------------------+------------------------------+---------------+-----------------------------------------------------------+----------+------------+ ++------------------+-----------------+---------------+---------------------------+----------+------------+ +| src_type | src_id | dst_type | dst_id | rel_type | provenance | ++------------------+-----------------+---------------+---------------------------+----------+------------+ +| k8s.pod | uid-1 | k8s.container | uid-1,main | contains | attribute | +| service.instance | shop/api,inst-1 | service | shop/api | part_of | attribute | +| k8s.pod | uid-1 | k8s.workload | default,ReplicaSet,api-rs | part_of | attribute | +| k8s.pod | uid-2 | k8s.workload | default,ReplicaSet,api-rs | part_of | attribute | +| k8s.pod | uid-1 | k8s.node | node-a | runs_on | attribute | +| k8s.pod | uid-2 | k8s.node | node-a | runs_on | attribute | ++------------------+-----------------+---------------+---------------------------+----------+------------+ drop table kube_pod_info; @@ -627,3 +627,182 @@ drop table http_requests_total; Affected Rows: 0 +-- OTel conventions: the ingestion-synthesized greptime_otel_resource_info descriptor +-- (stamped signal_type=metric + source=opentelemetry + metric.type=info) gets implicit +-- declarations under the raw OTel column names. host/container identities are +-- the stable ids only, so rows with an empty host.id / container.id link no +-- infrastructure, and a row without an instance value declares no +-- service.instance. The same table name stamped source=prometheus is not +-- whitelisted. +create table greptime_otel_resource_info ( + greptime_timestamp timestamp(3) time index, + job string, + instance string, + "service.name" string, + "service.namespace" string, + "host.id" string, + "host.name" string, + "container.id" string, + "container.name" string, + greptime_value double, + primary key (job, instance, "service.name", "service.namespace", + "host.id", "host.name", "container.id", "container.name") +) with ( + 'greptime.semantic.signal_type' = 'metric', + 'greptime.semantic.source' = 'opentelemetry', + 'greptime.semantic.metric.type' = 'info' +); + +Affected Rows: 0 + +insert into greptime_otel_resource_info values + (now(), 'shop/api', 'inst-1', 'api', 'shop', 'h-1', 'node-a', 'c-1', 'api-ctr', 1), + (now(), 'shop/api', 'inst-2', 'api', 'shop', '', 'laptop', '', '', 1), + (now(), 'worker', '', 'worker', '', 'h-1', 'node-a', '', '', 1); + +Affected Rows: 3 + +-- SQLNESS PROTOCOL MYSQL +select entity_type, entity_id, source_tables +from greptime_private.semantic_entities +order by entity_type, entity_id, source_tables; + ++------------------+-----------------+----------------------------------------+ +| entity_type | entity_id | source_tables | ++------------------+-----------------+----------------------------------------+ +| container | c-1 | ["public.greptime_otel_resource_info"] | +| host | h-1 | ["public.greptime_otel_resource_info"] | +| service | shop/api | ["public.greptime_otel_resource_info"] | +| service | worker | ["public.greptime_otel_resource_info"] | +| service.instance | shop/api,inst-1 | ["public.greptime_otel_resource_info"] | +| service.instance | shop/api,inst-2 | ["public.greptime_otel_resource_info"] | ++------------------+-----------------+----------------------------------------+ + +-- SQLNESS PROTOCOL MYSQL +select src_type, src_id, dst_type, dst_id, rel_type, provenance +from greptime_private.semantic_relationships +order by rel_type, src_id, dst_id; + ++------------------+-----------------+-----------+----------+----------+------------+ +| src_type | src_id | dst_type | dst_id | rel_type | provenance | ++------------------+-----------------+-----------+----------+----------+------------+ +| service.instance | shop/api,inst-1 | service | shop/api | part_of | attribute | +| service.instance | shop/api,inst-2 | service | shop/api | part_of | attribute | +| container | c-1 | host | h-1 | runs_on | attribute | +| service.instance | shop/api,inst-1 | container | c-1 | runs_on | attribute | +| service.instance | shop/api,inst-1 | host | h-1 | runs_on | attribute | ++------------------+-----------------+-----------+----------+----------+------------+ + +drop table greptime_otel_resource_info; + +Affected Rows: 0 + +-- The descriptor whitelist is gated on source=opentelemetry: the same table +-- shape stamped as a prometheus source contributes nothing. +create table greptime_otel_resource_info ( + greptime_timestamp timestamp(3) time index, + job string, + instance string, + "host.id" string, + greptime_value double, + primary key (job, instance, "host.id") +) with ( + 'greptime.semantic.signal_type' = 'metric', + 'greptime.semantic.source' = 'prometheus' +); + +Affected Rows: 0 + +insert into greptime_otel_resource_info values + (now(), 'shop/api', 'inst-1', 'h-1', 1); + +Affected Rows: 1 + +-- SQLNESS PROTOCOL MYSQL +select entity_type, entity_id +from greptime_private.semantic_entities +order by entity_type, entity_id; + +affected_rows: 0 + +drop table greptime_otel_resource_info; + +Affected Rows: 0 + +-- Cross-signal identity: the same service reaches the graph as trace spans and +-- as an ingestion-synthesized descriptor. The two sources name their identity +-- columns differently and only the metric side pre-composes the namespace into +-- job, so this asserts one entity per service and per instance, sourced from +-- both tables. +create table graph_xsignal_traces ( + "timestamp" timestamp(9) time index, + trace_id string, + span_id string, + parent_span_id string, + span_kind string, + span_status_code string, + service_name string, + duration_nano bigint unsigned, + "resource_attributes.service.namespace" string, + "resource_attributes.service.instance.id" string, + primary key (service_name, "resource_attributes.service.namespace", + "resource_attributes.service.instance.id") +) with ('table_data_model' = 'greptime_trace_v1', 'append_mode' = 'true'); + +Affected Rows: 0 + +create table greptime_otel_resource_info ( + greptime_timestamp timestamp(3) time index, + job string, + instance string, + "service.name" string, + "service.namespace" string, + greptime_value double, + primary key (job, instance, "service.name", "service.namespace") +) with ( + 'greptime.semantic.signal_type' = 'metric', + 'greptime.semantic.source' = 'opentelemetry', + 'greptime.semantic.metric.type' = 'info' +); + +Affected Rows: 0 + +insert into graph_xsignal_traces values + (now(), 't1', 's1', NULL, 'SPAN_KIND_SERVER', 'STATUS_CODE_UNSET', 'api', 100, 'shop', 'inst-1'), + (now(), 't2', 's2', NULL, 'SPAN_KIND_SERVER', 'STATUS_CODE_UNSET', 'worker', 100, '', 'inst-2'); + +Affected Rows: 2 + +insert into greptime_otel_resource_info values + (now(), 'shop/api', 'inst-1', 'api', 'shop', 1), + (now(), 'worker', 'inst-2', 'worker', '', 1); + +Affected Rows: 2 + +-- SQLNESS PROTOCOL MYSQL +select entity_type, entity_id, source_tables +from greptime_private.semantic_entities +where entity_type in ('service', 'service.instance') +order by entity_type, entity_id, source_tables; + ++------------------+-----------------+----------------------------------------+ +| entity_type | entity_id | source_tables | ++------------------+-----------------+----------------------------------------+ +| service | shop/api | ["public.graph_xsignal_traces"] | +| service | shop/api | ["public.greptime_otel_resource_info"] | +| service | worker | ["public.graph_xsignal_traces"] | +| service | worker | ["public.greptime_otel_resource_info"] | +| service.instance | shop/api,inst-1 | ["public.graph_xsignal_traces"] | +| service.instance | shop/api,inst-1 | ["public.greptime_otel_resource_info"] | +| service.instance | worker,inst-2 | ["public.graph_xsignal_traces"] | +| service.instance | worker,inst-2 | ["public.greptime_otel_resource_info"] | ++------------------+-----------------+----------------------------------------+ + +drop table graph_xsignal_traces; + +Affected Rows: 0 + +drop table greptime_otel_resource_info; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/system/semantic_graph.sql b/tests/cases/standalone/common/system/semantic_graph.sql index bfc9ea8fd7b..60d95630d91 100644 --- a/tests/cases/standalone/common/system/semantic_graph.sql +++ b/tests/cases/standalone/common/system/semantic_graph.sql @@ -387,3 +387,122 @@ drop table kube_service_info; drop table target_info; drop table http_requests_total; + +-- OTel conventions: the ingestion-synthesized greptime_otel_resource_info descriptor +-- (stamped signal_type=metric + source=opentelemetry + metric.type=info) gets implicit +-- declarations under the raw OTel column names. host/container identities are +-- the stable ids only, so rows with an empty host.id / container.id link no +-- infrastructure, and a row without an instance value declares no +-- service.instance. The same table name stamped source=prometheus is not +-- whitelisted. +create table greptime_otel_resource_info ( + greptime_timestamp timestamp(3) time index, + job string, + instance string, + "service.name" string, + "service.namespace" string, + "host.id" string, + "host.name" string, + "container.id" string, + "container.name" string, + greptime_value double, + primary key (job, instance, "service.name", "service.namespace", + "host.id", "host.name", "container.id", "container.name") +) with ( + 'greptime.semantic.signal_type' = 'metric', + 'greptime.semantic.source' = 'opentelemetry', + 'greptime.semantic.metric.type' = 'info' +); + +insert into greptime_otel_resource_info values + (now(), 'shop/api', 'inst-1', 'api', 'shop', 'h-1', 'node-a', 'c-1', 'api-ctr', 1), + (now(), 'shop/api', 'inst-2', 'api', 'shop', '', 'laptop', '', '', 1), + (now(), 'worker', '', 'worker', '', 'h-1', 'node-a', '', '', 1); + +-- SQLNESS PROTOCOL MYSQL +select entity_type, entity_id, source_tables +from greptime_private.semantic_entities +order by entity_type, entity_id, source_tables; + +-- SQLNESS PROTOCOL MYSQL +select src_type, src_id, dst_type, dst_id, rel_type, provenance +from greptime_private.semantic_relationships +order by rel_type, src_id, dst_id; + +drop table greptime_otel_resource_info; + +-- The descriptor whitelist is gated on source=opentelemetry: the same table +-- shape stamped as a prometheus source contributes nothing. +create table greptime_otel_resource_info ( + greptime_timestamp timestamp(3) time index, + job string, + instance string, + "host.id" string, + greptime_value double, + primary key (job, instance, "host.id") +) with ( + 'greptime.semantic.signal_type' = 'metric', + 'greptime.semantic.source' = 'prometheus' +); + +insert into greptime_otel_resource_info values + (now(), 'shop/api', 'inst-1', 'h-1', 1); + +-- SQLNESS PROTOCOL MYSQL +select entity_type, entity_id +from greptime_private.semantic_entities +order by entity_type, entity_id; + +drop table greptime_otel_resource_info; + +-- Cross-signal identity: the same service reaches the graph as trace spans and +-- as an ingestion-synthesized descriptor. The two sources name their identity +-- columns differently and only the metric side pre-composes the namespace into +-- job, so this asserts one entity per service and per instance, sourced from +-- both tables. +create table graph_xsignal_traces ( + "timestamp" timestamp(9) time index, + trace_id string, + span_id string, + parent_span_id string, + span_kind string, + span_status_code string, + service_name string, + duration_nano bigint unsigned, + "resource_attributes.service.namespace" string, + "resource_attributes.service.instance.id" string, + primary key (service_name, "resource_attributes.service.namespace", + "resource_attributes.service.instance.id") +) with ('table_data_model' = 'greptime_trace_v1', 'append_mode' = 'true'); + +create table greptime_otel_resource_info ( + greptime_timestamp timestamp(3) time index, + job string, + instance string, + "service.name" string, + "service.namespace" string, + greptime_value double, + primary key (job, instance, "service.name", "service.namespace") +) with ( + 'greptime.semantic.signal_type' = 'metric', + 'greptime.semantic.source' = 'opentelemetry', + 'greptime.semantic.metric.type' = 'info' +); + +insert into graph_xsignal_traces values + (now(), 't1', 's1', NULL, 'SPAN_KIND_SERVER', 'STATUS_CODE_UNSET', 'api', 100, 'shop', 'inst-1'), + (now(), 't2', 's2', NULL, 'SPAN_KIND_SERVER', 'STATUS_CODE_UNSET', 'worker', 100, '', 'inst-2'); + +insert into greptime_otel_resource_info values + (now(), 'shop/api', 'inst-1', 'api', 'shop', 1), + (now(), 'worker', 'inst-2', 'worker', '', 1); + +-- SQLNESS PROTOCOL MYSQL +select entity_type, entity_id, source_tables +from greptime_private.semantic_entities +where entity_type in ('service', 'service.instance') +order by entity_type, entity_id, source_tables; + +drop table graph_xsignal_traces; + +drop table greptime_otel_resource_info; diff --git a/tests/conf/frontend-test.toml.template b/tests/conf/frontend-test.toml.template index 632947a3cde..5a879e04882 100644 --- a/tests/conf/frontend-test.toml.template +++ b/tests/conf/frontend-test.toml.template @@ -7,3 +7,6 @@ parallelism = 4 [event_recorder] event_types = [] + +[otlp] +experimental_enable_resource_info = true diff --git a/tests/conf/standalone-test.toml.template b/tests/conf/standalone-test.toml.template index ebe25f82487..8d7e307d511 100644 --- a/tests/conf/standalone-test.toml.template +++ b/tests/conf/standalone-test.toml.template @@ -53,3 +53,6 @@ retry_delay = "500ms" [event_recorder] event_types = [] + +[otlp] +experimental_enable_resource_info = true diff --git a/tests/perf/AGENTS.md b/tests/perf/AGENTS.md index 1f13dea58dc..822b99080f9 100644 --- a/tests/perf/AGENTS.md +++ b/tests/perf/AGENTS.md @@ -3,9 +3,23 @@ - Keep GitHub Actions YAML thin. Put non-trivial control flow, case expansion, report generation, and metadata writing in scripts under `.github/scripts/`; workflow steps should mostly invoke those scripts. +- Runner lifecycle: the default path provisions one ephemeral Aliyun ECS + instance per run via `.github/scripts/aliyun-ecs-runner-provision.py` and + always releases it via `aliyun-ecs-runner-teardown.py`; a scheduled janitor + workflow sweeps leftovers. Build caches live on that instance's system disk + and are discarded with the VM. Runs do not share a workflow concurrency + group. The ECS custom image is built from the runner Dockerfile by + `.github/runner-scale-sets/query-regression/ecs-image/build-ecs-image.py`; + keep the Dockerfile the single source of the tool contract. Dispatching with + any other `runner` value treats it as a literal self-hosted runner label + (see `ecs-image/bootstrap-runner-host.sh` for preparing such a host). - Query regression PR runs should build base/candidate binaries once, then run the default case set. Do not hard-code a single case such as `promql_pushdown_7913` into the workflow path. +- Scheduled nightly comparison lives in `query-regression-nightly.yml`: it + waits for a successful Nightly Build, then calls `query-regression.yml` + with the previous vs current nightly SHAs. Keep SHA selection in + `.github/scripts/query-regression-nightly-refs.py`. - The case DSL is not required to keep compatibility inside this PR. When the DSL changes, update TOML cases, the outer lifecycle script, Rust helpers, and docs together. @@ -19,8 +33,10 @@ - Keep the direct-SST generator generic. Issue-specific behavior belongs in case files and thresholds, not in Rust generator logic. - Before pushing perf harness changes, run at least: - - the Python tests in the `Test query regression tooling` step of - `.github/workflows/query-regression.yml` + - the Python tests in the `test-tooling` job of + `.github/workflows/query-regression.yml` (ubuntu-latest, not the ECS runner). + The Checks workflow runs the same tests on ordinary PRs so they are not + gated on the `query-regression` / `heavy-regression` labels. - `cargo fmt --all -- --check` - `cargo build -p cmd --bin query_perf_fixture --features dev-tools` - `cargo build -p cmd --bin query_regression_runner --features dev-tools` diff --git a/tests/perf/README.md b/tests/perf/README.md index cbe03b51794..47495ca531f 100644 --- a/tests/perf/README.md +++ b/tests/perf/README.md @@ -255,11 +255,12 @@ gh workflow run query-regression.yml \ -f candidate_ref= \ -f cargo_profile=nightly \ -f http_timeout=300 \ - -f runner=perf-regression-8-cores + -f runner=aliyun-ecs ``` -The selected ARC scale set must already be deployed with the runner-image -digest built from the current query-regression Dockerfile. +The `aliyun-ecs` path provisions a fresh ECS instance per run from the custom +image built from the current query-regression Dockerfile; see +`.github/runner-scale-sets/query-regression/README.md` for its configuration. ## Generator contract diff --git a/tests/perf/test_aliyun_ecs_runner_scripts.py b/tests/perf/test_aliyun_ecs_runner_scripts.py new file mode 100644 index 00000000000..581fc88abd0 --- /dev/null +++ b/tests/perf/test_aliyun_ecs_runner_scripts.py @@ -0,0 +1,122 @@ +#!/usr/bin/env python3 +# Copyright 2023 Greptime Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Coverage for the pure parts of the Aliyun ECS runner provision/teardown scripts.""" + +import base64 +import importlib.util +import sys +import unittest +from datetime import datetime, timedelta, timezone +from pathlib import Path + + +SCRIPTS_DIR = Path(__file__).parents[2] / ".github/scripts" + + +def load_module(name: str, filename: str): + spec = importlib.util.spec_from_file_location(name, SCRIPTS_DIR / filename) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +provision = load_module("aliyun_ecs_runner_provision_under_test", "aliyun-ecs-runner-provision.py") +teardown = load_module("aliyun_ecs_runner_teardown_under_test", "aliyun-ecs-runner-teardown.py") + + +class ProvisionNamingTest(unittest.TestCase): + def test_runner_name_and_label_derive_from_run_id(self) -> None: + self.assertEqual(provision.runner_name_for_run("12345"), "qreg-ecs-12345") + self.assertEqual(provision.runner_label_for_run("12345"), "query-regression-ecs-12345") + + +class ProvisionUserDataTest(unittest.TestCase): + def render(self) -> str: + return provision.render_user_data( + runner_name="qreg-ecs-12345", + runner_label="query-regression-ecs-12345", + runner_token="TOKEN", + repo="GreptimeTeam/greptimedb", + ) + + def test_user_data_creates_cache_paths_on_the_system_disk(self) -> None: + script = self.render() + self.assertNotIn("DISK_SERIAL", script) + self.assertNotIn("mount --bind", script) + self.assertNotIn("mkfs.ext4", script) + for destination in provision.CACHE_PATHS: + self.assertIn(f'"{destination}"', script) + + def test_user_data_wires_runner_registration(self) -> None: + script = self.render() + self.assertIn("RUNNER_NAME=qreg-ecs-12345", script) + self.assertIn("RUNNER_LABELS=query-regression-ecs-12345", script) + self.assertIn("RUNNER_TOKEN=TOKEN", script) + self.assertIn("REPO_URL=https://github.com/GreptimeTeam/greptimedb", script) + self.assertIn("PATH=/opt/cargo/bin:", script) + self.assertIn("systemctl restart --no-block ephemeral-github-runner.service", script) + + def test_encode_user_data_round_trips(self) -> None: + script = self.render() + self.assertEqual( + base64.b64decode(provision.encode_user_data(script)).decode("utf-8"), script + ) + + def test_user_data_enables_swap_and_masks_oomd(self) -> None: + script = self.render() + self.assertIn("systemctl mask systemd-oomd.socket systemd-oomd.service", script) + self.assertIn("systemctl mask unattended-upgrades.service apt-daily.timer apt-daily-upgrade.timer", script) + self.assertIn('APT::Periodic::Unattended-Upgrade "0"', script) + self.assertIn(f'fallocate --length {provision.SWAP_SIZE_GIB}G "{provision.SWAP_FILE}"', script) + self.assertIn(f'swapon "{provision.SWAP_FILE}"', script) + self.assertIn("sysctl --write vm.swappiness=10", script) + self.assertIn("OOMPolicy=continue", script) + self.assertNotIn("OOMScoreAdjust", script) + self.assertLess(script.index("swapon"), script.index("systemctl restart --no-block ephemeral-github-runner.service")) + + +class TeardownExpiryTest(unittest.TestCase): + NOW = datetime(2026, 8, 17, 6, 0, tzinfo=timezone.utc) + TTL = timedelta(hours=4) + + def test_parse_creation_time_formats(self) -> None: + self.assertEqual( + teardown.parse_creation_time("2026-08-17T01:02:03Z"), + datetime(2026, 8, 17, 1, 2, 3, tzinfo=timezone.utc), + ) + self.assertEqual( + teardown.parse_creation_time("2026-08-17T01:02Z"), + datetime(2026, 8, 17, 1, 2, tzinfo=timezone.utc), + ) + with self.assertRaises(ValueError): + teardown.parse_creation_time("not-a-time") + + def test_expired_instance_names_selects_only_old_instances(self) -> None: + instances = [ + ("i-old", "qreg-ecs-1", "2026-08-17T01:00Z"), # 5h old: expired + ("i-edge", "qreg-ecs-2", "2026-08-17T02:00Z"), # exactly TTL: expired + ("i-fresh", "qreg-ecs-3", "2026-08-17T05:30Z"), # 30m old: kept + ] + self.assertEqual( + teardown.expired_instance_names(instances, self.NOW, self.TTL), + [("i-old", "qreg-ecs-1"), ("i-edge", "qreg-ecs-2")], + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/perf/test_query_regression_nightly_refs.py b/tests/perf/test_query_regression_nightly_refs.py new file mode 100644 index 00000000000..74b168bb464 --- /dev/null +++ b/tests/perf/test_query_regression_nightly_refs.py @@ -0,0 +1,120 @@ +#!/usr/bin/env python3 +# Copyright 2023 Greptime Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for consecutive Nightly Build SHA selection.""" + +import importlib.util +import sys +import unittest +from pathlib import Path + + +SCRIPTS_DIR = Path(__file__).parents[2] / ".github/scripts" + + +def load_module(): + spec = importlib.util.spec_from_file_location( + "query_regression_nightly_refs_under_test", + SCRIPTS_DIR / "query-regression-nightly-refs.py", + ) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +refs = load_module() + + +def run(*, run_id: int, sha: str, branch: str = "main", created_at: str = "2026-08-24T00:00:00Z"): + return refs.WorkflowRun( + id=run_id, + head_sha=sha, + head_branch=branch, + html_url=f"https://github.com/example/run/{run_id}", + created_at=created_at, + conclusion="success", + event="schedule", + ) + + +class SelectNightlyRefsTest(unittest.TestCase): + def test_picks_newest_and_previous_on_same_branch(self) -> None: + pair = refs.select_base_and_candidate( + [ + run(run_id=3, sha="ccc", created_at="2026-08-22T00:00:00Z"), + run(run_id=2, sha="bbb", created_at="2026-08-21T00:00:00Z"), + run(run_id=1, sha="aaa", created_at="2026-08-20T00:00:00Z"), + ] + ) + self.assertFalse(pair.skip) + assert pair.candidate is not None and pair.base is not None + self.assertEqual(pair.candidate.head_sha, "ccc") + self.assertEqual(pair.base.head_sha, "bbb") + + def test_candidate_run_id_uses_that_run_and_the_next_older(self) -> None: + pair = refs.select_base_and_candidate( + [ + run(run_id=3, sha="ccc", created_at="2026-08-22T00:00:00Z"), + run(run_id=2, sha="bbb", created_at="2026-08-21T00:00:00Z"), + run(run_id=1, sha="aaa", created_at="2026-08-20T00:00:00Z"), + ], + candidate_run_id=2, + ) + self.assertFalse(pair.skip) + assert pair.candidate is not None and pair.base is not None + self.assertEqual(pair.candidate.id, 2) + self.assertEqual(pair.base.id, 1) + + def test_skips_when_consecutive_nightlies_share_a_sha(self) -> None: + pair = refs.select_base_and_candidate( + [ + run(run_id=2, sha="same", created_at="2026-08-22T00:00:00Z"), + run(run_id=1, sha="same", created_at="2026-08-21T00:00:00Z"), + ] + ) + self.assertTrue(pair.skip) + self.assertIn("matches candidate", pair.reason) + + def test_skips_other_branches_when_picking_previous(self) -> None: + pair = refs.select_base_and_candidate( + [ + run(run_id=3, sha="ccc", branch="main", created_at="2026-08-22T00:00:00Z"), + run(run_id=2, sha="other", branch="feat", created_at="2026-08-21T12:00:00Z"), + run(run_id=1, sha="aaa", branch="main", created_at="2026-08-21T00:00:00Z"), + ] + ) + self.assertFalse(pair.skip) + assert pair.base is not None + self.assertEqual(pair.base.head_sha, "aaa") + + def test_skips_when_only_one_nightly_exists(self) -> None: + pair = refs.select_base_and_candidate( + [run(run_id=1, sha="aaa")], + ) + self.assertTrue(pair.skip) + self.assertIn("no previous", pair.reason) + + def test_explicit_refs_bypass_github(self) -> None: + pair = refs.override_pair("base-sha", "cand-sha") + self.assertFalse(pair.skip) + assert pair.base is not None and pair.candidate is not None + self.assertEqual(pair.base.head_sha, "base-sha") + self.assertEqual(pair.candidate.head_sha, "cand-sha") + + +if __name__ == "__main__": + unittest.main()