diff --git a/.github/workflows/continuous-integration.yml b/.github/workflows/continuous-integration.yml index c9d1479c..f384396a 100644 --- a/.github/workflows/continuous-integration.yml +++ b/.github/workflows/continuous-integration.yml @@ -184,7 +184,6 @@ jobs: run: scripts/container-scripts/test rust-crate: name: Test rust crate - if: ${{ false }} # FIXME: turn back on before v0.10 release runs-on: ubuntu-latest needs: - changes @@ -215,15 +214,29 @@ jobs: persist-credentials: false - uses: Swatinem/rust-cache@e18b497796c12c097a38f9edb9d0641fb99eee32 # v2 - uses: astral-sh/setup-uv@08807647e7069bb48b6ef5acd8ec9567f424441b # v8.1.0 - - name: Install pypgstac - working-directory: src/pypgstac + - name: Build pgstac test databases run: | - uv venv /tmp/ci-venv - uv pip install --python /tmp/ci-venv/bin/python .[psycopg] - echo "/tmp/ci-venv/bin" >> "$GITHUB_PATH" - - name: Migrate pgstac - run: pypgstac migrate - - name: Set search_path - run: psql -c "ALTER ROLE username SET search_path TO pgstac, public;" + scripts/pgstac-rs-test-db + scripts/pgstac-rs-test-db pgstac_rs_ingest_template + scripts/pgstac-rs-test-rich-db + - name: Compile-check the feature matrix + run: | + set -euo pipefail + for feat in \ + "--no-default-features" \ + "--no-default-features --features pool" \ + "--no-default-features --features export" \ + "--features cli" \ + "--all-features"; do + echo "::group::cargo check $feat" + cargo check $feat --all-targets --manifest-path src/pgstac-rs/Cargo.toml + echo "::endgroup::" + done + - name: Check docs (deny warnings) + env: + RUSTDOCFLAGS: -D warnings + run: cargo doc --all-features --no-deps --manifest-path src/pgstac-rs/Cargo.toml - name: Test - run: cargo test -p pgstac --all-features --manifest-path src/pgstac-rs/Cargo.toml + env: + PGSTAC_RS_TEST_DB: postgresql://username:password@localhost:5439/pgstac_rs_test_rich + run: cargo test --features cli --manifest-path src/pgstac-rs/Cargo.toml diff --git a/CHANGELOG.md b/CHANGELOG.md index 746d9ec0..a406a18c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -15,9 +15,16 @@ and this project adheres to [Semantic Versioning](http://semver.org/). - Add deterministic SHA-256 `content_hash` to STAC items to track data changes across migrations. - Add `pgstac_updated_at` column to items table as part of separating STAC property updates from database metadata updates. - Deterministic Planetary Computer benchmark fixture manifest + fetch tooling for `naip`, `sentinel-2-l2a`, and `landsat-c2-l2` (1000 items per collection), plus CI/manual benchmark workflows that emit JSON/CSV/Markdown artifacts and branch comparison reports. +- `src/pgstac-rs` read/streaming engine: a Rust-side keyset search that drives `search_plan`/`collection_search_plan` and does the band stepping, hydration, token minting, and the STAC `fields` include/exclude projection client-side (page-equivalent to SQL `search()`); a flat-memory streaming iterator (`query_raw` portal on one pooled connection + per-new-fragment fetch on a parallel connection, so memory stays flat regardless of result size); byte-identical Rust hydration (EWKB→GeoJSON, serialize-time fragment-asset merge that never deep-parses large asset arrays, `bbox` emitted from raw bytes to preserve numeric precision); collection search, item/collection getters, and queryables; a pooled hydration-invariant cache; and a parallel-context primitive. +- `pgstac` CLI (`src/pgstac-rs`, `cli` feature): `dump` a pgstac instance to a directory / `*.tar(.zst)` / S3 (object-store) / stdout as fully-hydrated stac-geoparquet (one file per partition) + collection/queryables/settings JSON + a sha256'd manifest, with collection/datetime/bbox prefilters, parallel and consistent-snapshot modes, and dry-run; `search` to stream results off the keyset engine, selecting the format with rustac's `stac_io::Format` spelling — `--format ndjson` (default), `json` (one ItemCollection page), or `geoparquet[]` (compression carried in the format string, e.g. `geoparquet[snappy]`). +- `pypgstac_rs` Python wheel (`src/pgstac-rs`, `python` feature, built with maturin): a pyo3 extension exposing the async read API (search, collection search, getters, queryables, parallel match count) over a startup-created connection pool, for use in `stac-fastapi-pgstac`. +- New [Rust client & CLI](https://stac-utils.github.io/pgstac/pgstac-rs/) reference documenting the crate features, the full read/pool API, every `pgstac dump` / `pgstac search` option, and all connection + test environment variables. +- `src/pgstac-rs`: a `store` cargo feature that gates the `object_store`-backed remote dump sink (local fs + S3/GCS/Azure), so the remote sink can be built without the full `cli` feature. ### Changed +- `search_plan`: bake the collection clamp as a literal so the datetime-band query is parameterized only by `$1`/`$2`/`$3` (band low/high, limit) and can be prepared once by a streaming client; fix the non-datetime branch's datetime clamp to be exclusive of the next month (`< months[last] + 1 month`) so items dated after the start of the final month are not dropped. +- `fields_to_itemcols`: emit excluded heavy columns as `NULL::type AS ` (named) instead of an unaliased `NULL::type`, so a client preparing the `search_plan` query keeps a stable result schema across any `fields` value while still never transferring the heavy value. It also nulls `fragment_id` in the projection when the requested fields need no shared fragment (`needs_fragment` is false), so a client driving the `search_plan` query skips the per-row `item_fragments` lookup entirely — the fragment-skip rides in the returned projection rather than a separate flag. - Replaced expensive row-based trigger for item inserts with optimized SQL/PLPGSQL hydration strategies to improve ingestion throughput. - Update pypgstac loaders to dynamically generate hashes during ingestion where required, avoiding trigger recalculation. - Add tombstone table `items_deleted_log` and `pgstac_updated_at` metadata column to items table. diff --git a/docs/mkdocs.yml b/docs/mkdocs.yml index fb863e31..0eec206c 100644 --- a/docs/mkdocs.yml +++ b/docs/mkdocs.yml @@ -21,6 +21,7 @@ nav: - PgSTAC: "pgstac.md" - Promoted Fields: "promoted-fields.md" - pyPgSTAC: "pypgstac.md" + - "Rust client & CLI": "pgstac-rs.md" - Performance: - item_size_analysis.ipynb - Development - Contributing: "contributing.md" diff --git a/docs/src/pgstac-rs.md b/docs/src/pgstac-rs.md new file mode 100644 index 00000000..0bbf7666 --- /dev/null +++ b/docs/src/pgstac-rs.md @@ -0,0 +1,160 @@ +# Rust client & CLI (`pgstac-rs`) + +`src/pgstac-rs` is a Rust crate (`pgstac`) and a command-line tool (`pgstac`) that talk to a PgSTAC +database. It provides: + +- a **read engine** that drives `search_plan` / `collection_search_plan` and does the band stepping, + hydration (EWKB→GeoJSON, fragment merge), keyset token minting, and the STAC `fields` include/exclude + projection **in Rust** — page-equivalent to SQL `search()` but with flat memory under streaming and + the CPU moved off the database. When a search's `fields` need no shared fragment, `search_plan` nulls + `fragment_id` in the projection it returns, so the engine skips the per-row `item_fragments` lookup; +- a **connection pool** (`PgstacPool`, the `pool` feature) with the async read API; +- a **dump/export** library + the `pgstac` **CLI** (`export` / `cli` features); +- a **Python wheel** (`pypgstac_rs`, the `python` feature, built with [maturin](https://github.com/PyO3/maturin)). + +## Cargo features + +| Feature | Adds | +| ----------- | ---- | +| *(default)* | The `Pgstac` trait over any `tokio_postgres::GenericClient`. | +| `pool` | `PgstacPool`: pooled async read API (rustls TLS, PgBouncer-safe). | +| `export` | The dump library (scan partitions → stac-geoparquet + manifest). | +| `cli` | The `pgstac` binary (implies `export`). | +| `python` | The `pypgstac_rs` pyo3 extension module. | + +## Environment variables + +The library and CLI resolve a connection from a single DSN if one is given, otherwise from the standard +**libpq** environment variables. The connection `search_path` is **always** set to `pgstac, public`. + +### Connection + +| Variable | Purpose | +| -------- | ------- | +| `PGSTAC_DSN` | Full Postgres connection string (URL or key/value). The CLI `--dsn` flag defaults to this. | +| `DATABASE_URL` | Fallback DSN if `PGSTAC_DSN` is unset. | +| `PGHOST` | Server host. | +| `PGPORT` | Server port. | +| `PGDATABASE` | Database name. | +| `PGUSER` | Username. | +| `PGPASSWORD` | Password. | +| `PGOPTIONS` | Extra startup options. | +| `PGAPPNAME` | `application_name` for the connection. | +| `PGCONNECT_TIMEOUT` | Connection timeout, in seconds. | +| `PGSSLMODE` | TLS mode (`disable`, `prefer`, `require`, `verify-ca`, `verify-full`). | +| `PGSSLROOTCERT` | Path to the CA certificate. | +| `PGSSLCERT` | Path to the client certificate. | +| `PGSSLKEY` | Path to the client key. | + +A DSN (`PGSTAC_DSN` / `--dsn`) takes precedence; the individual `PG*` variables fill in any field the +DSN does not set. + +### Test fixtures + +The integration tests skip (rather than fail) when their database is unreachable. Point them at your +fixtures with: + +| Variable | Default | Used by | +| -------- | ------- | ------- | +| `PGSTAC_RS_TEST_DB` | `postgresql://username:password@localhost:5439/postgis` | clean-template tests | +| `PGSTAC_RS_TEST_RICH_DB` | `…/pgstac_rs_test_rich` | search/stream/hydration/collection parity tests | +| `PGSTAC_RS_TEST_TEMPLATE` | `pgstac_rs_test_template` | pool clone-from-template tests | +| `PGSTAC_PARITY_DB_010` | `…/a_parity010` | 0.10 dump end-to-end tests | +| `PGSTAC_PARITY_DB_0911` | `…/a_parity0911` | 0.9.11 dump end-to-end tests | + +## Library + +```rust,no_run +use pgstac::{ConnectConfig, PgstacPool}; +use futures::StreamExt; +use serde_json::json; + +# tokio_test::block_on(async { +// Pool from the environment (PGSTAC_DSN or the PG* vars). Create once at startup. +let pool = PgstacPool::connect(ConnectConfig::from_env()).await.unwrap(); + +// A bounded page, with keyset next/prev tokens. +let page = pool.search_page(&json!({"collections": ["landsat-c2-l2"]}), None, 10).await.unwrap(); + +// Stream every match with flat memory. +let mut items = Box::pin(pool.search_items(json!({"collections": ["landsat-c2-l2"]}), None, None)); +while let Some(item) = items.next().await { let _ = item.unwrap(); } +# }) +``` + +Pool methods: `search_page`, `search_items` (stream), `search_collect_items`, `stream_ndjson`, +`search_matched`, `collection_search`, `get_item`, `get_collection`, `get_queryables`. + +## CLI + +```text +pgstac + dump Dump a pgstac instance (or a subset) to a directory, tar, S3, or stdout + search Stream a search/CQL2 result as NDJSON / ItemCollection / geoparquet (0.10 only) +``` + +Both subcommands read `--dsn` from `$PGSTAC_DSN` when the flag is omitted. + +### `pgstac dump` + +Writes fully-hydrated items as stac-geoparquet (one file per partition, ordered by `datetime, id`) plus +`collection.json` / `queryables.json` / `settings.json` and a sha256'd `manifest.json` (written last — +its presence marks a complete dump). + +| Option | Default | Description | +| ------ | ------- | ----------- | +| `--dsn ` | `$PGSTAC_DSN` / local dev | Postgres connection string. | +| `-o, --out ` | *(required)* | Destination: a directory, `s3://bucket/prefix` (or other object-store URL), a `*.tar` / `*.tar.zst` file, or `-` for stdout. | +| `-c, --collection ` | *(all)* | Restrict to these collection ids (repeatable). | +| `--datetime-start ` | — | Inclusive datetime prefilter start (requires `--datetime-end`). | +| `--datetime-end ` | — | Exclusive datetime prefilter end (requires `--datetime-start`). | +| `--bbox ` | — | Bbox prefilter (EPSG:4326). | +| `--compression ` | `zstd` | Parquet codec: `zstd`, `snappy`, `uncompressed`. | +| `--skip-errors` | off | Continue past per-item errors, recording skips in the report. | +| `--memory-budget ` | ~25% RAM | Memory budget for the buffered geoparquet path. | +| `--concurrency ` | `1` | Max partitions dumped concurrently (>1 opens one extra connection per worker). | +| `--consistent` | off | Dump the whole instance under one repeatable-read snapshot (implies a parallel run). | +| `--tar-zstd` | off | Write a compressed `.tar.zst` when `--out` is a `.tar` path. | +| `--dry-run` | off | Report the plan without writing data. | + +### `pgstac search` + +Streams a 0.10 search off the keyset engine. + +| Option | Default | Description | +| ------ | ------- | ----------- | +| `--dsn ` | `$PGSTAC_DSN` / local dev | Postgres connection string. | +| `-o, --out ` | `-` (stdout) | A file path or `-` for stdout. | +| `--format ` | `ndjson` | `ndjson` (streams all pages), `itemcollection` (one SQL-faithful page with next/prev + context), or `geoparquet`. | +| `-c, --collection ` | *(all)* | Restrict to these collection ids (repeatable). | +| `--id ` | — | Restrict to these item ids (repeatable). | +| `--bbox ` | — | Bbox (EPSG:4326). | +| `--datetime
` | — | Datetime / interval (STAC datetime syntax). | +| `--filter ` | — | CQL2-JSON filter (a JSON object). | +| `--limit ` | server default | Page size (the keyset limit). | +| `--token ` | — | Continuation token (`next:…` / `prev:…`). With `itemcollection`, returns exactly that one page. | +| `--max-items ` | *(unbounded)* | Cap total items streamed (`ndjson` / `geoparquet`). | +| `--compression ` | `zstd` | Parquet codec (`geoparquet` only). | + +## Python wheel + +```sh +maturin build -m src/pgstac-rs/Cargo.toml # produces the pypgstac_rs wheel +``` + +```python +import asyncio, orjson, pypgstac_rs + +async def main(): + pool = await pypgstac_rs.Pgstac.connect() # libpq env, or pass a dsn string + body = orjson.dumps({"collections": ["landsat-c2-l2"], "limit": 10}).decode() + fc = orjson.loads(await pool.search(body)) + print(fc["numberReturned"]) + +asyncio.run(main()) +``` + +`Pgstac` methods (each `async`, returning a JSON string): `connect(dsn=None)`, `search(search, +token=None, limit=10)`, `search_collect(search, max_items=None)`, `search_matched(search)`, +`collection_search(search, token=None)`, `get_item(collection_id, item_id)`, +`get_collection(collection_id)`, `get_queryables(collection_id=None)`. diff --git a/scripts/container-scripts/test b/scripts/container-scripts/test index 6d2d0256..717f6743 100755 --- a/scripts/container-scripts/test +++ b/scripts/container-scripts/test @@ -35,6 +35,7 @@ Options: --pypgstac Run Python tests. --migrations Run migration-path validation. --pgdump Run pg_dump / pg_restore validation. + --rust Run the pgstac-rs Rust tests (cargo test --features cli). --nomigrations Run everything except migration and pg_dump tests. --v Raise client_min_messages to notice. --vv Raise client_min_messages to debug1. @@ -368,6 +369,30 @@ EOSQL export PGDATABASE="${_prev_pgdatabase}" } +function test_rust(){ + # The Rust tests clone per-test databases from three templates (pgstac_rs_test_template / + # _ingest_template / _rich) built by the builder scripts. Passing an explicit PGHOST makes those + # builders target the dev database (pgstac:5432 in-container) rather than .env's host. + local host="${PGHOST:-pgstac}" port="${PGPORT:-5432}" user="${PGUSER:-username}" pass="${PGPASSWORD:-password}" + local scripts_dir base + scripts_dir="$(cd "$SRCDIR/../scripts" && pwd)" + base="postgresql://${user}:${pass}@${host}:${port}" + + echo "=== Building pgstac-rs test templates on ${host}:${port} ===" + PGHOST="$host" PGPORT="$port" PGUSER="$user" PGPASSWORD="$pass" "$scripts_dir/pgstac-rs-test-db" + PGHOST="$host" PGPORT="$port" PGUSER="$user" PGPASSWORD="$pass" "$scripts_dir/pgstac-rs-test-db" pgstac_rs_ingest_template + PGHOST="$host" PGPORT="$port" PGUSER="$user" PGPASSWORD="$pass" "$scripts_dir/pgstac-rs-test-rich-db" + + echo "=== Running pgstac-rs cargo tests ===" + ( + cd "$SRCDIR/pgstac-rs" + PGSTAC_RS_TEST_BASE="$base" \ + PGSTAC_RS_TEST_DB="${base}/pgstac_rs_test_rich" \ + cargo test --features cli + ) + echo "pgstac-rs Rust tests passed!" +} + FORMATTING=0 SETUPDB=0 PGTAP=0 @@ -378,6 +403,7 @@ PGDUMP=0 MESSAGENOTICE=0 MESSAGELOG=0 CREATEBASICSQLOUT=0 +RUST=0 while [[ $# -gt 0 ]] do @@ -450,6 +476,11 @@ while [[ $# -gt 0 ]] shift ;; + --rust) + RUST=1 + shift + ;; + *) # unknown option usage exit 1; @@ -463,7 +494,7 @@ CLIENTMESSAGES='warning' [[ $MESSAGEDEBUG -eq 1 ]] && CLIENTMESSAGES='debug1' echo $CLIENTMESSAGES -if [[ ($FORMATTING -eq 0) && ($SETUPDB -eq 0) && ($PGTAP -eq 0) && ($BASICSQL -eq 0) && ($PYPGSTAC -eq 0) && ($MIGRATIONS -eq 0) && ($PGDUMP -eq 0) ]] +if [[ ($FORMATTING -eq 0) && ($SETUPDB -eq 0) && ($PGTAP -eq 0) && ($BASICSQL -eq 0) && ($PYPGSTAC -eq 0) && ($MIGRATIONS -eq 0) && ($PGDUMP -eq 0) && ($RUST -eq 0) ]] then FORMATTING=1 SETUPDB=1 @@ -473,6 +504,7 @@ then # run-all block while Python loader updates are pending on this branch. # Use --pypgstac or --migrations flags to run them explicitly. PGDUMP=1 + RUST=1 fi [ $FORMATTING -eq 1 ] && test_formatting @@ -486,5 +518,6 @@ fi [ $PYPGSTAC -eq 1 ] && test_pypgstac [ $MIGRATIONS -eq 1 ] && test_migrations [ $PGDUMP -eq 1 ] && test_pgdump +[ $RUST -eq 1 ] && test_rust exit 0 diff --git a/scripts/pgstac-rs-test-db b/scripts/pgstac-rs-test-db new file mode 100755 index 00000000..d630b949 --- /dev/null +++ b/scripts/pgstac-rs-test-db @@ -0,0 +1,44 @@ +#!/bin/bash +# Build the clean, empty pgstac install used as the clone template by the pgstac-rs +# (`src/pgstac-rs`) Rust tests. Each Rust test clones a fresh database from this template, so it +# must be a pristine install (zero collections/items). Run this once against a running dev database +# (see `scripts/server`); re-run after schema changes to refresh the template. +# +# Usage: scripts/pgstac-rs-test-db [TEMPLATE_NAME] +# TEMPLATE_NAME defaults to pgstac_rs_test_template (matches PGSTAC_RS_TEST_TEMPLATE in the tests). +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd ) +source "$SCRIPT_DIR/pgstacenv" # cd to repo root, ensure .env + +# Connection defaults (overridable via env / .env). An explicitly-set PGHOST (e.g. from the container +# test runner) wins over .env, so .env is only sourced when no host was provided. +[[ -z "${PGHOST:-}" && -f .env ]] && set -a && source .env && set +a +PGHOST=${PGHOST:-localhost} +PGPORT=${PGPORT:-5439} +PGUSER=${PGUSER:-username} +export PGPASSWORD=${PGPASSWORD:-password} + +TEMPLATE=${1:-${PGSTAC_RS_TEST_TEMPLATE:-pgstac_rs_test_template}} +PGSTAC_SQL="src/pgstac/pgstac.sql" + +if [[ ! -f "$PGSTAC_SQL" ]]; then + echo "ERROR: $PGSTAC_SQL not found (run from a repo checkout)." >&2 + exit 1 +fi + +echo "Building clean pgstac template '$TEMPLATE' on $PGHOST:$PGPORT ..." + +# (Re)create the template database from a maintenance connection (postgres/postgis), then load the +# assembled pgstac.sql into it. Terminating sessions first lets a re-run drop an in-use template. +psql -h "$PGHOST" -p "$PGPORT" -U "$PGUSER" -d postgres -v ON_ERROR_STOP=1 -q < /dev/null && pwd ) +source "$SCRIPT_DIR/pgstacenv" # cd to repo root, ensure .env + +# An explicitly-set PGHOST (e.g. from the container test runner) wins over .env. +[[ -z "${PGHOST:-}" && -f .env ]] && set -a && source .env && set +a +PGHOST=${PGHOST:-localhost} +PGPORT=${PGPORT:-5439} +PGUSER=${PGUSER:-username} +export PGPASSWORD=${PGPASSWORD:-password} + +TEMPLATE=${1:-${PGSTAC_RS_TEST_RICH_TEMPLATE:-pgstac_rs_test_rich}} +PGSTAC_SQL="src/pgstac/pgstac.sql" +FIXROOT="src/pgstac/tests/testdata/planetary-computer" +FIXTURES=(landsat-c2-l2 sentinel-2-l2a naip) + +for required in "$PGSTAC_SQL" "$FIXROOT/collections.ndjson"; do + [[ -f "$required" ]] || { echo "ERROR: $required not found (run from a repo checkout)." >&2; exit 1; } +done + +psql() { command psql -h "$PGHOST" -p "$PGPORT" -U "$PGUSER" "$@"; } + +echo "Building rich pgstac template '$TEMPLATE' on $PGHOST:$PGPORT ..." + +psql -d postgres -v ON_ERROR_STOP=1 -q <&2; exit 1; } + psql -d "$TEMPLATE" -v ON_ERROR_STOP=1 -q \ + -c "\\copy items_staging (content) FROM '$items' $JSON_COPY" +done + +# The bulk items_staging path does not populate partition_stats; analyze_items() fills the per-month +# histogram so the search() band engine and partition_bounds work. (analyze_items is a procedure that +# commits internally, so it must run on its own outside a transaction block.) +psql -d "$TEMPLATE" -v ON_ERROR_STOP=1 -q -c "CALL analyze_items();" + +psql -d "$TEMPLATE" -tAc "SET search_path=pgstac,public; + SELECT format( + 'rich template ready: %s collections (%s with fragment_config), %s items, %s fragments', + (SELECT count(*) FROM collections), + (SELECT count(*) FROM collections WHERE fragment_config IS NOT NULL), + (SELECT count(*) FROM items), + (SELECT count(*) FROM item_fragments));" diff --git a/scripts/test b/scripts/test index 2ddcef0a..66860caf 100755 --- a/scripts/test +++ b/scripts/test @@ -14,9 +14,21 @@ Options: --no-strict Allow stale-image shortcuts and watch fallback. -h, --help Show this help text. +Test suite flags (passed through to the in-container runner; no suite flag = all except +migrations/pypgstac): + --formatting ruff + ty checks. + --pgtap PGTap SQL tests. + --basicsql Basic SQL output comparison tests. + --rust pgstac-rs Rust tests (cargo test --features cli). + --pgdump pg_dump / pg_restore round-trip. + --pypgstac Python (pypgstac) tests. + --migrations Full migration-chain validation. + --nomigrations Everything except migrations + pg_dump. + Examples: $(basename "$0") $(basename "$0") --fast + $(basename "$0") --rust $(basename "$0") --pypgstac --build-policy always EOF } diff --git a/src/pgstac-rs/Cargo.lock b/src/pgstac-rs/Cargo.lock index 53ad26fb..0a6576c1 100644 --- a/src/pgstac-rs/Cargo.lock +++ b/src/pgstac-rs/Cargo.lock @@ -2,6 +2,12 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "adler2" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" + [[package]] name = "ahash" version = "0.8.12" @@ -9,6 +15,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5a15f179cd60c4584b8a8c596927aadc462e27f2ca70c04e0071964a73ba7a75" dependencies = [ "cfg-if", + "const-random", "getrandom 0.3.4", "once_cell", "serde", @@ -25,6 +32,21 @@ dependencies = [ "memchr", ] +[[package]] +name = "alloc-no-stdlib" +version = "2.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc7bb162ec39d46ab1ca8c77bf72e890535becd1751bb45f64c597edb4c8c6b3" + +[[package]] +name = "alloc-stdlib" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0e76a019e91224d279006ff972f1e984179a6e9feb050adba6ce8274aef23195" +dependencies = [ + "alloc-no-stdlib", +] + [[package]] name = "allocator-api2" version = "0.2.21" @@ -40,6 +62,56 @@ dependencies = [ "libc", ] +[[package]] +name = "anstream" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "824a212faf96e9acacdbd09febd34438f8f711fb84e09a8916013cd7815ca28d" +dependencies = [ + "anstyle", + "anstyle-parse", + "anstyle-query", + "anstyle-wincon", + "colorchoice", + "is_terminal_polyfill", + "utf8parse", +] + +[[package]] +name = "anstyle" +version = "1.0.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "940b3a0ca603d1eade50a4846a2afffd5ef57a9feac2c0e2ec2e14f9ead76000" + +[[package]] +name = "anstyle-parse" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52ce7f38b242319f7cabaa6813055467063ecdc9d355bbb4ce0c68908cd8130e" +dependencies = [ + "utf8parse", +] + +[[package]] +name = "anstyle-query" +version = "1.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "anstyle-wincon" +version = "3.0.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" +dependencies = [ + "anstyle", + "once_cell_polyfill", + "windows-sys 0.61.2", +] + [[package]] name = "anyhow" version = "1.0.102" @@ -64,6 +136,157 @@ dependencies = [ "object", ] +[[package]] +name = "arrow-arith" +version = "58.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a0ab212d2c1886e802f51c5212d78ebbcbb0bec980fff9dadc1eb8d45cd0b738" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "chrono", + "num-traits", +] + +[[package]] +name = "arrow-array" +version = "58.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cfd33d3e92f207444098c75b42de99d329562be0cf686b307b097cc52b4e999e" +dependencies = [ + "ahash", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "chrono", + "chrono-tz", + "half", + "hashbrown 0.17.1", + "num-complex", + "num-integer", + "num-traits", +] + +[[package]] +name = "arrow-buffer" +version = "58.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c6cd424c2693bcdbc150d843dc9d4d137dd2de4782ce6df491ad11a3a0416c0" +dependencies = [ + "bytes", + "half", + "num-bigint", + "num-traits", +] + +[[package]] +name = "arrow-cast" +version = "58.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4c5aefb56a2c02e9e2b30746241058b85f8983f0fcff2ba0c6d09006e1cded7f" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-ord", + "arrow-schema", + "arrow-select", + "atoi", + "base64", + "chrono", + "half", + "lexical-core", + "num-traits", + "ryu", +] + +[[package]] +name = "arrow-data" +version = "58.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3c88210023a2bfee1896af366309a3028fc3bcbd6515fa29a7990ee1baa08ee0" +dependencies = [ + "arrow-buffer", + "arrow-schema", + "half", + "num-integer", + "num-traits", +] + +[[package]] +name = "arrow-ipc" +version = "58.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "238438f0834483703d88896db6fe5a7138b2230debc31b34c0336c2996e3c64f" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "arrow-select", + "flatbuffers", +] + +[[package]] +name = "arrow-json" +version = "58.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "205ca2119e6d679d5c133c6f30e68f027738d95ed948cf77677ea69c7800036b" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-cast", + "arrow-ord", + "arrow-schema", + "arrow-select", + "chrono", + "half", + "indexmap 2.14.0", + "itoa", + "lexical-core", + "memchr", + "num-traits", + "ryu", + "serde_core", + "serde_json", + "simdutf8", +] + +[[package]] +name = "arrow-ord" +version = "58.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1bffd8fd2579286a5d63bac898159873e5094a79009940bcb42bbfce4f19f1d0" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "arrow-select", +] + +[[package]] +name = "arrow-schema" +version = "58.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f633dbfdf39c039ada1bf9e34c694816eb71fbb7dc78f613993b7245e078a1ed" + +[[package]] +name = "arrow-select" +version = "58.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8cd065c54172ac787cf3f2f8d4107e0d3fdc26edba76fdf4f4cc170258942222" +dependencies = [ + "ahash", + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "num-traits", +] + [[package]] name = "as-slice" version = "0.1.5" @@ -76,6 +299,28 @@ dependencies = [ "stable_deref_trait", ] +[[package]] +name = "async-stream" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b5a71a6f37880a80d1d7f19efd781e4b5de42c88f0722cc13bcb6cc2cfe8476" +dependencies = [ + "async-stream-impl", + "futures-core", + "pin-project-lite", +] + +[[package]] +name = "async-stream-impl" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7c24de15d275a1ecfd47a380fb4d5ec9bfe0933f309ed5e705b775596a3574d" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "async-trait" version = "0.1.89" @@ -87,6 +332,15 @@ dependencies = [ "syn", ] +[[package]] +name = "atoi" +version = "2.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f28d99ec8bfea296261ca1af174f24225171fea9664ba9003cbebee704810528" +dependencies = [ + "num-traits", +] + [[package]] name = "atomic-polyfill" version = "1.0.3" @@ -96,18 +350,52 @@ dependencies = [ "critical-section", ] +[[package]] +name = "atomic-waker" +version = "1.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0" + [[package]] name = "autocfg" version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8" +[[package]] +name = "aws-lc-rs" +version = "1.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5ec2f1fc3ec205783a5da9a7e6c1509cc69dedf09a1949e412c1e18469326d00" +dependencies = [ + "aws-lc-sys", + "zeroize", +] + +[[package]] +name = "aws-lc-sys" +version = "0.41.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1a2f9779ce85b93ab6170dd940ad0169b5766ff848247aff13bb788b832fe3f4" +dependencies = [ + "cc", + "cmake", + "dunce", + "fs_extra", +] + [[package]] name = "base64" version = "0.22.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" +[[package]] +name = "base64ct" +version = "1.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2af50177e190e07a26ab74f8b1efbfe2ef87da2116221318cb1c2e82baf7de06" + [[package]] name = "bit-set" version = "0.8.0" @@ -153,6 +441,36 @@ version = "0.2.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dc0b364ead1874514c8c2855ab558056ebfeb775653e7ae45ff72f28f8f3166c" +[[package]] +name = "brotli" +version = "8.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5cc91aac060a7a1e25823bdccbfb6af1875b88f17c6daac97894eed8207166b3" +dependencies = [ + "alloc-no-stdlib", + "alloc-stdlib", + "brotli-decompressor", +] + +[[package]] +name = "brotli-decompressor" +version = "5.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3a32acac15fe1967bc3986b2a6347dffc965602354ea6f450ad07e8bfd253583" +dependencies = [ + "alloc-no-stdlib", + "alloc-stdlib", +] + +[[package]] +name = "bs58" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bf88ba1141d185c399bee5288d850d63b8369520c1eafc32a0430b5b6c287bf4" +dependencies = [ + "tinyvec", +] + [[package]] name = "bumpalo" version = "3.20.2" @@ -184,6 +502,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a1dce859f0832a7d088c4f1119888ab94ef4b5d6795d1ce05afb7fe159d79f98" dependencies = [ "find-msvc-tools", + "jobserver", + "libc", "shlex", ] @@ -193,6 +513,12 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" +[[package]] +name = "cfg_aliases" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" + [[package]] name = "chacha20" version = "0.10.0" @@ -218,18 +544,139 @@ dependencies = [ "windows-link", ] +[[package]] +name = "chrono-tz" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6139a8597ed92cf816dfb33f5dd6cf0bb93a6adc938f11039f371bc5bcd26c3" +dependencies = [ + "chrono", + "phf 0.12.1", +] + +[[package]] +name = "clap" +version = "4.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ddb117e43bbf7dacf0a4190fef4d345b9bad68dfc649cb349e7d17d28428e51" +dependencies = [ + "clap_builder", + "clap_derive", +] + +[[package]] +name = "clap_builder" +version = "4.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "714a53001bf66416adb0e2ef5ac857140e7dc3a0c48fb28b2f10762fc4b5069f" +dependencies = [ + "anstream", + "anstyle", + "clap_lex", + "strsim", +] + +[[package]] +name = "clap_derive" +version = "4.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2ce8604710f6733aa641a2b3731eaa1e8b3d9973d5e3565da11800813f997a9" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "clap_lex" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8d4a3bb8b1e0c1050499d1815f5ab16d04f0959b233085fb31653fbfc9d98f9" + +[[package]] +name = "cmake" +version = "0.1.58" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c0f78a02292a74a88ac736019ab962ece0bc380e3f977bf72e376c5d78ff0678" +dependencies = [ + "cc", +] + [[package]] name = "cmov" version = "0.5.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3f88a43d011fc4a6876cb7344703e297c71dda42494fee094d5f7c76bf13f746" +[[package]] +name = "colorchoice" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d07550c9036bf2ae0c684c4297d503f838287c83c53686d05370d0e139ae570" + +[[package]] +name = "combine" +version = "4.6.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba5a308b75df32fe02788e748662718f03fde005016435c444eea572398219fd" +dependencies = [ + "bytes", + "memchr", +] + +[[package]] +name = "const-oid" +version = "0.9.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2459377285ad874054d797f3ccebf984978aa39129f6eafde5cdc8315b612f8" + [[package]] name = "const-oid" version = "0.10.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a6ef517f0926dd24a1582492c791b6a4818a4d94e789a334894aa15b0d12f55c" +[[package]] +name = "const-random" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87e00182fe74b066627d63b85fd550ac2998d4b0bd86bfed477a0ae4c7c71359" +dependencies = [ + "const-random-macro", +] + +[[package]] +name = "const-random-macro" +version = "0.1.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9d839f2a20b0aee515dc581a6172f2321f96cab76c1a38a4c584a194955390e" +dependencies = [ + "getrandom 0.2.17", + "once_cell", + "tiny-keccak", +] + +[[package]] +name = "core-foundation" +version = "0.9.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91e195e091a93c46f7102ec7818a2aa394e1e1771c3ab4825963fa03e45afb8f" +dependencies = [ + "core-foundation-sys", + "libc", +] + +[[package]] +name = "core-foundation" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b2a6cd9ae233e7f62ba4e9353e81a88df7fc8a5987b8d445b4d90c879bd156f6" +dependencies = [ + "core-foundation-sys", + "libc", +] + [[package]] name = "core-foundation-sys" version = "0.8.7" @@ -311,6 +758,12 @@ version = "0.8.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28" +[[package]] +name = "crunchy" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" + [[package]] name = "crypto-common" version = "0.1.6" @@ -340,8 +793,109 @@ dependencies = [ ] [[package]] -name = "digest" -version = "0.10.7" +name = "darling" +version = "0.23.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "25ae13da2f202d56bd7f91c25fba009e7717a1e4a1cc98a76d844b65ae912e9d" +dependencies = [ + "darling_core", + "darling_macro", +] + +[[package]] +name = "darling_core" +version = "0.23.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9865a50f7c335f53564bb694ef660825eb8610e0a53d3e11bf1b0d3df31e03b0" +dependencies = [ + "ident_case", + "proc-macro2", + "quote", + "strsim", + "syn", +] + +[[package]] +name = "darling_macro" +version = "0.23.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac3984ec7bd6cfa798e62b4a642426a5be0e68f9401cfc2a01e3fa9ea2fcdb8d" +dependencies = [ + "darling_core", + "quote", + "syn", +] + +[[package]] +name = "deadpool" +version = "0.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0be2b1d1d6ec8d846f05e137292d0b89133caf95ef33695424c09568bdd39b1b" +dependencies = [ + "deadpool-runtime", + "lazy_static", + "num_cpus", + "tokio", +] + +[[package]] +name = "deadpool-postgres" +version = "0.14.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3d697d376cbfa018c23eb4caab1fd1883dd9c906a8c034e8d9a3cb06a7e0bef9" +dependencies = [ + "async-trait", + "deadpool", + "getrandom 0.2.17", + "tokio", + "tokio-postgres", + "tracing", +] + +[[package]] +name = "deadpool-runtime" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "092966b41edc516079bdf31ec78a2e0588d1d0c08f78b91d8307215928642b2b" +dependencies = [ + "tokio", +] + +[[package]] +name = "der" +version = "0.7.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7c1832837b905bbfb5101e07cc24c8deddf52f93225eee6ead5f4d63d53ddcb" +dependencies = [ + "const-oid 0.9.6", + "der_derive", + "flagset", + "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", +] + +[[package]] +name = "deranged" +version = "0.5.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7cd812cc2bc1d69d4764bd80df88b4317eaef9e773c75226407d9bc0876b211c" +dependencies = [ + "serde_core", +] + +[[package]] +name = "digest" +version = "0.10.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" dependencies = [ @@ -356,7 +910,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f1dd6dbb5841937940781866fa1281a1ff7bd3bf827091440879f9994983d5c2" dependencies = [ "block-buffer 0.12.0", - "const-oid", + "const-oid 0.10.2", "crypto-common 0.2.1", "ctutils", ] @@ -372,6 +926,18 @@ dependencies = [ "syn", ] +[[package]] +name = "dunce" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92773504d58c093f6de2459af4af33faa518c13451eb8f2b5698ed3d36e7c813" + +[[package]] +name = "dyn-clone" +version = "1.0.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d0881ea181b1df73ff77ffaaf9c7544ecc11e82fba9b5f27b262a3c73a332555" + [[package]] name = "earcut" version = "0.4.5" @@ -396,12 +962,31 @@ dependencies = [ "serde", ] +[[package]] +name = "encoding_rs" +version = "0.8.35" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "75030f3c4f45dafd7586dd6780965a8c7e8e285a5ecb86713e63a79c5b2766f3" +dependencies = [ + "cfg-if", +] + [[package]] name = "equivalent" version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" +[[package]] +name = "errno" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + [[package]] name = "fallible-iterator" version = "0.2.0" @@ -419,12 +1004,54 @@ dependencies = [ "regex-syntax", ] +[[package]] +name = "fastrand" +version = "2.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f1f227452a390804cdb637b74a86990f2a7d7ba4b7d5693aac9b4dd6defd8d6" + +[[package]] +name = "filetime" +version = "0.2.29" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c287a33c7f0a620c38e641e7f60827713987b3c0f26e8ddc9462cc69cf75759" +dependencies = [ + "cfg-if", + "libc", +] + [[package]] name = "find-msvc-tools" version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" +[[package]] +name = "flagset" +version = "0.4.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b7ac824320a75a52197e8f2d787f6a38b6718bb6897a35142d749af3c0e8f4fe" + +[[package]] +name = "flatbuffers" +version = "25.12.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "35f6839d7b3b98adde531effaf34f0c2badc6f4735d26fe74709d8e513a96ef3" +dependencies = [ + "bitflags", + "rustc_version", +] + +[[package]] +name = "flate2" +version = "1.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "843fba2746e448b37e26a819579957415c8cef339bf08564fe8b7ddbd959573c" +dependencies = [ + "miniz_oxide", + "zlib-rs", +] + [[package]] name = "float_next_after" version = "2.0.0" @@ -442,6 +1069,12 @@ dependencies = [ "serde", ] +[[package]] +name = "fnv" +version = "1.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" + [[package]] name = "foldhash" version = "0.1.5" @@ -473,6 +1106,27 @@ dependencies = [ "num", ] +[[package]] +name = "fs_extra" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c" + +[[package]] +name = "futures" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b147ee9d1f6d097cef9ce628cd2ee62288d963e16fb287bd9286455b241382d" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + [[package]] name = "futures-channel" version = "0.3.32" @@ -489,6 +1143,23 @@ version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7e3450815272ef58cec6d564423f6e755e25379b217b0bc688e295ba24df6b1d" +[[package]] +name = "futures-executor" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "baf29c38818342a3b26b5b923639e7b1f4a61fc5e76102d4b1981c6dc7a7579d" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-io" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cecba35d7ad927e23624b22ad55235f2239cfa44fd10428eecbeba6d6a717718" + [[package]] name = "futures-macro" version = "0.3.32" @@ -524,10 +1195,13 @@ version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "389ca41296e6190b48053de0321d02a77f32f8a5d2461dd38762c0593805c6d6" dependencies = [ + "futures-channel", "futures-core", + "futures-io", "futures-macro", "futures-sink", "futures-task", + "memchr", "pin-project-lite", "slab", ] @@ -608,6 +1282,35 @@ dependencies = [ "spade", ] +[[package]] +name = "geoarrow-array" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dafe7b7de3fab1a8b7099fd6a6434ca955fa65065f9c19f0f8a133693f3c2b0e" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-schema", + "geo-traits", + "geoarrow-schema", + "num-traits", + "wkb", + "wkt 0.14.0", +] + +[[package]] +name = "geoarrow-schema" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4d4a7edb2a1d87024a93805332a9c8184a0354836271d42c0d18cf628a5e3cd0" +dependencies = [ + "arrow-schema", + "geo-traits", + "serde", + "serde_json", + "thiserror 1.0.69", +] + [[package]] name = "geographiclib-rs" version = "0.2.7" @@ -644,6 +1347,29 @@ dependencies = [ "tinyvec", ] +[[package]] +name = "geoparquet" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a64b758b2b1fc749c5eb212215afa6bc14e1fca93884dac339ab7097ffc20ce1" +dependencies = [ + "arrow-arith", + "arrow-array", + "arrow-buffer", + "arrow-ord", + "arrow-schema", + "geo-traits", + "geo-types", + "geoarrow-array", + "geoarrow-schema", + "indexmap 2.14.0", + "parquet", + "serde", + "serde_json", + "serde_with", + "wkt 0.14.0", +] + [[package]] name = "geozero" version = "0.14.0" @@ -653,11 +1379,25 @@ dependencies = [ "geo-types", "geojson 0.24.2", "log", + "scroll", "serde_json", "thiserror 1.0.69", "wkt 0.11.1", ] +[[package]] +name = "getrandom" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff2abc00be7fca6ebc474524697ae276ad847ad0a6b3faa4bcb027e9a4614ad0" +dependencies = [ + "cfg-if", + "js-sys", + "libc", + "wasi 0.11.1+wasi-snapshot-preview1", + "wasm-bindgen", +] + [[package]] name = "getrandom" version = "0.3.4" @@ -665,9 +1405,11 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" dependencies = [ "cfg-if", + "js-sys", "libc", "r-efi 5.3.0", "wasip2", + "wasm-bindgen", ] [[package]] @@ -690,6 +1432,37 @@ version = "0.3.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0cc23270f6e1808e30a928bdc84dea0b9b4136a8bc82338574f23baf47bbd280" +[[package]] +name = "h2" +version = "0.4.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6cb093c84e8bd9b188d4c4a8cb6579fc016968d14c99882163cd3ff402a4f155" +dependencies = [ + "atomic-waker", + "bytes", + "fnv", + "futures-core", + "futures-sink", + "http", + "indexmap 2.14.0", + "slab", + "tokio", + "tokio-util", + "tracing", +] + +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "cfg-if", + "crunchy", + "num-traits", + "zerocopy", +] + [[package]] name = "hash32" version = "0.1.1" @@ -717,6 +1490,12 @@ dependencies = [ "byteorder", ] +[[package]] +name = "hashbrown" +version = "0.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8a9ee70c43aaf417c914396645a0fa852624801b24ebb7ae78fe8272889ac888" + [[package]] name = "hashbrown" version = "0.15.5" @@ -784,6 +1563,18 @@ version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" +[[package]] +name = "hermit-abi" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc0fef456e4baa96da950455cd02c081ca953b141298e41db3fc7e36b1da849c" + +[[package]] +name = "hex" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70" + [[package]] name = "hmac" version = "0.13.0" @@ -794,57 +1585,164 @@ dependencies = [ ] [[package]] -name = "hybrid-array" -version = "0.4.12" +name = "http" +version = "1.4.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9155a582abd142abc056962c29e3ce5ff2ad5469f4246b537ed42c5deba857da" +checksum = "6970f50e31d6fc17d3fa27329444bfa74e196cf62e95052a3f6fee181dba6425" dependencies = [ - "typenum", + "bytes", + "itoa", ] [[package]] -name = "i_float" -version = "1.16.0" +name = "http-body" +version = "1.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "813145bb0ad5b60f55cbbf3c74cdceda1c0a9d253b35c4cc36ae0df7887cb78f" +checksum = "1efedce1fb8e6913f23e0c92de8e62cd5b772a67e7b3946df930a62566c93184" dependencies = [ - "libm", + "bytes", + "http", ] [[package]] -name = "i_key_sort" -version = "0.10.3" +name = "http-body-util" +version = "0.1.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d73d122b937fca067feb0ad74f62388920272b27c356d4df2d0cfdd59e044cf0" +checksum = "b021d93e26becf5dc7e1b75b1bed1fd93124b374ceb73f43d4d4eafec896a64a" dependencies = [ - "rayon", + "bytes", + "futures-core", + "http", + "http-body", + "pin-project-lite", ] [[package]] -name = "i_overlay" -version = "4.5.1" +name = "httparse" +version = "1.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "934cba666ad1bf60436a190af6eaa3f58f57b5bad28268ce3fb45d9185862232" -dependencies = [ - "i_float", - "i_key_sort", - "i_shape", - "i_tree", - "rayon", -] +checksum = "6dbf3de79e51f3d586ab4cb9d5c3e2c14aa28ed23d180cf89b4df0454a69cc87" [[package]] -name = "i_shape" -version = "1.18.0" +name = "humantime" +version = "2.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bfa9eac533d7509a8ab87672b60ac610c17240f9ea4851d26227689fdfe349c8" -dependencies = [ - "i_float", -] +checksum = "135b12329e5e3ce057a9f972339ea52bc954fe1e9358ef27f95e89716fbc5424" [[package]] -name = "i_tree" -version = "0.18.0" +name = "hybrid-array" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9155a582abd142abc056962c29e3ce5ff2ad5469f4246b537ed42c5deba857da" +dependencies = [ + "typenum", +] + +[[package]] +name = "hyper" +version = "1.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "55281c53a1894c864990125767da440a4e630446785086f52523b20033b74498" +dependencies = [ + "atomic-waker", + "bytes", + "futures-channel", + "futures-core", + "h2", + "http", + "http-body", + "httparse", + "itoa", + "pin-project-lite", + "smallvec", + "tokio", + "want", +] + +[[package]] +name = "hyper-rustls" +version = "0.27.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "33ca68d021ef39cf6463ab54c1d0f5daf03377b70561305bb89a8f83aab66e0f" +dependencies = [ + "http", + "hyper", + "hyper-util", + "rustls", + "rustls-native-certs", + "tokio", + "tokio-rustls", + "tower-service", +] + +[[package]] +name = "hyper-util" +version = "0.1.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96547c2556ec9d12fb1578c4eaf448b04993e7fb79cbaad930a656880a6bdfa0" +dependencies = [ + "base64", + "bytes", + "futures-channel", + "futures-util", + "http", + "http-body", + "hyper", + "ipnet", + "libc", + "percent-encoding", + "pin-project-lite", + "socket2", + "system-configuration", + "tokio", + "tower-service", + "tracing", + "windows-registry", +] + +[[package]] +name = "i_float" +version = "1.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "813145bb0ad5b60f55cbbf3c74cdceda1c0a9d253b35c4cc36ae0df7887cb78f" +dependencies = [ + "libm", +] + +[[package]] +name = "i_key_sort" +version = "0.10.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d73d122b937fca067feb0ad74f62388920272b27c356d4df2d0cfdd59e044cf0" +dependencies = [ + "rayon", +] + +[[package]] +name = "i_overlay" +version = "4.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "934cba666ad1bf60436a190af6eaa3f58f57b5bad28268ce3fb45d9185862232" +dependencies = [ + "i_float", + "i_key_sort", + "i_shape", + "i_tree", + "rayon", +] + +[[package]] +name = "i_shape" +version = "1.18.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bfa9eac533d7509a8ab87672b60ac610c17240f9ea4851d26227689fdfe349c8" +dependencies = [ + "i_float", +] + +[[package]] +name = "i_tree" +version = "0.18.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4804bdc1dc124eb7e1aa9e144ecc04096bcf787a10a15fa44af682b51f0f6cce" @@ -960,6 +1858,12 @@ version = "2.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3d3067d79b975e8844ca9eb072e16b31c3c1c36928edf9c6789548c524d0d954" +[[package]] +name = "ident_case" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9e0384b61958566e926dc50660321d12159025e767c18e043daf26b70104c39" + [[package]] name = "idna" version = "1.1.0" @@ -981,6 +1885,17 @@ dependencies = [ "icu_properties", ] +[[package]] +name = "indexmap" +version = "1.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bd070e393353796e801d209ad339e89596eb4c8d430d18ede6a1cced8fafbd99" +dependencies = [ + "autocfg", + "hashbrown 0.12.3", + "serde", +] + [[package]] name = "indexmap" version = "2.14.0" @@ -993,6 +1908,51 @@ dependencies = [ "serde_core", ] +[[package]] +name = "indoc" +version = "2.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "79cf5c93f93228cf8efb3ba362535fb11199ac548a09ce117c9b1adc3030d706" +dependencies = [ + "rustversion", +] + +[[package]] +name = "integer-encoding" +version = "3.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8bb03732005da905c88227371639bf1ad885cc712789c011c31c5fb3ab3ccf02" + +[[package]] +name = "ipnet" +version = "2.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d98f6fed1fde3f8c21bc40a1abb88dd75e67924f9cffc3ef95607bad8017f8e2" + +[[package]] +name = "is_terminal_polyfill" +version = "1.70.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6cb138bb79a146c1bd460005623e142ef0181e3d0219cb493e02f7d08a35695" + +[[package]] +name = "itertools" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "413ee7dfc52ee1a4949ceeb7dbc8a33f2d6c088194d9f922fb8318faf1f01186" +dependencies = [ + "either", +] + +[[package]] +name = "itertools" +version = "0.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b192c782037fadd9cfa75548310488aabdbf3d2da73885b31bd0abd03351285" +dependencies = [ + "either", +] + [[package]] name = "itoa" version = "1.0.18" @@ -1011,7 +1971,7 @@ dependencies = [ "portable-atomic", "portable-atomic-util", "serde_core", - "windows-sys", + "windows-sys 0.61.2", ] [[package]] @@ -1040,6 +2000,65 @@ dependencies = [ "jiff-tzdb", ] +[[package]] +name = "jni" +version = "0.22.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5efd9a482cf3a427f00d6b35f14332adc7902ce91efb778580e180ff90fa3498" +dependencies = [ + "cfg-if", + "combine", + "jni-macros", + "jni-sys", + "log", + "simd_cesu8", + "thiserror 2.0.18", + "walkdir", + "windows-link", +] + +[[package]] +name = "jni-macros" +version = "0.22.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a00109accc170f0bdb141fed3e393c565b6f5e072365c3bd58f5b062591560a3" +dependencies = [ + "proc-macro2", + "quote", + "rustc_version", + "simd_cesu8", + "syn", +] + +[[package]] +name = "jni-sys" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c6377a88cb3910bee9b0fa88d4f42e1d2da8e79915598f65fb0c7ee14c878af2" +dependencies = [ + "jni-sys-macros", +] + +[[package]] +name = "jni-sys-macros" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "38c0b942f458fe50cdac086d2f946512305e5631e720728f2a61aabcd47a6264" +dependencies = [ + "quote", + "syn", +] + +[[package]] +name = "jobserver" +version = "0.1.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9afb3de4395d6b3e67a780b6de64b51c978ecf11cb9a462c66be7d4ca9039d33" +dependencies = [ + "getrandom 0.3.4", + "libc", +] + [[package]] name = "js-sys" version = "0.3.98" @@ -1102,6 +2121,63 @@ version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "09edd9e8b54e49e587e4f6295a7d29c3ea94d469cb40ab8ca70b288248a81db2" +[[package]] +name = "lexical-core" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7d8d125a277f807e55a77304455eb7b1cb52f2b18c143b60e766c120bd64a594" +dependencies = [ + "lexical-parse-float", + "lexical-parse-integer", + "lexical-util", + "lexical-write-float", + "lexical-write-integer", +] + +[[package]] +name = "lexical-parse-float" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52a9f232fbd6f550bc0137dcb5f99ab674071ac2d690ac69704593cb4abbea56" +dependencies = [ + "lexical-parse-integer", + "lexical-util", +] + +[[package]] +name = "lexical-parse-integer" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a7a039f8fb9c19c996cd7b2fcce303c1b2874fe1aca544edc85c4a5f8489b34" +dependencies = [ + "lexical-util", +] + +[[package]] +name = "lexical-util" +version = "1.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2604dd126bb14f13fb5d1bd6a66155079cb9fa655b37f875b3a742c705dbed17" + +[[package]] +name = "lexical-write-float" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "50c438c87c013188d415fbabbb1dceb44249ab81664efbd31b14ae55dabb6361" +dependencies = [ + "lexical-util", + "lexical-write-integer", +] + +[[package]] +name = "lexical-write-integer" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "409851a618475d2d5796377cad353802345cba92c867d9fbcde9cf4eac4e14df" +dependencies = [ + "lexical-util", +] + [[package]] name = "libc" version = "0.2.186" @@ -1129,6 +2205,12 @@ version = "0.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fc7281e4b2b1a1fae03463a7c49dd21464de50251a450f6da9715c40c7b21a70" +[[package]] +name = "linux-raw-sys" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53" + [[package]] name = "litemap" version = "0.8.2" @@ -1150,6 +2232,31 @@ version = "0.4.29" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897" +[[package]] +name = "lru-slab" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "112b39cec0b298b6c1999fee3e31427f74f676e4cb9879ed1a121b43661a4154" + +[[package]] +name = "lz4_flex" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7ef0d4ed8669f8f8826eb00dc878084aa8f253506c4fd5e8f58f5bce72ddb97e" +dependencies = [ + "twox-hash", +] + +[[package]] +name = "md-5" +version = "0.10.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d89e7ee0cfbedfc4da3340218492196241d89eefb6dab27de5df917a6d2e78cf" +dependencies = [ + "cfg-if", + "digest 0.10.7", +] + [[package]] name = "md-5" version = "0.11.0" @@ -1166,12 +2273,31 @@ version = "2.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f8ca58f447f06ed17d5fc4043ce1b10dd205e060fb3ce5b979b8ed8e59ff3f79" +[[package]] +name = "memoffset" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "488016bfae457b036d996092f6cb448677611ce4449e970ceaf42695203f218a" +dependencies = [ + "autocfg", +] + [[package]] name = "mime" version = "0.3.17" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6877bb514081ee2a7ff5ef9de3281f14a4dd4bceac4c09388074a6b5df8a139a" +[[package]] +name = "miniz_oxide" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fa76a2c86f704bdb222d66965fb3d63269ce38518b83cb0575fca855ebb6316" +dependencies = [ + "adler2", + "simd-adler32", +] + [[package]] name = "mio" version = "1.2.0" @@ -1180,7 +2306,7 @@ checksum = "50b7e5b27aa02a74bac8c3f23f448f8d87ff11f92d3aac1a6ed369ee08cc56c1" dependencies = [ "libc", "wasi 0.11.1+wasi-snapshot-preview1", - "windows-sys", + "windows-sys 0.61.2", ] [[package]] @@ -1222,6 +2348,12 @@ dependencies = [ "num-traits", ] +[[package]] +name = "num-conv" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "521739c6d2bac4aa25192232afe6841231376b2b26d4d9fae5ecf8ca5772e441" + [[package]] name = "num-integer" version = "0.1.46" @@ -1263,6 +2395,38 @@ dependencies = [ "libm", ] +[[package]] +name = "num_cpus" +version = "1.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91df4bbde75afed763b708b7eee1e8e7651e02d97f6d5dd763e89367e957b23b" +dependencies = [ + "hermit-abi", + "libc", +] + +[[package]] +name = "num_enum" +version = "0.7.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5d0bca838442ec211fa11de3a8b0e0e8f3a4522575b5c4c06ed722e005036f26" +dependencies = [ + "num_enum_derive", + "rustversion", +] + +[[package]] +name = "num_enum_derive" +version = "0.7.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "680998035259dcfcafe653688bf2aa6d3e2dc05e98be6ab46afb089dc84f1df8" +dependencies = [ + "proc-macro-crate", + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "objc2-core-foundation" version = "0.3.2" @@ -1291,17 +2455,93 @@ dependencies = [ ] [[package]] -name = "once_cell" -version = "1.21.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" - -[[package]] -name = "outref" -version = "0.5.2" +name = "object_store" +version = "0.11.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1a80800c0488c3a21695ea981a54918fbb37abf04f4d0720c453632255e2ff0e" - +checksum = "3cfccb68961a56facde1163f9319e0d15743352344e7808a11795fb99698dcaf" +dependencies = [ + "async-trait", + "base64", + "bytes", + "chrono", + "futures", + "humantime", + "hyper", + "itertools 0.13.0", + "md-5 0.10.6", + "parking_lot", + "percent-encoding", + "quick-xml", + "rand 0.8.6", + "reqwest 0.12.28", + "ring", + "serde", + "serde_json", + "snafu", + "tokio", + "tracing", + "url", + "walkdir", +] + +[[package]] +name = "object_store" +version = "0.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "622acbc9100d3c10e2ee15804b0caa40e55c933d5aa53814cd520805b7958a49" +dependencies = [ + "async-trait", + "bytes", + "chrono", + "futures-channel", + "futures-core", + "futures-util", + "http", + "humantime", + "itertools 0.14.0", + "parking_lot", + "percent-encoding", + "thiserror 2.0.18", + "tokio", + "tracing", + "url", + "wasm-bindgen-futures", + "web-time", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "once_cell_polyfill" +version = "1.70.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe" + +[[package]] +name = "openssl-probe" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7c87def4c32ab89d880effc9e097653c8da5d6ef28e6b539d313baaacfbafcbe" + +[[package]] +name = "ordered-float" +version = "2.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68f19d67e5a2795c94e73e0bb1cc1a7edeb2e28efd39e2e1c9b7a40c1108b11c" +dependencies = [ + "num-traits", +] + +[[package]] +name = "outref" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1a80800c0488c3a21695ea981a54918fbb37abf04f4d0720c453632255e2ff0e" + [[package]] name = "parking_lot" version = "0.12.5" @@ -1325,6 +2565,48 @@ dependencies = [ "windows-link", ] +[[package]] +name = "parquet" +version = "58.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5dafa7d01085b62a47dd0c1829550a0a36710ea9c4fe358a05a85477cec8a908" +dependencies = [ + "ahash", + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-ipc", + "arrow-schema", + "arrow-select", + "base64", + "brotli", + "bytes", + "chrono", + "flate2", + "futures", + "half", + "hashbrown 0.17.1", + "lz4_flex", + "num-bigint", + "num-integer", + "num-traits", + "object_store 0.13.2", + "paste", + "seq-macro", + "simdutf8", + "snap", + "thrift", + "tokio", + "twox-hash", + "zstd", +] + +[[package]] +name = "paste" +version = "1.0.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" + [[package]] name = "pdqselect" version = "0.1.0" @@ -1393,16 +2675,41 @@ dependencies = [ name = "pgstac" version = "0.9.11-dev" dependencies = [ + "arrow-array", + "async-stream", + "async-trait", + "base64", + "bytes", + "chrono", + "clap", + "deadpool-postgres", + "futures", "geojson 1.0.0", + "geoparquet", + "geozero", + "object_store 0.11.2", + "parquet", + "pyo3", + "pyo3-async-runtimes", "rstest", + "rustls", + "rustls-pemfile", "serde", "serde_json", + "sha2 0.10.9", "stac", + "stac-io", + "tar", + "tempfile", "thiserror 2.0.18", "tokio", "tokio-postgres", + "tokio-postgres-rustls", "tokio-test", "tracing", + "url", + "webpki-roots 0.26.11", + "zstd", ] [[package]] @@ -1415,6 +2722,15 @@ dependencies = [ "phf_shared 0.11.3", ] +[[package]] +name = "phf" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "913273894cec178f401a31ec4b656318d95473527be05c0752cc41cdc32be8b7" +dependencies = [ + "phf_shared 0.12.1", +] + [[package]] name = "phf" version = "0.13.1" @@ -1457,6 +2773,15 @@ dependencies = [ "siphasher", ] +[[package]] +name = "phf_shared" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06005508882fb681fd97892ecff4b7fd0fee13ef1aa569f8695dae7ab9099981" +dependencies = [ + "siphasher", +] + [[package]] name = "phf_shared" version = "0.13.1" @@ -1472,6 +2797,12 @@ version = "0.2.17" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" +[[package]] +name = "pkg-config" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "19f132c84eca552bf34cab8ec81f1c1dcc229b811638f9d283dceabe58c5569e" + [[package]] name = "portable-atomic" version = "1.13.1" @@ -1498,7 +2829,7 @@ dependencies = [ "bytes", "fallible-iterator", "hmac", - "md-5", + "md-5 0.11.0", "memchr", "rand 0.10.1", "sha2 0.11.0", @@ -1512,6 +2843,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8dc729a129e682e8d24170cd30ae1aa01b336b096cbb56df6d534ffec133d186" dependencies = [ "bytes", + "chrono", "fallible-iterator", "postgres-protocol", "serde_core", @@ -1527,6 +2859,21 @@ dependencies = [ "zerovec", ] +[[package]] +name = "powerfmt" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "439ee305def115ba05938db6eb1644ff94165c5ab5e9420d1c1bcedbba909391" + +[[package]] +name = "ppv-lite86" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85eae3c4ed2f50dcfe72643da4befc30deadb458a9b590d720cde2f2b1e97da9" +dependencies = [ + "zerocopy", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -1565,6 +2912,148 @@ dependencies = [ "cc", ] +[[package]] +name = "pyo3" +version = "0.23.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7778bffd85cf38175ac1f545509665d0b9b92a198ca7941f131f85f7a4f9a872" +dependencies = [ + "cfg-if", + "indoc", + "libc", + "memoffset", + "once_cell", + "portable-atomic", + "pyo3-build-config", + "pyo3-ffi", + "pyo3-macros", + "unindent", +] + +[[package]] +name = "pyo3-async-runtimes" +version = "0.23.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "977dc837525cfd22919ba6a831413854beb7c99a256c03bf8624ad707e45810e" +dependencies = [ + "futures", + "once_cell", + "pin-project-lite", + "pyo3", + "tokio", +] + +[[package]] +name = "pyo3-build-config" +version = "0.23.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94f6cbe86ef3bf18998d9df6e0f3fc1050a8c5efa409bf712e661a4366e010fb" +dependencies = [ + "once_cell", + "target-lexicon", +] + +[[package]] +name = "pyo3-ffi" +version = "0.23.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e9f1b4c431c0bb1c8fb0a338709859eed0d030ff6daa34368d3b152a63dfdd8d" +dependencies = [ + "libc", + "pyo3-build-config", +] + +[[package]] +name = "pyo3-macros" +version = "0.23.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fbc2201328f63c4710f68abdf653c89d8dbc2858b88c5d88b0ff38a75288a9da" +dependencies = [ + "proc-macro2", + "pyo3-macros-backend", + "quote", + "syn", +] + +[[package]] +name = "pyo3-macros-backend" +version = "0.23.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fca6726ad0f3da9c9de093d6f116a93c1a38e417ed73bf138472cf4064f72028" +dependencies = [ + "heck", + "proc-macro2", + "pyo3-build-config", + "quote", + "syn", +] + +[[package]] +name = "quick-xml" +version = "0.37.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "331e97a1af0bf59823e6eadffe373d7b27f485be8748f71471c662c1f269b7fb" +dependencies = [ + "memchr", + "serde", +] + +[[package]] +name = "quinn" +version = "0.11.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c1a41e437b6bbd489372cd4971de128e85c855f56c57f283d20ff016cf7c0a8" +dependencies = [ + "bytes", + "cfg_aliases", + "pin-project-lite", + "quinn-proto", + "quinn-udp", + "rustc-hash", + "rustls", + "socket2", + "thiserror 2.0.18", + "tokio", + "tracing", + "web-time", +] + +[[package]] +name = "quinn-proto" +version = "0.11.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fcb935c5bec503c2f0e306bdd3e58bb9029dcb14fa8d9ac76e3a5256ac0763e" +dependencies = [ + "aws-lc-rs", + "bytes", + "getrandom 0.3.4", + "lru-slab", + "rand 0.9.4", + "ring", + "rustc-hash", + "rustls", + "rustls-pki-types", + "slab", + "thiserror 2.0.18", + "tinyvec", + "tracing", + "web-time", +] + +[[package]] +name = "quinn-udp" +version = "0.5.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "addec6a0dcad8a8d96a771f815f0eaf55f9d1805756410b39f5fa81332574cbd" +dependencies = [ + "cfg_aliases", + "libc", + "once_cell", + "socket2", + "tracing", + "windows-sys 0.52.0", +] + [[package]] name = "quote" version = "1.0.45" @@ -1592,9 +3081,21 @@ version = "0.8.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5ca0ecfa931c29007047d1bc58e623ab12e5590e8c7cc53200d5202b69266d8a" dependencies = [ + "libc", + "rand_chacha 0.3.1", "rand_core 0.6.4", ] +[[package]] +name = "rand" +version = "0.9.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "44c5af06bb1b7d3216d91932aed5265164bf384dc89cd6ba05cf59a35f5f76ea" +dependencies = [ + "rand_chacha 0.9.0", + "rand_core 0.9.5", +] + [[package]] name = "rand" version = "0.10.1" @@ -1606,11 +3107,43 @@ dependencies = [ "rand_core 0.10.1", ] +[[package]] +name = "rand_chacha" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" +dependencies = [ + "ppv-lite86", + "rand_core 0.6.4", +] + +[[package]] +name = "rand_chacha" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" +dependencies = [ + "ppv-lite86", + "rand_core 0.9.5", +] + [[package]] name = "rand_core" version = "0.6.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" +dependencies = [ + "getrandom 0.2.17", +] + +[[package]] +name = "rand_core" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76afc826de14238e6e8c374ddcc1fa19e374fd8dd986b0d2af0d02377261d83c" +dependencies = [ + "getrandom 0.3.4", +] [[package]] name = "rand_core" @@ -1746,32 +3279,131 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ba39f3699c378cd8970968dcbff9c43159ea4cfbd88d43c00b22f2ef10a435d2" [[package]] -name = "robust" -version = "1.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4e27ee8bb91ca0adcf0ecb116293afa12d393f9c2b9b9cd54d33e8078fe19839" - -[[package]] -name = "rstar" -version = "0.8.4" +name = "reqwest" +version = "0.12.28" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3a45c0e8804d37e4d97e55c6f258bc9ad9c5ee7b07437009dd152d764949a27c" +checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147" dependencies = [ - "heapless 0.6.1", - "num-traits", - "pdqselect", + "base64", + "bytes", + "futures-core", + "futures-util", + "h2", + "http", + "http-body", + "http-body-util", + "hyper", + "hyper-rustls", + "hyper-util", + "js-sys", + "log", + "percent-encoding", + "pin-project-lite", + "quinn", + "rustls", + "rustls-native-certs", + "rustls-pki-types", "serde", - "smallvec", + "serde_json", + "serde_urlencoded", + "sync_wrapper", + "tokio", + "tokio-rustls", + "tokio-util", + "tower", + "tower-http", + "tower-service", + "url", + "wasm-bindgen", + "wasm-bindgen-futures", + "wasm-streams", + "web-sys", ] [[package]] -name = "rstar" -version = "0.9.3" +name = "reqwest" +version = "0.13.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b40f1bfe5acdab44bc63e6699c28b74f75ec43afb59f3eda01e145aff86a25fa" +checksum = "219c5811de6525e5416c7d5d53bb656d3afdbc6c5af816e0802bcfa42dbdc1c3" dependencies = [ - "heapless 0.7.17", - "num-traits", + "base64", + "bytes", + "encoding_rs", + "futures-channel", + "futures-core", + "futures-util", + "h2", + "http", + "http-body", + "http-body-util", + "hyper", + "hyper-rustls", + "hyper-util", + "js-sys", + "log", + "mime", + "percent-encoding", + "pin-project-lite", + "quinn", + "rustls", + "rustls-pki-types", + "rustls-platform-verifier", + "serde", + "serde_json", + "serde_urlencoded", + "sync_wrapper", + "tokio", + "tokio-rustls", + "tower", + "tower-http", + "tower-service", + "url", + "wasm-bindgen", + "wasm-bindgen-futures", + "web-sys", +] + +[[package]] +name = "ring" +version = "0.17.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a4689e6c2294d81e88dc6261c768b63bc4fcdb852be6d1352498b114f61383b7" +dependencies = [ + "cc", + "cfg-if", + "getrandom 0.2.17", + "libc", + "untrusted", + "windows-sys 0.52.0", +] + +[[package]] +name = "robust" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4e27ee8bb91ca0adcf0ecb116293afa12d393f9c2b9b9cd54d33e8078fe19839" + +[[package]] +name = "rstar" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3a45c0e8804d37e4d97e55c6f258bc9ad9c5ee7b07437009dd152d764949a27c" +dependencies = [ + "heapless 0.6.1", + "num-traits", + "pdqselect", + "serde", + "smallvec", +] + +[[package]] +name = "rstar" +version = "0.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b40f1bfe5acdab44bc63e6699c28b74f75ec43afb59f3eda01e145aff86a25fa" +dependencies = [ + "heapless 0.7.17", + "num-traits", "serde", "smallvec", ] @@ -1841,6 +3473,12 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "rustc-hash" +version = "2.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94300abf3f1ae2e2b8ffb7b58043de3d399c73fa6f4b73826402a5c457614dbe" + [[package]] name = "rustc_version" version = "0.4.1" @@ -1850,6 +3488,105 @@ dependencies = [ "semver", ] +[[package]] +name = "rustix" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6fe4565b9518b83ef4f91bb47ce29620ca828bd32cb7e408f0062e9930ba190" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys 0.61.2", +] + +[[package]] +name = "rustls" +version = "0.23.41" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b92b125634d9b795e7beca796cc790df15a7fb38323bf3196fda83292d06b1f" +dependencies = [ + "aws-lc-rs", + "log", + "once_cell", + "ring", + "rustls-pki-types", + "rustls-webpki", + "subtle", + "zeroize", +] + +[[package]] +name = "rustls-native-certs" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dab5152771c58876a2146916e53e35057e1a4dfa2b9df0f0305b07f611fdea4d" +dependencies = [ + "openssl-probe", + "rustls-pki-types", + "schannel", + "security-framework", +] + +[[package]] +name = "rustls-pemfile" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dce314e5fee3f39953d46bb63bb8a46d40c2f8fb7cc5a3b6cab2bde9721d6e50" +dependencies = [ + "rustls-pki-types", +] + +[[package]] +name = "rustls-pki-types" +version = "1.14.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "30a7197ae7eb376e574fe940d068c30fe0462554a3ddbe4eca7838e049c937a9" +dependencies = [ + "web-time", + "zeroize", +] + +[[package]] +name = "rustls-platform-verifier" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "26d1e2536ce4f35f4846aa13bff16bd0ff40157cdb14cc056c7b14ba41233ba0" +dependencies = [ + "core-foundation 0.10.1", + "core-foundation-sys", + "jni", + "log", + "once_cell", + "rustls", + "rustls-native-certs", + "rustls-platform-verifier-android", + "rustls-webpki", + "security-framework", + "security-framework-sys", + "webpki-root-certs", + "windows-sys 0.61.2", +] + +[[package]] +name = "rustls-platform-verifier-android" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f87165f0995f63a9fbeea62b64d10b4d9d8e78ec6d7d51fb2125fda7bb36788f" + +[[package]] +name = "rustls-webpki" +version = "0.103.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "61c429a8649f110dddef65e2a5ad240f747e85f7758a6bccc7e5777bd33f756e" +dependencies = [ + "aws-lc-rs", + "ring", + "rustls-pki-types", + "untrusted", +] + [[package]] name = "rustversion" version = "1.0.22" @@ -1862,18 +3599,95 @@ version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" +[[package]] +name = "same-file" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93fc1dc3aaa9bfed95e02e6eadabb4baf7e3078b0bd1b4d7b6b0b68378900502" +dependencies = [ + "winapi-util", +] + +[[package]] +name = "schannel" +version = "0.1.29" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91c1b7e4904c873ef0710c1f407dde2e6287de2bebc1bbbf7d430bb7cbffd939" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "schemars" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4cd191f9397d57d581cddd31014772520aa448f65ef991055d7f61582c65165f" +dependencies = [ + "dyn-clone", + "ref-cast", + "serde", + "serde_json", +] + +[[package]] +name = "schemars" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2b42f36aa1cd011945615b92222f6bf73c599a102a300334cd7f8dbeec726cc" +dependencies = [ + "dyn-clone", + "ref-cast", + "serde", + "serde_json", +] + [[package]] name = "scopeguard" version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" +[[package]] +name = "scroll" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "04c565b551bafbef4157586fa379538366e4385d42082f255bfd96e4fe8519da" + +[[package]] +name = "security-framework" +version = "3.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b7f4bc775c73d9a02cde8bf7b2ec4c9d12743edf609006c7facc23998404cd1d" +dependencies = [ + "bitflags", + "core-foundation 0.10.1", + "core-foundation-sys", + "libc", + "security-framework-sys", +] + +[[package]] +name = "security-framework-sys" +version = "2.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ce2691df843ecc5d231c0b14ece2acc3efb62c0a398c7e1d875f3983ce020e3" +dependencies = [ + "core-foundation-sys", + "libc", +] + [[package]] name = "semver" version = "1.0.28" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8a7852d02fc848982e0c167ef163aaff9cd91dc640ba85e263cb1ce46fae51cd" +[[package]] +name = "seq-macro" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1bc711410fbe7399f390ca1c3b60ad0f53f80e95c5eb935e52268a0e2cd49acc" + [[package]] name = "serde" version = "1.0.228" @@ -1910,7 +3724,7 @@ version = "1.0.149" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "83fc039473c5595ace860d8c4fafa220ff474b3fc6bfdb4293327f1a37e94d86" dependencies = [ - "indexmap", + "indexmap 2.14.0", "itoa", "memchr", "serde", @@ -1930,6 +3744,38 @@ dependencies = [ "serde", ] +[[package]] +name = "serde_with" +version = "3.21.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76a5c54c7310e7b8b9577c286d7e399ddd876c3e12b3ed917a8aabc4b96e9e8c" +dependencies = [ + "base64", + "bs58", + "chrono", + "hex", + "indexmap 1.9.3", + "indexmap 2.14.0", + "schemars 0.9.0", + "schemars 1.2.1", + "serde_core", + "serde_json", + "serde_with_macros", + "time", +] + +[[package]] +name = "serde_with_macros" +version = "3.21.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "84d57bc0c8b9a17920c178daa6bb924850d54a9c97ab45194bb8c17ad66bb660" +dependencies = [ + "darling", + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "sha2" version = "0.10.9" @@ -1964,6 +3810,28 @@ version = "0.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d7f45b8998ced5134fb1d75732c77842a3e888f19c1ff98481822e8fbfbf930b" +[[package]] +name = "simd-adler32" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "703d5c7ef118737c72f1af64ad2f6f8c5e1921f818cdcb97b8fe6fc69bf66214" + +[[package]] +name = "simd_cesu8" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94f90157bb87cddf702797c5dadfa0be7d266cdf49e22da2fcaa32eff75b2c33" +dependencies = [ + "rustc_version", + "simdutf8", +] + +[[package]] +name = "simdutf8" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3a9fe34e3e7a50316060351f37187a3f546bce95496156754b601a5fa71b76e" + [[package]] name = "siphasher" version = "1.0.3" @@ -1982,6 +3850,33 @@ version = "1.15.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "67b1b7a3b5fe4f1376887184045fcf45c69e92af734b7aaddc05fb777b6fbd03" +[[package]] +name = "snafu" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e84b3f4eacbf3a1ce05eac6763b4d629d60cbc94d632e4092c54ade71f1e1a2" +dependencies = [ + "snafu-derive", +] + +[[package]] +name = "snafu-derive" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c1c97747dbf44bb1ca44a561ece23508e99cb592e862f22222dcf42f51d1e451" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "snap" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1b6b67fb9a61334225b5b790716f609cd58395f895b3fe8b328786812a40bc3b" + [[package]] name = "socket2" version = "0.6.3" @@ -1989,7 +3884,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3a766e1110788c36f4fa1c2b71b387a7815aa65f88ce0229841826633d93723e" dependencies = [ "libc", - "windows-sys", + "windows-sys 0.61.2", ] [[package]] @@ -2013,6 +3908,16 @@ dependencies = [ "lock_api", ] +[[package]] +name = "spki" +version = "0.7.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d91ed6c858b01f942cd56b37a94b3e0a1798290327d1236e4d9cf4eaca44d29d" +dependencies = [ + "base64ct", + "der", +] + [[package]] name = "sqlparser" version = "0.58.0" @@ -2043,17 +3948,30 @@ checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596" [[package]] name = "stac" -version = "0.17.0" +version = "0.17.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a069e4454aec3e4cf9340d2b2b242764086e23c452f1ff3deff29b30b2bf0b60" +checksum = "79e4b3d417e58a0eeae1022484652d91bb92b8a8e69daabae126452397f15622" dependencies = [ + "arrow-array", + "arrow-cast", + "arrow-json", + "arrow-schema", + "async-stream", "bytes", "chrono", "cql2", + "futures", + "futures-core", + "geo-traits", + "geo-types", + "geoarrow-array", + "geoarrow-schema", "geojson 1.0.0", - "indexmap", + "geoparquet", + "indexmap 2.14.0", "log", "mime", + "parquet", "serde", "serde_json", "serde_urlencoded", @@ -2061,6 +3979,7 @@ dependencies = [ "thiserror 2.0.18", "tracing", "url", + "wkb", ] [[package]] @@ -2073,6 +3992,27 @@ dependencies = [ "syn", ] +[[package]] +name = "stac-io" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9962a05ecbefe29748d9213baf932d054fad0c07952faf2258e451364c39f42" +dependencies = [ + "async-stream", + "bytes", + "futures", + "http", + "parquet", + "reqwest 0.13.4", + "serde", + "serde_json", + "stac", + "thiserror 2.0.18", + "tokio", + "tracing", + "url", +] + [[package]] name = "stacker" version = "0.1.24" @@ -2083,7 +4023,7 @@ dependencies = [ "cfg-if", "libc", "psm", - "windows-sys", + "windows-sys 0.61.2", ] [[package]] @@ -2097,6 +4037,18 @@ dependencies = [ "unicode-properties", ] +[[package]] +name = "strsim" +version = "0.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" + +[[package]] +name = "subtle" +version = "2.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292" + [[package]] name = "syn" version = "2.0.117" @@ -2108,6 +4060,15 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "sync_wrapper" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0bf256ce5efdfa370213c1dabab5935a12e49f2c58d15e9eac2870d3b4f27263" +dependencies = [ + "futures-core", +] + [[package]] name = "synstructure" version = "0.13.2" @@ -2120,43 +4081,144 @@ dependencies = [ ] [[package]] -name = "thiserror" -version = "1.0.69" +name = "system-configuration" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a13f3d0daba03132c0aa9767f98351b3488edc2c100cda2d2ec2b04f3d8d3c8b" +dependencies = [ + "bitflags", + "core-foundation 0.9.4", + "system-configuration-sys", +] + +[[package]] +name = "system-configuration-sys" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e1d1b10ced5ca923a1fcb8d03e96b8d3268065d724548c0211415ff6ac6bac4" +dependencies = [ + "core-foundation-sys", + "libc", +] + +[[package]] +name = "tar" +version = "0.4.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f6221d9a6003c78398e3b239969f352578258df48c8eb051caadae0015bc840" +dependencies = [ + "filetime", + "libc", + "xattr", +] + +[[package]] +name = "target-lexicon" +version = "0.12.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "61c41af27dd6d1e27b1b16b489db798443478cef1f06a660c96db617ba5de3b1" + +[[package]] +name = "tempfile" +version = "3.27.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" +dependencies = [ + "fastrand", + "getrandom 0.4.2", + "once_cell", + "rustix", + "windows-sys 0.61.2", +] + +[[package]] +name = "thiserror" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" +dependencies = [ + "thiserror-impl 1.0.69", +] + +[[package]] +name = "thiserror" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" +dependencies = [ + "thiserror-impl 2.0.18", +] + +[[package]] +name = "thiserror-impl" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "thrift" +version = "0.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7e54bc85fc7faa8bc175c4bab5b92ba8d9a3ce893d0e9f42cc455c8ab16a9e09" +dependencies = [ + "byteorder", + "integer-encoding", + "ordered-float", +] + +[[package]] +name = "time" +version = "0.3.51" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" +checksum = "85c17d80feb7334b40c484e45ed1a5273dfd8bfda537c3be2e74a06a6686f327" dependencies = [ - "thiserror-impl 1.0.69", + "deranged", + "num-conv", + "powerfmt", + "serde_core", + "time-core", + "time-macros", ] [[package]] -name = "thiserror" -version = "2.0.18" +name = "time-core" +version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" -dependencies = [ - "thiserror-impl 2.0.18", -] +checksum = "9e1c906769ad99c88eaa54e728060edef082f8e358ff32030cb7c7d315e81109" [[package]] -name = "thiserror-impl" -version = "1.0.69" +name = "time-macros" +version = "0.2.30" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" +checksum = "dcef1a61bdb119096e153208ec5cbec23944ce8bca13be5c7f60c634f7403935" dependencies = [ - "proc-macro2", - "quote", - "syn", + "num-conv", + "time-core", ] [[package]] -name = "thiserror-impl" -version = "2.0.18" +name = "tiny-keccak" +version = "2.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" +checksum = "2c9d3793400a45f954c52e73d068316d76b6f4e36977e3fcebb13a2721e80237" dependencies = [ - "proc-macro2", - "quote", - "syn", + "crunchy", ] [[package]] @@ -2185,6 +4247,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", +] + [[package]] name = "tokio" version = "1.52.3" @@ -2197,7 +4280,7 @@ dependencies = [ "pin-project-lite", "socket2", "tokio-macros", - "windows-sys", + "windows-sys 0.61.2", ] [[package]] @@ -2237,6 +4320,31 @@ dependencies = [ "whoami", ] +[[package]] +name = "tokio-postgres-rustls" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "27d684bad428a0f2481f42241f821db42c54e2dc81d8c00db8536c506b0a0144" +dependencies = [ + "const-oid 0.9.6", + "ring", + "rustls", + "tokio", + "tokio-postgres", + "tokio-rustls", + "x509-cert", +] + +[[package]] +name = "tokio-rustls" +version = "0.26.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1729aa945f29d91ba541258c8df89027d5792d85a8841fb65e8bf0f4ede4ef61" +dependencies = [ + "rustls", + "tokio", +] + [[package]] name = "tokio-stream" version = "0.1.18" @@ -2287,7 +4395,7 @@ version = "0.25.11+spec-1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0b59c4d22ed448339746c59b905d24568fcbb3ab65a500494f7b8c3e97739f2b" dependencies = [ - "indexmap", + "indexmap 2.14.0", "toml_datetime", "toml_parser", "winnow", @@ -2302,6 +4410,51 @@ dependencies = [ "winnow", ] +[[package]] +name = "tower" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebe5ef63511595f1344e2d5cfa636d973292adc0eec1f0ad45fae9f0851ab1d4" +dependencies = [ + "futures-core", + "futures-util", + "pin-project-lite", + "sync_wrapper", + "tokio", + "tower-layer", + "tower-service", +] + +[[package]] +name = "tower-http" +version = "0.6.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4cfcf7e2740e6fc6d4d688b4ef00650406bb94adf4731e43c096c3a19fe40840" +dependencies = [ + "bitflags", + "bytes", + "futures-util", + "http", + "http-body", + "pin-project-lite", + "tower", + "tower-layer", + "tower-service", + "url", +] + +[[package]] +name = "tower-layer" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "121c2a6cda46980bb0fcd1647ffaf6cd3fc79a013de288782836f6df9c48780e" + +[[package]] +name = "tower-service" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8df9b6e13f2d32c91b9bd719c00d1958837bc7dec474d94952798cc8e69eeec3" + [[package]] name = "tracing" version = "0.1.44" @@ -2333,6 +4486,18 @@ dependencies = [ "once_cell", ] +[[package]] +name = "try-lock" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b" + +[[package]] +name = "twox-hash" +version = "2.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ea3136b675547379c4bd395ca6b938e5ad3c3d20fad76e7fe85f9e0d011419c" + [[package]] name = "typenum" version = "1.20.0" @@ -2387,6 +4552,18 @@ version = "0.2.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853" +[[package]] +name = "unindent" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7264e107f553ccae879d21fbea1d6724ac785e8c3bfc762137959b5802826ef3" + +[[package]] +name = "untrusted" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ecb6da28b8a351d773b68d5825ac39017e680750f980f3a1a85cd8dd28a47c1" + [[package]] name = "url" version = "2.5.8" @@ -2406,6 +4583,12 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" +[[package]] +name = "utf8parse" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" + [[package]] name = "uuid" version = "1.23.1" @@ -2439,6 +4622,25 @@ version = "0.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5c3082ca00d5a5ef149bb8b555a72ae84c9c59f7250f013ac822ac2e49b19c64" +[[package]] +name = "walkdir" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29790946404f91d9c5d06f9874efddea1dc06c5efe94541a7d6863108e3a5e4b" +dependencies = [ + "same-file", + "winapi-util", +] + +[[package]] +name = "want" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bfa7760aed19e106de2c7c0b581b509f2f25d3dacaf737cb82ac61bc6d760b0e" +dependencies = [ + "try-lock", +] + [[package]] name = "wasi" version = "0.11.1+wasi-snapshot-preview1" @@ -2494,6 +4696,16 @@ dependencies = [ "wasm-bindgen-shared", ] +[[package]] +name = "wasm-bindgen-futures" +version = "0.4.71" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96492d0d3ffba25305a7dc88720d250b1401d7edca02cc3bcd50633b424673b8" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + [[package]] name = "wasm-bindgen-macro" version = "0.2.121" @@ -2543,11 +4755,24 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bb0e353e6a2fbdc176932bbaab493762eb1255a7900fe0fea1a2f96c296cc909" dependencies = [ "anyhow", - "indexmap", + "indexmap 2.14.0", "wasm-encoder", "wasmparser", ] +[[package]] +name = "wasm-streams" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "15053d8d85c7eccdbefef60f06769760a563c7f0a9d6902a13d35c7800b0ad65" +dependencies = [ + "futures-util", + "js-sys", + "wasm-bindgen", + "wasm-bindgen-futures", + "web-sys", +] + [[package]] name = "wasmparser" version = "0.244.0" @@ -2556,7 +4781,7 @@ checksum = "47b807c72e1bac69382b3a6fb3dbe8ea4c0ed87ff5629b8685ae6b9a611028fe" dependencies = [ "bitflags", "hashbrown 0.15.5", - "indexmap", + "indexmap 2.14.0", "semver", ] @@ -2570,6 +4795,43 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "web-time" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a6580f308b1fad9207618087a65c04e7a10bc77e02c8e84e9b00dd4b12fa0bb" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + +[[package]] +name = "webpki-root-certs" +version = "1.0.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d46a5a140e6f7afeccd8eae97eff335163939eac8b929834875168b29b3d267" +dependencies = [ + "rustls-pki-types", +] + +[[package]] +name = "webpki-roots" +version = "0.26.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "521bc38abb08001b01866da9f51eb7c5d647a19260e00054a8c7fd5f9e57f7a9" +dependencies = [ + "webpki-roots 1.0.8", +] + +[[package]] +name = "webpki-roots" +version = "1.0.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bf85cb06032201fa7c6f829d7db5a7e5aa45bcc0655327713065f6f0576731bf" +dependencies = [ + "rustls-pki-types", +] + [[package]] name = "whoami" version = "2.1.2" @@ -2583,6 +4845,15 @@ dependencies = [ "web-sys", ] +[[package]] +name = "winapi-util" +version = "0.1.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" +dependencies = [ + "windows-sys 0.61.2", +] + [[package]] name = "windows-core" version = "0.62.2" @@ -2624,6 +4895,17 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" +[[package]] +name = "windows-registry" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "02752bf7fbdcce7f2a27a742f798510f3e5ad88dbe84871e5168e2120c3d5720" +dependencies = [ + "windows-link", + "windows-result", + "windows-strings", +] + [[package]] name = "windows-result" version = "0.4.1" @@ -2642,6 +4924,15 @@ dependencies = [ "windows-link", ] +[[package]] +name = "windows-sys" +version = "0.52.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "282be5f36a8ce781fad8c8ae18fa3f9beff57ec1b52cb3de0789201425d9a33d" +dependencies = [ + "windows-targets", +] + [[package]] name = "windows-sys" version = "0.61.2" @@ -2651,6 +4942,70 @@ dependencies = [ "windows-link", ] +[[package]] +name = "windows-targets" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" +dependencies = [ + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", + "windows_i686_gnullvm", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", +] + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" + +[[package]] +name = "windows_i686_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" + +[[package]] +name = "windows_i686_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" + +[[package]] +name = "windows_i686_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" + [[package]] name = "winnow" version = "1.0.2" @@ -2694,7 +5049,7 @@ checksum = "b7c566e0f4b284dd6561c786d9cb0142da491f46a9fbed79ea69cdad5db17f21" dependencies = [ "anyhow", "heck", - "indexmap", + "indexmap 2.14.0", "prettyplease", "syn", "wasm-metadata", @@ -2725,7 +5080,7 @@ checksum = "9d66ea20e9553b30172b5e831994e35fbde2d165325bec84fc43dbf6f4eb9cb2" dependencies = [ "anyhow", "bitflags", - "indexmap", + "indexmap 2.14.0", "log", "serde", "serde_derive", @@ -2744,7 +5099,7 @@ checksum = "ecc8ac4bc1dc3381b7f59c34f00b67e18f910c2c0f50015669dde7def656a736" dependencies = [ "anyhow", "id-arena", - "indexmap", + "indexmap 2.14.0", "log", "semver", "serde", @@ -2754,6 +5109,18 @@ dependencies = [ "wasmparser", ] +[[package]] +name = "wkb" +version = "0.9.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a120b336c7ad17749026d50427c23d838ecb50cd64aaea6254b5030152f890a9" +dependencies = [ + "byteorder", + "geo-traits", + "num_enum", + "thiserror 1.0.69", +] + [[package]] name = "wkt" version = "0.11.1" @@ -2785,6 +5152,28 @@ version = "0.6.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1ffae5123b2d3fc086436f8834ae3ab053a283cfac8fe0a0b8eaae044768a4c4" +[[package]] +name = "x509-cert" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1301e935010a701ae5f8655edc0ad17c44bad3ac5ce8c39185f75453b720ae94" +dependencies = [ + "const-oid 0.9.6", + "der", + "spki", + "tls_codec", +] + +[[package]] +name = "xattr" +version = "1.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32e45ad4206f6d2479085147f02bc2ef834ac85886624a23575ae137c8aa8156" +dependencies = [ + "libc", + "rustix", +] + [[package]] name = "yoke" version = "0.8.2" @@ -2849,6 +5238,26 @@ dependencies = [ "synstructure", ] +[[package]] +name = "zeroize" +version = "1.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e13c156562582aa81c60cb29407084cdb54c4164760106ab78e6c5b0858cf64e" +dependencies = [ + "zeroize_derive", +] + +[[package]] +name = "zeroize_derive" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3c50655cbb0fe3fc43170059e702f1ce5e19b84cec58dc87b037a09935c2f328" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "zerotrie" version = "0.2.4" @@ -2882,8 +5291,42 @@ dependencies = [ "syn", ] +[[package]] +name = "zlib-rs" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "977347db8caa080403f6b6b7c1cda9479a8e869316f7e13a59b19076a40f94e3" + [[package]] name = "zmij" version = "1.0.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa" + +[[package]] +name = "zstd" +version = "0.13.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e91ee311a569c327171651566e07972200e76fcfe2242a4fa446149a3881c08a" +dependencies = [ + "zstd-safe", +] + +[[package]] +name = "zstd-safe" +version = "7.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f49c4d5f0abb602a93fb8736af2a4f4dd9512e36f7f570d66e65ff867ed3b9d" +dependencies = [ + "zstd-sys", +] + +[[package]] +name = "zstd-sys" +version = "2.0.16+zstd.1.5.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91e19ebc2adc8f83e43039e79776e3fda8ca919132d68a1fed6a5faca2683748" +dependencies = [ + "cc", + "pkg-config", +] diff --git a/src/pgstac-rs/Cargo.toml b/src/pgstac-rs/Cargo.toml index 82895714..a7f39571 100644 --- a/src/pgstac-rs/Cargo.toml +++ b/src/pgstac-rs/Cargo.toml @@ -11,14 +11,103 @@ repository = "https://github.com/stac-utils/pgstac" license = "MIT OR Apache-2.0" rust-version = "1.88" +[lib] +crate-type = ["rlib", "cdylib"] + +[features] +default = [] +# A Python extension module (built with maturin) exposing the read + geoparquet-export API over a +# connection pool. +python = ["pool", "export", "dep:pyo3", "dep:pyo3-async-runtimes"] +migrations = [] +# A `deadpool`-backed connection pool ([`PgstacPool`]) with rustls TLS. +pool = [ + "dep:deadpool-postgres", + "dep:tokio-postgres-rustls", + "dep:rustls", + "dep:rustls-pemfile", + "dep:webpki-roots", + "dep:async-stream", +] +# Export/dump library: row sources, hydration glue, formats, sinks, manifest. Pulls in the stac +# geoparquet writer + checksum/spill helpers, plus the parquet/geoparquet/arrow crates the parallel +# loader decode drives directly (same versions stac/geoparquet resolves transitively). +export = [ + "stac/geoparquet", + "dep:async-trait", + "dep:tempfile", + "dep:arrow-array", + "dep:parquet", + "dep:geoparquet", + "tokio/sync", + "tokio/macros", +] +# The `object_store`-backed remote dump sink (local fs + S3/GCS/Azure), usable without the whole CLI. +store = [ + "export", + "dep:object_store", + "dep:url", + "tokio/fs", + "tokio/io-util", +] +# The `pgstac` binary (CLI). Implies `export` + `store` (dump/search) + `pool` (load/restore via the Rust loader). +cli = [ + "export", + "pool", + "store", + "dep:clap", + "dep:stac-io", + "dep:tar", + "dep:zstd", + "tokio/macros", + "tokio/rt-multi-thread", + "tokio/fs", +] + [dependencies] +async-stream = { version = "0.3", optional = true } +base64 = "0.22" +bytes = "1" +chrono = { version = "0.4", default-features = false, features = ["clock"] } +deadpool-postgres = { version = "0.14", optional = true } +futures = "0.3" +geozero = { version = "0.14", default-features = false, features = ["with-wkb", "with-geojson"] } +rustls = { version = "0.23", optional = true } +rustls-pemfile = { version = "2.2", optional = true } serde = "1.0" -serde_json = "1.0" -stac = "0.17.0" +serde_json = { version = "1.0", features = ["raw_value", "float_roundtrip"] } +sha2 = "0.10" +stac = { version = "0.17.0", features = ["async"] } thiserror = "2.0" -tokio = { version = "1.44", features = ["rt"] } -tokio-postgres = { version = "0.7.12", features = ["with-serde_json-1"] } +tokio = { version = "1.44", features = ["rt", "sync"] } +tokio-postgres = { version = "0.7.12", features = ["with-serde_json-1", "with-chrono-0_4"] } +tokio-postgres-rustls = { version = "0.13", optional = true } tracing = "0.1.40" +webpki-roots = { version = "0.26", optional = true } + +# export feature (sha2 is a non-optional dep above — canonical.rs always needs it) +async-trait = { version = "0.1", optional = true } +tempfile = { version = "3", optional = true } +# export feature: parallel stac-geoparquet decode reads row groups directly. Versions match what +# stac's `geoparquet` feature pulls in transitively (parquet 58, geoparquet 0.8, arrow-array 58). +arrow-array = { version = "58", optional = true } +parquet = { version = "58", optional = true } +geoparquet = { version = "0.8", optional = true } + +# store feature +object_store = { version = "0.11", optional = true, features = ["aws"] } +url = { version = "2", optional = true } + +# cli feature +clap = { version = "4", optional = true, features = ["derive", "env"] } +# rustac's format enum (also folds in geoparquet compression) for the CLI's --format. +stac-io = { version = "0.3", optional = true, features = ["geoparquet"] } +tar = { version = "0.4", optional = true } +zstd = { version = "0.13", optional = true } + +# python feature +pyo3 = { version = "0.23", optional = true, features = ["macros"] } +pyo3-async-runtimes = { version = "0.23", optional = true, features = ["tokio-runtime"] } [dev-dependencies] geojson = "1.0" @@ -26,6 +115,11 @@ rstest = "0.26.1" tokio = { version = "1.44", features = ["rt-multi-thread", "macros"] } tokio-test = "0.4.4" +[[bin]] +name = "pgstac" +path = "src/bin/pgstac.rs" +required-features = ["cli"] + [package.metadata.docs.rs] all-features = true rustdoc-args = ["--cfg", "docsrs"] diff --git a/src/pgstac-rs/README.md b/src/pgstac-rs/README.md index 6df92d00..30816d6a 100644 --- a/src/pgstac-rs/README.md +++ b/src/pgstac-rs/README.md @@ -16,6 +16,68 @@ pgstac = "*" See the [documentation](https://docs.rs/pgstac) for more. +### Cargo features + +| Feature | What it adds | +| --------- | ------------ | +| *(default)* | The `Pgstac` trait over any `tokio_postgres::GenericClient`. | +| `pool` | [`PgstacPool`] — a `deadpool` connection pool (rustls TLS, PgBouncer-safe) with the read API: `search_page`, the flat-memory streaming iterator `search_items` / `stream_ndjson`, `collection_search`, `get_item`/`get_collection`/`get_queryables`, and `search_matched`. | +| `export` | The export/dump library: scan partitions and write fully-hydrated stac-geoparquet + a sha256'd manifest. | +| `cli` | The `pgstac` binary (implies `export`): `pgstac dump` and `pgstac search`. | +| `python` | A pyo3 extension module (`pypgstac_rs`) built with [maturin]. | + +The read API drives `search_plan` / `collection_search_plan` and does the band-stepping, +hydration (EWKB→GeoJSON, fragment merge), and keyset token minting **in Rust** — page-equivalent +to SQL `search()` but with flat memory under streaming and the work moved off the database. + +### Pooled read API + +```rust,no_run +use pgstac::{ConnectConfig, PgstacPool}; +use futures::StreamExt; +use serde_json::json; + +# tokio_test::block_on(async { +let pool = PgstacPool::connect(ConnectConfig::from_env()).await.unwrap(); + +// A bounded page (keyset tokens included). +let page = pool.search_page(&json!({"collections": ["landsat-c2-l2"]}), None, 10).await.unwrap(); + +// Or stream every match with flat memory (one row + the fragment cache at a time). +let mut items = Box::pin(pool.search_items(json!({"collections": ["landsat-c2-l2"]}), None, None)); +while let Some(item) = items.next().await { + let _item = item.unwrap(); +} +# }) +``` + +### CLI (`cli` feature) + +```sh +# Dump an instance to fully-hydrated stac-geoparquet + manifest. +pgstac dump --dsn "$PGSTAC_DSN" --out ./dump # or s3://…, foo.tar.zst, - + +# Stream a search as NDJSON / a FeatureCollection page / geoparquet. +pgstac search --dsn "$PGSTAC_DSN" -c landsat-c2-l2 --datetime 2024-01-01/.. --format ndjson +``` + +### Python wheel (`python` feature) + +```sh +maturin build -m src/pgstac-rs/Cargo.toml # builds the pypgstac_rs wheel +``` + +```python +import asyncio, orjson, pypgstac_rs + +async def main(): + pool = await pypgstac_rs.Pgstac.connect() # libpq env, or pass a dsn + fc = orjson.loads(await pool.search(orjson.dumps({"collections": ["landsat-c2-l2"], "limit": 10}).decode())) + print(fc["numberReturned"]) + +asyncio.run(main()) +``` + ## Testing **pgstac** needs a blank **pgstac** database for testing, so is not part of the default workspace build. @@ -45,3 +107,5 @@ PGSTAC_RS_TEST_DB=postgresql://otherusername:otherpassword@otherhost:7822/otherd ## Other info This crate used to be part of the [rustac](https://github.com/stac-utils/rustac) monorepo, but was moved here in May 2026. + +[maturin]: https://github.com/PyO3/maturin diff --git a/src/pgstac-rs/pyproject.toml b/src/pgstac-rs/pyproject.toml new file mode 100644 index 00000000..3d8c11d5 --- /dev/null +++ b/src/pgstac-rs/pyproject.toml @@ -0,0 +1,18 @@ +[build-system] +requires = ["maturin>=1.7,<2.0"] +build-backend = "maturin" + +[project] +name = "pypgstac-rs" +description = "Rust-accelerated pgstac read API (search, collection search, getters, queryables) for stac-fastapi-pgstac" +requires-python = ">=3.9" +license = { text = "MIT" } +dynamic = ["version"] +classifiers = [ + "Programming Language :: Rust", + "Programming Language :: Python :: Implementation :: CPython", +] + +[tool.maturin] +features = ["python", "pyo3/extension-module"] +module-name = "pypgstac_rs" diff --git a/src/pgstac-rs/src/api/mod.rs b/src/pgstac-rs/src/api/mod.rs new file mode 100644 index 00000000..42c6b487 --- /dev/null +++ b/src/pgstac-rs/src/api/mod.rs @@ -0,0 +1,4 @@ +//! The rustac-native `stac::api` client-trait implementations over [`PgstacPool`](crate::PgstacPool). + +#[cfg(feature = "pool")] +pub(crate) mod stac_api; diff --git a/src/pgstac-rs/src/api/stac_api.rs b/src/pgstac-rs/src/api/stac_api.rs new file mode 100644 index 00000000..a24c22eb --- /dev/null +++ b/src/pgstac-rs/src/api/stac_api.rs @@ -0,0 +1,72 @@ +//! Native [`stac::api`] client-trait impls for [`PgstacPool`], so a pgstac database is a first-class +//! rustac backend (search / collections / transaction). Every method routes through the same engine +//! the rest of the crate uses — the keyset search in [`crate::search()`] and the Rust loader for writes — +//! so these traits are the rustac-native API surface over that engine, not a second implementation. + +use crate::{Error, PgstacPool}; +use futures::{Stream, StreamExt}; +use stac::api::{ + CollectionsClient, ItemCollection, ItemsClient, Search, StreamItemsClient, TransactionClient, +}; +use stac::{Collection, Item}; + +impl ItemsClient for PgstacPool { + type Error = Error; + + async fn search(&self, search: Search) -> Result { + self.client().await?.search(search).await + } + + async fn item(&self, collection_id: &str, item_id: &str) -> Result, Error> { + match self.get_item(collection_id, item_id).await? { + Some(value) => Ok(Some(serde_json::from_value(value)?)), + None => Ok(None), + } + } +} + +impl CollectionsClient for PgstacPool { + type Error = Error; + + async fn collections(&self) -> Result, Error> { + self.client().await?.collections().await + } + + async fn collection(&self, id: &str) -> Result, Error> { + match self.get_collection(id).await? { + Some(value) => Ok(Some(serde_json::from_value(value)?)), + None => Ok(None), + } + } +} + +impl TransactionClient for PgstacPool { + type Error = Error; + + async fn add_collection(&mut self, collection: Collection) -> Result<(), Error> { + let value = serde_json::to_value(collection)?; + self.create_collection(&value).await + } + + async fn add_item(&mut self, item: Item) -> Result<(), Error> { + let value = serde_json::to_value(item)?; + self.create_item(value).await.map(|_| ()) + } +} + +impl StreamItemsClient for PgstacPool { + type Error = Error; + + /// Streams every matching item across all keyset pages with flat memory, by mapping the crate's + /// portal-based [`search_items`](PgstacPool::search_items) stream into `stac::api::Item`s. This is + /// the rustac-native, constant-memory alternative to page-by-page pagination. + async fn search_stream( + &self, + search: Search, + ) -> Result> + Send, Error> { + let search = serde_json::to_value(search)?; + let items = self.search_items(search, None, None); + Ok(items + .map(|result| result.and_then(|value| serde_json::from_value(value).map_err(Error::from)))) + } +} diff --git a/src/pgstac-rs/src/bin/pgstac.rs b/src/pgstac-rs/src/bin/pgstac.rs new file mode 100644 index 00000000..fc025bd1 --- /dev/null +++ b/src/pgstac-rs/src/bin/pgstac.rs @@ -0,0 +1,1212 @@ +//! The `pgstac` CLI binary (built with `--features cli`). +//! +//! Thin wrapper over the `pgstac` library: one binary, subcommand per feature +//! (`dump`/`search`/`load`/`restore`). The CLI is pure arg-parsing + wiring; all +//! logic lives in the lib. `load`/`restore` route through the Rust loader +//! ([`PgstacPool::create_items`] → `load_items`), so dehydration + fragment +//! splitting + the binary COPY all happen in Rust. + +use clap::{Parser, Subcommand}; +use pgstac::export::format::default_compression; +use pgstac::export::plan::Prefilter; +use pgstac::export::sink::{DirSink, StdoutSink, TarSink}; +use pgstac::export::{DumpConfig, DumpPlanner}; +use pgstac::ingest::ConflictPolicy; +use pgstac::{ConnectConfig, PgstacPool, PoolOptions}; +use serde_json::Value; +use stac::geoparquet::Compression; +use stac_io::Format; +use std::path::{Path, PathBuf}; +use std::process::ExitCode; +use std::time::Instant; +use tokio_postgres::NoTls; + +#[derive(Parser, Debug)] +#[command( + name = "pgstac", + version, + about = "pgstac tooling: dump, search, load, restore" +)] +struct Cli { + /// Postgres connection string. Defaults to $PGSTAC_DSN or a local dev DSN. Shared by every + /// subcommand. + #[arg(long, env = "PGSTAC_DSN", global = true)] + dsn: Option, + + #[command(subcommand)] + command: Command, +} + +#[derive(Subcommand, Debug)] +enum Command { + /// Dump a pgstac instance (or a subset) to a directory, tar, S3, or stdout. + Dump(DumpArgs), + /// Stream a search/CQL2 result as NDJSON / ItemCollection / geoparquet + /// (0.10 only — rides the keyset search engine). + Search(SearchArgs), + /// Load STAC items + collections (stac-geoparquet, NDJSON, or JSON) through + /// the Rust loader: dehydration, fragment splitting, and the binary COPY all + /// run in Rust. Collections are created before items. + Load(LoadArgs), + /// Restore a `pgstac dump` directory/tar (collections + partition geoparquet) + /// through the same Rust loader. + Restore(RestoreArgs), + /// Run the async partition-stats maintenance: + /// recompute exact bounds + counts for partitions ingest left dirty. + Maintain(MaintainArgs), + /// Delete an item (with --item) or a whole collection (without --item). + Delete(DeleteArgs), +} + +#[derive(clap::Args, Debug)] +struct DeleteArgs { + /// Collection id (required). + #[arg(long, short)] + collection: String, + + /// Item id. Given, deletes that item; omitted, deletes the whole collection (and its items). + #[arg(long)] + item: Option, + + /// Skip the interactive confirmation prompt. + #[arg(long)] + yes: bool, +} + +#[derive(clap::Args, Debug)] +struct MaintainArgs { + /// Cap the number of dirty partitions tightened this run (oldest first). + /// Omit to tighten all dirty partitions. + #[arg(long)] + limit: Option, +} + +#[derive(clap::Args, Debug)] +struct LoadArgs { + /// Inputs to load: stac-geoparquet (`.parquet`/`.geoparquet`), NDJSON + /// (`.ndjson`), JSON (item or ItemCollection), or collection JSON. Directories + /// are scanned recursively. Collection files are loaded before items. + #[arg(required = true)] + inputs: Vec, + + /// Items per binary-COPY batch (one loader transaction per batch). + #[arg(long, default_value_t = DEFAULT_BATCH_SIZE)] + batch_size: usize, + + /// Number of batches loaded concurrently (each uses one pooled connection). + #[arg(long, default_value_t = default_ingest_parallelism())] + concurrency: usize, + + /// Conflict policy when an item id already exists. + #[arg(long, value_enum, default_value_t = ConflictPolicy::Upsert)] + policy: ConflictPolicy, + + /// Pool size. Defaults to max(concurrency, 4). + #[arg(long)] + pool_size: Option, + + /// Cap the total number of items loaded (across all sources). For benchmarking / sampling. + #[arg(long)] + limit: Option, + + /// Skip items already present-and-current, loading only new + changed (a per-partition precheck for + /// re-ingest / sync). `upsert`/`delsert` skip ids whose content hash matches; `ignore` skips by id alone. + #[arg(long)] + skip_unchanged: bool, +} + +/// Items per binary-COPY batch by default. +const DEFAULT_BATCH_SIZE: usize = 5_000; +/// Ceiling on auto-detected ingest parallelism (decode threads + concurrent load batches). +const MAX_INGEST_PARALLELISM: usize = 8; + +/// Default ingest parallelism: CPU count capped at `MAX_INGEST_PARALLELISM` (a safe ceiling for +/// decode/connection fan-out). Powers both `--concurrency` (parallel load batches) and the ndjson +/// decode-thread default, so the loader parallelizes out of the box. +fn default_ingest_parallelism() -> usize { + std::thread::available_parallelism() + .map(|n| n.get()) + .unwrap_or(4) + .min(MAX_INGEST_PARALLELISM) +} + +#[derive(clap::Args, Debug)] +struct RestoreArgs { + /// A `pgstac dump` directory (collections.ndjson + per-partition geoparquet). + src: PathBuf, + + /// Items per binary-COPY batch. + #[arg(long, default_value_t = DEFAULT_BATCH_SIZE)] + batch_size: usize, + + /// Number of partition files restored concurrently. + #[arg(long, default_value_t = default_ingest_parallelism())] + concurrency: usize, + + /// Pool size. Defaults to max(concurrency, 4). + #[arg(long)] + pool_size: Option, +} + +#[derive(clap::Args, Debug)] +struct SearchArgs { + /// Output destination: a file path or `-` for stdout (default). + #[arg(long, short, default_value = "-")] + out: String, + + /// Output format: `json` (one ItemCollection page), `ndjson` (default, stream all), or + /// `geoparquet[]` — rustac's format spelling, which also carries the geoparquet + /// compression codec. + #[arg(long, default_value = "ndjson", value_parser = str::parse::)] + format: Format, + + /// Restrict to these collection ids (repeatable). + #[arg(long = "collection", short = 'c')] + collections: Vec, + + /// Restrict to these item ids (repeatable). + #[arg(long = "id")] + ids: Vec, + + /// Bbox: west,south,east,north (EPSG:4326). + #[arg(long, value_delimiter = ',', num_args = 4, allow_hyphen_values = true)] + bbox: Option>, + + /// Datetime / interval (RFC3339, STAC datetime syntax). + #[arg(long)] + datetime: Option, + + /// CQL2-JSON filter (a JSON object). + #[arg(long)] + filter: Option, + + /// Full STAC search body as JSON (sortby / fields / intersects / query / conf / ...). When set it + /// overrides the individual filter flags above, letting the CLI drive the whole search spec — handy + /// for benchmarking the engine on sort + fields (full-vs-slim) + complex geometry. + #[arg(long)] + search_json: Option, + + /// Page size (the keyset limit). Omit for the server default. + #[arg(long)] + limit: Option, + + /// Continuation token (`next:` / `prev:`). With this set the + /// ItemCollection format returns exactly that one page. + #[arg(long)] + token: Option, + + /// Max rows per parquet row-group (geoparquet only). Smaller values bound peak memory when + /// streaming out parquet. Defaults to $PGSTAC_PARQUET_ROW_GROUP_SIZE, else the parquet default. + #[arg(long, env = "PGSTAC_PARQUET_ROW_GROUP_SIZE")] + row_group_size: Option, + + /// Cap the total items streamed (NDJSON / geoparquet). + #[arg(long)] + max_items: Option, +} + +#[derive(clap::Args, Debug)] +struct DumpArgs { + /// Output destination: a directory path, `s3://bucket/prefix` (or other + /// object-store URL), a `*.tar` / `*.tar.zst` file, or `-` for stdout. + #[arg(long, short)] + out: String, + + /// Restrict to these collection ids (repeatable). Omit for a full dump. + #[arg(long = "collection", short = 'c')] + collections: Vec, + + /// Datetime prefilter start (RFC3339, inclusive). + #[arg(long, requires = "datetime_end")] + datetime_start: Option, + + /// Datetime prefilter end (RFC3339, exclusive). + #[arg(long, requires = "datetime_start")] + datetime_end: Option, + + /// Bbox prefilter: west,south,east,north (EPSG:4326). + #[arg(long, value_delimiter = ',', num_args = 4, allow_hyphen_values = true)] + bbox: Option>, + + /// Parquet compression codec in parquet's spelling (e.g. `snappy`, `uncompressed`, `zstd(15)`). + /// Omit for the fast default (zstd). + #[arg(long, value_parser = str::parse::)] + compression: Option, + + /// Continue past per-item errors, recording skips in the report. + #[arg(long)] + skip_errors: bool, + + /// Memory budget for the buffered geoparquet path (bytes). Default ~25% RAM. + #[arg(long)] + memory_budget: Option, + + /// Max partitions dumped concurrently. 1 = sequential. Values > 1 + /// open one extra Postgres connection per worker. + #[arg(long, default_value_t = default_ingest_parallelism())] + concurrency: usize, + + /// Dump the whole instance under one repeatable-read snapshot so + /// concurrent ingest cannot tear the output. Implies a parallel run. + #[arg(long)] + consistent: bool, + + /// Write a `.tar.zst` (compressed) tar when --out is a `.tar` path. + #[arg(long)] + tar_zstd: bool, + + /// Report the plan without writing data. + #[arg(long)] + dry_run: bool, +} + +const DEFAULT_DSN: &str = "postgresql://username:password@localhost:5432/postgis"; + +/// Builds a [`ConnectConfig`] from the environment, with an explicit `--dsn` (or `$PGSTAC_DSN`) taking +/// precedence over ambient `PG*` / `DATABASE_URL` values. +fn connect_config(dsn: Option<&str>) -> ConnectConfig { + let mut config = ConnectConfig::from_env(); + if let Some(dsn) = dsn { + config.dsn = Some(dsn.to_string()); + } + config +} + +/// Confirms a destructive delete. `--yes` proceeds without asking; on a TTY it prompts (a `y`/`yes` +/// answer proceeds, anything else aborts). When stdin is not a TTY and `--yes` was not given it errors +/// rather than hang waiting for input that will never arrive. +fn confirm_delete(target: &str, yes: bool) -> Result> { + use std::io::{IsTerminal, Write}; + if yes { + return Ok(true); + } + if !std::io::stdin().is_terminal() { + return Err( + "refusing to delete without confirmation: stdin is not a TTY; pass --yes to confirm".into(), + ); + } + eprint!("delete {target}? this cannot be undone [y/N]: "); + std::io::stderr().flush()?; + let mut answer = String::new(); + let _ = std::io::stdin().read_line(&mut answer)?; + Ok(matches!( + answer.trim().to_ascii_lowercase().as_str(), + "y" | "yes" + )) +} + +fn main() -> ExitCode { + let cli = Cli::parse(); + let runtime = match tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build() + { + Ok(rt) => rt, + Err(e) => { + eprintln!("error: failed to start runtime: {e}"); + return ExitCode::FAILURE; + } + }; + let Cli { dsn, command } = cli; + let result = match command { + Command::Dump(args) => runtime.block_on(run_dump(args, dsn)), + Command::Search(args) => runtime.block_on(run_search(args, dsn)), + Command::Load(args) => runtime.block_on(run_load(args, dsn)), + Command::Restore(args) => runtime.block_on(run_restore(args, dsn)), + Command::Maintain(args) => runtime.block_on(run_maintain(args, dsn)), + Command::Delete(args) => runtime.block_on(run_delete(args, dsn)), + }; + match result { + Ok(()) => ExitCode::SUCCESS, + Err(e) => { + eprintln!("error: {e}"); + // Walk the source chain: tokio_postgres::Error renders as a terse `db error`, so a Postgres + // failure otherwise hides its SQLSTATE + server message. Surface them (plus detail/hint). + let mut src = e.source(); + while let Some(s) = src { + if let Some(db) = s.downcast_ref::() { + eprintln!(" SQLSTATE {}: {}", db.code().code(), db.message()); + if let Some(detail) = db.detail() { + eprintln!(" detail: {detail}"); + } + if let Some(hint) = db.hint() { + eprintln!(" hint: {hint}"); + } + } else if s.downcast_ref::().is_none() { + // tokio_postgres::Error renders as a terse "db error"; skip it — its DbError source + // (handled on the next hop) carries the real SQLSTATE + message. + eprintln!(" caused by: {s}"); + } + src = s.source(); + } + ExitCode::FAILURE + } + } +} + +async fn run_dump(args: DumpArgs, dsn: Option) -> Result<(), Box> { + let dsn = dsn.unwrap_or_else(|| DEFAULT_DSN.to_string()); + let (client, connection) = tokio_postgres::connect(&dsn, NoTls).await?; + tokio::spawn(async move { + if let Err(e) = connection.await { + eprintln!("connection error: {e}"); + } + }); + client + .batch_execute("SET search_path = pgstac, public;") + .await?; + + let datetime = match (&args.datetime_start, &args.datetime_end) { + (Some(s), Some(e)) => Some((s.clone(), e.clone())), + _ => None, + }; + let bbox = match &args.bbox { + Some(b) if b.len() == 4 => Some([b[0], b[1], b[2], b[3]]), + _ => None, + }; + let prefilter = Prefilter { datetime, bbox }; + + // Resume: if a prior _checkpoint.json sits next to a directory dump, + // load + verify it so already-completed partition files are skipped. + let resume_completed = + if args.out != "-" && !is_tar_path(&args.out) && object_store_url(&args.out).is_none() { + load_resume_set(&args.out) + } else { + Default::default() + }; + + let config = DumpConfig { + collection_ids: if args.collections.is_empty() { + None + } else { + Some(args.collections.clone()) + }, + prefilter, + compression: args.compression.unwrap_or_else(default_compression), + skip_errors: args.skip_errors, + resume_completed, + memory_budget: args.memory_budget, + tool_version: env!("CARGO_PKG_VERSION").to_string(), + server: None, + concurrency: args.concurrency, + consistent: args.consistent, + }; + + if args.dry_run { + return dry_run(&client, &config).await; + } + + // A parallel run is requested when concurrency > 1 or --consistent is set. + // The TAR/stdout sinks serialize internally but still work; fs/object_store + // sinks write concurrently. + let parallel = config.concurrency > 1 || config.consistent; + let planner = DumpPlanner::new(config); + + // Choose a sink from --out (format ⟂ sink; here output is always the + // backup geoparquet layout). + let report = if parallel { + use pgstac::export::DsnFactory; + use std::sync::Arc; + let factory = DsnFactory::new(dsn.clone()); + if args.out == "-" { + planner + .run_parallel(&client, Arc::new(StdoutSink), &factory) + .await? + } else if is_tar_path(&args.out) { + let zstd = args.tar_zstd || args.out.ends_with(".tar.zst"); + let sink = Arc::new(TarSink::new(&args.out, zstd)?); + planner.run_parallel(&client, sink, &factory).await? + } else if let Some(url) = object_store_url(&args.out) { + let sink = Arc::new(pgstac::export::sink::ObjectStoreSink::from_url_opts( + &url, + std::iter::empty(), + )?); + planner.run_parallel(&client, sink, &factory).await? + } else { + let sink = Arc::new(DirSink::new(&args.out)?); + planner.run_parallel(&client, sink, &factory).await? + } + } else if args.out == "-" { + planner.run(&client, &StdoutSink).await? + } else if is_tar_path(&args.out) { + let zstd = args.tar_zstd || args.out.ends_with(".tar.zst"); + let sink = TarSink::new(&args.out, zstd)?; + planner.run(&client, &sink).await? + } else if let Some(url) = object_store_url(&args.out) { + let sink = pgstac::export::sink::ObjectStoreSink::from_url_opts(&url, std::iter::empty())?; + planner.run(&client, &sink).await? + } else { + let sink = DirSink::new(&args.out)?; + planner.run(&client, &sink).await? + }; + + eprintln!( + "dump complete: {} collections, {} partition files, {} items{}", + report.collections, + report.partitions, + report.items, + if report.skipped > 0 { + format!(", {} skipped", report.skipped) + } else { + String::new() + } + ); + Ok(()) +} + +/// Opens the output sink for a search export: a file, or stdout for `-`. +fn make_sink(out: &str) -> Result, Box> { + Ok(if out == "-" { + Box::new(std::io::BufWriter::new(std::io::stdout().lock())) + } else { + Box::new(std::io::BufWriter::new(std::fs::File::create(out)?)) + }) +} + +/// Search-driven export (0.10 only): stream a search / CQL2 result as NDJSON, a single ItemCollection +/// page, or stac-geoparquet. Every page is built in Rust from `search_plan` (Rust band-stepping + +/// keyset minting); the SQL `search()` function is never called. +async fn run_search(args: SearchArgs, dsn: Option) -> Result<(), Box> { + use stac::api::{ItemsClient, Search}; + use std::io::Write as _; + + let pool = PgstacPool::connect(connect_config(dsn.as_deref())).await?; + + // Build the Search from CLI args. + let mut search = Search { + collections: args.collections.clone(), + ids: args.ids.clone(), + ..Default::default() + }; + if let Some(b) = &args.bbox + && b.len() == 4 + { + search = search.bbox(stac::Bbox::new(b[0], b[1], b[2], b[3])); + } + if let Some(dt) = &args.datetime { + search = search.datetime(dt); + } + if let Some(limit) = args.limit { + search = search.limit(limit); + } + if let Some(filter) = &args.filter { + let filter_json: serde_json::Value = serde_json::from_str(filter)?; + search.items.filter = Some(stac::api::Filter::Cql2Json(serde_json::from_value( + filter_json, + )?)); + } + // A full JSON body overrides the typed args (whole-spec searches: sortby, fields, intersects, query, conf). + if let Some(body) = &args.search_json { + search = serde_json::from_str(body)?; + } + + let max_items = args.max_items.map(|m| m as i64); + match args.format { + Format::NdJson => { + // Flat-memory streaming: rows arrive from a server-side portal and are written one at a + // time with the shared fragment merged at serialize time — never a buffered page. + let mut out = make_sink(&args.out)?; + let body = serde_json::to_value(&search)?; + let n = pool + .stream_ndjson(body, args.token.as_deref(), max_items, &mut out) + .await?; + out.flush()?; + eprintln!("search: streamed {n} item(s) as NDJSON"); + } + Format::Json(_) => { + // A single keyset page assembled in Rust via ItemsClient::search (search_plan + Rust + // band-stepping): features, next/prev tokens, links, and context all minted client-side. + let mut out = make_sink(&args.out)?; + if let Some(token) = &args.token { + let _ = search + .additional_fields + .insert("token".into(), serde_json::Value::String(token.clone())); + } + let item_collection = ItemsClient::search(&pool, search).await?; + let mut fc = serde_json::to_value(&item_collection)?; + // ItemCollection serializes next/prev only into links; also surface them as top-level + // fields (as before) so token-paging callers need not re-parse links. + if let serde_json::Value::Object(map) = &mut fc { + if let Some(t) = item_collection.next.as_ref().and_then(|m| m.get("token")) { + let _ = map.insert("next".into(), t.clone()); + } + if let Some(t) = item_collection.prev.as_ref().and_then(|m| m.get("token")) { + let _ = map.insert("prev".into(), t.clone()); + } + } + serde_json::to_writer(&mut out, &fc)?; + out.flush()?; + let n = item_collection + .number_returned + .unwrap_or(item_collection.items.len() as u64); + eprintln!("search: wrote 1 ItemCollection page ({n} item(s))"); + } + Format::Geoparquet(writer_options) => { + // Flat-memory streaming: items stream from the search portal into row-group batches, so the + // item working set is one batch (0.10 registry schema), not the whole result. + let body = serde_json::to_value(&search)?; + // --format geoparquet[] carries the compression; --row-group-size stays a flag. + let compression = writer_options.compression.unwrap_or_else(default_compression); + let row_group_size = args.row_group_size; + let n = if args.out == "-" { + // The geoparquet writer needs a Send sink and stdout's lock is not Send: encode into a + // buffer, then write it out. + let mut buf: Vec = Vec::new(); + let n = pool + .stream_geoparquet(body, max_items, compression, row_group_size, &mut buf) + .await?; + std::io::stdout().write_all(&buf)?; + n + } else { + let file = std::fs::File::create(&args.out)?; + pool.stream_geoparquet(body, max_items, compression, row_group_size, file) + .await? + }; + eprintln!("search: wrote geoparquet ({n} item(s))"); + } + } + Ok(()) +} + +/// Load STAC items + collections through the Rust loader. Collections are +/// created first (items reference them); then each item source is read into +/// memory, chunked into `batch_size` batches, and loaded via +/// [`PgstacPool::create_items`] (`load_items`) — so dehydration, fragment +/// splitting, and the binary COPY all run in Rust. `concurrency` batches run in +/// parallel, each on its own pooled connection. +async fn run_load(args: LoadArgs, dsn: Option) -> Result<(), Box> { + let (collection_files, item_sources) = classify_inputs(&args.inputs)?; + if collection_files.is_empty() && item_sources.is_empty() { + return Err("no loadable inputs found (.parquet/.geoparquet/.ndjson/.json)".into()); + } + + let config = connect_config(dsn.as_deref()); + let pool_size = args.pool_size.unwrap_or_else(|| args.concurrency.max(4)); + let pool = PgstacPool::connect_with( + config, + PoolOptions { + max_size: pool_size, + ..Default::default() + }, + ) + .await?; + let policy = args.policy; + if args.skip_unchanged && matches!(policy, ConflictPolicy::Error) { + return Err( + "--skip-unchanged is incompatible with --policy error (error must fail on an existing id, \ + not skip it)" + .into(), + ); + } + + for cf in &collection_files { + let raw: Value = serde_json::from_slice(&std::fs::read(cf)?)?; + // A `pgstac dump` writes `{"stac": , "pgstac": {partition_trunc, fragment_config, + // private}}`; a plain collection file is the bare STAC Collection. Unwrap the dump envelope and + // apply its partition_trunc so a restore recreates the same partition layout (fragment_config is + // re-derived by create_collection from item_assets). + let (collection, partition_trunc) = match raw.get("stac") { + Some(stac) => ( + stac.clone(), + raw.get("pgstac") + .and_then(|p| p.get("partition_trunc")) + .and_then(|t| t.as_str()) + .map(str::to_string), + ), + None => (raw, None), + }; + let id = collection + .get("id") + .and_then(|v| v.as_str()) + .unwrap_or("?") + .to_string(); + // Upsert (not bare create) so load/restore are idempotent: a re-run updates the collection instead + // of failing on the duplicate id — which is what makes `pgstac restore` resumable after an interrupt. + // partition_trunc is applied on insert and preserved (same value) on a re-run. + pool.upsert_collection(&collection, partition_trunc.as_deref()) + .await?; + eprintln!("collection: {id}"); + } + + let start = Instant::now(); + let mut total = 0u64; + let mut remaining = args.limit.unwrap_or(usize::MAX); + for src in &item_sources { + if remaining == 0 { + break; + } + load_source_streaming( + &pool, + src, + &StreamingOptions { + batch_size: args.batch_size, + concurrency: args.concurrency, + policy, + skip_unchanged: args.skip_unchanged, + }, + &mut remaining, + &mut total, + ) + .await?; + eprintln!(" {} done", src.display()); + } + let secs = start.elapsed().as_secs_f64(); + eprintln!( + "load complete: {total} items in {secs:.2}s ({:.0} items/s, streaming decode+load)", + total as f64 / secs.max(1e-9), + ); + pool.close(); + Ok(()) +} + +/// Restore a `pgstac dump` directory through the Rust loader. The dump layout +/// (per-collection `collection.json` + partition geoparquet) is exactly what +/// [`run_load`] consumes, so this is `load` over the dump directory with an +/// upsert policy. +async fn run_restore(args: RestoreArgs, dsn: Option) -> Result<(), Box> { + if !args.src.is_dir() { + return Err(format!( + "restore source must be a dump directory: {} (extract a .tar dump first)", + args.src.display() + ) + .into()); + } + run_load( + LoadArgs { + inputs: vec![args.src], + batch_size: args.batch_size, + concurrency: args.concurrency, + policy: ConflictPolicy::Upsert, + pool_size: args.pool_size, + limit: None, + skip_unchanged: false, + }, + dsn, + ) + .await +} + +/// Run the async partition-stats maintenance: +/// recompute exact bounds + row counts for partitions ingest left dirty. Safe + +/// optional — a generous (un-tightened) envelope only over-includes a partition +/// in search, never loses rows; tightening restores tight pruning + honest +/// counts. Operators usually schedule this off-hours (pg_cron). +async fn run_maintain( + args: MaintainArgs, + dsn: Option, +) -> Result<(), Box> { + let dsn = dsn.unwrap_or_else(|| DEFAULT_DSN.to_string()); + let (client, connection) = tokio_postgres::connect(&dsn, NoTls).await?; + tokio::spawn(async move { + if let Err(e) = connection.await { + eprintln!("connection error: {e}"); + } + }); + client + .batch_execute("SET search_path = pgstac, public;") + .await?; + let row = client + .query_one( + "SELECT tighten_dirty_partition_stats($1::int)", + &[&args.limit], + ) + .await?; + let tightened: i32 = row.get(0); + eprintln!("maintain: tightened {tightened} dirty partition(s)"); + Ok(()) +} + +/// Delete an item (with `--item`) or a whole collection (without). Both go through the pool's +/// SECURITY DEFINER delete functions. +async fn run_delete(args: DeleteArgs, dsn: Option) -> Result<(), Box> { + let target = match &args.item { + Some(item) => format!("item {}/{}", args.collection, item), + None => format!("collection {} and all its items", args.collection), + }; + if !confirm_delete(&target, args.yes)? { + eprintln!("aborted"); + return Ok(()); + } + let pool = PgstacPool::connect(connect_config(dsn.as_deref())).await?; + match &args.item { + Some(item) => { + pool.delete_item(&args.collection, item).await?; + eprintln!("deleted item {}/{}", args.collection, item); + } + None => { + pool.delete_collection(&args.collection).await?; + eprintln!("deleted collection {}", args.collection); + } + } + pool.close(); + Ok(()) +} + +/// Expand the inputs (recursing into directories) and split into collection +/// files (loaded first) and item sources, each sorted for a deterministic order. +fn classify_inputs( + inputs: &[PathBuf], +) -> Result<(Vec, Vec), Box> { + let mut files = Vec::new(); + for input in inputs { + collect_files(input, &mut files)?; + } + let mut collections = Vec::new(); + let mut items = Vec::new(); + for f in files { + match classify_file(&f)? { + FileRole::Collection => collections.push(f), + FileRole::Items => items.push(f), + FileRole::Skip => {} // a dump's manifest/settings/queryables JSON, etc. + } + } + collections.sort(); + items.sort(); + Ok((collections, items)) +} + +/// How an input file is loaded. +enum FileRole { + Collection, + Items, + Skip, +} + +/// Recursively collect loadable files (by extension) under `path`. +fn collect_files(path: &Path, out: &mut Vec) -> std::io::Result<()> { + if path.is_dir() { + let mut paths: Vec = std::fs::read_dir(path)? + .collect::, _>>()? + .into_iter() + .map(|e| e.path()) + .collect(); + paths.sort(); + for p in paths { + collect_files(&p, out)?; + } + } else if path.is_file() { + let ext = path.extension().and_then(|e| e.to_str()).unwrap_or(""); + if matches!(ext, "parquet" | "geoparquet" | "ndjson" | "json") { + out.push(path.to_path_buf()); + } + } + Ok(()) +} + +/// Classify a file. Parquet/NDJSON are always item sources. A `.json` is routed by its STAC `type` +/// (with a `collection.json` name fast-path): `Collection` → a collection, `Feature`/`FeatureCollection` +/// → items, anything else → skipped — so a dump's `manifest.json` / `settings.json` / `queryables.json` +/// are ignored rather than mis-loaded as items. +fn classify_file(path: &Path) -> Result> { + let ext = path.extension().and_then(|e| e.to_str()).unwrap_or(""); + if matches!(ext, "parquet" | "geoparquet" | "ndjson") { + return Ok(FileRole::Items); + } + if path.file_name().and_then(|n| n.to_str()) == Some("collection.json") { + return Ok(FileRole::Collection); + } + let value: Value = serde_json::from_slice(&std::fs::read(path)?)?; + Ok(match value.get("type").and_then(|t| t.as_str()) { + Some("Collection") => FileRole::Collection, + Some("Feature") | Some("FeatureCollection") => FileRole::Items, + _ => FileRole::Skip, + }) +} + +/// A streaming batch iterator over an item source. Each yielded `Vec` is one chunk of item JSON values: +/// a stac-geoparquet **row group** (via [`stac::geoparquet::from_reader_iter`], so a multi-GB file is +/// never held in memory), a slice of NDJSON lines, or a whole `.json` file (item / ItemCollection). +type ItemBatches = Box, Box>>>; + +fn item_batches(path: &Path) -> Result> { + let ext = path.extension().and_then(|e| e.to_str()).unwrap_or(""); + match ext { + "parquet" | "geoparquet" => { + let file = std::fs::File::open(path)?; + let iter = stac::geoparquet::from_reader_iter(file)?.map(|batch| { + batch? + .into_iter() + .map(|item| serde_json::to_value(item).map_err(Into::into)) + .collect::, Box>>() + }); + Ok(Box::new(iter)) + } + "ndjson" => { + use std::io::BufRead; + const NDJSON_BATCH: usize = 10_000; + let mut lines = std::io::BufReader::new(std::fs::File::open(path)?).lines(); + let iter = std::iter::from_fn(move || { + let mut batch = Vec::with_capacity(NDJSON_BATCH); + for line in lines.by_ref() { + let line = match line { + Ok(line) => line, + Err(e) => return Some(Err(e.into())), + }; + if line.trim().is_empty() { + continue; + } + match serde_json::from_str(&line) { + Ok(value) => batch.push(value), + Err(e) => return Some(Err(e.into())), + } + if batch.len() >= NDJSON_BATCH { + return Some(Ok(batch)); + } + } + (!batch.is_empty()).then_some(Ok(batch)) + }); + Ok(Box::new(iter)) + } + "json" => { + let value: Value = serde_json::from_slice(&std::fs::read(path)?)?; + let items = match value.get("features").and_then(|f| f.as_array()) { + Some(features) => features.clone(), + None => vec![value], + }; + Ok(Box::new(std::iter::once(Ok(items)))) + } + _ => Err(format!("unsupported input extension: {}", path.display()).into()), + } +} + +/// Parallel ndjson decode: split the file into `n` byte ranges and spawn one decoder thread per range. +/// Each thread parses only the lines whose START offset falls in its range — boundary-aligned (a thread +/// discards the partial line straddling its start, UNLESS its start sits exactly on a line boundary, and +/// fully consumes the line straddling its end) so no line is dropped or double-decoded. Small files use a +/// single range (no alignment). All threads feed the shared `tx`. This parallelizes the serde parse, the +/// single-thread decode ceiling. Order is not preserved (fine for load). +fn spawn_ndjson_range_decoders( + path: &Path, + n: usize, + tx: tokio::sync::mpsc::Sender, usize), String>>, +) -> std::io::Result>> { + use std::io::{BufRead, BufReader, Read, Seek, SeekFrom}; + const BATCH: usize = 10_000; + // Also flush at a byte budget, not just an item count: the channel holds up to `concurrency` batches and + // each decoder builds one, so the in-flight set is bounded by (decoders + channel slots) * BYTE_CAP + // regardless of item size OR file size. Without this, big STAC items (~60 KB each as a parsed Value) at + // 10k/batch across an 8-slot channel sustained ~12 GB on a full multi-GB load and OOM-killed it. + const BYTE_CAP: usize = 8 * 1024 * 1024; + let size = std::fs::metadata(path)?.len(); + let n = (n.max(1)) as u64; + // Sharding a small file buys nothing (a sub-MB decode is instant) and only invites the line-boundary + // edge cases below, so cap to ~one decoder per MB; tiny files fall back to a single range. + let n = n.min((size / (1u64 << 20)).max(1)); + let mut handles = Vec::with_capacity(n as usize); + for i in 0..n { + let start = i * size / n; + let end = (i + 1) * size / n; + let path = path.to_path_buf(); + let tx = tx.clone(); + handles.push(std::thread::spawn(move || { + let f = match std::fs::File::open(&path) { + Ok(f) => f, + Err(e) => { + let _ = tx.blocking_send(Err(e.to_string())); + return; + } + }; + let mut reader = BufReader::with_capacity(256 * 1024, f); + let mut consumed = start; + if i > 0 { + // Align to a line boundary. The line straddling `start` belongs to the previous range, so + // discard it — UNLESS `start` already sits on a line boundary (the byte before it is '\n'), + // in which case the line AT `start` is ours and skipping it would silently drop a record (a + // shard boundary landing exactly on a newline). + if let Err(e) = reader.seek(SeekFrom::Start(start - 1)) { + let _ = tx.blocking_send(Err(e.to_string())); + return; + } + let mut prev = [0u8; 1]; + let on_boundary = match reader.read_exact(&mut prev) { + Ok(()) => prev[0] == b'\n', + Err(e) => { + let _ = tx.blocking_send(Err(e.to_string())); + return; + } + }; + if !on_boundary { + let mut junk = Vec::new(); + match reader.read_until(b'\n', &mut junk) { + Ok(nr) => consumed += nr as u64, + Err(e) => { + let _ = tx.blocking_send(Err(e.to_string())); + return; + } + } + } + } + let mut batch = Vec::with_capacity(BATCH); + let mut batch_bytes = 0usize; + let mut line = String::new(); + loop { + if consumed >= end { + break; // the next range owns lines starting at/after `end` + } + line.clear(); + let nr = match reader.read_line(&mut line) { + Ok(nr) => nr, + Err(e) => { + let _ = tx.blocking_send(Err(e.to_string())); + return; + } + }; + if nr == 0 { + break; // EOF + } + consumed += nr as u64; + let trimmed = line.trim(); + if !trimmed.is_empty() { + match serde_json::from_str::(trimmed) { + Ok(v) => batch.push(v), + Err(e) => { + let _ = tx.blocking_send(Err(e.to_string())); + return; + } + } + batch_bytes += trimmed.len(); + if batch.len() >= BATCH || batch_bytes >= BYTE_CAP { + let sent = std::mem::replace(&mut batch, Vec::with_capacity(BATCH)); + if tx.blocking_send(Ok((sent, batch_bytes))).is_err() { + return; // receiver gone + } + batch_bytes = 0; + } + } + } + if !batch.is_empty() { + let _ = tx.blocking_send(Ok((batch, batch_bytes))); + } + })); + } + Ok(handles) +} + +/// Stream one source into the loader. Decoding (the blocking parquet/geoarrow work) runs on a dedicated +/// thread feeding a **bounded channel**; the async side re-chunks to `batch_size` and keeps up to +/// `concurrency` COPY batches in flight (each on its own pooled connection). So CPU decode overlaps with +/// the DB loads, and the bounded channel + bounded in-flight set keep memory flat regardless of file +/// size (decode blocks when the loaders fall behind). Stops after `remaining` items; adds rows to `total`. +struct StreamingOptions { + batch_size: usize, + concurrency: usize, + policy: ConflictPolicy, + skip_unchanged: bool, +} + +async fn load_source_streaming( + pool: &PgstacPool, + src: &Path, + opts: &StreamingOptions, + remaining: &mut usize, + total: &mut u64, +) -> Result<(), Box> { + let batch_size = opts.batch_size.max(1); + let concurrency = opts.concurrency.max(1); + let policy = opts.policy; + let skip_unchanged = opts.skip_unchanged; + // Cap the decoded-item bytes held per in-flight load task. Large STAC items (e.g. landsat with many asset + // bands) cost ~6x their text size once parsed + dehydrated + COPY-encoded, so item-count batching alone + // (batch_size * concurrency) can balloon to tens of GB and OOM a multi-GB ndjson load. Tunable via + // PGSTAC_LOAD_BYTES (MB); only ndjson decoders report byte sizes, so other formats batch by count. + let load_bytes = std::env::var("PGSTAC_LOAD_BYTES") + .ok() + .and_then(|s| s.parse::().ok()) + .unwrap_or(64) + .saturating_mul(1024 * 1024); + + // Decoder thread(s): build + drain the (non-Send) batch iterator here, sending only the Send + // `Vec` batches across. The bounded channel applies backpressure (decode pauses when loaders + // are behind). PGSTAC_DECODE_THREADS>1 shards an ndjson file across N parallel range decoders (by + // byte range) or a stac-geoparquet file across N decoders (by row group) to lift the single-thread + // decode ceiling; other formats (plain .json) keep the single decoder. + let (tx, mut rx) = + tokio::sync::mpsc::channel::, usize), String>>(concurrency.max(2)); + let path = src.to_path_buf(); + let decode_threads = std::env::var("PGSTAC_DECODE_THREADS") + .ok() + .and_then(|s| s.parse::().ok()) + .unwrap_or_else(default_ingest_parallelism) + .max(1); + let ext = path.extension().and_then(|e| e.to_str()); + let is_ndjson = ext == Some("ndjson"); + let is_parquet = matches!(ext, Some("parquet") | Some("geoparquet")); + let decoders: Vec> = if is_ndjson && decode_threads > 1 { + spawn_ndjson_range_decoders(&path, decode_threads, tx)? + } else if is_parquet && decode_threads > 1 { + // Parquet decode was the single-thread ingest bottleneck (much slower than ndjson). Shard the + // file's row groups across N decoder threads; each reports per-batch byte estimates so the + // load loop's PGSTAC_LOAD_BYTES budget applies to parquet too. + pgstac::parquet_decode::spawn_parquet_decoders(&path, decode_threads, tx)? + } else { + vec![std::thread::spawn(move || match item_batches(&path) { + Ok(batches) => { + for batch in batches { + // Non-ndjson decoders don't track byte size; send 0 so the load loop batches them by + // item count only (the byte budget is an ndjson-streaming safeguard). + let msg = batch.map(|b| (b, 0usize)).map_err(|e| e.to_string()); + if tx.blocking_send(msg).is_err() { + break; // receiver gone (early stop / error) + } + } + } + Err(e) => { + let _ = tx.blocking_send(Err(e.to_string())); + } + })] + }; + + // Single ingest path: create_items (staging + flush), which resolves every conflict policy (a raw + // COPY can't). + + let mut set: tokio::task::JoinSet> = tokio::task::JoinSet::new(); + let mut pending: Vec = Vec::with_capacity(batch_size); + let mut pending_bytes = 0usize; + loop { + let msg = rx.recv().await; + let Some(msg) = msg else { break }; + if *remaining == 0 { + break; + } + let (mut batch, bytes) = msg.map_err(|e| -> Box { e.into() })?; + if batch.len() > *remaining { + batch.truncate(*remaining); + } + *remaining -= batch.len(); + pending.append(&mut batch); + pending_bytes += bytes; + // Flush a task when pending reaches batch_size items OR the byte budget — the byte budget bounds the + // in-flight working set (concurrency * chunk) for large items, which otherwise OOMs a big ndjson load. + if pending.len() >= batch_size || (pending_bytes >= load_bytes && !pending.is_empty()) { + while set.len() >= concurrency { + let r = set.join_next().await.expect("set is non-empty"); + *total += r??; + } + let chunk = std::mem::replace(&mut pending, Vec::with_capacity(batch_size)); + pending_bytes = 0; + let pool = pool.clone(); + set.spawn(async move { + if skip_unchanged { + pool.upsert_items_precheck(chunk, policy) + .await + .map(|(_skipped, loaded)| loaded) + } else { + pool.create_items(chunk, policy).await + } + }); + } + } + if !pending.is_empty() { + let pool = pool.clone(); + set.spawn(async move { + if skip_unchanged { + pool.upsert_items_precheck(pending, policy) + .await + .map(|(_skipped, loaded)| loaded) + } else { + pool.create_items(pending, policy).await + } + }); + } + while let Some(res) = set.join_next().await { + *total += res??; + } + // Unblock the decoder(s) if we stopped early (channel full), then join them. + rx.close(); + for decoder in decoders { + let _ = decoder.join(); + } + Ok(()) +} + +/// Dry-run: report the plan (collections, partitions, est. items) +/// without writing data. +async fn dry_run( + client: &tokio_postgres::Client, + config: &DumpConfig, +) -> Result<(), Box> { + let plans = pgstac::export::plan::plan_collections( + client, + config.collection_ids.as_deref(), + &config.prefilter, + ) + .await?; + let mut total_parts = 0usize; + let mut est_items = 0f64; + println!("dry-run plan:"); + for plan in &plans { + let parts = plan.partitions.len(); + total_parts += parts; + let coll_est: f64 = plan.partitions.iter().map(|p| p.reltuples as f64).sum(); + est_items += coll_est; + println!( + " {} ({:?}): {} partitions, ~{} items", + plan.id, plan.partition_trunc, parts, coll_est as u64 + ); + } + println!( + "total: {} collections, {} partitions, ~{} items (estimated from reltuples)", + plans.len(), + total_parts, + est_items as u64 + ); + Ok(()) +} + +fn is_tar_path(out: &str) -> bool { + out.ends_with(".tar") || out.ends_with(".tar.zst") +} + +/// Returns an object_store URL if `out` looks like a remote/object-store URL. +fn object_store_url(out: &str) -> Option { + for scheme in ["s3://", "gs://", "az://", "azure://", "file://"] { + if out.starts_with(scheme) { + return Some(out.to_string()); + } + } + None +} + +/// Loads + verifies a prior `_checkpoint.json` from a directory dump for resume +///. Returns the verified completed entries keyed by relative path: a file +/// is included only if it exists on disk AND matches its recorded SHA-256; a +/// missing or mismatched file is left out (it is redone whole). Absent checkpoint +/// => empty map (fresh dump). A finished dump (manifest present, no checkpoint) +/// also yields an empty map, so a re-run is idempotent. +fn load_resume_set( + dir: &str, +) -> std::collections::HashMap { + use pgstac::export::manifest::{Checkpoint, sha256_hex}; + let root = std::path::Path::new(dir); + let cp_path = root.join("_checkpoint.json"); + let Ok(bytes) = std::fs::read(&cp_path) else { + return Default::default(); + }; + let Ok(cp): Result = serde_json::from_slice(&bytes) else { + eprintln!("warning: ignoring unreadable _checkpoint.json"); + return Default::default(); + }; + let mut verified = std::collections::HashMap::new(); + let mut skipped = 0usize; + for entry in &cp.completed { + let p = root.join(&entry.file); + match std::fs::read(&p) { + Ok(b) if sha256_hex(&b) == entry.sha256 => { + let _ = verified.insert(entry.file.clone(), entry.clone()); + } + _ => skipped += 1, + } + } + if !verified.is_empty() || skipped > 0 { + eprintln!( + "resume: {} completed partition file(s) verified and will be skipped{}", + verified.len(), + if skipped > 0 { + format!("; {skipped} missing/mismatched will be redone") + } else { + String::new() + } + ); + } + verified +} diff --git a/src/pgstac-rs/src/client.rs b/src/pgstac-rs/src/client.rs deleted file mode 100644 index 6286216c..00000000 --- a/src/pgstac-rs/src/client.rs +++ /dev/null @@ -1,119 +0,0 @@ -use crate::{Error, Pgstac}; -use serde_json::Map; -use stac::api::{CollectionsClient, ItemCollection, ItemsClient, Search, TransactionClient}; -use stac::{Collection, Item}; -use std::ops::{Deref, DerefMut}; -use tokio_postgres::GenericClient; - -/// A newtype wrapper around a [`GenericClient`] that implements the STAC -/// client traits ([`ItemsClient`], [`CollectionsClient`], and -/// [`TransactionClient`]). -/// -/// This wrapper allows any [`tokio_postgres`] client or transaction to be used -/// with the STAC client trait abstractions. Use [`Deref`] to access the -/// underlying [`Pgstac`] methods directly. -/// -/// # Examples -/// -/// ```no_run -/// use pgstac::Client; -/// use stac::api::ItemsClient; -/// use tokio_postgres::NoTls; -/// -/// # tokio_test::block_on(async { -/// let (pg_client, connection) = tokio_postgres::connect( -/// "postgresql://username:password@localhost:5432/postgis", -/// NoTls, -/// ).await.unwrap(); -/// tokio::spawn(async move { -/// if let Err(e) = connection.await { -/// eprintln!("connection error: {}", e); -/// } -/// }); -/// let client = Client(pg_client); -/// let item_collection = client.search(Default::default()).await.unwrap(); -/// # }) -/// ``` -#[derive(Debug)] -pub struct Client(pub C); - -impl Deref for Client { - type Target = C; - - fn deref(&self) -> &C { - &self.0 - } -} - -impl DerefMut for Client { - fn deref_mut(&mut self) -> &mut C { - &mut self.0 - } -} - -impl ItemsClient for Client { - type Error = Error; - - async fn search(&self, search: Search) -> Result { - let page = Pgstac::search(&self.0, search).await?; - let next_token = page.next_token(); - let prev_token = page.prev_token(); - let mut item_collection = ItemCollection::new(page.features)?; - if let Some(next_token) = next_token { - let mut next = Map::new(); - let _ = next.insert("token".into(), next_token.into()); - item_collection.next = Some(next); - } - if let Some(prev_token) = prev_token { - let mut prev = Map::new(); - let _ = prev.insert("token".into(), prev_token.into()); - item_collection.prev = Some(prev); - } - item_collection.context = page.context; - Ok(item_collection) - } - - async fn item(&self, collection_id: &str, item_id: &str) -> Result, Error> { - let value = Pgstac::item(&self.0, item_id, Some(collection_id)).await?; - value - .map(serde_json::from_value) - .transpose() - .map_err(Error::from) - } -} - -impl CollectionsClient for Client { - type Error = Error; - - async fn collections(&self) -> Result, Error> { - let values = Pgstac::collections(&self.0).await?; - values - .into_iter() - .map(|v| serde_json::from_value(v).map_err(Error::from)) - .collect() - } - - async fn collection(&self, id: &str) -> Result, Error> { - let value = Pgstac::collection(&self.0, id).await?; - value - .map(serde_json::from_value) - .transpose() - .map_err(Error::from) - } -} - -impl TransactionClient for Client { - type Error = Error; - - async fn add_collection(&mut self, collection: Collection) -> Result<(), Error> { - Pgstac::add_collection(&self.0, collection).await - } - - async fn add_item(&mut self, item: Item) -> Result<(), Error> { - Pgstac::add_item(&self.0, item).await - } - - async fn add_items(&mut self, items: Vec) -> Result<(), Error> { - Pgstac::add_items(&self.0, &items).await - } -} diff --git a/src/pgstac-rs/src/db/call.rs b/src/pgstac-rs/src/db/call.rs new file mode 100644 index 00000000..fb66c72d --- /dev/null +++ b/src/pgstac-rs/src/db/call.rs @@ -0,0 +1,53 @@ +//! Low-level helpers for invoking pgstac SQL functions over any connection. + +use crate::{Error, Result}; +use tokio_postgres::types::ToSql; +use tokio_postgres::{GenericClient, Row}; + +/// Calls `pgstac.()` and returns the single result row. +async fn row( + client: &impl GenericClient, + function: &str, + params: &[&(dyn ToSql + Sync)], +) -> std::result::Result { + let param_string = (0..params.len()) + .map(|i| format!("${}", i + 1)) + .collect::>() + .join(", "); + let query = format!("SELECT * from pgstac.{function}({param_string})"); + client.query_one(&query, params).await +} + +/// Calls a pgstac function returning a single `text` column named for the function. +pub(crate) async fn string( + client: &impl GenericClient, + function: &str, + params: &[&(dyn ToSql + Sync)], +) -> Result { + row(client, function, params) + .await? + .try_get(function) + .map_err(Error::from) +} + +/// Calls a pgstac function returning a single `boolean` column named for the function. +pub(crate) async fn boolean( + client: &impl GenericClient, + function: &str, + params: &[&(dyn ToSql + Sync)], +) -> Result { + row(client, function, params) + .await? + .try_get(function) + .map_err(Error::from) +} + +/// Calls a pgstac function purely for its side effects, discarding the result. +pub(crate) async fn void( + client: &impl GenericClient, + function: &str, + params: &[&(dyn ToSql + Sync)], +) -> Result<()> { + let _ = row(client, function, params).await?; + Ok(()) +} diff --git a/src/pgstac-rs/src/db/client.rs b/src/pgstac-rs/src/db/client.rs new file mode 100644 index 00000000..88d189d9 --- /dev/null +++ b/src/pgstac-rs/src/db/client.rs @@ -0,0 +1,295 @@ +use crate::dehydrate::DehydrateSchema; +use crate::ingest::{ConflictPolicy, load_items}; +use crate::db::call; +use crate::read::collections; +use crate::search::{SearchPage, search_page_with}; +use crate::source::CachedHydration; +use crate::Error; +use serde::Serialize; +use serde_json::Value; +use stac::api::{CollectionsClient, ItemCollection, ItemsClient, Search, TransactionClient}; +use stac::{Collection, Item}; +use std::ops::{Deref, DerefMut}; +use std::sync::Arc; +use tokio::sync::OnceCell; + +/// Bridges the concrete connection types the engine runs on to a `&tokio_postgres::Client`. Implemented for +/// a bare [`tokio_postgres::Client`] and (with the `pool` feature) a pooled `deadpool_postgres::Client`, so a +/// single [`Client`] serves both direct and pooled connections. +pub trait PgConn { + /// The underlying tokio-postgres client, for reads. + fn pg(&self) -> &tokio_postgres::Client; + /// The underlying tokio-postgres client, for the loader's `&mut` transaction. + fn pg_mut(&mut self) -> &mut tokio_postgres::Client; +} + +impl PgConn for tokio_postgres::Client { + fn pg(&self) -> &tokio_postgres::Client { + self + } + fn pg_mut(&mut self) -> &mut tokio_postgres::Client { + self + } +} + +#[cfg(feature = "pool")] +impl PgConn for deadpool_postgres::Client { + fn pg(&self) -> &tokio_postgres::Client { + self + } + fn pg_mut(&mut self) -> &mut tokio_postgres::Client { + self + } +} + +/// A wrapper around a connection (a bare [`tokio_postgres::Client`] or, with the `pool` feature, a pooled +/// `deadpool_postgres::Client`) that implements the STAC client traits ([`ItemsClient`], +/// [`CollectionsClient`], and [`TransactionClient`]) on the Rust pgstac engine. +/// +/// `Client` caches the database's hydration invariants (storage model + promoted schema) on first use and +/// reuses them across every read on this connection, saving the catalog round trips a fresh detection would +/// cost per call. The pgstac-specific operations that are not STAC CRUD — version, settings, and collection +/// maintenance — are inherent methods; [`Deref`] reaches the underlying [`tokio_postgres::Client`] for raw +/// SQL. +/// +/// # Examples +/// +/// ```no_run +/// use pgstac::Client; +/// use stac::api::ItemsClient; +/// use tokio_postgres::NoTls; +/// +/// # tokio_test::block_on(async { +/// let (pg_client, connection) = tokio_postgres::connect( +/// "postgresql://username:password@localhost:5432/postgis", +/// NoTls, +/// ).await.unwrap(); +/// tokio::spawn(async move { +/// if let Err(e) = connection.await { +/// eprintln!("connection error: {}", e); +/// } +/// }); +/// let client = Client::new(pg_client); +/// let item_collection = client.search(Default::default()).await.unwrap(); +/// # }) +/// ``` +#[derive(Debug)] +pub struct Client { + client: C, + hydration: Arc>, +} + +impl Client { + /// Wraps a connection, caching the hydration invariants across calls. + pub fn new(client: C) -> Self { + Client::with_cache(client, Arc::new(OnceCell::new())) + } + + /// Wraps a connection with a hydration cache shared with other clients (every client a + /// [`crate::PgstacPool`] hands out shares the pool's cache, so detection happens once pool-wide). + pub(crate) fn with_cache(client: C, hydration: Arc>) -> Self { + Client { client, hydration } + } +} + +impl Deref for Client { + type Target = tokio_postgres::Client; + + fn deref(&self) -> &tokio_postgres::Client { + self.client.pg() + } +} + +impl DerefMut for Client { + fn deref_mut(&mut self) -> &mut tokio_postgres::Client { + self.client.pg_mut() + } +} + +impl Client { + /// The hydration invariants for this connection, detected once and cached. + async fn cached_hydration(&self) -> Result { + self.hydration + .get_or_try_init(|| CachedHydration::detect(self.client.pg())) + .await + .cloned() + } + + /// Returns the **pgstac** version. + pub async fn pgstac_version(&self) -> Result { + call::string(self.client.pg(), "get_version", &[]).await + } + + /// Returns whether the **pgstac** database is readonly. + pub async fn readonly(&self) -> Result { + call::boolean(self.client.pg(), "readonly", &[]).await + } + + /// Returns the value of the `context` **pgstac** setting (defaults to "off"). + pub async fn context(&self) -> Result { + call::string(self.client.pg(), "get_setting", &[&"context"]) + .await + .map(|value| value == "on") + } + + /// Sets the value of a **pgstac** setting. + pub async fn set_pgstac_setting(&self, key: &str, value: &str) -> Result<(), Error> { + let _ = self + .client + .pg() + .execute( + "INSERT INTO pgstac_settings (name, value) VALUES ($1, $2) ON CONFLICT ON CONSTRAINT pgstac_settings_pkey DO UPDATE SET value = excluded.value;", + &[&key, &value], + ) + .await?; + Ok(()) + } + + /// Adds or updates a collection. + pub async fn upsert_collection(&self, collection: T) -> Result<(), Error> { + let collection = serde_json::to_value(collection)?; + call::void(self.client.pg(), "upsert_collection", &[&collection]).await + } + + /// Updates a collection. + pub async fn update_collection(&self, collection: T) -> Result<(), Error> { + let collection = serde_json::to_value(collection)?; + call::void(self.client.pg(), "update_collection", &[&collection]).await + } + + /// Updates all collection extents. + pub async fn update_collection_extents(&self) -> Result<(), Error> { + call::void(self.client.pg(), "update_collection_extents", &[]).await + } + + /// Deletes a collection. + pub async fn delete_collection(&self, id: &str) -> Result<(), Error> { + call::void(self.client.pg(), "delete_collection", &[&id]).await + } + + /// Deletes an item. + pub async fn delete_item(&self, id: &str, collection: Option<&str>) -> Result<(), Error> { + call::void(self.client.pg(), "delete_item", &[&id, &collection]).await + } +} + +impl ItemsClient for Client { + type Error = Error; + + async fn search(&self, search: Search) -> Result { + let hydration = self.cached_hydration().await?; + let search = search.into_cql2_json()?; + let body = serde_json::to_value(search)?; + let token = body + .get("token") + .and_then(Value::as_str) + .map(str::to_string); + let limit = body.get("limit").and_then(Value::as_i64).unwrap_or(10); + let page: SearchPage = match search_page_with( + self.client.pg(), + hydration.model, + hydration.schema.as_deref(), + &body, + token.as_deref(), + limit, + ) + .await + { + Ok(page) => page, + // A malformed token surfaces as a `keyset_decode` error in SQL; map it to a clear + // InvalidToken rather than a raw DB error (and never an offset-style fallback). + Err(Error::TokioPostgres(e)) => { + let in_keyset = e + .as_db_error() + .and_then(|db| db.where_()) + .is_some_and(|w| w.contains("keyset_decode") || w.contains("keyset_where")); + let msg = e.to_string(); + if in_keyset + || msg.contains("token") + || msg.contains("keyset") + || msg.contains("base64") + { + return Err(Error::InvalidToken(msg)); + } + return Err(Error::TokioPostgres(e)); + } + Err(other) => return Err(other), + }; + ItemCollection::try_from(page) + } + + async fn item(&self, collection_id: &str, item_id: &str) -> Result, Error> { + let hydration = self.cached_hydration().await?; + let body = + serde_json::json!({"collections": [collection_id], "ids": [item_id], "limit": 1}); + let value = search_page_with( + self.client.pg(), + hydration.model, + hydration.schema.as_deref(), + &body, + None, + 1, + ) + .await? + .features + .into_iter() + .next(); + value + .map(serde_json::from_value) + .transpose() + .map_err(Error::from) + } +} + +impl CollectionsClient for Client { + type Error = Error; + + async fn collections(&self) -> Result, Error> { + collections::all(self.client.pg()) + .await? + .into_iter() + .map(|v| serde_json::from_value(v).map_err(Error::from)) + .collect() + } + + async fn collection(&self, id: &str) -> Result, Error> { + collections::get_collection(self.client.pg(), id) + .await? + .map(serde_json::from_value) + .transpose() + .map_err(Error::from) + } +} + +/// Writes go through the **Rust loader**, not the SQL `create_item`/`create_items` functions: +/// `add_items` dehydrates, splits fragments, and binary-COPYs entirely in Rust (see [`load_items`]). +/// The loader needs a `&mut tokio_postgres::Client` to open its own transaction for the binary COPY; it +/// reaches it via `pg_mut()`, so this works for any `PgConn` — a direct or a pooled connection. +impl TransactionClient for Client { + type Error = Error; + + /// Creates a collection via the SQL `create_collection` (which derives the `fragment_config` from + /// `item_assets`). Collections are not subject to the item-write restriction and are not a bulk path. + async fn add_collection(&mut self, collection: Collection) -> Result<(), Error> { + let collection = serde_json::to_value(collection)?; + call::void(self.client.pg(), "create_collection", &[&collection]).await + } + + /// Creates a single item through the Rust loader, erroring if its id already exists. + async fn add_item(&mut self, item: Item) -> Result<(), Error> { + self.add_items(vec![item]).await + } + + /// Creates items through the Rust loader (dehydrate + fragment split + binary COPY, all in Rust), + /// erroring if any id already exists. Use [`crate::PgstacPool::create_items`] to choose an + /// upsert/ignore [`ConflictPolicy`]. + async fn add_items(&mut self, items: Vec) -> Result<(), Error> { + let values = items + .into_iter() + .map(serde_json::to_value) + .collect::, _>>()?; + let schema = DehydrateSchema::load(self.client.pg()).await?; + let _ = load_items(self.client.pg_mut(), values, &schema, ConflictPolicy::Error).await?; + Ok(()) + } +} diff --git a/src/pgstac-rs/src/db/connect.rs b/src/pgstac-rs/src/db/connect.rs new file mode 100644 index 00000000..af58b658 --- /dev/null +++ b/src/pgstac-rs/src/db/connect.rs @@ -0,0 +1,375 @@ +//! Connection configuration for pgstac. +//! +//! A connection string (DSN) is the base configuration; the standard libpq environment variables +//! (`PGHOST`, `PGPORT`, `PGDATABASE`, `PGUSER`, `PGPASSWORD`, `PGOPTIONS`, `PGAPPNAME`, +//! `PGCONNECT_TIMEOUT`, `PGSSLMODE`) fill only the parameters the DSN leaves unset. Because the DSN is +//! the base, an explicit DSN always wins over ambient `PG*` values. +//! +//! Critically, it guarantees the pgstac `search_path` is applied at connection **startup** (via the +//! libpq `options` startup parameter) rather than a runtime `SET`. A runtime `SET search_path` does +//! not survive a transaction-pooling proxy such as PgBouncer, so setting it at startup is required for +//! correctness behind a pooler. + +use crate::Result; +use std::time::Duration; +use tokio_postgres::Config; +use tokio_postgres::config::SslMode; + +/// The default pgstac search path, applied at connection startup. +pub const DEFAULT_SEARCH_PATH: &str = "pgstac,public"; + +/// The default application name reported to Postgres. +pub const DEFAULT_APPLICATION_NAME: &str = "pgstac"; + +/// Resolved connection configuration for a pgstac database. +/// +/// The [`dsn`](Self::dsn) (or an empty [`Config`] when there is none) is the base; the libpq `PG*` +/// environment fills only the parameters the DSN leaves unset, so a DSN always wins. Build one with +/// [`ConnectConfig::from_env`] (reads `PGSTAC_DSN` + the TLS cert paths) or [`ConnectConfig::default`], +/// set/override any fields, then turn it into a [`tokio_postgres::Config`] with +/// [`ConnectConfig::to_pg_config`]. +/// +/// The `tokio_postgres::Config`-native connection parameters (host, port, database, user, password, +/// options, application name, connect timeout, ssl mode) are not mirrored here; set them through the +/// DSN or the `PG*` environment. Only what `Config` cannot carry stays: the startup search path and the +/// TLS certificate paths. +/// +/// # Examples +/// +/// ``` +/// use pgstac::ConnectConfig; +/// +/// let config = ConnectConfig { +/// dsn: Some("postgresql://username:password@localhost:5432/postgis".to_string()), +/// ..Default::default() +/// }; +/// let pg_config = config.to_pg_config().unwrap(); +/// assert_eq!(pg_config.get_dbname(), Some("postgis")); +/// ``` +#[derive(Debug, Clone)] +pub struct ConnectConfig { + /// A full connection string (libpq keyword/value or `postgresql://` URL). It is the base + /// configuration; the `PG*` environment fills only the parameters it leaves unset. + pub dsn: Option, + /// Path to a trusted CA certificate (`PGSSLROOTCERT`); consumed by the TLS connector. + pub sslrootcert: Option, + /// Path to a client certificate (`PGSSLCERT`); consumed by the TLS connector. + pub sslcert: Option, + /// Path to a client private key (`PGSSLKEY`); consumed by the TLS connector. + pub sslkey: Option, + /// The search path applied at connection startup. Defaults to [`DEFAULT_SEARCH_PATH`]. + pub search_path: String, +} + +impl Default for ConnectConfig { + fn default() -> Self { + ConnectConfig { + dsn: None, + sslrootcert: None, + sslcert: None, + sslkey: None, + search_path: DEFAULT_SEARCH_PATH.to_string(), + } + } +} + +impl ConnectConfig { + /// Resolves a configuration from the process environment: `PGSTAC_DSN` (or `DATABASE_URL`) for the + /// DSN, and `PGSSLROOTCERT`/`PGSSLCERT`/`PGSSLKEY` for the TLS certificate paths. The remaining libpq + /// `PG*` connection parameters are resolved at [`to_pg_config`](Self::to_pg_config) time, filling only + /// what the DSN leaves unset. + /// + /// # Examples + /// + /// ``` + /// use pgstac::ConnectConfig; + /// + /// let config = ConnectConfig::from_env(); + /// ``` + pub fn from_env() -> Self { + Self::from_getter(|key| std::env::var(key).ok()) + } + + /// Resolves a configuration using a caller-supplied environment getter. + /// + /// This is the testable core of [`ConnectConfig::from_env`]; tests pass a closure over a fixed map + /// instead of mutating the process environment. + pub fn from_getter(get: F) -> Self + where + F: Fn(&str) -> Option, + { + ConnectConfig { + dsn: get("PGSTAC_DSN").or_else(|| get("DATABASE_URL")), + sslrootcert: get("PGSSLROOTCERT"), + sslcert: get("PGSSLCERT"), + sslkey: get("PGSSLKEY"), + search_path: DEFAULT_SEARCH_PATH.to_string(), + } + } + + /// Builds a [`tokio_postgres::Config`], guaranteeing `search_path` is set at startup via `options`. + /// + /// The [`dsn`](Self::dsn) (or an empty [`Config`] when there is none) is the base; each libpq `PG*` + /// connection variable fills its parameter only when the DSN left it unset, so the DSN wins. The + /// search path is merged into the `options` startup string unless one is already present there. + /// + /// # Examples + /// + /// ``` + /// use pgstac::ConnectConfig; + /// + /// let pg_config = ConnectConfig::default().to_pg_config().unwrap(); + /// // The pgstac search_path is applied at startup via the options parameter, so it survives + /// // transaction-mode poolers like PgBouncer. + /// assert!(pg_config.get_options().unwrap().contains("search_path=pgstac,public")); + /// ``` + pub fn to_pg_config(&self) -> Result { + self.to_pg_config_from(|key| std::env::var(key).ok()) + } + + /// The testable core of [`to_pg_config`](Self::to_pg_config): the same resolution against a + /// caller-supplied environment getter instead of the process environment. + fn to_pg_config_from(&self, get: F) -> Result + where + F: Fn(&str) -> Option, + { + let mut config = match &self.dsn { + Some(dsn) => dsn.parse::()?, + None => Config::new(), + }; + + // Fill each connection parameter from the environment only when the DSN left it unset. + if config.get_hosts().is_empty() + && let Some(host) = get("PGHOST") + { + let _ = config.host(&host); + } + if config.get_ports().is_empty() + && let Some(port) = get("PGPORT").and_then(|value| value.parse().ok()) + { + let _ = config.port(port); + } + if config.get_dbname().is_none() + && let Some(dbname) = get("PGDATABASE") + { + let _ = config.dbname(&dbname); + } + if config.get_user().is_none() + && let Some(user) = get("PGUSER") + { + let _ = config.user(&user); + } + if config.get_password().is_none() + && let Some(password) = get("PGPASSWORD") + { + let _ = config.password(&password); + } + if config.get_connect_timeout().is_none() + && let Some(timeout) = get("PGCONNECT_TIMEOUT").and_then(|value| value.parse().ok()) + { + let _ = config.connect_timeout(Duration::from_secs(timeout)); + } + // Application name: the DSN value wins, then PGAPPNAME, then the pgstac default. + if config.get_application_name().is_none() { + let app = get("PGAPPNAME"); + let _ = config.application_name(app.as_deref().unwrap_or(DEFAULT_APPLICATION_NAME)); + } + // SSL mode: PGSSLMODE applies only when the DSN did not set one (a parsed DSN defaults to + // Prefer, so a Prefer here means "unset"). + if matches!(config.get_ssl_mode(), SslMode::Prefer) + && let Some(sslmode) = get("PGSSLMODE") + { + let _ = config.ssl_mode(parse_ssl_mode(&sslmode)); + } + + // Merge the search path into the startup options (pgbouncer-safe), without duplicating one that + // is already present (from the DSN or PGOPTIONS). DSN options win over PGOPTIONS. + let base_options = config + .get_options() + .map(str::to_string) + .or_else(|| get("PGOPTIONS")) + .unwrap_or_default(); + let merged = merge_search_path(&base_options, &self.search_path); + let _ = config.options(&merged); + + Ok(config) + } +} + +/// Maps a libpq `sslmode` string to a [`SslMode`]. +/// +/// `tokio_postgres` only distinguishes disable / prefer / require; the `verify-ca` and `verify-full` +/// modes map to `require` and rely on the TLS connector for certificate verification. Unknown values +/// fall back to `prefer` (the libpq default). +fn parse_ssl_mode(value: &str) -> SslMode { + match value.to_ascii_lowercase().as_str() { + "disable" => SslMode::Disable, + "require" | "verify-ca" | "verify-full" => SslMode::Require, + _ => SslMode::Prefer, + } +} + +/// Appends `-c search_path=` to a libpq `options` string unless a `search_path` is +/// already set there. +fn merge_search_path(base_options: &str, search_path: &str) -> String { + if base_options.contains("search_path") { + return base_options.to_string(); + } + let setting = format!("-c search_path={search_path}"); + if base_options.is_empty() { + setting + } else { + format!("{base_options} {setting}") + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::collections::HashMap; + + /// Builds a `get` closure over a fixed environment map. + fn env(pairs: &[(&str, &str)]) -> HashMap { + pairs + .iter() + .map(|(k, v)| (k.to_string(), v.to_string())) + .collect() + } + + #[test] + fn default_injects_search_path_at_startup() { + let config = ConnectConfig::default() + .to_pg_config_from(|_| None) + .unwrap(); + let options = config.get_options().unwrap_or_default(); + assert!( + options.contains("search_path=pgstac,public"), + "options should carry the startup search_path, got {options:?}" + ); + assert_eq!(config.get_application_name(), Some("pgstac")); + } + + #[test] + fn env_fills_unset_connection_params() { + let map = env(&[ + ("PGHOST", "db.example.com"), + ("PGPORT", "6432"), + ("PGDATABASE", "stac"), + ("PGUSER", "reader"), + ("PGPASSWORD", "secret"), + ("PGAPPNAME", "myapp"), + ("PGCONNECT_TIMEOUT", "7"), + ]); + let pg = ConnectConfig::default() + .to_pg_config_from(|key| map.get(key).cloned()) + .unwrap(); + assert_eq!(pg.get_ports(), &[6432]); + assert_eq!(pg.get_dbname(), Some("stac")); + assert_eq!(pg.get_user(), Some("reader")); + assert_eq!(pg.get_password(), Some(&b"secret"[..])); + assert_eq!(pg.get_application_name(), Some("myapp")); + assert_eq!(pg.get_connect_timeout(), Some(&Duration::from_secs(7))); + assert!( + pg.get_options().unwrap_or_default().contains("search_path"), + "search_path must be present even with env-var config" + ); + } + + #[test] + fn dsn_is_base_and_gets_search_path() { + let cfg = ConnectConfig { + dsn: Some("postgresql://u:p@host:5432/db".to_string()), + ..Default::default() + }; + let pg = cfg.to_pg_config_from(|_| None).unwrap(); + assert_eq!(pg.get_dbname(), Some("db")); + assert_eq!(pg.get_user(), Some("u")); + assert_eq!(pg.get_ports(), &[5432]); + assert!(pg.get_options().unwrap_or_default().contains("search_path")); + } + + #[test] + fn dsn_wins_over_ambient_env() { + // An explicit DSN must be authoritative: ambient `PGDATABASE`/`PGHOST`/`PGPORT` must not + // override the target the DSN names. Regression for the CI failure where a `--dsn ` + // load connected to `PGDATABASE=postgis` (no pgstac) instead of the clone. + let map = env(&[ + ("PGDATABASE", "postgis"), + ("PGHOST", "envhost"), + ("PGPORT", "1111"), + ]); + let cfg = ConnectConfig { + dsn: Some("postgresql://u:p@dsnhost:5432/target".to_string()), + ..Default::default() + }; + let pg = cfg.to_pg_config_from(|key| map.get(key).cloned()).unwrap(); + assert_eq!(pg.get_dbname(), Some("target")); + assert_eq!(pg.get_ports(), &[5432]); + } + + #[test] + fn pgoptions_preserved_and_search_path_appended() { + let map = env(&[("PGOPTIONS", "-c statement_timeout=5000")]); + let pg = ConnectConfig::default() + .to_pg_config_from(|key| map.get(key).cloned()) + .unwrap(); + let options = pg.get_options().unwrap_or_default(); + assert!( + options.contains("statement_timeout=5000"), + "got {options:?}" + ); + assert!( + options.contains("search_path=pgstac,public"), + "got {options:?}" + ); + } + + #[test] + fn existing_search_path_not_duplicated() { + let merged = merge_search_path("-c search_path=custom", "pgstac,public"); + assert_eq!(merged, "-c search_path=custom"); + assert_eq!(merged.matches("search_path").count(), 1); + } + + #[test] + fn default_ssl_mode_is_prefer() { + let pg = ConnectConfig::default() + .to_pg_config_from(|_| None) + .unwrap(); + assert!(matches!(pg.get_ssl_mode(), SslMode::Prefer)); + } + + #[test] + fn ssl_mode_resolves_from_pgsslmode() { + let disable = env(&[("PGSSLMODE", "disable")]); + assert!(matches!( + ConnectConfig::default() + .to_pg_config_from(|key| disable.get(key).cloned()) + .unwrap() + .get_ssl_mode(), + SslMode::Disable + )); + + let verify = env(&[("PGSSLMODE", "verify-full")]); + assert!(matches!( + ConnectConfig::default() + .to_pg_config_from(|key| verify.get(key).cloned()) + .unwrap() + .get_ssl_mode(), + SslMode::Require + )); + } + + #[test] + fn from_getter_reads_ssl_cert_paths() { + let map = env(&[ + ("PGSSLROOTCERT", "/etc/ssl/root.crt"), + ("PGSSLCERT", "/etc/ssl/client.crt"), + ("PGSSLKEY", "/etc/ssl/client.key"), + ]); + let cfg = ConnectConfig::from_getter(|key| map.get(key).cloned()); + assert_eq!(cfg.sslrootcert.as_deref(), Some("/etc/ssl/root.crt")); + assert_eq!(cfg.sslcert.as_deref(), Some("/etc/ssl/client.crt")); + assert_eq!(cfg.sslkey.as_deref(), Some("/etc/ssl/client.key")); + } +} diff --git a/src/pgstac-rs/src/db/mod.rs b/src/pgstac-rs/src/db/mod.rs new file mode 100644 index 00000000..9d07bd11 --- /dev/null +++ b/src/pgstac-rs/src/db/mod.rs @@ -0,0 +1,9 @@ +//! Connection configuration, TLS, the pooled client, and the connection pool. + +pub(crate) mod call; +pub(crate) mod client; +pub(crate) mod connect; +#[cfg(feature = "pool")] +pub(crate) mod pool; +#[cfg(feature = "pool")] +pub(crate) mod tls; diff --git a/src/pgstac-rs/src/db/pool.rs b/src/pgstac-rs/src/db/pool.rs new file mode 100644 index 00000000..485117ae --- /dev/null +++ b/src/pgstac-rs/src/db/pool.rs @@ -0,0 +1,478 @@ +//! A [`deadpool`](deadpool_postgres)-backed connection pool for pgstac. + +use crate::source::CachedHydration; +use crate::tls::make_tls_connect; +use crate::{ConnectConfig, Result}; +use deadpool_postgres::{Client, Manager, ManagerConfig, Pool, RecyclingMethod}; +use serde_json::Value; +use std::sync::Arc; +use tokio::sync::OnceCell; + +/// The default maximum number of connections in a [`PgstacPool`]. +pub const DEFAULT_POOL_SIZE: usize = 4; + +/// How the pool talks to the database, to stay compatible with external connection poolers. +/// +/// The difference matters only behind a pooler such as PgBouncer. In either mode the pgstac +/// `search_path` is applied at connection startup (see [`ConnectConfig`]). +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum PoolerMode { + /// Use the extended query protocol with server-side prepared statements (fastest). Safe for a + /// direct connection or a **session**-mode pooler. The default. + #[default] + Session, + /// Use the simple query protocol with no named prepared statements, so the connection is safe to + /// reuse across a **transaction**-mode pooler. Slightly slower; argument values are inlined. + Transaction, +} + +/// Options controlling how a [`PgstacPool`] is built. +#[derive(Debug, Clone)] +pub struct PoolOptions { + /// The maximum number of connections in the pool. Defaults to [`DEFAULT_POOL_SIZE`]. + pub max_size: usize, + /// The protocol mode used for queries. Defaults to [`PoolerMode::Session`]. + pub pooler_mode: PoolerMode, +} + +impl Default for PoolOptions { + fn default() -> Self { + PoolOptions { + max_size: DEFAULT_POOL_SIZE, + pooler_mode: PoolerMode::default(), + } + } +} + +/// A pool of connections to a pgstac database. +/// +/// Every connection applies the pgstac `search_path` at startup (see [`ConnectConfig`]), so the pool +/// is safe to use behind a transaction-mode connection pooler such as PgBouncer. +/// +/// # Examples +/// +/// ```no_run +/// use pgstac::{ConnectConfig, PgstacPool}; +/// +/// # tokio_test::block_on(async { +/// let pool = PgstacPool::connect(ConnectConfig::from_env()).await.unwrap(); +/// let version = pool.version().await.unwrap(); +/// # }) +/// ``` +#[derive(Debug, Clone)] +pub struct PgstacPool { + pool: Pool, + pooler_mode: PoolerMode, + /// Lazily-loaded, shared-across-clones hydration invariants (model + promoted schema). + hydration: Arc>, + /// Lazily-loaded, shared-across-clones dehydrate schema (the promoted-column registry) for the + /// write path — cached so repeated ingests don't re-query the catalog on every `create_items`. + dehydrate_schema: Arc>>, +} + +impl PgstacPool { + /// Builds a pool from a [`ConnectConfig`] with default [`PoolOptions`]. + /// + /// A connection is acquired during construction to fail fast if the database is unreachable. + /// + /// # Examples + /// + /// ```no_run + /// use pgstac::{ConnectConfig, PgstacPool}; + /// + /// # tokio_test::block_on(async { + /// let pool = PgstacPool::connect(ConnectConfig::from_env()).await.unwrap(); + /// # }) + /// ``` + pub async fn connect(config: ConnectConfig) -> Result { + PgstacPool::connect_with(config, PoolOptions::default()).await + } + + /// Builds a pool from a [`ConnectConfig`] and explicit [`PoolOptions`]. + /// + /// A connection is acquired during construction to fail fast if the database is unreachable. + pub async fn connect_with(config: ConnectConfig, options: PoolOptions) -> Result { + let tls = make_tls_connect(&config)?; + let pg_config = config.to_pg_config()?; + let manager_config = ManagerConfig { + recycling_method: RecyclingMethod::Fast, + }; + let manager = Manager::from_config(pg_config, tls, manager_config); + let pool = Pool::builder(manager).max_size(options.max_size).build()?; + let pgstac_pool = PgstacPool { + pool, + pooler_mode: options.pooler_mode, + hydration: Arc::new(OnceCell::new()), + dehydrate_schema: Arc::new(OnceCell::new()), + }; + // Acquire (and immediately release) a connection so a bad DSN or unreachable database fails + // here, at startup, rather than on the first request. + let _client = pgstac_pool.get().await?; + Ok(pgstac_pool) + } + + /// The cached hydration invariants (storage model + promoted schema), loaded once on first use and + /// shared across every search on this pool (and its clones). + pub(crate) async fn cached_hydration(&self) -> Result { + let cached = self + .hydration + .get_or_try_init(|| async { + let client = self.get().await?; + CachedHydration::detect(&**client).await + }) + .await?; + Ok(cached.clone()) + } + + /// The cached dehydrate schema (promoted-column registry) for the write path, loaded once on the + /// first ingest and reused across `create_items` calls + pool clones — so repeated ingests (e.g. + /// many small stac-fastapi POSTs) skip the per-call catalog round trip. + pub(crate) async fn cached_dehydrate_schema( + &self, + ) -> Result> { + let cached = self + .dehydrate_schema + .get_or_try_init(|| async { + let client = self.get().await?; + Ok::<_, crate::Error>(Arc::new( + crate::dehydrate::DehydrateSchema::load(&**client).await?, + )) + }) + .await?; + Ok(cached.clone()) + } + + /// The [`PoolerMode`] this pool uses for queries. + pub fn pooler_mode(&self) -> PoolerMode { + self.pooler_mode + } + + /// Acquires a connection from the pool. + /// + /// The returned client dereferences to a [`tokio_postgres::Client`]; use [`client`](Self::client) for + /// the STAC client traits with cached hydration. + pub async fn get(&self) -> Result { + self.pool.get().await.map_err(Into::into) + } + + /// A [`crate::Client`] backed by one pooled connection, sharing the pool's cached hydration invariants + /// (so the detection round trip happens once pool-wide, not per checked-out client). This is the + /// idiomatic rustac write/read surface — call the STAC client traits ([`stac::api::ItemsClient`], + /// [`stac::api::CollectionsClient`], [`stac::api::TransactionClient`]) on the returned `Client`. + pub async fn client(&self) -> Result> { + Ok(crate::Client::with_cache( + self.get().await?, + self.hydration.clone(), + )) + } + + /// Returns the pgstac version reported by the database. + pub async fn version(&self) -> Result { + self.client().await?.pgstac_version().await + } + + /// Runs one page of a search and returns the hydrated features plus the keyset pagination tokens. + /// + /// This drives the search client-side via [`search_page`](crate::search::search_page): it walks + /// the keyset query, hydrates the rows in Rust, and mints tokens via SQL `keyset_encode` (so they + /// are byte-identical to the SQL `search()` tokens). `token` is a `next:`/`prev:` continuation + /// token (or `None` for the first page); `limit` is the page size. + /// + /// Currently uses the extended query protocol, so it is safe for a direct connection or a + /// session-mode pooler. Transaction-mode-safe search (simple-protocol routing per + /// [`PoolerMode::Transaction`]) is a follow-up. + /// + /// # Examples + /// + /// ```no_run + /// use pgstac::{ConnectConfig, PgstacPool}; + /// use serde_json::json; + /// + /// # tokio_test::block_on(async { + /// let pool = PgstacPool::connect(ConnectConfig::from_env()).await.unwrap(); + /// let page = pool.search_page(&json!({"limit": 10}), None, 10).await.unwrap(); + /// # }) + /// ``` + pub async fn search_page( + &self, + search: &Value, + token: Option<&str>, + limit: i64, + ) -> Result { + let hydration = self.cached_hydration().await?; + let client = self.get().await?; + crate::search::search_page_with( + &**client, + hydration.model, + hydration.schema.as_deref(), + search, + token, + limit, + ) + .await + } + + /// Returns the total match count (`numberMatched`) for a search, independent of fetching a page. + /// + /// This is the parallel-context primitive: a caller that needs the count can run this on one pooled + /// connection **concurrently** with [`search_page`](Self::search_page) (with context disabled in the + /// search body) on another, overlapping the count latency with the data fetch. Returns `None` when + /// context counting is off for the search. + /// + /// # Examples + /// + /// ```no_run + /// use pgstac::{ConnectConfig, PgstacPool}; + /// use serde_json::json; + /// + /// # tokio_test::block_on(async { + /// let pool = PgstacPool::connect(ConnectConfig::from_env()).await.unwrap(); + /// let body = json!({"collections": ["landsat-c2-l2"]}); + /// // Fetch the page and the count concurrently. + /// let (page, matched) = tokio::join!( + /// pool.search_page(&body, None, 10), + /// pool.search_matched(&body), + /// ); + /// let (_page, _matched) = (page.unwrap(), matched.unwrap()); + /// # }) + /// ``` + pub async fn search_matched(&self, search: &Value) -> Result> { + let client = self.get().await?; + let row = client + .query_one( + "SELECT context_count, ctx_query FROM search_plan($1::jsonb, NULL::text, NULL::int)", + &[search], + ) + .await?; + if let Some(count) = row.try_get::<_, Option>("context_count")? { + return Ok(Some(count)); + } + let Some(ctx_query) = row.try_get::<_, Option>("ctx_query")? else { + return Ok(None); + }; + let count_row = client.query_one(&ctx_query, &[]).await?; + Ok(count_row.try_get::<_, Option>(0)?) + } + + /// Collects every matching item (up to `max_items`) into one [`SearchPage`](crate::search::SearchPage), + /// paginating internally over the keyset tokens. + pub async fn search_collect( + &self, + search: &Value, + max_items: Option, + ) -> Result { + let client = self.get().await?; + crate::search::search_collect(&**client, search, max_items).await + } + + /// Streams every matching item (up to `max_items`) as newline-delimited JSON to `write`. + /// + /// One connection is held for the whole stream and only one page is buffered at a time, so memory + /// stays flat regardless of the result size. Returns the number of items written. + pub async fn search_stream( + &self, + search: &Value, + write: &mut W, + max_items: Option, + ) -> Result { + let client = self.get().await?; + crate::search::search_stream(&**client, search, write, max_items).await + } + + /// Fetches a single item by id from a collection. + pub async fn get_item(&self, collection_id: &str, item_id: &str) -> Result> { + let client = self.get().await?; + crate::search::get_item(&**client, collection_id, item_id).await + } + + /// Runs one page of a collection search, returning the matching collections + keyset tokens. + pub async fn collection_search( + &self, + search: &Value, + token: Option<&str>, + ) -> Result { + let client = self.get().await?; + crate::collections::collection_search(&**client, search, token).await + } + + /// Fetches a single collection by id. + pub async fn get_collection(&self, collection_id: &str) -> Result> { + let client = self.get().await?; + crate::collections::get_collection(&**client, collection_id).await + } + + /// Fetches the queryables document for a collection (or catalog-wide when `collection_id` is None). + pub async fn get_queryables(&self, collection_id: Option<&str>) -> Result { + let client = self.get().await?; + crate::collections::get_queryables(&**client, collection_id).await + } + + /// Closes the pool, preventing it from handing out further connections. + pub fn close(&self) { + self.pool.close(); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + use std::sync::atomic::{AtomicU32, Ordering}; + use tokio_postgres::NoTls; + + const LOCAL_BASE: &str = "postgresql://username:password@localhost:5439"; + + fn test_dsn() -> String { + std::env::var("PGSTAC_RS_TEST_DB").unwrap_or_else(|_| format!("{LOCAL_BASE}/postgis")) + } + + /// A disposable database cloned from the clean test template, dropped on `Drop`. + /// + /// Search calls need a clean pgstac install (the shared dev database can carry ambiguous + /// multi-version functions), and they must not hold a connection to the template itself (that + /// would block other tests from cloning it). Each instance clones an isolated database. + struct CloneDb { + name: String, + } + + impl CloneDb { + async fn create() -> CloneDb { + static COUNTER: AtomicU32 = AtomicU32::new(0); + let name = format!( + "pgstac_rs_pool_test_{}_{}", + std::process::id(), + COUNTER.fetch_add(1, Ordering::Relaxed) + ); + let template = std::env::var("PGSTAC_RS_TEST_TEMPLATE") + .unwrap_or_else(|_| "pgstac_rs_test_template".to_string()); + let (client, connection) = + tokio_postgres::connect(&format!("{LOCAL_BASE}/postgres"), NoTls) + .await + .unwrap(); + let handle = tokio::spawn(connection); + let _ = client + .execute(&format!("CREATE DATABASE {name} TEMPLATE {template}"), &[]) + .await + .unwrap(); + handle.abort(); + CloneDb { name } + } + + fn dsn(&self) -> String { + format!("{LOCAL_BASE}/{}", self.name) + } + } + + impl Drop for CloneDb { + fn drop(&mut self) { + let name = self.name.clone(); + std::thread::scope(|scope| { + let _ = scope.spawn(|| { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + runtime.block_on(async move { + let (client, connection) = + tokio_postgres::connect(&format!("{LOCAL_BASE}/postgres"), NoTls) + .await + .unwrap(); + let handle = tokio::spawn(connection); + let _ = client + .execute( + "SELECT pg_terminate_backend(pid) FROM pg_stat_activity \ + WHERE datname = $1", + &[&name], + ) + .await; + let _ = client + .execute(&format!("DROP DATABASE IF EXISTS {name}"), &[]) + .await; + handle.abort(); + }); + }); + }); + } + } + + async fn clone_pool(pooler_mode: PoolerMode) -> (CloneDb, PgstacPool) { + let db = CloneDb::create().await; + let config = ConnectConfig { + dsn: Some(db.dsn()), + ..Default::default() + }; + let pool = PgstacPool::connect_with( + config, + PoolOptions { + pooler_mode, + ..Default::default() + }, + ) + .await + .unwrap(); + (db, pool) + } + + #[tokio::test] + async fn search_page_through_pool() { + // Exercises the pool -> search_page wiring end to end. The clone is an empty clean install, + // so the page has no features and no continuation token; search-result fidelity on real data + // is covered by tests/search_page.rs. + let (_db, pool) = clone_pool(PoolerMode::Session).await; + let page = pool + .search_page(&json!({"limit": 5}), None, 5) + .await + .unwrap(); + assert_eq!(page.number_returned, 0); + assert!(page.features.is_empty()); + assert!(page.next_token.is_none()); + pool.close(); + } + + #[test] + fn default_pooler_mode_is_session() { + assert_eq!(PoolOptions::default().pooler_mode, PoolerMode::Session); + } + + #[tokio::test] + async fn pool_reports_version_through_startup_search_path() { + let config = ConnectConfig { + dsn: Some(test_dsn()), + ..Default::default() + }; + let pool = PgstacPool::connect(config).await.unwrap(); + // `get_version()` lives in the pgstac schema, so a non-empty result also proves the startup + // `search_path` reached the server. + let version = pool.version().await.unwrap(); + assert!(!version.is_empty(), "expected a pgstac version, got empty"); + } + + #[tokio::test] + async fn pool_connects_with_sslmode_disable() { + // sslmode=disable skips the TLS handshake entirely; the connector is still built but unused. + let config = ConnectConfig { + dsn: Some(format!("{}?sslmode=disable", test_dsn())), + ..Default::default() + }; + let pool = PgstacPool::connect(config).await.unwrap(); + assert!(!pool.version().await.unwrap().is_empty()); + } + + #[tokio::test] + async fn pool_respects_max_size() { + let config = ConnectConfig { + dsn: Some(test_dsn()), + ..Default::default() + }; + let pool = PgstacPool::connect_with( + config, + PoolOptions { + max_size: 2, + ..Default::default() + }, + ) + .await + .unwrap(); + assert_eq!(pool.pool.status().max_size, 2); + } +} diff --git a/src/pgstac-rs/src/db/tls.rs b/src/pgstac-rs/src/db/tls.rs new file mode 100644 index 00000000..2b6e5a49 --- /dev/null +++ b/src/pgstac-rs/src/db/tls.rs @@ -0,0 +1,74 @@ +//! rustls TLS connector construction for the connection [pool](crate::PgstacPool). +//! +//! A single [`MakeRustlsConnect`] is always used; whether TLS is actually negotiated is governed by +//! the connection's `sslmode` (see [`ConnectConfig`](crate::ConnectConfig)). The trust store is the +//! bundled Mozilla root set plus any `PGSSLROOTCERT`, and client certificates are loaded from +//! `PGSSLCERT`/`PGSSLKEY` when both are present. + +use crate::{ConnectConfig, Error, Result}; +use rustls::RootCertStore; +use rustls::pki_types::{CertificateDer, PrivateKeyDer}; +use std::sync::Arc; +use tokio_postgres_rustls::MakeRustlsConnect; + +/// Builds the rustls TLS connector for the given configuration. +pub(crate) fn make_tls_connect(config: &ConnectConfig) -> Result { + let mut roots = RootCertStore::empty(); + roots.extend(webpki_roots::TLS_SERVER_ROOTS.iter().cloned()); + if let Some(path) = &config.sslrootcert { + for cert in load_certs(path)? { + roots.add(cert)?; + } + } + + let provider = Arc::new(rustls::crypto::aws_lc_rs::default_provider()); + let builder = rustls::ClientConfig::builder_with_provider(provider) + .with_safe_default_protocol_versions()? + .with_root_certificates(roots); + + let client_config = match (&config.sslcert, &config.sslkey) { + (Some(cert_path), Some(key_path)) => { + let certs = load_certs(cert_path)?; + let key = load_key(key_path)?; + builder.with_client_auth_cert(certs, key)? + } + _ => builder.with_no_client_auth(), + }; + + Ok(MakeRustlsConnect::new(client_config)) +} + +/// Loads all PEM certificates from a file. +fn load_certs(path: &str) -> Result>> { + let data = std::fs::read(path)?; + rustls_pemfile::certs(&mut data.as_slice()) + .collect::>>() + .map_err(Error::from) +} + +/// Loads the first PEM private key from a file. +fn load_key(path: &str) -> Result> { + let data = std::fs::read(path)?; + rustls_pemfile::private_key(&mut data.as_slice())? + .ok_or_else(|| Error::Tls(format!("no private key found in {path}"))) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn builds_default_connector() { + // The default config (Mozilla roots, no client auth) must yield a valid rustls connector. + let _connector = make_tls_connect(&ConnectConfig::default()).unwrap(); + } + + #[test] + fn missing_root_cert_is_an_error() { + let config = ConnectConfig { + sslrootcert: Some("/nonexistent/path/to/root.crt".to_string()), + ..Default::default() + }; + assert!(make_tls_connect(&config).is_err()); + } +} diff --git a/src/pgstac-rs/src/export/budget.rs b/src/pgstac-rs/src/export/budget.rs new file mode 100644 index 00000000..261b7b5b --- /dev/null +++ b/src/pgstac-rs/src/export/budget.rs @@ -0,0 +1,214 @@ +//! Global memory budget + spill coordination for the buffered-widening +//! geoparquet path (0.9.11). +//! +//! Only the buffered-widening path consumes the budget; the 0.10 stream-write +//! path does not buffer. The budget defaults to ~25% of the available memory, +//! read **cgroup-aware** so it behaves in containers. A partition that would +//! exceed its share spills its buffered items to a temp file before final encode. + +use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; + +/// Default fraction of detected memory to use as the global buffer budget. +pub const DEFAULT_BUDGET_FRACTION: f64 = 0.25; + +/// A global, shared memory budget across all partition workers. +/// +/// Workers reserve an estimated number of bytes before buffering a partition; if +/// the reservation would exceed the budget the worker spills to disk instead. +#[derive(Debug, Clone)] +pub struct MemoryBudget { + total: u64, + used: Arc, +} + +impl MemoryBudget { + /// Creates a budget with an explicit byte cap. + pub fn with_bytes(total: u64) -> Self { + MemoryBudget { + total, + used: Arc::new(AtomicU64::new(0)), + } + } + + /// Creates a budget sized to a fraction of detected (cgroup-aware) memory. + pub fn from_fraction(fraction: f64) -> Self { + let detected = detect_available_memory_bytes(); + let total = ((detected as f64) * fraction).max(1.0) as u64; + Self::with_bytes(total) + } + + /// Creates the default budget (~25% of detected memory). + pub fn default_budget() -> Self { + Self::from_fraction(DEFAULT_BUDGET_FRACTION) + } + + /// Total budget in bytes. + pub fn total(&self) -> u64 { + self.total + } + + /// Bytes currently reserved. + pub fn used(&self) -> u64 { + self.used.load(Ordering::Relaxed) + } + + /// Tries to reserve `bytes`. Returns a [`BudgetGuard`] that releases on drop, + /// or `None` if the reservation would exceed the budget (caller should spill). + pub fn try_reserve(&self, bytes: u64) -> Option { + let mut current = self.used.load(Ordering::Relaxed); + loop { + let next = current.saturating_add(bytes); + if next > self.total { + return None; + } + match self.used.compare_exchange_weak( + current, + next, + Ordering::AcqRel, + Ordering::Relaxed, + ) { + Ok(_) => { + return Some(BudgetGuard { + used: self.used.clone(), + bytes, + }); + } + Err(observed) => current = observed, + } + } + } +} + +/// A reservation against a [`MemoryBudget`]; releases its bytes on drop. +#[derive(Debug)] +pub struct BudgetGuard { + used: Arc, + bytes: u64, +} + +impl BudgetGuard { + /// Reserved bytes. + pub fn bytes(&self) -> u64 { + self.bytes + } +} + +impl Drop for BudgetGuard { + fn drop(&mut self) { + let _ = self.used.fetch_sub(self.bytes, Ordering::AcqRel); + } +} + +/// Detects available memory in bytes, preferring cgroup v2/v1 limits over the +/// raw machine total so the budget stays correct inside containers. +/// +/// Falls back to a conservative 2 GiB if nothing can be read. +pub fn detect_available_memory_bytes() -> u64 { + const FALLBACK: u64 = 2 * 1024 * 1024 * 1024; + + if let Some(limit) = cgroup_memory_limit() { + return limit; + } + if let Some(total) = meminfo_total() { + return total; + } + FALLBACK +} + +/// cgroup v2 `memory.max`, then v1 `memory.limit_in_bytes`. A "max" / very large +/// value (unlimited) is treated as no cgroup limit. +/// +/// Linux only; other platforms have no cgroup files, so [`detect_available_memory_bytes`] falls back to +/// the RAM-fraction default. +#[cfg(target_os = "linux")] +fn cgroup_memory_limit() -> Option { + // cgroup v2 + if let Ok(s) = std::fs::read_to_string("/sys/fs/cgroup/memory.max") { + let s = s.trim(); + if s != "max" + && let Ok(v) = s.parse::() + && is_real_limit(v) + { + return Some(v); + } + } + // cgroup v1 + if let Ok(s) = std::fs::read_to_string("/sys/fs/cgroup/memory/memory.limit_in_bytes") + && let Ok(v) = s.trim().parse::() + && is_real_limit(v) + { + return Some(v); + } + None +} + +/// Non-Linux fallback: no cgroup files, so detection uses the RAM-fraction default. +#[cfg(not(target_os = "linux"))] +fn cgroup_memory_limit() -> Option { + None +} + +/// A cgroup limit at/near the unsigned max (or page-aligned PAGE_COUNTER_MAX) +/// means "unlimited"; ignore it. +#[cfg(target_os = "linux")] +fn is_real_limit(v: u64) -> bool { + // Common "unlimited" sentinels are within a small factor of u64::MAX. + v < (u64::MAX / 2) +} + +/// `MemTotal` from `/proc/meminfo`, in bytes. +fn meminfo_total() -> Option { + let s = std::fs::read_to_string("/proc/meminfo").ok()?; + for line in s.lines() { + if let Some(rest) = line.strip_prefix("MemTotal:") { + let kb: u64 = rest.split_whitespace().next()?.parse().ok()?; + return Some(kb * 1024); + } + } + None +} + +/// Creates a named temp file for spilling a partition's buffered items. +/// +/// The file is removed when the returned handle is dropped. +pub fn spill_file() -> std::io::Result { + tempfile::Builder::new() + .prefix("pgstac-dump-spill-") + .tempfile() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn reserve_and_release() { + let budget = MemoryBudget::with_bytes(100); + let g1 = budget.try_reserve(60).expect("first fits"); + assert_eq!(budget.used(), 60); + assert!(budget.try_reserve(50).is_none(), "would exceed -> spill"); + let g2 = budget.try_reserve(40).expect("exact remainder fits"); + assert_eq!(budget.used(), 100); + drop(g1); + assert_eq!(budget.used(), 40); + drop(g2); + assert_eq!(budget.used(), 0); + } + + #[test] + fn detection_is_positive() { + assert!(detect_available_memory_bytes() > 0); + } + + #[test] + fn spill_file_roundtrip() { + use std::io::{Read, Seek, SeekFrom, Write}; + let mut f = spill_file().unwrap(); + f.write_all(b"hello").unwrap(); + let _ = f.as_file_mut().seek(SeekFrom::Start(0)).unwrap(); + let mut s = String::new(); + let _ = f.as_file_mut().read_to_string(&mut s).unwrap(); + assert_eq!(s, "hello"); + } +} diff --git a/src/pgstac-rs/src/export/format.rs b/src/pgstac-rs/src/export/format.rs new file mode 100644 index 00000000..9af5dfdd --- /dev/null +++ b/src/pgstac-rs/src/export/format.rs @@ -0,0 +1,484 @@ +//! Output formats for a dump unit: NDJSON, ItemCollection JSON, and +//! stac-geoparquet (both write modes). +//! +//! A [`Format`] is paired with a [`crate::export::sink::Sink`] (format ⟂ sink). +//! The format owns *encoding* hydrated items into bytes; the sink owns *where* +//! the bytes go. NDJSON / ItemCollection are pure serde. Geoparquet delegates to +//! the published `stac::geoparquet` writer: +//! +//! * **Buffered-widening** ([`GeoparquetMode::Buffered`], default; required for +//! 0.9.11): buffer the whole partition, then encode once so the Arrow schema is +//! complete by construction (no statistical sampling can drop a column). An +//! oversized partition spills to a temp file before encode (see +//! [`crate::export::budget`]). +//! * **Stream-write** ([`GeoparquetMode::Stream`], 0.10): the 0.10 registry +//! supplies a complete schema up front, so items are written batch-by-batch +//! with no full buffering; the first batch fixes the schema and every later +//! batch (registry-complete) conforms. + +use crate::{Error, Result}; +use serde_json::Value; +use stac::Item; +use stac::geoparquet::{Compression, WriterBuilder, WriterOptions, ZstdLevel}; + +/// The default (fast) parquet compression codec: zstd at parquet's default level. +pub fn default_compression() -> Compression { + Compression::ZSTD(ZstdLevel::default()) +} + +/// A short lowercase codec label for the dump manifest. parquet's clean codec name +/// (`Compression::codec_to_string`) is `pub(crate)`, and its public `Display` is the Debug repr (e.g. +/// `ZSTD(ZstdLevel(1))`), so there is no public upstream way to get the plain "zstd" — map the label here. +pub fn compression_label(compression: Compression) -> &'static str { + match compression { + Compression::UNCOMPRESSED => "uncompressed", + Compression::SNAPPY => "snappy", + Compression::GZIP(_) => "gzip", + Compression::LZO => "lzo", + Compression::BROTLI(_) => "brotli", + Compression::LZ4 => "lz4", + Compression::ZSTD(_) => "zstd", + Compression::LZ4_RAW => "lz4_raw", + } +} + +/// Which geoparquet write strategy to use. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum GeoparquetMode { + /// Buffer the whole partition then encode once (schema complete by + /// construction). Required for 0.9.11. + Buffered, + /// Stream batches with a registry-complete schema (0.10). + Stream, +} + +/// The output format for a dump unit. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Format { + /// Newline-delimited JSON, one item per line. + Ndjson, + /// A single STAC `FeatureCollection` JSON object. + ItemCollectionJson, + /// stac-geoparquet. + Geoparquet { + /// Compression codec. + compression: Compression, + /// Write mode. + mode: GeoparquetMode, + /// Max rows per parquet row-group (`None` = the stac/parquet default). Smaller values bound the + /// encoder's in-memory working set when streaming, at the cost of more, smaller row-groups. + max_row_group_row_count: Option, + }, +} + +impl Format { + /// The default partition geoparquet format for a storage model: buffered for + /// 0.9.11 (no registry), stream for 0.10. + pub fn geoparquet_for(model: crate::hydrate::HydrationModel) -> Format { + let mode = match model { + crate::hydrate::HydrationModel::BaseItem => GeoparquetMode::Buffered, + crate::hydrate::HydrationModel::Fragment => GeoparquetMode::Stream, + }; + Format::Geoparquet { + compression: default_compression(), + mode, + max_row_group_row_count: None, + } + } +} + +/// Converts a hydrated item [`Value`] into a [`stac::Item`]. +fn value_to_item(value: Value) -> Result { + serde_json::from_value(value).map_err(Error::from) +} + +/// Encodes a complete set of hydrated items (one dump unit) to bytes. +/// +/// For NDJSON / ItemCollection this serializes directly. For geoparquet it uses +/// the buffered path (encode once from the full vec) — the simplest correct +/// strategy and the one required for 0.9.11. The streaming geoparquet path is +/// [`GeoparquetStreamWriter`]. +pub fn encode_all(format: Format, items: Vec) -> Result> { + match format { + Format::Ndjson => encode_ndjson(&items), + Format::ItemCollectionJson => encode_item_collection(items), + Format::Geoparquet { + compression, + max_row_group_row_count, + .. + } => encode_geoparquet(items, compression, max_row_group_row_count), + } +} + +fn encode_ndjson(items: &[Value]) -> Result> { + let mut buf = Vec::new(); + for item in items { + let line = serde_json::to_vec(item)?; + buf.extend_from_slice(&line); + buf.push(b'\n'); + } + Ok(buf) +} + +fn encode_item_collection(items: Vec) -> Result> { + let fc = serde_json::json!({ + "type": "FeatureCollection", + "features": items, + }); + serde_json::to_vec(&fc).map_err(Error::from) +} + +fn encode_geoparquet( + mut items: Vec, + compression: Compression, + max_row_group_row_count: Option, +) -> Result> { + // Buffered-widening completeness guarantee: the whole partition is in hand, so resolve any + // property whose scalar JSON type conflicts across items into a single Arrow-representable type + // before encoding. Without this, a column that is e.g. a string in 99 items and an integer in 1 + // makes the Arrow encoder fail and the dump would lose the partition. + coerce_conflicting_properties(&mut items); + + let items: Result> = items.into_iter().map(value_to_item).collect(); + let items = items?; + let mut buf = Vec::new(); + let mut options = WriterOptions::new().with_compression(compression); + if let Some(n) = max_row_group_row_count { + options = options.with_max_row_group_row_count(n); + } + WriterBuilder::new(&mut buf) + .writer_options(options) + .build(items)? + .finish()?; + Ok(buf) +} + +/// A coarse JSON scalar class used to detect cross-item type conflicts. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +enum ScalarClass { + Null, + Bool, + Number, + String, + /// Object or array — treated as non-scalar; conflicts among these or with + /// scalars are left to the encoder (nested handling is out of scope here). + Compound, +} + +fn classify(v: &Value) -> ScalarClass { + match v { + Value::Null => ScalarClass::Null, + Value::Bool(_) => ScalarClass::Bool, + Value::Number(_) => ScalarClass::Number, + Value::String(_) => ScalarClass::String, + Value::Array(_) | Value::Object(_) => ScalarClass::Compound, + } +} + +/// Detects `properties.` columns whose scalar type conflicts across the +/// buffered items and coerces every value of those columns to a string (numbers +/// and bools become their JSON text). This is the lossless-completeness default: +/// the field is always captured and the partition always encodes; the int/string +/// distinction for a conflicted column is normalized to string. +/// +/// Only top-level `properties` scalar keys are normalized (the common conflict +/// surface); nested/compound conflicts are left to the encoder. +fn coerce_conflicting_properties(items: &mut [Value]) { + use std::collections::HashMap; + // First pass: collect the set of scalar classes seen per property key. + let mut classes: HashMap> = HashMap::new(); + for item in items.iter() { + if let Some(Value::Object(props)) = item.get("properties") { + for (k, v) in props { + let c = classify(v); + if c == ScalarClass::Null { + continue; // null is compatible with any column (nullable) + } + let _ = classes.entry(k.clone()).or_default().insert(c); + } + } + } + // Conflicting = more than one non-null scalar class, and all involved are + // scalar (we only safely stringify scalars). + let conflicted: std::collections::HashSet = classes + .into_iter() + .filter(|(_, set)| set.len() > 1 && set.iter().all(|c| *c != ScalarClass::Compound)) + .map(|(k, _)| k) + .collect(); + if conflicted.is_empty() { + return; + } + // Second pass: stringify the conflicted keys in place. + for item in items.iter_mut() { + if let Some(Value::Object(props)) = item.get_mut("properties") { + for key in &conflicted { + if let Some(v) = props.get_mut(key) + && !v.is_null() + { + let s = match v { + Value::String(_) => continue, + Value::Bool(b) => b.to_string(), + Value::Number(n) => n.to_string(), + _ => continue, + }; + *v = Value::String(s); + } + } + } + } +} + +/// A streaming geoparquet writer for the 0.10 path: write batches as they +/// arrive into a caller-owned sink `W`, finalize at the end. The first batch +/// fixes the schema (registry-complete on 0.10), so later batches conform +/// without buffering the whole partition — the *item* working set is one batch. +/// +/// The caller owns the sink `W` (a file, a `&mut Vec`, etc.) and reads/streams +/// the bytes after [`finish`](Self::finish). This keeps the encoded-file bytes off +/// the memory budget when `W` is file-backed. +#[allow(missing_debug_implementations)] +pub struct GeoparquetStreamWriter { + inner: WriterState, + compression: Compression, + max_row_group_row_count: Option, +} + +enum WriterState { + /// Not yet started: holds the pending sink. + Pending(Option), + /// Started: holds the live writer (boxed — the parquet writer is large). + Active(Box>), +} + +impl GeoparquetStreamWriter { + /// Creates a new streaming writer that encodes into `sink`. + pub fn new(sink: W, compression: Compression) -> Self { + GeoparquetStreamWriter { + inner: WriterState::Pending(Some(sink)), + compression, + max_row_group_row_count: None, + } + } + + /// Sets the max rows per parquet row-group. Smaller values flush row-groups sooner, bounding the + /// in-memory working set while streaming (at the cost of more, smaller row-groups). + pub fn with_max_row_group_row_count(mut self, max_row_group_row_count: usize) -> Self { + self.max_row_group_row_count = Some(max_row_group_row_count); + self + } + + /// Writes one batch of hydrated items. The first non-empty batch fixes the + /// schema. + pub fn write_batch(&mut self, items: Vec) -> Result<()> { + if items.is_empty() { + return Ok(()); + } + let items: Result> = items.into_iter().map(value_to_item).collect(); + let items = items?; + match &mut self.inner { + WriterState::Pending(sink) => { + let sink = sink.take().expect("sink present until started"); + let mut options = WriterOptions::new().with_compression(self.compression); + if let Some(n) = self.max_row_group_row_count { + options = options.with_max_row_group_row_count(n); + } + let writer = WriterBuilder::new(sink) + .writer_options(options) + .build(items)?; + self.inner = WriterState::Active(Box::new(writer)); + } + WriterState::Active(writer) => { + writer.write(items)?; + } + } + Ok(()) + } + + /// Whether any batch has been written (i.e. a file would be produced). + pub fn has_data(&self) -> bool { + matches!(self.inner, WriterState::Active(_)) + } + + /// Finalizes the file, flushing the parquet footer into the sink. Returns + /// `false` if no batches were written (no file should be produced). + pub fn finish(self) -> Result { + match self.inner { + WriterState::Active(writer) => { + (*writer).finish()?; + Ok(true) + } + WriterState::Pending(_) => Ok(false), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + fn sample_item(id: &str) -> Value { + json!({ + "type": "Feature", + "stac_version": "1.0.0", + "id": id, + "collection": "c", + "geometry": {"type": "Point", "coordinates": [1.0, 2.0]}, + "bbox": [1.0, 2.0, 1.0, 2.0], + "properties": {"datetime": "2020-01-01T00:00:00Z"}, + "assets": {}, + "links": [] + }) + } + + #[test] + fn ndjson_roundtrip() { + let items = vec![sample_item("a"), sample_item("b")]; + let bytes = encode_all(Format::Ndjson, items).unwrap(); + let text = String::from_utf8(bytes).unwrap(); + let lines: Vec<&str> = text.lines().collect(); + assert_eq!(lines.len(), 2); + let first: Value = serde_json::from_str(lines[0]).unwrap(); + assert_eq!(first["id"], "a"); + } + + #[test] + fn item_collection_shape() { + let items = vec![sample_item("a")]; + let bytes = encode_all(Format::ItemCollectionJson, items).unwrap(); + let fc: Value = serde_json::from_slice(&bytes).unwrap(); + assert_eq!(fc["type"], "FeatureCollection"); + assert_eq!(fc["features"][0]["id"], "a"); + } + + #[test] + fn compression_label_is_clean_codec() { + assert_eq!(compression_label(default_compression()), "zstd"); + assert_eq!(compression_label(Compression::SNAPPY), "snappy"); + assert_eq!(compression_label(Compression::UNCOMPRESSED), "uncompressed"); + } + + /// Reads geoparquet bytes back via a temp file and returns the item count. + fn read_geoparquet(bytes: &[u8]) -> usize { + use std::io::{Seek, SeekFrom, Write}; + let mut tf = tempfile::tempfile().unwrap(); + tf.write_all(bytes).unwrap(); + let _ = tf.seek(SeekFrom::Start(0)).unwrap(); + stac::geoparquet::from_reader(tf).unwrap().items.len() + } + + #[test] + fn geoparquet_buffered_readable() { + let items = vec![sample_item("a"), sample_item("b")]; + let bytes = encode_all( + Format::Geoparquet { + compression: default_compression(), + mode: GeoparquetMode::Buffered, + max_row_group_row_count: None, + }, + items, + ) + .unwrap(); + assert_eq!(read_geoparquet(&bytes), 2); + } + + #[test] + fn geoparquet_stream_readable() { + let mut buf: Vec = Vec::new(); + let mut writer = GeoparquetStreamWriter::new(&mut buf, default_compression()); + writer.write_batch(vec![sample_item("a")]).unwrap(); + writer.write_batch(vec![sample_item("b")]).unwrap(); + assert!(writer.has_data()); + assert!(writer.finish().unwrap()); + assert_eq!(read_geoparquet(&buf), 2); + } + + #[test] + fn coerce_mixed_type_property() { + let mut items = vec![ + json!({"type":"Feature","id":"a","geometry":{"type":"Point","coordinates":[0.0,0.0]}, + "properties":{"datetime":"2020-01-01T00:00:00Z","naip:year":"2011"}}), + json!({"type":"Feature","id":"b","geometry":{"type":"Point","coordinates":[0.0,0.0]}, + "properties":{"datetime":"2020-01-01T00:00:00Z","naip:year":2013}}), + ]; + coerce_conflicting_properties(&mut items); + // Both naip:year are now strings. + assert_eq!(items[0]["properties"]["naip:year"], "2011"); + assert_eq!(items[1]["properties"]["naip:year"], "2013"); + // Now it encodes without error. + let bytes = encode_all( + Format::Geoparquet { + compression: default_compression(), + mode: GeoparquetMode::Buffered, + max_row_group_row_count: None, + }, + items, + ) + .unwrap(); + assert_eq!(read_geoparquet(&bytes), 2); + } + + #[test] + fn coerce_leaves_consistent_property() { + let mut items = vec![ + json!({"properties":{"gsd":10}}), + json!({"properties":{"gsd":20}}), + ]; + coerce_conflicting_properties(&mut items); + // No conflict -> numbers stay numbers. + assert_eq!(items[0]["properties"]["gsd"], 10); + assert!(items[1]["properties"]["gsd"].is_number()); + } + + #[test] + fn geoparquet_stream_empty_no_file() { + let mut buf: Vec = Vec::new(); + let writer = GeoparquetStreamWriter::new(&mut buf, default_compression()); + assert!(!writer.has_data()); + assert!(!writer.finish().unwrap(), "no batches -> no file"); + assert!(buf.is_empty()); + } + + /// Number of parquet row-groups in the encoded bytes. + fn num_row_groups(bytes: &[u8]) -> usize { + use parquet::file::reader::{FileReader, SerializedFileReader}; + use std::io::{Seek, SeekFrom, Write}; + let mut tf = tempfile::tempfile().unwrap(); + tf.write_all(bytes).unwrap(); + let _ = tf.seek(SeekFrom::Start(0)).unwrap(); + SerializedFileReader::new(tf) + .unwrap() + .metadata() + .num_row_groups() + } + + #[test] + fn geoparquet_stream_row_group_size_takes_effect() { + let items: Vec = (0..100) + .map(|i| sample_item(&format!("item-{i}"))) + .collect(); + + // Cap at 10 rows/group, fed in 10-item batches -> 10 row-groups. + let mut capped: Vec = Vec::new(); + let mut w = GeoparquetStreamWriter::new(&mut capped, default_compression()) + .with_max_row_group_row_count(10); + for chunk in items.chunks(10) { + w.write_batch(chunk.to_vec()).unwrap(); + } + assert!(w.finish().unwrap()); + + // Default (no cap): the same 100 items land in a single row-group. + let mut default: Vec = Vec::new(); + let mut w2 = GeoparquetStreamWriter::new(&mut default, default_compression()); + w2.write_batch(items.clone()).unwrap(); + assert!(w2.finish().unwrap()); + + assert_eq!(read_geoparquet(&capped), 100, "all items round-trip"); + assert_eq!(num_row_groups(&default), 1, "no cap => single row-group"); + assert_eq!( + num_row_groups(&capped), + 10, + "cap=10 over 100 items => 10 row-groups" + ); + } +} diff --git a/src/pgstac-rs/src/export/manifest.rs b/src/pgstac-rs/src/export/manifest.rs new file mode 100644 index 00000000..74c53750 --- /dev/null +++ b/src/pgstac-rs/src/export/manifest.rs @@ -0,0 +1,248 @@ +//! Dump manifest, checkpoint, and root-metadata models, plus SHA-256 helpers. +//! +//! The manifest is the cross-PR contract the ingest side reads, written **last**; its presence signals a +//! complete dump. + +use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256}; +use std::collections::BTreeMap; + +/// The manifest format version. Import refuses an unknown major. +pub const MANIFEST_VERSION: &str = "1"; + +/// Computes the lowercase hex SHA-256 of a byte slice. +pub fn sha256_hex(bytes: &[u8]) -> String { + let mut hasher = Sha256::new(); + hasher.update(bytes); + hex_encode(&hasher.finalize()) +} + +/// A running SHA-256 over streamed bytes (for files written incrementally). +#[derive(Debug, Default, Clone)] +pub struct Sha256Writer { + hasher: Sha256, + bytes: u64, +} + +impl Sha256Writer { + /// Creates a new running hasher. + pub fn new() -> Self { + Self::default() + } + + /// Feeds bytes into the hash. + pub fn update(&mut self, bytes: &[u8]) { + self.hasher.update(bytes); + self.bytes += bytes.len() as u64; + } + + /// Number of bytes hashed so far. + pub fn bytes(&self) -> u64 { + self.bytes + } + + /// Finalizes and returns the hex digest. + pub fn finalize_hex(self) -> String { + hex_encode(&self.hasher.finalize()) + } +} + +fn hex_encode(bytes: &[u8]) -> String { + use std::fmt::Write; + let mut s = String::with_capacity(bytes.len() * 2); + for b in bytes { + let _ = write!(s, "{b:02x}"); + } + s +} + +/// A half-open `[start, end)` datetime range. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct DatetimeRange { + /// Inclusive start, STAC text (`...Z`). + pub start: String, + /// Exclusive end, STAC text (`...Z`). + pub end: String, +} + +/// A recorded file: path relative to the dump root, plus integrity metadata. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct FileEntry { + /// Path relative to the dump root. + pub path: String, + /// Lowercase hex SHA-256 of the file bytes. + pub sha256: String, + /// File size in bytes. + pub bytes: u64, + /// Row count, when meaningful (e.g. queryables/settings/parquet). + #[serde(skip_serializing_if = "Option::is_none")] + pub count: Option, +} + +/// One partition file in the manifest. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct PartitionEntry { + /// Source partition name (informational/debug only). + pub name: String, + /// Path to the parquet file relative to the dump root. + pub file: String, + /// Lowercase hex SHA-256. + pub sha256: String, + /// File size in bytes. + pub bytes: u64, + /// Exact rows written. + pub item_count: u64, + /// `datetime` footprint of the written rows. + pub datetime_range: Option, + /// `end_datetime` footprint of the written rows. + pub end_datetime_range: Option, + /// Spatial extent `[w, s, e, n]` (optional). + #[serde(skip_serializing_if = "Option::is_none")] + pub bbox: Option>, +} + +/// One collection in the manifest. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct CollectionEntry { + /// Collection id (authoritative, from `collection.json`). + pub id: String, + /// The `collection.json` file entry. + pub collection_file: FileEntry, + /// `collections.partition_trunc` (`month`/`year`/null). + pub partition_trunc: Option, + /// Total items dumped for this collection. + pub item_count: u64, + /// Partition files. + pub partitions: Vec, +} + +/// The exporter tool block. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Tool { + /// Tool name (`pgstac`). + pub name: String, + /// Command (`dump`). + pub command: String, + /// Exporter binary version. + pub version: String, +} + +/// The source-instance block. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Source { + /// `get_version()` of the dumped instance. + pub pgstac_version: String, + /// Postgres server version (informational). + #[serde(skip_serializing_if = "Option::is_none")] + pub postgres_version: Option, + /// Server description, no credentials (informational). + #[serde(skip_serializing_if = "Option::is_none")] + pub server: Option, +} + +/// Output options that affect how files were written. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Options { + /// Always true (decision: items are always hydrated). + pub hydrated: bool, + /// Parquet codec used. + pub compression: String, + /// Always `["datetime", "id"]`. + pub ordering: Vec, +} + +/// Partial-dump filter description. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Filter { + /// Selected collection ids. + pub collection_ids: Vec, + /// Optional datetime prefilter. + #[serde(skip_serializing_if = "Option::is_none")] + pub datetime: Option, + /// Optional bbox prefilter `[w, s, e, n]`. + #[serde(skip_serializing_if = "Option::is_none")] + pub bbox: Option>, +} + +/// The dump manifest, written last. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Manifest { + /// Manifest format version. + pub manifest_version: String, + /// RFC3339 creation timestamp. + pub created_at: String, + /// Exporter tool info. + pub tool: Tool, + /// Source instance info. + pub source: Source, + /// `"full"` or `"partial"`. + pub dump_type: String, + /// `null` for full; the filter for partial. + pub filter: Option, + /// Whether `--consistent` was used. + pub consistent: bool, + /// Informational snapshot block. + #[serde(skip_serializing_if = "Option::is_none")] + pub snapshot: Option, + /// Output options. + pub options: Options, + /// Root metadata files (queryables, settings). + pub metadata_files: BTreeMap, + /// Per-collection entries. + pub collections: Vec, +} + +/// A `_checkpoint.json` entry: one fully-written partition file. +/// +/// Carries the full [`PartitionEntry`] and its owning collection id so a resumed +/// run can place a skipped partition back into the manifest without re-scanning +/// it (the manifest is rebuilt completely on every run, including resumes). The +/// flat `file`/`sha256`/`item_count` mirror the entry for quick human/diff +/// reading and to verify the on-disk file before skipping. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct CheckpointEntry { + /// Path relative to the dump root. + pub file: String, + /// Lowercase hex SHA-256. + pub sha256: String, + /// Rows written. + pub item_count: u64, + /// Owning collection id (to re-place the entry on resume). + pub collection_id: String, + /// The full manifest entry for this partition file. + pub entry: PartitionEntry, +} + +/// The `_checkpoint.json` progress file; present only mid-dump. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Checkpoint { + /// Manifest format version. + pub manifest_version: String, + /// RFC3339 start timestamp. + pub started_at: String, + /// Completed partition files. + pub completed: Vec, +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn sha256_known_vector() { + // SHA-256("abc") + assert_eq!( + sha256_hex(b"abc"), + "ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad" + ); + } + + #[test] + fn streaming_matches_oneshot() { + let mut w = Sha256Writer::new(); + w.update(b"ab"); + w.update(b"c"); + assert_eq!(w.bytes(), 3); + assert_eq!(w.finalize_hex(), sha256_hex(b"abc")); + } +} diff --git a/src/pgstac-rs/src/export/metadata.rs b/src/pgstac-rs/src/export/metadata.rs new file mode 100644 index 00000000..d802f4b4 --- /dev/null +++ b/src/pgstac-rs/src/export/metadata.rs @@ -0,0 +1,108 @@ +//! Root-metadata fetchers: queryables.json, settings.json, collection.json. +//! +//! These shapes are the cross-PR contract the ingest side reads. +//! Generated columns (`queryables.id`, `collections.base_item`/geometry/…) are +//! never dumped — the target recomputes them. + +use crate::Result; +use crate::export::plan::CollectionPlan; +use serde_json::{Value, json}; +use tokio_postgres::GenericClient; + +/// Builds the `queryables.json` body: all queryables rows verbatim, minus the +/// GENERATED `id`. Each row carries its own `collection_ids` (null = global), so +/// global + per-collection queryables share one file. +pub async fn queryables_json(client: &C) -> Result<(Value, u64)> { + let rows = client + .query( + "SELECT name, collection_ids, definition, property_path, \ + property_wrapper, property_index_type \ + FROM queryables ORDER BY name, collection_ids NULLS FIRST", + &[], + ) + .await?; + let queryables: Vec = rows + .iter() + .map(|r| { + json!({ + "name": r.get::<_, String>("name"), + "collection_ids": r.get::<_, Option>>("collection_ids"), + "definition": r.get::<_, Option>("definition"), + "property_path": r.get::<_, Option>("property_path"), + "property_wrapper": r.get::<_, Option>("property_wrapper"), + "property_index_type": r.get::<_, Option>("property_index_type"), + }) + }) + .collect(); + let count = queryables.len() as u64; + Ok((json!({ "queryables": queryables }), count)) +} + +/// Builds the `settings.json` body: all `pgstac_settings` rows verbatim. Import decides what to apply. +pub async fn settings_json(client: &C) -> Result<(Value, u64)> { + let rows = client + .query("SELECT name, value FROM pgstac_settings ORDER BY name", &[]) + .await?; + let settings: Vec = rows + .iter() + .map(|r| { + json!({ + "name": r.get::<_, String>("name"), + "value": r.get::<_, Option>("value"), + }) + }) + .collect(); + let count = settings.len() as u64; + Ok((json!({ "settings": settings }), count)) +} + +/// Builds the `collection.json` body for a collection: the STAC content plus the +/// pgstac restore metadata (partition_trunc, private, fragment_config). +pub fn collection_json(plan: &CollectionPlan) -> Value { + json!({ + "stac": plan.content, + "pgstac": { + "partition_trunc": plan.partition_trunc.as_manifest_str(), + "private": plan.private, + "fragment_config": plan.fragment_config, + } + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::export::plan::PartitionTrunc; + + #[test] + fn collection_json_shape() { + let plan = CollectionPlan { + id: "c".into(), + partition_trunc: PartitionTrunc::Month, + content: json!({"id": "c", "type": "Collection"}), + private: Some(json!({"secret": 1})), + fragment_config: Some(vec!["[\"assets\"]".into()]), + partitions: vec![], + }; + let j = collection_json(&plan); + assert_eq!(j["stac"]["id"], "c"); + assert_eq!(j["pgstac"]["partition_trunc"], "month"); + assert_eq!(j["pgstac"]["private"]["secret"], 1); + assert_eq!(j["pgstac"]["fragment_config"][0], "[\"assets\"]"); + } + + #[test] + fn collection_json_null_trunc() { + let plan = CollectionPlan { + id: "c".into(), + partition_trunc: PartitionTrunc::None, + content: json!({"id": "c"}), + private: None, + fragment_config: None, + partitions: vec![], + }; + let j = collection_json(&plan); + assert!(j["pgstac"]["partition_trunc"].is_null()); + assert!(j["pgstac"]["fragment_config"].is_null()); + } +} diff --git a/src/pgstac-rs/src/export/mod.rs b/src/pgstac-rs/src/export/mod.rs new file mode 100644 index 00000000..92204b4c --- /dev/null +++ b/src/pgstac-rs/src/export/mod.rs @@ -0,0 +1,122 @@ +//! Data export (dump) module for pgstac. +//! +//! This module turns a pgstac instance (0.9.11 or 0.10) into a self-describing, +//! restorable dump: fully hydrated items as stac-geoparquet (one file per +//! partition) plus collection/queryables/settings JSON and a manifest. +//! +//! It is library code with no CLI dependencies; the `pgstac` binary (behind the +//! `cli` feature) is a thin wrapper. + +pub mod budget; +pub mod format; +pub mod manifest; +pub mod metadata; +pub mod parallel; +pub mod plan; +pub mod planner; +pub mod sink; + +pub use parallel::{ClientFactory, DsnFactory, Snapshot}; +pub use planner::{DumpConfig, DumpPlanner, DumpReport}; + +use crate::Result; +use crate::hydrate::{CollectionContext, HydrationModel}; +use serde_json::Value; +use tokio_postgres::GenericClient; + +/// Detects the storage model of a pgstac instance from its version string. +/// +/// 0.9.x and earlier use the per-collection `base_item` model; 0.10+ uses the +/// fragment model. A version that does not parse as `0.9.x` (e.g. the dev +/// `"unreleased"` build, or `1.x`) is treated as the fragment model. +/// +/// Prefer [`detect_hydration_model`], which inspects the actual schema and is +/// robust to non-numeric version strings; this helper is exposed for callers +/// that already have a release version in hand. +pub fn hydration_model_for_version(version: &str) -> HydrationModel { + // Versions are `MAJOR.MINOR.PATCH`. Only a clean `0.<10` is base_item. + let mut parts = version.split('.'); + let major: Option = parts.next().and_then(|p| p.parse().ok()); + let minor: Option = parts.next().and_then(|p| p.parse().ok()); + match (major, minor) { + (Some(0), Some(minor)) if minor < 10 => HydrationModel::BaseItem, + _ => HydrationModel::Fragment, + } +} + +/// Detects the storage model by inspecting the live schema (robust to dev +/// version strings like `"unreleased"`). +/// +/// The fragment model is identified by the presence of the `items.fragment_id` +/// column (0.10); otherwise the `collections.base_item` column (0.9.11) means the +/// base_item model. +pub async fn detect_hydration_model(client: &C) -> Result { + let has_fragment_id: bool = client + .query_one( + "SELECT EXISTS (\ + SELECT 1 FROM information_schema.columns \ + WHERE table_schema = 'pgstac' \ + AND table_name = 'items' \ + AND column_name = 'fragment_id'\ + ) AS present", + &[], + ) + .await? + .get("present"); + if has_fragment_id { + Ok(HydrationModel::Fragment) + } else { + Ok(HydrationModel::BaseItem) + } +} + +/// Loads the [`CollectionContext`] (base_item for 0.9.11) for a collection. +/// +/// For the fragment model the context carries no per-collection data, but we +/// still confirm the collection exists. +pub async fn load_collection_context( + client: &C, + model: HydrationModel, + collection_id: &str, +) -> Result { + match model { + HydrationModel::BaseItem => { + let rows = client + .query( + "SELECT base_item FROM collections WHERE id = $1", + &[&collection_id], + ) + .await?; + let base_item = rows + .first() + .and_then(|row| row.get::<_, Option>("base_item")); + Ok(CollectionContext { base_item }) + } + HydrationModel::Fragment => Ok(CollectionContext::default()), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn model_detection() { + assert_eq!( + hydration_model_for_version("0.9.11"), + HydrationModel::BaseItem + ); + assert_eq!( + hydration_model_for_version("0.9.15"), + HydrationModel::BaseItem + ); + assert_eq!( + hydration_model_for_version("0.10.0"), + HydrationModel::Fragment + ); + assert_eq!( + hydration_model_for_version("1.0.0"), + HydrationModel::Fragment + ); + } +} diff --git a/src/pgstac-rs/src/export/parallel.rs b/src/pgstac-rs/src/export/parallel.rs new file mode 100644 index 00000000..69f4c8c9 --- /dev/null +++ b/src/pgstac-rs/src/export/parallel.rs @@ -0,0 +1,150 @@ +//! Parallel partition fan-out for the dump and the optional single +//! repeatable-read snapshot. +//! +//! True parallelism needs one Postgres connection per concurrent scan (a single +//! `tokio_postgres::Client` serializes its queries on one connection), so the +//! parallel path takes a [`ClientFactory`] that mints fresh worker connections. +//! Each worker owns its connection and pulls partition jobs from a shared queue; +//! a bounded pool of `concurrency` workers runs partitions concurrently. The +//! object_store / fs sinks accept concurrent writes; the TAR sink serializes +//! internally via its own mutex, so it is safe under fan-out too (just not +//! faster). +//! +//! `--consistent`: a coordinator connection opens one +//! `REPEATABLE READ` transaction and `pg_export_snapshot()`s it; every worker +//! `SET TRANSACTION SNAPSHOT`s onto that id, so all partitions see one +//! point-in-time view even while concurrent ingest continues. The coordinator +//! holds its transaction open for the whole dump. + +use crate::{Error, Result}; +use std::future::Future; +use std::pin::Pin; +use tokio_postgres::Client; + +/// Mints fresh worker connections for the parallel dump. +/// +/// Each call must return a brand-new connection on `search_path = pgstac, +/// public`; the caller drives it concurrently with the others. +pub trait ClientFactory: Send + Sync { + /// Opens a new connection. + fn connect(&self) -> Pin> + Send + '_>>; +} + +/// A [`ClientFactory`] backed by a DSN over `NoTls` (the CLI default). +/// +/// Spawns each connection's driver task and sets the pgstac search_path before +/// handing the client back. +#[derive(Debug, Clone)] +pub struct DsnFactory { + dsn: String, +} + +impl DsnFactory { + /// Creates a factory for the given connection string. + pub fn new(dsn: impl Into) -> Self { + DsnFactory { dsn: dsn.into() } + } +} + +impl ClientFactory for DsnFactory { + fn connect(&self) -> Pin> + Send + '_>> { + let dsn = self.dsn.clone(); + Box::pin(async move { + let (client, connection) = tokio_postgres::connect(&dsn, tokio_postgres::NoTls).await?; + // The driver task runs for the connection's lifetime; we detach it + // (the client owns the connection's liveness for our purposes). + let _driver = tokio::spawn(async move { + if let Err(e) = connection.await { + tracing::error!("pgstac worker connection error: {e}"); + } + }); + client + .batch_execute("SET search_path = pgstac, public;") + .await?; + Ok(client) + }) + } +} + +/// A held repeatable-read snapshot for a `--consistent` dump. +/// +/// Owns the coordinator connection with an open transaction; dropping it (or +/// calling [`Snapshot::finish`]) ends the transaction and releases the snapshot. +pub struct Snapshot { + client: Client, + id: String, +} + +impl std::fmt::Debug for Snapshot { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Snapshot").field("id", &self.id).finish() + } +} + +impl Snapshot { + /// Opens a coordinator transaction and exports its snapshot id. + pub async fn export(factory: &dyn ClientFactory) -> Result { + let client = factory.connect().await?; + // A read-only repeatable-read transaction pins one MVCC snapshot for the + // whole dump. It must stay open (the snapshot is only valid while the + // exporting transaction is alive). + client + .batch_execute("BEGIN ISOLATION LEVEL REPEATABLE READ READ ONLY;") + .await?; + let id: String = client + .query_one("SELECT pg_export_snapshot() AS s", &[]) + .await? + .get("s"); + Ok(Snapshot { client, id }) + } + + /// The exported snapshot id (e.g. `00000003-0000001B-1`). + pub fn id(&self) -> &str { + &self.id + } + + /// Ends the coordinator transaction, releasing the snapshot. + pub async fn finish(self) -> Result<()> { + self.client.batch_execute("COMMIT;").await?; + Ok(()) + } +} + +/// Prepares a freshly minted worker connection to read under a shared snapshot. +/// +/// Opens a repeatable-read transaction on the worker and binds it to the +/// coordinator's exported snapshot id so the worker sees the same point in time. +pub async fn bind_worker_to_snapshot(client: &Client, snapshot_id: &str) -> Result<()> { + client + .batch_execute("BEGIN ISOLATION LEVEL REPEATABLE READ READ ONLY;") + .await?; + // The snapshot id is generated by Postgres (pg_export_snapshot); it is a + // fixed shape (digits and dashes). Validate before inlining to be safe. + if !snapshot_id + .bytes() + .all(|b| b.is_ascii_hexdigit() || b == b'-') + || snapshot_id.is_empty() + { + return Err(Error::Export(format!( + "refusing to bind to malformed snapshot id {snapshot_id:?}" + ))); + } + client + .batch_execute(&format!("SET TRANSACTION SNAPSHOT '{snapshot_id}';")) + .await?; + Ok(()) +} + +#[cfg(test)] +mod tests { + #[test] + fn rejects_malformed_snapshot_inline() { + // The validation rule (used by bind_worker_to_snapshot) rejects anything + // outside [0-9A-Fa-f-]; a quote-injection attempt must be refused. + let bad = "00000003'; DROP TABLE items; --"; + let ok = bad.bytes().all(|b| b.is_ascii_hexdigit() || b == b'-'); + assert!(!ok, "injection-shaped id must be rejected"); + let good = "00000003-0000001B-1"; + assert!(good.bytes().all(|b| b.is_ascii_hexdigit() || b == b'-')); + } +} diff --git a/src/pgstac-rs/src/export/plan.rs b/src/pgstac-rs/src/export/plan.rs new file mode 100644 index 00000000..1b862a49 --- /dev/null +++ b/src/pgstac-rs/src/export/plan.rs @@ -0,0 +1,249 @@ +//! Work-unit enumeration for a dump: collections and their partitions, with +//! datetime/bbox prefilter pruning (no keyset search engine, so this runs on +//! both 0.9.11 and 0.10). + +use crate::Result; +use serde_json::Value; +use tokio_postgres::GenericClient; + +/// The partitioning granularity of a collection (`collections.partition_trunc`). +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum PartitionTrunc { + /// One partition per year-month; file `.parquet`. + Month, + /// One partition per year; file `.parquet`. + Year, + /// No datetime sub-partitioning; one file `items.parquet`. + None, +} + +impl PartitionTrunc { + /// Parses the `collections.partition_trunc` text value. + pub fn parse(value: Option<&str>) -> PartitionTrunc { + match value { + Some("month") => PartitionTrunc::Month, + Some("year") => PartitionTrunc::Year, + _ => PartitionTrunc::None, + } + } + + /// The text form for the manifest (`"month"`/`"year"`/`null`). + pub fn as_manifest_str(&self) -> Option<&'static str> { + match self { + PartitionTrunc::Month => Some("month"), + PartitionTrunc::Year => Some("year"), + PartitionTrunc::None => None, + } + } +} + +/// One partition to dump: its source name + temporal footprint. +#[derive(Debug, Clone)] +pub struct PartitionPlan { + /// Source partition relation name (e.g. `_items_2_202401`). + pub name: String, + /// Constraint datetime range `[start, end)` as STAC text, if bounded. + pub dtrange: Option<(String, String)>, + /// Estimated row count (`reltuples`). + pub reltuples: f32, + /// The parquet file name within the collection dir, per MANIFEST.md naming. + pub file_name: String, +} + +/// One collection to dump: its id, partition_trunc, and partitions. +#[derive(Debug, Clone)] +pub struct CollectionPlan { + /// Collection id (authoritative). + pub id: String, + /// Partitioning granularity. + pub partition_trunc: PartitionTrunc, + /// The collection's STAC content (`collections.content`). + pub content: Value, + /// The collection's `private` jsonb, if any. + pub private: Option, + /// The collection's `fragment_config` (0.10), if any. + pub fragment_config: Option>, + /// Partitions to dump (already pruned by any prefilter). + pub partitions: Vec, +} + +/// A datetime/bbox prefilter for a partial dump. +#[derive(Debug, Clone, Default)] +pub struct Prefilter { + /// Inclusive-start, exclusive-end datetime window (STAC text). + pub datetime: Option<(String, String)>, + /// Spatial bbox `[w, s, e, n]`. + pub bbox: Option<[f64; 4]>, +} + +/// Derives the partition file name from the truncation and the partition's +/// datetime range start (per MANIFEST.md naming table). +fn file_name_for(trunc: PartitionTrunc, dtrange_start: Option<&str>) -> String { + match trunc { + PartitionTrunc::None => "items.parquet".to_string(), + PartitionTrunc::Month => match dtrange_start { + // `2024-01-01...` -> `202401.parquet` + Some(s) if s.len() >= 7 => { + format!("{}{}.parquet", &s[0..4], &s[5..7]) + } + _ => "items.parquet".to_string(), + }, + PartitionTrunc::Year => match dtrange_start { + Some(s) if s.len() >= 4 => format!("{}.parquet", &s[0..4]), + _ => "items.parquet".to_string(), + }, + } +} + +/// Lower bound text of a `tstzrange`, or `None` for `-infinity`/empty. +fn range_bound_text(range_lower: Option<&str>) -> Option { + match range_lower { + Some(s) if s != "-infinity" && s != "infinity" && !s.is_empty() => { + // Postgres renders tstz as `2024-01-01 00:00:00+00`; normalize to a + // sortable date prefix for naming. Keep the full text for the range. + Some(s.to_string()) + } + _ => None, + } +} + +/// Plans the collections + partitions to dump. +/// +/// `collection_ids` restricts to a subset (partial dump); `None` = all +/// collections. `prefilter` prunes partitions whose datetime footprint does not +/// overlap the requested window. bbox pruning happens per-item at scan time (the +/// partition footprint is temporal only), so a bbox-only filter does not prune +/// partitions here. +pub async fn plan_collections( + client: &C, + collection_ids: Option<&[String]>, + prefilter: &Prefilter, +) -> Result> { + // `fragment_config` exists only on 0.10. Detect it once and build the + // projection accordingly (a literal column reference is parsed even inside an + // unreached CASE branch, so it must be conditionally omitted, not guarded). + let has_fragment_config: bool = client + .query_one( + "SELECT EXISTS (SELECT 1 FROM information_schema.columns \ + WHERE table_schema='pgstac' AND table_name='collections' \ + AND column_name='fragment_config') AS present", + &[], + ) + .await? + .get("present"); + let fragment_expr = if has_fragment_config { + "to_jsonb(fragment_config)" + } else { + "NULL::jsonb" + }; + + let collection_rows = match collection_ids { + Some(ids) => { + let q = format!( + "SELECT id, content, private, partition_trunc, {fragment_expr} AS fragment_config \ + FROM collections WHERE id = ANY($1) ORDER BY id" + ); + client.query(&q, &[&ids]).await? + } + None => { + let q = format!( + "SELECT id, content, private, partition_trunc, {fragment_expr} AS fragment_config \ + FROM collections ORDER BY id" + ); + client.query(&q, &[]).await? + } + }; + + let mut plans = Vec::with_capacity(collection_rows.len()); + for crow in &collection_rows { + let id: String = crow.get("id"); + let content: Value = crow.get("content"); + let private: Option = crow.get("private"); + let trunc = + PartitionTrunc::parse(crow.get::<_, Option>("partition_trunc").as_deref()); + let fragment_config: Option> = crow + .get::<_, Option>("fragment_config") + .and_then(|v| serde_json::from_value(v).ok()); + + // Partition metadata + bounds rendered as text for naming/manifest. + let prows = client + .query( + "SELECT partition, reltuples, \ + lower(constraint_dtrange)::text AS dt_lower, \ + upper(constraint_dtrange)::text AS dt_upper \ + FROM partition_sys_meta WHERE collection = $1 ORDER BY partition", + &[&id], + ) + .await?; + + let mut partitions = Vec::with_capacity(prows.len()); + for prow in &prows { + let name: String = prow.get("partition"); + let reltuples: f32 = prow.get("reltuples"); + let dt_lower: Option = prow.get("dt_lower"); + let dt_upper: Option = prow.get("dt_upper"); + let dtrange = match (range_bound_text(dt_lower.as_deref()), &dt_upper) { + (Some(lo), Some(up)) if up != "infinity" => Some((lo, up.clone())), + _ => None, + }; + + // Prune by datetime prefilter: skip a partition whose bounded + // footprint does not overlap the requested window. Unbounded + // (NULL-trunc / infinite) partitions are never pruned (scanned with + // a per-item predicate instead). + if let Some((want_lo, want_hi)) = &prefilter.datetime + && let Some((p_lo, p_hi)) = &dtrange + && (p_hi.as_str() <= want_lo.as_str() || p_lo.as_str() >= want_hi.as_str()) + { + continue; + } + + let file_name = file_name_for(trunc, range_bound_text(dt_lower.as_deref()).as_deref()); + partitions.push(PartitionPlan { + name, + dtrange, + reltuples, + file_name, + }); + } + + plans.push(CollectionPlan { + id, + partition_trunc: trunc, + content, + private, + fragment_config, + partitions, + }); + } + Ok(plans) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn trunc_parse_and_str() { + assert_eq!(PartitionTrunc::parse(Some("month")), PartitionTrunc::Month); + assert_eq!(PartitionTrunc::parse(Some("year")), PartitionTrunc::Year); + assert_eq!(PartitionTrunc::parse(None), PartitionTrunc::None); + assert_eq!(PartitionTrunc::Month.as_manifest_str(), Some("month")); + assert_eq!(PartitionTrunc::None.as_manifest_str(), None); + } + + #[test] + fn file_naming() { + assert_eq!( + file_name_for(PartitionTrunc::Month, Some("2024-01-01 00:00:00+00")), + "202401.parquet" + ); + assert_eq!( + file_name_for(PartitionTrunc::Year, Some("2021-01-01 00:00:00+00")), + "2021.parquet" + ); + assert_eq!(file_name_for(PartitionTrunc::None, None), "items.parquet"); + // Month with no bound start -> falls back. + assert_eq!(file_name_for(PartitionTrunc::Month, None), "items.parquet"); + } +} diff --git a/src/pgstac-rs/src/export/planner.rs b/src/pgstac-rs/src/export/planner.rs new file mode 100644 index 00000000..cfbe02b1 --- /dev/null +++ b/src/pgstac-rs/src/export/planner.rs @@ -0,0 +1,1221 @@ +//! `DumpPlanner` — the dump orchestrator (backup). +//! +//! Enumerates collections + partitions, scans each partition (cursor, ordered), +//! hydrates + encodes geoparquet to a temp file, streams it to the [`Sink`], +//! records per-file SHA-256 + counts + temporal/spatial footprint, checkpoints +//! after each completed partition, writes the root metadata files, and writes +//! `manifest.json` **last** (its presence = a complete dump). +//! +//! Failure policy: default fail-fast; [`DumpConfig::skip_errors`] records a +//! per-item skip and continues. Resume: a partition whose checkpointed +//! file verifies by SHA-256 is skipped; a missing/partial one is redone whole. + +use crate::export::budget::MemoryBudget; +use crate::export::format::{ + Format, GeoparquetMode, GeoparquetStreamWriter, compression_label, default_compression, +}; +use stac::geoparquet::Compression; +use crate::export::manifest::{ + Checkpoint, CheckpointEntry, CollectionEntry, DatetimeRange, FileEntry, Filter, Manifest, + Options, PartitionEntry, Source, Tool, sha256_hex, +}; +use crate::export::metadata::{collection_json, queryables_json, settings_json}; +use crate::export::plan::{Prefilter, plan_collections}; +use crate::export::sink::Sink; +use crate::hydrate::HydrationModel; +use crate::source::{ScanFilter, detect_hydration_model, load_collection_context, scan_partition}; +use crate::{Error, Result}; +use serde_json::Value; +use std::collections::BTreeMap; +use tokio_postgres::GenericClient; + +/// Configuration for a dump run. +#[derive(Debug, Clone)] +pub struct DumpConfig { + /// Restrict to these collection ids (partial dump). `None` = all (full). + pub collection_ids: Option>, + /// Datetime/bbox prefilter (partial dump). + pub prefilter: Prefilter, + /// Parquet compression. + pub compression: Compression, + /// Continue past per-item errors, recording them. Default false. + pub skip_errors: bool, + /// Resume from a prior run: partition files already completed (and + /// verified by SHA-256), keyed by relative path. The caller loads + verifies + /// a prior `_checkpoint.json` (see the CLI's `load_resume_set`). A keyed + /// partition is *not* re-dumped; its prior entry is carried straight into the + /// new manifest, so a resumed dump produces the same complete manifest as an + /// uninterrupted run. + pub resume_completed: std::collections::HashMap, + /// Memory budget for the buffered geoparquet path; `None` = default ~25% RAM. + pub memory_budget: Option, + /// Exporter binary version string (for the manifest `tool.version`). + pub tool_version: String, + /// Optional informational server description (no credentials). + pub server: Option, + /// Max partitions processed concurrently in [`DumpPlanner::run_parallel`] + ///. `1` = effectively sequential. The sequential [`DumpPlanner::run`] + /// ignores it. Concurrency is additionally bounded by the global memory + /// budget on the buffered (0.9.11) path. + pub concurrency: usize, + /// Run the whole dump under one repeatable-read snapshot. Default off. + /// Honored by [`DumpPlanner::run_parallel`]: a coordinator transaction + /// exports its snapshot and every worker `SET TRANSACTION SNAPSHOT`s onto it, + /// so all partitions see one point in time. The plain [`DumpPlanner::run`] + /// ignores it (single connection sees its own view already). + pub consistent: bool, +} + +impl Default for DumpConfig { + fn default() -> Self { + DumpConfig { + collection_ids: None, + prefilter: Prefilter::default(), + compression: default_compression(), + skip_errors: false, + resume_completed: std::collections::HashMap::new(), + memory_budget: None, + tool_version: env!("CARGO_PKG_VERSION").to_string(), + server: None, + concurrency: 1, + consistent: false, + } + } +} + +/// Summary returned by a completed dump. +#[derive(Debug, Clone)] +pub struct DumpReport { + /// Number of collections dumped. + pub collections: usize, + /// Number of partition files written. + pub partitions: usize, + /// Total items written. + pub items: u64, + /// Items skipped (only non-zero with `skip_errors`). + pub skipped: u64, +} + +/// The dump orchestrator. +#[derive(Debug)] +pub struct DumpPlanner { + config: DumpConfig, + budget: MemoryBudget, +} + +impl DumpPlanner { + /// Creates a planner with the given config. + pub fn new(config: DumpConfig) -> Self { + let budget = match config.memory_budget { + Some(bytes) => MemoryBudget::with_bytes(bytes), + None => MemoryBudget::default_budget(), + }; + DumpPlanner { config, budget } + } + + /// Runs the dump against `client`, writing to `sink`. + /// + /// `client` must be on `search_path = pgstac, public`. Works on 0.9.11 + /// (base_item) and 0.10 (fragment). + pub async fn run(&self, client: &C, sink: &S) -> Result { + let model = detect_hydration_model(client).await?; + let pgstac_version: String = client + .query_one("SELECT get_version() AS v", &[]) + .await? + .get("v"); + + let plans = plan_collections( + client, + self.config.collection_ids.as_deref(), + &self.config.prefilter, + ) + .await?; + + // Checkpoint: records each completed partition file as it lands; + // written after every partition and removed implicitly when the manifest + // (the completion marker) is written last. Resume consults + // `config.resume_completed` (file paths already verified by the caller). + let mut checkpoint = Checkpoint { + manifest_version: crate::export::manifest::MANIFEST_VERSION.to_string(), + started_at: now_rfc3339(), + completed: Vec::new(), + }; + let already_done = &self.config.resume_completed; + + let geoparquet_mode = match model { + HydrationModel::BaseItem => GeoparquetMode::Buffered, + HydrationModel::Fragment => GeoparquetMode::Stream, + }; + + let mut collection_entries: Vec = Vec::new(); + let mut total_items: u64 = 0; + let mut total_skipped: u64 = 0; + let mut total_partition_files: usize = 0; + + for plan in &plans { + let ctx = load_collection_context(client, model, &plan.id).await?; + let coll_dir = format!("collections/{}", encode_collection_dir(&plan.id)); + + // collection.json + let coll_json = collection_json(plan); + let coll_bytes = serde_json::to_vec_pretty(&coll_json)?; + let coll_file = sink.put(&format!("{coll_dir}/collection.json"), &coll_bytes).await?; + + let mut partition_entries: Vec = Vec::new(); + let mut coll_item_count: u64 = 0; + + for partition in &plan.partitions { + let rel_file = format!("{coll_dir}/{}", partition.file_name); + + // Resume: a caller-verified completed file is NOT + // re-dumped; its prior entry is carried into the new manifest so + // the resumed dump's manifest is identical to an uninterrupted + // run (no missing partitions). + if let Some(prior) = already_done.get(&rel_file) { + let entry = prior.entry.clone(); + coll_item_count += entry.item_count; + total_items += entry.item_count; + total_partition_files += 1; + checkpoint.completed.push(prior.clone()); + let cp_bytes = serde_json::to_vec_pretty(&checkpoint)?; + let _ = sink.put("_checkpoint.json", &cp_bytes).await?; + partition_entries.push(entry); + continue; + } + + let scan_filter = scan_filter_for(&self.config.prefilter); + let outcome = dump_one_partition( + client, + sink, + model, + geoparquet_mode, + &partition.name, + &plan.id, + &rel_file, + &ctx, + &scan_filter, + self.config.compression, + self.config.skip_errors, + &self.budget, + ) + .await?; + + total_skipped += outcome.skipped; + + if let Some(entry) = outcome.entry { + coll_item_count += entry.item_count; + total_items += entry.item_count; + total_partition_files += 1; + checkpoint.completed.push(CheckpointEntry { + file: entry.file.clone(), + sha256: entry.sha256.clone(), + item_count: entry.item_count, + collection_id: plan.id.clone(), + entry: entry.clone(), + }); + // Persist checkpoint after each completed partition. + let cp_bytes = serde_json::to_vec_pretty(&checkpoint)?; + let _ = sink.put("_checkpoint.json", &cp_bytes).await?; + partition_entries.push(entry); + } + } + + collection_entries.push(CollectionEntry { + id: plan.id.clone(), + collection_file: FileEntry { + path: format!("{coll_dir}/collection.json"), + sha256: coll_file.sha256, + bytes: coll_file.bytes, + count: None, + }, + partition_trunc: plan.partition_trunc.as_manifest_str().map(str::to_string), + item_count: coll_item_count, + partitions: partition_entries, + }); + } + + self.write_metadata_and_manifest(client, sink, collection_entries, pgstac_version, None) + .await?; + + Ok(DumpReport { + collections: plans.len(), + partitions: total_partition_files, + items: total_items, + skipped: total_skipped, + }) + } + + /// Parallel partition fan-out with optional single snapshot. + /// + /// `primary` is one connection used for the cheap sequential work (model + /// detection, planning, collection.json / queryables / settings / manifest). + /// Worker connections for the concurrent partition scans come from `factory`; + /// up to [`DumpConfig::concurrency`] run at once, each owning its own + /// connection. `sink` is shared (`Arc`) across workers — fs/object_store + /// sinks write concurrently; the TAR sink serializes internally. + /// + /// With [`DumpConfig::consistent`] a coordinator transaction exports one + /// repeatable-read snapshot and every worker binds to it, so all partitions + /// see one point in time. + /// + /// Resume: a partition file already in + /// [`DumpConfig::resume_completed`] is skipped, exactly as in [`Self::run`]. + pub async fn run_parallel( + &self, + primary: &C, + sink: std::sync::Arc, + factory: &F, + ) -> Result + where + C: GenericClient, + S: Sink + 'static, + F: crate::export::parallel::ClientFactory, + { + use crate::export::parallel::{Snapshot, bind_worker_to_snapshot}; + + let model = detect_hydration_model(primary).await?; + let pgstac_version: String = primary + .query_one("SELECT get_version() AS v", &[]) + .await? + .get("v"); + let plans = plan_collections( + primary, + self.config.collection_ids.as_deref(), + &self.config.prefilter, + ) + .await?; + + let geoparquet_mode = match model { + HydrationModel::BaseItem => GeoparquetMode::Buffered, + HydrationModel::Fragment => GeoparquetMode::Stream, + }; + let scan_filter = scan_filter_for(&self.config.prefilter); + + // Optional shared snapshot. + let snapshot = if self.config.consistent { + Some(Snapshot::export(factory).await?) + } else { + None + }; + let snapshot_id = snapshot.as_ref().map(|s| s.id().to_string()); + + // Pre-seed totals + checkpoint with resumed (already-complete) partitions + // so the new manifest is complete and resumed files are not re-dumped. + let mut total_items: u64 = 0; + let mut total_partition_files: usize = 0; + let mut checkpoint = Checkpoint { + manifest_version: crate::export::manifest::MANIFEST_VERSION.to_string(), + started_at: now_rfc3339(), + completed: Vec::new(), + }; + + // Write each collection.json up front (sequential, cheap) and build the + // flat job list. Per-collection contexts are cloned into the jobs. + let mut collection_meta: Vec = Vec::with_capacity(plans.len()); + let mut jobs: Vec = Vec::new(); + for (ci, plan) in plans.iter().enumerate() { + let ctx = load_collection_context(primary, model, &plan.id).await?; + let coll_dir = format!("collections/{}", encode_collection_dir(&plan.id)); + let coll_json = collection_json(plan); + let coll_bytes = serde_json::to_vec_pretty(&coll_json)?; + let coll_file = sink.put(&format!("{coll_dir}/collection.json"), &coll_bytes).await?; + collection_meta.push(CollectionMeta { + id: plan.id.clone(), + coll_dir: coll_dir.clone(), + collection_file: FileEntry { + path: format!("{coll_dir}/collection.json"), + sha256: coll_file.sha256, + bytes: coll_file.bytes, + count: None, + }, + partition_trunc: plan.partition_trunc.as_manifest_str().map(str::to_string), + partitions: Vec::new(), + item_count: 0, + }); + let ctx = std::sync::Arc::new(ctx); + for partition in &plan.partitions { + let rel_file = format!("{coll_dir}/{}", partition.file_name); + // Resume: carry a verified prior entry straight into the + // manifest instead of re-dumping it. + if let Some(prior) = self.config.resume_completed.get(&rel_file) { + let entry = prior.entry.clone(); + total_items += entry.item_count; + total_partition_files += 1; + let meta = &mut collection_meta[ci]; + meta.item_count += entry.item_count; + meta.partitions.push(entry); + checkpoint.completed.push(prior.clone()); + continue; + } + jobs.push(PartitionJob { + collection_index: ci, + collection_id: plan.id.clone(), + partition_name: partition.name.clone(), + rel_file, + ctx: ctx.clone(), + }); + } + } + + // Shared job queue + worker pool. Workers own one connection each and + // pull jobs until the queue drains. + let concurrency = self.config.concurrency.max(1).min(jobs.len().max(1)); + let queue = std::sync::Arc::new(tokio::sync::Mutex::new( + jobs.into_iter().collect::>(), + )); + let compression = self.config.compression; + let skip_errors = self.config.skip_errors; + let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::>(); + + let budget = self.budget.clone(); + let mut workers = tokio::task::JoinSet::new(); + for _ in 0..concurrency { + let queue = queue.clone(); + let sink = sink.clone(); + let budget = budget.clone(); + let tx = tx.clone(); + let snap_id = snapshot_id.clone(); + let scan_filter = scan_filter.clone(); + let factory_fut = factory.connect(); + // Connect each worker up front so a connection failure is reported + // before the worker loop starts. + let client = factory_fut.await?; + if let Some(id) = &snap_id { + bind_worker_to_snapshot(&client, id).await?; + } + let _ = workers.spawn(async move { + loop { + let job = { + let mut q = queue.lock().await; + q.pop_front() + }; + let Some(job) = job else { break }; + let outcome = dump_one_partition( + &client, + sink.as_ref(), + model, + geoparquet_mode, + &job.partition_name, + &job.collection_id, + &job.rel_file, + job.ctx.as_ref(), + &scan_filter, + compression, + skip_errors, + &budget, + ) + .await; + let msg = outcome.map(|o| JobResult { + collection_index: job.collection_index, + outcome: o, + }); + let failed = msg.is_err(); + // If the receiver is gone the dump is aborting; stop. + if tx.send(msg).is_err() || failed { + break; + } + } + }); + } + drop(tx); // workers hold the remaining senders; rx closes when all finish + + // Collect results as they arrive. Checkpoint (pre-seeded with resumed + // files above) is updated incrementally; its order is non-deterministic + // under fan-out but only used for resume (set membership), so order does + // not matter. + let mut total_skipped: u64 = 0; + let mut first_error: Option = None; + + while let Some(msg) = rx.recv().await { + match msg { + Ok(JobResult { + collection_index, + outcome, + }) => { + total_skipped += outcome.skipped; + if let Some(entry) = outcome.entry { + total_items += entry.item_count; + total_partition_files += 1; + let meta = &mut collection_meta[collection_index]; + checkpoint.completed.push(CheckpointEntry { + file: entry.file.clone(), + sha256: entry.sha256.clone(), + item_count: entry.item_count, + collection_id: meta.id.clone(), + entry: entry.clone(), + }); + let cp_bytes = serde_json::to_vec_pretty(&checkpoint)?; + let _ = sink.put("_checkpoint.json", &cp_bytes).await?; + meta.item_count += entry.item_count; + meta.partitions.push(entry); + } + } + Err(e) => { + // Fail-fast: keep the first error, stop accepting more. + if first_error.is_none() { + first_error = Some(e); + } + break; + } + } + } + + // Drain/await workers (they stop when the queue empties or rx drops). + while workers.join_next().await.is_some() {} + + if let Some(e) = first_error { + // Leave _checkpoint.json in place so a resume can pick up. + if let Some(s) = snapshot { + let _ = s.finish().await; + } + return Err(e); + } + + // Per-partition order within a collection is non-deterministic under + // fan-out; sort by file name so the manifest is deterministic. + let mut collection_entries: Vec = + Vec::with_capacity(collection_meta.len()); + for meta in collection_meta { + let mut partitions = meta.partitions; + partitions.sort_by(|a, b| a.file.cmp(&b.file)); + collection_entries.push(CollectionEntry { + id: meta.id, + collection_file: meta.collection_file, + partition_trunc: meta.partition_trunc, + item_count: meta.item_count, + partitions, + }); + } + + self.write_metadata_and_manifest( + primary, + sink.as_ref(), + collection_entries, + pgstac_version, + snapshot_id, + ) + .await?; + + if let Some(s) = snapshot { + s.finish().await?; + } + + Ok(DumpReport { + collections: plans.len(), + partitions: total_partition_files, + items: total_items, + skipped: total_skipped, + }) + } + + /// Writes the root metadata files (queryables/settings), the manifest LAST, + /// its sibling sha256, and drops the checkpoint. `snapshot_id` is recorded + /// (informational) when the dump ran under `--consistent`. + async fn write_metadata_and_manifest( + &self, + client: &C, + sink: &S, + collection_entries: Vec, + pgstac_version: String, + snapshot_id: Option, + ) -> Result<()> { + // Root metadata files. + let (q_json, q_count) = queryables_json(client).await?; + let q_bytes = serde_json::to_vec_pretty(&q_json)?; + let q_file = sink.put("queryables.json", &q_bytes).await?; + + let (s_json, s_count) = settings_json(client).await?; + let s_bytes = serde_json::to_vec_pretty(&s_json)?; + let s_file = sink.put("settings.json", &s_bytes).await?; + + let mut metadata_files = BTreeMap::new(); + let _ = metadata_files.insert( + "queryables".to_string(), + FileEntry { + path: "queryables.json".to_string(), + sha256: q_file.sha256, + bytes: q_file.bytes, + count: Some(q_count), + }, + ); + let _ = metadata_files.insert( + "settings".to_string(), + FileEntry { + path: "settings.json".to_string(), + sha256: s_file.sha256, + bytes: s_file.bytes, + count: Some(s_count), + }, + ); + + // Manifest (written LAST). + let dump_type = if self.is_partial() { "partial" } else { "full" }; + let snapshot = snapshot_id + .as_ref() + .map(|id| serde_json::json!({ "exported_snapshot": id })); + let manifest = Manifest { + manifest_version: crate::export::manifest::MANIFEST_VERSION.to_string(), + created_at: now_rfc3339(), + tool: Tool { + name: "pgstac".to_string(), + command: "dump".to_string(), + version: self.config.tool_version.clone(), + }, + source: Source { + pgstac_version, + postgres_version: None, + server: self.config.server.clone(), + }, + dump_type: dump_type.to_string(), + filter: self.filter_for_manifest(), + consistent: snapshot_id.is_some(), + snapshot, + options: Options { + hydrated: true, + compression: compression_label(self.config.compression).to_string(), + ordering: vec!["datetime".to_string(), "id".to_string()], + }, + metadata_files, + collections: collection_entries, + }; + let manifest_bytes = serde_json::to_vec_pretty(&manifest)?; + let _ = sink.put("manifest.json", &manifest_bytes).await?; + // Sibling integrity for the manifest itself. + let _ = sink + .put( + "manifest.json.sha256", + format!("{}\n", sha256_hex(&manifest_bytes)).as_bytes(), + ) + .await?; + + // Manifest is the completion marker; drop the in-progress checkpoint. + sink.remove("_checkpoint.json").await?; + + sink.finalize().await?; + Ok(()) + } + + fn is_partial(&self) -> bool { + self.config.collection_ids.is_some() + || self.config.prefilter.datetime.is_some() + || self.config.prefilter.bbox.is_some() + } + + fn filter_for_manifest(&self) -> Option { + if !self.is_partial() { + return None; + } + Some(Filter { + collection_ids: self.config.collection_ids.clone().unwrap_or_default(), + datetime: self + .config + .prefilter + .datetime + .as_ref() + .map(|(s, e)| DatetimeRange { + start: s.clone(), + end: e.clone(), + }), + bbox: self.config.prefilter.bbox.map(|b| b.to_vec()), + }) + } +} + +/// Budget-gated in-memory item buffer for the 0.9.11 buffered-widening path, +/// with spill-to-disk for an oversized single partition. +/// +/// Items accumulate in RAM only while a [`BudgetGuard`] reservation holds. When +/// the running estimate would exceed the global budget the buffer **spills**: +/// it serializes its accumulated items as NDJSON to a temp file and releases the +/// reservation, so later batches no longer count against RAM. At encode time a +/// spilled buffer replays its NDJSON back into the encoder. Either way the whole +/// partition is encoded once, schema complete by construction. +/// +/// [`BudgetGuard`]: crate::export::budget::BudgetGuard +struct BufferedItems<'b> { + budget: &'b MemoryBudget, + /// Items held in memory (not yet spilled). + items: Vec, + /// Bytes currently reserved against the budget for `items`. + reserved: u64, + /// The live reservation, dropped (released) on spill or at end. + guard: Option, + /// Spill file (NDJSON of already-spilled items), created on first spill. + spill: Option>, + spill_handle: Option, + /// Whether anything (memory or spill) has been written. + nonempty: bool, +} + +impl<'b> BufferedItems<'b> { + fn new(budget: &'b MemoryBudget) -> Self { + BufferedItems { + budget, + items: Vec::new(), + reserved: 0, + guard: None, + spill: None, + spill_handle: None, + nonempty: false, + } + } + + /// Rough per-item heap estimate used for budget accounting. + fn estimate(item: &Value) -> u64 { + // A serialized-length proxy; cheap and good enough to bound RAM. The + // in-memory `Value` is larger than its text, so scale up a little. + (serde_json::to_string(item).map(|s| s.len()).unwrap_or(256) as u64) * 2 + } + + /// Adds a batch, reserving budget; spills to disk if the reservation fails. + fn extend(&mut self, batch: Vec) -> Result<()> { + if batch.is_empty() { + return Ok(()); + } + self.nonempty = true; + let need: u64 = batch.iter().map(Self::estimate).sum(); + + // Try to grow the reservation to cover the new batch. + if self.try_reserve_more(need) { + self.items.extend(batch); + return Ok(()); + } + + // Could not reserve: spill what we have, then keep the new batch in + // memory under a fresh (smaller) reservation if possible, else spill it + // too. This bounds peak RAM to roughly one batch over the budget. + self.spill_current()?; + if self.try_reserve_more(need) { + self.items.extend(batch); + } else { + // Even one batch exceeds the budget: write it straight to spill. + self.spill_batch(&batch)?; + } + Ok(()) + } + + /// Attempts to extend the reservation by `extra` bytes. + fn try_reserve_more(&mut self, extra: u64) -> bool { + match self.budget.try_reserve(self.reserved + extra) { + Some(g) => { + // Replace the old guard (releases it) with the larger one. + self.guard = Some(g); + self.reserved += extra; + true + } + None => false, + } + } + + /// Spills the in-memory items to the NDJSON spill file and releases the + /// reservation. + fn spill_current(&mut self) -> Result<()> { + if !self.items.is_empty() { + let batch = std::mem::take(&mut self.items); + self.spill_batch(&batch)?; + } + self.reserved = 0; + self.guard = None; // release the reservation + Ok(()) + } + + /// Appends a batch of items as NDJSON to the spill file (creating it lazily). + fn spill_batch(&mut self, batch: &[Value]) -> Result<()> { + use std::io::Write; + if self.spill.is_none() { + let handle = crate::export::budget::spill_file()?; + let file = handle.reopen()?; + self.spill_handle = Some(handle); + self.spill = Some(std::io::BufWriter::new(file)); + } + let w = self.spill.as_mut().expect("spill writer present"); + for item in batch { + serde_json::to_writer(&mut *w, item)?; + w.write_all(b"\n")?; + } + Ok(()) + } + + /// Encodes all buffered items (spilled + in-memory) to a geoparquet file at + /// `out`. Returns `false` if there were no items (no file should be written). + fn encode_to(mut self, out: &std::path::Path, compression: Compression) -> Result { + if !self.nonempty { + return Ok(false); + } + // Gather: replay any spilled NDJSON, then append in-memory items. + let mut all: Vec = Vec::new(); + if let Some(w) = self.spill.take() { + // Flush + reopen the spill file for reading. + let _ = w.into_inner().map_err(std::io::Error::other)?; + if let Some(handle) = &self.spill_handle { + use std::io::{BufRead, BufReader}; + let f = handle.reopen()?; + for line in BufReader::new(f).lines() { + let line = line?; + if line.is_empty() { + continue; + } + all.push(serde_json::from_str(&line)?); + } + } + } + all.append(&mut self.items); + if all.is_empty() { + return Ok(false); + } + let bytes = crate::export::format::encode_all( + Format::Geoparquet { + compression, + mode: GeoparquetMode::Buffered, + max_row_group_row_count: None, + }, + all, + )?; + std::fs::write(out, &bytes)?; + Ok(true) + } +} + +/// Dumps one partition to a temp geoparquet file, streams it to the sink, and +/// returns the manifest entry (None if the partition has no items). +/// +/// Free function (not a method) so the parallel fan-out can run it in +/// concurrent worker tasks each owning their own connection. `budget` gates the +/// buffered (0.9.11) path: items accumulate in memory only while a reservation +/// holds; an oversized partition spills its buffered items to a temp file before +/// the final encode so peak memory stays under the global budget. +#[allow(clippy::too_many_arguments)] +async fn dump_one_partition( + client: &C, + sink: &S, + model: HydrationModel, + mode: GeoparquetMode, + partition_name: &str, + collection: &str, + rel_file: &str, + ctx: &crate::hydrate::CollectionContext, + scan_filter: &ScanFilter, + compression: Compression, + skip_errors: bool, + budget: &MemoryBudget, +) -> Result { + // Footprint accumulators. + let mut footprint = Footprint::default(); + let mut item_count: u64 = 0; + let mut skipped: u64 = 0; + + // Temp file for the encoded geoparquet -> streamed to sink. The + // NamedTempFile owns the path + cleanup; we read/write via fresh handles + // on its path so no borrow outlives the encode (the stream writer needs + // exclusive ownership of its File for the whole scan). + let spill = crate::export::budget::spill_file()?; + let spill_path = spill.path().to_path_buf(); + + { + // Buffered (0.9.11) path: accumulate items, spilling to a temp file when + // the in-memory working set would exceed the budget. Stream (0.10) path + // never buffers, so the budget does not apply. + let mut buffer = BufferedItems::new(budget); + let mut stream_writer: Option> = match mode { + GeoparquetMode::Stream => { + let file = std::fs::File::create(&spill_path)?; + Some(GeoparquetStreamWriter::new(file, compression)) + } + GeoparquetMode::Buffered => None, + }; + + // Scan + per-batch handling. The scan callback is sync; we collect + // footprint and feed the encoder. + let scan_result = scan_partition( + client, + model, + partition_name, + collection, + ctx, + scan_filter, + |batch| { + let mut accepted: Vec = Vec::with_capacity(batch.len()); + for item in batch { + match validate_item(&item) { + Ok(()) => { + footprint.observe(&item); + item_count += 1; + accepted.push(item); + } + Err(e) => { + if skip_errors { + skipped += 1; + } else { + return Err(e); + } + } + } + } + match mode { + GeoparquetMode::Buffered => buffer.extend(accepted)?, + GeoparquetMode::Stream => { + stream_writer + .as_mut() + .expect("stream writer present") + .write_batch(accepted)?; + } + } + Ok(()) + }, + ) + .await; + + let _scanned = scan_result?; + + // Finalize the encoder. + let produced = match mode { + GeoparquetMode::Buffered => buffer.encode_to(&spill_path, compression)?, + GeoparquetMode::Stream => stream_writer + .take() + .expect("stream writer present") + .finish()?, + }; + + if !produced { + // Empty partition -> no file. + return Ok(PartitionOutcome { + entry: None, + skipped, + }); + } + } + + // Stream the finished temp file to the sink (also yields sha256). `spill` + // still owns the path until it drops at end of scope. + let written = sink.put_from_path(rel_file, &spill_path).await?; + drop(spill); + + Ok(PartitionOutcome { + entry: Some(PartitionEntry { + name: partition_name.to_string(), + file: rel_file.to_string(), + sha256: written.sha256, + bytes: written.bytes, + item_count, + datetime_range: footprint.datetime_range(), + end_datetime_range: footprint.end_datetime_range(), + bbox: footprint.bbox(), + }), + skipped, + }) +} + +struct PartitionOutcome { + entry: Option, + skipped: u64, +} + +/// One partition's worth of work for the parallel pool. +struct PartitionJob { + /// Index into `collection_meta` (which collection this partition belongs to). + collection_index: usize, + /// Collection id (for pre-loading the collection's fragments before the scan). + collection_id: String, + /// Source partition relation name. + partition_name: String, + /// Output path relative to the dump root. + rel_file: String, + /// Per-collection hydration context (base_item for 0.9.11), shared (`Arc`). + ctx: std::sync::Arc, +} + +/// A worker's result for one partition. +struct JobResult { + collection_index: usize, + outcome: PartitionOutcome, +} + +/// Mutable per-collection accumulator for the parallel path (partition entries +/// arrive out of order, so we collect them here and assemble at the end). +struct CollectionMeta { + id: String, + #[allow(dead_code)] + coll_dir: String, + collection_file: FileEntry, + partition_trunc: Option, + partitions: Vec, + item_count: u64, +} + +/// Validates an item enough to encode it (must be a Feature object with an id). +fn validate_item(item: &Value) -> Result<()> { + if !item.is_object() { + return Err(Error::Export("item is not a JSON object".into())); + } + if item.get("id").and_then(Value::as_str).is_none() { + return Err(Error::Export("item missing string id".into())); + } + Ok(()) +} + +/// Accumulates per-partition temporal + spatial footprint from streamed items. +#[derive(Debug, Default)] +struct Footprint { + dt_min: Option, + dt_max: Option, + edt_min: Option, + edt_max: Option, + bbox: Option<[f64; 4]>, +} + +impl Footprint { + fn observe(&mut self, item: &Value) { + let props = &item["properties"]; + let dt = props + .get("datetime") + .and_then(Value::as_str) + .or_else(|| props.get("start_datetime").and_then(Value::as_str)); + if let Some(dt) = dt { + observe_min(&mut self.dt_min, dt); + observe_max(&mut self.dt_max, dt); + } + let edt = props + .get("end_datetime") + .and_then(Value::as_str) + .or_else(|| props.get("datetime").and_then(Value::as_str)); + if let Some(edt) = edt { + observe_min(&mut self.edt_min, edt); + observe_max(&mut self.edt_max, edt); + } + if let Some(bbox) = item.get("bbox").and_then(Value::as_array) + && bbox.len() >= 4 + { + let vals: Vec = bbox.iter().filter_map(Value::as_f64).collect(); + if vals.len() >= 4 { + // bbox may be 4 or 6 (3D): west/south are [0],[1]; east/north sit at len/2 and + // len/2+1 (indices 2,3 for a 4-element bbox; 3,4 for a 6-element 3D bbox). + let mid = vals.len() / 2; + let (w, s) = (vals[0], vals[1]); + let (e, n) = (vals[mid], vals[mid + 1]); + self.merge_bbox(w, s, e, n); + } + } + } + + fn merge_bbox(&mut self, w: f64, s: f64, e: f64, n: f64) { + match &mut self.bbox { + Some(b) => { + b[0] = b[0].min(w); + b[1] = b[1].min(s); + b[2] = b[2].max(e); + b[3] = b[3].max(n); + } + None => self.bbox = Some([w, s, e, n]), + } + } + + fn datetime_range(&self) -> Option { + match (&self.dt_min, &self.dt_max) { + (Some(lo), Some(hi)) => Some(DatetimeRange { + start: lo.clone(), + end: hi.clone(), + }), + _ => None, + } + } + + fn end_datetime_range(&self) -> Option { + match (&self.edt_min, &self.edt_max) { + (Some(lo), Some(hi)) => Some(DatetimeRange { + start: lo.clone(), + end: hi.clone(), + }), + _ => None, + } + } + + fn bbox(&self) -> Option> { + self.bbox.map(|b| b.to_vec()) + } +} + +fn observe_min(slot: &mut Option, v: &str) { + match slot { + Some(cur) if cur.as_str() <= v => {} + _ => *slot = Some(v.to_string()), + } +} + +fn observe_max(slot: &mut Option, v: &str) { + match slot { + Some(cur) if cur.as_str() >= v => {} + _ => *slot = Some(v.to_string()), + } +} + +/// Translates a plan-level [`Prefilter`] into a per-item [`ScanFilter`]. +fn scan_filter_for(prefilter: &Prefilter) -> ScanFilter { + ScanFilter { + datetime: prefilter.datetime.clone(), + bbox: prefilter.bbox, + } +} + +/// Percent-encodes path-unsafe characters in a collection id for use as a +/// directory name. The authoritative id is always the JSON +/// `id`, never parsed from the dir name. +fn encode_collection_dir(id: &str) -> String { + let mut out = String::with_capacity(id.len()); + for b in id.bytes() { + let safe = b.is_ascii_alphanumeric() || matches!(b, b'-' | b'_' | b'.'); + if safe { + out.push(b as char); + } else { + out.push('%'); + out.push_str(&format!("{b:02X}")); + } + } + out +} + +/// RFC3339 (UTC, second precision) timestamp via Postgres-independent clock. +fn now_rfc3339() -> String { + // Avoid a chrono dep: format from SystemTime. + let now = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_secs(); + // Days/time arithmetic for a UTC RFC3339 timestamp. + let days = now / 86_400; + let secs_of_day = now % 86_400; + let (h, m, s) = ( + secs_of_day / 3600, + (secs_of_day % 3600) / 60, + secs_of_day % 60, + ); + let (y, mo, d) = civil_from_days(days as i64); + format!("{y:04}-{mo:02}-{d:02}T{h:02}:{m:02}:{s:02}Z") +} + +/// Days-since-epoch -> (year, month, day), Howard Hinnant's civil algorithm. +fn civil_from_days(z: i64) -> (i64, u32, u32) { + let z = z + 719_468; + let era = if z >= 0 { z } else { z - 146_096 } / 146_097; + let doe = z - era * 146_097; + let yoe = (doe - doe / 1460 + doe / 36_524 - doe / 146_096) / 365; + let y = yoe + era * 400; + let doy = doe - (365 * yoe + yoe / 4 - yoe / 100); + let mp = (5 * doy + 2) / 153; + let d = (doy - (153 * mp + 2) / 5 + 1) as u32; + let m = (if mp < 10 { mp + 3 } else { mp - 9 }) as u32; + (if m <= 2 { y + 1 } else { y }, m, d) +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + fn feat(id: &str) -> Value { + json!({ + "type": "Feature", + "stac_version": "1.0.0", + "id": id, + "collection": "c", + "geometry": {"type": "Point", "coordinates": [1.0, 2.0]}, + "bbox": [1.0, 2.0, 1.0, 2.0], + "properties": {"datetime": "2020-01-01T00:00:00Z"}, + "assets": {}, + "links": [] + }) + } + + fn read_geoparquet_ids(path: &std::path::Path) -> Vec { + let f = std::fs::File::open(path).unwrap(); + let ic = stac::geoparquet::from_reader(f).unwrap(); + ic.items + .iter() + .map(|it| { + serde_json::to_value(it).unwrap()["id"] + .as_str() + .unwrap() + .to_string() + }) + .collect() + } + + /// Spill: a budget far smaller than the data forces the buffered + /// path to spill to disk, yet every item is still encoded (no drops). The + /// encode replays the spill + in-memory tail, so all ids round-trip. + #[test] + fn buffered_items_spills_under_tiny_budget_and_loses_nothing() { + // A 1-byte budget: every batch fails to reserve -> everything spills. + let budget = MemoryBudget::with_bytes(1); + let mut buf = BufferedItems::new(&budget); + let ids: Vec = (0..50).map(|i| format!("item-{i:03}")).collect(); + // Feed in several batches. + for chunk in ids.chunks(7) { + let batch: Vec = chunk.iter().map(|id| feat(id)).collect(); + buf.extend(batch).unwrap(); + } + // Budget fully released after spilling (no reservation held). + assert_eq!(budget.used(), 0, "spill must release the reservation"); + let out = crate::export::budget::spill_file().unwrap(); + let produced = buf.encode_to(out.path(), default_compression()).unwrap(); + assert!(produced, "non-empty buffer must produce a file"); + let mut got = read_geoparquet_ids(out.path()); + got.sort(); + let mut want = ids.clone(); + want.sort(); + assert_eq!(got, want, "spill+encode must preserve every item"); + } + + /// With a generous budget nothing spills; items stay in memory and encode. + #[test] + fn buffered_items_in_memory_when_budget_fits() { + let budget = MemoryBudget::with_bytes(64 * 1024 * 1024); + let mut buf = BufferedItems::new(&budget); + buf.extend(vec![feat("a"), feat("b")]).unwrap(); + assert!(budget.used() > 0, "in-memory buffer holds a reservation"); + let out = crate::export::budget::spill_file().unwrap(); + assert!(buf.encode_to(out.path(), default_compression()).unwrap()); + let mut got = read_geoparquet_ids(out.path()); + got.sort(); + assert_eq!(got, vec!["a".to_string(), "b".to_string()]); + } + + #[test] + fn buffered_items_empty_produces_no_file() { + let budget = MemoryBudget::with_bytes(1024); + let buf = BufferedItems::new(&budget); + let out = crate::export::budget::spill_file().unwrap(); + assert!(!buf.encode_to(out.path(), default_compression()).unwrap()); + } + + #[test] + fn dir_encoding() { + assert_eq!(encode_collection_dir("landsat-c2-l2"), "landsat-c2-l2"); + assert_eq!(encode_collection_dir("a b/c"), "a%20b%2Fc"); + assert_eq!(encode_collection_dir("ok.id_1"), "ok.id_1"); + } + + #[test] + fn footprint_min_max() { + let mut fp = Footprint::default(); + fp.observe(&serde_json::json!({ + "properties": {"datetime": "2024-01-10T00:00:00Z"}, + "bbox": [-10.0, -5.0, 10.0, 5.0] + })); + fp.observe(&serde_json::json!({ + "properties": {"datetime": "2024-01-05T00:00:00Z"}, + "bbox": [-20.0, -1.0, 2.0, 8.0] + })); + let dr = fp.datetime_range().unwrap(); + assert_eq!(dr.start, "2024-01-05T00:00:00Z"); + assert_eq!(dr.end, "2024-01-10T00:00:00Z"); + assert_eq!(fp.bbox().unwrap(), vec![-20.0, -5.0, 10.0, 8.0]); + } + + #[test] + fn civil_epoch() { + // 1970-01-01 + assert_eq!(civil_from_days(0), (1970, 1, 1)); + // 2000-01-01 = 10957 days after epoch + assert_eq!(civil_from_days(10_957), (2000, 1, 1)); + } +} diff --git a/src/pgstac-rs/src/export/sink.rs b/src/pgstac-rs/src/export/sink.rs new file mode 100644 index 00000000..449c43f7 --- /dev/null +++ b/src/pgstac-rs/src/export/sink.rs @@ -0,0 +1,453 @@ +//! Output sinks: where a dump's encoded bytes go. Format ⟂ Sink. +//! +//! A [`Sink`] accepts whole small files ([`Sink::put`], for the JSON metadata) +//! and finished large files streamed from a local temp file +//! ([`Sink::put_from_path`], for geoparquet). Streaming the large file from disk +//! keeps its bytes off the memory budget and lets object_store use S3 +//! multipart. +//! +//! Each write returns a [`FileWritten`] (SHA-256 + byte count) so the manifest +//! can record integrity without a second pass. +//! +//! Implementations: +//! * [`ObjectStoreSink`] (feature `store`) — local dir + S3/GCS/Azure via +//! `object_store` (`parse_url_opts`); large files use multipart. +//! * [`TarSink`] — a sequential `.tar` (default) or `.tar.zst` archive; the dump +//! tree packed into one file. Sequential, so writes to it serialize. +//! * [`StdoutSink`] — passthrough for single-stream output (search/NDJSON). + +#[cfg(any(feature = "store", feature = "cli"))] +use crate::Error; +use crate::Result; +use crate::export::manifest::{Sha256Writer, sha256_hex}; +use async_trait::async_trait; +use std::io::{Read, Write}; +use std::path::{Path, PathBuf}; +#[cfg(feature = "cli")] +use std::sync::Mutex; + +/// Integrity metadata for a written file. +#[derive(Debug, Clone)] +pub struct FileWritten { + /// Lowercase hex SHA-256 of the bytes. + pub sha256: String, + /// Number of bytes written. + pub bytes: u64, +} + +/// A dump output destination. +/// +/// The methods take a *relative path within the dump root* (e.g. +/// `collections/landsat-c2-l2/202401.parquet`). The trait is async so the remote +/// [`ObjectStoreSink`] drives `object_store` directly; the local sinks run their +/// synchronous file IO inline. +#[async_trait] +pub trait Sink: Send + Sync { + /// Writes a whole small file from an in-memory buffer. + async fn put(&self, rel_path: &str, bytes: &[u8]) -> Result; + + /// Writes a finished large file by streaming it from a local path. + async fn put_from_path(&self, rel_path: &str, src: &Path) -> Result; + + /// Removes a previously written file if the sink supports it (used to drop + /// `_checkpoint.json` on success). No-op for append-only sinks (tar/stdout). + async fn remove(&self, _rel_path: &str) -> Result<()> { + Ok(()) + } + + /// Flushes/finalizes the sink (e.g. closes a tar archive). Idempotent. + async fn finalize(&self) -> Result<()> { + Ok(()) + } +} + +// --------------------------------------------------------------------------- +// Stdout sink (single-stream passthrough) +// --------------------------------------------------------------------------- + +/// Writes all bytes to stdout (single-stream output; ignores `rel_path`). Useful +/// for `pgstac dump --format ndjson -` style piping. +#[derive(Debug, Default)] +pub struct StdoutSink; + +#[async_trait] +impl Sink for StdoutSink { + async fn put(&self, _rel_path: &str, bytes: &[u8]) -> Result { + let mut out = std::io::stdout().lock(); + out.write_all(bytes)?; + out.flush()?; + Ok(FileWritten { + sha256: sha256_hex(bytes), + bytes: bytes.len() as u64, + }) + } + + async fn put_from_path(&self, _rel_path: &str, src: &Path) -> Result { + let mut f = std::fs::File::open(src)?; + let mut out = std::io::stdout().lock(); + let mut hasher = Sha256Writer::new(); + let mut buf = vec![0u8; 1 << 16]; + loop { + let n = f.read(&mut buf)?; + if n == 0 { + break; + } + hasher.update(&buf[..n]); + out.write_all(&buf[..n])?; + } + out.flush()?; + Ok(FileWritten { + bytes: hasher.bytes(), + sha256: hasher.finalize_hex(), + }) + } +} + +// --------------------------------------------------------------------------- +// object_store sink (local dir + S3/GCS/Azure) +// --------------------------------------------------------------------------- + +#[cfg(feature = "store")] +mod object_store_sink { + use super::*; + use object_store::{ObjectStore, PutPayload, WriteMultipart, path::Path as ObjPath}; + use std::sync::Arc; + use url::Url; + + /// Streams to any `object_store` backend (local fs, S3, GCS, Azure). + /// + /// Construct from a base URL (`file:///abs/dir`, `s3://bucket/prefix`, …); + /// relative dump paths are joined onto the base prefix. Large files use + /// multipart upload. + pub struct ObjectStoreSink { + store: Arc, + base_prefix: ObjPath, + /// Multipart chunk size (bytes). + chunk_size: usize, + } + + impl std::fmt::Debug for ObjectStoreSink { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ObjectStoreSink") + .field("base_prefix", &self.base_prefix) + .field("chunk_size", &self.chunk_size) + .finish() + } + } + + impl ObjectStoreSink { + /// Builds a sink from a base URL and `object_store` options + /// (region/endpoint/credentials). + pub fn from_url_opts( + base_url: &str, + options: impl IntoIterator, + ) -> Result { + let url = Url::parse(base_url) + .map_err(|e| Error::Export(format!("invalid sink url {base_url}: {e}")))?; + let (store, base_prefix) = object_store::parse_url_opts(&url, options)?; + Ok(ObjectStoreSink { + store: Arc::from(store), + base_prefix, + chunk_size: 8 * 1024 * 1024, + }) + } + + fn child(&self, rel_path: &str) -> ObjPath { + let mut p = self.base_prefix.clone(); + for part in rel_path.split('/').filter(|s| !s.is_empty()) { + p = p.child(part); + } + p + } + } + + #[async_trait] + impl Sink for ObjectStoreSink { + async fn put(&self, rel_path: &str, bytes: &[u8]) -> Result { + let location = self.child(rel_path); + let payload = PutPayload::from(bytes.to_vec()); + let _ = self.store.put(&location, payload).await?; + Ok(FileWritten { + sha256: sha256_hex(bytes), + bytes: bytes.len() as u64, + }) + } + + async fn put_from_path(&self, rel_path: &str, src: &Path) -> Result { + use tokio::io::AsyncReadExt; + let location = self.child(rel_path); + let upload = self.store.put_multipart(&location).await?; + let mut writer = WriteMultipart::new_with_chunk_size(upload, self.chunk_size); + let mut f = tokio::fs::File::open(src).await?; + let mut hasher = Sha256Writer::new(); + let mut buf = vec![0u8; self.chunk_size]; + loop { + let n = f.read(&mut buf).await?; + if n == 0 { + break; + } + hasher.update(&buf[..n]); + writer.write(&buf[..n]); + } + let _ = writer.finish().await?; + Ok(FileWritten { + bytes: hasher.bytes(), + sha256: hasher.finalize_hex(), + }) + } + + async fn remove(&self, rel_path: &str) -> Result<()> { + let location = self.child(rel_path); + match self.store.delete(&location).await { + Ok(()) => Ok(()), + Err(object_store::Error::NotFound { .. }) => Ok(()), + Err(e) => Err(e.into()), + } + } + } +} + +#[cfg(feature = "store")] +pub use object_store_sink::ObjectStoreSink; + +// --------------------------------------------------------------------------- +// Local directory sink (no object_store dependency; available with `export`) +// --------------------------------------------------------------------------- + +/// Writes the dump tree into a local directory. Available without the `store` +/// feature (no `object_store`); the planner uses this for `--out `. +#[derive(Debug)] +pub struct DirSink { + root: PathBuf, +} + +impl DirSink { + /// Creates a directory sink rooted at `root` (created if missing). + pub fn new(root: impl Into) -> Result { + let root = root.into(); + std::fs::create_dir_all(&root)?; + Ok(DirSink { root }) + } + + fn target(&self, rel_path: &str) -> PathBuf { + let mut p = self.root.clone(); + for part in rel_path.split('/').filter(|s| !s.is_empty()) { + p.push(part); + } + p + } +} + +#[async_trait] +impl Sink for DirSink { + async fn put(&self, rel_path: &str, bytes: &[u8]) -> Result { + let target = self.target(rel_path); + if let Some(parent) = target.parent() { + std::fs::create_dir_all(parent)?; + } + std::fs::write(&target, bytes)?; + Ok(FileWritten { + sha256: sha256_hex(bytes), + bytes: bytes.len() as u64, + }) + } + + async fn put_from_path(&self, rel_path: &str, src: &Path) -> Result { + let target = self.target(rel_path); + if let Some(parent) = target.parent() { + std::fs::create_dir_all(parent)?; + } + // Copy + hash in one pass. + let mut f = std::fs::File::open(src)?; + let mut out = std::fs::File::create(&target)?; + let mut hasher = Sha256Writer::new(); + let mut buf = vec![0u8; 1 << 16]; + loop { + let n = f.read(&mut buf)?; + if n == 0 { + break; + } + hasher.update(&buf[..n]); + out.write_all(&buf[..n])?; + } + out.flush()?; + Ok(FileWritten { + bytes: hasher.bytes(), + sha256: hasher.finalize_hex(), + }) + } + + async fn remove(&self, rel_path: &str) -> Result<()> { + let target = self.target(rel_path); + match std::fs::remove_file(&target) { + Ok(()) => Ok(()), + Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(()), + Err(e) => Err(e.into()), + } + } +} + +// --------------------------------------------------------------------------- +// tar sink (sequential archive; .tar default, .tar.zst optional) +// --------------------------------------------------------------------------- + +#[cfg(feature = "cli")] +mod tar_sink { + use super::*; + use std::fs::File; + + /// Writes the dump tree into a single sequential `.tar` archive (optionally + /// zstd-compressed). Sequential => writes serialize (the planner accepts this + /// for the tar sink; fs/object_store sinks stay concurrent). + pub struct TarSink { + inner: Mutex, + } + + enum TarInner { + Plain(tar::Builder), + Zstd(tar::Builder>), + Finalized, + } + + impl std::fmt::Debug for TarSink { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("TarSink").finish() + } + } + + impl TarSink { + /// Creates a tar sink writing to `path`. `zstd` toggles outer + /// compression (default: plain `.tar`). + pub fn new(path: impl AsRef, zstd_compress: bool) -> Result { + let file = File::create(path.as_ref())?; + let inner = if zstd_compress { + let enc = zstd::Encoder::new(file, 0)?; + TarInner::Zstd(tar::Builder::new(enc)) + } else { + TarInner::Plain(tar::Builder::new(file)) + }; + Ok(TarSink { + inner: Mutex::new(inner), + }) + } + + fn append(&self, rel_path: &str, bytes: &[u8]) -> Result<()> { + let mut header = tar::Header::new_gnu(); + header.set_size(bytes.len() as u64); + header.set_mode(0o644); + header.set_cksum(); + let mut guard = self + .inner + .lock() + .map_err(|_| Error::Export("tar sink poisoned".into()))?; + match &mut *guard { + TarInner::Plain(b) => b.append_data(&mut header, rel_path, bytes)?, + TarInner::Zstd(b) => b.append_data(&mut header, rel_path, bytes)?, + TarInner::Finalized => { + return Err(Error::Export("tar sink already finalized".into())); + } + } + Ok(()) + } + } + + #[async_trait] + impl Sink for TarSink { + async fn put(&self, rel_path: &str, bytes: &[u8]) -> Result { + self.append(rel_path, bytes)?; + Ok(FileWritten { + sha256: sha256_hex(bytes), + bytes: bytes.len() as u64, + }) + } + + async fn put_from_path(&self, rel_path: &str, src: &Path) -> Result { + // tar needs the size up front; read the (already-finished) file in. + // Hash as we read so we don't traverse twice. + let mut f = File::open(src)?; + let mut bytes = Vec::new(); + let _ = f.read_to_end(&mut bytes)?; + let written = FileWritten { + sha256: sha256_hex(&bytes), + bytes: bytes.len() as u64, + }; + self.append(rel_path, &bytes)?; + Ok(written) + } + + async fn finalize(&self) -> Result<()> { + let mut guard = self + .inner + .lock() + .map_err(|_| Error::Export("tar sink poisoned".into()))?; + let inner = std::mem::replace(&mut *guard, TarInner::Finalized); + match inner { + TarInner::Plain(b) => { + let _ = b.into_inner()?; + } + TarInner::Zstd(b) => { + let enc = b.into_inner()?; + let _ = enc.finish()?; + } + TarInner::Finalized => {} + } + Ok(()) + } + } +} + +#[cfg(feature = "cli")] +pub use tar_sink::TarSink; + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn dir_sink_put_and_put_from_path() { + let dir = tempfile::tempdir().unwrap(); + let sink = DirSink::new(dir.path()).unwrap(); + let w = sink.put("a/b/meta.json", b"{\"k\":1}").await.unwrap(); + assert_eq!(w.bytes, 7); + assert_eq!(w.sha256, sha256_hex(b"{\"k\":1}")); + let written = std::fs::read(dir.path().join("a/b/meta.json")).unwrap(); + assert_eq!(written, b"{\"k\":1}"); + + // put_from_path streams + hashes. + let mut tf = tempfile::NamedTempFile::new().unwrap(); + Write::write_all(&mut tf, b"parquet-bytes").unwrap(); + let w2 = sink.put_from_path("c/items.parquet", tf.path()).await.unwrap(); + assert_eq!(w2.bytes, 13); + assert_eq!(w2.sha256, sha256_hex(b"parquet-bytes")); + let copied = std::fs::read(dir.path().join("c/items.parquet")).unwrap(); + assert_eq!(copied, b"parquet-bytes"); + } + + #[cfg(feature = "cli")] + #[tokio::test] + async fn tar_sink_roundtrip() { + let dir = tempfile::tempdir().unwrap(); + let tar_path = dir.path().join("dump.tar"); + { + let sink = TarSink::new(&tar_path, false).unwrap(); + let _ = sink.put("manifest.json", b"{}").await.unwrap(); + let mut tf = tempfile::NamedTempFile::new().unwrap(); + Write::write_all(&mut tf, b"col").unwrap(); + let _ = sink + .put_from_path("collections/c/items.parquet", tf.path()) + .await + .unwrap(); + sink.finalize().await.unwrap(); + } + // Read back the tar entries. + let f = std::fs::File::open(&tar_path).unwrap(); + let mut ar = tar::Archive::new(f); + let mut names = Vec::new(); + for entry in ar.entries().unwrap() { + let entry = entry.unwrap(); + names.push(entry.path().unwrap().to_string_lossy().into_owned()); + } + assert!(names.contains(&"manifest.json".to_string())); + assert!(names.contains(&"collections/c/items.parquet".to_string())); + } +} diff --git a/src/pgstac-rs/src/json/canonical.rs b/src/pgstac-rs/src/json/canonical.rs new file mode 100644 index 00000000..7cbf5958 --- /dev/null +++ b/src/pgstac-rs/src/json/canonical.rs @@ -0,0 +1,467 @@ +//! Deterministic content hash — produces bytes identical to SQL `pgstac.jsonb_canonical_hash` +//! (parity-gated in `tests/canonical_parity.rs`), so an item hashed here and the same item hashed by the +//! SQL ingest path get the same `item_hash`. +//! +//! The document is flattened to leaf `(path, value)` rows, each rendered as `path value`, sorted by +//! path (byte order), joined by ``, and SHA-256'd. Numbers are encoded as their `float8` form and +//! RFC3339 timestamps as a canonical UTC instant, so the digest matches a value after the `float8` / +//! `timestamptz` column round-trip pgstac storage performs. Separators are the C0 control chars US/RS/FS +//! (`0x1F`/`0x1E`/`0x1D`), which do not occur in STAC JSON keys or strings. + +use serde_json::{Number, Value}; +use sha2::{Digest, Sha256}; +use std::fmt::Write as _; + +/// Path segment separator (Unit Separator). +const US: char = '\u{1F}'; +/// Path/value separator (Record Separator is used between leaves; this splits a leaf's path from its value). +const FS: char = '\u{1E}'; +/// Between-leaf separator. +const RS: char = '\u{1D}'; + +/// The 32-byte SHA-256 identity digest of `value` — equal to SQL `pgstac.jsonb_canonical_hash(value)`. +/// +/// Errors (loudly, never silently) if `value` contains a string shaped exactly like an RFC3339 timestamp +/// that is not a valid instant — a malformed datetime must surface, not be reinterpreted as a plain string. +/// +/// # Examples +/// +/// ``` +/// use serde_json::json; +/// let h = pgstac::canonical::jsonb_canonical_hash(&json!({"b": 1, "a": 2})).unwrap(); +/// assert_eq!(h.len(), 32); +/// // key order is irrelevant; datetime spelling is irrelevant +/// assert_eq!( +/// pgstac::canonical::jsonb_canonical_hash(&json!({"a": 2, "b": 1})).unwrap(), +/// h, +/// ); +/// ``` +pub fn jsonb_canonical_hash(value: &Value) -> crate::Result<[u8; 32]> { + // (path, encoded-value) for every leaf; borrow the input (paths are the only allocation). + let mut leaves: Vec<(String, String)> = Vec::new(); + collect_leaves(&mut leaves, String::new(), value)?; + // Sort by path in raw byte order (matches SQL `ORDER BY path COLLATE "C"`). Paths are unique per tree. + leaves.sort_by(|a, b| a.0.as_bytes().cmp(b.0.as_bytes())); + + let mut buf = String::new(); + for (i, (path, encoded)) in leaves.iter().enumerate() { + if i > 0 { + buf.push(RS); + } + buf.push_str(path); + buf.push(FS); + buf.push_str(encoded); + } + Ok(Sha256::digest(buf.as_bytes()).into()) +} + +/// Walks `value`, appending a `(path, encoded-value)` row for every leaf (scalar or empty container). +/// A path step is `` then `k` (object member) or `i` (0-based array index); the `k`/`i` +/// tags keep object key `"0"` distinct from array index 0, and `` keeps nesting distinct from a key +/// that itself contains a separator character. +fn collect_leaves(out: &mut Vec<(String, String)>, path: String, value: &Value) -> crate::Result<()> { + match value { + Value::Object(map) if !map.is_empty() => { + for (key, child) in map { + let mut child_path = path.clone(); + child_path.push(US); + child_path.push('k'); + child_path.push_str(key); + collect_leaves(out, child_path, child)?; + } + } + Value::Array(items) if !items.is_empty() => { + for (idx, child) in items.iter().enumerate() { + let mut child_path = path.clone(); + child_path.push(US); + child_path.push('i'); + let _ = write!(child_path, "{idx}"); + collect_leaves(out, child_path, child)?; + } + } + // Leaf: scalar, or an empty object/array. + _ => out.push((path, encode_value(value)?)), + } + Ok(()) +} + +/// Encodes a leaf value with a 1-char type tag (matching the SQL `CASE jsonb_typeof(...)`): +/// `n` / `b`true|false / `z` (null) / `t` / `s` / +/// `e` (empty object) / `a` (empty array). +fn encode_value(value: &Value) -> crate::Result { + Ok(match value { + Value::Number(n) => { + let mut s = String::from("n"); + write_pg_float8(&mut s, n); + s + } + Value::Bool(true) => "btrue".to_string(), + Value::Bool(false) => "bfalse".to_string(), + Value::Null => "z".to_string(), + Value::String(s) => match canonical_stac_datetime(s)? { + Some(instant) => format!("t{instant}"), + None => format!("s{s}"), + }, + // Only empty containers reach here (non-empty ones recurse in `collect_leaves`). + Value::Object(_) => "e".to_string(), + Value::Array(_) => "a".to_string(), + }) +} + +/// Returns the canonical instant of `s` (UTC, fixed 6-digit microseconds, `Z`) when `s` is datetime-shaped, +/// `Ok(None)` when it is not (a plain string), and an **error** when it is datetime-shaped but not a valid +/// instant. The error case is deliberately loud — a malformed datetime is never silently reinterpreted as a +/// string. Mirrors the SQL string branch: the same shape prefilter and a `::timestamptz` cast pinned to +/// `SET TIMEZONE='UTC'` (matching pgstac's `to_tstz` ingest), which raises on a shaped-but-invalid value. +/// +/// Forgiving like `to_tstz`: `YYYY-MM-DD`, an optional `T`/`t`/space time, an optional fractional part, and +/// an optional offset; offset-less values are assumed UTC. For ≤6 fractional digits the output is +/// byte-identical to the SQL side; >6-digit fractions are truncated here but rounded by PostgreSQL (out of +/// STAC contract) — the one known edge where the two could differ. +fn canonical_stac_datetime(s: &str) -> crate::Result> { + if !is_datetime_shaped(s.as_bytes()) { + return Ok(None); // not a datetime — a plain string, unchanged + } + // Rewrite to a strict RFC3339 string (uppercase `T`/`Z`, `±HH:MM` offset, `Z` when the offset is + // absent), then parse once — the same instant `to_tstz`'s UTC-pinned `::timestamptz` cast produces. + let parsed = chrono::DateTime::parse_from_rfc3339(&to_rfc3339_utc_assumed(s)).map_err(|e| { + crate::Error::Dehydrate(format!("timestamp-shaped item value {s:?} is not a valid instant: {e}")) + })?; + Ok(Some( + parsed + .with_timezone(&chrono::Utc) + .format("%Y-%m-%dT%H:%M:%S%.6fZ") + .to_string(), + )) +} + +/// Rewrites a datetime-shaped string (per [`is_datetime_shaped`]) into a strict RFC3339 string chrono can +/// parse: date-only → midnight; a `T`/`t`/space separator → `T`; a missing offset → `Z` (UTC assumed, as +/// pgstac's UTC-pinned ingest does); a `±HHMM` offset → `±HH:MM`; a lowercase `z` → `Z`. +fn to_rfc3339_utc_assumed(s: &str) -> String { + if s.len() == 10 { + return format!("{s}T00:00:00Z"); // date-only → midnight UTC + } + // Normalize the date/time separator (byte 10) to `T`. + let mut out = String::with_capacity(s.len() + 1); + out.push_str(&s[..10]); + out.push('T'); + out.push_str(&s[11..]); + // Fix the offset on the tail (everything after the `T`). + let tail = &out[11..]; + if tail.ends_with(['Z', 'z']) { + out.truncate(out.len() - 1); // drop the `Z`/`z` (ASCII, 1 byte) + out.push('Z'); + out + } else if tail.contains(['+', '-']) { + insert_offset_colon(&out) // `±HHMM` → `±HH:MM` (no-op if already `±HH:MM`) + } else { + out.push('Z'); // no offset → assume UTC + out + } +} + +/// Cheap byte check for `^\d{4}-\d{2}-\d{2}([Tt ]\d{2}:\d{2}:\d{2}(\.\d+)?([Zz]|[+-]\d{2}:?\d{2})?)?$` — the +/// same shape the SQL side uses to prefilter datetime strings. Date-only and offset-less forms are allowed +/// (offset-less is assumed UTC downstream); the dash anchor excludes bare years, compact dates, and the +/// volatile `now`/`today` specials that a bare `::timestamptz` cast would otherwise accept. +fn is_datetime_shaped(b: &[u8]) -> bool { + let digit = |i: usize| b.get(i).is_some_and(u8::is_ascii_digit); + // YYYY-MM-DD + if b.len() < 10 + || !(digit(0) && digit(1) && digit(2) && digit(3)) + || b[4] != b'-' + || !(digit(5) && digit(6)) + || b[7] != b'-' + || !(digit(8) && digit(9)) + { + return false; + } + if b.len() == 10 { + return true; // date-only + } + // separator (T/t/space) + HH:MM:SS + if !matches!(b[10], b'T' | b't' | b' ') + || !(digit(11) && digit(12)) + || b.get(13) != Some(&b':') + || !(digit(14) && digit(15)) + || b.get(16) != Some(&b':') + || !(digit(17) && digit(18)) + { + return false; + } + let mut i = 19; + if b.get(i) == Some(&b'.') { + i += 1; + let start = i; + while b.get(i).is_some_and(u8::is_ascii_digit) { + i += 1; + } + if i == start { + return false; // '.' with no digits + } + } + match b.get(i) { + None => true, // offset-less (assumed UTC) + Some(&b'Z') | Some(&b'z') => i + 1 == b.len(), + Some(&b'+') | Some(&b'-') => { + i += 1; + if !(digit(i) && digit(i + 1)) { + return false; + } + i += 2; + if b.get(i) == Some(&b':') { + i += 1; + } + digit(i) && digit(i + 1) && i + 2 == b.len() + } + _ => false, + } +} + +/// Inserts a colon into a trailing `±HHMM` numeric offset (`+0500` -> `+05:00`) so chrono's RFC3339 parser +/// accepts the offset-without-colon variant the SQL regex allows. Returns `s` unchanged otherwise. +fn insert_offset_colon(s: &str) -> String { + let b = s.as_bytes(); + let n = b.len(); + if n >= 5 + && (b[n - 5] == b'+' || b[n - 5] == b'-') + && b[n - 4].is_ascii_digit() + && b[n - 3].is_ascii_digit() + && b[n - 2].is_ascii_digit() + && b[n - 1].is_ascii_digit() + { + format!("{}:{}", &s[..n - 2], &s[n - 2..]) + } else { + s.to_string() + } +} + +/// Appends a JSON number as PostgreSQL `float8::text`. The canonical value for a number is +/// `(v)::float8::text`, so the number is cast through `float8` and printed with PostgreSQL's shortest +/// round-trip formatting (fixed-point in the usual range, scientific notation outside it). This is required +/// (not merely convenient): promoted numbers are stored in `float8` columns, so this is the value the +/// hydrated/reproducible form carries. +fn write_pg_float8(out: &mut String, n: &Number) { + let Some(f) = n.as_f64() else { + // A JSON number is always f64-representable via serde_json; defensively emit the raw text. + out.push_str(&n.to_string()); + return; + }; + write_float8(out, f); +} + +/// `format_float8` writing directly into `out` (the hot path; no result-String allocation). +fn write_float8(out: &mut String, f: f64) { + if f == 0.0 { + // jsonb numbers pass through `numeric`, which has no signed zero, so -0.0 becomes "0". + out.push('0'); + return; + } + if f.is_nan() { + out.push_str("NaN"); + return; + } + if f.is_infinite() { + out.push_str(if f < 0.0 { "-Infinity" } else { "Infinity" }); + return; + } + + let negative = f < 0.0; + // Rust's `{:e}` yields the shortest round-trip mantissa + base-10 exponent. + let sci = format!("{:e}", f.abs()); + let (mantissa, exp_str) = sci + .split_once('e') + .expect("scientific notation contains 'e'"); + let decexp: i32 = exp_str.parse().expect("valid base-10 exponent"); + let digits: String = mantissa.chars().filter(|&c| c != '.').collect(); + let ndigits = digits.len() as i32; + + if negative { + out.push('-'); + } + + // Fixed point for -4 <= decexp < 15; scientific outside that range (matches PostgreSQL). + if !(-4..15).contains(&decexp) { + out.push_str(&digits[..1]); + if ndigits > 1 { + out.push('.'); + out.push_str(&digits[1..]); + } + out.push('e'); + out.push(if decexp >= 0 { '+' } else { '-' }); + let abs_exp = decexp.unsigned_abs(); + if abs_exp < 10 { + out.push('0'); + } + let _ = write!(out, "{abs_exp}"); + } else if decexp >= 0 { + let int_digits = decexp + 1; + if ndigits <= int_digits { + out.push_str(&digits); + for _ in 0..(int_digits - ndigits) { + out.push('0'); + } + } else { + out.push_str(&digits[..int_digits as usize]); + out.push('.'); + out.push_str(&digits[int_digits as usize..]); + } + } else { + out.push_str("0."); + for _ in 0..(-decexp - 1) { + out.push('0'); + } + out.push_str(&digits); + } +} + +/// Formats `f` exactly as PostgreSQL `float8out` does (shortest round-trip, `%g`-style fixed/scientific +/// selection). Validated against the database in `tests/canonical_parity.rs`. Thin wrapper over +/// [`write_float8`] (the hot path writes into a shared buffer; this is for tests wanting a `String`). +#[cfg(test)] +fn format_float8(f: f64) -> String { + let mut out = String::new(); + write_float8(&mut out, f); + out +} + +#[cfg(test)] +mod tests { + use super::{ + canonical_stac_datetime, format_float8, is_datetime_shaped, jsonb_canonical_hash, + to_rfc3339_utc_assumed, + }; + use serde_json::json; + + fn hash(v: &serde_json::Value) -> [u8; 32] { + jsonb_canonical_hash(v).unwrap() + } + + #[test] + fn key_order_and_number_scale_are_irrelevant() { + assert_eq!(hash(&json!({"b": 1, "a": {"y": 1, "x": 2}})), hash(&json!({"a": {"x": 2, "y": 1}, "b": 1}))); + // 1.0 and 1 both canonicalize to float8 "1". + assert_eq!(hash(&json!({"n": 1.0})), hash(&json!({"n": 1}))); + } + + #[test] + fn datetime_forgiving_spellings_are_one_instant() { + // Every forgiving spelling of 2020-01-01T00:00:00Z hashes identically (mirrors to_tstz + UTC). + let base = hash(&json!({"d": "2020-01-01T00:00:00Z"})); + for form in [ + "2020-01-01", // date-only -> midnight UTC + "2020-01-01 00:00:00", // space separator, offset-less -> UTC + "2020-01-01T00:00:00", // T separator, offset-less -> UTC + "2020-01-01t00:00:00z", // lowercase t/z + "2020-01-01T00:00:00.000000Z", // explicit zero micros + "2020-01-01T05:00:00+05:00", // offset respected + "2020-01-01T05:00:00+0500", // offset without a colon + "2019-12-31T19:00:00-05:00", // offset that crosses midnight + ] { + assert_eq!(hash(&json!({ "d": form })), base, "{form:?}"); + } + } + + #[test] + fn classifier_accepts_forgiving_rejects_specials() { + for ok in [ + "2020-01-01", + "2020-01-01 00:00:00", + "2020-01-01T00:00:00", + "2020-01-01T00:00:00Z", + "2020-01-01t00:00:00z", + "2020-01-01T05:00:00+05:00", + "2020-01-01T05:00:00+0500", + "2020-06-15T12:34:56.789Z", + ] { + assert!(is_datetime_shaped(ok.as_bytes()), "{ok:?} should be datetime-shaped"); + } + for no in [ + "2020", // bare year (PostgreSQL rejects it too) + "2020-01", // year-month + "now", // volatile special + "today", // volatile special + "epoch", // special + "20200101", // compact (no dashes) + "hello", // plain string + "http://x/2020-01-01", // not anchored at the start + "2020-01-01T00:00", // time without seconds + "2020-01-01Txx:00:00", // non-numeric time + "2020-01-01T00:00:00+5", // truncated offset + ] { + assert!(!is_datetime_shaped(no.as_bytes()), "{no:?} should not be datetime-shaped"); + } + } + + #[test] + fn to_rfc3339_normalizes_to_parseable_utc() { + assert_eq!(to_rfc3339_utc_assumed("2020-01-01"), "2020-01-01T00:00:00Z"); + assert_eq!(to_rfc3339_utc_assumed("2020-01-01 12:00:00"), "2020-01-01T12:00:00Z"); + assert_eq!(to_rfc3339_utc_assumed("2020-01-01T12:00:00"), "2020-01-01T12:00:00Z"); + assert_eq!(to_rfc3339_utc_assumed("2020-01-01t12:00:00z"), "2020-01-01T12:00:00Z"); + assert_eq!(to_rfc3339_utc_assumed("2020-01-01T12:00:00Z"), "2020-01-01T12:00:00Z"); + assert_eq!(to_rfc3339_utc_assumed("2020-01-01T12:00:00+0500"), "2020-01-01T12:00:00+05:00"); + assert_eq!(to_rfc3339_utc_assumed("2020-01-01T12:00:00+05:00"), "2020-01-01T12:00:00+05:00"); + assert_eq!(to_rfc3339_utc_assumed("2020-01-01T12:00:00.5-05:00"), "2020-01-01T12:00:00.5-05:00"); + } + + #[test] + fn rejected_datetime_shaped_strings_stay_raw() { + // "now"/"today" must NOT normalize to the current time (non-deterministic); they stay raw strings. + assert_ne!(hash(&json!({"d": "now"})), hash(&json!({"d": "today"}))); + // a bare year stays raw, distinct from the normalized full date + assert_ne!(hash(&json!({"d": "2020"})), hash(&json!({"d": "2020-01-01"}))); + } + + #[test] + fn canonical_stac_datetime_classifies() { + assert_eq!(canonical_stac_datetime("hello").unwrap(), None); // not a datetime + assert_eq!( + canonical_stac_datetime("2020-01-01").unwrap(), + Some("2020-01-01T00:00:00.000000Z".to_string()) + ); + assert!(canonical_stac_datetime("2020-13-45").is_err()); // shaped but invalid -> loud + } + + #[test] + fn paths_do_not_collide() { + // object nesting vs a key containing '/' + assert_ne!(hash(&json!({"a/b": 1})), hash(&json!({"a": {"b": 1}}))); + // object key "0" vs array index 0 + assert_ne!(hash(&json!({"0": 1})), hash(&json!([1]))); + // empty object vs empty array vs absent + assert_ne!(hash(&json!({"a": {}})), hash(&json!({"a": []}))); + } + + #[test] + fn shaped_but_invalid_timestamp_errors_loudly() { + for bad in [ + "2020-13-45", // month 13, day 45 (date-only) + "2020-02-30", // Feb 30 + "2020-13-45T99:99:99Z", // every field out of range + "2020-01-01T25:00:00Z", // hour 25 + ] { + assert!(jsonb_canonical_hash(&json!({ "d": bad })).is_err(), "{bad:?} should error"); + } + // a non-datetime string is fine (classification, not an error) + assert!(jsonb_canonical_hash(&json!({"id": "item-2020"})).is_ok()); + } + + #[test] + fn float8_fixed_and_scientific() { + assert_eq!(format_float8(0.0), "0"); + assert_eq!(format_float8(-0.0), "0"); + assert_eq!(format_float8(1.0), "1"); + assert_eq!(format_float8(100.0), "100"); + assert_eq!(format_float8(0.1), "0.1"); + assert_eq!(format_float8(123.456), "123.456"); + assert_eq!(format_float8(0.0001), "0.0001"); + assert_eq!(format_float8(0.00001), "1e-05"); + assert_eq!(format_float8(1e14), "100000000000000"); + assert_eq!(format_float8(1e15), "1e+15"); + assert_eq!(format_float8(1e16), "1e+16"); + assert_eq!(format_float8(123456789012345.6), "123456789012345.6"); + assert_eq!(format_float8(-105.1019), "-105.1019"); + } +} diff --git a/src/pgstac-rs/src/json/geom.rs b/src/pgstac-rs/src/json/geom.rs new file mode 100644 index 00000000..bc3bc4f0 --- /dev/null +++ b/src/pgstac-rs/src/json/geom.rs @@ -0,0 +1,137 @@ +//! Reading the raw `items.geometry` column (EWKB on the wire) and converting it client-side. +//! +//! pgstac stores geometry as a PostGIS `geometry`; `search_plan`'s projection selects the column +//! unwrapped (`i.geometry`), so the binary protocol hands Rust the raw EWKB bytes. We convert to +//! GeoJSON **here, only for JSON output** — Arrow/Parquet output keeps the WKB and never touches this. + +use crate::Result; +use bytes::BytesMut; +use geozero::geojson::GeoJson; +use geozero::wkb::Ewkb as GeozeroEwkb; +use geozero::{CoordDimensions, ToJson, ToWkb}; +use serde_json::Value; +use std::error::Error; +use tokio_postgres::types::{FromSql, IsNull, ToSql, Type, to_sql_checked}; + +/// Encodes a GeoJSON geometry [`Value`] to EWKB bytes with SRID 4326 — the form the `items.geometry` +/// column expects on a binary COPY. Borrows `geometry` and serializes it once (no deep clone). +/// +/// # Examples +/// +/// ``` +/// use pgstac::geom::{geojson_to_ewkb, Ewkb}; +/// use serde_json::json; +/// let ewkb = geojson_to_ewkb(&json!({"type": "Point", "coordinates": [1.0, 2.0]})).unwrap(); +/// let round_trip = Ewkb(ewkb).to_geojson().unwrap(); +/// assert_eq!(round_trip["type"], "Point"); +/// assert_eq!(round_trip["coordinates"], json!([1, 2])); +/// ``` +pub fn geojson_to_ewkb(geometry: &Value) -> Result> { + let geojson = serde_json::to_string(geometry)?; + Ok(GeoJson(&geojson).to_ewkb(CoordDimensions::xy(), Some(4326))?) +} + +/// The EWKB bytes of an `items.geometry` value. +/// +/// The PostGIS `geometry` binary wire format *is* EWKB, so one type serves both directions: it reads +/// the raw bytes off the binary protocol ([`FromSql`]) — keeping them lets the dump path forward WKB to +/// Arrow/Parquet untouched while the JSON path converts to GeoJSON via [`Ewkb::to_geojson`] — and writes +/// them back verbatim on a binary COPY ([`ToSql`]). +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Ewkb(pub Vec); + +impl Ewkb { + /// Converts the EWKB to a GeoJSON geometry [`Value`]. + /// + /// # Examples + /// + /// ``` + /// use pgstac::geom::Ewkb; + /// // EWKB for POINT(1 2) (little-endian, no SRID flag). + /// let ewkb = Ewkb(hex_to_bytes("0101000000000000000000f03f0000000000000040")); + /// let gj = ewkb.to_geojson().unwrap(); + /// assert_eq!(gj["type"], "Point"); + /// assert_eq!(gj["coordinates"], serde_json::json!([1, 2])); + /// # fn hex_to_bytes(s: &str) -> Vec { + /// # (0..s.len()).step_by(2).map(|i| u8::from_str_radix(&s[i..i + 2], 16).unwrap()).collect() + /// # } + /// ``` + pub fn to_geojson(&self) -> Result { + let json = GeozeroEwkb(&self.0).to_json()?; + Ok(serde_json::from_str(&json)?) + } +} + +impl<'a> FromSql<'a> for Ewkb { + fn from_sql( + _ty: &Type, + raw: &'a [u8], + ) -> std::result::Result> { + Ok(Ewkb(raw.to_vec())) + } + + fn accepts(ty: &Type) -> bool { + ty.name() == "geometry" + } +} + +impl ToSql for Ewkb { + fn to_sql( + &self, + _ty: &Type, + out: &mut BytesMut, + ) -> std::result::Result> { + out.extend_from_slice(&self.0); + Ok(IsNull::No) + } + + fn accepts(ty: &Type) -> bool { + ty.name() == "geometry" + } + + to_sql_checked!(); +} + +#[cfg(test)] +mod tests { + use super::*; + + fn hex(s: &str) -> Vec { + (0..s.len()) + .step_by(2) + .map(|i| u8::from_str_radix(&s[i..i + 2], 16).unwrap()) + .collect() + } + + #[test] + fn point_ewkb_to_geojson() { + // POINT(1 2), little-endian WKB. + let geom = Ewkb(hex("0101000000000000000000f03f0000000000000040")); + let gj = geom.to_geojson().unwrap(); + assert_eq!(gj["type"], "Point"); + assert_eq!(gj["coordinates"], serde_json::json!([1, 2])); + } + + #[test] + fn ewkb_with_srid_to_geojson() { + // SRID=4326;POINT(19 10) — EWKB with the SRID flag (0x20000000) set. + let geom = Ewkb(hex("0101000020e610000000000000000033400000000000002440")); + let gj = geom.to_geojson().unwrap(); + assert_eq!(gj["type"], "Point"); + assert_eq!(gj["coordinates"], serde_json::json!([19, 10])); + } + + #[test] + fn geojson_to_ewkb_round_trips_with_srid() { + let geom = serde_json::json!({"type": "Point", "coordinates": [19.0, 10.0]}); + let ewkb = geojson_to_ewkb(&geom).unwrap(); + // SRID=4326 -> the EWKB encodes the SRID and round-trips to the same geometry. + assert_eq!( + ewkb, + hex("0101000020e610000000000000000033400000000000002440") + ); + let gj = Ewkb(ewkb).to_geojson().unwrap(); + assert_eq!(gj["type"], "Point"); + assert_eq!(gj["coordinates"], serde_json::json!([19, 10])); + } +} diff --git a/src/pgstac-rs/src/json/mod.rs b/src/pgstac-rs/src/json/mod.rs new file mode 100644 index 00000000..0a804353 --- /dev/null +++ b/src/pgstac-rs/src/json/mod.rs @@ -0,0 +1,6 @@ +//! Low-level JSON and column-conversion helpers (canonical hashing, raw jsonb, geometry, datetime). + +pub mod canonical; +pub mod geom; +pub mod rawjson; +pub mod temporal; diff --git a/src/pgstac-rs/src/json/rawjson.rs b/src/pgstac-rs/src/json/rawjson.rs new file mode 100644 index 00000000..2fef5603 --- /dev/null +++ b/src/pgstac-rs/src/json/rawjson.rs @@ -0,0 +1,78 @@ +//! Reading a `jsonb` column as raw bytes (no parse) for the byte-assembly output path. +//! +//! Parsing `assets`/`bbox` jsonb into a [`serde_json::Value`] is the dominant per-item cost for rich +//! items (the big `classification:bitfields` arrays) and it loses PostgreSQL `numeric` precision on +//! `bbox` (f64 round-trip). [`RawJson`] keeps the original JSON text, so the byte path can splice it in +//! verbatim and merge only where it must (see [`crate::feature`]). + +use serde_json::Value; +use serde_json::value::RawValue; +use std::error::Error; +use tokio_postgres::types::{FromSql, Type}; + +/// The raw JSON text of a `jsonb` value, kept unparsed. +#[derive(Debug, Clone)] +pub struct RawJson(pub Box); + +impl RawJson { + /// The raw JSON text. + pub fn get(&self) -> &str { + self.0.get() + } + + /// Parses the raw text into a [`Value`] (jsonb is always valid JSON; `Null` on the impossible + /// failure). + pub fn to_value(&self) -> Value { + serde_json::from_str(self.0.get()).unwrap_or(Value::Null) + } + + /// Builds a [`RawJson`] from a [`Value`] — used by tests and the `from_value` round trip. + /// + /// # Examples + /// + /// ``` + /// use pgstac::rawjson::RawJson; + /// use serde_json::json; + /// + /// let raw = RawJson::from_value(&json!({"a": 1})); + /// assert_eq!(raw.to_value(), json!({"a": 1})); + /// ``` + pub fn from_value(value: &Value) -> RawJson { + RawJson(RawValue::from_string(value.to_string()).expect("a Value serializes to valid JSON")) + } +} + +impl<'a> FromSql<'a> for RawJson { + fn from_sql(_ty: &Type, raw: &'a [u8]) -> Result> { + // jsonb wire format: a 1-byte version header (currently `1`) followed by the UTF-8 JSON text. + // `json` has no header. Strip a leading version byte only for jsonb. + let text = match raw.first() { + Some(1) if raw.len() > 1 => std::str::from_utf8(&raw[1..])?, + _ => std::str::from_utf8(raw)?, + }; + Ok(RawJson(RawValue::from_string(text.to_string())?)) + } + + fn accepts(ty: &Type) -> bool { + matches!(*ty, Type::JSONB | Type::JSON) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn round_trips_through_value() { + let raw = RawJson::from_value(&json!({"b": [1, 2, 3], "a": "x"})); + assert_eq!(raw.to_value(), json!({"b": [1, 2, 3], "a": "x"})); + } + + #[test] + fn preserves_high_precision_number_text() { + // A number beyond f64 precision survives verbatim (the bbox-precision motivation). + let raw = RawJson(RawValue::from_string("[-22.650799880000001]".to_string()).unwrap()); + assert_eq!(raw.get(), "[-22.650799880000001]"); + } +} diff --git a/src/pgstac-rs/src/json/temporal.rs b/src/pgstac-rs/src/json/temporal.rs new file mode 100644 index 00000000..7b574eeb --- /dev/null +++ b/src/pgstac-rs/src/json/temporal.rs @@ -0,0 +1,56 @@ +//! Rendering the raw `timestamptz` columns to STAC datetime text, client-side. +//! +//! pgstac's `tstz_to_stac_text` formats `YYYY-MM-DDTHH:MM:SS[.ffffff]Z` in UTC, trimming trailing +//! fractional zeros (and the dot when the fraction is all zeros). `search_plan`'s projection now hands +//! Rust the raw `timestamptz`, so we reproduce that rendering here instead of asking the server. + +use chrono::{DateTime, Utc}; + +/// Formats a `timestamptz` exactly like pgstac's `tstz_to_stac_text`. +/// +/// # Examples +/// +/// ``` +/// use chrono::{TimeZone, Utc}; +/// use pgstac::temporal::tstz_to_stac_text; +/// +/// let dt = Utc.with_ymd_and_hms(2013, 4, 19, 12, 34, 56).unwrap(); +/// assert_eq!(tstz_to_stac_text(dt), "2013-04-19T12:34:56Z"); +/// ``` +pub fn tstz_to_stac_text(value: DateTime) -> String { + // `%.6f` always emits 6 fractional digits; trim trailing zeros and a bare dot to match + // `trim(trailing '.' from trim(trailing '0' from to_char(..., 'US')))`. + let base = value.format("%Y-%m-%dT%H:%M:%S%.6f").to_string(); + let trimmed = base.trim_end_matches('0').trim_end_matches('.'); + format!("{trimmed}Z") +} + +#[cfg(test)] +mod tests { + use super::*; + use chrono::TimeZone; + + #[test] + fn whole_second_drops_fraction() { + let dt = Utc.with_ymd_and_hms(2013, 4, 19, 0, 0, 0).unwrap(); + assert_eq!(tstz_to_stac_text(dt), "2013-04-19T00:00:00Z"); + } + + #[test] + fn trailing_zeros_trimmed() { + // .120000 -> .12 + let dt = Utc + .with_ymd_and_hms(2013, 4, 19, 12, 34, 56) + .unwrap() + .with_timezone(&Utc) + + chrono::Duration::microseconds(120_000); + assert_eq!(tstz_to_stac_text(dt), "2013-04-19T12:34:56.12Z"); + } + + #[test] + fn full_microseconds_kept() { + let dt = Utc.with_ymd_and_hms(2013, 4, 19, 12, 34, 56).unwrap() + + chrono::Duration::microseconds(123_456); + assert_eq!(tstz_to_stac_text(dt), "2013-04-19T12:34:56.123456Z"); + } +} diff --git a/src/pgstac-rs/src/lib.rs b/src/pgstac-rs/src/lib.rs index e76c3961..9314daad 100644 --- a/src/pgstac-rs/src/lib.rs +++ b/src/pgstac-rs/src/lib.rs @@ -1,50 +1,56 @@ +#![deny(rustdoc::broken_intra_doc_links)] //! Rust interface for [pgstac](https://github.com/stac-utils/pgstac). //! //! # Examples //! -//! [Pgstac] is a trait to query a **pgstac** database. -//! It is implemented for anything that implements [tokio_postgres::GenericClient]: +//! [Client] wraps a connection and talks to a **pgstac** database, caching the database's hydration +//! invariants across calls: //! //! ```no_run -//! use pgstac::Pgstac; +//! use pgstac::Client; //! use tokio_postgres::NoTls; //! //! # tokio_test::block_on(async { //! let config = "postgresql://username:password@localhost:5432/postgis"; -//! let (client, connection) = tokio_postgres::connect(config, NoTls).await.unwrap(); +//! let (connection, conn) = tokio_postgres::connect(config, NoTls).await.unwrap(); //! tokio::spawn(async move { -//! if let Err(e) = connection.await { +//! if let Err(e) = conn.await { //! eprintln!("connection error: {}", e); //! } //! }); +//! let client = Client::new(connection); //! println!("{}", client.pgstac_version().await.unwrap()); //! # }) //! ``` //! -//! If you want to work in a transaction, you can do that too: +//! Reads and writes both go through the same `Client` via the [`stac::api`] client traits: //! //! ```no_run -//! use pgstac::Pgstac; +//! use pgstac::Client; //! use stac::Collection; +//! use stac::api::{ItemsClient, TransactionClient}; //! use tokio_postgres::NoTls; //! //! # tokio_test::block_on(async { //! let config = "postgresql://username:password@localhost:5432/postgis"; -//! let (mut client, connection) = tokio_postgres::connect(config, NoTls).await.unwrap(); +//! let (connection, conn) = tokio_postgres::connect(config, NoTls).await.unwrap(); //! tokio::spawn(async move { -//! if let Err(e) = connection.await { +//! if let Err(e) = conn.await { //! eprintln!("connection error: {}", e); //! } //! }); -//! let transaction = client.transaction().await.unwrap(); -//! transaction.add_collection(Collection::new("an-id", "a description")).await.unwrap(); -//! transaction.commit().await.unwrap(); +//! let mut client = Client::new(connection); +//! client.add_collection(Collection::new("an-id", "a description")).await.unwrap(); +//! let items = client.search(Default::default()).await.unwrap(); //! # }) //! ``` //! //! # Features //! -//! - `tls`: provide a function to create an unverified tls provider, which can be useful in some circumstances (see ) +//! - `pool`: a `deadpool`-backed connection pool (`PgstacPool`) with rustls TLS. +//! - `export`: the stac-geoparquet dump/export library. +//! - `cli`: the `pgstac` binary (implies `export` + `pool`). +//! - `python`: the Python extension module built with maturin (implies `pool` + `export`). #![deny( elided_lifetimes_in_paths, @@ -75,14 +81,34 @@ )] #![warn(missing_docs)] -mod client; -mod page; - -pub use client::Client; -pub use page::Page; -use serde::{Serialize, de::DeserializeOwned}; -use stac::api::{ItemCollection, Search}; -use tokio_postgres::{GenericClient, NoTls, Row, types::ToSql}; +mod api; +mod db; +#[cfg(feature = "export")] +pub mod export; +mod json; +mod load; +#[cfg(feature = "python")] +mod python; +mod read; + +// The modules above live in api/ db/ json/ load/ python/ read/ subdirectories; they are re-exported at +// the crate root so existing `pgstac::…` and internal `crate::…` paths are unchanged — public modules +// stay public, previously-private ones stay crate-internal. +pub use json::{canonical, geom, rawjson, temporal}; +pub use load::{dehydrate, fragment, ingest}; +#[cfg(feature = "export")] +pub use load::parquet_decode; +pub use read::{collections, feature, fields, hydrate, keyset, search, source}; +pub(crate) use load::field_registry; +#[cfg(feature = "pool")] +pub(crate) use db::tls; + +pub use db::client::Client; +pub use db::connect::{ConnectConfig, DEFAULT_APPLICATION_NAME, DEFAULT_SEARCH_PATH}; +#[cfg(feature = "pool")] +pub use db::pool::{DEFAULT_POOL_SIZE, PgstacPool, PoolOptions, PoolerMode}; +use stac::api::{ItemCollection, ItemsClient, Search}; +use tokio_postgres::NoTls; /// Crate-specific error enum. #[derive(Debug, thiserror::Error)] @@ -96,6 +122,38 @@ pub enum Error { #[error(transparent)] Stac(#[from] stac::Error), + /// [geozero::error::GeozeroError] + #[error(transparent)] + Geozero(#[from] geozero::error::GeozeroError), + + /// A malformed item could not be dehydrated. + #[error("dehydrate error: {0}")] + Dehydrate(String), + + /// [std::io::Error] + #[error(transparent)] + Io(#[from] std::io::Error), + + /// [deadpool_postgres::PoolError] + #[cfg(feature = "pool")] + #[error(transparent)] + Pool(#[from] deadpool_postgres::PoolError), + + /// [rustls::Error] + #[cfg(feature = "pool")] + #[error(transparent)] + Rustls(#[from] rustls::Error), + + /// A TLS configuration error. + #[cfg(feature = "pool")] + #[error("tls configuration error: {0}")] + Tls(String), + + /// [deadpool_postgres::BuildError] + #[cfg(feature = "pool")] + #[error(transparent)] + PoolBuild(#[from] deadpool_postgres::BuildError), + /// [tokio_postgres::Error] #[error(transparent)] TokioPostgres(#[from] tokio_postgres::Error), @@ -103,8 +161,34 @@ pub enum Error { /// [std::num::TryFromIntError] #[error(transparent)] TryFromInt(#[from] std::num::TryFromIntError), + + /// A malformed search / keyset pagination token. + #[error("invalid search token: {0}")] + InvalidToken(String), + + /// An export/dump error. + #[cfg(feature = "export")] + #[error("export error: {0}")] + Export(String), + + /// [object_store::Error] + #[cfg(feature = "store")] + #[error(transparent)] + ObjectStore(#[from] object_store::Error), + + /// [url::ParseError] + #[cfg(feature = "store")] + #[error(transparent)] + UrlParse(#[from] url::ParseError), } +// `clap` and `stac-io` are used only by the `pgstac` binary (a separate crate target), so the library +// would otherwise flag them as unused dependencies. +#[cfg(feature = "cli")] +use clap as _; +#[cfg(feature = "cli")] +use stac_io as _; + /// Crate-specific result type. pub type Result = std::result::Result; @@ -130,7 +214,13 @@ pub async fn search( mut search: Search, max_items: Option, ) -> Result { - let (client, connection) = tokio_postgres::connect(connection_string, NoTls).await?; + // Route through ConnectConfig so the pgstac startup search_path (and any PG* env gap-fill) applies, + // the same as the pool and CLI paths. + let config = ConnectConfig { + dsn: Some(connection_string.to_string()), + ..Default::default() + }; + let (client, connection) = config.to_pg_config()?.connect(NoTls).await?; let task = tokio::spawn(async move { if let Err(e) = connection.await { tracing::error!("pgstac connection error: {}", e); @@ -152,17 +242,23 @@ pub async fn search( search.items.limit = Some(max_items.try_into()?); } + let client = Client::new(client); loop { tracing::info!("Fetching page"); let page = client.search(search.clone()).await?; - let next_token = page.next_token(); + let next_token = page + .next + .as_ref() + .and_then(|token_map| token_map.get("token")) + .and_then(|token| token.as_str()) + .map(str::to_string); let has_next_token = next_token.is_some(); if let Some(token) = next_token { let _ = search .additional_fields .insert("token".into(), token.into()); } - for item in page.features { + for item in page.items { all_items.push(item); if let Some(max_items) = max_items && all_items.len() >= max_items @@ -186,236 +282,18 @@ pub async fn search( Ok(ItemCollection::new(all_items)?) } -/// Methods for working with **pgstac**. -#[allow(async_fn_in_trait)] -pub trait Pgstac: GenericClient { - /// Returns the **pgstac** version. - async fn pgstac_version(&self) -> Result { - self.pgstac_string("get_version", &[]).await - } - - /// Returns whether the **pgstac** database is readonly. - async fn readonly(&self) -> Result { - self.pgstac_bool("readonly", &[]).await - } - - /// Returns the value of the `context` **pgstac** setting. - /// - /// This setting defaults to "off". See [the **pgstac** - /// docs](https://github.com/stac-utils/pgstac/blob/main/docs/src/pgstac.md#pgstac-settings) - /// for more information on the settings and their meaning. - async fn context(&self) -> Result { - self.pgstac_string("get_setting", &[&"context"]) - .await - .map(|value| value == "on") - } - - /// Sets the value of a **pgstac** setting. - async fn set_pgstac_setting(&self, key: &str, value: &str) -> Result<()> { - self.execute( - "INSERT INTO pgstac_settings (name, value) VALUES ($1, $2) ON CONFLICT ON CONSTRAINT pgstac_settings_pkey DO UPDATE SET value = excluded.value;", - &[&key, &value], - ).await.map(|_| ()).map_err(Error::from) - } - - /// Fetches all collections. - async fn collections(&self) -> Result> { - self.pgstac_vec("all_collections", &[]).await - } - - /// Fetches a collection by id. - async fn collection(&self, id: &str) -> Result> { - self.pgstac_opt("get_collection", &[&id]).await - } - - /// Adds a collection. - async fn add_collection(&self, collection: T) -> Result<()> - where - T: Serialize, - { - let collection = serde_json::to_value(collection)?; - self.pgstac_void("create_collection", &[&collection]).await - } - - /// Adds or updates a collection. - async fn upsert_collection(&self, collection: T) -> Result<()> - where - T: Serialize, - { - let collection = serde_json::to_value(collection)?; - self.pgstac_void("upsert_collection", &[&collection]).await - } - - /// Updates all collection extents. - async fn update_collection_extents(&self) -> Result<()> { - self.pgstac_void("update_collection_extents", &[]).await - } - - /// Updates a collection. - async fn update_collection(&self, collection: T) -> Result<()> - where - T: Serialize, - { - let collection = serde_json::to_value(collection)?; - self.pgstac_void("update_collection", &[&collection]).await - } - - /// Deletes a collection. - async fn delete_collection(&self, id: &str) -> Result<()> { - self.pgstac_void("delete_collection", &[&id]).await - } - - /// Fetches an item. - async fn item(&self, id: &str, collection: Option<&str>) -> Result> { - self.pgstac_opt("get_item", &[&id, &collection]).await - } - - /// Adds an item. - async fn add_item(&self, item: T) -> Result<()> - where - T: Serialize, - { - let item = serde_json::to_value(item)?; - self.pgstac_void("create_item", &[&item]).await - } - - /// Adds items. - async fn add_items(&self, items: &[T]) -> Result<()> - where - T: Serialize, - { - let items = serde_json::to_value(items)?; - self.pgstac_void("create_items", &[&items]).await - } - - /// Updates an item. - async fn update_item(&self, item: T) -> Result<()> - where - T: Serialize, - { - let item = serde_json::to_value(item)?; - self.pgstac_void("update_item", &[&item]).await - } - - /// Upserts an item. - async fn upsert_item(&self, item: T) -> Result<()> - where - T: Serialize, - { - let item = serde_json::to_value(item)?; - self.pgstac_void("upsert_item", &[&item]).await - } - - /// Upserts items. - /// - /// To avoid having to iterate the entire slice to serialize, these items - /// must all be [serde_json::Value]. - async fn upsert_items(&self, items: &[T]) -> Result<()> - where - T: Serialize, - { - let items = serde_json::to_value(items)?; - self.pgstac_void("upsert_items", &[&items]).await - } - - /// Deletes an item. - async fn delete_item(&self, id: &str, collection: Option<&str>) -> Result<()> { - self.pgstac_void("delete_item", &[&id, &collection]).await - } - - /// Searches for items. - async fn search(&self, search: Search) -> Result { - let search = search.into_cql2_json()?; - let search = serde_json::to_value(search)?; - self.pgstac_value("search", &[&search]).await - } - - /// Runs a pgstac function. - async fn pgstac( - &self, - function: &str, - params: &[&(dyn ToSql + Sync)], - ) -> std::result::Result { - let param_string = (0..params.len()) - .map(|i| format!("${}", i + 1)) - .collect::>() - .join(", "); - let query = format!("SELECT * from pgstac.{function}({param_string})"); - self.query_one(&query, params).await - } - - /// Returns a string result from a pgstac function. - async fn pgstac_string( - &self, - function: &str, - params: &[&(dyn ToSql + Sync)], - ) -> Result { - let row = self.pgstac(function, params).await?; - row.try_get(function).map_err(Error::from) - } - - /// Returns a bool result from a pgstac function. - async fn pgstac_bool(&self, function: &str, params: &[&(dyn ToSql + Sync)]) -> Result { - let row = self.pgstac(function, params).await?; - row.try_get(function).map_err(Error::from) - } - - /// Returns a vector from a pgstac function. - async fn pgstac_vec(&self, function: &str, params: &[&(dyn ToSql + Sync)]) -> Result> - where - T: DeserializeOwned, - { - if let Some(value) = self.pgstac_opt(function, params).await? { - Ok(value) - } else { - Ok(Vec::new()) - } - } - - /// Returns an optional value from a pgstac function. - async fn pgstac_opt( - &self, - function: &str, - params: &[&(dyn ToSql + Sync)], - ) -> Result> - where - T: DeserializeOwned, - { - let row = self.pgstac(function, params).await?; - let option: Option = row.try_get(function)?; - let option = option.map(|v| serde_json::from_value(v)).transpose()?; - Ok(option) - } - - /// Returns a deserializable value from a pgstac function. - async fn pgstac_value(&self, function: &str, params: &[&(dyn ToSql + Sync)]) -> Result - where - T: DeserializeOwned, - { - let row = self.pgstac(function, params).await?; - let value = row.try_get(function)?; - serde_json::from_value(value).map_err(Error::from) - } - - /// Returns nothing from a pgstac function. - async fn pgstac_void(&self, function: &str, params: &[&(dyn ToSql + Sync)]) -> Result<()> { - let _ = self.pgstac(function, params).await?; - Ok(()) - } -} - -impl Pgstac for T where T: GenericClient {} - #[cfg(test)] pub(crate) mod tests { - use super::Pgstac; + use super::Client as PgstacClient; use geojson::Geometry; use rstest::{fixture, rstest}; use serde_json::{Map, json}; - use stac::api::{Fields, Filter, Search, Sortby}; + use stac::api::{ + CollectionsClient, Fields, Filter, ItemsClient, Search, Sortby, TransactionClient, + }; use stac::{Collection, Item}; use std::{ - ops::Deref, + ops::{Deref, DerefMut}, sync::{LazyLock, atomic::AtomicU16}, }; use tokio::sync::Mutex; @@ -425,7 +303,7 @@ pub(crate) mod tests { static MUTEX: LazyLock> = LazyLock::new(|| Mutex::new(())); struct TestClient { - client: Client, + client: PgstacClient, config: Config, dbname: String, } @@ -437,6 +315,17 @@ pub(crate) mod tests { .unwrap() } + /// Name of the clean, empty pgstac install used as the clone template for each test. + /// + /// Tests assume a *fresh* pgstac install (e.g. zero collections), so they must clone from a + /// known-clean template rather than from whatever database the maintenance connection + /// ([`config`]) points at — in local dev that database is often populated. Build the template + /// once with `scripts/pgstac-rs-test-db`, or override the name via `PGSTAC_RS_TEST_TEMPLATE`. + fn template() -> String { + std::env::var("PGSTAC_RS_TEST_TEMPLATE") + .unwrap_or_else(|_| "pgstac_rs_test_template".to_string()) + } + impl TestClient { async fn new(id: u16) -> TestClient { let dbname = format!("pgstac_test_{id}"); @@ -445,13 +334,30 @@ pub(crate) mod tests { let _mutex = MUTEX.lock().await; let (client, connection) = config.connect(NoTls).await.unwrap(); let _handle = tokio::spawn(async move { connection.await.unwrap() }); + // Clone from the clean template (not the maintenance db, which may be populated). + // Connecting via `config` (a different database) means the template has no active + // session, so CREATE DATABASE ... TEMPLATE succeeds. + let _ = client + .execute( + &format!("CREATE DATABASE {} TEMPLATE {}", dbname, template()), + &[], + ) + .await + .unwrap_or_else(|e| { + panic!( + "failed to create test db {dbname} from template {}: {e}. \ + Build the template first with `scripts/pgstac-rs-test-db` \ + (or set PGSTAC_RS_TEST_TEMPLATE).", + template() + ) + }); + // `CREATE DATABASE ... TEMPLATE` copies the schema but NOT the template's per-database + // settings (pg_db_role_setting), so the clone loses the template's + // `search_path = pgstac, public`. Re-apply it, or the unqualified refs inside the pgstac + // functions resolve against `public` only and every query fails. let _ = client .execute( - &format!( - "CREATE DATABASE {} TEMPLATE {}", - dbname, - config.get_dbname().unwrap() - ), + &format!("ALTER DATABASE {dbname} SET search_path TO pgstac, public"), &[], ) .await @@ -461,12 +367,71 @@ pub(crate) mod tests { let (client, connection) = test_config.dbname(&dbname).connect(NoTls).await.unwrap(); let _handle = tokio::spawn(async move { connection.await.unwrap() }); TestClient { - client, + client: PgstacClient::new(client), config, dbname, } } + /// Seeds items through the **Rust loader** (the single write path) on a fresh connection — + /// the test analogue of the removed SQL `create_item`/`upsert_item`. `policy` resolves id + /// collisions (Error = "add", Upsert = "upsert/update"). + async fn load( + &self, + items: Vec, + policy: crate::ingest::ConflictPolicy, + ) -> crate::Result<()> { + let mut config = self.config.clone(); + let (mut client, connection) = + config.dbname(&self.dbname).connect(NoTls).await.unwrap(); + let handle = tokio::spawn(connection); + client + .batch_execute("SET search_path TO pgstac, public") + .await?; + let schema = crate::dehydrate::DehydrateSchema::load(&client).await?; + let _ = crate::ingest::load_items(&mut client, items, &schema, policy).await?; + handle.abort(); + Ok(()) + } + + async fn add_item(&self, item: T) -> crate::Result<()> { + self.load( + vec![serde_json::to_value(item)?], + crate::ingest::ConflictPolicy::Error, + ) + .await + } + + async fn add_items(&self, items: &[T]) -> crate::Result<()> { + let values = items + .iter() + .map(serde_json::to_value) + .collect::, _>>()?; + self.load(values, crate::ingest::ConflictPolicy::Error) + .await + } + + async fn upsert_item(&self, item: T) -> crate::Result<()> { + self.load( + vec![serde_json::to_value(item)?], + crate::ingest::ConflictPolicy::Upsert, + ) + .await + } + + async fn upsert_items(&self, items: &[T]) -> crate::Result<()> { + let values = items + .iter() + .map(serde_json::to_value) + .collect::, _>>()?; + self.load(values, crate::ingest::ConflictPolicy::Upsert) + .await + } + + async fn update_item(&self, item: T) -> crate::Result<()> { + self.upsert_item(item).await + } + async fn terminate(&mut self) { let (client, connection) = self.config.connect(NoTls).await.unwrap(); let _handle = tokio::spawn(async move { connection.await.unwrap() }); @@ -499,16 +464,30 @@ pub(crate) mod tests { } impl Deref for TestClient { - type Target = Client; + type Target = PgstacClient; fn deref(&self) -> &Self::Target { &self.client } } + impl DerefMut for TestClient { + fn deref_mut(&mut self) -> &mut Self::Target { + &mut self.client + } + } + fn longmont() -> Geometry { Geometry::new_point(vec![-105.1019, 40.1672]) } + /// Extracts the keyset token string from an `ItemCollection` `next`/`prev` pagination map. + fn token(link: &Option>) -> Option { + link.as_ref() + .and_then(|map| map.get("token")) + .and_then(|value| value.as_str()) + .map(str::to_string) + } + #[fixture] fn id() -> u16 { static COUNTER: AtomicU16 = AtomicU16::new(0); @@ -547,7 +526,7 @@ pub(crate) mod tests { #[rstest] #[tokio::test] - async fn collections(#[future(awt)] client: TestClient) { + async fn collections(#[future(awt)] mut client: TestClient) { assert!(client.collections().await.unwrap().is_empty()); client .add_collection(Collection::new("an-id", "a description")) @@ -558,7 +537,7 @@ pub(crate) mod tests { #[rstest] #[tokio::test] - async fn add_collection_duplicate(#[future(awt)] client: TestClient) { + async fn add_collection_duplicate(#[future(awt)] mut client: TestClient) { assert!(client.collections().await.unwrap().is_empty()); let collection = Collection::new("an-id", "a description"); client.add_collection(collection.clone()).await.unwrap(); @@ -574,14 +553,14 @@ pub(crate) mod tests { collection.title = Some("a title".to_string()); client.upsert_collection(collection).await.unwrap(); assert_eq!( - client.collection("an-id").await.unwrap().unwrap()["title"], - "a title" + client.collection("an-id").await.unwrap().unwrap().title, + Some("a title".to_string()) ); } #[rstest] #[tokio::test] - async fn update_collection(#[future(awt)] client: TestClient) { + async fn update_collection(#[future(awt)] mut client: TestClient) { let mut collection = Collection::new("an-id", "a description"); client.add_collection(collection.clone()).await.unwrap(); assert!( @@ -590,15 +569,15 @@ pub(crate) mod tests { .await .unwrap() .unwrap() - .get("title") + .title .is_none() ); collection.title = Some("a title".to_string()); client.update_collection(collection).await.unwrap(); assert_eq!(client.collections().await.unwrap().len(), 1); assert_eq!( - client.collection("an-id").await.unwrap().unwrap()["title"], - "a title" + client.collection("an-id").await.unwrap().unwrap().title, + Some("a title".to_string()) ); } @@ -617,7 +596,7 @@ pub(crate) mod tests { #[rstest] #[tokio::test] - async fn delete_collection(#[future(awt)] client: TestClient) { + async fn delete_collection(#[future(awt)] mut client: TestClient) { let collection = Collection::new("an-id", "a description"); client.add_collection(collection.clone()).await.unwrap(); assert!(client.collection("an-id").await.unwrap().is_some()); @@ -633,10 +612,10 @@ pub(crate) mod tests { #[rstest] #[tokio::test] - async fn item(#[future(awt)] client: TestClient) { + async fn item(#[future(awt)] mut client: TestClient) { assert!( client - .item("an-id", Some("collection-id")) + .item("collection-id", "an-id") .await .unwrap() .is_none() @@ -650,17 +629,26 @@ pub(crate) mod tests { .additional_fields .insert("type".into(), "Feature".into()); client.add_item(item.clone()).await.unwrap(); - assert_eq!( - client - .item("an-id", Some("collection-id")) - .await - .unwrap() - .unwrap(), - serde_json::to_value(item).unwrap(), - ); + let got = client + .item("collection-id", "an-id") + .await + .unwrap() + .expect("item present"); + let got = serde_json::to_value(got).unwrap(); + // v0.10 re-renders the item from its columns rather than echoing the stored JSON: the + // timestamptz `datetime` is microsecond precision (the input nanoseconds are truncated) and an + // empty `assets` object is dropped, so assert the round-tripped identity + geometry rather than + // byte-equality with the input. + assert_eq!(got["id"], "an-id"); + assert_eq!(got["collection"], "collection-id"); + assert_eq!(got["type"], "Feature"); + assert_eq!(got["geometry"], serde_json::to_value(longmont()).unwrap()); + assert!(got["properties"].get("datetime").is_some()); client.update_collection_extents().await.unwrap(); } + // The Rust loader (now the only write path) errors on an item whose collection does not exist, + // rather than the legacy SQL `create_item` behavior of silently dropping it. #[rstest] #[tokio::test] async fn item_without_collection(#[future(awt)] client: TestClient) { @@ -670,7 +658,7 @@ pub(crate) mod tests { #[rstest] #[tokio::test] - async fn update_item(#[future(awt)] client: TestClient) { + async fn update_item(#[future(awt)] mut client: TestClient) { let collection = Collection::new("collection-id", "a description"); client.add_collection(collection).await.unwrap(); let mut item = Item::new("an-id"); @@ -682,19 +670,20 @@ pub(crate) mod tests { .additional_fields .insert("foo".into(), "bar".into()); client.update_item(item).await.unwrap(); - assert_eq!( + let got = serde_json::to_value( client - .item("an-id", Some("collection-id")) + .item("collection-id", "an-id") .await .unwrap() - .unwrap()["properties"]["foo"], - "bar" - ); + .unwrap(), + ) + .unwrap(); + assert_eq!(got["properties"]["foo"], "bar"); } #[rstest] #[tokio::test] - async fn delete_item(#[future(awt)] client: TestClient) { + async fn delete_item(#[future(awt)] mut client: TestClient) { let collection = Collection::new("collection-id", "a description"); client.add_collection(collection).await.unwrap(); let mut item = Item::new("an-id"); @@ -706,14 +695,14 @@ pub(crate) mod tests { .await .unwrap(); assert_eq!( - client.item("an-id", Some("collection-id")).await.unwrap(), + client.item("collection-id", "an-id").await.unwrap(), None, ); } #[rstest] #[tokio::test] - async fn upsert_item(#[future(awt)] client: TestClient) { + async fn upsert_item(#[future(awt)] mut client: TestClient) { let collection = Collection::new("collection-id", "a description"); client.add_collection(collection).await.unwrap(); let mut item = Item::new("an-id"); @@ -725,7 +714,7 @@ pub(crate) mod tests { #[rstest] #[tokio::test] - async fn add_items(#[future(awt)] client: TestClient) { + async fn add_items(#[future(awt)] mut client: TestClient) { let collection = Collection::new("collection-id", "a description"); client.add_collection(collection).await.unwrap(); let mut item = Item::new("an-id"); @@ -736,14 +725,14 @@ pub(crate) mod tests { client.add_items(&[item, other_item]).await.unwrap(); assert!( client - .item("an-id", Some("collection-id")) + .item("collection-id", "an-id") .await .unwrap() .is_some() ); assert!( client - .item("other-id", Some("collection-id")) + .item("collection-id", "other-id") .await .unwrap() .is_some() @@ -752,7 +741,7 @@ pub(crate) mod tests { #[rstest] #[tokio::test] - async fn upsert_items(#[future(awt)] client: TestClient) { + async fn upsert_items(#[future(awt)] mut client: TestClient) { let collection = Collection::new("collection-id", "a description"); client.add_collection(collection).await.unwrap(); let mut item = Item::new("an-id"); @@ -767,13 +756,13 @@ pub(crate) mod tests { #[rstest] #[tokio::test] - async fn search_everything(#[future(awt)] client: TestClient) { + async fn search_everything(#[future(awt)] mut client: TestClient) { assert!( client .search(Search::default()) .await .unwrap() - .features + .items .is_empty() ); let collection = Collection::new("collection-id", "a description"); @@ -782,15 +771,17 @@ pub(crate) mod tests { item.collection = Some("collection-id".to_string()); item.geometry = Some(longmont()); client.add_item(item.clone()).await.unwrap(); - assert_eq!( - client.search(Search::default()).await.unwrap().features[0], - *serde_json::to_value(item).unwrap().as_object().unwrap() - ); + // See `item`: v0.10 re-renders from columns, so assert identity + geometry, not byte-equality. + let page = client.search(Search::default()).await.unwrap(); + let got = serde_json::to_value(&page.items[0]).unwrap(); + assert_eq!(got["id"], "an-id"); + assert_eq!(got["collection"], "collection-id"); + assert_eq!(got["geometry"], serde_json::to_value(longmont()).unwrap()); } #[rstest] #[tokio::test] - async fn search_ids(#[future(awt)] client: TestClient) { + async fn search_ids(#[future(awt)] mut client: TestClient) { let collection = Collection::new("collection-id", "a description"); client.add_collection(collection).await.unwrap(); let mut item = Item::new("an-id"); @@ -801,17 +792,17 @@ pub(crate) mod tests { ids: vec!["an-id".to_string()], ..Default::default() }; - assert_eq!(client.search(search).await.unwrap().features.len(), 1); + assert_eq!(client.search(search).await.unwrap().items.len(), 1); let search = Search { ids: vec!["not-an-id".to_string()], ..Default::default() }; - assert!(client.search(search).await.unwrap().features.is_empty()); + assert!(client.search(search).await.unwrap().items.is_empty()); } #[rstest] #[tokio::test] - async fn search_collections(#[future(awt)] client: TestClient) { + async fn search_collections(#[future(awt)] mut client: TestClient) { let collection = Collection::new("collection-id", "a description"); client.add_collection(collection).await.unwrap(); let mut item = Item::new("an-id"); @@ -822,17 +813,17 @@ pub(crate) mod tests { collections: vec!["collection-id".to_string()], ..Default::default() }; - assert_eq!(client.search(search).await.unwrap().features.len(), 1); + assert_eq!(client.search(search).await.unwrap().items.len(), 1); let search = Search { collections: vec!["not-an-id".to_string()], ..Default::default() }; - assert!(client.search(search).await.unwrap().features.is_empty()); + assert!(client.search(search).await.unwrap().items.is_empty()); } #[rstest] #[tokio::test] - async fn search_limit(#[future(awt)] client: TestClient) { + async fn search_limit(#[future(awt)] mut client: TestClient) { let collection = Collection::new("collection-id", "a description"); client.add_collection(collection).await.unwrap(); let mut item = Item::new("an-id"); @@ -844,7 +835,7 @@ pub(crate) mod tests { let mut search = Search::default(); search.items.limit = Some(1); let page = client.search(search).await.unwrap(); - assert_eq!(page.features.len(), 1); + assert_eq!(page.items.len(), 1); if let Some(context) = page.context { // v0.8 assert_eq!(context.limit.unwrap(), 1); @@ -856,7 +847,7 @@ pub(crate) mod tests { #[rstest] #[tokio::test] - async fn search_bbox(#[future(awt)] client: TestClient) { + async fn search_bbox(#[future(awt)] mut client: TestClient) { let collection = Collection::new("collection-id", "a description"); client.add_collection(collection).await.unwrap(); let mut item = Item::new("an-id"); @@ -866,16 +857,16 @@ pub(crate) mod tests { let mut search = Search::default(); search.items.bbox = Some(vec![-106., 40., -105., 41.].try_into().unwrap()); assert_eq!( - client.search(search.clone()).await.unwrap().features.len(), + client.search(search.clone()).await.unwrap().items.len(), 1 ); search.items.bbox = Some(vec![-106., 41., -105., 42.].try_into().unwrap()); - assert!(client.search(search).await.unwrap().features.is_empty()); + assert!(client.search(search).await.unwrap().items.is_empty()); } #[rstest] #[tokio::test] - async fn search_datetime(#[future(awt)] client: TestClient) { + async fn search_datetime(#[future(awt)] mut client: TestClient) { let collection = Collection::new("collection-id", "a description"); client.add_collection(collection).await.unwrap(); let mut item = Item::new("an-id"); @@ -886,16 +877,16 @@ pub(crate) mod tests { let mut search = Search::default(); search.items.datetime = Some("2023-01-07T00:00:00Z".to_string()); assert_eq!( - client.search(search.clone()).await.unwrap().features.len(), + client.search(search.clone()).await.unwrap().items.len(), 1 ); search.items.datetime = Some("2023-01-08T00:00:00Z".to_string()); - assert!(client.search(search).await.unwrap().features.is_empty()); + assert!(client.search(search).await.unwrap().items.is_empty()); } #[rstest] #[tokio::test] - async fn search_intersects(#[future(awt)] client: TestClient) { + async fn search_intersects(#[future(awt)] mut client: TestClient) { let collection = Collection::new("collection-id", "a description"); client.add_collection(collection).await.unwrap(); let mut item = Item::new("an-id"); @@ -918,7 +909,7 @@ pub(crate) mod tests { ), ..Default::default() }; - assert_eq!(client.search(search).await.unwrap().features.len(), 1); + assert_eq!(client.search(search).await.unwrap().items.len(), 1); let search = Search { intersects: Some( serde_json::from_value( @@ -935,12 +926,12 @@ pub(crate) mod tests { ), ..Default::default() }; - assert!(client.search(search).await.unwrap().features.is_empty()); + assert!(client.search(search).await.unwrap().items.is_empty()); } #[rstest] #[tokio::test] - async fn pagination(#[future(awt)] client: TestClient) { + async fn pagination(#[future(awt)] mut client: TestClient) { let collection = Collection::new("collection-id", "a description"); client.add_collection(collection).await.unwrap(); let mut item = Item::new("an-id"); @@ -954,22 +945,26 @@ pub(crate) mod tests { let mut search = Search::default(); search.items.limit = Some(1); let page = client.search(search.clone()).await.unwrap(); - assert_eq!(page.features[0]["id"], "an-id"); + assert_eq!(serde_json::to_value(&page.items[0]).unwrap()["id"], "an-id"); + // Page with the real keyset token the server minted; the old "collection-id:an-id" offset-style + // token format no longer exists in v0.10 (keyset_decode rejects it). + let next = token(&page.next).expect("next token"); let _ = search .additional_fields - .insert("token".to_string(), "next:collection-id:an-id".into()); + .insert("token".to_string(), next.into()); let page = client.search(search.clone()).await.unwrap(); - assert_eq!(page.features[0]["id"], "another-id"); + assert_eq!(serde_json::to_value(&page.items[0]).unwrap()["id"], "another-id"); + let prev = token(&page.prev).expect("prev token"); let _ = search .additional_fields - .insert("token".to_string(), "prev:collection-id:another-id".into()); + .insert("token".to_string(), prev.into()); let page = client.search(search).await.unwrap(); - assert_eq!(page.features[0]["id"], "an-id"); + assert_eq!(serde_json::to_value(&page.items[0]).unwrap()["id"], "an-id"); } #[rstest] #[tokio::test] - async fn base_url(#[future(awt)] client: TestClient) { + async fn base_url(#[future(awt)] mut client: TestClient) { client .set_pgstac_setting("base_url", "http://pgstac.test") .await @@ -994,7 +989,7 @@ pub(crate) mod tests { #[rstest] #[tokio::test] - async fn fields(#[future(awt)] client: TestClient) { + async fn fields(#[future(awt)] mut client: TestClient) { let collection = Collection::new("collection-id", "a description"); client.add_collection(collection).await.unwrap(); let mut item = Item::new("an-id"); @@ -1015,14 +1010,14 @@ pub(crate) mod tests { exclude: vec!["properties.bar".to_string()], }); let page = client.search(search).await.unwrap(); - let item = &page.features[0]; + let item = serde_json::to_value(&page.items[0]).unwrap(); assert!(item["properties"].as_object().unwrap().get("foo").is_some()); assert!(item["properties"].as_object().unwrap().get("bar").is_none()); } #[rstest] #[tokio::test] - async fn sortby(#[future(awt)] client: TestClient) { + async fn sortby(#[future(awt)] mut client: TestClient) { let collection = Collection::new("collection-id", "a description"); client.add_collection(collection).await.unwrap(); let mut item = Item::new("a"); @@ -1034,18 +1029,18 @@ pub(crate) mod tests { let mut search = Search::default(); search.items.sortby = vec![Sortby::asc("id")]; let page = client.search(search.clone()).await.unwrap(); - assert_eq!(page.features[0]["id"], "a"); - assert_eq!(page.features[1]["id"], "b"); + assert_eq!(serde_json::to_value(&page.items[0]).unwrap()["id"], "a"); + assert_eq!(serde_json::to_value(&page.items[1]).unwrap()["id"], "b"); search.items.sortby = vec![Sortby::desc("id")]; let page = client.search(search).await.unwrap(); - assert_eq!(page.features[0]["id"], "b"); - assert_eq!(page.features[1]["id"], "a"); + assert_eq!(serde_json::to_value(&page.items[0]).unwrap()["id"], "b"); + assert_eq!(serde_json::to_value(&page.items[1]).unwrap()["id"], "a"); } #[rstest] #[tokio::test] - async fn filter(#[future(awt)] client: TestClient) { + async fn filter(#[future(awt)] mut client: TestClient) { let collection = Collection::new("collection-id", "a description"); client.add_collection(collection).await.unwrap(); let mut item = Item::new("a"); @@ -1068,12 +1063,12 @@ pub(crate) mod tests { let mut search = Search::default(); search.items.filter = Some(Filter::Cql2Json(filter)); let page = client.search(search).await.unwrap(); - assert_eq!(page.features.len(), 1); + assert_eq!(page.items.len(), 1); } #[rstest] #[tokio::test] - async fn query(#[future(awt)] client: TestClient) { + async fn query(#[future(awt)] mut client: TestClient) { let collection = Collection::new("collection-id", "a description"); client.add_collection(collection).await.unwrap(); let mut item = Item::new("a"); @@ -1095,6 +1090,6 @@ pub(crate) mod tests { let mut search = Search::default(); search.items.query = Some(query); let page = client.search(search).await.unwrap(); - assert_eq!(page.features.len(), 1); + assert_eq!(page.items.len(), 1); } } diff --git a/src/pgstac-rs/src/load/dehydrate.rs b/src/pgstac-rs/src/load/dehydrate.rs new file mode 100644 index 00000000..a927043c --- /dev/null +++ b/src/pgstac-rs/src/load/dehydrate.rs @@ -0,0 +1,410 @@ +//! Rust-side dehydration of a full STAC item into pgstac's split `items`-row columns. +//! +//! This is the inverse of [`hydrate`](crate::hydrate): it takes a self-contained STAC item and produces +//! the column values pgstac stores — promoted queryable columns, the residual `properties`/`extra`, the +//! EWKB geometry, the temporal columns, and the canonical [`item_hash`](crate::canonical). It must match +//! the SQL `content_dehydrate` (003a_items.sql) so an item ingested through Rust is byte-identical to one +//! ingested through the SQL path; that is gated by `tests/dehydrate_parity.rs`. +//! +//! Fragment extraction (factoring shared content into `item_fragments`) is layered on top separately; +//! this module produces the pre-fragment row (`fragment_id = None`, full `assets`/`properties`/`links`), +//! exactly like `content_dehydrate`. +//! +//! Allocation discipline: the input item is consumed and its sub-values are **moved** out of the JSON +//! map (`Map::remove`); per-collection context (the promoted-column [`DehydrateSchema`]) is shared by +//! reference / `Arc`, never deep-cloned per item. + +use crate::canonical::jsonb_canonical_hash; +use crate::geom::geojson_to_ewkb; +use crate::{Error, Result}; +use chrono::{DateTime, Utc}; +use serde_json::{Map, Value}; +use tokio_postgres::GenericClient; + +/// Top-level item keys that map to dedicated `items` columns and so are removed before the remainder +/// becomes `extra`. Mirrors the `content - '{...}'` set in `content_dehydrate`. +const KNOWN_TOP_LEVEL: &[&str] = &[ + "id", + "geometry", + "collection", + "type", + "bbox", + "links", + "assets", + "properties", + "stac_version", + "stac_extensions", +]; + +/// The Postgres storage type of a promoted column, controlling how its property value is coerced. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum PromotedKind { + /// `text`. + Text, + /// `float8` / `float4` / `numeric`. + Float, + /// `int4` / `int2`. + Int, + /// `int8`. + BigInt, + /// `text[]`. + TextArray, + /// `jsonb` (stored verbatim). + Jsonb, + /// `timestamptz`. + Timestamptz, +} + +impl PromotedKind { + fn from_udt(udt: &str) -> PromotedKind { + match udt { + "text" | "varchar" | "bpchar" => PromotedKind::Text, + "float8" | "float4" | "numeric" => PromotedKind::Float, + "int4" | "int2" => PromotedKind::Int, + "int8" => PromotedKind::BigInt, + "_text" | "_varchar" => PromotedKind::TextArray, + "timestamptz" | "timestamp" => PromotedKind::Timestamptz, + _ => PromotedKind::Jsonb, + } + } +} + +/// One promoted queryable column: its STAC property name, the `items` column, and the column type. +#[derive(Debug, Clone)] +pub struct PromotedColumn { + /// STAC property name in `properties`, e.g. `eo:cloud_cover`. + pub stac_name: String, + /// `items` column name, e.g. `eo_cloud_cover`. + pub column: String, + /// The column's storage type. + pub kind: PromotedKind, +} + +/// The promoted-column schema, derived at runtime from `promoted_item_property_defs()` joined with the +/// `items` column types — the ordered `(stac_name, column, kind)` mapping the dehydrate promotes out of +/// `properties`. Loaded once per ingest and shared (by reference / `Arc`) across all items. +#[derive(Debug, Clone)] +pub struct DehydrateSchema { + columns: Vec, +} + +impl DehydrateSchema { + /// Loads the promoted-column schema from the live database. + pub async fn load(client: &C) -> Result { + let type_rows = client + .query( + "SELECT column_name, udt_name FROM information_schema.columns \ + WHERE table_schema = 'pgstac' AND table_name = 'items'", + &[], + ) + .await?; + let mut types = std::collections::HashMap::new(); + for row in &type_rows { + let name: String = row.get("column_name"); + let udt: String = row.get("udt_name"); + let _ = types.insert(name, udt); + } + + let def_rows = client + .query( + "SELECT name, property_path FROM promoted_item_property_defs() \ + WITH ORDINALITY AS d(name, definition, property_path, ord) ORDER BY ord", + &[], + ) + .await?; + let mut columns = Vec::with_capacity(def_rows.len()); + for row in &def_rows { + let stac_name: String = row.get("name"); + let column: String = row.get("property_path"); + let kind = types + .get(&column) + .map(|u| PromotedKind::from_udt(u)) + .unwrap_or(PromotedKind::Jsonb); + columns.push(PromotedColumn { + stac_name, + column, + kind, + }); + } + Ok(DehydrateSchema { columns }) + } + + /// The promoted columns in `promoted_item_property_defs()` order. + pub fn columns(&self) -> &[PromotedColumn] { + &self.columns + } +} + +/// A coerced promoted-column value, ready for a typed binary COPY into `items`. +#[derive(Debug, Clone, PartialEq)] +pub enum PromotedValue { + /// A `text` value (`None` = SQL NULL). + Text(Option), + /// A `float8` value. + Float(Option), + /// An `int4` value. + Int(Option), + /// An `int8` value. + BigInt(Option), + /// A `text[]` value. + TextArray(Option>), + /// A `jsonb` value (stored verbatim). + Jsonb(Option), + /// A `timestamptz` value. + Timestamptz(Option>), +} + +/// A fully dehydrated `items` row, ready for a binary COPY. Jsonb columns are kept as +/// [`Value`](serde_json::Value); `geometry` is EWKB (SRID 4326). +#[derive(Debug, Clone)] +pub struct DehydratedRow { + /// `items.id`. + pub id: String, + /// `items.collection`. + pub collection: String, + /// `items.geometry` as EWKB bytes (SRID 4326). + pub geometry: Vec, + /// `items.datetime` (start_datetime for ranges). + pub datetime: DateTime, + /// `items.end_datetime`. + pub end_datetime: DateTime, + /// `items.datetime_is_range`. + pub datetime_is_range: bool, + /// `items.stac_version`. + pub stac_version: Option, + /// `items.stac_extensions` (defaults to `[]`). + pub stac_extensions: Value, + /// `items.item_hash` — the canonical content hash. + pub item_hash: [u8; 32], + /// `items.fragment_id` — always `None` from this (pre-fragment) dehydrate. + pub fragment_id: Option, + /// `items.bbox`. + pub bbox: Option, + /// `items.links` (`None` when absent or empty `[]`). + pub links: Option, + /// `items.assets` (`None` when absent or empty `{}`). + pub assets: Option, + /// `items.properties` — the residual after promoting columns and removing temporal keys. + pub properties: Value, + /// `items.extra` — top-level keys outside the known column set. + pub extra: Value, + /// `items.link_hrefs` — the per-link `href`s (NULL entries preserved). + pub link_hrefs: Option>>, + /// Promoted queryable columns, in schema order, aligned with `schema.columns()`. + pub promoted: Vec, +} + +/// Dehydrates one self-contained STAC item into a [`DehydratedRow`], consuming `item`. +/// +/// Matches SQL `content_dehydrate(item)`: computes the canonical `item_hash` over the whole item, splits +/// the geometry to EWKB, derives the temporal columns (`stac_daterange` semantics), promotes the +/// queryable columns out of `properties`, and routes the remaining top-level keys to `extra`. +pub fn dehydrate(item: Value, schema: &DehydrateSchema) -> Result { + // The hash is over the full item, so compute it before deconstructing (borrow, no clone). + let item_hash = jsonb_canonical_hash(&item)?; + + let mut map = match item { + Value::Object(map) => map, + _ => return Err(Error::Dehydrate("item is not a JSON object".to_string())), + }; + + let id = + take_string(&mut map, "id").ok_or_else(|| Error::Dehydrate("item has no id".into()))?; + let collection = take_string(&mut map, "collection") + .ok_or_else(|| Error::Dehydrate(format!("item {id} has no collection")))?; + + let geometry_value = map + .remove("geometry") + .filter(|v| !v.is_null()) + .ok_or_else(|| Error::Dehydrate(format!("item {id} has no geometry")))?; + let geometry = geojson_to_ewkb(&geometry_value)?; + + let stac_version = take_string(&mut map, "stac_version"); + let stac_extensions = map + .remove("stac_extensions") + .filter(|v| !v.is_null()) + .unwrap_or_else(|| Value::Array(Vec::new())); + let bbox = map.remove("bbox").filter(|v| !v.is_null()); + let links_raw = map.remove("links"); + let assets = map + .remove("assets") + .filter(|v| !v.is_null() && v.as_object().is_none_or(|o| !o.is_empty())); + let _ = map.remove("type"); + + // `properties` is taken out as a map we can mutate (promote + strip temporal keys). + let mut properties = match map.remove("properties") { + Some(Value::Object(p)) => p, + _ => Map::new(), + }; + + // Temporal columns: stac_daterange semantics over `properties`. + let (datetime, end_datetime, datetime_is_range) = temporal(&properties, &id)?; + + // Promote each queryable column out of `properties` (removing it), coercing per the column type. + let mut promoted = Vec::with_capacity(schema.columns.len()); + for col in &schema.columns { + let raw = properties.remove(&col.stac_name).filter(|v| !v.is_null()); + promoted.push(coerce_promoted(raw, col.kind, &col.stac_name, &id)?); + } + // strip_promoted_properties also removes the temporal keys. + let _ = properties.remove("datetime"); + let _ = properties.remove("start_datetime"); + let _ = properties.remove("end_datetime"); + + // links column: NULL when absent or empty `[]`. Compute link_hrefs (borrow) before moving links_raw. + let link_hrefs = link_hrefs_of(links_raw.as_ref()); + let links = links_raw.filter(|v| v.as_array().is_some_and(|a| !a.is_empty())); + + // Everything left at the top level is `extra`. + for key in KNOWN_TOP_LEVEL { + let _ = map.remove(*key); + } + let extra = Value::Object(map); + + Ok(DehydratedRow { + id, + collection, + geometry, + datetime, + end_datetime, + datetime_is_range, + stac_version, + stac_extensions, + item_hash, + fragment_id: None, + bbox, + links, + assets, + properties: Value::Object(properties), + extra, + link_hrefs, + promoted, + }) +} + +/// Removes a string-valued key from `map`, returning the owned `String` (or `None` if absent/non-string). +fn take_string(map: &mut Map, key: &str) -> Option { + match map.remove(key) { + Some(Value::String(s)) => Some(s), + _ => None, + } +} + +/// Computes `(datetime, end_datetime, datetime_is_range)` from `properties`, matching `stac_daterange` +/// + `content_dehydrate`'s `datetime_is_range`. +fn temporal(props: &Map, id: &str) -> Result<(DateTime, DateTime, bool)> { + let datetime_present = props.get("datetime").is_some_and(|v| !v.is_null()); + let start = props.get("start_datetime").filter(|v| !v.is_null()); + let end = props.get("end_datetime").filter(|v| !v.is_null()); + + let (dt_v, edt_v) = if start.is_some() && end.is_some() { + (start, end) + } else { + let dt = props.get("datetime").filter(|v| !v.is_null()); + (dt, dt) + }; + let dt = parse_stac_datetime(dt_v, id)?; + let edt = parse_stac_datetime(edt_v, id)?; + + let datetime_is_range = if datetime_present { + false + } else { + start.is_some() || end.is_some() + }; + Ok((dt, edt, datetime_is_range)) +} + +/// Parses a STAC datetime [`Value`] (an RFC 3339 string) into a UTC timestamp. +fn parse_stac_datetime(v: Option<&Value>, id: &str) -> Result> { + let s = v.and_then(Value::as_str).ok_or_else(|| { + Error::Dehydrate(format!( + "item {id}: datetime, or both start_datetime and end_datetime, must be set" + )) + })?; + DateTime::parse_from_rfc3339(s) + .map(|dt| dt.with_timezone(&Utc)) + .map_err(|e| Error::Dehydrate(format!("item {id}: invalid datetime {s:?}: {e}"))) +} + +/// Coerces a promoted property value to its column type, matching `content_dehydrate`'s per-column casts. +fn coerce_promoted( + raw: Option, + kind: PromotedKind, + stac_name: &str, + id: &str, +) -> Result { + let bad = + |what: &str| Error::Dehydrate(format!("item {id}: property {stac_name} is not {what}")); + Ok(match (kind, raw) { + (PromotedKind::Text, Some(v)) => PromotedValue::Text(Some(value_to_text(v))), + (PromotedKind::Text, None) => PromotedValue::Text(None), + (PromotedKind::Float, Some(v)) => { + PromotedValue::Float(Some(value_to_f64(&v).ok_or_else(|| bad("a number"))?)) + } + (PromotedKind::Float, None) => PromotedValue::Float(None), + (PromotedKind::Int, Some(v)) => PromotedValue::Int(Some( + i32::try_from(value_to_i64(&v).ok_or_else(|| bad("an integer"))?) + .map_err(|_| bad("an integer within int32 range"))?, + )), + (PromotedKind::Int, None) => PromotedValue::Int(None), + (PromotedKind::BigInt, Some(v)) => { + PromotedValue::BigInt(Some(value_to_i64(&v).ok_or_else(|| bad("an integer"))?)) + } + (PromotedKind::BigInt, None) => PromotedValue::BigInt(None), + (PromotedKind::TextArray, Some(v)) => { + PromotedValue::TextArray(Some(value_to_text_array(v))) + } + (PromotedKind::TextArray, None) => PromotedValue::TextArray(None), + (PromotedKind::Timestamptz, Some(v)) => { + PromotedValue::Timestamptz(Some(parse_stac_datetime(Some(&v), id)?)) + } + (PromotedKind::Timestamptz, None) => PromotedValue::Timestamptz(None), + (PromotedKind::Jsonb, raw) => PromotedValue::Jsonb(raw), + }) +} + +/// `->>` semantics: a JSON string yields its contents; any other scalar yields its JSON text. +fn value_to_text(v: Value) -> String { + match v { + Value::String(s) => s, + other => other.to_string(), + } +} + +fn value_to_f64(v: &Value) -> Option { + match v { + Value::Number(n) => n.as_f64(), + Value::String(s) => s.parse().ok(), + _ => None, + } +} + +fn value_to_i64(v: &Value) -> Option { + match v { + Value::Number(n) => n.as_i64().or_else(|| n.as_f64().map(|f| f as i64)), + Value::String(s) => s.parse().ok(), + _ => None, + } +} + +/// `to_text_array`: an array yields each element as text; a scalar yields a one-element array. +fn value_to_text_array(v: Value) -> Vec { + match v { + Value::Array(items) => items.into_iter().map(value_to_text).collect(), + other => vec![value_to_text(other)], + } +} + +/// Extracts each link's `href` (NULL when absent), matching `stac_links_href_array`. Returns `None` when +/// there are no links. +fn link_hrefs_of(links: Option<&Value>) -> Option>> { + let arr = links?.as_array()?; + if arr.is_empty() { + return None; + } + Some( + arr.iter() + .map(|link| link.get("href").and_then(Value::as_str).map(str::to_string)) + .collect(), + ) +} diff --git a/src/pgstac-rs/src/load/field_registry.rs b/src/pgstac-rs/src/load/field_registry.rs new file mode 100644 index 00000000..00b43372 --- /dev/null +++ b/src/pgstac-rs/src/load/field_registry.rs @@ -0,0 +1,184 @@ +//! Client-side field registry: the per-collection set of JSON paths and value kinds the loader maintains. +//! +//! The loader walks EVERY ingested item's content — no sampling — to collect the COMPLETE set of JSON paths +//! and their value kinds per collection, then UPSERTs them ADD-ONLY into `item_field_registry` so it is +//! always a *superset* of the data (too-wide is fine; expiring stale paths is deferred async maintenance). +//! The walk mirrors the SQL `jsonb_field_rows` exactly (`.`-joined paths, leaf flag, `jsonb_typeof` value +//! kind, depth cap 20) so a Rust-loaded registry matches one the SQL path would produce. + +use serde_json::{Map, Value}; +use std::collections::{BTreeMap, BTreeSet}; + +/// Matches `jsonb_field_rows`' `max_depth` guard against pathologically nested documents. +const MAX_DEPTH: u32 = 20; + +/// The accumulated distinct paths (and their merged leaf flag + value kinds) observed across the items of a +/// single collection. `is_leaf` is AND-merged (a path is a leaf only if it never gains children); value kinds +/// only accumulate — so merging more items only ever widens, never narrows. +#[derive(Default)] +pub(crate) struct FieldRegistry { + paths: BTreeMap, +} + +#[derive(Default)] +struct Entry { + is_leaf: bool, + kinds: BTreeSet<&'static str>, +} + +impl FieldRegistry { + /// Walk one item's content, merging its paths into the registry. + pub(crate) fn observe(&mut self, content: &Value) { + let mut path = String::new(); + walk( + content, + &mut path, + MAX_DEPTH, + &mut |p, is_leaf, kind| match self.paths.get_mut(p) { + Some(e) => { + e.is_leaf = e.is_leaf && is_leaf; + let _ = e.kinds.insert(kind); + } + None => { + let mut kinds = BTreeSet::new(); + let _ = kinds.insert(kind); + let _ = self.paths.insert(p.to_string(), Entry { is_leaf, kinds }); + } + }, + ); + } + + pub(crate) fn is_empty(&self) -> bool { + self.paths.is_empty() + } + + /// Render as the JSONB array the loader's `item_field_registry` UPSERT expects: + /// `[{"path": text, "is_leaf": bool, "value_kinds": [text, ...]}, ...]`. + pub(crate) fn to_entries(&self) -> Value { + Value::Array( + self.paths + .iter() + .map(|(path, e)| { + let mut o = Map::new(); + let _ = o.insert("path".into(), Value::String(path.clone())); + let _ = o.insert("is_leaf".into(), Value::Bool(e.is_leaf)); + let _ = o.insert( + "value_kinds".into(), + Value::Array(e.kinds.iter().map(|k| Value::String((*k).into())).collect()), + ); + Value::Object(o) + }) + .collect(), + ) + } +} + +/// `jsonb_typeof` for a serde_json value. +fn kind(v: &Value) -> &'static str { + match v { + Value::Object(_) => "object", + Value::Array(_) => "array", + Value::String(_) => "string", + Value::Number(_) => "number", + Value::Bool(_) => "boolean", + Value::Null => "null", + } +} + +/// Recursively emit one `(path, is_leaf, value_kind)` per field path, mirroring SQL `jsonb_field_rows`: +/// objects emit each key (a container key emits its own non-leaf row AND recurses); arrays recurse into +/// object elements under the SAME path (arrays of scalars are already covered by the container row above); +/// `path` is a reused buffer (push the `.key` segment, recurse, truncate back). +fn walk( + value: &Value, + path: &mut String, + depth: u32, + emit: &mut impl FnMut(&str, bool, &'static str), +) { + if depth == 0 { + return; + } + match value { + Value::Object(map) => { + for (k, v) in map { + let len = path.len(); + if !path.is_empty() { + path.push('.'); + } + path.push_str(k); + match v { + Value::Object(_) | Value::Array(_) => { + emit(path, false, kind(v)); + walk(v, path, depth - 1, emit); + } + _ => emit(path, true, kind(v)), + } + path.truncate(len); + } + } + Value::Array(arr) => { + for v in arr { + if matches!(v, Value::Object(_)) { + walk(v, path, depth - 1, emit); + } + } + } + _ => {} + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn walk_matches_jsonb_field_rows_semantics() { + // object keys, nested object, array-of-objects (recurses under the same path), array-of-scalars + // (only the container row), and a scalar leaf. + let item = json!({ + "id": "x", + "geometry": {"type": "Point", "coordinates": [1, 2]}, + "properties": {"datetime": "2020-01-01T00:00:00Z", "eo:cloud_cover": 10}, + "assets": [{"href": "a"}, {"href": "b", "roles": ["data"]}] + }); + let mut reg = FieldRegistry::default(); + reg.observe(&item); + let paths: Vec<&str> = reg.paths.keys().map(String::as_str).collect(); + // leaf scalars + assert!(paths.contains(&"id")); + assert!(paths.contains(&"properties.datetime")); + assert!(paths.contains(&"properties.eo:cloud_cover")); + // containers emit their own (non-leaf) row + recurse + assert!(paths.contains(&"geometry")); + assert!(paths.contains(&"geometry.type")); + assert!(paths.contains(&"geometry.coordinates")); // array -> non-leaf row + // array of objects: recurse under the same path, merged across elements + assert!(paths.contains(&"assets")); // the array container + assert!(paths.contains(&"assets.href")); + assert!(paths.contains(&"assets.roles")); + // leaf/kind merges + assert!(reg.paths["id"].is_leaf); + assert!(!reg.paths["geometry"].is_leaf); + assert_eq!( + reg.paths["properties.eo:cloud_cover"] + .kinds + .iter() + .copied() + .collect::>(), + vec!["number"] + ); + } + + #[test] + fn merge_only_widens() { + let mut reg = FieldRegistry::default(); + reg.observe(&json!({"a": 1})); + reg.observe(&json!({"a": "s", "b": true})); // a gains a second kind; b is new + assert_eq!( + reg.paths["a"].kinds.iter().copied().collect::>(), + vec!["number", "string"] + ); + assert!(reg.paths.contains_key("b")); + } +} diff --git a/src/pgstac-rs/src/load/fragment.rs b/src/pgstac-rs/src/load/fragment.rs new file mode 100644 index 00000000..6287cdd7 --- /dev/null +++ b/src/pgstac-rs/src/load/fragment.rs @@ -0,0 +1,189 @@ +//! Fragment split for collections with a `fragment_config` (the 0.10 asset-dedup model). +//! +//! Ports the SQL `extract_fragment` / `strip_fragment_col` (003a_items.sql) so the Rust loader can split +//! an item the same way `items_staging_dehydrate` does: the configured paths (stable per-collection asset +//! metadata) move into a shared `item_fragments` row, and the per-item columns keep only what differs. +//! +//! This module is the pure, parity-tested core. Wiring it into the loader (dedup → `ensure_fragments` → +//! stamp `fragment_id`, plus links-template / stac_version / stac_extensions handling) is layered on top. + +use crate::Result; +use serde_json::{Map, Value}; + +/// A collection's fragment paths, parsed from `collections.fragment_config`. +/// +/// The stored form is a `text[]` whose elements are JSON arrays of path segments, e.g. +/// `["assets","B1","type"]` (see SQL `fragment_path_text`); the array form avoids ambiguity for keys +/// containing dots. +/// +/// # Examples +/// +/// ``` +/// use pgstac::fragment::FragmentConfig; +/// let config = FragmentConfig::parse(&[r#"["assets","B1","type"]"#.to_string()]).unwrap(); +/// assert!(!config.is_empty()); +/// ``` +#[derive(Debug, Clone, Default)] +pub struct FragmentConfig { + paths: Vec>, +} + +impl FragmentConfig { + /// Parses the `text[]` config — each element a JSON array of path segments. + pub fn parse(config: &[String]) -> Result { + let mut paths = Vec::with_capacity(config.len()); + for entry in config { + paths.push(serde_json::from_str(entry)?); + } + Ok(FragmentConfig { paths }) + } + + /// Whether the collection fragments nothing (no shared metadata to extract). + pub fn is_empty(&self) -> bool { + self.paths.is_empty() + } + + /// The parsed path arrays. + pub fn paths(&self) -> &[Vec] { + &self.paths + } + + /// Whether the config contains exactly this path (e.g. `["stac_version"]`), used to decide whether a + /// top-level column is fragmented wholesale. + pub fn has_path(&self, path: &[&str]) -> bool { + self.paths + .iter() + .any(|p| p.len() == path.len() && p.iter().zip(path).all(|(a, b)| a == b)) + } +} + +/// Gets the value at a root-relative object path — mirrors `content #> path` for object keys (fragment +/// paths are always object keys, never array indices). +fn get_path<'a>(content: &'a Value, path: &[String]) -> Option<&'a Value> { + let mut cur = content; + for key in path { + cur = cur.as_object()?.get(key)?; + } + Some(cur) +} + +/// Sets `val` at `path` in `result`, creating intermediate objects as needed (mirrors `jsonb_set_nested`, +/// so depth-3+ paths sharing intermediate keys merge rather than overwrite). +fn set_path(result: &mut Map, path: &[String], val: Value) { + let (last, parents) = match path.split_last() { + Some(parts) => parts, + None => return, + }; + let mut cur = result; + for key in parents { + let child = cur + .entry(key.clone()) + .or_insert_with(|| Value::Object(Map::new())); + if !child.is_object() { + *child = Value::Object(Map::new()); + } + cur = child.as_object_mut().expect("just ensured object"); + } + let _ = cur.insert(last.clone(), val); +} + +/// Removes a nested object path, mirroring the `#-` operator for object keys. +fn remove_path(value: &mut Value, path: &[String]) { + let Some((first, rest)) = path.split_first() else { + return; + }; + let Some(obj) = value.as_object_mut() else { + return; + }; + if rest.is_empty() { + let _ = obj.remove(first); + } else if let Some(child) = obj.get_mut(first) { + remove_path(child, rest); + } +} + +/// The sparse fragment overlay: every configured path present in `content`, reassembled into a nested +/// object. Returns `None` when no configured path is present (matches SQL `extract_fragment` returning +/// NULL for an empty result). +/// +/// Clones the extracted values into the overlay (the same values are also stripped from the per-item +/// columns by [`strip_fragment_col`]); this is the deferred dedup path, not the per-item hot path. +pub fn extract_fragment(content: &Value, config: &FragmentConfig) -> Option { + if config.is_empty() { + return None; + } + let mut result = Map::new(); + for path in config.paths() { + if path.is_empty() { + continue; + } + if let Some(val) = get_path(content, path) { + set_path(&mut result, path, val.clone()); + } + } + if result.is_empty() { + None + } else { + Some(Value::Object(result)) + } +} + +/// Removes the fragment-owned sub-keys for `col_name` from `col_value` (mirrors SQL `strip_fragment_col`): +/// a depth-1 path equal to `col_name` zeroes the whole column (`{}`); deeper paths remove the nested key. +pub fn strip_fragment_col(mut col_value: Value, col_name: &str, config: &FragmentConfig) -> Value { + for path in config.paths() { + if path.first().map(String::as_str) != Some(col_name) { + continue; + } + if path.len() == 1 { + return Value::Object(Map::new()); + } + remove_path(&mut col_value, &path[1..]); + } + col_value +} + +/// Builds the fragment payload for an item — `{content?, links_template?}` — that `ensure_fragments` +/// hashes + dedups (mirrors the SQL `items_staging_dehydrate` `fragmented` CTE's +/// `jsonb_strip_nulls(jsonb_build_object('content', NULLIF(frag_content,'{}'), 'links_template', ...))`). +/// +/// `content` is [`extract_fragment`] (already `None` for an empty overlay); `links_template` is +/// [`strip_link_hrefs`] of the item's `links`. Returns `None` when both are absent (the item has no +/// fragment — `fragment_id` stays NULL). Only the present keys appear, matching `jsonb_strip_nulls`. +pub fn build_fragment_payload(item: &Value, config: &FragmentConfig) -> Option { + let content = extract_fragment(item, config); + let links_template = item.get("links").and_then(strip_link_hrefs); + if content.is_none() && links_template.is_none() { + return None; + } + let mut payload = Map::new(); + if let Some(content) = content { + let _ = payload.insert("content".to_string(), content); + } + if let Some(links_template) = links_template { + let _ = payload.insert("links_template".to_string(), links_template); + } + Some(Value::Object(payload)) +} + +/// The links template: each link with its per-item `href` removed, preserving order (mirrors SQL +/// `stac_links_strip_hrefs`). Returns `None` for a null/non-array/empty `links` value. The template is the +/// shared part of an item's links that goes into the fragment; the per-item hrefs stay in `link_hrefs`. +pub fn strip_link_hrefs(links: &Value) -> Option { + let array = links.as_array()?; + if array.is_empty() { + return None; + } + let stripped = array + .iter() + .map(|link| match link.as_object() { + Some(object) => { + let mut object = object.clone(); + let _ = object.remove("href"); + Value::Object(object) + } + None => link.clone(), + }) + .collect(); + Some(Value::Array(stripped)) +} diff --git a/src/pgstac-rs/src/load/ingest.rs b/src/pgstac-rs/src/load/ingest.rs new file mode 100644 index 00000000..045fb600 --- /dev/null +++ b/src/pgstac-rs/src/load/ingest.rs @@ -0,0 +1,942 @@ +//! Loading dehydrated item rows into pgstac via binary COPY. +//! +//! The Rust loader binary-COPYs fully-dehydrated rows ([`DehydratedRow`]) into a session-local staging +//! table shaped like `items` (created by the SQL `make_binary_staging`), then calls +//! `flush_items_staging_binary` to move them into `items` (the only path that writes `items` — direct writes +//! are revoked from `pgstac_ingest`). This module owns the client-side binary encoding of a row. +//! +//! Geometry is sent as raw EWKB through [`Ewkb`](crate::geom::Ewkb): the PostGIS `geometry` binary +//! wire format *is* EWKB, so the bytes [`dehydrate`](crate::dehydrate) already produced go out untouched. + +use crate::Result; +#[cfg(feature = "pool")] +use crate::canonical::jsonb_canonical_hash; +use crate::dehydrate::{DehydrateSchema, DehydratedRow, PromotedValue, dehydrate}; +use crate::field_registry::FieldRegistry; +use crate::fragment::{FragmentConfig, build_fragment_payload, strip_fragment_col}; +use crate::geom::Ewkb; +use chrono::{DateTime, Datelike, TimeZone, Utc}; +use futures::pin_mut; +use serde_json::{Map, Value}; +use std::collections::{HashMap, HashSet}; +#[cfg(feature = "pool")] +use std::sync::atomic::{AtomicU64, Ordering}; +use tokio_postgres::binary_copy::BinaryCopyInWriter; +use tokio_postgres::types::{ToSql, Type}; + +impl PromotedValue { + /// Boxes the inner value as a trait object for the binary COPY writer. `Send` so the loader's futures + /// (which hold these across the COPY await) are `Send` and can be `tokio::spawn`ed under concurrency. + fn into_boxed(self) -> Box { + match self { + PromotedValue::Text(v) => Box::new(v), + PromotedValue::Float(v) => Box::new(v), + PromotedValue::Int(v) => Box::new(v), + PromotedValue::BigInt(v) => Box::new(v), + PromotedValue::TextArray(v) => Box::new(v), + PromotedValue::Jsonb(v) => Box::new(v), + PromotedValue::Timestamptz(v) => Box::new(v), + } + } +} + +/// The fixed (non-promoted) staging columns, in the exact order [`row_boxes`] pushes them. A staging +/// column not listed here and not a promoted column COPYs as SQL NULL (the sentinel slot). +const FIXED_COLUMNS: [&str; 18] = [ + "id", + "geometry", + "collection", + "datetime", + "end_datetime", + "datetime_is_range", + "stac_version", + "stac_extensions", + "pgstac_updated_at", + "item_hash", + "fragment_id", + "bbox", + "links", + "assets", + "properties", + "extra", + "link_hrefs", + "private", +]; + +/// Marshals a [`DehydratedRow`] into binary-COPY trait objects in a **fixed positional order** (the 18 +/// [`FIXED_COLUMNS`], then the promoted columns in `schema` order, then one trailing NULL sentinel), +/// consuming the row (fields are moved into the boxes — no deep clones, and no per-row map or key +/// strings; [`copy_rows`] reorders these into staging-column order via a permutation computed once). +fn row_boxes(row: DehydratedRow, n_promoted: usize) -> Vec> { + let mut v: Vec> = + Vec::with_capacity(FIXED_COLUMNS.len() + n_promoted + 1); + v.push(Box::new(row.id)); + v.push(Box::new(Ewkb(row.geometry))); + v.push(Box::new(row.collection)); + v.push(Box::new(row.datetime)); + v.push(Box::new(row.end_datetime)); + v.push(Box::new(row.datetime_is_range)); + v.push(Box::new(row.stac_version)); + v.push(Box::new(row.stac_extensions)); + v.push(Box::new(Utc::now())); + v.push(Box::new(row.item_hash.to_vec())); + v.push(Box::new(row.fragment_id)); + v.push(Box::new(row.bbox)); + v.push(Box::new(row.links)); + v.push(Box::new(row.assets)); + v.push(Box::new(row.properties)); + v.push(Box::new(row.extra)); + v.push(Box::new(row.link_hrefs)); + v.push(Box::new(Option::::None)); // private + for value in row.promoted { + v.push(value.into_boxed()); + } + v.push(Box::new(Option::::None)); // NULL sentinel for unmapped staging columns + v +} + +/// Binary-COPYs `rows` into the staging table `staging` (a session-local table shaped like `items`). +/// +/// Reads the staging table's column names + types once, then writes each row in column order. Returns +/// the number of rows written. Takes a transaction because the staging table is session-local +/// (`ON COMMIT DROP`), so the COPY must share the transaction that created it. +pub async fn copy_rows( + tx: &tokio_postgres::Transaction<'_>, + staging: &str, + rows: Vec, + schema: &DehydrateSchema, +) -> Result { + let describe = tx + .prepare(&format!("SELECT * FROM {staging} WHERE false")) + .await?; + let types: Vec = describe + .columns() + .iter() + .map(|c| c.type_().clone()) + .collect(); + + // Permutation computed ONCE: for each staging column, the index into the positional `row_boxes` + // vector that fills it (fixed columns 0..18, promoted 18.., NULL sentinel last). Unmapped staging + // columns point at the sentinel and COPY as SQL NULL. No per-row map or column-name strings. + let n_promoted = schema.columns().len(); + let null_idx = FIXED_COLUMNS.len() + n_promoted; + let mut name_to_box: HashMap<&str, usize> = HashMap::with_capacity(null_idx); + for (i, name) in FIXED_COLUMNS.iter().enumerate() { + let _ = name_to_box.insert(name, i); + } + for (j, col) in schema.columns().iter().enumerate() { + let _ = name_to_box.insert(col.column.as_str(), FIXED_COLUMNS.len() + j); + } + let perm: Vec = describe + .columns() + .iter() + .map(|c| name_to_box.get(c.name()).copied().unwrap_or(null_idx)) + .collect(); + + let sink = tx + .copy_in(&format!("COPY {staging} FROM STDIN WITH (FORMAT BINARY)")) + .await?; + let writer = BinaryCopyInWriter::new(sink, &types); + pin_mut!(writer); + + for row in rows { + let boxes = row_boxes(row, n_promoted); + // Reorder into staging-column order via the precomputed permutation. Drop the `Send` bound for + // the writer's `&[&(dyn ToSql + Sync)]` signature (implicit coercion). + let params: Vec<&(dyn ToSql + Sync)> = perm + .iter() + .map(|&i| -> &(dyn ToSql + Sync) { &*boxes[i] }) + .collect(); + writer.as_mut().write(¶ms).await?; + } + Ok(writer.finish().await?) +} + +/// Whether a failed DB call is worth retrying: transient contention (deadlock, serialization failure, +/// lock-not-available, object-in-use) from many writers racing to CREATE + widen the same partitions. +/// Determinate errors (bad data, unique violation, missing collection, …) are NOT retried. +fn is_retryable_db(err: &tokio_postgres::Error) -> bool { + use tokio_postgres::error::SqlState; + match err.as_db_error() { + Some(db) => matches!( + *db.code(), + SqlState::T_R_DEADLOCK_DETECTED + | SqlState::T_R_SERIALIZATION_FAILURE + | SqlState::LOCK_NOT_AVAILABLE + | SqlState::OBJECT_IN_USE + ), + None => false, + } +} + +/// How `flush_items_staging_binary` resolves an id collision with an existing item. +/// +/// ORPHAN CAVEAT: `Ignore` and `Upsert` only touch the partition the *incoming* item routes to. If a +/// datetime change moves an item to a DIFFERENT partition, the old row in its former partition is left +/// behind as an orphan (a duplicate `(collection, id)` across partitions). Use `Delsert` when updates may +/// move an item between partitions; it deletes the old row wherever it lives. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[cfg_attr(feature = "cli", derive(clap::ValueEnum))] +pub enum ConflictPolicy { + /// Fail the whole load if any item id already exists. + Error, + /// Replace a changed item IN THE PARTITION IT ROUTES TO (fast: the delete is window-pruned to that one + /// partition). A datetime change that moves the item to another partition orphans the old row — use + /// `Delsert` for that case. + Upsert, + /// Replace a changed item wherever it lives: delete the old row CROSS-partition (by collection+id) first, + /// then insert. Move-safe, but the delete probes every partition for the id. + Delsert, + /// Skip items whose id already exists (insert-if-absent). Like `Upsert`, a cross-partition move orphans + /// the old row. + Ignore, +} + +impl ConflictPolicy { + /// The SQL token `flush_items_staging_binary` expects. + fn as_sql(self) -> &'static str { + match self { + ConflictPolicy::Error => "error", + ConflictPolicy::Upsert => "upsert", + ConflictPolicy::Delsert => "delsert", + ConflictPolicy::Ignore => "ignore", + } + } +} + +/// Loads a batch of full STAC items into pgstac via the Rust dehydrate + binary-COPY path. +/// +/// For each item: dehydrate it; for a collection with a `fragment_config`, split out the shared fragment +/// (upserted + deduped via `ensure_fragments`) and strip the fragment-owned keys from the row; ensure its +/// partition exists and its stats cover it (`ensure_partitions`, committed up front); then +/// binary-COPY all rows into a session-local staging table and `flush_items_staging_binary` them into +/// `items` resolving id collisions by `policy`. Matches `items_staging_dehydrate` for both the no-fragment +/// and fragment paths. Returns the number of rows flushed. Collections must already exist. +pub async fn load_items( + client: &mut tokio_postgres::Client, + items: Vec, + schema: &DehydrateSchema, + policy: ConflictPolicy, +) -> Result { + if items.is_empty() { + return Ok(0); + } + + // Per-collection fragment_config for any fragment-enabled collection in the batch (empty/absent = + // the no-fragment fast path). + let configs = fetch_fragment_configs(client, &items).await?; + + // Dehydrate each item; for a fragment collection, compute its fragment payload from the ORIGINAL item + // (before dehydrate consumes it) and strip the fragment-owned keys from the dehydrated row. + let mut rows: Vec = Vec::with_capacity(items.len()); + let mut payloads: Vec> = Vec::with_capacity(items.len()); + let mut field_registry: HashMap = HashMap::new(); + for item in items { + let collection = item.get("collection").and_then(Value::as_str); + // Field registry: walk the FULL item content (every item, no sampling) into the per-collection + // registry, flushed below as an add-only widen so it stays a superset of every field the items carry. + if let Some(coll) = collection { + field_registry + .entry(coll.to_string()) + .or_default() + .observe(&item); + } + let config = collection + .and_then(|collection| configs.get(collection)) + .filter(|config| !config.is_empty()); + let payload = config.and_then(|config| build_fragment_payload(&item, config)); + let mut row = dehydrate(item, schema)?; + if let Some(config) = config { + apply_fragment_split(&mut row, config); + } + rows.push(row); + payloads.push(payload); + } + + // Create + widen every partition the batch lands in BEFORE the load. Bucket client-side + // by (collection, partition window) and call prepare_partition_for_load once per partition — per-partition + // metadata, O(#partitions), NOT the per-item batch-length arrays the old `ensure_partitions` shipped just + // to group them server-side. Each call widens dt/edt + the real SPATIAL envelope (`ensure_partitions` + // passed NULL) AND bumps n by this partition's staged count, so the flush writes only `items`, never + // partition_stats. Over-counting n is safe (ignore/upsert may insert fewer); the async tightener resets it. + let mut seen_collections: HashSet<&str> = HashSet::new(); + let distinct_collections: Vec<&str> = rows + .iter() + .map(|r| r.collection.as_str()) + .filter(|c| seen_collections.insert(*c)) + .collect(); + let truncs: HashMap> = client + .query( + "SELECT id, partition_trunc FROM collections WHERE id = ANY($1)", + &[&distinct_collections], + ) + .await? + .into_iter() + .map(|r| (r.get::<_, String>(0), r.get::<_, Option>(1))) + .collect(); + let mut aggs: HashMap<(String, i64), PartitionAgg> = HashMap::new(); + for row in &rows { + let trunc = truncs.get(&row.collection).and_then(Option::as_deref); + let (xmin, ymin, xmax, ymax) = bbox_envelope(&row.bbox); + let agg = aggs + .entry((row.collection.clone(), window_key(row.datetime, trunc))) + .or_insert_with(|| PartitionAgg { + collection: row.collection.clone(), + dt_lo: row.datetime, + dt_hi: row.datetime, + edt_lo: row.end_datetime, + edt_hi: row.end_datetime, + xmin, + ymin, + xmax, + ymax, + count: 0, + }); + agg.dt_lo = agg.dt_lo.min(row.datetime); + agg.dt_hi = agg.dt_hi.max(row.datetime); + agg.edt_lo = agg.edt_lo.min(row.end_datetime); + agg.edt_hi = agg.edt_hi.max(row.end_datetime); + agg.xmin = agg.xmin.min(xmin); + agg.ymin = agg.ymin.min(ymin); + agg.xmax = agg.xmax.max(xmax); + agg.ymax = agg.ymax.max(ymax); + agg.count += 1; + } + for agg in aggs.values() { + // Many writers race to CREATE + widen the same partitions; concurrent partition DDL can transiently + // deadlock / serialization-fail / hit lock-not-available. Retry those a few times with a small backoff; + // determinate errors propagate. + let mut attempt = 0u32; + loop { + match client + .execute( + "SELECT prepare_partition_for_load($1,$2,$3,$4,$5,$6,$7,$8,$9,$10)", + &[ + &agg.collection, + &agg.dt_lo, + &agg.dt_hi, + &agg.edt_lo, + &agg.edt_hi, + &agg.xmin, + &agg.ymin, + &agg.xmax, + &agg.ymax, + &agg.count, + ], + ) + .await + { + Ok(_) => break, + Err(e) if attempt < 5 && is_retryable_db(&e) => { + attempt += 1; + tokio::time::sleep(std::time::Duration::from_millis(20 * u64::from(attempt))) + .await; + } + Err(e) => return Err(e.into()), + } + } + } + + // Add-only UPSERT of the COMPLETE per-collection path set walked above, so item_field_registry stays a + // superset of every field the items carry. Direct INSERT ... ON CONFLICT (no SECURITY DEFINER function) — + // pgstac_ingest keeps INSERT/UPDATE on item_field_registry (DELETE/TRUNCATE revoked, so this can only + // widen, never narrow). Retried on transient contention, since concurrent writers UPSERT the same rows. + for (collection, registry) in &field_registry { + if registry.is_empty() { + continue; + } + let entries = registry.to_entries(); + let mut attempt = 0u32; + loop { + match client + .execute( + "INSERT INTO item_field_registry (collection, path, is_leaf, value_kinds, first_seen, last_seen) \ + SELECT $1, e->>'path', (e->>'is_leaf')::boolean, \ + ARRAY(SELECT jsonb_array_elements_text(e->'value_kinds')), now(), now() \ + FROM jsonb_array_elements($2::jsonb) AS e \ + ON CONFLICT (collection, path) DO UPDATE SET \ + is_leaf = item_field_registry.is_leaf AND EXCLUDED.is_leaf, \ + value_kinds = (SELECT array_agg(DISTINCT v) \ + FROM unnest(item_field_registry.value_kinds || EXCLUDED.value_kinds) t(v)), \ + last_seen = now()", + &[collection, &entries], + ) + .await + { + Ok(_) => break, + Err(e) if attempt < 5 && is_retryable_db(&e) => { + attempt += 1; + tokio::time::sleep(std::time::Duration::from_millis(20 * u64::from(attempt))).await; + } + Err(e) => return Err(e.into()), + } + } + } + + // Upsert + dedup the fragments and stamp each row's fragment_id, committed before the load (the + // items COPY references fragment_id; the fragments must already exist). + stamp_fragment_ids(client, &mut rows, &payloads).await?; + + // Staging + COPY + flush in one transaction (the staging table is ON COMMIT DROP). + let tx = client.transaction().await?; + let staging: String = tx + .query_one("SELECT make_binary_staging()", &[]) + .await? + .get(0); + let _ = copy_rows(&tx, &staging, rows, schema).await?; + let flushed: i64 = tx + .query_one( + "SELECT flush_items_staging_binary($1, $2)", + &[&staging, &policy.as_sql()], + ) + .await? + .get(0); + tx.commit().await?; + + Ok(u64::try_from(flushed)?) +} + +/// Per-partition aggregate accumulated in pass 1 (extent + count), used to widen partition_stats BEFORE +/// the COPY (golden rule: stats are at least as wide as the data before any data lands). +struct PartitionAgg { + collection: String, + dt_lo: DateTime, + dt_hi: DateTime, + edt_lo: DateTime, + edt_hi: DateTime, + xmin: f64, + ymin: f64, + xmax: f64, + ymax: f64, + count: i64, +} + +/// The partition-window start (epoch seconds) a datetime falls in, for client-side bucketing. Matches the +/// server's UTC `date_trunc` (the cluster runs UTC). A NULL `partition_trunc` collapses to one bucket per +/// collection (sentinel `i64::MIN`). +fn window_key(dt: DateTime, trunc: Option<&str>) -> i64 { + let start = match trunc { + Some("month") => Utc + .with_ymd_and_hms(dt.year(), dt.month(), 1, 0, 0, 0) + .single(), + Some("year") => Utc.with_ymd_and_hms(dt.year(), 1, 1, 0, 0, 0).single(), + _ => return i64::MIN, + }; + start.map(|d| d.timestamp()).unwrap_or(i64::MIN) +} + +/// The `(xmin, ymin, xmax, ymax)` envelope of a STAC bbox (2D `[w,s,e,n]` or 3D `[w,s,zmin,e,n,zmax]`). +/// A missing/odd bbox falls back to the whole world (over-wide is the safe direction for the widen). +fn bbox_envelope(bbox: &Option) -> (f64, f64, f64, f64) { + let nums: Option> = bbox + .as_ref() + .and_then(|b| b.as_array()) + .map(|a| a.iter().filter_map(Value::as_f64).collect()); + match nums.as_deref() { + Some([xmin, ymin, xmax, ymax]) => (*xmin, *ymin, *xmax, *ymax), + Some([xmin, ymin, _zmin, xmax, ymax, _zmax]) => (*xmin, *ymin, *xmax, *ymax), + _ => (-180.0, -90.0, 180.0, 90.0), + } +} + +/// Per-collection `partition_trunc` for the collections referenced in `items`. +#[cfg(feature = "pool")] +async fn fetch_partition_truncs( + client: &tokio_postgres::Client, + items: &[Value], +) -> Result>> { + let mut seen = HashSet::new(); + let collections: Vec<&str> = items + .iter() + .filter_map(|item| item.get("collection").and_then(Value::as_str)) + .filter(|c| seen.insert(*c)) + .collect(); + let rows = client + .query( + "SELECT id, partition_trunc FROM collections WHERE id = ANY($1)", + &[&collections], + ) + .await?; + Ok(rows + .into_iter() + .map(|r| (r.get::<_, String>(0), r.get::<_, Option>(1))) + .collect()) +} + +/// The partition window `[start, end)` a datetime lands in, or `None` for a non-partitioned (single +/// partition) collection. Mirrors [`window_key`] / the server's UTC `date_trunc`, so the client can prune a +/// precheck probe to exactly one partition. +#[cfg(feature = "pool")] +fn window_bounds(dt: DateTime, trunc: Option<&str>) -> Option<(DateTime, DateTime)> { + match trunc { + Some("month") => { + let start = Utc + .with_ymd_and_hms(dt.year(), dt.month(), 1, 0, 0, 0) + .single()?; + let (ey, em) = if dt.month() == 12 { + (dt.year() + 1, 1) + } else { + (dt.year(), dt.month() + 1) + }; + let end = Utc.with_ymd_and_hms(ey, em, 1, 0, 0, 0).single()?; + Some((start, end)) + } + Some("year") => { + let start = Utc.with_ymd_and_hms(dt.year(), 1, 1, 0, 0, 0).single()?; + let end = Utc + .with_ymd_and_hms(dt.year() + 1, 1, 1, 0, 0, 0) + .single()?; + Some((start, end)) + } + _ => None, + } +} + +/// A STAC item's nominal datetime (`properties.datetime`, falling back to `start_datetime`) as UTC. +#[cfg(feature = "pool")] +fn item_datetime(item: &Value) -> Result> { + let props = item.get("properties"); + let raw = props + .and_then(|p| p.get("datetime")) + .and_then(Value::as_str) + .filter(|s| !s.is_empty()) + .or_else(|| { + props + .and_then(|p| p.get("start_datetime")) + .and_then(Value::as_str) + }) + .ok_or_else(|| crate::Error::Dehydrate("item missing datetime".to_string()))?; + Ok(DateTime::parse_from_rfc3339(raw) + .map_err(|e| crate::Error::Dehydrate(format!("unparseable datetime {raw:?}: {e}")))? + .with_timezone(&Utc)) +} + +/// One partition's worth of precheck input: the original item indices, ids and canonical hashes that landed +/// in this partition window, a representative datetime (to locate the partition), and the window bounds used +/// to prune the probe to that one partition. +#[cfg(feature = "pool")] +struct PrecheckBucket { + collection: String, + repr_dt: DateTime, + bounds: Option<(DateTime, DateTime)>, + ords: Vec, + ids: Vec, + hashes: Vec>, +} + +#[cfg(feature = "pool")] +static PRECHECK_TEMP_SEQ: AtomicU64 = AtomicU64::new(0); + +/// Loads `items`, loading only those not already present-and-current and skipping the rest, via +/// [`crate::PgstacPool::create_items`]. A cheap pass (id + datetime, plus a content hash for upsert/delsert) +/// buckets items by partition window; `precheck_one_partition` classifies each partition's bucket in +/// parallel, and the survivors take the normal load (resolving same-partition conflicts per `policy`). +/// `ignore` skips every existing id; `upsert`/`delsert` skip only ids whose content is unchanged. +/// +/// Matching is by id: a datetime change within a partition is caught, but a cross-partition move reads as +/// new (the old row lives in another partition) — `Upsert` leaves that old row in place, `Delsert` removes +/// it. Returns `(unchanged_skipped, loaded)`. +#[cfg(feature = "pool")] +pub async fn precheck_upsert( + pool: &crate::PgstacPool, + items: Vec, + policy: ConflictPolicy, +) -> Result<(u64, u64)> { + if items.is_empty() { + return Ok((0, 0)); + } + // ignore/error only care whether the id already exists in the partition, not whether content changed — + // skip the content hash entirely (cheaper pass) and skip ALL existing ids, not just unchanged ones. + let use_hash = matches!(policy, ConflictPolicy::Upsert | ConflictPolicy::Delsert); + let meta = pool.get().await?; + let truncs = fetch_partition_truncs(&meta, &items).await?; + drop(meta); + + // Cheap pass: (id, canonical hash, datetime) per item, bucketed by (collection, partition window). + let mut buckets: HashMap<(String, i64), PrecheckBucket> = HashMap::new(); + for (i, item) in items.iter().enumerate() { + let collection = item + .get("collection") + .and_then(Value::as_str) + .ok_or_else(|| crate::Error::Dehydrate("item missing collection".to_string()))?; + let id = item + .get("id") + .and_then(Value::as_str) + .ok_or_else(|| crate::Error::Dehydrate("item missing id".to_string()))?; + let dt = item_datetime(item)?; + let trunc = truncs.get(collection).and_then(Option::as_deref); + let bucket = buckets + .entry((collection.to_string(), window_key(dt, trunc))) + .or_insert_with(|| PrecheckBucket { + collection: collection.to_string(), + repr_dt: dt, + bounds: window_bounds(dt, trunc), + ords: Vec::new(), + ids: Vec::new(), + hashes: Vec::new(), + }); + bucket + .ords + .push(i32::try_from(i).expect("batch index fits i32")); + bucket.ids.push(id.to_string()); + if use_hash { + bucket.hashes.push(jsonb_canonical_hash(item)?.to_vec()); + } + } + + // Classify each partition's bucket in parallel -> (ords to load, unchanged count). + let mut set: tokio::task::JoinSet, u64)>> = tokio::task::JoinSet::new(); + for (_key, bucket) in buckets { + let pool = pool.clone(); + let _ = set.spawn(async move { precheck_one_partition(&pool, bucket, use_hash).await }); + } + let mut load_ords: HashSet = HashSet::new(); + let mut unchanged: u64 = 0; + while let Some(joined) = set.join_next().await { + let (ords, unch) = joined.expect("precheck task panicked")?; + load_ords.extend(ords); + unchanged += unch; + } + + // Load only new + changed through the normal path (which resolves same-partition conflicts per policy). + let to_load: Vec = items + .into_iter() + .enumerate() + .filter_map(|(i, item)| { + load_ords + .contains(&i32::try_from(i).expect("batch index fits i32")) + .then_some(item) + }) + .collect(); + let loaded = pool.create_items(to_load, policy).await?; + Ok((unchanged, loaded)) +} + +/// Classify one partition's bucket: returns the ords (original item indices) to LOAD and the count SKIPPED. +/// With `use_hash` (upsert/delsert) an existing id is skipped only when its content hash matches; without it +/// (ignore) an existing id is skipped regardless of content. Picks the cheaper side: empty/just-created +/// partition -> all new (no probe); batch bigger than the partition -> pull the partition's ids (+ item_hash +/// when hashing) and compare in memory; otherwise binary-COPY the (smaller) batch into a TEMP table and JOIN +/// this one partition (window-pruned). No per-item SQL-function arguments either way. +#[cfg(feature = "pool")] +async fn precheck_one_partition( + pool: &crate::PgstacPool, + bucket: PrecheckBucket, + use_hash: bool, +) -> Result<(Vec, u64)> { + // Pre-load row count for the partition that holds repr_dt (named by the server's own partition_name). + // No row / 0 => empty or just-created => every item is new; skip the probe entirely. ONE connection for + // the whole task (the n-query, the pull, and the probe's temp-table tx all reuse it) — holding a second + // would deadlock the pool once there are more parallel partition tasks than pool/2. + let mut client = pool.get().await?; + let n: i64 = client + .query_one( + "SELECT COALESCE((SELECT n FROM pgstac.partition_stats \ + WHERE partition = (pgstac.partition_name($1, $2::timestamptz)).partition_name), 0)", + &[&bucket.collection, &bucket.repr_dt], + ) + .await? + .get(0); + if n == 0 { + return Ok((bucket.ords, 0)); + } + + let batch = bucket.ords.len() as i64; + if batch > n { + // PULL: the partition is the smaller side. Stream its ids (+ item_hash for upsert); compare in memory. + if use_hash { + let rows = match bucket.bounds { + Some((ws, we)) => { + client + .query( + "SELECT id, item_hash FROM pgstac.items \ + WHERE collection = $1 AND datetime >= $2 AND datetime < $3", + &[&bucket.collection, &ws, &we], + ) + .await? + } + None => { + client + .query( + "SELECT id, item_hash FROM pgstac.items WHERE collection = $1", + &[&bucket.collection], + ) + .await? + } + }; + let mut stored: HashMap>> = HashMap::with_capacity(rows.len()); + for row in &rows { + let _ = stored.insert(row.get(0), row.get(1)); + } + let mut to_load = Vec::new(); + let mut unchanged = 0u64; + for ((ord, id), hash) in bucket.ords.iter().zip(&bucket.ids).zip(&bucket.hashes) { + match stored.get(id) { + Some(Some(stored_hash)) if stored_hash == hash => unchanged += 1, + _ => to_load.push(*ord), + } + } + return Ok((to_load, unchanged)); + } + // ignore/error: only id existence matters — skip ALL existing ids (no hash). + let rows = match bucket.bounds { + Some((ws, we)) => { + client + .query( + "SELECT id FROM pgstac.items \ + WHERE collection = $1 AND datetime >= $2 AND datetime < $3", + &[&bucket.collection, &ws, &we], + ) + .await? + } + None => { + client + .query( + "SELECT id FROM pgstac.items WHERE collection = $1", + &[&bucket.collection], + ) + .await? + } + }; + let stored: HashSet = rows.iter().map(|row| row.get(0)).collect(); + let mut to_load = Vec::new(); + let mut unchanged = 0u64; + for (ord, id) in bucket.ords.iter().zip(&bucket.ids) { + if stored.contains(id) { + unchanged += 1; + } else { + to_load.push(*ord); + } + } + return Ok((to_load, unchanged)); + } + + // PROBE: the batch is the smaller side. Binary-COPY it into a TEMP table, JOIN this one partition + // (window-pruned). Reuses the task's single connection (the n-query above) — see the pool-deadlock note. + let tx = client.transaction().await?; + let tmp = format!("_pc_{}", PRECHECK_TEMP_SEQ.fetch_add(1, Ordering::Relaxed)); + if use_hash { + // (ord, id, hash); status: 0 = new (no match), 1 = unchanged (hash equal), 2 = changed. Load 0 + 2. + tx.batch_execute(&format!( + "CREATE TEMP TABLE {tmp} (ord int4, id text, hash bytea) ON COMMIT DROP" + )) + .await?; + { + let sink = tx + .copy_in(&format!("COPY {tmp} FROM STDIN WITH (FORMAT BINARY)")) + .await?; + let writer = BinaryCopyInWriter::new(sink, &[Type::INT4, Type::TEXT, Type::BYTEA]); + pin_mut!(writer); + for ((ord, id), hash) in bucket.ords.iter().zip(&bucket.ids).zip(&bucket.hashes) { + let params: [&(dyn ToSql + Sync); 3] = [ord, id, hash]; + writer.as_mut().write(¶ms).await?; + } + let _ = writer.finish().await?; + } + let classify = "CASE WHEN p.id IS NULL THEN 0 \ + WHEN t.hash IS NOT DISTINCT FROM p.item_hash THEN 1 ELSE 2 END"; + let rows = match bucket.bounds { + Some((ws, we)) => { + tx.query( + &format!( + "SELECT t.ord, {classify} FROM {tmp} t \ + LEFT JOIN pgstac.items p \ + ON p.collection = $1 AND p.datetime >= $2 AND p.datetime < $3 AND p.id = t.id" + ), + &[&bucket.collection, &ws, &we], + ) + .await? + } + None => { + tx.query( + &format!( + "SELECT t.ord, {classify} FROM {tmp} t \ + LEFT JOIN pgstac.items p ON p.collection = $1 AND p.id = t.id" + ), + &[&bucket.collection], + ) + .await? + } + }; + let mut to_load = Vec::new(); + let mut unchanged = 0u64; + for row in &rows { + let ord: i32 = row.get(0); + let status: i32 = row.get(1); + if status == 1 { + unchanged += 1; + } else { + to_load.push(ord); + } + } + tx.commit().await?; + return Ok((to_load, unchanged)); + } + // ignore/error: (ord, id) only — an existing id is skipped, a new id is loaded. + tx.batch_execute(&format!( + "CREATE TEMP TABLE {tmp} (ord int4, id text) ON COMMIT DROP" + )) + .await?; + { + let sink = tx + .copy_in(&format!("COPY {tmp} FROM STDIN WITH (FORMAT BINARY)")) + .await?; + let writer = BinaryCopyInWriter::new(sink, &[Type::INT4, Type::TEXT]); + pin_mut!(writer); + for (ord, id) in bucket.ords.iter().zip(&bucket.ids) { + let params: [&(dyn ToSql + Sync); 2] = [ord, id]; + writer.as_mut().write(¶ms).await?; + } + let _ = writer.finish().await?; + } + let rows = match bucket.bounds { + Some((ws, we)) => { + tx.query( + &format!( + "SELECT t.ord, (p.id IS NOT NULL) FROM {tmp} t \ + LEFT JOIN pgstac.items p \ + ON p.collection = $1 AND p.datetime >= $2 AND p.datetime < $3 AND p.id = t.id" + ), + &[&bucket.collection, &ws, &we], + ) + .await? + } + None => { + tx.query( + &format!( + "SELECT t.ord, (p.id IS NOT NULL) FROM {tmp} t \ + LEFT JOIN pgstac.items p ON p.collection = $1 AND p.id = t.id" + ), + &[&bucket.collection], + ) + .await? + } + }; + let mut to_load = Vec::new(); + let mut unchanged = 0u64; + for row in &rows { + let ord: i32 = row.get(0); + let exists: bool = row.get(1); + if exists { + unchanged += 1; + } else { + to_load.push(ord); + } + } + tx.commit().await?; + Ok((to_load, unchanged)) +} + +/// Loads the `fragment_config` for the distinct collections referenced in `items`. Only collections with +/// a non-null config appear; everything else takes the no-fragment path. +async fn fetch_fragment_configs( + client: &tokio_postgres::Client, + items: &[Value], +) -> Result> { + let mut seen = HashSet::new(); + let collections: Vec<&str> = items + .iter() + .filter_map(|item| item.get("collection").and_then(Value::as_str)) + .filter(|collection| seen.insert(*collection)) + .collect(); + let mut configs = HashMap::new(); + if collections.is_empty() { + return Ok(configs); + } + let rows = client + .query( + "SELECT id, fragment_config FROM collections \ + WHERE id = ANY($1) AND fragment_config IS NOT NULL", + &[&collections], + ) + .await?; + for row in rows { + let id: String = row.get("id"); + let config: Vec = row.get("fragment_config"); + let _ = configs.insert(id, FragmentConfig::parse(&config)?); + } + Ok(configs) +} + +/// Strips a collection's fragment-owned keys from a dehydrated row (mirrors the `enriched` CTE of +/// `items_staging_dehydrate`). `config` is non-empty. `fragment_id` is stamped later by +/// [`stamp_fragment_ids`]; a row whose item has no fragment payload keeps `fragment_id = None`. +fn apply_fragment_split(row: &mut DehydratedRow, config: &FragmentConfig) { + let assets = row + .assets + .take() + .unwrap_or_else(|| Value::Object(Map::new())); + row.assets = Some(strip_fragment_col(assets, "assets", config)); + let properties = std::mem::replace(&mut row.properties, Value::Null); + row.properties = strip_fragment_col(properties, "properties", config); + if config.has_path(&["stac_version"]) { + row.stac_version = None; + } + if config.has_path(&["stac_extensions"]) { + row.stac_extensions = Value::Array(Vec::new()); + } + // When link_hrefs carries the per-item hrefs, the link structure lives in the fragment's + // links_template, so the row's links column is nulled (matches items_staging_dehydrate). + if row + .link_hrefs + .as_ref() + .is_some_and(|hrefs| !hrefs.is_empty()) + { + row.links = None; + } +} + +/// Dedups the per-collection fragment payloads, upserts them via `ensure_fragments`, and stamps each +/// contributing row's `fragment_id`. Local dedup (by serialized payload) just shrinks the array sent; +/// `ensure_fragments` is itself dup-safe and assigns ids by hash server-side. +async fn stamp_fragment_ids( + client: &tokio_postgres::Client, + rows: &mut [DehydratedRow], + payloads: &[Option], +) -> Result<()> { + // Group contributing item indices by collection (owned key — we mutate rows below). + let mut by_collection: HashMap> = HashMap::new(); + for (index, payload) in payloads.iter().enumerate() { + if payload.is_some() { + by_collection + .entry(rows[index].collection.clone()) + .or_default() + .push(index); + } + } + + for (collection, indices) in &by_collection { + // Distinct payloads (preserving first-seen order) + each item's slot in that distinct list. + let mut distinct: Vec = Vec::new(); + let mut key_to_slot: HashMap = HashMap::new(); + let mut slot_per_item: Vec<(usize, usize)> = Vec::with_capacity(indices.len()); + for &index in indices { + let payload = payloads[index] + .as_ref() + .expect("contributing item has a payload"); + let key = serde_json::to_string(payload)?; + let slot = *key_to_slot.entry(key).or_insert_with(|| { + distinct.push(payload.clone()); + distinct.len() - 1 + }); + slot_per_item.push((index, slot)); + } + + // ensure_fragments returns (ord, frag_id) where ord is the 1-based input position. + let returned = client + .query( + "SELECT ord, frag_id FROM ensure_fragments($1, $2)", + &[collection, &distinct], + ) + .await?; + let mut ids = vec![0i64; distinct.len()]; + for row in returned { + let ord: i32 = row.get("ord"); + ids[usize::try_from(ord - 1)?] = row.get("frag_id"); + } + for (index, slot) in slot_per_item { + rows[index].fragment_id = Some(ids[slot]); + } + } + Ok(()) +} diff --git a/src/pgstac-rs/src/load/mod.rs b/src/pgstac-rs/src/load/mod.rs new file mode 100644 index 00000000..ffef0e80 --- /dev/null +++ b/src/pgstac-rs/src/load/mod.rs @@ -0,0 +1,11 @@ +//! The ingest/write path: dehydration, fragment splitting, the field registry, and the Rust loader. + +pub mod dehydrate; +pub(crate) mod field_registry; +pub mod fragment; +pub mod ingest; +#[cfg(feature = "export")] +pub mod parquet_decode; +#[cfg(feature = "pool")] +pub(crate) mod pool_ingest; +pub(crate) mod queryables; diff --git a/src/pgstac-rs/src/load/parquet_decode.rs b/src/pgstac-rs/src/load/parquet_decode.rs new file mode 100644 index 00000000..ec18f0c9 --- /dev/null +++ b/src/pgstac-rs/src/load/parquet_decode.rs @@ -0,0 +1,329 @@ +//! Parallel stac-geoparquet decode for the loader. +//! +//! A stac-geoparquet file is a stack of independent **row groups**; decoding one needs nothing from +//! any other. [`spawn_parquet_decoders`] shards a file's row groups across N threads, each decoding its +//! shard to [`Vec`](serde_json::Value) batches and feeding the shared bounded channel the loader +//! already drains for the NDJSON path. This lifts the single-thread parquet→GeoJSON ceiling that made +//! the parquet ingest path decode-bound and far slower than NDJSON. +//! +//! The single-file reader mirrors [`stac::geoparquet::from_reader_iter`] (same geoarrow schema, same +//! per-batch decode), restricted to a row-group subset via +//! [`ParquetRecordBatchReaderBuilder::with_row_groups`]. Row-group sharding is exact: every row group +//! is assigned to exactly one shard, so the union over shards decodes every row exactly once — no +//! dropped or duplicated items. + +use arrow_array::RecordBatch; +use geoparquet::reader::{GeoParquetReaderBuilder, GeoParquetRecordBatchReader}; +use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; +use serde_json::Value; +use std::path::Path; +use tokio::sync::mpsc::Sender; + +/// Distributes the row-group indices `0..num_row_groups` across at most `n` shards. +/// +/// Indices are dealt **round-robin** (`rg % shards`) rather than in contiguous ranges, so when row +/// groups differ in size — common in stac-geoparquet, whose last group is a remainder — each shard +/// still gets a balanced spread of small and large groups instead of one shard inheriting all the big +/// ones. The number of shards is capped at `num_row_groups` (an empty file yields no shards), and every +/// index `0..num_row_groups` appears in exactly one returned shard. +/// +/// # Examples +/// +/// ``` +/// use pgstac::parquet_decode::shard_row_groups; +/// +/// // 5 row groups over 2 threads: round-robin split. +/// assert_eq!(shard_row_groups(5, 2), vec![vec![0, 2, 4], vec![1, 3]]); +/// // More threads than row groups: one row group per shard, no empties. +/// assert_eq!(shard_row_groups(2, 8), vec![vec![0], vec![1]]); +/// // Empty file: no shards. +/// assert!(shard_row_groups(0, 4).is_empty()); +/// ``` +pub fn shard_row_groups(num_row_groups: usize, n: usize) -> Vec> { + let shards = n.max(1).min(num_row_groups); + let mut out: Vec> = vec![Vec::new(); shards]; + for rg in 0..num_row_groups { + out[rg % shards].push(rg); + } + out +} + +/// Decodes a row-group subset of a stac-geoparquet file, invoking `on_batch` once per record batch +/// with the decoded items and an estimated in-memory byte size. +/// +/// Mirrors [`stac::geoparquet::from_reader_iter`] — same geoarrow schema, same per-batch +/// `items_from_record_batch` decode — but restricted to `row_groups` via +/// [`ParquetRecordBatchReaderBuilder::with_row_groups`]. The byte estimate is the row groups' +/// uncompressed columnar size prorated per item (a cheap, monotonic proxy for the JSON text size the +/// NDJSON path reports), so the loader's `PGSTAC_LOAD_BYTES` budget bounds the in-flight set for +/// parquet just as it does for NDJSON. Errors are returned as `String` to match the loader channel. +fn decode_row_groups(path: &Path, row_groups: &[usize], mut on_batch: F) -> Result<(), String> +where + F: FnMut(Vec, usize) -> Result<(), String>, +{ + let file = std::fs::File::open(path).map_err(|e| e.to_string())?; + let builder = ParquetRecordBatchReaderBuilder::try_new(file).map_err(|e| e.to_string())?; + + // Per-item byte proxy: the uncompressed columnar size of this shard's row groups, divided by their + // row count. Reported per batch as `batch.len() * per_item_bytes` so the loader's byte budget sees a + // sensible, monotonic size (matching the NDJSON path's summed item-text length closely enough to + // bound memory) without a second serialization pass over every item. + let meta = builder.metadata(); + let (shard_bytes, shard_rows) = row_groups.iter().fold((0i64, 0i64), |(b, r), &rg| { + let g = meta.row_group(rg); + (b + g.total_byte_size(), r + g.num_rows()) + }); + let per_item_bytes = if shard_rows > 0 { + (shard_bytes / shard_rows).max(0) as usize + } else { + 0 + }; + + // Build the same geoarrow schema from_reader_iter would, then restrict to this shard's row groups. + let geo_metadata = builder + .geoparquet_metadata() + .transpose() + .map_err(|e| e.to_string())? + .ok_or_else(|| "missing geoparquet metadata".to_string())?; + let geoarrow_schema = builder + .geoarrow_schema(&geo_metadata, true, Default::default()) + .map_err(|e| e.to_string())?; + let reader = builder + .with_row_groups(row_groups.to_vec()) + .build() + .map_err(|e| e.to_string())?; + let reader = + GeoParquetRecordBatchReader::try_new(reader, geoarrow_schema).map_err(|e| e.to_string())?; + + for record_batch in reader { + let record_batch: RecordBatch = record_batch.map_err(|e| e.to_string())?; + let items = + stac::geoarrow::items_from_record_batch(record_batch).map_err(|e| e.to_string())?; + let batch: Vec = items + .into_iter() + .map(serde_json::to_value) + .collect::>() + .map_err(|e| e.to_string())?; + let bytes = batch.len().saturating_mul(per_item_bytes); + on_batch(batch, bytes)?; + } + Ok(()) +} + +/// Parallel stac-geoparquet decode: shard one file's row groups over `n` threads, each feeding `tx`. +/// +/// Reads the file footer once to count row groups, shards them with [`shard_row_groups`] (round-robin, +/// capped at the row-group count), and spawns one decoder thread per shard. Each thread opens its own +/// file handle and decodes only its row groups — so every row group is decoded exactly once and the +/// union loads exactly the items a single-thread pass would, in no particular order (fine for load). +/// Every batch is sent as `(Vec, byte_estimate)`; the bounded channel applies backpressure +/// (decode pauses when the loaders fall behind). A decode error is sent as `Err(String)` and that +/// thread stops. Returns the spawned thread handles for the caller to join. +/// +/// # Examples +/// +/// ```no_run +/// use std::path::Path; +/// use serde_json::Value; +/// +/// # tokio_test::block_on(async { +/// let (tx, mut rx) = tokio::sync::mpsc::channel::, usize), String>>(4); +/// let handles = pgstac::parquet_decode::spawn_parquet_decoders(Path::new("items.parquet"), 4, tx) +/// .unwrap(); +/// while let Some(msg) = rx.recv().await { +/// let (items, _bytes) = msg.unwrap(); +/// println!("{} items", items.len()); +/// } +/// for h in handles { +/// let _ = h.join(); +/// } +/// # }) +/// ``` +pub fn spawn_parquet_decoders( + path: &Path, + n: usize, + tx: Sender, usize), String>>, +) -> std::io::Result>> { + // Read the footer once to learn the row-group count, then shard. Opening the file again per thread + // (below) re-reads the footer, which is cheap relative to decoding the data pages. + let num_row_groups = { + let file = std::fs::File::open(path)?; + let builder = ParquetRecordBatchReaderBuilder::try_new(file) + .map_err(|e| std::io::Error::other(e.to_string()))?; + builder.metadata().num_row_groups() + }; + + let shards = shard_row_groups(num_row_groups, n); + // A file with no row groups spawns no decoders: `tx` is dropped on return, so the loader's channel + // closes immediately and it sees an empty load rather than hanging. + if shards.is_empty() { + return Ok(Vec::new()); + } + + let mut handles = Vec::with_capacity(shards.len()); + for shard in shards { + let path = path.to_path_buf(); + let tx = tx.clone(); + handles.push(std::thread::spawn(move || { + let result: Result<(), String> = decode_row_groups(&path, &shard, |batch, bytes| { + // `Err` here only means the receiver is gone (early stop / downstream error); treat it + // as a stop signal, not a decode failure, so we don't also send a spurious Err. + if tx.blocking_send(Ok((batch, bytes))).is_err() { + Err("receiver closed".to_string()) + } else { + Ok(()) + } + }); + if let Err(e) = result + && e != "receiver closed" + { + let _ = tx.blocking_send(Err(e)); + } + })); + } + Ok(handles) +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + use stac::ItemCollection; + use std::collections::BTreeMap; + use std::collections::BTreeSet; + + /// Round-robin sharding covers every row group exactly once, never empties a shard while indices + /// remain, and caps shard count at the row-group count. + #[test] + fn shard_row_groups_partitions_exactly() { + for num in 0..40usize { + for n in 1..12usize { + let shards = shard_row_groups(num, n); + // Cap: at most min(n, num) shards, and no empty shards when any exist. + assert!(shards.len() <= n.min(num).max(if num == 0 { 0 } else { 1 })); + assert_eq!(shards.len(), n.min(num)); + for s in &shards { + assert!(!s.is_empty(), "no empty shards (num={num}, n={n})"); + } + // Exact partition: union == 0..num, no duplicates. + let mut seen: BTreeSet = BTreeSet::new(); + for s in &shards { + for &rg in s { + assert!( + seen.insert(rg), + "row group {rg} assigned twice (num={num}, n={n})" + ); + } + } + let expected: BTreeSet = (0..num).collect(); + assert_eq!(seen, expected, "every row group covered (num={num}, n={n})"); + } + } + } + + /// Writes `count` distinct items as stac-geoparquet with a small row-group size, forcing several + /// row groups so the parallel decode actually shards. + fn write_multi_rowgroup_parquet(count: usize, rows_per_group: usize) -> Vec { + use stac::IntoGeoparquet; + let items: Vec = (0..count) + .map(|i| { + let mut item = stac::Item::new(format!("item-{i:05}")); + item.collection = Some("c".to_string()); + item.geometry = Some( + serde_json::from_value(json!({ + "type": "Point", + "coordinates": [(-180.0 + i as f64 * 0.001), 40.0], + })) + .unwrap(), + ); + item.properties.datetime = Some("2023-01-01T00:00:00Z".parse().unwrap()); + let _ = item + .properties + .additional_fields + .insert("seq".into(), (i as i64).into()); + item + }) + .collect(); + let options = + stac::geoparquet::WriterOptions::new().with_max_row_group_row_count(rows_per_group); + ItemCollection::from(items) + .into_geoparquet_vec(options) + .unwrap() + } + + /// Collects every item the parallel decoder yields for a file, keyed by id, with the total count. + fn decode_parallel_all(path: &Path, n: usize) -> (usize, BTreeMap) { + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(2) + .enable_all() + .build() + .unwrap(); + runtime.block_on(async move { + let (tx, mut rx) = tokio::sync::mpsc::channel::, usize), String>>(4); + let handles = spawn_parquet_decoders(path, n, tx).unwrap(); + let mut count = 0usize; + let mut by_id: BTreeMap = BTreeMap::new(); + while let Some(msg) = rx.recv().await { + let (items, _bytes) = msg.unwrap(); + count += items.len(); + for item in items { + let id = item["id"].as_str().unwrap().to_string(); + assert!( + by_id.insert(id.clone(), item).is_none(), + "duplicate id {id}" + ); + } + } + for h in handles { + h.join().unwrap(); + } + (count, by_id) + }) + } + + /// The parallel decode loads exactly the same items (count + content) as the single-thread + /// `from_reader_iter` path, with no dropped or duplicated rows — across several shard counts and a + /// file with multiple row groups. + #[test] + fn parallel_decode_matches_single_thread() { + let count = 2_500usize; + let bytes = write_multi_rowgroup_parquet(count, 300); // ~9 row groups + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("items.parquet"); + std::fs::write(&path, &bytes).unwrap(); + + // Single-thread baseline via the same API the loader uses today. + let mut single: BTreeMap = BTreeMap::new(); + for batch in + stac::geoparquet::from_reader_iter(std::fs::File::open(&path).unwrap()).unwrap() + { + for item in batch.unwrap() { + let v = serde_json::to_value(item).unwrap(); + let id = v["id"].as_str().unwrap().to_string(); + let _ = single.insert(id, v); + } + } + assert_eq!(single.len(), count, "baseline sees every item"); + + for n in [1usize, 2, 3, 4, 8, 16] { + let (parallel_count, parallel) = decode_parallel_all(&path, n); + assert_eq!(parallel_count, count, "n={n}: item count matches"); + assert_eq!(parallel.len(), count, "n={n}: no duplicate ids"); + assert_eq!(parallel, single, "n={n}: identical items to single-thread"); + } + } + + /// A single-row-group file still decodes correctly when more threads are requested than row groups + /// (the shard count caps at 1). + #[test] + fn parallel_decode_single_rowgroup() { + let count = 50usize; + let bytes = write_multi_rowgroup_parquet(count, 150_000); // one row group + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("items.parquet"); + std::fs::write(&path, &bytes).unwrap(); + let (parallel_count, parallel) = decode_parallel_all(&path, 8); + assert_eq!(parallel_count, count); + assert_eq!(parallel.len(), count); + } +} diff --git a/src/pgstac-rs/src/load/pool_ingest.rs b/src/pgstac-rs/src/load/pool_ingest.rs new file mode 100644 index 00000000..da155a66 --- /dev/null +++ b/src/pgstac-rs/src/load/pool_ingest.rs @@ -0,0 +1,145 @@ +//! Write methods on [`PgstacPool`]: the Rust ingest path. +//! +//! Item writes dehydrate in Rust and binary-COPY through the SECURITY DEFINER staging/flush functions +//! (direct writes to `items` are revoked from `pgstac_ingest`). Collection writes use the SQL +//! `create_collection` / `update_collection` / `delete_collection` functions. + +use crate::PgstacPool; +use crate::Result; +use crate::ingest::{ConflictPolicy, load_items}; +use serde_json::Value; + +impl PgstacPool { + /// Loads a batch of full STAC items via the Rust dehydrate + binary-COPY pipeline, resolving id + /// collisions by `policy`. Returns the number of rows flushed. + /// + /// # Examples + /// + /// ```no_run + /// use pgstac::{ConnectConfig, PgstacPool}; + /// use pgstac::ingest::ConflictPolicy; + /// use serde_json::json; + /// + /// # tokio_test::block_on(async { + /// let pool = PgstacPool::connect(ConnectConfig::from_env()).await.unwrap(); + /// let item = json!({"type": "Feature", "id": "a", "collection": "c1", /* ... */ }); + /// pool.create_items(vec![item], ConflictPolicy::Upsert).await.unwrap(); + /// # }) + /// ``` + pub async fn create_items(&self, items: Vec, policy: ConflictPolicy) -> Result { + let schema = self.cached_dehydrate_schema().await?; + let mut client = self.get().await?; + load_items(&mut client, items, &schema, policy).await + } + + /// Loads `items`, skipping any whose content is unchanged from what's already stored (a per-partition + /// content-hash pre-filter) and loading only new + changed items via `create_items(policy)`. Returns + /// `(unchanged_skipped, loaded)`. + pub async fn upsert_items_precheck( + &self, + items: Vec, + policy: ConflictPolicy, + ) -> Result<(u64, u64)> { + crate::ingest::precheck_upsert(self, items, policy).await + } + + /// Inserts a single item, erroring if its id already exists in the collection. + pub async fn create_item(&self, item: Value) -> Result { + self.create_items(vec![item], ConflictPolicy::Error).await + } + + /// Upserts a single item, replacing it if its content changed. + pub async fn upsert_item(&self, item: Value) -> Result { + self.create_items(vec![item], ConflictPolicy::Upsert).await + } + + /// Deletes an item by id from a collection. + pub async fn delete_item(&self, collection_id: &str, item_id: &str) -> Result<()> { + let client = self.get().await?; + let _ = client + .execute("SELECT delete_item($1, $2)", &[&item_id, &collection_id]) + .await?; + Ok(()) + } + + /// Creates a collection from its STAC JSON (deriving `fragment_config` from `item_assets`). + pub async fn create_collection(&self, collection: &Value) -> Result<()> { + let client = self.get().await?; + let _ = client + .execute("SELECT create_collection($1::jsonb)", &[collection]) + .await?; + Ok(()) + } + + /// Inserts or updates a collection (idempotent), setting `partition_trunc` on insert and preserving it + /// (same value) on conflict. This is what makes `load`/`restore` re-runnable: a second pass updates the + /// collection content instead of failing on the duplicate id, so an interrupted restore resumes cleanly. + pub async fn upsert_collection( + &self, + collection: &Value, + partition_trunc: Option<&str>, + ) -> Result<()> { + let client = self.get().await?; + let _ = client + .execute( + "SELECT upsert_collection($1::jsonb, $2::text)", + &[collection, &partition_trunc], + ) + .await?; + Ok(()) + } + + /// Runs the async partition-stats maintenance: recompute exact + /// bounds + row counts for partitions ingest left dirty (oldest first). `limit` caps the batch + /// (`None` = all dirty). Returns the number of partitions tightened. Safe + optional — a generous + /// (un-tightened) envelope only over-includes a partition in search, never loses rows. + pub async fn tighten_dirty_stats(&self, limit: Option) -> Result { + let client = self.get().await?; + let row = client + .query_one("SELECT tighten_dirty_partition_stats($1::int)", &[&limit]) + .await?; + Ok(row.get(0)) + } + + /// Sets a collection's `partition_trunc` (`Some("year")`/`Some("month")`, or `None` for a single + /// partition). The collections trigger repartitions existing items; on a freshly-created (empty) + /// collection it just records the setting, so set it before loading items (e.g. on restore, to + /// recreate the dumped partition layout). + pub async fn set_partition_trunc( + &self, + collection_id: &str, + partition_trunc: Option<&str>, + ) -> Result<()> { + let client = self.get().await?; + let _ = client + .execute( + "UPDATE collections SET partition_trunc = $2 WHERE id = $1", + &[&collection_id, &partition_trunc], + ) + .await?; + Ok(()) + } + + /// Replaces a collection's content (preserving its operator-configured partitioning/fragmenting). + pub async fn update_collection(&self, collection: &Value) -> Result<()> { + let client = self.get().await?; + let _ = client + .execute("SELECT update_collection($1::jsonb)", &[collection]) + .await?; + Ok(()) + } + + /// Deletes a collection (and its items) by id. + pub async fn delete_collection(&self, collection_id: &str) -> Result<()> { + let client = self.get().await?; + let _ = client + .execute("SELECT delete_collection($1)", &[&collection_id]) + .await?; + Ok(()) + } +} + +// The rustac `stac::api::TransactionClient` (+ `ItemsClient` / `CollectionsClient`) live on [`crate::Client`]: +// call `pool.client().await?` for the idiomatic per-request client (it shares the pool's hydration cache). +// The pool keeps the native `create_items` / `create_collection` / `upsert_collection` write methods above, +// which are what those traits delegate to. diff --git a/src/pgstac-rs/src/load/queryables.rs b/src/pgstac-rs/src/load/queryables.rs new file mode 100644 index 00000000..9417d377 --- /dev/null +++ b/src/pgstac-rs/src/load/queryables.rs @@ -0,0 +1,5 @@ +#![allow(dead_code)] +// Stub for queryables loading +pub(crate) async fn load_queryables() { + unimplemented!("load_queryables mapping to pgstac queryables"); +} diff --git a/src/pgstac-rs/src/page.rs b/src/pgstac-rs/src/page.rs deleted file mode 100644 index c4bf8947..00000000 --- a/src/pgstac-rs/src/page.rs +++ /dev/null @@ -1,75 +0,0 @@ -use serde::{Deserialize, Serialize}; -use serde_json::{Map, Value}; -use stac::Link; -use stac::api::{Context, Item}; - -/// A page of search results. -#[derive(Debug, Deserialize, Serialize)] -pub struct Page { - /// These are the out features, usually STAC items, but maybe not legal STAC - /// items if fields are excluded. - pub features: Vec, - - /// The next id. - #[serde(skip_serializing_if = "Option::is_none")] - pub next: Option, - - /// The previous id. - #[serde(skip_serializing_if = "Option::is_none")] - pub prev: Option, - - /// The search context. - /// - /// This was removed in pgstac v0.9 - #[serde(skip_serializing_if = "Option::is_none")] - pub context: Option, - - /// The number of values returned. - /// - /// Added in pgstac v0.9 - #[serde(rename = "numberReturned", skip_serializing_if = "Option::is_none")] - pub number_returned: Option, - - /// Links - /// - /// Added in pgstac v0.9 - #[serde(default, skip_serializing_if = "Vec::is_empty")] - pub links: Vec, - - /// Additional fields. - #[serde(flatten)] - pub additional_fields: Map, -} - -impl Page { - /// Returns this page's next token, if it has one. - pub fn next_token(&self) -> Option { - if let Some(next) = &self.next { - return Some(format!("next:{next}")); - } - - self.links - .iter() - .find(|link| link.rel == "next") - .and_then(|link| extract_token_from_href(&link.href)) - } - - /// Returns this page's prev token, if it has one. - pub fn prev_token(&self) -> Option { - if let Some(prev) = &self.prev { - return Some(format!("prev:{prev}")); - } - - self.links - .iter() - .find(|link| link.rel == "prev") - .and_then(|link| extract_token_from_href(&link.href)) - } -} - -fn extract_token_from_href(href: &str) -> Option { - href.split("token=") - .nth(1) - .and_then(|token_part| token_part.split('&').next()) - .map(|token| token.to_string()) -} diff --git a/src/pgstac-rs/src/python/mod.rs b/src/pgstac-rs/src/python/mod.rs new file mode 100644 index 00000000..6cc2bda8 --- /dev/null +++ b/src/pgstac-rs/src/python/mod.rs @@ -0,0 +1,394 @@ +//! Python extension module (built with maturin) exposing the pgstac **read + write** API over a +//! connection pool, for use in `stac-fastapi-pgstac` and similar. +//! +//! The pool is created once (`Pgstac.connect`) at application startup and shared. Every method is +//! `async` (a Python awaitable via `pyo3-async-runtimes`). Reads return JSON strings (the caller does +//! `orjson.loads`); writes go through the **Rust loader** (dehydration + fragment splitting + the +//! binary COPY all in Rust), driving the same engine as the rest of the crate. + +#![cfg(feature = "python")] + +use crate::ingest::ConflictPolicy; +use crate::{ConnectConfig, DEFAULT_SEARCH_PATH, PgstacPool}; +use pyo3::prelude::*; +use pyo3_async_runtimes::tokio::future_into_py; +use serde_json::{Value, json}; + +/// Serializes a value to a JSON string for the Python boundary. +fn encode(value: &T) -> PyResult { + serde_json::to_string(value).map_err(pyerr) +} + +fn pyerr(error: E) -> PyErr { + pyo3::exceptions::PyRuntimeError::new_err(error.to_string()) +} + +fn parse_search(search: &str) -> PyResult { + serde_json::from_str(search).map_err(pyerr) +} + +/// Parse the write input into a list of item values: a JSON array, a +/// FeatureCollection (`features`), or a single item object. +fn parse_items(items: &str) -> PyResult> { + match serde_json::from_str(items).map_err(pyerr)? { + Value::Array(items) => Ok(items), + Value::Object(mut map) => match map.remove("features") { + Some(Value::Array(features)) => Ok(features), + _ => Ok(vec![Value::Object(map)]), + }, + other => Ok(vec![other]), + } +} + +/// Map the conflict-policy string ("upsert" default / "ignore" / "error"). +fn parse_policy(policy: Option<&str>) -> PyResult { + Ok(match policy.unwrap_or("upsert") { + "upsert" => ConflictPolicy::Upsert, + "ignore" => ConflictPolicy::Ignore, + "error" => ConflictPolicy::Error, + other => return Err(pyerr(format!("unknown conflict policy: {other}"))), + }) +} + +/// A pooled pgstac client exposing the read API to Python. +#[pyclass] +struct Pgstac { + pool: PgstacPool, +} + +#[pymethods] +impl Pgstac { + /// Connects a pool. `dsn` defaults to the standard libpq environment (`PGHOST`, `PGDATABASE`, …). + /// + /// Returns an awaitable resolving to a `Pgstac`. Call once at startup and reuse. + #[staticmethod] + #[pyo3(signature = (dsn=None))] + fn connect(py: Python<'_>, dsn: Option) -> PyResult> { + future_into_py(py, async move { + let config = match dsn { + Some(dsn) => ConnectConfig { + dsn: Some(dsn), + search_path: DEFAULT_SEARCH_PATH.to_string(), + ..Default::default() + }, + None => ConnectConfig::from_env(), + }; + let pool = PgstacPool::connect(config).await.map_err(pyerr)?; + Ok(Pgstac { pool }) + }) + } + + /// One page of an item search. Returns a FeatureCollection (with `next`/`prev` tokens, + /// `numberReturned`, `numberMatched`) as JSON bytes. + #[pyo3(signature = (search, token=None, limit=10))] + fn search<'py>( + &self, + py: Python<'py>, + search: &str, + token: Option, + limit: i64, + ) -> PyResult> { + let pool = self.pool.clone(); + let search = parse_search(search)?; + // The search body's `limit` (the STAC contract) wins over the method default. + let limit = search.get("limit").and_then(Value::as_i64).unwrap_or(limit); + future_into_py(py, async move { + let page = pool + .search_page(&search, token.as_deref(), limit) + .await + .map_err(pyerr)?; + encode(&json!({ + "type": "FeatureCollection", + "features": page.features, + "numberReturned": page.number_returned, + "numberMatched": page.number_matched, + "next": page.next_token, + "prev": page.prev_token, + })) + }) + } + + /// Collects every matching item (up to `max_items`) into one FeatureCollection (JSON bytes), + /// draining the flat-memory streaming iterator. + #[pyo3(signature = (search, max_items=None))] + fn search_collect<'py>( + &self, + py: Python<'py>, + search: &str, + max_items: Option, + ) -> PyResult> { + let pool = self.pool.clone(); + let search = parse_search(search)?; + future_into_py(py, async move { + let items = pool + .search_collect_items(search, max_items) + .await + .map_err(pyerr)?; + let number_returned = items.len(); + encode(&json!({ + "type": "FeatureCollection", + "features": items, + "numberReturned": number_returned, + })) + }) + } + + /// The total match count (`numberMatched`) for a search, run independently (use concurrently with + /// `search` for the parallel-context pattern). `None` when context counting is off. + fn search_matched<'py>(&self, py: Python<'py>, search: &str) -> PyResult> { + let pool = self.pool.clone(); + let search = parse_search(search)?; + future_into_py(py, async move { + pool.search_matched(&search).await.map_err(pyerr) + }) + } + + /// Streams every matching item (up to `max_items`) into stac-geoparquet, returning the file bytes. + /// + /// `row_group_size` caps rows per parquet row-group (smaller = lower peak memory while encoding; + /// `None` = the parquet default). `compression` is a parquet codec in parquet's spelling (e.g. + /// `snappy`, `uncompressed`, `zstd(15)`); `None` (default) uses the fast default codec. + #[pyo3(signature = (search, max_items=None, row_group_size=None, compression=None))] + fn search_to_geoparquet<'py>( + &self, + py: Python<'py>, + search: &str, + max_items: Option, + row_group_size: Option, + compression: Option, + ) -> PyResult> { + let pool = self.pool.clone(); + let search = parse_search(search)?; + let compression = match compression { + Some(name) => name + .parse::() + .map_err(pyerr)?, + None => crate::export::format::default_compression(), + }; + future_into_py(py, async move { + let mut buf: Vec = Vec::new(); + let _ = pool + .stream_geoparquet(search, max_items, compression, row_group_size, &mut buf) + .await + .map_err(pyerr)?; + Ok(buf) + }) + } + + /// One page of a collection search. Returns the collections + tokens as JSON bytes. + #[pyo3(signature = (search, token=None))] + fn collection_search<'py>( + &self, + py: Python<'py>, + search: &str, + token: Option, + ) -> PyResult> { + let pool = self.pool.clone(); + let search = parse_search(search)?; + future_into_py(py, async move { + let page = pool + .collection_search(&search, token.as_deref()) + .await + .map_err(pyerr)?; + encode(&json!({ + "collections": page.features, + "numberReturned": page.number_returned, + "numberMatched": page.number_matched, + "next": page.next_token, + "prev": page.prev_token, + })) + }) + } + + /// Fetches one item by id; JSON bytes, or `None` if it does not exist. + fn get_item<'py>( + &self, + py: Python<'py>, + collection_id: String, + item_id: String, + ) -> PyResult> { + let pool = self.pool.clone(); + future_into_py(py, async move { + let item = pool + .get_item(&collection_id, &item_id) + .await + .map_err(pyerr)?; + item.map(|v| encode(&v)).transpose() + }) + } + + /// Fetches one collection by id; JSON bytes, or `None` if it does not exist. + fn get_collection<'py>( + &self, + py: Python<'py>, + collection_id: String, + ) -> PyResult> { + let pool = self.pool.clone(); + future_into_py(py, async move { + let collection = pool.get_collection(&collection_id).await.map_err(pyerr)?; + collection.map(|v| encode(&v)).transpose() + }) + } + + /// The queryables document for a collection (or catalog-wide when `collection_id` is `None`). + #[pyo3(signature = (collection_id=None))] + fn get_queryables<'py>( + &self, + py: Python<'py>, + collection_id: Option, + ) -> PyResult> { + let pool = self.pool.clone(); + future_into_py(py, async move { + let queryables = pool + .get_queryables(collection_id.as_deref()) + .await + .map_err(pyerr)?; + encode(&queryables) + }) + } + + /// Loads items through the Rust loader (dehydrate + fragment split + binary COPY). `items` is a + /// JSON array, a FeatureCollection, or a single item. `policy` is "upsert" (default), "ignore", or + /// "error". Returns the number of rows written. Use for both single-item (stac-fastapi POST) and + /// bulk ingest — the same pipeline at different batch sizes. + #[pyo3(signature = (items, policy=None))] + fn create_items<'py>( + &self, + py: Python<'py>, + items: &str, + policy: Option, + ) -> PyResult> { + let pool = self.pool.clone(); + let items = parse_items(items)?; + let policy = parse_policy(policy.as_deref())?; + future_into_py(py, async move { + pool.create_items(items, policy).await.map_err(pyerr) + }) + } + + /// Creates a collection from its STAC JSON (deriving `fragment_config` from `item_assets`). + fn create_collection<'py>( + &self, + py: Python<'py>, + collection: &str, + ) -> PyResult> { + let pool = self.pool.clone(); + let collection = parse_search(collection)?; + future_into_py(py, async move { + pool.create_collection(&collection).await.map_err(pyerr) + }) + } + + /// Replaces a collection's content (preserving its operator-configured partitioning/fragmenting). + fn update_collection<'py>( + &self, + py: Python<'py>, + collection: &str, + ) -> PyResult> { + let pool = self.pool.clone(); + let collection = parse_search(collection)?; + future_into_py(py, async move { + pool.update_collection(&collection).await.map_err(pyerr) + }) + } + + /// Deletes a collection (and its items) by id. + fn delete_collection<'py>( + &self, + py: Python<'py>, + collection_id: String, + ) -> PyResult> { + let pool = self.pool.clone(); + future_into_py(py, async move { + pool.delete_collection(&collection_id).await.map_err(pyerr) + }) + } + + /// Sets a collection's partition truncation: "year", "month", or `None` for a single partition. + /// On a collection with items this repartitions; set it before loading to choose the layout. + #[pyo3(signature = (collection_id, partition_trunc=None))] + fn set_partition_trunc<'py>( + &self, + py: Python<'py>, + collection_id: String, + partition_trunc: Option, + ) -> PyResult> { + let pool = self.pool.clone(); + future_into_py(py, async move { + pool.set_partition_trunc(&collection_id, partition_trunc.as_deref()) + .await + .map_err(pyerr) + }) + } + + /// Deletes an item by id from a collection. + fn delete_item<'py>( + &self, + py: Python<'py>, + collection_id: String, + item_id: String, + ) -> PyResult> { + let pool = self.pool.clone(); + future_into_py(py, async move { + pool.delete_item(&collection_id, &item_id) + .await + .map_err(pyerr) + }) + } + + /// Runs the async partition-stats maintenance: recompute exact + /// bounds + row counts for partitions ingest left dirty. `limit` caps the batch (oldest first); + /// `None` tightens all. Returns the number of partitions tightened. Schedule off-hours. + #[pyo3(signature = (limit=None))] + fn maintain<'py>(&self, py: Python<'py>, limit: Option) -> PyResult> { + let pool = self.pool.clone(); + future_into_py(py, async move { + pool.tighten_dirty_stats(limit).await.map_err(pyerr) + }) + } + + #[pyo3(signature = (_file, _collection_ids=None, _delete_missing=None, _index_fields=None))] + fn load_queryables<'py>( + &self, + py: Python<'py>, + _file: String, + _collection_ids: Option>, + _delete_missing: Option, + _index_fields: Option>, + ) -> PyResult> { + future_into_py(py, async move { + Err::<(), _>(pyo3::exceptions::PyNotImplementedError::new_err("load_queryables is stubbed in rust")) + }) + } + + #[pyo3(signature = (_toversion=None))] + fn migrate<'py>( + &self, + py: Python<'py>, + _toversion: Option, + ) -> PyResult> { + future_into_py(py, async move { + Err::(pyo3::exceptions::PyNotImplementedError::new_err("migrate is stubbed in rust")) + }) + } + + fn runqueue<'py>(&self, py: Python<'py>) -> PyResult> { + future_into_py(py, async move { + Err::(pyo3::exceptions::PyNotImplementedError::new_err("runqueue is stubbed in rust")) + }) + } + + fn pgready<'py>(&self, py: Python<'py>) -> PyResult> { + future_into_py(py, async move { + Err::<(), _>(pyo3::exceptions::PyNotImplementedError::new_err("pgready is stubbed in rust")) + }) + } +} + + +/// The `pypgstac_rs` extension module. +#[pymodule] +fn pypgstac_rs(m: &Bound<'_, PyModule>) -> PyResult<()> { + m.add_class::()?; + Ok(()) +} diff --git a/src/pgstac-rs/src/read/collections.rs b/src/pgstac-rs/src/read/collections.rs new file mode 100644 index 00000000..0d8b8d4e --- /dev/null +++ b/src/pgstac-rs/src/read/collections.rs @@ -0,0 +1,144 @@ +//! Client-side collection search and lookups, driven by `collection_search_plan`. +//! +//! Collections are not split-stored, so there is no hydration and no datetime band engine: the plan +//! returns a query yielding `(content, keys)` rows. This walks that query, applies keyset pagination, +//! and mints tokens **in Rust** ([`crate::keyset::keyset_encode`]) from the per-row `keys`. + +use crate::Result; +use crate::keyset::keyset_encode; +use crate::search::SearchPage; +use serde_json::{Value, json}; +use tokio_postgres::GenericClient; + +/// Runs one page of a collection search and returns the matching collections plus keyset tokens. +/// +/// Mirrors the SQL `collection_search` paging contract (including the `prev` reversal) but does the +/// work client-side off `collection_search_plan`. The returned [`SearchPage`] holds the collections in +/// `features`. +pub async fn collection_search( + client: &C, + search: &Value, + token: Option<&str>, +) -> Result { + let is_prev = token.is_some_and(|t| t.starts_with("prev:")); + // A non-empty keyset (the token minus its next:/prev: prefix) means this is not the first page. + let has_keyset = token + .map(|t| t.trim_start_matches("next:").trim_start_matches("prev:")) + .is_some_and(|k| !k.is_empty()); + let limit = search + .get("limit") + .and_then(Value::as_i64) + .filter(|&n| n > 0) + .unwrap_or(10); + + let plan = client + .query_one( + "SELECT query, ctx_query FROM collection_search_plan($1::jsonb, $2)", + &[search, &token], + ) + .await?; + let query: String = plan.try_get("query")?; + let ctx_query: Option = plan.try_get("ctx_query")?; + + let number_matched = match ctx_query { + Some(ctx) if !ctx.is_empty() => { + Some(client.query_one(&ctx, &[]).await?.try_get::<_, i64>(0)?) + } + _ => None, + }; + + // The plan query takes the page size (+1 to detect a further page) as a bound parameter and + // returns (content jsonb, keys text[]). `LIMIT $1` infers $1 as bigint. + let over_limit: i64 = limit + 1; + let rows = client.query(&query, &[&over_limit]).await?; + let has_more = rows.len() as i64 > limit; + + let mut features: Vec = Vec::new(); + let mut first_keys: Option>> = None; + let mut last_keys: Option>> = None; + for (i, row) in rows.iter().enumerate() { + let keys: Vec> = row.try_get(1)?; + if i == 0 { + first_keys = Some(keys.clone()); + } + if (i as i64) < limit { + features.push(row.try_get(0)?); + last_keys = Some(keys); + } + } + + let (fwd_first, fwd_last, next_present, prev_present) = if is_prev { + features.reverse(); + (last_keys, first_keys, has_keyset, has_more) + } else { + (first_keys, last_keys, has_more, has_keyset) + }; + + let mut next_token = None; + let mut prev_token = None; + if !features.is_empty() { + if let (true, Some(keys)) = (next_present, &fwd_last) { + next_token = Some(format!("next:{}", keyset_encode(keys))); + } + if let (true, Some(keys)) = (prev_present, &fwd_first) { + prev_token = Some(format!("prev:{}", keyset_encode(keys))); + } + } + + Ok(SearchPage { + number_returned: features.len(), + features, + next_token, + prev_token, + number_matched, + }) +} + +/// Fetches every collection, paging [`collection_search`] to completion. +pub(crate) async fn all(client: &C) -> Result> { + let mut all = Vec::new(); + let mut token: Option = None; + loop { + let page = collection_search(client, &json!({}), token.as_deref()).await?; + let next = page.next_token; + all.extend(page.features); + match next { + Some(next) => token = Some(next), + None => break, + } + } + Ok(all) +} + +/// Fetches a single collection by id (wraps [`collection_search`] with an id filter). +pub async fn get_collection( + client: &C, + collection_id: &str, +) -> Result> { + let body = json!({"ids": [collection_id], "limit": 1}); + let page = collection_search(client, &body, None).await?; + Ok(page.features.into_iter().next()) +} + +/// Fetches the queryables document for a collection, or the catalog-wide queryables when +/// `collection_id` is `None`. +pub async fn get_queryables( + client: &C, + collection_id: Option<&str>, +) -> Result { + let row = match collection_id { + Some(id) => { + client + .query_one("SELECT get_queryables($1::text) AS q", &[&id]) + .await? + } + None => { + // `get_queryables()` with no args is ambiguous between the text[] and text overloads; + // pin the text overload with an explicit NULL for the catalog-wide queryables. + client + .query_one("SELECT get_queryables(NULL::text) AS q", &[]) + .await? + } + }; + row.try_get("q").map_err(Into::into) +} diff --git a/src/pgstac-rs/src/read/feature.rs b/src/pgstac-rs/src/read/feature.rs new file mode 100644 index 00000000..0f82e8d6 --- /dev/null +++ b/src/pgstac-rs/src/read/feature.rs @@ -0,0 +1,269 @@ +//! Byte-assembly output for hydrated 0.10 items. +//! +//! [`DehydratedItem::write_fragment_feature`] writes a fully-hydrated STAC item straight to a writer +//! without ever materializing the merged `assets` as a [`Value`]: the item's `assets` are read as raw +//! JSON, the +//! shared fragment's `assets` are borrowed (by `&`, never cloned per item), and the two are merged +//! **at serialize time** — descending only where both sides are objects at a key, and emitting arrays, +//! scalars, and one-sided keys **verbatim** from the raw bytes. `bbox` is likewise emitted from raw +//! bytes (preserving PostgreSQL `numeric` precision that an f64 round-trip would lose). Everything else +//! reuses `hydrate_fragment_core`. +//! +//! The merge reproduces SQL `jsonb_merge_recursive(frag, item)` exactly (see the parallel +//! implementation in [`crate::hydrate`]); validated by parsing the output back and comparing to +//! `content_hydrate`. + +use crate::Result; +use crate::hydrate::{DehydratedItem, FragmentContext, hydrate_fragment_core}; +use serde::Serialize; +use serde::ser::{SerializeMap, Serializer}; +use serde_json::value::RawValue; +use serde_json::{Map, Value}; +use std::collections::BTreeMap; +use std::io::Write; + +impl DehydratedItem { + /// Writes this fully-hydrated 0.10 item as JSON to `write`. + pub fn write_fragment_feature( + self, + fragment: Option<&FragmentContext>, + write: &mut W, + ) -> Result<()> { + let frag_assets = fragment + .and_then(|f| f.content.as_ref()) + .and_then(|c| c.get("assets")); + let (core, assets_raw, bbox_raw) = hydrate_fragment_core(self, fragment); + let feature = FeatureWriter { + core: &core, + bbox: bbox_raw.as_ref().map(|r| r.0.as_ref()), + frag_assets, + item_assets: assets_raw.as_ref().map(|r| r.0.as_ref()), + }; + serde_json::to_writer(write, &feature)?; + Ok(()) + } +} + +/// Parses a raw JSON object one level into `key -> raw value`, or `None` when it is not an object. +fn raw_object(raw: &RawValue) -> Option> { + serde_json::from_str(raw.get()).ok() +} + +/// Whether a raw object has at least one key. +fn raw_object_nonempty(raw: &RawValue) -> bool { + raw_object(raw).is_some_and(|m| !m.is_empty()) +} + +/// The top-level feature: core keys, then raw `bbox`, then the serialize-time-merged `assets`. +struct FeatureWriter<'a> { + core: &'a Map, + bbox: Option<&'a RawValue>, + frag_assets: Option<&'a Value>, + item_assets: Option<&'a RawValue>, +} + +impl Serialize for FeatureWriter<'_> { + fn serialize(&self, serializer: S) -> std::result::Result { + let mut map = serializer.serialize_map(None)?; + for (k, v) in self.core { + map.serialize_entry(k, v)?; + } + if let Some(bbox) = self.bbox + && bbox.get().trim() != "null" + { + map.serialize_entry("bbox", bbox)?; + } + // assets := jsonb_merge_recursive(frag.assets, item.assets); emitted unless it would be {}. + let frag_nonempty = self + .frag_assets + .and_then(Value::as_object) + .is_some_and(|o| !o.is_empty()); + let item_nonempty = self.item_assets.is_some_and(raw_object_nonempty); + if frag_nonempty || item_nonempty { + map.serialize_entry( + "assets", + &Merge { + frag: self.frag_assets, + item: self.item_assets, + }, + )?; + } + map.end() + } +} + +/// Serializes `jsonb_merge_recursive(frag, item)` where `frag` is a parsed [`Value`] (shared, by ref) +/// and `item` is raw JSON (never deep-parsed). +struct Merge<'a> { + frag: Option<&'a Value>, + item: Option<&'a RawValue>, +} + +impl Serialize for Merge<'_> { + fn serialize(&self, serializer: S) -> std::result::Result { + match (self.frag, self.item) { + // frag NULL -> COALESCE(item, {}) + (None, Some(item)) => item.serialize(serializer), + (None, None) => serializer.serialize_map(Some(0))?.end(), + // item NULL -> frag + (Some(frag), None) => frag.serialize(serializer), + (Some(frag), Some(item)) => { + let item_obj = raw_object(item); + // item = {} -> frag (regardless of frag's type) + if item_obj.as_ref().is_some_and(|m| m.is_empty()) { + return frag.serialize(serializer); + } + match (frag.as_object(), item_obj) { + // both objects -> FULL JOIN merge + (Some(frag_obj), Some(item_obj)) => { + serialize_object_merge(frag_obj, &item_obj, serializer) + } + // else -> item wins, verbatim + _ => item.serialize(serializer), + } + } + } + } +} + +/// The both-objects FULL JOIN: frag keys (then item-only keys), with per-key merge rules matching +/// `jsonb_merge_recursive`'s inner branch. +fn serialize_object_merge( + frag: &Map, + item: &BTreeMap, + serializer: S, +) -> std::result::Result { + let mut map = serializer.serialize_map(None)?; + for (key, fv) in frag { + match item.get(key.as_str()) { + // frag-only + None => map.serialize_entry(key, fv)?, + Some(iv) => match (fv.as_object(), raw_object(iv)) { + // both objects: disjoint -> concat (f || i); else recurse + (Some(fvo), Some(ivo)) => { + let disjoint = !fvo.keys().any(|fk| ivo.contains_key(fk.as_str())); + if disjoint { + map.serialize_entry( + key, + &Concat { + frag: fvo, + item: &ivo, + }, + )?; + } else { + map.serialize_entry( + key, + &Merge { + frag: Some(fv), + item: Some(iv), + }, + )?; + } + } + // else -> item wins, verbatim + _ => map.serialize_entry(key, iv)?, + }, + } + } + for (key, iv) in item { + if !frag.contains_key(key.as_str()) { + map.serialize_entry(key, iv)?; + } + } + map.end() +} + +/// `frag.value || item.value` shallow concat for the disjoint-objects case (no key overlap, so order +/// is frag's keys then item's). +struct Concat<'a> { + frag: &'a Map, + item: &'a BTreeMap, +} + +impl Serialize for Concat<'_> { + fn serialize(&self, serializer: S) -> std::result::Result { + let mut map = serializer.serialize_map(None)?; + for (k, v) in self.frag { + map.serialize_entry(k, v)?; + } + for (k, v) in self.item { + if !self.frag.contains_key(k.as_str()) { + map.serialize_entry(k, v)?; + } + } + map.end() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::hydrate::{ + CollectionContext, DehydratedItem, FragmentContext, HydrationModel, Hydrator, + PromotedProperties, + }; + use crate::rawjson::RawJson; + use serde_json::json; + + fn temporal(dt: &str) -> Map { + serde_json::from_value(json!({ "datetime": dt })).unwrap() + } + + #[test] + fn byte_path_equals_value_path_and_preserves_bbox_precision() { + // Exercises every merge branch: recurse (data), disjoint concat (data.nested), + // item-wins-on-leaf (data.type), frag-only (thumbnail), item-only (item_only). + let make_item = || DehydratedItem { + id: "i1".into(), + collection: "c1".into(), + geometry: Some(json!({"type": "Point", "coordinates": [1, 2]})), + // Built from raw text (NOT json!, which would round through f64 first) so the high-precision + // coordinate reaches the byte path verbatim. + bbox: Some(RawJson( + RawValue::from_string("[-22.650799880000001,1.0,2.0,3.0]".to_string()).unwrap(), + )), + assets: Some(RawJson::from_value(&json!({ + "data": {"href": "s3://x", "type": "override", "nested": {"i": 1}}, + "item_only": {"k": [1, 2, 3]} + }))), + stac_extensions: Some(json!(["https://ext/v1"])), + promoted: PromotedProperties { + temporal: temporal("2020-01-01T00:00:00Z"), + fields: vec![("platform", json!("l8"))], + owned_fields: vec![], + }, + fragment_id: Some(7), + ..Default::default() + }; + let frag = FragmentContext { + content: Some(json!({ + "assets": { + "data": {"type": "image/tiff", "nested": {"j": 2}}, + "thumbnail": {"href": "t", "big": [0, 0, 0]} + }, + "stac_version": "1.0.0" + })), + links_template: None, + }; + + let h = Hydrator::new(HydrationModel::Fragment); + let value = h.hydrate(make_item(), &CollectionContext::default(), Some(&frag)); + + let mut buf: Vec = Vec::new(); + make_item().write_fragment_feature(Some(&frag), &mut buf).unwrap(); + let byte_value: Value = serde_json::from_slice(&buf).unwrap(); + + assert_eq!(byte_value, value, "byte path != value path"); + assert_eq!(byte_value["assets"]["data"]["type"], "override"); // item wins on leaf + assert_eq!( + byte_value["assets"]["data"]["nested"], + json!({"j": 2, "i": 1}) + ); // disjoint concat + assert!(byte_value["assets"]["thumbnail"]["big"].is_array()); // frag-only, verbatim + // bbox numeric precision preserved verbatim in the byte output. + assert!( + String::from_utf8_lossy(&buf).contains("-22.650799880000001"), + "bbox precision lost" + ); + } +} diff --git a/src/pgstac-rs/src/read/fields.rs b/src/pgstac-rs/src/read/fields.rs new file mode 100644 index 00000000..77bc4db7 --- /dev/null +++ b/src/pgstac-rs/src/read/fields.rs @@ -0,0 +1,193 @@ +//! STAC `fields` include/exclude projection, replicating pgstac's `jsonb_fields`. +//! +//! `jsonb_fields(j, f) = jsonb_exclude(jsonb_include(j, f), f)`: +//! * **include** — when `fields.include` is a non-empty array, the output keeps only those dot-paths +//! (plus the always-required `id` and `collection`); otherwise the item passes through unchanged. +//! * **exclude** — every `fields.exclude` dot-path is then removed. +//! +//! Paths are split on `.` (e.g. `properties.eo:cloud_cover`); numeric segments index arrays. + +use serde_json::{Map, Value}; + +/// Applies a STAC `fields` object (`{ "include": [...], "exclude": [...] }`) to a hydrated item, +/// matching pgstac's `jsonb_fields`. +/// +/// # Examples +/// +/// ``` +/// use pgstac::fields::apply_fields; +/// use serde_json::json; +/// +/// let item = json!({ +/// "id": "x", "collection": "c", +/// "assets": {"data": {"href": "h"}}, +/// "properties": {"datetime": "2020-01-01T00:00:00Z", "eo:cloud_cover": 5} +/// }); +/// +/// // Exclude drops the listed dot-paths. +/// let trimmed = apply_fields(item.clone(), &json!({"exclude": ["assets", "properties.eo:cloud_cover"]})); +/// assert!(trimmed.get("assets").is_none()); +/// assert_eq!(trimmed["properties"]["datetime"], "2020-01-01T00:00:00Z"); +/// +/// // Include keeps only the listed paths, plus the always-required id/collection. +/// let narrowed = apply_fields(item, &json!({"include": ["properties.datetime"]})); +/// assert_eq!(narrowed["id"], "x"); +/// assert!(narrowed.get("assets").is_none()); +/// ``` +pub fn apply_fields(item: Value, fields: &Value) -> Value { + exclude(include(item, fields), fields) +} + +/// Splits a `fields` list into dot-path segment vectors. +fn dotpaths(list: Option<&Value>) -> Vec> { + list.and_then(Value::as_array) + .map(|array| { + array + .iter() + .filter_map(Value::as_str) + .map(|s| s.split('.').map(String::from).collect()) + .collect() + }) + .unwrap_or_default() +} + +/// `jsonb_include`: keep only the included dot-paths (+ `id`/`collection`), or pass through when there +/// is no include list. +fn include(item: Value, fields: &Value) -> Value { + let mut paths = dotpaths(fields.get("include")); + if paths.is_empty() { + return item; + } + paths.push(vec!["id".to_string()]); + if item.get("collection").is_some() { + paths.push(vec!["collection".to_string()]); + } + let mut out = Value::Object(Map::new()); + for path in &paths { + let value = get_path(&item, path).cloned().unwrap_or(Value::Null); + set_nested(&mut out, path, value); + } + out +} + +/// `jsonb_exclude`: remove each excluded dot-path. +fn exclude(mut item: Value, fields: &Value) -> Value { + for path in dotpaths(fields.get("exclude")) { + remove_path(&mut item, &path); + } + item +} + +/// Reads the value at a dot-path (object key, or array index for a numeric segment). +fn get_path<'a>(value: &'a Value, path: &[String]) -> Option<&'a Value> { + let mut current = value; + for segment in path { + current = match current { + Value::Object(map) => map.get(segment)?, + Value::Array(array) => array.get(segment.parse::().ok()?)?, + _ => return None, + }; + } + Some(current) +} + +/// Sets `value` at a dot-path, creating intermediate objects (`jsonb_set_nested`). +fn set_nested(out: &mut Value, path: &[String], value: Value) { + if !out.is_object() { + *out = Value::Object(Map::new()); + } + let mut current = out; + for (i, segment) in path.iter().enumerate() { + let map = current.as_object_mut().expect("intermediate is an object"); + if i == path.len() - 1 { + let _ = map.insert(segment.clone(), value); + return; + } + let entry = map + .entry(segment.clone()) + .or_insert_with(|| Value::Object(Map::new())); + if !entry.is_object() { + *entry = Value::Object(Map::new()); + } + current = entry; + } +} + +/// Removes the value at a dot-path (`#-`). +fn remove_path(value: &mut Value, path: &[String]) { + let Some((last, parents)) = path.split_last() else { + return; + }; + let mut current = value; + for segment in parents { + current = match current { + Value::Object(map) => match map.get_mut(segment) { + Some(child) => child, + None => return, + }, + Value::Array(array) => { + match segment.parse::().ok().and_then(|i| array.get_mut(i)) { + Some(child) => child, + None => return, + } + } + _ => return, + }; + } + match current { + Value::Object(map) => { + let _ = map.remove(last); + } + Value::Array(array) => { + if let Ok(i) = last.parse::() + && i < array.len() + { + let _ = array.remove(i); + } + } + _ => {} + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + fn item() -> Value { + json!({ + "id": "i", "collection": "c", "type": "Feature", + "geometry": {"type": "Point", "coordinates": [1, 2]}, + "assets": {"data": {"href": "x"}}, + "properties": {"datetime": "2020-01-01T00:00:00Z", "eo:cloud_cover": 5, "platform": "l8"} + }) + } + + #[test] + fn exclude_removes_paths() { + let out = apply_fields( + item(), + &json!({"exclude": ["assets", "properties.eo:cloud_cover"]}), + ); + assert!(out.get("assets").is_none()); + assert!(out["properties"].get("eo:cloud_cover").is_none()); + assert_eq!(out["properties"]["platform"], "l8"); + } + + #[test] + fn include_keeps_only_paths_plus_id_collection() { + let out = apply_fields(item(), &json!({"include": ["properties.datetime"]})); + let keys: Vec<&String> = out.as_object().unwrap().keys().collect(); + assert!(keys.contains(&&"id".to_string())); + assert!(keys.contains(&&"collection".to_string())); + assert!(keys.contains(&&"properties".to_string())); + assert!(out.get("geometry").is_none()); + assert_eq!(out["properties"]["datetime"], "2020-01-01T00:00:00Z"); + assert!(out["properties"].get("platform").is_none()); + } + + #[test] + fn empty_fields_pass_through() { + assert_eq!(apply_fields(item(), &json!({})), item()); + } +} diff --git a/src/pgstac-rs/src/read/hydrate.rs b/src/pgstac-rs/src/read/hydrate.rs new file mode 100644 index 00000000..cf46e885 --- /dev/null +++ b/src/pgstac-rs/src/read/hydrate.rs @@ -0,0 +1,686 @@ +//! Rust-side hydration of dehydrated pgstac item rows back into full STAC items. +//! +//! pgstac stores items in a "dehydrated" form: shared content is factored out into +//! a per-collection `base_item` (pgstac 0.9.11) or shared `item_fragments` +//! (pgstac 0.10), and a fixed set of queryable properties is promoted to dedicated +//! columns. [content_hydrate] reassembles the full self-contained STAC item. +//! +//! This module reproduces that reassembly client-side so a dump can emit fully +//! hydrated items without a per-row server round trip. It targets the **export +//! path only**, where the requested fields are always the full item (`fields = +//! {}`), so the include/exclude projection that `content_hydrate(_, fields)` +//! applies is the identity and is intentionally omitted here. +//! +//! Two storage models are supported: +//! +//! * [`HydrationModel::BaseItem`] — pgstac 0.9.11. The dehydrated row carries the +//! item `content` jsonb; hydration deep-merges it over the collection +//! `base_item` using the same precedence and deletion-sentinel rules as the SQL +//! `merge_jsonb`. +//! * [`HydrationModel::Fragment`] — pgstac 0.10. Promoted columns are folded back +//! into `properties`, optional shared fragment content is recursively merged in +//! (item wins), and per-item link hrefs are spliced into the fragment's link +//! template. +//! +//! The output must equal the SQL `content_hydrate` over the same rows; this is +//! gated by parity tests in `tests/hydrate_parity.rs`. +//! +//! [content_hydrate]: https://github.com/stac-utils/pgstac/blob/main/src/pgstac/sql/003a_items.sql + +use crate::rawjson::RawJson; +use serde_json::{Map, Value, json}; + +/// Which pgstac storage model a dehydrated row came from. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum HydrationModel { + /// pgstac 0.9.11: per-collection `base_item` deep-merge. + BaseItem, + /// pgstac 0.10: promoted columns + shared `item_fragments`. + Fragment, +} + +/// Per-collection context needed to hydrate an item. +/// +/// For [`HydrationModel::BaseItem`] only `base_item` is used. For +/// [`HydrationModel::Fragment`] the fragment (looked up by the row's +/// `fragment_id`) is supplied per row, so this carries nothing model-specific. +#[derive(Debug, Default, Clone)] +pub struct CollectionContext { + /// The collection's `base_item` jsonb (0.9.11). `None`/`null` on 0.10. + pub base_item: Option, +} + +/// The shared fragment content for a 0.10 item, looked up by `fragment_id`. +#[derive(Debug, Default, Clone)] +pub struct FragmentContext { + /// `item_fragments.content` — shared assets/properties/etc. + pub content: Option, + /// `item_fragments.links_template` — link list with hrefs templated out. + pub links_template: Option, +} + +/// A dehydrated item row plus everything needed to hydrate it. +/// +/// All jsonb-typed columns are kept as [`serde_json::Value`] for byte-fidelity +/// with the SQL path. `geometry` is the already-computed `ST_AsGeoJSON(geom, +/// 20)::jsonb` value (the export query computes it server-side so Rust never +/// re-encodes WKB and cannot drift on coordinate precision). +#[derive(Debug, Default, Clone)] +pub struct DehydratedItem { + // --- always present identity / geometry --- + /// `items.id`. + pub id: String, + /// `items.collection`. + pub collection: String, + /// `ST_AsGeoJSON(geometry, 20)::jsonb`, or `None` if the geometry is null. + pub geometry: Option, + + // --- 0.9.11 base_item model --- + /// `items.content` jsonb (0.9.11 only). + pub content: Option, + + // --- 0.10 fragment model: split columns --- + /// `items.bbox`, kept as raw JSON so the byte path can emit it verbatim (preserving PG numeric + /// precision) and the Value path parses it on demand. + pub bbox: Option, + /// `items.links` (per-item link list when there is no fragment). + pub links: Option, + /// `items.assets`, kept as raw JSON so the byte path can merge it with the fragment without + /// deep-parsing the large per-asset arrays; the Value path parses it on demand. + pub assets: Option, + /// `items.properties` (non-promoted, non-fragment properties). + pub properties: Option, + /// `items.extra` (top-level keys outside the known set). + pub extra: Option, + /// `items.stac_version`. + pub stac_version: Option, + /// `items.stac_extensions`. + pub stac_extensions: Option, + /// `items.link_hrefs` — per-item hrefs spliced into a fragment template. + pub link_hrefs: Option>>, + /// `items.fragment_id` — set when the item shares a fragment. + pub fragment_id: Option, + /// Promoted queryable columns (0.10), in `promoted_properties_from_item` order. + pub promoted: PromotedProperties, +} + +/// The promoted queryable columns, in the exact STAC-name order that +/// `promoted_properties_from_item` rebuilds them. +/// +/// Each non-temporal column value is already rendered as a JSON-ready +/// [`serde_json::Value`] (timestamps as STAC text, arrays as arrays, etc.). A +/// `Value::Null` (or absent) column is dropped, matching the SQL's +/// `jsonb_strip_nulls`. +/// +/// The temporal keys are handled separately because the range case emits an +/// **explicit** `"datetime": null` which `jsonb_strip_nulls` must *not* remove +/// (it is produced by `jsonb_build_object`, outside the stripped block — see +/// `temporal_properties_from_item`). +#[derive(Debug, Default, Clone)] +pub struct PromotedProperties { + /// The temporal block (`temporal_properties_from_item`): + /// + /// * instant item: `{ "datetime": "" }` + /// * range item: `{ "datetime": null, "start_datetime": "", + /// "end_datetime": "" }` (start/end stripped if null). + /// + /// Built by the source so the precise tstz-text formatting matches the SQL. + pub temporal: Map, + /// Ordered `(stac_name, value)` pairs for non-temporal promoted columns with + /// **static** names (used by unit tests / hardcoded callers). The order is + /// load-bearing for byte-parity. + pub fields: Vec<(&'static str, Value)>, + /// Ordered `(stac_name, value)` pairs for non-temporal promoted columns with + /// **owned** names (the runtime-derived path; see `PromotedSchema`). Emitted + /// after [`fields`](Self::fields). The order is load-bearing for byte-parity. + pub owned_fields: Vec<(String, Value)>, +} + +impl PromotedProperties { + /// Builds the promoted-properties object: the temporal block first (with its + /// explicit `datetime` key preserved), then the non-temporal columns with + /// null values dropped (`jsonb_strip_nulls`). + fn to_object(&self) -> Map { + let mut out = self.temporal.clone(); + for (name, value) in &self.fields { + if value.is_null() { + continue; + } + let _ = out.insert((*name).to_string(), value.clone()); + } + for (name, value) in &self.owned_fields { + if value.is_null() { + continue; + } + let _ = out.insert(name.clone(), value.clone()); + } + out + } +} + +/// Hydrates dehydrated rows back into full STAC items. +#[derive(Debug, Clone, Copy)] +pub struct Hydrator { + model: HydrationModel, +} + +impl Hydrator { + /// Creates a hydrator for the given storage model. + pub fn new(model: HydrationModel) -> Self { + Hydrator { model } + } + + /// The storage model this hydrator targets. + pub fn model(&self) -> HydrationModel { + self.model + } + + /// Hydrates one dehydrated row into a full STAC item object. + /// + /// `fragment` is the shared fragment content for the row's `fragment_id` + /// (0.10 only); pass `None` for 0.9.11 or for items without a fragment. + pub fn hydrate( + &self, + item: DehydratedItem, + ctx: &CollectionContext, + fragment: Option<&FragmentContext>, + ) -> Value { + match self.model { + HydrationModel::BaseItem => hydrate_base_item(item, ctx), + HydrationModel::Fragment => hydrate_fragment(item, fragment), + } + } +} + +// --------------------------------------------------------------------------- +// 0.9.11 base_item model +// --------------------------------------------------------------------------- + +/// The pgstac 0.9.11 deletion sentinel (`"𒍟※"`): a base_item key with this value +/// in the per-item content marks the key as removed in the hydrated output. +const DELETION_SENTINEL: &str = "𒍟※"; + +/// 0.9.11: `merge_jsonb(item_jsonb || content, base_item)` where +/// `item_jsonb = {id, geometry, collection, type:Feature}`. +fn hydrate_base_item(item: DehydratedItem, ctx: &CollectionContext) -> Value { + let DehydratedItem { + id, + collection, + geometry, + content, + .. + } = item; + let mut item_obj = Map::new(); + let _ = item_obj.insert("id".to_string(), Value::String(id)); + let _ = item_obj.insert("geometry".to_string(), geometry.unwrap_or(Value::Null)); + let _ = item_obj.insert("collection".to_string(), Value::String(collection)); + let _ = item_obj.insert("type".to_string(), Value::String("Feature".to_string())); + + // `{id,geometry,collection,type} || content` — jsonb concat, content wins. + if let Some(Value::Object(content)) = content { + for (k, v) in content { + let _ = item_obj.insert(k, v); + } + } + + let base = ctx.base_item.clone().unwrap_or(Value::Null); + merge_jsonb(Value::Object(item_obj), base) +} + +/// Port of the SQL `merge_jsonb(_a, _b)`: `_a` is the per-item value, `_b` the +/// base_item value. `_a` takes precedence; objects deep-merge; equal-length +/// arrays merge element-wise; the deletion sentinel in `_a` removes the key. +fn merge_jsonb(a: Value, b: Value) -> Value { + // WHEN _a = sentinel THEN NULL + if let Value::String(s) = &a + && s == DELETION_SENTINEL + { + return Value::Null; + } + // WHEN _a IS NULL OR jsonb_typeof(_a) = 'null' THEN _b + if a.is_null() { + return b; + } + match (&a, &b) { + (Value::Object(ao), Value::Object(bo)) => { + // jsonb_strip_nulls(jsonb_object_agg(key, merge_jsonb(a.value, b.value))) + // over a FULL JOIN on key. + let mut keys: Vec = Vec::new(); + let mut seen = std::collections::HashSet::new(); + for k in ao.keys().chain(bo.keys()) { + if seen.insert(k.clone()) { + keys.push(k.clone()); + } + } + let mut out = Map::new(); + for k in keys { + let av = ao.get(&k).cloned(); + let bv = bo.get(&k).cloned(); + let merged = match (av, bv) { + (Some(av), Some(bv)) => merge_jsonb(av, bv), + // FULL JOIN: a missing side is SQL NULL. + (Some(av), None) => merge_jsonb(av, Value::Null), + (None, Some(bv)) => merge_jsonb(Value::Null, bv), + (None, None) => Value::Null, + }; + // jsonb_strip_nulls drops null *values* at every object level. + if !merged.is_null() { + let _ = out.insert(k, merged); + } + } + Value::Object(out) + } + (Value::Array(aa), Value::Array(ba)) if aa.len() == ba.len() => { + // Element-wise merge of equal-length arrays. + let merged: Vec = aa + .iter() + .zip(ba.iter()) + .map(|(x, y)| merge_jsonb(x.clone(), y.clone())) + .collect(); + Value::Array(merged) + } + // ELSE _a + _ => a, + } +} + +// --------------------------------------------------------------------------- +// 0.10 fragment model +// --------------------------------------------------------------------------- + +/// 0.10: reassemble from split columns + promoted columns + optional fragment. +/// +/// Consumes `item` so its (potentially large) split-column values move into the output instead of +/// being cloned per row; only the shared `fragment` parts are cloned. +fn hydrate_fragment(item: DehydratedItem, fragment: Option<&FragmentContext>) -> Value { + let (mut output, assets_raw, bbox_raw) = hydrate_fragment_core(item, fragment); + + // bbox (parsed for the Value path). + if let Some(bbox) = bbox_raw.map(|r| r.to_value()) + && !bbox.is_null() + { + let _ = output.insert("bbox".to_string(), bbox); + } + + // merged_assets := jsonb_merge_recursive(frag.assets, COALESCE(item.assets, {})) + let frag_assets = fragment + .and_then(|f| f.content.as_ref()) + .and_then(|c| c.get("assets")) + .cloned(); + let item_assets = assets_raw + .map(|r| r.to_value()) + .unwrap_or_else(|| json!({})); + let merged_assets = + jsonb_merge_recursive(frag_assets, Some(item_assets)).unwrap_or_else(|| json!({})); + if merged_assets != json!({}) { + let _ = output.insert("assets".to_string(), merged_assets); + } + + Value::Object(output) +} + +/// The fragment-model item assembled for **every key except `assets` and `bbox`**, returned alongside +/// the item's raw `assets`/`bbox`. +/// +/// The Value path ([`hydrate_fragment`]) parses + merges those into the map; the byte path +/// ([`crate::feature`]) emits the raw `bbox` verbatim and merges `assets` at serialize time. Object key +/// order is irrelevant to STAC semantics (and to the [`Value`] parity check), so the two callers add +/// `assets`/`bbox` wherever is convenient. +pub(crate) fn hydrate_fragment_core( + item: DehydratedItem, + fragment: Option<&FragmentContext>, +) -> (Map, Option, Option) { + let DehydratedItem { + id, + collection, + geometry, + bbox, + links, + assets, + properties, + extra, + stac_version, + stac_extensions, + link_hrefs, + fragment_id, + promoted, + .. + } = item; + + let frag_content = fragment.and_then(|f| f.content.as_ref()); + let frag_links_template = fragment.and_then(|f| f.links_template.as_ref()); + let has_fragment = fragment_id.is_some() && fragment.is_some(); + + // merged_properties := jsonb_merge_recursive(frag.properties, COALESCE(item.properties, {})) + let frag_properties = frag_content.and_then(|c| c.get("properties")).cloned(); + let item_properties = properties.unwrap_or_else(|| json!({})); + let merged_properties = jsonb_merge_recursive(frag_properties, Some(item_properties)); + // merged_properties := promoted_properties_from_item || COALESCE(merged, {}) + let merged_properties = concat_objects( + Value::Object(promoted.to_object()), + merged_properties.unwrap_or_else(|| json!({})), + ); + + // hydrated_stac_version := COALESCE(item.stac_version, frag.stac_version) + let hydrated_stac_version = stac_version.or_else(|| { + frag_content + .and_then(|c| c.get("stac_version")) + .and_then(|v| v.as_str()) + .map(str::to_string) + }); + + // hydrated_stac_extensions: + // item.stac_extensions when non-null and != [] + // else COALESCE(frag.stac_extensions, item.stac_extensions) + let hydrated_stac_extensions = match stac_extensions { + Some(e) if !e.is_null() && !is_empty_array(&e) => Some(e), + other => frag_content + .and_then(|c| c.get("stac_extensions")) + .cloned() + .or(other), + }; + + // links + let hydrated_links = if has_fragment { + stac_links_hydrate(frag_links_template, link_hrefs.as_deref()) + } else { + links.unwrap_or_else(|| json!([])) + }; + + let mut output = Map::new(); + let _ = output.insert("id".to_string(), Value::String(id)); + let _ = output.insert("geometry".to_string(), geometry.unwrap_or(Value::Null)); + let _ = output.insert("collection".to_string(), Value::String(collection)); + let _ = output.insert("type".to_string(), Value::String("Feature".to_string())); + if let Some(v) = hydrated_stac_version { + let _ = output.insert("stac_version".to_string(), Value::String(v)); + } + if let Some(exts) = hydrated_stac_extensions + && !exts.is_null() + && !is_empty_array(&exts) + { + let _ = output.insert("stac_extensions".to_string(), exts); + } + // hydrated_links IS NOT NULL is always true here (we coalesce to []). + let _ = output.insert("links".to_string(), hydrated_links); + // merged_properties IS NOT NULL is always true (we coalesce). + let _ = output.insert("properties".to_string(), merged_properties); + + // output := output || item.extra + if let Some(Value::Object(extra)) = extra { + for (k, v) in extra { + let _ = output.insert(k, v); + } + } + + (output, assets, bbox) +} + +/// `_a || _b` jsonb concatenation for two objects (right side wins on key +/// collision). Used for `promoted_properties_from_item || merged_properties`. +fn concat_objects(a: Value, b: Value) -> Value { + match (a, b) { + (Value::Object(mut ao), Value::Object(bo)) => { + for (k, v) in bo { + let _ = ao.insert(k, v); + } + Value::Object(ao) + } + // promoted is always an object; merged_properties coalesced to {}. + (_, b) => b, + } +} + +/// Port of `jsonb_merge_recursive(frag, item)`: deep merge with *item* +/// precedence. Object keys merge recursively; for non-object collisions the item +/// value wins. Disjoint object key sets short-circuit to a shallow concat +/// (`frag || item`), which the SQL does for performance and is byte-equivalent. +fn jsonb_merge_recursive(frag: Option, item: Option) -> Option { + match (frag, item) { + // WHEN frag IS NULL THEN COALESCE(item, {}) + (None, item) => Some(item.unwrap_or_else(|| json!({}))), + (Some(frag), item) => { + let item = match item { + // WHEN item IS NULL OR item = {} THEN frag + None => return Some(frag), + Some(i) if i == json!({}) => return Some(frag), + Some(i) => i, + }; + match (&frag, &item) { + (Value::Object(fo), Value::Object(io)) => { + // FULL JOIN on key; per-key: + // i null -> f ; f null -> i + // both objects -> disjoint? f||i : recurse + // else -> i + let mut keys: Vec = Vec::new(); + let mut seen = std::collections::HashSet::new(); + for k in fo.keys().chain(io.keys()) { + if seen.insert(k.clone()) { + keys.push(k.clone()); + } + } + let mut out = Map::new(); + for k in keys { + let fv = fo.get(&k); + let iv = io.get(&k); + let merged = match (fv, iv) { + (Some(fv), None) => fv.clone(), + (None, Some(iv)) => iv.clone(), + (Some(fv), Some(iv)) => { + if let (Value::Object(fvo), Value::Object(ivo)) = (fv, iv) { + let disjoint = !fvo.keys().any(|fk| ivo.contains_key(fk)); + if disjoint { + // f.value || i.value + concat_objects(fv.clone(), iv.clone()) + } else { + jsonb_merge_recursive(Some(fv.clone()), Some(iv.clone())) + .unwrap_or_else(|| json!({})) + } + } else { + iv.clone() + } + } + (None, None) => continue, + }; + let _ = out.insert(k, merged); + } + Some(Value::Object(out)) + } + // ELSE item + _ => Some(item), + } + } + } +} + +/// Port of `stac_links_hydrate(links_template, link_hrefs)`: splice per-item +/// hrefs back into the fragment link template by ordinal position. +fn stac_links_hydrate(template: Option<&Value>, hrefs: Option<&[Option]>) -> Value { + let template = match template { + Some(Value::Array(a)) => a, + // links_template null or not an array -> [] + _ => return json!([]), + }; + + // No hrefs -> return the template as-is. + let hrefs = match hrefs { + Some(h) if !h.is_empty() => h, + _ => return Value::Array(template.clone()), + }; + + let out: Vec = template + .iter() + .enumerate() + .map(|(idx, link)| match link { + Value::Object(obj) => { + // ordinality is 1-based in SQL; link_hrefs[ord]. + let href = hrefs.get(idx).and_then(|h| h.as_ref()); + let mut new_obj = obj.clone(); + let _ = new_obj.remove("href"); + if let Some(href) = href { + let _ = new_obj.insert("href".to_string(), Value::String(href.clone())); + } + Value::Object(new_obj) + } + other => other.clone(), + }) + .collect(); + Value::Array(out) +} + +fn is_empty_array(v: &Value) -> bool { + matches!(v, Value::Array(a) if a.is_empty()) +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + fn temporal_instant(dt: &str) -> Map { + let mut m = Map::new(); + let _ = m.insert("datetime".to_string(), Value::String(dt.to_string())); + m + } + + #[test] + fn merge_jsonb_item_wins() { + let a = json!({"a": 1, "b": 2}); + let b = json!({"b": 99, "c": 3}); + assert_eq!( + merge_jsonb(a, b), + json!({"a": 1, "b": 2, "c": 3}), + "item value wins on collision; base-only keys included" + ); + } + + #[test] + fn merge_jsonb_deletion_sentinel() { + let a = json!({"a": DELETION_SENTINEL, "b": 2}); + let b = json!({"a": 1, "b": 99}); + // a -> sentinel removes key (merge to null, then stripped); b kept from a. + assert_eq!(merge_jsonb(a, b), json!({"b": 2})); + } + + #[test] + fn merge_jsonb_strip_nulls() { + // base-only key whose value is null is stripped. + let a = json!({"a": 1}); + let b = json!({"a": 1, "n": null}); + assert_eq!(merge_jsonb(a, b), json!({"a": 1})); + } + + #[test] + fn merge_jsonb_equal_arrays() { + let a = json!([{"x": 1}, {"y": 2}]); + let b = json!([{"x": 0, "z": 9}, {"y": 0, "w": 8}]); + assert_eq!( + merge_jsonb(a, b), + json!([{"x": 1, "z": 9}, {"y": 2, "w": 8}]) + ); + } + + #[test] + fn merge_jsonb_unequal_arrays_a_wins() { + let a = json!([1, 2]); + let b = json!([3, 4, 5]); + assert_eq!(merge_jsonb(a, b), json!([1, 2])); + } + + #[test] + fn base_item_basic() { + let h = Hydrator::new(HydrationModel::BaseItem); + let item = DehydratedItem { + id: "i1".into(), + collection: "c1".into(), + geometry: Some(json!({"type": "Point", "coordinates": [1.0, 2.0]})), + content: Some(json!({ + "properties": {"datetime": "2020-01-01T00:00:00Z", "eo:cloud_cover": 5}, + "assets": {"data": {"href": "s3://x"}} + })), + ..Default::default() + }; + let ctx = CollectionContext { + base_item: Some(json!({ + "stac_version": "1.0.0", + "properties": {"platform": "sentinel-2"}, + "assets": {"data": {"type": "image/tiff"}} + })), + }; + let out = h.hydrate(item, &ctx, None); + assert_eq!(out["id"], "i1"); + assert_eq!(out["type"], "Feature"); + assert_eq!(out["stac_version"], "1.0.0"); + assert_eq!(out["properties"]["platform"], "sentinel-2"); + assert_eq!(out["properties"]["eo:cloud_cover"], 5); + assert_eq!(out["assets"]["data"]["href"], "s3://x"); + assert_eq!(out["assets"]["data"]["type"], "image/tiff"); + } + + #[test] + fn fragment_promoted_properties() { + let h = Hydrator::new(HydrationModel::Fragment); + let item = DehydratedItem { + id: "i1".into(), + collection: "c1".into(), + geometry: Some(json!({"type": "Point", "coordinates": [1.0, 2.0]})), + promoted: PromotedProperties { + temporal: temporal_instant("2020-01-01T00:00:00Z"), + fields: vec![ + ("platform", json!("landsat-8")), + ("eo:cloud_cover", json!(12.5)), + ("nullcol", Value::Null), + ], + owned_fields: vec![], + }, + ..Default::default() + }; + let out = h.hydrate(item, &CollectionContext::default(), None); + assert_eq!(out["properties"]["datetime"], "2020-01-01T00:00:00Z"); + assert_eq!(out["properties"]["platform"], "landsat-8"); + assert_eq!(out["properties"]["eo:cloud_cover"], 12.5); + assert!(out["properties"].get("nullcol").is_none()); + // links always present (coalesced to []). + assert_eq!(out["links"], json!([])); + } + + #[test] + fn fragment_merge_and_links() { + let h = Hydrator::new(HydrationModel::Fragment); + let item = DehydratedItem { + id: "i1".into(), + collection: "c1".into(), + geometry: Some(json!({"type": "Point", "coordinates": [0.0, 0.0]})), + fragment_id: Some(7), + link_hrefs: Some(vec![Some("https://a/self".into()), None]), + properties: Some(json!({"eo:cloud_cover": 1})), + promoted: PromotedProperties { + temporal: temporal_instant("2020-01-01T00:00:00Z"), + fields: vec![], + owned_fields: vec![], + }, + ..Default::default() + }; + let frag = FragmentContext { + content: Some(json!({ + "properties": {"platform": "x"}, + "assets": {"thumb": {"href": "t"}} + })), + links_template: Some(json!([ + {"rel": "self", "type": "application/json"}, + {"rel": "root", "href": "https://static/root"} + ])), + }; + let out = h.hydrate(item, &CollectionContext::default(), Some(&frag)); + assert_eq!(out["properties"]["platform"], "x"); + assert_eq!(out["properties"]["eo:cloud_cover"], 1); + assert_eq!(out["properties"]["datetime"], "2020-01-01T00:00:00Z"); + assert_eq!(out["assets"]["thumb"]["href"], "t"); + // link 0: href spliced from link_hrefs[0] + assert_eq!(out["links"][0]["href"], "https://a/self"); + assert_eq!(out["links"][0]["rel"], "self"); + // link 1: href None -> href removed from template + assert_eq!(out["links"][1]["rel"], "root"); + assert!(out["links"][1].get("href").is_none()); + } +} diff --git a/src/pgstac-rs/src/read/keyset.rs b/src/pgstac-rs/src/read/keyset.rs new file mode 100644 index 00000000..d8334d77 --- /dev/null +++ b/src/pgstac-rs/src/read/keyset.rs @@ -0,0 +1,155 @@ +//! Minting keyset pagination tokens in Rust, mirroring SQL `keyset_encode` / `keyset_sortkeys`. +//! +//! A token is the base64 of the boundary row's sort-key values, joined by `chr(31)` with NULL encoded +//! as `chr(30)`. The values come from the **hydrated feature** (the sort keys are item fields), so no +//! extra server round trip is needed. On the next request the server `keyset_decode`s the token and +//! `keyset_where` casts each value back to its column type, so the rendering only has to round-trip — +//! it does not have to be byte-identical to what SQL `search()` would emit. + +use base64::Engine; +use serde_json::{Value, json}; +use std::collections::HashSet; + +/// The unit separator (`chr(31)`) joining keyset values. +const UNIT_SEP: &str = "\u{1f}"; +/// The record separator (`chr(30)`) standing in for a NULL value. +const NULL_SENTINEL: &str = "\u{1e}"; + +/// The ordered sort-key fields for a search: `sortby` fields, then appended `id` and `collection`, +/// deduped keeping the first occurrence. Mirrors `keyset_sortkeys` (direction is irrelevant to the +/// encoded value, so only the field order matters here). +/// +/// # Examples +/// +/// ``` +/// use pgstac::keyset::sort_key_fields; +/// use serde_json::json; +/// +/// // Default sort is by datetime, then the id/collection tiebreaks are appended. +/// assert_eq!(sort_key_fields(&json!({})), ["datetime", "id", "collection"]); +/// // An explicit id sort dedupes the appended id. +/// assert_eq!( +/// sort_key_fields(&json!({"sortby": [{"field": "id", "direction": "asc"}]})), +/// ["id", "collection"] +/// ); +/// ``` +pub fn sort_key_fields(search: &Value) -> Vec { + let default = json!([{"field": "datetime", "direction": "desc"}]); + let sortby = search + .get("sortby") + .filter(|v| v.as_array().is_some_and(|a| !a.is_empty())) + .unwrap_or(&default); + + let mut fields: Vec = sortby + .as_array() + .into_iter() + .flatten() + .filter_map(|e| e.get("field").and_then(Value::as_str).map(String::from)) + .collect(); + fields.push("id".to_string()); + fields.push("collection".to_string()); + + let mut seen = HashSet::new(); + fields.retain(|f| seen.insert(f.clone())); + fields +} + +/// The keyset value text for one sort-key field, pulled from a hydrated feature. +fn keyset_value(feature: &Value, field: &str) -> Option { + let props = &feature["properties"]; + match field { + "id" => feature["id"].as_str().map(String::from), + "collection" => feature["collection"].as_str().map(String::from), + // The keyset uses the `datetime` column, which is `start_datetime` for range items (whose + // `properties.datetime` is null). + "datetime" => props["datetime"] + .as_str() + .or_else(|| props["start_datetime"].as_str()) + .map(String::from), + "end_datetime" => props["end_datetime"] + .as_str() + .or_else(|| props["datetime"].as_str()) + .map(String::from), + other => json_text(&props[other]), + } +} + +/// Renders a JSON value to the text Postgres `->>` would produce (for property sort keys). +fn json_text(value: &Value) -> Option { + match value { + Value::Null => None, + Value::String(s) => Some(s.clone()), + Value::Bool(b) => Some(b.to_string()), + Value::Number(n) => Some(n.to_string()), + other => Some(other.to_string()), + } +} + +/// Encodes keyset values into a token, matching SQL `keyset_encode`. +/// +/// # Examples +/// +/// ``` +/// use pgstac::keyset::keyset_encode; +/// +/// // base64("a" + chr(31) + "b") +/// assert_eq!(keyset_encode(&[Some("a".into()), Some("b".into())]), "YR9i"); +/// // A NULL becomes chr(30): base64("a" + chr(31) + chr(30)) +/// assert_eq!(keyset_encode(&[Some("a".into()), None]), "YR8e"); +/// ``` +pub fn keyset_encode(values: &[Option]) -> String { + let joined = values + .iter() + .map(|v| v.as_deref().unwrap_or(NULL_SENTINEL)) + .collect::>() + .join(UNIT_SEP); + base64::engine::general_purpose::STANDARD.encode(joined.as_bytes()) +} + +/// Mints the keyset token for a hydrated feature given the search's ordered sort-key `fields`. +pub fn mint_token(feature: &Value, fields: &[String]) -> String { + let values: Vec> = fields.iter().map(|f| keyset_value(feature, f)).collect(); + keyset_encode(&values) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn encode_matches_sql_keyset_encode() { + assert_eq!(keyset_encode(&[Some("a".into()), Some("b".into())]), "YR9i"); + assert_eq!(keyset_encode(&[Some("a".into()), None]), "YR8e"); + } + + #[test] + fn mint_uses_datetime_then_id_collection() { + let feature = json!({ + "id": "item-1", + "collection": "coll-a", + "properties": {"datetime": "2013-04-19T00:00:00Z"} + }); + let fields = sort_key_fields(&json!({})); + let token = mint_token(&feature, &fields); + // datetime + chr(31) + id + chr(31) + collection + let expected = keyset_encode(&[ + Some("2013-04-19T00:00:00Z".into()), + Some("item-1".into()), + Some("coll-a".into()), + ]); + assert_eq!(token, expected); + } + + #[test] + fn range_item_uses_start_datetime_for_keyset() { + let feature = json!({ + "id": "r1", + "collection": "c", + "properties": {"datetime": null, "start_datetime": "2020-01-01T00:00:00Z", "end_datetime": "2020-01-02T00:00:00Z"} + }); + assert_eq!( + keyset_value(&feature, "datetime").as_deref(), + Some("2020-01-01T00:00:00Z") + ); + } +} diff --git a/src/pgstac-rs/src/read/mod.rs b/src/pgstac-rs/src/read/mod.rs new file mode 100644 index 00000000..69938d7d --- /dev/null +++ b/src/pgstac-rs/src/read/mod.rs @@ -0,0 +1,12 @@ +//! The client-side read path: keyset search, flat-memory streaming, hydration, projection, and pages. + +pub mod collections; +pub mod feature; +pub mod fields; +pub mod hydrate; +pub mod keyset; +pub(crate) mod page; +pub mod search; +pub mod source; +#[cfg(feature = "pool")] +pub(crate) mod stream; diff --git a/src/pgstac-rs/src/read/page.rs b/src/pgstac-rs/src/read/page.rs new file mode 100644 index 00000000..3a3ff859 --- /dev/null +++ b/src/pgstac-rs/src/read/page.rs @@ -0,0 +1,42 @@ +//! Adapts the keyset engine's [`SearchPage`] into the rustac-native [`stac::api::ItemCollection`]. + +use crate::Error; +use crate::search::SearchPage; +use serde_json::{Map, Value}; +use stac::api::{Context, Item, ItemCollection}; + +/// Adapts the engine's [`SearchPage`] into a [`stac::api::ItemCollection`]: the hydrated feature values +/// parse into [`Item`]s, the `next:`/`prev:`-prefixed keyset tokens become `{"token": …}` pagination +/// maps, and the match count (when the search counted one) surfaces as both `number_matched` and a +/// [`Context`]. This is the single page-to-ItemCollection path the `stac::api` client-trait impls route +/// through. +impl TryFrom for ItemCollection { + type Error = Error; + + fn try_from(page: SearchPage) -> Result { + let items = page + .features + .into_iter() + .map(serde_json::from_value) + .collect::, _>>()?; + let token_map = |token: String| { + let mut map = Map::new(); + let _ = map.insert("token".into(), Value::String(token)); + map + }; + let context = page.number_matched.map(|matched| Context { + returned: page.number_returned as u64, + limit: None, + matched: Some(matched as u64), + additional_fields: Map::new(), + }); + let number_matched = context.as_ref().and_then(|context| context.matched); + let mut item_collection = ItemCollection::new(items)?; + item_collection.number_matched = number_matched; + item_collection.number_returned = Some(page.number_returned as u64); + item_collection.context = context; + item_collection.next = page.next_token.map(token_map); + item_collection.prev = page.prev_token.map(token_map); + Ok(item_collection) + } +} diff --git a/src/pgstac-rs/src/read/search.rs b/src/pgstac-rs/src/read/search.rs new file mode 100644 index 00000000..3c491ec0 --- /dev/null +++ b/src/pgstac-rs/src/read/search.rs @@ -0,0 +1,451 @@ +//! Client-side search driven by `search_plan`. +//! +//! `search_plan` returns a ready-to-prepare SELECT (projection, `WHERE`, keyset seek and `ORDER BY` +//! all baked in), the per-band datetime histogram, the lead direction, and the context count/query. +//! This module does the rest in Rust: +//! +//! 1. call `search_plan` **once** and `prepare()` its query, so stepping multiple histogram bands +//! reuses one plan (no replanning); +//! 2. step the bands, binding `$1`/`$2` (band low/high) and `$3` (limit) per band, until the page is +//! full — or, for a non-datetime-leading sort, run the single `$1`=limit query; +//! 3. read each row's **raw** columns (EWKB geometry, raw `timestamptz`) and **hydrate in Rust** with +//! the [`Hydrator`](crate::hydrate::Hydrator); +//! 4. **mint the keyset tokens in Rust** ([`crate::keyset`]) from the boundary features — no extra +//! round trip. Tokens are validated by page-equivalence (the server casts each value back). + +use crate::Result; +use crate::hydrate::{HydrationModel, Hydrator}; +use crate::keyset; +use crate::source::{ + PromotedSchema, detect_hydration_model, fetch_fragments, load_collection_context, + read_base_item_row, read_fragment_row, +}; +use chrono::{DateTime, TimeZone, Utc}; +use serde_json::{Value, json}; +use std::collections::HashMap; +use tokio_postgres::types::ToSql; +use tokio_postgres::{GenericClient, Row}; + +/// One page of a search: hydrated features plus the keyset pagination tokens and context. +#[derive(Debug, Clone, Default)] +pub struct SearchPage { + /// The hydrated STAC items. + pub features: Vec, + /// The `next:`-prefixed continuation token, if there is a next page. + pub next_token: Option, + /// The `prev:`-prefixed continuation token, if there is a previous page. + pub prev_token: Option, + /// The number of features returned in this page. + pub number_returned: usize, + /// The total number of matches, when context counting is enabled. + pub number_matched: Option, +} + +/// The pieces of `search_plan` the client needs to step + tokenize a page. +pub(crate) struct Plan { + /// The ready-to-prepare SELECT. + pub(crate) query: String, + /// The per-month `[{m, n}]` histogram, when the sort is datetime-leading. + pub(crate) histogram: Option, + /// Whether the effective lead direction is descending (already accounts for `prev`). + pub(crate) lead_desc: bool, + /// Whether the sort leads with datetime (so the query is `$1,$2,$3` band-stepped). + pub(crate) datetime_leading: bool, + /// The inlined context count, when fresh. + pub(crate) context_count: Option, + /// The query to run for the context count on a cache miss. + pub(crate) ctx_query: Option, + /// The search carries a property filter (`filter`/`query`), so its matches can be sparse across + /// datetime bands. [`step_bands`] then scans the whole range in one query (like SQL `search()`) + /// instead of geometric-doubling through sparse leading bands (many round trips). + pub(crate) selective: bool, +} + +/// A high sentinel for the most recent band's exclusive upper bound (no items exist past the last +/// histogram month, so any value above them works). +fn far_future() -> DateTime { + Utc.with_ymd_and_hms(9999, 12, 31, 23, 59, 59).unwrap() +} + +/// Fetches the plan for a search in a single round trip. +pub(crate) async fn fetch_plan( + client: &C, + search: &Value, + token: Option<&str>, + limit: i64, +) -> Result { + let limit = i32::try_from(limit)?; + let row = client + .query_one( + "SELECT query, histogram, lead_desc, datetime_leading, context_count, ctx_query \ + FROM search_plan($1::jsonb, $2, $3)", + &[search, &token, &limit], + ) + .await?; + Ok(Plan { + query: row.try_get("query")?, + histogram: row.try_get("histogram")?, + lead_desc: row.try_get("lead_desc")?, + datetime_leading: row.try_get("datetime_leading")?, + context_count: row.try_get("context_count")?, + ctx_query: row.try_get("ctx_query")?, + selective: search.get("filter").is_some() || search.get("query").is_some(), + }) +} + +/// Parses the histogram into contiguous `(low, high)` datetime bands (ascending), one per month, that +/// together cover the whole histogram range up to [`far_future`]. +/// +/// Every month is kept — bands must NOT be dropped by count. The histogram counts are an estimate +/// (`partition_bounds` prorates a partition's row count uniformly across its dtrange months), so a sparse +/// collection — few items over a long span — rounds every month to 0 even though the data is there. +/// Dropping zero-count months would leave datetime gaps and silently lose those rows (SQL `search()` +/// scans the whole range regardless). [`step_bands`] sizes its scan windows by geometric doubling, not by +/// these counts, so keeping the empty months costs only a few extra (cheap) empty-window probes. +pub(crate) fn histogram_bands(histogram: &Value) -> Result, DateTime)>> { + let entries: Vec<(DateTime, i64)> = histogram + .as_array() + .map(Vec::as_slice) + .unwrap_or(&[]) + .iter() + .map(|e| { + let m = e["m"] + .as_str() + .and_then(|s| DateTime::parse_from_rfc3339(s).ok()) + .map(|dt| dt.with_timezone(&Utc)); + let n = e["n"].as_i64().unwrap_or(0); + (m, n) + }) + .filter_map(|(m, n)| m.map(|m| (m, n))) + .collect(); + + let mut bands = Vec::new(); + for i in 0..entries.len() { + let (low, _n) = entries[i]; + let high = entries.get(i + 1).map(|e| e.0).unwrap_or_else(far_future); + bands.push((low, high)); + } + Ok(bands) +} + +/// The ordered band ranges to walk for a plan: the datetime histogram bands in lead order, or a single +/// `None` band for a non-datetime sort (the query has no band parameters). Used by the streaming +/// iterator (the buffered [`step_bands`] inlines an equivalent walk with its geometric windowing). +#[cfg(feature = "pool")] +#[allow(clippy::type_complexity)] +pub(crate) fn band_ranges(plan: &Plan) -> Result, DateTime)>>> { + if !plan.datetime_leading { + return Ok(vec![None]); + } + let mut bands = match &plan.histogram { + Some(h) => histogram_bands(h)?, + None => Vec::new(), + }; + if plan.lead_desc { + bands.reverse(); + } + Ok(bands.into_iter().map(Some).collect()) +} + +/// Runs the prepared query across the datetime histogram bands until `want` rows are collected. +async fn step_bands(client: &C, plan: &Plan, want: i64) -> Result> { + let stmt = client.prepare(&plan.query).await?; + + if !plan.datetime_leading { + // Non-datetime sort: a single `$1`=limit query covers the whole range. + return Ok(client.query(&stmt, &[&want]).await?); + } + + let mut bands = match &plan.histogram { + Some(h) => histogram_bands(h)?, + None => Vec::new(), + }; + if plan.lead_desc { + bands.reverse(); + } + + // Geometric-window band walk (non-lossy): query the + // next `window` contiguous bands as ONE statement with `LIMIT = remaining`, doubling the window each + // step. A window that returns `< remaining` rows is exhausted (advance past it); one that fills + // `remaining` ends the walk. Starting at a single band keeps time-to-first-row low and avoids the + // wide-`Append` planning tax of scanning every partition up front; the doubling keeps a deep page (or + // a token that has consumed the leading bands) down to ~log(bands) round trips instead of one per band. + let mut rows: Vec = Vec::new(); + let mut remaining = want; + let mut i = 0usize; + // Unfiltered datetime sorts fill from the first band(s), so start narrow (low time-to-first-row) and + // double. A property filter can be sparse across bands, so scan the whole range in one query (like + // SQL `search()`) rather than walking sparse leading bands one round trip at a time. + let mut window = if plan.selective { + bands.len().max(1) + } else { + 1 + }; + while i < bands.len() && remaining > 0 { + let end = (i + window).min(bands.len()); + // The window spans contiguous bands, so its datetime range is [min low, max high). + let low = bands[i..end] + .iter() + .map(|b| b.0) + .min() + .expect("non-empty window"); + let high = bands[i..end] + .iter() + .map(|b| b.1) + .max() + .expect("non-empty window"); + let params: [&(dyn ToSql + Sync); 3] = [&low, &high, &remaining]; + let band_rows = client.query(&stmt, ¶ms).await?; + let n = band_rows.len() as i64; + rows.extend(band_rows); + if n >= remaining { + break; // filled the page from this window + } + remaining -= n; + i = end; + window *= 2; + } + Ok(rows) +} + +/// The total match count (`numberMatched`), or `None` when context counting is off. +async fn resolve_context(client: &C, plan: &Plan) -> Result> { + if let Some(count) = plan.context_count { + return Ok(Some(count)); + } + let Some(ctx_query) = &plan.ctx_query else { + return Ok(None); + }; + let row = client.query_one(ctx_query, &[]).await?; + Ok(row.try_get::<_, Option>(0)?) +} + +/// Runs one page of a search: prepares `search_plan`'s query, steps the histogram bands, hydrates the +/// rows in Rust, and mints the pagination tokens in Rust. +/// +/// Detects the hydration model (and loads the promoted schema) on every call; callers that already +/// know these — e.g. [`PgstacPool`](crate::PgstacPool), which caches them — use +/// [`search_page_with`] to skip those catalog round trips. +pub async fn search_page( + client: &C, + search: &Value, + token: Option<&str>, + limit: i64, +) -> Result { + let model = detect_hydration_model(client).await?; + let promoted = match model { + HydrationModel::Fragment => Some(PromotedSchema::load(client).await?), + HydrationModel::BaseItem => None, + }; + search_page_with(client, model, promoted.as_ref(), search, token, limit).await +} + +/// Like [`search_page`] but with the hydration model + promoted schema supplied (so they aren't +/// re-detected per call). +pub async fn search_page_with( + client: &C, + model: HydrationModel, + promoted: Option<&PromotedSchema>, + search: &Value, + token: Option<&str>, + limit: i64, +) -> Result { + let is_prev = token.is_some_and(|t| t.starts_with("prev:")); + let plan = fetch_plan(client, search, token, limit).await?; + + // Fetch one extra row across the bands to detect whether a further page exists. + let rows = step_bands(client, &plan, limit + 1).await?; + let mut items = match (&model, &promoted) { + (HydrationModel::Fragment, Some(schema)) => rows + .iter() + .map(|row| read_fragment_row(row, schema)) + .collect::>>()?, + _ => rows + .iter() + .map(read_base_item_row) + .collect::>>()?, + }; + + let has_more = items.len() as i64 > limit; + items.truncate(limit as usize); + + // When paginating backwards the rows came back in reverse order; flip them to forward order. + if is_prev { + items.reverse(); + } + + // Hydrate: load each collection's context once and the referenced fragments. + let mut contexts: HashMap = HashMap::new(); + for item in &items { + if !contexts.contains_key(&item.collection) { + let ctx = load_collection_context(client, model, &item.collection).await?; + let _ = contexts.insert(item.collection.clone(), ctx); + } + } + let frag_ids: Vec = items.iter().filter_map(|i| i.fragment_id).collect(); + let fragments = fetch_fragments(client, &frag_ids).await?; + let hydrator = Hydrator::new(model); + let features: Vec = items + .into_iter() + .map(|item| { + let ctx = &contexts[&item.collection]; + let fragment = item.fragment_id.and_then(|id| fragments.get(&id)); + hydrator.hydrate(item, ctx, fragment) + }) + .collect(); + + // Mint tokens in Rust from the boundary features' sort-key values. + let fields = keyset::sort_key_fields(search); + let (mut next_token, mut prev_token) = (None, None); + let next_present = if is_prev { token.is_some() } else { has_more }; + let prev_present = if is_prev { has_more } else { token.is_some() }; + if let Some(last) = features.last() + && next_present + { + next_token = Some(format!("next:{}", keyset::mint_token(last, &fields))); + } + if let Some(first) = features.first() + && prev_present + { + prev_token = Some(format!("prev:{}", keyset::mint_token(first, &fields))); + } + + let number_matched = resolve_context(client, &plan).await?; + + // Apply the STAC `fields` projection last, so token minting above still saw the full sort-key + // values even when `fields` excludes them. + let features = project_fields(features, search); + + Ok(SearchPage { + number_returned: features.len(), + features, + next_token, + prev_token, + number_matched, + }) +} + +/// Applies the search's STAC `fields` include/exclude projection to each feature (a no-op when the +/// search has no `fields`). +pub(crate) fn project_fields(features: Vec, search: &Value) -> Vec { + match search.get("fields") { + Some(fields) if fields.is_object() => features + .into_iter() + .map(|feature| crate::fields::apply_fields(feature, fields)) + .collect(), + _ => features, + } +} + +/// The page size used to paginate `search_collect` / `search_stream` when the search body has no +/// explicit `limit`. +const DEFAULT_BATCH: i64 = 1000; + +/// The per-page batch size for a search, bounded by `max_items` so the last fetch doesn't overshoot. +fn batch_size(search: &Value, collected: usize, max_items: Option) -> i64 { + let base = search + .get("limit") + .and_then(Value::as_i64) + .filter(|&n| n > 0) + .unwrap_or(DEFAULT_BATCH); + match max_items { + Some(max) => { + let remaining = max.saturating_sub(collected) as i64; + remaining.min(base) + } + None => base, + } +} + +/// Collects every matching item (up to `max_items`) into one page by following the keyset `next` +/// tokens. The returned [`SearchPage`] carries the full set of features (and the match count); it has +/// no continuation tokens. +pub async fn search_collect( + client: &C, + search: &Value, + max_items: Option, +) -> Result { + let mut features: Vec = Vec::new(); + let mut token: Option = None; + let mut number_matched = None; + loop { + let limit = batch_size(search, features.len(), max_items); + if limit <= 0 { + break; + } + let page = search_page(client, search, token.as_deref(), limit).await?; + if number_matched.is_none() { + number_matched = page.number_matched; + } + let next = page.next_token; + let empty = page.features.is_empty(); + let take = match max_items { + Some(max) => max - features.len(), + None => usize::MAX, + }; + features.extend(page.features.into_iter().take(take)); + match next { + Some(next) if !empty => token = Some(next), + _ => break, + } + } + Ok(SearchPage { + number_returned: features.len(), + features, + next_token: None, + prev_token: None, + number_matched, + }) +} + +/// Streams every matching item (up to `max_items`) as newline-delimited JSON to `write`, following +/// the keyset `next` tokens. Memory stays flat: one page is held at a time. Returns the number of +/// items written. +pub async fn search_stream( + client: &C, + search: &Value, + write: &mut W, + max_items: Option, +) -> Result +where + C: GenericClient, + W: std::io::Write, +{ + let mut written = 0usize; + let mut token: Option = None; + loop { + let limit = batch_size(search, written, max_items); + if limit <= 0 { + break; + } + let page = search_page(client, search, token.as_deref(), limit).await?; + let next = page.next_token; + for feature in &page.features { + serde_json::to_writer(&mut *write, feature)?; + write.write_all(b"\n")?; + written += 1; + if max_items.is_some_and(|max| written >= max) { + return Ok(written); + } + } + match next { + Some(next) if !page.features.is_empty() => token = Some(next), + _ => break, + } + } + Ok(written) +} + +/// Fetches a single item by id from a collection. +/// +/// A getter is just a one-row search with the collection and id pinned, so it shares the same +/// hydration path as [`search_page`]. +pub async fn get_item( + client: &C, + collection_id: &str, + item_id: &str, +) -> Result> { + let body = json!({"collections": [collection_id], "ids": [item_id], "limit": 1}); + let page = search_page(client, &body, None, 1).await?; + Ok(page.features.into_iter().next()) +} diff --git a/src/pgstac-rs/src/read/source.rs b/src/pgstac-rs/src/read/source.rs new file mode 100644 index 00000000..4484f54f --- /dev/null +++ b/src/pgstac-rs/src/read/source.rs @@ -0,0 +1,720 @@ +//! Row sources: fetch dehydrated item rows from a partition (or arbitrary +//! `WHERE`) and turn them into [`DehydratedItem`]s ready for the [`Hydrator`]. +//! +//! Columns are read **raw/binary** and converted in Rust: `geometry` arrives as +//! EWKB (the column is selected unwrapped) and is turned into GeoJSON via +//! [`Ewkb`](crate::geom::Ewkb); `timestamptz` columns are rendered +//! to STAC text with [`tstz_to_stac_text`](crate::temporal::tstz_to_stac_text). +//! Nothing is pre-formatted server-side. +//! +//! [`Hydrator`]: crate::hydrate::Hydrator + +use crate::Result; +use crate::geom::Ewkb; +use crate::hydrate::HydrationModel; +use crate::hydrate::{ + CollectionContext, DehydratedItem, FragmentContext, Hydrator, PromotedProperties, +}; +use crate::temporal::tstz_to_stac_text; +use chrono::{DateTime, Utc}; +use futures::StreamExt; +use futures::pin_mut; +use serde_json::{Map, Value}; +use std::collections::HashMap; +use tokio_postgres::GenericClient; +use tokio_postgres::types::ToSql; + +/// Reads the raw `geometry` column (EWKB) into a GeoJSON [`Value`], or `None` when null. +fn read_geometry(row: &tokio_postgres::Row, col: &str) -> Result> { + row.try_get::<_, Option>(col)? + .map(|g| g.to_geojson()) + .transpose() +} + +/// How a promoted column maps to a JSON value, derived from its Postgres type. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum PromotedKind { + Text, + Float, + Int, + BigInt, + TextArray, + Jsonb, + /// `timestamptz` rendered to STAC text by the source query. + TstzText, +} + +impl PromotedKind { + /// Maps a Postgres `udt_name`/`data_type` to a kind. + fn from_pg_type(udt_name: &str) -> PromotedKind { + match udt_name { + "text" | "varchar" | "bpchar" => PromotedKind::Text, + "float8" | "float4" | "numeric" => PromotedKind::Float, + "int4" | "int2" => PromotedKind::Int, + "int8" => PromotedKind::BigInt, + "_text" | "_varchar" => PromotedKind::TextArray, + "jsonb" | "json" => PromotedKind::Jsonb, + "timestamptz" | "timestamp" => PromotedKind::TstzText, + // Unknown -> jsonb is the safest passthrough (PostGIS geom etc. don't + // appear in the promoted set). + _ => PromotedKind::Jsonb, + } + } +} + +/// One promoted column: its STAC property name, SQL column, and value kind. +#[derive(Debug, Clone)] +struct PromotedColumn { + stac_name: String, + column: String, + kind: PromotedKind, +} + +/// The promoted-column schema for an instance, derived at runtime from +/// `promoted_item_property_defs()` (STAC name + ordered column list) joined with +/// the items table column types. +/// +/// Deriving this at runtime — rather than hardcoding — keeps the source portable +/// across pgstac schema revisions, which differ in their promoted-column set +/// (e.g. `eo:bands` vs `bands`, `proj:epsg` vs `proj:code`). +#[derive(Debug, Clone)] +pub struct PromotedSchema { + columns: Vec, +} + +/// Hydration invariants for a pgstac database: the storage model + the fragment-model promoted schema. +/// These don't change between calls, so a stateful client/pool detects them once and caches them, +/// saving the catalog round trips (`detect_hydration_model` + `PromotedSchema::load`) on every search. +#[derive(Debug, Clone)] +pub(crate) struct CachedHydration { + pub(crate) model: HydrationModel, + pub(crate) schema: Option>, +} + +impl CachedHydration { + /// Detects the storage model and loads the promoted schema (fragment model only). + pub(crate) async fn detect(client: &C) -> Result { + let model = detect_hydration_model(client).await?; + let schema = match model { + HydrationModel::Fragment => { + Some(std::sync::Arc::new(PromotedSchema::load(client).await?)) + } + HydrationModel::BaseItem => None, + }; + Ok(CachedHydration { model, schema }) + } +} + +impl PromotedSchema { + /// Loads the promoted-column schema from the live instance. + /// + /// Uses `promoted_item_property_defs()` for the ordered `(stac_name, column)` + /// mapping (the order is load-bearing for byte-parity) and + /// `information_schema.columns` for each column's type. + pub async fn load(client: &C) -> Result { + // Column -> udt_name for the items table. + let type_rows = client + .query( + "SELECT column_name, udt_name FROM information_schema.columns \ + WHERE table_schema = 'pgstac' AND table_name = 'items'", + &[], + ) + .await?; + let mut types = HashMap::new(); + for row in &type_rows { + let name: String = row.get("column_name"); + let udt: String = row.get("udt_name"); + let _ = types.insert(name, udt); + } + + // Ordered (stac_name, column) from the metadata function. Its row type is + // (name, definition, property_path) where property_path is the column. + let def_rows = client + .query( + "SELECT name, property_path FROM promoted_item_property_defs() \ + WITH ORDINALITY AS d(name, definition, property_path, ord) ORDER BY ord", + &[], + ) + .await?; + + let mut columns = Vec::with_capacity(def_rows.len()); + for row in &def_rows { + let stac_name: String = row.get("name"); + let column: String = row.get("property_path"); + let udt = types.get(&column).map(String::as_str).unwrap_or("jsonb"); + columns.push(PromotedColumn { + stac_name, + kind: PromotedKind::from_pg_type(udt), + column, + }); + } + Ok(PromotedSchema { columns }) + } + + /// SELECT expressions for the promoted columns — the raw column names. The + /// `timestamptz` columns are rendered to STAC text in Rust, not the server. + fn select_exprs(&self) -> Vec { + self.columns.iter().map(|c| c.column.clone()).collect() + } +} + +/// Builds the `SELECT` projection for the 0.10 (fragment) model. +/// +/// Columns are raw: `geometry` is the unwrapped EWKB column and `datetime`/`end_datetime` are raw +/// `timestamptz` (converted in Rust). This is used by the export/partition-scan path; the search path +/// runs `search_plan`'s own (equivalent, fields-aware) projection. +pub(crate) fn fragment_select_columns(schema: &PromotedSchema) -> Vec { + let mut cols = vec![ + "id".to_string(), + "collection".to_string(), + "geometry".to_string(), + "datetime_is_range".to_string(), + "datetime".to_string(), + "end_datetime".to_string(), + "bbox".to_string(), + "links".to_string(), + "assets".to_string(), + "properties".to_string(), + "extra".to_string(), + "stac_version".to_string(), + "stac_extensions".to_string(), + "link_hrefs".to_string(), + "fragment_id".to_string(), + ]; + cols.extend(schema.select_exprs()); + cols +} + +/// Reads one [`tokio_postgres::Row`] (fragment model) into a [`DehydratedItem`]. +/// +/// Reads the raw columns and does the conversions in Rust: EWKB→GeoJSON geometry and +/// `timestamptz`→STAC text. +pub(crate) fn read_fragment_row( + row: &tokio_postgres::Row, + schema: &PromotedSchema, +) -> Result { + let datetime_is_range: bool = row.try_get("datetime_is_range")?; + let dt_text: Option = row + .try_get::<_, Option>>("datetime")? + .map(tstz_to_stac_text); + let edt_text: Option = row + .try_get::<_, Option>>("end_datetime")? + .map(tstz_to_stac_text); + + let mut promoted = PromotedProperties { + temporal: build_temporal(datetime_is_range, dt_text, edt_text), + fields: Vec::new(), + owned_fields: Vec::with_capacity(schema.columns.len()), + }; + for col in &schema.columns { + let value = read_promoted(row, &col.column, col.kind)?; + promoted.owned_fields.push((col.stac_name.clone(), value)); + } + + Ok(DehydratedItem { + id: row.try_get("id")?, + collection: row.try_get("collection")?, + geometry: read_geometry(row, "geometry")?, + content: None, + bbox: row.try_get("bbox")?, + links: row.try_get("links")?, + assets: row.try_get("assets")?, + properties: row.try_get("properties")?, + extra: row.try_get("extra")?, + stac_version: row.try_get("stac_version")?, + stac_extensions: row.try_get("stac_extensions")?, + link_hrefs: row.try_get("link_hrefs")?, + fragment_id: row.try_get("fragment_id")?, + promoted, + }) +} + +/// Reads a single promoted column value into a JSON-ready [`Value`]. +fn read_promoted(row: &tokio_postgres::Row, col: &str, kind: PromotedKind) -> Result { + let value = match kind { + PromotedKind::Text => row + .try_get::<_, Option>(col)? + .map(Value::String) + .unwrap_or(Value::Null), + PromotedKind::TstzText => row + .try_get::<_, Option>>(col)? + .map(|dt| Value::String(tstz_to_stac_text(dt))) + .unwrap_or(Value::Null), + PromotedKind::Float => row + .try_get::<_, Option>(col)? + .map(float8_to_json) + .unwrap_or(Value::Null), + PromotedKind::Int => row + .try_get::<_, Option>(col)? + .map(|v| Value::Number(v.into())) + .unwrap_or(Value::Null), + PromotedKind::BigInt => row + .try_get::<_, Option>(col)? + .map(|v| Value::Number(v.into())) + .unwrap_or(Value::Null), + PromotedKind::TextArray => match row.try_get::<_, Option>>(col)? { + Some(items) => Value::Array(items.into_iter().map(Value::String).collect()), + None => Value::Null, + }, + PromotedKind::Jsonb => row.try_get::<_, Option>(col)?.unwrap_or(Value::Null), + }; + Ok(value) +} + +/// Converts a `float8` to a JSON number matching Postgres `to_jsonb(float8)`. +/// +/// Postgres jsonb stores numbers as `numeric`, so an **integral** float renders +/// without a decimal point (`10.0` -> `10`); non-integral floats keep their +/// fractional part. Mirroring that keeps the promoted float columns byte-equal +/// to the SQL `content_hydrate` output (whose values come through jsonb). +fn float8_to_json(v: f64) -> Value { + if v.is_finite() && v.fract() == 0.0 && v.abs() < (i64::MAX as f64) { + // Integral and i64-representable -> integer, matching jsonb. + Value::Number((v as i64).into()) + } else { + serde_json::Number::from_f64(v) + .map(Value::Number) + .unwrap_or(Value::Null) + } +} + +/// Builds the temporal block matching `temporal_properties_from_item`. +fn build_temporal( + datetime_is_range: bool, + dt_text: Option, + edt_text: Option, +) -> Map { + let mut m = Map::new(); + if datetime_is_range { + // jsonb_build_object('datetime', NULL) || strip_nulls({start, end}) + let _ = m.insert("datetime".to_string(), Value::Null); + if let Some(dt) = dt_text { + let _ = m.insert("start_datetime".to_string(), Value::String(dt)); + } + if let Some(edt) = edt_text { + let _ = m.insert("end_datetime".to_string(), Value::String(edt)); + } + } else { + // jsonb_build_object('datetime', tstz_to_stac_text(datetime)) + let _ = m.insert( + "datetime".to_string(), + dt_text.map(Value::String).unwrap_or(Value::Null), + ); + } + m +} + +/// Reads one [`tokio_postgres::Row`] (base_item model, 0.9.11) into a +/// [`DehydratedItem`]. +pub(crate) fn read_base_item_row(row: &tokio_postgres::Row) -> Result { + Ok(DehydratedItem { + id: row.try_get("id")?, + collection: row.try_get("collection")?, + geometry: read_geometry(row, "geometry")?, + content: row.try_get("content")?, + ..Default::default() + }) +} + +/// The `SELECT` projection for the 0.9.11 (base_item) model. Geometry is the raw EWKB column. +pub(crate) fn base_item_select_columns() -> Vec { + vec![ + "id".to_string(), + "collection".to_string(), + "geometry".to_string(), + "content".to_string(), + ] +} + +/// Fetches dehydrated rows for the given model from `from_clause` with an +/// optional `where_clause`, applying `params`. Used by parity tests and the +/// partition source. +/// +/// `from_clause` is a table or partition name (or subquery). `where_clause` is +/// the predicate body (without `WHERE`); pass `None` for all rows. `limit` +/// caps the number of rows (None = no limit). Results are ordered +/// `(datetime, id)` to match the dump contract. +pub async fn fetch_dehydrated( + client: &C, + model: HydrationModel, + from_clause: &str, + where_clause: Option<&str>, + params: &[&(dyn ToSql + Sync)], +) -> Result> { + fetch_dehydrated_limited(client, model, from_clause, where_clause, params, None).await +} + +/// Like [`fetch_dehydrated`] but with an optional row `limit`. +pub async fn fetch_dehydrated_limited( + client: &C, + model: HydrationModel, + from_clause: &str, + where_clause: Option<&str>, + params: &[&(dyn ToSql + Sync)], + limit: Option, +) -> Result> { + // The fragment model needs the runtime-derived promoted-column schema. + let promoted_schema = match model { + HydrationModel::Fragment => Some(PromotedSchema::load(client).await?), + HydrationModel::BaseItem => None, + }; + let columns = match model { + HydrationModel::BaseItem => base_item_select_columns(), + HydrationModel::Fragment => fragment_select_columns( + promoted_schema + .as_ref() + .expect("fragment model loads a promoted schema"), + ), + }; + let where_sql = match where_clause { + Some(w) => format!("WHERE {w}"), + None => String::new(), + }; + let limit_sql = match limit { + Some(n) => format!("LIMIT {n}"), + None => String::new(), + }; + let query = format!( + "SELECT {} FROM {} {} ORDER BY datetime, id {}", + columns.join(", "), + from_clause, + where_sql, + limit_sql, + ); + let rows = client.query(&query, params).await?; + match model { + HydrationModel::BaseItem => rows.iter().map(read_base_item_row).collect(), + HydrationModel::Fragment => { + let schema = promoted_schema + .as_ref() + .expect("fragment model loads a promoted schema"); + rows.iter() + .map(|row| read_fragment_row(row, schema)) + .collect() + } + } +} + +/// Loads a single fragment by id (0.10) — used by the streaming iterator's per-new-id parallel fetch. +pub async fn fetch_fragment( + client: &C, + id: i64, +) -> Result> { + let rows = client + .query( + "SELECT content, links_template FROM item_fragments WHERE id = $1", + &[&id], + ) + .await?; + match rows.first() { + Some(row) => Ok(Some(FragmentContext { + content: row.try_get("content")?, + links_template: row.try_get("links_template")?, + })), + None => Ok(None), + } +} + +/// Loads the shared fragment content for a set of fragment ids (0.10). +pub async fn fetch_fragments( + client: &C, + ids: &[i64], +) -> Result> { + let mut map = HashMap::new(); + if ids.is_empty() { + return Ok(map); + } + let rows = client + .query( + "SELECT id, content, links_template FROM item_fragments WHERE id = ANY($1)", + &[&ids], + ) + .await?; + for row in rows { + let id: i64 = row.try_get("id")?; + let _ = map.insert( + id, + FragmentContext { + content: row.try_get("content")?, + links_template: row.try_get("links_template")?, + }, + ); + } + Ok(map) +} + +/// Loads **all** of a collection's fragments into a cache keyed by fragment id. Fragments are the +/// deduplicated shared part of item content, so a collection has only a handful regardless of item +/// count — cheap to load whole. Used to pre-populate the fragment cache before a streaming partition +/// scan so hydration never issues a query while the scan's portal is open (see [`scan_partition`]). +pub async fn fetch_collection_fragments( + client: &C, + collection: &str, +) -> Result> { + let mut map = HashMap::new(); + let rows = client + .query( + "SELECT id, content, links_template FROM item_fragments WHERE collection = $1", + &[&collection], + ) + .await?; + for row in rows { + let id: i64 = row.try_get("id")?; + let _ = map.insert( + id, + FragmentContext { + content: row.try_get("content")?, + links_template: row.try_get("links_template")?, + }, + ); + } + Ok(map) +} + +/// A per-item prefilter applied during a partition scan. +#[derive(Debug, Clone, Default)] +pub struct ScanFilter { + /// Inclusive-start, exclusive-end datetime window (STAC text). Applied as + /// `datetime >= start AND datetime < end` (datetime = pgstac `datetime` + /// column = start_datetime for ranges, matching the dump ordering key). + pub datetime: Option<(String, String)>, + /// Spatial bbox `[w, s, e, n]`; applied as `ST_Intersects(geometry, env)`. + pub bbox: Option<[f64; 4]>, +} + +/// Builds the `WHERE` body and bound params for a [`ScanFilter`]. +/// +/// Returns `(where_body, owned_param_values)`. Bbox is inlined as numeric +/// literals (finite f64 -> safe), datetime as bound `$N` params. +fn scan_where(filter: &ScanFilter) -> (String, Vec) { + let mut clauses: Vec = Vec::new(); + let mut params: Vec = Vec::new(); + if let Some((lo, hi)) = &filter.datetime { + // Bind as text, then cast in SQL, so the param's Rust type (String) + // matches what the driver sends ($N::text::timestamptz). + params.push(lo.clone()); + clauses.push(format!("datetime >= ${}::text::timestamptz", params.len())); + params.push(hi.clone()); + clauses.push(format!("datetime < ${}::text::timestamptz", params.len())); + } + if let Some([w, s, e, n]) = &filter.bbox { + // Coordinates are finite f64; format as plain decimals (no injection). + clauses.push(format!( + "ST_Intersects(geometry, ST_MakeEnvelope({w}, {s}, {e}, {n}, 4326))" + )); + } + (clauses.join(" AND "), params) +} + +/// Streams a single partition's items, hydrated, in `(datetime, id)` order, +/// using a server-side cursor (`query_raw` portal) so server memory stays flat. +/// +/// `callback` is invoked once per batch of hydrated items (batch size chosen by +/// the driver's portal); returning `Err` aborts the scan. Fragments referenced +/// by the batch are looked up lazily and cached across batches. +/// +/// Works on both 0.9.11 (base_item) and 0.10 (fragment); the only server-side +/// state is the cursor, so it does not need the keyset search engine. +pub async fn scan_partition( + client: &C, + model: HydrationModel, + partition: &str, + collection: &str, + ctx: &CollectionContext, + filter: &ScanFilter, + mut callback: F, +) -> Result +where + C: GenericClient, + F: FnMut(Vec) -> Result<()>, +{ + let hydrator = Hydrator::new(model); + let promoted_schema = match model { + HydrationModel::Fragment => Some(PromotedSchema::load(client).await?), + HydrationModel::BaseItem => None, + }; + let columns = match model { + HydrationModel::BaseItem => base_item_select_columns(), + HydrationModel::Fragment => fragment_select_columns( + promoted_schema + .as_ref() + .expect("fragment model loads a promoted schema"), + ), + }; + let (where_body, where_params) = scan_where(filter); + let where_sql = if where_body.is_empty() { + String::new() + } else { + format!("WHERE {where_body}") + }; + let query = format!( + "SELECT {} FROM {} {} ORDER BY datetime, id", + columns.join(", "), + partition, + where_sql, + ); + + // Bind params as &dyn ToSql (coercion, not a cast — keeps trivial_casts happy). + let params: Vec<&(dyn ToSql + Sync)> = where_params + .iter() + .map(|p| { + let r: &(dyn ToSql + Sync) = p; + r + }) + .collect(); + + // Pre-load ALL of this collection's (deduplicated, few) fragments BEFORE opening the streaming + // portal below. Fetching fragments lazily inside the scan loop issued a query on this same + // connection while the portal's unread rows filled the connection buffer, deadlocking on any + // non-trivial partition (the server blocks sending the fragment result, the client blocks not + // draining the portal). Pre-loading keeps the portal scan query-free. `flush_batch` keeps a lazy + // fetch only as a fallback for the rare case of a fragment committed after this load. + let mut frag_cache: HashMap = match model { + HydrationModel::Fragment => fetch_collection_fragments(client, collection).await?, + HydrationModel::BaseItem => HashMap::new(), + }; + + // query_raw streams via a portal (flat server memory). + let stream = client.query_raw(&query, params).await?; + pin_mut!(stream); + + // Buffer rows into modest batches so the callback amortizes work while + // memory stays bounded. + const BATCH: usize = 1024; + let mut row_batch: Vec = Vec::with_capacity(BATCH); + let mut total: u64 = 0; + + while let Some(row) = stream.next().await { + row_batch.push(row?); + if row_batch.len() >= BATCH { + total += flush_batch( + client, + &hydrator, + model, + promoted_schema.as_ref(), + ctx, + &mut frag_cache, + &mut row_batch, + &mut callback, + ) + .await?; + } + } + if !row_batch.is_empty() { + total += flush_batch( + client, + &hydrator, + model, + promoted_schema.as_ref(), + ctx, + &mut frag_cache, + &mut row_batch, + &mut callback, + ) + .await?; + } + Ok(total) +} + +/// Hydrates a batch of raw rows and hands the items to the callback. +#[allow(clippy::too_many_arguments)] +async fn flush_batch( + client: &C, + hydrator: &Hydrator, + model: HydrationModel, + promoted_schema: Option<&PromotedSchema>, + ctx: &CollectionContext, + frag_cache: &mut HashMap, + row_batch: &mut Vec, + callback: &mut F, +) -> Result +where + C: GenericClient, + F: FnMut(Vec) -> Result<()>, +{ + // Parse rows into dehydrated items. + let mut items: Vec = Vec::with_capacity(row_batch.len()); + for row in row_batch.iter() { + let item = match model { + HydrationModel::BaseItem => read_base_item_row(row)?, + HydrationModel::Fragment => read_fragment_row( + row, + promoted_schema.expect("fragment model loads a promoted schema"), + )?, + }; + items.push(item); + } + row_batch.clear(); + + // Resolve any new fragment ids not yet cached. + if model == HydrationModel::Fragment { + let missing: Vec = items + .iter() + .filter_map(|i| i.fragment_id) + .filter(|id| !frag_cache.contains_key(id)) + .collect(); + if !missing.is_empty() { + let fetched = fetch_fragments(client, &missing).await?; + frag_cache.extend(fetched); + } + } + + let count = items.len() as u64; + let hydrated: Vec = items + .into_iter() + .map(|item| { + let fragment = item.fragment_id.and_then(|id| frag_cache.get(&id)); + hydrator.hydrate(item, ctx, fragment) + }) + .collect(); + callback(hydrated)?; + Ok(count) +} + +/// Detects which storage [`HydrationModel`] a database uses. +/// +/// A `pgstac.items.fragment_id` column means the 0.10 fragment model; its absence means the 0.9.11 +/// `base_item` model. +pub async fn detect_hydration_model(client: &C) -> Result { + let has_fragment_id: bool = client + .query_one( + "SELECT EXISTS (\ + SELECT 1 FROM information_schema.columns \ + WHERE table_schema = 'pgstac' \ + AND table_name = 'items' \ + AND column_name = 'fragment_id'\ + ) AS present", + &[], + ) + .await? + .get("present"); + if has_fragment_id { + Ok(HydrationModel::Fragment) + } else { + Ok(HydrationModel::BaseItem) + } +} + +/// Loads the per-collection hydration context (the `base_item` for the 0.9.11 model; empty for the +/// 0.10 fragment model, whose shared content lives in `item_fragments`). +pub async fn load_collection_context( + client: &C, + model: HydrationModel, + collection_id: &str, +) -> Result { + match model { + HydrationModel::BaseItem => { + let rows = client + .query( + "SELECT base_item FROM collections WHERE id = $1", + &[&collection_id], + ) + .await?; + let base_item = rows + .first() + .and_then(|row| row.get::<_, Option>("base_item")); + Ok(CollectionContext { base_item }) + } + HydrationModel::Fragment => Ok(CollectionContext::default()), + } +} diff --git a/src/pgstac-rs/src/read/stream.rs b/src/pgstac-rs/src/read/stream.rs new file mode 100644 index 00000000..0d818093 --- /dev/null +++ b/src/pgstac-rs/src/read/stream.rs @@ -0,0 +1,329 @@ +//! A true streaming search iterator with flat memory. +//! +//! Item rows stream from a server **portal** (`query_raw`) on one pooled connection — a whole page is +//! never buffered. As each row surfaces a `fragment_id` not yet seen, that **single** fragment is +//! fetched synchronously on a **second, parallel** pooled connection and cached. With high fragment +//! sharing (the split-storage design point) the shared fragment crosses the wire once and peak memory +//! stays at ~one row plus the small fragment cache, independent of result size. + +#[cfg(feature = "export")] +use crate::export::format::{Format, GeoparquetMode, GeoparquetStreamWriter, encode_all}; +#[cfg(feature = "export")] +use stac::geoparquet::Compression; +use crate::hydrate::{CollectionContext, FragmentContext, HydrationModel, Hydrator}; +use crate::search::{band_ranges, fetch_plan}; +use crate::source::{ + fetch_fragment, load_collection_context, read_base_item_row, read_fragment_row, +}; +use crate::{PgstacPool, Result}; +use async_stream::try_stream; +use futures::{Stream, StreamExt}; +use serde_json::Value; +use std::collections::HashMap; +use std::sync::Arc; +use tokio_postgres::GenericClient; +use tokio_postgres::types::ToSql; + +/// The `_limit` passed to `search_plan` for an unbounded stream — affects only the histogram/context +/// estimate, never the streamed result. +const PLAN_LIMIT_HINT: i64 = 10_000; + +/// Fetches `fid` once: returns the cached `Arc` if present, otherwise fetches it on `conn` (the +/// parallel connection) and caches it. The `Arc` is cloned out before any await so no borrow of the +/// cache is held across the fetch. +async fn cached_fragment( + cache: &mut HashMap>, + conn: &C, + fid: i64, +) -> Result>> { + if let Some(fragment) = cache.get(&fid) { + return Ok(Some(Arc::clone(fragment))); + } + let fetched = fetch_fragment(conn, fid).await?.map(Arc::new); + if let Some(fragment) = &fetched { + let _ = cache.insert(fid, Arc::clone(fragment)); + } + Ok(fetched) +} + +impl PgstacPool { + /// Streams hydrated STAC items one at a time with flat memory. + /// + /// `limit` caps the number of items (`None` = unbounded). The returned stream owns two pooled + /// connections for its lifetime: one streams the row portal, the other serves the per-new-fragment + /// fetches in parallel. + /// + /// # Examples + /// + /// ```no_run + /// use pgstac::{ConnectConfig, PgstacPool}; + /// use futures::StreamExt; + /// use serde_json::json; + /// + /// # tokio_test::block_on(async { + /// let pool = PgstacPool::connect(ConnectConfig::from_env()).await.unwrap(); + /// let mut items = Box::pin(pool.search_items(json!({}), None, Some(5))); + /// while let Some(item) = items.next().await { + /// println!("{}", item.unwrap()["id"]); + /// } + /// # }) + /// ``` + pub fn search_items( + &self, + search: Value, + token: Option, + limit: Option, + ) -> impl Stream> + 'static { + let pool = self.clone(); + try_stream! { + // Model + promoted schema come from the pool's shared cache (loaded once, not per stream). + let hydration = pool.cached_hydration().await?; + let model = hydration.model; + let schema = hydration.schema; + let hydrator = Hydrator::new(model); + + let data = pool.get().await?; + let frag_conn = pool.get().await?; + + let plan = fetch_plan(&**data, &search, token.as_deref(), limit.unwrap_or(PLAN_LIMIT_HINT)).await?; + let stmt = data.prepare(&plan.query).await?; + let bands = band_ranges(&plan)?; + let fields = search.get("fields").filter(|f| f.is_object()).cloned(); + + let mut contexts: HashMap = HashMap::new(); + let mut frag_cache: HashMap> = HashMap::new(); + let mut remaining = limit; + + 'bands: for band in bands { + if remaining == Some(0) { + break; + } + let lim = remaining.unwrap_or(i64::MAX); + let row_stream = match &band { + Some((low, high)) => { + let params: Vec<&(dyn ToSql + Sync)> = vec![low, high, &lim]; + data.query_raw(&stmt, params).await? + } + None => { + let params: Vec<&(dyn ToSql + Sync)> = vec![&lim]; + data.query_raw(&stmt, params).await? + } + }; + futures::pin_mut!(row_stream); + + while let Some(row) = row_stream.next().await { + let row = row?; + let item = match (&model, &schema) { + (HydrationModel::Fragment, Some(s)) => read_fragment_row(&row, s)?, + _ => read_base_item_row(&row)?, + }; + + // Per-collection context: fragment model is empty (no query); base_item queries on + // the parallel connection so the data portal is never interrupted. + if !contexts.contains_key(&item.collection) { + let ctx = load_collection_context(&**frag_conn, model, &item.collection).await?; + let _ = contexts.insert(item.collection.clone(), ctx); + } + + let fragment = match item.fragment_id { + None => None, + Some(fid) => cached_fragment(&mut frag_cache, &**frag_conn, fid).await?, + }; + + let ctx = &contexts[&item.collection]; + let feature = hydrator.hydrate(item, ctx, fragment.as_deref()); + let feature = match &fields { + Some(f) => crate::fields::apply_fields(feature, f), + None => feature, + }; + yield feature; + + if let Some(ref mut r) = remaining { + *r -= 1; + if *r == 0 { + break 'bands; + } + } + } + } + } + } + + /// Collects every matching item (up to `max_items`) by draining [`search_items`](Self::search_items). + pub async fn search_collect_items( + &self, + search: Value, + max_items: Option, + ) -> Result> { + let mut out = Vec::new(); + let stream = self.search_items(search, None, max_items); + futures::pin_mut!(stream); + while let Some(item) = stream.next().await { + out.push(item?); + } + Ok(out) + } + + /// Streams every matching item (up to `max_items`) into a stac-geoparquet file written to `sink`. + /// + /// On the 0.10 fragment model the schema is registry-complete, so items are written in row-group + /// batches as they stream — the item working set is one batch, never the whole result. The legacy + /// base-item model has no registry, so its schema must be widened over the full set; that case + /// buffers (the same split the dump makes). Returns the number of items written. + #[cfg(feature = "export")] + pub async fn stream_geoparquet( + &self, + search: Value, + max_items: Option, + compression: Compression, + row_group_size: Option, + sink: W, + ) -> Result { + const BATCH: usize = 1024; + let model = self.cached_hydration().await?.model; + match model { + HydrationModel::Fragment => { + let mut writer = GeoparquetStreamWriter::new(sink, compression); + if let Some(n) = row_group_size { + writer = writer.with_max_row_group_row_count(n); + } + let mut written = 0usize; + let mut batch: Vec = Vec::with_capacity(BATCH); + let stream = self.search_items(search, None, max_items); + futures::pin_mut!(stream); + while let Some(item) = stream.next().await { + batch.push(item?); + written += 1; + if batch.len() >= BATCH { + writer.write_batch(std::mem::take(&mut batch))?; + } + } + if !batch.is_empty() { + writer.write_batch(batch)?; + } + let _ = writer.finish()?; + Ok(written) + } + HydrationModel::BaseItem => { + // No registry to complete the schema, so widen over the full set: buffer + encode once. + let mut sink = sink; + let items = self.search_collect_items(search, max_items).await?; + let written = items.len(); + let bytes = encode_all( + Format::Geoparquet { + compression, + mode: GeoparquetMode::Buffered, + max_row_group_row_count: row_group_size, + }, + items, + )?; + sink.write_all(&bytes)?; + Ok(written) + } + } + } + + /// Streams every matching item (up to `max_items`) as newline-delimited JSON to `write`, holding + /// only one item at a time. Begins at `token` (`None` streams from the start). Returns the number + /// of items written. + /// + /// For the 0.10 fragment model this uses the byte path + /// ([`DehydratedItem::write_fragment_feature`](crate::hydrate::DehydratedItem::write_fragment_feature)): + /// each item is serialized + /// straight to `write` with the shared fragment merged at serialize time and `bbox` emitted from raw + /// bytes — no per-item [`Value`] is ever materialized. A `fields` projection runs on the value + /// path instead (it needs the structured feature). + pub async fn stream_ndjson( + &self, + search: Value, + token: Option<&str>, + max_items: Option, + write: &mut W, + ) -> Result { + let hydration = self.cached_hydration().await?; + let model = hydration.model; + let schema = hydration.schema; + let hydrator = Hydrator::new(model); + + let data = self.get().await?; + let frag_conn = self.get().await?; + let plan = fetch_plan( + &**data, + &search, + token, + max_items.unwrap_or(PLAN_LIMIT_HINT), + ) + .await?; + let stmt = data.prepare(&plan.query).await?; + let bands = band_ranges(&plan)?; + let fields = search.get("fields").filter(|f| f.is_object()).cloned(); + + let mut contexts: HashMap = HashMap::new(); + let mut frag_cache: HashMap> = HashMap::new(); + let mut written = 0usize; + let mut remaining = max_items; + + 'bands: for band in bands { + if remaining == Some(0) { + break; + } + let lim = remaining.unwrap_or(i64::MAX); + let row_stream = match &band { + Some((low, high)) => { + let params: Vec<&(dyn ToSql + Sync)> = vec![low, high, &lim]; + data.query_raw(&stmt, params).await? + } + None => { + let params: Vec<&(dyn ToSql + Sync)> = vec![&lim]; + data.query_raw(&stmt, params).await? + } + }; + futures::pin_mut!(row_stream); + + while let Some(row) = row_stream.next().await { + let row = row?; + let item = match (&model, &schema) { + (HydrationModel::Fragment, Some(s)) => read_fragment_row(&row, s)?, + _ => read_base_item_row(&row)?, + }; + + if !contexts.contains_key(&item.collection) { + let ctx = + load_collection_context(&**frag_conn, model, &item.collection).await?; + let _ = contexts.insert(item.collection.clone(), ctx); + } + let fragment = match item.fragment_id { + None => None, + Some(fid) => cached_fragment(&mut frag_cache, &**frag_conn, fid).await?, + }; + + match model { + // Byte path (no per-item Value materialized) — only when no field projection is + // requested; `fields` needs the structured feature, so fall back to the value path. + HydrationModel::Fragment if fields.is_none() => { + item.write_fragment_feature(fragment.as_deref(), &mut *write)?; + } + _ => { + let ctx = &contexts[&item.collection]; + let feature = hydrator.hydrate(item, ctx, fragment.as_deref()); + let feature = match &fields { + Some(f) => crate::fields::apply_fields(feature, f), + None => feature, + }; + serde_json::to_writer(&mut *write, &feature)?; + } + } + write.write_all(b"\n")?; + written += 1; + + if let Some(ref mut r) = remaining { + *r -= 1; + if *r == 0 { + break 'bands; + } + } + } + } + Ok(written) + } +} diff --git a/src/pgstac-rs/tests/canonical_parity.rs b/src/pgstac-rs/tests/canonical_parity.rs new file mode 100644 index 00000000..b8a0e8b1 --- /dev/null +++ b/src/pgstac-rs/tests/canonical_parity.rs @@ -0,0 +1,184 @@ +//! Parity gate: Rust `canonical::jsonb_canonical_hash` must equal SQL `pgstac.jsonb_canonical_hash` +//! byte-for-byte, so an item dehydrated in Rust gets the same `item_hash` as one ingested through the SQL +//! path. +//! +//! Stands up its own throwaway database from `src/pgstac/pgstac.sql` (dropped on `Drop`), so +//! `jsonb_canonical_hash` is guaranteed present regardless of any ambient test database. Override the +//! maintenance connection base (no database) via `PGSTAC_RS_TEST_BASE`. + +use pgstac::canonical; +use serde_json::{Value, json}; +use std::sync::atomic::{AtomicU32, Ordering}; +use tokio_postgres::{Client, NoTls}; + +/// The maintenance connection base (no database), used to create + drop the throwaway database. +fn base() -> String { + std::env::var("PGSTAC_RS_TEST_BASE") + .unwrap_or_else(|_| "postgresql://username:password@localhost:5439".to_string()) +} + +/// The assembled pgstac schema, loaded into the fresh database. +const PGSTAC_SQL: &str = include_str!("../../pgstac/pgstac.sql"); + +/// A throwaway database built from `pgstac.sql`, dropped on `Drop`. +struct FreshDb { + name: String, +} + +impl FreshDb { + async fn create() -> FreshDb { + static COUNTER: AtomicU32 = AtomicU32::new(0); + let name = format!( + "pgstac_rs_canonical_test_{}_{}", + std::process::id(), + COUNTER.fetch_add(1, Ordering::Relaxed) + ); + let (client, connection) = tokio_postgres::connect(&format!("{}/postgres", base()), NoTls) + .await + .unwrap(); + let handle = tokio::spawn(connection); + // CREATE DATABASE cannot run inside a transaction, so issue each statement on its own. + let _ = client + .execute(&format!("CREATE DATABASE {name}"), &[]) + .await + .unwrap(); + let _ = client + .execute( + &format!("ALTER DATABASE {name} SET search_path TO pgstac, public"), + &[], + ) + .await + .unwrap(); + handle.abort(); + FreshDb { name } + } + + fn dsn(&self) -> String { + format!("{}/{}", base(), self.name) + } +} + +impl Drop for FreshDb { + fn drop(&mut self) { + let name = self.name.clone(); + std::thread::scope(|scope| { + let _ = scope.spawn(|| { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + runtime.block_on(async move { + let (client, connection) = + tokio_postgres::connect(&format!("{}/postgres", base()), NoTls) + .await + .unwrap(); + let handle = tokio::spawn(connection); + let _ = client + .execute( + "SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname = $1", + &[&name], + ) + .await; + let _ = client + .execute(&format!("DROP DATABASE IF EXISTS {name}"), &[]) + .await; + handle.abort(); + }); + }); + }); + } +} + +/// Connects to the fresh database and loads the assembled `pgstac.sql` into it. +async fn connect_and_install(db: &FreshDb) -> Client { + let (client, connection) = tokio_postgres::connect(&db.dsn(), NoTls).await.unwrap(); + tokio::spawn(async move { + let _ = connection.await; + }); + client + .batch_execute("SET search_path TO pgstac, public;") + .await + .unwrap(); + client.batch_execute(PGSTAC_SQL).await.unwrap(); + client +} + +/// A battery of values exercising key ordering, nesting, unicode, escaping, and the number formatter +/// (fixed point, the scientific boundaries, integers, negative zero, large/small magnitudes). +fn cases() -> Vec { + vec![ + json!({"b": 1, "a": 2, "Z": 3, "aa": 4, "A": 5}), + json!([3, 1, 2, [4, 5], {"k": 6}]), + json!("hello \"world\"\n\t/path\\x"), + json!({"café": "naïve", "emoji": "🚀", "Δ": "δ"}), + // Datetime normalization (UTC, fixed 6-digit microseconds) — same instant, many spellings. + json!({"datetime": "2023-01-07T00:00:00Z"}), + json!({"datetime": "2023-01-07T00:00:00.000000Z"}), + json!({"datetime": "2023-01-07T12:34:56.789Z"}), + json!({"datetime": "2023-01-07T12:34:56.789000Z"}), + json!({"datetime": "2023-01-07T05:00:00+05:00"}), + json!({"datetime": "2023-01-07T00:00:00-00:00"}), + // Forgiving forms pgstac's to_tstz accepts (date-only, space, offset-less->UTC, lowercase, ±HHMM). + json!({"datetime": "2023-01-07"}), + json!({"datetime": "2023-01-07 00:00:00"}), + json!({"datetime": "2023-01-07T00:00:00"}), + json!({"datetime": "2023-01-07t00:00:00z"}), + json!({"datetime": "2023-01-07T05:00:00+0500"}), + json!({"start_datetime": "2019-12-31T19:00:00-05:00", "end_datetime": "2023-06-15T12:34:56.789Z"}), + // Datetime-shaped-but-not-a-datetime strings that must stay raw on BOTH sides (equal either way). + json!({"suffix": "2023-01-07T00:00:00Z-suffix", "bare": "2020", "compact": "20200101", "kw": "now"}), + json!({"nested": {"y": [1, {"x": true, "w": null}], "datetime": "2023-01-07T00:00:00Z"}}), + json!(42), + json!(42.0), + json!(-1), + json!(0), + json!(-0.0), + json!(0.1), + json!(0.5), + json!(100.0), + json!(1000000.0), + json!(123.456), + json!(-105.1019), + json!(40.1672), + json!(0.0001), + json!(0.00001), + json!(0.000001), + json!(1e15), + json!(1e16), + json!(1e20), + json!(1e21), + json!(1.5e-10), + json!(6.022e23), + json!(123456789012345_i64), + json!(9999999999999999_i64), + json!(2.5), + json!(98765.4321), + json!(0.30000000000000004), + json!({ + "type": "Feature", + "stac_version": "1.0.0", + "id": "item-1", + "geometry": {"type": "Point", "coordinates": [-105.1019, 40.1672]}, + "bbox": [-105.1019, 40.1672, -105.1019, 40.1672], + "properties": {"eo:cloud_cover": 12.5, "gsd": 30.0, "datetime": "2023-01-07T00:00:00Z"} + }), + ] +} + +#[tokio::test] +async fn canonical_and_hash_match_sql() { + let db = FreshDb::create().await; + let client = connect_and_install(&db).await; + for v in cases() { + let sql_hash: Vec = client + .query_one("SELECT pgstac.jsonb_canonical_hash($1::jsonb)", &[&v]) + .await + .unwrap() + .get(0); + assert_eq!( + canonical::jsonb_canonical_hash(&v).unwrap().to_vec(), + sql_hash, + "jsonb_canonical_hash mismatch for {v}" + ); + } +} diff --git a/src/pgstac-rs/tests/cli.rs b/src/pgstac-rs/tests/cli.rs new file mode 100644 index 00000000..50f2f7f7 --- /dev/null +++ b/src/pgstac-rs/tests/cli.rs @@ -0,0 +1,210 @@ +#![cfg(feature = "cli")] +//! End-to-end tests for the `pgstac` CLI binary: drive the actual built binary (via +//! `CARGO_BIN_EXE_pgstac`) through load -> maintain -> delete against a fresh clone of the ingest +//! template, asserting the database state over a side connection. Loads route through the Rust loader. + +use std::process::{Command, Output}; +use std::sync::atomic::{AtomicU32, Ordering}; +use tokio_postgres::NoTls; + +fn base() -> String { + std::env::var("PGSTAC_RS_TEST_BASE") + .unwrap_or_else(|_| "postgresql://username:password@localhost:5439".to_string()) +} + +fn template() -> String { + std::env::var("PGSTAC_RS_INGEST_TEMPLATE") + .unwrap_or_else(|_| "pgstac_rs_ingest_template".to_string()) +} + +/// A disposable database cloned from the ingest template, dropped on `Drop`. +struct CloneDb { + name: String, +} + +impl CloneDb { + async fn create() -> CloneDb { + static COUNTER: AtomicU32 = AtomicU32::new(0); + let name = format!( + "pgstac_rs_cli_test_{}_{}", + std::process::id(), + COUNTER.fetch_add(1, Ordering::Relaxed) + ); + let (client, connection) = tokio_postgres::connect(&format!("{}/postgres", base()), NoTls) + .await + .unwrap(); + let handle = tokio::spawn(connection); + client + .execute( + &format!("CREATE DATABASE {name} TEMPLATE {}", template()), + &[], + ) + .await + .unwrap(); + handle.abort(); + CloneDb { name } + } + + fn dsn(&self) -> String { + format!("{}/{}", base(), self.name) + } + + async fn count_items(&self) -> i64 { + let (client, connection) = tokio_postgres::connect(&self.dsn(), NoTls).await.unwrap(); + let handle = tokio::spawn(connection); + client + .batch_execute("SET search_path TO pgstac, public") + .await + .unwrap(); + let n: i64 = client + .query_one("SELECT count(*) FROM items", &[]) + .await + .unwrap() + .get(0); + handle.abort(); + n + } +} + +impl Drop for CloneDb { + fn drop(&mut self) { + let name = self.name.clone(); + std::thread::scope(|s| { + s.spawn(|| { + tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build() + .unwrap() + .block_on(async { + let (client, connection) = + tokio_postgres::connect(&format!("{}/postgres", base()), NoTls).await.unwrap(); + let handle = tokio::spawn(connection); + let _ = client + .execute( + "SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname = $1", + &[&name], + ) + .await; + let _ = client.execute(&format!("DROP DATABASE IF EXISTS {name} WITH (force)"), &[]).await; + handle.abort(); + }); + }); + }); + } +} + +fn pgstac(args: &[&str]) -> Output { + Command::new(env!("CARGO_BIN_EXE_pgstac")) + .args(args) + .output() + .expect("run pgstac binary") +} + +const COLLECTION: &str = r#"{"id":"clitest","type":"Collection","stac_version":"1.0.0","description":"cli test","license":"proprietary","extent":{"spatial":{"bbox":[[-180,-90,180,90]]},"temporal":{"interval":[[null,null]]}},"links":[]}"#; + +fn item(id: &str, datetime: &str) -> String { + format!( + r#"{{"type":"Feature","stac_version":"1.0.0","id":"{id}","collection":"clitest","geometry":{{"type":"Point","coordinates":[-105.1,40.1]}},"bbox":[-105.1,40.1,-105.1,40.1],"properties":{{"datetime":"{datetime}"}},"assets":{{}},"links":[]}}"# + ) +} + +/// load (collection + NDJSON items via the Rust loader) -> maintain (tighten) -> delete item -> delete +/// collection, asserting the database state at each step over a side connection. +#[tokio::test] +async fn cli_load_maintain_delete() { + let db = CloneDb::create().await; + let dsn = db.dsn(); + let dir = tempfile::tempdir().unwrap(); + std::fs::write(dir.path().join("collection.json"), COLLECTION).unwrap(); + std::fs::write( + dir.path().join("items.ndjson"), + format!( + "{}\n{}\n", + item("i1", "2023-01-05T00:00:00Z"), + item("i2", "2023-02-05T00:00:00Z") + ), + ) + .unwrap(); + + let out = pgstac(&["load", "--dsn", &dsn, dir.path().to_str().unwrap()]); + assert!( + out.status.success(), + "load failed: {}", + String::from_utf8_lossy(&out.stderr) + ); + assert_eq!(db.count_items().await, 2, "both items loaded via the CLI"); + + let out = pgstac(&["maintain", "--dsn", &dsn]); + assert!( + out.status.success(), + "maintain failed: {}", + String::from_utf8_lossy(&out.stderr) + ); + + let out = pgstac(&[ + "delete", + "--dsn", + &dsn, + "--collection", + "clitest", + "--item", + "i1", + "--yes", + ]); + assert!( + out.status.success(), + "delete item failed: {}", + String::from_utf8_lossy(&out.stderr) + ); + assert_eq!(db.count_items().await, 1, "one item left after deleting i1"); + + let out = pgstac(&["delete", "--dsn", &dsn, "--collection", "clitest", "--yes"]); + assert!( + out.status.success(), + "delete collection failed: {}", + String::from_utf8_lossy(&out.stderr) + ); + assert_eq!( + db.count_items().await, + 0, + "collection delete cascaded its items" + ); +} + +/// `load --policy error` rejects a duplicate id (the loader's Error conflict policy), and the CLI exits +/// non-zero. +#[tokio::test] +async fn cli_load_error_policy_rejects_duplicate() { + let db = CloneDb::create().await; + let dsn = db.dsn(); + let dir = tempfile::tempdir().unwrap(); + std::fs::write(dir.path().join("collection.json"), COLLECTION).unwrap(); + std::fs::write( + dir.path().join("items.ndjson"), + format!("{}\n", item("dup", "2023-01-05T00:00:00Z")), + ) + .unwrap(); + + let out = pgstac(&["load", "--dsn", &dsn, dir.path().to_str().unwrap()]); + assert!( + out.status.success(), + "first load: {}", + String::from_utf8_lossy(&out.stderr) + ); + + // Re-load the same id with the error policy (items only) -> must fail. + let items = dir.path().join("items.ndjson"); + let out = pgstac(&[ + "load", + "--dsn", + &dsn, + "--policy", + "error", + items.to_str().unwrap(), + ]); + assert!( + !out.status.success(), + "duplicate load with --policy error should fail" + ); + assert_eq!(db.count_items().await, 1, "duplicate not added"); +} diff --git a/src/pgstac-rs/tests/client.rs b/src/pgstac-rs/tests/client.rs new file mode 100644 index 00000000..d3ff5f50 --- /dev/null +++ b/src/pgstac-rs/tests/client.rs @@ -0,0 +1,64 @@ +//! The cached [`Client`] wrapper: the upstream `stac::api` traits (ItemsClient / CollectionsClient / +//! TransactionClient) backed by the Rust engine, caching the hydration invariants across calls. +//! Targets the rich fragment instance (default `pgstac_rs_test_rich`); skips if unreachable. + +#![allow(unused_crate_dependencies)] + +use pgstac::Client; +use stac::api::{CollectionsClient, ItemsClient, Search}; +use tokio_postgres::NoTls; + +async fn connect() -> Option> { + let dsn = std::env::var("PGSTAC_RS_TEST_RICH_DB").unwrap_or_else(|_| { + "postgresql://username:password@localhost:5439/pgstac_rs_test_rich".into() + }); + match tokio_postgres::connect(&dsn, NoTls).await { + Ok((client, connection)) => { + tokio::spawn(connection); + client + .batch_execute("SET search_path = pgstac, public;") + .await + .ok()?; + Some(Client::new(client)) + } + Err(e) => { + eprintln!("[client] skipping: {e}"); + None + } + } +} + +#[tokio::test] +async fn cached_client_implements_stac_traits() { + let Some(client) = connect().await else { + return; + }; + + // A (collection, id) sample, fetched via Deref to the wrapped connection. + let row = client + .query_one( + "SELECT collection, id FROM items WHERE collection = 'landsat-c2-l2' LIMIT 1", + &[], + ) + .await + .unwrap(); + let (collection, id): (String, String) = (row.get("collection"), row.get("id")); + + // ItemsClient::item — the first read detects + caches the hydration invariants. + let item = ItemsClient::item(&client, &collection, &id).await.unwrap(); + assert_eq!(item.map(|i| i.id), Some(id)); + + // ItemsClient::search — reuses the cached hydration (no re-detection round trips). + let mut search = Search::default(); + search.items.limit = Some(5); + let items = ItemsClient::search(&client, search).await.unwrap(); + assert_eq!(items.items.len(), 5); + + // CollectionsClient, also on the engine. + let collections = CollectionsClient::collections(&client).await.unwrap(); + assert!(!collections.is_empty()); + let one = CollectionsClient::collection(&client, &collection) + .await + .unwrap(); + assert_eq!(one.map(|c| c.id), Some(collection)); +} diff --git a/src/pgstac-rs/tests/collections.rs b/src/pgstac-rs/tests/collections.rs new file mode 100644 index 00000000..c2e00e6d --- /dev/null +++ b/src/pgstac-rs/tests/collections.rs @@ -0,0 +1,161 @@ +//! Fidelity gate for the Rust collection search + getters against their SQL counterparts. +//! +//! Targets a clean fragment instance (default `pgstac_rs_test_rich`); skips if unreachable. + +#![allow(unused_crate_dependencies)] + +use pgstac::collections::{collection_search, get_collection, get_queryables}; +use pgstac::search::get_item; +use serde_json::{Value, json}; +use tokio_postgres::{Client, NoTls}; + +async fn connect() -> Option { + let dsn = std::env::var("PGSTAC_RS_TEST_RICH_DB").unwrap_or_else(|_| { + "postgresql://username:password@localhost:5439/pgstac_rs_test_rich".into() + }); + match tokio_postgres::connect(&dsn, NoTls).await { + Ok((client, connection)) => { + tokio::spawn(connection); + client + .batch_execute("SET search_path = pgstac, public;") + .await + .ok()?; + Some(client) + } + Err(e) => { + eprintln!("[collections] skipping: cannot connect ({dsn}): {e}"); + None + } + } +} + +async fn sql_collection_search(client: &Client, body: &Value) -> Value { + client + .query_one("SELECT collection_search($1::jsonb) AS cs", &[body]) + .await + .expect("sql collection_search") + .get("cs") +} + +fn link_token(doc: &Value, rel: &str) -> Option { + doc["links"] + .as_array()? + .iter() + .find(|l| l["rel"] == rel) + .and_then(|l| l["href"].as_str()) + .and_then(|href| href.split("token=").nth(1)) + .map(str::to_string) +} + +fn id_of(v: &Value) -> String { + v["id"].as_str().unwrap_or_default().to_string() +} + +#[tokio::test] +async fn collection_search_paginates_and_matches_sql() { + let Some(client) = connect().await else { + return; + }; + // One collection per page; walk all of them via next tokens and compare to SQL collection_search. + let body = json!({"limit": 1, "sortby": [{"field": "id", "direction": "asc"}]}); + let mut token: Option = None; + let mut ids = Vec::new(); + let mut guard = 0; + loop { + guard += 1; + assert!(guard < 100, "collection pagination did not terminate"); + + let page = collection_search(&client, &body, token.as_deref()) + .await + .unwrap(); + + let mut sql_body = body.clone(); + if let Some(t) = &token { + let _ = sql_body + .as_object_mut() + .unwrap() + .insert("token".into(), Value::String(t.clone())); + } + let sql = sql_collection_search(&client, &sql_body).await; + let sql_collections = sql["collections"].as_array().cloned().unwrap_or_default(); + + assert_eq!( + page.features, sql_collections, + "collections mismatch at token {token:?}" + ); + assert_eq!( + page.next_token, + link_token(&sql, "next"), + "next token mismatch" + ); + assert_eq!( + page.prev_token, + link_token(&sql, "prev"), + "prev token mismatch" + ); + + ids.extend(page.features.iter().map(id_of)); + match page.next_token { + Some(next) if !page.features.is_empty() => token = Some(next), + _ => break, + } + } + assert_eq!( + ids.len(), + 3, + "expected to walk all 3 collections, got {ids:?}" + ); + eprintln!( + "[collections] paginated {} collections, matches collection_search()", + ids.len() + ); +} + +#[tokio::test] +async fn getters_return_the_right_objects() { + let Some(client) = connect().await else { + return; + }; + + // get_collection + let coll = get_collection(&client, "landsat-c2-l2").await.unwrap(); + assert_eq!(coll.as_ref().map(id_of), Some("landsat-c2-l2".to_string())); + assert!( + get_collection(&client, "does-not-exist") + .await + .unwrap() + .is_none() + ); + + // get_item: pick a real (collection, id) from the instance. + let row = client + .query_one( + "SELECT collection, id FROM items WHERE collection = 'landsat-c2-l2' LIMIT 1", + &[], + ) + .await + .unwrap(); + let (collection, id): (String, String) = (row.get("collection"), row.get("id")); + let item = get_item(&client, &collection, &id) + .await + .unwrap() + .expect("item"); + assert_eq!(item["id"], id); + // Fully hydrated: a landsat item carries its assets. + assert!(item["assets"].is_object() && !item["assets"].as_object().unwrap().is_empty()); + assert!( + get_item(&client, &collection, "no-such-item") + .await + .unwrap() + .is_none() + ); + + // get_queryables, catalog-wide and collection-scoped, are JSON objects. + let q = get_queryables(&client, None).await.unwrap(); + assert!(q.is_object() && q.get("properties").is_some()); + let qc = get_queryables(&client, Some("landsat-c2-l2")) + .await + .unwrap(); + assert!(qc.is_object() && qc.get("properties").is_some()); + eprintln!("[collections] getters + queryables ok"); +} diff --git a/src/pgstac-rs/tests/concurrency.rs b/src/pgstac-rs/tests/concurrency.rs new file mode 100644 index 00000000..4f7ab803 --- /dev/null +++ b/src/pgstac-rs/tests/concurrency.rs @@ -0,0 +1,517 @@ +//! High-concurrency ingest correctness tests — especially many writers hammering a SINGLE partition, +//! which is where the partition-creation / partition_stats / fragment locking has to be right. +//! +//! Gated on `PGSTAC_TEST_DSN` (a DSN to a pgstac DB that already has the pgstac schema); the test is a +//! no-op skip when it's unset, so `cargo test` stays green without a database. Each test uses a unique +//! collection id so runs don't collide. +//! +//! Run: `PGSTAC_TEST_DSN=postgresql://username:password@localhost:5439/ cargo test --features pool \ +//! --test concurrency -- --nocapture` + +#![cfg(feature = "pool")] + +use pgstac::ingest::ConflictPolicy; +use pgstac::{ConnectConfig, PgstacPool}; +use serde_json::{Value, json}; + +/// Flatten an error + its source chain into one string, so a failed assertion shows the real Postgres +/// message (SQLSTATE / server text) instead of tokio_postgres's terse top-line `db error`. +fn full_error(e: &dyn std::error::Error) -> String { + let mut s = e.to_string(); + let mut src = e.source(); + while let Some(inner) = src { + s.push_str(" | "); + s.push_str(&inner.to_string()); + src = inner.source(); + } + s +} + +fn dsn() -> Option { + std::env::var("PGSTAC_TEST_DSN").ok() +} + +async fn connect(dsn: &str) -> PgstacPool { + PgstacPool::connect(ConnectConfig { + dsn: Some(dsn.to_string()), + ..Default::default() + }) + .await + .expect("pool connect") +} + +fn collection_json(id: &str) -> Value { + json!({ + "type": "Collection", + "stac_version": "1.0.0", + "id": id, + "description": "concurrency test", + "license": "proprietary", + "extent": { + "spatial": {"bbox": [[-180.0, -90.0, 180.0, 90.0]]}, + "temporal": {"interval": [["2020-01-01T00:00:00Z", null]]} + }, + "links": [] + }) +} + +/// A minimal self-contained STAC item with the given id, in the given UTC month/day (so a `month` +/// partition_trunc routes it to one partition window). +fn make_item(collection: &str, id: &str, year: i32, month: u32, day: u32) -> Value { + let lon = -100.0 + (f64::from(day) * 0.01); + let lat = 40.0; + json!({ + "type": "Feature", + "stac_version": "1.0.0", + "id": id, + "collection": collection, + "geometry": {"type": "Polygon", "coordinates": [[ + [lon, lat], [lon + 0.1, lat], [lon + 0.1, lat + 0.1], [lon, lat + 0.1], [lon, lat] + ]]}, + "bbox": [lon, lat, lon + 0.1, lat + 0.1], + "properties": {"datetime": format!("{year}-{month:02}-{day:02}T00:00:00Z")}, + "assets": {}, + "links": [] + }) +} + +/// Create a fresh month-partitioned collection, returning a pooled connection's executor for assertions. +async fn setup_collection(pool: &PgstacPool, collection: &str) { + pool.create_collection(&collection_json(collection)) + .await + .expect("create_collection"); + let client = pool.get().await.expect("get"); + client + .execute( + "UPDATE pgstac.collections SET partition_trunc='month' WHERE id=$1", + &[&collection], + ) + .await + .expect("set month"); +} + +async fn count(pool: &PgstacPool, sql: &str, collection: &str) -> i64 { + let client = pool.get().await.expect("get"); + client + .query_one(sql, &[&collection]) + .await + .expect("count query") + .get(0) +} + +async fn cleanup(pool: &PgstacPool, collection: &str) { + if let Ok(client) = pool.get().await { + let _ = client + .execute("SELECT pgstac.delete_collection($1)", &[&collection]) + .await; + } +} + +/// Many concurrent bulk loads, ALL landing in the SAME month partition. Asserts no task errors, no lost or +/// duplicated rows, exactly one partition, and a partition_stats.n that covers the data (golden rule). +#[tokio::test(flavor = "multi_thread", worker_threads = 8)] +async fn concurrent_writers_single_partition() { + let Some(dsn) = dsn() else { + eprintln!("skip concurrent_writers_single_partition: PGSTAC_TEST_DSN unset"); + return; + }; + let pool = connect(&dsn).await; + let collection = format!("conc_single_{}", std::process::id()); + cleanup(&pool, &collection).await; + setup_collection(&pool, &collection).await; + + let writers = 8usize; + let per_writer = 1500usize; + let total = writers * per_writer; + + // Every item is in 2020-07 -> one partition; ids are globally unique across writers. + let mut batches: Vec> = (0..writers) + .map(|_| Vec::with_capacity(per_writer)) + .collect(); + for i in 0..total { + let day = 1 + u32::try_from(i % 28).unwrap(); + batches[i % writers].push(make_item(&collection, &format!("item-{i}"), 2020, 7, day)); + } + + let mut set = tokio::task::JoinSet::new(); + for batch in batches { + let pool = pool.clone(); + let _ = set.spawn(async move { pool.create_items(batch, ConflictPolicy::Error).await }); + } + let mut errors = Vec::new(); + while let Some(joined) = set.join_next().await { + match joined.expect("task panicked") { + Ok(_) => {} + Err(e) => errors.push(full_error(&e)), + } + } + assert!(errors.is_empty(), "concurrent loads errored: {errors:?}"); + + let n_items = count( + &pool, + "SELECT count(*) FROM pgstac.items WHERE collection=$1", + &collection, + ) + .await; + let n_distinct = count( + &pool, + "SELECT count(DISTINCT id) FROM pgstac.items WHERE collection=$1", + &collection, + ) + .await; + let n_parts = count( + &pool, + "SELECT count(*) FROM pgstac.partition_stats WHERE collection=$1", + &collection, + ) + .await; + let stats_n = count( + &pool, + "SELECT COALESCE(sum(n),0)::bigint FROM pgstac.partition_stats WHERE collection=$1", + &collection, + ) + .await; + + cleanup(&pool, &collection).await; + + assert_eq!( + n_items, total as i64, + "lost or duplicated rows under concurrency" + ); + assert_eq!(n_distinct, total as i64, "duplicate ids under concurrency"); + assert_eq!(n_parts, 1, "expected exactly one month partition"); + assert!( + stats_n >= total as i64, + "partition_stats.n ({stats_n}) under-counts the data ({total}) — golden rule violated" + ); +} + +/// Concurrent bulk loads spread across MANY months (so writers race on different + shared partitions and +/// on parent-partition creation). Asserts no errors, no lost/duplicated rows, and the right partition count. +#[tokio::test(flavor = "multi_thread", worker_threads = 8)] +async fn concurrent_writers_many_partitions() { + let Some(dsn) = dsn() else { + eprintln!("skip concurrent_writers_many_partitions: PGSTAC_TEST_DSN unset"); + return; + }; + let pool = connect(&dsn).await; + let collection = format!("conc_many_{}", std::process::id()); + cleanup(&pool, &collection).await; + setup_collection(&pool, &collection).await; + + let writers = 8usize; + let months = 12u32; + let per_writer = 1200usize; + let total = writers * per_writer; + + // Each writer's items span all 12 months -> every writer touches every partition (max contention on + // shared parent creation + each month's stats row). + let mut batches: Vec> = (0..writers) + .map(|_| Vec::with_capacity(per_writer)) + .collect(); + for i in 0..total { + let month = 1 + u32::try_from(i).unwrap() % months; + batches[i % writers].push(make_item( + &collection, + &format!("item-{i}"), + 2021, + month, + 15, + )); + } + + let mut set = tokio::task::JoinSet::new(); + for batch in batches { + let pool = pool.clone(); + let _ = set.spawn(async move { pool.create_items(batch, ConflictPolicy::Error).await }); + } + let mut errors = Vec::new(); + while let Some(joined) = set.join_next().await { + match joined.expect("task panicked") { + Ok(_) => {} + Err(e) => errors.push(full_error(&e)), + } + } + assert!( + errors.is_empty(), + "concurrent multi-partition loads errored: {errors:?}" + ); + + let n_items = count( + &pool, + "SELECT count(*) FROM pgstac.items WHERE collection=$1", + &collection, + ) + .await; + let n_distinct = count( + &pool, + "SELECT count(DISTINCT id) FROM pgstac.items WHERE collection=$1", + &collection, + ) + .await; + let n_parts = count( + &pool, + "SELECT count(*) FROM pgstac.partition_stats WHERE collection=$1", + &collection, + ) + .await; + + cleanup(&pool, &collection).await; + + assert_eq!( + n_items, total as i64, + "lost or duplicated rows under concurrency" + ); + assert_eq!(n_distinct, total as i64, "duplicate ids under concurrency"); + assert_eq!( + n_parts, months as i64, + "expected one partition per month touched" + ); +} + +/// Precheck-driven upsert: skips unchanged items on re-ingest, detects a content change, and correctly +/// handles a datetime change that MOVES an item to a different partition (no duplicate across partitions — +/// the case per-partition id uniqueness would miss). +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn precheck_skip_unchanged_and_partition_move() { + let Some(dsn) = dsn() else { + eprintln!("skip precheck_skip_unchanged_and_partition_move: PGSTAC_TEST_DSN unset"); + return; + }; + let pool = connect(&dsn).await; + let collection = format!("precheck_{}", std::process::id()); + cleanup(&pool, &collection).await; + setup_collection(&pool, &collection).await; + + let n = 1000usize; + let items: Vec = (0..n) + .map(|i| { + make_item( + &collection, + &format!("item-{i}"), + 2020, + 7, + 1 + u32::try_from(i % 28).unwrap(), + ) + }) + .collect(); + let count_sql = "SELECT count(*) FROM pgstac.items WHERE collection=$1"; + + // initial upsert: everything is new. + let (unchanged0, loaded0) = pool + .upsert_items_precheck(items.clone(), ConflictPolicy::Upsert) + .await + .unwrap(); + assert_eq!( + (unchanged0, loaded0), + (0, n as u64), + "initial load: all new" + ); + + // re-upsert identical items: all unchanged, nothing loaded (the skip-unchanged / re-ingest win). + let (unchanged1, loaded1) = pool + .upsert_items_precheck(items.clone(), ConflictPolicy::Upsert) + .await + .unwrap(); + assert_eq!( + (unchanged1, loaded1), + (n as u64, 0), + "re-ingest of identical items skips everything" + ); + + // change one item's content -> exactly one changed/loaded, the rest skipped, no duplicate. + let mut changed = items.clone(); + changed[0]["properties"]["pgstac:test"] = json!("modified"); + let (unchanged2, loaded2) = pool + .upsert_items_precheck(changed, ConflictPolicy::Upsert) + .await + .unwrap(); + assert_eq!( + (unchanged2, loaded2), + ((n - 1) as u64, 1), + "one content change -> one load" + ); + assert_eq!( + count(&pool, count_sql, &collection).await, + n as i64, + "content change must not duplicate" + ); + + // move one item to a different month (datetime change -> different partition). The precheck reports it + // 'new' (the new partition holds no such id), so we load with DELSERT, whose cross-partition delete + // removes the old row from its former partition -> exactly one row. (UPSERT would orphan it by design — + // it only touches the partition the new row routes to.) + let mut moved = items.clone(); + moved[0]["properties"]["pgstac:test"] = json!("modified"); // item-0 already stored as modified -> unchanged + moved[1] = make_item(&collection, "item-1", 2020, 8, 15); // item-1: 2020-07 -> 2020-08 + let (unchanged3, loaded3) = pool + .upsert_items_precheck(moved, ConflictPolicy::Delsert) + .await + .unwrap(); + assert_eq!( + (unchanged3, loaded3), + ((n - 1) as u64, 1), + "only the moved item is loaded" + ); + assert_eq!( + count(&pool, count_sql, &collection).await, + n as i64, + "partition move must not duplicate" + ); + let item1_rows = count( + &pool, + "SELECT count(*) FROM pgstac.items WHERE collection=$1 AND id='item-1'", + &collection, + ) + .await; + assert_eq!( + item1_rows, 1, + "moved item must exist exactly once (not in both partitions)" + ); + + cleanup(&pool, &collection).await; +} + +/// Ignore + precheck (id-only): every existing id is skipped regardless of content, so a changed re-ingest +/// under `Ignore` loads nothing and leaves the stored rows untouched; only genuinely new ids are loaded. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn precheck_ignore_skips_existing_by_id() { + let Some(dsn) = dsn() else { + eprintln!("skip precheck_ignore_skips_existing_by_id: PGSTAC_TEST_DSN unset"); + return; + }; + let pool = connect(&dsn).await; + let collection = format!("precheck_ign_{}", std::process::id()); + cleanup(&pool, &collection).await; + setup_collection(&pool, &collection).await; + + let n = 1000usize; + let items: Vec = (0..n) + .map(|i| { + make_item( + &collection, + &format!("item-{i}"), + 2020, + 7, + 1 + u32::try_from(i % 28).unwrap(), + ) + }) + .collect(); + let count_sql = "SELECT count(*) FROM pgstac.items WHERE collection=$1"; + + // initial ignore load: everything is new. + let (skipped0, loaded0) = pool + .upsert_items_precheck(items.clone(), ConflictPolicy::Ignore) + .await + .unwrap(); + assert_eq!( + (skipped0, loaded0), + (0, n as u64), + "initial ignore load: all new" + ); + + // re-ingest with content CHANGED on every id (same ids). id-only ignore skips them ALL by id — a + // hash-based precheck would instead call them 'changed' and skip none. + let mut changed = items.clone(); + for it in &mut changed { + it["properties"]["pgstac:test"] = json!("modified"); + } + let (skipped1, loaded1) = pool + .upsert_items_precheck(changed, ConflictPolicy::Ignore) + .await + .unwrap(); + assert_eq!( + (skipped1, loaded1), + (n as u64, 0), + "ignore skips every existing id (changed or not); nothing loaded" + ); + assert_eq!( + count(&pool, count_sql, &collection).await, + n as i64, + "ignore must not add or duplicate rows" + ); + + // a batch mixing the existing ids with genuinely new ones -> only the new ids load. + let new_ids = 50usize; + let mut mixed = items.clone(); + for j in 0..new_ids { + mixed.push(make_item( + &collection, + &format!("new-{j}"), + 2020, + 7, + 1 + u32::try_from(j % 28).unwrap(), + )); + } + let (skipped2, loaded2) = pool + .upsert_items_precheck(mixed, ConflictPolicy::Ignore) + .await + .unwrap(); + assert_eq!( + (skipped2, loaded2), + (n as u64, new_ids as u64), + "ignore loads only the new ids, skips all existing" + ); + assert_eq!( + count(&pool, count_sql, &collection).await, + (n + new_ids) as i64, + "exactly the new ids were added" + ); + + cleanup(&pool, &collection).await; +} + +/// Locks in the upsert vs delsert semantics on a cross-partition move: `Upsert` only touches the partition +/// the new row routes to, so it ORPHANS the old row; `Delsert` deletes cross-partition, so it does not. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn upsert_orphans_on_move_but_delsert_does_not() { + let Some(dsn) = dsn() else { + eprintln!("skip upsert_orphans_on_move_but_delsert_does_not: PGSTAC_TEST_DSN unset"); + return; + }; + let pool = connect(&dsn).await; + let collection = format!("orphan_{}", std::process::id()); + cleanup(&pool, &collection).await; + setup_collection(&pool, &collection).await; + let id_rows = "SELECT count(*) FROM pgstac.items WHERE collection=$1 AND id='x'"; + + // initial: x in 2020-07. + let _ = pool + .create_items( + vec![make_item(&collection, "x", 2020, 7, 1)], + ConflictPolicy::Ignore, + ) + .await + .unwrap(); + assert_eq!(count(&pool, id_rows, &collection).await, 1, "initial load"); + + // UPSERT a move 2020-07 -> 2020-08: same-partition only, so the 2020-07 row is left behind -> 2 rows. + let _ = pool + .create_items( + vec![make_item(&collection, "x", 2020, 8, 1)], + ConflictPolicy::Upsert, + ) + .await + .unwrap(); + assert_eq!( + count(&pool, id_rows, &collection).await, + 2, + "Upsert on a cross-partition move orphans the old row (by design)" + ); + + // DELSERT a move -> 2020-09: the cross-partition delete removes BOTH prior rows -> 1 row. + let _ = pool + .create_items( + vec![make_item(&collection, "x", 2020, 9, 1)], + ConflictPolicy::Delsert, + ) + .await + .unwrap(); + assert_eq!( + count(&pool, id_rows, &collection).await, + 1, + "Delsert removes the old rows cross-partition" + ); + + cleanup(&pool, &collection).await; +} diff --git a/src/pgstac-rs/tests/dehydrate_parity.rs b/src/pgstac-rs/tests/dehydrate_parity.rs new file mode 100644 index 00000000..d197ef51 --- /dev/null +++ b/src/pgstac-rs/tests/dehydrate_parity.rs @@ -0,0 +1,188 @@ +//! Parity gate: the Rust `dehydrate` must produce the same `items` row as SQL `content_dehydrate`, so an +//! item ingested through Rust is byte-identical to one ingested through the SQL path. +//! +//! Connects read-only to a pgstac database (`PGSTAC_RS_TEST_DB`, default the dev `postgis` db). + +use chrono::{DateTime, Utc}; +use pgstac::dehydrate::{DehydrateSchema, PromotedColumn, PromotedKind, PromotedValue, dehydrate}; +use pgstac::geom::Ewkb; +use serde_json::{Value, json}; +use tokio_postgres::{Client, NoTls, Row}; + +fn dsn() -> String { + std::env::var("PGSTAC_RS_TEST_DB") + .unwrap_or_else(|_| "postgresql://username:password@localhost:5439/postgis".to_string()) +} + +async fn connect() -> Client { + let (client, connection) = tokio_postgres::connect(&dsn(), NoTls).await.unwrap(); + tokio::spawn(async move { + let _ = connection.await; + }); + client + .batch_execute("SET search_path TO pgstac, public;") + .await + .unwrap(); + client +} + +/// Representative items: singular datetime, datetime range, promoted props (text/float/int/array/jsonb/ +/// timestamptz), empty/absent assets+links, extra top-level keys. +fn items() -> Vec { + vec![ + json!({ + "type": "Feature", + "stac_version": "1.0.0", + "stac_extensions": ["https://stac-extensions.github.io/eo/v1.1.0/schema.json"], + "id": "item-singular", + "collection": "c1", + "geometry": {"type": "Point", "coordinates": [-105.1019, 40.1672]}, + "bbox": [-105.1019, 40.1672, -105.1019, 40.1672], + "properties": { + "datetime": "2023-01-07T00:00:00Z", + "platform": "landsat-8", + "instruments": ["oli", "tirs"], + "eo:cloud_cover": 12.5, + "gsd": 30.0, + "sat:relative_orbit": 42, + "created": "2023-01-08T01:02:03Z", + "proj:bbox": [1.0, 2.0, 3.0, 4.0], + "custom:keep": "stays-in-properties" + }, + "assets": {"thumbnail": {"href": "https://x/t.png", "type": "image/png"}}, + "links": [{"rel": "self", "href": "https://x/item-1"}, {"rel": "root", "href": "https://x/"}], + "extra_top": {"a": 1} + }), + json!({ + "type": "Feature", + "stac_version": "1.1.0", + "id": "item-range", + "collection": "c1", + "geometry": {"type": "Polygon", "coordinates": [[[-85.3, 30.9], [-85.3, 31.0], [-85.2, 31.0], [-85.2, 30.9], [-85.3, 30.9]]]}, + "properties": { + "start_datetime": "2020-01-01T00:00:00Z", + "end_datetime": "2020-06-01T00:00:00Z", + "datetime": null, + "eo:cloud_cover": 0.0 + }, + "assets": {}, + "links": [] + }), + json!({ + "type": "Feature", + "id": "item-minimal", + "collection": "c2", + "geometry": {"type": "Point", "coordinates": [0.0, 0.0]}, + "properties": {"datetime": "2024-06-26T12:00:00.500Z"} + }), + ] +} + +fn read_sql_promoted(row: &Row, col: &PromotedColumn) -> PromotedValue { + let c = col.column.as_str(); + match col.kind { + PromotedKind::Text => PromotedValue::Text(row.get(c)), + PromotedKind::Float => PromotedValue::Float(row.get(c)), + PromotedKind::Int => PromotedValue::Int(row.get(c)), + PromotedKind::BigInt => PromotedValue::BigInt(row.get(c)), + PromotedKind::TextArray => PromotedValue::TextArray(row.get(c)), + PromotedKind::Jsonb => PromotedValue::Jsonb(row.get(c)), + PromotedKind::Timestamptz => PromotedValue::Timestamptz(row.get(c)), + } +} + +#[tokio::test] +async fn dehydrate_matches_content_dehydrate() { + let client = connect().await; + let schema = DehydrateSchema::load(&client).await.unwrap(); + + for item in items() { + let row = client + .query_one("SELECT * FROM content_dehydrate($1::jsonb)", &[&item]) + .await + .unwrap(); + let id_for_msg = item["id"].as_str().unwrap_or("?").to_string(); + let rust = dehydrate(item, &schema).unwrap(); + + assert_eq!(rust.id, row.get::<_, String>("id"), "id ({id_for_msg})"); + assert_eq!( + rust.collection, + row.get::<_, String>("collection"), + "collection ({id_for_msg})" + ); + assert_eq!( + rust.datetime, + row.get::<_, DateTime>("datetime"), + "datetime ({id_for_msg})" + ); + assert_eq!( + rust.end_datetime, + row.get::<_, DateTime>("end_datetime"), + "end_datetime ({id_for_msg})" + ); + assert_eq!( + rust.datetime_is_range, + row.get::<_, bool>("datetime_is_range"), + "datetime_is_range ({id_for_msg})" + ); + assert_eq!( + rust.stac_version, + row.get::<_, Option>("stac_version"), + "stac_version ({id_for_msg})" + ); + assert_eq!( + rust.stac_extensions, + row.get::<_, Value>("stac_extensions"), + "stac_extensions ({id_for_msg})" + ); + assert_eq!( + rust.item_hash.to_vec(), + row.get::<_, Vec>("item_hash"), + "item_hash ({id_for_msg})" + ); + assert_eq!( + rust.bbox, + row.get::<_, Option>("bbox"), + "bbox ({id_for_msg})" + ); + assert_eq!( + rust.links, + row.get::<_, Option>("links"), + "links ({id_for_msg})" + ); + assert_eq!( + rust.assets, + row.get::<_, Option>("assets"), + "assets ({id_for_msg})" + ); + assert_eq!( + rust.properties, + row.get::<_, Value>("properties"), + "properties ({id_for_msg})" + ); + assert_eq!( + rust.extra, + row.get::<_, Value>("extra"), + "extra ({id_for_msg})" + ); + assert_eq!( + rust.link_hrefs, + row.get::<_, Option>>>("link_hrefs"), + "link_hrefs ({id_for_msg})" + ); + assert_eq!( + rust.geometry, + row.get::<_, Ewkb>("geometry").0, + "geometry ({id_for_msg})" + ); + + for (i, col) in schema.columns().iter().enumerate() { + assert_eq!( + rust.promoted[i], + read_sql_promoted(&row, col), + "promoted {} ({id_for_msg})", + col.column + ); + } + } +} diff --git a/src/pgstac-rs/tests/dump_e2e.rs b/src/pgstac-rs/tests/dump_e2e.rs new file mode 100644 index 00000000..1141e065 --- /dev/null +++ b/src/pgstac-rs/tests/dump_e2e.rs @@ -0,0 +1,231 @@ +//! End-to-end dump: DumpPlanner -> DirSink, full + partial, +//! on a 0.9.11 (NULL-trunc, base_item) and a 0.10 (month, fragment) instance. +//! +//! Validates the on-disk layout + manifest (MANIFEST.md), per-file SHA-256, item +//! counts vs SQL, and that every partition geoparquet is readable and hydrated. +//! Skips if the DB is unreachable. + +#![cfg(feature = "export")] +#![allow(unused_crate_dependencies)] + +use pgstac::export::manifest::Manifest; +use pgstac::export::plan::Prefilter; +use pgstac::export::sink::DirSink; +use pgstac::export::{DumpConfig, DumpPlanner}; +use std::collections::HashMap; +use std::path::Path; +use tokio_postgres::{Client, NoTls}; + +async fn connect(env: &str, default: &str) -> Option { + let conn = std::env::var(env).unwrap_or_else(|_| default.to_string()); + match tokio_postgres::connect(&conn, NoTls).await { + Ok((client, connection)) => { + tokio::spawn(async move { + let _ = connection.await; + }); + if client + .batch_execute("SET search_path = pgstac, public;") + .await + .is_err() + { + return None; + } + Some(client) + } + Err(e) => { + eprintln!("[dump] skipping: cannot connect via {env} ({default}): {e}"); + None + } + } +} + +/// SQL item counts per collection. +async fn sql_counts(client: &Client) -> HashMap { + let rows = client + .query( + "SELECT collection, count(*) AS c FROM items GROUP BY collection", + &[], + ) + .await + .unwrap(); + rows.iter() + .map(|r| (r.get::<_, String>("collection"), r.get::<_, i64>("c"))) + .collect() +} + +/// Reads + sha256-verifies the manifest against the on-disk files, and confirms +/// every partition geoparquet is readable with the recorded item count. +fn verify_dump(root: &Path, sql_counts: &HashMap, expect_full: bool) { + let manifest_path = root.join("manifest.json"); + assert!( + manifest_path.exists(), + "manifest.json must exist (complete)" + ); + // Manifest presence is the completion marker; the checkpoint is removed. + assert!( + !root.join("_checkpoint.json").exists(), + "_checkpoint.json should be removed on success" + ); + + let manifest_bytes = std::fs::read(&manifest_path).unwrap(); + let manifest: Manifest = serde_json::from_slice(&manifest_bytes).unwrap(); + assert_eq!(manifest.manifest_version, "1"); + assert!(manifest.options.hydrated); + assert_eq!(manifest.options.ordering, vec!["datetime", "id"]); + if expect_full { + assert_eq!(manifest.dump_type, "full"); + assert!(manifest.filter.is_none()); + } + + // Root metadata files exist + sha256 match. + for fe in manifest.metadata_files.values() { + let p = root.join(&fe.path); + let bytes = std::fs::read(&p).unwrap_or_else(|_| panic!("missing {}", fe.path)); + assert_eq!( + pgstac::export::manifest::sha256_hex(&bytes), + fe.sha256, + "sha256 mismatch for {}", + fe.path + ); + assert_eq!(bytes.len() as u64, fe.bytes); + } + + for coll in &manifest.collections { + // collection.json + let cp = root.join(&coll.collection_file.path); + let cbytes = std::fs::read(&cp).unwrap(); + assert_eq!( + pgstac::export::manifest::sha256_hex(&cbytes), + coll.collection_file.sha256 + ); + + // Sum of partition item counts == collection count == SQL count. + let mut sum = 0u64; + for part in &coll.partitions { + let pp = root.join(&part.file); + let pbytes = std::fs::read(&pp).unwrap_or_else(|_| panic!("missing {}", part.file)); + assert_eq!( + pgstac::export::manifest::sha256_hex(&pbytes), + part.sha256, + "parquet sha256 mismatch for {}", + part.file + ); + assert_eq!(pbytes.len() as u64, part.bytes); + + // The geoparquet must be readable with the recorded count. + let read = read_geoparquet_count(&pbytes); + assert_eq!( + read as u64, part.item_count, + "geoparquet row count != manifest for {}", + part.file + ); + sum += part.item_count; + } + assert_eq!(sum, coll.item_count, "partition sum != collection count"); + if let Some(expected) = sql_counts.get(&coll.id) { + assert_eq!( + coll.item_count as i64, *expected, + "dumped count != SQL count for {}", + coll.id + ); + } + } + eprintln!( + "[dump] verified {} collections, manifest + sha256 + geoparquet readable", + manifest.collections.len() + ); +} + +fn read_geoparquet_count(bytes: &[u8]) -> usize { + use std::io::{Seek, SeekFrom, Write}; + let mut tf = tempfile::tempfile().unwrap(); + tf.write_all(bytes).unwrap(); + let _ = tf.seek(SeekFrom::Start(0)).unwrap(); + let ic = stac::geoparquet::from_reader(tf).unwrap(); + // Sanity: first item is hydrated (has type Feature + properties). + if let Some(first) = ic.items.first() { + let v = serde_json::to_value(first).unwrap(); + assert_eq!(v["type"], "Feature"); + assert!(v.get("properties").is_some()); + } + ic.items.len() +} + +#[tokio::test] +async fn dump_full_0911_nulltrunc() { + let Some(client) = connect( + "PGSTAC_PARITY_DB_0911", + "postgresql://username:password@localhost:5439/a_parity0911", + ) + .await + else { + return; + }; + let dir = tempfile::tempdir().unwrap(); + let sink = DirSink::new(dir.path()).unwrap(); + let planner = DumpPlanner::new(DumpConfig::default()); + let report = planner.run(&client, &sink).await.unwrap(); + eprintln!( + "[dump] 0911 full: {} collections, {} partitions, {} items", + report.collections, report.partitions, report.items + ); + assert!(report.items > 0); + let counts = sql_counts(&client).await; + verify_dump(dir.path(), &counts, true); +} + +#[tokio::test] +async fn dump_full_010_month() { + let Some(client) = connect( + "PGSTAC_PARITY_DB_010", + "postgresql://username:password@localhost:5439/a_parity010", + ) + .await + else { + return; + }; + let dir = tempfile::tempdir().unwrap(); + let sink = DirSink::new(dir.path()).unwrap(); + let planner = DumpPlanner::new(DumpConfig::default()); + let report = planner.run(&client, &sink).await.unwrap(); + eprintln!( + "[dump] 010 full: {} collections, {} partitions, {} items", + report.collections, report.partitions, report.items + ); + assert!(report.items > 0); + assert!(report.partitions >= 2, "month-partitioned -> many files"); + let counts = sql_counts(&client).await; + verify_dump(dir.path(), &counts, true); +} + +#[tokio::test] +async fn dump_partial_010_one_collection() { + let Some(client) = connect( + "PGSTAC_PARITY_DB_010", + "postgresql://username:password@localhost:5439/a_parity010", + ) + .await + else { + return; + }; + // Partial: a single collection. + let dir = tempfile::tempdir().unwrap(); + let sink = DirSink::new(dir.path()).unwrap(); + let config = DumpConfig { + collection_ids: Some(vec!["bench_uniform".to_string()]), + prefilter: Prefilter::default(), + ..Default::default() + }; + let planner = DumpPlanner::new(config); + let report = planner.run(&client, &sink).await.unwrap(); + assert_eq!(report.collections, 1); + assert!(report.items > 0); + + let manifest: Manifest = + serde_json::from_slice(&std::fs::read(dir.path().join("manifest.json")).unwrap()).unwrap(); + assert_eq!(manifest.dump_type, "partial"); + let filter = manifest.filter.expect("partial has a filter"); + assert_eq!(filter.collection_ids, vec!["bench_uniform"]); + assert_eq!(manifest.collections.len(), 1); + eprintln!("[dump] 010 partial bench_uniform: {} items", report.items); +} diff --git a/src/pgstac-rs/tests/dump_followups.rs b/src/pgstac-rs/tests/dump_followups.rs new file mode 100644 index 00000000..81bf581a --- /dev/null +++ b/src/pgstac-rs/tests/dump_followups.rs @@ -0,0 +1,378 @@ +//! Follow-up integration tests: parallel fan-out, --consistent snapshot, and resume. +//! +//! * Parallel parity: `run_parallel` produces the same item counts + per-file +//! sha256 set as the sequential `run` on the same instance. +//! * --consistent: a concurrent insert during a consistent dump is NOT captured. +//! * Resume: a dump with a pre-seeded checkpoint skips the completed partition +//! yet still lists it in the manifest (no missing partitions, no re-dump). +//! +//! All tests skip if the DB is unreachable, so the suite stays green without +//! fixtures. Point them at month-partitioned 0.10 clones via the env vars below. + +#![cfg(feature = "export")] +#![allow(unused_crate_dependencies)] + +use pgstac::export::manifest::Manifest; +use pgstac::export::parallel::DsnFactory; +use pgstac::export::sink::DirSink; +use pgstac::export::{DumpConfig, DumpPlanner}; +use std::collections::BTreeMap; +use std::sync::Arc; +use tokio_postgres::{Client, NoTls}; + +const DEFAULT_010: &str = "postgresql://username:password@localhost:5439/a_parity010"; + +async fn connect(dsn: &str) -> Option { + match tokio_postgres::connect(dsn, NoTls).await { + Ok((client, connection)) => { + tokio::spawn(async move { + let _ = connection.await; + }); + if client + .batch_execute("SET search_path = pgstac, public;") + .await + .is_err() + { + return None; + } + Some(client) + } + Err(e) => { + eprintln!("[followups] skipping: cannot connect ({dsn}): {e}"); + None + } + } +} + +fn dsn_010() -> String { + std::env::var("PGSTAC_PARITY_DB_010").unwrap_or_else(|_| DEFAULT_010.to_string()) +} + +/// Maps every partition file in a manifest to its (sha256, item_count). +fn file_map(m: &Manifest) -> BTreeMap { + let mut map = BTreeMap::new(); + for c in &m.collections { + for p in &c.partitions { + let _ = map.insert(p.file.clone(), (p.sha256.clone(), p.item_count)); + } + } + map +} + +fn read_manifest(dir: &std::path::Path) -> Manifest { + serde_json::from_slice(&std::fs::read(dir.join("manifest.json")).unwrap()).unwrap() +} + +/// Parallel fan-out yields the same files + counts as the sequential run. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn parallel_matches_sequential_010() { + let dsn = dsn_010(); + let Some(client) = connect(&dsn).await else { + return; + }; + + // Restrict to one collection to keep the test quick (still many partitions). + let only = first_collection(&client).await; + + // Sequential. + let seq_dir = tempfile::tempdir().unwrap(); + let seq_report = DumpPlanner::new(DumpConfig { + collection_ids: only.clone(), + ..Default::default() + }) + .run(&client, &DirSink::new(seq_dir.path()).unwrap()) + .await + .unwrap(); + + // Parallel (4 workers). + let par_dir = tempfile::tempdir().unwrap(); + let factory = DsnFactory::new(dsn.clone()); + let par_report = DumpPlanner::new(DumpConfig { + collection_ids: only.clone(), + concurrency: 4, + ..Default::default() + }) + .run_parallel( + &client, + Arc::new(DirSink::new(par_dir.path()).unwrap()), + &factory, + ) + .await + .unwrap(); + + assert_eq!(seq_report.items, par_report.items, "item totals differ"); + assert_eq!( + seq_report.partitions, par_report.partitions, + "partition counts differ" + ); + + let seq = file_map(&read_manifest(seq_dir.path())); + let par = file_map(&read_manifest(par_dir.path())); + // Same set of files, identical per-file item counts. (Geoparquet bytes are + // deterministic for the same rows, so sha256 should match too.) + assert_eq!( + seq.keys().collect::>(), + par.keys().collect::>(), + "different partition file sets" + ); + for (file, (seq_sha, seq_n)) in &seq { + let (par_sha, par_n) = par.get(file).unwrap(); + assert_eq!(seq_n, par_n, "item_count differs for {file}"); + assert_eq!(seq_sha, par_sha, "sha256 differs for {file}"); + } + eprintln!( + "[followups] parallel parity: {} files, {} items match sequential", + seq.len(), + par_report.items + ); +} + +async fn first_collection(client: &Client) -> Option> { + let row = client + .query_opt( + "SELECT id FROM collections WHERE partition_trunc IS NOT NULL ORDER BY id LIMIT 1", + &[], + ) + .await + .unwrap(); + row.map(|r| vec![r.get::<_, String>("id")]) +} + +/// A concurrent insert during a --consistent dump must NOT appear in the +/// output (the dump reads one repeatable-read snapshot taken before the insert). +/// +/// This test mutates the DB, so it only runs against an explicitly opt-in +/// disposable DSN (`PGSTAC_CONSISTENT_DB`) to avoid corrupting shared fixtures. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn consistent_excludes_concurrent_insert() { + let Ok(dsn) = std::env::var("PGSTAC_CONSISTENT_DB") else { + eprintln!( + "[followups] skipping consistent test: set PGSTAC_CONSISTENT_DB to a disposable DB" + ); + return; + }; + let Some(client) = connect(&dsn).await else { + return; + }; + + // Pick a target collection + a fresh item id that does not yet exist. + let Some(coll) = first_collection(&client).await else { + eprintln!("[followups] no partitioned collection; skipping"); + return; + }; + let coll_id = &coll[0]; + + let before: i64 = client + .query_one("SELECT count(*) AS c FROM items", &[]) + .await + .unwrap() + .get("c"); + + // Start a consistent dump on one connection; while it holds its snapshot, + // insert a new item on a second connection. The snapshot is taken at dump + // start, so the insert must not be captured. + // + // We force the insert to land mid-dump by inserting right after the planner + // exports the snapshot. Since run_parallel exports the snapshot first thing, + // a small spawn-then-insert races safely after we kick off the dump. + let dir = tempfile::tempdir().unwrap(); + let factory = DsnFactory::new(dsn.clone()); + let planner = DumpPlanner::new(DumpConfig { + collection_ids: Some(coll.clone()), + concurrency: 2, + consistent: true, + ..Default::default() + }); + + // Insert concurrently shortly after the dump begins. + let insert_dsn = dsn.clone(); + let insert_coll = coll_id.clone(); + let inserter = tokio::spawn(async move { + tokio::time::sleep(std::time::Duration::from_millis(150)).await; + let Some(mut c2) = connect(&insert_dsn).await else { + return false; + }; + let item = serde_json::json!({ + "type": "Feature", + "stac_version": "1.0.0", + "id": "consistency-probe-item", + "collection": insert_coll, + "geometry": {"type": "Point", "coordinates": [0.0, 0.0]}, + "bbox": [0.0, 0.0, 0.0, 0.0], + "properties": {"datetime": "2024-01-15T00:00:00Z"}, + "assets": {}, + "links": [] + }); + // Insert through the Rust loader (the only write path from Rust), not the SQL create_item. + let Ok(schema) = pgstac::dehydrate::DehydrateSchema::load(&c2).await else { + return false; + }; + pgstac::ingest::load_items( + &mut c2, + vec![item], + &schema, + pgstac::ingest::ConflictPolicy::Error, + ) + .await + .is_ok() + }); + + let report = planner + .run_parallel( + &client, + Arc::new(DirSink::new(dir.path()).unwrap()), + &factory, + ) + .await + .unwrap(); + let inserted = inserter.await.unwrap(); + assert!( + inserted, + "concurrent insert did not succeed; test inconclusive" + ); + + // The dump's item total reflects only the snapshot before the insert. + let manifest = read_manifest(dir.path()); + let dumped: u64 = manifest.collections.iter().map(|c| c.item_count).sum(); + assert_eq!(report.items, dumped); + // No partition file should contain the probe item id. Read each parquet. + let mut found_probe = false; + for c in &manifest.collections { + for p in &c.partitions { + let bytes = std::fs::read(dir.path().join(&p.file)).unwrap(); + if geoparquet_has_id(&bytes, "consistency-probe-item") { + found_probe = true; + } + } + } + assert!( + !found_probe, + "consistent dump captured a concurrently-inserted item (snapshot not honored)" + ); + assert!( + manifest.consistent, + "manifest should record consistent=true" + ); + assert!( + manifest.snapshot.is_some(), + "manifest should record the snapshot id" + ); + + // Clean up the probe item so the disposable DB can be reused. + let _ = client + .execute( + "SELECT delete_item('consistency-probe-item', $1)", + &[&coll_id], + ) + .await; + + eprintln!( + "[followups] consistent: dumped {} items (before={before}); concurrent insert correctly excluded", + report.items + ); +} + +fn geoparquet_has_id(bytes: &[u8], id: &str) -> bool { + use std::io::{Seek, SeekFrom, Write}; + let mut tf = tempfile::tempfile().unwrap(); + tf.write_all(bytes).unwrap(); + let _ = tf.seek(SeekFrom::Start(0)).unwrap(); + let ic = stac::geoparquet::from_reader(tf).unwrap(); + ic.items.iter().any(|it| { + serde_json::to_value(it) + .ok() + .and_then(|v| v.get("id").and_then(|x| x.as_str()).map(|s| s == id)) + .unwrap_or(false) + }) +} + +/// A resumed dump skips a completed partition (does not re-dump it) but +/// still lists it in the manifest, producing the same complete manifest. +#[tokio::test] +async fn resume_skips_completed_but_keeps_in_manifest() { + let dsn = dsn_010(); + let Some(client) = connect(&dsn).await else { + return; + }; + let only = first_collection(&client).await; + + // Full dump first. + let dir1 = tempfile::tempdir().unwrap(); + let _ = DumpPlanner::new(DumpConfig { + collection_ids: only.clone(), + ..Default::default() + }) + .run(&client, &DirSink::new(dir1.path()).unwrap()) + .await + .unwrap(); + let full = read_manifest(dir1.path()); + let full_files = file_map(&full); + assert!( + full_files.len() >= 2, + "need multiple partitions to test resume" + ); + + // Build a resume set from the full manifest: pretend the first partition was + // already completed (point its checkpoint entry at the real file on disk). + let (first_file, (first_sha, first_n)) = full_files.iter().next().unwrap(); + let entry = full + .collections + .iter() + .flat_map(|c| c.partitions.iter()) + .find(|p| &p.file == first_file) + .unwrap() + .clone(); + let cp_entry = pgstac::export::manifest::CheckpointEntry { + file: first_file.clone(), + sha256: first_sha.clone(), + item_count: *first_n, + collection_id: full.collections[0].id.clone(), + entry, + }; + // Copy that one parquet into a fresh dump dir so resume can verify it. + let dir2 = tempfile::tempdir().unwrap(); + let src = dir1.path().join(first_file); + let dst = dir2.path().join(first_file); + std::fs::create_dir_all(dst.parent().unwrap()).unwrap(); + std::fs::copy(&src, &dst).unwrap(); + + let mut resume = std::collections::HashMap::new(); + let _ = resume.insert(first_file.clone(), cp_entry); + + let report = DumpPlanner::new(DumpConfig { + collection_ids: only.clone(), + resume_completed: resume, + ..Default::default() + }) + .run(&client, &DirSink::new(dir2.path()).unwrap()) + .await + .unwrap(); + + let resumed = read_manifest(dir2.path()); + let resumed_files = file_map(&resumed); + // Same set of partition files + identical per-file counts as the full dump. + assert_eq!( + full_files.keys().collect::>(), + resumed_files.keys().collect::>(), + "resumed manifest is missing partitions" + ); + for (f, (_sha, n)) in &full_files { + assert_eq!(&resumed_files.get(f).unwrap().1, n, "count differs for {f}"); + } + assert_eq!( + report.items, + full_files.values().map(|(_, n)| n).sum::() + ); + // The carried-over file's bytes match the original (it was not re-dumped). + let carried = std::fs::read(dir2.path().join(first_file)).unwrap(); + assert_eq!( + pgstac::export::manifest::sha256_hex(&carried), + *first_sha, + "carried-over partition file changed" + ); + eprintln!( + "[followups] resume: carried {} (skipped re-dump), manifest complete with {} files", + first_file, + resumed_files.len() + ); +} diff --git a/src/pgstac-rs/tests/fragment_parity.rs b/src/pgstac-rs/tests/fragment_parity.rs new file mode 100644 index 00000000..342ef4ad --- /dev/null +++ b/src/pgstac-rs/tests/fragment_parity.rs @@ -0,0 +1,176 @@ +//! Parity gate: the Rust `extract_fragment` / `strip_fragment_col` must match the SQL functions of the +//! same name (003a_items.sql), so the Rust loader splits fragment-config collections exactly like +//! `items_staging_dehydrate`. Connects read-only to a pgstac database (`PGSTAC_RS_TEST_DB`). + +use pgstac::fragment::{ + FragmentConfig, build_fragment_payload, extract_fragment, strip_fragment_col, strip_link_hrefs, +}; +use serde_json::{Value, json}; +use tokio_postgres::{Client, NoTls}; + +fn dsn() -> String { + std::env::var("PGSTAC_RS_TEST_DB") + .unwrap_or_else(|_| "postgresql://username:password@localhost:5439/postgis".to_string()) +} + +async fn connect() -> Client { + let (client, connection) = tokio_postgres::connect(&dsn(), NoTls).await.unwrap(); + tokio::spawn(async move { + let _ = connection.await; + }); + client + .batch_execute("SET search_path TO pgstac, public;") + .await + .unwrap(); + client +} + +/// A representative item with nested asset metadata + properties to fragment. +fn content() -> Value { + json!({ + "type": "Feature", + "stac_version": "1.0.0", + "id": "frag-item", + "collection": "c1", + "geometry": {"type": "Point", "coordinates": [0, 0]}, + "properties": { + "datetime": "2023-01-07T00:00:00Z", + "platform": "landsat-8", + "constellation": "landsat" + }, + "assets": { + "B1": { + "href": "https://x/B1.tif", + "type": "image/tiff; application=geotiff", + "title": "Coastal/Aerosol", + "roles": ["data"], + "eo:bands": [{"name": "B1", "common_name": "coastal"}] + }, + "thumbnail": { + "href": "https://x/t.png", + "type": "image/png", + "roles": ["thumbnail"] + } + } + }) +} + +/// Fragment configs covering depth-2, depth-3, a whole-column path, and a top-level scalar. +fn configs() -> Vec> { + vec![ + vec![ + r#"["assets","B1","type"]"#.into(), + r#"["assets","B1","roles"]"#.into(), + r#"["assets","B1","eo:bands"]"#.into(), + r#"["assets","thumbnail","type"]"#.into(), + ], + vec![ + r#"["properties","platform"]"#.into(), + r#"["properties","constellation"]"#.into(), + ], + vec![r#"["assets","thumbnail"]"#.into()], + vec![ + r#"["assets","B1","title"]"#.into(), + r#"["properties","missing-key"]"#.into(), + ], + ] +} + +#[tokio::test] +async fn extract_and_strip_match_sql() { + let client = connect().await; + let item = content(); + + for config in configs() { + let rust_config = FragmentConfig::parse(&config).unwrap(); + + // extract_fragment(content, config) + let sql_extract: Option = client + .query_one( + "SELECT extract_fragment($1::jsonb, $2::text[])", + &[&item, &config], + ) + .await + .unwrap() + .get(0); + let rust_extract = extract_fragment(&item, &rust_config); + assert_eq!( + rust_extract, sql_extract, + "extract_fragment mismatch for {config:?}" + ); + + // strip_fragment_col(col, col_name, config) for assets + properties + for col_name in ["assets", "properties"] { + let col = item[col_name].clone(); + let sql_strip: Value = client + .query_one( + "SELECT strip_fragment_col($1::jsonb, $2::text, $3::text[])", + &[&col, &col_name, &config], + ) + .await + .unwrap() + .get(0); + let rust_strip = strip_fragment_col(col.clone(), col_name, &rust_config); + assert_eq!( + rust_strip, sql_strip, + "strip_fragment_col({col_name}) mismatch for {config:?}" + ); + } + } +} + +#[tokio::test] +async fn strip_link_hrefs_matches_sql() { + let client = connect().await; + let cases = vec![ + json!([{"rel": "self", "href": "https://x/1"}, {"rel": "root", "href": "https://x/"}]), + json!([]), + json!([{"rel": "license", "title": "L"}]), + json!(null), + json!([{"rel": "self", "href": "https://x/1"}, "not-an-object"]), + ]; + for links in cases { + let sql: Option = client + .query_one("SELECT stac_links_strip_hrefs($1::jsonb)", &[&links]) + .await + .unwrap() + .get(0); + let rust = strip_link_hrefs(&links); + assert_eq!(rust, sql, "strip_link_hrefs mismatch for {links:?}"); + } +} + +#[tokio::test] +async fn build_fragment_payload_matches_sql() { + let client = connect().await; + // content() has assets+properties (no links); add a links-bearing item to exercise links_template. + let mut with_links = content(); + with_links["links"] = json!([ + {"rel": "self", "href": "https://x/frag-item"}, + {"rel": "license", "title": "L"} + ]); + let items = vec![content(), with_links]; + + for item in &items { + for config in configs() { + let rust = build_fragment_payload(item, &FragmentConfig::parse(&config).unwrap()) + // SQL jsonb_strip_nulls yields {} (not NULL) when there is no fragment; map None -> {}. + .unwrap_or_else(|| json!({})); + let sql: Value = client + .query_one( + "SELECT jsonb_strip_nulls(jsonb_build_object(\ + 'content', NULLIF(extract_fragment($1::jsonb, $2::text[]), '{}'::jsonb), \ + 'links_template', stac_links_strip_hrefs($1::jsonb -> 'links')))", + &[item, &config], + ) + .await + .unwrap() + .get(0); + assert_eq!( + rust, sql, + "payload mismatch for id={} config={config:?}", + item["id"] + ); + } + } +} diff --git a/src/pgstac-rs/tests/hydrate_parity.rs b/src/pgstac-rs/tests/hydrate_parity.rs new file mode 100644 index 00000000..58e902db --- /dev/null +++ b/src/pgstac-rs/tests/hydrate_parity.rs @@ -0,0 +1,182 @@ +//! Hydration parity gate (ARCHITECTURE A2, §8). +//! +//! For sampled items, asserts the Rust `Hydrator` output equals the SQL +//! `content_hydrate(i)` output, on **both** a 0.9.11 (base_item) and a 0.10 +//! (fragment) instance. This is the top-risk component and the gate for Phase 1. +//! +//! These tests connect to disposable clones. By default they target: +//! +//! * 0.9.11: `PGSTAC_PARITY_DB_0911` (default `...:5439/a_parity0911`) +//! * 0.10: `PGSTAC_PARITY_DB_010` (default `...:5439/a_parity010`) +//! +//! Each test skips (passes with a logged note) if its DB is unreachable, so the +//! suite stays green on machines without the fixtures; CI/overnight runs point +//! the env vars at real clones. +//! +//! Comparison is **semantic** (`serde_json::Value` equality is order-independent +//! for objects), which is the correct criterion: jsonb has no inherent key order +//! and `content_hydrate` makes no order guarantee. + +#![allow(unused_crate_dependencies)] + +use pgstac::hydrate::{HydrationModel, Hydrator}; +use pgstac::source::{detect_hydration_model, load_collection_context}; +use pgstac::source::{fetch_dehydrated_limited, fetch_fragments}; +use serde_json::Value; +use std::collections::HashMap; +use tokio_postgres::{Client, NoTls}; + +async fn connect(env: &str, default: &str) -> Option { + let conn = std::env::var(env).unwrap_or_else(|_| default.to_string()); + match tokio_postgres::connect(&conn, NoTls).await { + Ok((client, connection)) => { + tokio::spawn(async move { + let _ = connection.await; + }); + if client + .batch_execute("SET search_path = pgstac, public;") + .await + .is_err() + { + return None; + } + Some(client) + } + Err(e) => { + eprintln!("[parity] skipping: cannot connect via {env} ({default}): {e}"); + None + } + } +} + +/// Lists collection ids in the instance. +async fn list_collections(client: &Client) -> Vec { + let rows = client + .query("SELECT id FROM collections ORDER BY id", &[]) + .await + .expect("list collections"); + rows.iter().map(|r| r.get::<_, String>("id")).collect() +} + +/// Fetches the SQL `content_hydrate(i)` output for each item id in a collection, +/// limited to `sample_limit` ordered `(datetime, id)`. +async fn sql_hydrated( + client: &Client, + collection: &str, + sample_limit: i64, +) -> HashMap { + let rows = client + .query( + "SELECT id, content_hydrate(i) AS h \ + FROM items i WHERE collection = $1 \ + ORDER BY datetime, id LIMIT $2", + &[&collection, &sample_limit], + ) + .await + .expect("sql content_hydrate"); + rows.iter() + .map(|r| (r.get::<_, String>("id"), r.get::<_, Value>("h"))) + .collect() +} + +/// Runs the parity comparison for one instance. +async fn run_parity(client: &Client, sample_limit: i64) { + let version: String = client + .query_one("SELECT get_version() AS v", &[]) + .await + .expect("get_version") + .get("v"); + let model = detect_hydration_model(client).await.expect("detect model"); + let hydrator = Hydrator::new(model); + eprintln!("[parity] instance version {version} -> model {model:?}"); + + let collections = list_collections(client).await; + assert!(!collections.is_empty(), "instance has no collections"); + + let mut total = 0usize; + for collection in &collections { + let ctx = load_collection_context(client, model, collection) + .await + .expect("load context"); + + let where_clause = "collection = $1"; + let items = fetch_dehydrated_limited( + client, + model, + "items", + Some(where_clause), + &[&collection], + Some(sample_limit), + ) + .await + .expect("fetch dehydrated"); + if items.is_empty() { + continue; + } + + // Load fragments for any items that reference one (0.10). + let frag_ids: Vec = items.iter().filter_map(|i| i.fragment_id).collect(); + let fragments = fetch_fragments(client, &frag_ids) + .await + .expect("fetch fragments"); + + let sql = sql_hydrated(client, collection, sample_limit).await; + + for item in &items { + let fragment = item.fragment_id.and_then(|id| fragments.get(&id)); + let rust = hydrator.hydrate(item.clone(), &ctx, fragment); + let expected = sql + .get(&item.id) + .unwrap_or_else(|| panic!("no SQL hydrate for id={}", item.id)); + assert_eq!( + &rust, + expected, + "hydration mismatch for id={} collection={}\n--- rust ---\n{}\n--- sql ---\n{}", + item.id, + item.collection, + serde_json::to_string_pretty(&rust).unwrap(), + serde_json::to_string_pretty(expected).unwrap(), + ); + total += 1; + } + eprintln!("[parity] {collection}: {} items matched", items.len()); + } + assert!(total > 0, "compared zero items"); + eprintln!("[parity] {model:?}: total {total} items, all match"); +} + +#[tokio::test] +async fn parity_base_item_0911() { + let Some(client) = connect( + "PGSTAC_PARITY_DB_0911", + "postgresql://username:password@localhost:5439/pgstac_v0911_bench", + ) + .await + else { + return; + }; + assert_eq!( + detect_hydration_model(&client).await.unwrap(), + HydrationModel::BaseItem, + "0911 fixture should be base_item model" + ); + run_parity(&client, 2000).await; +} + +#[tokio::test] +async fn parity_fragment_010() { + let Some(client) = connect( + "PGSTAC_PARITY_DB_010", + "postgresql://username:password@localhost:5439/pgstac_rs_test_rich", + ) + .await + else { + return; + }; + assert_eq!( + detect_hydration_model(&client).await.unwrap(), + HydrationModel::Fragment, + "010 fixture should be fragment model" + ); + run_parity(&client, 2000).await; +} diff --git a/src/pgstac-rs/tests/ingest_copy.rs b/src/pgstac-rs/tests/ingest_copy.rs new file mode 100644 index 00000000..5071895b --- /dev/null +++ b/src/pgstac-rs/tests/ingest_copy.rs @@ -0,0 +1,88 @@ +//! Round-trip gate for the binary-COPY loader: dehydrate an item, binary-COPY it into a `LIKE items` +//! staging table, and assert the stored row equals SQL `content_dehydrate` of the same item (every +//! column except `pgstac_updated_at`, which is a load-time `now()`). Validates the per-column binary +//! encoding — geometry EWKB, promoted typed columns, jsonb, `text[]` — against the database. +//! +//! Connects to a pgstac database (`PGSTAC_RS_TEST_DB`, default the dev `postgis` db). Uses a manual +//! `CREATE TEMP TABLE … (LIKE items)` rather than the SECURITY DEFINER `make_binary_staging`, so it runs on +//! any pgstac install (the SECURITY DEFINER flush + the revoked direct writes are covered by the SQL gate). + +use pgstac::dehydrate::{DehydrateSchema, dehydrate}; +use pgstac::ingest::copy_rows; +use serde_json::{Value, json}; +use tokio_postgres::{Client, NoTls}; + +fn dsn() -> String { + std::env::var("PGSTAC_RS_TEST_DB") + .unwrap_or_else(|_| "postgresql://username:password@localhost:5439/postgis".to_string()) +} + +async fn connect() -> Client { + let (client, connection) = tokio_postgres::connect(&dsn(), NoTls).await.unwrap(); + tokio::spawn(async move { + let _ = connection.await; + }); + client + .batch_execute("SET search_path TO pgstac, public;") + .await + .unwrap(); + client +} + +fn item() -> Value { + json!({ + "type": "Feature", + "stac_version": "1.0.0", + "stac_extensions": ["https://stac-extensions.github.io/eo/v1.1.0/schema.json"], + "id": "copy-item", + "collection": "c1", + "geometry": {"type": "Point", "coordinates": [-105.1019, 40.1672]}, + "bbox": [-105.1019, 40.1672, -105.1019, 40.1672], + "properties": { + "datetime": "2023-01-07T00:00:00Z", + "platform": "landsat-8", + "instruments": ["oli", "tirs"], + "eo:cloud_cover": 12.5, + "gsd": 30.0, + "sat:relative_orbit": 42, + "created": "2023-01-08T01:02:03Z", + "proj:bbox": [1.0, 2.0, 3.0, 4.0], + "custom:keep": "stays" + }, + "assets": {"thumbnail": {"href": "https://x/t.png", "type": "image/png", "roles": ["thumbnail"]}}, + "links": [{"rel": "self", "href": "https://x/copy-item"}], + "extra_top": {"a": 1} + }) +} + +#[tokio::test] +async fn copy_round_trips_against_content_dehydrate() { + let mut client = connect().await; + let schema = DehydrateSchema::load(&client).await.unwrap(); + + // copy_rows shares the transaction that owns the (ON COMMIT DROP) staging table. + let tx = client.transaction().await.unwrap(); + tx.batch_execute("CREATE TEMP TABLE _ingest_copy_staging (LIKE pgstac.items) ON COMMIT DROP") + .await + .unwrap(); + + let it = item(); + let row = dehydrate(it.clone(), &schema).unwrap(); + let n = copy_rows(&tx, "_ingest_copy_staging", vec![row], &schema) + .await + .unwrap(); + assert_eq!(n, 1, "one row copied"); + + // The staged row must equal content_dehydrate of the same item, ignoring the load-time timestamp. + let matches: bool = tx + .query_one( + "SELECT (to_jsonb(s) - 'pgstac_updated_at') = (to_jsonb(d) - 'pgstac_updated_at') \ + FROM _ingest_copy_staging s, content_dehydrate($1::jsonb) d", + &[&it], + ) + .await + .unwrap() + .get(0); + assert!(matches, "staged row must equal content_dehydrate output"); + tx.rollback().await.unwrap(); +} diff --git a/src/pgstac-rs/tests/ingest_load.rs b/src/pgstac-rs/tests/ingest_load.rs new file mode 100644 index 00000000..bd42681f --- /dev/null +++ b/src/pgstac-rs/tests/ingest_load.rs @@ -0,0 +1,215 @@ +//! End-to-end gate for the Rust ingest pipeline: `load_items` (dehydrate -> check_partition -> +//! make_binary_staging -> binary COPY -> flush_items_staging_binary) stores items that equal SQL +//! `content_dehydrate`, and a `pgstac_ingest` role can ingest only through the SECURITY DEFINER path +//! (direct item writes are denied). +//! +//! Needs a database with THIS branch's schema (the SECURITY DEFINER functions + revoked direct writes), so it clones a dedicated +//! template `pgstac_rs_ingest_template` (build with: +//! `psql ... -c 'CREATE DATABASE pgstac_rs_ingest_template'; psql ... pgstac_rs_ingest_template -f src/pgstac/pgstac.sql`). +//! Override the template/base via `PGSTAC_RS_INGEST_TEMPLATE` / `PGSTAC_RS_TEST_BASE`. + +use pgstac::dehydrate::DehydrateSchema; +use pgstac::ingest::{ConflictPolicy, load_items}; +use serde_json::{Value, json}; +use std::sync::atomic::{AtomicU32, Ordering}; +use tokio_postgres::{Client, NoTls}; + +fn base() -> String { + std::env::var("PGSTAC_RS_TEST_BASE") + .unwrap_or_else(|_| "postgresql://username:password@localhost:5439".to_string()) +} + +fn template() -> String { + std::env::var("PGSTAC_RS_INGEST_TEMPLATE") + .unwrap_or_else(|_| "pgstac_rs_ingest_template".to_string()) +} + +/// A disposable database cloned from the ingest template, dropped on `Drop`. +struct CloneDb { + name: String, +} + +impl CloneDb { + async fn create() -> CloneDb { + static COUNTER: AtomicU32 = AtomicU32::new(0); + let name = format!( + "pgstac_rs_ingest_test_{}_{}", + std::process::id(), + COUNTER.fetch_add(1, Ordering::Relaxed) + ); + let (client, connection) = tokio_postgres::connect(&format!("{}/postgres", base()), NoTls) + .await + .unwrap(); + let handle = tokio::spawn(connection); + let _ = client + .execute( + &format!("CREATE DATABASE {name} TEMPLATE {}", template()), + &[], + ) + .await + .unwrap(); + handle.abort(); + CloneDb { name } + } + + fn dsn(&self) -> String { + format!("{}/{}", base(), self.name) + } +} + +impl Drop for CloneDb { + fn drop(&mut self) { + let name = self.name.clone(); + std::thread::scope(|scope| { + let _ = scope.spawn(|| { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + runtime.block_on(async move { + let (client, connection) = + tokio_postgres::connect(&format!("{}/postgres", base()), NoTls) + .await + .unwrap(); + let handle = tokio::spawn(connection); + let _ = client + .execute( + "SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname = $1", + &[&name], + ) + .await; + let _ = client + .execute(&format!("DROP DATABASE IF EXISTS {name}"), &[]) + .await; + handle.abort(); + }); + }); + }); + } +} + +async fn connect(dsn: &str) -> Client { + let (client, connection) = tokio_postgres::connect(dsn, NoTls).await.unwrap(); + tokio::spawn(async move { + let _ = connection.await; + }); + client + .batch_execute("SET search_path TO pgstac, public;") + .await + .unwrap(); + client +} + +fn collection() -> Value { + json!({ + "id": "c1", + "type": "Collection", + "stac_version": "1.0.0", + "description": "test", + "license": "proprietary", + "extent": {"spatial": {"bbox": [[-180, -90, 180, 90]]}, "temporal": {"interval": [[null, null]]}}, + "links": [] + }) +} + +fn item(id: &str, datetime: &str) -> Value { + json!({ + "type": "Feature", + "stac_version": "1.0.0", + "id": id, + "collection": "c1", + "geometry": {"type": "Point", "coordinates": [-105.1019, 40.1672]}, + "bbox": [-105.1019, 40.1672, -105.1019, 40.1672], + "properties": {"datetime": datetime, "platform": "landsat-8", "eo:cloud_cover": 12.5}, + "assets": {"thumbnail": {"href": "https://x/t.png", "type": "image/png"}} + }) +} + +/// Each stored item must equal content_dehydrate of the same item (ignoring the load-time timestamp). +async fn assert_stored_equals_content_dehydrate(client: &Client, items: &[Value]) { + for it in items { + let id = it["id"].as_str().unwrap(); + let coll = it["collection"].as_str().unwrap(); + let matches: bool = client + .query_one( + "SELECT (to_jsonb(i) - 'pgstac_updated_at') = (to_jsonb(d) - 'pgstac_updated_at') \ + FROM items i, content_dehydrate($1::jsonb) d WHERE i.id = $2 AND i.collection = $3", + &[it, &id, &coll], + ) + .await + .unwrap() + .get(0); + assert!( + matches, + "stored item {id} must equal content_dehydrate output" + ); + } +} + +#[tokio::test] +async fn load_items_stores_items_matching_content_dehydrate() { + let db = CloneDb::create().await; + let mut client = connect(&db.dsn()).await; + let schema = DehydrateSchema::load(&client).await.unwrap(); + + let _ = client + .execute("SELECT create_collection($1::jsonb)", &[&collection()]) + .await + .unwrap(); + + let items = vec![ + item("a", "2023-01-07T00:00:00Z"), + item("b", "2023-02-15T00:00:00Z"), + ]; + let n = load_items(&mut client, items.clone(), &schema, ConflictPolicy::Ignore) + .await + .unwrap(); + assert_eq!(n, 2, "two items flushed"); + + let count: i64 = client + .query_one("SELECT count(*) FROM items WHERE collection = 'c1'", &[]) + .await + .unwrap() + .get(0); + assert_eq!(count, 2, "two items stored"); + + assert_stored_equals_content_dehydrate(&client, &items).await; +} + +#[tokio::test] +async fn load_items_works_as_pgstac_ingest_under_the_wall() { + let db = CloneDb::create().await; + let mut client = connect(&db.dsn()).await; + let schema = DehydrateSchema::load(&client).await.unwrap(); + let _ = client + .execute("SELECT create_collection($1::jsonb)", &[&collection()]) + .await + .unwrap(); + + // Act as a role inheriting pgstac_ingest: direct item writes must be denied... + client + .batch_execute("SET ROLE pgstac_ingest;") + .await + .unwrap(); + let direct = client + .execute( + "INSERT INTO items (id, collection, geometry, datetime, end_datetime, datetime_is_range, item_hash) \ + VALUES ('x', 'c1', ST_SetSRID(ST_MakePoint(0,0),4326), now(), now(), false, '\\x00')", + &[], + ) + .await; + assert!( + direct.is_err(), + "direct INSERT INTO items must be denied for pgstac_ingest" + ); + + // ...but the Rust pipeline (via the SECURITY DEFINER flush) ingests fine. + let items = vec![item("a", "2023-01-07T00:00:00Z")]; + let n = load_items(&mut client, items.clone(), &schema, ConflictPolicy::Ignore) + .await + .unwrap(); + assert_eq!(n, 1, "one item flushed via the SECURITY DEFINER path"); + + client.batch_execute("RESET ROLE;").await.unwrap(); + assert_stored_equals_content_dehydrate(&client, &items).await; +} diff --git a/src/pgstac-rs/tests/pool_ingest.rs b/src/pgstac-rs/tests/pool_ingest.rs new file mode 100644 index 00000000..281038c2 --- /dev/null +++ b/src/pgstac-rs/tests/pool_ingest.rs @@ -0,0 +1,813 @@ +#![cfg(feature = "pool")] +//! End-to-end gate for [`PgstacPool`] write methods: create_collection -> create_items -> get_item +//! round-trip, the `Error` conflict policy rejects a duplicate id, and delete_item removes an item. +//! Runs against a clone of this branch's ingest template (see tests/ingest_load.rs for how to build it). + +use pgstac::ingest::ConflictPolicy; +use pgstac::{ConnectConfig, PgstacPool}; +use serde_json::{Value, json}; +use std::sync::atomic::{AtomicU32, Ordering}; +use tokio_postgres::NoTls; + +fn base() -> String { + std::env::var("PGSTAC_RS_TEST_BASE") + .unwrap_or_else(|_| "postgresql://username:password@localhost:5439".to_string()) +} + +fn template() -> String { + std::env::var("PGSTAC_RS_INGEST_TEMPLATE") + .unwrap_or_else(|_| "pgstac_rs_ingest_template".to_string()) +} + +/// A disposable database cloned from the ingest template, dropped on `Drop`. +struct CloneDb { + name: String, +} + +impl CloneDb { + async fn create() -> CloneDb { + static COUNTER: AtomicU32 = AtomicU32::new(0); + let name = format!( + "pgstac_rs_pool_ingest_test_{}_{}", + std::process::id(), + COUNTER.fetch_add(1, Ordering::Relaxed) + ); + let (client, connection) = tokio_postgres::connect(&format!("{}/postgres", base()), NoTls) + .await + .unwrap(); + let handle = tokio::spawn(connection); + let _ = client + .execute( + &format!("CREATE DATABASE {name} TEMPLATE {}", template()), + &[], + ) + .await + .unwrap(); + handle.abort(); + CloneDb { name } + } + + fn dsn(&self) -> String { + format!("{}/{}", base(), self.name) + } +} + +impl Drop for CloneDb { + fn drop(&mut self) { + let name = self.name.clone(); + std::thread::scope(|scope| { + let _ = scope.spawn(|| { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + runtime.block_on(async move { + let (client, connection) = + tokio_postgres::connect(&format!("{}/postgres", base()), NoTls) + .await + .unwrap(); + let handle = tokio::spawn(connection); + let _ = client + .execute( + "SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname = $1", + &[&name], + ) + .await; + let _ = client + .execute(&format!("DROP DATABASE IF EXISTS {name}"), &[]) + .await; + handle.abort(); + }); + }); + }); + } +} + +fn collection() -> Value { + json!({ + "id": "c1", + "type": "Collection", + "stac_version": "1.0.0", + "description": "test", + "license": "proprietary", + "extent": {"spatial": {"bbox": [[-180, -90, 180, 90]]}, "temporal": {"interval": [[null, null]]}}, + "links": [] + }) +} + +fn item(id: &str, datetime: &str) -> Value { + json!({ + "type": "Feature", + "stac_version": "1.0.0", + "id": id, + "collection": "c1", + "geometry": {"type": "Point", "coordinates": [-105.1019, 40.1672]}, + "bbox": [-105.1019, 40.1672, -105.1019, 40.1672], + "properties": {"datetime": datetime, "platform": "landsat-8", "eo:cloud_cover": 12.5}, + "assets": {"thumbnail": {"href": "https://x/t.png", "type": "image/png"}} + }) +} + +async fn pool(db: &CloneDb) -> PgstacPool { + PgstacPool::connect(ConnectConfig { + dsn: Some(db.dsn()), + ..Default::default() + }) + .await + .unwrap() +} + +#[tokio::test] +async fn pool_create_items_then_get_item() { + let db = CloneDb::create().await; + let pool = pool(&db).await; + pool.create_collection(&collection()).await.unwrap(); + + let n = pool + .create_items( + vec![ + item("a", "2023-01-07T00:00:00Z"), + item("b", "2023-02-15T00:00:00Z"), + ], + ConflictPolicy::Ignore, + ) + .await + .unwrap(); + assert_eq!(n, 2, "two items created"); + + let got = pool.get_item("c1", "a").await.unwrap(); + assert_eq!( + got.as_ref().map(|i| &i["id"]), + Some(&json!("a")), + "get_item round-trips the created item" + ); + + pool.close(); +} + +#[tokio::test] +async fn pool_search_finds_sparse_items_across_a_long_span() { + // Regression: 3 items years apart in one (NULL-trunc) partition. partition_stats.n is tiny relative + // to the month span, so partition_bounds prorates every month's count to 0. The keyset search must + // still return all 3 — bands must cover the whole range, never be dropped by (estimated) count. + let db = CloneDb::create().await; + let pool = pool(&db).await; + pool.create_collection(&collection()).await.unwrap(); + let items = vec![ + item("s2010", "2010-06-01T00:00:00Z"), + item("s2015", "2015-06-01T00:00:00Z"), + item("s2023", "2023-06-01T00:00:00Z"), + ]; + let n = pool + .create_items(items, ConflictPolicy::Ignore) + .await + .unwrap(); + assert_eq!(n, 3, "three items created"); + + let page = pool + .search_collect(&json!({"collections": ["c1"], "limit": 100}), None) + .await + .unwrap(); + assert_eq!( + page.features.len(), + 3, + "search must return all sparse items, none dropped" + ); + pool.close(); +} + +#[tokio::test] +async fn pool_create_item_rejects_duplicate_then_deletes() { + let db = CloneDb::create().await; + let pool = pool(&db).await; + pool.create_collection(&collection()).await.unwrap(); + + assert_eq!( + pool.create_item(item("a", "2023-01-07T00:00:00Z")) + .await + .unwrap(), + 1 + ); + + // The Error policy must reject a second item with the same id. + let duplicate = pool.create_item(item("a", "2023-01-07T00:00:00Z")).await; + assert!( + duplicate.is_err(), + "create_item must error on an existing id" + ); + + // delete_item removes it. + pool.delete_item("c1", "a").await.unwrap(); + assert!( + pool.get_item("c1", "a").await.unwrap().is_none(), + "item deleted" + ); + + pool.close(); +} + +/// A collection with `item_assets` so `create_collection` derives a non-empty `fragment_config`. +fn frag_collection(id: &str) -> Value { + json!({ + "id": id, + "type": "Collection", + "stac_version": "1.0.0", + "description": "fragment test", + "license": "proprietary", + "extent": {"spatial": {"bbox": [[-180, -90, 180, 90]]}, "temporal": {"interval": [[null, null]]}}, + "links": [], + "item_assets": { + "B1": {"type": "image/tiff; application=geotiff", "title": "Coastal", "roles": ["data"], "eo:bands": [{"name": "B1", "common_name": "coastal"}]}, + "thumbnail": {"type": "image/png", "roles": ["thumbnail"]} + } + }) +} + +fn frag_item(collection: &str, id: &str) -> Value { + json!({ + "type": "Feature", + "stac_version": "1.0.0", + "id": id, + "collection": collection, + "geometry": {"type": "Point", "coordinates": [-105.1, 40.1]}, + "bbox": [-105.1, 40.1, -105.1, 40.1], + "properties": {"datetime": "2023-01-07T00:00:00Z", "platform": "landsat-8"}, + "assets": { + "B1": {"href": format!("https://x/{id}/B1.tif"), "type": "image/tiff; application=geotiff", "title": "Coastal", "roles": ["data"], "eo:bands": [{"name": "B1", "common_name": "coastal"}]}, + "thumbnail": {"href": format!("https://x/{id}/t.png"), "type": "image/png", "roles": ["thumbnail"]} + }, + "links": [{"rel": "self", "href": format!("https://x/{id}")}] + }) +} + +/// The Rust loader's fragment split must produce the SAME stored rows + fragments as the SQL +/// items_staging path. Load identical items (modulo collection) via both paths into two fragment-config +/// collections and assert the per-item columns + the fragment they point to are byte-identical. +#[tokio::test] +async fn fragment_load_matches_sql_items_staging() { + let db = CloneDb::create().await; + let pool = pool(&db).await; + pool.create_collection(&frag_collection("c_rust")) + .await + .unwrap(); + pool.create_collection(&frag_collection("c_sql")) + .await + .unwrap(); + + let (client, connection) = tokio_postgres::connect(&db.dsn(), NoTls).await.unwrap(); + tokio::spawn(async move { + let _ = connection.await; + }); + client + .batch_execute("SET search_path TO pgstac, public;") + .await + .unwrap(); + + // Both collections must actually fragment (else the test is vacuous). + let with_config: i64 = client + .query_one( + "SELECT count(*) FROM collections WHERE id IN ('c_rust','c_sql') \ + AND fragment_config IS NOT NULL AND cardinality(fragment_config) > 0", + &[], + ) + .await + .unwrap() + .get(0); + assert_eq!( + with_config, 2, + "both collections need a non-empty fragment_config" + ); + + // Rust loader -> c_rust. + let rust_items: Vec = ["a", "b", "c"] + .iter() + .map(|id| frag_item("c_rust", id)) + .collect(); + pool.create_items(rust_items, ConflictPolicy::Ignore) + .await + .unwrap(); + + // SQL items_staging -> c_sql. + for id in ["a", "b", "c"] { + client + .execute( + "INSERT INTO items_staging (content) VALUES ($1::jsonb)", + &[&frag_item("c_sql", id)], + ) + .await + .unwrap(); + } + + // Rust must have actually created fragments + stamped fragment_id. + let frags: i64 = client + .query_one( + "SELECT count(*) FROM item_fragments WHERE collection='c_rust'", + &[], + ) + .await + .unwrap() + .get(0); + assert!(frags >= 1, "Rust load must create fragments"); + let stamped: i64 = client + .query_one( + "SELECT count(*) FROM items WHERE collection='c_rust' AND fragment_id IS NOT NULL", + &[], + ) + .await + .unwrap() + .get(0); + assert_eq!(stamped, 3, "all Rust items must reference a fragment"); + + // Per-item columns + the fragment they point to must match between the two paths (symmetric diff = 0). + let diff: i64 = client + .query_one( + "WITH cols AS (\ + SELECT i.collection, i.id, i.assets, i.properties, i.stac_version, i.stac_extensions, \ + i.links, i.link_hrefs, f.content AS frag_content, f.links_template \ + FROM items i LEFT JOIN item_fragments f ON i.fragment_id = f.id \ + WHERE i.collection IN ('c_rust','c_sql')), \ + r AS (SELECT id, assets, properties, stac_version, stac_extensions, links, link_hrefs, frag_content, links_template FROM cols WHERE collection='c_rust'), \ + s AS (SELECT id, assets, properties, stac_version, stac_extensions, links, link_hrefs, frag_content, links_template FROM cols WHERE collection='c_sql') \ + SELECT count(*) FROM ((SELECT * FROM r EXCEPT SELECT * FROM s) UNION ALL (SELECT * FROM s EXCEPT SELECT * FROM r)) d", + &[], + ) + .await + .unwrap() + .get(0); + assert_eq!( + diff, 0, + "Rust fragment split must match SQL items_staging row-for-row" + ); + pool.close(); +} + +/// The Rust loader must NOT silently drop an item whose collection doesn't exist (the SQL items_staging +/// path does — see the ignored lib test tests::item_without_collection). ensure_partitions -> +/// check_partition raises for an unknown collection. +#[tokio::test] +async fn create_items_errors_on_absent_collection() { + let db = CloneDb::create().await; + let pool = pool(&db).await; + // No collection created. + let result = pool + .create_item(item("orphan", "2023-01-07T00:00:00Z")) + .await; + assert!( + result.is_err(), + "loading an item for a nonexistent collection must error, not silently drop" + ); + pool.close(); +} + +/// The Upsert policy must rewrite an item only when its content changed (flush's content-aware +/// DELETE-if-distinct + INSERT), leaving an unchanged re-upsert untouched. +#[tokio::test] +async fn upsert_only_rewrites_changed_items() { + let db = CloneDb::create().await; + let pool = pool(&db).await; + pool.create_collection(&collection()).await.unwrap(); + let (client, connection) = tokio_postgres::connect(&db.dsn(), NoTls).await.unwrap(); + tokio::spawn(async move { + let _ = connection.await; + }); + client + .batch_execute("SET search_path TO pgstac, public;") + .await + .unwrap(); + + let v1 = item("u", "2023-01-07T00:00:00Z"); + pool.upsert_item(v1.clone()).await.unwrap(); + let t1: chrono::DateTime = client + .query_one("SELECT pgstac_updated_at FROM items WHERE id='u'", &[]) + .await + .unwrap() + .get(0); + + // Re-upsert identical content -> not rewritten (pgstac_updated_at unchanged). + pool.upsert_item(v1.clone()).await.unwrap(); + let t2: chrono::DateTime = client + .query_one("SELECT pgstac_updated_at FROM items WHERE id='u'", &[]) + .await + .unwrap() + .get(0); + assert_eq!(t1, t2, "unchanged upsert must not rewrite the row"); + + // Upsert changed content -> rewritten with the new value. + let mut v2 = v1; + v2["properties"]["eo:cloud_cover"] = json!(99.0); + pool.upsert_item(v2).await.unwrap(); + let (cc, t3): (f64, chrono::DateTime) = { + let row = client + .query_one( + "SELECT eo_cloud_cover, pgstac_updated_at FROM items WHERE id='u'", + &[], + ) + .await + .unwrap(); + (row.get(0), row.get(1)) + }; + assert_eq!(cc, 99.0, "changed upsert must store the new content"); + assert!( + t3 > t1, + "changed upsert must rewrite (newer pgstac_updated_at)" + ); + pool.close(); +} + +/// After a load, partition_stats covers the data (row count raised, dtrange generous + dirty); +/// tighten_partition_stats then narrows to the EXACT extent and clears dirty. +#[tokio::test] +async fn tighten_narrows_widen_now_stats_to_exact() { + let db = CloneDb::create().await; + let pool = pool(&db).await; + pool.create_collection(&collection()).await.unwrap(); + pool.create_items( + vec![ + item("a", "2023-01-10T00:00:00Z"), + item("b", "2023-01-20T00:00:00Z"), + ], + ConflictPolicy::Ignore, + ) + .await + .unwrap(); + let (client, connection) = tokio_postgres::connect(&db.dsn(), NoTls).await.unwrap(); + tokio::spawn(async move { + let _ = connection.await; + }); + client + .batch_execute("SET search_path TO pgstac, public;") + .await + .unwrap(); + + // Stats now cover the load: row count >= the 2 rows, dtrange COVERS the data span, dirty set. + let row = client + .query_one( + "SELECT n, dtrange @> tstzrange('2023-01-10','2023-01-20','[]') AS covers, dirty \ + FROM partition_stats WHERE collection='c1'", + &[], + ) + .await + .unwrap(); + let n: i64 = row.get("n"); + let covers: bool = row.get("covers"); + let dirty: bool = row.get("dirty"); + assert!(n >= 2, "row count covers the loaded rows"); + assert!(covers, "dtrange must cover the data span"); + assert!( + dirty, + "freshly-loaded partition is dirty (awaiting tighten)" + ); + + // Tighten -> exact extent + exact count + dirty cleared. + client + .execute( + "SELECT tighten_partition_stats(partition) FROM partition_stats WHERE collection='c1'", + &[], + ) + .await + .unwrap(); + let row = client + .query_one( + "SELECT n, dtrange = tstzrange('2023-01-10','2023-01-20','[]') AS exact, dirty \ + FROM partition_stats WHERE collection='c1'", + &[], + ) + .await + .unwrap(); + let n: i64 = row.get("n"); + let exact: bool = row.get("exact"); + let dirty: bool = row.get("dirty"); + assert_eq!(n, 2, "tighten sets the exact row count"); + assert!(exact, "tighten narrows dtrange to the exact data extent"); + assert!(!dirty, "tighten clears dirty"); + pool.close(); +} + +/// High-concurrency single-item ingest: many concurrent create_item +/// calls to the same collection/partition must all succeed (idempotent check_partition + covered-guard +/// widen + ON CONFLICT) and every item must land. +#[tokio::test] +async fn concurrent_single_item_creates_all_land() { + let db = CloneDb::create().await; + let pool = pool(&db).await; + pool.create_collection(&collection()).await.unwrap(); + + const N: usize = 24; + let mut handles = Vec::with_capacity(N); + for i in 0..N { + let pool = pool.clone(); + handles.push(tokio::spawn(async move { + pool.create_item(item(&format!("c{i:03}"), "2023-01-15T00:00:00Z")) + .await + })); + } + let mut ok = 0; + for handle in handles { + handle + .await + .unwrap() + .expect("concurrent create_item must succeed"); + ok += 1; + } + assert_eq!(ok, N, "all concurrent creates succeeded"); + + let (client, connection) = tokio_postgres::connect(&db.dsn(), NoTls).await.unwrap(); + tokio::spawn(async move { + let _ = connection.await; + }); + client + .batch_execute("SET search_path TO pgstac, public;") + .await + .unwrap(); + let count: i64 = client + .query_one("SELECT count(*) FROM items WHERE collection='c1'", &[]) + .await + .unwrap() + .get(0); + assert_eq!(count, N as i64, "every concurrently-created item landed"); + pool.close(); +} + +/// The async tighten sweep clears every dirty partition (oldest first), exercising the loop over a +/// multi-partition (month-truncated) collection. +#[tokio::test] +async fn tighten_sweep_clears_all_dirty_partitions() { + let db = CloneDb::create().await; + let pool = pool(&db).await; + let (client, connection) = tokio_postgres::connect(&db.dsn(), NoTls).await.unwrap(); + tokio::spawn(async move { + let _ = connection.await; + }); + client + .batch_execute("SET search_path TO pgstac, public;") + .await + .unwrap(); + // Month-truncated collection -> items in different months land in different partitions. + client + .execute( + "SELECT create_collection($1::jsonb, 'month')", + &[&collection()], + ) + .await + .unwrap(); + + pool.create_items( + vec![ + item("a", "2023-01-10T00:00:00Z"), + item("b", "2023-02-20T00:00:00Z"), + item("c", "2023-03-05T00:00:00Z"), + ], + ConflictPolicy::Ignore, + ) + .await + .unwrap(); + + let dirty_before: i64 = client + .query_one("SELECT count(*) FROM partition_stats WHERE dirty", &[]) + .await + .unwrap() + .get(0); + assert!( + dirty_before >= 3, + "each month partition is dirty after load (got {dirty_before})" + ); + + let tightened: i32 = client + .query_one("SELECT tighten_dirty_partition_stats(NULL)", &[]) + .await + .unwrap() + .get(0); + assert_eq!( + i64::from(tightened), + dirty_before, + "sweep tightens every dirty partition" + ); + + let dirty_after: i64 = client + .query_one("SELECT count(*) FROM partition_stats WHERE dirty", &[]) + .await + .unwrap() + .get(0); + assert_eq!(dirty_after, 0, "no dirty partitions remain after the sweep"); + pool.close(); +} + +/// The typed `stac::api::TransactionClient` routes writes through the Rust loader (not SQL +/// `create_item`): a pooled `Client` (`pool.client()`) lands items via `add_items`, a `Client` +/// lands one via `add_item`, and `add_item` to an absent collection errors (the loader behavior, vs the +/// legacy SQL silent-drop). +#[tokio::test] +async fn transaction_client_routes_through_loader() { + use stac::api::TransactionClient; + + let db = CloneDb::create().await; + let pool = pool(&db).await; + pool.create_collection(&collection()).await.unwrap(); + + // A pooled Client as TransactionClient: typed add_items -> loader. + let items: Vec = vec![ + serde_json::from_value(item("ta", "2023-01-05T00:00:00Z")).unwrap(), + serde_json::from_value(item("tb", "2023-02-05T00:00:00Z")).unwrap(), + ]; + let mut pooled = pool.client().await.unwrap(); + TransactionClient::add_items(&mut pooled, items) + .await + .unwrap(); + + let (client, connection) = tokio_postgres::connect(&db.dsn(), NoTls).await.unwrap(); + tokio::spawn(async move { + let _ = connection.await; + }); + client + .batch_execute("SET search_path TO pgstac, public;") + .await + .unwrap(); + let n: i64 = client + .query_one("SELECT count(*) FROM items WHERE collection='c1'", &[]) + .await + .unwrap() + .get(0); + assert_eq!(n, 2, "typed add_items landed both items through the loader"); + + // Client as TransactionClient: typed add_item -> loader. + let (raw, conn) = tokio_postgres::connect(&db.dsn(), NoTls).await.unwrap(); + tokio::spawn(async move { + let _ = conn.await; + }); + raw.batch_execute("SET search_path TO pgstac, public;") + .await + .unwrap(); + let mut wrapped = pgstac::Client::new(raw); + let third: stac::Item = serde_json::from_value(item("tc", "2023-04-05T00:00:00Z")).unwrap(); + TransactionClient::add_item(&mut wrapped, third) + .await + .unwrap(); + let n: i64 = client + .query_one("SELECT count(*) FROM items WHERE collection='c1'", &[]) + .await + .unwrap() + .get(0); + assert_eq!(n, 3, "Client add_item landed via the loader"); + + // add_item to an absent collection errors (loader), not a silent drop. + let mut orphan: stac::Item = + serde_json::from_value(item("orphan", "2023-03-05T00:00:00Z")).unwrap(); + orphan.collection = Some("does-not-exist".to_string()); + assert!( + TransactionClient::add_item(&mut pooled, orphan) + .await + .is_err(), + "add_item errors on an absent collection" + ); + + pool.close(); +} + +/// The loader widens `item_field_registry` on every load: after a load the registry must be populated +/// and a SUPERSET of every field path in the loaded items (cross-checked against SQL `jsonb_field_rows`). +/// Guards against the widen call being silently dropped, which would leave the registry empty. +#[tokio::test] +async fn load_widens_field_registry() { + let db = CloneDb::create().await; + let pool = pool(&db).await; + pool.create_collection(&collection()).await.unwrap(); + pool.create_items( + vec![ + item("fr-a", "2023-01-05T00:00:00Z"), + item("fr-b", "2023-02-05T00:00:00Z"), + ], + ConflictPolicy::Error, + ) + .await + .unwrap(); + + let (client, connection) = tokio_postgres::connect(&db.dsn(), NoTls).await.unwrap(); + tokio::spawn(async move { + let _ = connection.await; + }); + client + .batch_execute("SET search_path TO pgstac, public;") + .await + .unwrap(); + + let paths: i64 = client + .query_one( + "SELECT count(*) FROM item_field_registry WHERE collection='c1'", + &[], + ) + .await + .unwrap() + .get(0); + assert!(paths > 0, "the load populated the field registry"); + + // Every path the items actually carried is in the registry (superset). `content_hydrate` injects an + // empty top-level `links` array that the loaded items didn't have — a hydrate artifact, not a loaded data + // field — so exclude it; the registry tracks what was loaded. + let missing: i64 = client + .query_one( + "WITH sql_paths AS (SELECT DISTINCT r.path FROM items i \ + CROSS JOIN LATERAL jsonb_field_rows(content_hydrate(i)) r \ + WHERE i.collection='c1' AND r.path <> 'links') \ + SELECT count(*) FROM sql_paths s \ + WHERE NOT EXISTS (SELECT 1 FROM item_field_registry reg WHERE reg.collection='c1' AND reg.path=s.path)", + &[], + ) + .await + .unwrap() + .get(0); + assert_eq!( + missing, 0, + "field registry is a superset of every loaded item field" + ); + + pool.close(); +} + +/// The Rust loader populates the stored-row metadata the SQL `create_item` path did — `item_hash` +/// (32-byte sha256), `pgstac_updated_at`, `fragment_id` for a fragment collection, and the promoted +/// `datetime`/`geometry` — an upsert with changed content refreshes `item_hash`, and `delete_item` +/// writes a tombstone to `items_deleted_log`. This preserves, on the Rust loader, the coverage the +/// pgtap `003_items` create_item/update_item/delete_item assertions provided. +#[tokio::test] +async fn loader_stored_row_invariants() { + let db = CloneDb::create().await; + let pool = pool(&db).await; + pool.create_collection(&frag_collection("cf")) + .await + .unwrap(); + pool.create_items(vec![frag_item("cf", "f1")], ConflictPolicy::Error) + .await + .unwrap(); + + let (client, connection) = tokio_postgres::connect(&db.dsn(), NoTls).await.unwrap(); + tokio::spawn(async move { + let _ = connection.await; + }); + client + .batch_execute("SET search_path TO pgstac, public;") + .await + .unwrap(); + + let row = client + .query_one( + "SELECT octet_length(item_hash) AS hash_len, \ + pgstac_updated_at IS NOT NULL AS has_ts, \ + fragment_id IS NOT NULL AS has_frag, \ + datetime IS NOT NULL AS has_dt, \ + geometry IS NOT NULL AS has_geom \ + FROM items WHERE id='f1' AND collection='cf'", + &[], + ) + .await + .unwrap(); + assert_eq!( + row.get::<_, i32>("hash_len"), + 32, + "item_hash is a 32-byte sha256" + ); + assert!(row.get::<_, bool>("has_ts"), "pgstac_updated_at populated"); + assert!( + row.get::<_, bool>("has_frag"), + "fragment_id assigned for a fragment collection" + ); + assert!( + row.get::<_, bool>("has_dt"), + "datetime promoted to its column" + ); + assert!(row.get::<_, bool>("has_geom"), "geometry stored"); + + let hash_before: Vec = client + .query_one("SELECT item_hash FROM items WHERE id='f1'", &[]) + .await + .unwrap() + .get(0); + let mut changed = frag_item("cf", "f1"); + changed["properties"]["platform"] = json!("landsat-9"); + pool.create_items(vec![changed], ConflictPolicy::Upsert) + .await + .unwrap(); + let hash_after: Vec = client + .query_one("SELECT item_hash FROM items WHERE id='f1'", &[]) + .await + .unwrap() + .get(0); + assert_ne!( + hash_before, hash_after, + "upsert with changed content refreshes item_hash" + ); + + pool.delete_item("cf", "f1").await.unwrap(); + let tombstones: i64 = client + .query_one( + "SELECT count(*) FROM items_deleted_log WHERE item_id='f1'", + &[], + ) + .await + .unwrap() + .get(0); + assert!( + tombstones >= 1, + "delete_item writes a tombstone to items_deleted_log" + ); + pool.close(); +} diff --git a/src/pgstac-rs/tests/profile_search.rs b/src/pgstac-rs/tests/profile_search.rs new file mode 100644 index 00000000..c31cacd3 --- /dev/null +++ b/src/pgstac-rs/tests/profile_search.rs @@ -0,0 +1,83 @@ +#![cfg(feature = "pool")] +//! Manual profiling harness for the Rust search engine ([`PgstacPool::search_page`]) vs the SQL +//! `search()` function, warm-cache, on a populated database. Ignored by default; run with: +//! PGSTAC_RS_PROFILE_DB=postgresql://…/pgstac_profile \ +//! cargo test --features pool --test profile_search -- --ignored --nocapture + +use pgstac::{ConnectConfig, PgstacPool}; +use serde_json::{Value, json}; +use std::time::{Duration, Instant}; +use tokio_postgres::NoTls; + +async fn median(iters: usize, mut run: F) -> Duration +where + F: FnMut() -> Fut, + Fut: std::future::Future, +{ + for _ in 0..3 { + run().await; // warm + } + let mut times = Vec::with_capacity(iters); + for _ in 0..iters { + let t = Instant::now(); + run().await; + times.push(t.elapsed()); + } + times.sort(); + times[times.len() / 2] +} + +#[tokio::test] +#[ignore = "manual profiling; set PGSTAC_RS_PROFILE_DB to a populated database"] +async fn profile_search_page() { + let Ok(dsn) = std::env::var("PGSTAC_RS_PROFILE_DB") else { + eprintln!("skip: set PGSTAC_RS_PROFILE_DB"); + return; + }; + let pool = PgstacPool::connect(ConnectConfig { + dsn: Some(dsn.clone()), + ..Default::default() + }) + .await + .unwrap(); + let (client, connection) = tokio_postgres::connect(&dsn, NoTls).await.unwrap(); + tokio::spawn(connection); + client + .batch_execute("SET search_path TO pgstac, public") + .await + .unwrap(); + + let queries: [(&str, Value); 4] = [ + ("broad limit=100", json!({"limit": 100})), + ( + "bbox (CONUS)", + json!({"limit": 100, "bbox": [-125.0, 25.0, -65.0, 50.0]}), + ), + ( + "datetime 1mo", + json!({"limit": 100, "datetime": "2020-07-01T00:00:00Z/2020-07-31T23:59:59Z"}), + ), + ( + "cql2 cloud<20", + json!({"limit": 100, "filter-lang": "cql2-json", "filter": {"op": "<", "args": [{"property": "eo:cloud_cover"}, 20]}}), + ), + ]; + + println!("\n=== search_page (Rust engine) vs SQL search() — median of 20 warm runs ==="); + for (name, body) in &queries { + let rust = median(20, || async { + let _ = pool.search_page(body, None, 100).await.unwrap(); + }) + .await; + let sql = median(20, || async { + let _ = client.query_one("SELECT search($1)", &[body]).await; + }) + .await; + println!( + " {name:<18} rust {:>7.1} ms sql {:>7.1} ms", + rust.as_secs_f64() * 1000.0, + sql.as_secs_f64() * 1000.0, + ); + } + pool.close(); +} diff --git a/src/pgstac-rs/tests/search_dimensions.rs b/src/pgstac-rs/tests/search_dimensions.rs new file mode 100644 index 00000000..8b2a2f75 --- /dev/null +++ b/src/pgstac-rs/tests/search_dimensions.rs @@ -0,0 +1,156 @@ +//! Search-dimension parity: the Rust engine must return the same items as SQL `search()` across every +//! search dimension (ids, bbox, datetime, CQL2 filter, non-datetime sort, and `fields` projection), +//! mirroring the functionalities pypgstac / the upstream client test. Targets the rich fragment +//! instance; skips if unreachable. + +#![allow(unused_crate_dependencies)] + +use pgstac::search::search_page; +use serde_json::{Value, json}; +use tokio_postgres::{Client, NoTls}; + +async fn connect() -> Option { + let dsn = std::env::var("PGSTAC_RS_TEST_RICH_DB").unwrap_or_else(|_| { + "postgresql://username:password@localhost:5439/pgstac_rs_test_rich".into() + }); + match tokio_postgres::connect(&dsn, NoTls).await { + Ok((client, connection)) => { + tokio::spawn(connection); + client + .batch_execute("SET search_path = pgstac, public;") + .await + .ok()?; + Some(client) + } + Err(e) => { + eprintln!("[dims] skipping: {e}"); + None + } + } +} + +async fn sql_features(client: &Client, body: &Value) -> Vec { + let fc: Value = client + .query_one("SELECT search($1::jsonb)::jsonb AS fc", &[body]) + .await + .expect("sql search") + .get("fc"); + fc["features"].as_array().cloned().unwrap_or_default() +} + +async fn rust_features(client: &Client, body: &Value) -> Vec { + let limit = body["limit"].as_i64().unwrap_or(10); + search_page(client, body, None, limit) + .await + .expect("rust search_page") + .features +} + +#[tokio::test] +async fn dimensions_match_sql() { + let Some(client) = connect().await else { + return; + }; + let coll = "landsat-c2-l2"; + // A real item's id + spatial envelope, so the geometry cases are *selective* (a small real footprint), + // not the trivial world polygon — this exercises the actual spatial WHERE clause + partition pruning. + let sample = client + .query_one( + "SELECT id, ST_AsGeoJSON(ST_Envelope(geometry)) AS env, ST_XMin(geometry) AS xmin, \ + ST_YMin(geometry) AS ymin, ST_XMax(geometry) AS xmax, ST_YMax(geometry) AS ymax \ + FROM items WHERE collection = 'landsat-c2-l2' ORDER BY datetime LIMIT 1", + &[], + ) + .await + .unwrap(); + let sample_id: String = sample.get("id"); + let env: Value = + serde_json::from_str(&sample.get::<_, String>("env")).expect("envelope geojson"); + let (xmin, ymin, xmax, ymax): (f64, f64, f64, f64) = ( + sample.get("xmin"), + sample.get("ymin"), + sample.get("xmax"), + sample.get("ymax"), + ); + + let bodies = vec![ + ( + "ids", + json!({"collections": [coll], "ids": [sample_id], "limit": 5}), + ), + ( + "bbox", + json!({"collections": [coll], "bbox": [-180, -90, 180, 90], "limit": 5}), + ), + ( + "selective_bbox", + json!({"collections": [coll], "bbox": [xmin, ymin, xmax, ymax], "limit": 5}), + ), + ( + "datetime", + json!({"collections": [coll], "datetime": "2024-01-01T00:00:00Z/..", "limit": 5}), + ), + ( + "cql2_filter", + json!({"collections": [coll], "filter": {"op": "<", "args": [{"property": "eo:cloud_cover"}, 25]}, "limit": 5}), + ), + ( + "query_ext", + json!({"collections": [coll], "query": {"eo:cloud_cover": {"lt": 25}}, "limit": 5}), + ), + ( + "intersects", + json!({"collections": [coll], "intersects": {"type": "Polygon", "coordinates": [[[-180, -90], [180, -90], [180, 90], [-180, 90], [-180, -90]]]}, "limit": 5}), + ), + ( + "selective_intersects", + json!({"collections": [coll], "intersects": env.clone(), "limit": 5}), + ), + // Adversarial: a selective footprint + a wide datetime window + a non-datetime sort all at once — the + // spatial filter prunes partitions, the sort drives a non-keyset order; both engines must still agree. + ( + "intersects_datetime_sort", + json!({"collections": [coll], "intersects": env, "datetime": "2013-01-01T00:00:00Z/..", "sortby": [{"field": "eo:cloud_cover", "direction": "asc"}], "limit": 5}), + ), + // A footprint over open ocean: both engines must return zero rows (the empty-result path). + ( + "intersects_no_match", + json!({"collections": [coll], "intersects": {"type": "Polygon", "coordinates": [[[0.0, 0.0], [0.01, 0.0], [0.01, 0.01], [0.0, 0.01], [0.0, 0.0]]]}, "limit": 5}), + ), + ( + "nondatetime_sort", + json!({"collections": [coll], "sortby": [{"field": "eo:cloud_cover", "direction": "asc"}], "limit": 5}), + ), + ( + "nondatetime_sort_desc", + json!({"collections": [coll], "sortby": [{"field": "eo:cloud_cover", "direction": "desc"}], "limit": 5}), + ), + ]; + for (name, body) in bodies { + let rust = rust_features(&client, &body).await; + let sql = sql_features(&client, &body).await; + assert_eq!(rust, sql, "{name} mismatch (body={body})"); + eprintln!("[dims] {name}: {} items match SQL", rust.len()); + } +} + +#[tokio::test] +async fn fields_projection_matches_sql() { + let Some(client) = connect().await else { + return; + }; + let coll = "landsat-c2-l2"; + for fields in [ + json!({"exclude": ["assets", "properties.eo:cloud_cover"]}), + json!({"include": ["id", "properties.datetime"]}), + // include id+geometry only: search_plan nulls fragment_id in the projection (no shared + // fragment needed), so the engine skips the per-row fragment fetch -- must still match SQL. + json!({"include": ["id", "geometry"]}), + ] { + let body = json!({"collections": [coll], "fields": fields, "limit": 3}); + let rust = rust_features(&client, &body).await; + let sql = sql_features(&client, &body).await; + assert_eq!(rust, sql, "fields mismatch (fields={fields})"); + eprintln!("[dims] fields {fields}: match SQL"); + } +} diff --git a/src/pgstac-rs/tests/search_page.rs b/src/pgstac-rs/tests/search_page.rs new file mode 100644 index 00000000..b1ee39f8 --- /dev/null +++ b/src/pgstac-rs/tests/search_page.rs @@ -0,0 +1,230 @@ +//! Fidelity + pagination gate for the Rust `search_page` against SQL `search()`. +//! +//! For datetime-ascending, datetime-descending, and a non-datetime sort (`id`), this paginates an +//! entire collection page-by-page through the `next` tokens and asserts, at every page, that +//! `search_page` matches `search()` byte-for-byte (hydrated features + keyset tokens). It also checks +//! the full traversal has no duplicates, no gaps (covers every item exactly once), is correctly +//! ordered, and that `prev` tokens round-trip. +//! +//! Targets a clean fragment instance (default `pgstac_rs_test_rich`); skips if unreachable. + +#![allow(unused_crate_dependencies)] + +use pgstac::search::search_page; +use serde_json::{Value, json}; +use std::collections::HashSet; +use tokio_postgres::{Client, NoTls}; + +async fn connect() -> Option { + let dsn = std::env::var("PGSTAC_RS_TEST_RICH_DB").unwrap_or_else(|_| { + "postgresql://username:password@localhost:5439/pgstac_rs_test_rich".into() + }); + match tokio_postgres::connect(&dsn, NoTls).await { + Ok((client, connection)) => { + tokio::spawn(connection); + client + .batch_execute("SET search_path = pgstac, public;") + .await + .ok()?; + Some(client) + } + Err(e) => { + eprintln!("[search_page] skipping: cannot connect ({dsn}): {e}"); + None + } + } +} + +async fn sql_search(client: &Client, body: &Value) -> Value { + client + .query_one("SELECT search($1::jsonb) AS fc", &[body]) + .await + .expect("sql search") + .get("fc") +} + +fn features(fc: &Value) -> &Vec { + fc["features"].as_array().expect("features array") +} + +fn id_of(feature: &Value) -> String { + feature["id"].as_str().unwrap_or_default().to_string() +} + +async fn count_in(client: &Client, collection: &str) -> usize { + let n: i64 = client + .query_one( + "SELECT count(*) AS n FROM items WHERE collection = $1", + &[&collection], + ) + .await + .unwrap() + .get("n"); + n as usize +} + +/// Paginates `body` fully through the Rust `next` tokens; at each page asserts the Rust page's +/// features equal SQL `search()` fed the same (Rust-minted) token. Returns the ordered list of feature +/// ids and the per-page `prev` tokens. +async fn paginate(client: &Client, body: &Value, limit: i64) -> (Vec, Vec>) { + let mut token: Option = None; + let mut ids = Vec::new(); + let mut prev_tokens = Vec::new(); + let mut guard = 0; + loop { + guard += 1; + assert!(guard < 1000, "pagination did not terminate for {body}"); + + let page = search_page(client, body, token.as_deref(), limit) + .await + .unwrap(); + + let mut sql_body = body.clone(); + if let Some(t) = &token { + let _ = sql_body + .as_object_mut() + .unwrap() + .insert("token".into(), Value::String(t.clone())); + } + let fc = sql_search(client, &sql_body).await; + + // Page-equivalence: feeding the Rust-minted token to SQL `search()` returns exactly the + // items the Rust page returned. This proves the token round-trips (the server casts each + // decoded keyset value back to its column type), without byte-identical token strings. + assert_eq!( + &page.features, + features(&fc), + "features mismatch (body={body}, token={token:?})" + ); + + ids.extend(page.features.iter().map(id_of)); + prev_tokens.push(page.prev_token.clone()); + + match page.next_token { + Some(next) if !page.features.is_empty() => token = Some(next), + _ => break, + } + } + (ids, prev_tokens) +} + +#[tokio::test] +async fn pagination_matches_search_for_asc_desc_and_non_datetime() { + let Some(client) = connect().await else { + return; + }; + let collection = "landsat-c2-l2"; + let total = count_in(&client, collection).await; + let limit = 137; // a prime-ish page size so the last page is partial + + for sort in [ + json!([{"field": "datetime", "direction": "desc"}]), + json!([{"field": "datetime", "direction": "asc"}]), + json!([{"field": "id", "direction": "asc"}]), // non-datetime: exercises the non-band path + json!([{"field": "id", "direction": "desc"}]), + ] { + let body = json!({"collections": [collection], "sortby": sort, "limit": limit}); + let (ids, _prev) = paginate(&client, &body, limit).await; + + // Covered every item exactly once (no gaps, no duplicates). + let unique: HashSet<&String> = ids.iter().collect(); + assert_eq!(unique.len(), ids.len(), "duplicate ids paginating {sort}"); + assert_eq!( + ids.len(), + total, + "did not cover all {total} items for {sort}" + ); + eprintln!( + "[search_page] sort {sort}: {} items over {} pages, matches search()", + ids.len(), + ids.len().div_ceil(limit as usize) + ); + } +} + +#[tokio::test] +async fn prev_token_round_trips_to_previous_page() { + let Some(client) = connect().await else { + return; + }; + let body = json!({"collections": ["landsat-c2-l2"], "limit": 50}); + + // Page 1, then page 2 via its next token. + let page1 = search_page(&client, &body, None, 50).await.unwrap(); + let next = page1.next_token.clone().expect("next token"); + let page2 = search_page(&client, &body, Some(&next), 50).await.unwrap(); + + // Page 2's prev token must return exactly page 1. + let prev = page2.prev_token.clone().expect("prev token on page 2"); + let back = search_page(&client, &body, Some(&prev), 50).await.unwrap(); + assert_eq!( + back.features.iter().map(id_of).collect::>(), + page1.features.iter().map(id_of).collect::>(), + "prev token did not round-trip to page 1" + ); + eprintln!("[search_page] prev token round-trip ok"); +} + +fn ndjson_ids(buf: &[u8]) -> Vec { + buf.split(|&b| b == b'\n') + .filter(|line| !line.is_empty()) + .map(|line| { + let v: Value = serde_json::from_slice(line).unwrap(); + v["id"].as_str().unwrap_or_default().to_string() + }) + .collect() +} + +#[tokio::test] +async fn collect_and_stream_cover_all_items_in_the_same_order() { + let Some(client) = connect().await else { + return; + }; + let collection = "landsat-c2-l2"; + let total = count_in(&client, collection).await; + let body = json!({"collections": [collection], "limit": 200}); + + // collect-all paginates internally and returns the whole set. + let collected = pgstac::search::search_collect(&client, &body, None) + .await + .unwrap(); + assert_eq!(collected.number_returned, total); + assert!(collected.next_token.is_none()); + let collected_ids: Vec = collected.features.iter().map(id_of).collect(); + let unique: HashSet<&String> = collected_ids.iter().collect(); + assert_eq!(unique.len(), total, "collect produced duplicates"); + + // stream emits the same items, in the same order, as NDJSON. + let mut buf: Vec = Vec::new(); + let written = pgstac::search::search_stream(&client, &body, &mut buf, None) + .await + .unwrap(); + assert_eq!(written, total); + assert_eq!( + ndjson_ids(&buf), + collected_ids, + "stream != collect order/set" + ); + eprintln!("[search_page] collect + stream cover all {total} items identically"); +} + +#[tokio::test] +async fn collect_and_stream_respect_max_items() { + let Some(client) = connect().await else { + return; + }; + let body = json!({"collections": ["landsat-c2-l2"], "limit": 7}); + + let collected = pgstac::search::search_collect(&client, &body, Some(10)) + .await + .unwrap(); + assert_eq!(collected.number_returned, 10); + + let mut buf: Vec = Vec::new(); + let written = pgstac::search::search_stream(&client, &body, &mut buf, Some(10)) + .await + .unwrap(); + assert_eq!(written, 10); + assert_eq!(ndjson_ids(&buf).len(), 10); + eprintln!("[search_page] collect + stream respect max_items"); +} diff --git a/src/pgstac-rs/tests/stream.rs b/src/pgstac-rs/tests/stream.rs new file mode 100644 index 00000000..518d2656 --- /dev/null +++ b/src/pgstac-rs/tests/stream.rs @@ -0,0 +1,117 @@ +//! The flat-memory streaming iterator: must yield the same items, in the same order, with the same +//! byte-for-byte hydration as the buffered `search_page`. Targets the rich fragment instance; skips if +//! unreachable. Requires the `pool` feature. + +#![cfg(feature = "pool")] +#![allow(unused_crate_dependencies)] + +use futures::StreamExt; +use pgstac::{ConnectConfig, DEFAULT_SEARCH_PATH, PgstacPool}; +use serde_json::{Value, json}; + +async fn pool() -> Option { + let dsn = std::env::var("PGSTAC_RS_TEST_RICH_DB").unwrap_or_else(|_| { + "postgresql://username:password@localhost:5439/pgstac_rs_test_rich".into() + }); + let config = ConnectConfig { + dsn: Some(dsn), + search_path: DEFAULT_SEARCH_PATH.to_string(), + ..Default::default() + }; + match PgstacPool::connect(config).await { + Ok(p) => Some(p), + Err(e) => { + eprintln!("[stream] skipping: cannot connect: {e}"); + None + } + } +} + +fn id_of(v: &Value) -> String { + v["id"].as_str().unwrap_or_default().to_string() +} + +#[tokio::test] +async fn stream_matches_search_page_order_and_hydration() { + let Some(pool) = pool().await else { + return; + }; + let body = json!({"collections": ["landsat-c2-l2"], "limit": 1000}); + + // The streaming iterator (flat memory). + let streamed: Vec = { + let stream = pool.search_items(body.clone(), None, Some(1000)); + futures::pin_mut!(stream); + let mut out = Vec::new(); + while let Some(item) = stream.next().await { + out.push(item.unwrap()); + } + out + }; + + // The buffered reference, on a pooled connection. + let client = pool.get().await.unwrap(); + let page = pgstac::search::search_page(&**client, &body, None, 1000) + .await + .unwrap(); + + assert_eq!(streamed.len(), 1000, "stream returned {}", streamed.len()); + assert_eq!( + streamed.iter().map(id_of).collect::>(), + page.features.iter().map(id_of).collect::>(), + "stream order/set differs from search_page", + ); + // Full hydration must be byte-identical (same Hydrator, same fragment content). + assert_eq!( + streamed, page.features, + "stream hydration differs from search_page" + ); + eprintln!("[stream] 1000 items match search_page byte-for-byte"); +} + +#[tokio::test] +async fn search_matched_equals_page_number_matched() { + let Some(pool) = pool().await else { + return; + }; + let body = json!({"collections": ["landsat-c2-l2"]}); + // The standalone count (parallel-context primitive) must agree with the page's numberMatched. + let (page, matched) = tokio::join!( + pool.search_page(&body, None, 10), + pool.search_matched(&body) + ); + assert_eq!( + matched.unwrap(), + page.unwrap().number_matched, + "search_matched != search_page numberMatched" + ); + eprintln!("[stream] search_matched agrees with search_page numberMatched"); +} + +#[tokio::test] +async fn collect_and_ndjson_helpers() { + let Some(pool) = pool().await else { + return; + }; + let body = json!({"collections": ["landsat-c2-l2"]}); + + let collected = pool + .search_collect_items(body.clone(), Some(50)) + .await + .unwrap(); + assert_eq!(collected.len(), 50); + + let mut buf: Vec = Vec::new(); + let n = pool.stream_ndjson(body, None, Some(50), &mut buf).await.unwrap(); + assert_eq!(n, 50); + assert_eq!(buf.iter().filter(|&&b| b == b'\n').count(), 50); + // The byte-path NDJSON (serialize-time merge) must equal the Value path feature-for-feature on real + // landsat items (25-asset structure + shared fragment) — transitively == SQL content_hydrate. + let line_features: Vec = buf + .split(|&b| b == b'\n') + .filter(|l| !l.is_empty()) + .map(|l| serde_json::from_slice::(l).unwrap()) + .collect(); + assert_eq!(line_features, collected, "byte-path NDJSON != Value path"); + eprintln!("[stream] byte-path NDJSON matches the Value path for 50 landsat items"); +} diff --git a/src/pgstac/migrations/pgstac--0.9.11--unreleased.sql b/src/pgstac/migrations/pgstac--0.9.11--unreleased.sql index 58a7d7e4..8eb47b2a 100644 --- a/src/pgstac/migrations/pgstac--0.9.11--unreleased.sql +++ b/src/pgstac/migrations/pgstac--0.9.11--unreleased.sql @@ -45,6 +45,7 @@ DO $$ $$; + GRANT pgstac_admin TO current_user; -- Function to make sure pgstac_admin is the owner of items @@ -102,10 +103,8 @@ $$ LANGUAGE PLPGSQL; SELECT pgstac_admin_owns(); CREATE SCHEMA IF NOT EXISTS pgstac AUTHORIZATION pgstac_admin; - -GRANT ALL ON ALL FUNCTIONS IN SCHEMA pgstac to pgstac_admin; -GRANT ALL ON ALL TABLES IN SCHEMA pgstac to pgstac_admin; -GRANT ALL ON ALL SEQUENCES IN SCHEMA pgstac to pgstac_admin; +-- pgstac_admin owns the schema and all objects in it (pgstac_admin_owns() above + +-- AUTHORIZATION), so it already has every privilege — no explicit GRANT ALL needed. ALTER ROLE pgstac_admin SET SEARCH_PATH TO pgstac, public; ALTER ROLE pgstac_read SET SEARCH_PATH TO pgstac, public; @@ -197,41 +196,61 @@ RETURNS timestamptz AS $$ END ; $$ LANGUAGE SQL IMMUTABLE STRICT; -create sequence "pgstac"."item_fragments_id_seq"; - -drop trigger if exists "items_after_delete_trigger" on "pgstac"."items"; - -drop trigger if exists "items_after_update_trigger" on "pgstac"."items"; - -drop function if exists "pgstac"."content_hydrate"(_item items, _collection collections, fields jsonb); - -drop function if exists "pgstac"."content_hydrate"(_item jsonb, _base_item jsonb, fields jsonb); - -drop function if exists "pgstac"."content_nonhydrated"(_item items, fields jsonb); - -drop function if exists "pgstac"."content_slim"(_item jsonb); - -drop function if exists "pgstac"."create_collection"(data jsonb); - -drop function if exists "pgstac"."format_item"(_item items, _fields jsonb, _hydrated boolean); - - -drop function if exists "pgstac"."search_rows"(_where text, _orderby text, partitions text[], _limit integer); - -drop function if exists "pgstac"."update_collection"(data jsonb); -drop function if exists "pgstac"."upsert_collection"(data jsonb); - -drop function if exists "pgstac"."where_stats"(inwhere text, updatestats boolean, conf jsonb); +-- Drop objects superseded by the current partition_stats model. +DROP MATERIALIZED VIEW IF EXISTS partitions CASCADE; +DROP MATERIALIZED VIEW IF EXISTS partition_steps; +DROP VIEW IF EXISTS partition_steps; + +-- Drop function signatures whose argument lists changed (CREATE OR REPLACE cannot alter them) +DROP FUNCTION IF EXISTS chunker(pred_envelope); +DROP FUNCTION IF EXISTS search_bands(pred_envelope, boolean, integer, integer); +DROP FUNCTION IF EXISTS search_rows(jsonb, integer, text, boolean); +DROP FUNCTION IF EXISTS search_sql(jsonb, text); +DROP FUNCTION IF EXISTS search_cursor(jsonb, text, refcursor); +DROP FUNCTION IF EXISTS search_cursor(jsonb); +DROP FUNCTION IF EXISTS fields_to_columns(jsonb); +DROP FUNCTION IF EXISTS fields_to_rowjsonb(jsonb, text[]); +DROP FUNCTION IF EXISTS fields_to_rowjsonb(jsonb); +DROP FUNCTION IF EXISTS collection_search_rows(jsonb); +DROP FUNCTION IF EXISTS collection_fragments_properties(text); +DROP FUNCTION IF EXISTS format_item(items, jsonb); +DROP FUNCTION IF EXISTS geometrysearch(geometry, text, jsonb, integer, integer, interval, boolean, boolean); +DROP FUNCTION IF EXISTS geojsonsearch(jsonb, text, jsonb, integer, integer, interval, boolean, boolean); +DROP FUNCTION IF EXISTS xyzsearch(integer, integer, integer, text, jsonb, integer, integer, interval, boolean, boolean); +DROP FUNCTION IF EXISTS search(jsonb); +DROP FUNCTION IF EXISTS search_page(jsonb, integer, text, boolean); +DROP FUNCTION IF EXISTS search_plan(jsonb, text); +DROP FUNCTION IF EXISTS fields_to_itemcols(jsonb); +DROP FUNCTION IF EXISTS search_query(jsonb, boolean, jsonb); +DROP FUNCTION IF EXISTS where_stats(text, text, boolean, jsonb); +DROP FUNCTION IF EXISTS keyset_sortkeys(jsonb); +DROP FUNCTION IF EXISTS paging_dtrange(jsonb); +DROP FUNCTION IF EXISTS paging_collections(jsonb); +DROP FUNCTION IF EXISTS get_token_filter(jsonb, items, boolean, boolean); +DROP FUNCTION IF EXISTS get_token_record(text); +DROP FUNCTION IF EXISTS get_token_val_str(text, items); +DROP FUNCTION IF EXISTS sort_dir_to_op(text, boolean); +DROP FUNCTION IF EXISTS sort_sqlorderby(jsonb, boolean); +DROP FUNCTION IF EXISTS parse_sort_dir(text, boolean); +DROP FUNCTION IF EXISTS cql2_envelope_safe(jsonb); +DROP TABLE IF EXISTS search_wheres; +drop materialized view if exists "pgstac"."partition_steps"; + +drop materialized view if exists "pgstac"."partitions"; drop view if exists "pgstac"."collections_asitems"; -drop function if exists "pgstac"."content_hydrate"(_item items, fields jsonb); +drop view if exists "pgstac"."partition_sys_meta"; + +drop view if exists "pgstac"."partitions_view"; alter table "pgstac"."format_item_cache" drop constraint "format_item_cache_pkey"; alter table "pgstac"."search_wheres" drop constraint "search_wheres_pkey"; +drop index if exists "pgstac"."partitions_partition_idx"; + drop index if exists "pgstac"."format_item_cache_lastused_idx"; drop index if exists "pgstac"."format_item_cache_pkey"; @@ -242,195 +261,169 @@ drop index if exists "pgstac"."search_wheres_pkey"; drop index if exists "pgstac"."search_wheres_where"; -drop table "pgstac"."format_item_cache"; - -drop table "pgstac"."search_wheres"; - -create table "pgstac"."item_field_registry" ( - "collection" text not null, - "path" text not null, - "is_leaf" boolean default true, - "value_kinds" text[] default '{}'::text[], - "first_seen" timestamp with time zone not null default now(), - "last_seen" timestamp with time zone not null default now() -); +create sequence "pgstac"."item_fragments_id_seq"; +drop trigger if exists "items_after_delete_trigger" on "pgstac"."items"; -create table "pgstac"."item_fragments" ( - "id" bigint not null default nextval('item_fragments_id_seq'::regclass), - "collection" text not null, - "hash" text not null, - "content" jsonb not null, - "created_at" timestamp with time zone not null default now() -); +drop trigger if exists "items_after_insert_trigger" on "pgstac"."items"; +drop trigger if exists "items_after_update_trigger" on "pgstac"."items"; -create table "pgstac"."items_deleted_log" ( - "id" bigint generated always as identity not null, - "item_id" text not null, - "collection" text not null, - "partition" text, - "datetime" timestamp with time zone, - "end_datetime" timestamp with time zone, - "content_hash" text not null default ''::text, - "deleted_at" timestamp with time zone not null default now() +create type "pgstac"."pred_envelope" as ( + "colls" _text, + "dt" tstzmultirange, + "edt" tstzmultirange, + "geom" geometry ); +set check_function_bodies = off; -alter table "pgstac"."collections" drop column "base_item"; - -alter table "pgstac"."collections" add column "fragment_config" text[]; - -alter table "pgstac"."items" drop column "content"; - -alter table "pgstac"."items" drop column "private"; - -alter table "pgstac"."items" add column "assets" jsonb default '{}'::jsonb; - -alter table "pgstac"."items" add column "bbox" jsonb; - -alter table "pgstac"."items" add column "constellation" text; - -alter table "pgstac"."items" add column "content_hash" text not null default ''::text; - -alter table "pgstac"."items" add column "created" timestamp with time zone; - -alter table "pgstac"."items" add column "datetime_is_range" boolean not null default false; - -alter table "pgstac"."items" add column "eo_bands" jsonb; - -alter table "pgstac"."items" add column "eo_cloud_cover" double precision; - -alter table "pgstac"."items" add column "eo_snow_cover" double precision; - -alter table "pgstac"."items" add column "extra" jsonb; - -alter table "pgstac"."items" add column "file_byte_order" text; - -alter table "pgstac"."items" add column "file_checksum" text; - -alter table "pgstac"."items" add column "file_header_size" bigint; - -alter table "pgstac"."items" add column "file_size" bigint; - -alter table "pgstac"."items" add column "file_values_regex" text; - -alter table "pgstac"."items" add column "fragment_id" bigint; - -alter table "pgstac"."items" add column "gsd" double precision; - -alter table "pgstac"."items" add column "instruments" text[]; - -alter table "pgstac"."items" add column "links" jsonb default '[]'::jsonb; - -alter table "pgstac"."items" add column "mission" text; - -alter table "pgstac"."items" add column "pgstac_updated_at" timestamp with time zone not null default now(); - -alter table "pgstac"."items" add column "platform" text; - -alter table "pgstac"."items" add column "proj_bbox" jsonb; - -alter table "pgstac"."items" add column "proj_centroid" jsonb; - -alter table "pgstac"."items" add column "proj_epsg" integer; - -alter table "pgstac"."items" add column "proj_projjson" jsonb; - -alter table "pgstac"."items" add column "proj_shape" jsonb; - -alter table "pgstac"."items" add column "proj_transform" jsonb; - -alter table "pgstac"."items" add column "proj_wkt2" text; - -alter table "pgstac"."items" add column "properties" jsonb default '{}'::jsonb; - -alter table "pgstac"."items" add column "sat_absolute_orbit" integer; - -alter table "pgstac"."items" add column "sat_orbit_state" text; - -alter table "pgstac"."items" add column "sat_relative_orbit" integer; - -alter table "pgstac"."items" add column "sci_citation" text; - -alter table "pgstac"."items" add column "sci_doi" text; - -alter table "pgstac"."items" add column "sci_publications" jsonb; - -alter table "pgstac"."items" add column "stac_extensions" jsonb default '[]'::jsonb; - -alter table "pgstac"."items" add column "stac_version" text; - -alter table "pgstac"."items" add column "updated" timestamp with time zone; - -alter table "pgstac"."items" add column "view_azimuth" double precision; - -alter table "pgstac"."items" add column "view_incidence_angle" double precision; - -alter table "pgstac"."items" add column "view_off_nadir" double precision; - -alter table "pgstac"."items" add column "view_sun_azimuth" double precision; - -alter table "pgstac"."items" add column "view_sun_elevation" double precision; - -alter table "pgstac"."searches" add column "context_count" bigint; - -alter table "pgstac"."searches" add column "created_at" timestamp with time zone default now(); - -alter table "pgstac"."searches" add column "name" text; - -alter table "pgstac"."searches" add column "pinned" boolean not null default false; - -alter table "pgstac"."searches" add column "statslastupdated" timestamp with time zone; - -alter table "pgstac"."searches" alter column "hash" drop expression; - -alter sequence "pgstac"."item_fragments_id_seq" owned by "pgstac"."item_fragments"."id"; - -CREATE INDEX item_field_registry_path_idx ON pgstac.item_field_registry USING btree (path); - -CREATE UNIQUE INDEX item_field_registry_pkey ON pgstac.item_field_registry USING btree (collection, path); - -CREATE UNIQUE INDEX item_fragments_collection_hash_key ON pgstac.item_fragments USING btree (collection, hash); - -CREATE INDEX item_fragments_collection_idx ON pgstac.item_fragments USING btree (collection); - -CREATE UNIQUE INDEX item_fragments_pkey ON pgstac.item_fragments USING btree (id); - -CREATE INDEX items_deleted_log_deleted_at_idx ON pgstac.items_deleted_log USING btree (deleted_at); - -CREATE UNIQUE INDEX items_deleted_log_pkey ON pgstac.items_deleted_log USING btree (id); - -CREATE INDEX items_fragment_id_idx ON ONLY pgstac.items USING btree (fragment_id) WHERE (fragment_id IS NOT NULL); - -CREATE INDEX searches_lastused_anon_idx ON pgstac.searches USING btree (lastused) WHERE ((name IS NULL) AND (NOT pinned)); - -CREATE UNIQUE INDEX searches_name_key ON pgstac.searches USING btree (name); - -alter table "pgstac"."item_field_registry" add constraint "item_field_registry_pkey" PRIMARY KEY using index "item_field_registry_pkey"; - -alter table "pgstac"."item_fragments" add constraint "item_fragments_pkey" PRIMARY KEY using index "item_fragments_pkey"; +CREATE OR REPLACE FUNCTION pgstac.build_pending_indexes(_limit integer DEFAULT NULL::integer) + RETURNS integer + LANGUAGE plpgsql + SECURITY DEFINER +AS $function$ +DECLARE + _part text; + _q text; + _count int := 0; +BEGIN + FOR _part IN + SELECT partition FROM pgstac.partition_stats + WHERE indexes_pending + ORDER BY last_updated NULLS FIRST + LIMIT _limit + LOOP + FOR _q IN SELECT * FROM pgstac.maintain_partition_queries(_part) LOOP + EXECUTE _q; + END LOOP; + UPDATE pgstac.partition_stats SET indexes_pending = false WHERE partition = _part; + _count := _count + 1; + END LOOP; + RETURN _count; +END; +$function$ +; -alter table "pgstac"."items_deleted_log" add constraint "items_deleted_log_pkey" PRIMARY KEY using index "items_deleted_log_pkey"; +CREATE OR REPLACE FUNCTION pgstac.check_partition(_collection text, _dtrange tstzrange, _edtrange tstzrange, _spatial geometry DEFAULT NULL::geometry) + RETURNS text + LANGUAGE plpgsql + SECURITY DEFINER +AS $function$ +DECLARE + c RECORD; + _partition_name text; + _parent_name text; + _partition_range tstzrange; +BEGIN + SELECT * INTO c FROM pgstac.collections WHERE id=_collection; + IF NOT FOUND THEN + RAISE EXCEPTION 'Collection % does not exist', _collection USING ERRCODE = 'foreign_key_violation', HINT = 'Make sure collection exists before adding items'; + END IF; -alter table "pgstac"."item_field_registry" add constraint "item_field_registry_collection_fkey" FOREIGN KEY ("collection") REFERENCES "pgstac"."collections"("id") ON DELETE CASCADE NOT VALID; + IF c.partition_trunc IS NOT NULL THEN + _partition_range := tstzrange( + date_trunc(c.partition_trunc, lower(_dtrange)), + date_trunc(c.partition_trunc, lower(_dtrange)) + (concat('1 ', c.partition_trunc))::interval, + '[)' + ); + ELSE + _partition_range := '[-infinity, infinity]'::tstzrange; + END IF; -alter table "pgstac"."item_field_registry" validate constraint "item_field_registry_collection_fkey"; + IF NOT _partition_range @> _dtrange THEN + RAISE EXCEPTION 'dtrange % spans more than the % partition window for collection %', _dtrange, c.partition_trunc, _collection; + END IF; -alter table "pgstac"."item_fragments" add constraint "item_fragments_collection_fkey" FOREIGN KEY ("collection") REFERENCES "pgstac"."collections"("id") ON DELETE CASCADE NOT VALID; + _parent_name := format('_items_%s', c.key); + IF c.partition_trunc = 'year' THEN + _partition_name := format('%s_%s', _parent_name, to_char(lower(_partition_range),'YYYY')); + ELSIF c.partition_trunc = 'month' THEN + _partition_name := format('%s_%s', _parent_name, to_char(lower(_partition_range),'YYYYMM')); + ELSE + _partition_name := _parent_name; + END IF; -alter table "pgstac"."item_fragments" validate constraint "item_fragments_collection_fkey"; + -- Create the collection-level PARENT partition (_items_) first, for sub-partitioned collections. + -- It is shared across every child window, so concurrent setup of different children would race on its + -- CREATE TABLE. Guard it with a parent-scoped advisory lock, taken before any child lock and only when + -- the parent is missing: parent-before-child is one lock order (parent < child) that can't deadlock + -- with ensure_partitions' sorted child locks, and skipping it once the parent exists keeps steady-state + -- ingest off the parent lock. + IF c.partition_trunc IS NOT NULL AND to_regclass(format('pgstac.%I', _parent_name)) IS NULL THEN + PERFORM pg_advisory_xact_lock(hashtext('pgstac.check_partition'), hashtext(_parent_name)); + IF to_regclass(format('pgstac.%I', _parent_name)) IS NULL THEN + EXECUTE format( + 'CREATE TABLE IF NOT EXISTS %I partition OF items FOR VALUES IN (%L) PARTITION BY RANGE (datetime)', + _parent_name, _collection + ); + END IF; + END IF; -alter table "pgstac"."item_fragments" add constraint "item_fragments_collection_hash_key" UNIQUE using index "item_fragments_collection_hash_key"; + -- Serialize concurrent setup of THIS (leaf) partition with a partition-scoped advisory lock. Different + -- partitions hash to different keys, so this never serializes unrelated ingest; ensure_partitions + -- acquires these in sorted order, so concurrent multi-partition batches cannot deadlock. The lock + -- releases at commit (setup is fast + idempotent). + PERFORM pg_advisory_xact_lock(hashtext('pgstac.check_partition'), hashtext(_partition_name)); + + -- Create the leaf partition if missing. Parent-inherited indexes only (id PK here; datetime/geometry + -- come from the items parent); no CHECK constraints; queryable indexes are deferred via + -- indexes_pending. A SELECT grant lets read/ingest query it; writes reach it only through the + -- SECURITY DEFINER write functions (the privilege wall in 998_idempotent_post). + -- Skip the DDL when it already exists: re-running CREATE/GRANT takes a relation lock that deadlocks + -- with concurrent INSERTs. The existence check is race-safe under the advisory lock. + IF to_regclass(format('pgstac.%I', _partition_name)) IS NULL THEN + BEGIN + IF c.partition_trunc IS NULL THEN + EXECUTE format( + $q$ + CREATE TABLE IF NOT EXISTS %I partition OF items FOR VALUES IN (%L); + CREATE UNIQUE INDEX IF NOT EXISTS %I ON %I (id); + GRANT SELECT ON %I TO pgstac_read, pgstac_ingest; + $q$, + _partition_name, _collection, + concat(_partition_name, '_pk'), _partition_name, + _partition_name + ); + ELSE + EXECUTE format( + $q$ + CREATE TABLE IF NOT EXISTS %I partition OF %I FOR VALUES FROM (%L) TO (%L); + CREATE UNIQUE INDEX IF NOT EXISTS %I ON %I (id); + GRANT SELECT ON %I TO pgstac_read, pgstac_ingest; + $q$, + _partition_name, _parent_name, lower(_partition_range), upper(_partition_range), + concat(_partition_name, '_pk'), _partition_name, + _partition_name + ); + END IF; + EXCEPTION + -- A concurrent creator that finished between our checks: benign, it exists now. + WHEN duplicate_table THEN + RAISE DEBUG 'Partition % already exists.', _partition_name; + -- Do NOT swallow other errors: a failed creation must propagate so the caller never goes on to + -- write a partition that does not exist (invariant: check_partition succeeds before use). + END; + END IF; -alter table "pgstac"."searches" add constraint "searches_name_key" UNIQUE using index "searches_name_key"; + -- Seed the partition_stats row (collection only — the data ranges start NULL), then cover this batch + -- via the shared widen guard, which fills dtrange/edtrange. ON CONFLICT keeps an existing row (so a + -- sweep that cleared indexes_pending/dirty is not reset on a later check_partition for the same + -- partition). + INSERT INTO partition_stats (partition, collection, dirty, indexes_pending) + VALUES (_partition_name, _collection, true, true) + ON CONFLICT (partition) DO NOTHING; + PERFORM widen_partition_stats(_partition_name, _dtrange, _edtrange, _spatial, _partition_range); -set check_function_bodies = off; + RETURN _partition_name; +END; +$function$ +; CREATE OR REPLACE FUNCTION pgstac.collection_fragment_config_default(content jsonb) RETURNS text[] LANGUAGE plpgsql - STABLE SECURITY DEFINER + STABLE AS $function$ DECLARE paths text[] := ARRAY[]::text[]; @@ -474,35 +467,364 @@ END; $function$ ; -CREATE OR REPLACE FUNCTION pgstac.create_collection(data jsonb, _partition_trunc text DEFAULT NULL::text, _fragment_config text[] DEFAULT NULL::text[]) - RETURNS void - LANGUAGE sql +CREATE OR REPLACE FUNCTION pgstac.collection_search_plan(_search jsonb DEFAULT '{}'::jsonb, _token text DEFAULT NULL::text, OUT query text, OUT ctx_query text) + RETURNS record + LANGUAGE plpgsql + STABLE SECURITY DEFINER SET search_path TO 'pgstac', 'public' AS $function$ - INSERT INTO collections (content, fragment_config, partition_trunc) - VALUES ( - data, - COALESCE(_fragment_config, collection_fragment_config_default(data)), - _partition_trunc - ) - ; +DECLARE + is_prev boolean := _token LIKE 'prev:%'; + keyset text := nullif(regexp_replace(coalesce(_token, ''), '^(next|prev):', ''), ''); + _fields jsonb := coalesce(_search->'fields', '{}'::jsonb); + -- collections default to id ASC unless the caller supplied an explicit sortby. + _eff jsonb := CASE WHEN _search ? 'sortby' THEN _search + ELSE _search || '{"sortby":[{"field":"id","direction":"asc"}]}'::jsonb END; + _cql2 jsonb := search_to_cql2(_search); + _where text := coalesce(nullif(btrim(cql2_query(_cql2)), ''), 'TRUE'); + orderby_str text; keys_proj text; keyset_w text; full_where text; +BEGIN + -- effective ORDER BY (reversed for prev) + the key projection, from the keyset keys. + SELECT string_agg(expr || ' ' || CASE WHEN is_prev THEN (CASE dir WHEN 'ASC' THEN 'DESC' ELSE 'ASC' END) ELSE dir END, ', ' ORDER BY ord), + 'ARRAY[' || string_agg(format('(%s)::text', expr), ',' ORDER BY ord) || ']::text[]' + INTO orderby_str, keys_proj + FROM keyset_sortkeys(_eff); + + IF keyset IS NOT NULL THEN + keyset_w := keyset_where(_eff, keyset_decode(keyset), is_prev); + END IF; + full_where := concat_ws(' AND ', _where, keyset_w); + IF full_where IS NULL OR btrim(full_where) = '' THEN full_where := 'TRUE'; END IF; + + query := format( + $q$ SELECT jsonb_fields(collectionjson, %L) AS content, %s AS keys + FROM collections_asitems WHERE %s ORDER BY %s LIMIT $1 $q$, + _fields, keys_proj, full_where, orderby_str); + ctx_query := format('SELECT count(*) FROM collections_asitems WHERE %s', _where); +END; $function$ ; -CREATE OR REPLACE FUNCTION pgstac.extract_fragment(content jsonb, fragment_paths text[]) +CREATE OR REPLACE FUNCTION pgstac.content_hydrate(_item items, fields jsonb DEFAULT '{}'::jsonb, _skip_fragment boolean DEFAULT false) RETURNS jsonb LANGUAGE plpgsql - IMMUTABLE PARALLEL SAFE SECURITY DEFINER + STABLE PARALLEL SAFE AS $function$ DECLARE - result jsonb := '{}'::jsonb; - p text; - pth text[]; - val jsonb; + geom jsonb; + output jsonb; + frag_content jsonb; + frag_links_template jsonb; + merged_assets jsonb; + merged_properties jsonb; + hydrated_stac_version text; + hydrated_stac_extensions jsonb; + hydrated_links jsonb; BEGIN - IF content IS NULL OR fragment_paths IS NULL OR cardinality(fragment_paths) = 0 THEN - RETURN NULL; - END IF; + IF include_field('geometry', fields) THEN + geom := ST_ASGeoJson(_item.geometry, 20)::jsonb; + END IF; + + -- Fetch shared fragment content (NULL when item has no fragment). _skip_fragment lets a caller + -- that has already determined the requested fields are satisfiable from item columns alone + -- (via needs_fragment) avoid this per-row lookup entirely. + IF _item.fragment_id IS NOT NULL AND NOT _skip_fragment THEN + SELECT content, links_template + INTO frag_content, frag_links_template + FROM item_fragments + WHERE id = _item.fragment_id; + END IF; + + -- Merge: fragment provides shared asset/property values; per-item provides individual values. + merged_assets := jsonb_merge_recursive(frag_content->'assets', COALESCE(_item.assets, '{}'::jsonb)); + merged_properties := jsonb_merge_recursive(frag_content->'properties', COALESCE(_item.properties, '{}'::jsonb)); + merged_properties := promoted_properties_from_item(_item) || COALESCE(merged_properties, '{}'::jsonb); + hydrated_stac_version := COALESCE(_item.stac_version, frag_content->>'stac_version'); + hydrated_stac_extensions := CASE + WHEN _item.stac_extensions IS NOT NULL AND _item.stac_extensions <> '[]'::jsonb THEN _item.stac_extensions + ELSE COALESCE(frag_content->'stac_extensions', _item.stac_extensions) + END; + IF _item.fragment_id IS NOT NULL AND NOT _skip_fragment THEN + hydrated_links := stac_links_hydrate(frag_links_template, _item.link_hrefs); + ELSE + hydrated_links := COALESCE(_item.links, '[]'::jsonb); + END IF; + + output := jsonb_build_object( + 'id', _item.id, + 'geometry', geom, + 'collection', _item.collection, + 'type', 'Feature' + ); + IF _item.bbox IS NOT NULL THEN + output := output || jsonb_build_object('bbox', _item.bbox); + END IF; + IF hydrated_stac_version IS NOT NULL THEN + output := output || jsonb_build_object('stac_version', hydrated_stac_version); + END IF; + IF hydrated_stac_extensions IS NOT NULL AND hydrated_stac_extensions <> '[]'::jsonb THEN + output := output || jsonb_build_object('stac_extensions', hydrated_stac_extensions); + END IF; + IF hydrated_links IS NOT NULL THEN + output := output || jsonb_build_object('links', hydrated_links); + END IF; + IF merged_assets != '{}'::jsonb THEN + output := output || jsonb_build_object('assets', merged_assets); + END IF; + IF merged_properties IS NOT NULL THEN + output := output || jsonb_build_object('properties', merged_properties); + END IF; + IF _item.extra IS NOT NULL THEN + output := output || _item.extra; + END IF; + + RETURN jsonb_fields(output, fields); +END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.cql2_collection_set(op text, args jsonb) + RETURNS text[] + LANGUAGE plpgsql + STABLE +AS $function$ +DECLARE lst jsonb; +BEGIN + IF op IN ('=','eq') AND jsonb_typeof(args->1) = 'string' THEN + RETURN ARRAY[args->>1]; + ELSIF op = 'in' THEN + lst := CASE WHEN jsonb_typeof(args->1)='object' AND args->1 ? 'list' THEN args->1->'list' ELSE args->1 END; + IF jsonb_typeof(lst) <> 'array' THEN RETURN NULL; END IF; + RETURN to_text_array(lst); + ELSIF jsonb_typeof(args->1) = 'string' THEN + IF op IN ('<','lt') THEN RETURN coalesce((SELECT array_agg(id) FROM collections WHERE id < (args->>1)), '{}'); + ELSIF op IN ('<=','lte') THEN RETURN coalesce((SELECT array_agg(id) FROM collections WHERE id <= (args->>1)), '{}'); + ELSIF op IN ('>','gt') THEN RETURN coalesce((SELECT array_agg(id) FROM collections WHERE id > (args->>1)), '{}'); + ELSIF op IN ('>=','gte') THEN RETURN coalesce((SELECT array_agg(id) FROM collections WHERE id >= (args->>1)), '{}'); + ELSIF op IN ('like') THEN RETURN coalesce((SELECT array_agg(id) FROM collections WHERE id LIKE (args->>1)), '{}'); + ELSIF op IN ('ilike') THEN RETURN coalesce((SELECT array_agg(id) FROM collections WHERE id ILIKE (args->>1)), '{}'); + END IF; + ELSIF op = 'between' AND jsonb_typeof(args->1)='string' AND jsonb_typeof(args->2)='string' THEN + RETURN coalesce((SELECT array_agg(id) FROM collections WHERE id BETWEEN (args->>1) AND (args->>2)), '{}'); + END IF; + RETURN NULL; +END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.cql2_envelope(j jsonb) + RETURNS pred_envelope + LANGUAGE plpgsql + STABLE +AS $function$ +DECLARE op text; args jsonb; child jsonb; acc pred_envelope; r tstzrange; prop text; v timestamptz; g geometry; +BEGIN + IF j IS NULL OR jsonb_typeof(j) <> 'object' OR NOT j ? 'op' THEN RETURN env_full(); END IF; + op := lower(j->>'op'); args := j->'args'; + IF op = 'and' THEN + acc := env_full(); + FOR child IN SELECT * FROM jsonb_array_elements(args) LOOP acc := env_and(acc, cql2_envelope(child)); END LOOP; + RETURN acc; + ELSIF op = 'or' THEN + acc := NULL; + FOR child IN SELECT * FROM jsonb_array_elements(args) LOOP + acc := CASE WHEN acc IS NULL THEN cql2_envelope(child) ELSE env_or(acc, cql2_envelope(child)) END; + END LOOP; + RETURN coalesce(acc, env_full()); + ELSIF op = 'not' THEN RETURN env_full(); + ELSIF op ILIKE 't_%' OR op = 'anyinteracts' THEN + r := parse_dtrange(args->1); acc := env_full(); + IF op IN ('t_intersects','anyinteracts','t_during','t_equals','t_starts','t_finishes','t_overlaps') THEN + acc.dt := tstzmultirange(tstzrange('-infinity', upper(r), '(]')); + acc.edt := tstzmultirange(tstzrange(lower(r), 'infinity', '[)')); + ELSIF op IN ('t_before','t_meets') THEN + acc.edt := tstzmultirange(tstzrange('-infinity', lower(r), '()')); + ELSIF op IN ('t_after','t_metby') THEN + acc.dt := tstzmultirange(tstzrange(upper(r), 'infinity', '()')); + END IF; + RETURN acc; + ELSIF op ILIKE 's_%' OR op = 'intersects' THEN + BEGIN + g := ST_GeomFromGeoJSON(args->1); + EXCEPTION WHEN others THEN + RAISE EXCEPTION 'Invalid GeoJSON geometry: %', args->1 USING ERRCODE = '22P02'; + END; + acc := env_full(); acc.geom := ST_Envelope(g); RETURN acc; + ELSIF op IN ('=','<','<=','>','>=','between','eq','lt','lte','gt','gte','in','like','ilike') + AND jsonb_typeof(args)='array' AND args->0 ? 'property' THEN + prop := args->0->>'property'; acc := env_full(); + IF prop IN ('datetime','end_datetime') THEN + IF op IN ('in','like','ilike') THEN RETURN env_full(); END IF; + IF op = 'between' THEN r := tstzrange(cql2_ts(args->1), cql2_ts(args->2), '[]'); + ELSIF op IN ('<','<=','lt','lte') THEN r := tstzrange('-infinity', cql2_ts(args->1), '(]'); + ELSIF op IN ('>','>=','gt','gte') THEN r := tstzrange(cql2_ts(args->1), 'infinity', '[)'); + ELSE v := cql2_ts(args->1); r := tstzrange(v, v, '[]'); END IF; + IF prop = 'datetime' THEN acc.dt := tstzmultirange(r); ELSE acc.edt := tstzmultirange(r); END IF; + RETURN acc; + ELSIF prop = 'collection' THEN + acc.colls := cql2_collection_set(op, args); RETURN acc; + ELSE RETURN env_full(); END IF; + ELSE RETURN env_full(); + END IF; +END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.cql2_ts(v jsonb) + RETURNS timestamp with time zone + LANGUAGE sql + IMMUTABLE +AS $function$ + SELECT coalesce(v->>'timestamp', v->>'date', v#>>'{}')::timestamptz; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.create_collection(data jsonb, _partition_trunc text DEFAULT NULL::text, _fragment_config text[] DEFAULT NULL::text[]) + RETURNS void + LANGUAGE sql + SET search_path TO 'pgstac', 'public' +AS $function$ + INSERT INTO collections (content, fragment_config, partition_trunc) + VALUES ( + data, + COALESCE(_fragment_config, collection_fragment_config_default(data)), + _partition_trunc + ) + ; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.ensure_fragments(_collection text, _fragments jsonb[]) + RETURNS TABLE(ord integer, frag_id bigint) + LANGUAGE sql + SECURITY DEFINER +AS $function$ + WITH input AS ( + SELECT + o::int AS ord, + pgstac_hash_fragment( + jsonb_strip_nulls(jsonb_build_object( + 'content', NULLIF(f->'content', '{}'::jsonb), + 'links_template', f->'links_template' + )) + ) AS hash, + COALESCE(f->'content', '{}'::jsonb) AS content, + f->'links_template' AS links_template + FROM unnest(_fragments) WITH ORDINALITY AS u(f, o) + ), + distinct_hashes AS ( + SELECT DISTINCT ON (hash) hash, content, links_template + FROM input + ORDER BY hash + ), + upserted AS ( + INSERT INTO item_fragments (collection, hash, content, links_template) + SELECT _collection, hash, content, links_template FROM distinct_hashes + -- DO UPDATE (a no-op write) rather than DO NOTHING: it locks + RETURNS the conflicting row even when + -- a CONCURRENT transaction committed the same (collection, hash) mid-statement. DO NOTHING returns + -- nothing on conflict, and the old "UNION the existing rows via a separate SELECT" ran on the + -- statement-start snapshot, so it could MISS such a concurrent insert -> that hash dropped out of the + -- final join and the item was stamped with NO fragment_id. DO UPDATE returns every distinct hash + -- exactly once (newly inserted or pre-existing), making the stamp concurrency-safe. + ON CONFLICT (collection, hash) DO UPDATE SET content = item_fragments.content + RETURNING id, hash + ) + SELECT input.ord, upserted.id + FROM input JOIN upserted ON input.hash = upserted.hash + ORDER BY input.ord; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.ensure_partitions(_collections text[], _datetimes timestamp with time zone[], _end_datetimes timestamp with time zone[]) + RETURNS void + LANGUAGE plpgsql + SECURITY DEFINER +AS $function$ +DECLARE + r RECORD; +BEGIN + -- Call check_partition per distinct (collection, partition window) in a DETERMINISTIC order: each + -- check_partition takes that partition's advisory lock, so locking in a consistent order prevents two + -- concurrent multi-partition batches from advisory-deadlocking on the locks. (flush locks in the same + -- sorted order.) A PLPGSQL FOR loop guarantees the order; a set-returning SELECT would not. + FOR r IN + WITH t AS ( + SELECT + unnest(_collections) AS collection, + unnest(_datetimes) AS dt, + unnest(_end_datetimes) AS edt + ), + j AS ( + SELECT t.collection, t.dt, t.edt, c.partition_trunc + FROM t JOIN pgstac.collections c ON t.collection = c.id + ) + SELECT + collection, + tstzrange(min(dt), max(dt), '[]') AS dtrange, + tstzrange(min(edt), max(edt), '[]') AS edtrange + FROM j + GROUP BY collection, COALESCE(date_trunc(partition_trunc::text, dt), '-infinity'::timestamptz) + ORDER BY collection, COALESCE(date_trunc(partition_trunc::text, dt), '-infinity'::timestamptz) + LOOP + PERFORM pgstac.check_partition(r.collection, r.dtrange, r.edtrange); + END LOOP; +END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.env_and(a pred_envelope, b pred_envelope) + RETURNS pred_envelope + LANGUAGE sql + IMMUTABLE +AS $function$ + SELECT ( + CASE WHEN a.colls IS NULL THEN b.colls WHEN b.colls IS NULL THEN a.colls + ELSE array_intersection(a.colls, b.colls) END, + a.dt * b.dt, a.edt * b.edt, + CASE WHEN a.geom IS NULL THEN b.geom WHEN b.geom IS NULL THEN a.geom + ELSE ST_Envelope(ST_Intersection(a.geom, b.geom)) END)::pred_envelope; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.env_full() + RETURNS pred_envelope + LANGUAGE sql + IMMUTABLE +AS $function$ + SELECT (NULL::text[], + tstzmultirange(tstzrange('-infinity','infinity','[]')), + tstzmultirange(tstzrange('-infinity','infinity','[]')), NULL)::pred_envelope; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.env_or(a pred_envelope, b pred_envelope) + RETURNS pred_envelope + LANGUAGE sql + IMMUTABLE +AS $function$ + SELECT ( + CASE WHEN a.colls IS NULL OR b.colls IS NULL THEN NULL + ELSE ARRAY(SELECT DISTINCT unnest(a.colls || b.colls)) END, + a.dt + b.dt, a.edt + b.edt, + CASE WHEN a.geom IS NULL OR b.geom IS NULL THEN NULL + ELSE ST_Envelope(ST_Collect(a.geom, b.geom)) END)::pred_envelope; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.extract_fragment(content jsonb, fragment_paths text[]) + RETURNS jsonb + LANGUAGE plpgsql + IMMUTABLE PARALLEL SAFE +AS $function$ +DECLARE + result jsonb := '{}'::jsonb; + p text; + pth text[]; + val jsonb; +BEGIN + IF content IS NULL OR fragment_paths IS NULL OR cardinality(fragment_paths) = 0 THEN + RETURN NULL; + END IF; FOREACH p IN ARRAY fragment_paths LOOP pth := fragment_path_array(p); @@ -523,13 +845,89 @@ END; $function$ ; -CREATE OR REPLACE FUNCTION pgstac.format_item(_item items, _fields jsonb DEFAULT '{}'::jsonb) - RETURNS jsonb +CREATE OR REPLACE FUNCTION pgstac.field_included(_field text, _includes text[], _excludes text[]) + RETURNS boolean + LANGUAGE sql + IMMUTABLE +AS $function$ + SELECT CASE + WHEN _field = ANY(_excludes) THEN false + WHEN array_length(_includes, 1) IS NOT NULL THEN _field = ANY(_includes) + ELSE true END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.fields_to_itemcols(fields jsonb DEFAULT '{}'::jsonb, _skip_fragment boolean DEFAULT false) + RETURNS text + LANGUAGE plpgsql + STABLE +AS $function$ +DECLARE + includes text[] := ARRAY(SELECT jsonb_array_elements_text(fields->'include')); + excludes text[] := ARRAY(SELECT jsonb_array_elements_text(fields->'exclude')); + cols text; +BEGIN + SELECT string_agg( + CASE + WHEN a.attname = 'fragment_id' AND _skip_fragment + THEN format('NULL::%s AS %I', format_type(a.atttypid, a.atttypmod), a.attname) + WHEN a.attname = ANY (ARRAY['geometry','bbox','assets','links','link_hrefs', + 'extra','properties','stac_version','stac_extensions']) + AND NOT field_included(CASE WHEN a.attname = 'link_hrefs' THEN 'links' ELSE a.attname END, + includes, excludes) + THEN format('NULL::%s AS %I', format_type(a.atttypid, a.atttypmod), a.attname) + ELSE format('i.%I', a.attname) + END, ', ' ORDER BY a.attnum) + INTO cols + FROM pg_attribute a + WHERE a.attrelid = 'items'::regclass AND a.attnum > 0 AND NOT a.attisdropped; + RETURN cols; +END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.flush_items_staging_binary(_staging text, _policy text DEFAULT 'ignore'::text) + RETURNS bigint LANGUAGE plpgsql SECURITY DEFINER AS $function$ +DECLARE + nrows bigint; BEGIN - RETURN content_hydrate(_item, _fields); + IF _policy = 'ignore' THEN + EXECUTE format('INSERT INTO items SELECT * FROM %1$I ON CONFLICT DO NOTHING', _staging); + ELSIF _policy = 'error' THEN + EXECUTE format('INSERT INTO items SELECT * FROM %1$I', _staging); + ELSIF _policy = 'upsert' THEN + -- Fast SAME-partition upsert: delete a changed row only in the partition the incoming row routes to. + -- The window range [date_trunc(trunc, s.datetime), + 1 trunc) is derived per staged row, so the + -- planner runtime-prunes `items` to that one partition (no cross-partition scan). A datetime change + -- within the partition is applied; one that MOVES the item to another partition isn't seen here, so + -- the old row orphans (use 'delsert'). NULL partition_trunc => the collection's single partition. + EXECUTE format($q$ + DELETE FROM items i USING %1$I s JOIN collections c ON c.id = s.collection + WHERE i.collection = s.collection AND i.id = s.id + AND i.datetime >= (CASE WHEN c.partition_trunc IS NULL THEN '-infinity'::timestamptz + ELSE date_trunc(c.partition_trunc::text, s.datetime) END) + AND i.datetime < (CASE WHEN c.partition_trunc IS NULL THEN 'infinity'::timestamptz + ELSE date_trunc(c.partition_trunc::text, s.datetime) + + ('1 ' || c.partition_trunc::text)::interval END) + AND ( %2$s ) + $q$, _staging, items_content_distinct_sql('i', 's')); + EXECUTE format('INSERT INTO items SELECT * FROM %1$I ON CONFLICT DO NOTHING', _staging); + ELSIF _policy = 'delsert' THEN + -- Move-safe CROSS-partition upsert: delete the old row wherever it lives (by collection+id, no + -- datetime bound, so it probes every partition), then insert. + EXECUTE format($q$ + DELETE FROM items i USING %1$I s + WHERE i.id = s.id AND i.collection = s.collection AND ( %2$s ) + $q$, _staging, items_content_distinct_sql('i', 's')); + EXECUTE format('INSERT INTO items SELECT * FROM %1$I ON CONFLICT DO NOTHING', _staging); + ELSE + RAISE EXCEPTION 'unknown conflict policy % (expected ignore | upsert | delsert | error)', _policy; + END IF; + GET DIAGNOSTICS nrows = ROW_COUNT; + RETURN nrows; END; $function$ ; @@ -537,21 +935,12 @@ $function$ CREATE OR REPLACE FUNCTION pgstac.fragment_path_array(_path_text text) RETURNS text[] LANGUAGE plpgsql - IMMUTABLE PARALLEL SAFE STRICT SECURITY DEFINER + IMMUTABLE PARALLEL SAFE STRICT AS $function$ BEGIN - IF _path_text ~ '^\s*\[' THEN - RETURN ARRAY( - SELECT jsonb_array_elements_text(_path_text::jsonb) - ); - END IF; - - -- Legacy support for pre-JSON serialization. - RETURN string_to_array(_path_text, '.'); -EXCEPTION - WHEN others THEN - -- Be permissive for existing operator-configured values. - RETURN string_to_array(_path_text, '.'); + RETURN ARRAY( + SELECT jsonb_array_elements_text(_path_text::jsonb) + ); END; $function$ ; @@ -559,36 +948,12 @@ $function$ CREATE OR REPLACE FUNCTION pgstac.fragment_path_text(_path text[]) RETURNS text LANGUAGE sql - IMMUTABLE PARALLEL SAFE STRICT SECURITY DEFINER + IMMUTABLE PARALLEL SAFE STRICT AS $function$ SELECT to_jsonb(_path)::text; $function$ ; -CREATE OR REPLACE FUNCTION pgstac.gc_anonymous_searches(retention_interval interval DEFAULT NULL::interval, conf jsonb DEFAULT NULL::jsonb) - RETURNS bigint - LANGUAGE sql - SECURITY DEFINER -AS $function$ - WITH effective_retention AS ( - SELECT COALESCE( - retention_interval, - search_gc_retention_interval(conf) - ) AS i - ), - deleted AS ( - DELETE FROM searches - USING effective_retention - WHERE - name IS NULL - AND NOT pinned - AND lastused < now() - effective_retention.i - RETURNING 1 - ) - SELECT count(*)::bigint FROM deleted; -$function$ -; - CREATE OR REPLACE FUNCTION pgstac.gc_deleted_items_log(retention_interval interval DEFAULT '30 days'::interval) RETURNS bigint LANGUAGE sql @@ -680,22 +1045,10 @@ AS $function$ $function$ ; -CREATE OR REPLACE FUNCTION pgstac.gc_search_caches(retention_interval interval DEFAULT NULL::interval, conf jsonb DEFAULT NULL::jsonb) - RETURNS jsonb - LANGUAGE sql - SECURITY DEFINER -AS $function$ - SELECT jsonb_build_object( - 'removed_searches', - gc_anonymous_searches(retention_interval, conf) - ); -$function$ -; - CREATE OR REPLACE FUNCTION pgstac.items_content_changed(left_item items, right_item items) RETURNS boolean LANGUAGE plpgsql - IMMUTABLE PARALLEL SAFE SECURITY DEFINER + IMMUTABLE PARALLEL SAFE AS $function$ DECLARE changed boolean; @@ -712,7 +1065,7 @@ $function$ CREATE OR REPLACE FUNCTION pgstac.items_content_distinct_sql(left_ref text, right_ref text) RETURNS text LANGUAGE plpgsql - IMMUTABLE PARALLEL SAFE SECURITY DEFINER + IMMUTABLE PARALLEL SAFE AS $function$ DECLARE clauses text[]; @@ -721,15 +1074,16 @@ BEGIN format('%s.datetime_is_range IS DISTINCT FROM %s.datetime_is_range', left_ref, right_ref), format('%s.datetime IS DISTINCT FROM %s.datetime', left_ref, right_ref), format('%s.end_datetime IS DISTINCT FROM %s.end_datetime', left_ref, right_ref), + format('%s.item_hash IS DISTINCT FROM %s.item_hash', left_ref, right_ref), format('%s.geometry IS DISTINCT FROM %s.geometry', left_ref, right_ref), format('%s.bbox IS DISTINCT FROM %s.bbox', left_ref, right_ref), format('%s.links IS DISTINCT FROM %s.links', left_ref, right_ref), + format('%s.link_hrefs IS DISTINCT FROM %s.link_hrefs', left_ref, right_ref), format('%s.assets IS DISTINCT FROM %s.assets', left_ref, right_ref), format('%s.properties IS DISTINCT FROM %s.properties', left_ref, right_ref), format('%s.extra IS DISTINCT FROM %s.extra', left_ref, right_ref), format('%s.stac_version IS DISTINCT FROM %s.stac_version', left_ref, right_ref), - format('%s.stac_extensions IS DISTINCT FROM %s.stac_extensions', left_ref, right_ref), - format('%s.fragment_id IS DISTINCT FROM %s.fragment_id', left_ref, right_ref) + format('%s.stac_extensions IS DISTINCT FROM %s.stac_extensions', left_ref, right_ref) ]; clauses := clauses || ARRAY( @@ -754,7 +1108,7 @@ BEGIN partition, datetime, end_datetime, - content_hash + item_hash ) SELECT old_rows.id, @@ -762,7 +1116,7 @@ BEGIN (partition_name(old_rows.collection, old_rows.datetime)).partition_name, old_rows.datetime, old_rows.end_datetime, - old_rows.content_hash + old_rows.item_hash FROM old_rows; RETURN NULL; @@ -770,6 +1124,142 @@ END; $function$ ; +CREATE OR REPLACE FUNCTION pgstac.items_staging_dehydrate(_contents jsonb[]) + RETURNS SETOF items + LANGUAGE sql +AS $function$ + WITH raw AS MATERIALIZED ( + SELECT + n.content AS orig_content, + d.* + FROM unnest(_contents) AS n(content) + CROSS JOIN LATERAL content_dehydrate(n.content) d + ), + fragmented_base AS MATERIALIZED ( + SELECT + r.*, + c.fragment_config, + stac_links_strip_hrefs(r.orig_content->'links') AS links_template, + extract_fragment(r.orig_content, c.fragment_config) AS frag_content + FROM raw r + JOIN collections c ON c.id = r.collection + ), + fragmented AS MATERIALIZED ( + SELECT + fb.*, + CASE + WHEN fb.frag_content IS NULL + AND fb.links_template IS NULL THEN NULL + ELSE pgstac_hash_fragment( + jsonb_strip_nulls( + jsonb_build_object( + 'content', NULLIF(fb.frag_content, '{}'::jsonb), + 'links_template', fb.links_template + ) + ) + ) + END AS frag_hash + FROM fragmented_base fb + ), + fragments AS MATERIALIZED ( + SELECT DISTINCT ON (collection, frag_hash) + collection, + frag_hash, + frag_content, + fragment_config, + links_template + FROM fragmented + WHERE frag_hash IS NOT NULL + ORDER BY collection, frag_hash + ), + insert_fragments AS ( + INSERT INTO item_fragments (collection, hash, content, links_template) + SELECT collection, frag_hash, COALESCE(frag_content, '{}'::jsonb), links_template + FROM fragments + ON CONFLICT (collection, hash) DO NOTHING + RETURNING id, collection, hash + ), + all_fragments AS ( + SELECT id, collection, hash FROM insert_fragments + UNION ALL + SELECT f.id, f.collection, f.hash + FROM item_fragments f + JOIN fragments p ON f.collection = p.collection AND f.hash = p.frag_hash + ), + enriched AS MATERIALIZED ( + SELECT + r.id, + r.geometry, + r.collection, + r.datetime, + r.end_datetime, + r.datetime_is_range, + CASE + WHEN fragment_path_text(ARRAY['stac_version']) = ANY(f.fragment_config) THEN NULL + ELSE r.stac_version + END AS stac_version, + CASE + WHEN fragment_path_text(ARRAY['stac_extensions']) = ANY(f.fragment_config) THEN '[]'::jsonb + ELSE r.stac_extensions + END AS stac_extensions, + r.pgstac_updated_at, + r.item_hash, + af.id AS fragment_id, + r.bbox, + CASE + WHEN r.link_hrefs IS NOT NULL AND array_length(r.link_hrefs, 1) > 0 THEN NULL + ELSE r.links + END AS links, + strip_fragment_col(COALESCE(r.assets, '{}'::jsonb), 'assets', f.fragment_config) AS assets, + strip_fragment_col(COALESCE(r.properties, '{}'::jsonb), 'properties', f.fragment_config) AS properties, + r.extra, + r.created, + r.updated, + r.platform, + r.instruments, + r.constellation, + r.mission, + r.eo_cloud_cover, + r.bands, + r.eo_snow_cover, + r.gsd, + r.proj_code, + r.proj_geometry, + r.proj_wkt2, + r.proj_projjson, + r.proj_bbox, + r.proj_centroid, + r.proj_shape, + r.proj_transform, + r.sci_doi, + r.sci_citation, + r.sci_publications, + r.view_off_nadir, + r.view_incidence_angle, + r.view_azimuth, + r.view_sun_azimuth, + r.view_sun_elevation, + r.view_moon_azimuth, + r.view_moon_elevation, + r.file_size, + r.file_header_size, + r.file_checksum, + r.file_byte_order, + r.sat_orbit_state, + r.sat_relative_orbit, + r.sat_absolute_orbit, + r.sat_platform_international_designator, + r.sat_anx_datetime, + r.link_hrefs, + NULL::jsonb AS private + FROM fragmented r + LEFT JOIN fragments f ON f.collection = r.collection AND f.frag_hash = r.frag_hash + LEFT JOIN all_fragments af ON af.collection = f.collection AND af.hash = f.frag_hash + ) + SELECT * FROM enriched; +$function$ +; + CREATE OR REPLACE FUNCTION pgstac.items_touch_triggerfunc() RETURNS trigger LANGUAGE plpgsql @@ -781,12 +1271,48 @@ BEGIN END IF; NEW.pgstac_updated_at := now(); - NEW.content_hash := pgstac_item_hash(content_hydrate(NEW)); RETURN NEW; END; $function$ ; +CREATE OR REPLACE FUNCTION pgstac.jsonb_canonical_hash(j jsonb) + RETURNS bytea + LANGUAGE sql + IMMUTABLE PARALLEL SAFE STRICT + SET "TimeZone" TO 'UTC' +AS $function$ + WITH RECURSIVE t(path, value) AS ( + SELECT ''::text, j + UNION ALL + SELECT t.path || E'\x1F' || e.seg, e.value + FROM t CROSS JOIN LATERAL ( + SELECT 'k' || key AS seg, value + FROM jsonb_each(t.value) WHERE jsonb_typeof(t.value) = 'object' + UNION ALL + SELECT 'i' || (ord - 1)::text AS seg, value + FROM jsonb_array_elements(t.value) WITH ORDINALITY x(value, ord) + WHERE jsonb_typeof(t.value) = 'array' + ) e + ) + SELECT sha256(convert_to(COALESCE(string_agg( + t.path || E'\x1E' || CASE jsonb_typeof(t.value) + WHEN 'number' THEN 'n' || (t.value #>> '{}')::float8::text + WHEN 'boolean' THEN 'b' || (t.value #>> '{}') + WHEN 'null' THEN 'z' + WHEN 'string' THEN CASE + WHEN (t.value #>> '{}') ~ '^\d{4}-\d{2}-\d{2}([Tt ]\d{2}:\d{2}:\d{2}(\.\d+)?([Zz]|[+-]\d{2}:?\d{2})?)?$' + THEN 't' || to_char((t.value #>> '{}')::timestamptz AT TIME ZONE 'UTC', 'YYYY-MM-DD"T"HH24:MI:SS.US') || 'Z' + ELSE 's' || (t.value #>> '{}') END + WHEN 'object' THEN 'e' + WHEN 'array' THEN 'a' + ELSE 'c' END, + E'\x1D' ORDER BY t.path COLLATE "C"), ''), 'UTF8')) + FROM t + WHERE jsonb_typeof(t.value) NOT IN ('object', 'array') OR t.value = '{}'::jsonb OR t.value = '[]'::jsonb; +$function$ +; + CREATE OR REPLACE FUNCTION pgstac.jsonb_common_paths_final(docs jsonb[]) RETURNS text[] LANGUAGE sql @@ -839,7 +1365,7 @@ $function$ CREATE OR REPLACE FUNCTION pgstac.jsonb_common_values(left_doc jsonb, right_doc jsonb) RETURNS jsonb LANGUAGE sql - IMMUTABLE PARALLEL SAFE SECURITY DEFINER + IMMUTABLE PARALLEL SAFE AS $function$ SELECT CASE WHEN left_doc IS NULL OR right_doc IS NULL THEN NULL @@ -908,7 +1434,7 @@ $function$ CREATE OR REPLACE FUNCTION pgstac.jsonb_leaf_rows(data jsonb, parent_path text) RETURNS TABLE(path text, value jsonb) LANGUAGE plpgsql - IMMUTABLE PARALLEL SAFE SECURITY DEFINER + IMMUTABLE PARALLEL SAFE AS $function$ DECLARE key_text text; @@ -944,168 +1470,416 @@ END; $function$ ; -CREATE OR REPLACE FUNCTION pgstac.jsonb_merge_level1(frag jsonb, item jsonb) +CREATE OR REPLACE FUNCTION pgstac.jsonb_merge_recursive(frag jsonb, item jsonb) RETURNS jsonb LANGUAGE sql - IMMUTABLE PARALLEL SAFE SECURITY DEFINER + IMMUTABLE PARALLEL SAFE AS $function$ - SELECT COALESCE( - (SELECT jsonb_object_agg( - COALESCE(f.key, i.key), - CASE - WHEN i.value IS NULL THEN f.value - WHEN f.value IS NULL THEN i.value - WHEN jsonb_typeof(f.value) = 'object' AND jsonb_typeof(i.value) = 'object' - THEN f.value || i.value - ELSE i.value - END - ) - FROM - jsonb_each(COALESCE(frag, '{}')) f - FULL JOIN jsonb_each(COALESCE(item, '{}')) i USING (key) - ), - '{}'::jsonb - ); + SELECT CASE + WHEN frag IS NULL THEN COALESCE(item, '{}'::jsonb) + WHEN item IS NULL OR item = '{}'::jsonb THEN frag + WHEN jsonb_typeof(frag) = 'object' AND jsonb_typeof(item) = 'object' THEN + COALESCE( + ( + SELECT jsonb_object_agg( + key, + CASE + WHEN i.value IS NULL THEN f.value + WHEN f.value IS NULL THEN i.value + WHEN jsonb_typeof(f.value) = 'object' AND jsonb_typeof(i.value) = 'object' THEN + CASE + WHEN NOT EXISTS ( + SELECT 1 FROM jsonb_object_keys(f.value) k WHERE i.value ? k + ) THEN f.value || i.value + ELSE jsonb_merge_recursive(f.value, i.value) + END + ELSE i.value + END + ) + FROM jsonb_each(frag) f + FULL JOIN jsonb_each(item) i USING (key) + ), + '{}'::jsonb + ) + ELSE item + END; $function$ ; -CREATE OR REPLACE FUNCTION pgstac.name_search(_search jsonb, _name text, _metadata jsonb DEFAULT '{}'::jsonb) - RETURNS searches +CREATE OR REPLACE FUNCTION pgstac.keyset_decode(token text) + RETURNS text[] LANGUAGE plpgsql - SECURITY DEFINER + IMMUTABLE AS $function$ -DECLARE - named searches%ROWTYPE; BEGIN - named := search_query(_search, false, _metadata); - UPDATE searches - SET - name = _name, - lastused = now(), - usecount = searches.usecount + 1 - WHERE hash = named.hash - RETURNING * INTO named; - - IF named IS NULL THEN - RAISE EXCEPTION 'Could not name search for input: %', _search; - END IF; - - RETURN named; + IF token IS NULL OR token = '' THEN RETURN NULL; END IF; + RETURN array_replace( + string_to_array(convert_from(decode(token,'base64'),'UTF8'), chr(31)), + chr(30), NULL); +EXCEPTION WHEN others THEN + -- A non-empty token that does not decode is a client error, not an empty page. + RAISE EXCEPTION 'Invalid pagination token: %', token USING ERRCODE = '22P02'; END; $function$ ; -CREATE OR REPLACE FUNCTION pgstac.pgstac_hash(data text) +CREATE OR REPLACE FUNCTION pgstac.keyset_encode(vals text[]) RETURNS text LANGUAGE sql - IMMUTABLE PARALLEL SAFE STRICT + IMMUTABLE AS $function$ - SELECT encode(sha256(convert_to(data, 'UTF8')), 'hex'); + SELECT encode(convert_to(array_to_string(vals, chr(31), chr(30)), 'UTF8'), 'base64'); $function$ ; -CREATE OR REPLACE FUNCTION pgstac.pgstac_hash_fragment(fragment jsonb) +CREATE OR REPLACE FUNCTION pgstac.keyset_orderby(_search jsonb, _prev boolean DEFAULT false) RETURNS text LANGUAGE sql - IMMUTABLE PARALLEL SAFE SECURITY DEFINER + STABLE AS $function$ -SELECT pgstac_hash(fragment::text); + SELECT string_agg( + expr || ' ' || CASE WHEN _prev + THEN (CASE dir WHEN 'ASC' THEN 'DESC' ELSE 'ASC' END) ELSE dir END, + ', ' ORDER BY ord) + FROM keyset_sortkeys(_search); $function$ ; -CREATE OR REPLACE FUNCTION pgstac.pgstac_item_hash(item_json jsonb) - RETURNS text +CREATE OR REPLACE FUNCTION pgstac.keyset_sortkeys(_search jsonb) + RETURNS TABLE(ord integer, field text, expr text, dir text, "notnull" boolean) LANGUAGE sql - IMMUTABLE PARALLEL SAFE STRICT SECURITY DEFINER + STABLE AS $function$ - SELECT encode(sha256( - (SELECT jsonb_object_agg(key, value ORDER BY key) - FROM jsonb_each(item_json))::text::bytea - ), 'hex'); + WITH base AS ( + SELECT coalesce(_search->'sortby','[{"field":"datetime","direction":"desc"}]'::jsonb) AS s + ), + app AS ( + SELECT s + || jsonb_build_object('field','id','direction', s->0->>'direction') + || jsonb_build_object('field','collection','direction', s->0->>'direction') AS s + FROM base + ), + rows AS ( + SELECT value->>'field' AS field, get_sort_dir(value) AS dir, o + FROM app, jsonb_array_elements(s) WITH ORDINALITY AS t(value, o) + ), + firsts AS ( + SELECT DISTINCT ON (field) field, dir, o FROM rows ORDER BY field, o + ) + SELECT (row_number() OVER (ORDER BY o))::int, field, (queryable(field)).expression, dir, + field IN ('id', 'collection', 'datetime', 'end_datetime') + FROM firsts ORDER BY o; $function$ ; -CREATE OR REPLACE FUNCTION pgstac.pin_search(_name text) - RETURNS searches +CREATE OR REPLACE FUNCTION pgstac.keyset_where(_search jsonb, _values text[], prev boolean DEFAULT false) + RETURNS text + LANGUAGE plpgsql + STABLE +AS $function$ +DECLARE + k record; vlit text; orterm text; + andfilters text[] := '{}'::text[]; + orfilters text[] := '{}'::text[]; +BEGIN + IF _values IS NULL THEN RETURN NULL; END IF; + FOR k IN SELECT * FROM keyset_sortkeys(_search) ORDER BY ord LOOP + vlit := CASE WHEN _values[k.ord] IS NULL THEN NULL ELSE quote_literal(_values[k.ord]) END; + orterm := NULL; + IF vlit IS NOT NULL AND ((prev AND k.dir='ASC') OR (NOT prev AND k.dir='DESC')) THEN + orterm := format('(%s < %s)', k.expr, vlit); + ELSIF vlit IS NULL AND ((prev AND k.dir='ASC') OR (NOT prev AND k.dir='DESC')) THEN + orterm := format('(%s IS NOT NULL)', k.expr); + ELSIF vlit IS NULL THEN + orterm := NULL; + ELSIF k.notnull THEN + orterm := format('(%s > %s)', k.expr, vlit); + ELSE + orterm := format('((%s > %s) OR (%s IS NULL))', k.expr, vlit, k.expr); + END IF; + IF orterm IS NOT NULL THEN + IF array_length(andfilters,1) IS NULL THEN + orfilters := orfilters || orterm; + ELSE + orfilters := orfilters || format('(%s AND %s)', array_to_string(andfilters,' AND '), orterm); + END IF; + END IF; + andfilters := andfilters || CASE WHEN vlit IS NULL + THEN format('(%s IS NULL)', k.expr) ELSE format('(%s = %s)', k.expr, vlit) END; + END LOOP; + IF array_length(orfilters,1) IS NULL THEN RETURN NULL; END IF; + RETURN '(' || array_to_string(orfilters, ' OR ') || ')'; +END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.make_binary_staging() + RETURNS text LANGUAGE plpgsql SECURITY DEFINER AS $function$ DECLARE - pinned_search searches%ROWTYPE; + _name text := format('_staging_%s', replace(gen_random_uuid()::text, '-', '')); BEGIN - UPDATE searches - SET - pinned = true, - lastused = now(), - usecount = searches.usecount + 1 - WHERE name = _name - RETURNING * INTO pinned_search; + EXECUTE format('CREATE TEMP TABLE %I (LIKE items) ON COMMIT DROP', _name); + EXECUTE format('GRANT INSERT ON pg_temp.%I TO pgstac_ingest', _name); + RETURN _name; +END; +$function$ +; - IF pinned_search IS NULL THEN - RAISE EXCEPTION 'Named search % not found', _name; +CREATE OR REPLACE FUNCTION pgstac.needs_fragment(fields jsonb, _colls text[]) + RETURNS boolean + LANGUAGE plpgsql + STABLE PARALLEL SAFE +AS $function$ +DECLARE + includes text[] := ARRAY(SELECT jsonb_array_elements_text(fields->'include')); + promoted text[] := ARRAY['datetime','start_datetime','end_datetime'] + || ARRAY(SELECT name FROM promoted_item_property_defs()); + inc text; propkey text; prop_paths text[] := ARRAY[]::text[]; +BEGIN + IF array_length(includes, 1) IS NULL THEN RETURN true; END IF; + IF _colls IS NULL THEN RETURN true; END IF; + FOREACH inc IN ARRAY includes LOOP + IF inc IN ('assets','links','stac_version','stac_extensions','properties') + OR inc LIKE 'assets.%' OR inc LIKE 'links.%' THEN + RETURN true; + ELSIF inc LIKE 'properties.%' THEN + propkey := substring(inc FROM 'properties\\.(.*)'); + IF NOT (propkey = ANY (promoted)) THEN RETURN true; END IF; + prop_paths := prop_paths || propkey; + END IF; + END LOOP; + IF array_length(prop_paths, 1) IS NOT NULL AND EXISTS ( + SELECT 1 FROM collections c, unnest(c.fragment_config) cfg(p), + LATERAL (SELECT fragment_path_array(cfg.p) AS parr) z + WHERE c.id = ANY (_colls) + AND z.parr[1] = 'properties' + AND (array_length(z.parr, 1) = 1 OR z.parr[2] = ANY (prop_paths)) + ) THEN RETURN true; END IF; + RETURN false; +END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.next_band(_counts bigint[], _cursor_idx integer, _target numeric, _cap_months integer, _descending boolean DEFAULT false, OUT band_start_idx integer, OUT band_end_idx integer, OUT scanned bigint, OUT next_cursor_idx integer, OUT done boolean) + RETURNS record + LANGUAGE plpgsql + STABLE +AS $function$ +DECLARE + idx int; + n int; + cumulative bigint := 0; +BEGIN + done := false; scanned := 0; + band_start_idx := NULL; band_end_idx := NULL; next_cursor_idx := _cursor_idx; + + IF _counts IS NULL OR array_length(_counts, 1) IS NULL OR _cursor_idx IS NULL THEN + done := true; -- no histogram / no cursor => nothing to walk + RETURN; + END IF; + n := array_length(_counts, 1); + + IF _descending THEN + -- Walk downward (newest -> oldest). The high bound is the cursor (clamped into range). + IF _cursor_idx < 1 THEN done := true; RETURN; END IF; + idx := LEAST(_cursor_idx, n); + FOR i IN REVERSE idx..GREATEST(1, idx - _cap_months + 1) LOOP + cumulative := cumulative + _counts[i]; + scanned := scanned + _counts[i]; + IF cumulative >= _target THEN + band_start_idx := i; -- low (older) month bound + band_end_idx := idx; -- high (newer) month bound + next_cursor_idx := i - 1; + IF next_cursor_idx < 1 THEN done := true; END IF; + RETURN; + END IF; + END LOOP; + -- Target not reached within the cap: end the band at the cap boundary. + band_start_idx := GREATEST(1, idx - _cap_months + 1); + band_end_idx := idx; + next_cursor_idx := band_start_idx - 1; + done := (band_start_idx <= 1); + RETURN; END IF; - RETURN pinned_search; + -- Ascending (oldest -> newest). + idx := NULL; + FOR i IN 1..n LOOP + IF i >= _cursor_idx THEN idx := i; EXIT; END IF; + END LOOP; + IF idx IS NULL THEN done := true; next_cursor_idx := _cursor_idx; RETURN; END IF; + + FOR i IN idx..n LOOP + IF i > idx + _cap_months - 1 THEN EXIT; END IF; + cumulative := cumulative + _counts[i]; + scanned := scanned + _counts[i]; + IF cumulative >= _target THEN + band_start_idx := idx; + band_end_idx := i; + next_cursor_idx := i + 1; + IF next_cursor_idx > n THEN done := true; END IF; + RETURN; + END IF; + END LOOP; + + -- Target not reached within the cap: end the band at the cap boundary (not the array end), + -- so the cap actually limits band width. done only when we've consumed the whole histogram. + band_start_idx := idx; + band_end_idx := LEAST(idx + _cap_months - 1, n); + next_cursor_idx := band_end_idx + 1; + done := (band_end_idx >= n); +END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.partition_bounds(_env pred_envelope, OUT months timestamp with time zone[], OUT counts bigint[], OUT collections text[], OUT total_count bigint) + RETURNS record + LANGUAGE sql + STABLE +AS $function$ + WITH cand AS ( + SELECT ps.collection, + lower(ps.dtrange) AS lo, + upper(ps.dtrange) AS hi, + coalesce(ps.n, 0) AS n + FROM partition_stats ps + WHERE ((_env).colls IS NULL OR ps.collection = ANY((_env).colls)) + AND (_env).dt && COALESCE(ps.dtrange, tstzrange('-infinity','infinity','[]')) + AND (_env).edt && COALESCE(ps.edtrange, tstzrange('-infinity','infinity','[]')) + AND ((_env).geom IS NULL OR ps.spatial IS NULL OR ps.spatial && (_env).geom) + ), + monthly AS ( + SELECT date_trunc('month', gs) AS month_start, + CASE + WHEN c.hi <= c.lo THEN c.n::numeric + ELSE c.n * ( + extract(epoch FROM ( + LEAST(c.hi, date_trunc('month', gs) + interval '1 month') + - GREATEST(c.lo, date_trunc('month', gs)))) + / extract(epoch FROM (c.hi - c.lo))) + END AS pn + FROM cand c, + generate_series(date_trunc('month', c.lo), + date_trunc('month', GREATEST(c.lo, c.hi - interval '1 microsecond')), + interval '1 month') AS gs + ), + buckets AS ( + SELECT month_start, round(sum(pn))::bigint AS n + FROM monthly + GROUP BY month_start + ) + SELECT + (SELECT array_agg(month_start ORDER BY month_start) FROM buckets), + (SELECT array_agg(n ORDER BY month_start) FROM buckets), + (SELECT array_agg(DISTINCT collection) FROM cand), + (SELECT coalesce(sum(n), 0) FROM cand); +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.pgstac_hash(data text) + RETURNS text + LANGUAGE sql + IMMUTABLE PARALLEL SAFE STRICT +AS $function$ + SELECT encode(sha256(convert_to(data, 'UTF8')), 'hex'); +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.pgstac_hash_fragment(fragment jsonb) + RETURNS bytea + LANGUAGE sql + IMMUTABLE PARALLEL SAFE +AS $function$ +SELECT sha256(convert_to(fragment::text, 'UTF8')); +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.prepare_partition_for_load(_collection text, _dt_lo timestamp with time zone, _dt_hi timestamp with time zone, _edt_lo timestamp with time zone, _edt_hi timestamp with time zone, _xmin double precision, _ymin double precision, _xmax double precision, _ymax double precision, _n_add bigint, OUT partition_name text, OUT pre_load_n bigint) + RETURNS record + LANGUAGE plpgsql + SECURITY DEFINER +AS $function$ +BEGIN + partition_name := pgstac.check_partition( + _collection, + tstzrange(_dt_lo, _dt_hi, '[]'), + tstzrange(_edt_lo, _edt_hi, '[]'), + st_setsrid(st_makeenvelope(_xmin, _ymin, _xmax, _ymax), 4326) + ); + SELECT COALESCE(n, 0) INTO pre_load_n + FROM pgstac.partition_stats WHERE partition = partition_name; + -- Over-estimating n is the safe direction; the async tightener computes the exact count + extent off the + -- hot path. Single-row atomic UPDATE, so concurrent loads into the same partition serialize on the row. + UPDATE pgstac.partition_stats + SET n = COALESCE(n, 0) + _n_add, dirty = true, last_updated = now() + WHERE partition = partition_name; END; $function$ ; CREATE OR REPLACE FUNCTION pgstac.promoted_item_property_defs() - RETURNS TABLE(name text, definition jsonb, property_path text, property_wrapper text) + RETURNS TABLE(name text, definition jsonb, property_path text) LANGUAGE sql IMMUTABLE PARALLEL SAFE AS $function$ SELECT * FROM (VALUES - ('created', '{"description": "Metadata creation timestamp","type": "string","format": "date-time","title": "Created"}'::jsonb, 'created', 'to_tstz'), - ('updated', '{"description": "Metadata update timestamp","type": "string","format": "date-time","title": "Updated"}'::jsonb, 'updated', 'to_tstz'), - ('platform', '{"description": "Platform name","type": "string","title": "Platform"}'::jsonb, 'platform', 'to_text'), - ('instruments', '{"description": "Instrument names","type": "array","title": "Instruments"}'::jsonb, 'instruments', 'to_text_array'), - ('constellation', '{"description": "Constellation name","type": "string","title": "Constellation"}'::jsonb, 'constellation', 'to_text'), - ('mission', '{"description": "Mission name","type": "string","title": "Mission"}'::jsonb, 'mission', 'to_text'), - ('eo:cloud_cover', '{"description": "EO cloud cover percentage","type": "number","title": "Cloud Cover"}'::jsonb, 'eo_cloud_cover', 'to_float'), - ('eo:bands', '{"description": "EO band metadata","type": "array","title": "EO Bands"}'::jsonb, 'eo_bands', 'to_text'), - ('eo:snow_cover', '{"description": "EO snow cover percentage","type": "number","title": "Snow Cover"}'::jsonb, 'eo_snow_cover', 'to_float'), - ('gsd', '{"description": "Ground sample distance","type": "number","title": "Ground Sample Distance"}'::jsonb, 'gsd', 'to_float'), - ('proj:epsg', '{"description": "EPSG code","type": "integer","title": "Projection EPSG"}'::jsonb, 'proj_epsg', 'to_int'), - ('proj:wkt2', '{"description": "WKT2 CRS definition","type": "string","title": "Projection WKT2"}'::jsonb, 'proj_wkt2', 'to_text'), - ('proj:projjson', '{"description": "PROJJSON CRS definition","type": ["object", "string"],"title": "Projection PROJJSON"}'::jsonb, 'proj_projjson', 'to_text'), - ('proj:bbox', '{"description": "Projection bbox","type": "array","title": "Projection BBOX"}'::jsonb, 'proj_bbox', 'to_text'), - ('proj:centroid', '{"description": "Projection centroid","type": "object","title": "Projection Centroid"}'::jsonb, 'proj_centroid', 'to_text'), - ('proj:shape', '{"description": "Projection shape","type": "array","title": "Projection Shape"}'::jsonb, 'proj_shape', 'to_text'), - ('proj:transform', '{"description": "Projection affine transform","type": "array","title": "Projection Transform"}'::jsonb, 'proj_transform', 'to_text'), - ('sci:doi', '{"description": "Scientific DOI","type": "string","title": "Scientific DOI"}'::jsonb, 'sci_doi', 'to_text'), - ('sci:citation', '{"description": "Scientific citation","type": "string","title": "Scientific Citation"}'::jsonb, 'sci_citation', 'to_text'), - ('sci:publications', '{"description": "Scientific publications","type": "array","title": "Scientific Publications"}'::jsonb, 'sci_publications', 'to_text'), - ('view:off_nadir', '{"description": "Viewing angle off nadir","type": "number","title": "View Off Nadir"}'::jsonb, 'view_off_nadir', 'to_float'), - ('view:incidence_angle','{"description": "View incidence angle","type": "number","title": "View Incidence Angle"}'::jsonb, 'view_incidence_angle', 'to_float'), - ('view:azimuth', '{"description": "View azimuth angle","type": "number","title": "View Azimuth"}'::jsonb, 'view_azimuth', 'to_float'), - ('view:sun_azimuth', '{"description": "Sun azimuth angle","type": "number","title": "View Sun Azimuth"}'::jsonb, 'view_sun_azimuth', 'to_float'), - ('view:sun_elevation', '{"description": "Sun elevation angle","type": "number","title": "View Sun Elevation"}'::jsonb, 'view_sun_elevation', 'to_float'), - ('file:size', '{"description": "File size in bytes","type": "integer","title": "File Size"}'::jsonb, 'file_size', 'to_int'), - ('file:header_size', '{"description": "File header size in bytes","type": "integer","title": "File Header Size"}'::jsonb, 'file_header_size', 'to_int'), - ('file:checksum', '{"description": "File checksum","type": "string","title": "File Checksum"}'::jsonb, 'file_checksum', 'to_text'), - ('file:byte_order', '{"description": "File byte order","type": "string","title": "File Byte Order"}'::jsonb, 'file_byte_order', 'to_text'), - ('file:values_regex', '{"description": "File values regex","type": "string","title": "File Values Regex"}'::jsonb, 'file_values_regex', 'to_text'), - ('sat:orbit_state', '{"description": "Satellite orbit state","type": "string","title": "Orbit State"}'::jsonb, 'sat_orbit_state', 'to_text'), - ('sat:relative_orbit', '{"description": "Satellite relative orbit","type": "integer","title": "Relative Orbit"}'::jsonb, 'sat_relative_orbit', 'to_int'), - ('sat:absolute_orbit', '{"description": "Satellite absolute orbit","type": "integer","title": "Absolute Orbit"}'::jsonb, 'sat_absolute_orbit', 'to_int') - ) AS t(name, definition, property_path, property_wrapper); + ('created', '{"description": "Metadata creation timestamp","type": "string","format": "date-time","title": "Created"}'::jsonb, 'created'), + ('updated', '{"description": "Metadata update timestamp","type": "string","format": "date-time","title": "Updated"}'::jsonb, 'updated'), + ('platform', '{"description": "Platform name","type": "string","title": "Platform"}'::jsonb, 'platform'), + ('instruments', '{"description": "Instrument names","type": "array","title": "Instruments"}'::jsonb, 'instruments'), + ('constellation', '{"description": "Constellation name","type": "string","title": "Constellation"}'::jsonb, 'constellation'), + ('mission', '{"description": "Mission name","type": "string","title": "Mission"}'::jsonb, 'mission'), + ('eo:cloud_cover', '{"description": "EO cloud cover percentage","type": "number","title": "Cloud Cover"}'::jsonb, 'eo_cloud_cover'), + ('bands', '{"description": "Bands metadata (STAC 1.1 common bands, successor of eo:bands)","type": "array","title": "Bands"}'::jsonb, 'bands'), + ('eo:snow_cover', '{"description": "EO snow cover percentage","type": "number","title": "Snow Cover"}'::jsonb, 'eo_snow_cover'), + ('gsd', '{"description": "Ground sample distance","type": "number","title": "Ground Sample Distance"}'::jsonb, 'gsd'), + ('proj:code', '{"description": "Authority and code of the CRS, e.g. EPSG:32659","type": "string","title": "Projection Code"}'::jsonb, 'proj_code'), + ('proj:geometry', '{"description": "Footprint of the Item in its native CRS","type": "object","title": "Projection Geometry"}'::jsonb, 'proj_geometry'), + ('proj:wkt2', '{"description": "WKT2 CRS definition","type": "string","title": "Projection WKT2"}'::jsonb, 'proj_wkt2'), + ('proj:projjson', '{"description": "PROJJSON CRS definition","type": ["object", "string"],"title": "Projection PROJJSON"}'::jsonb, 'proj_projjson'), + ('proj:bbox', '{"description": "Projection bbox","type": "array","title": "Projection BBOX"}'::jsonb, 'proj_bbox'), + ('proj:centroid', '{"description": "Projection centroid","type": "object","title": "Projection Centroid"}'::jsonb, 'proj_centroid'), + ('proj:shape', '{"description": "Projection shape","type": "array","title": "Projection Shape"}'::jsonb, 'proj_shape'), + ('proj:transform', '{"description": "Projection affine transform","type": "array","title": "Projection Transform"}'::jsonb, 'proj_transform'), + ('sci:doi', '{"description": "Scientific DOI","type": "string","title": "Scientific DOI"}'::jsonb, 'sci_doi'), + ('sci:citation', '{"description": "Scientific citation","type": "string","title": "Scientific Citation"}'::jsonb, 'sci_citation'), + ('sci:publications', '{"description": "Scientific publications","type": "array","title": "Scientific Publications"}'::jsonb, 'sci_publications'), + ('view:off_nadir', '{"description": "Viewing angle off nadir","type": "number","title": "View Off Nadir"}'::jsonb, 'view_off_nadir'), + ('view:incidence_angle','{"description": "View incidence angle","type": "number","title": "View Incidence Angle"}'::jsonb, 'view_incidence_angle'), + ('view:azimuth', '{"description": "View azimuth angle","type": "number","title": "View Azimuth"}'::jsonb, 'view_azimuth'), + ('view:sun_azimuth', '{"description": "Sun azimuth angle","type": "number","title": "View Sun Azimuth"}'::jsonb, 'view_sun_azimuth'), + ('view:sun_elevation', '{"description": "Sun elevation angle","type": "number","title": "View Sun Elevation"}'::jsonb, 'view_sun_elevation'), + ('view:moon_azimuth', '{"description": "Moon azimuth angle","type": "number","title": "View Moon Azimuth"}'::jsonb, 'view_moon_azimuth'), + ('view:moon_elevation', '{"description": "Moon elevation angle","type": "number","title": "View Moon Elevation"}'::jsonb, 'view_moon_elevation'), + ('file:size', '{"description": "File size in bytes","type": "integer","title": "File Size"}'::jsonb, 'file_size'), + ('file:header_size', '{"description": "File header size in bytes","type": "integer","title": "File Header Size"}'::jsonb, 'file_header_size'), + ('file:checksum', '{"description": "File checksum","type": "string","title": "File Checksum"}'::jsonb, 'file_checksum'), + ('file:byte_order', '{"description": "File byte order","type": "string","title": "File Byte Order"}'::jsonb, 'file_byte_order'), + ('sat:orbit_state', '{"description": "Satellite orbit state","type": "string","title": "Orbit State"}'::jsonb, 'sat_orbit_state'), + ('sat:relative_orbit', '{"description": "Satellite relative orbit","type": "integer","title": "Relative Orbit"}'::jsonb, 'sat_relative_orbit'), + ('sat:absolute_orbit', '{"description": "Satellite absolute orbit","type": "integer","title": "Absolute Orbit"}'::jsonb, 'sat_absolute_orbit'), + ('sat:platform_international_designator', '{"description": "Platform International Designator","type": "string","title": "Platform International Designator"}'::jsonb, 'sat_platform_international_designator'), + ('sat:anx_datetime', '{"description": "Ascending node crossing time","type": "string","format": "date-time","title": "ANX Datetime"}'::jsonb, 'sat_anx_datetime') + ) AS t(name, definition, property_path); $function$ ; CREATE OR REPLACE FUNCTION pgstac.promoted_items_column_list() RETURNS text[] LANGUAGE sql - IMMUTABLE PARALLEL SAFE SECURITY DEFINER + IMMUTABLE PARALLEL SAFE AS $function$ SELECT ARRAY[ 'created', 'updated', 'platform', 'instruments', 'constellation', 'mission', - 'eo_cloud_cover', 'eo_bands', 'eo_snow_cover', 'gsd', - 'proj_epsg', 'proj_wkt2', 'proj_projjson', 'proj_bbox', 'proj_centroid', 'proj_shape', 'proj_transform', + 'eo_cloud_cover', 'bands', 'eo_snow_cover', 'gsd', + 'proj_code', 'proj_geometry', 'proj_wkt2', 'proj_projjson', 'proj_bbox', 'proj_centroid', 'proj_shape', 'proj_transform', 'sci_doi', 'sci_citation', 'sci_publications', - 'view_off_nadir', 'view_incidence_angle', 'view_azimuth', 'view_sun_azimuth', 'view_sun_elevation', - 'file_size', 'file_header_size', 'file_checksum', 'file_byte_order', 'file_values_regex', - 'sat_orbit_state', 'sat_relative_orbit', 'sat_absolute_orbit' + 'view_off_nadir', 'view_incidence_angle', 'view_azimuth', 'view_sun_azimuth', 'view_sun_elevation', 'view_moon_azimuth', 'view_moon_elevation', + 'file_size', 'file_header_size', 'file_checksum', 'file_byte_order', + 'sat_orbit_state', 'sat_relative_orbit', 'sat_absolute_orbit', 'sat_platform_international_designator', 'sat_anx_datetime' ]::text[]; $function$ ; @@ -1113,55 +1887,47 @@ $function$ CREATE OR REPLACE FUNCTION pgstac.promoted_properties_from_item(_item items) RETURNS jsonb LANGUAGE sql - IMMUTABLE PARALLEL SAFE SECURITY DEFINER -AS $function$ - WITH promoted_values(property_path, value) AS ( - VALUES - ('created', CASE WHEN _item.created IS NULL THEN NULL ELSE to_jsonb(tstz_to_stac_text(_item.created)) END), - ('updated', CASE WHEN _item.updated IS NULL THEN NULL ELSE to_jsonb(tstz_to_stac_text(_item.updated)) END), - ('platform', to_jsonb(_item.platform)), - ('instruments', to_jsonb(_item.instruments)), - ('constellation', to_jsonb(_item.constellation)), - ('mission', to_jsonb(_item.mission)), - ('eo_cloud_cover', to_jsonb(_item.eo_cloud_cover)), - ('eo_bands', to_jsonb(_item.eo_bands)), - ('eo_snow_cover', to_jsonb(_item.eo_snow_cover)), - ('gsd', to_jsonb(_item.gsd)), - ('proj_epsg', to_jsonb(_item.proj_epsg)), - ('proj_wkt2', to_jsonb(_item.proj_wkt2)), - ('proj_projjson', to_jsonb(_item.proj_projjson)), - ('proj_bbox', to_jsonb(_item.proj_bbox)), - ('proj_centroid', to_jsonb(_item.proj_centroid)), - ('proj_shape', to_jsonb(_item.proj_shape)), - ('proj_transform', to_jsonb(_item.proj_transform)), - ('sci_doi', to_jsonb(_item.sci_doi)), - ('sci_citation', to_jsonb(_item.sci_citation)), - ('sci_publications', to_jsonb(_item.sci_publications)), - ('view_off_nadir', to_jsonb(_item.view_off_nadir)), - ('view_incidence_angle', to_jsonb(_item.view_incidence_angle)), - ('view_azimuth', to_jsonb(_item.view_azimuth)), - ('view_sun_azimuth', to_jsonb(_item.view_sun_azimuth)), - ('view_sun_elevation', to_jsonb(_item.view_sun_elevation)), - ('file_size', to_jsonb(_item.file_size)), - ('file_header_size', to_jsonb(_item.file_header_size)), - ('file_checksum', to_jsonb(_item.file_checksum)), - ('file_byte_order', to_jsonb(_item.file_byte_order)), - ('file_values_regex', to_jsonb(_item.file_values_regex)), - ('sat_orbit_state', to_jsonb(_item.sat_orbit_state)), - ('sat_relative_orbit', to_jsonb(_item.sat_relative_orbit)), - ('sat_absolute_orbit', to_jsonb(_item.sat_absolute_orbit)) - ) - SELECT temporal_properties_from_item(_item) - || COALESCE( - ( - SELECT jsonb_object_agg(defs.name, promoted_values.value ORDER BY defs.name) - FROM promoted_item_property_defs() defs - JOIN promoted_values USING (property_path) - WHERE promoted_values.value IS NOT NULL - AND promoted_values.value <> 'null'::jsonb - ), - '{}'::jsonb - ); + IMMUTABLE PARALLEL SAFE +AS $function$ + SELECT temporal_properties_from_item(_item) || jsonb_strip_nulls(jsonb_build_object( + 'created', CASE WHEN _item.created IS NULL THEN NULL ELSE tstz_to_stac_text(_item.created) END, + 'updated', CASE WHEN _item.updated IS NULL THEN NULL ELSE tstz_to_stac_text(_item.updated) END, + 'platform', _item.platform, + 'instruments', _item.instruments, + 'constellation', _item.constellation, + 'mission', _item.mission, + 'eo:cloud_cover', _item.eo_cloud_cover, + 'bands', _item.bands, + 'eo:snow_cover', _item.eo_snow_cover, + 'gsd', _item.gsd, + 'proj:code', _item.proj_code, + 'proj:geometry', _item.proj_geometry, + 'proj:wkt2', _item.proj_wkt2, + 'proj:projjson', _item.proj_projjson, + 'proj:bbox', _item.proj_bbox, + 'proj:centroid', _item.proj_centroid, + 'proj:shape', _item.proj_shape, + 'proj:transform', _item.proj_transform, + 'sci:doi', _item.sci_doi, + 'sci:citation', _item.sci_citation, + 'sci:publications', _item.sci_publications, + 'view:off_nadir', _item.view_off_nadir, + 'view:incidence_angle', _item.view_incidence_angle, + 'view:azimuth', _item.view_azimuth, + 'view:sun_azimuth', _item.view_sun_azimuth, + 'view:sun_elevation', _item.view_sun_elevation, + 'view:moon_azimuth', _item.view_moon_azimuth, + 'view:moon_elevation', _item.view_moon_elevation, + 'file:size', _item.file_size, + 'file:header_size', _item.file_header_size, + 'file:checksum', _item.file_checksum, + 'file:byte_order', _item.file_byte_order, + 'sat:orbit_state', _item.sat_orbit_state, + 'sat:relative_orbit', _item.sat_relative_orbit, + 'sat:absolute_orbit', _item.sat_absolute_orbit, + 'sat:platform_international_designator', _item.sat_platform_international_designator, + 'sat:anx_datetime', CASE WHEN _item.sat_anx_datetime IS NULL THEN NULL ELSE tstz_to_stac_text(_item.sat_anx_datetime) END + )); $function$ ; @@ -1171,15 +1937,31 @@ CREATE OR REPLACE FUNCTION pgstac.promoted_queryables_defaults() IMMUTABLE PARALLEL SAFE AS $function$ SELECT * FROM (VALUES - ('stac_version', '{"description": "STAC specification version","type": "string","title": "STAC Version"}'::jsonb, 'stac_version', 'to_text'), - ('stac_extensions', '{"description": "List of STAC extension schema URIs","type": "array","title": "STAC Extensions"}'::jsonb, 'stac_extensions', 'to_text') + ('stac_version', '{"description": "STAC specification version","type": "string","title": "STAC Version"}'::jsonb, 'stac_version', NULL), + ('stac_extensions', '{"description": "List of STAC extension schema URIs","type": "array","title": "STAC Extensions"}'::jsonb, 'stac_extensions', NULL) ) AS top_level(name, definition, property_path, property_wrapper) UNION ALL - SELECT p.name, p.definition, p.property_path, p.property_wrapper + SELECT p.name, p.definition, p.property_path, NULL::text FROM promoted_item_property_defs() p; $function$ ; +CREATE OR REPLACE FUNCTION pgstac.q_op_query(args jsonb) + RETURNS text + LANGUAGE sql + STABLE +AS $function$ + SELECT format( + $q$( + to_tsvector('english', coalesce(properties->>'description', '')) || + to_tsvector('english', coalesce(properties->>'title', '')) || + to_tsvector('english', coalesce(properties->>'keywords', '')) + ) @@ %L$q$, + q_to_tsquery(args) + ); +$function$ +; + CREATE OR REPLACE FUNCTION pgstac.queryable_index_field(q queryables) RETURNS text LANGUAGE sql @@ -1218,45 +2000,66 @@ AS $function$ $function$ ; -CREATE OR REPLACE FUNCTION pgstac.rename_search(_old_name text, _new_name text) +CREATE OR REPLACE FUNCTION pgstac.register_search(search searches) RETURNS searches LANGUAGE plpgsql SECURITY DEFINER + SET search_path TO 'pgstac', 'public' AS $function$ -DECLARE - renamed searches%ROWTYPE; +DECLARE cached_search searches%ROWTYPE; BEGIN - -- Serialize rename-pair operations to avoid deadlocks on concurrent name swaps. - PERFORM pg_advisory_xact_lock( - hashtext( - least(_old_name, _new_name) - || '|' - || greatest(_old_name, _new_name) - ) - ); - - UPDATE searches - SET - name = _new_name, - lastused = now(), - usecount = searches.usecount + 1 - WHERE name = _old_name - RETURNING * INTO renamed; - - IF renamed IS NULL THEN - RAISE EXCEPTION 'Named search % not found', _old_name; + IF pgstac.readonly() THEN RETURN search; END IF; + UPDATE searches SET lastused = now(), usecount = searches.usecount + 1 + WHERE ctid = (SELECT ctid FROM searches WHERE hash = search.hash FOR UPDATE SKIP LOCKED LIMIT 1) + RETURNING * INTO cached_search; + IF cached_search IS NULL THEN + IF pg_try_advisory_xact_lock(hashtext(search.hash)) THEN + INSERT INTO searches (hash, search, _where, orderby, lastused, usecount, metadata) + VALUES (search.hash, search.search, search._where, search.orderby, now(), 1, search.metadata) + ON CONFLICT (hash) DO UPDATE SET lastused = EXCLUDED.lastused, usecount = searches.usecount + 1 + RETURNING * INTO cached_search; + END IF; + IF cached_search IS NULL THEN + SELECT * INTO cached_search FROM searches WHERE hash = search.hash; + END IF; END IF; - - RETURN renamed; + IF cached_search IS NOT NULL THEN + cached_search._where := search._where; + cached_search.orderby := search.orderby; + RETURN cached_search; + END IF; + RETURN search; END; $function$ ; -CREATE OR REPLACE FUNCTION pgstac.search_gc_retention_interval(conf jsonb DEFAULT NULL::jsonb) - RETURNS interval +CREATE OR REPLACE FUNCTION pgstac.search_envelope(j jsonb) + RETURNS pred_envelope LANGUAGE sql + STABLE +AS $function$ + SELECT cql2_envelope(search_to_cql2(j)); +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.search_from_json(_search jsonb DEFAULT '{}'::jsonb, _metadata jsonb DEFAULT '{}'::jsonb) + RETURNS TABLE(hash text, metadata jsonb) + LANGUAGE plpgsql + SECURITY DEFINER + SET search_path TO 'pgstac', 'public' AS $function$ - SELECT pgstac.get_setting('search_gc_retention_interval', conf)::interval; +DECLARE search searches%ROWTYPE; +BEGIN + search.search := _search; + search.metadata := _metadata; + search._where := stac_search_to_where(_search); + search.hash := search_hash_from_where(search._where, search.metadata); + search.orderby := keyset_orderby(_search); + search.lastused := now(); + search.usecount := 1; + search := register_search(search); + RETURN QUERY SELECT search.hash, search.metadata; +END; $function$ ; @@ -1265,10 +2068,7 @@ CREATE OR REPLACE FUNCTION pgstac.search_hash(_search jsonb, _metadata jsonb DEF LANGUAGE sql STABLE PARALLEL SAFE AS $function$ - SELECT search_hash_from_where( - stac_search_to_where(_search), - _metadata - ); + SELECT search_hash_from_where(stac_search_to_where(_search), _metadata); $function$ ; @@ -1277,239 +2077,559 @@ CREATE OR REPLACE FUNCTION pgstac.search_hash_from_where(_where text, _metadata LANGUAGE sql IMMUTABLE PARALLEL SAFE AS $function$ - SELECT pgstac_hash( - format( - '%s|%s', - _where, - coalesce(_metadata, '{}'::jsonb)::text - ) - ); + SELECT pgstac_hash(format('%s|%s', _where, coalesce(_metadata, '{}'::jsonb)::text)); $function$ ; -CREATE OR REPLACE FUNCTION pgstac.search_rows(_where text DEFAULT 'TRUE'::text, _orderby text DEFAULT 'datetime DESC, id DESC'::text, _limit integer DEFAULT 10) - RETURNS SETOF items +CREATE OR REPLACE FUNCTION pgstac.search_page(_search jsonb DEFAULT '{}'::jsonb, _limit integer DEFAULT 100, _token text DEFAULT NULL::text, _prev boolean DEFAULT false, _fields jsonb DEFAULT '{}'::jsonb, OUT features json, OUT number_returned integer, OUT next_token text, OUT prev_token text, OUT number_matched bigint) + RETURNS record LANGUAGE plpgsql + SECURITY DEFINER SET search_path TO 'pgstac', 'public' AS $function$ DECLARE - base_query text; - query text; - sdate timestamptz; - edate timestamptz; - n int; - records_left int := _limit; - timer timestamptz := clock_timestamp(); - full_timer timestamptz := clock_timestamp(); + band_margin CONSTANT numeric := 3.0; + band_safety CONSTANT numeric := 1.5; + _where text; _hash text; total_count bigint; + keyset_w text; full_where text; orderby_str text; keys_proj text; + lead_field text; eff_lead_dir text; datetime_leading boolean; + acc json := '[]'::json; cnt bigint := 0; has_more boolean := false; + first_k text[]; last_k text[]; fwd_first_k text[]; fwd_last_k text[]; + have_row boolean := false; next_tok text; prev_tok text; + next_present boolean; prev_present boolean; + _env pred_envelope; _cql2 jsonb; + bnds record; clamp text; clamped_where text; + band_cap_months CONSTANT int := 18; + page_rows items[] := '{}'::items[]; chunk_rows items[]; + target int := _limit + 1; got int := 0; got_band int; + is_asc boolean; band record; + band_target numeric; obs_sel numeric; band_where text; + guard int := 0; cum_scanned bigint := 0; + proj_expr text; mo interval := interval '1 month'; + cursor_idx int; BEGIN -IF _where IS NULL OR trim(_where) = '' THEN - _where = ' TRUE '; -END IF; -RAISE NOTICE 'Getting chunks for % %', _where, _orderby; + -- The requested STAC `fields` live in the search request; honor them over the (defaulted) + -- parameter so include/exclude projection is actually applied for search(). + _fields := coalesce(_search->'fields', _fields, '{}'::jsonb); + _cql2 := search_to_cql2(_search); + _where := cql2_query(_cql2); + IF _where IS NULL OR btrim(_where) = '' THEN _where := ' TRUE '; END IF; + _hash := search_hash_from_where(_where, '{}'::jsonb); + _env := cql2_envelope(_cql2); + + SELECT * INTO bnds FROM partition_bounds(_env); + IF bnds.collections IS NOT NULL THEN + clamp := format('i.collection = ANY (%L::text[])', bnds.collections); + END IF; -base_query := $q$ - SELECT * FROM items - WHERE - datetime >= %L AND datetime < %L - AND (%s) - ORDER BY %s - LIMIT %L -$q$; - -IF _orderby ILIKE 'datetime d%' THEN - FOR sdate, edate IN SELECT * FROM chunker(_where) ORDER BY 1 DESC LOOP - RAISE NOTICE 'Running Query for % to %. %', sdate, edate, age_ms(full_timer); - query := format( - base_query, - sdate, - edate, - _where, - _orderby, - records_left - ); - RAISE DEBUG 'QUERY: %', query; - timer := clock_timestamp(); - RETURN QUERY EXECUTE query; - - GET DIAGNOSTICS n = ROW_COUNT; - records_left := records_left - n; - RAISE NOTICE 'Returned %/% Rows From % to %. % to go. Time: %ms', n, _limit, sdate, edate, records_left, age_ms(timer); - timer := clock_timestamp(); - IF records_left <= 0 THEN - RAISE NOTICE 'SEARCH_ROWS TOOK %ms', age_ms(full_timer); - RETURN; - END IF; - END LOOP; -ELSIF _orderby ILIKE 'datetime a%' THEN - FOR sdate, edate IN SELECT * FROM chunker(_where) ORDER BY 1 ASC LOOP - RAISE NOTICE 'Running Query for % to %. %', sdate, edate, age_ms(full_timer); - query := format( - base_query, - sdate, - edate, - _where, - _orderby, - records_left - ); - RAISE DEBUG 'QUERY: %', query; - timer := clock_timestamp(); - RETURN QUERY EXECUTE query; - - GET DIAGNOSTICS n = ROW_COUNT; - records_left := records_left - n; - RAISE NOTICE 'Returned %/% Rows From % to %. % to go. Time: %ms', n, _limit, sdate, edate, records_left, age_ms(timer); - timer := clock_timestamp(); - IF records_left <= 0 THEN - RAISE NOTICE 'SEARCH_ROWS TOOK %ms', age_ms(full_timer); - RETURN; - END IF; - END LOOP; -ELSE - query := format($q$ - SELECT * FROM items - WHERE %s - ORDER BY %s - LIMIT %L - $q$, _where, _orderby, _limit - ); - RAISE DEBUG 'QUERY: %', query; - timer := clock_timestamp(); - RETURN QUERY EXECUTE query; - RAISE NOTICE 'FULL QUERY TOOK %ms', age_ms(timer); -END IF; -RAISE NOTICE 'SEARCH_ROWS TOOK %ms', age_ms(full_timer); -RETURN; + SELECT string_agg(expr||' '||CASE WHEN _prev THEN (CASE dir WHEN 'ASC' THEN 'DESC' ELSE 'ASC' END) ELSE dir END, ', ' ORDER BY ord), + 'ARRAY['||string_agg(format('(%s)::text', expr), ',' ORDER BY ord)||']::text[]', + (array_agg(field ORDER BY ord))[1], + (array_agg(CASE WHEN _prev THEN (CASE dir WHEN 'ASC' THEN 'DESC' ELSE 'ASC' END) ELSE dir END ORDER BY ord))[1] + INTO orderby_str, keys_proj, lead_field, eff_lead_dir + FROM keyset_sortkeys(_search); + datetime_leading := (lead_field = 'datetime'); + is_asc := (eff_lead_dir = 'ASC'); + + IF context(_search->'conf') <> 'off' THEN + DECLARE s searches%ROWTYPE; BEGIN + s.search := _search; s.metadata := '{}'::jsonb; s._where := _where; s.hash := _hash; + s.orderby := keyset_orderby(_search); s.lastused := now(); s.usecount := 1; + PERFORM register_search(s); + END; + total_count := (where_stats(_hash, _where, false, _search->'conf', clamp)).context_count; + END IF; + + IF _token IS NOT NULL THEN + keyset_w := keyset_where(_search, keyset_decode(_token), _prev); + END IF; + full_where := concat_ws(' AND ', _where, keyset_w); + IF full_where IS NULL OR btrim(full_where) = '' THEN full_where := 'TRUE'; END IF; + clamped_where := concat_ws(' AND ', clamp, full_where); + IF clamped_where IS NULL OR btrim(clamped_where) = '' THEN clamped_where := 'TRUE'; END IF; + + -- Per-row projection. When the requested fields can be satisfied from item columns alone + -- (needs_fragment, evaluated once for the whole query), tell content_hydrate to skip the shared + -- item_fragments lookup entirely; otherwise it fetches and merges the fragment per row. + IF needs_fragment(_fields, bnds.collections) THEN + proj_expr := format('content_hydrate(i, %L::jsonb)', _fields); + ELSE + proj_expr := format('content_hydrate(i, %L::jsonb, true)', _fields); + END IF; + + IF datetime_leading AND array_length(bnds.months, 1) IS NOT NULL THEN + -- Start the band walk at the leading edge for the sort direction: oldest month for ascending, + -- most recent month for descending. (A descending walk that started at the oldest band would + -- collect the oldest items and never reach the recent months.) + cursor_idx := CASE WHEN is_asc THEN 1 ELSE array_length(bnds.months, 1) END; + band_target := target * band_margin; + WHILE got < target AND guard < 80 LOOP + guard := guard + 1; + SELECT * INTO band FROM next_band(bnds.counts, cursor_idx, band_target, band_cap_months, NOT is_asc); + -- a valid band must be processed even when next_band also flags done (it consumed the + -- last bucket); only stop when there is no band at all. + EXIT WHEN band.band_start_idx IS NULL; + band_where := format('i.datetime >= %L AND i.datetime < %L AND (%s)', + bnds.months[band.band_start_idx], bnds.months[band.band_end_idx] + mo, full_where); + EXECUTE format('SELECT array_agg(i ORDER BY %s) FROM (SELECT * FROM items i WHERE %s ORDER BY %s LIMIT %s) i', + orderby_str, band_where, orderby_str, target - got) INTO chunk_rows; + got_band := coalesce(array_length(chunk_rows, 1), 0); + IF chunk_rows IS NOT NULL THEN page_rows := page_rows || chunk_rows; got := got + got_band; END IF; + cum_scanned := cum_scanned + band.scanned; + cursor_idx := band.next_cursor_idx; + EXIT WHEN got >= target; + obs_sel := GREATEST(got::numeric, 0.5) / GREATEST(cum_scanned, 1); + band_target := ((target - got) / obs_sel) * band_safety; + END LOOP; + ELSE + EXECUTE format('SELECT array_agg(i ORDER BY %s) FROM (SELECT * FROM items i WHERE %s ORDER BY %s LIMIT %s) i', + orderby_str, clamped_where, orderby_str, target) INTO page_rows; + END IF; + + EXECUTE format($q$ + WITH page AS ( + SELECT %1$s AS content, + CASE WHEN row_number() OVER (ORDER BY %2$s) IN (1, %3$s) THEN %4$s END AS keys, + row_number() OVER (ORDER BY %2$s) AS rn + FROM unnest($1::items[]) i + ), + counts AS (SELECT count(*) AS n FROM page) + SELECT + coalesce(json_agg(content ORDER BY rn) FILTER (WHERE rn <= %3$s), '[]'::json), + (SELECT n FROM counts), + (SELECT keys FROM page WHERE rn = 1), + (SELECT keys FROM page WHERE rn = LEAST(%3$s, (SELECT n FROM counts)::int)) + FROM page + $q$, proj_expr, orderby_str, _limit, keys_proj) + USING page_rows INTO acc, cnt, first_k, last_k; + + IF acc IS NULL THEN acc := '[]'::json; END IF; + has_more := cnt > _limit; + have_row := cnt > 0; + + IF _prev THEN + acc := (SELECT coalesce(json_agg(e ORDER BY ord DESC), '[]'::json) FROM json_array_elements(acc) WITH ORDINALITY t(e, ord)); + fwd_first_k := last_k; fwd_last_k := first_k; + next_present := (_token IS NOT NULL); + prev_present := has_more; + ELSE + fwd_first_k := first_k; fwd_last_k := last_k; + next_present := has_more; + prev_present := (_token IS NOT NULL); + END IF; + IF have_row AND next_present THEN next_tok := keyset_encode(fwd_last_k); END IF; + IF have_row AND prev_present THEN prev_tok := keyset_encode(fwd_first_k); END IF; + + features := acc; + number_returned := json_array_length(acc); + next_token := next_tok; + prev_token := prev_tok; + number_matched := total_count; END; $function$ ; -CREATE OR REPLACE FUNCTION pgstac.strip_fragment_col(col_value jsonb, col_name text, fragment_paths text[]) - RETURNS jsonb +CREATE OR REPLACE FUNCTION pgstac.search_plan(_search jsonb DEFAULT '{}'::jsonb, _token text DEFAULT NULL::text, _limit integer DEFAULT NULL::integer, OUT query text, OUT histogram jsonb, OUT min_datetime timestamp with time zone, OUT max_datetime timestamp with time zone, OUT max_count bigint, OUT lead_desc boolean, OUT ctx_query text, OUT datetime_leading boolean, OUT context_count bigint) + RETURNS record LANGUAGE plpgsql - IMMUTABLE PARALLEL SAFE SECURITY DEFINER + SECURITY DEFINER + SET search_path TO 'pgstac', 'public' AS $function$ DECLARE - result jsonb := col_value; - p text; - pth text[]; - n int; + _cql2 jsonb := search_to_cql2(_search); + _where text := cql2_query(_cql2); + is_prev boolean := _token LIKE 'prev:%'; + keyset text := nullif(regexp_replace(coalesce(_token, ''), '^(next|prev):', ''), ''); + keyset_w text; full_where text; orderby_str text; collist text; + lead_field text; eff_dir text; _env pred_envelope; _hash text; + bnds record; coll_clamp text := ''; clamp text; BEGIN - IF col_value IS NULL OR fragment_paths IS NULL THEN RETURN col_value; END IF; - - FOREACH p IN ARRAY fragment_paths LOOP - pth := fragment_path_array(p); - n := cardinality(pth); - IF pth IS NULL OR n = 0 OR pth[1] <> col_name THEN CONTINUE; END IF; + IF _where IS NULL OR btrim(_where) = '' THEN _where := ' TRUE '; END IF; + orderby_str := keyset_orderby(_search, is_prev); + SELECT (array_agg(field ORDER BY ord))[1], + (array_agg(CASE WHEN is_prev THEN (CASE dir WHEN 'ASC' THEN 'DESC' ELSE 'ASC' END) ELSE dir END ORDER BY ord))[1] + INTO lead_field, eff_dir FROM keyset_sortkeys(_search); + datetime_leading := (lead_field = 'datetime'); + lead_desc := (eff_dir = 'DESC'); + + IF keyset IS NOT NULL THEN + keyset_w := keyset_where(_search, keyset_decode(keyset), is_prev); + END IF; + full_where := concat_ws(' AND ', _where, keyset_w); + IF full_where IS NULL OR btrim(full_where) = '' THEN full_where := 'TRUE'; END IF; + + _env := cql2_envelope(_cql2); + SELECT * INTO bnds FROM partition_bounds(_env); + min_datetime := bnds.months[1]; + max_datetime := bnds.months[array_length(bnds.months, 1)]; + max_count := bnds.total_count; + IF bnds.collections IS NOT NULL THEN + clamp := format('i.collection = ANY(%L::text[])', bnds.collections); + END IF; + IF bnds.collections IS NOT NULL THEN + coll_clamp := format('i.collection = ANY(%L::text[]) AND ', bnds.collections); + END IF; - IF n = 1 THEN - RETURN '{}'::jsonb; -- entire column goes to fragment - ELSE - -- Remove the nested sub-path from the column value at any depth. - -- #- handles depth-2 (removes a top-level key) through depth-N (removes a - -- nested key) using the path tail pth[2:n]. - result := result #- pth[2:n]; - END IF; - END LOOP; + -- Build the projection now that the resolved collections are known: when the requested fields are + -- satisfiable from item columns alone (needs_fragment is false), fragment_id is nulled in the + -- projection so a client driving this query never looks up item_fragments -- the fragment-skip + -- rides in the returned query's SELECT list, not a separate flag. + collist := fields_to_itemcols( + coalesce(_search->'fields', '{}'::jsonb), + NOT needs_fragment(coalesce(_search->'fields', '{}'::jsonb), bnds.collections)); + + IF datetime_leading AND array_length(bnds.months, 1) IS NOT NULL THEN + -- collections baked in as a literal so the query is parameterized only by $1 (band low), + -- $2 (band high), $3 (limit): the client prepares it once and binds per histogram band. + query := format( + 'SELECT %s FROM items i WHERE %si.datetime >= $1 AND i.datetime < $2 AND (%s) ORDER BY %s LIMIT $3', + collist, coll_clamp, full_where, orderby_str); + -- serialize the per-month histogram (months[]/counts[]) to jsonb for the streaming client + histogram := ( + SELECT jsonb_agg(jsonb_build_object('m', m, 'n', n) ORDER BY ord) + FROM unnest(bnds.months, bnds.counts) WITH ORDINALITY AS h(m, n, ord) + ); + ELSE + DECLARE dt_clamp text := ''; BEGIN + IF array_length(bnds.months, 1) IS NOT NULL THEN + -- months[] holds month *starts*; the last bucket spans [start, start + 1 month), so the + -- upper bound must be exclusive of the next month, not <= the last month's start (which + -- would drop items dated after the start of the final month). + dt_clamp := format('i.datetime >= %L AND i.datetime < %L AND ', bnds.months[1], bnds.months[array_length(bnds.months, 1)] + interval '1 month'); + END IF; + query := format( + 'SELECT %s FROM items i WHERE %s%s(%s) ORDER BY %s LIMIT $1', + collist, coll_clamp, dt_clamp, full_where, orderby_str); + histogram := NULL; + END; + END IF; - RETURN result; + IF context(_search->'conf') <> 'off' THEN + _hash := search_hash_from_where(_where, '{}'::jsonb); + DECLARE s searches%ROWTYPE; BEGIN + s.search := _search; s.metadata := '{}'::jsonb; s._where := _where; s.hash := _hash; + s.orderby := keyset_orderby(_search); s.lastused := now(); s.usecount := 1; + PERFORM register_search(s); + END; + -- Inline the cached count when stats are fresh (same rule as where_stats), so the client + -- can skip ctx_query on a cache hit. NULL => miss/stale => client races ctx_query. + SELECT s2.context_count INTO context_count + FROM searches s2 + WHERE s2.hash = _hash + AND s2.statslastupdated IS NOT NULL + AND s2.context_count IS NOT NULL + AND now() - s2.statslastupdated <= context_stats_ttl(_search->'conf'); + ctx_query := format('SELECT (where_stats(%L, %L, false, %L, %L)).context_count', + _hash, _where, _search->'conf', clamp); + ELSE + ctx_query := NULL; + context_count := NULL; + END IF; END; $function$ ; -CREATE OR REPLACE FUNCTION pgstac.strip_promoted_properties(props jsonb) - RETURNS jsonb - LANGUAGE sql - IMMUTABLE PARALLEL SAFE SECURITY DEFINER +CREATE OR REPLACE FUNCTION pgstac.search_query(_search jsonb DEFAULT '{}'::jsonb, _metadata jsonb DEFAULT '{}'::jsonb) + RETURNS TABLE(hash text, metadata jsonb) + LANGUAGE plpgsql + SECURITY DEFINER + SET search_path TO 'pgstac', 'public' AS $function$ - SELECT COALESCE(props, '{}'::jsonb) - COALESCE( - (SELECT array_agg(name ORDER BY name) FROM promoted_item_property_defs()), - '{}'::text[] - ) - ARRAY['datetime', 'start_datetime', 'end_datetime']; +DECLARE search searches%ROWTYPE; +BEGIN + search.search := _search; + search.metadata := _metadata; + search._where := stac_search_to_where(_search); + search.hash := search_hash_from_where(search._where, search.metadata); + search.orderby := keyset_orderby(_search); + search.lastused := now(); + search.usecount := 1; + search := register_search(search); + RETURN QUERY SELECT search.hash, search.metadata; +END; $function$ ; -CREATE OR REPLACE FUNCTION pgstac.temporal_properties_from_item(_item items) +CREATE OR REPLACE FUNCTION pgstac.search_to_cql2(j jsonb) RETURNS jsonb - LANGUAGE sql - IMMUTABLE PARALLEL SAFE SECURITY DEFINER + LANGUAGE plpgsql + STABLE AS $function$ - SELECT CASE - WHEN _item.datetime_is_range THEN jsonb_build_object('datetime', NULL) - || jsonb_strip_nulls(jsonb_build_object( - 'start_datetime', tstz_to_stac_text(_item.datetime), - 'end_datetime', tstz_to_stac_text(_item.end_datetime) - )) - ELSE jsonb_build_object('datetime', tstz_to_stac_text(_item.datetime)) - END; -$function$ -; +DECLARE + parts jsonb := '[]'::jsonb; + fil jsonb; + filterlang text; + g geometry; +BEGIN + IF j ? 'ids' THEN + parts := parts || jsonb_build_object('op', 'in', + 'args', jsonb_build_array(jsonb_build_object('property', 'id'), j->'ids')); + END IF; + IF j ? 'collections' THEN + parts := parts || jsonb_build_object('op', 'in', + 'args', jsonb_build_array(jsonb_build_object('property', 'collection'), j->'collections')); + END IF; + IF j ? 'datetime' THEN + parts := parts || jsonb_build_object('op', 'anyinteracts', + 'args', jsonb_build_array(jsonb_build_object('property', 'datetime'), j->'datetime')); + END IF; + g := stac_geom(j); + IF g IS NOT NULL THEN + parts := parts || jsonb_build_object('op', 's_intersects', + 'args', jsonb_build_array(jsonb_build_object('property', 'geometry'), ST_AsGeoJSON(g)::jsonb)); + END IF; + IF j ? 'q' THEN parts := parts || jsonb_build_object('op', 'q', 'args', j->'q'); END IF; -CREATE OR REPLACE FUNCTION pgstac.tstz_to_stac_text(value timestamp with time zone) - RETURNS text + IF j ? 'query' AND j ? 'filter' THEN + RAISE EXCEPTION 'Can only use either query or filter at one time.'; + END IF; + IF j ? 'query' THEN fil := query_to_cql2(j->'query'); + ELSIF j ? 'filter' THEN + filterlang := COALESCE(j->>'filter-lang', get_setting('default_filter_lang', j->'conf')); + IF NOT (j->'filter') @? '$.**.op' OR filterlang = 'cql-json' THEN + fil := cql1_to_cql2(j->'filter'); + ELSE fil := j->'filter'; + END IF; + END IF; + IF fil IS NOT NULL THEN parts := parts || fil; END IF; + + IF jsonb_array_length(parts) = 0 THEN RETURN NULL; + ELSIF jsonb_array_length(parts) = 1 THEN RETURN parts->0; + END IF; + RETURN jsonb_build_object('op', 'and', 'args', parts); +END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.stac_links_href_array(links jsonb) + RETURNS text[] LANGUAGE sql - IMMUTABLE PARALLEL SAFE SECURITY DEFINER + IMMUTABLE PARALLEL SAFE AS $function$ SELECT CASE - WHEN value IS NULL THEN NULL - ELSE trim(trailing '.' FROM trim(trailing '0' FROM to_char( - value AT TIME ZONE 'UTC', - 'YYYY-MM-DD"T"HH24:MI:SS.US' - ))) || 'Z' + WHEN links IS NULL OR jsonb_typeof(links) <> 'array' OR jsonb_array_length(links) = 0 THEN NULL + ELSE ARRAY( + SELECT link->>'href' + FROM jsonb_array_elements(links) WITH ORDINALITY AS elements(link, ord) + ORDER BY ord + ) END; $function$ ; -CREATE OR REPLACE FUNCTION pgstac.unname_search(_name text) - RETURNS searches +CREATE OR REPLACE FUNCTION pgstac.stac_links_hydrate(links_template jsonb, link_hrefs text[]) + RETURNS jsonb + LANGUAGE sql + IMMUTABLE PARALLEL SAFE +AS $function$ + SELECT CASE + WHEN links_template IS NULL OR jsonb_typeof(links_template) <> 'array' THEN '[]'::jsonb + WHEN link_hrefs IS NULL OR array_length(link_hrefs, 1) IS NULL THEN + COALESCE( + (SELECT jsonb_agg(link ORDER BY ord) + FROM jsonb_array_elements(links_template) WITH ORDINALITY AS elements(link, ord)), + '[]'::jsonb + ) + ELSE COALESCE( + ( + SELECT jsonb_agg( + CASE + WHEN jsonb_typeof(link) = 'object' AND hrefs.href IS NOT NULL THEN + jsonb_set(link - 'href', '{href}', to_jsonb(hrefs.href), true) + WHEN jsonb_typeof(link) = 'object' THEN link - 'href' + ELSE link + END + ORDER BY ord + ) + FROM jsonb_array_elements(links_template) WITH ORDINALITY AS elements(link, ord) + LEFT JOIN LATERAL ( + SELECT link_hrefs[ord] AS href + ) AS hrefs ON TRUE + ), + '[]'::jsonb + ) + END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.stac_links_strip_hrefs(links jsonb) + RETURNS jsonb + LANGUAGE sql + IMMUTABLE PARALLEL SAFE +AS $function$ + SELECT CASE + WHEN links IS NULL OR jsonb_typeof(links) <> 'array' OR jsonb_array_length(links) = 0 THEN NULL + ELSE + ( + SELECT jsonb_agg( + CASE + WHEN jsonb_typeof(link) = 'object' THEN link - 'href' + ELSE link + END + ORDER BY ord + ) + FROM jsonb_array_elements(links) WITH ORDINALITY AS elements(link, ord) + ) + END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.strip_fragment_col(col_value jsonb, col_name text, fragment_paths text[]) + RETURNS jsonb LANGUAGE plpgsql - SECURITY DEFINER + IMMUTABLE PARALLEL SAFE AS $function$ DECLARE - unnamed searches%ROWTYPE; + result jsonb := col_value; + p text; + pth text[]; + n int; BEGIN - UPDATE searches - SET - name = NULL, - pinned = false, - lastused = now(), - usecount = searches.usecount + 1 - WHERE name = _name - RETURNING * INTO unnamed; + IF col_value IS NULL OR fragment_paths IS NULL THEN RETURN col_value; END IF; - IF unnamed IS NULL THEN - RAISE EXCEPTION 'Named search % not found', _name; - END IF; + FOREACH p IN ARRAY fragment_paths LOOP + pth := fragment_path_array(p); + n := cardinality(pth); + IF pth IS NULL OR n = 0 OR pth[1] <> col_name THEN CONTINUE; END IF; + + IF n = 1 THEN + RETURN '{}'::jsonb; -- entire column goes to fragment + ELSE + -- Remove the nested sub-path from the column value at any depth. + -- #- handles depth-2 (removes a top-level key) through depth-N (removes a + -- nested key) using the path tail pth[2:n]. + result := result #- pth[2:n]; + END IF; + END LOOP; - RETURN unnamed; + RETURN result; END; $function$ ; -CREATE OR REPLACE FUNCTION pgstac.unpin_search(_name text) - RETURNS searches +CREATE OR REPLACE FUNCTION pgstac.strip_promoted_properties(props jsonb) + RETURNS jsonb + LANGUAGE sql + IMMUTABLE PARALLEL SAFE +AS $function$ + SELECT COALESCE(props, '{}'::jsonb) - COALESCE( + (SELECT array_agg(name ORDER BY name) FROM promoted_item_property_defs()), + '{}'::text[] + ) - ARRAY['datetime', 'start_datetime', 'end_datetime']; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.temporal_properties_from_item(_item items) + RETURNS jsonb + LANGUAGE sql + IMMUTABLE PARALLEL SAFE +AS $function$ + SELECT CASE + WHEN _item.datetime_is_range THEN jsonb_build_object('datetime', NULL) + || jsonb_strip_nulls(jsonb_build_object( + 'start_datetime', tstz_to_stac_text(_item.datetime), + 'end_datetime', tstz_to_stac_text(_item.end_datetime) + )) + ELSE jsonb_build_object('datetime', tstz_to_stac_text(_item.datetime)) + END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.tighten_dirty_partition_stats(_limit integer DEFAULT NULL::integer) + RETURNS integer LANGUAGE plpgsql SECURITY DEFINER AS $function$ DECLARE - unpinned_search searches%ROWTYPE; + _part text; + _count int := 0; BEGIN - UPDATE searches - SET - pinned = false, - lastused = now(), - usecount = searches.usecount + 1 - WHERE name = _name - RETURNING * INTO unpinned_search; + FOR _part IN + SELECT partition FROM pgstac.partition_stats + WHERE dirty + ORDER BY last_updated NULLS FIRST + LIMIT _limit + LOOP + PERFORM pgstac.tighten_partition_stats(_part); + _count := _count + 1; + END LOOP; + RETURN _count; +END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.tighten_partition_stats(_partition text) + RETURNS void + LANGUAGE plpgsql + SECURITY DEFINER +AS $function$ +DECLARE + _n bigint; + _dtmin timestamptz; _dtmax timestamptz; + _edtmin timestamptz; _edtmax timestamptz; + _spatial geometry; + _dtrange tstzrange; + _edtrange tstzrange; + _collection text; +BEGIN + -- Hold the partition's advisory lock across the scan + write: tighten narrows dtrange to the scanned + -- extent and clears dirty, so without the lock a concurrent ingest committing a row outside that + -- extent would be left uncovered (search would prune + miss it). Same lock check_partition uses, so + -- tighten serializes with ingest into this partition only. + PERFORM pg_advisory_xact_lock(hashtext('pgstac.check_partition'), hashtext(_partition)); + + EXECUTE format( + $q$ + SELECT count(*), min(datetime), max(datetime), min(end_datetime), max(end_datetime), + st_setsrid(st_extent(geometry)::geometry, 4326) + FROM %I + $q$, + _partition + ) INTO _n, _dtmin, _dtmax, _edtmin, _edtmax, _spatial; - IF unpinned_search IS NULL THEN - RAISE EXCEPTION 'Named search % not found', _name; + IF _n = 0 THEN + _dtrange := 'empty'::tstzrange; + _edtrange := 'empty'::tstzrange; + _spatial := NULL; + ELSE + _dtrange := tstzrange(_dtmin, _dtmax, '[]'); + _edtrange := tstzrange(_edtmin, _edtmax, '[]'); END IF; - RETURN unpinned_search; + SELECT pv.collection + INTO _collection + FROM partitions_view pv WHERE partition = _partition; + + INSERT INTO partition_stats + (partition, collection, dtrange, edtrange, spatial, n, dirty, last_updated) + VALUES (_partition, _collection, _dtrange, _edtrange, _spatial, _n, false, now()) + ON CONFLICT (partition) DO UPDATE SET + collection = EXCLUDED.collection, + dtrange = EXCLUDED.dtrange, + edtrange = EXCLUDED.edtrange, + spatial = EXCLUDED.spatial, + n = EXCLUDED.n, + dirty = false, + last_updated = now(); END; $function$ ; +CREATE OR REPLACE FUNCTION pgstac.tstz_to_stac_text(value timestamp with time zone) + RETURNS text + LANGUAGE sql + IMMUTABLE PARALLEL SAFE +AS $function$ + SELECT CASE + WHEN value IS NULL THEN NULL + ELSE trim(trailing '.' FROM trim(trailing '0' FROM to_char( + value AT TIME ZONE 'UTC', + 'YYYY-MM-DD"T"HH24:MI:SS.US' + ))) || 'Z' + END; +$function$ +; + CREATE OR REPLACE FUNCTION pgstac.update_collection(data jsonb, _partition_trunc text DEFAULT NULL::text, _fragment_config text[] DEFAULT NULL::text[]) RETURNS void LANGUAGE plpgsql @@ -1665,296 +2785,3421 @@ AS $function$ $function$ ; -CREATE OR REPLACE FUNCTION pgstac.where_stats(inhash text, inwhere text, updatestats boolean DEFAULT false, conf jsonb DEFAULT NULL::jsonb) +CREATE OR REPLACE FUNCTION pgstac.where_stats(inhash text, inwhere text, updatestats boolean DEFAULT false, conf jsonb DEFAULT NULL::jsonb, inclamp text DEFAULT NULL::text) RETURNS searches LANGUAGE plpgsql SECURITY DEFINER AS $function$ DECLARE - t timestamptz; - i interval; - explain_json jsonb; - sw searches%ROWTYPE; - sw_statslastupdated timestamptz; - sw_estimated_count bigint; - sw_estimated_cost float; + t timestamptz; i interval; explain_json jsonb; + sw searches%ROWTYPE; sw_statslastupdated timestamptz; + sw_estimated_count bigint; sw_estimated_cost float; _context text := lower(context(conf)); _stats_ttl interval := context_stats_ttl(conf); _estimated_cost_threshold float := context_estimated_cost(conf); _estimated_count_threshold int := context_estimated_count(conf); ro bool := pgstac.readonly(conf); BEGIN - -- If updatestats is true then set ttl to 0 - IF updatestats THEN - RAISE DEBUG 'Updatestats set to TRUE, setting TTL to 0'; - _stats_ttl := '0'::interval; - END IF; - - -- If we don't need to calculate context, just return - IF _context = 'off' THEN - RETURN sw; - END IF; - - -- Read current stats state without holding row locks during expensive - -- estimate/count operations. + IF updatestats THEN _stats_ttl := '0'::interval; END IF; + IF _context = 'off' THEN RETURN sw; END IF; SELECT * INTO sw FROM searches WHERE hash = inhash; - - IF sw IS NULL THEN - -- In read-only mode, searches may not be persisted. Continue with - -- non-persistent estimate/count calculation so context can still be - -- returned to callers. - sw.hash := inhash; - sw._where := inwhere; - sw_statslastupdated := NULL; - ELSE - sw_statslastupdated := sw.statslastupdated; - END IF; - - -- If there is a cached row, figure out if we need to update - IF - sw IS NOT NULL - AND sw.statslastupdated IS NOT NULL - AND sw.context_count IS NOT NULL - AND now() - sw.statslastupdated <= _stats_ttl - THEN - -- We have a cached row with data that is within our ttl. - RAISE DEBUG 'Stats present in table and lastupdated within ttl: %', sw; - RAISE DEBUG 'Returning cached counts. %', sw; - RETURN sw; - END IF; - - -- Calculate estimated cost and rows - -- Use explain to get estimated count/cost - RAISE DEBUG 'Calculating estimated stats'; + IF sw IS NULL THEN sw.hash := inhash; sw._where := inwhere; sw_statslastupdated := NULL; + ELSE sw_statslastupdated := sw.statslastupdated; END IF; + IF sw IS NOT NULL AND sw.statslastupdated IS NOT NULL AND sw.context_count IS NOT NULL + AND now() - sw_statslastupdated <= _stats_ttl THEN RETURN sw; END IF; t := clock_timestamp(); - EXECUTE format('EXPLAIN (format json) SELECT 1 FROM items WHERE %s', inwhere) - INTO explain_json; - RAISE DEBUG 'Time for just the explain: %', clock_timestamp() - t; + EXECUTE format('EXPLAIN (format json) SELECT 1 FROM items i WHERE %s', + concat_ws(' AND ', inclamp, inwhere)) INTO explain_json; i := clock_timestamp() - t; - sw_estimated_count := (explain_json->0->'Plan'->>'Plan Rows')::bigint; sw_estimated_cost := (explain_json->0->'Plan'->>'Total Cost')::float; - - RAISE DEBUG 'ESTIMATED_COUNT: %, THRESHOLD %', sw_estimated_count, _estimated_count_threshold; - RAISE DEBUG 'ESTIMATED_COST: %, THRESHOLD %', sw_estimated_cost, _estimated_cost_threshold; - - -- If context is set to auto and the costs are within the threshold return the estimated costs - IF - _context = 'auto' - AND sw_estimated_count >= _estimated_count_threshold - AND sw_estimated_cost >= _estimated_cost_threshold - THEN + IF _context = 'auto' AND sw_estimated_count >= _estimated_count_threshold + AND sw_estimated_cost >= _estimated_cost_threshold THEN sw.context_count := sw_estimated_count; IF NOT ro THEN - UPDATE searches SET - statslastupdated = now(), - context_count = sw.context_count - WHERE - hash = inhash - AND statslastupdated IS NOT DISTINCT FROM sw_statslastupdated + UPDATE searches SET statslastupdated = now(), context_count = sw.context_count + WHERE hash = inhash AND statslastupdated IS NOT DISTINCT FROM sw_statslastupdated RETURNING * INTO sw; - - IF sw IS NULL THEN - SELECT * INTO sw FROM searches WHERE hash = inhash; - END IF; + IF sw IS NULL THEN SELECT * INTO sw FROM searches WHERE hash = inhash; END IF; END IF; - RAISE DEBUG 'Estimates are within thresholds, returning estimates. %', sw; RETURN sw; END IF; - - -- Calculate Actual Count t := clock_timestamp(); - RAISE NOTICE 'Calculating actual count...'; - EXECUTE format( - 'SELECT count(*) FROM items WHERE %s', - inwhere - ) INTO sw.context_count; + EXECUTE format('SELECT count(*) FROM items i WHERE %s', concat_ws(' AND ', inclamp, inwhere)) + INTO sw.context_count; i := clock_timestamp() - t; - RAISE NOTICE 'Actual Count: % -- %', sw.context_count, i; - IF NOT ro THEN - UPDATE searches SET - statslastupdated = now(), - context_count = sw.context_count - WHERE - hash = inhash - AND statslastupdated IS NOT DISTINCT FROM sw_statslastupdated + UPDATE searches SET statslastupdated = now(), context_count = sw.context_count + WHERE hash = inhash AND statslastupdated IS NOT DISTINCT FROM sw_statslastupdated RETURNING * INTO sw; - - IF sw IS NULL THEN - SELECT * INTO sw FROM searches WHERE hash = inhash; - END IF; + IF sw IS NULL THEN SELECT * INTO sw FROM searches WHERE hash = inhash; END IF; END IF; - RAISE DEBUG 'Returning with actual count. %', sw; RETURN sw; END; $function$ ; -CREATE OR REPLACE FUNCTION pgstac.collection_base_item(cid text) - RETURNS jsonb - LANGUAGE sql - STABLE PARALLEL SAFE +CREATE OR REPLACE FUNCTION pgstac.widen_partition_stats(_partition text, _dtrange tstzrange, _edtrange tstzrange, _spatial geometry DEFAULT NULL::geometry, _constraint_dtrange tstzrange DEFAULT NULL::tstzrange) + RETURNS void + LANGUAGE plpgsql + SECURITY DEFINER AS $function$ - SELECT collection_base_item(content) - FROM collections - WHERE id = cid - LIMIT 1; +DECLARE + cur RECORD; + is_unbounded boolean; + tail interval; + buf interval := COALESCE(get_setting('partition_stats_widen_buffer'), '1 month')::interval; + target_dtrange tstzrange; + target_edtrange tstzrange; + new_dtrange tstzrange; + new_edtrange tstzrange; + new_spatial geometry; + dt_covered boolean; + edt_covered boolean; + spatial_covered boolean; +BEGIN + SELECT dtrange, edtrange, spatial + INTO cur + FROM partition_stats WHERE partition = _partition; + IF NOT FOUND THEN + RAISE EXCEPTION 'partition_stats row for % does not exist; call check_partition first', _partition; + END IF; + + dt_covered := COALESCE(cur.dtrange @> _dtrange, false); + edt_covered := COALESCE(cur.edtrange @> _edtrange, false); + spatial_covered := cur.spatial IS NULL OR _spatial IS NULL OR ST_Covers(cur.spatial, _spatial); + IF dt_covered AND edt_covered AND spatial_covered THEN + RETURN; -- already covered: no write, no lock + END IF; + + -- Clamp the widen to the partition's structural bound (_constraint_dtrange). A sub-partitioned + -- (month/year) partition's bound is finite: cover the whole partition so dtrange never misses and + -- never spills into a neighbour. An unbounded partition (NULL-partition_trunc or NULL bound) has + -- nothing to clamp to, so pad the batch range by the widen buffer each side. + is_unbounded := _constraint_dtrange IS NULL + OR (lower(_constraint_dtrange) = '-infinity'::timestamptz + AND upper(_constraint_dtrange) = 'infinity'::timestamptz); + tail := GREATEST(upper(_edtrange) - upper(_dtrange), '0'::interval); + IF is_unbounded THEN + target_dtrange := tstzrange(lower(_dtrange) - buf, upper(_dtrange) + buf, '[]'); + target_edtrange := tstzrange(lower(_dtrange) - buf, upper(_edtrange) + buf, '[]'); + ELSE + target_dtrange := _constraint_dtrange; + target_edtrange := tstzrange(lower(_constraint_dtrange), upper(_constraint_dtrange) + tail, '[]'); + END IF; + + new_dtrange := range_merge(COALESCE(cur.dtrange, target_dtrange), target_dtrange); + new_edtrange := range_merge(COALESCE(cur.edtrange, target_edtrange), target_edtrange); + new_spatial := CASE WHEN spatial_covered THEN cur.spatial ELSE NULL END; + + UPDATE partition_stats + SET dtrange = new_dtrange, + edtrange = new_edtrange, + spatial = new_spatial, + dirty = true, + last_updated = now() + WHERE partition = _partition; +END; $function$ ; -create or replace view "pgstac"."collections_asitems" as SELECT id, - geometry, - 'collections'::text AS collection, - datetime, - end_datetime, - content - '{links,assets,stac_version,stac_extensions}'::text AS properties, - jsonb_build_object('properties', content - '{links,assets,stac_version,stac_extensions}'::text, 'links', content -> 'links'::text, 'assets', content -> 'assets'::text, 'stac_version', content -> 'stac_version'::text, 'stac_extensions', content -> 'stac_extensions'::text) AS content, - content AS collectionjson - FROM collections; - +CREATE OR REPLACE FUNCTION pgstac.additional_properties() + RETURNS boolean + LANGUAGE sql +AS $function$ + SELECT pgstac.get_setting_bool('additional_properties'); +$function$ +; -CREATE OR REPLACE FUNCTION pgstac.content_dehydrate(content jsonb) - RETURNS items - LANGUAGE plpgsql - STABLE +CREATE OR REPLACE FUNCTION pgstac.age_ms(a timestamp with time zone, b timestamp with time zone DEFAULT clock_timestamp()) + RETURNS double precision + LANGUAGE sql + IMMUTABLE PARALLEL SAFE AS $function$ -DECLARE - out items; - props jsonb; -BEGIN - out.id := content->>'id'; - out.geometry := stac_geom(content); - out.collection := content->>'collection'; - props := content->'properties'; - out.datetime := stac_datetime(content); - out.end_datetime := stac_end_datetime(content); - out.datetime_is_range := CASE - WHEN props->'datetime' IS NOT NULL AND props->'datetime' <> 'null'::jsonb THEN FALSE - ELSE ( - (props->'start_datetime' IS NOT NULL AND props->'start_datetime' <> 'null'::jsonb) - OR (props->'end_datetime' IS NOT NULL AND props->'end_datetime' <> 'null'::jsonb) - ) - END; - out.stac_version := content->>'stac_version'; - out.stac_extensions := COALESCE(content->'stac_extensions', '[]'::jsonb); - out.pgstac_updated_at := now(); - out.content_hash := pgstac_item_hash(content); - - -- Split columns: dedicated storage for standard top-level STAC fields. - -- These enable index-only scans on promoted queryables and JSONB-free hot paths. - out.bbox := content->'bbox'; - out.links := COALESCE(content->'links', '[]'::jsonb); - out.assets := COALESCE(content->'assets', '{}'::jsonb); - out.properties := strip_promoted_properties(props); - out.extra := content - '{id,geometry,collection,type,bbox,links,assets,properties,stac_version,stac_extensions}'::text[]; - - out.created := (props->>'created')::timestamptz; - out.updated := (props->>'updated')::timestamptz; - out.platform := props->>'platform'; - out.instruments := to_text_array(props->'instruments'); - out.constellation := props->>'constellation'; - out.mission := props->>'mission'; - out.eo_cloud_cover := (props->>'eo:cloud_cover')::float8; - out.eo_bands := props->'eo:bands'; - out.eo_snow_cover := (props->>'eo:snow_cover')::float8; - out.gsd := (props->>'gsd')::float8; - out.proj_epsg := (props->>'proj:epsg')::integer; - out.proj_wkt2 := props->>'proj:wkt2'; - out.proj_projjson := props->'proj:projjson'; - out.proj_bbox := props->'proj:bbox'; - out.proj_centroid := props->'proj:centroid'; - out.proj_shape := props->'proj:shape'; - out.proj_transform := props->'proj:transform'; - out.sci_doi := props->>'sci:doi'; - out.sci_citation := props->>'sci:citation'; - out.sci_publications := props->'sci:publications'; - out.view_off_nadir := (props->>'view:off_nadir')::float8; - out.view_incidence_angle := (props->>'view:incidence_angle')::float8; - out.view_azimuth := (props->>'view:azimuth')::float8; - out.view_sun_azimuth := (props->>'view:sun_azimuth')::float8; - out.view_sun_elevation := (props->>'view:sun_elevation')::float8; - out.file_size := (props->>'file:size')::bigint; - out.file_header_size := (props->>'file:header_size')::bigint; - out.file_checksum := props->>'file:checksum'; - out.file_byte_order := props->>'file:byte_order'; - out.file_values_regex := props->>'file:values_regex'; - out.sat_orbit_state := props->>'sat:orbit_state'; - out.sat_relative_orbit := (props->>'sat:relative_orbit')::integer; - out.sat_absolute_orbit := (props->>'sat:absolute_orbit')::integer; - - -- fragment_id is NULL on initial dehydration; assigned by the staging trigger. - out.fragment_id := NULL; - RETURN out; -END; + SELECT abs(extract(epoch from age(a,b)) * 1000); $function$ ; -CREATE OR REPLACE FUNCTION pgstac.content_hydrate(_item items, fields jsonb DEFAULT '{}'::jsonb) +CREATE OR REPLACE FUNCTION pgstac.all_collections() RETURNS jsonb - LANGUAGE plpgsql - STABLE PARALLEL SAFE + LANGUAGE sql + SET search_path TO 'pgstac', 'public' AS $function$ + SELECT coalesce(jsonb_agg(content), '[]'::jsonb) FROM collections; +$function$ +; + +CREATE OR REPLACE PROCEDURE pgstac.analyze_items() + LANGUAGE plpgsql +AS $procedure$ DECLARE - geom jsonb; - output jsonb; - frag_content jsonb; - merged_assets jsonb; - merged_properties jsonb; - hydrated_stac_version text; - hydrated_stac_extensions jsonb; + q text; + timeout_ts timestamptz; BEGIN - IF include_field('geometry', fields) THEN - geom := ST_ASGeoJson(_item.geometry, 20)::jsonb; - END IF; + timeout_ts := statement_timestamp() + queue_timeout(); + WHILE clock_timestamp() < timeout_ts LOOP + SELECT format('ANALYZE (VERBOSE, SKIP_LOCKED) %I;', relname) INTO q + FROM pg_stat_user_tables + WHERE relname like '_item%' AND (n_mod_since_analyze>0 OR last_analyze IS NULL) LIMIT 1; + IF NOT FOUND THEN + EXIT; + END IF; + RAISE NOTICE '%', q; + EXECUTE q; + COMMIT; + END LOOP; +END; +$procedure$ +; - -- Fetch shared fragment content (NULL when item has no fragment). - IF _item.fragment_id IS NOT NULL THEN - SELECT content INTO frag_content FROM item_fragments WHERE id = _item.fragment_id; - END IF; +CREATE OR REPLACE FUNCTION pgstac.array_intersection(_a anyarray, _b anyarray) + RETURNS anyarray + LANGUAGE sql + IMMUTABLE +AS $function$ + SELECT ARRAY ( SELECT unnest(_a) INTERSECT SELECT UNNEST(_b) ); +$function$ +; - -- Merge: fragment provides shared asset/property values; per-item provides individual values. - merged_assets := jsonb_merge_level1(frag_content->'assets', _item.assets); - merged_properties := jsonb_merge_level1(frag_content->'properties', _item.properties); - merged_properties := promoted_properties_from_item(_item) || COALESCE(merged_properties, '{}'::jsonb); - hydrated_stac_version := COALESCE(_item.stac_version, frag_content->>'stac_version'); - hydrated_stac_extensions := CASE - WHEN _item.stac_extensions IS NOT NULL AND _item.stac_extensions <> '[]'::jsonb THEN _item.stac_extensions - ELSE COALESCE(frag_content->'stac_extensions', _item.stac_extensions) +CREATE OR REPLACE FUNCTION pgstac.array_map_ident(_a text[]) + RETURNS text[] + LANGUAGE sql + IMMUTABLE PARALLEL SAFE +AS $function$ + SELECT array_agg(quote_ident(v)) FROM unnest(_a) v; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.array_map_literal(_a text[]) + RETURNS text[] + LANGUAGE sql + IMMUTABLE PARALLEL SAFE +AS $function$ + SELECT array_agg(quote_literal(v)) FROM unnest(_a) v; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.array_reverse(anyarray) + RETURNS anyarray + LANGUAGE sql + IMMUTABLE STRICT +AS $function$ +SELECT ARRAY( + SELECT $1[i] + FROM generate_subscripts($1,1) AS s(i) + ORDER BY i DESC +); +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.array_to_path(arr text[]) + RETURNS text + LANGUAGE sql + IMMUTABLE STRICT +AS $function$ + SELECT string_agg( + quote_literal(v), + '->' + ) FROM unnest(arr) v; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.base_url(conf jsonb DEFAULT NULL::jsonb) + RETURNS text + LANGUAGE sql +AS $function$ + SELECT COALESCE(pgstac.get_setting('base_url', conf), '.'); +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.bbox_geom(_bbox jsonb) + RETURNS geometry + LANGUAGE sql + IMMUTABLE PARALLEL SAFE STRICT +AS $function$ +SELECT CASE jsonb_array_length(_bbox) + WHEN 4 THEN + ST_SetSRID(ST_MakeEnvelope( + (_bbox->>0)::float, + (_bbox->>1)::float, + (_bbox->>2)::float, + (_bbox->>3)::float + ),4326) + WHEN 6 THEN + ST_SetSRID(ST_3DMakeBox( + ST_MakePoint( + (_bbox->>0)::float, + (_bbox->>1)::float, + (_bbox->>2)::float + ), + ST_MakePoint( + (_bbox->>3)::float, + (_bbox->>4)::float, + (_bbox->>5)::float + ) + ),4326) + ELSE null END; +; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.check_pgstac_settings(_sysmem text DEFAULT NULL::text) + RETURNS void + LANGUAGE plpgsql + SET search_path TO 'pgstac', 'public' + SET client_min_messages TO 'notice' +AS $function$ +DECLARE + settingval text; + sysmem bigint := pg_size_bytes(_sysmem); + effective_cache_size bigint := pg_size_bytes(current_setting('effective_cache_size', TRUE)); + shared_buffers bigint := pg_size_bytes(current_setting('shared_buffers', TRUE)); + work_mem bigint := pg_size_bytes(current_setting('work_mem', TRUE)); + max_connections int := current_setting('max_connections', TRUE); + maintenance_work_mem bigint := pg_size_bytes(current_setting('maintenance_work_mem', TRUE)); + seq_page_cost float := current_setting('seq_page_cost', TRUE); + random_page_cost float := current_setting('random_page_cost', TRUE); + temp_buffers bigint := pg_size_bytes(current_setting('temp_buffers', TRUE)); + r record; +BEGIN + IF _sysmem IS NULL THEN + RAISE NOTICE 'Call function with the size of your system memory `SELECT check_pgstac_settings(''4GB'')` to get pg system setting recommendations.'; + ELSE + IF effective_cache_size < (sysmem * 0.5) THEN + RAISE WARNING 'effective_cache_size of % is set low for a system with %. Recomended value between % and %', pg_size_pretty(effective_cache_size), pg_size_pretty(sysmem), pg_size_pretty(sysmem * 0.5), pg_size_pretty(sysmem * 0.75); + ELSIF effective_cache_size > (sysmem * 0.75) THEN + RAISE WARNING 'effective_cache_size of % is set high for a system with %. Recomended value between % and %', pg_size_pretty(effective_cache_size), pg_size_pretty(sysmem), pg_size_pretty(sysmem * 0.5), pg_size_pretty(sysmem * 0.75); + ELSE + RAISE NOTICE 'effective_cache_size of % is set appropriately for a system with %', pg_size_pretty(effective_cache_size), pg_size_pretty(sysmem); + END IF; + + IF shared_buffers < (sysmem * 0.2) THEN + RAISE WARNING 'shared_buffers of % is set low for a system with %. Recomended value between % and %', pg_size_pretty(shared_buffers), pg_size_pretty(sysmem), pg_size_pretty(sysmem * 0.2), pg_size_pretty(sysmem * 0.3); + ELSIF shared_buffers > (sysmem * 0.3) THEN + RAISE WARNING 'shared_buffers of % is set high for a system with %. Recomended value between % and %', pg_size_pretty(shared_buffers), pg_size_pretty(sysmem), pg_size_pretty(sysmem * 0.2), pg_size_pretty(sysmem * 0.3); + ELSE + RAISE NOTICE 'shared_buffers of % is set appropriately for a system with %', pg_size_pretty(shared_buffers), pg_size_pretty(sysmem); + END IF; + shared_buffers = sysmem * 0.3; + IF maintenance_work_mem < (sysmem * 0.2) THEN + RAISE WARNING 'maintenance_work_mem of % is set low for shared_buffers of %. Recomended value between % and %', pg_size_pretty(maintenance_work_mem), pg_size_pretty(shared_buffers), pg_size_pretty(shared_buffers * 0.2), pg_size_pretty(shared_buffers * 0.3); + ELSIF maintenance_work_mem > (shared_buffers * 0.3) THEN + RAISE WARNING 'maintenance_work_mem of % is set high for shared_buffers of %. Recomended value between % and %', pg_size_pretty(maintenance_work_mem), pg_size_pretty(shared_buffers), pg_size_pretty(shared_buffers * 0.2), pg_size_pretty(shared_buffers * 0.3); + ELSE + RAISE NOTICE 'maintenance_work_mem of % is set appropriately for shared_buffers of %', pg_size_pretty(shared_buffers), pg_size_pretty(shared_buffers); + END IF; + + IF work_mem * max_connections > shared_buffers THEN + RAISE WARNING 'work_mem setting of % is set high for % max_connections please reduce work_mem to % or decrease max_connections to %', pg_size_pretty(work_mem), max_connections, pg_size_pretty(shared_buffers/max_connections), floor(shared_buffers/work_mem); + ELSIF work_mem * max_connections < (shared_buffers * 0.75) THEN + RAISE WARNING 'work_mem setting of % is set low for % max_connections you may consider raising work_mem to % or increasing max_connections to %', pg_size_pretty(work_mem), max_connections, pg_size_pretty(shared_buffers/max_connections), floor(shared_buffers/work_mem); + ELSE + RAISE NOTICE 'work_mem setting of % and max_connections of % are adequate for shared_buffers of %', pg_size_pretty(work_mem), max_connections, pg_size_pretty(shared_buffers); + END IF; + + IF random_page_cost / seq_page_cost != 1.1 THEN + RAISE WARNING 'random_page_cost (%) /seq_page_cost (%) should be set to 1.1 for SSD. Change random_page_cost to %', random_page_cost, seq_page_cost, 1.1 * seq_page_cost; + ELSE + RAISE NOTICE 'random_page_cost and seq_page_cost set appropriately for SSD'; + END IF; + + IF temp_buffers < greatest(pg_size_bytes('128MB'),(maintenance_work_mem / 2)) THEN + RAISE WARNING 'pgstac makes heavy use of temp tables, consider raising temp_buffers from % to %', pg_size_pretty(temp_buffers), greatest('128MB', pg_size_pretty((shared_buffers / 16))); + END IF; + END IF; + + RAISE NOTICE 'VALUES FOR PGSTAC VARIABLES'; + RAISE NOTICE 'These can be set either as GUC system variables or by setting in the pgstac_settings table.'; + + FOR r IN SELECT name, get_setting(name) as setting, CASE WHEN current_setting(concat('pgstac.',name), TRUE) IS NOT NULL THEN concat('pgstac.',name, ' GUC') WHEN value IS NOT NULL THEN 'pgstac_settings table' ELSE 'Not Set' END as loc FROM pgstac_settings LOOP + RAISE NOTICE '% is set to % from the %', r.name, r.setting, r.loc; + END LOOP; + + SELECT installed_version INTO settingval from pg_available_extensions WHERE name = 'pg_cron'; + IF NOT FOUND OR settingval IS NULL THEN + RAISE NOTICE 'Consider intalling pg_cron which can be used to automate tasks'; + ELSE + RAISE NOTICE 'pg_cron % is installed', settingval; + END IF; + + SELECT installed_version INTO settingval from pg_available_extensions WHERE name = 'pgstattuple'; + IF NOT FOUND OR settingval IS NULL THEN + RAISE NOTICE 'Consider installing the pgstattuple extension which can be used to help maintain tables and indexes.'; + ELSE + RAISE NOTICE 'pgstattuple % is installed', settingval; + END IF; + + SELECT installed_version INTO settingval from pg_available_extensions WHERE name = 'pg_stat_statements'; + IF NOT FOUND OR settingval IS NULL THEN + RAISE NOTICE 'Consider installing the pg_stat_statements extension which is very helpful for tracking the types of queries on the system'; + ELSE + RAISE NOTICE 'pg_stat_statements % is installed', settingval; + IF current_setting('pg_stat_statements.track_statements', TRUE) IS DISTINCT FROM 'all' THEN + RAISE WARNING 'SET pg_stat_statements.track_statements TO ''all''; --In order to track statements within functions.'; + END IF; + END IF; + +END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.collection_bbox(id text) + RETURNS jsonb + LANGUAGE sql + IMMUTABLE PARALLEL SAFE + SET search_path TO 'pgstac', 'public' +AS $function$ + SELECT (replace(replace(replace(st_extent(geometry)::text,'BOX(','[['),')',']]'),' ',','))::jsonb + FROM items WHERE collection=$1; + ; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.collection_datetime(content jsonb) + RETURNS timestamp with time zone + LANGUAGE sql + IMMUTABLE STRICT +AS $function$ + SELECT + CASE + WHEN + (content->'extent'->'temporal'->'interval'->0->>0) IS NULL + THEN '-infinity'::timestamptz + ELSE + (content->'extent'->'temporal'->'interval'->0->>0)::timestamptz + END + ; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.collection_delete_trigger_func() + RETURNS trigger + LANGUAGE plpgsql +AS $function$ +DECLARE + collection_base_partition text := concat('_items_', OLD.key); +BEGIN + EXECUTE format($q$ + DELETE FROM partition_stats WHERE partition IN ( + SELECT partition FROM partition_sys_meta + WHERE collection=%L + ); + DROP TABLE IF EXISTS %I CASCADE; + $q$, + OLD.id, + collection_base_partition + ); + RETURN OLD; +END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.collection_enddatetime(content jsonb) + RETURNS timestamp with time zone + LANGUAGE sql + IMMUTABLE STRICT +AS $function$ + SELECT + CASE + WHEN + (content->'extent'->'temporal'->'interval'->0->>1) IS NULL + THEN 'infinity'::timestamptz + ELSE + (content->'extent'->'temporal'->'interval'->0->>1)::timestamptz + END + ; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.collection_extent(_collection text, runupdate boolean DEFAULT false) + RETURNS jsonb + LANGUAGE plpgsql +AS $function$ +DECLARE + geom_extent geometry; + mind timestamptz; + maxd timestamptz; + extent jsonb; +BEGIN + IF runupdate THEN + PERFORM tighten_partition_stats(partition) + FROM partitions_view WHERE collection=_collection; + END IF; + SELECT + min(lower(dtrange)), + max(upper(edtrange)), + st_extent(spatial) + INTO + mind, + maxd, + geom_extent + FROM partitions_view + WHERE collection=_collection; + + IF geom_extent IS NOT NULL AND mind IS NOT NULL AND maxd IS NOT NULL THEN + extent := jsonb_build_object( + 'spatial', jsonb_build_object( + 'bbox', to_jsonb(array[array[st_xmin(geom_extent), st_ymin(geom_extent), st_xmax(geom_extent), st_ymax(geom_extent)]]) + ), + 'temporal', jsonb_build_object( + 'interval', to_jsonb(array[array[mind, maxd]]) + ) + ); + RETURN extent; + END IF; + RETURN NULL; +END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.collection_geom(content jsonb) + RETURNS geometry + LANGUAGE sql + IMMUTABLE STRICT +AS $function$ + WITH box AS (SELECT content->'extent'->'spatial'->'bbox'->0 as box) + SELECT + st_makeenvelope( + (box->>0)::float, + (box->>1)::float, + (box->>2)::float, + (box->>3)::float, + 4326 + ) + FROM box; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.collection_search(_search jsonb DEFAULT '{}'::jsonb) + RETURNS jsonb + LANGUAGE plpgsql + STABLE PARALLEL SAFE +AS $function$ +DECLARE + _limit int := coalesce((_search->>'limit')::float::int, 10); + _token text := _search->>'token'; + is_prev boolean := _token LIKE 'prev:%'; + keyset text := nullif(regexp_replace(coalesce(_token, ''), '^(next|prev):', ''), ''); + _q text; _ctx text; + number_matched bigint; + number_returned bigint; + acc jsonb := '[]'::jsonb; cnt int := 0; + first_k text[]; last_k text[]; fwd_first_k text[]; fwd_last_k text[]; + has_more boolean; next_present boolean; prev_present boolean; + next_tok text; prev_tok text; + links jsonb := '[]'::jsonb; + burl text := concat(rtrim(base_url(_search->'conf'), '/'), '/collections'); + -- typed loop targets (not a `record`) so plpgsql_check can analyze the dynamic EXECUTE. + _content jsonb; _keys text[]; +BEGIN + SELECT query, ctx_query INTO _q, _ctx FROM collection_search_plan(_search, _token); + EXECUTE _ctx INTO number_matched; -- numberMatched, always + + -- the plan query returns (content, keys); +1 over _limit detects a further page. + FOR _content, _keys IN EXECUTE _q USING (_limit + 1) + LOOP + cnt := cnt + 1; + IF cnt = 1 THEN first_k := _keys; END IF; + IF cnt <= _limit THEN acc := acc || _content; last_k := _keys; END IF; + END LOOP; + + has_more := cnt > _limit; + IF is_prev THEN + acc := flip_jsonb_array(acc); + fwd_first_k := last_k; fwd_last_k := first_k; + next_present := (keyset IS NOT NULL); -- a prev page always has the origin ahead + prev_present := has_more; -- ... and a further-back page iff more remain + ELSE + fwd_first_k := first_k; fwd_last_k := last_k; + next_present := has_more; + prev_present := (keyset IS NOT NULL); + END IF; + IF cnt > 0 AND next_present THEN next_tok := keyset_encode(fwd_last_k); END IF; + IF cnt > 0 AND prev_present THEN prev_tok := keyset_encode(fwd_first_k); END IF; + + number_returned := jsonb_array_length(acc); + + IF next_tok IS NOT NULL THEN + links := links || jsonb_build_object('rel', 'next', 'type', 'application/json', + 'method', 'GET', 'href', burl || '?token=next:' || next_tok); + END IF; + IF prev_tok IS NOT NULL THEN + links := links || jsonb_build_object('rel', 'prev', 'type', 'application/json', + 'method', 'GET', 'href', burl || '?token=prev:' || prev_tok); + END IF; + + RETURN jsonb_build_object( + 'collections', acc, + 'numberMatched', number_matched, + 'numberReturned', number_returned, + 'links', links + ); +END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.collection_temporal_extent(id text) + RETURNS jsonb + LANGUAGE sql + IMMUTABLE PARALLEL SAFE + SET search_path TO 'pgstac', 'public' +AS $function$ + SELECT to_jsonb(array[array[min(datetime), max(datetime)]]) + FROM items WHERE collection=$1; +; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.collections_trigger_func() + RETURNS trigger + LANGUAGE plpgsql +AS $function$ +DECLARE + q text; + partition_name text := format('_items_%s', NEW.key); + partition_exists boolean := false; + partition_empty boolean := true; + err_context text; + loadtemp boolean := FALSE; +BEGIN + RAISE NOTICE 'Collection Trigger. % %', NEW.id, NEW.key; + IF TG_OP = 'UPDATE' AND NEW.partition_trunc IS DISTINCT FROM OLD.partition_trunc THEN + PERFORM repartition(NEW.id, NEW.partition_trunc, TRUE); + END IF; + RETURN NEW; +END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.constraint_tstzrange(expr text) + RETURNS tstzrange + LANGUAGE sql + IMMUTABLE PARALLEL SAFE STRICT +AS $function$ + WITH t AS ( + SELECT regexp_matches( + expr, + E'\\(''\([0-9 :+-]*\)''\\).*\\(''\([0-9 :+-]*\)''\\)' + ) AS m + ) SELECT tstzrange(m[1]::timestamptz, m[2]::timestamptz) FROM t + ; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.content_dehydrate(content jsonb) + RETURNS items + LANGUAGE sql + STABLE +AS $function$ + SELECT + content->>'id' AS id, + stac_geom(content) AS geometry, + content->>'collection' AS collection, + stac_datetime(content) AS datetime, + stac_end_datetime(content) AS end_datetime, + CASE + WHEN (content->'properties')->'datetime' IS NOT NULL AND (content->'properties')->'datetime' <> 'null'::jsonb THEN FALSE + ELSE ( + ((content->'properties')->'start_datetime' IS NOT NULL AND (content->'properties')->'start_datetime' <> 'null'::jsonb) + OR ((content->'properties')->'end_datetime' IS NOT NULL AND (content->'properties')->'end_datetime' <> 'null'::jsonb) + ) + END AS datetime_is_range, + content->>'stac_version' AS stac_version, + COALESCE(content->'stac_extensions', '[]'::jsonb) AS stac_extensions, + now() AS pgstac_updated_at, + pgstac.jsonb_canonical_hash(content) AS item_hash, + NULL::bigint AS fragment_id, + content->'bbox' AS bbox, + CASE WHEN content->'links' IS NOT NULL AND content->'links' <> '[]'::jsonb THEN content->'links' END AS links, + CASE WHEN content->'assets' IS NOT NULL AND content->'assets' <> '{}'::jsonb THEN content->'assets' END AS assets, + strip_promoted_properties(content->'properties') AS properties, + content - '{id,geometry,collection,type,bbox,links,assets,properties,stac_version,stac_extensions}'::text[] AS extra, + + ((content->'properties')->>'created')::timestamptz AS created, + ((content->'properties')->>'updated')::timestamptz AS updated, + ((content->'properties')->>'platform') AS platform, + to_text_array((content->'properties')->'instruments') AS instruments, + ((content->'properties')->>'constellation') AS constellation, + ((content->'properties')->>'mission') AS mission, + ((content->'properties')->>'eo:cloud_cover')::float8 AS eo_cloud_cover, + ((content->'properties')->'bands') AS bands, + ((content->'properties')->>'eo:snow_cover')::float8 AS eo_snow_cover, + ((content->'properties')->>'gsd')::float8 AS gsd, + + ((content->'properties')->>'proj:code') AS proj_code, + ((content->'properties')->'proj:geometry') AS proj_geometry, + ((content->'properties')->>'proj:wkt2') AS proj_wkt2, + ((content->'properties')->'proj:projjson') AS proj_projjson, + ((content->'properties')->'proj:bbox') AS proj_bbox, + ((content->'properties')->'proj:centroid') AS proj_centroid, + ((content->'properties')->'proj:shape') AS proj_shape, + ((content->'properties')->'proj:transform') AS proj_transform, + + ((content->'properties')->>'sci:doi') AS sci_doi, + ((content->'properties')->>'sci:citation') AS sci_citation, + ((content->'properties')->'sci:publications') AS sci_publications, + + ((content->'properties')->>'view:off_nadir')::float8 AS view_off_nadir, + ((content->'properties')->>'view:incidence_angle')::float8 AS view_incidence_angle, + ((content->'properties')->>'view:azimuth')::float8 AS view_azimuth, + ((content->'properties')->>'view:sun_azimuth')::float8 AS view_sun_azimuth, + ((content->'properties')->>'view:sun_elevation')::float8 AS view_sun_elevation, + ((content->'properties')->>'view:moon_azimuth')::float8 AS view_moon_azimuth, + ((content->'properties')->>'view:moon_elevation')::float8 AS view_moon_elevation, + + ((content->'properties')->>'file:size')::bigint AS file_size, + ((content->'properties')->>'file:header_size')::bigint AS file_header_size, + ((content->'properties')->>'file:checksum') AS file_checksum, + ((content->'properties')->>'file:byte_order') AS file_byte_order, + + ((content->'properties')->>'sat:orbit_state') AS sat_orbit_state, + ((content->'properties')->>'sat:relative_orbit')::integer AS sat_relative_orbit, + ((content->'properties')->>'sat:absolute_orbit')::integer AS sat_absolute_orbit, + ((content->'properties')->>'sat:platform_international_designator') AS sat_platform_international_designator, + ((content->'properties')->>'sat:anx_datetime')::timestamptz AS sat_anx_datetime, + stac_links_href_array(content->'links') AS link_hrefs, + -- private is operator-managed metadata outside the STAC item; always NULL from ingest + NULL::jsonb AS private; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.context(conf jsonb DEFAULT NULL::jsonb) + RETURNS text + LANGUAGE sql +AS $function$ + SELECT pgstac.get_setting('context', conf); +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.context_estimated_cost(conf jsonb DEFAULT NULL::jsonb) + RETURNS double precision + LANGUAGE sql +AS $function$ + SELECT pgstac.get_setting('context_estimated_cost', conf)::float; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.context_estimated_count(conf jsonb DEFAULT NULL::jsonb) + RETURNS integer + LANGUAGE sql +AS $function$ + SELECT pgstac.get_setting('context_estimated_count', conf)::int; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.context_stats_ttl(conf jsonb DEFAULT NULL::jsonb) + RETURNS interval + LANGUAGE sql +AS $function$ + SELECT pgstac.get_setting('context_stats_ttl', conf)::interval; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.cql1_to_cql2(j jsonb) + RETURNS jsonb + LANGUAGE plpgsql + IMMUTABLE STRICT +AS $function$ +DECLARE + args jsonb; + ret jsonb; +BEGIN + RAISE NOTICE 'CQL1_TO_CQL2: %', j; + IF j ? 'filter' THEN + RETURN cql1_to_cql2(j->'filter'); + END IF; + IF j ? 'property' THEN + RETURN j; + END IF; + IF jsonb_typeof(j) = 'array' THEN + SELECT jsonb_agg(cql1_to_cql2(el)) INTO args FROM jsonb_array_elements(j) el; + RETURN args; + END IF; + IF jsonb_typeof(j) = 'number' THEN + RETURN j; + END IF; + IF jsonb_typeof(j) = 'string' THEN + RETURN j; + END IF; + + IF jsonb_typeof(j) = 'object' THEN + -- GeoJSON geometry args (Point/Polygon/.../GeometryCollection) are not cql expressions; + -- pass them through unchanged so spatial ops keep their geometry intact. + IF j ? 'type' AND (j ? 'coordinates' OR j ? 'geometries') THEN + RETURN j; + END IF; + SELECT jsonb_build_object( + 'op', key, + 'args', cql1_to_cql2(value) + ) INTO ret + FROM jsonb_each(j) + WHERE j IS NOT NULL; + RETURN ret; + END IF; + RETURN NULL; +END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.cql2_query(j jsonb, wrapper text DEFAULT NULL::text) + RETURNS text + LANGUAGE plpgsql + STABLE +AS $function$ +#variable_conflict use_variable +DECLARE + args jsonb := j->'args'; + arg jsonb; + op text := lower(j->>'op'); + cql2op RECORD; + literal text; + _wrapper text; + leftarg text; + rightarg text; + prop text; + extra_props bool := pgstac.additional_properties(); + queryable_row RECORD; +BEGIN + IF j IS NULL OR (op IS NOT NULL AND args IS NULL) THEN + RETURN NULL; + END IF; + RAISE NOTICE 'CQL2_QUERY: %', j; + + -- check if all properties are represented in the queryables + IF NOT extra_props THEN + FOR prop IN + SELECT DISTINCT p->>0 + FROM jsonb_path_query(j, 'strict $.**.property') p + WHERE p->>0 NOT IN ('id', 'datetime', 'geometry', 'end_datetime', 'collection') + LOOP + IF (queryable(prop)).nulled_wrapper IS NULL THEN + RAISE EXCEPTION 'Term % is not found in queryables.', prop; + END IF; + END LOOP; + END IF; + + IF j ? 'filter' THEN + RETURN cql2_query(j->'filter'); + END IF; + + IF j ? 'upper' THEN + RETURN cql2_query(jsonb_build_object('op', 'upper', 'args', j->'upper')); + END IF; + + IF j ? 'lower' THEN + RETURN cql2_query(jsonb_build_object('op', 'lower', 'args', j->'lower')); + END IF; + + -- Temporal Query + IF op ilike 't_%' or op = 'anyinteracts' THEN + RETURN temporal_op_query(op, args); + END IF; + + -- If property is a timestamp convert it to text to use with + -- general operators + IF j ? 'timestamp' THEN + RETURN format('%L::timestamptz', to_tstz(j->'timestamp')); + END IF; + IF j ? 'interval' THEN + RAISE EXCEPTION 'Please use temporal operators when using intervals.'; + END IF; + + -- Spatial Query + IF op ilike 's_%' or op = 'intersects' THEN + RETURN spatial_op_query(op, args); + END IF; + + -- Full-text Query (pgstac `q` operator) + IF op = 'q' THEN + RETURN q_op_query(args); + END IF; + + IF op IN ('a_equals','a_contains','a_contained_by','a_overlaps') THEN + IF args->0 ? 'property' THEN + leftarg := format('to_text_array(%s)', (queryable(args->0->>'property')).path); + END IF; + IF args->1 ? 'property' THEN + rightarg := format('to_text_array(%s)', (queryable(args->1->>'property')).path); + END IF; + RETURN FORMAT( + '%s %s %s', + COALESCE(leftarg, quote_literal(to_text_array(args->0))), + CASE op + WHEN 'a_equals' THEN '=' + WHEN 'a_contains' THEN '@>' + WHEN 'a_contained_by' THEN '<@' + WHEN 'a_overlaps' THEN '&&' + END, + COALESCE(rightarg, quote_literal(to_text_array(args->1))) + ); + END IF; + + IF op = 'in' THEN + RAISE NOTICE 'IN : % % %', args, jsonb_build_array(args->0), args->1; + args := jsonb_build_array(args->0) || (args->1); + RAISE NOTICE 'IN2 : %', args; + END IF; + + + + IF op = 'between' THEN + args = jsonb_build_array( + args->0, + args->1, + args->2 + ); + END IF; + + -- Make sure that args is an array and run cql2_query on + -- each element of the array + RAISE NOTICE 'ARGS PRE: %', args; + IF j ? 'args' THEN + IF jsonb_typeof(args) != 'array' THEN + args := jsonb_build_array(args); + END IF; + + IF jsonb_path_exists(args, '$[*] ? (@.property == "id" || @.property == "datetime" || @.property == "end_datetime" || @.property == "collection")') THEN + wrapper := NULL; + ELSE + -- if any of the arguments are a property, try to get the property_wrapper + FOR arg IN SELECT jsonb_path_query(args, '$[*] ? (@.property != null)') LOOP + RAISE NOTICE 'Arg: %', arg; + wrapper := (queryable(arg->>'property')).nulled_wrapper; + RAISE NOTICE 'Property: %, Wrapper: %', arg, wrapper; + IF wrapper IS NOT NULL THEN + EXIT; + END IF; + END LOOP; + + -- if the property was not in queryables, see if any args were numbers + IF + wrapper IS NULL + AND jsonb_path_exists(args, '$[*] ? (@.type()=="number")') + THEN + wrapper := 'to_float'; + END IF; + wrapper := coalesce(wrapper, 'to_text'); + END IF; + + SELECT jsonb_agg(cql2_query(a, wrapper)) + INTO args + FROM jsonb_array_elements(args) a; + END IF; + RAISE NOTICE 'ARGS: %', args; + + IF op IN ('and', 'or') THEN + RETURN + format( + '(%s)', + array_to_string(to_text_array(args), format(' %s ', upper(op))) + ); + END IF; + + IF op = 'in' THEN + RAISE NOTICE 'IN -- % %', args->0, to_text(args->0); + RETURN format( + '%s IN (%s)', + to_text(args->0), + array_to_string((to_text_array(args))[2:], ',') + ); + END IF; + + -- Look up template from cql2_ops + IF j ? 'op' THEN + SELECT * INTO cql2op FROM cql2_ops WHERE cql2_ops.op ilike op; + IF FOUND THEN + -- If specific index set in queryables for a property cast other arguments to that type + + RETURN format( + cql2op.template, + VARIADIC (to_text_array(args)) + ); + ELSE + RAISE EXCEPTION 'Operator % Not Supported.', op; + END IF; + END IF; + + + IF wrapper IS NOT NULL THEN + RAISE NOTICE 'Wrapping % with %', j, wrapper; + IF j ? 'property' THEN + SELECT * INTO queryable_row FROM queryable(j->>'property'); + -- For native promoted columns (expression = path, no JSONB extraction), + -- the column's type already matches; applying a cast wrapper like to_int() + -- is redundant and prevents index-only scans. Return the bare expression. + IF + wrapper = ANY (ARRAY['to_int', 'to_float', 'to_tstz', 'to_text', 'to_text_array']) + AND queryable_row.expression = queryable_row.path + THEN + RETURN queryable_row.expression; + END IF; + RETURN format('%I(%s)', wrapper, queryable_row.path); + ELSE + RETURN format('%I(%L)', wrapper, j); + END IF; + ELSIF j ? 'property' THEN + RETURN quote_ident(j->>'property'); + END IF; + + RETURN quote_literal(to_text(j)); +END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.create_item(data jsonb) + RETURNS void + LANGUAGE sql + SET search_path TO 'pgstac', 'public' +AS $function$ + INSERT INTO items_staging (content) VALUES (data); +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.create_items(data jsonb) + RETURNS void + LANGUAGE sql + SET search_path TO 'pgstac', 'public' +AS $function$ + INSERT INTO items_staging (content) + SELECT * FROM jsonb_array_elements(data); +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.delete_collection(_id text) + RETURNS void + LANGUAGE plpgsql + SET search_path TO 'pgstac', 'public' +AS $function$ +DECLARE + out collections%ROWTYPE; +BEGIN + DELETE FROM collections WHERE id = _id RETURNING * INTO STRICT out; +END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.delete_item(_id text, _collection text DEFAULT NULL::text) + RETURNS void + LANGUAGE plpgsql + SECURITY DEFINER +AS $function$ +DECLARE +out items%ROWTYPE; +BEGIN + DELETE FROM items WHERE id = _id AND (_collection IS NULL OR collection=_collection) RETURNING * INTO STRICT out; +END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.empty_arr(anyarray) + RETURNS boolean + LANGUAGE sql + IMMUTABLE PARALLEL SAFE +AS $function$ +SELECT CASE + WHEN $1 IS NULL THEN TRUE + WHEN cardinality($1)<1 THEN TRUE +ELSE FALSE +END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.explode_dotpaths(j jsonb) + RETURNS SETOF text[] + LANGUAGE sql + IMMUTABLE PARALLEL SAFE +AS $function$ + SELECT string_to_array(p, '.') as e FROM jsonb_array_elements_text(j) p; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.explode_dotpaths_recurse(j jsonb) + RETURNS SETOF text[] + LANGUAGE sql + IMMUTABLE PARALLEL SAFE +AS $function$ + WITH RECURSIVE t AS ( + SELECT e FROM explode_dotpaths(j) e + UNION ALL + SELECT e[1:cardinality(e)-1] + FROM t + WHERE cardinality(e)>1 + ) SELECT e FROM t; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.first_notnull_sfunc(anyelement, anyelement) + RETURNS anyelement + LANGUAGE sql + IMMUTABLE PARALLEL SAFE +AS $function$ + SELECT COALESCE($1,$2); +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.flip_jsonb_array(j jsonb) + RETURNS jsonb + LANGUAGE sql + IMMUTABLE PARALLEL SAFE +AS $function$ + SELECT jsonb_agg(value) FROM (SELECT value FROM jsonb_array_elements(j) WITH ORDINALITY ORDER BY ordinality DESC) as t; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.ftime() + RETURNS interval + LANGUAGE sql +AS $function$ +SELECT age(clock_timestamp(), transaction_timestamp()); +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.geojsonsearch(geojson jsonb, queryhash text, fields jsonb DEFAULT NULL::jsonb, _scanlimit integer DEFAULT 10000, _limit integer DEFAULT 100, _timelimit interval DEFAULT '00:00:05'::interval, exitwhenfull boolean DEFAULT true, skipcovered boolean DEFAULT true) + RETURNS json + LANGUAGE sql +AS $function$ + SELECT * FROM geometrysearch( + st_geomfromgeojson(geojson), + queryhash, + fields, + _scanlimit, + _limit, + _timelimit, + exitwhenfull, + skipcovered + ); +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.geom_bbox(_geom geometry) + RETURNS jsonb + LANGUAGE sql + IMMUTABLE STRICT +AS $function$ + SELECT jsonb_build_array( + st_xmin(_geom), + st_ymin(_geom), + st_xmax(_geom), + st_ymax(_geom) + ); +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.geometrysearch(geom geometry, queryhash text, fields jsonb DEFAULT NULL::jsonb, _scanlimit integer DEFAULT 10000, _limit integer DEFAULT 100, _timelimit interval DEFAULT '00:00:05'::interval, exitwhenfull boolean DEFAULT true, skipcovered boolean DEFAULT true) + RETURNS json + LANGUAGE plpgsql +AS $function$ +DECLARE + -- Same adaptive band sizing as search_page: first band ~one page (band_margin headroom), each + -- subsequent band re-sized from the band's observed spatial hit-rate (band_safety headroom). + band_margin CONSTANT numeric := 3.0; + band_safety CONSTANT numeric := 1.5; + band_cap_months CONSTANT int := 18; + search searches%ROWTYPE; + curs refcursor; + _where text; query text; proj_expr text; + iter_record items%ROWTYPE; + page_rows items[] := '{}'::items[]; -- coverage-passing rows, hydrated once at the end + features json; + exit_flag boolean := FALSE; + counter int := 1; scancounter int := 1; remaining_limit int := _scanlimit; + tilearea float; unionedgeom geometry; clippedgeom geometry; + unionedgeom_area float := 0; prev_area float := 0; + _env pred_envelope; bnds record; + lead_field text; eff_dir text; datetime_leading boolean; is_asc boolean; orderby_str text; + cursor_idx int; mo interval := interval '1 month'; band record; band_target numeric; obs_sel numeric; + band_fetched int; guard int := 0; cum_scanned bigint := 0; +BEGIN + -- If the passed in geometry is not an area, coverage tests are meaningless. + IF ST_GeometryType(geom) !~* 'polygon' THEN + skipcovered := FALSE; exitwhenfull := FALSE; + END IF; + -- skipcovered implies exitwhenfull (once covered, nothing new can show through). + IF skipcovered THEN exitwhenfull := TRUE; END IF; + + search := search_fromhash(queryhash); + IF search IS NULL THEN + RAISE EXCEPTION 'Search with Query Hash % Not Found', queryhash; + END IF; + + tilearea := st_area(geom); + _where := format('%s AND st_intersects(geometry, %L::geometry)', search._where, geom); + + -- Discovery envelope from the registered search AND the tile geometry, then the single + -- partition_bounds read (unified per-month histogram) that next_band walks. + _env := search_envelope(search.search); + _env.geom := CASE + WHEN _env.geom IS NULL THEN ST_Envelope(geom) + ELSE ST_Envelope(ST_Intersection(_env.geom, ST_Envelope(geom))) + END; + SELECT * INTO bnds FROM partition_bounds(_env); + + -- Leading sort field/direction + the registered forward orderby. Bands are datetime windows walked + -- newest-first unless the search sorts datetime ascending; within a band rows follow search.orderby. + SELECT (array_agg(field ORDER BY ord))[1], (array_agg(dir ORDER BY ord))[1] + INTO lead_field, eff_dir FROM keyset_sortkeys(search.search); + datetime_leading := (lead_field = 'datetime'); + is_asc := (datetime_leading AND eff_dir = 'ASC'); + orderby_str := search.orderby; + + -- Per-row projection: skip the shared item_fragments lookup when the requested fields are + -- satisfiable from item columns alone (needs_fragment, evaluated once for the whole query). + IF needs_fragment(coalesce(fields, '{}'::jsonb), bnds.collections) THEN + proj_expr := format('content_hydrate(i, %L::jsonb)', coalesce(fields, '{}'::jsonb)); + ELSE + proj_expr := format('content_hydrate(i, %L::jsonb, true)', coalesce(fields, '{}'::jsonb)); + END IF; + + IF array_length(bnds.months, 1) IS NOT NULL THEN + -- Walk bands in the SEARCH's datetime direction: newest month first for a descending search, + -- oldest first when sorting datetime ascending. Mirrors search_page (004_search). Walking the + -- wrong direction returns older items before newer ones for a descending limit/early-exit search, + -- and (because band grouping shifts with the partition_stats histogram) makes the result depend on + -- stats — a bug, since stats must only affect performance, never which rows come back. + cursor_idx := CASE WHEN is_asc THEN 1 ELSE array_length(bnds.months, 1) END; + band_target := (_limit + 1) * band_margin; + <> + WHILE NOT exit_flag AND guard < 80 LOOP + guard := guard + 1; + SELECT * INTO band FROM next_band(bnds.counts, cursor_idx, band_target, band_cap_months, NOT is_asc); + -- process a valid band even when next_band also flags done; stop only on no band. + EXIT bands WHEN band.band_start_idx IS NULL; + query := format( + 'SELECT * FROM items i WHERE i.collection = ANY(%L::text[]) AND i.datetime >= %L AND i.datetime < %L AND %s ORDER BY %s LIMIT %L', + bnds.collections, bnds.months[band.band_start_idx], bnds.months[band.band_end_idx] + mo, _where, orderby_str, remaining_limit); + band_fetched := 0; + OPEN curs FOR EXECUTE query; + LOOP + FETCH curs INTO iter_record; + EXIT WHEN NOT FOUND; + band_fetched := band_fetched + 1; + IF exitwhenfull OR skipcovered THEN -- skip expensive geometry ops when neither is on + clippedgeom := st_intersection(geom, iter_record.geometry); + IF unionedgeom IS NULL THEN + unionedgeom := clippedgeom; + ELSE + unionedgeom := st_union(unionedgeom, clippedgeom); + END IF; + unionedgeom_area := st_area(unionedgeom); + IF skipcovered AND prev_area = unionedgeom_area THEN + scancounter := scancounter + 1; + CONTINUE; + END IF; + prev_area := unionedgeom_area; + END IF; + -- Collect the row; hydration happens once at the end (same content_hydrate path as search_page). + page_rows := page_rows || iter_record; + IF counter >= _limit + OR scancounter > _scanlimit + OR ftime() > _timelimit + OR (exitwhenfull AND unionedgeom_area >= tilearea) + THEN + exit_flag := TRUE; + EXIT; + END IF; + counter := counter + 1; + scancounter := scancounter + 1; + END LOOP; + CLOSE curs; + cum_scanned := cum_scanned + band.scanned; + cursor_idx := band.next_cursor_idx; + EXIT bands WHEN exit_flag; + -- LEARN: size the next band from this band's spatial hit-rate (fetched / rows scanned). + obs_sel := GREATEST(band_fetched::numeric, 0.5) / GREATEST(cum_scanned, 1); + band_target := ((_limit + 1 - counter) / obs_sel) * band_safety; + remaining_limit := _scanlimit - scancounter; + END LOOP; + END IF; + + -- Hydrate the collected rows once (content_hydrate + needs_fragment skip, json output -- 1GB text + -- ceiling, no 256MB jsonb-array limit), exactly like search_page. + EXECUTE format('SELECT coalesce(json_agg(%s), ''[]''::json) FROM unnest($1::items[]) i', proj_expr) + USING page_rows INTO features; + + RETURN json_build_object( + 'type', 'FeatureCollection', + 'features', coalesce(features, '[]'::json) + ); +END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.get_collection(id text) + RETURNS jsonb + LANGUAGE sql + SET search_path TO 'pgstac', 'public' +AS $function$ + SELECT content FROM collections + WHERE id=$1 + ; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.get_item(_id text, _collection text DEFAULT NULL::text) + RETURNS jsonb + LANGUAGE sql + STABLE + SET search_path TO 'pgstac', 'public' +AS $function$ + SELECT content_hydrate(items) FROM items WHERE id=_id AND (_collection IS NULL OR collection=_collection); +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.get_partition_name(relid regclass) + RETURNS text + LANGUAGE sql + STABLE STRICT +AS $function$ + SELECT (parse_ident(relid::text))[cardinality(parse_ident(relid::text))]; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.get_queryables() + RETURNS jsonb + LANGUAGE sql +AS $function$ + SELECT get_queryables(NULL::text[]); +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.get_queryables(_collection text DEFAULT NULL::text) + RETURNS jsonb + LANGUAGE sql +AS $function$ + SELECT + CASE + WHEN _collection IS NULL THEN get_queryables(NULL::text[]) + ELSE get_queryables(ARRAY[_collection]) + END + ; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.get_queryables(_collection_ids text[] DEFAULT NULL::text[]) + RETURNS jsonb + LANGUAGE plpgsql + STABLE +AS $function$ +DECLARE +BEGIN + -- Build up queryables if the input contains valid collection ids or is empty + IF EXISTS ( + SELECT 1 FROM collections + WHERE + _collection_ids IS NULL + OR cardinality(_collection_ids) = 0 + OR id = ANY(_collection_ids) + ) + THEN + RETURN ( + WITH base AS ( + SELECT + unnest(collection_ids) as collection_id, + name, + coalesce(definition, '{"type":"string"}'::jsonb) as definition + FROM queryables + WHERE + _collection_ids IS NULL OR + _collection_ids = '{}'::text[] OR + _collection_ids && collection_ids + UNION ALL + SELECT null, name, coalesce(definition, '{"type":"string"}'::jsonb) as definition + FROM queryables WHERE collection_ids IS NULL OR collection_ids = '{}'::text[] + ), g AS ( + SELECT + name, + first_notnull(definition) as definition, + jsonb_array_unique_merge(definition->'enum') as enum, + jsonb_min(definition->'minimum') as minimum, + jsonb_min(definition->'maxiumn') as maximum + FROM base + GROUP BY 1 + ) + SELECT + jsonb_build_object( + '$schema', 'http://json-schema.org/draft-07/schema#', + '$id', '', + 'type', 'object', + 'title', 'STAC Queryables.', + 'properties', jsonb_object_agg( + name, + definition + || + jsonb_strip_nulls(jsonb_build_object( + 'enum', enum, + 'minimum', minimum, + 'maximum', maximum + )) + ), + 'additionalProperties', pgstac.additional_properties() + ) + FROM g + ); + ELSE + RETURN NULL; + END IF; +END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.get_setting(_setting text, conf jsonb DEFAULT NULL::jsonb) + RETURNS text + LANGUAGE sql +AS $function$ +SELECT COALESCE( + nullif(conf->>_setting, ''), + nullif(current_setting(concat('pgstac.',_setting), TRUE),''), + nullif((SELECT value FROM pgstac.pgstac_settings WHERE name=_setting),'') +); +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.get_setting_bool(_setting text, conf jsonb DEFAULT NULL::jsonb) + RETURNS boolean + LANGUAGE sql +AS $function$ +SELECT COALESCE( + nullif(conf->>_setting, ''), + nullif(current_setting(concat('pgstac.',_setting), TRUE),''), + nullif((SELECT value FROM pgstac.pgstac_settings WHERE name=_setting),''), + 'FALSE' +)::boolean; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.get_sort_dir(sort_item jsonb) + RETURNS text + LANGUAGE sql + IMMUTABLE PARALLEL SAFE +AS $function$ + SELECT CASE WHEN sort_item->>'direction' ILIKE 'desc%' THEN 'DESC' ELSE 'ASC' END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.get_tstz_constraint(reloid oid, colname text) + RETURNS tstzrange + LANGUAGE plpgsql + STABLE STRICT +AS $function$ +DECLARE + expr text := NULL; + m text[]; + ts_lower timestamptz := NULL; + ts_upper timestamptz := NULL; + lower_inclusive text := '['; + upper_inclusive text := ']'; + ts timestamptz; +BEGIN + SELECT INTO expr + string_agg(def, ' AND ') + FROM pg_constraint JOIN LATERAL pg_get_constraintdef(oid) AS def ON TRUE + WHERE + conrelid = reloid + AND contype = 'c' + AND def LIKE '%' || colname || '%' + ; + + IF expr IS NULL THEN + RETURN NULL; + END IF; + + RAISE DEBUG 'Constraint expression for % on %: %', colname, reloid::regclass, expr; + -- collect all constraints for the specified column + FOR m IN SELECT regexp_matches(expr, '[ (]' || colname || $expr$\s*([<>=]{1,2})\s*'([0-9 :.+\-]+)'$expr$, 'g') LOOP + ts := m[2]::timestamptz; + IF m[1] IN ('>', '>=') + THEN + IF ts_lower IS NULL OR ts > ts_lower OR (ts = ts_lower AND m[1] = '>') THEN + ts_lower := ts; + lower_inclusive := CASE WHEN m[1] = '>' THEN '(' ELSE '[' END; + END IF; + ELSIF m[1] IN ('<', '<=') + THEN + IF ts_upper IS NULL OR ts < ts_upper OR (ts = ts_upper AND m[1] = '<') THEN + ts_upper := ts; + upper_inclusive := CASE WHEN m[1] = '<' THEN ')' ELSE ']' END; + END IF; + END IF; + END LOOP; + RAISE DEBUG 'Constraint % for %: % %', colname, reloid::regclass, ts_lower, ts_upper; + RETURN tstzrange(ts_lower, ts_upper, lower_inclusive || upper_inclusive); +END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.get_version() + RETURNS text + LANGUAGE sql +AS $function$ + SELECT version FROM pgstac.migrations ORDER BY datetime DESC, version DESC LIMIT 1; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.include_field(f text, fields jsonb DEFAULT '{}'::jsonb) + RETURNS boolean + LANGUAGE plpgsql + IMMUTABLE +AS $function$ +DECLARE + includes jsonb := fields->'include'; + excludes jsonb := fields->'exclude'; +BEGIN + IF f IS NULL THEN + RETURN NULL; + END IF; + + + IF + jsonb_typeof(excludes) = 'array' + AND jsonb_array_length(excludes)>0 + AND excludes ? f + THEN + RETURN FALSE; + END IF; + + IF + ( + jsonb_typeof(includes) = 'array' + AND jsonb_array_length(includes) > 0 + AND includes ? f + ) OR + ( + includes IS NULL + OR jsonb_typeof(includes) = 'null' + OR jsonb_array_length(includes) = 0 + ) + THEN + RETURN TRUE; + END IF; + + RETURN FALSE; +END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.indexdef(q queryables) + RETURNS text + LANGUAGE plpgsql + IMMUTABLE +AS $function$ + DECLARE + out text; + BEGIN + IF q.name = 'id' THEN + out := 'CREATE UNIQUE INDEX ON %I USING btree (id)'; + ELSIF q.name = 'datetime' THEN + out := 'CREATE INDEX ON %I USING btree (datetime DESC, end_datetime)'; + ELSIF q.name = 'geometry' THEN + out := 'CREATE INDEX ON %I USING gist (geometry)'; + ELSIF q.property_path IS NOT NULL AND queryable_uses_native_path(q.property_path) THEN + -- Native promoted column: index the column directly, no type-cast wrapper needed. + out := format( + 'CREATE INDEX ON %%I USING %s (%s)', + lower(COALESCE(q.property_index_type, 'BTREE')), + q.property_path + ); + ELSE + out := format($q$CREATE INDEX ON %%I USING %s (%s(((content -> 'properties'::text) -> %L::text)))$q$, + lower(COALESCE(q.property_index_type, 'BTREE')), + lower(COALESCE(q.property_wrapper, 'to_text')), + q.name + ); + END IF; + RETURN btrim(out, ' \n\t'); + END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.item_by_id(_id text, _collection text DEFAULT NULL::text) + RETURNS items + LANGUAGE plpgsql + STABLE + SET search_path TO 'pgstac', 'public' +AS $function$ +DECLARE + i items%ROWTYPE; +BEGIN + SELECT * INTO i FROM items WHERE id=_id AND (_collection IS NULL OR collection=_collection) LIMIT 1; + RETURN i; +END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.items_staging_triggerfunc() + RETURNS trigger + LANGUAGE plpgsql + SECURITY DEFINER +AS $function$ +DECLARE + part text; + ts timestamptz := clock_timestamp(); + nrows int; + batch jsonb[]; +BEGIN + RAISE NOTICE 'Creating Partitions. %', clock_timestamp() - ts; + + -- Fail loudly on items whose collection does not exist instead of silently dropping them + -- in the collections JOINs below (matches the Rust loader, which also errors on this). + IF EXISTS ( + SELECT 1 FROM newdata n + WHERE NOT EXISTS ( + SELECT 1 FROM collections c WHERE c.id = n.content->>'collection' + ) + ) THEN + RAISE EXCEPTION 'cannot load item(s) into nonexistent collection(s): %', + (SELECT string_agg(DISTINCT coalesce(n.content->>'collection', ''), ', ') + FROM newdata n + WHERE NOT EXISTS ( + SELECT 1 FROM collections c WHERE c.id = n.content->>'collection' + )); + END IF; + + FOR part IN WITH t AS ( + SELECT + n.content->>'collection' as collection, + stac_daterange(n.content->'properties') as dtr, + partition_trunc + FROM newdata n JOIN collections ON (n.content->>'collection'=collections.id) + ), p AS ( + SELECT + collection, + COALESCE(date_trunc(partition_trunc::text, lower(dtr)),'-infinity') as d, + tstzrange(min(lower(dtr)),max(lower(dtr)),'[]') as dtrange, + tstzrange(min(upper(dtr)),max(upper(dtr)),'[]') as edtrange + FROM t + GROUP BY 1,2 + ) SELECT check_partition(collection, dtrange, edtrange) FROM p LOOP + RAISE NOTICE 'Partition %', part; + END LOOP; + + batch := ARRAY(SELECT content FROM newdata); + + RAISE NOTICE 'Doing the insert. %', clock_timestamp() - ts; + IF TG_TABLE_NAME = 'items_staging' THEN + INSERT INTO items SELECT * FROM items_staging_dehydrate(batch); + GET DIAGNOSTICS nrows = ROW_COUNT; + RAISE NOTICE 'Inserted % rows to items. %', nrows, clock_timestamp() - ts; + ELSIF TG_TABLE_NAME = 'items_staging_ignore' THEN + INSERT INTO items SELECT * FROM items_staging_dehydrate(batch) + ON CONFLICT DO NOTHING; + GET DIAGNOSTICS nrows = ROW_COUNT; + RAISE NOTICE 'Inserted % rows to items. %', nrows, clock_timestamp() - ts; + ELSIF TG_TABLE_NAME = 'items_staging_upsert' THEN + -- Delete existing rows whose stored content actually changed, then insert. + EXECUTE format( + $sql$ + DELETE FROM items i USING ( + SELECT d.* + FROM newdata s + CROSS JOIN LATERAL content_dehydrate(s.content) d + ) s + WHERE + i.id = s.id + AND i.collection = s.collection + AND ( + %s + ) + $sql$, + items_content_distinct_sql('i', 's') + ); + GET DIAGNOSTICS nrows = ROW_COUNT; + RAISE NOTICE 'Deleted % rows from items. %', nrows, clock_timestamp() - ts; + INSERT INTO items SELECT * FROM items_staging_dehydrate(batch) + ON CONFLICT DO NOTHING; + GET DIAGNOSTICS nrows = ROW_COUNT; + RAISE NOTICE 'Inserted % rows to items. %', nrows, clock_timestamp() - ts; + END IF; + + -- Bump each partition's row-count estimate so search stepping sees the new rows. n drives + -- partition_bounds' histogram and the Rust keyset search_page skips zero-count bands, so a partition + -- left at n=0 would hide its freshly-ingested items from that path (SQL search() scans regardless). + -- Mirrors flush_items_staging_binary on the Rust loader. Over-count is safe (ignore/upsert may insert + -- fewer than staged); the async tightener resets n exactly. check_partition above already widened the + -- temporal envelope; spatial is left to the tightener on this fallback path. (n only affects search + -- stepping performance, never which rows/order come back — see geometrysearch band direction.) + UPDATE partition_stats ps + SET n = COALESCE(ps.n, 0) + agg.c, dirty = true, last_updated = now() + FROM ( + SELECT (partition_name(nd.content->>'collection', + lower(stac_daterange(nd.content->'properties')))).partition_name AS partition, + count(*) AS c + FROM newdata nd + GROUP BY 1 + ) agg + WHERE ps.partition = agg.partition; + + RAISE NOTICE 'Deleting data from staging table. %', clock_timestamp() - ts; + EXECUTE format('DELETE FROM %I', TG_TABLE_NAME); + RAISE NOTICE 'Done. %', clock_timestamp() - ts; + + RETURN NULL; +END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.jsonb_array_unique(j jsonb) + RETURNS jsonb + LANGUAGE sql + IMMUTABLE +AS $function$ + SELECT nullif_jsonbnullempty(jsonb_agg(DISTINCT a)) v FROM jsonb_array_elements(j) a; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.jsonb_concat_ignorenull(a jsonb, b jsonb) + RETURNS jsonb + LANGUAGE sql + IMMUTABLE +AS $function$ + SELECT coalesce(a,'[]'::jsonb) || coalesce(b,'[]'::jsonb); +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.jsonb_exclude(j jsonb, f jsonb) + RETURNS jsonb + LANGUAGE plpgsql + IMMUTABLE +AS $function$ +DECLARE + excludes jsonb := f-> 'exclude'; + outj jsonb := j; + path text[]; +BEGIN + IF + excludes IS NULL + OR jsonb_array_length(excludes) = 0 + THEN + RETURN j; + ELSE + FOR path IN SELECT explode_dotpaths(excludes) LOOP + outj := outj #- path; + END LOOP; + END IF; + RETURN outj; +END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.jsonb_fields(j jsonb, f jsonb DEFAULT '{"fields": []}'::jsonb) + RETURNS jsonb + LANGUAGE sql + IMMUTABLE +AS $function$ + SELECT jsonb_exclude(jsonb_include(j, f), f); +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.jsonb_greatest(a jsonb, b jsonb) + RETURNS jsonb + LANGUAGE sql + IMMUTABLE +AS $function$ + SELECT nullif_jsonbnullempty(greatest(a, b)); +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.jsonb_include(j jsonb, f jsonb) + RETURNS jsonb + LANGUAGE plpgsql + IMMUTABLE +AS $function$ +DECLARE + includes jsonb := f-> 'include'; + outj jsonb := '{}'::jsonb; + path text[]; +BEGIN + IF + includes IS NULL + OR jsonb_array_length(includes) = 0 + THEN + RETURN j; + ELSE + includes := includes || ( + CASE WHEN j ? 'collection' THEN + '["id","collection"]' + ELSE + '["id"]' + END)::jsonb; + FOR path IN SELECT explode_dotpaths(includes) LOOP + outj := jsonb_set_nested(outj, path, j #> path); + END LOOP; + END IF; + RETURN outj; +END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.jsonb_least(a jsonb, b jsonb) + RETURNS jsonb + LANGUAGE sql + IMMUTABLE +AS $function$ + SELECT nullif_jsonbnullempty(least(nullif_jsonbnullempty(a), nullif_jsonbnullempty(b))); +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.jsonb_set_nested(j jsonb, path text[], val jsonb) + RETURNS jsonb + LANGUAGE plpgsql + IMMUTABLE +AS $function$ +DECLARE +BEGIN + IF cardinality(path) > 1 THEN + FOR i IN 1..(cardinality(path)-1) LOOP + IF j #> path[:i] IS NULL THEN + j := jsonb_set_lax(j, path[:i], '{}', TRUE); + END IF; + END LOOP; + END IF; + RETURN jsonb_set_lax(j, path, val, true); + +END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.maintain_index(indexname text, queryable_idx text, dropindexes boolean DEFAULT false, rebuildindexes boolean DEFAULT false, idxconcurrently boolean DEFAULT false) + RETURNS void + LANGUAGE plpgsql + SECURITY DEFINER +AS $function$ +DECLARE +BEGIN + IF indexname IS NOT NULL THEN + IF dropindexes OR queryable_idx IS NOT NULL THEN + EXECUTE format('DROP INDEX IF EXISTS %I;', indexname); + ELSIF rebuildindexes THEN + IF idxconcurrently THEN + EXECUTE format('REINDEX INDEX CONCURRENTLY %I;', indexname); + ELSE + EXECUTE format('REINDEX INDEX CONCURRENTLY %I;', indexname); + END IF; + END IF; + END IF; + IF queryable_idx IS NOT NULL THEN + IF idxconcurrently THEN + EXECUTE replace(queryable_idx, 'INDEX', 'INDEX CONCURRENTLY'); + ELSE EXECUTE queryable_idx; + END IF; + END IF; +END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.maintain_partition_queries(part text DEFAULT 'items'::text, dropindexes boolean DEFAULT false, rebuildindexes boolean DEFAULT false, idxconcurrently boolean DEFAULT false) + RETURNS SETOF text + LANGUAGE plpgsql +AS $function$ +DECLARE + rec record; + q text; +BEGIN + FOR rec IN ( + SELECT * FROM queryable_indexes(part,true) + ) LOOP + q := format( + 'SELECT maintain_index( + %L,%L,%L,%L,%L + );', + rec.indexname, + rec.queryable_idx, + dropindexes, + rebuildindexes, + idxconcurrently + ); + RAISE NOTICE 'Q: %', q; + RETURN NEXT q; + END LOOP; + RETURN; +END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.maintain_partitions(part text DEFAULT 'items'::text, dropindexes boolean DEFAULT false, rebuildindexes boolean DEFAULT false) + RETURNS void + LANGUAGE sql +AS $function$ + WITH t AS ( + SELECT run_or_queue(q) FROM maintain_partition_queries(part, dropindexes, rebuildindexes) q + ) SELECT count(*) FROM t; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.missing_queryables(_collection text, _tablesample double precision DEFAULT 5, minrows double precision DEFAULT 10) + RETURNS TABLE(collection text, name text, definition jsonb, property_wrapper text) + LANGUAGE plpgsql +AS $function$ +DECLARE + q text; + _partition text; + explain_json json; + psize float; + estrows float; +BEGIN + SELECT format('_items_%s', key) INTO _partition FROM collections WHERE id=_collection; + + EXECUTE format('EXPLAIN (format json) SELECT 1 FROM %I;', _partition) + INTO explain_json; + psize := explain_json->0->'Plan'->'Plan Rows'; + estrows := _tablesample * .01 * psize; + IF estrows < minrows THEN + _tablesample := least(100,greatest(_tablesample, (estrows / psize) / 100)); + RAISE NOTICE '%', (psize / estrows) / 100; + END IF; + RAISE NOTICE 'Using tablesample % to find missing queryables from % % that has ~% rows estrows: %', _tablesample, _collection, _partition, psize, estrows; + + q := format( + $q$ + WITH q AS ( + SELECT * FROM queryables + WHERE + collection_ids IS NULL + OR %L = ANY(collection_ids) + ), t AS ( + SELECT + properties + FROM + %I + TABLESAMPLE SYSTEM(%L) + ), p AS ( + SELECT DISTINCT ON (key) + key, + value, + s.definition + FROM t + JOIN LATERAL jsonb_each(properties) ON TRUE + LEFT JOIN q ON (q.name=key) + LEFT JOIN stac_extension_queryables s ON (s.name=key) + WHERE q.definition IS NULL + ) + SELECT + %L, + key, + COALESCE(definition, jsonb_build_object('type',jsonb_typeof(value))) as definition, + CASE + WHEN definition->>'type' = 'integer' THEN 'to_int' + WHEN COALESCE(definition->>'type', jsonb_typeof(value)) = 'number' THEN 'to_float' + WHEN COALESCE(definition->>'type', jsonb_typeof(value)) = 'array' THEN 'to_text_array' + ELSE 'to_text' + END + FROM p; + $q$, + _collection, + _partition, + _tablesample, + _collection + ); + RETURN QUERY EXECUTE q; +END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.missing_queryables(_tablesample double precision DEFAULT 5) + RETURNS TABLE(collection_ids text[], name text, definition jsonb, property_wrapper text) + LANGUAGE sql +AS $function$ + SELECT + array_agg(collection), + name, + definition, + property_wrapper + FROM + collections + JOIN LATERAL + missing_queryables(id, _tablesample) c + ON TRUE + GROUP BY + 2,3,4 + ORDER BY 2,1 + ; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.normalize_indexdef(def text) + RETURNS text + LANGUAGE plpgsql + IMMUTABLE PARALLEL SAFE STRICT +AS $function$ +DECLARE +BEGIN + def := btrim(def, ' \n\t'); + def := regexp_replace(def, '^CREATE (UNIQUE )?INDEX ([^ ]* )?ON (ONLY )?([^ ]* )?', '', 'i'); + RETURN def; +END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.notice(VARIADIC text[]) + RETURNS boolean + LANGUAGE plpgsql +AS $function$ +DECLARE +debug boolean := current_setting('pgstac.debug', true); +BEGIN + IF debug THEN + RAISE NOTICE 'NOTICE FROM FUNC: % >>>>> %', concat_ws(' | ', $1), clock_timestamp(); + RETURN TRUE; + END IF; + RETURN FALSE; +END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.nullif_jsonbnullempty(j jsonb) + RETURNS jsonb + LANGUAGE sql + IMMUTABLE PARALLEL SAFE STRICT +AS $function$ + SELECT nullif(nullif(nullif(j,'null'::jsonb),'{}'::jsonb),'[]'::jsonb); +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.parse_dtrange(_indate jsonb, relative_base timestamp with time zone DEFAULT date_trunc('hour'::text, CURRENT_TIMESTAMP)) + RETURNS tstzrange + LANGUAGE plpgsql + STABLE PARALLEL SAFE STRICT + SET "TimeZone" TO 'UTC' +AS $function$ +DECLARE + timestrs text[]; + s timestamptz; + e timestamptz; +BEGIN + timestrs := + CASE + WHEN _indate ? 'timestamp' THEN + ARRAY[_indate->>'timestamp'] + WHEN _indate ? 'interval' THEN + to_text_array(_indate->'interval') + WHEN jsonb_typeof(_indate) = 'array' THEN + to_text_array(_indate) + ELSE + regexp_split_to_array( + _indate->>0, + '/' + ) + END; + RAISE NOTICE 'TIMESTRS %', timestrs; + IF cardinality(timestrs) = 1 THEN + IF timestrs[1] ILIKE 'P%' THEN + RETURN tstzrange(relative_base - upper(timestrs[1])::interval, relative_base, '[)'); + END IF; + s := timestrs[1]::timestamptz; + RETURN tstzrange(s, s, '[]'); + END IF; + + IF cardinality(timestrs) != 2 THEN + RAISE EXCEPTION 'Timestamp cannot have more than 2 values'; + END IF; + + IF timestrs[1] = '..' OR timestrs[1] = '' THEN + s := '-infinity'::timestamptz; + e := timestrs[2]::timestamptz; + RETURN tstzrange(s,e,'[)'); + END IF; + + IF timestrs[2] = '..' OR timestrs[2] = '' THEN + s := timestrs[1]::timestamptz; + e := 'infinity'::timestamptz; + RETURN tstzrange(s,e,'[)'); + END IF; + + IF timestrs[1] ILIKE 'P%' AND timestrs[2] NOT ILIKE 'P%' THEN + e := timestrs[2]::timestamptz; + s := e - upper(timestrs[1])::interval; + RETURN tstzrange(s,e,'[)'); + END IF; + + IF timestrs[2] ILIKE 'P%' AND timestrs[1] NOT ILIKE 'P%' THEN + s := timestrs[1]::timestamptz; + e := s + upper(timestrs[2])::interval; + RETURN tstzrange(s,e,'[)'); + END IF; + + s := timestrs[1]::timestamptz; + e := timestrs[2]::timestamptz; + + RETURN tstzrange(s,e,'[)'); + + RETURN NULL; + +END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.parse_dtrange(_indate text, relative_base timestamp with time zone DEFAULT CURRENT_TIMESTAMP) + RETURNS tstzrange + LANGUAGE sql + STABLE PARALLEL SAFE STRICT +AS $function$ + SELECT parse_dtrange(to_jsonb(_indate), relative_base); +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.partition_name(collection text, dt timestamp with time zone, OUT partition_name text, OUT partition_range tstzrange) + RETURNS record + LANGUAGE plpgsql + STABLE +AS $function$ +DECLARE + c RECORD; + parent_name text; +BEGIN + SELECT * INTO c FROM pgstac.collections WHERE id=collection; + IF NOT FOUND THEN + RAISE EXCEPTION 'Collection % does not exist', collection USING ERRCODE = 'foreign_key_violation', HINT = 'Make sure collection exists before adding items'; + END IF; + parent_name := format('_items_%s', c.key); + + + IF c.partition_trunc = 'year' THEN + partition_name := format('%s_%s', parent_name, to_char(dt,'YYYY')); + ELSIF c.partition_trunc = 'month' THEN + partition_name := format('%s_%s', parent_name, to_char(dt,'YYYYMM')); + ELSE + partition_name := parent_name; + partition_range := tstzrange('-infinity'::timestamptz, 'infinity'::timestamptz, '[]'); + END IF; + IF partition_range IS NULL THEN + partition_range := tstzrange( + date_trunc(c.partition_trunc::text, dt), + date_trunc(c.partition_trunc::text, dt) + concat('1 ', c.partition_trunc)::interval + ); + END IF; + RETURN; + +END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.q_to_tsquery(jinput jsonb) + RETURNS tsquery + LANGUAGE plpgsql +AS $function$ +DECLARE + input text; + processed_text text; + temp_text text; + quote_array text[]; + placeholder text := '@QUOTE@'; +BEGIN + IF jsonb_typeof(jinput) = 'string' THEN + input := jinput->>0; + ELSIF jsonb_typeof(jinput) = 'array' THEN + input := array_to_string( + array(select jsonb_array_elements_text(jinput)), + ' OR ' + ); + ELSE + RAISE EXCEPTION 'Input must be a string or an array of strings.'; + END IF; + -- Extract all quoted phrases and store in array + quote_array := regexp_matches(input, '"[^"]*"', 'g'); + + -- Replace each quoted part with a unique placeholder if there are any quoted phrases + IF array_length(quote_array, 1) IS NOT NULL THEN + processed_text := input; + FOR i IN array_lower(quote_array, 1) .. array_upper(quote_array, 1) LOOP + processed_text := replace(processed_text, quote_array[i], placeholder || i || placeholder); + END LOOP; + ELSE + processed_text := input; + END IF; + + -- Replace non-quoted text using regular expressions + + -- , -> | + processed_text := regexp_replace(processed_text, ',(?=(?:[^"]*"[^"]*")*[^"]*$)', ' | ', 'g'); + + -- and -> & + processed_text := regexp_replace(processed_text, '\s+AND\s+', ' & ', 'gi'); + + -- or -> | + processed_text := regexp_replace(processed_text, '\s+OR\s+', ' | ', 'gi'); + + -- + -> + processed_text := regexp_replace(processed_text, '^\s*\+([a-zA-Z0-9_]+)', '\1', 'g'); -- +term at start + processed_text := regexp_replace(processed_text, '\s*\+([a-zA-Z0-9_]+)', ' & \1', 'g'); -- +term elsewhere + + -- - -> ! + processed_text := regexp_replace(processed_text, '^\s*\-([a-zA-Z0-9_]+)', '! \1', 'g'); -- -term at start + processed_text := regexp_replace(processed_text, '\s*\-([a-zA-Z0-9_]+)', ' & ! \1', 'g'); -- -term elsewhere + + -- terms separated with spaces are assumed to represent adjacent terms. loop through these + -- occurrences and replace them with the adjacency operator (<->) + LOOP + temp_text := regexp_replace(processed_text, '([a-zA-Z0-9_]+)\s+([a-zA-Z0-9_]+)(?!\s*[&|<>])', '\1 <-> \2', 'g'); + IF temp_text = processed_text THEN + EXIT; -- No more replacements were made + END IF; + processed_text := temp_text; + END LOOP; + + + -- Replace placeholders back with quoted phrases if there were any + IF array_length(quote_array, 1) IS NOT NULL THEN + FOR i IN array_lower(quote_array, 1) .. array_upper(quote_array, 1) LOOP + processed_text := replace(processed_text, placeholder || i || placeholder, '''' || substring(quote_array[i] from 2 for length(quote_array[i]) - 2) || ''''); + END LOOP; + END IF; + + RETURN to_tsquery('english', processed_text); +END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.query_to_cql2(q jsonb) + RETURNS jsonb + LANGUAGE sql + IMMUTABLE STRICT +AS $function$ +-- Translates anything passed in through the deprecated "query" into equivalent CQL2 +WITH t AS ( + SELECT key as property, value as ops + FROM jsonb_each(q) +), t2 AS ( + SELECT property, (jsonb_each(ops)).* + FROM t WHERE jsonb_typeof(ops) = 'object' + UNION ALL + SELECT property, 'eq', ops + FROM t WHERE jsonb_typeof(ops) != 'object' +) +SELECT + jsonb_strip_nulls(jsonb_build_object( + 'op', 'and', + 'args', jsonb_agg( + jsonb_build_object( + 'op', key, + 'args', jsonb_build_array( + jsonb_build_object('property',property), + value + ) + ) + ) + ) +) as qcql FROM t2 +; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.queryable(dotpath text, OUT path text, OUT expression text, OUT wrapper text, OUT nulled_wrapper text) + RETURNS record + LANGUAGE plpgsql + STABLE STRICT +AS $function$ +DECLARE + q RECORD; + path_elements text[]; +BEGIN + dotpath := replace(dotpath, 'properties.', ''); + IF dotpath = 'start_datetime' THEN + dotpath := 'datetime'; + END IF; + IF dotpath IN ('id', 'geometry', 'datetime', 'end_datetime', 'collection') THEN + path := dotpath; + expression := dotpath; + wrapper := NULL; + RETURN; + END IF; + + SELECT * INTO q FROM queryables + WHERE + name=dotpath + OR name = 'properties.' || dotpath + OR name = replace(dotpath, 'properties.', '') + ; + IF q.property_wrapper IS NULL THEN + IF q.definition->>'type' = 'number' THEN + wrapper := 'to_float'; + nulled_wrapper := wrapper; + ELSIF q.definition->>'format' = 'date-time' THEN + wrapper := 'to_tstz'; + nulled_wrapper := wrapper; + ELSE + nulled_wrapper := NULL; + wrapper := 'to_text'; + END IF; + ELSE + wrapper := q.property_wrapper; + nulled_wrapper := wrapper; + END IF; + IF q.property_path IS NOT NULL THEN + path := q.property_path; + ELSE + path_elements := string_to_array(dotpath, '.'); + IF path_elements[1] IN ('links', 'assets', 'stac_version', 'stac_extensions') THEN + -- links, assets, stac_version, stac_extensions are now split columns. + IF array_length(path_elements, 1) = 1 THEN + path := path_elements[1]; + ELSE + path := format('%I->%s', path_elements[1], array_to_path(path_elements[2:])); + END IF; + ELSIF path_elements[1] = 'properties' THEN + -- properties is a split JSONB column; generate properties->... path. + IF array_length(path_elements, 1) = 1 THEN + path := 'properties'; + ELSE + path := format('properties->%s', array_to_path(path_elements[2:])); + END IF; + ELSE + -- Non-prefixed queryable names are assumed to live in properties. + path := format($F$properties->%s$F$, array_to_path(path_elements)); + END IF; + END IF; + IF queryable_uses_native_path(path) THEN + IF q.definition->>'type' IN ('number', 'integer') OR q.property_wrapper IN ('to_int', 'to_float') THEN + wrapper := 'to_float'; + nulled_wrapper := wrapper; + ELSIF q.definition->>'format' = 'date-time' THEN + wrapper := 'to_tstz'; + nulled_wrapper := wrapper; + ELSIF q.property_wrapper IS NULL THEN + wrapper := 'to_text'; + nulled_wrapper := NULL; + END IF; + END IF; + IF wrapper IS NULL OR queryable_uses_native_path(path) THEN + expression := path; + ELSE + expression := format('%I(%s)', wrapper, path); + END IF; + RETURN; +END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.queryable_indexes(treeroot text DEFAULT 'items'::text, changes boolean DEFAULT false, OUT collection text, OUT partition text, OUT field text, OUT indexname text, OUT existing_idx text, OUT queryable_idx text) + RETURNS SETOF record + LANGUAGE sql +AS $function$ +WITH p AS ( + SELECT + relid::text as partition, + replace(replace( + CASE + WHEN parentrelid::regclass::text='items' THEN pg_get_expr(c.relpartbound, c.oid) + ELSE pg_get_expr(parent.relpartbound, parent.oid) + END, + 'FOR VALUES IN (''',''), ''')', + '' + ) AS collection + FROM pg_partition_tree(treeroot) + JOIN pg_class c ON (relid::regclass = c.oid) + JOIN pg_class parent ON (parentrelid::regclass = parent.oid AND isleaf) + ), i AS ( + SELECT + partition, + indexname, + regexp_replace(btrim(replace(replace(indexdef, indexname, ''),'pgstac.',''),' \t\n'), '[ ]+', ' ', 'g') as iidx, + COALESCE( + substring(indexdef FROM '\(([a-zA-Z0-9_]+)\)'), + substring(indexdef FROM '\(content -> ''properties''::text\) -> ''([a-zA-Z0-9\:\_-]+)''::text'), + CASE WHEN indexdef ~* '\(datetime desc, end_datetime\)' THEN 'datetime' ELSE NULL END + ) AS field + FROM + pg_indexes + JOIN p ON (tablename=partition) + ), q AS ( + SELECT + queryable_index_field(queryables) AS field, + collection, + partition, + format(indexdef(queryables), partition) as qidx + FROM queryables, unnest_collection(queryables.collection_ids) collection + JOIN p USING (collection) + WHERE property_index_type IS NOT NULL OR name IN ('datetime','geometry','id') + ) + SELECT + collection, + partition, + field, + indexname, + iidx as existing_idx, + qidx as queryable_idx + FROM i FULL JOIN q USING (field, partition) + WHERE CASE WHEN changes THEN lower(iidx) IS DISTINCT FROM lower(qidx) ELSE TRUE END; +; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.queryable_signature(n text, c text[]) + RETURNS text + LANGUAGE sql + IMMUTABLE PARALLEL SAFE +AS $function$ + SELECT concat(n, c); +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.queryables_constraint_triggerfunc() + RETURNS trigger + LANGUAGE plpgsql +AS $function$ +DECLARE + allcollections text[]; +BEGIN + RAISE NOTICE 'Making sure that name/collection is unique for queryables %', NEW; + IF NEW.collection_ids IS NOT NULL THEN + IF EXISTS ( + SELECT 1 + FROM unnest(NEW.collection_ids) c + LEFT JOIN + collections + ON (collections.id = c) + WHERE collections.id IS NULL + ) THEN + RAISE foreign_key_violation USING MESSAGE = format( + 'One or more collections in %s do not exist.', NEW.collection_ids + ); + RETURN NULL; + END IF; + END IF; + IF TG_OP = 'INSERT' THEN + IF EXISTS ( + SELECT 1 FROM queryables q + WHERE + q.name = NEW.name + AND ( + q.collection_ids && NEW.collection_ids + OR + q.collection_ids IS NULL + OR + NEW.collection_ids IS NULL + ) + ) THEN + RAISE unique_violation USING MESSAGE = format( + 'There is already a queryable for %s for a collection in %s: %s', + NEW.name, + NEW.collection_ids, + (SELECT json_agg(row_to_json(q)) FROM queryables q WHERE + q.name = NEW.name + AND ( + q.collection_ids && NEW.collection_ids + OR + q.collection_ids IS NULL + OR + NEW.collection_ids IS NULL + )) + ); + RETURN NULL; + END IF; + END IF; + IF TG_OP = 'UPDATE' THEN + IF EXISTS ( + SELECT 1 FROM queryables q + WHERE + q.id != NEW.id + AND + q.name = NEW.name + AND ( + q.collection_ids && NEW.collection_ids + OR + q.collection_ids IS NULL + OR + NEW.collection_ids IS NULL + ) + ) THEN + RAISE unique_violation + USING MESSAGE = format( + 'There is already a queryable for %s for a collection in %s', + NEW.name, + NEW.collection_ids + ); + RETURN NULL; + END IF; + END IF; + + RETURN NEW; +END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.queryables_trigger_func() + RETURNS trigger + LANGUAGE plpgsql + SECURITY DEFINER +AS $function$ +DECLARE +BEGIN + -- Queryable definitions changed, so every partition's queryable indexes may be stale. Flag them and let + -- the async sweep (build_pending_indexes) (re)build off the hot path, instead of rebuilding every + -- partition synchronously inside this trigger. + UPDATE pgstac.partition_stats SET indexes_pending = true; + RETURN NULL; +END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.queue_timeout() + RETURNS interval + LANGUAGE sql +AS $function$ + SELECT t2s(coalesce( + get_setting('queue_timeout'), + '1h' + ))::interval; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.readonly(conf jsonb DEFAULT NULL::jsonb) + RETURNS boolean + LANGUAGE sql +AS $function$ + SELECT pgstac.get_setting_bool('readonly', conf); +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.repartition(_collection text, _partition_trunc text, triggered boolean DEFAULT false) + RETURNS text + LANGUAGE plpgsql + SECURITY DEFINER +AS $function$ +DECLARE + c RECORD; + q text; + from_trunc text; +BEGIN + SELECT * INTO c FROM pgstac.collections WHERE id=_collection; + IF NOT FOUND THEN + RAISE EXCEPTION 'Collection % does not exist', _collection USING ERRCODE = 'foreign_key_violation', HINT = 'Make sure collection exists before adding items'; + END IF; + IF triggered THEN + RAISE NOTICE 'Converting % to % partitioning via Trigger', _collection, _partition_trunc; + ELSE + RAISE NOTICE 'Converting % from using % to % partitioning', _collection, c.partition_trunc, _partition_trunc; + IF c.partition_trunc IS NOT DISTINCT FROM _partition_trunc THEN + RAISE NOTICE 'Collection % already set to use partition by %', _collection, _partition_trunc; + RETURN _collection; + END IF; + END IF; + + IF EXISTS (SELECT 1 FROM partitions_view WHERE collection=_collection LIMIT 1) THEN + EXECUTE format( + $q$ + CREATE TEMP TABLE changepartitionstaging ON COMMIT DROP AS SELECT * FROM %I; + DROP TABLE IF EXISTS %I CASCADE; + WITH p AS ( + SELECT + collection, + CASE + WHEN %L IS NULL THEN '-infinity'::timestamptz + ELSE date_trunc(%L, datetime) + END as d, + tstzrange(min(datetime),max(datetime),'[]') as dtrange, + tstzrange(min(end_datetime),max(end_datetime),'[]') as edtrange + FROM changepartitionstaging + GROUP BY 1,2 + ) SELECT check_partition(collection, dtrange, edtrange) FROM p; + INSERT INTO items SELECT * FROM changepartitionstaging; + DROP TABLE changepartitionstaging; + $q$, + concat('_items_', c.key), + concat('_items_', c.key), + c.partition_trunc, + c.partition_trunc + ); + END IF; + RETURN _collection; +END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.run_or_queue(query text) + RETURNS void + LANGUAGE plpgsql +AS $function$ +DECLARE + use_queue text := COALESCE(get_setting('use_queue'), 'FALSE')::boolean; +BEGIN + IF get_setting_bool('debug') THEN + RAISE NOTICE '%', query; + END IF; + IF use_queue THEN + INSERT INTO query_queue (query) VALUES (query) ON CONFLICT DO NOTHING; + ELSE + EXECUTE query; + END IF; + RETURN; +END; +$function$ +; + +CREATE OR REPLACE PROCEDURE pgstac.run_queued_queries() + LANGUAGE plpgsql +AS $procedure$ +DECLARE + qitem query_queue%ROWTYPE; + timeout_ts timestamptz; + error text; + cnt int := 0; +BEGIN + timeout_ts := statement_timestamp() + queue_timeout(); + WHILE clock_timestamp() < timeout_ts LOOP + DELETE FROM query_queue WHERE query = (SELECT query FROM query_queue ORDER BY added DESC LIMIT 1 FOR UPDATE SKIP LOCKED) RETURNING * INTO qitem; + IF NOT FOUND THEN + EXIT; + END IF; + cnt := cnt + 1; + BEGIN + RAISE NOTICE 'RUNNING QUERY: %', qitem.query; + EXECUTE qitem.query; + EXCEPTION WHEN others THEN + error := format('%s | %s', SQLERRM, SQLSTATE); + END; + INSERT INTO query_queue_history (query, added, finished, error) + VALUES (qitem.query, qitem.added, clock_timestamp(), error); + COMMIT; + END LOOP; +END; +$procedure$ +; + +CREATE OR REPLACE FUNCTION pgstac.run_queued_queries_intransaction() + RETURNS integer + LANGUAGE plpgsql +AS $function$ +DECLARE + qitem query_queue%ROWTYPE; + timeout_ts timestamptz; + error text; + cnt int := 0; +BEGIN + timeout_ts := statement_timestamp() + queue_timeout(); + WHILE clock_timestamp() < timeout_ts LOOP + DELETE FROM query_queue WHERE query = (SELECT query FROM query_queue ORDER BY added DESC LIMIT 1 FOR UPDATE SKIP LOCKED) RETURNING * INTO qitem; + IF NOT FOUND THEN + RETURN cnt; + END IF; + cnt := cnt + 1; + BEGIN + qitem.query := regexp_replace(qitem.query, 'CONCURRENTLY', ''); + RAISE NOTICE 'RUNNING QUERY: %', qitem.query; + + EXECUTE qitem.query; + EXCEPTION WHEN others THEN + error := format('%s | %s', SQLERRM, SQLSTATE); + RAISE WARNING '%', error; + END; + INSERT INTO query_queue_history (query, added, finished, error) + VALUES (qitem.query, qitem.added, clock_timestamp(), error); + END LOOP; + RETURN cnt; +END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.schema_qualify_refs(url text, j jsonb) + RETURNS jsonb + LANGUAGE sql + IMMUTABLE PARALLEL SAFE STRICT +AS $function$ + SELECT regexp_replace(j::text, '"\$ref": "#', concat('"$ref": "', url, '#'), 'g')::jsonb; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.search(_search jsonb DEFAULT '{}'::jsonb) + RETURNS json + LANGUAGE plpgsql +AS $function$ +DECLARE + -- caller-provided limit/token come from the search body; default to the configured page size. + _limit int := coalesce((_search->>'limit')::int, + nullif(get_setting('default_page_size', _search->'conf'), '')::int, 10); + _token text := _search->>'token'; + -- The keyset is the token minus its next/prev prefix; an empty token means no keyset (first + -- page). A non-empty keyset that does not decode raises in keyset_decode downstream. + keyset text := nullif(regexp_replace(coalesce(_token,''), '^(next|prev):', ''), ''); + is_prev boolean := (_token LIKE 'prev:%') AND keyset IS NOT NULL; + pg record; + burl text := rtrim(coalesce(base_url(_search->'conf'), ''), '/'); + links jsonb := '[]'::jsonb; + out json; +BEGIN + SELECT * INTO pg FROM search_page(_search, _limit, keyset, is_prev); + links := links + || jsonb_build_object('rel','root','type','application/json','href', burl) + || jsonb_build_object('rel','self','type','application/json','href',burl||'/search'); + IF pg.next_token IS NOT NULL THEN + links := links || jsonb_build_object('rel','next','type','application/geo+json','method','GET', + 'href', burl||'/search?token=next:'||pg.next_token); + END IF; + IF pg.prev_token IS NOT NULL THEN + links := links || jsonb_build_object('rel','prev','type','application/geo+json','method','GET', + 'href', burl||'/search?token=prev:'||pg.prev_token); + END IF; + IF pg.number_matched IS NOT NULL THEN + out := json_build_object( + 'type','FeatureCollection', + 'features', coalesce(pg.features,'[]'::json), + 'links', links, + 'numberReturned', pg.number_returned, + 'numberMatched', pg.number_matched); + ELSE + out := json_build_object( + 'type','FeatureCollection', + 'features', coalesce(pg.features,'[]'::json), + 'links', links, + 'numberReturned', pg.number_returned); + END IF; + RETURN out; +END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.search_fromhash(_hash text) + RETURNS searches + LANGUAGE sql + STRICT +AS $function$ + SELECT * FROM searches WHERE hash = _hash LIMIT 1; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.set_version(text) + RETURNS text + LANGUAGE sql +AS $function$ + INSERT INTO pgstac.migrations (version) VALUES ($1) + ON CONFLICT DO NOTHING + RETURNING version; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.spatial_op_query(op text, args jsonb) + RETURNS text + LANGUAGE plpgsql +AS $function$ +DECLARE + geom text; + j jsonb := args->1; +BEGIN + op := lower(op); + RAISE NOTICE 'Constructing spatial query OP: %, ARGS: %', op, args; + IF op NOT IN ('s_equals','s_disjoint','s_touches','s_within','s_overlaps','s_crosses','s_intersects','intersects','s_contains') THEN + RAISE EXCEPTION 'Spatial Operator % Not Supported', op; + END IF; + op := regexp_replace(op, '^s_', 'st_'); + IF op = 'intersects' THEN + op := 'st_intersects'; + END IF; + -- Convert geometry to WKB string + IF j ? 'type' AND j ? 'coordinates' THEN + geom := st_geomfromgeojson(j)::text; + ELSIF jsonb_typeof(j) = 'array' THEN + geom := bbox_geom(j)::text; + END IF; + + RETURN format('%s(geometry, %L::geometry)', op, geom); +END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.stac_daterange(value jsonb) + RETURNS tstzrange + LANGUAGE plpgsql + IMMUTABLE PARALLEL SAFE + SET "TimeZone" TO 'UTC' +AS $function$ +DECLARE + props jsonb := value; + dt timestamptz; + edt timestamptz; +BEGIN + IF props ? 'properties' THEN + props := props->'properties'; + END IF; + IF + props ? 'start_datetime' + AND props->>'start_datetime' IS NOT NULL + AND props ? 'end_datetime' + AND props->>'end_datetime' IS NOT NULL + THEN + dt := props->>'start_datetime'; + edt := props->>'end_datetime'; + IF dt > edt THEN + RAISE EXCEPTION 'start_datetime must be < end_datetime'; + END IF; + ELSE + dt := props->>'datetime'; + edt := props->>'datetime'; + END IF; + IF dt is NULL OR edt IS NULL THEN + RAISE NOTICE 'DT: %, EDT: %', dt, edt; + RAISE EXCEPTION 'Either datetime (%) or both start_datetime (%) and end_datetime (%) must be set.', props->>'datetime',props->>'start_datetime',props->>'end_datetime'; + END IF; + RETURN tstzrange(dt, edt, '[]'); +END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.stac_datetime(value jsonb) + RETURNS timestamp with time zone + LANGUAGE sql + IMMUTABLE PARALLEL SAFE + SET "TimeZone" TO 'UTC' +AS $function$ + SELECT lower(stac_daterange(value)); +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.stac_end_datetime(value jsonb) + RETURNS timestamp with time zone + LANGUAGE sql + IMMUTABLE PARALLEL SAFE + SET "TimeZone" TO 'UTC' +AS $function$ + SELECT upper(stac_daterange(value)); +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.stac_geom(value jsonb) + RETURNS geometry + LANGUAGE sql + IMMUTABLE PARALLEL SAFE +AS $function$ +SELECT + CASE + WHEN value ? 'intersects' THEN + ST_GeomFromGeoJSON(value->>'intersects') + WHEN value ? 'geometry' THEN + ST_GeomFromGeoJSON(value->>'geometry') + WHEN value ? 'bbox' THEN + pgstac.bbox_geom(value->'bbox') + ELSE NULL + END as geometry +; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.stac_search_to_where(j jsonb) + RETURNS text + LANGUAGE plpgsql + STABLE +AS $function$ +DECLARE _where text := cql2_query(search_to_cql2(j)); +BEGIN + IF _where IS NULL OR btrim(_where) = '' THEN RETURN ' TRUE '; END IF; + RETURN _where; +END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.t2s(text) + RETURNS text + LANGUAGE sql + IMMUTABLE PARALLEL SAFE STRICT +AS $function$ + SELECT extract(epoch FROM $1::interval)::text || ' s'; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.table_empty(text) + RETURNS boolean + LANGUAGE plpgsql +AS $function$ +DECLARE + retval boolean; +BEGIN + EXECUTE format($q$ + SELECT NOT EXISTS (SELECT 1 FROM %I LIMIT 1) + $q$, + $1 + ) INTO retval; + RETURN retval; +END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.temporal_op_query(op text, args jsonb) + RETURNS text + LANGUAGE plpgsql + STABLE STRICT +AS $function$ +DECLARE + ll text := 'datetime'; + lh text := 'end_datetime'; + rrange tstzrange; + rl text; + rh text; + outq text; +BEGIN + rrange := parse_dtrange(args->1); + RAISE NOTICE 'Constructing temporal query OP: %, ARGS: %, RRANGE: %', op, args, rrange; + op := lower(op); + rl := format('%L::timestamptz', lower(rrange)); + rh := format('%L::timestamptz', upper(rrange)); + outq := CASE op + WHEN 't_before' THEN 'lh < rl' + WHEN 't_after' THEN 'll > rh' + WHEN 't_meets' THEN 'lh = rl' + WHEN 't_metby' THEN 'll = rh' + WHEN 't_overlaps' THEN 'll < rl AND rl < lh < rh' + WHEN 't_overlappedby' THEN 'rl < ll < rh AND lh > rh' + WHEN 't_starts' THEN 'll = rl AND lh < rh' + WHEN 't_startedby' THEN 'll = rl AND lh > rh' + WHEN 't_during' THEN 'll > rl AND lh < rh' + WHEN 't_contains' THEN 'll < rl AND lh > rh' + WHEN 't_finishes' THEN 'll > rl AND lh = rh' + WHEN 't_finishedby' THEN 'll < rl AND lh = rh' + WHEN 't_equals' THEN 'll = rl AND lh = rh' + WHEN 't_disjoint' THEN 'NOT (ll <= rh AND lh >= rl)' + WHEN 't_intersects' THEN 'll <= rh AND lh >= rl' + WHEN 'anyinteracts' THEN 'll <= rh AND lh >= rl' + END; + outq := regexp_replace(outq, '\mll\M', ll); + outq := regexp_replace(outq, '\mlh\M', lh); + outq := regexp_replace(outq, '\mrl\M', rl); + outq := regexp_replace(outq, '\mrh\M', rh); + outq := format('(%s)', outq); + RETURN outq; +END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.tileenvelope(zoom integer, x integer, y integer) + RETURNS geometry + LANGUAGE sql + IMMUTABLE PARALLEL SAFE +AS $function$ +WITH t AS ( + SELECT + 20037508.3427892 as merc_max, + -20037508.3427892 as merc_min, + (2 * 20037508.3427892) / (2 ^ zoom) as tile_size +) +SELECT st_makeenvelope( + merc_min + (tile_size * x), + merc_max - (tile_size * (y + 1)), + merc_min + (tile_size * (x + 1)), + merc_max - (tile_size * y), + 3857 +) FROM t; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.to_float(jsonb) + RETURNS double precision + LANGUAGE sql + IMMUTABLE PARALLEL SAFE STRICT COST 5000 +AS $function$ + SELECT ($1->>0)::float; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.to_int(jsonb) + RETURNS integer + LANGUAGE sql + IMMUTABLE PARALLEL SAFE STRICT COST 5000 +AS $function$ + SELECT floor(($1->>0)::float)::int; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.to_text(jsonb) + RETURNS text + LANGUAGE sql + IMMUTABLE PARALLEL SAFE STRICT COST 5000 +AS $function$ + SELECT CASE WHEN jsonb_typeof($1) IN ('array','object') THEN $1::text ELSE $1->>0 END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.to_text_array(jsonb) + RETURNS text[] + LANGUAGE sql + IMMUTABLE PARALLEL SAFE STRICT COST 5000 +AS $function$ + SELECT + CASE jsonb_typeof($1) + WHEN 'array' THEN ARRAY(SELECT jsonb_array_elements_text($1)) + ELSE ARRAY[$1->>0] + END + ; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.to_tstz(jsonb) + RETURNS timestamp with time zone + LANGUAGE sql + IMMUTABLE PARALLEL SAFE STRICT COST 5000 + SET "TimeZone" TO 'UTC' +AS $function$ + SELECT ($1->>0)::timestamptz; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.unnest_collection(collection_ids text[] DEFAULT NULL::text[]) + RETURNS SETOF text + LANGUAGE plpgsql + STABLE +AS $function$ + DECLARE + BEGIN + IF collection_ids IS NULL THEN + RETURN QUERY SELECT id FROM collections; + END IF; + RETURN QUERY SELECT unnest(collection_ids); + END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.update_collection_extents() + RETURNS void + LANGUAGE sql +AS $function$ +UPDATE collections + SET content = jsonb_set_lax( + content, + '{extent}'::text[], + collection_extent(id, FALSE), + true, + 'use_json_null' + ) +; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.update_item(content jsonb) + RETURNS void + LANGUAGE plpgsql + SET search_path TO 'pgstac', 'public' +AS $function$ +DECLARE + old items %ROWTYPE; + out items%ROWTYPE; +BEGIN + PERFORM delete_item(content->>'id', content->>'collection'); + PERFORM create_item(content); +END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.upsert_item(data jsonb) + RETURNS void + LANGUAGE sql + SET search_path TO 'pgstac', 'public' +AS $function$ + INSERT INTO items_staging_upsert (content) VALUES (data); +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.upsert_items(data jsonb) + RETURNS void + LANGUAGE sql + SET search_path TO 'pgstac', 'public' +AS $function$ + INSERT INTO items_staging_upsert (content) + SELECT * FROM jsonb_array_elements(data); +$function$ +; + +CREATE OR REPLACE PROCEDURE pgstac.validate_constraints() + LANGUAGE plpgsql +AS $procedure$ +DECLARE + q text; +BEGIN + FOR q IN + SELECT + FORMAT( + 'ALTER TABLE %I.%I VALIDATE CONSTRAINT %I;', + nsp.nspname, + cls.relname, + con.conname + ) + + FROM pg_constraint AS con + JOIN pg_class AS cls + ON con.conrelid = cls.oid + JOIN pg_namespace AS nsp + ON cls.relnamespace = nsp.oid + WHERE convalidated = FALSE AND contype in ('c','f') + AND nsp.nspname = 'pgstac' + LOOP + RAISE NOTICE '%', q; + PERFORM run_or_queue(q); + COMMIT; + END LOOP; +END; +$procedure$ +; + +CREATE OR REPLACE FUNCTION pgstac.xyzsearch(_x integer, _y integer, _z integer, queryhash text, fields jsonb DEFAULT NULL::jsonb, _scanlimit integer DEFAULT 10000, _limit integer DEFAULT 100, _timelimit interval DEFAULT '00:00:05'::interval, exitwhenfull boolean DEFAULT true, skipcovered boolean DEFAULT true) + RETURNS json + LANGUAGE sql +AS $function$ + SELECT * FROM geometrysearch( + st_transform(tileenvelope(_z, _x, _y), 4326), + queryhash, + fields, + _scanlimit, + _limit, + _timelimit, + exitwhenfull, + skipcovered + ); +$function$ +; + +CREATE OR REPLACE FUNCTION public.pgstac_admin_owns() + RETURNS void + LANGUAGE plpgsql +AS $function$ +DECLARE + f RECORD; +BEGIN + FOR f IN ( + SELECT + concat( + oid::regproc::text, + '(', + coalesce(pg_get_function_identity_arguments(oid),''), + ')' + ) AS name, + CASE prokind WHEN 'f' THEN 'FUNCTION' WHEN 'p' THEN 'PROCEDURE' WHEN 'a' THEN 'AGGREGATE' END as typ + FROM pg_proc + WHERE + pronamespace=to_regnamespace('pgstac') + AND proowner != to_regrole('pgstac_admin') + AND proname NOT LIKE 'pg_stat%' + ) + LOOP + BEGIN + EXECUTE format('ALTER %s %s OWNER TO pgstac_admin;', f.typ, f.name); + EXCEPTION WHEN others THEN + RAISE NOTICE '%, skipping', SQLERRM USING ERRCODE = SQLSTATE; + END; + END LOOP; + FOR f IN ( + SELECT + oid::regclass::text as name, + CASE relkind + WHEN 'i' THEN 'INDEX' + WHEN 'I' THEN 'INDEX' + WHEN 'p' THEN 'TABLE' + WHEN 'r' THEN 'TABLE' + WHEN 'v' THEN 'VIEW' + WHEN 'S' THEN 'SEQUENCE' + ELSE NULL + END as typ + FROM pg_class + WHERE relnamespace=to_regnamespace('pgstac') and relowner != to_regrole('pgstac_admin') AND relkind IN ('r','p','v','S') AND relname NOT LIKE 'pg_stat' + ) + LOOP + BEGIN + EXECUTE format('ALTER %s %s OWNER TO pgstac_admin;', f.typ, f.name); + EXCEPTION WHEN others THEN + RAISE NOTICE '%, skipping', SQLERRM USING ERRCODE = SQLSTATE; END; + END LOOP; + RETURN; +END; +$function$ +; + +drop table "pgstac"."format_item_cache"; + +drop table "pgstac"."search_wheres"; + +create table "pgstac"."item_field_registry" ( + "collection" text not null, + "path" text not null, + "is_leaf" boolean default true, + "value_kinds" text[] default '{}'::text[], + "first_seen" timestamp with time zone not null default now(), + "last_seen" timestamp with time zone not null default now() +); + + +create table "pgstac"."item_fragments" ( + "id" bigint not null default nextval('item_fragments_id_seq'::regclass), + "collection" text not null, + "hash" bytea not null, + "content" jsonb, + "links_template" jsonb, + "created_at" timestamp with time zone not null default now() +); + + +create table "pgstac"."items_deleted_log" ( + "id" bigint generated always as identity not null, + "item_id" text not null, + "collection" text not null, + "partition" text, + "datetime" timestamp with time zone, + "end_datetime" timestamp with time zone, + "item_hash" bytea not null default '\x'::bytea, + "deleted_at" timestamp with time zone not null default now() +); + + +alter table "pgstac"."collections" drop column "base_item"; + +alter table "pgstac"."collections" add column "fragment_config" text[]; + +alter table "pgstac"."items" drop column "content"; + +alter table "pgstac"."items" add column "assets" jsonb; + +alter table "pgstac"."items" add column "bands" jsonb; + +alter table "pgstac"."items" add column "bbox" jsonb; + +alter table "pgstac"."items" add column "constellation" text; + +alter table "pgstac"."items" add column "created" timestamp with time zone; + +alter table "pgstac"."items" add column "datetime_is_range" boolean not null default false; + +alter table "pgstac"."items" add column "eo_cloud_cover" double precision; + +alter table "pgstac"."items" add column "eo_snow_cover" double precision; + +alter table "pgstac"."items" add column "extra" jsonb; + +alter table "pgstac"."items" add column "file_byte_order" text; + +alter table "pgstac"."items" add column "file_checksum" text; + +alter table "pgstac"."items" add column "file_header_size" bigint; + +alter table "pgstac"."items" add column "file_size" bigint; + +alter table "pgstac"."items" add column "fragment_id" bigint; + +alter table "pgstac"."items" add column "gsd" double precision; + +alter table "pgstac"."items" add column "instruments" text[]; + +alter table "pgstac"."items" add column "item_hash" bytea not null default '\x'::bytea; + +alter table "pgstac"."items" add column "link_hrefs" text[]; + +alter table "pgstac"."items" add column "links" jsonb; + +alter table "pgstac"."items" add column "mission" text; + +alter table "pgstac"."items" add column "pgstac_updated_at" timestamp with time zone not null default now(); + +alter table "pgstac"."items" add column "platform" text; + +alter table "pgstac"."items" add column "proj_bbox" jsonb; + +alter table "pgstac"."items" add column "proj_centroid" jsonb; + +alter table "pgstac"."items" add column "proj_code" text; + +alter table "pgstac"."items" add column "proj_geometry" jsonb; + +alter table "pgstac"."items" add column "proj_projjson" jsonb; + +alter table "pgstac"."items" add column "proj_shape" jsonb; + +alter table "pgstac"."items" add column "proj_transform" jsonb; + +alter table "pgstac"."items" add column "proj_wkt2" text; - output := jsonb_build_object( - 'id', _item.id, - 'geometry', geom, - 'collection', _item.collection, - 'type', 'Feature' - ); - IF _item.bbox IS NOT NULL THEN - output := output || jsonb_build_object('bbox', _item.bbox); - END IF; - IF hydrated_stac_version IS NOT NULL THEN - output := output || jsonb_build_object('stac_version', hydrated_stac_version); - END IF; - IF hydrated_stac_extensions IS NOT NULL AND hydrated_stac_extensions <> '[]'::jsonb THEN - output := output || jsonb_build_object('stac_extensions', hydrated_stac_extensions); - END IF; - IF _item.links IS NOT NULL THEN - output := output || jsonb_build_object('links', _item.links); - END IF; - IF merged_assets != '{}'::jsonb THEN - output := output || jsonb_build_object('assets', merged_assets); - END IF; - IF merged_properties IS NOT NULL THEN - output := output || jsonb_build_object('properties', merged_properties); - END IF; - IF _item.extra IS NOT NULL THEN - output := output || _item.extra; - END IF; +alter table "pgstac"."items" add column "properties" jsonb; + +alter table "pgstac"."items" add column "sat_absolute_orbit" integer; + +alter table "pgstac"."items" add column "sat_anx_datetime" timestamp with time zone; + +alter table "pgstac"."items" add column "sat_orbit_state" text; + +alter table "pgstac"."items" add column "sat_platform_international_designator" text; + +alter table "pgstac"."items" add column "sat_relative_orbit" integer; + +alter table "pgstac"."items" add column "sci_citation" text; + +alter table "pgstac"."items" add column "sci_doi" text; + +alter table "pgstac"."items" add column "sci_publications" jsonb; + +alter table "pgstac"."items" add column "stac_extensions" jsonb default '[]'::jsonb; + +alter table "pgstac"."items" add column "stac_version" text; + +alter table "pgstac"."items" add column "updated" timestamp with time zone; + +alter table "pgstac"."items" add column "view_azimuth" double precision; + +alter table "pgstac"."items" add column "view_incidence_angle" double precision; + +alter table "pgstac"."items" add column "view_moon_azimuth" double precision; + +alter table "pgstac"."items" add column "view_moon_elevation" double precision; + +alter table "pgstac"."items" add column "view_off_nadir" double precision; + +alter table "pgstac"."items" add column "view_sun_azimuth" double precision; + +alter table "pgstac"."items" add column "view_sun_elevation" double precision; + +alter table "pgstac"."partition_stats" drop column "keys"; + +alter table "pgstac"."partition_stats" add column "collection" text; + +alter table "pgstac"."partition_stats" add column "dirty" boolean not null default false; + +alter table "pgstac"."partition_stats" add column "indexes_pending" boolean not null default false; + +alter table "pgstac"."partition_stats" add column "n" bigint; + +alter table "pgstac"."searches" add column "context_count" bigint; + +alter table "pgstac"."searches" add column "created_at" timestamp with time zone default now(); + +alter table "pgstac"."searches" add column "statslastupdated" timestamp with time zone; + +alter table "pgstac"."searches" alter column "hash" drop expression; + +alter sequence "pgstac"."item_fragments_id_seq" owned by "pgstac"."item_fragments"."id"; + +drop function if exists "pgstac"."check_partition"(_collection text, _dtrange tstzrange, _edtrange tstzrange); + +drop function if exists "pgstac"."chunker"(_where text, OUT s timestamp with time zone, OUT e timestamp with time zone); + +drop function if exists "pgstac"."collection_base_item"(cid text); + +drop function if exists "pgstac"."collection_base_item"(content jsonb); + +drop function if exists "pgstac"."collection_search_matched"(_search jsonb, OUT matched bigint); + +drop function if exists "pgstac"."collection_search_rows"(_search jsonb); + +drop function if exists "pgstac"."content_hydrate"(_item items, _collection collections, fields jsonb); + +drop function if exists "pgstac"."content_hydrate"(_item items, fields jsonb); + +drop function if exists "pgstac"."content_hydrate"(_item jsonb, _base_item jsonb, fields jsonb); + +drop function if exists "pgstac"."content_nonhydrated"(_item items, fields jsonb); + +drop function if exists "pgstac"."content_slim"(_item jsonb); + +drop function if exists "pgstac"."create_collection"(data jsonb); + +drop function if exists "pgstac"."create_table_constraints"(t text, _dtrange tstzrange, _edtrange tstzrange); + +drop function if exists "pgstac"."drop_table_constraints"(t text); + +drop function if exists "pgstac"."format_item"(_item items, _fields jsonb, _hydrated boolean); + +drop function if exists "pgstac"."get_token_filter"(_sortby jsonb, token_item items, prev boolean, inclusive boolean); + +drop function if exists "pgstac"."get_token_record"(_token text, OUT prev boolean, OUT item items); + +drop function if exists "pgstac"."get_token_val_str"(_field text, _item items); + +drop function if exists "pgstac"."merge_jsonb"(_a jsonb, _b jsonb); + +drop function if exists "pgstac"."paging_collections"(j jsonb); + +drop function if exists "pgstac"."paging_dtrange"(j jsonb); + +drop function if exists "pgstac"."parse_sort_dir"(_dir text, reverse boolean); + +drop function if exists "pgstac"."partition_after_triggerfunc"(); + +drop function if exists "pgstac"."partition_queries"(_where text, _orderby text, partitions text[]); + +drop function if exists "pgstac"."partition_query_view"(_where text, _orderby text, _limit integer); + +drop function if exists "pgstac"."search_cursor"(_search jsonb); + + +drop function if exists "pgstac"."search_query"(_search jsonb, updatestats boolean, _metadata jsonb); + +drop function if exists "pgstac"."search_rows"(_where text, _orderby text, partitions text[], _limit integer); + +drop function if exists "pgstac"."sort_dir_to_op"(_dir text, prev boolean); + +drop function if exists "pgstac"."sort_sqlorderby"(_search jsonb, reverse boolean); + +drop function if exists "pgstac"."strip_jsonb"(_a jsonb, _b jsonb); + +drop function if exists "pgstac"."update_collection"(data jsonb); + +drop function if exists "pgstac"."update_partition_stats"(_partition text, istrigger boolean); + +drop function if exists "pgstac"."update_partition_stats_q"(_partition text, istrigger boolean); + +drop function if exists "pgstac"."upsert_collection"(data jsonb); + +drop function if exists "pgstac"."where_stats"(inwhere text, updatestats boolean, conf jsonb); + +drop function if exists "pgstac"."parse_dtrange"(_indate text, relative_base timestamp with time zone); + +create or replace view "pgstac"."collections_asitems" as SELECT id, + geometry, + 'collections'::text AS collection, + datetime, + end_datetime, + content - '{links,assets,stac_version,stac_extensions}'::text AS properties, + jsonb_build_object('properties', content - '{links,assets,stac_version,stac_extensions}'::text, 'links', content -> 'links'::text, 'assets', content -> 'assets'::text, 'stac_version', content -> 'stac_version'::text, 'stac_extensions', content -> 'stac_extensions'::text) AS content, + content AS collectionjson + FROM collections; + + +create or replace view "pgstac"."partition_sys_meta" as SELECT partition.partition, + replace(replace( + CASE + WHEN pg_partition_tree.level = 1 THEN partition_expr.partition_expr + ELSE parent_partition_expr.parent_partition_expr + END, 'FOR VALUES IN ('''::text, ''::text), ''')'::text, ''::text) AS collection, + pg_partition_tree.level, + c.reltuples, + c.relhastriggers, + COALESCE(get_tstz_constraint(c.oid, 'datetime'::text), constraint_tstzrange(pg_get_expr(c.relpartbound, c.oid)), inf_range.inf_range) AS constraint_dtrange, + COALESCE(get_tstz_constraint(c.oid, 'end_datetime'::text), inf_range.inf_range) AS constraint_edtrange + FROM pg_partition_tree('items'::regclass) pg_partition_tree(relid, parentrelid, isleaf, level) + JOIN pg_class c ON pg_partition_tree.relid::oid = c.oid + JOIN pg_class parent ON pg_partition_tree.parentrelid::oid = parent.oid AND pg_partition_tree.isleaf + LEFT JOIN pg_constraint edt ON edt.conrelid = c.oid AND edt.contype = 'c'::"char" + JOIN LATERAL get_partition_name(pg_partition_tree.relid) partition(partition) ON true + JOIN LATERAL pg_get_expr(c.relpartbound, c.oid) partition_expr(partition_expr) ON true + JOIN LATERAL pg_get_expr(parent.relpartbound, parent.oid) parent_partition_expr(parent_partition_expr) ON true + JOIN LATERAL tstzrange('-infinity'::timestamp with time zone, 'infinity'::timestamp with time zone, '[]'::text) inf_range(inf_range) ON true + JOIN LATERAL get_tstz_constraint(c.oid, 'datetime'::text) datetime_constraint(datetime_constraint) ON true + JOIN LATERAL get_tstz_constraint(c.oid, 'end_datetime'::text) end_datetime_constraint(end_datetime_constraint) ON true + WHERE pg_partition_tree.isleaf; + + +create or replace view "pgstac"."partitions_view" as SELECT (parse_ident(pg_partition_tree.relid::text))[cardinality(parse_ident(pg_partition_tree.relid::text))] AS partition, + replace(replace( + CASE + WHEN pg_partition_tree.level = 1 THEN partition_expr.partition_expr + ELSE parent_partition_expr.parent_partition_expr + END, 'FOR VALUES IN ('''::text, ''::text), ''')'::text, ''::text) AS collection, + pg_partition_tree.level, + c.reltuples, + c.relhastriggers, + COALESCE(get_tstz_constraint(c.oid, 'datetime'::text), constraint_tstzrange(pg_get_expr(c.relpartbound, c.oid)), inf_range.inf_range) AS constraint_dtrange, + COALESCE(get_tstz_constraint(c.oid, 'end_datetime'::text), inf_range.inf_range) AS constraint_edtrange, + ps.dtrange, + ps.edtrange, + ps.spatial, + ps.last_updated + FROM pg_partition_tree('items'::regclass) pg_partition_tree(relid, parentrelid, isleaf, level) + JOIN pg_class c ON pg_partition_tree.relid::oid = c.oid + JOIN pg_class parent ON pg_partition_tree.parentrelid::oid = parent.oid AND pg_partition_tree.isleaf + LEFT JOIN pg_constraint edt ON edt.conrelid = c.oid AND edt.contype = 'c'::"char" + JOIN LATERAL get_partition_name(pg_partition_tree.relid) partition(partition) ON true + JOIN LATERAL pg_get_expr(c.relpartbound, c.oid) partition_expr(partition_expr) ON true + JOIN LATERAL pg_get_expr(parent.relpartbound, parent.oid) parent_partition_expr(parent_partition_expr) ON true + JOIN LATERAL tstzrange('-infinity'::timestamp with time zone, 'infinity'::timestamp with time zone, '[]'::text) inf_range(inf_range) ON true + JOIN LATERAL get_tstz_constraint(c.oid, 'datetime'::text) datetime_constraint(datetime_constraint) ON true + JOIN LATERAL get_tstz_constraint(c.oid, 'end_datetime'::text) end_datetime_constraint(end_datetime_constraint) ON true + LEFT JOIN ( SELECT partition_stats.partition, + partition_stats.dtrange, + partition_stats.edtrange, + partition_stats.spatial, + partition_stats.last_updated + FROM partition_stats) ps USING (partition) + WHERE pg_partition_tree.isleaf; - RETURN jsonb_fields(output, fields); -END; -$function$ -; create or replace view "pgstac"."pgstac_indexes" as SELECT schemaname, tablename, @@ -1996,160 +6241,276 @@ CREATE TRIGGER items_before_update_trigger BEFORE UPDATE ON pgstac.items FOR EAC CREATE TRIGGER items_delete_log_after_delete_trigger AFTER DELETE ON pgstac.items REFERENCING OLD TABLE AS old_rows FOR EACH STATEMENT EXECUTE FUNCTION items_delete_log_trigger(); -CREATE TRIGGER items_after_delete_trigger AFTER DELETE ON pgstac.items REFERENCING OLD TABLE AS newdata FOR EACH STATEMENT EXECUTE FUNCTION partition_after_triggerfunc(); +CREATE INDEX item_field_registry_path_idx ON pgstac.item_field_registry USING btree (path); -CREATE TRIGGER items_after_update_trigger AFTER UPDATE ON pgstac.items REFERENCING NEW TABLE AS newdata FOR EACH STATEMENT EXECUTE FUNCTION partition_after_triggerfunc(); +CREATE UNIQUE INDEX item_field_registry_pkey ON pgstac.item_field_registry USING btree (collection, path); -CREATE OR REPLACE FUNCTION pgstac.collection_base_item(cid text) - RETURNS jsonb - LANGUAGE sql - STABLE PARALLEL SAFE -AS $function$ - SELECT collection_base_item(content) - FROM collections - WHERE id = cid - LIMIT 1; -$function$ -; +CREATE UNIQUE INDEX item_fragments_collection_hash_key ON pgstac.item_fragments USING btree (collection, hash); -CREATE OR REPLACE FUNCTION pgstac.content_dehydrate(content jsonb) - RETURNS items +CREATE INDEX item_fragments_collection_idx ON pgstac.item_fragments USING btree (collection); + +CREATE UNIQUE INDEX item_fragments_pkey ON pgstac.item_fragments USING btree (id); + +CREATE INDEX items_deleted_log_deleted_at_idx ON pgstac.items_deleted_log USING btree (deleted_at); + +CREATE UNIQUE INDEX items_deleted_log_pkey ON pgstac.items_deleted_log USING btree (id); + +CREATE INDEX items_fragment_id_idx ON ONLY pgstac.items USING btree (fragment_id) WHERE (fragment_id IS NOT NULL); + +CREATE INDEX partition_stats_collection_idx ON pgstac.partition_stats USING btree (collection); + +CREATE INDEX partition_stats_dirty_idx ON pgstac.partition_stats USING btree (partition) WHERE dirty; + +CREATE INDEX partition_stats_indexes_pending_idx ON pgstac.partition_stats USING btree (partition) WHERE indexes_pending; + +CREATE INDEX partition_stats_spatial_idx ON pgstac.partition_stats USING gist (spatial) WHERE (spatial IS NOT NULL); + +CREATE INDEX searches_lastused_anon_idx ON pgstac.searches USING btree (lastused) WHERE (metadata IS NOT NULL); + +alter table "pgstac"."item_field_registry" add constraint "item_field_registry_pkey" PRIMARY KEY using index "item_field_registry_pkey"; + +alter table "pgstac"."item_fragments" add constraint "item_fragments_pkey" PRIMARY KEY using index "item_fragments_pkey"; + +alter table "pgstac"."items_deleted_log" add constraint "items_deleted_log_pkey" PRIMARY KEY using index "items_deleted_log_pkey"; + +alter table "pgstac"."item_field_registry" add constraint "item_field_registry_collection_fkey" FOREIGN KEY ("collection") REFERENCES "pgstac"."collections"("id") ON DELETE CASCADE NOT VALID; + +alter table "pgstac"."item_field_registry" validate constraint "item_field_registry_collection_fkey"; + +alter table "pgstac"."item_fragments" add constraint "item_fragments_collection_fkey" FOREIGN KEY ("collection") REFERENCES "pgstac"."collections"("id") ON DELETE CASCADE NOT VALID; + +alter table "pgstac"."item_fragments" validate constraint "item_fragments_collection_fkey"; + +alter table "pgstac"."item_fragments" add constraint "item_fragments_collection_hash_key" UNIQUE using index "item_fragments_collection_hash_key"; + +CREATE OR REPLACE FUNCTION pgstac.collection_extent(_collection text, runupdate boolean DEFAULT false) + RETURNS jsonb LANGUAGE plpgsql - STABLE AS $function$ DECLARE - out items; - props jsonb; -BEGIN - out.id := content->>'id'; - out.geometry := stac_geom(content); - out.collection := content->>'collection'; - props := content->'properties'; - out.datetime := stac_datetime(content); - out.end_datetime := stac_end_datetime(content); - out.datetime_is_range := CASE - WHEN props->'datetime' IS NOT NULL AND props->'datetime' <> 'null'::jsonb THEN FALSE - ELSE ( - (props->'start_datetime' IS NOT NULL AND props->'start_datetime' <> 'null'::jsonb) - OR (props->'end_datetime' IS NOT NULL AND props->'end_datetime' <> 'null'::jsonb) - ) - END; - out.stac_version := content->>'stac_version'; - out.stac_extensions := COALESCE(content->'stac_extensions', '[]'::jsonb); - out.pgstac_updated_at := now(); - out.content_hash := pgstac_item_hash(content); - - -- Split columns: dedicated storage for standard top-level STAC fields. - -- These enable index-only scans on promoted queryables and JSONB-free hot paths. - out.bbox := content->'bbox'; - out.links := COALESCE(content->'links', '[]'::jsonb); - out.assets := COALESCE(content->'assets', '{}'::jsonb); - out.properties := strip_promoted_properties(props); - out.extra := content - '{id,geometry,collection,type,bbox,links,assets,properties,stac_version,stac_extensions}'::text[]; - - out.created := (props->>'created')::timestamptz; - out.updated := (props->>'updated')::timestamptz; - out.platform := props->>'platform'; - out.instruments := to_text_array(props->'instruments'); - out.constellation := props->>'constellation'; - out.mission := props->>'mission'; - out.eo_cloud_cover := (props->>'eo:cloud_cover')::float8; - out.eo_bands := props->'eo:bands'; - out.eo_snow_cover := (props->>'eo:snow_cover')::float8; - out.gsd := (props->>'gsd')::float8; - out.proj_epsg := (props->>'proj:epsg')::integer; - out.proj_wkt2 := props->>'proj:wkt2'; - out.proj_projjson := props->'proj:projjson'; - out.proj_bbox := props->'proj:bbox'; - out.proj_centroid := props->'proj:centroid'; - out.proj_shape := props->'proj:shape'; - out.proj_transform := props->'proj:transform'; - out.sci_doi := props->>'sci:doi'; - out.sci_citation := props->>'sci:citation'; - out.sci_publications := props->'sci:publications'; - out.view_off_nadir := (props->>'view:off_nadir')::float8; - out.view_incidence_angle := (props->>'view:incidence_angle')::float8; - out.view_azimuth := (props->>'view:azimuth')::float8; - out.view_sun_azimuth := (props->>'view:sun_azimuth')::float8; - out.view_sun_elevation := (props->>'view:sun_elevation')::float8; - out.file_size := (props->>'file:size')::bigint; - out.file_header_size := (props->>'file:header_size')::bigint; - out.file_checksum := props->>'file:checksum'; - out.file_byte_order := props->>'file:byte_order'; - out.file_values_regex := props->>'file:values_regex'; - out.sat_orbit_state := props->>'sat:orbit_state'; - out.sat_relative_orbit := (props->>'sat:relative_orbit')::integer; - out.sat_absolute_orbit := (props->>'sat:absolute_orbit')::integer; - - -- fragment_id is NULL on initial dehydration; assigned by the staging trigger. - out.fragment_id := NULL; - RETURN out; + geom_extent geometry; + mind timestamptz; + maxd timestamptz; + extent jsonb; +BEGIN + IF runupdate THEN + PERFORM tighten_partition_stats(partition) + FROM partitions_view WHERE collection=_collection; + END IF; + SELECT + min(lower(dtrange)), + max(upper(edtrange)), + st_extent(spatial) + INTO + mind, + maxd, + geom_extent + FROM partitions_view + WHERE collection=_collection; + + IF geom_extent IS NOT NULL AND mind IS NOT NULL AND maxd IS NOT NULL THEN + extent := jsonb_build_object( + 'spatial', jsonb_build_object( + 'bbox', to_jsonb(array[array[st_xmin(geom_extent), st_ymin(geom_extent), st_xmax(geom_extent), st_ymax(geom_extent)]]) + ), + 'temporal', jsonb_build_object( + 'interval', to_jsonb(array[array[mind, maxd]]) + ) + ); + RETURN extent; + END IF; + RETURN NULL; END; $function$ ; -CREATE OR REPLACE FUNCTION pgstac.content_hydrate(_item items, fields jsonb DEFAULT '{}'::jsonb) +CREATE OR REPLACE FUNCTION pgstac.collection_search(_search jsonb DEFAULT '{}'::jsonb) RETURNS jsonb LANGUAGE plpgsql STABLE PARALLEL SAFE AS $function$ DECLARE - geom jsonb; - output jsonb; - frag_content jsonb; - merged_assets jsonb; - merged_properties jsonb; - hydrated_stac_version text; - hydrated_stac_extensions jsonb; + _limit int := coalesce((_search->>'limit')::float::int, 10); + _token text := _search->>'token'; + is_prev boolean := _token LIKE 'prev:%'; + keyset text := nullif(regexp_replace(coalesce(_token, ''), '^(next|prev):', ''), ''); + _q text; _ctx text; + number_matched bigint; + number_returned bigint; + acc jsonb := '[]'::jsonb; cnt int := 0; + first_k text[]; last_k text[]; fwd_first_k text[]; fwd_last_k text[]; + has_more boolean; next_present boolean; prev_present boolean; + next_tok text; prev_tok text; + links jsonb := '[]'::jsonb; + burl text := concat(rtrim(base_url(_search->'conf'), '/'), '/collections'); + -- typed loop targets (not a `record`) so plpgsql_check can analyze the dynamic EXECUTE. + _content jsonb; _keys text[]; BEGIN - IF include_field('geometry', fields) THEN - geom := ST_ASGeoJson(_item.geometry, 20)::jsonb; - END IF; + SELECT query, ctx_query INTO _q, _ctx FROM collection_search_plan(_search, _token); + EXECUTE _ctx INTO number_matched; -- numberMatched, always + + -- the plan query returns (content, keys); +1 over _limit detects a further page. + FOR _content, _keys IN EXECUTE _q USING (_limit + 1) + LOOP + cnt := cnt + 1; + IF cnt = 1 THEN first_k := _keys; END IF; + IF cnt <= _limit THEN acc := acc || _content; last_k := _keys; END IF; + END LOOP; - -- Fetch shared fragment content (NULL when item has no fragment). - IF _item.fragment_id IS NOT NULL THEN - SELECT content INTO frag_content FROM item_fragments WHERE id = _item.fragment_id; + has_more := cnt > _limit; + IF is_prev THEN + acc := flip_jsonb_array(acc); + fwd_first_k := last_k; fwd_last_k := first_k; + next_present := (keyset IS NOT NULL); -- a prev page always has the origin ahead + prev_present := has_more; -- ... and a further-back page iff more remain + ELSE + fwd_first_k := first_k; fwd_last_k := last_k; + next_present := has_more; + prev_present := (keyset IS NOT NULL); END IF; + IF cnt > 0 AND next_present THEN next_tok := keyset_encode(fwd_last_k); END IF; + IF cnt > 0 AND prev_present THEN prev_tok := keyset_encode(fwd_first_k); END IF; - -- Merge: fragment provides shared asset/property values; per-item provides individual values. - merged_assets := jsonb_merge_level1(frag_content->'assets', _item.assets); - merged_properties := jsonb_merge_level1(frag_content->'properties', _item.properties); - merged_properties := promoted_properties_from_item(_item) || COALESCE(merged_properties, '{}'::jsonb); - hydrated_stac_version := COALESCE(_item.stac_version, frag_content->>'stac_version'); - hydrated_stac_extensions := CASE - WHEN _item.stac_extensions IS NOT NULL AND _item.stac_extensions <> '[]'::jsonb THEN _item.stac_extensions - ELSE COALESCE(frag_content->'stac_extensions', _item.stac_extensions) - END; + number_returned := jsonb_array_length(acc); - output := jsonb_build_object( - 'id', _item.id, - 'geometry', geom, - 'collection', _item.collection, - 'type', 'Feature' - ); - IF _item.bbox IS NOT NULL THEN - output := output || jsonb_build_object('bbox', _item.bbox); + IF next_tok IS NOT NULL THEN + links := links || jsonb_build_object('rel', 'next', 'type', 'application/json', + 'method', 'GET', 'href', burl || '?token=next:' || next_tok); END IF; - IF hydrated_stac_version IS NOT NULL THEN - output := output || jsonb_build_object('stac_version', hydrated_stac_version); + IF prev_tok IS NOT NULL THEN + links := links || jsonb_build_object('rel', 'prev', 'type', 'application/json', + 'method', 'GET', 'href', burl || '?token=prev:' || prev_tok); END IF; - IF hydrated_stac_extensions IS NOT NULL AND hydrated_stac_extensions <> '[]'::jsonb THEN - output := output || jsonb_build_object('stac_extensions', hydrated_stac_extensions); + + RETURN jsonb_build_object( + 'collections', acc, + 'numberMatched', number_matched, + 'numberReturned', number_returned, + 'links', links + ); +END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.content_dehydrate(content jsonb) + RETURNS items + LANGUAGE sql + STABLE +AS $function$ + SELECT + content->>'id' AS id, + stac_geom(content) AS geometry, + content->>'collection' AS collection, + stac_datetime(content) AS datetime, + stac_end_datetime(content) AS end_datetime, + CASE + WHEN (content->'properties')->'datetime' IS NOT NULL AND (content->'properties')->'datetime' <> 'null'::jsonb THEN FALSE + ELSE ( + ((content->'properties')->'start_datetime' IS NOT NULL AND (content->'properties')->'start_datetime' <> 'null'::jsonb) + OR ((content->'properties')->'end_datetime' IS NOT NULL AND (content->'properties')->'end_datetime' <> 'null'::jsonb) + ) + END AS datetime_is_range, + content->>'stac_version' AS stac_version, + COALESCE(content->'stac_extensions', '[]'::jsonb) AS stac_extensions, + now() AS pgstac_updated_at, + pgstac.jsonb_canonical_hash(content) AS item_hash, + NULL::bigint AS fragment_id, + content->'bbox' AS bbox, + CASE WHEN content->'links' IS NOT NULL AND content->'links' <> '[]'::jsonb THEN content->'links' END AS links, + CASE WHEN content->'assets' IS NOT NULL AND content->'assets' <> '{}'::jsonb THEN content->'assets' END AS assets, + strip_promoted_properties(content->'properties') AS properties, + content - '{id,geometry,collection,type,bbox,links,assets,properties,stac_version,stac_extensions}'::text[] AS extra, + + ((content->'properties')->>'created')::timestamptz AS created, + ((content->'properties')->>'updated')::timestamptz AS updated, + ((content->'properties')->>'platform') AS platform, + to_text_array((content->'properties')->'instruments') AS instruments, + ((content->'properties')->>'constellation') AS constellation, + ((content->'properties')->>'mission') AS mission, + ((content->'properties')->>'eo:cloud_cover')::float8 AS eo_cloud_cover, + ((content->'properties')->'bands') AS bands, + ((content->'properties')->>'eo:snow_cover')::float8 AS eo_snow_cover, + ((content->'properties')->>'gsd')::float8 AS gsd, + + ((content->'properties')->>'proj:code') AS proj_code, + ((content->'properties')->'proj:geometry') AS proj_geometry, + ((content->'properties')->>'proj:wkt2') AS proj_wkt2, + ((content->'properties')->'proj:projjson') AS proj_projjson, + ((content->'properties')->'proj:bbox') AS proj_bbox, + ((content->'properties')->'proj:centroid') AS proj_centroid, + ((content->'properties')->'proj:shape') AS proj_shape, + ((content->'properties')->'proj:transform') AS proj_transform, + + ((content->'properties')->>'sci:doi') AS sci_doi, + ((content->'properties')->>'sci:citation') AS sci_citation, + ((content->'properties')->'sci:publications') AS sci_publications, + + ((content->'properties')->>'view:off_nadir')::float8 AS view_off_nadir, + ((content->'properties')->>'view:incidence_angle')::float8 AS view_incidence_angle, + ((content->'properties')->>'view:azimuth')::float8 AS view_azimuth, + ((content->'properties')->>'view:sun_azimuth')::float8 AS view_sun_azimuth, + ((content->'properties')->>'view:sun_elevation')::float8 AS view_sun_elevation, + ((content->'properties')->>'view:moon_azimuth')::float8 AS view_moon_azimuth, + ((content->'properties')->>'view:moon_elevation')::float8 AS view_moon_elevation, + + ((content->'properties')->>'file:size')::bigint AS file_size, + ((content->'properties')->>'file:header_size')::bigint AS file_header_size, + ((content->'properties')->>'file:checksum') AS file_checksum, + ((content->'properties')->>'file:byte_order') AS file_byte_order, + + ((content->'properties')->>'sat:orbit_state') AS sat_orbit_state, + ((content->'properties')->>'sat:relative_orbit')::integer AS sat_relative_orbit, + ((content->'properties')->>'sat:absolute_orbit')::integer AS sat_absolute_orbit, + ((content->'properties')->>'sat:platform_international_designator') AS sat_platform_international_designator, + ((content->'properties')->>'sat:anx_datetime')::timestamptz AS sat_anx_datetime, + stac_links_href_array(content->'links') AS link_hrefs, + -- private is operator-managed metadata outside the STAC item; always NULL from ingest + NULL::jsonb AS private; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.cql1_to_cql2(j jsonb) + RETURNS jsonb + LANGUAGE plpgsql + IMMUTABLE STRICT +AS $function$ +DECLARE + args jsonb; + ret jsonb; +BEGIN + RAISE NOTICE 'CQL1_TO_CQL2: %', j; + IF j ? 'filter' THEN + RETURN cql1_to_cql2(j->'filter'); END IF; - IF _item.links IS NOT NULL THEN - output := output || jsonb_build_object('links', _item.links); + IF j ? 'property' THEN + RETURN j; END IF; - IF merged_assets != '{}'::jsonb THEN - output := output || jsonb_build_object('assets', merged_assets); + IF jsonb_typeof(j) = 'array' THEN + SELECT jsonb_agg(cql1_to_cql2(el)) INTO args FROM jsonb_array_elements(j) el; + RETURN args; END IF; - IF merged_properties IS NOT NULL THEN - output := output || jsonb_build_object('properties', merged_properties); + IF jsonb_typeof(j) = 'number' THEN + RETURN j; END IF; - IF _item.extra IS NOT NULL THEN - output := output || _item.extra; + IF jsonb_typeof(j) = 'string' THEN + RETURN j; END IF; - RETURN jsonb_fields(output, fields); + IF jsonb_typeof(j) = 'object' THEN + -- GeoJSON geometry args (Point/Polygon/.../GeometryCollection) are not cql expressions; + -- pass them through unchanged so spatial ops keep their geometry intact. + IF j ? 'type' AND (j ? 'coordinates' OR j ? 'geometries') THEN + RETURN j; + END IF; + SELECT jsonb_build_object( + 'op', key, + 'args', cql1_to_cql2(value) + ) INTO ret + FROM jsonb_each(j) + WHERE j IS NOT NULL; + RETURN ret; + END IF; + RETURN NULL; END; $function$ ; @@ -2215,7 +6576,6 @@ BEGIN END IF; IF j ? 'interval' THEN RAISE EXCEPTION 'Please use temporal operators when using intervals.'; - RETURN NONE; END IF; -- Spatial Query @@ -2223,6 +6583,11 @@ BEGIN RETURN spatial_op_query(op, args); END IF; + -- Full-text Query (pgstac `q` operator) + IF op = 'q' THEN + RETURN q_op_query(args); + END IF; + IF op IN ('a_equals','a_contains','a_contained_by','a_overlaps') THEN IF args->0 ? 'property' THEN leftarg := format('to_text_array(%s)', (queryable(args->0->>'property')).path); @@ -2355,6 +6720,175 @@ END; $function$ ; +CREATE OR REPLACE FUNCTION pgstac.delete_item(_id text, _collection text DEFAULT NULL::text) + RETURNS void + LANGUAGE plpgsql + SECURITY DEFINER +AS $function$ +DECLARE +out items%ROWTYPE; +BEGIN + DELETE FROM items WHERE id = _id AND (_collection IS NULL OR collection=_collection) RETURNING * INTO STRICT out; +END; +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.geojsonsearch(geojson jsonb, queryhash text, fields jsonb DEFAULT NULL::jsonb, _scanlimit integer DEFAULT 10000, _limit integer DEFAULT 100, _timelimit interval DEFAULT '00:00:05'::interval, exitwhenfull boolean DEFAULT true, skipcovered boolean DEFAULT true) + RETURNS json + LANGUAGE sql +AS $function$ + SELECT * FROM geometrysearch( + st_geomfromgeojson(geojson), + queryhash, + fields, + _scanlimit, + _limit, + _timelimit, + exitwhenfull, + skipcovered + ); +$function$ +; + +CREATE OR REPLACE FUNCTION pgstac.geometrysearch(geom geometry, queryhash text, fields jsonb DEFAULT NULL::jsonb, _scanlimit integer DEFAULT 10000, _limit integer DEFAULT 100, _timelimit interval DEFAULT '00:00:05'::interval, exitwhenfull boolean DEFAULT true, skipcovered boolean DEFAULT true) + RETURNS json + LANGUAGE plpgsql +AS $function$ +DECLARE + -- Same adaptive band sizing as search_page: first band ~one page (band_margin headroom), each + -- subsequent band re-sized from the band's observed spatial hit-rate (band_safety headroom). + band_margin CONSTANT numeric := 3.0; + band_safety CONSTANT numeric := 1.5; + band_cap_months CONSTANT int := 18; + search searches%ROWTYPE; + curs refcursor; + _where text; query text; proj_expr text; + iter_record items%ROWTYPE; + page_rows items[] := '{}'::items[]; -- coverage-passing rows, hydrated once at the end + features json; + exit_flag boolean := FALSE; + counter int := 1; scancounter int := 1; remaining_limit int := _scanlimit; + tilearea float; unionedgeom geometry; clippedgeom geometry; + unionedgeom_area float := 0; prev_area float := 0; + _env pred_envelope; bnds record; + lead_field text; eff_dir text; datetime_leading boolean; is_asc boolean; orderby_str text; + cursor_idx int; mo interval := interval '1 month'; band record; band_target numeric; obs_sel numeric; + band_fetched int; guard int := 0; cum_scanned bigint := 0; +BEGIN + -- If the passed in geometry is not an area, coverage tests are meaningless. + IF ST_GeometryType(geom) !~* 'polygon' THEN + skipcovered := FALSE; exitwhenfull := FALSE; + END IF; + -- skipcovered implies exitwhenfull (once covered, nothing new can show through). + IF skipcovered THEN exitwhenfull := TRUE; END IF; + + search := search_fromhash(queryhash); + IF search IS NULL THEN + RAISE EXCEPTION 'Search with Query Hash % Not Found', queryhash; + END IF; + + tilearea := st_area(geom); + _where := format('%s AND st_intersects(geometry, %L::geometry)', search._where, geom); + + -- Discovery envelope from the registered search AND the tile geometry, then the single + -- partition_bounds read (unified per-month histogram) that next_band walks. + _env := search_envelope(search.search); + _env.geom := CASE + WHEN _env.geom IS NULL THEN ST_Envelope(geom) + ELSE ST_Envelope(ST_Intersection(_env.geom, ST_Envelope(geom))) + END; + SELECT * INTO bnds FROM partition_bounds(_env); + + -- Leading sort field/direction + the registered forward orderby. Bands are datetime windows walked + -- newest-first unless the search sorts datetime ascending; within a band rows follow search.orderby. + SELECT (array_agg(field ORDER BY ord))[1], (array_agg(dir ORDER BY ord))[1] + INTO lead_field, eff_dir FROM keyset_sortkeys(search.search); + datetime_leading := (lead_field = 'datetime'); + is_asc := (datetime_leading AND eff_dir = 'ASC'); + orderby_str := search.orderby; + + -- Per-row projection: skip the shared item_fragments lookup when the requested fields are + -- satisfiable from item columns alone (needs_fragment, evaluated once for the whole query). + IF needs_fragment(coalesce(fields, '{}'::jsonb), bnds.collections) THEN + proj_expr := format('content_hydrate(i, %L::jsonb)', coalesce(fields, '{}'::jsonb)); + ELSE + proj_expr := format('content_hydrate(i, %L::jsonb, true)', coalesce(fields, '{}'::jsonb)); + END IF; + + IF array_length(bnds.months, 1) IS NOT NULL THEN + -- Walk bands in the SEARCH's datetime direction: newest month first for a descending search, + -- oldest first when sorting datetime ascending. Mirrors search_page (004_search). Walking the + -- wrong direction returns older items before newer ones for a descending limit/early-exit search, + -- and (because band grouping shifts with the partition_stats histogram) makes the result depend on + -- stats — a bug, since stats must only affect performance, never which rows come back. + cursor_idx := CASE WHEN is_asc THEN 1 ELSE array_length(bnds.months, 1) END; + band_target := (_limit + 1) * band_margin; + <> + WHILE NOT exit_flag AND guard < 80 LOOP + guard := guard + 1; + SELECT * INTO band FROM next_band(bnds.counts, cursor_idx, band_target, band_cap_months, NOT is_asc); + -- process a valid band even when next_band also flags done; stop only on no band. + EXIT bands WHEN band.band_start_idx IS NULL; + query := format( + 'SELECT * FROM items i WHERE i.collection = ANY(%L::text[]) AND i.datetime >= %L AND i.datetime < %L AND %s ORDER BY %s LIMIT %L', + bnds.collections, bnds.months[band.band_start_idx], bnds.months[band.band_end_idx] + mo, _where, orderby_str, remaining_limit); + band_fetched := 0; + OPEN curs FOR EXECUTE query; + LOOP + FETCH curs INTO iter_record; + EXIT WHEN NOT FOUND; + band_fetched := band_fetched + 1; + IF exitwhenfull OR skipcovered THEN -- skip expensive geometry ops when neither is on + clippedgeom := st_intersection(geom, iter_record.geometry); + IF unionedgeom IS NULL THEN + unionedgeom := clippedgeom; + ELSE + unionedgeom := st_union(unionedgeom, clippedgeom); + END IF; + unionedgeom_area := st_area(unionedgeom); + IF skipcovered AND prev_area = unionedgeom_area THEN + scancounter := scancounter + 1; + CONTINUE; + END IF; + prev_area := unionedgeom_area; + END IF; + -- Collect the row; hydration happens once at the end (same content_hydrate path as search_page). + page_rows := page_rows || iter_record; + IF counter >= _limit + OR scancounter > _scanlimit + OR ftime() > _timelimit + OR (exitwhenfull AND unionedgeom_area >= tilearea) + THEN + exit_flag := TRUE; + EXIT; + END IF; + counter := counter + 1; + scancounter := scancounter + 1; + END LOOP; + CLOSE curs; + cum_scanned := cum_scanned + band.scanned; + cursor_idx := band.next_cursor_idx; + EXIT bands WHEN exit_flag; + -- LEARN: size the next band from this band's spatial hit-rate (fetched / rows scanned). + obs_sel := GREATEST(band_fetched::numeric, 0.5) / GREATEST(cum_scanned, 1); + band_target := ((_limit + 1 - counter) / obs_sel) * band_safety; + remaining_limit := _scanlimit - scancounter; + END LOOP; + END IF; + + -- Hydrate the collected rows once (content_hydrate + needs_fragment skip, json output -- 1GB text + -- ceiling, no 256MB jsonb-array limit), exactly like search_page. + EXECUTE format('SELECT coalesce(json_agg(%s), ''[]''::json) FROM unnest($1::items[]) i', proj_expr) + USING page_rows INTO features; + + RETURN json_build_object( + 'type', 'FeatureCollection', + 'features', coalesce(features, '[]'::json) + ); +END; +$function$ +; + CREATE OR REPLACE FUNCTION pgstac.indexdef(q queryables) RETURNS text LANGUAGE plpgsql @@ -2391,14 +6925,32 @@ $function$ CREATE OR REPLACE FUNCTION pgstac.items_staging_triggerfunc() RETURNS trigger LANGUAGE plpgsql + SECURITY DEFINER AS $function$ DECLARE part text; ts timestamptz := clock_timestamp(); nrows int; + batch jsonb[]; BEGIN RAISE NOTICE 'Creating Partitions. %', clock_timestamp() - ts; + -- Fail loudly on items whose collection does not exist instead of silently dropping them + -- in the collections JOINs below (matches the Rust loader, which also errors on this). + IF EXISTS ( + SELECT 1 FROM newdata n + WHERE NOT EXISTS ( + SELECT 1 FROM collections c WHERE c.id = n.content->>'collection' + ) + ) THEN + RAISE EXCEPTION 'cannot load item(s) into nonexistent collection(s): %', + (SELECT string_agg(DISTINCT coalesce(n.content->>'collection', ''), ', ') + FROM newdata n + WHERE NOT EXISTS ( + SELECT 1 FROM collections c WHERE c.id = n.content->>'collection' + )); + END IF; + FOR part IN WITH t AS ( SELECT n.content->>'collection' as collection, @@ -2417,117 +6969,27 @@ BEGIN RAISE NOTICE 'Partition %', part; END LOOP; - RAISE NOTICE 'Creating temp table with data to be added. %', clock_timestamp() - ts; - -- TEMP tables are session-scoped. Reusing the tmpdata name here is safe even - -- with concurrent ingest sessions because each session gets its own temp schema. - DROP TABLE IF EXISTS tmpdata; - CREATE TEMP TABLE tmpdata ON COMMIT DROP AS - SELECT - -- orig_content stores the full STAC JSON so we can extract fragment keys later. - -- It is NOT a column in items; we use an explicit column list on INSERT below. - n.content AS orig_content, - (content_dehydrate(n.content)).* - FROM newdata n; - GET DIAGNOSTICS nrows = ROW_COUNT; - RAISE NOTICE 'Added % rows to tmpdata. %', nrows, clock_timestamp() - ts; - - -- Batch fragment dedup: compute the configured fragment payload per row using - -- fragment_config from the collection row, insert unique fragments, then assign - -- fragment_id and strip the fragment-covered keys from per-item assets/properties. - RAISE NOTICE 'Batch inserting fragments. %', clock_timestamp() - ts; - INSERT INTO item_fragments (collection, hash, content) - SELECT DISTINCT ON (collection, pgstac_hash_fragment(fragment_content)) - collection, - pgstac_hash_fragment(fragment_content) AS hash, - fragment_content - FROM ( - SELECT - t.collection, - extract_fragment(t.orig_content, c.fragment_config) AS fragment_content - FROM tmpdata t - JOIN collections c ON c.id = t.collection - ) fragments - WHERE fragment_content IS NOT NULL AND fragment_content != '{}'::jsonb - ON CONFLICT (collection, hash) DO NOTHING; - - RAISE NOTICE 'Assigning fragment_id. %', clock_timestamp() - ts; - UPDATE tmpdata t - SET - fragment_id = f.id, - -- Strip the fragment-covered keys from per-item columns so items.assets/properties - -- only contain per-item-specific values; fragment provides the shared baseline. - stac_version = CASE - WHEN 'stac_version' = ANY(c.fragment_config) THEN NULL - ELSE t.stac_version - END, - stac_extensions = CASE - WHEN 'stac_extensions' = ANY(c.fragment_config) THEN '[]'::jsonb - ELSE t.stac_extensions - END, - assets = strip_fragment_col(t.assets, 'assets', c.fragment_config), - properties = strip_fragment_col(t.properties, 'properties', c.fragment_config) - FROM collections c, - item_fragments f - WHERE c.id = t.collection - AND f.collection = t.collection - AND c.fragment_config IS NOT NULL - AND f.hash = pgstac_hash_fragment(extract_fragment(t.orig_content, c.fragment_config)); - - -- Queue registry sampling per collection (async via run_or_queue so it does not - -- block the ingest transaction). One queued call per distinct collection in the batch. - PERFORM run_or_queue(format('SELECT update_field_registry_from_items(%L);', c)) - FROM (SELECT DISTINCT collection FROM tmpdata) AS cte(c); - - -- Explicit column list excludes the orig_content extra column we added to tmpdata. + batch := ARRAY(SELECT content FROM newdata); + RAISE NOTICE 'Doing the insert. %', clock_timestamp() - ts; IF TG_TABLE_NAME = 'items_staging' THEN - INSERT INTO items (id, geometry, collection, datetime, end_datetime, pgstac_updated_at, - datetime_is_range, stac_version, stac_extensions, content_hash, fragment_id, bbox, links, assets, properties, extra, - created, updated, platform, instruments, constellation, mission, - eo_cloud_cover, eo_bands, eo_snow_cover, gsd, - proj_epsg, proj_wkt2, proj_projjson, proj_bbox, proj_centroid, proj_shape, proj_transform, - sci_doi, sci_citation, sci_publications, - view_off_nadir, view_incidence_angle, view_azimuth, view_sun_azimuth, view_sun_elevation, - file_size, file_header_size, file_checksum, file_byte_order, file_values_regex, - sat_orbit_state, sat_relative_orbit, sat_absolute_orbit) - SELECT id, geometry, collection, datetime, end_datetime, pgstac_updated_at, - datetime_is_range, stac_version, stac_extensions, content_hash, fragment_id, bbox, links, assets, properties, extra, - created, updated, platform, instruments, constellation, mission, - eo_cloud_cover, eo_bands, eo_snow_cover, gsd, - proj_epsg, proj_wkt2, proj_projjson, proj_bbox, proj_centroid, proj_shape, proj_transform, - sci_doi, sci_citation, sci_publications, - view_off_nadir, view_incidence_angle, view_azimuth, view_sun_azimuth, view_sun_elevation, - file_size, file_header_size, file_checksum, file_byte_order, file_values_regex, - sat_orbit_state, sat_relative_orbit, sat_absolute_orbit - FROM tmpdata; + INSERT INTO items SELECT * FROM items_staging_dehydrate(batch); GET DIAGNOSTICS nrows = ROW_COUNT; RAISE NOTICE 'Inserted % rows to items. %', nrows, clock_timestamp() - ts; ELSIF TG_TABLE_NAME = 'items_staging_ignore' THEN - INSERT INTO items (id, geometry, collection, datetime, end_datetime, pgstac_updated_at, - datetime_is_range, stac_version, stac_extensions, content_hash, fragment_id, bbox, links, assets, properties, extra, - created, updated, platform, instruments, constellation, mission, - eo_cloud_cover, eo_bands, eo_snow_cover, gsd, - proj_epsg, proj_wkt2, proj_projjson, proj_bbox, proj_centroid, proj_shape, proj_transform, - sci_doi, sci_citation, sci_publications, - view_off_nadir, view_incidence_angle, view_azimuth, view_sun_azimuth, view_sun_elevation, - file_size, file_header_size, file_checksum, file_byte_order, file_values_regex, - sat_orbit_state, sat_relative_orbit, sat_absolute_orbit) - SELECT id, geometry, collection, datetime, end_datetime, pgstac_updated_at, - datetime_is_range, stac_version, stac_extensions, content_hash, fragment_id, bbox, links, assets, properties, extra, - created, updated, platform, instruments, constellation, mission, - eo_cloud_cover, eo_bands, eo_snow_cover, gsd, - proj_epsg, proj_wkt2, proj_projjson, proj_bbox, proj_centroid, proj_shape, proj_transform, - sci_doi, sci_citation, sci_publications, - view_off_nadir, view_incidence_angle, view_azimuth, view_sun_azimuth, view_sun_elevation, - file_size, file_header_size, file_checksum, file_byte_order, file_values_regex, - sat_orbit_state, sat_relative_orbit, sat_absolute_orbit - FROM tmpdata ON CONFLICT DO NOTHING; + INSERT INTO items SELECT * FROM items_staging_dehydrate(batch) + ON CONFLICT DO NOTHING; GET DIAGNOSTICS nrows = ROW_COUNT; RAISE NOTICE 'Inserted % rows to items. %', nrows, clock_timestamp() - ts; ELSIF TG_TABLE_NAME = 'items_staging_upsert' THEN + -- Delete existing rows whose stored content actually changed, then insert. EXECUTE format( $sql$ - DELETE FROM items i USING tmpdata s + DELETE FROM items i USING ( + SELECT d.* + FROM newdata s + CROSS JOIN LATERAL content_dehydrate(s.content) d + ) s WHERE i.id = s.id AND i.collection = s.collection @@ -2539,85 +7001,39 @@ BEGIN ); GET DIAGNOSTICS nrows = ROW_COUNT; RAISE NOTICE 'Deleted % rows from items. %', nrows, clock_timestamp() - ts; - INSERT INTO items (id, geometry, collection, datetime, end_datetime, pgstac_updated_at, - datetime_is_range, stac_version, stac_extensions, content_hash, fragment_id, bbox, links, assets, properties, extra, - created, updated, platform, instruments, constellation, mission, - eo_cloud_cover, eo_bands, eo_snow_cover, gsd, - proj_epsg, proj_wkt2, proj_projjson, proj_bbox, proj_centroid, proj_shape, proj_transform, - sci_doi, sci_citation, sci_publications, - view_off_nadir, view_incidence_angle, view_azimuth, view_sun_azimuth, view_sun_elevation, - file_size, file_header_size, file_checksum, file_byte_order, file_values_regex, - sat_orbit_state, sat_relative_orbit, sat_absolute_orbit) - SELECT id, geometry, collection, datetime, end_datetime, pgstac_updated_at, - datetime_is_range, stac_version, stac_extensions, content_hash, fragment_id, bbox, links, assets, properties, extra, - created, updated, platform, instruments, constellation, mission, - eo_cloud_cover, eo_bands, eo_snow_cover, gsd, - proj_epsg, proj_wkt2, proj_projjson, proj_bbox, proj_centroid, proj_shape, proj_transform, - sci_doi, sci_citation, sci_publications, - view_off_nadir, view_incidence_angle, view_azimuth, view_sun_azimuth, view_sun_elevation, - file_size, file_header_size, file_checksum, file_byte_order, file_values_regex, - sat_orbit_state, sat_relative_orbit, sat_absolute_orbit - FROM tmpdata ON CONFLICT DO NOTHING; + INSERT INTO items SELECT * FROM items_staging_dehydrate(batch) + ON CONFLICT DO NOTHING; GET DIAGNOSTICS nrows = ROW_COUNT; RAISE NOTICE 'Inserted % rows to items. %', nrows, clock_timestamp() - ts; END IF; + -- Bump each partition's row-count estimate so search stepping sees the new rows. n drives + -- partition_bounds' histogram and the Rust keyset search_page skips zero-count bands, so a partition + -- left at n=0 would hide its freshly-ingested items from that path (SQL search() scans regardless). + -- Mirrors flush_items_staging_binary on the Rust loader. Over-count is safe (ignore/upsert may insert + -- fewer than staged); the async tightener resets n exactly. check_partition above already widened the + -- temporal envelope; spatial is left to the tightener on this fallback path. (n only affects search + -- stepping performance, never which rows/order come back — see geometrysearch band direction.) + UPDATE partition_stats ps + SET n = COALESCE(ps.n, 0) + agg.c, dirty = true, last_updated = now() + FROM ( + SELECT (partition_name(nd.content->>'collection', + lower(stac_daterange(nd.content->'properties')))).partition_name AS partition, + count(*) AS c + FROM newdata nd + GROUP BY 1 + ) agg + WHERE ps.partition = agg.partition; + RAISE NOTICE 'Deleting data from staging table. %', clock_timestamp() - ts; - -- Use TG_TABLE_NAME so the correct staging table is cleared. - -- The previous hard-coded 'DELETE FROM items_staging' was a bug that left - -- items_staging_ignore and items_staging_upsert un-cleared after processing. EXECUTE format('DELETE FROM %I', TG_TABLE_NAME); RAISE NOTICE 'Done. %', clock_timestamp() - ts; RETURN NULL; - END; $function$ ; -CREATE OR REPLACE FUNCTION pgstac.merge_jsonb(_a jsonb, _b jsonb) - RETURNS jsonb - LANGUAGE sql - IMMUTABLE -AS $function$ - SELECT - CASE - WHEN _a = '"𒍟※"'::jsonb THEN NULL - WHEN _a IS NULL THEN _b - WHEN jsonb_typeof(_a) = 'null' THEN 'null'::jsonb - WHEN jsonb_typeof(_a) = 'object' AND jsonb_typeof(_b) = 'object' THEN - ( - SELECT coalesce(jsonb_object_agg(sub.key, sub.val), '{}'::jsonb) - FROM ( - SELECT key, merge_jsonb(a.value, b.value) AS val - FROM - jsonb_each(coalesce(_a,'{}'::jsonb)) as a - FULL JOIN - jsonb_each(coalesce(_b,'{}'::jsonb)) as b - USING (key) - ) sub - WHERE sub.val IS NOT NULL - ) - WHEN - jsonb_typeof(_a) = 'array' - AND jsonb_typeof(_b) = 'array' - AND jsonb_array_length(_a) = jsonb_array_length(_b) - THEN - ( - SELECT jsonb_agg(m) FROM - ( SELECT - merge_jsonb( - jsonb_array_elements(_a), - jsonb_array_elements(_b) - ) as m - ) as l - ) - ELSE _a - END - ; -$function$ -; - CREATE OR REPLACE FUNCTION pgstac.missing_queryables(_collection text, _tablesample double precision DEFAULT 5, minrows double precision DEFAULT 10) RETURNS TABLE(collection text, name text, definition jsonb, property_wrapper text) LANGUAGE plpgsql @@ -2687,24 +7103,78 @@ END; $function$ ; -CREATE OR REPLACE FUNCTION pgstac.partition_after_triggerfunc() - RETURNS trigger +CREATE OR REPLACE FUNCTION pgstac.q_to_tsquery(jinput jsonb) + RETURNS tsquery LANGUAGE plpgsql - SECURITY DEFINER AS $function$ DECLARE - p text; - t timestamptz := clock_timestamp(); + input text; + processed_text text; + temp_text text; + quote_array text[]; + placeholder text := '@QUOTE@'; BEGIN - RAISE NOTICE 'Updating partition stats %', t; - FOR p IN SELECT DISTINCT partition - FROM newdata n JOIN partition_sys_meta p - ON (n.collection=p.collection AND n.datetime <@ p.partition_dtrange) + IF jsonb_typeof(jinput) = 'string' THEN + input := jinput->>0; + ELSIF jsonb_typeof(jinput) = 'array' THEN + input := array_to_string( + array(select jsonb_array_elements_text(jinput)), + ' OR ' + ); + ELSE + RAISE EXCEPTION 'Input must be a string or an array of strings.'; + END IF; + -- Extract all quoted phrases and store in array + quote_array := regexp_matches(input, '"[^"]*"', 'g'); + + -- Replace each quoted part with a unique placeholder if there are any quoted phrases + IF array_length(quote_array, 1) IS NOT NULL THEN + processed_text := input; + FOR i IN array_lower(quote_array, 1) .. array_upper(quote_array, 1) LOOP + processed_text := replace(processed_text, quote_array[i], placeholder || i || placeholder); + END LOOP; + ELSE + processed_text := input; + END IF; + + -- Replace non-quoted text using regular expressions + + -- , -> | + processed_text := regexp_replace(processed_text, ',(?=(?:[^"]*"[^"]*")*[^"]*$)', ' | ', 'g'); + + -- and -> & + processed_text := regexp_replace(processed_text, '\s+AND\s+', ' & ', 'gi'); + + -- or -> | + processed_text := regexp_replace(processed_text, '\s+OR\s+', ' | ', 'gi'); + + -- + -> + processed_text := regexp_replace(processed_text, '^\s*\+([a-zA-Z0-9_]+)', '\1', 'g'); -- +term at start + processed_text := regexp_replace(processed_text, '\s*\+([a-zA-Z0-9_]+)', ' & \1', 'g'); -- +term elsewhere + + -- - -> ! + processed_text := regexp_replace(processed_text, '^\s*\-([a-zA-Z0-9_]+)', '! \1', 'g'); -- -term at start + processed_text := regexp_replace(processed_text, '\s*\-([a-zA-Z0-9_]+)', ' & ! \1', 'g'); -- -term elsewhere + + -- terms separated with spaces are assumed to represent adjacent terms. loop through these + -- occurrences and replace them with the adjacency operator (<->) LOOP - PERFORM run_or_queue(format('SELECT update_partition_stats(%L, %L);', p, true)); + temp_text := regexp_replace(processed_text, '([a-zA-Z0-9_]+)\s+([a-zA-Z0-9_]+)(?!\s*[&|<>])', '\1 <-> \2', 'g'); + IF temp_text = processed_text THEN + EXIT; -- No more replacements were made + END IF; + processed_text := temp_text; END LOOP; - RAISE NOTICE 't: % %', t, clock_timestamp() - t; - RETURN NULL; + + + -- Replace placeholders back with quoted phrases if there were any + IF array_length(quote_array, 1) IS NOT NULL THEN + FOR i IN array_lower(quote_array, 1) .. array_upper(quote_array, 1) LOOP + processed_text := replace(processed_text, placeholder || i || placeholder, '''' || substring(quote_array[i] from 2 for length(quote_array[i]) - 2) || ''''); + END LOOP; + END IF; + + RETURN to_tsquery('english', processed_text); END; $function$ ; @@ -2849,177 +7319,67 @@ WITH p AS ( $function$ ; +CREATE OR REPLACE FUNCTION pgstac.queryables_trigger_func() + RETURNS trigger + LANGUAGE plpgsql + SECURITY DEFINER +AS $function$ +DECLARE +BEGIN + -- Queryable definitions changed, so every partition's queryable indexes may be stale. Flag them and let + -- the async sweep (build_pending_indexes) (re)build off the hot path, instead of rebuilding every + -- partition synchronously inside this trigger. + UPDATE pgstac.partition_stats SET indexes_pending = true; + RETURN NULL; +END; +$function$ +; + CREATE OR REPLACE FUNCTION pgstac.search(_search jsonb DEFAULT '{}'::jsonb) - RETURNS jsonb + RETURNS json LANGUAGE plpgsql AS $function$ DECLARE - searches searches%ROWTYPE; - _where text; - orderby text; - search_where searches%ROWTYPE; - total_count bigint; - token record; - token_prev boolean; - token_item items%ROWTYPE; - token_where text; - full_where text; - init_ts timestamptz := clock_timestamp(); - timer timestamptz := clock_timestamp(); - prev text; - next text; - collection jsonb; - out_records jsonb; - out_len int; - _limit int := coalesce((_search->>'limit')::int, 10); - _querylimit int; - _fields jsonb := coalesce(_search->'fields', '{}'::jsonb); - has_prev boolean := FALSE; - has_next boolean := FALSE; + -- caller-provided limit/token come from the search body; default to the configured page size. + _limit int := coalesce((_search->>'limit')::int, + nullif(get_setting('default_page_size', _search->'conf'), '')::int, 10); + _token text := _search->>'token'; + -- The keyset is the token minus its next/prev prefix; an empty token means no keyset (first + -- page). A non-empty keyset that does not decode raises in keyset_decode downstream. + keyset text := nullif(regexp_replace(coalesce(_token,''), '^(next|prev):', ''), ''); + is_prev boolean := (_token LIKE 'prev:%') AND keyset IS NOT NULL; + pg record; + burl text := rtrim(coalesce(base_url(_search->'conf'), ''), '/'); links jsonb := '[]'::jsonb; - base_url text:= concat(rtrim(base_url(_search->'conf'),'/')); -BEGIN - searches := search_query(_search); - _where := searches._where; - orderby := searches.orderby; - search_where := where_stats(searches.hash, _where, false, _search->'conf'); - total_count := search_where.context_count; - RAISE NOTICE 'SEARCH:TOKEN: %', _search->>'token'; - token := get_token_record(_search->>'token'); - RAISE NOTICE '***TOKEN: %', token; - _querylimit := _limit + 1; - IF token IS NOT NULL THEN - token_prev := token.prev; - token_item := token.item; - token_where := get_token_filter(_search->'sortby', token_item, token_prev, FALSE); - RAISE DEBUG 'TOKEN_WHERE: % (%ms from search start)', token_where, age_ms(timer); - IF token_prev THEN -- if we are using a prev token, we know has_next is true - RAISE DEBUG 'There is a previous token, so automatically setting has_next to true'; - has_next := TRUE; - orderby := sort_sqlorderby(_search, TRUE); - ELSE - RAISE DEBUG 'There is a next token, so automatically setting has_prev to true'; - has_prev := TRUE; - - END IF; - ELSE -- if there was no token, we know there is no prev - RAISE DEBUG 'There is no token, so we know there is no prev. setting has_prev to false'; - has_prev := FALSE; - END IF; - - full_where := concat_ws(' AND ', _where, token_where); - RAISE NOTICE 'FULL WHERE CLAUSE: %', full_where; - RAISE NOTICE 'Time to get counts and build query %', age_ms(timer); - timer := clock_timestamp(); - - RAISE NOTICE 'Getting hydrated data.'; - RAISE NOTICE 'Time to set hydration/formatting %', age_ms(timer); - timer := clock_timestamp(); - SELECT jsonb_agg(format_item(i, _fields)) INTO out_records - FROM search_rows( - full_where, - orderby, - _querylimit - ) as i; - - RAISE NOTICE 'Time to fetch rows %', age_ms(timer); - timer := clock_timestamp(); - - - IF token_prev THEN - out_records := flip_jsonb_array(out_records); - END IF; - - RAISE NOTICE 'Query returned % records.', jsonb_array_length(out_records); - RAISE DEBUG 'TOKEN: % %', token_item.id, token_item.collection; - RAISE DEBUG 'RECORD_1: % %', out_records->0->>'id', out_records->0->>'collection'; - RAISE DEBUG 'RECORD-1: % %', out_records->-1->>'id', out_records->-1->>'collection'; - - -- REMOVE records that were from our token - IF out_records->0->>'id' = token_item.id AND out_records->0->>'collection' = token_item.collection THEN - out_records := out_records - 0; - ELSIF out_records->-1->>'id' = token_item.id AND out_records->-1->>'collection' = token_item.collection THEN - out_records := out_records - -1; - END IF; - - out_len := jsonb_array_length(out_records); - - IF out_len = _limit + 1 THEN - IF token_prev THEN - has_prev := TRUE; - out_records := out_records - 0; - ELSE - has_next := TRUE; - out_records := out_records - -1; - END IF; - END IF; - - - links := links || jsonb_build_object( - 'rel', 'root', - 'type', 'application/json', - 'href', base_url - ) || jsonb_build_object( - 'rel', 'self', - 'type', 'application/json', - 'href', concat(base_url, '/search') - ); - - IF has_next THEN - next := concat(out_records->-1->>'collection', ':', out_records->-1->>'id'); - RAISE NOTICE 'HAS NEXT | %', next; - links := links || jsonb_build_object( - 'rel', 'next', - 'type', 'application/geo+json', - 'method', 'GET', - 'href', concat(base_url, '/search?token=next:', next) - ); + out json; +BEGIN + SELECT * INTO pg FROM search_page(_search, _limit, keyset, is_prev); + links := links + || jsonb_build_object('rel','root','type','application/json','href', burl) + || jsonb_build_object('rel','self','type','application/json','href',burl||'/search'); + IF pg.next_token IS NOT NULL THEN + links := links || jsonb_build_object('rel','next','type','application/geo+json','method','GET', + 'href', burl||'/search?token=next:'||pg.next_token); END IF; - - IF has_prev THEN - prev := concat(out_records->0->>'collection', ':', out_records->0->>'id'); - RAISE NOTICE 'HAS PREV | %', prev; - links := links || jsonb_build_object( - 'rel', 'prev', - 'type', 'application/geo+json', - 'method', 'GET', - 'href', concat(base_url, '/search?token=prev:', prev) - ); + IF pg.prev_token IS NOT NULL THEN + links := links || jsonb_build_object('rel','prev','type','application/geo+json','method','GET', + 'href', burl||'/search?token=prev:'||pg.prev_token); END IF; - - RAISE NOTICE 'Time to get prev/next %', age_ms(timer); - timer := clock_timestamp(); - - - collection := jsonb_build_object( - 'type', 'FeatureCollection', - 'features', coalesce(out_records, '[]'::jsonb), - 'links', links - ); - - - - IF context(_search->'conf') != 'off' THEN - collection := collection || jsonb_strip_nulls(jsonb_build_object( - 'numberMatched', total_count, - 'numberReturned', coalesce(jsonb_array_length(out_records), 0) - )); + IF pg.number_matched IS NOT NULL THEN + out := json_build_object( + 'type','FeatureCollection', + 'features', coalesce(pg.features,'[]'::json), + 'links', links, + 'numberReturned', pg.number_returned, + 'numberMatched', pg.number_matched); ELSE - collection := collection || jsonb_strip_nulls(jsonb_build_object( - 'numberReturned', coalesce(jsonb_array_length(out_records), 0) - )); + out := json_build_object( + 'type','FeatureCollection', + 'features', coalesce(pg.features,'[]'::json), + 'links', links, + 'numberReturned', pg.number_returned); END IF; - - IF get_setting_bool('timing', _search->'conf') THEN - collection = collection || jsonb_build_object('timing', age_ms(init_ts)); - END IF; - - RAISE NOTICE 'Time to build final json %', age_ms(timer); - timer := clock_timestamp(); - - RAISE NOTICE 'Total Time: %', age_ms(current_timestamp); - RAISE NOTICE 'RETURNING % records. NEXT: %. PREV: %', collection->>'numberReturned', collection->>'next', collection->>'prev'; - RETURN collection; + RETURN out; END; $function$ ; @@ -3033,238 +7393,33 @@ AS $function$ $function$ ; -CREATE OR REPLACE FUNCTION pgstac.search_query(_search jsonb DEFAULT '{}'::jsonb, updatestats boolean DEFAULT false, _metadata jsonb DEFAULT '{}'::jsonb) - RETURNS searches - LANGUAGE plpgsql - SECURITY DEFINER -AS $function$ -DECLARE - search searches%ROWTYPE; - cached_search searches%ROWTYPE; - search_where searches%ROWTYPE; - ro boolean := pgstac.readonly(); -BEGIN - RAISE NOTICE 'SEARCH: %', _search; - -- Calculate hash, where clause, and order by statement - search.search := _search; - search.metadata := _metadata; - search._where := stac_search_to_where(_search); - search.hash := search_hash_from_where(search._where, search.metadata); - search.orderby := sort_sqlorderby(_search); - search.lastused := now(); - search.usecount := 1; - - -- If we are in read only mode, directly return search - IF ro THEN - RETURN search; - END IF; - - -- Cache bookkeeping is best-effort and non-blocking. We always return - -- canonical hash + where, even if cache touch cannot be acquired quickly. - UPDATE searches - SET - lastused = now(), - usecount = searches.usecount + 1 - WHERE ctid = ( - SELECT ctid - FROM searches - WHERE hash = search.hash - FOR UPDATE SKIP LOCKED - LIMIT 1 - ) - RETURNING * INTO cached_search; - - IF cached_search IS NULL THEN - IF pg_try_advisory_xact_lock(hashtext(search.hash)) THEN - INSERT INTO searches (hash, search, _where, orderby, lastused, usecount, metadata) - VALUES (search.hash, search.search, search._where, search.orderby, now(), 1, search.metadata) - ON CONFLICT (hash) DO UPDATE SET - lastused = EXCLUDED.lastused, - usecount = searches.usecount + 1 - RETURNING * INTO cached_search; - END IF; - - IF cached_search IS NULL THEN - SELECT * INTO cached_search FROM searches WHERE hash = search.hash; - END IF; - END IF; - - IF cached_search IS NOT NULL THEN - cached_search._where = search._where; - cached_search.orderby = search.orderby; - IF updatestats THEN - search_where := where_stats( - cached_search.hash, - cached_search._where, - true, - _search->'conf' - ); - cached_search.context_count := search_where.context_count; - cached_search.statslastupdated := search_where.statslastupdated; - END IF; - RETURN cached_search; - END IF; - - IF updatestats THEN - search_where := where_stats( - search.hash, - search._where, - true, - _search->'conf' - ); - search.context_count := search_where.context_count; - search.statslastupdated := search_where.statslastupdated; - END IF; - - RETURN search; - -END; -$function$ -; - CREATE OR REPLACE FUNCTION pgstac.stac_search_to_where(j jsonb) RETURNS text LANGUAGE plpgsql STABLE AS $function$ -DECLARE - where_segments text[]; - _where text; - dtrange tstzrange; - collections text[]; - geom geometry; - sdate timestamptz; - edate timestamptz; - filterlang text; - filter jsonb := j->'filter'; - ft_query tsquery; +DECLARE _where text := cql2_query(search_to_cql2(j)); BEGIN - IF j ? 'ids' THEN - where_segments := where_segments || format('id = ANY (%L) ', to_text_array(j->'ids')); - END IF; - - IF j ? 'collections' THEN - collections := to_text_array(j->'collections'); - where_segments := where_segments || format('collection = ANY (%L) ', collections); - END IF; - - IF j ? 'datetime' THEN - dtrange := parse_dtrange(j->'datetime'); - sdate := lower(dtrange); - edate := upper(dtrange); - - where_segments := where_segments || format(' datetime <= %L::timestamptz AND end_datetime >= %L::timestamptz ', - edate, - sdate - ); - END IF; - - IF j ? 'q' THEN - ft_query := q_to_tsquery(j->'q'); - where_segments := where_segments || format( - $quote$ - ( - -- Use the split properties column directly (v0.10 schema). - -- Previously read from content->'properties'->>'description' etc. - to_tsvector('english', properties->>'description') || - to_tsvector('english', coalesce(properties->>'title', '')) || - to_tsvector('english', coalesce(properties->>'keywords', '')) - ) @@ %L - $quote$, - ft_query - ); - END IF; - - geom := stac_geom(j); - IF geom IS NOT NULL THEN - where_segments := where_segments || format('st_intersects(geometry, %L)',geom); - END IF; - - filterlang := COALESCE( - j->>'filter-lang', - get_setting('default_filter_lang', j->'conf') - ); - IF NOT filter @? '$.**.op' THEN - filterlang := 'cql-json'; - END IF; - - IF filterlang NOT IN ('cql-json','cql2-json') AND j ? 'filter' THEN - RAISE EXCEPTION '% is not a supported filter-lang. Please use cql-json or cql2-json.', filterlang; - END IF; - - IF j ? 'query' AND j ? 'filter' THEN - RAISE EXCEPTION 'Can only use either query or filter at one time.'; - END IF; - - IF j ? 'query' THEN - filter := query_to_cql2(j->'query'); - ELSIF filterlang = 'cql-json' THEN - filter := cql1_to_cql2(filter); - END IF; - RAISE NOTICE 'FILTER: %', filter; - where_segments := where_segments || cql2_query(filter); - IF cardinality(where_segments) < 1 THEN - RETURN ' TRUE '; - END IF; - - _where := array_to_string(array_remove(where_segments, NULL), ' AND '); - - IF _where IS NULL OR BTRIM(_where) = '' THEN - RETURN ' TRUE '; - END IF; + IF _where IS NULL OR btrim(_where) = '' THEN RETURN ' TRUE '; END IF; RETURN _where; - END; $function$ ; -CREATE OR REPLACE FUNCTION pgstac.strip_jsonb(_a jsonb, _b jsonb) - RETURNS jsonb +CREATE OR REPLACE FUNCTION pgstac.xyzsearch(_x integer, _y integer, _z integer, queryhash text, fields jsonb DEFAULT NULL::jsonb, _scanlimit integer DEFAULT 10000, _limit integer DEFAULT 100, _timelimit interval DEFAULT '00:00:05'::interval, exitwhenfull boolean DEFAULT true, skipcovered boolean DEFAULT true) + RETURNS json LANGUAGE sql - IMMUTABLE AS $function$ - -- strip_jsonb: RETAINED FOR USE BY MIGRATION SCRIPTS ONLY. - -- Must not be called from any ingest, hydrate, or search code path after v0.10. - -- Will be removed once the migration path has been finalized and tested. - SELECT - CASE - - WHEN _a IS NULL AND _b IS NOT NULL AND jsonb_typeof(_b) != 'null' THEN '"𒍟※"'::jsonb - WHEN _a IS NULL THEN NULL - WHEN _a = _b AND jsonb_typeof(_a) = 'object' THEN '{}'::jsonb - WHEN _a = _b THEN NULL - WHEN jsonb_typeof(_a) = 'null' THEN 'null'::jsonb - WHEN _b IS NULL THEN _a - WHEN jsonb_typeof(_a) = 'object' AND jsonb_typeof(_b) = 'object' THEN - ( - SELECT coalesce(jsonb_object_agg(sub.key, sub.val), '{}'::jsonb) - FROM ( - SELECT key, strip_jsonb(a.value, b.value) AS val - FROM - jsonb_each(_a) as a - FULL JOIN - jsonb_each(_b) as b - USING (key) - ) sub - WHERE sub.val IS NOT NULL - ) - WHEN - jsonb_typeof(_a) = 'array' - AND jsonb_typeof(_b) = 'array' - AND jsonb_array_length(_a) = jsonb_array_length(_b) - THEN - ( - SELECT jsonb_agg(m) FROM - ( SELECT - strip_jsonb( - jsonb_array_elements(_a), - jsonb_array_elements(_b) - ) as m - ) as l - ) - ELSE _a - END - ; + SELECT * FROM geometrysearch( + st_transform(tileenvelope(_z, _x, _y), 4326), + queryhash, + fields, + _scanlimit, + _limit, + _timelimit, + exitwhenfull, + skipcovered + ); $function$ ; DO $$ @@ -3294,7 +7449,7 @@ DO $$ END $$; --- Register promoted native-column queryables (v0.10 split schema). +-- Register promoted native-column queryables. -- Each entry maps a STAC property name to the promoted items column via property_path. -- CQL2 queries and auto-created indexes will use the native column, not JSONB extraction. -- The seed data lives in promoted_queryables_defaults() (002a_queryables.sql) so it @@ -3308,15 +7463,15 @@ WHERE NOT EXISTS ( SELECT 1 FROM queryables q WHERE q.name = p.name ); --- Pass 2: backfill property_path/property_wrapper on rows that were created by an --- older install (pre-v0.10) and therefore have property_path=NULL. +-- Pass 2: backfill property_path on older rows and normalize promoted wrappers +-- to the defaults (NULL for native promoted columns). UPDATE queryables q SET property_path = CASE WHEN q.property_index_type IS NULL THEN COALESCE(q.property_path, p.property_path) ELSE q.property_path END, property_wrapper = CASE - WHEN q.property_index_type IS NULL THEN COALESCE(q.property_wrapper, p.property_wrapper) + WHEN q.property_index_type IS NULL THEN p.property_wrapper ELSE q.property_wrapper END, definition = COALESCE(q.definition, p.definition) @@ -3332,7 +7487,6 @@ INSERT INTO pgstac_settings (name, value) VALUES ('context_estimated_count', '100000'), ('context_estimated_cost', '100000'), ('context_stats_ttl', '1 day'), - ('search_gc_retention_interval', '7 days'), ('default_filter_lang', 'cql2-json'), ('additional_properties', 'true'), ('use_queue', 'false'), @@ -3386,62 +7540,33 @@ ALTER FUNCTION to_int COST 5000; ALTER FUNCTION to_tstz COST 5000; ALTER FUNCTION to_text_array COST 5000; -ALTER FUNCTION update_partition_stats SECURITY DEFINER; -ALTER FUNCTION partition_after_triggerfunc SECURITY DEFINER; -ALTER FUNCTION drop_table_constraints SECURITY DEFINER; -ALTER FUNCTION create_table_constraints SECURITY DEFINER; -ALTER FUNCTION check_partition SECURITY DEFINER; -ALTER FUNCTION repartition SECURITY DEFINER; -ALTER FUNCTION where_stats(text, text, boolean, jsonb) SECURITY DEFINER; -ALTER FUNCTION search_query SECURITY DEFINER; -ALTER FUNCTION name_search SECURITY DEFINER; -ALTER FUNCTION rename_search SECURITY DEFINER; -ALTER FUNCTION unname_search SECURITY DEFINER; -ALTER FUNCTION pin_search SECURITY DEFINER; -ALTER FUNCTION unpin_search SECURITY DEFINER; -ALTER FUNCTION gc_anonymous_searches(interval, jsonb) SECURITY DEFINER; -ALTER FUNCTION gc_search_caches(interval, jsonb) SECURITY DEFINER; -ALTER FUNCTION gc_deleted_items_log_batch(interval, integer) SECURITY DEFINER; -ALTER FUNCTION gc_deleted_items_log(interval, integer) SECURITY DEFINER; -ALTER FUNCTION gc_deleted_items_log(interval) SECURITY DEFINER; -ALTER FUNCTION format_item SECURITY DEFINER; -ALTER FUNCTION maintain_index SECURITY DEFINER; -ALTER FUNCTION pgstac_item_hash(jsonb) SECURITY DEFINER; -ALTER FUNCTION promoted_items_column_list() SECURITY DEFINER; -ALTER FUNCTION items_content_distinct_sql(text, text) SECURITY DEFINER; -ALTER FUNCTION items_content_changed(items, items) SECURITY DEFINER; -ALTER FUNCTION items_touch_triggerfunc SECURITY DEFINER; -ALTER FUNCTION items_delete_log_trigger SECURITY DEFINER; -ALTER FUNCTION strip_promoted_properties(jsonb) SECURITY DEFINER; -ALTER FUNCTION tstz_to_stac_text(timestamptz) SECURITY DEFINER; -ALTER FUNCTION temporal_properties_from_item(items) SECURITY DEFINER; -ALTER FUNCTION promoted_properties_from_item(items) SECURITY DEFINER; -ALTER FUNCTION extract_fragment(jsonb, text[]) SECURITY DEFINER; -ALTER FUNCTION pgstac_hash_fragment(jsonb) SECURITY DEFINER; -ALTER FUNCTION gc_fragments(text, interval) SECURITY DEFINER; -ALTER FUNCTION strip_fragment_col(jsonb, text, text[]) SECURITY DEFINER; -ALTER FUNCTION update_field_registry_from_sample(text, jsonb[]) SECURITY DEFINER; -ALTER FUNCTION update_field_registry_from_items(text) SECURITY DEFINER; -ALTER FUNCTION refresh_field_registry(text, interval) SECURITY DEFINER; -ALTER FUNCTION collection_fragment_config_default(jsonb) SECURITY DEFINER; -ALTER FUNCTION jsonb_leaf_rows(jsonb, text) SECURITY DEFINER; -ALTER FUNCTION jsonb_common_values(jsonb, jsonb) SECURITY DEFINER; -ALTER FUNCTION jsonb_merge_level1(jsonb, jsonb) SECURITY DEFINER; -ALTER FUNCTION fragment_path_text(text[]) SECURITY DEFINER; -ALTER FUNCTION fragment_path_array(text) SECURITY DEFINER; +-- SECURITY DEFINER is declared INLINE in each function's CREATE (the single source of truth), +-- not re-applied here. Functions that create partitions/indexes/constraints declare it inline so +-- the created objects are owned by pgstac_admin; functions that write the search cache from the +-- read path declare it inline too. Pure helpers stay SECURITY INVOKER. Keeping a separate ALTER +-- list here only let it drift from the definitions (stale/duplicate/wrong-signature entries). -GRANT USAGE ON SCHEMA pgstac to pgstac_read; -GRANT ALL ON SCHEMA pgstac to pgstac_ingest; -GRANT ALL ON SCHEMA pgstac to pgstac_admin; +-- Schema USAGE for pgstac_read / pgstac_ingest is granted in 000_idempotent_pre.sql; pgstac_admin +-- owns the schema. Not re-granted here. --- pgstac_read role limited to using function apis +-- pgstac_read API surface. Functions are EXECUTE-able by PUBLIC by default, so these grants are not +-- required for access today; they document the intended top-level read API (and would be the point to +-- enforce from if EXECUTE were ever revoked from PUBLIC). Internal helpers (keyset_*, partition_bounds, +-- cql2_*, next_band, ...) are deliberately NOT listed — read reaches them only inside these entry points. GRANT EXECUTE ON FUNCTION search TO pgstac_read; GRANT EXECUTE ON FUNCTION search_query TO pgstac_read; GRANT EXECUTE ON FUNCTION item_by_id TO pgstac_read; GRANT EXECUTE ON FUNCTION get_item TO pgstac_read; -GRANT EXECUTE ON FUNCTION format_item TO pgstac_read; GRANT EXECUTE ON FUNCTION content_hydrate TO pgstac_read; -GRANT EXECUTE ON FUNCTION pgstac_item_hash TO pgstac_read; +GRANT EXECUTE ON FUNCTION search_page TO pgstac_read; +GRANT EXECUTE ON FUNCTION search_plan TO pgstac_read; +GRANT EXECUTE ON FUNCTION collection_search_plan TO pgstac_read; +GRANT EXECUTE ON FUNCTION collection_search TO pgstac_read; +GRANT EXECUTE ON FUNCTION geometrysearch TO pgstac_read; +GRANT EXECUTE ON FUNCTION geojsonsearch TO pgstac_read; +GRANT EXECUTE ON FUNCTION xyzsearch TO pgstac_read; +GRANT EXECUTE ON FUNCTION search_from_json(jsonb, jsonb) TO pgstac_read; +-- Tables are NOT readable by PUBLIC; read needs an explicit SELECT grant. GRANT SELECT ON ALL TABLES IN SCHEMA pgstac TO pgstac_read; @@ -3449,6 +7574,21 @@ GRANT EXECUTE ON ALL FUNCTIONS IN SCHEMA pgstac to pgstac_ingest; GRANT ALL ON ALL TABLES IN SCHEMA pgstac to pgstac_ingest; GRANT USAGE ON ALL SEQUENCES IN SCHEMA pgstac to pgstac_ingest; +-- Privilege wall (INV-1): the connecting role (inheriting pgstac_ingest) may READ these tables but may +-- only MUTATE them through the SECURITY DEFINER write functions (check_partition, widen/tighten_partition_stats, +-- ensure_fragments, make_binary_staging, flush_items_staging_binary, the items_staging trigger, create/ +-- update/upsert/delete_item, the field-registry/fragment helpers). This makes an un-widened or +-- envelope-narrowing direct write structurally impossible. New partitions are created SELECT-only for +-- pgstac_ingest by check_partition; the items parent is revoked here. Staging tables (items_staging*) stay +-- writable so the SQL-only ingest path can COPY into them. +REVOKE INSERT, UPDATE, DELETE, TRUNCATE ON items, partition_stats, item_fragments FROM pgstac_ingest; +-- item_field_registry is the one exception to the wall: the loader WIDENS it directly with an add-only +-- INSERT ... ON CONFLICT DO UPDATE (no SD function), so INSERT + UPDATE stay granted. DELETE/TRUNCATE remain +-- revoked, so even a direct write cannot NARROW the registry — INV-1 (registry is a superset of the data) +-- holds structurally against narrowing, while the cheap widen stays on the hot load path. +REVOKE DELETE, TRUNCATE ON item_field_registry FROM pgstac_ingest; + + REVOKE ALL PRIVILEGES ON PROCEDURE run_queued_queries FROM public; GRANT ALL ON PROCEDURE run_queued_queries TO pgstac_admin; @@ -3460,6 +7600,7 @@ GRANT ALL ON PROCEDURE gc_deleted_items_log_committed(interval, integer) TO pgst RESET ROLE; -SET ROLE pgstac_ingest; -SELECT update_partition_stats_q(partition) FROM partitions_view; +-- (No install-time stats seeding: a fresh install has no partitions, and partition_stats rows are now +-- seeded by check_partition as partitions are created. Exact stats come from tighten_partition_stats via +-- the maintenance sweeps.) SELECT set_version('unreleased'); diff --git a/src/pgstac/migrations/pgstac--unreleased.sql b/src/pgstac/migrations/pgstac--unreleased.sql index f54715d1..77894b32 100644 --- a/src/pgstac/migrations/pgstac--unreleased.sql +++ b/src/pgstac/migrations/pgstac--unreleased.sql @@ -44,6 +44,7 @@ DO $$ $$; + GRANT pgstac_admin TO current_user; -- Function to make sure pgstac_admin is the owner of items @@ -101,10 +102,8 @@ $$ LANGUAGE PLPGSQL; SELECT pgstac_admin_owns(); CREATE SCHEMA IF NOT EXISTS pgstac AUTHORIZATION pgstac_admin; - -GRANT ALL ON ALL FUNCTIONS IN SCHEMA pgstac to pgstac_admin; -GRANT ALL ON ALL TABLES IN SCHEMA pgstac to pgstac_admin; -GRANT ALL ON ALL SEQUENCES IN SCHEMA pgstac to pgstac_admin; +-- pgstac_admin owns the schema and all objects in it (pgstac_admin_owns() above + +-- AUTHORIZATION), so it already has every privilege — no explicit GRANT ALL needed. ALTER ROLE pgstac_admin SET SEARCH_PATH TO pgstac, public; ALTER ROLE pgstac_read SET SEARCH_PATH TO pgstac, public; @@ -197,7 +196,12 @@ RETURNS timestamptz AS $$ ; $$ LANGUAGE SQL IMMUTABLE STRICT; --- v0.10 search: drop old function signatures +-- Drop objects superseded by the current partition_stats model. +DROP MATERIALIZED VIEW IF EXISTS partitions CASCADE; +DROP MATERIALIZED VIEW IF EXISTS partition_steps; +DROP VIEW IF EXISTS partition_steps; + +-- Drop function signatures whose argument lists changed (CREATE OR REPLACE cannot alter them) DROP FUNCTION IF EXISTS chunker(pred_envelope); DROP FUNCTION IF EXISTS search_bands(pred_envelope, boolean, integer, integer); DROP FUNCTION IF EXISTS search_rows(jsonb, integer, text, boolean); @@ -216,6 +220,7 @@ DROP FUNCTION IF EXISTS xyzsearch(integer, integer, integer, text, jsonb, intege DROP FUNCTION IF EXISTS search(jsonb); DROP FUNCTION IF EXISTS search_page(jsonb, integer, text, boolean); DROP FUNCTION IF EXISTS search_plan(jsonb, text); +DROP FUNCTION IF EXISTS fields_to_itemcols(jsonb); DROP FUNCTION IF EXISTS search_query(jsonb, boolean, jsonb); DROP FUNCTION IF EXISTS where_stats(text, text, boolean, jsonb); DROP FUNCTION IF EXISTS keyset_sortkeys(jsonb); @@ -315,10 +320,6 @@ CREATE OR REPLACE FUNCTION context_stats_ttl(conf jsonb DEFAULT NULL) RETURNS in SELECT pgstac.get_setting('context_stats_ttl', conf)::interval; $$ LANGUAGE SQL; -CREATE OR REPLACE FUNCTION search_gc_retention_interval(conf jsonb DEFAULT NULL) RETURNS interval AS $$ - SELECT pgstac.get_setting('search_gc_retention_interval', conf)::interval; -$$ LANGUAGE SQL; - CREATE OR REPLACE FUNCTION t2s(text) RETURNS text AS $$ SELECT extract(epoch FROM $1::interval)::text || ' s'; $$ LANGUAGE SQL IMMUTABLE PARALLEL SAFE STRICT; @@ -646,63 +647,51 @@ CREATE OR REPLACE FUNCTION explode_dotpaths_recurse(IN j jsonb) RETURNS SETOF te $$ LANGUAGE SQL IMMUTABLE PARALLEL SAFE; --- jsonb_canonical: RFC 8785 (JSON Canonicalization Scheme)-aligned serialization. --- Produces a deterministic, key-order-independent text encoding that an external --- client can reproduce byte-for-byte. NOTE: do NOT use `jsonb::text` for hashing — --- PostgreSQL re-normalizes object key order to length-then-bytewise and inserts --- ": " / ", " separators, so `jsonb::text` is neither alphabetical nor compact. --- --- Canonical rules (must match the external recipe below): --- * object keys sorted by Unicode code point (== UTF-8 byte order, COLLATE "C"), --- * compact separators: ',' between members, ':' between key and value, --- * strings: standard JSON escaping, NON-ASCII left as UTF-8 (no \uXXXX), --- * numbers: IEEE-754 double, shortest round-trip form (Ryu) — matches --- PostgreSQL float8 output and ECMAScript Number::toString for in-range --- values. (STAC numbers are physical quantities; integers beyond 2^53 are --- out of contract, as in RFC 8785.) --- * true / false / null as literals. +-- jsonb_canonical_hash: a deterministic 32-byte SHA-256 identity digest of a JSONB document, reproducible +-- outside PostgreSQL (pgstac-rs `canonical::jsonb_canonical_hash` produces identical bytes). -- --- External equivalents: --- Python: an RFC 8785 canonicalizer, or the rule-for-rule reference: --- def canon(v): --- if isinstance(v, bool): return 'true' if v else 'false' --- if v is None: return 'null' --- if isinstance(v, dict): --- return '{'+','.join(json.dumps(k,ensure_ascii=False)+':'+canon(v[k]) --- for k in sorted(v))+'}' --- if isinstance(v, list): return '['+','.join(canon(x) for x in v)+']' --- if isinstance(v,(int,float)): --- f=float(v); return str(int(f)) if f==int(f) and abs(f)<1e16 else repr(f) --- return json.dumps(v, ensure_ascii=False) --- Rust: the `rfc8785` crate (serde_jcs) over serde_json::Value. -CREATE OR REPLACE FUNCTION jsonb_canonical(j jsonb) RETURNS text AS $$ - SELECT CASE jsonb_typeof(j) - WHEN 'object' THEN COALESCE(( - SELECT '{' || string_agg( - to_json(kv.key)::text || ':' || jsonb_canonical(kv.value), - ',' ORDER BY kv.key COLLATE "C" - ) || '}' - FROM jsonb_each(j) kv - ), '{}') - WHEN 'array' THEN COALESCE(( - SELECT '[' || string_agg(jsonb_canonical(e.value), ',' ORDER BY e.ord) || ']' - FROM jsonb_array_elements(j) WITH ORDINALITY e(value, ord) - ), '[]') - WHEN 'number' THEN (j #>> '{}')::float8::text - ELSE j::text -- string (JSON-escaped, UTF-8 preserved), 'true' / 'false' / 'null' - END; -$$ LANGUAGE SQL IMMUTABLE PARALLEL SAFE STRICT; - --- jsonb_hash: raw 32-byte sha256 of the canonical (RFC 8785-aligned) JSON form. --- Returns bytea so callers store the compact binary digest directly (32 B vs --- 64-char hex). Use encode(jsonb_hash(j), 'hex') when a printable string is --- needed for display or external comparison. --- Externally reproducible: sha256(utf8_bytes(jsonb_canonical(j))). --- The private jsonb column on items/collections is intentionally excluded — it --- stores operator metadata outside the STAC item identity contract. -CREATE OR REPLACE FUNCTION jsonb_hash(j jsonb) RETURNS bytea AS $$ - SELECT sha256(convert_to(jsonb_canonical(j), 'UTF8')); -$$ LANGUAGE SQL IMMUTABLE PARALLEL SAFE STRICT; +-- The document is flattened to leaf (path, value) rows, each rendered as `path value`, sorted by path +-- (byte order), joined by , and hashed. Separators are C0 control chars US/RS/FS (0x1F/0x1E/0x1D), +-- which do not occur in STAC JSON keys or strings. +-- * path — per step, then 'k'||key (object) or 'i'||idx (0-based array). The 'k'/'i' tags keep an +-- object key "0" distinct from array index 0; keeps nesting distinct from a key with '/'. +-- * number — 'n' || (v)::float8::text, so the digest matches a value stored in a float8 (promoted) column +-- and read back. +-- * string — datetime-shaped (YYYY-MM-DD, optional T/space time, optional offset): 't' || the instant +-- normalized to UTC / 6-digit microseconds / 'Z'. Offset-less is assumed UTC (the SET TIMEZONE +-- below pins the cast), matching pgstac's to_tstz ingest; the cast raises on a shaped-but- +-- invalid timestamp. Other strings: 's' || the raw string. +-- * boolean — 'b'||'true'/'false'; null — 'z'; empty {} — 'e', empty [] — 'a'. +CREATE OR REPLACE FUNCTION jsonb_canonical_hash(j jsonb) RETURNS bytea AS $$ + WITH RECURSIVE t(path, value) AS ( + SELECT ''::text, j + UNION ALL + SELECT t.path || E'\x1F' || e.seg, e.value + FROM t CROSS JOIN LATERAL ( + SELECT 'k' || key AS seg, value + FROM jsonb_each(t.value) WHERE jsonb_typeof(t.value) = 'object' + UNION ALL + SELECT 'i' || (ord - 1)::text AS seg, value + FROM jsonb_array_elements(t.value) WITH ORDINALITY x(value, ord) + WHERE jsonb_typeof(t.value) = 'array' + ) e + ) + SELECT sha256(convert_to(COALESCE(string_agg( + t.path || E'\x1E' || CASE jsonb_typeof(t.value) + WHEN 'number' THEN 'n' || (t.value #>> '{}')::float8::text + WHEN 'boolean' THEN 'b' || (t.value #>> '{}') + WHEN 'null' THEN 'z' + WHEN 'string' THEN CASE + WHEN (t.value #>> '{}') ~ '^\d{4}-\d{2}-\d{2}([Tt ]\d{2}:\d{2}:\d{2}(\.\d+)?([Zz]|[+-]\d{2}:?\d{2})?)?$' + THEN 't' || to_char((t.value #>> '{}')::timestamptz AT TIME ZONE 'UTC', 'YYYY-MM-DD"T"HH24:MI:SS.US') || 'Z' + ELSE 's' || (t.value #>> '{}') END + WHEN 'object' THEN 'e' + WHEN 'array' THEN 'a' + ELSE 'c' END, + E'\x1D' ORDER BY t.path COLLATE "C"), ''), 'UTF8')) + FROM t + WHERE jsonb_typeof(t.value) NOT IN ('object', 'array') OR t.value = '{}'::jsonb OR t.value = '[]'::jsonb; +$$ LANGUAGE SQL IMMUTABLE PARALLEL SAFE STRICT SET TIMEZONE = 'UTC'; -- jsonb_field_rows: Recursively walk a JSONB document and emit one row per field path. -- max_depth guards against runaway recursion on pathologically nested documents. @@ -1746,10 +1735,13 @@ $$ LANGUAGE SQL; CREATE OR REPLACE FUNCTION queryables_trigger_func() RETURNS TRIGGER AS $$ DECLARE BEGIN - PERFORM maintain_partitions(); + -- Queryable definitions changed, so every partition's queryable indexes may be stale. Flag them and let + -- the async sweep (build_pending_indexes) (re)build off the hot path, instead of rebuilding every + -- partition synchronously inside this trigger. + UPDATE pgstac.partition_stats SET indexes_pending = true; RETURN NULL; END; -$$ LANGUAGE PLPGSQL; +$$ LANGUAGE PLPGSQL SECURITY DEFINER; CREATE TRIGGER queryables_trigger AFTER INSERT OR UPDATE ON queryables FOR EACH STATEMENT EXECUTE PROCEDURE queryables_trigger_func(); @@ -2153,6 +2145,98 @@ BEGIN END; $$ LANGUAGE PLPGSQL; +-- q_to_tsquery: parse the STAC `q` free-text parameter (a string or array of strings) into a +-- tsquery, honoring quoted phrases, AND/OR, +/- prefixes, commas (OR) and adjacency. +CREATE OR REPLACE FUNCTION q_to_tsquery (jinput jsonb) + RETURNS tsquery + AS $$ +DECLARE + input text; + processed_text text; + temp_text text; + quote_array text[]; + placeholder text := '@QUOTE@'; +BEGIN + IF jsonb_typeof(jinput) = 'string' THEN + input := jinput->>0; + ELSIF jsonb_typeof(jinput) = 'array' THEN + input := array_to_string( + array(select jsonb_array_elements_text(jinput)), + ' OR ' + ); + ELSE + RAISE EXCEPTION 'Input must be a string or an array of strings.'; + END IF; + -- Extract all quoted phrases and store in array + quote_array := regexp_matches(input, '"[^"]*"', 'g'); + + -- Replace each quoted part with a unique placeholder if there are any quoted phrases + IF array_length(quote_array, 1) IS NOT NULL THEN + processed_text := input; + FOR i IN array_lower(quote_array, 1) .. array_upper(quote_array, 1) LOOP + processed_text := replace(processed_text, quote_array[i], placeholder || i || placeholder); + END LOOP; + ELSE + processed_text := input; + END IF; + + -- Replace non-quoted text using regular expressions + + -- , -> | + processed_text := regexp_replace(processed_text, ',(?=(?:[^"]*"[^"]*")*[^"]*$)', ' | ', 'g'); + + -- and -> & + processed_text := regexp_replace(processed_text, '\s+AND\s+', ' & ', 'gi'); + + -- or -> | + processed_text := regexp_replace(processed_text, '\s+OR\s+', ' | ', 'gi'); + + -- + -> + processed_text := regexp_replace(processed_text, '^\s*\+([a-zA-Z0-9_]+)', '\1', 'g'); -- +term at start + processed_text := regexp_replace(processed_text, '\s*\+([a-zA-Z0-9_]+)', ' & \1', 'g'); -- +term elsewhere + + -- - -> ! + processed_text := regexp_replace(processed_text, '^\s*\-([a-zA-Z0-9_]+)', '! \1', 'g'); -- -term at start + processed_text := regexp_replace(processed_text, '\s*\-([a-zA-Z0-9_]+)', ' & ! \1', 'g'); -- -term elsewhere + + -- terms separated with spaces are assumed to represent adjacent terms. loop through these + -- occurrences and replace them with the adjacency operator (<->) + LOOP + temp_text := regexp_replace(processed_text, '([a-zA-Z0-9_]+)\s+([a-zA-Z0-9_]+)(?!\s*[&|<>])', '\1 <-> \2', 'g'); + IF temp_text = processed_text THEN + EXIT; -- No more replacements were made + END IF; + processed_text := temp_text; + END LOOP; + + + -- Replace placeholders back with quoted phrases if there were any + IF array_length(quote_array, 1) IS NOT NULL THEN + FOR i IN array_lower(quote_array, 1) .. array_upper(quote_array, 1) LOOP + processed_text := replace(processed_text, placeholder || i || placeholder, '''' || substring(quote_array[i] from 2 for length(quote_array[i]) - 2) || ''''); + END LOOP; + END IF; + + RETURN to_tsquery('english', processed_text); +END; +$$ +LANGUAGE plpgsql; + +-- q_op_query: SQL predicate for the pgstac `q` full-text operator. args is the search term(s) +-- (a string or array of strings, the STAC `q` parameter); it is matched against a tsvector +-- built from the row's description/title/keywords. Modeled on spatial_op_query / +-- temporal_op_query so full-text is a first-class CQL2 op. +CREATE OR REPLACE FUNCTION q_op_query(args jsonb) RETURNS text AS $$ + SELECT format( + $q$( + to_tsvector('english', coalesce(properties->>'description', '')) || + to_tsvector('english', coalesce(properties->>'title', '')) || + to_tsvector('english', coalesce(properties->>'keywords', '')) + ) @@ %L$q$, + q_to_tsquery(args) + ); +$$ LANGUAGE SQL STABLE; + CREATE OR REPLACE FUNCTION query_to_cql2(q jsonb) RETURNS jsonb AS $$ -- Translates anything passed in through the deprecated "query" into equivalent CQL2 WITH t AS ( @@ -2207,6 +2291,11 @@ BEGIN END IF; IF jsonb_typeof(j) = 'object' THEN + -- GeoJSON geometry args (Point/Polygon/.../GeometryCollection) are not cql expressions; + -- pass them through unchanged so spatial ops keep their geometry intact. + IF j ? 'type' AND (j ? 'coordinates' OR j ? 'geometries') THEN + RETURN j; + END IF; SELECT jsonb_build_object( 'op', key, 'args', cql1_to_cql2(value) @@ -2317,7 +2406,6 @@ BEGIN END IF; IF j ? 'interval' THEN RAISE EXCEPTION 'Please use temporal operators when using intervals.'; - RETURN NONE; END IF; -- Spatial Query @@ -2325,6 +2413,11 @@ BEGIN RETURN spatial_op_query(op, args); END IF; + -- Full-text Query (pgstac `q` operator) + IF op = 'q' THEN + RETURN q_op_query(args); + END IF; + IF op IN ('a_equals','a_contains','a_contained_by','a_overlaps') THEN IF args->0 ? 'property' THEN leftarg := format('to_text_array(%s)', (queryable(args->0->>'property')).path); @@ -2457,87 +2550,6 @@ END; $$ LANGUAGE PLPGSQL STABLE; -CREATE OR REPLACE FUNCTION paging_dtrange( - j jsonb -) RETURNS tstzrange AS $$ -DECLARE - op text; - filter jsonb := j->'filter'; - dtrange tstzrange := tstzrange('-infinity'::timestamptz,'infinity'::timestamptz); - sdate timestamptz := '-infinity'::timestamptz; - edate timestamptz := 'infinity'::timestamptz; - jpitem jsonb; -BEGIN - - IF j ? 'datetime' THEN - dtrange := parse_dtrange(j->'datetime'); - sdate := lower(dtrange); - edate := upper(dtrange); - END IF; - IF NOT (filter @? '$.**.op ? (@ == "or" || @ == "not")') THEN - FOR jpitem IN SELECT j FROM jsonb_path_query(filter,'strict $.** ? (@.args[*].property == "datetime")'::jsonpath) j LOOP - op := lower(jpitem->>'op'); - dtrange := parse_dtrange(jpitem->'args'->1); - IF op IN ('<=', 'lt', 'lte', '<', 'le', 't_before') THEN - sdate := greatest(sdate,'-infinity'); - edate := least(edate, upper(dtrange)); - ELSIF op IN ('>=', '>', 'gt', 'gte', 'ge', 't_after') THEN - edate := least(edate, 'infinity'); - sdate := greatest(sdate, lower(dtrange)); - ELSIF op IN ('=', 'eq') THEN - edate := least(edate, upper(dtrange)); - sdate := greatest(sdate, lower(dtrange)); - END IF; - RAISE NOTICE '2 OP: %, ARGS: %, DTRANGE: %, SDATE: %, EDATE: %', op, jpitem->'args'->1, dtrange, sdate, edate; - END LOOP; - END IF; - IF sdate > edate THEN - RETURN 'empty'::tstzrange; - END IF; - RETURN tstzrange(sdate,edate, '[]'); -END; -$$ LANGUAGE PLPGSQL STABLE STRICT SET TIME ZONE 'UTC'; - -CREATE OR REPLACE FUNCTION paging_collections( - IN j jsonb -) RETURNS text[] AS $$ -DECLARE - filter jsonb := j->'filter'; - jpitem jsonb; - op text; - args jsonb; - arg jsonb; - collections text[]; -BEGIN - IF j ? 'collections' THEN - collections := to_text_array(j->'collections'); - END IF; - IF NOT (filter @? '$.**.op ? (@ == "or" || @ == "not")') THEN - FOR jpitem IN SELECT j FROM jsonb_path_query(filter,'strict $.** ? (@.args[*].property == "collection")'::jsonpath) j LOOP - RAISE NOTICE 'JPITEM: %', jpitem; - op := jpitem->>'op'; - args := jpitem->'args'; - IF op IN ('=', 'eq', 'in') THEN - FOR arg IN SELECT a FROM jsonb_array_elements(args) a LOOP - IF jsonb_typeof(arg) IN ('string', 'array') THEN - RAISE NOTICE 'arg: %, collections: %', arg, collections; - IF collections IS NULL OR collections = '{}'::text[] THEN - collections := to_text_array(arg); - ELSE - collections := array_intersection(collections, to_text_array(arg)); - END IF; - END IF; - END LOOP; - END IF; - END LOOP; - END IF; - IF collections = '{}'::text[] THEN - RETURN NULL; - END IF; - RETURN collections; -END; -$$ LANGUAGE PLPGSQL STABLE STRICT; - -- coerce a cql2 scalar (string | {"timestamp":..} | {"date":..}) to timestamptz CREATE OR REPLACE FUNCTION cql2_ts(v jsonb) RETURNS timestamptz LANGUAGE sql IMMUTABLE AS $$ SELECT coalesce(v->>'timestamp', v->>'date', v#>>'{}')::timestamptz; @@ -2643,7 +2655,11 @@ BEGIN END IF; RETURN acc; ELSIF op ILIKE 's_%' OR op = 'intersects' THEN - g := ST_GeomFromGeoJSON(args->1); + BEGIN + g := ST_GeomFromGeoJSON(args->1); + EXCEPTION WHEN others THEN + RAISE EXCEPTION 'Invalid GeoJSON geometry: %', args->1 USING ERRCODE = '22P02'; + END; acc := env_full(); acc.geom := ST_Envelope(g); RETURN acc; ELSIF op IN ('=','<','<=','>','>=','between','eq','lt','lte','gt','gte','in','like','ilike') AND jsonb_typeof(args)='array' AND args->0 ? 'property' THEN @@ -2711,22 +2727,20 @@ BEGIN END; $$ LANGUAGE PLPGSQL STABLE; - - - - -- search_envelope: convert a STAC search JSON to a pred_envelope for partition pruning. -- Used by tilesearch and external callers that have raw search JSON. -CREATE OR REPLACE FUNCTION search_envelope(j jsonb) RETURNS pred_envelope LANGUAGE plpgsql STABLE AS $$ +CREATE OR REPLACE FUNCTION search_envelope(j jsonb) RETURNS pred_envelope LANGUAGE sql STABLE AS $$ SELECT cql2_envelope(search_to_cql2(j)); $$; --- partition_bounds: read partition_stats once using an envelope. Returns candidate --- collections, per-month row counts, and total count. _asc controls output array --- order: ASC for datetime-ASC sorts, DESC for datetime-DESC sorts. +-- partition_bounds: read partition_stats once using an envelope. Returns the candidate +-- collections, the per-month row-count histogram as aligned ascending arrays (months[] + +-- counts[]), and the total candidate count. next_band walks counts[] by index; callers map +-- indices back to timestamps via months[]. Each candidate partition's row estimate is prorated +-- across the calendar months its [lo,hi) data extent spans (an instant lo=hi puts all n in its +-- month), then summed per month. CREATE OR REPLACE FUNCTION partition_bounds( _env pred_envelope, - _asc boolean DEFAULT true, OUT months timestamptz[], OUT counts bigint[], OUT collections text[], @@ -2734,8 +2748,8 @@ CREATE OR REPLACE FUNCTION partition_bounds( ) RETURNS record LANGUAGE sql STABLE AS $$ WITH cand AS ( SELECT ps.collection, - lower(coalesce(ps.dtrange, ps.partition_dtrange)) AS lo, - upper(coalesce(ps.dtrange, ps.partition_dtrange)) AS hi, + lower(ps.dtrange) AS lo, + upper(ps.dtrange) AS hi, coalesce(ps.n, 0) AS n FROM partition_stats ps WHERE ((_env).colls IS NULL OR ps.collection = ANY((_env).colls)) @@ -2745,31 +2759,44 @@ CREATE OR REPLACE FUNCTION partition_bounds( ), monthly AS ( SELECT date_trunc('month', gs) AS month_start, - CASE WHEN c.hi <= c.lo THEN c.n::numeric - ELSE c.n * ( - extract(epoch FROM (LEAST(c.hi, date_trunc('month', gs) + interval '1 month') - GREATEST(c.lo, date_trunc('month', gs))) + CASE + WHEN c.hi <= c.lo THEN c.n::numeric + ELSE c.n * ( + extract(epoch FROM ( + LEAST(c.hi, date_trunc('month', gs) + interval '1 month') + - GREATEST(c.lo, date_trunc('month', gs)))) / extract(epoch FROM (c.hi - c.lo))) END AS pn FROM cand c, generate_series(date_trunc('month', c.lo), date_trunc('month', GREATEST(c.lo, c.hi - interval '1 microsecond')), interval '1 month') AS gs + ), + buckets AS ( + SELECT month_start, round(sum(pn))::bigint AS n + FROM monthly + GROUP BY month_start ) SELECT - ARRAY(SELECT month_start FROM monthly ORDER BY month_start CASE WHEN _asc THEN ASC ELSE DESC END), - ARRAY(SELECT round(pn)::bigint FROM monthly ORDER BY month_start CASE WHEN _asc THEN ASC ELSE DESC END), + (SELECT array_agg(month_start ORDER BY month_start) FROM buckets), + (SELECT array_agg(n ORDER BY month_start) FROM buckets), (SELECT array_agg(DISTINCT collection) FROM cand), (SELECT coalesce(sum(n), 0) FROM cand); $$; -- next_band: walk a histogram using array indexes. Given per-month counts, a -- cursor position, and a target row count, returns the next band's index range. --- The caller converts band indexes back to timestamps for SQL queries. +-- The caller converts band indexes back to timestamps for SQL queries. The band +-- index range [band_start_idx, band_end_idx] is always low..high regardless of +-- direction; only the walk order differs. For a descending search (_descending) +-- the cursor starts at the most recent month and walks toward older months, so +-- the most recent items are collected first. CREATE OR REPLACE FUNCTION next_band( _counts bigint[], _cursor_idx int, _target numeric, _cap_months int, + _descending boolean DEFAULT false, OUT band_start_idx int, OUT band_end_idx int, OUT scanned bigint, @@ -2778,22 +2805,49 @@ CREATE OR REPLACE FUNCTION next_band( ) RETURNS record LANGUAGE plpgsql STABLE AS $$ DECLARE idx int; + n int; cumulative bigint := 0; BEGIN done := false; scanned := 0; band_start_idx := NULL; band_end_idx := NULL; next_cursor_idx := _cursor_idx; IF _counts IS NULL OR array_length(_counts, 1) IS NULL OR _cursor_idx IS NULL THEN + done := true; -- no histogram / no cursor => nothing to walk + RETURN; + END IF; + n := array_length(_counts, 1); + + IF _descending THEN + -- Walk downward (newest -> oldest). The high bound is the cursor (clamped into range). + IF _cursor_idx < 1 THEN done := true; RETURN; END IF; + idx := LEAST(_cursor_idx, n); + FOR i IN REVERSE idx..GREATEST(1, idx - _cap_months + 1) LOOP + cumulative := cumulative + _counts[i]; + scanned := scanned + _counts[i]; + IF cumulative >= _target THEN + band_start_idx := i; -- low (older) month bound + band_end_idx := idx; -- high (newer) month bound + next_cursor_idx := i - 1; + IF next_cursor_idx < 1 THEN done := true; END IF; + RETURN; + END IF; + END LOOP; + -- Target not reached within the cap: end the band at the cap boundary. + band_start_idx := GREATEST(1, idx - _cap_months + 1); + band_end_idx := idx; + next_cursor_idx := band_start_idx - 1; + done := (band_start_idx <= 1); RETURN; END IF; + -- Ascending (oldest -> newest). idx := NULL; - FOR i IN 1..array_length(_counts, 1) LOOP + FOR i IN 1..n LOOP IF i >= _cursor_idx THEN idx := i; EXIT; END IF; END LOOP; IF idx IS NULL THEN done := true; next_cursor_idx := _cursor_idx; RETURN; END IF; - FOR i IN idx..array_length(_counts, 1) LOOP + FOR i IN idx..n LOOP IF i > idx + _cap_months - 1 THEN EXIT; END IF; cumulative := cumulative + _counts[i]; scanned := scanned + _counts[i]; @@ -2801,15 +2855,17 @@ BEGIN band_start_idx := idx; band_end_idx := i; next_cursor_idx := i + 1; - IF next_cursor_idx > array_length(_counts, 1) THEN done := true; END IF; + IF next_cursor_idx > n THEN done := true; END IF; RETURN; END IF; END LOOP; + -- Target not reached within the cap: end the band at the cap boundary (not the array end), + -- so the cap actually limits band width. done only when we've consumed the whole histogram. band_start_idx := idx; - band_end_idx := array_length(_counts, 1); - next_cursor_idx := array_length(_counts, 1) + 1; - done := true; + band_end_idx := LEAST(idx + _cap_months - 1, n); + next_cursor_idx := band_end_idx + 1; + done := (band_end_idx >= n); END; $$; -- END FRAGMENT: 002c_envelope.sql @@ -2820,12 +2876,20 @@ CREATE OR REPLACE FUNCTION keyset_encode(vals text[]) RETURNS text AS $$ SELECT encode(convert_to(array_to_string(vals, chr(31), chr(30)), 'UTF8'), 'base64'); $$ LANGUAGE sql IMMUTABLE; --- Decode a base64 keyset token back to sort key values. +-- Decode a base64 keyset token back to sort key values. An empty/NULL token returns NULL +-- ("no keyset" => first page). A non-empty token that is not a valid base64 keyset (e.g. a +-- stale/old-style token) raises 22P02 rather than silently returning the first page. CREATE OR REPLACE FUNCTION keyset_decode(token text) RETURNS text[] AS $$ - SELECT array_replace( +BEGIN + IF token IS NULL OR token = '' THEN RETURN NULL; END IF; + RETURN array_replace( string_to_array(convert_from(decode(token,'base64'),'UTF8'), chr(31)), chr(30), NULL); -$$ LANGUAGE sql IMMUTABLE; +EXCEPTION WHEN others THEN + -- A non-empty token that does not decode is a client error, not an empty page. + RAISE EXCEPTION 'Invalid pagination token: %', token USING ERRCODE = '22P02'; +END; +$$ LANGUAGE plpgsql IMMUTABLE; -- Resolve sortby + id/collection tiebreaks into ordered sort keys with SQL -- expressions and directions for a unique total row order. @@ -2952,7 +3016,7 @@ CREATE TABLE items ( stac_version text, stac_extensions jsonb DEFAULT '[]'::jsonb, pgstac_updated_at timestamptz NOT NULL DEFAULT now(), - -- 32-byte sha256 of the canonical (RFC 8785-aligned) STAC item JSON set at + -- 32-byte sha256 of the canonical (jsonb_canonical) STAC item JSON set at -- ingest time. Allows external clients to detect unchanged items without a -- full fetch. Does NOT include the private column (operator metadata). item_hash bytea NOT NULL DEFAULT '\x'::bytea, @@ -3042,47 +3106,15 @@ CREATE STATISTICS datetime_stats (dependencies) on datetime, end_datetime from i ALTER TABLE items ADD CONSTRAINT items_collections_fk FOREIGN KEY (collection) REFERENCES collections(id) ON DELETE CASCADE DEFERRABLE; --- partition_after_triggerfunc: After-statement trigger on items. --- Updates partition statistics for every partition touched by the current batch, --- using run_or_queue() so the work is deferred rather than blocking the ingest --- transaction. -CREATE OR REPLACE FUNCTION partition_after_triggerfunc() RETURNS TRIGGER AS $$ -DECLARE - p text; - t timestamptz := clock_timestamp(); -BEGIN - RAISE NOTICE 'Updating partition stats %', t; - FOR p IN SELECT DISTINCT partition - FROM newdata n JOIN partition_sys_meta p - ON (n.collection=p.collection AND n.datetime <@ p.partition_dtrange) - LOOP - PERFORM run_or_queue(format('SELECT update_partition_stats(%L, %L);', p, true)); - END LOOP; - RAISE NOTICE 't: % %', t, clock_timestamp() - t; - RETURN NULL; -END; -$$ LANGUAGE PLPGSQL SECURITY DEFINER; - +-- The old per-statement stats trigger (partition_after_triggerfunc) is intentionally GONE. Under the +-- widen-now / tighten-async model, partition_stats coverage is maintained by the ingest path itself +-- (check_partition seeds + widen_partition_stats covers, in their own transactions) and tightened +-- off the hot path by tighten_partition_stats; a per-insert trigger that re-scanned partitions and +-- refreshed matviews was the lock pathology this model removes. DROP TRIGGER IF EXISTS items_after_insert_trigger ON items; -CREATE TRIGGER items_after_insert_trigger -AFTER INSERT ON items -REFERENCING NEW TABLE AS newdata -FOR EACH STATEMENT -EXECUTE FUNCTION partition_after_triggerfunc(); - DROP TRIGGER IF EXISTS items_after_update_trigger ON items; -CREATE TRIGGER items_after_update_trigger -AFTER UPDATE ON items -REFERENCING NEW TABLE AS newdata -FOR EACH STATEMENT -EXECUTE FUNCTION partition_after_triggerfunc(); - DROP TRIGGER IF EXISTS items_after_delete_trigger ON items; -CREATE TRIGGER items_after_delete_trigger -AFTER DELETE ON items -REFERENCING OLD TABLE AS newdata -FOR EACH STATEMENT -EXECUTE FUNCTION partition_after_triggerfunc(); +DROP FUNCTION IF EXISTS partition_after_triggerfunc(); -- items_content_distinct_sql / items_content_changed: detect whether a direct -- UPDATE actually changed the stored item content. Used by the touch trigger @@ -3133,7 +3165,7 @@ $$ LANGUAGE PLPGSQL IMMUTABLE PARALLEL SAFE; -- items_touch_triggerfunc: refresh pgstac_updated_at when a direct UPDATE changes -- the stored item content. It deliberately does NOT recompute item_hash: --- item_hash is the canonical (RFC 8785-aligned) hash of the item *as ingested* +-- item_hash is the canonical (jsonb_canonical) hash of the item *as ingested* -- through create_item / upsert_item / update_item (set once in content_dehydrate), -- so it stays externally reproducible by a client hashing its own copy. -- A raw `UPDATE items SET ...` that bypasses the staging path leaves item_hash @@ -3346,7 +3378,7 @@ CREATE OR REPLACE FUNCTION content_dehydrate(content jsonb) RETURNS items AS $$ content->>'stac_version' AS stac_version, COALESCE(content->'stac_extensions', '[]'::jsonb) AS stac_extensions, now() AS pgstac_updated_at, - pgstac.jsonb_hash(content) AS item_hash, + pgstac.jsonb_canonical_hash(content) AS item_hash, NULL::bigint AS fragment_id, content->'bbox' AS bbox, CASE WHEN content->'links' IS NOT NULL AND content->'links' <> '[]'::jsonb THEN content->'links' END AS links, @@ -3401,6 +3433,9 @@ CREATE OR REPLACE FUNCTION content_dehydrate(content jsonb) RETURNS items AS $$ NULL::jsonb AS private; $$ LANGUAGE SQL STABLE; +-- include_field: STAC fields include/exclude decision over a fields jsonb (used by content_hydrate); +-- the jsonb-form of field_included(). Same rule: exclude wins, an explicit include list restricts to +-- its members, otherwise everything is included. NULL field returns NULL. CREATE OR REPLACE FUNCTION include_field(f text, fields jsonb DEFAULT '{}'::jsonb) RETURNS boolean AS $$ DECLARE includes jsonb := fields->'include'; @@ -3438,13 +3473,12 @@ BEGIN END; $$ LANGUAGE PLPGSQL IMMUTABLE; --- content_hydrate: Reassemble a full STAC item JSON from the split columns --- and the shared fragment content. This is the single hydrate function; --- the old content_nonhydrated wrapper and 3-arg _collection parameter have --- been removed. +-- content_hydrate: reassemble a full STAC item JSON from the split columns and the shared fragment +-- content. The single hydrate function. CREATE OR REPLACE FUNCTION content_hydrate( _item items, - fields jsonb DEFAULT '{}'::jsonb + fields jsonb DEFAULT '{}'::jsonb, + _skip_fragment boolean DEFAULT false ) RETURNS jsonb AS $$ DECLARE geom jsonb; @@ -3461,8 +3495,10 @@ BEGIN geom := ST_ASGeoJson(_item.geometry, 20)::jsonb; END IF; - -- Fetch shared fragment content (NULL when item has no fragment). - IF _item.fragment_id IS NOT NULL THEN + -- Fetch shared fragment content (NULL when item has no fragment). _skip_fragment lets a caller + -- that has already determined the requested fields are satisfiable from item columns alone + -- (via needs_fragment) avoid this per-row lookup entirely. + IF _item.fragment_id IS NOT NULL AND NOT _skip_fragment THEN SELECT content, links_template INTO frag_content, frag_links_template FROM item_fragments @@ -3478,7 +3514,7 @@ BEGIN WHEN _item.stac_extensions IS NOT NULL AND _item.stac_extensions <> '[]'::jsonb THEN _item.stac_extensions ELSE COALESCE(frag_content->'stac_extensions', _item.stac_extensions) END; - IF _item.fragment_id IS NOT NULL THEN + IF _item.fragment_id IS NOT NULL AND NOT _skip_fragment THEN hydrated_links := stac_links_hydrate(frag_links_template, _item.link_hrefs); ELSE hydrated_links := COALESCE(_item.links, '[]'::jsonb); @@ -3533,7 +3569,7 @@ CREATE UNLOGGED TABLE items_staging_upsert ( -- item_fragments rows via ON CONFLICT), assigns fragment_id, and strips -- fragment-covered keys. Returns the fully-enriched rows as the items rowtype so -- each staging trigger branch is a single INSERT differing only in conflict --- policy. The enriched column list lives here once (previously duplicated 3x). +-- policy. The enriched column list lives here once. CREATE OR REPLACE FUNCTION items_staging_dehydrate(_contents jsonb[]) RETURNS SETOF items AS $$ WITH raw AS MATERIALIZED ( SELECT @@ -3680,6 +3716,22 @@ DECLARE BEGIN RAISE NOTICE 'Creating Partitions. %', clock_timestamp() - ts; + -- Fail loudly on items whose collection does not exist instead of silently dropping them + -- in the collections JOINs below (matches the Rust loader, which also errors on this). + IF EXISTS ( + SELECT 1 FROM newdata n + WHERE NOT EXISTS ( + SELECT 1 FROM collections c WHERE c.id = n.content->>'collection' + ) + ) THEN + RAISE EXCEPTION 'cannot load item(s) into nonexistent collection(s): %', + (SELECT string_agg(DISTINCT coalesce(n.content->>'collection', ''), ', ') + FROM newdata n + WHERE NOT EXISTS ( + SELECT 1 FROM collections c WHERE c.id = n.content->>'collection' + )); + END IF; + FOR part IN WITH t AS ( SELECT n.content->>'collection' as collection, @@ -3736,13 +3788,31 @@ BEGIN RAISE NOTICE 'Inserted % rows to items. %', nrows, clock_timestamp() - ts; END IF; + -- Bump each partition's row-count estimate so search stepping sees the new rows. n drives + -- partition_bounds' histogram and the Rust keyset search_page skips zero-count bands, so a partition + -- left at n=0 would hide its freshly-ingested items from that path (SQL search() scans regardless). + -- Mirrors flush_items_staging_binary on the Rust loader. Over-count is safe (ignore/upsert may insert + -- fewer than staged); the async tightener resets n exactly. check_partition above already widened the + -- temporal envelope; spatial is left to the tightener on this fallback path. (n only affects search + -- stepping performance, never which rows/order come back — see geometrysearch band direction.) + UPDATE partition_stats ps + SET n = COALESCE(ps.n, 0) + agg.c, dirty = true, last_updated = now() + FROM ( + SELECT (partition_name(nd.content->>'collection', + lower(stac_daterange(nd.content->'properties')))).partition_name AS partition, + count(*) AS c + FROM newdata nd + GROUP BY 1 + ) agg + WHERE ps.partition = agg.partition; + RAISE NOTICE 'Deleting data from staging table. %', clock_timestamp() - ts; EXECUTE format('DELETE FROM %I', TG_TABLE_NAME); RAISE NOTICE 'Done. %', clock_timestamp() - ts; RETURN NULL; END; -$$ LANGUAGE PLPGSQL; +$$ LANGUAGE PLPGSQL SECURITY DEFINER; DROP TRIGGER IF EXISTS items_staging_insert_trigger ON items_staging; @@ -3778,7 +3848,7 @@ out items%ROWTYPE; BEGIN DELETE FROM items WHERE id = _id AND (_collection IS NULL OR collection=_collection) RETURNING * INTO STRICT out; END; -$$ LANGUAGE PLPGSQL; +$$ LANGUAGE PLPGSQL SECURITY DEFINER; --/* CREATE OR REPLACE FUNCTION create_item(data jsonb) RETURNS VOID AS $$ @@ -3823,156 +3893,10 @@ CREATE OR REPLACE FUNCTION collection_temporal_extent(id text) RETURNS jsonb AS ; $$ LANGUAGE SQL IMMUTABLE PARALLEL SAFE SET SEARCH_PATH TO pgstac, public; -CREATE OR REPLACE FUNCTION update_collection_extents() RETURNS VOID AS $$ -UPDATE collections - SET content = jsonb_set_lax( - content, - '{extent}'::text[], - collection_extent(id, FALSE), - true, - 'use_json_null' - ) -; -$$ LANGUAGE SQL; - --- --------------------------------------------------------------------------- --- Field Registry: walks JSONB item content to track which paths exist in each --- collection. Used to auto-populate queryables and support schema inference. --- jsonb_field_rows is defined in 001a_jsonutils.sql (loaded first). --- --------------------------------------------------------------------------- - --- update_field_registry_from_sample: UPSERT registry rows from a pre-selected array of --- raw item content JSONBs. Callers supply the sample to decouple sampling strategy --- from the registry write; merge value_kinds to accumulate observed types over time. -CREATE OR REPLACE FUNCTION update_field_registry_from_sample( - _collection text, - item_contents jsonb[] -) RETURNS void AS $$ - INSERT INTO item_field_registry (collection, path, is_leaf, value_kinds, first_seen, last_seen) - SELECT - _collection, - r.path, - bool_and(r.is_leaf) AS is_leaf, - array_agg(DISTINCT r.value_kind) FILTER (WHERE r.value_kind IS NOT NULL) AS value_kinds, - now(), - now() - FROM unnest(item_contents) AS item(content) - CROSS JOIN LATERAL jsonb_field_rows(item.content) AS r(path, is_leaf, value_kind) - GROUP BY r.path - ON CONFLICT (collection, path) DO UPDATE SET - is_leaf = EXCLUDED.is_leaf, - value_kinds = ( - SELECT array_agg(DISTINCT v) - FROM unnest(item_field_registry.value_kinds || EXCLUDED.value_kinds) t(v) - ), - last_seen = now() - ; -$$ LANGUAGE SQL VOLATILE; --- update_field_registry_from_items: Sample a live collection and UPSERT registry rows. --- Uses TABLESAMPLE BERNOULLI(5) for large collections (>10k rows by pg_class estimate) --- and LIMIT 1000 for smaller ones to avoid a full seq-scan for tiny collections. --- pg_class.reltuples is an estimate (may be stale); its only role is threshold selection. --- Returns (registered_paths, rows_processed) for observability. -CREATE OR REPLACE FUNCTION update_field_registry_from_items( - _collection text -) RETURNS TABLE (registered_paths int, rows_processed int) AS $$ -DECLARE - est_rows bigint; - nrows int; - npaths int; -BEGIN - -- Sum reltuples across the registered item partitions for this collection. - -- reltuples can be -1 (never analyzed); treat negative values as zero. - SELECT COALESCE(sum(GREATEST(c.reltuples::bigint, 0)), 0) INTO est_rows - FROM partitions_view p - JOIN pg_class c ON c.relname = p.partition - JOIN pg_namespace n ON n.oid = c.relnamespace - WHERE p.collection = _collection - AND n.nspname = 'pgstac' - AND c.relkind = 'r'; - IF est_rows > 10000 THEN - -- Large collection: use statistical sampling to avoid full seq-scan. - WITH sampled AS ( - SELECT content_hydrate(i) AS content FROM items i TABLESAMPLE BERNOULLI(5) WHERE i.collection = _collection - ), - upserted AS ( - INSERT INTO item_field_registry (collection, path, is_leaf, value_kinds, first_seen, last_seen) - SELECT - _collection, - r.path, - bool_and(r.is_leaf) AS is_leaf, - array_agg(DISTINCT r.value_kind) FILTER (WHERE r.value_kind IS NOT NULL) AS value_kinds, - now(), now() - FROM sampled - CROSS JOIN LATERAL jsonb_field_rows(content) AS r(path, is_leaf, value_kind) - GROUP BY r.path - ON CONFLICT (collection, path) DO UPDATE SET - is_leaf = EXCLUDED.is_leaf, - value_kinds = ( - SELECT array_agg(DISTINCT v) - FROM unnest(item_field_registry.value_kinds || EXCLUDED.value_kinds) t(v) - ), - last_seen = now() - RETURNING 1 - ) - SELECT - (SELECT count(*)::int FROM upserted), - (SELECT count(*)::int FROM sampled) - INTO npaths, nrows; - ELSE - -- Small collection: process up to 1000 rows to avoid BERNOULLI returning 0 rows. - WITH sampled AS ( - SELECT content_hydrate(i) AS content FROM items i WHERE i.collection = _collection LIMIT 1000 - ), - upserted AS ( - INSERT INTO item_field_registry (collection, path, is_leaf, value_kinds, first_seen, last_seen) - SELECT - _collection, - r.path, - bool_and(r.is_leaf) AS is_leaf, - array_agg(DISTINCT r.value_kind) FILTER (WHERE r.value_kind IS NOT NULL) AS value_kinds, - now(), now() - FROM sampled - CROSS JOIN LATERAL jsonb_field_rows(content) AS r(path, is_leaf, value_kind) - GROUP BY r.path - ON CONFLICT (collection, path) DO UPDATE SET - is_leaf = EXCLUDED.is_leaf, - value_kinds = ( - SELECT array_agg(DISTINCT v) - FROM unnest(item_field_registry.value_kinds || EXCLUDED.value_kinds) t(v) - ), - last_seen = now() - RETURNING 1 - ) - SELECT - (SELECT count(*)::int FROM upserted), - (SELECT count(*)::int FROM sampled) - INTO npaths, nrows; - END IF; - RETURN QUERY SELECT npaths, nrows; -END; -$$ LANGUAGE PLPGSQL VOLATILE SECURITY DEFINER; --- refresh_field_registry: Expire stale registry entries that haven't been seen recently. --- Intended for scheduled maintenance (e.g. pg_cron daily job). --- Returns (collection, expired_paths) for each collection affected. -CREATE OR REPLACE FUNCTION refresh_field_registry( - _collection text DEFAULT NULL, - retention_interval interval DEFAULT '90 days' -) RETURNS TABLE (collection_id text, expired_paths int) AS $$ - WITH deleted AS ( - DELETE FROM item_field_registry - WHERE (_collection IS NULL OR collection = _collection) - AND last_seen < now() - retention_interval - RETURNING collection - ) - SELECT collection, count(*)::int - FROM deleted - GROUP BY collection; -$$ LANGUAGE SQL VOLATILE; -- Item Fragment Management functions @@ -4029,38 +3953,6 @@ CREATE OR REPLACE FUNCTION pgstac_hash_fragment(fragment jsonb) RETURNS bytea AS SELECT sha256(convert_to(fragment::text, 'UTF8')); $$ LANGUAGE SQL IMMUTABLE PARALLEL SAFE; --- gc_fragments: Garbage collect orphaned fragments using a single set-based DELETE. --- Replaces the previous per-collection FOR LOOP with a single statement that lets --- the planner choose the optimal join/anti-join strategy across all collections. --- The NOT EXISTS sub-select is evaluated per fragment; with an index on items.fragment_id --- this is an efficient anti-join rather than a full seq-scan. --- --- Operational note: because items.fragment_id is intentionally unmanaged by an FK --- (partitioned-items incremental NOT VALID FKs are not supported), gc_fragments has a --- small race window with concurrent inserts. A fragment that is unreferenced at the time --- the DELETE snapshot is taken but becomes referenced by a later insert could be removed. --- The retention_interval guard makes this unlikely for normal ingest, but operators should --- still run gc_fragments during low-ingest periods or with a sufficiently conservative --- retention interval. This is a documented operational tradeoff, not a silent invariant. -CREATE OR REPLACE FUNCTION gc_fragments( - _collection text DEFAULT NULL, - retention_interval interval DEFAULT '90 days' -) RETURNS TABLE ( - collection_id text, - fragments_removed int -) AS $$ - WITH deleted AS ( - DELETE FROM item_fragments f - WHERE - (_collection IS NULL OR f.collection = _collection) - AND f.created_at < now() - retention_interval - AND NOT EXISTS (SELECT 1 FROM items i WHERE i.fragment_id = f.id) - RETURNING f.collection - ) - SELECT collection, count(*)::int - FROM deleted - GROUP BY collection; -$$ LANGUAGE SQL VOLATILE PARALLEL UNSAFE; -- strip_fragment_col: Remove fragment-owned sub-keys from a split column value. -- col_name is the top-level STAC key that this column represents (e.g. 'assets' or 'properties'). @@ -4110,14 +4002,31 @@ $$ LANGUAGE PLPGSQL IMMUTABLE PARALLEL SAFE; -- BEGIN FRAGMENT: 003b_partitions.sql CREATE TABLE partition_stats ( partition text PRIMARY KEY, + collection text, + -- dtrange/edtrange are the data bounds used for search pruning. widen_partition_stats fills them + -- (generously) on ingest; tighten_partition_stats narrows them to the exact extent. The partition's + -- structural range lives in its own constraint, read via partitions_view.constraint_dtrange. dtrange tstzrange, edtrange tstzrange, spatial geometry, last_updated timestamptz, - keys text[] -) WITH (FILLFACTOR=90); + n bigint, + -- Deferred-maintenance flags: + -- dirty: the stored envelope may be WIDER than the actual data; an async tightener + -- should recompute exact min/max/extent and clear the flag. Search correctness + -- never depends on this (a wide envelope only over-includes a partition). + -- indexes_pending: the partition currently carries only the parent-inherited indexes (id PK, + -- datetime, geometry); its queryable-defined indexes have not been built yet. + dirty boolean NOT NULL DEFAULT false, + indexes_pending boolean NOT NULL DEFAULT false +) WITH (FILLFACTOR=70); CREATE INDEX partitions_range_idx ON partition_stats USING GIST(dtrange); +CREATE INDEX partition_stats_collection_idx ON partition_stats (collection); +CREATE INDEX partition_stats_spatial_idx ON partition_stats USING GIST(spatial) WHERE spatial IS NOT NULL; +-- Work-queue indexes: the maintenance sweeps find pending partitions without scanning every row. +CREATE INDEX partition_stats_dirty_idx ON partition_stats (partition) WHERE dirty; +CREATE INDEX partition_stats_indexes_pending_idx ON partition_stats (partition) WHERE indexes_pending; CREATE OR REPLACE FUNCTION constraint_tstzrange(expr text) RETURNS tstzrange AS $$ @@ -4195,10 +4104,9 @@ SELECT level, c.reltuples, c.relhastriggers, - partition_dtrange, COALESCE( get_tstz_constraint(c.oid, 'datetime'), - partition_dtrange, + constraint_tstzrange(pg_get_expr(c.relpartbound, c.oid)), inf_range ) as constraint_dtrange, COALESCE( @@ -4214,7 +4122,6 @@ FROM JOIN LATERAL pg_get_expr(c.relpartbound, c.oid) as partition_expr ON TRUE JOIN LATERAL pg_get_expr(parent.relpartbound, parent.oid) as parent_partition_expr ON TRUE JOIN LATERAL tstzrange('-infinity', 'infinity','[]') as inf_range ON TRUE - JOIN LATERAL COALESCE(constraint_tstzrange(pg_get_expr(c.relpartbound, c.oid)), inf_range) as partition_dtrange ON TRUE JOIN LATERAL get_tstz_constraint(c.oid, 'datetime') as datetime_constraint ON TRUE JOIN LATERAL get_tstz_constraint(c.oid, 'end_datetime') as end_datetime_constraint ON TRUE WHERE isleaf @@ -4235,10 +4142,9 @@ SELECT level, c.reltuples, c.relhastriggers, - partition_dtrange, COALESCE( get_tstz_constraint(c.oid, 'datetime'), - partition_dtrange, + constraint_tstzrange(pg_get_expr(c.relpartbound, c.oid)), inf_range ) as constraint_dtrange, COALESCE( @@ -4258,107 +4164,17 @@ FROM JOIN LATERAL pg_get_expr(c.relpartbound, c.oid) as partition_expr ON TRUE JOIN LATERAL pg_get_expr(parent.relpartbound, parent.oid) as parent_partition_expr ON TRUE JOIN LATERAL tstzrange('-infinity', 'infinity','[]') as inf_range ON TRUE - JOIN LATERAL COALESCE(constraint_tstzrange(pg_get_expr(c.relpartbound, c.oid)), inf_range) as partition_dtrange ON TRUE JOIN LATERAL get_tstz_constraint(c.oid, 'datetime') as datetime_constraint ON TRUE JOIN LATERAL get_tstz_constraint(c.oid, 'end_datetime') as end_datetime_constraint ON TRUE - LEFT JOIN pgstac.partition_stats USING (partition) + -- the view computes its own collection + constraint_dtrange from the live tree; pull only the + -- data-extent columns from partition_stats to avoid colliding with those names. + LEFT JOIN ( + SELECT partition, dtrange, edtrange, spatial, last_updated + FROM pgstac.partition_stats + ) ps USING (partition) WHERE isleaf ; -CREATE MATERIALIZED VIEW partitions AS -SELECT * FROM partitions_view; -CREATE UNIQUE INDEX ON partitions (partition); - -CREATE MATERIALIZED VIEW partition_steps AS -SELECT - partition as name, - date_trunc('month',lower(partition_dtrange)) as sdate, - date_trunc('month', upper(partition_dtrange)) + '1 month'::interval as edate - FROM partitions_view WHERE partition_dtrange IS NOT NULL AND partition_dtrange != 'empty'::tstzrange - ORDER BY dtrange ASC -; - - -CREATE OR REPLACE FUNCTION update_partition_stats_q(_partition text, istrigger boolean default false) RETURNS VOID AS $$ -DECLARE -BEGIN - PERFORM run_or_queue( - format('SELECT update_partition_stats(%L, %L);', _partition, istrigger) - ); -END; -$$ LANGUAGE PLPGSQL; - -CREATE OR REPLACE FUNCTION update_partition_stats(_partition text, istrigger boolean default false) RETURNS VOID AS $$ -DECLARE - dtrange tstzrange; - edtrange tstzrange; - cdtrange tstzrange; - cedtrange tstzrange; - extent geometry; - collection text; -BEGIN - RAISE NOTICE 'Updating stats for %.', _partition; - EXECUTE format( - $q$ - SELECT - tstzrange(min(datetime), max(datetime),'[]'), - tstzrange(min(end_datetime), max(end_datetime), '[]') - FROM %I - $q$, - _partition - ) INTO dtrange, edtrange; - EXECUTE format('ANALYZE %I;', _partition); - extent := st_estimatedextent('pgstac', _partition, 'geometry'); - RAISE DEBUG 'Estimated Extent: %', extent; - INSERT INTO partition_stats (partition, dtrange, edtrange, spatial, last_updated) - SELECT _partition, dtrange, edtrange, extent, now() - ON CONFLICT (partition) DO - UPDATE SET - dtrange=EXCLUDED.dtrange, - edtrange=EXCLUDED.edtrange, - spatial=EXCLUDED.spatial, - last_updated=EXCLUDED.last_updated - ; - - SELECT - constraint_dtrange, constraint_edtrange, pv.collection - INTO cdtrange, cedtrange, collection - FROM partitions_view pv WHERE partition = _partition; - - RAISE NOTICE 'Checking if we need to modify constraints...'; - RAISE NOTICE 'cdtrange: % dtrange: % cedtrange: % edtrange: %',cdtrange, dtrange, cedtrange, edtrange; - IF - (cdtrange IS DISTINCT FROM dtrange OR edtrange IS DISTINCT FROM cedtrange) - AND NOT istrigger - THEN - RAISE NOTICE 'Modifying Constraints'; - RAISE NOTICE 'Existing % %', cdtrange, cedtrange; - RAISE NOTICE 'New % %', dtrange, edtrange; - PERFORM drop_table_constraints(_partition); - PERFORM create_table_constraints(_partition, dtrange, edtrange); - END IF; - REFRESH MATERIALIZED VIEW partitions; - REFRESH MATERIALIZED VIEW partition_steps; - RAISE NOTICE 'Checking if we need to update collection extents.'; - IF get_setting_bool('update_collection_extent') THEN - RAISE NOTICE 'updating collection extent for %', collection; - PERFORM run_or_queue(format($q$ - UPDATE collections - SET content = jsonb_set_lax( - content, - '{extent}'::text[], - collection_extent(%L, FALSE), - true, - 'use_json_null' - ) WHERE id=%L - ; - $q$, collection, collection)); - ELSE - RAISE NOTICE 'Not updating collection extent for %', collection; - END IF; - -END; -$$ LANGUAGE PLPGSQL STRICT SECURITY DEFINER; CREATE OR REPLACE FUNCTION partition_name( IN collection text, IN dt timestamptz, OUT partition_name text, OUT partition_range tstzrange) AS $$ @@ -4393,107 +4209,82 @@ END; $$ LANGUAGE PLPGSQL STABLE; -CREATE OR REPLACE FUNCTION drop_table_constraints(t text) RETURNS text AS $$ -DECLARE - q text; -BEGIN - IF NOT EXISTS (SELECT 1 FROM partitions_view WHERE partition=t) THEN - RETURN NULL; - END IF; - FOR q IN SELECT FORMAT( - $q$ - ALTER TABLE %I DROP CONSTRAINT IF EXISTS %I; - $q$, - t, - conname - ) FROM pg_constraint - WHERE conrelid=t::regclass::oid AND contype='c' - LOOP - EXECUTE q; - END LOOP; - RETURN t; -END; -$$ LANGUAGE PLPGSQL SECURITY DEFINER; +-- Partitions carry no _dt CHECK constraints; pruning is partition_stats-driven. get_tstz_constraint / +-- constraint_tstzrange remain so the views' constraint_* columns resolve to inf_range when a partition +-- has no datetime constraint. + -CREATE OR REPLACE FUNCTION create_table_constraints(t text, _dtrange tstzrange, _edtrange tstzrange) RETURNS text AS $$ +-- widen_partition_stats: widen a partition's stored envelope to cover a batch. A no-op (snapshot read, no +-- write, no lock) when the envelope already covers the batch; otherwise widens and sets dirty=true: +-- * datetime : sub-partitioned collections widen to the partition's datetime bound (covers anything +-- that can land there); a NULL-partition_trunc partition pads the batch range by +-- partition_stats_widen_buffer (default 1 month) each side. +-- * end_datetime: the datetime target extended by the batch's max (end_datetime - datetime) tail. +-- * spatial : NULL means "always a search candidate"; a spatial miss resets spatial to NULL until +-- the tightener computes the real extent. +-- Requires the partition_stats row to exist (check_partition seeds it); raises if it does not. +CREATE OR REPLACE FUNCTION widen_partition_stats( + _partition text, + _dtrange tstzrange, + _edtrange tstzrange, + _spatial geometry DEFAULT NULL, + _constraint_dtrange tstzrange DEFAULT NULL +) RETURNS void AS $$ DECLARE - q text; + cur RECORD; + is_unbounded boolean; + tail interval; + buf interval := COALESCE(get_setting('partition_stats_widen_buffer'), '1 month')::interval; + target_dtrange tstzrange; + target_edtrange tstzrange; + new_dtrange tstzrange; + new_edtrange tstzrange; + new_spatial geometry; + dt_covered boolean; + edt_covered boolean; + spatial_covered boolean; BEGIN - IF NOT EXISTS (SELECT 1 FROM partitions_view WHERE partition=t) THEN - RETURN NULL; - END IF; - RAISE NOTICE 'Creating Table Constraints for % % %', t, _dtrange, _edtrange; - IF _dtrange = 'empty' AND _edtrange = 'empty' THEN - q :=format( - $q$ - DO $block$ - BEGIN - ALTER TABLE %I DROP CONSTRAINT IF EXISTS %I; - ALTER TABLE %I - ADD CONSTRAINT %I - CHECK (((datetime IS NULL) AND (end_datetime IS NULL))) NOT VALID - ; - ALTER TABLE %I - VALIDATE CONSTRAINT %I - ; - - - - EXCEPTION WHEN others THEN - RAISE WARNING '%%, Issue Altering Constraints. Please run update_partition_stats(%I)', SQLERRM USING ERRCODE = SQLSTATE; - END; - $block$; - $q$, - t, - format('%s_dt', t), - t, - format('%s_dt', t), - t, - format('%s_dt', t), - t - ); + SELECT dtrange, edtrange, spatial + INTO cur + FROM partition_stats WHERE partition = _partition; + IF NOT FOUND THEN + RAISE EXCEPTION 'partition_stats row for % does not exist; call check_partition first', _partition; + END IF; + + dt_covered := COALESCE(cur.dtrange @> _dtrange, false); + edt_covered := COALESCE(cur.edtrange @> _edtrange, false); + spatial_covered := cur.spatial IS NULL OR _spatial IS NULL OR ST_Covers(cur.spatial, _spatial); + IF dt_covered AND edt_covered AND spatial_covered THEN + RETURN; -- already covered: no write, no lock + END IF; + + -- Clamp the widen to the partition's structural bound (_constraint_dtrange). A sub-partitioned + -- (month/year) partition's bound is finite: cover the whole partition so dtrange never misses and + -- never spills into a neighbour. An unbounded partition (NULL-partition_trunc or NULL bound) has + -- nothing to clamp to, so pad the batch range by the widen buffer each side. + is_unbounded := _constraint_dtrange IS NULL + OR (lower(_constraint_dtrange) = '-infinity'::timestamptz + AND upper(_constraint_dtrange) = 'infinity'::timestamptz); + tail := GREATEST(upper(_edtrange) - upper(_dtrange), '0'::interval); + IF is_unbounded THEN + target_dtrange := tstzrange(lower(_dtrange) - buf, upper(_dtrange) + buf, '[]'); + target_edtrange := tstzrange(lower(_dtrange) - buf, upper(_edtrange) + buf, '[]'); ELSE - q :=format( - $q$ - DO $block$ - BEGIN - - ALTER TABLE %I DROP CONSTRAINT IF EXISTS %I; - ALTER TABLE %I - ADD CONSTRAINT %I - CHECK ( - (datetime >= %L) - AND (datetime <= %L) - AND (end_datetime >= %L) - AND (end_datetime <= %L) - ) NOT VALID - ; - ALTER TABLE %I - VALIDATE CONSTRAINT %I - ; - - - - EXCEPTION WHEN others THEN - RAISE WARNING '%%, Issue Altering Constraints. Please run update_partition_stats(%I)', SQLERRM USING ERRCODE = SQLSTATE; - END; - $block$; - $q$, - t, - format('%s_dt', t), - t, - format('%s_dt', t), - lower(_dtrange), - upper(_dtrange), - lower(_edtrange), - upper(_edtrange), - t, - format('%s_dt', t), - t - ); + target_dtrange := _constraint_dtrange; + target_edtrange := tstzrange(lower(_constraint_dtrange), upper(_constraint_dtrange) + tail, '[]'); END IF; - PERFORM run_or_queue(q); - RETURN t; + + new_dtrange := range_merge(COALESCE(cur.dtrange, target_dtrange), target_dtrange); + new_edtrange := range_merge(COALESCE(cur.edtrange, target_edtrange), target_edtrange); + new_spatial := CASE WHEN spatial_covered THEN cur.spatial ELSE NULL END; + + UPDATE partition_stats + SET dtrange = new_dtrange, + edtrange = new_edtrange, + spatial = new_spatial, + dirty = true, + last_updated = now() + WHERE partition = _partition; END; $$ LANGUAGE PLPGSQL SECURITY DEFINER; @@ -4501,18 +4292,14 @@ $$ LANGUAGE PLPGSQL SECURITY DEFINER; CREATE OR REPLACE FUNCTION check_partition( _collection text, _dtrange tstzrange, - _edtrange tstzrange + _edtrange tstzrange, + _spatial geometry DEFAULT NULL ) RETURNS text AS $$ DECLARE c RECORD; - pm RECORD; _partition_name text; - _partition_dtrange tstzrange; - _constraint_dtrange tstzrange; - _constraint_edtrange tstzrange; - q text; - deferrable_q text; - err_context text; + _parent_name text; + _partition_range tstzrange; BEGIN SELECT * INTO c FROM pgstac.collections WHERE id=_collection; IF NOT FOUND THEN @@ -4520,115 +4307,99 @@ BEGIN END IF; IF c.partition_trunc IS NOT NULL THEN - _partition_dtrange := tstzrange( + _partition_range := tstzrange( date_trunc(c.partition_trunc, lower(_dtrange)), date_trunc(c.partition_trunc, lower(_dtrange)) + (concat('1 ', c.partition_trunc))::interval, '[)' ); ELSE - _partition_dtrange := '[-infinity, infinity]'::tstzrange; + _partition_range := '[-infinity, infinity]'::tstzrange; END IF; - IF NOT _partition_dtrange @> _dtrange THEN - RAISE EXCEPTION 'dtrange % is greater than the partition size % for collection %', _dtrange, c.partition_trunc, _collection; + IF NOT _partition_range @> _dtrange THEN + RAISE EXCEPTION 'dtrange % spans more than the % partition window for collection %', _dtrange, c.partition_trunc, _collection; END IF; - + _parent_name := format('_items_%s', c.key); IF c.partition_trunc = 'year' THEN - _partition_name := format('_items_%s_%s', c.key, to_char(lower(_partition_dtrange),'YYYY')); + _partition_name := format('%s_%s', _parent_name, to_char(lower(_partition_range),'YYYY')); ELSIF c.partition_trunc = 'month' THEN - _partition_name := format('_items_%s_%s', c.key, to_char(lower(_partition_dtrange),'YYYYMM')); + _partition_name := format('%s_%s', _parent_name, to_char(lower(_partition_range),'YYYYMM')); ELSE - _partition_name := format('_items_%s', c.key); - END IF; - - SELECT * INTO pm FROM partition_sys_meta WHERE collection=_collection AND partition_dtrange @> _dtrange; - IF FOUND THEN - RAISE NOTICE '% % %', _edtrange, _dtrange, pm; - _constraint_edtrange := - tstzrange( - least( - lower(_edtrange), - nullif(lower(pm.constraint_edtrange), '-infinity') - ), - greatest( - upper(_edtrange), - nullif(upper(pm.constraint_edtrange), 'infinity') - ), - '[]' - ); - _constraint_dtrange := - tstzrange( - least( - lower(_dtrange), - nullif(lower(pm.constraint_dtrange), '-infinity') - ), - greatest( - upper(_dtrange), - nullif(upper(pm.constraint_dtrange), 'infinity') - ), - '[]' + _partition_name := _parent_name; + END IF; + + -- Create the collection-level PARENT partition (_items_) first, for sub-partitioned collections. + -- It is shared across every child window, so concurrent setup of different children would race on its + -- CREATE TABLE. Guard it with a parent-scoped advisory lock, taken before any child lock and only when + -- the parent is missing: parent-before-child is one lock order (parent < child) that can't deadlock + -- with ensure_partitions' sorted child locks, and skipping it once the parent exists keeps steady-state + -- ingest off the parent lock. + IF c.partition_trunc IS NOT NULL AND to_regclass(format('pgstac.%I', _parent_name)) IS NULL THEN + PERFORM pg_advisory_xact_lock(hashtext('pgstac.check_partition'), hashtext(_parent_name)); + IF to_regclass(format('pgstac.%I', _parent_name)) IS NULL THEN + EXECUTE format( + 'CREATE TABLE IF NOT EXISTS %I partition OF items FOR VALUES IN (%L) PARTITION BY RANGE (datetime)', + _parent_name, _collection ); - - IF pm.constraint_edtrange @> _edtrange AND pm.constraint_dtrange @> _dtrange THEN - RETURN pm.partition; - ELSE - PERFORM drop_table_constraints(_partition_name); END IF; - ELSE - _constraint_edtrange := _edtrange; - _constraint_dtrange := _dtrange; END IF; - RAISE NOTICE 'EXISTING CONSTRAINTS % %, NEW % %', pm.constraint_dtrange, pm.constraint_edtrange, _constraint_dtrange, _constraint_edtrange; - RAISE NOTICE 'Creating partition % %', _partition_name, _partition_dtrange; - IF c.partition_trunc IS NULL THEN - q := format( - $q$ - CREATE TABLE IF NOT EXISTS %I partition OF items FOR VALUES IN (%L); - CREATE UNIQUE INDEX IF NOT EXISTS %I ON %I (id); - GRANT ALL ON %I to pgstac_ingest; - $q$, - _partition_name, - _collection, - concat(_partition_name,'_pk'), - _partition_name, - _partition_name - ); - ELSE - q := format( - $q$ - CREATE TABLE IF NOT EXISTS %I partition OF items FOR VALUES IN (%L) PARTITION BY RANGE (datetime); - CREATE TABLE IF NOT EXISTS %I partition OF %I FOR VALUES FROM (%L) TO (%L); - CREATE UNIQUE INDEX IF NOT EXISTS %I ON %I (id); - GRANT ALL ON %I TO pgstac_ingest; - $q$, - format('_items_%s', c.key), - _collection, - _partition_name, - format('_items_%s', c.key), - lower(_partition_dtrange), - upper(_partition_dtrange), - format('%s_pk', _partition_name), - _partition_name, - _partition_name - ); + + -- Serialize concurrent setup of THIS (leaf) partition with a partition-scoped advisory lock. Different + -- partitions hash to different keys, so this never serializes unrelated ingest; ensure_partitions + -- acquires these in sorted order, so concurrent multi-partition batches cannot deadlock. The lock + -- releases at commit (setup is fast + idempotent). + PERFORM pg_advisory_xact_lock(hashtext('pgstac.check_partition'), hashtext(_partition_name)); + + -- Create the leaf partition if missing. Parent-inherited indexes only (id PK here; datetime/geometry + -- come from the items parent); no CHECK constraints; queryable indexes are deferred via + -- indexes_pending. A SELECT grant lets read/ingest query it; writes reach it only through the + -- SECURITY DEFINER write functions (the privilege wall in 998_idempotent_post). + -- Skip the DDL when it already exists: re-running CREATE/GRANT takes a relation lock that deadlocks + -- with concurrent INSERTs. The existence check is race-safe under the advisory lock. + IF to_regclass(format('pgstac.%I', _partition_name)) IS NULL THEN + BEGIN + IF c.partition_trunc IS NULL THEN + EXECUTE format( + $q$ + CREATE TABLE IF NOT EXISTS %I partition OF items FOR VALUES IN (%L); + CREATE UNIQUE INDEX IF NOT EXISTS %I ON %I (id); + GRANT SELECT ON %I TO pgstac_read, pgstac_ingest; + $q$, + _partition_name, _collection, + concat(_partition_name, '_pk'), _partition_name, + _partition_name + ); + ELSE + EXECUTE format( + $q$ + CREATE TABLE IF NOT EXISTS %I partition OF %I FOR VALUES FROM (%L) TO (%L); + CREATE UNIQUE INDEX IF NOT EXISTS %I ON %I (id); + GRANT SELECT ON %I TO pgstac_read, pgstac_ingest; + $q$, + _partition_name, _parent_name, lower(_partition_range), upper(_partition_range), + concat(_partition_name, '_pk'), _partition_name, + _partition_name + ); + END IF; + EXCEPTION + -- A concurrent creator that finished between our checks: benign, it exists now. + WHEN duplicate_table THEN + RAISE DEBUG 'Partition % already exists.', _partition_name; + -- Do NOT swallow other errors: a failed creation must propagate so the caller never goes on to + -- write a partition that does not exist (invariant: check_partition succeeds before use). + END; END IF; - BEGIN - EXECUTE q; - EXCEPTION - WHEN duplicate_table THEN - RAISE NOTICE 'Partition % already exists.', _partition_name; - WHEN others THEN - GET STACKED DIAGNOSTICS err_context = PG_EXCEPTION_CONTEXT; - RAISE INFO 'Error Name:%',SQLERRM; - RAISE INFO 'Error State:%', SQLSTATE; - RAISE INFO 'Error Context:%', err_context; - END; - PERFORM maintain_partitions(_partition_name); - PERFORM update_partition_stats_q(_partition_name, true); - REFRESH MATERIALIZED VIEW partitions; - REFRESH MATERIALIZED VIEW partition_steps; + -- Seed the partition_stats row (collection only — the data ranges start NULL), then cover this batch + -- via the shared widen guard, which fills dtrange/edtrange. ON CONFLICT keeps an existing row (so a + -- sweep that cleared indexes_pending/dirty is not reset on a later check_partition for the same + -- partition). + INSERT INTO partition_stats (partition, collection, dirty, indexes_pending) + VALUES (_partition_name, _collection, true, true) + ON CONFLICT (partition) DO NOTHING; + PERFORM widen_partition_stats(_partition_name, _dtrange, _edtrange, _spatial, _partition_range); + RETURN _partition_name; END; $$ LANGUAGE PLPGSQL SECURITY DEFINER; @@ -4709,6 +4480,219 @@ UPDATE ON collections FOR EACH ROW EXECUTE FUNCTION collections_trigger_func(); -- END FRAGMENT: 003b_partitions.sql +-- BEGIN FRAGMENT: 003c_ingest.sql +-- --------------------------------------------------------------------------- +-- Rust-first ingest: SECURITY DEFINER staging + flush + fragment helpers. +-- +-- These are the server-side seam for the Rust loader's binary-COPY path. The connecting role never +-- writes the real tables directly; it binary-COPYs fully-dehydrated rows into a session-local TEMP +-- table (make_binary_staging) and calls flush_items_staging_binary, which is the only thing that writes +-- `items`. Partition existence + stats (extent + n) and registry/fragment coverage are established BEFORE +-- the load in their own transactions (prepare_partition_for_load / ensure_fragments); the flush writes only +-- `items` and never touches partition_stats. +-- --------------------------------------------------------------------------- + +-- ensure_fragments: upsert per-collection fragment payloads and return, for each input position, the +-- item_fragments id the Rust loader stamps into items.fragment_id. +-- +-- The caller (the Rust loader) deduplicates fragments locally and sends only the DISTINCT set (one input +-- element per unique fragment), then maps each item -> its fragment via the returned `ord`. Sending one +-- fragment per item would ship millions of duplicates when a collection has only a handful of distinct +-- fragments. The function is nonetheless dup-safe (a DISTINCT ON guards the insert), so a caller that +-- passes duplicates still gets the correct id for every position. +-- +-- Each input element is {"content": , "links_template": }; the canonical +-- hash matches pgstac_hash_fragment (so it dedups identically to the SQL ingest path). ON CONFLICT keeps +-- existing rows; pre-existing ids come from item_fragments (snapshot), new ids from the INSERT's RETURNING. +CREATE OR REPLACE FUNCTION ensure_fragments( + _collection text, + _fragments jsonb[] +) RETURNS TABLE (ord int, frag_id bigint) AS $$ + WITH input AS ( + SELECT + o::int AS ord, + pgstac_hash_fragment( + jsonb_strip_nulls(jsonb_build_object( + 'content', NULLIF(f->'content', '{}'::jsonb), + 'links_template', f->'links_template' + )) + ) AS hash, + COALESCE(f->'content', '{}'::jsonb) AS content, + f->'links_template' AS links_template + FROM unnest(_fragments) WITH ORDINALITY AS u(f, o) + ), + distinct_hashes AS ( + SELECT DISTINCT ON (hash) hash, content, links_template + FROM input + ORDER BY hash + ), + upserted AS ( + INSERT INTO item_fragments (collection, hash, content, links_template) + SELECT _collection, hash, content, links_template FROM distinct_hashes + -- DO UPDATE (a no-op write) rather than DO NOTHING: it locks + RETURNS the conflicting row even when + -- a CONCURRENT transaction committed the same (collection, hash) mid-statement. DO NOTHING returns + -- nothing on conflict, and the old "UNION the existing rows via a separate SELECT" ran on the + -- statement-start snapshot, so it could MISS such a concurrent insert -> that hash dropped out of the + -- final join and the item was stamped with NO fragment_id. DO UPDATE returns every distinct hash + -- exactly once (newly inserted or pre-existing), making the stamp concurrency-safe. + ON CONFLICT (collection, hash) DO UPDATE SET content = item_fragments.content + RETURNING id, hash + ) + SELECT input.ord, upserted.id + FROM input JOIN upserted ON input.hash = upserted.hash + ORDER BY input.ord; +$$ LANGUAGE SQL SECURITY DEFINER; + + +-- ensure_partitions: create + widen every partition a batch will land in, BEFORE the load (widen-now). +-- The Rust loader passes parallel arrays of (collection, datetime, end_datetime) — one element per item. +-- Grouping happens HERE, server-side, so the month/year truncation uses the server's date_trunc (correct +-- timezone) rather than re-deriving partition boundaries in the client. One call per batch replaces the +-- loader's per-item check_partition round trips: it groups by (collection, partition window) and calls +-- check_partition once per distinct partition with that group's datetime range. Runs as its own committed +-- statement before the load transaction, so partitions exist + their stats cover the batch by the time +-- the binary COPY + flush run. +CREATE OR REPLACE FUNCTION ensure_partitions( + _collections text[], + _datetimes timestamptz[], + _end_datetimes timestamptz[] +) RETURNS void AS $$ +DECLARE + r RECORD; +BEGIN + -- Call check_partition per distinct (collection, partition window) in a DETERMINISTIC order: each + -- check_partition takes that partition's advisory lock, so locking in a consistent order prevents two + -- concurrent multi-partition batches from advisory-deadlocking on the locks. (flush locks in the same + -- sorted order.) A PLPGSQL FOR loop guarantees the order; a set-returning SELECT would not. + FOR r IN + WITH t AS ( + SELECT + unnest(_collections) AS collection, + unnest(_datetimes) AS dt, + unnest(_end_datetimes) AS edt + ), + j AS ( + SELECT t.collection, t.dt, t.edt, c.partition_trunc + FROM t JOIN pgstac.collections c ON t.collection = c.id + ) + SELECT + collection, + tstzrange(min(dt), max(dt), '[]') AS dtrange, + tstzrange(min(edt), max(edt), '[]') AS edtrange + FROM j + GROUP BY collection, COALESCE(date_trunc(partition_trunc::text, dt), '-infinity'::timestamptz) + ORDER BY collection, COALESCE(date_trunc(partition_trunc::text, dt), '-infinity'::timestamptz) + LOOP + PERFORM pgstac.check_partition(r.collection, r.dtrange, r.edtrange); + END LOOP; +END; +$$ LANGUAGE PLPGSQL SECURITY DEFINER; + + +-- prepare_partition_for_load: per-partition metadata for the Rust direct / precheck load paths. ONE small +-- self-contained txn per partition, run BEFORE any COPY: create + widen the partition to cover this batch +-- (dt + edt + the real SPATIAL envelope, unlike ensure_partitions which passes NULL) and bump n, so +-- partition_stats is at least as wide as the data (golden rule) and search treats the partition as non-empty +-- before the data lands. Returns the partition name + the pre-load row count (n BEFORE the bump) so the +-- loader can choose its adaptive precheck path: empty -> skip the precheck; batch > n -> pull the partition's +-- (id,item_hash) to the client; n >= batch -> COPY the batch (id,hash) to a temp table + JOIN this partition. +CREATE OR REPLACE FUNCTION prepare_partition_for_load( + _collection text, + _dt_lo timestamptz, _dt_hi timestamptz, + _edt_lo timestamptz, _edt_hi timestamptz, + _xmin float8, _ymin float8, _xmax float8, _ymax float8, + _n_add bigint, + OUT partition_name text, + OUT pre_load_n bigint +) AS $$ +BEGIN + partition_name := pgstac.check_partition( + _collection, + tstzrange(_dt_lo, _dt_hi, '[]'), + tstzrange(_edt_lo, _edt_hi, '[]'), + st_setsrid(st_makeenvelope(_xmin, _ymin, _xmax, _ymax), 4326) + ); + SELECT COALESCE(n, 0) INTO pre_load_n + FROM pgstac.partition_stats WHERE partition = partition_name; + -- Over-estimating n is the safe direction; the async tightener computes the exact count + extent off the + -- hot path. Single-row atomic UPDATE, so concurrent loads into the same partition serialize on the row. + UPDATE pgstac.partition_stats + SET n = COALESCE(n, 0) + _n_add, dirty = true, last_updated = now() + WHERE partition = partition_name; +END; +$$ LANGUAGE PLPGSQL SECURITY DEFINER; + + +-- make_binary_staging: create a session-local TEMP staging table shaped exactly like `items` +-- (ON COMMIT DROP) and grant INSERT to pgstac_ingest so the connecting role can binary-COPY into it. +-- Created inside this SECURITY DEFINER function, so the temp table is owned by pgstac_admin (which also +-- owns `items`), letting flush_items_staging_binary read it. Returns the generated table name. +CREATE OR REPLACE FUNCTION make_binary_staging() RETURNS text AS $$ +DECLARE + _name text := format('_staging_%s', replace(gen_random_uuid()::text, '-', '')); +BEGIN + EXECUTE format('CREATE TEMP TABLE %I (LIKE items) ON COMMIT DROP', _name); + EXECUTE format('GRANT INSERT ON pg_temp.%I TO pgstac_ingest', _name); + RETURN _name; +END; +$$ LANGUAGE PLPGSQL SECURITY DEFINER; + + +-- flush_items_staging_binary: move fully-dehydrated rows from a TEMP staging table into `items` with the +-- conflict policy. partition_stats (extent + n + dirty) is set entirely by prepare_partition_for_load in +-- the preflight, so this writes only `items`: +-- ignore -> ON CONFLICT DO NOTHING (idempotent; orphans the old row on a cross-partition move) +-- upsert -> delete a changed row IN THE PARTITION IT ROUTES TO (window-pruned), then insert. A +-- datetime change that stays in the partition is applied; one that moves the item to a +-- different partition orphans the old row (use 'delsert'). +-- delsert -> delete a changed row CROSS-partition (by collection+id, move-safe), then insert +-- error -> plain INSERT (raises on any duplicate) +-- Returns the number of rows inserted. +CREATE OR REPLACE FUNCTION flush_items_staging_binary( + _staging text, + _policy text DEFAULT 'ignore' +) RETURNS bigint AS $$ +DECLARE + nrows bigint; +BEGIN + IF _policy = 'ignore' THEN + EXECUTE format('INSERT INTO items SELECT * FROM %1$I ON CONFLICT DO NOTHING', _staging); + ELSIF _policy = 'error' THEN + EXECUTE format('INSERT INTO items SELECT * FROM %1$I', _staging); + ELSIF _policy = 'upsert' THEN + -- Fast SAME-partition upsert: delete a changed row only in the partition the incoming row routes to. + -- The window range [date_trunc(trunc, s.datetime), + 1 trunc) is derived per staged row, so the + -- planner runtime-prunes `items` to that one partition (no cross-partition scan). A datetime change + -- within the partition is applied; one that MOVES the item to another partition isn't seen here, so + -- the old row orphans (use 'delsert'). NULL partition_trunc => the collection's single partition. + EXECUTE format($q$ + DELETE FROM items i USING %1$I s JOIN collections c ON c.id = s.collection + WHERE i.collection = s.collection AND i.id = s.id + AND i.datetime >= (CASE WHEN c.partition_trunc IS NULL THEN '-infinity'::timestamptz + ELSE date_trunc(c.partition_trunc::text, s.datetime) END) + AND i.datetime < (CASE WHEN c.partition_trunc IS NULL THEN 'infinity'::timestamptz + ELSE date_trunc(c.partition_trunc::text, s.datetime) + + ('1 ' || c.partition_trunc::text)::interval END) + AND ( %2$s ) + $q$, _staging, items_content_distinct_sql('i', 's')); + EXECUTE format('INSERT INTO items SELECT * FROM %1$I ON CONFLICT DO NOTHING', _staging); + ELSIF _policy = 'delsert' THEN + -- Move-safe CROSS-partition upsert: delete the old row wherever it lives (by collection+id, no + -- datetime bound, so it probes every partition), then insert. + EXECUTE format($q$ + DELETE FROM items i USING %1$I s + WHERE i.id = s.id AND i.collection = s.collection AND ( %2$s ) + $q$, _staging, items_content_distinct_sql('i', 's')); + EXECUTE format('INSERT INTO items SELECT * FROM %1$I ON CONFLICT DO NOTHING', _staging); + ELSE + RAISE EXCEPTION 'unknown conflict policy % (expected ignore | upsert | delsert | error)', _policy; + END IF; + GET DIAGNOSTICS nrows = ROW_COUNT; + RETURN nrows; +END; +$$ LANGUAGE PLPGSQL SECURITY DEFINER; +-- END FRAGMENT: 003c_ingest.sql + -- BEGIN FRAGMENT: 004_search.sql -- Search hashing @@ -4859,7 +4843,7 @@ BEGIN search := register_search(search); RETURN QUERY SELECT search.hash, search.metadata; END; -$$ LANGUAGE PLPGSQL SECURITY DEFINER; +$$ LANGUAGE PLPGSQL SECURITY DEFINER SET search_path TO pgstac, public; -- search_fromhash: lookup a cached search by its hash. CREATE OR REPLACE FUNCTION search_fromhash(_hash text) RETURNS searches AS $$ @@ -4884,14 +4868,18 @@ BEGIN search := register_search(search); RETURN QUERY SELECT search.hash, search.metadata; END; -$$ LANGUAGE PLPGSQL SECURITY DEFINER; +$$ LANGUAGE PLPGSQL SECURITY DEFINER SET search_path TO pgstac, public; --- field_included: canonical STAC fields include/exclude decision. +-- field_included: STAC fields include/exclude decision over text[] arrays (the array-form used by +-- the column projector fields_to_itemcols). include_field() is the equivalent over a fields jsonb +-- (used by content_hydrate); both apply the same rule: exclude wins, then an explicit include list +-- restricts to its members, otherwise everything is included. CREATE OR REPLACE FUNCTION field_included(_field text, _includes text[], _excludes text[]) RETURNS boolean LANGUAGE sql IMMUTABLE AS $$ - SELECT CASE WHEN array_length(_includes, 1) IS NOT NULL - THEN _field = ANY(_includes) - ELSE NOT (_field = ANY(_excludes)) END; + SELECT CASE + WHEN _field = ANY(_excludes) THEN false + WHEN array_length(_includes, 1) IS NOT NULL THEN _field = ANY(_includes) + ELSE true END; $$; -- needs_fragment: determine whether satisfying the requested fields for the @@ -4930,8 +4918,16 @@ $$ LANGUAGE PLPGSQL STABLE PARALLEL SAFE; -- fields_to_itemcols: produce a SELECT list of item columns in attnum order. -- Heavy columns (geometry, bbox, assets, links, properties, extra) are emitted as --- NULL::type when their controlling field is excluded/not-included. -CREATE OR REPLACE FUNCTION fields_to_itemcols(fields jsonb DEFAULT '{}'::jsonb) RETURNS text AS $$ +-- NULL::type AS when their controlling field is excluded/not-included: the +-- column is kept (named) so the result schema is stable across any `fields` value +-- -- a client preparing this query once can step bands without re-binding columns -- +-- while the heavy value itself is never transferred. +-- When _skip_fragment is true (the caller determined the requested fields are +-- satisfiable from item columns alone, via needs_fragment), fragment_id is emitted +-- as NULL too: a client driving this query then sees no fragment_id and skips the +-- per-row item_fragments lookup entirely -- the fragment-skip is carried in the +-- projection itself, not a separate flag. +CREATE OR REPLACE FUNCTION fields_to_itemcols(fields jsonb DEFAULT '{}'::jsonb, _skip_fragment boolean DEFAULT false) RETURNS text AS $$ DECLARE includes text[] := ARRAY(SELECT jsonb_array_elements_text(fields->'include')); excludes text[] := ARRAY(SELECT jsonb_array_elements_text(fields->'exclude')); @@ -4939,11 +4935,13 @@ DECLARE BEGIN SELECT string_agg( CASE + WHEN a.attname = 'fragment_id' AND _skip_fragment + THEN format('NULL::%s AS %I', format_type(a.atttypid, a.atttypmod), a.attname) WHEN a.attname = ANY (ARRAY['geometry','bbox','assets','links','link_hrefs', 'extra','properties','stac_version','stac_extensions']) AND NOT field_included(CASE WHEN a.attname = 'link_hrefs' THEN 'links' ELSE a.attname END, includes, excludes) - THEN format('NULL::%s', format_type(a.atttypid, a.atttypmod)) + THEN format('NULL::%s AS %I', format_type(a.atttypid, a.atttypmod), a.attname) ELSE format('i.%I', a.attname) END, ', ' ORDER BY a.attnum) INTO cols @@ -4984,12 +4982,15 @@ DECLARE band_cap_months CONSTANT int := 18; page_rows items[] := '{}'::items[]; chunk_rows items[]; target int := _limit + 1; got int := 0; got_band int; - is_asc boolean; cursor_ts timestamptz; band record; + is_asc boolean; band record; band_target numeric; obs_sel numeric; band_where text; guard int := 0; cum_scanned bigint := 0; proj_expr text; mo interval := interval '1 month'; cursor_idx int; BEGIN + -- The requested STAC `fields` live in the search request; honor them over the (defaulted) + -- parameter so include/exclude projection is actually applied for search(). + _fields := coalesce(_search->'fields', _fields, '{}'::jsonb); _cql2 := search_to_cql2(_search); _where := cql2_query(_cql2); IF _where IS NULL OR btrim(_where) = '' THEN _where := ' TRUE '; END IF; @@ -5027,20 +5028,27 @@ BEGIN clamped_where := concat_ws(' AND ', clamp, full_where); IF clamped_where IS NULL OR btrim(clamped_where) = '' THEN clamped_where := 'TRUE'; END IF; - IF needs_fragment(coalesce(_fields, '{}'::jsonb), bnds.collections) THEN - proj_expr := format('content_hydrate(i, %L::jsonb, (SELECT f FROM item_fragments f WHERE f.id = i.fragment_id))', _fields); + -- Per-row projection. When the requested fields can be satisfied from item columns alone + -- (needs_fragment, evaluated once for the whole query), tell content_hydrate to skip the shared + -- item_fragments lookup entirely; otherwise it fetches and merges the fragment per row. + IF needs_fragment(_fields, bnds.collections) THEN + proj_expr := format('content_hydrate(i, %L::jsonb)', _fields); ELSE - proj_expr := format('content_hydrate(i, %L::jsonb, NULL, true)', _fields); + proj_expr := format('content_hydrate(i, %L::jsonb, true)', _fields); END IF; IF datetime_leading AND array_length(bnds.months, 1) IS NOT NULL THEN - cursor_ts := CASE WHEN is_asc THEN bnds.months[1] ELSE bnds.months[array_length(bnds.months, 1)] + mo END; - cursor_idx := 1; + -- Start the band walk at the leading edge for the sort direction: oldest month for ascending, + -- most recent month for descending. (A descending walk that started at the oldest band would + -- collect the oldest items and never reach the recent months.) + cursor_idx := CASE WHEN is_asc THEN 1 ELSE array_length(bnds.months, 1) END; band_target := target * band_margin; WHILE got < target AND guard < 80 LOOP guard := guard + 1; - SELECT * INTO band FROM next_band(bnds.counts, cursor_idx, band_target, band_cap_months); - EXIT WHEN band.done; + SELECT * INTO band FROM next_band(bnds.counts, cursor_idx, band_target, band_cap_months, NOT is_asc); + -- a valid band must be processed even when next_band also flags done (it consumed the + -- last bucket); only stop when there is no band at all. + EXIT WHEN band.band_start_idx IS NULL; band_where := format('i.datetime >= %L AND i.datetime < %L AND (%s)', bnds.months[band.band_start_idx], bnds.months[band.band_end_idx] + mo, full_where); EXECUTE format('SELECT array_agg(i ORDER BY %s) FROM (SELECT * FROM items i WHERE %s ORDER BY %s LIMIT %s) i', @@ -5103,12 +5111,20 @@ $$; -- search: FeatureCollection API wrapper CREATE OR REPLACE FUNCTION search(_search jsonb DEFAULT '{}'::jsonb) RETURNS json AS $$ DECLARE + -- caller-provided limit/token come from the search body; default to the configured page size. + _limit int := coalesce((_search->>'limit')::int, + nullif(get_setting('default_page_size', _search->'conf'), '')::int, 10); + _token text := _search->>'token'; + -- The keyset is the token minus its next/prev prefix; an empty token means no keyset (first + -- page). A non-empty keyset that does not decode raises in keyset_decode downstream. + keyset text := nullif(regexp_replace(coalesce(_token,''), '^(next|prev):', ''), ''); + is_prev boolean := (_token LIKE 'prev:%') AND keyset IS NOT NULL; pg record; burl text := rtrim(coalesce(base_url(_search->'conf'), ''), '/'); links jsonb := '[]'::jsonb; out json; BEGIN - SELECT * INTO pg FROM search_page(_search); + SELECT * INTO pg FROM search_page(_search, _limit, keyset, is_prev); links := links || jsonb_build_object('rel','root','type','application/json','href', burl) || jsonb_build_object('rel','self','type','application/json','href',burl||'/search'); @@ -5149,7 +5165,9 @@ CREATE OR REPLACE FUNCTION search_plan( OUT max_datetime timestamptz, OUT max_count bigint, OUT lead_desc boolean, - OUT ctx_query text + OUT ctx_query text, + OUT datetime_leading boolean, + OUT context_count bigint ) RETURNS record LANGUAGE plpgsql VOLATILE SECURITY DEFINER SET search_path TO pgstac, public AS $$ DECLARE _cql2 jsonb := search_to_cql2(_search); @@ -5157,11 +5175,10 @@ DECLARE is_prev boolean := _token LIKE 'prev:%'; keyset text := nullif(regexp_replace(coalesce(_token, ''), '^(next|prev):', ''), ''); keyset_w text; full_where text; orderby_str text; collist text; - lead_field text; eff_dir text; datetime_leading boolean; _env pred_envelope; _hash text; + lead_field text; eff_dir text; _env pred_envelope; _hash text; bnds record; coll_clamp text := ''; clamp text; BEGIN IF _where IS NULL OR btrim(_where) = '' THEN _where := ' TRUE '; END IF; - collist := fields_to_itemcols(coalesce(_search->'fields', '{}'::jsonb)); orderby_str := keyset_orderby(_search, is_prev); SELECT (array_agg(field ORDER BY ord))[1], (array_agg(CASE WHEN is_prev THEN (CASE dir WHEN 'ASC' THEN 'DESC' ELSE 'ASC' END) ELSE dir END ORDER BY ord))[1] @@ -5187,15 +5204,32 @@ BEGIN coll_clamp := format('i.collection = ANY(%L::text[]) AND ', bnds.collections); END IF; + -- Build the projection now that the resolved collections are known: when the requested fields are + -- satisfiable from item columns alone (needs_fragment is false), fragment_id is nulled in the + -- projection so a client driving this query never looks up item_fragments -- the fragment-skip + -- rides in the returned query's SELECT list, not a separate flag. + collist := fields_to_itemcols( + coalesce(_search->'fields', '{}'::jsonb), + NOT needs_fragment(coalesce(_search->'fields', '{}'::jsonb), bnds.collections)); + IF datetime_leading AND array_length(bnds.months, 1) IS NOT NULL THEN + -- collections baked in as a literal so the query is parameterized only by $1 (band low), + -- $2 (band high), $3 (limit): the client prepares it once and binds per histogram band. query := format( - 'SELECT %s FROM items i WHERE i.collection = ANY($4) AND i.datetime >= $1 AND i.datetime < $2 AND (%s) ORDER BY %s LIMIT $3', - collist, full_where, orderby_str); - histogram := bnds.histogram; + 'SELECT %s FROM items i WHERE %si.datetime >= $1 AND i.datetime < $2 AND (%s) ORDER BY %s LIMIT $3', + collist, coll_clamp, full_where, orderby_str); + -- serialize the per-month histogram (months[]/counts[]) to jsonb for the streaming client + histogram := ( + SELECT jsonb_agg(jsonb_build_object('m', m, 'n', n) ORDER BY ord) + FROM unnest(bnds.months, bnds.counts) WITH ORDINALITY AS h(m, n, ord) + ); ELSE DECLARE dt_clamp text := ''; BEGIN IF array_length(bnds.months, 1) IS NOT NULL THEN - dt_clamp := format('i.datetime >= %L AND i.datetime <= %L AND ', bnds.months[1], bnds.months[array_length(bnds.months, 1)]); + -- months[] holds month *starts*; the last bucket spans [start, start + 1 month), so the + -- upper bound must be exclusive of the next month, not <= the last month's start (which + -- would drop items dated after the start of the final month). + dt_clamp := format('i.datetime >= %L AND i.datetime < %L AND ', bnds.months[1], bnds.months[array_length(bnds.months, 1)] + interval '1 month'); END IF; query := format( 'SELECT %s FROM items i WHERE %s%s(%s) ORDER BY %s LIMIT $1', @@ -5211,10 +5245,19 @@ BEGIN s.orderby := keyset_orderby(_search); s.lastused := now(); s.usecount := 1; PERFORM register_search(s); END; + -- Inline the cached count when stats are fresh (same rule as where_stats), so the client + -- can skip ctx_query on a cache hit. NULL => miss/stale => client races ctx_query. + SELECT s2.context_count INTO context_count + FROM searches s2 + WHERE s2.hash = _hash + AND s2.statslastupdated IS NOT NULL + AND s2.context_count IS NOT NULL + AND now() - s2.statslastupdated <= context_stats_ttl(_search->'conf'); ctx_query := format('SELECT (where_stats(%L, %L, false, %L, %L)).context_count', _hash, _where, _search->'conf', clamp); ELSE ctx_query := NULL; + context_count := NULL; END IF; END; $$; @@ -5245,28 +5288,6 @@ SELECT FROM collections; -CREATE OR REPLACE FUNCTION collection_search_matched( - IN _search jsonb DEFAULT '{}'::jsonb, - OUT matched bigint -) RETURNS bigint AS $$ -DECLARE - _where text := stac_search_to_where(_search); -BEGIN - EXECUTE format( - $query$ - SELECT - count(*) - FROM - collections_asitems - WHERE %s - ; - $query$, - _where - ) INTO matched; - RETURN; -END; -$$ LANGUAGE PLPGSQL STABLE PARALLEL SAFE; - -- collection_search_plan: the collection counterpart of search_plan -- the CLIENT-STREAMING entry -- for collections. Returns the data query (collection content + keyset keys) the client PREPAREs and -- the always-on numberMatched query. Built from the SAME building blocks as search_plan/search_page @@ -5314,7 +5335,7 @@ END; $$ LANGUAGE plpgsql STABLE SECURITY DEFINER SET search_path TO pgstac, public; --- collection_search: keyset-paginated collection listing (v0.10 — replaces the prior OFFSET paging). +-- collection_search: keyset-paginated collection listing. -- Builds its data + numberMatched queries via collection_search_plan (one source of truth, shared -- with the client-streaming path), fetches _limit+1 to detect a further page, and links next/prev as -- opaque keyset tokens. Collections default to id ASC. numberMatched is ALWAYS returned (small table). @@ -5446,7 +5467,7 @@ DECLARE unionedgeom_area float := 0; prev_area float := 0; _env pred_envelope; bnds record; lead_field text; eff_dir text; datetime_leading boolean; is_asc boolean; orderby_str text; - cursor_ts timestamptz; band record; band_target numeric; obs_sel numeric; + cursor_idx int; mo interval := interval '1 month'; band record; band_target numeric; obs_sel numeric; band_fetched int; guard int := 0; cum_scanned bigint := 0; BEGIN -- If the passed in geometry is not an area, coverage tests are meaningless. @@ -5481,28 +5502,31 @@ BEGIN is_asc := (datetime_leading AND eff_dir = 'ASC'); orderby_str := search.orderby; - -- Same projection choice as search_page: skip the fragment when the fields allow it. + -- Per-row projection: skip the shared item_fragments lookup when the requested fields are + -- satisfiable from item columns alone (needs_fragment, evaluated once for the whole query). IF needs_fragment(coalesce(fields, '{}'::jsonb), bnds.collections) THEN - proj_expr := format( - 'content_hydrate(i, %L::jsonb, (SELECT f FROM item_fragments f WHERE f.id = i.fragment_id))', - coalesce(fields, '{}'::jsonb)); + proj_expr := format('content_hydrate(i, %L::jsonb)', coalesce(fields, '{}'::jsonb)); ELSE - proj_expr := format('content_hydrate(i, %L::jsonb, NULL, true)', coalesce(fields, '{}'::jsonb)); + proj_expr := format('content_hydrate(i, %L::jsonb, true)', coalesce(fields, '{}'::jsonb)); END IF; - IF bnds.min_datetime IS NOT NULL THEN - cursor_ts := CASE WHEN is_asc THEN bnds.min_datetime ELSE bnds.max_datetime + interval '1 microsecond' END; + IF array_length(bnds.months, 1) IS NOT NULL THEN + -- Walk bands in the SEARCH's datetime direction: newest month first for a descending search, + -- oldest first when sorting datetime ascending. Mirrors search_page (004_search). Walking the + -- wrong direction returns older items before newer ones for a descending limit/early-exit search, + -- and (because band grouping shifts with the partition_stats histogram) makes the result depend on + -- stats — a bug, since stats must only affect performance, never which rows come back. + cursor_idx := CASE WHEN is_asc THEN 1 ELSE array_length(bnds.months, 1) END; band_target := (_limit + 1) * band_margin; <> WHILE NOT exit_flag AND guard < 80 LOOP guard := guard + 1; - SELECT * INTO band FROM next_band( - bnds.histogram, bnds.min_datetime, bnds.max_datetime, - cursor_ts, band_target, band_cap_months, is_asc); - EXIT bands WHEN band.done; + SELECT * INTO band FROM next_band(bnds.counts, cursor_idx, band_target, band_cap_months, NOT is_asc); + -- process a valid band even when next_band also flags done; stop only on no band. + EXIT bands WHEN band.band_start_idx IS NULL; query := format( 'SELECT * FROM items i WHERE i.collection = ANY(%L::text[]) AND i.datetime >= %L AND i.datetime < %L AND %s ORDER BY %s LIMIT %L', - band.colls, band.qlo, band.qhi, _where, orderby_str, remaining_limit); + bnds.collections, bnds.months[band.band_start_idx], bnds.months[band.band_end_idx] + mo, _where, orderby_str, remaining_limit); band_fetched := 0; OPEN curs FOR EXECUTE query; LOOP @@ -5538,10 +5562,8 @@ BEGIN END LOOP; CLOSE curs; cum_scanned := cum_scanned + band.scanned; - cursor_ts := band.next_cursor; - EXIT bands WHEN exit_flag - OR (is_asc AND cursor_ts > bnds.max_datetime) - OR (NOT is_asc AND cursor_ts <= bnds.min_datetime); + cursor_idx := band.next_cursor_idx; + EXIT bands WHEN exit_flag; -- LEARN: size the next band from this band's spatial hit-rate (fetched / rows scanned). obs_sel := GREATEST(band_fetched::numeric, 0.5) / GREATEST(cum_scanned, 1); band_target := ((_limit + 1 - counter) / obs_sel) * band_safety; @@ -5612,6 +5634,66 @@ $$ LANGUAGE SQL; -- BEGIN FRAGMENT: 997_maintenance.sql +-- tighten_dirty_partition_stats: recompute the exact envelope + row count for dirty partitions (oldest +-- first), clearing dirty. Run off-hours (pg_cron or the maintenance CLI). Optional: a wide envelope only +-- over-includes a partition in search, so skipping it never loses rows. `_limit` caps the batch (NULL = +-- all dirty); returns the number of partitions tightened. +-- +-- pg_cron example (operators install this themselves): +-- SELECT cron.schedule('pgstac-tighten', '*/15 * * * *', +-- $$SELECT pgstac.tighten_dirty_partition_stats(200)$$); +CREATE OR REPLACE FUNCTION tighten_dirty_partition_stats(_limit int DEFAULT NULL) +RETURNS int AS $$ +DECLARE + _part text; + _count int := 0; +BEGIN + FOR _part IN + SELECT partition FROM pgstac.partition_stats + WHERE dirty + ORDER BY last_updated NULLS FIRST + LIMIT _limit + LOOP + PERFORM pgstac.tighten_partition_stats(_part); + _count := _count + 1; + END LOOP; + RETURN _count; +END; +$$ LANGUAGE PLPGSQL SECURITY DEFINER; + + +-- build_pending_indexes: build the queryable indexes for partitions flagged `indexes_pending` (new +-- partitions are created index-light for fast ingest; changing a queryable flags every partition via the +-- queryables trigger), then clear the flag. Builds the DDL directly (not via the queue), so schedule it +-- off-hours like the tighten sweep. `_limit` caps the batch (NULL = all pending); returns the count built. +-- +-- pg_cron example (operators install this themselves): +-- SELECT cron.schedule('pgstac-build-indexes', '*/30 * * * *', +-- $$SELECT pgstac.build_pending_indexes(50)$$); +CREATE OR REPLACE FUNCTION build_pending_indexes(_limit int DEFAULT NULL) +RETURNS int AS $$ +DECLARE + _part text; + _q text; + _count int := 0; +BEGIN + FOR _part IN + SELECT partition FROM pgstac.partition_stats + WHERE indexes_pending + ORDER BY last_updated NULLS FIRST + LIMIT _limit + LOOP + FOR _q IN SELECT * FROM pgstac.maintain_partition_queries(_part) LOOP + EXECUTE _q; + END LOOP; + UPDATE pgstac.partition_stats SET indexes_pending = false WHERE partition = _part; + _count := _count + 1; + END LOOP; + RETURN _count; +END; +$$ LANGUAGE PLPGSQL SECURITY DEFINER; + + CREATE OR REPLACE PROCEDURE analyze_items() AS $$ DECLARE q text; @@ -5670,7 +5752,7 @@ DECLARE extent jsonb; BEGIN IF runupdate THEN - PERFORM update_partition_stats_q(partition) + PERFORM tighten_partition_stats(partition) FROM partitions_view WHERE collection=_collection; END IF; SELECT @@ -5762,85 +5844,240 @@ BEGIN END; $$ LANGUAGE PLPGSQL; --- gc_anonymous_searches: clean up searches where metadata IS NULL --- Controlled by search_gc_anonymous_retention_interval setting --- Set to '0' or '-1' to disable (never run) -CREATE OR REPLACE FUNCTION gc_anonymous_searches( - retention_interval interval DEFAULT NULL, - conf jsonb DEFAULT NULL -) RETURNS bigint AS $$ -DECLARE effective_interval interval; - result bigint; -BEGIN - effective_interval := coalesce( - retention_interval, - get_setting('search_gc_anonymous_retention_interval', conf)::interval - ); - - IF effective_interval <= '0'::interval THEN - RETURN 0; - END IF; - WITH deleted AS ( - DELETE FROM searches - WHERE metadata IS NULL - AND lastused < now() - effective_interval - RETURNING 1 +-- update_collection_extents: recompute every collection's extent from partitions_view. +CREATE OR REPLACE FUNCTION update_collection_extents() RETURNS VOID AS $$ +UPDATE collections + SET content = jsonb_set_lax( + content, + '{extent}'::text[], + collection_extent(id, FALSE), + true, + 'use_json_null' ) - SELECT count(*)::bigint INTO result FROM deleted; +; +$$ LANGUAGE SQL; - RETURN result; -END; -$$ LANGUAGE PLPGSQL SECURITY DEFINER; --- gc_searches_with_metadata: clean up searches where metadata IS NOT NULL --- Controlled by search_gc_metadata_retention_interval setting --- Set to '0' or '-1' to disable (never run) -CREATE OR REPLACE FUNCTION gc_searches_with_metadata( - retention_interval interval DEFAULT NULL, - conf jsonb DEFAULT NULL -) RETURNS bigint AS $$ -DECLARE effective_interval interval; - result bigint; +-- --------------------------------------------------------------------------- +-- Field registry maintenance: track which paths (and value types) exist per +-- collection. Used for schema inference (e.g. the geoparquet export schema). +-- jsonb_field_rows is in 001a_jsonutils.sql. +-- --------------------------------------------------------------------------- + +-- update_field_registry_from_sample: UPSERT registry rows from an array of raw item content JSONBs (the +-- caller picks the sample); value_kinds accumulate observed types over time. +CREATE OR REPLACE FUNCTION update_field_registry_from_sample( + _collection text, + item_contents jsonb[] +) RETURNS void AS $$ + INSERT INTO item_field_registry (collection, path, is_leaf, value_kinds, first_seen, last_seen) + SELECT + _collection, + r.path, + bool_and(r.is_leaf) AS is_leaf, + array_agg(DISTINCT r.value_kind) FILTER (WHERE r.value_kind IS NOT NULL) AS value_kinds, + now(), + now() + FROM unnest(item_contents) AS item(content) + CROSS JOIN LATERAL jsonb_field_rows(item.content) AS r(path, is_leaf, value_kind) + GROUP BY r.path + ON CONFLICT (collection, path) DO UPDATE SET + is_leaf = EXCLUDED.is_leaf, + value_kinds = ( + SELECT array_agg(DISTINCT v) + FROM unnest(item_field_registry.value_kinds || EXCLUDED.value_kinds) t(v) + ), + last_seen = now() + ; +$$ LANGUAGE SQL VOLATILE SECURITY DEFINER; + + +-- update_field_registry_from_items: sample a live collection and UPSERT registry rows (TABLESAMPLE +-- BERNOULLI(5) above ~10k rows by pg_class estimate, else LIMIT 1000). Returns (registered_paths, +-- rows_processed). +CREATE OR REPLACE FUNCTION update_field_registry_from_items( + _collection text +) RETURNS TABLE (registered_paths int, rows_processed int) AS $$ +DECLARE + est_rows bigint; + nrows int; + npaths int; BEGIN - effective_interval := coalesce( - retention_interval, - get_setting('search_gc_metadata_retention_interval', conf)::interval - ); + -- Sum reltuples across the registered item partitions for this collection. + -- reltuples can be -1 (never analyzed); treat negative values as zero. + SELECT COALESCE(sum(GREATEST(c.reltuples::bigint, 0)), 0) INTO est_rows + FROM partitions_view p + JOIN pg_class c ON c.relname = p.partition + JOIN pg_namespace n ON n.oid = c.relnamespace + WHERE p.collection = _collection + AND n.nspname = 'pgstac' + AND c.relkind = 'r'; - IF effective_interval <= '0'::interval THEN - RETURN 0; + IF est_rows > 10000 THEN + -- Large collection: use statistical sampling to avoid full seq-scan. + WITH sampled AS ( + SELECT content_hydrate(i) AS content FROM items i TABLESAMPLE BERNOULLI(5) WHERE i.collection = _collection + ), + upserted AS ( + INSERT INTO item_field_registry (collection, path, is_leaf, value_kinds, first_seen, last_seen) + SELECT + _collection, + r.path, + bool_and(r.is_leaf) AS is_leaf, + array_agg(DISTINCT r.value_kind) FILTER (WHERE r.value_kind IS NOT NULL) AS value_kinds, + now(), now() + FROM sampled + CROSS JOIN LATERAL jsonb_field_rows(content) AS r(path, is_leaf, value_kind) + GROUP BY r.path + ON CONFLICT (collection, path) DO UPDATE SET + is_leaf = EXCLUDED.is_leaf, + value_kinds = ( + SELECT array_agg(DISTINCT v) + FROM unnest(item_field_registry.value_kinds || EXCLUDED.value_kinds) t(v) + ), + last_seen = now() + RETURNING 1 + ) + SELECT + (SELECT count(*)::int FROM upserted), + (SELECT count(*)::int FROM sampled) + INTO npaths, nrows; + ELSE + -- Small collection: process up to 1000 rows to avoid BERNOULLI returning 0 rows. + WITH sampled AS ( + SELECT content_hydrate(i) AS content FROM items i WHERE i.collection = _collection LIMIT 1000 + ), + upserted AS ( + INSERT INTO item_field_registry (collection, path, is_leaf, value_kinds, first_seen, last_seen) + SELECT + _collection, + r.path, + bool_and(r.is_leaf) AS is_leaf, + array_agg(DISTINCT r.value_kind) FILTER (WHERE r.value_kind IS NOT NULL) AS value_kinds, + now(), now() + FROM sampled + CROSS JOIN LATERAL jsonb_field_rows(content) AS r(path, is_leaf, value_kind) + GROUP BY r.path + ON CONFLICT (collection, path) DO UPDATE SET + is_leaf = EXCLUDED.is_leaf, + value_kinds = ( + SELECT array_agg(DISTINCT v) + FROM unnest(item_field_registry.value_kinds || EXCLUDED.value_kinds) t(v) + ), + last_seen = now() + RETURNING 1 + ) + SELECT + (SELECT count(*)::int FROM upserted), + (SELECT count(*)::int FROM sampled) + INTO npaths, nrows; END IF; + RETURN QUERY SELECT npaths, nrows; +END; +$$ LANGUAGE PLPGSQL VOLATILE SECURITY DEFINER; + + +-- refresh_field_registry: expire registry entries not seen within retention_interval (scheduled +-- maintenance). Returns (collection, expired_paths) per affected collection. +CREATE OR REPLACE FUNCTION refresh_field_registry( + _collection text DEFAULT NULL, + retention_interval interval DEFAULT '90 days' +) RETURNS TABLE (collection_id text, expired_paths int) AS $$ WITH deleted AS ( - DELETE FROM searches - WHERE metadata IS NOT NULL - AND lastused < now() - effective_interval - RETURNING 1 + DELETE FROM item_field_registry + WHERE (_collection IS NULL OR collection = _collection) + AND last_seen < now() - retention_interval + RETURNING collection ) - SELECT count(*)::bigint INTO result FROM deleted; + SELECT collection, count(*)::int + FROM deleted + GROUP BY collection; +$$ LANGUAGE SQL VOLATILE SECURITY DEFINER; - RETURN result; -END; -$$ LANGUAGE PLPGSQL SECURITY DEFINER; --- gc_search_caches: wrapper that calls both GC functions -CREATE OR REPLACE FUNCTION gc_search_caches( - retention_interval interval DEFAULT NULL, - conf jsonb DEFAULT NULL -) RETURNS jsonb AS $$ -DECLARE anon_count bigint; - meta_count bigint; +-- gc_fragments: garbage-collect orphaned item_fragments with a single set-based DELETE (NOT EXISTS +-- anti-join against items.fragment_id). items.fragment_id has no FK (partitioned-items incremental +-- NOT VALID FKs aren't supported), so a fragment unreferenced at the DELETE snapshot but referenced by a +-- concurrent insert can be removed; the retention_interval guard makes this unlikely. Run during +-- low-ingest periods. +CREATE OR REPLACE FUNCTION gc_fragments( + _collection text DEFAULT NULL, + retention_interval interval DEFAULT '90 days' +) RETURNS TABLE ( + collection_id text, + fragments_removed int +) AS $$ + WITH deleted AS ( + DELETE FROM item_fragments f + WHERE + (_collection IS NULL OR f.collection = _collection) + AND f.created_at < now() - retention_interval + AND NOT EXISTS (SELECT 1 FROM items i WHERE i.fragment_id = f.id) + RETURNING f.collection + ) + SELECT collection, count(*)::int + FROM deleted + GROUP BY collection; +$$ LANGUAGE SQL VOLATILE PARALLEL UNSAFE SECURITY DEFINER; + + +-- tighten_partition_stats: recompute the exact envelope (datetime/end_datetime min-max + spatial extent) +-- and row count for one partition, then clear `dirty`. The only function that narrows a partition_stats +-- envelope; the maintenance sweeps drive it over partition_stats WHERE dirty. An empty partition tightens +-- to an 'empty' range so search prunes it out. +CREATE OR REPLACE FUNCTION tighten_partition_stats(_partition text) RETURNS void AS $$ +DECLARE + _n bigint; + _dtmin timestamptz; _dtmax timestamptz; + _edtmin timestamptz; _edtmax timestamptz; + _spatial geometry; + _dtrange tstzrange; + _edtrange tstzrange; + _collection text; BEGIN - anon_count := gc_anonymous_searches(retention_interval, conf); - meta_count := gc_searches_with_metadata(retention_interval, conf); + -- Hold the partition's advisory lock across the scan + write: tighten narrows dtrange to the scanned + -- extent and clears dirty, so without the lock a concurrent ingest committing a row outside that + -- extent would be left uncovered (search would prune + miss it). Same lock check_partition uses, so + -- tighten serializes with ingest into this partition only. + PERFORM pg_advisory_xact_lock(hashtext('pgstac.check_partition'), hashtext(_partition)); - RETURN jsonb_build_object( - 'removed_anonymous', anon_count, - 'removed_with_metadata', meta_count - ); + EXECUTE format( + $q$ + SELECT count(*), min(datetime), max(datetime), min(end_datetime), max(end_datetime), + st_setsrid(st_extent(geometry)::geometry, 4326) + FROM %I + $q$, + _partition + ) INTO _n, _dtmin, _dtmax, _edtmin, _edtmax, _spatial; + + IF _n = 0 THEN + _dtrange := 'empty'::tstzrange; + _edtrange := 'empty'::tstzrange; + _spatial := NULL; + ELSE + _dtrange := tstzrange(_dtmin, _dtmax, '[]'); + _edtrange := tstzrange(_edtmin, _edtmax, '[]'); + END IF; + + SELECT pv.collection + INTO _collection + FROM partitions_view pv WHERE partition = _partition; + + INSERT INTO partition_stats + (partition, collection, dtrange, edtrange, spatial, n, dirty, last_updated) + VALUES (_partition, _collection, _dtrange, _edtrange, _spatial, _n, false, now()) + ON CONFLICT (partition) DO UPDATE SET + collection = EXCLUDED.collection, + dtrange = EXCLUDED.dtrange, + edtrange = EXCLUDED.edtrange, + spatial = EXCLUDED.spatial, + n = EXCLUDED.n, + dirty = false, + last_updated = now(); END; -$$ LANGUAGE SQL SECURITY DEFINER; +$$ LANGUAGE PLPGSQL SECURITY DEFINER; -- END FRAGMENT: 997_maintenance.sql -- BEGIN FRAGMENT: 998_idempotent_post.sql @@ -5871,7 +6108,7 @@ DO $$ END $$; --- Register promoted native-column queryables (v0.10 split schema). +-- Register promoted native-column queryables. -- Each entry maps a STAC property name to the promoted items column via property_path. -- CQL2 queries and auto-created indexes will use the native column, not JSONB extraction. -- The seed data lives in promoted_queryables_defaults() (002a_queryables.sql) so it @@ -5909,7 +6146,6 @@ INSERT INTO pgstac_settings (name, value) VALUES ('context_estimated_count', '100000'), ('context_estimated_cost', '100000'), ('context_stats_ttl', '1 day'), - ('search_gc_retention_interval', '7 days'), ('default_filter_lang', 'cql2-json'), ('additional_properties', 'true'), ('use_queue', 'false'), @@ -5963,61 +6199,33 @@ ALTER FUNCTION to_int COST 5000; ALTER FUNCTION to_tstz COST 5000; ALTER FUNCTION to_text_array COST 5000; -ALTER FUNCTION update_partition_stats SECURITY DEFINER; -ALTER FUNCTION partition_after_triggerfunc SECURITY DEFINER; -ALTER FUNCTION drop_table_constraints SECURITY DEFINER; -ALTER FUNCTION create_table_constraints SECURITY DEFINER; -ALTER FUNCTION check_partition SECURITY DEFINER; -ALTER FUNCTION repartition SECURITY DEFINER; -ALTER FUNCTION where_stats(text, text, boolean, jsonb) SECURITY DEFINER; -ALTER FUNCTION search_query SECURITY DEFINER; -ALTER FUNCTION name_search SECURITY DEFINER; -ALTER FUNCTION rename_search SECURITY DEFINER; -ALTER FUNCTION unname_search SECURITY DEFINER; -ALTER FUNCTION pin_search SECURITY DEFINER; -ALTER FUNCTION unpin_search SECURITY DEFINER; -ALTER FUNCTION gc_anonymous_searches(interval, jsonb) SECURITY DEFINER; -ALTER FUNCTION gc_search_caches(interval, jsonb) SECURITY DEFINER; -ALTER FUNCTION gc_deleted_items_log_batch(interval, integer) SECURITY DEFINER; -ALTER FUNCTION gc_deleted_items_log(interval, integer) SECURITY DEFINER; -ALTER FUNCTION gc_deleted_items_log(interval) SECURITY DEFINER; -ALTER FUNCTION format_item SECURITY DEFINER; -ALTER FUNCTION maintain_index SECURITY DEFINER; -ALTER FUNCTION pgstac.jsonb_hash(jsonb) SECURITY DEFINER; -ALTER FUNCTION promoted_items_column_list() SECURITY DEFINER; -ALTER FUNCTION items_content_distinct_sql(text, text) SECURITY DEFINER; -ALTER FUNCTION items_content_changed(items, items) SECURITY DEFINER; -ALTER FUNCTION items_touch_triggerfunc SECURITY DEFINER; -ALTER FUNCTION items_delete_log_trigger SECURITY DEFINER; -ALTER FUNCTION strip_promoted_properties(jsonb) SECURITY DEFINER; -ALTER FUNCTION tstz_to_stac_text(timestamptz) SECURITY DEFINER; -ALTER FUNCTION temporal_properties_from_item(items) SECURITY DEFINER; -ALTER FUNCTION promoted_properties_from_item(items) SECURITY DEFINER; -ALTER FUNCTION extract_fragment(jsonb, text[]) SECURITY DEFINER; -ALTER FUNCTION pgstac_hash_fragment(jsonb) SECURITY DEFINER; -ALTER FUNCTION gc_fragments(text, interval) SECURITY DEFINER; -ALTER FUNCTION strip_fragment_col(jsonb, text, text[]) SECURITY DEFINER; -ALTER FUNCTION update_field_registry_from_sample(text, jsonb[]) SECURITY DEFINER; -ALTER FUNCTION update_field_registry_from_items(text) SECURITY DEFINER; -ALTER FUNCTION refresh_field_registry(text, interval) SECURITY DEFINER; -ALTER FUNCTION collection_fragment_config_default(jsonb) SECURITY DEFINER; -ALTER FUNCTION jsonb_leaf_rows(jsonb, text) SECURITY DEFINER; -ALTER FUNCTION jsonb_common_values(jsonb, jsonb) SECURITY DEFINER; -ALTER FUNCTION fragment_path_text(text[]) SECURITY DEFINER; -ALTER FUNCTION fragment_path_array(text) SECURITY DEFINER; +-- SECURITY DEFINER is declared INLINE in each function's CREATE (the single source of truth), +-- not re-applied here. Functions that create partitions/indexes/constraints declare it inline so +-- the created objects are owned by pgstac_admin; functions that write the search cache from the +-- read path declare it inline too. Pure helpers stay SECURITY INVOKER. Keeping a separate ALTER +-- list here only let it drift from the definitions (stale/duplicate/wrong-signature entries). -GRANT USAGE ON SCHEMA pgstac to pgstac_read; -GRANT ALL ON SCHEMA pgstac to pgstac_ingest; -GRANT ALL ON SCHEMA pgstac to pgstac_admin; +-- Schema USAGE for pgstac_read / pgstac_ingest is granted in 000_idempotent_pre.sql; pgstac_admin +-- owns the schema. Not re-granted here. --- pgstac_read role limited to using function apis +-- pgstac_read API surface. Functions are EXECUTE-able by PUBLIC by default, so these grants are not +-- required for access today; they document the intended top-level read API (and would be the point to +-- enforce from if EXECUTE were ever revoked from PUBLIC). Internal helpers (keyset_*, partition_bounds, +-- cql2_*, next_band, ...) are deliberately NOT listed — read reaches them only inside these entry points. GRANT EXECUTE ON FUNCTION search TO pgstac_read; GRANT EXECUTE ON FUNCTION search_query TO pgstac_read; GRANT EXECUTE ON FUNCTION item_by_id TO pgstac_read; GRANT EXECUTE ON FUNCTION get_item TO pgstac_read; -GRANT EXECUTE ON FUNCTION format_item TO pgstac_read; GRANT EXECUTE ON FUNCTION content_hydrate TO pgstac_read; -GRANT EXECUTE ON FUNCTION pgstac.jsonb_hash(jsonb) TO pgstac_read; +GRANT EXECUTE ON FUNCTION search_page TO pgstac_read; +GRANT EXECUTE ON FUNCTION search_plan TO pgstac_read; +GRANT EXECUTE ON FUNCTION collection_search_plan TO pgstac_read; +GRANT EXECUTE ON FUNCTION collection_search TO pgstac_read; +GRANT EXECUTE ON FUNCTION geometrysearch TO pgstac_read; +GRANT EXECUTE ON FUNCTION geojsonsearch TO pgstac_read; +GRANT EXECUTE ON FUNCTION xyzsearch TO pgstac_read; +GRANT EXECUTE ON FUNCTION search_from_json(jsonb, jsonb) TO pgstac_read; +-- Tables are NOT readable by PUBLIC; read needs an explicit SELECT grant. GRANT SELECT ON ALL TABLES IN SCHEMA pgstac TO pgstac_read; @@ -6025,6 +6233,21 @@ GRANT EXECUTE ON ALL FUNCTIONS IN SCHEMA pgstac to pgstac_ingest; GRANT ALL ON ALL TABLES IN SCHEMA pgstac to pgstac_ingest; GRANT USAGE ON ALL SEQUENCES IN SCHEMA pgstac to pgstac_ingest; +-- Privilege wall (INV-1): the connecting role (inheriting pgstac_ingest) may READ these tables but may +-- only MUTATE them through the SECURITY DEFINER write functions (check_partition, widen/tighten_partition_stats, +-- ensure_fragments, make_binary_staging, flush_items_staging_binary, the items_staging trigger, create/ +-- update/upsert/delete_item, the field-registry/fragment helpers). This makes an un-widened or +-- envelope-narrowing direct write structurally impossible. New partitions are created SELECT-only for +-- pgstac_ingest by check_partition; the items parent is revoked here. Staging tables (items_staging*) stay +-- writable so the SQL-only ingest path can COPY into them. +REVOKE INSERT, UPDATE, DELETE, TRUNCATE ON items, partition_stats, item_fragments FROM pgstac_ingest; +-- item_field_registry is the one exception to the wall: the loader WIDENS it directly with an add-only +-- INSERT ... ON CONFLICT DO UPDATE (no SD function), so INSERT + UPDATE stay granted. DELETE/TRUNCATE remain +-- revoked, so even a direct write cannot NARROW the registry — INV-1 (registry is a superset of the data) +-- holds structurally against narrowing, while the cheap widen stays on the hot load path. +REVOKE DELETE, TRUNCATE ON item_field_registry FROM pgstac_ingest; + + REVOKE ALL PRIVILEGES ON PROCEDURE run_queued_queries FROM public; GRANT ALL ON PROCEDURE run_queued_queries TO pgstac_admin; @@ -6036,38 +6259,9 @@ GRANT ALL ON PROCEDURE gc_deleted_items_log_committed(interval, integer) TO pgst RESET ROLE; -SET ROLE pgstac_ingest; -SELECT update_partition_stats_q(partition) FROM partitions_view; - --- v0.10 search: grants for new functions -ALTER FUNCTION where_stats(text, text, boolean, jsonb, text) SECURITY DEFINER; -ALTER FUNCTION search_plan(jsonb, text, integer) SECURITY DEFINER; -ALTER FUNCTION collection_search_plan(jsonb) SECURITY DEFINER; -ALTER FUNCTION gc_anonymous_searches(interval, jsonb) SECURITY DEFINER; -ALTER FUNCTION gc_searches_with_metadata(interval, jsonb) SECURITY DEFINER; - -GRANT EXECUTE ON FUNCTION keyset_encode(text[]) TO pgstac_read; -GRANT EXECUTE ON FUNCTION keyset_decode(text) TO pgstac_read; -GRANT EXECUTE ON FUNCTION keyset_sortkeys(jsonb) TO pgstac_read; -GRANT EXECUTE ON FUNCTION keyset_where(jsonb, text[], boolean) TO pgstac_read; -GRANT EXECUTE ON FUNCTION search_page(jsonb, integer, text, boolean) TO pgstac_read; -GRANT EXECUTE ON FUNCTION search_plan(jsonb, text, integer) TO pgstac_read; -GRANT EXECUTE ON FUNCTION search_from_json(jsonb, jsonb) TO pgstac_read; -GRANT EXECUTE ON FUNCTION collection_search_plan(jsonb) TO pgstac_read; -GRANT EXECUTE ON FUNCTION search_to_cql2(jsonb) TO pgstac_read; -GRANT EXECUTE ON FUNCTION search_envelope(jsonb) TO pgstac_read; -GRANT EXECUTE ON FUNCTION cql2_envelope(jsonb) TO pgstac_read; -GRANT EXECUTE ON FUNCTION cql2_collection_set(text, jsonb) TO pgstac_read; -GRANT EXECUTE ON FUNCTION q_op_query(jsonb) TO pgstac_read; -GRANT EXECUTE ON FUNCTION partition_bounds(pred_envelope, boolean) TO pgstac_read; -GRANT EXECUTE ON FUNCTION next_band(bigint[], integer, numeric, integer) TO pgstac_read; -GRANT EXECUTE ON FUNCTION needs_fragment(jsonb, text[]) TO pgstac_read; - --- GC retention settings -INSERT INTO pgstac_settings (name, value) VALUES - ('search_gc_anonymous_retention_interval', '1 day'), - ('search_gc_metadata_retention_interval', '30 days') -ON CONFLICT (name) DO UPDATE SET value = EXCLUDED.value; +-- (No install-time stats seeding: a fresh install has no partitions, and partition_stats rows are now +-- seeded by check_partition as partitions are created. Exact stats come from tighten_partition_stats via +-- the maintenance sweeps.) -- END FRAGMENT: 998_idempotent_post.sql -- BEGIN FRAGMENT: 999_version.sql diff --git a/src/pgstac/pgstac.sql b/src/pgstac/pgstac.sql index 0383cb94..77894b32 100644 --- a/src/pgstac/pgstac.sql +++ b/src/pgstac/pgstac.sql @@ -44,6 +44,7 @@ DO $$ $$; + GRANT pgstac_admin TO current_user; -- Function to make sure pgstac_admin is the owner of items @@ -195,6 +196,11 @@ RETURNS timestamptz AS $$ ; $$ LANGUAGE SQL IMMUTABLE STRICT; +-- Drop objects superseded by the current partition_stats model. +DROP MATERIALIZED VIEW IF EXISTS partitions CASCADE; +DROP MATERIALIZED VIEW IF EXISTS partition_steps; +DROP VIEW IF EXISTS partition_steps; + -- Drop function signatures whose argument lists changed (CREATE OR REPLACE cannot alter them) DROP FUNCTION IF EXISTS chunker(pred_envelope); DROP FUNCTION IF EXISTS search_bands(pred_envelope, boolean, integer, integer); @@ -214,6 +220,7 @@ DROP FUNCTION IF EXISTS xyzsearch(integer, integer, integer, text, jsonb, intege DROP FUNCTION IF EXISTS search(jsonb); DROP FUNCTION IF EXISTS search_page(jsonb, integer, text, boolean); DROP FUNCTION IF EXISTS search_plan(jsonb, text); +DROP FUNCTION IF EXISTS fields_to_itemcols(jsonb); DROP FUNCTION IF EXISTS search_query(jsonb, boolean, jsonb); DROP FUNCTION IF EXISTS where_stats(text, text, boolean, jsonb); DROP FUNCTION IF EXISTS keyset_sortkeys(jsonb); @@ -640,63 +647,51 @@ CREATE OR REPLACE FUNCTION explode_dotpaths_recurse(IN j jsonb) RETURNS SETOF te $$ LANGUAGE SQL IMMUTABLE PARALLEL SAFE; --- jsonb_canonical: RFC 8785 (JSON Canonicalization Scheme)-aligned serialization. --- Produces a deterministic, key-order-independent text encoding that an external --- client can reproduce byte-for-byte. NOTE: do NOT use `jsonb::text` for hashing — --- PostgreSQL re-normalizes object key order to length-then-bytewise and inserts --- ": " / ", " separators, so `jsonb::text` is neither alphabetical nor compact. +-- jsonb_canonical_hash: a deterministic 32-byte SHA-256 identity digest of a JSONB document, reproducible +-- outside PostgreSQL (pgstac-rs `canonical::jsonb_canonical_hash` produces identical bytes). -- --- Canonical rules (must match the external recipe below): --- * object keys sorted by Unicode code point (== UTF-8 byte order, COLLATE "C"), --- * compact separators: ',' between members, ':' between key and value, --- * strings: standard JSON escaping, NON-ASCII left as UTF-8 (no \uXXXX), --- * numbers: IEEE-754 double, shortest round-trip form (Ryu) — matches --- PostgreSQL float8 output and ECMAScript Number::toString for in-range --- values. (STAC numbers are physical quantities; integers beyond 2^53 are --- out of contract, as in RFC 8785.) --- * true / false / null as literals. --- --- External equivalents: --- Python: an RFC 8785 canonicalizer, or the rule-for-rule reference: --- def canon(v): --- if isinstance(v, bool): return 'true' if v else 'false' --- if v is None: return 'null' --- if isinstance(v, dict): --- return '{'+','.join(json.dumps(k,ensure_ascii=False)+':'+canon(v[k]) --- for k in sorted(v))+'}' --- if isinstance(v, list): return '['+','.join(canon(x) for x in v)+']' --- if isinstance(v,(int,float)): --- f=float(v); return str(int(f)) if f==int(f) and abs(f)<1e16 else repr(f) --- return json.dumps(v, ensure_ascii=False) --- Rust: the `rfc8785` crate (serde_jcs) over serde_json::Value. -CREATE OR REPLACE FUNCTION jsonb_canonical(j jsonb) RETURNS text AS $$ - SELECT CASE jsonb_typeof(j) - WHEN 'object' THEN COALESCE(( - SELECT '{' || string_agg( - to_json(kv.key)::text || ':' || jsonb_canonical(kv.value), - ',' ORDER BY kv.key COLLATE "C" - ) || '}' - FROM jsonb_each(j) kv - ), '{}') - WHEN 'array' THEN COALESCE(( - SELECT '[' || string_agg(jsonb_canonical(e.value), ',' ORDER BY e.ord) || ']' - FROM jsonb_array_elements(j) WITH ORDINALITY e(value, ord) - ), '[]') - WHEN 'number' THEN (j #>> '{}')::float8::text - ELSE j::text -- string (JSON-escaped, UTF-8 preserved), 'true' / 'false' / 'null' - END; -$$ LANGUAGE SQL IMMUTABLE PARALLEL SAFE STRICT; - --- jsonb_hash: raw 32-byte sha256 of the canonical (RFC 8785-aligned) JSON form. --- Returns bytea so callers store the compact binary digest directly (32 B vs --- 64-char hex). Use encode(jsonb_hash(j), 'hex') when a printable string is --- needed for display or external comparison. --- Externally reproducible: sha256(utf8_bytes(jsonb_canonical(j))). --- The private jsonb column on items/collections is intentionally excluded — it --- stores operator metadata outside the STAC item identity contract. -CREATE OR REPLACE FUNCTION jsonb_hash(j jsonb) RETURNS bytea AS $$ - SELECT sha256(convert_to(jsonb_canonical(j), 'UTF8')); -$$ LANGUAGE SQL IMMUTABLE PARALLEL SAFE STRICT; +-- The document is flattened to leaf (path, value) rows, each rendered as `path value`, sorted by path +-- (byte order), joined by , and hashed. Separators are C0 control chars US/RS/FS (0x1F/0x1E/0x1D), +-- which do not occur in STAC JSON keys or strings. +-- * path — per step, then 'k'||key (object) or 'i'||idx (0-based array). The 'k'/'i' tags keep an +-- object key "0" distinct from array index 0; keeps nesting distinct from a key with '/'. +-- * number — 'n' || (v)::float8::text, so the digest matches a value stored in a float8 (promoted) column +-- and read back. +-- * string — datetime-shaped (YYYY-MM-DD, optional T/space time, optional offset): 't' || the instant +-- normalized to UTC / 6-digit microseconds / 'Z'. Offset-less is assumed UTC (the SET TIMEZONE +-- below pins the cast), matching pgstac's to_tstz ingest; the cast raises on a shaped-but- +-- invalid timestamp. Other strings: 's' || the raw string. +-- * boolean — 'b'||'true'/'false'; null — 'z'; empty {} — 'e', empty [] — 'a'. +CREATE OR REPLACE FUNCTION jsonb_canonical_hash(j jsonb) RETURNS bytea AS $$ + WITH RECURSIVE t(path, value) AS ( + SELECT ''::text, j + UNION ALL + SELECT t.path || E'\x1F' || e.seg, e.value + FROM t CROSS JOIN LATERAL ( + SELECT 'k' || key AS seg, value + FROM jsonb_each(t.value) WHERE jsonb_typeof(t.value) = 'object' + UNION ALL + SELECT 'i' || (ord - 1)::text AS seg, value + FROM jsonb_array_elements(t.value) WITH ORDINALITY x(value, ord) + WHERE jsonb_typeof(t.value) = 'array' + ) e + ) + SELECT sha256(convert_to(COALESCE(string_agg( + t.path || E'\x1E' || CASE jsonb_typeof(t.value) + WHEN 'number' THEN 'n' || (t.value #>> '{}')::float8::text + WHEN 'boolean' THEN 'b' || (t.value #>> '{}') + WHEN 'null' THEN 'z' + WHEN 'string' THEN CASE + WHEN (t.value #>> '{}') ~ '^\d{4}-\d{2}-\d{2}([Tt ]\d{2}:\d{2}:\d{2}(\.\d+)?([Zz]|[+-]\d{2}:?\d{2})?)?$' + THEN 't' || to_char((t.value #>> '{}')::timestamptz AT TIME ZONE 'UTC', 'YYYY-MM-DD"T"HH24:MI:SS.US') || 'Z' + ELSE 's' || (t.value #>> '{}') END + WHEN 'object' THEN 'e' + WHEN 'array' THEN 'a' + ELSE 'c' END, + E'\x1D' ORDER BY t.path COLLATE "C"), ''), 'UTF8')) + FROM t + WHERE jsonb_typeof(t.value) NOT IN ('object', 'array') OR t.value = '{}'::jsonb OR t.value = '[]'::jsonb; +$$ LANGUAGE SQL IMMUTABLE PARALLEL SAFE STRICT SET TIMEZONE = 'UTC'; -- jsonb_field_rows: Recursively walk a JSONB document and emit one row per field path. -- max_depth guards against runaway recursion on pathologically nested documents. @@ -1740,10 +1735,13 @@ $$ LANGUAGE SQL; CREATE OR REPLACE FUNCTION queryables_trigger_func() RETURNS TRIGGER AS $$ DECLARE BEGIN - PERFORM maintain_partitions(); + -- Queryable definitions changed, so every partition's queryable indexes may be stale. Flag them and let + -- the async sweep (build_pending_indexes) (re)build off the hot path, instead of rebuilding every + -- partition synchronously inside this trigger. + UPDATE pgstac.partition_stats SET indexes_pending = true; RETURN NULL; END; -$$ LANGUAGE PLPGSQL; +$$ LANGUAGE PLPGSQL SECURITY DEFINER; CREATE TRIGGER queryables_trigger AFTER INSERT OR UPDATE ON queryables FOR EACH STATEMENT EXECUTE PROCEDURE queryables_trigger_func(); @@ -2293,7 +2291,7 @@ BEGIN END IF; IF jsonb_typeof(j) = 'object' THEN - -- GeoJSON geometry args (Point/Polygon/.../GeometryCollection) are not cql1 expressions; + -- GeoJSON geometry args (Point/Polygon/.../GeometryCollection) are not cql expressions; -- pass them through unchanged so spatial ops keep their geometry intact. IF j ? 'type' AND (j ? 'coordinates' OR j ? 'geometries') THEN RETURN j; @@ -2750,8 +2748,8 @@ CREATE OR REPLACE FUNCTION partition_bounds( ) RETURNS record LANGUAGE sql STABLE AS $$ WITH cand AS ( SELECT ps.collection, - lower(coalesce(ps.dtrange, ps.partition_dtrange)) AS lo, - upper(coalesce(ps.dtrange, ps.partition_dtrange)) AS hi, + lower(ps.dtrange) AS lo, + upper(ps.dtrange) AS hi, coalesce(ps.n, 0) AS n FROM partition_stats ps WHERE ((_env).colls IS NULL OR ps.collection = ANY((_env).colls)) @@ -2788,12 +2786,17 @@ $$; -- next_band: walk a histogram using array indexes. Given per-month counts, a -- cursor position, and a target row count, returns the next band's index range. --- The caller converts band indexes back to timestamps for SQL queries. +-- The caller converts band indexes back to timestamps for SQL queries. The band +-- index range [band_start_idx, band_end_idx] is always low..high regardless of +-- direction; only the walk order differs. For a descending search (_descending) +-- the cursor starts at the most recent month and walks toward older months, so +-- the most recent items are collected first. CREATE OR REPLACE FUNCTION next_band( _counts bigint[], _cursor_idx int, _target numeric, _cap_months int, + _descending boolean DEFAULT false, OUT band_start_idx int, OUT band_end_idx int, OUT scanned bigint, @@ -2802,6 +2805,7 @@ CREATE OR REPLACE FUNCTION next_band( ) RETURNS record LANGUAGE plpgsql STABLE AS $$ DECLARE idx int; + n int; cumulative bigint := 0; BEGIN done := false; scanned := 0; @@ -2811,14 +2815,39 @@ BEGIN done := true; -- no histogram / no cursor => nothing to walk RETURN; END IF; + n := array_length(_counts, 1); + + IF _descending THEN + -- Walk downward (newest -> oldest). The high bound is the cursor (clamped into range). + IF _cursor_idx < 1 THEN done := true; RETURN; END IF; + idx := LEAST(_cursor_idx, n); + FOR i IN REVERSE idx..GREATEST(1, idx - _cap_months + 1) LOOP + cumulative := cumulative + _counts[i]; + scanned := scanned + _counts[i]; + IF cumulative >= _target THEN + band_start_idx := i; -- low (older) month bound + band_end_idx := idx; -- high (newer) month bound + next_cursor_idx := i - 1; + IF next_cursor_idx < 1 THEN done := true; END IF; + RETURN; + END IF; + END LOOP; + -- Target not reached within the cap: end the band at the cap boundary. + band_start_idx := GREATEST(1, idx - _cap_months + 1); + band_end_idx := idx; + next_cursor_idx := band_start_idx - 1; + done := (band_start_idx <= 1); + RETURN; + END IF; + -- Ascending (oldest -> newest). idx := NULL; - FOR i IN 1..array_length(_counts, 1) LOOP + FOR i IN 1..n LOOP IF i >= _cursor_idx THEN idx := i; EXIT; END IF; END LOOP; IF idx IS NULL THEN done := true; next_cursor_idx := _cursor_idx; RETURN; END IF; - FOR i IN idx..array_length(_counts, 1) LOOP + FOR i IN idx..n LOOP IF i > idx + _cap_months - 1 THEN EXIT; END IF; cumulative := cumulative + _counts[i]; scanned := scanned + _counts[i]; @@ -2826,7 +2855,7 @@ BEGIN band_start_idx := idx; band_end_idx := i; next_cursor_idx := i + 1; - IF next_cursor_idx > array_length(_counts, 1) THEN done := true; END IF; + IF next_cursor_idx > n THEN done := true; END IF; RETURN; END IF; END LOOP; @@ -2834,9 +2863,9 @@ BEGIN -- Target not reached within the cap: end the band at the cap boundary (not the array end), -- so the cap actually limits band width. done only when we've consumed the whole histogram. band_start_idx := idx; - band_end_idx := LEAST(idx + _cap_months - 1, array_length(_counts, 1)); + band_end_idx := LEAST(idx + _cap_months - 1, n); next_cursor_idx := band_end_idx + 1; - done := (band_end_idx >= array_length(_counts, 1)); + done := (band_end_idx >= n); END; $$; -- END FRAGMENT: 002c_envelope.sql @@ -2987,7 +3016,7 @@ CREATE TABLE items ( stac_version text, stac_extensions jsonb DEFAULT '[]'::jsonb, pgstac_updated_at timestamptz NOT NULL DEFAULT now(), - -- 32-byte sha256 of the canonical (RFC 8785-aligned) STAC item JSON set at + -- 32-byte sha256 of the canonical (jsonb_canonical) STAC item JSON set at -- ingest time. Allows external clients to detect unchanged items without a -- full fetch. Does NOT include the private column (operator metadata). item_hash bytea NOT NULL DEFAULT '\x'::bytea, @@ -3077,47 +3106,15 @@ CREATE STATISTICS datetime_stats (dependencies) on datetime, end_datetime from i ALTER TABLE items ADD CONSTRAINT items_collections_fk FOREIGN KEY (collection) REFERENCES collections(id) ON DELETE CASCADE DEFERRABLE; --- partition_after_triggerfunc: After-statement trigger on items. --- Updates partition statistics for every partition touched by the current batch, --- using run_or_queue() so the work is deferred rather than blocking the ingest --- transaction. -CREATE OR REPLACE FUNCTION partition_after_triggerfunc() RETURNS TRIGGER AS $$ -DECLARE - p text; - t timestamptz := clock_timestamp(); -BEGIN - RAISE NOTICE 'Updating partition stats %', t; - FOR p IN SELECT DISTINCT partition - FROM newdata n JOIN partition_sys_meta p - ON (n.collection=p.collection AND n.datetime <@ p.partition_dtrange) - LOOP - PERFORM run_or_queue(format('SELECT update_partition_stats(%L, %L);', p, true)); - END LOOP; - RAISE NOTICE 't: % %', t, clock_timestamp() - t; - RETURN NULL; -END; -$$ LANGUAGE PLPGSQL SECURITY DEFINER; - +-- The old per-statement stats trigger (partition_after_triggerfunc) is intentionally GONE. Under the +-- widen-now / tighten-async model, partition_stats coverage is maintained by the ingest path itself +-- (check_partition seeds + widen_partition_stats covers, in their own transactions) and tightened +-- off the hot path by tighten_partition_stats; a per-insert trigger that re-scanned partitions and +-- refreshed matviews was the lock pathology this model removes. DROP TRIGGER IF EXISTS items_after_insert_trigger ON items; -CREATE TRIGGER items_after_insert_trigger -AFTER INSERT ON items -REFERENCING NEW TABLE AS newdata -FOR EACH STATEMENT -EXECUTE FUNCTION partition_after_triggerfunc(); - DROP TRIGGER IF EXISTS items_after_update_trigger ON items; -CREATE TRIGGER items_after_update_trigger -AFTER UPDATE ON items -REFERENCING NEW TABLE AS newdata -FOR EACH STATEMENT -EXECUTE FUNCTION partition_after_triggerfunc(); - DROP TRIGGER IF EXISTS items_after_delete_trigger ON items; -CREATE TRIGGER items_after_delete_trigger -AFTER DELETE ON items -REFERENCING OLD TABLE AS newdata -FOR EACH STATEMENT -EXECUTE FUNCTION partition_after_triggerfunc(); +DROP FUNCTION IF EXISTS partition_after_triggerfunc(); -- items_content_distinct_sql / items_content_changed: detect whether a direct -- UPDATE actually changed the stored item content. Used by the touch trigger @@ -3168,7 +3165,7 @@ $$ LANGUAGE PLPGSQL IMMUTABLE PARALLEL SAFE; -- items_touch_triggerfunc: refresh pgstac_updated_at when a direct UPDATE changes -- the stored item content. It deliberately does NOT recompute item_hash: --- item_hash is the canonical (RFC 8785-aligned) hash of the item *as ingested* +-- item_hash is the canonical (jsonb_canonical) hash of the item *as ingested* -- through create_item / upsert_item / update_item (set once in content_dehydrate), -- so it stays externally reproducible by a client hashing its own copy. -- A raw `UPDATE items SET ...` that bypasses the staging path leaves item_hash @@ -3381,7 +3378,7 @@ CREATE OR REPLACE FUNCTION content_dehydrate(content jsonb) RETURNS items AS $$ content->>'stac_version' AS stac_version, COALESCE(content->'stac_extensions', '[]'::jsonb) AS stac_extensions, now() AS pgstac_updated_at, - pgstac.jsonb_hash(content) AS item_hash, + pgstac.jsonb_canonical_hash(content) AS item_hash, NULL::bigint AS fragment_id, content->'bbox' AS bbox, CASE WHEN content->'links' IS NOT NULL AND content->'links' <> '[]'::jsonb THEN content->'links' END AS links, @@ -3476,10 +3473,8 @@ BEGIN END; $$ LANGUAGE PLPGSQL IMMUTABLE; --- content_hydrate: Reassemble a full STAC item JSON from the split columns --- and the shared fragment content. This is the single hydrate function; --- the old content_nonhydrated wrapper and 3-arg _collection parameter have --- been removed. +-- content_hydrate: reassemble a full STAC item JSON from the split columns and the shared fragment +-- content. The single hydrate function. CREATE OR REPLACE FUNCTION content_hydrate( _item items, fields jsonb DEFAULT '{}'::jsonb, @@ -3721,6 +3716,22 @@ DECLARE BEGIN RAISE NOTICE 'Creating Partitions. %', clock_timestamp() - ts; + -- Fail loudly on items whose collection does not exist instead of silently dropping them + -- in the collections JOINs below (matches the Rust loader, which also errors on this). + IF EXISTS ( + SELECT 1 FROM newdata n + WHERE NOT EXISTS ( + SELECT 1 FROM collections c WHERE c.id = n.content->>'collection' + ) + ) THEN + RAISE EXCEPTION 'cannot load item(s) into nonexistent collection(s): %', + (SELECT string_agg(DISTINCT coalesce(n.content->>'collection', ''), ', ') + FROM newdata n + WHERE NOT EXISTS ( + SELECT 1 FROM collections c WHERE c.id = n.content->>'collection' + )); + END IF; + FOR part IN WITH t AS ( SELECT n.content->>'collection' as collection, @@ -3777,13 +3788,31 @@ BEGIN RAISE NOTICE 'Inserted % rows to items. %', nrows, clock_timestamp() - ts; END IF; + -- Bump each partition's row-count estimate so search stepping sees the new rows. n drives + -- partition_bounds' histogram and the Rust keyset search_page skips zero-count bands, so a partition + -- left at n=0 would hide its freshly-ingested items from that path (SQL search() scans regardless). + -- Mirrors flush_items_staging_binary on the Rust loader. Over-count is safe (ignore/upsert may insert + -- fewer than staged); the async tightener resets n exactly. check_partition above already widened the + -- temporal envelope; spatial is left to the tightener on this fallback path. (n only affects search + -- stepping performance, never which rows/order come back — see geometrysearch band direction.) + UPDATE partition_stats ps + SET n = COALESCE(ps.n, 0) + agg.c, dirty = true, last_updated = now() + FROM ( + SELECT (partition_name(nd.content->>'collection', + lower(stac_daterange(nd.content->'properties')))).partition_name AS partition, + count(*) AS c + FROM newdata nd + GROUP BY 1 + ) agg + WHERE ps.partition = agg.partition; + RAISE NOTICE 'Deleting data from staging table. %', clock_timestamp() - ts; EXECUTE format('DELETE FROM %I', TG_TABLE_NAME); RAISE NOTICE 'Done. %', clock_timestamp() - ts; RETURN NULL; END; -$$ LANGUAGE PLPGSQL; +$$ LANGUAGE PLPGSQL SECURITY DEFINER; DROP TRIGGER IF EXISTS items_staging_insert_trigger ON items_staging; @@ -3819,7 +3848,7 @@ out items%ROWTYPE; BEGIN DELETE FROM items WHERE id = _id AND (_collection IS NULL OR collection=_collection) RETURNING * INTO STRICT out; END; -$$ LANGUAGE PLPGSQL; +$$ LANGUAGE PLPGSQL SECURITY DEFINER; --/* CREATE OR REPLACE FUNCTION create_item(data jsonb) RETURNS VOID AS $$ @@ -3864,156 +3893,10 @@ CREATE OR REPLACE FUNCTION collection_temporal_extent(id text) RETURNS jsonb AS ; $$ LANGUAGE SQL IMMUTABLE PARALLEL SAFE SET SEARCH_PATH TO pgstac, public; -CREATE OR REPLACE FUNCTION update_collection_extents() RETURNS VOID AS $$ -UPDATE collections - SET content = jsonb_set_lax( - content, - '{extent}'::text[], - collection_extent(id, FALSE), - true, - 'use_json_null' - ) -; -$$ LANGUAGE SQL; - --- --------------------------------------------------------------------------- --- Field Registry: walks JSONB item content to track which paths exist in each --- collection. Used to auto-populate queryables and support schema inference. --- jsonb_field_rows is defined in 001a_jsonutils.sql (loaded first). --- --------------------------------------------------------------------------- - --- update_field_registry_from_sample: UPSERT registry rows from a pre-selected array of --- raw item content JSONBs. Callers supply the sample to decouple sampling strategy --- from the registry write; merge value_kinds to accumulate observed types over time. -CREATE OR REPLACE FUNCTION update_field_registry_from_sample( - _collection text, - item_contents jsonb[] -) RETURNS void AS $$ - INSERT INTO item_field_registry (collection, path, is_leaf, value_kinds, first_seen, last_seen) - SELECT - _collection, - r.path, - bool_and(r.is_leaf) AS is_leaf, - array_agg(DISTINCT r.value_kind) FILTER (WHERE r.value_kind IS NOT NULL) AS value_kinds, - now(), - now() - FROM unnest(item_contents) AS item(content) - CROSS JOIN LATERAL jsonb_field_rows(item.content) AS r(path, is_leaf, value_kind) - GROUP BY r.path - ON CONFLICT (collection, path) DO UPDATE SET - is_leaf = EXCLUDED.is_leaf, - value_kinds = ( - SELECT array_agg(DISTINCT v) - FROM unnest(item_field_registry.value_kinds || EXCLUDED.value_kinds) t(v) - ), - last_seen = now() - ; -$$ LANGUAGE SQL VOLATILE; --- update_field_registry_from_items: Sample a live collection and UPSERT registry rows. --- Uses TABLESAMPLE BERNOULLI(5) for large collections (>10k rows by pg_class estimate) --- and LIMIT 1000 for smaller ones to avoid a full seq-scan for tiny collections. --- pg_class.reltuples is an estimate (may be stale); its only role is threshold selection. --- Returns (registered_paths, rows_processed) for observability. -CREATE OR REPLACE FUNCTION update_field_registry_from_items( - _collection text -) RETURNS TABLE (registered_paths int, rows_processed int) AS $$ -DECLARE - est_rows bigint; - nrows int; - npaths int; -BEGIN - -- Sum reltuples across the registered item partitions for this collection. - -- reltuples can be -1 (never analyzed); treat negative values as zero. - SELECT COALESCE(sum(GREATEST(c.reltuples::bigint, 0)), 0) INTO est_rows - FROM partitions_view p - JOIN pg_class c ON c.relname = p.partition - JOIN pg_namespace n ON n.oid = c.relnamespace - WHERE p.collection = _collection - AND n.nspname = 'pgstac' - AND c.relkind = 'r'; - IF est_rows > 10000 THEN - -- Large collection: use statistical sampling to avoid full seq-scan. - WITH sampled AS ( - SELECT content_hydrate(i) AS content FROM items i TABLESAMPLE BERNOULLI(5) WHERE i.collection = _collection - ), - upserted AS ( - INSERT INTO item_field_registry (collection, path, is_leaf, value_kinds, first_seen, last_seen) - SELECT - _collection, - r.path, - bool_and(r.is_leaf) AS is_leaf, - array_agg(DISTINCT r.value_kind) FILTER (WHERE r.value_kind IS NOT NULL) AS value_kinds, - now(), now() - FROM sampled - CROSS JOIN LATERAL jsonb_field_rows(content) AS r(path, is_leaf, value_kind) - GROUP BY r.path - ON CONFLICT (collection, path) DO UPDATE SET - is_leaf = EXCLUDED.is_leaf, - value_kinds = ( - SELECT array_agg(DISTINCT v) - FROM unnest(item_field_registry.value_kinds || EXCLUDED.value_kinds) t(v) - ), - last_seen = now() - RETURNING 1 - ) - SELECT - (SELECT count(*)::int FROM upserted), - (SELECT count(*)::int FROM sampled) - INTO npaths, nrows; - ELSE - -- Small collection: process up to 1000 rows to avoid BERNOULLI returning 0 rows. - WITH sampled AS ( - SELECT content_hydrate(i) AS content FROM items i WHERE i.collection = _collection LIMIT 1000 - ), - upserted AS ( - INSERT INTO item_field_registry (collection, path, is_leaf, value_kinds, first_seen, last_seen) - SELECT - _collection, - r.path, - bool_and(r.is_leaf) AS is_leaf, - array_agg(DISTINCT r.value_kind) FILTER (WHERE r.value_kind IS NOT NULL) AS value_kinds, - now(), now() - FROM sampled - CROSS JOIN LATERAL jsonb_field_rows(content) AS r(path, is_leaf, value_kind) - GROUP BY r.path - ON CONFLICT (collection, path) DO UPDATE SET - is_leaf = EXCLUDED.is_leaf, - value_kinds = ( - SELECT array_agg(DISTINCT v) - FROM unnest(item_field_registry.value_kinds || EXCLUDED.value_kinds) t(v) - ), - last_seen = now() - RETURNING 1 - ) - SELECT - (SELECT count(*)::int FROM upserted), - (SELECT count(*)::int FROM sampled) - INTO npaths, nrows; - END IF; - RETURN QUERY SELECT npaths, nrows; -END; -$$ LANGUAGE PLPGSQL VOLATILE SECURITY DEFINER; --- refresh_field_registry: Expire stale registry entries that haven't been seen recently. --- Intended for scheduled maintenance (e.g. pg_cron daily job). --- Returns (collection, expired_paths) for each collection affected. -CREATE OR REPLACE FUNCTION refresh_field_registry( - _collection text DEFAULT NULL, - retention_interval interval DEFAULT '90 days' -) RETURNS TABLE (collection_id text, expired_paths int) AS $$ - WITH deleted AS ( - DELETE FROM item_field_registry - WHERE (_collection IS NULL OR collection = _collection) - AND last_seen < now() - retention_interval - RETURNING collection - ) - SELECT collection, count(*)::int - FROM deleted - GROUP BY collection; -$$ LANGUAGE SQL VOLATILE; -- Item Fragment Management functions @@ -4070,37 +3953,6 @@ CREATE OR REPLACE FUNCTION pgstac_hash_fragment(fragment jsonb) RETURNS bytea AS SELECT sha256(convert_to(fragment::text, 'UTF8')); $$ LANGUAGE SQL IMMUTABLE PARALLEL SAFE; --- gc_fragments: Garbage collect orphaned fragments using a single set-based DELETE so the --- planner can choose the optimal join/anti-join strategy across all collections. --- The NOT EXISTS sub-select is evaluated per fragment; with an index on items.fragment_id --- this is an efficient anti-join rather than a full seq-scan. --- --- Operational note: because items.fragment_id is intentionally unmanaged by an FK --- (partitioned-items incremental NOT VALID FKs are not supported), gc_fragments has a --- small race window with concurrent inserts. A fragment that is unreferenced at the time --- the DELETE snapshot is taken but becomes referenced by a later insert could be removed. --- The retention_interval guard makes this unlikely for normal ingest, but operators should --- still run gc_fragments during low-ingest periods or with a sufficiently conservative --- retention interval. This is a documented operational tradeoff, not a silent invariant. -CREATE OR REPLACE FUNCTION gc_fragments( - _collection text DEFAULT NULL, - retention_interval interval DEFAULT '90 days' -) RETURNS TABLE ( - collection_id text, - fragments_removed int -) AS $$ - WITH deleted AS ( - DELETE FROM item_fragments f - WHERE - (_collection IS NULL OR f.collection = _collection) - AND f.created_at < now() - retention_interval - AND NOT EXISTS (SELECT 1 FROM items i WHERE i.fragment_id = f.id) - RETURNING f.collection - ) - SELECT collection, count(*)::int - FROM deleted - GROUP BY collection; -$$ LANGUAGE SQL VOLATILE PARALLEL UNSAFE; -- strip_fragment_col: Remove fragment-owned sub-keys from a split column value. -- col_name is the top-level STAC key that this column represents (e.g. 'assets' or 'properties'). @@ -4151,18 +4003,30 @@ $$ LANGUAGE PLPGSQL IMMUTABLE PARALLEL SAFE; CREATE TABLE partition_stats ( partition text PRIMARY KEY, collection text, - partition_dtrange tstzrange, + -- dtrange/edtrange are the data bounds used for search pruning. widen_partition_stats fills them + -- (generously) on ingest; tighten_partition_stats narrows them to the exact extent. The partition's + -- structural range lives in its own constraint, read via partitions_view.constraint_dtrange. dtrange tstzrange, edtrange tstzrange, spatial geometry, last_updated timestamptz, n bigint, - keys text[] -) WITH (FILLFACTOR=90); + -- Deferred-maintenance flags: + -- dirty: the stored envelope may be WIDER than the actual data; an async tightener + -- should recompute exact min/max/extent and clear the flag. Search correctness + -- never depends on this (a wide envelope only over-includes a partition). + -- indexes_pending: the partition currently carries only the parent-inherited indexes (id PK, + -- datetime, geometry); its queryable-defined indexes have not been built yet. + dirty boolean NOT NULL DEFAULT false, + indexes_pending boolean NOT NULL DEFAULT false +) WITH (FILLFACTOR=70); CREATE INDEX partitions_range_idx ON partition_stats USING GIST(dtrange); CREATE INDEX partition_stats_collection_idx ON partition_stats (collection); CREATE INDEX partition_stats_spatial_idx ON partition_stats USING GIST(spatial) WHERE spatial IS NOT NULL; +-- Work-queue indexes: the maintenance sweeps find pending partitions without scanning every row. +CREATE INDEX partition_stats_dirty_idx ON partition_stats (partition) WHERE dirty; +CREATE INDEX partition_stats_indexes_pending_idx ON partition_stats (partition) WHERE indexes_pending; CREATE OR REPLACE FUNCTION constraint_tstzrange(expr text) RETURNS tstzrange AS $$ @@ -4240,10 +4104,9 @@ SELECT level, c.reltuples, c.relhastriggers, - partition_dtrange, COALESCE( get_tstz_constraint(c.oid, 'datetime'), - partition_dtrange, + constraint_tstzrange(pg_get_expr(c.relpartbound, c.oid)), inf_range ) as constraint_dtrange, COALESCE( @@ -4259,7 +4122,6 @@ FROM JOIN LATERAL pg_get_expr(c.relpartbound, c.oid) as partition_expr ON TRUE JOIN LATERAL pg_get_expr(parent.relpartbound, parent.oid) as parent_partition_expr ON TRUE JOIN LATERAL tstzrange('-infinity', 'infinity','[]') as inf_range ON TRUE - JOIN LATERAL COALESCE(constraint_tstzrange(pg_get_expr(c.relpartbound, c.oid)), inf_range) as partition_dtrange ON TRUE JOIN LATERAL get_tstz_constraint(c.oid, 'datetime') as datetime_constraint ON TRUE JOIN LATERAL get_tstz_constraint(c.oid, 'end_datetime') as end_datetime_constraint ON TRUE WHERE isleaf @@ -4280,10 +4142,9 @@ SELECT level, c.reltuples, c.relhastriggers, - partition_dtrange, COALESCE( get_tstz_constraint(c.oid, 'datetime'), - partition_dtrange, + constraint_tstzrange(pg_get_expr(c.relpartbound, c.oid)), inf_range ) as constraint_dtrange, COALESCE( @@ -4303,10 +4164,9 @@ FROM JOIN LATERAL pg_get_expr(c.relpartbound, c.oid) as partition_expr ON TRUE JOIN LATERAL pg_get_expr(parent.relpartbound, parent.oid) as parent_partition_expr ON TRUE JOIN LATERAL tstzrange('-infinity', 'infinity','[]') as inf_range ON TRUE - JOIN LATERAL COALESCE(constraint_tstzrange(pg_get_expr(c.relpartbound, c.oid)), inf_range) as partition_dtrange ON TRUE JOIN LATERAL get_tstz_constraint(c.oid, 'datetime') as datetime_constraint ON TRUE JOIN LATERAL get_tstz_constraint(c.oid, 'end_datetime') as end_datetime_constraint ON TRUE - -- the view computes its own collection/partition_dtrange from the live tree; pull only the + -- the view computes its own collection + constraint_dtrange from the live tree; pull only the -- data-extent columns from partition_stats to avoid colliding with those names. LEFT JOIN ( SELECT partition, dtrange, edtrange, spatial, last_updated @@ -4315,109 +4175,6 @@ FROM WHERE isleaf ; -CREATE MATERIALIZED VIEW partitions AS -SELECT * FROM partitions_view; -CREATE UNIQUE INDEX ON partitions (partition); - -CREATE MATERIALIZED VIEW partition_steps AS -SELECT - partition as name, - date_trunc('month',lower(partition_dtrange)) as sdate, - date_trunc('month', upper(partition_dtrange)) + '1 month'::interval as edate - FROM partitions_view WHERE partition_dtrange IS NOT NULL AND partition_dtrange != 'empty'::tstzrange - ORDER BY dtrange ASC -; - - -CREATE OR REPLACE FUNCTION update_partition_stats_q(_partition text, istrigger boolean default false) RETURNS VOID AS $$ -DECLARE -BEGIN - PERFORM run_or_queue( - format('SELECT update_partition_stats(%L, %L);', _partition, istrigger) - ); -END; -$$ LANGUAGE PLPGSQL; - -CREATE OR REPLACE FUNCTION update_partition_stats(_partition text, istrigger boolean default false) RETURNS VOID AS $$ -DECLARE - dtrange tstzrange; - edtrange tstzrange; - cdtrange tstzrange; - cedtrange tstzrange; - extent geometry; - collection text; - _part_dtrange tstzrange; - _n bigint; -BEGIN - RAISE NOTICE 'Updating stats for %.', _partition; - EXECUTE format( - $q$ - SELECT - tstzrange(min(datetime), max(datetime),'[]'), - tstzrange(min(end_datetime), max(end_datetime), '[]') - FROM %I - $q$, - _partition - ) INTO dtrange, edtrange; - EXECUTE format('ANALYZE %I;', _partition); - extent := st_estimatedextent('pgstac', _partition, 'geometry'); - RAISE DEBUG 'Estimated Extent: %', extent; - - -- Per-partition metadata: collection + partition boundary from the live tree (partitions_view), - -- current constraint ranges (for the constraint check below), and the row estimate from pg_class. - SELECT pv.collection, pv.partition_dtrange, pv.constraint_dtrange, pv.constraint_edtrange - INTO collection, _part_dtrange, cdtrange, cedtrange - FROM partitions_view pv WHERE partition = _partition; - _n := (SELECT reltuples::bigint FROM pg_class WHERE oid = quote_ident(_partition)::regclass); - - INSERT INTO partition_stats - (partition, collection, partition_dtrange, dtrange, edtrange, spatial, n, last_updated) - SELECT _partition, collection, _part_dtrange, dtrange, edtrange, extent, _n, now() - ON CONFLICT (partition) DO - UPDATE SET - collection=EXCLUDED.collection, - partition_dtrange=EXCLUDED.partition_dtrange, - dtrange=EXCLUDED.dtrange, - edtrange=EXCLUDED.edtrange, - spatial=EXCLUDED.spatial, - n=EXCLUDED.n, - last_updated=EXCLUDED.last_updated - ; - - RAISE NOTICE 'Checking if we need to modify constraints...'; - RAISE NOTICE 'cdtrange: % dtrange: % cedtrange: % edtrange: %',cdtrange, dtrange, cedtrange, edtrange; - IF - (cdtrange IS DISTINCT FROM dtrange OR edtrange IS DISTINCT FROM cedtrange) - AND NOT istrigger - THEN - RAISE NOTICE 'Modifying Constraints'; - RAISE NOTICE 'Existing % %', cdtrange, cedtrange; - RAISE NOTICE 'New % %', dtrange, edtrange; - PERFORM drop_table_constraints(_partition); - PERFORM create_table_constraints(_partition, dtrange, edtrange); - END IF; - REFRESH MATERIALIZED VIEW partitions; - REFRESH MATERIALIZED VIEW partition_steps; - RAISE NOTICE 'Checking if we need to update collection extents.'; - IF get_setting_bool('update_collection_extent') THEN - RAISE NOTICE 'updating collection extent for %', collection; - PERFORM run_or_queue(format($q$ - UPDATE collections - SET content = jsonb_set_lax( - content, - '{extent}'::text[], - collection_extent(%L, FALSE), - true, - 'use_json_null' - ) WHERE id=%L - ; - $q$, collection, collection)); - ELSE - RAISE NOTICE 'Not updating collection extent for %', collection; - END IF; - -END; -$$ LANGUAGE PLPGSQL STRICT SECURITY DEFINER; CREATE OR REPLACE FUNCTION partition_name( IN collection text, IN dt timestamptz, OUT partition_name text, OUT partition_range tstzrange) AS $$ @@ -4452,107 +4209,82 @@ END; $$ LANGUAGE PLPGSQL STABLE; -CREATE OR REPLACE FUNCTION drop_table_constraints(t text) RETURNS text AS $$ -DECLARE - q text; -BEGIN - IF NOT EXISTS (SELECT 1 FROM partitions_view WHERE partition=t) THEN - RETURN NULL; - END IF; - FOR q IN SELECT FORMAT( - $q$ - ALTER TABLE %I DROP CONSTRAINT IF EXISTS %I; - $q$, - t, - conname - ) FROM pg_constraint - WHERE conrelid=t::regclass::oid AND contype='c' - LOOP - EXECUTE q; - END LOOP; - RETURN t; -END; -$$ LANGUAGE PLPGSQL SECURITY DEFINER; +-- Partitions carry no _dt CHECK constraints; pruning is partition_stats-driven. get_tstz_constraint / +-- constraint_tstzrange remain so the views' constraint_* columns resolve to inf_range when a partition +-- has no datetime constraint. -CREATE OR REPLACE FUNCTION create_table_constraints(t text, _dtrange tstzrange, _edtrange tstzrange) RETURNS text AS $$ + +-- widen_partition_stats: widen a partition's stored envelope to cover a batch. A no-op (snapshot read, no +-- write, no lock) when the envelope already covers the batch; otherwise widens and sets dirty=true: +-- * datetime : sub-partitioned collections widen to the partition's datetime bound (covers anything +-- that can land there); a NULL-partition_trunc partition pads the batch range by +-- partition_stats_widen_buffer (default 1 month) each side. +-- * end_datetime: the datetime target extended by the batch's max (end_datetime - datetime) tail. +-- * spatial : NULL means "always a search candidate"; a spatial miss resets spatial to NULL until +-- the tightener computes the real extent. +-- Requires the partition_stats row to exist (check_partition seeds it); raises if it does not. +CREATE OR REPLACE FUNCTION widen_partition_stats( + _partition text, + _dtrange tstzrange, + _edtrange tstzrange, + _spatial geometry DEFAULT NULL, + _constraint_dtrange tstzrange DEFAULT NULL +) RETURNS void AS $$ DECLARE - q text; + cur RECORD; + is_unbounded boolean; + tail interval; + buf interval := COALESCE(get_setting('partition_stats_widen_buffer'), '1 month')::interval; + target_dtrange tstzrange; + target_edtrange tstzrange; + new_dtrange tstzrange; + new_edtrange tstzrange; + new_spatial geometry; + dt_covered boolean; + edt_covered boolean; + spatial_covered boolean; BEGIN - IF NOT EXISTS (SELECT 1 FROM partitions_view WHERE partition=t) THEN - RETURN NULL; - END IF; - RAISE NOTICE 'Creating Table Constraints for % % %', t, _dtrange, _edtrange; - IF _dtrange = 'empty' AND _edtrange = 'empty' THEN - q :=format( - $q$ - DO $block$ - BEGIN - ALTER TABLE %I DROP CONSTRAINT IF EXISTS %I; - ALTER TABLE %I - ADD CONSTRAINT %I - CHECK (((datetime IS NULL) AND (end_datetime IS NULL))) NOT VALID - ; - ALTER TABLE %I - VALIDATE CONSTRAINT %I - ; - - - - EXCEPTION WHEN others THEN - RAISE WARNING '%%, Issue Altering Constraints. Please run update_partition_stats(%I)', SQLERRM USING ERRCODE = SQLSTATE; - END; - $block$; - $q$, - t, - format('%s_dt', t), - t, - format('%s_dt', t), - t, - format('%s_dt', t), - t - ); + SELECT dtrange, edtrange, spatial + INTO cur + FROM partition_stats WHERE partition = _partition; + IF NOT FOUND THEN + RAISE EXCEPTION 'partition_stats row for % does not exist; call check_partition first', _partition; + END IF; + + dt_covered := COALESCE(cur.dtrange @> _dtrange, false); + edt_covered := COALESCE(cur.edtrange @> _edtrange, false); + spatial_covered := cur.spatial IS NULL OR _spatial IS NULL OR ST_Covers(cur.spatial, _spatial); + IF dt_covered AND edt_covered AND spatial_covered THEN + RETURN; -- already covered: no write, no lock + END IF; + + -- Clamp the widen to the partition's structural bound (_constraint_dtrange). A sub-partitioned + -- (month/year) partition's bound is finite: cover the whole partition so dtrange never misses and + -- never spills into a neighbour. An unbounded partition (NULL-partition_trunc or NULL bound) has + -- nothing to clamp to, so pad the batch range by the widen buffer each side. + is_unbounded := _constraint_dtrange IS NULL + OR (lower(_constraint_dtrange) = '-infinity'::timestamptz + AND upper(_constraint_dtrange) = 'infinity'::timestamptz); + tail := GREATEST(upper(_edtrange) - upper(_dtrange), '0'::interval); + IF is_unbounded THEN + target_dtrange := tstzrange(lower(_dtrange) - buf, upper(_dtrange) + buf, '[]'); + target_edtrange := tstzrange(lower(_dtrange) - buf, upper(_edtrange) + buf, '[]'); ELSE - q :=format( - $q$ - DO $block$ - BEGIN - - ALTER TABLE %I DROP CONSTRAINT IF EXISTS %I; - ALTER TABLE %I - ADD CONSTRAINT %I - CHECK ( - (datetime >= %L) - AND (datetime <= %L) - AND (end_datetime >= %L) - AND (end_datetime <= %L) - ) NOT VALID - ; - ALTER TABLE %I - VALIDATE CONSTRAINT %I - ; - - - - EXCEPTION WHEN others THEN - RAISE WARNING '%%, Issue Altering Constraints. Please run update_partition_stats(%I)', SQLERRM USING ERRCODE = SQLSTATE; - END; - $block$; - $q$, - t, - format('%s_dt', t), - t, - format('%s_dt', t), - lower(_dtrange), - upper(_dtrange), - lower(_edtrange), - upper(_edtrange), - t, - format('%s_dt', t), - t - ); + target_dtrange := _constraint_dtrange; + target_edtrange := tstzrange(lower(_constraint_dtrange), upper(_constraint_dtrange) + tail, '[]'); END IF; - PERFORM run_or_queue(q); - RETURN t; + + new_dtrange := range_merge(COALESCE(cur.dtrange, target_dtrange), target_dtrange); + new_edtrange := range_merge(COALESCE(cur.edtrange, target_edtrange), target_edtrange); + new_spatial := CASE WHEN spatial_covered THEN cur.spatial ELSE NULL END; + + UPDATE partition_stats + SET dtrange = new_dtrange, + edtrange = new_edtrange, + spatial = new_spatial, + dirty = true, + last_updated = now() + WHERE partition = _partition; END; $$ LANGUAGE PLPGSQL SECURITY DEFINER; @@ -4560,18 +4292,14 @@ $$ LANGUAGE PLPGSQL SECURITY DEFINER; CREATE OR REPLACE FUNCTION check_partition( _collection text, _dtrange tstzrange, - _edtrange tstzrange + _edtrange tstzrange, + _spatial geometry DEFAULT NULL ) RETURNS text AS $$ DECLARE c RECORD; - pm RECORD; _partition_name text; - _partition_dtrange tstzrange; - _constraint_dtrange tstzrange; - _constraint_edtrange tstzrange; - q text; - deferrable_q text; - err_context text; + _parent_name text; + _partition_range tstzrange; BEGIN SELECT * INTO c FROM pgstac.collections WHERE id=_collection; IF NOT FOUND THEN @@ -4579,115 +4307,99 @@ BEGIN END IF; IF c.partition_trunc IS NOT NULL THEN - _partition_dtrange := tstzrange( + _partition_range := tstzrange( date_trunc(c.partition_trunc, lower(_dtrange)), date_trunc(c.partition_trunc, lower(_dtrange)) + (concat('1 ', c.partition_trunc))::interval, '[)' ); ELSE - _partition_dtrange := '[-infinity, infinity]'::tstzrange; + _partition_range := '[-infinity, infinity]'::tstzrange; END IF; - IF NOT _partition_dtrange @> _dtrange THEN - RAISE EXCEPTION 'dtrange % is greater than the partition size % for collection %', _dtrange, c.partition_trunc, _collection; + IF NOT _partition_range @> _dtrange THEN + RAISE EXCEPTION 'dtrange % spans more than the % partition window for collection %', _dtrange, c.partition_trunc, _collection; END IF; - + _parent_name := format('_items_%s', c.key); IF c.partition_trunc = 'year' THEN - _partition_name := format('_items_%s_%s', c.key, to_char(lower(_partition_dtrange),'YYYY')); + _partition_name := format('%s_%s', _parent_name, to_char(lower(_partition_range),'YYYY')); ELSIF c.partition_trunc = 'month' THEN - _partition_name := format('_items_%s_%s', c.key, to_char(lower(_partition_dtrange),'YYYYMM')); + _partition_name := format('%s_%s', _parent_name, to_char(lower(_partition_range),'YYYYMM')); ELSE - _partition_name := format('_items_%s', c.key); - END IF; - - SELECT * INTO pm FROM partition_sys_meta WHERE collection=_collection AND partition_dtrange @> _dtrange; - IF FOUND THEN - RAISE NOTICE '% % %', _edtrange, _dtrange, pm; - _constraint_edtrange := - tstzrange( - least( - lower(_edtrange), - nullif(lower(pm.constraint_edtrange), '-infinity') - ), - greatest( - upper(_edtrange), - nullif(upper(pm.constraint_edtrange), 'infinity') - ), - '[]' + _partition_name := _parent_name; + END IF; + + -- Create the collection-level PARENT partition (_items_) first, for sub-partitioned collections. + -- It is shared across every child window, so concurrent setup of different children would race on its + -- CREATE TABLE. Guard it with a parent-scoped advisory lock, taken before any child lock and only when + -- the parent is missing: parent-before-child is one lock order (parent < child) that can't deadlock + -- with ensure_partitions' sorted child locks, and skipping it once the parent exists keeps steady-state + -- ingest off the parent lock. + IF c.partition_trunc IS NOT NULL AND to_regclass(format('pgstac.%I', _parent_name)) IS NULL THEN + PERFORM pg_advisory_xact_lock(hashtext('pgstac.check_partition'), hashtext(_parent_name)); + IF to_regclass(format('pgstac.%I', _parent_name)) IS NULL THEN + EXECUTE format( + 'CREATE TABLE IF NOT EXISTS %I partition OF items FOR VALUES IN (%L) PARTITION BY RANGE (datetime)', + _parent_name, _collection ); - _constraint_dtrange := - tstzrange( - least( - lower(_dtrange), - nullif(lower(pm.constraint_dtrange), '-infinity') - ), - greatest( - upper(_dtrange), - nullif(upper(pm.constraint_dtrange), 'infinity') - ), - '[]' - ); - - IF pm.constraint_edtrange @> _edtrange AND pm.constraint_dtrange @> _dtrange THEN - RETURN pm.partition; - ELSE - PERFORM drop_table_constraints(_partition_name); END IF; - ELSE - _constraint_edtrange := _edtrange; - _constraint_dtrange := _dtrange; END IF; - RAISE NOTICE 'EXISTING CONSTRAINTS % %, NEW % %', pm.constraint_dtrange, pm.constraint_edtrange, _constraint_dtrange, _constraint_edtrange; - RAISE NOTICE 'Creating partition % %', _partition_name, _partition_dtrange; - IF c.partition_trunc IS NULL THEN - q := format( - $q$ - CREATE TABLE IF NOT EXISTS %I partition OF items FOR VALUES IN (%L); - CREATE UNIQUE INDEX IF NOT EXISTS %I ON %I (id); - GRANT ALL ON %I to pgstac_ingest; - $q$, - _partition_name, - _collection, - concat(_partition_name,'_pk'), - _partition_name, - _partition_name - ); - ELSE - q := format( - $q$ - CREATE TABLE IF NOT EXISTS %I partition OF items FOR VALUES IN (%L) PARTITION BY RANGE (datetime); - CREATE TABLE IF NOT EXISTS %I partition OF %I FOR VALUES FROM (%L) TO (%L); - CREATE UNIQUE INDEX IF NOT EXISTS %I ON %I (id); - GRANT ALL ON %I TO pgstac_ingest; - $q$, - format('_items_%s', c.key), - _collection, - _partition_name, - format('_items_%s', c.key), - lower(_partition_dtrange), - upper(_partition_dtrange), - format('%s_pk', _partition_name), - _partition_name, - _partition_name - ); + + -- Serialize concurrent setup of THIS (leaf) partition with a partition-scoped advisory lock. Different + -- partitions hash to different keys, so this never serializes unrelated ingest; ensure_partitions + -- acquires these in sorted order, so concurrent multi-partition batches cannot deadlock. The lock + -- releases at commit (setup is fast + idempotent). + PERFORM pg_advisory_xact_lock(hashtext('pgstac.check_partition'), hashtext(_partition_name)); + + -- Create the leaf partition if missing. Parent-inherited indexes only (id PK here; datetime/geometry + -- come from the items parent); no CHECK constraints; queryable indexes are deferred via + -- indexes_pending. A SELECT grant lets read/ingest query it; writes reach it only through the + -- SECURITY DEFINER write functions (the privilege wall in 998_idempotent_post). + -- Skip the DDL when it already exists: re-running CREATE/GRANT takes a relation lock that deadlocks + -- with concurrent INSERTs. The existence check is race-safe under the advisory lock. + IF to_regclass(format('pgstac.%I', _partition_name)) IS NULL THEN + BEGIN + IF c.partition_trunc IS NULL THEN + EXECUTE format( + $q$ + CREATE TABLE IF NOT EXISTS %I partition OF items FOR VALUES IN (%L); + CREATE UNIQUE INDEX IF NOT EXISTS %I ON %I (id); + GRANT SELECT ON %I TO pgstac_read, pgstac_ingest; + $q$, + _partition_name, _collection, + concat(_partition_name, '_pk'), _partition_name, + _partition_name + ); + ELSE + EXECUTE format( + $q$ + CREATE TABLE IF NOT EXISTS %I partition OF %I FOR VALUES FROM (%L) TO (%L); + CREATE UNIQUE INDEX IF NOT EXISTS %I ON %I (id); + GRANT SELECT ON %I TO pgstac_read, pgstac_ingest; + $q$, + _partition_name, _parent_name, lower(_partition_range), upper(_partition_range), + concat(_partition_name, '_pk'), _partition_name, + _partition_name + ); + END IF; + EXCEPTION + -- A concurrent creator that finished between our checks: benign, it exists now. + WHEN duplicate_table THEN + RAISE DEBUG 'Partition % already exists.', _partition_name; + -- Do NOT swallow other errors: a failed creation must propagate so the caller never goes on to + -- write a partition that does not exist (invariant: check_partition succeeds before use). + END; END IF; - BEGIN - EXECUTE q; - EXCEPTION - WHEN duplicate_table THEN - RAISE NOTICE 'Partition % already exists.', _partition_name; - WHEN others THEN - GET STACKED DIAGNOSTICS err_context = PG_EXCEPTION_CONTEXT; - RAISE INFO 'Error Name:%',SQLERRM; - RAISE INFO 'Error State:%', SQLSTATE; - RAISE INFO 'Error Context:%', err_context; - END; - PERFORM maintain_partitions(_partition_name); - PERFORM update_partition_stats_q(_partition_name, true); - REFRESH MATERIALIZED VIEW partitions; - REFRESH MATERIALIZED VIEW partition_steps; + -- Seed the partition_stats row (collection only — the data ranges start NULL), then cover this batch + -- via the shared widen guard, which fills dtrange/edtrange. ON CONFLICT keeps an existing row (so a + -- sweep that cleared indexes_pending/dirty is not reset on a later check_partition for the same + -- partition). + INSERT INTO partition_stats (partition, collection, dirty, indexes_pending) + VALUES (_partition_name, _collection, true, true) + ON CONFLICT (partition) DO NOTHING; + PERFORM widen_partition_stats(_partition_name, _dtrange, _edtrange, _spatial, _partition_range); + RETURN _partition_name; END; $$ LANGUAGE PLPGSQL SECURITY DEFINER; @@ -4768,6 +4480,219 @@ UPDATE ON collections FOR EACH ROW EXECUTE FUNCTION collections_trigger_func(); -- END FRAGMENT: 003b_partitions.sql +-- BEGIN FRAGMENT: 003c_ingest.sql +-- --------------------------------------------------------------------------- +-- Rust-first ingest: SECURITY DEFINER staging + flush + fragment helpers. +-- +-- These are the server-side seam for the Rust loader's binary-COPY path. The connecting role never +-- writes the real tables directly; it binary-COPYs fully-dehydrated rows into a session-local TEMP +-- table (make_binary_staging) and calls flush_items_staging_binary, which is the only thing that writes +-- `items`. Partition existence + stats (extent + n) and registry/fragment coverage are established BEFORE +-- the load in their own transactions (prepare_partition_for_load / ensure_fragments); the flush writes only +-- `items` and never touches partition_stats. +-- --------------------------------------------------------------------------- + +-- ensure_fragments: upsert per-collection fragment payloads and return, for each input position, the +-- item_fragments id the Rust loader stamps into items.fragment_id. +-- +-- The caller (the Rust loader) deduplicates fragments locally and sends only the DISTINCT set (one input +-- element per unique fragment), then maps each item -> its fragment via the returned `ord`. Sending one +-- fragment per item would ship millions of duplicates when a collection has only a handful of distinct +-- fragments. The function is nonetheless dup-safe (a DISTINCT ON guards the insert), so a caller that +-- passes duplicates still gets the correct id for every position. +-- +-- Each input element is {"content": , "links_template": }; the canonical +-- hash matches pgstac_hash_fragment (so it dedups identically to the SQL ingest path). ON CONFLICT keeps +-- existing rows; pre-existing ids come from item_fragments (snapshot), new ids from the INSERT's RETURNING. +CREATE OR REPLACE FUNCTION ensure_fragments( + _collection text, + _fragments jsonb[] +) RETURNS TABLE (ord int, frag_id bigint) AS $$ + WITH input AS ( + SELECT + o::int AS ord, + pgstac_hash_fragment( + jsonb_strip_nulls(jsonb_build_object( + 'content', NULLIF(f->'content', '{}'::jsonb), + 'links_template', f->'links_template' + )) + ) AS hash, + COALESCE(f->'content', '{}'::jsonb) AS content, + f->'links_template' AS links_template + FROM unnest(_fragments) WITH ORDINALITY AS u(f, o) + ), + distinct_hashes AS ( + SELECT DISTINCT ON (hash) hash, content, links_template + FROM input + ORDER BY hash + ), + upserted AS ( + INSERT INTO item_fragments (collection, hash, content, links_template) + SELECT _collection, hash, content, links_template FROM distinct_hashes + -- DO UPDATE (a no-op write) rather than DO NOTHING: it locks + RETURNS the conflicting row even when + -- a CONCURRENT transaction committed the same (collection, hash) mid-statement. DO NOTHING returns + -- nothing on conflict, and the old "UNION the existing rows via a separate SELECT" ran on the + -- statement-start snapshot, so it could MISS such a concurrent insert -> that hash dropped out of the + -- final join and the item was stamped with NO fragment_id. DO UPDATE returns every distinct hash + -- exactly once (newly inserted or pre-existing), making the stamp concurrency-safe. + ON CONFLICT (collection, hash) DO UPDATE SET content = item_fragments.content + RETURNING id, hash + ) + SELECT input.ord, upserted.id + FROM input JOIN upserted ON input.hash = upserted.hash + ORDER BY input.ord; +$$ LANGUAGE SQL SECURITY DEFINER; + + +-- ensure_partitions: create + widen every partition a batch will land in, BEFORE the load (widen-now). +-- The Rust loader passes parallel arrays of (collection, datetime, end_datetime) — one element per item. +-- Grouping happens HERE, server-side, so the month/year truncation uses the server's date_trunc (correct +-- timezone) rather than re-deriving partition boundaries in the client. One call per batch replaces the +-- loader's per-item check_partition round trips: it groups by (collection, partition window) and calls +-- check_partition once per distinct partition with that group's datetime range. Runs as its own committed +-- statement before the load transaction, so partitions exist + their stats cover the batch by the time +-- the binary COPY + flush run. +CREATE OR REPLACE FUNCTION ensure_partitions( + _collections text[], + _datetimes timestamptz[], + _end_datetimes timestamptz[] +) RETURNS void AS $$ +DECLARE + r RECORD; +BEGIN + -- Call check_partition per distinct (collection, partition window) in a DETERMINISTIC order: each + -- check_partition takes that partition's advisory lock, so locking in a consistent order prevents two + -- concurrent multi-partition batches from advisory-deadlocking on the locks. (flush locks in the same + -- sorted order.) A PLPGSQL FOR loop guarantees the order; a set-returning SELECT would not. + FOR r IN + WITH t AS ( + SELECT + unnest(_collections) AS collection, + unnest(_datetimes) AS dt, + unnest(_end_datetimes) AS edt + ), + j AS ( + SELECT t.collection, t.dt, t.edt, c.partition_trunc + FROM t JOIN pgstac.collections c ON t.collection = c.id + ) + SELECT + collection, + tstzrange(min(dt), max(dt), '[]') AS dtrange, + tstzrange(min(edt), max(edt), '[]') AS edtrange + FROM j + GROUP BY collection, COALESCE(date_trunc(partition_trunc::text, dt), '-infinity'::timestamptz) + ORDER BY collection, COALESCE(date_trunc(partition_trunc::text, dt), '-infinity'::timestamptz) + LOOP + PERFORM pgstac.check_partition(r.collection, r.dtrange, r.edtrange); + END LOOP; +END; +$$ LANGUAGE PLPGSQL SECURITY DEFINER; + + +-- prepare_partition_for_load: per-partition metadata for the Rust direct / precheck load paths. ONE small +-- self-contained txn per partition, run BEFORE any COPY: create + widen the partition to cover this batch +-- (dt + edt + the real SPATIAL envelope, unlike ensure_partitions which passes NULL) and bump n, so +-- partition_stats is at least as wide as the data (golden rule) and search treats the partition as non-empty +-- before the data lands. Returns the partition name + the pre-load row count (n BEFORE the bump) so the +-- loader can choose its adaptive precheck path: empty -> skip the precheck; batch > n -> pull the partition's +-- (id,item_hash) to the client; n >= batch -> COPY the batch (id,hash) to a temp table + JOIN this partition. +CREATE OR REPLACE FUNCTION prepare_partition_for_load( + _collection text, + _dt_lo timestamptz, _dt_hi timestamptz, + _edt_lo timestamptz, _edt_hi timestamptz, + _xmin float8, _ymin float8, _xmax float8, _ymax float8, + _n_add bigint, + OUT partition_name text, + OUT pre_load_n bigint +) AS $$ +BEGIN + partition_name := pgstac.check_partition( + _collection, + tstzrange(_dt_lo, _dt_hi, '[]'), + tstzrange(_edt_lo, _edt_hi, '[]'), + st_setsrid(st_makeenvelope(_xmin, _ymin, _xmax, _ymax), 4326) + ); + SELECT COALESCE(n, 0) INTO pre_load_n + FROM pgstac.partition_stats WHERE partition = partition_name; + -- Over-estimating n is the safe direction; the async tightener computes the exact count + extent off the + -- hot path. Single-row atomic UPDATE, so concurrent loads into the same partition serialize on the row. + UPDATE pgstac.partition_stats + SET n = COALESCE(n, 0) + _n_add, dirty = true, last_updated = now() + WHERE partition = partition_name; +END; +$$ LANGUAGE PLPGSQL SECURITY DEFINER; + + +-- make_binary_staging: create a session-local TEMP staging table shaped exactly like `items` +-- (ON COMMIT DROP) and grant INSERT to pgstac_ingest so the connecting role can binary-COPY into it. +-- Created inside this SECURITY DEFINER function, so the temp table is owned by pgstac_admin (which also +-- owns `items`), letting flush_items_staging_binary read it. Returns the generated table name. +CREATE OR REPLACE FUNCTION make_binary_staging() RETURNS text AS $$ +DECLARE + _name text := format('_staging_%s', replace(gen_random_uuid()::text, '-', '')); +BEGIN + EXECUTE format('CREATE TEMP TABLE %I (LIKE items) ON COMMIT DROP', _name); + EXECUTE format('GRANT INSERT ON pg_temp.%I TO pgstac_ingest', _name); + RETURN _name; +END; +$$ LANGUAGE PLPGSQL SECURITY DEFINER; + + +-- flush_items_staging_binary: move fully-dehydrated rows from a TEMP staging table into `items` with the +-- conflict policy. partition_stats (extent + n + dirty) is set entirely by prepare_partition_for_load in +-- the preflight, so this writes only `items`: +-- ignore -> ON CONFLICT DO NOTHING (idempotent; orphans the old row on a cross-partition move) +-- upsert -> delete a changed row IN THE PARTITION IT ROUTES TO (window-pruned), then insert. A +-- datetime change that stays in the partition is applied; one that moves the item to a +-- different partition orphans the old row (use 'delsert'). +-- delsert -> delete a changed row CROSS-partition (by collection+id, move-safe), then insert +-- error -> plain INSERT (raises on any duplicate) +-- Returns the number of rows inserted. +CREATE OR REPLACE FUNCTION flush_items_staging_binary( + _staging text, + _policy text DEFAULT 'ignore' +) RETURNS bigint AS $$ +DECLARE + nrows bigint; +BEGIN + IF _policy = 'ignore' THEN + EXECUTE format('INSERT INTO items SELECT * FROM %1$I ON CONFLICT DO NOTHING', _staging); + ELSIF _policy = 'error' THEN + EXECUTE format('INSERT INTO items SELECT * FROM %1$I', _staging); + ELSIF _policy = 'upsert' THEN + -- Fast SAME-partition upsert: delete a changed row only in the partition the incoming row routes to. + -- The window range [date_trunc(trunc, s.datetime), + 1 trunc) is derived per staged row, so the + -- planner runtime-prunes `items` to that one partition (no cross-partition scan). A datetime change + -- within the partition is applied; one that MOVES the item to another partition isn't seen here, so + -- the old row orphans (use 'delsert'). NULL partition_trunc => the collection's single partition. + EXECUTE format($q$ + DELETE FROM items i USING %1$I s JOIN collections c ON c.id = s.collection + WHERE i.collection = s.collection AND i.id = s.id + AND i.datetime >= (CASE WHEN c.partition_trunc IS NULL THEN '-infinity'::timestamptz + ELSE date_trunc(c.partition_trunc::text, s.datetime) END) + AND i.datetime < (CASE WHEN c.partition_trunc IS NULL THEN 'infinity'::timestamptz + ELSE date_trunc(c.partition_trunc::text, s.datetime) + + ('1 ' || c.partition_trunc::text)::interval END) + AND ( %2$s ) + $q$, _staging, items_content_distinct_sql('i', 's')); + EXECUTE format('INSERT INTO items SELECT * FROM %1$I ON CONFLICT DO NOTHING', _staging); + ELSIF _policy = 'delsert' THEN + -- Move-safe CROSS-partition upsert: delete the old row wherever it lives (by collection+id, no + -- datetime bound, so it probes every partition), then insert. + EXECUTE format($q$ + DELETE FROM items i USING %1$I s + WHERE i.id = s.id AND i.collection = s.collection AND ( %2$s ) + $q$, _staging, items_content_distinct_sql('i', 's')); + EXECUTE format('INSERT INTO items SELECT * FROM %1$I ON CONFLICT DO NOTHING', _staging); + ELSE + RAISE EXCEPTION 'unknown conflict policy % (expected ignore | upsert | delsert | error)', _policy; + END IF; + GET DIAGNOSTICS nrows = ROW_COUNT; + RETURN nrows; +END; +$$ LANGUAGE PLPGSQL SECURITY DEFINER; +-- END FRAGMENT: 003c_ingest.sql + -- BEGIN FRAGMENT: 004_search.sql -- Search hashing @@ -4993,8 +4918,16 @@ $$ LANGUAGE PLPGSQL STABLE PARALLEL SAFE; -- fields_to_itemcols: produce a SELECT list of item columns in attnum order. -- Heavy columns (geometry, bbox, assets, links, properties, extra) are emitted as --- NULL::type when their controlling field is excluded/not-included. -CREATE OR REPLACE FUNCTION fields_to_itemcols(fields jsonb DEFAULT '{}'::jsonb) RETURNS text AS $$ +-- NULL::type AS when their controlling field is excluded/not-included: the +-- column is kept (named) so the result schema is stable across any `fields` value +-- -- a client preparing this query once can step bands without re-binding columns -- +-- while the heavy value itself is never transferred. +-- When _skip_fragment is true (the caller determined the requested fields are +-- satisfiable from item columns alone, via needs_fragment), fragment_id is emitted +-- as NULL too: a client driving this query then sees no fragment_id and skips the +-- per-row item_fragments lookup entirely -- the fragment-skip is carried in the +-- projection itself, not a separate flag. +CREATE OR REPLACE FUNCTION fields_to_itemcols(fields jsonb DEFAULT '{}'::jsonb, _skip_fragment boolean DEFAULT false) RETURNS text AS $$ DECLARE includes text[] := ARRAY(SELECT jsonb_array_elements_text(fields->'include')); excludes text[] := ARRAY(SELECT jsonb_array_elements_text(fields->'exclude')); @@ -5002,11 +4935,13 @@ DECLARE BEGIN SELECT string_agg( CASE + WHEN a.attname = 'fragment_id' AND _skip_fragment + THEN format('NULL::%s AS %I', format_type(a.atttypid, a.atttypmod), a.attname) WHEN a.attname = ANY (ARRAY['geometry','bbox','assets','links','link_hrefs', 'extra','properties','stac_version','stac_extensions']) AND NOT field_included(CASE WHEN a.attname = 'link_hrefs' THEN 'links' ELSE a.attname END, includes, excludes) - THEN format('NULL::%s', format_type(a.atttypid, a.atttypmod)) + THEN format('NULL::%s AS %I', format_type(a.atttypid, a.atttypmod), a.attname) ELSE format('i.%I', a.attname) END, ', ' ORDER BY a.attnum) INTO cols @@ -5047,7 +4982,7 @@ DECLARE band_cap_months CONSTANT int := 18; page_rows items[] := '{}'::items[]; chunk_rows items[]; target int := _limit + 1; got int := 0; got_band int; - is_asc boolean; cursor_ts timestamptz; band record; + is_asc boolean; band record; band_target numeric; obs_sel numeric; band_where text; guard int := 0; cum_scanned bigint := 0; proj_expr text; mo interval := interval '1 month'; @@ -5103,12 +5038,14 @@ BEGIN END IF; IF datetime_leading AND array_length(bnds.months, 1) IS NOT NULL THEN - cursor_ts := CASE WHEN is_asc THEN bnds.months[1] ELSE bnds.months[array_length(bnds.months, 1)] + mo END; - cursor_idx := 1; + -- Start the band walk at the leading edge for the sort direction: oldest month for ascending, + -- most recent month for descending. (A descending walk that started at the oldest band would + -- collect the oldest items and never reach the recent months.) + cursor_idx := CASE WHEN is_asc THEN 1 ELSE array_length(bnds.months, 1) END; band_target := target * band_margin; WHILE got < target AND guard < 80 LOOP guard := guard + 1; - SELECT * INTO band FROM next_band(bnds.counts, cursor_idx, band_target, band_cap_months); + SELECT * INTO band FROM next_band(bnds.counts, cursor_idx, band_target, band_cap_months, NOT is_asc); -- a valid band must be processed even when next_band also flags done (it consumed the -- last bucket); only stop when there is no band at all. EXIT WHEN band.band_start_idx IS NULL; @@ -5242,7 +5179,6 @@ DECLARE bnds record; coll_clamp text := ''; clamp text; BEGIN IF _where IS NULL OR btrim(_where) = '' THEN _where := ' TRUE '; END IF; - collist := fields_to_itemcols(coalesce(_search->'fields', '{}'::jsonb)); orderby_str := keyset_orderby(_search, is_prev); SELECT (array_agg(field ORDER BY ord))[1], (array_agg(CASE WHEN is_prev THEN (CASE dir WHEN 'ASC' THEN 'DESC' ELSE 'ASC' END) ELSE dir END ORDER BY ord))[1] @@ -5268,10 +5204,20 @@ BEGIN coll_clamp := format('i.collection = ANY(%L::text[]) AND ', bnds.collections); END IF; + -- Build the projection now that the resolved collections are known: when the requested fields are + -- satisfiable from item columns alone (needs_fragment is false), fragment_id is nulled in the + -- projection so a client driving this query never looks up item_fragments -- the fragment-skip + -- rides in the returned query's SELECT list, not a separate flag. + collist := fields_to_itemcols( + coalesce(_search->'fields', '{}'::jsonb), + NOT needs_fragment(coalesce(_search->'fields', '{}'::jsonb), bnds.collections)); + IF datetime_leading AND array_length(bnds.months, 1) IS NOT NULL THEN + -- collections baked in as a literal so the query is parameterized only by $1 (band low), + -- $2 (band high), $3 (limit): the client prepares it once and binds per histogram band. query := format( - 'SELECT %s FROM items i WHERE i.collection = ANY($4) AND i.datetime >= $1 AND i.datetime < $2 AND (%s) ORDER BY %s LIMIT $3', - collist, full_where, orderby_str); + 'SELECT %s FROM items i WHERE %si.datetime >= $1 AND i.datetime < $2 AND (%s) ORDER BY %s LIMIT $3', + collist, coll_clamp, full_where, orderby_str); -- serialize the per-month histogram (months[]/counts[]) to jsonb for the streaming client histogram := ( SELECT jsonb_agg(jsonb_build_object('m', m, 'n', n) ORDER BY ord) @@ -5280,7 +5226,10 @@ BEGIN ELSE DECLARE dt_clamp text := ''; BEGIN IF array_length(bnds.months, 1) IS NOT NULL THEN - dt_clamp := format('i.datetime >= %L AND i.datetime <= %L AND ', bnds.months[1], bnds.months[array_length(bnds.months, 1)]); + -- months[] holds month *starts*; the last bucket spans [start, start + 1 month), so the + -- upper bound must be exclusive of the next month, not <= the last month's start (which + -- would drop items dated after the start of the final month). + dt_clamp := format('i.datetime >= %L AND i.datetime < %L AND ', bnds.months[1], bnds.months[array_length(bnds.months, 1)] + interval '1 month'); END IF; query := format( 'SELECT %s FROM items i WHERE %s%s(%s) ORDER BY %s LIMIT $1', @@ -5562,12 +5511,17 @@ BEGIN END IF; IF array_length(bnds.months, 1) IS NOT NULL THEN - cursor_idx := 1; + -- Walk bands in the SEARCH's datetime direction: newest month first for a descending search, + -- oldest first when sorting datetime ascending. Mirrors search_page (004_search). Walking the + -- wrong direction returns older items before newer ones for a descending limit/early-exit search, + -- and (because band grouping shifts with the partition_stats histogram) makes the result depend on + -- stats — a bug, since stats must only affect performance, never which rows come back. + cursor_idx := CASE WHEN is_asc THEN 1 ELSE array_length(bnds.months, 1) END; band_target := (_limit + 1) * band_margin; <> WHILE NOT exit_flag AND guard < 80 LOOP guard := guard + 1; - SELECT * INTO band FROM next_band(bnds.counts, cursor_idx, band_target, band_cap_months); + SELECT * INTO band FROM next_band(bnds.counts, cursor_idx, band_target, band_cap_months, NOT is_asc); -- process a valid band even when next_band also flags done; stop only on no band. EXIT bands WHEN band.band_start_idx IS NULL; query := format( @@ -5680,6 +5634,66 @@ $$ LANGUAGE SQL; -- BEGIN FRAGMENT: 997_maintenance.sql +-- tighten_dirty_partition_stats: recompute the exact envelope + row count for dirty partitions (oldest +-- first), clearing dirty. Run off-hours (pg_cron or the maintenance CLI). Optional: a wide envelope only +-- over-includes a partition in search, so skipping it never loses rows. `_limit` caps the batch (NULL = +-- all dirty); returns the number of partitions tightened. +-- +-- pg_cron example (operators install this themselves): +-- SELECT cron.schedule('pgstac-tighten', '*/15 * * * *', +-- $$SELECT pgstac.tighten_dirty_partition_stats(200)$$); +CREATE OR REPLACE FUNCTION tighten_dirty_partition_stats(_limit int DEFAULT NULL) +RETURNS int AS $$ +DECLARE + _part text; + _count int := 0; +BEGIN + FOR _part IN + SELECT partition FROM pgstac.partition_stats + WHERE dirty + ORDER BY last_updated NULLS FIRST + LIMIT _limit + LOOP + PERFORM pgstac.tighten_partition_stats(_part); + _count := _count + 1; + END LOOP; + RETURN _count; +END; +$$ LANGUAGE PLPGSQL SECURITY DEFINER; + + +-- build_pending_indexes: build the queryable indexes for partitions flagged `indexes_pending` (new +-- partitions are created index-light for fast ingest; changing a queryable flags every partition via the +-- queryables trigger), then clear the flag. Builds the DDL directly (not via the queue), so schedule it +-- off-hours like the tighten sweep. `_limit` caps the batch (NULL = all pending); returns the count built. +-- +-- pg_cron example (operators install this themselves): +-- SELECT cron.schedule('pgstac-build-indexes', '*/30 * * * *', +-- $$SELECT pgstac.build_pending_indexes(50)$$); +CREATE OR REPLACE FUNCTION build_pending_indexes(_limit int DEFAULT NULL) +RETURNS int AS $$ +DECLARE + _part text; + _q text; + _count int := 0; +BEGIN + FOR _part IN + SELECT partition FROM pgstac.partition_stats + WHERE indexes_pending + ORDER BY last_updated NULLS FIRST + LIMIT _limit + LOOP + FOR _q IN SELECT * FROM pgstac.maintain_partition_queries(_part) LOOP + EXECUTE _q; + END LOOP; + UPDATE pgstac.partition_stats SET indexes_pending = false WHERE partition = _part; + _count := _count + 1; + END LOOP; + RETURN _count; +END; +$$ LANGUAGE PLPGSQL SECURITY DEFINER; + + CREATE OR REPLACE PROCEDURE analyze_items() AS $$ DECLARE q text; @@ -5738,7 +5752,7 @@ DECLARE extent jsonb; BEGIN IF runupdate THEN - PERFORM update_partition_stats_q(partition) + PERFORM tighten_partition_stats(partition) FROM partitions_view WHERE collection=_collection; END IF; SELECT @@ -5830,6 +5844,240 @@ BEGIN END; $$ LANGUAGE PLPGSQL; + +-- update_collection_extents: recompute every collection's extent from partitions_view. +CREATE OR REPLACE FUNCTION update_collection_extents() RETURNS VOID AS $$ +UPDATE collections + SET content = jsonb_set_lax( + content, + '{extent}'::text[], + collection_extent(id, FALSE), + true, + 'use_json_null' + ) +; +$$ LANGUAGE SQL; + + +-- --------------------------------------------------------------------------- +-- Field registry maintenance: track which paths (and value types) exist per +-- collection. Used for schema inference (e.g. the geoparquet export schema). +-- jsonb_field_rows is in 001a_jsonutils.sql. +-- --------------------------------------------------------------------------- + +-- update_field_registry_from_sample: UPSERT registry rows from an array of raw item content JSONBs (the +-- caller picks the sample); value_kinds accumulate observed types over time. +CREATE OR REPLACE FUNCTION update_field_registry_from_sample( + _collection text, + item_contents jsonb[] +) RETURNS void AS $$ + INSERT INTO item_field_registry (collection, path, is_leaf, value_kinds, first_seen, last_seen) + SELECT + _collection, + r.path, + bool_and(r.is_leaf) AS is_leaf, + array_agg(DISTINCT r.value_kind) FILTER (WHERE r.value_kind IS NOT NULL) AS value_kinds, + now(), + now() + FROM unnest(item_contents) AS item(content) + CROSS JOIN LATERAL jsonb_field_rows(item.content) AS r(path, is_leaf, value_kind) + GROUP BY r.path + ON CONFLICT (collection, path) DO UPDATE SET + is_leaf = EXCLUDED.is_leaf, + value_kinds = ( + SELECT array_agg(DISTINCT v) + FROM unnest(item_field_registry.value_kinds || EXCLUDED.value_kinds) t(v) + ), + last_seen = now() + ; +$$ LANGUAGE SQL VOLATILE SECURITY DEFINER; + + +-- update_field_registry_from_items: sample a live collection and UPSERT registry rows (TABLESAMPLE +-- BERNOULLI(5) above ~10k rows by pg_class estimate, else LIMIT 1000). Returns (registered_paths, +-- rows_processed). +CREATE OR REPLACE FUNCTION update_field_registry_from_items( + _collection text +) RETURNS TABLE (registered_paths int, rows_processed int) AS $$ +DECLARE + est_rows bigint; + nrows int; + npaths int; +BEGIN + -- Sum reltuples across the registered item partitions for this collection. + -- reltuples can be -1 (never analyzed); treat negative values as zero. + SELECT COALESCE(sum(GREATEST(c.reltuples::bigint, 0)), 0) INTO est_rows + FROM partitions_view p + JOIN pg_class c ON c.relname = p.partition + JOIN pg_namespace n ON n.oid = c.relnamespace + WHERE p.collection = _collection + AND n.nspname = 'pgstac' + AND c.relkind = 'r'; + + IF est_rows > 10000 THEN + -- Large collection: use statistical sampling to avoid full seq-scan. + WITH sampled AS ( + SELECT content_hydrate(i) AS content FROM items i TABLESAMPLE BERNOULLI(5) WHERE i.collection = _collection + ), + upserted AS ( + INSERT INTO item_field_registry (collection, path, is_leaf, value_kinds, first_seen, last_seen) + SELECT + _collection, + r.path, + bool_and(r.is_leaf) AS is_leaf, + array_agg(DISTINCT r.value_kind) FILTER (WHERE r.value_kind IS NOT NULL) AS value_kinds, + now(), now() + FROM sampled + CROSS JOIN LATERAL jsonb_field_rows(content) AS r(path, is_leaf, value_kind) + GROUP BY r.path + ON CONFLICT (collection, path) DO UPDATE SET + is_leaf = EXCLUDED.is_leaf, + value_kinds = ( + SELECT array_agg(DISTINCT v) + FROM unnest(item_field_registry.value_kinds || EXCLUDED.value_kinds) t(v) + ), + last_seen = now() + RETURNING 1 + ) + SELECT + (SELECT count(*)::int FROM upserted), + (SELECT count(*)::int FROM sampled) + INTO npaths, nrows; + ELSE + -- Small collection: process up to 1000 rows to avoid BERNOULLI returning 0 rows. + WITH sampled AS ( + SELECT content_hydrate(i) AS content FROM items i WHERE i.collection = _collection LIMIT 1000 + ), + upserted AS ( + INSERT INTO item_field_registry (collection, path, is_leaf, value_kinds, first_seen, last_seen) + SELECT + _collection, + r.path, + bool_and(r.is_leaf) AS is_leaf, + array_agg(DISTINCT r.value_kind) FILTER (WHERE r.value_kind IS NOT NULL) AS value_kinds, + now(), now() + FROM sampled + CROSS JOIN LATERAL jsonb_field_rows(content) AS r(path, is_leaf, value_kind) + GROUP BY r.path + ON CONFLICT (collection, path) DO UPDATE SET + is_leaf = EXCLUDED.is_leaf, + value_kinds = ( + SELECT array_agg(DISTINCT v) + FROM unnest(item_field_registry.value_kinds || EXCLUDED.value_kinds) t(v) + ), + last_seen = now() + RETURNING 1 + ) + SELECT + (SELECT count(*)::int FROM upserted), + (SELECT count(*)::int FROM sampled) + INTO npaths, nrows; + END IF; + + RETURN QUERY SELECT npaths, nrows; +END; +$$ LANGUAGE PLPGSQL VOLATILE SECURITY DEFINER; + + +-- refresh_field_registry: expire registry entries not seen within retention_interval (scheduled +-- maintenance). Returns (collection, expired_paths) per affected collection. +CREATE OR REPLACE FUNCTION refresh_field_registry( + _collection text DEFAULT NULL, + retention_interval interval DEFAULT '90 days' +) RETURNS TABLE (collection_id text, expired_paths int) AS $$ + WITH deleted AS ( + DELETE FROM item_field_registry + WHERE (_collection IS NULL OR collection = _collection) + AND last_seen < now() - retention_interval + RETURNING collection + ) + SELECT collection, count(*)::int + FROM deleted + GROUP BY collection; +$$ LANGUAGE SQL VOLATILE SECURITY DEFINER; + + +-- gc_fragments: garbage-collect orphaned item_fragments with a single set-based DELETE (NOT EXISTS +-- anti-join against items.fragment_id). items.fragment_id has no FK (partitioned-items incremental +-- NOT VALID FKs aren't supported), so a fragment unreferenced at the DELETE snapshot but referenced by a +-- concurrent insert can be removed; the retention_interval guard makes this unlikely. Run during +-- low-ingest periods. +CREATE OR REPLACE FUNCTION gc_fragments( + _collection text DEFAULT NULL, + retention_interval interval DEFAULT '90 days' +) RETURNS TABLE ( + collection_id text, + fragments_removed int +) AS $$ + WITH deleted AS ( + DELETE FROM item_fragments f + WHERE + (_collection IS NULL OR f.collection = _collection) + AND f.created_at < now() - retention_interval + AND NOT EXISTS (SELECT 1 FROM items i WHERE i.fragment_id = f.id) + RETURNING f.collection + ) + SELECT collection, count(*)::int + FROM deleted + GROUP BY collection; +$$ LANGUAGE SQL VOLATILE PARALLEL UNSAFE SECURITY DEFINER; + + +-- tighten_partition_stats: recompute the exact envelope (datetime/end_datetime min-max + spatial extent) +-- and row count for one partition, then clear `dirty`. The only function that narrows a partition_stats +-- envelope; the maintenance sweeps drive it over partition_stats WHERE dirty. An empty partition tightens +-- to an 'empty' range so search prunes it out. +CREATE OR REPLACE FUNCTION tighten_partition_stats(_partition text) RETURNS void AS $$ +DECLARE + _n bigint; + _dtmin timestamptz; _dtmax timestamptz; + _edtmin timestamptz; _edtmax timestamptz; + _spatial geometry; + _dtrange tstzrange; + _edtrange tstzrange; + _collection text; +BEGIN + -- Hold the partition's advisory lock across the scan + write: tighten narrows dtrange to the scanned + -- extent and clears dirty, so without the lock a concurrent ingest committing a row outside that + -- extent would be left uncovered (search would prune + miss it). Same lock check_partition uses, so + -- tighten serializes with ingest into this partition only. + PERFORM pg_advisory_xact_lock(hashtext('pgstac.check_partition'), hashtext(_partition)); + + EXECUTE format( + $q$ + SELECT count(*), min(datetime), max(datetime), min(end_datetime), max(end_datetime), + st_setsrid(st_extent(geometry)::geometry, 4326) + FROM %I + $q$, + _partition + ) INTO _n, _dtmin, _dtmax, _edtmin, _edtmax, _spatial; + + IF _n = 0 THEN + _dtrange := 'empty'::tstzrange; + _edtrange := 'empty'::tstzrange; + _spatial := NULL; + ELSE + _dtrange := tstzrange(_dtmin, _dtmax, '[]'); + _edtrange := tstzrange(_edtmin, _edtmax, '[]'); + END IF; + + SELECT pv.collection + INTO _collection + FROM partitions_view pv WHERE partition = _partition; + + INSERT INTO partition_stats + (partition, collection, dtrange, edtrange, spatial, n, dirty, last_updated) + VALUES (_partition, _collection, _dtrange, _edtrange, _spatial, _n, false, now()) + ON CONFLICT (partition) DO UPDATE SET + collection = EXCLUDED.collection, + dtrange = EXCLUDED.dtrange, + edtrange = EXCLUDED.edtrange, + spatial = EXCLUDED.spatial, + n = EXCLUDED.n, + dirty = false, + last_updated = now(); +END; +$$ LANGUAGE PLPGSQL SECURITY DEFINER; -- END FRAGMENT: 997_maintenance.sql -- BEGIN FRAGMENT: 998_idempotent_post.sql @@ -5985,6 +6233,21 @@ GRANT EXECUTE ON ALL FUNCTIONS IN SCHEMA pgstac to pgstac_ingest; GRANT ALL ON ALL TABLES IN SCHEMA pgstac to pgstac_ingest; GRANT USAGE ON ALL SEQUENCES IN SCHEMA pgstac to pgstac_ingest; +-- Privilege wall (INV-1): the connecting role (inheriting pgstac_ingest) may READ these tables but may +-- only MUTATE them through the SECURITY DEFINER write functions (check_partition, widen/tighten_partition_stats, +-- ensure_fragments, make_binary_staging, flush_items_staging_binary, the items_staging trigger, create/ +-- update/upsert/delete_item, the field-registry/fragment helpers). This makes an un-widened or +-- envelope-narrowing direct write structurally impossible. New partitions are created SELECT-only for +-- pgstac_ingest by check_partition; the items parent is revoked here. Staging tables (items_staging*) stay +-- writable so the SQL-only ingest path can COPY into them. +REVOKE INSERT, UPDATE, DELETE, TRUNCATE ON items, partition_stats, item_fragments FROM pgstac_ingest; +-- item_field_registry is the one exception to the wall: the loader WIDENS it directly with an add-only +-- INSERT ... ON CONFLICT DO UPDATE (no SD function), so INSERT + UPDATE stay granted. DELETE/TRUNCATE remain +-- revoked, so even a direct write cannot NARROW the registry — INV-1 (registry is a superset of the data) +-- holds structurally against narrowing, while the cheap widen stays on the hot load path. +REVOKE DELETE, TRUNCATE ON item_field_registry FROM pgstac_ingest; + + REVOKE ALL PRIVILEGES ON PROCEDURE run_queued_queries FROM public; GRANT ALL ON PROCEDURE run_queued_queries TO pgstac_admin; @@ -5996,9 +6259,9 @@ GRANT ALL ON PROCEDURE gc_deleted_items_log_committed(interval, integer) TO pgst RESET ROLE; -SET ROLE pgstac_ingest; -SELECT update_partition_stats_q(partition) FROM partitions_view; -RESET ROLE; +-- (No install-time stats seeding: a fresh install has no partitions, and partition_stats rows are now +-- seeded by check_partition as partitions are created. Exact stats come from tighten_partition_stats via +-- the maintenance sweeps.) -- END FRAGMENT: 998_idempotent_post.sql -- BEGIN FRAGMENT: 999_version.sql diff --git a/src/pgstac/sql/000_idempotent_pre.sql b/src/pgstac/sql/000_idempotent_pre.sql index 86620c7a..de9996b1 100644 --- a/src/pgstac/sql/000_idempotent_pre.sql +++ b/src/pgstac/sql/000_idempotent_pre.sql @@ -39,6 +39,7 @@ DO $$ $$; + GRANT pgstac_admin TO current_user; -- Function to make sure pgstac_admin is the owner of items @@ -190,6 +191,11 @@ RETURNS timestamptz AS $$ ; $$ LANGUAGE SQL IMMUTABLE STRICT; +-- Drop objects superseded by the current partition_stats model. +DROP MATERIALIZED VIEW IF EXISTS partitions CASCADE; +DROP MATERIALIZED VIEW IF EXISTS partition_steps; +DROP VIEW IF EXISTS partition_steps; + -- Drop function signatures whose argument lists changed (CREATE OR REPLACE cannot alter them) DROP FUNCTION IF EXISTS chunker(pred_envelope); DROP FUNCTION IF EXISTS search_bands(pred_envelope, boolean, integer, integer); @@ -209,6 +215,7 @@ DROP FUNCTION IF EXISTS xyzsearch(integer, integer, integer, text, jsonb, intege DROP FUNCTION IF EXISTS search(jsonb); DROP FUNCTION IF EXISTS search_page(jsonb, integer, text, boolean); DROP FUNCTION IF EXISTS search_plan(jsonb, text); +DROP FUNCTION IF EXISTS fields_to_itemcols(jsonb); DROP FUNCTION IF EXISTS search_query(jsonb, boolean, jsonb); DROP FUNCTION IF EXISTS where_stats(text, text, boolean, jsonb); DROP FUNCTION IF EXISTS keyset_sortkeys(jsonb); diff --git a/src/pgstac/sql/001a_jsonutils.sql b/src/pgstac/sql/001a_jsonutils.sql index 0b2a9c44..4c7bba7c 100644 --- a/src/pgstac/sql/001a_jsonutils.sql +++ b/src/pgstac/sql/001a_jsonutils.sql @@ -79,63 +79,51 @@ CREATE OR REPLACE FUNCTION explode_dotpaths_recurse(IN j jsonb) RETURNS SETOF te $$ LANGUAGE SQL IMMUTABLE PARALLEL SAFE; --- jsonb_canonical: RFC 8785 (JSON Canonicalization Scheme)-aligned serialization. --- Produces a deterministic, key-order-independent text encoding that an external --- client can reproduce byte-for-byte. NOTE: do NOT use `jsonb::text` for hashing — --- PostgreSQL re-normalizes object key order to length-then-bytewise and inserts --- ": " / ", " separators, so `jsonb::text` is neither alphabetical nor compact. +-- jsonb_canonical_hash: a deterministic 32-byte SHA-256 identity digest of a JSONB document, reproducible +-- outside PostgreSQL (pgstac-rs `canonical::jsonb_canonical_hash` produces identical bytes). -- --- Canonical rules (must match the external recipe below): --- * object keys sorted by Unicode code point (== UTF-8 byte order, COLLATE "C"), --- * compact separators: ',' between members, ':' between key and value, --- * strings: standard JSON escaping, NON-ASCII left as UTF-8 (no \uXXXX), --- * numbers: IEEE-754 double, shortest round-trip form (Ryu) — matches --- PostgreSQL float8 output and ECMAScript Number::toString for in-range --- values. (STAC numbers are physical quantities; integers beyond 2^53 are --- out of contract, as in RFC 8785.) --- * true / false / null as literals. --- --- External equivalents: --- Python: an RFC 8785 canonicalizer, or the rule-for-rule reference: --- def canon(v): --- if isinstance(v, bool): return 'true' if v else 'false' --- if v is None: return 'null' --- if isinstance(v, dict): --- return '{'+','.join(json.dumps(k,ensure_ascii=False)+':'+canon(v[k]) --- for k in sorted(v))+'}' --- if isinstance(v, list): return '['+','.join(canon(x) for x in v)+']' --- if isinstance(v,(int,float)): --- f=float(v); return str(int(f)) if f==int(f) and abs(f)<1e16 else repr(f) --- return json.dumps(v, ensure_ascii=False) --- Rust: the `rfc8785` crate (serde_jcs) over serde_json::Value. -CREATE OR REPLACE FUNCTION jsonb_canonical(j jsonb) RETURNS text AS $$ - SELECT CASE jsonb_typeof(j) - WHEN 'object' THEN COALESCE(( - SELECT '{' || string_agg( - to_json(kv.key)::text || ':' || jsonb_canonical(kv.value), - ',' ORDER BY kv.key COLLATE "C" - ) || '}' - FROM jsonb_each(j) kv - ), '{}') - WHEN 'array' THEN COALESCE(( - SELECT '[' || string_agg(jsonb_canonical(e.value), ',' ORDER BY e.ord) || ']' - FROM jsonb_array_elements(j) WITH ORDINALITY e(value, ord) - ), '[]') - WHEN 'number' THEN (j #>> '{}')::float8::text - ELSE j::text -- string (JSON-escaped, UTF-8 preserved), 'true' / 'false' / 'null' - END; -$$ LANGUAGE SQL IMMUTABLE PARALLEL SAFE STRICT; - --- jsonb_hash: raw 32-byte sha256 of the canonical (RFC 8785-aligned) JSON form. --- Returns bytea so callers store the compact binary digest directly (32 B vs --- 64-char hex). Use encode(jsonb_hash(j), 'hex') when a printable string is --- needed for display or external comparison. --- Externally reproducible: sha256(utf8_bytes(jsonb_canonical(j))). --- The private jsonb column on items/collections is intentionally excluded — it --- stores operator metadata outside the STAC item identity contract. -CREATE OR REPLACE FUNCTION jsonb_hash(j jsonb) RETURNS bytea AS $$ - SELECT sha256(convert_to(jsonb_canonical(j), 'UTF8')); -$$ LANGUAGE SQL IMMUTABLE PARALLEL SAFE STRICT; +-- The document is flattened to leaf (path, value) rows, each rendered as `path value`, sorted by path +-- (byte order), joined by , and hashed. Separators are C0 control chars US/RS/FS (0x1F/0x1E/0x1D), +-- which do not occur in STAC JSON keys or strings. +-- * path — per step, then 'k'||key (object) or 'i'||idx (0-based array). The 'k'/'i' tags keep an +-- object key "0" distinct from array index 0; keeps nesting distinct from a key with '/'. +-- * number — 'n' || (v)::float8::text, so the digest matches a value stored in a float8 (promoted) column +-- and read back. +-- * string — datetime-shaped (YYYY-MM-DD, optional T/space time, optional offset): 't' || the instant +-- normalized to UTC / 6-digit microseconds / 'Z'. Offset-less is assumed UTC (the SET TIMEZONE +-- below pins the cast), matching pgstac's to_tstz ingest; the cast raises on a shaped-but- +-- invalid timestamp. Other strings: 's' || the raw string. +-- * boolean — 'b'||'true'/'false'; null — 'z'; empty {} — 'e', empty [] — 'a'. +CREATE OR REPLACE FUNCTION jsonb_canonical_hash(j jsonb) RETURNS bytea AS $$ + WITH RECURSIVE t(path, value) AS ( + SELECT ''::text, j + UNION ALL + SELECT t.path || E'\x1F' || e.seg, e.value + FROM t CROSS JOIN LATERAL ( + SELECT 'k' || key AS seg, value + FROM jsonb_each(t.value) WHERE jsonb_typeof(t.value) = 'object' + UNION ALL + SELECT 'i' || (ord - 1)::text AS seg, value + FROM jsonb_array_elements(t.value) WITH ORDINALITY x(value, ord) + WHERE jsonb_typeof(t.value) = 'array' + ) e + ) + SELECT sha256(convert_to(COALESCE(string_agg( + t.path || E'\x1E' || CASE jsonb_typeof(t.value) + WHEN 'number' THEN 'n' || (t.value #>> '{}')::float8::text + WHEN 'boolean' THEN 'b' || (t.value #>> '{}') + WHEN 'null' THEN 'z' + WHEN 'string' THEN CASE + WHEN (t.value #>> '{}') ~ '^\d{4}-\d{2}-\d{2}([Tt ]\d{2}:\d{2}:\d{2}(\.\d+)?([Zz]|[+-]\d{2}:?\d{2})?)?$' + THEN 't' || to_char((t.value #>> '{}')::timestamptz AT TIME ZONE 'UTC', 'YYYY-MM-DD"T"HH24:MI:SS.US') || 'Z' + ELSE 's' || (t.value #>> '{}') END + WHEN 'object' THEN 'e' + WHEN 'array' THEN 'a' + ELSE 'c' END, + E'\x1D' ORDER BY t.path COLLATE "C"), ''), 'UTF8')) + FROM t + WHERE jsonb_typeof(t.value) NOT IN ('object', 'array') OR t.value = '{}'::jsonb OR t.value = '[]'::jsonb; +$$ LANGUAGE SQL IMMUTABLE PARALLEL SAFE STRICT SET TIMEZONE = 'UTC'; -- jsonb_field_rows: Recursively walk a JSONB document and emit one row per field path. -- max_depth guards against runaway recursion on pathologically nested documents. diff --git a/src/pgstac/sql/002a_queryables.sql b/src/pgstac/sql/002a_queryables.sql index 66fd31f4..61c9c0db 100644 --- a/src/pgstac/sql/002a_queryables.sql +++ b/src/pgstac/sql/002a_queryables.sql @@ -446,10 +446,13 @@ $$ LANGUAGE SQL; CREATE OR REPLACE FUNCTION queryables_trigger_func() RETURNS TRIGGER AS $$ DECLARE BEGIN - PERFORM maintain_partitions(); + -- Queryable definitions changed, so every partition's queryable indexes may be stale. Flag them and let + -- the async sweep (build_pending_indexes) (re)build off the hot path, instead of rebuilding every + -- partition synchronously inside this trigger. + UPDATE pgstac.partition_stats SET indexes_pending = true; RETURN NULL; END; -$$ LANGUAGE PLPGSQL; +$$ LANGUAGE PLPGSQL SECURITY DEFINER; CREATE TRIGGER queryables_trigger AFTER INSERT OR UPDATE ON queryables FOR EACH STATEMENT EXECUTE PROCEDURE queryables_trigger_func(); diff --git a/src/pgstac/sql/002b_cql.sql b/src/pgstac/sql/002b_cql.sql index 0c01aa76..03883636 100644 --- a/src/pgstac/sql/002b_cql.sql +++ b/src/pgstac/sql/002b_cql.sql @@ -290,7 +290,7 @@ BEGIN END IF; IF jsonb_typeof(j) = 'object' THEN - -- GeoJSON geometry args (Point/Polygon/.../GeometryCollection) are not cql2 expressions; + -- GeoJSON geometry args (Point/Polygon/.../GeometryCollection) are not cql expressions; -- pass them through unchanged so spatial ops keep their geometry intact. IF j ? 'type' AND (j ? 'coordinates' OR j ? 'geometries') THEN RETURN j; diff --git a/src/pgstac/sql/002c_envelope.sql b/src/pgstac/sql/002c_envelope.sql index 05aeb971..309ad699 100644 --- a/src/pgstac/sql/002c_envelope.sql +++ b/src/pgstac/sql/002c_envelope.sql @@ -162,8 +162,8 @@ CREATE OR REPLACE FUNCTION partition_bounds( ) RETURNS record LANGUAGE sql STABLE AS $$ WITH cand AS ( SELECT ps.collection, - lower(coalesce(ps.dtrange, ps.partition_dtrange)) AS lo, - upper(coalesce(ps.dtrange, ps.partition_dtrange)) AS hi, + lower(ps.dtrange) AS lo, + upper(ps.dtrange) AS hi, coalesce(ps.n, 0) AS n FROM partition_stats ps WHERE ((_env).colls IS NULL OR ps.collection = ANY((_env).colls)) @@ -200,12 +200,17 @@ $$; -- next_band: walk a histogram using array indexes. Given per-month counts, a -- cursor position, and a target row count, returns the next band's index range. --- The caller converts band indexes back to timestamps for SQL queries. +-- The caller converts band indexes back to timestamps for SQL queries. The band +-- index range [band_start_idx, band_end_idx] is always low..high regardless of +-- direction; only the walk order differs. For a descending search (_descending) +-- the cursor starts at the most recent month and walks toward older months, so +-- the most recent items are collected first. CREATE OR REPLACE FUNCTION next_band( _counts bigint[], _cursor_idx int, _target numeric, _cap_months int, + _descending boolean DEFAULT false, OUT band_start_idx int, OUT band_end_idx int, OUT scanned bigint, @@ -214,6 +219,7 @@ CREATE OR REPLACE FUNCTION next_band( ) RETURNS record LANGUAGE plpgsql STABLE AS $$ DECLARE idx int; + n int; cumulative bigint := 0; BEGIN done := false; scanned := 0; @@ -223,14 +229,39 @@ BEGIN done := true; -- no histogram / no cursor => nothing to walk RETURN; END IF; + n := array_length(_counts, 1); + IF _descending THEN + -- Walk downward (newest -> oldest). The high bound is the cursor (clamped into range). + IF _cursor_idx < 1 THEN done := true; RETURN; END IF; + idx := LEAST(_cursor_idx, n); + FOR i IN REVERSE idx..GREATEST(1, idx - _cap_months + 1) LOOP + cumulative := cumulative + _counts[i]; + scanned := scanned + _counts[i]; + IF cumulative >= _target THEN + band_start_idx := i; -- low (older) month bound + band_end_idx := idx; -- high (newer) month bound + next_cursor_idx := i - 1; + IF next_cursor_idx < 1 THEN done := true; END IF; + RETURN; + END IF; + END LOOP; + -- Target not reached within the cap: end the band at the cap boundary. + band_start_idx := GREATEST(1, idx - _cap_months + 1); + band_end_idx := idx; + next_cursor_idx := band_start_idx - 1; + done := (band_start_idx <= 1); + RETURN; + END IF; + + -- Ascending (oldest -> newest). idx := NULL; - FOR i IN 1..array_length(_counts, 1) LOOP + FOR i IN 1..n LOOP IF i >= _cursor_idx THEN idx := i; EXIT; END IF; END LOOP; IF idx IS NULL THEN done := true; next_cursor_idx := _cursor_idx; RETURN; END IF; - FOR i IN idx..array_length(_counts, 1) LOOP + FOR i IN idx..n LOOP IF i > idx + _cap_months - 1 THEN EXIT; END IF; cumulative := cumulative + _counts[i]; scanned := scanned + _counts[i]; @@ -238,7 +269,7 @@ BEGIN band_start_idx := idx; band_end_idx := i; next_cursor_idx := i + 1; - IF next_cursor_idx > array_length(_counts, 1) THEN done := true; END IF; + IF next_cursor_idx > n THEN done := true; END IF; RETURN; END IF; END LOOP; @@ -246,8 +277,8 @@ BEGIN -- Target not reached within the cap: end the band at the cap boundary (not the array end), -- so the cap actually limits band width. done only when we've consumed the whole histogram. band_start_idx := idx; - band_end_idx := LEAST(idx + _cap_months - 1, array_length(_counts, 1)); + band_end_idx := LEAST(idx + _cap_months - 1, n); next_cursor_idx := band_end_idx + 1; - done := (band_end_idx >= array_length(_counts, 1)); + done := (band_end_idx >= n); END; $$; diff --git a/src/pgstac/sql/003a_items.sql b/src/pgstac/sql/003a_items.sql index 7763b4e2..e521b4f6 100644 --- a/src/pgstac/sql/003a_items.sql +++ b/src/pgstac/sql/003a_items.sql @@ -42,7 +42,7 @@ CREATE TABLE items ( stac_version text, stac_extensions jsonb DEFAULT '[]'::jsonb, pgstac_updated_at timestamptz NOT NULL DEFAULT now(), - -- 32-byte sha256 of the canonical (RFC 8785-aligned) STAC item JSON set at + -- 32-byte sha256 of the canonical (jsonb_canonical) STAC item JSON set at -- ingest time. Allows external clients to detect unchanged items without a -- full fetch. Does NOT include the private column (operator metadata). item_hash bytea NOT NULL DEFAULT '\x'::bytea, @@ -132,47 +132,15 @@ CREATE STATISTICS datetime_stats (dependencies) on datetime, end_datetime from i ALTER TABLE items ADD CONSTRAINT items_collections_fk FOREIGN KEY (collection) REFERENCES collections(id) ON DELETE CASCADE DEFERRABLE; --- partition_after_triggerfunc: After-statement trigger on items. --- Updates partition statistics for every partition touched by the current batch, --- using run_or_queue() so the work is deferred rather than blocking the ingest --- transaction. -CREATE OR REPLACE FUNCTION partition_after_triggerfunc() RETURNS TRIGGER AS $$ -DECLARE - p text; - t timestamptz := clock_timestamp(); -BEGIN - RAISE NOTICE 'Updating partition stats %', t; - FOR p IN SELECT DISTINCT partition - FROM newdata n JOIN partition_sys_meta p - ON (n.collection=p.collection AND n.datetime <@ p.partition_dtrange) - LOOP - PERFORM run_or_queue(format('SELECT update_partition_stats(%L, %L);', p, true)); - END LOOP; - RAISE NOTICE 't: % %', t, clock_timestamp() - t; - RETURN NULL; -END; -$$ LANGUAGE PLPGSQL SECURITY DEFINER; - +-- The old per-statement stats trigger (partition_after_triggerfunc) is intentionally GONE. Under the +-- widen-now / tighten-async model, partition_stats coverage is maintained by the ingest path itself +-- (check_partition seeds + widen_partition_stats covers, in their own transactions) and tightened +-- off the hot path by tighten_partition_stats; a per-insert trigger that re-scanned partitions and +-- refreshed matviews was the lock pathology this model removes. DROP TRIGGER IF EXISTS items_after_insert_trigger ON items; -CREATE TRIGGER items_after_insert_trigger -AFTER INSERT ON items -REFERENCING NEW TABLE AS newdata -FOR EACH STATEMENT -EXECUTE FUNCTION partition_after_triggerfunc(); - DROP TRIGGER IF EXISTS items_after_update_trigger ON items; -CREATE TRIGGER items_after_update_trigger -AFTER UPDATE ON items -REFERENCING NEW TABLE AS newdata -FOR EACH STATEMENT -EXECUTE FUNCTION partition_after_triggerfunc(); - DROP TRIGGER IF EXISTS items_after_delete_trigger ON items; -CREATE TRIGGER items_after_delete_trigger -AFTER DELETE ON items -REFERENCING OLD TABLE AS newdata -FOR EACH STATEMENT -EXECUTE FUNCTION partition_after_triggerfunc(); +DROP FUNCTION IF EXISTS partition_after_triggerfunc(); -- items_content_distinct_sql / items_content_changed: detect whether a direct -- UPDATE actually changed the stored item content. Used by the touch trigger @@ -223,7 +191,7 @@ $$ LANGUAGE PLPGSQL IMMUTABLE PARALLEL SAFE; -- items_touch_triggerfunc: refresh pgstac_updated_at when a direct UPDATE changes -- the stored item content. It deliberately does NOT recompute item_hash: --- item_hash is the canonical (RFC 8785-aligned) hash of the item *as ingested* +-- item_hash is the canonical (jsonb_canonical) hash of the item *as ingested* -- through create_item / upsert_item / update_item (set once in content_dehydrate), -- so it stays externally reproducible by a client hashing its own copy. -- A raw `UPDATE items SET ...` that bypasses the staging path leaves item_hash @@ -436,7 +404,7 @@ CREATE OR REPLACE FUNCTION content_dehydrate(content jsonb) RETURNS items AS $$ content->>'stac_version' AS stac_version, COALESCE(content->'stac_extensions', '[]'::jsonb) AS stac_extensions, now() AS pgstac_updated_at, - pgstac.jsonb_hash(content) AS item_hash, + pgstac.jsonb_canonical_hash(content) AS item_hash, NULL::bigint AS fragment_id, content->'bbox' AS bbox, CASE WHEN content->'links' IS NOT NULL AND content->'links' <> '[]'::jsonb THEN content->'links' END AS links, @@ -531,10 +499,8 @@ BEGIN END; $$ LANGUAGE PLPGSQL IMMUTABLE; --- content_hydrate: Reassemble a full STAC item JSON from the split columns --- and the shared fragment content. This is the single hydrate function; --- the old content_nonhydrated wrapper and 3-arg _collection parameter have --- been removed. +-- content_hydrate: reassemble a full STAC item JSON from the split columns and the shared fragment +-- content. The single hydrate function. CREATE OR REPLACE FUNCTION content_hydrate( _item items, fields jsonb DEFAULT '{}'::jsonb, @@ -776,6 +742,22 @@ DECLARE BEGIN RAISE NOTICE 'Creating Partitions. %', clock_timestamp() - ts; + -- Fail loudly on items whose collection does not exist instead of silently dropping them + -- in the collections JOINs below (matches the Rust loader, which also errors on this). + IF EXISTS ( + SELECT 1 FROM newdata n + WHERE NOT EXISTS ( + SELECT 1 FROM collections c WHERE c.id = n.content->>'collection' + ) + ) THEN + RAISE EXCEPTION 'cannot load item(s) into nonexistent collection(s): %', + (SELECT string_agg(DISTINCT coalesce(n.content->>'collection', ''), ', ') + FROM newdata n + WHERE NOT EXISTS ( + SELECT 1 FROM collections c WHERE c.id = n.content->>'collection' + )); + END IF; + FOR part IN WITH t AS ( SELECT n.content->>'collection' as collection, @@ -832,13 +814,31 @@ BEGIN RAISE NOTICE 'Inserted % rows to items. %', nrows, clock_timestamp() - ts; END IF; + -- Bump each partition's row-count estimate so search stepping sees the new rows. n drives + -- partition_bounds' histogram and the Rust keyset search_page skips zero-count bands, so a partition + -- left at n=0 would hide its freshly-ingested items from that path (SQL search() scans regardless). + -- Mirrors flush_items_staging_binary on the Rust loader. Over-count is safe (ignore/upsert may insert + -- fewer than staged); the async tightener resets n exactly. check_partition above already widened the + -- temporal envelope; spatial is left to the tightener on this fallback path. (n only affects search + -- stepping performance, never which rows/order come back — see geometrysearch band direction.) + UPDATE partition_stats ps + SET n = COALESCE(ps.n, 0) + agg.c, dirty = true, last_updated = now() + FROM ( + SELECT (partition_name(nd.content->>'collection', + lower(stac_daterange(nd.content->'properties')))).partition_name AS partition, + count(*) AS c + FROM newdata nd + GROUP BY 1 + ) agg + WHERE ps.partition = agg.partition; + RAISE NOTICE 'Deleting data from staging table. %', clock_timestamp() - ts; EXECUTE format('DELETE FROM %I', TG_TABLE_NAME); RAISE NOTICE 'Done. %', clock_timestamp() - ts; RETURN NULL; END; -$$ LANGUAGE PLPGSQL; +$$ LANGUAGE PLPGSQL SECURITY DEFINER; DROP TRIGGER IF EXISTS items_staging_insert_trigger ON items_staging; @@ -874,7 +874,7 @@ out items%ROWTYPE; BEGIN DELETE FROM items WHERE id = _id AND (_collection IS NULL OR collection=_collection) RETURNING * INTO STRICT out; END; -$$ LANGUAGE PLPGSQL; +$$ LANGUAGE PLPGSQL SECURITY DEFINER; --/* CREATE OR REPLACE FUNCTION create_item(data jsonb) RETURNS VOID AS $$ @@ -919,156 +919,10 @@ CREATE OR REPLACE FUNCTION collection_temporal_extent(id text) RETURNS jsonb AS ; $$ LANGUAGE SQL IMMUTABLE PARALLEL SAFE SET SEARCH_PATH TO pgstac, public; -CREATE OR REPLACE FUNCTION update_collection_extents() RETURNS VOID AS $$ -UPDATE collections - SET content = jsonb_set_lax( - content, - '{extent}'::text[], - collection_extent(id, FALSE), - true, - 'use_json_null' - ) -; -$$ LANGUAGE SQL; --- --------------------------------------------------------------------------- --- Field Registry: walks JSONB item content to track which paths exist in each --- collection. Used to auto-populate queryables and support schema inference. --- jsonb_field_rows is defined in 001a_jsonutils.sql (loaded first). --- --------------------------------------------------------------------------- - --- update_field_registry_from_sample: UPSERT registry rows from a pre-selected array of --- raw item content JSONBs. Callers supply the sample to decouple sampling strategy --- from the registry write; merge value_kinds to accumulate observed types over time. -CREATE OR REPLACE FUNCTION update_field_registry_from_sample( - _collection text, - item_contents jsonb[] -) RETURNS void AS $$ - INSERT INTO item_field_registry (collection, path, is_leaf, value_kinds, first_seen, last_seen) - SELECT - _collection, - r.path, - bool_and(r.is_leaf) AS is_leaf, - array_agg(DISTINCT r.value_kind) FILTER (WHERE r.value_kind IS NOT NULL) AS value_kinds, - now(), - now() - FROM unnest(item_contents) AS item(content) - CROSS JOIN LATERAL jsonb_field_rows(item.content) AS r(path, is_leaf, value_kind) - GROUP BY r.path - ON CONFLICT (collection, path) DO UPDATE SET - is_leaf = EXCLUDED.is_leaf, - value_kinds = ( - SELECT array_agg(DISTINCT v) - FROM unnest(item_field_registry.value_kinds || EXCLUDED.value_kinds) t(v) - ), - last_seen = now() - ; -$$ LANGUAGE SQL VOLATILE; - --- update_field_registry_from_items: Sample a live collection and UPSERT registry rows. --- Uses TABLESAMPLE BERNOULLI(5) for large collections (>10k rows by pg_class estimate) --- and LIMIT 1000 for smaller ones to avoid a full seq-scan for tiny collections. --- pg_class.reltuples is an estimate (may be stale); its only role is threshold selection. --- Returns (registered_paths, rows_processed) for observability. -CREATE OR REPLACE FUNCTION update_field_registry_from_items( - _collection text -) RETURNS TABLE (registered_paths int, rows_processed int) AS $$ -DECLARE - est_rows bigint; - nrows int; - npaths int; -BEGIN - -- Sum reltuples across the registered item partitions for this collection. - -- reltuples can be -1 (never analyzed); treat negative values as zero. - SELECT COALESCE(sum(GREATEST(c.reltuples::bigint, 0)), 0) INTO est_rows - FROM partitions_view p - JOIN pg_class c ON c.relname = p.partition - JOIN pg_namespace n ON n.oid = c.relnamespace - WHERE p.collection = _collection - AND n.nspname = 'pgstac' - AND c.relkind = 'r'; - - IF est_rows > 10000 THEN - -- Large collection: use statistical sampling to avoid full seq-scan. - WITH sampled AS ( - SELECT content_hydrate(i) AS content FROM items i TABLESAMPLE BERNOULLI(5) WHERE i.collection = _collection - ), - upserted AS ( - INSERT INTO item_field_registry (collection, path, is_leaf, value_kinds, first_seen, last_seen) - SELECT - _collection, - r.path, - bool_and(r.is_leaf) AS is_leaf, - array_agg(DISTINCT r.value_kind) FILTER (WHERE r.value_kind IS NOT NULL) AS value_kinds, - now(), now() - FROM sampled - CROSS JOIN LATERAL jsonb_field_rows(content) AS r(path, is_leaf, value_kind) - GROUP BY r.path - ON CONFLICT (collection, path) DO UPDATE SET - is_leaf = EXCLUDED.is_leaf, - value_kinds = ( - SELECT array_agg(DISTINCT v) - FROM unnest(item_field_registry.value_kinds || EXCLUDED.value_kinds) t(v) - ), - last_seen = now() - RETURNING 1 - ) - SELECT - (SELECT count(*)::int FROM upserted), - (SELECT count(*)::int FROM sampled) - INTO npaths, nrows; - ELSE - -- Small collection: process up to 1000 rows to avoid BERNOULLI returning 0 rows. - WITH sampled AS ( - SELECT content_hydrate(i) AS content FROM items i WHERE i.collection = _collection LIMIT 1000 - ), - upserted AS ( - INSERT INTO item_field_registry (collection, path, is_leaf, value_kinds, first_seen, last_seen) - SELECT - _collection, - r.path, - bool_and(r.is_leaf) AS is_leaf, - array_agg(DISTINCT r.value_kind) FILTER (WHERE r.value_kind IS NOT NULL) AS value_kinds, - now(), now() - FROM sampled - CROSS JOIN LATERAL jsonb_field_rows(content) AS r(path, is_leaf, value_kind) - GROUP BY r.path - ON CONFLICT (collection, path) DO UPDATE SET - is_leaf = EXCLUDED.is_leaf, - value_kinds = ( - SELECT array_agg(DISTINCT v) - FROM unnest(item_field_registry.value_kinds || EXCLUDED.value_kinds) t(v) - ), - last_seen = now() - RETURNING 1 - ) - SELECT - (SELECT count(*)::int FROM upserted), - (SELECT count(*)::int FROM sampled) - INTO npaths, nrows; - END IF; - RETURN QUERY SELECT npaths, nrows; -END; -$$ LANGUAGE PLPGSQL VOLATILE SECURITY DEFINER; - --- refresh_field_registry: Expire stale registry entries that haven't been seen recently. --- Intended for scheduled maintenance (e.g. pg_cron daily job). --- Returns (collection, expired_paths) for each collection affected. -CREATE OR REPLACE FUNCTION refresh_field_registry( - _collection text DEFAULT NULL, - retention_interval interval DEFAULT '90 days' -) RETURNS TABLE (collection_id text, expired_paths int) AS $$ - WITH deleted AS ( - DELETE FROM item_field_registry - WHERE (_collection IS NULL OR collection = _collection) - AND last_seen < now() - retention_interval - RETURNING collection - ) - SELECT collection, count(*)::int - FROM deleted - GROUP BY collection; -$$ LANGUAGE SQL VOLATILE; + + -- Item Fragment Management functions @@ -1125,37 +979,6 @@ CREATE OR REPLACE FUNCTION pgstac_hash_fragment(fragment jsonb) RETURNS bytea AS SELECT sha256(convert_to(fragment::text, 'UTF8')); $$ LANGUAGE SQL IMMUTABLE PARALLEL SAFE; --- gc_fragments: Garbage collect orphaned fragments using a single set-based DELETE so the --- planner can choose the optimal join/anti-join strategy across all collections. --- The NOT EXISTS sub-select is evaluated per fragment; with an index on items.fragment_id --- this is an efficient anti-join rather than a full seq-scan. --- --- Operational note: because items.fragment_id is intentionally unmanaged by an FK --- (partitioned-items incremental NOT VALID FKs are not supported), gc_fragments has a --- small race window with concurrent inserts. A fragment that is unreferenced at the time --- the DELETE snapshot is taken but becomes referenced by a later insert could be removed. --- The retention_interval guard makes this unlikely for normal ingest, but operators should --- still run gc_fragments during low-ingest periods or with a sufficiently conservative --- retention interval. This is a documented operational tradeoff, not a silent invariant. -CREATE OR REPLACE FUNCTION gc_fragments( - _collection text DEFAULT NULL, - retention_interval interval DEFAULT '90 days' -) RETURNS TABLE ( - collection_id text, - fragments_removed int -) AS $$ - WITH deleted AS ( - DELETE FROM item_fragments f - WHERE - (_collection IS NULL OR f.collection = _collection) - AND f.created_at < now() - retention_interval - AND NOT EXISTS (SELECT 1 FROM items i WHERE i.fragment_id = f.id) - RETURNING f.collection - ) - SELECT collection, count(*)::int - FROM deleted - GROUP BY collection; -$$ LANGUAGE SQL VOLATILE PARALLEL UNSAFE; -- strip_fragment_col: Remove fragment-owned sub-keys from a split column value. -- col_name is the top-level STAC key that this column represents (e.g. 'assets' or 'properties'). diff --git a/src/pgstac/sql/003b_partitions.sql b/src/pgstac/sql/003b_partitions.sql index 6f661398..157bbd4a 100644 --- a/src/pgstac/sql/003b_partitions.sql +++ b/src/pgstac/sql/003b_partitions.sql @@ -1,18 +1,30 @@ CREATE TABLE partition_stats ( partition text PRIMARY KEY, collection text, - partition_dtrange tstzrange, + -- dtrange/edtrange are the data bounds used for search pruning. widen_partition_stats fills them + -- (generously) on ingest; tighten_partition_stats narrows them to the exact extent. The partition's + -- structural range lives in its own constraint, read via partitions_view.constraint_dtrange. dtrange tstzrange, edtrange tstzrange, spatial geometry, last_updated timestamptz, n bigint, - keys text[] -) WITH (FILLFACTOR=90); + -- Deferred-maintenance flags: + -- dirty: the stored envelope may be WIDER than the actual data; an async tightener + -- should recompute exact min/max/extent and clear the flag. Search correctness + -- never depends on this (a wide envelope only over-includes a partition). + -- indexes_pending: the partition currently carries only the parent-inherited indexes (id PK, + -- datetime, geometry); its queryable-defined indexes have not been built yet. + dirty boolean NOT NULL DEFAULT false, + indexes_pending boolean NOT NULL DEFAULT false +) WITH (FILLFACTOR=70); CREATE INDEX partitions_range_idx ON partition_stats USING GIST(dtrange); CREATE INDEX partition_stats_collection_idx ON partition_stats (collection); CREATE INDEX partition_stats_spatial_idx ON partition_stats USING GIST(spatial) WHERE spatial IS NOT NULL; +-- Work-queue indexes: the maintenance sweeps find pending partitions without scanning every row. +CREATE INDEX partition_stats_dirty_idx ON partition_stats (partition) WHERE dirty; +CREATE INDEX partition_stats_indexes_pending_idx ON partition_stats (partition) WHERE indexes_pending; CREATE OR REPLACE FUNCTION constraint_tstzrange(expr text) RETURNS tstzrange AS $$ @@ -90,10 +102,9 @@ SELECT level, c.reltuples, c.relhastriggers, - partition_dtrange, COALESCE( get_tstz_constraint(c.oid, 'datetime'), - partition_dtrange, + constraint_tstzrange(pg_get_expr(c.relpartbound, c.oid)), inf_range ) as constraint_dtrange, COALESCE( @@ -109,7 +120,6 @@ FROM JOIN LATERAL pg_get_expr(c.relpartbound, c.oid) as partition_expr ON TRUE JOIN LATERAL pg_get_expr(parent.relpartbound, parent.oid) as parent_partition_expr ON TRUE JOIN LATERAL tstzrange('-infinity', 'infinity','[]') as inf_range ON TRUE - JOIN LATERAL COALESCE(constraint_tstzrange(pg_get_expr(c.relpartbound, c.oid)), inf_range) as partition_dtrange ON TRUE JOIN LATERAL get_tstz_constraint(c.oid, 'datetime') as datetime_constraint ON TRUE JOIN LATERAL get_tstz_constraint(c.oid, 'end_datetime') as end_datetime_constraint ON TRUE WHERE isleaf @@ -130,10 +140,9 @@ SELECT level, c.reltuples, c.relhastriggers, - partition_dtrange, COALESCE( get_tstz_constraint(c.oid, 'datetime'), - partition_dtrange, + constraint_tstzrange(pg_get_expr(c.relpartbound, c.oid)), inf_range ) as constraint_dtrange, COALESCE( @@ -153,10 +162,9 @@ FROM JOIN LATERAL pg_get_expr(c.relpartbound, c.oid) as partition_expr ON TRUE JOIN LATERAL pg_get_expr(parent.relpartbound, parent.oid) as parent_partition_expr ON TRUE JOIN LATERAL tstzrange('-infinity', 'infinity','[]') as inf_range ON TRUE - JOIN LATERAL COALESCE(constraint_tstzrange(pg_get_expr(c.relpartbound, c.oid)), inf_range) as partition_dtrange ON TRUE JOIN LATERAL get_tstz_constraint(c.oid, 'datetime') as datetime_constraint ON TRUE JOIN LATERAL get_tstz_constraint(c.oid, 'end_datetime') as end_datetime_constraint ON TRUE - -- the view computes its own collection/partition_dtrange from the live tree; pull only the + -- the view computes its own collection + constraint_dtrange from the live tree; pull only the -- data-extent columns from partition_stats to avoid colliding with those names. LEFT JOIN ( SELECT partition, dtrange, edtrange, spatial, last_updated @@ -165,109 +173,6 @@ FROM WHERE isleaf ; -CREATE MATERIALIZED VIEW partitions AS -SELECT * FROM partitions_view; -CREATE UNIQUE INDEX ON partitions (partition); - -CREATE MATERIALIZED VIEW partition_steps AS -SELECT - partition as name, - date_trunc('month',lower(partition_dtrange)) as sdate, - date_trunc('month', upper(partition_dtrange)) + '1 month'::interval as edate - FROM partitions_view WHERE partition_dtrange IS NOT NULL AND partition_dtrange != 'empty'::tstzrange - ORDER BY dtrange ASC -; - - -CREATE OR REPLACE FUNCTION update_partition_stats_q(_partition text, istrigger boolean default false) RETURNS VOID AS $$ -DECLARE -BEGIN - PERFORM run_or_queue( - format('SELECT update_partition_stats(%L, %L);', _partition, istrigger) - ); -END; -$$ LANGUAGE PLPGSQL; - -CREATE OR REPLACE FUNCTION update_partition_stats(_partition text, istrigger boolean default false) RETURNS VOID AS $$ -DECLARE - dtrange tstzrange; - edtrange tstzrange; - cdtrange tstzrange; - cedtrange tstzrange; - extent geometry; - collection text; - _part_dtrange tstzrange; - _n bigint; -BEGIN - RAISE NOTICE 'Updating stats for %.', _partition; - EXECUTE format( - $q$ - SELECT - tstzrange(min(datetime), max(datetime),'[]'), - tstzrange(min(end_datetime), max(end_datetime), '[]') - FROM %I - $q$, - _partition - ) INTO dtrange, edtrange; - EXECUTE format('ANALYZE %I;', _partition); - extent := st_estimatedextent('pgstac', _partition, 'geometry'); - RAISE DEBUG 'Estimated Extent: %', extent; - - -- Per-partition metadata: collection + partition boundary from the live tree (partitions_view), - -- current constraint ranges (for the constraint check below), and the row estimate from pg_class. - SELECT pv.collection, pv.partition_dtrange, pv.constraint_dtrange, pv.constraint_edtrange - INTO collection, _part_dtrange, cdtrange, cedtrange - FROM partitions_view pv WHERE partition = _partition; - _n := (SELECT reltuples::bigint FROM pg_class WHERE oid = quote_ident(_partition)::regclass); - - INSERT INTO partition_stats - (partition, collection, partition_dtrange, dtrange, edtrange, spatial, n, last_updated) - SELECT _partition, collection, _part_dtrange, dtrange, edtrange, extent, _n, now() - ON CONFLICT (partition) DO - UPDATE SET - collection=EXCLUDED.collection, - partition_dtrange=EXCLUDED.partition_dtrange, - dtrange=EXCLUDED.dtrange, - edtrange=EXCLUDED.edtrange, - spatial=EXCLUDED.spatial, - n=EXCLUDED.n, - last_updated=EXCLUDED.last_updated - ; - - RAISE NOTICE 'Checking if we need to modify constraints...'; - RAISE NOTICE 'cdtrange: % dtrange: % cedtrange: % edtrange: %',cdtrange, dtrange, cedtrange, edtrange; - IF - (cdtrange IS DISTINCT FROM dtrange OR edtrange IS DISTINCT FROM cedtrange) - AND NOT istrigger - THEN - RAISE NOTICE 'Modifying Constraints'; - RAISE NOTICE 'Existing % %', cdtrange, cedtrange; - RAISE NOTICE 'New % %', dtrange, edtrange; - PERFORM drop_table_constraints(_partition); - PERFORM create_table_constraints(_partition, dtrange, edtrange); - END IF; - REFRESH MATERIALIZED VIEW partitions; - REFRESH MATERIALIZED VIEW partition_steps; - RAISE NOTICE 'Checking if we need to update collection extents.'; - IF get_setting_bool('update_collection_extent') THEN - RAISE NOTICE 'updating collection extent for %', collection; - PERFORM run_or_queue(format($q$ - UPDATE collections - SET content = jsonb_set_lax( - content, - '{extent}'::text[], - collection_extent(%L, FALSE), - true, - 'use_json_null' - ) WHERE id=%L - ; - $q$, collection, collection)); - ELSE - RAISE NOTICE 'Not updating collection extent for %', collection; - END IF; - -END; -$$ LANGUAGE PLPGSQL STRICT SECURITY DEFINER; CREATE OR REPLACE FUNCTION partition_name( IN collection text, IN dt timestamptz, OUT partition_name text, OUT partition_range tstzrange) AS $$ @@ -302,107 +207,82 @@ END; $$ LANGUAGE PLPGSQL STABLE; -CREATE OR REPLACE FUNCTION drop_table_constraints(t text) RETURNS text AS $$ +-- Partitions carry no _dt CHECK constraints; pruning is partition_stats-driven. get_tstz_constraint / +-- constraint_tstzrange remain so the views' constraint_* columns resolve to inf_range when a partition +-- has no datetime constraint. + + +-- widen_partition_stats: widen a partition's stored envelope to cover a batch. A no-op (snapshot read, no +-- write, no lock) when the envelope already covers the batch; otherwise widens and sets dirty=true: +-- * datetime : sub-partitioned collections widen to the partition's datetime bound (covers anything +-- that can land there); a NULL-partition_trunc partition pads the batch range by +-- partition_stats_widen_buffer (default 1 month) each side. +-- * end_datetime: the datetime target extended by the batch's max (end_datetime - datetime) tail. +-- * spatial : NULL means "always a search candidate"; a spatial miss resets spatial to NULL until +-- the tightener computes the real extent. +-- Requires the partition_stats row to exist (check_partition seeds it); raises if it does not. +CREATE OR REPLACE FUNCTION widen_partition_stats( + _partition text, + _dtrange tstzrange, + _edtrange tstzrange, + _spatial geometry DEFAULT NULL, + _constraint_dtrange tstzrange DEFAULT NULL +) RETURNS void AS $$ DECLARE - q text; + cur RECORD; + is_unbounded boolean; + tail interval; + buf interval := COALESCE(get_setting('partition_stats_widen_buffer'), '1 month')::interval; + target_dtrange tstzrange; + target_edtrange tstzrange; + new_dtrange tstzrange; + new_edtrange tstzrange; + new_spatial geometry; + dt_covered boolean; + edt_covered boolean; + spatial_covered boolean; BEGIN - IF NOT EXISTS (SELECT 1 FROM partitions_view WHERE partition=t) THEN - RETURN NULL; + SELECT dtrange, edtrange, spatial + INTO cur + FROM partition_stats WHERE partition = _partition; + IF NOT FOUND THEN + RAISE EXCEPTION 'partition_stats row for % does not exist; call check_partition first', _partition; END IF; - FOR q IN SELECT FORMAT( - $q$ - ALTER TABLE %I DROP CONSTRAINT IF EXISTS %I; - $q$, - t, - conname - ) FROM pg_constraint - WHERE conrelid=t::regclass::oid AND contype='c' - LOOP - EXECUTE q; - END LOOP; - RETURN t; -END; -$$ LANGUAGE PLPGSQL SECURITY DEFINER; -CREATE OR REPLACE FUNCTION create_table_constraints(t text, _dtrange tstzrange, _edtrange tstzrange) RETURNS text AS $$ -DECLARE - q text; -BEGIN - IF NOT EXISTS (SELECT 1 FROM partitions_view WHERE partition=t) THEN - RETURN NULL; + dt_covered := COALESCE(cur.dtrange @> _dtrange, false); + edt_covered := COALESCE(cur.edtrange @> _edtrange, false); + spatial_covered := cur.spatial IS NULL OR _spatial IS NULL OR ST_Covers(cur.spatial, _spatial); + IF dt_covered AND edt_covered AND spatial_covered THEN + RETURN; -- already covered: no write, no lock END IF; - RAISE NOTICE 'Creating Table Constraints for % % %', t, _dtrange, _edtrange; - IF _dtrange = 'empty' AND _edtrange = 'empty' THEN - q :=format( - $q$ - DO $block$ - BEGIN - ALTER TABLE %I DROP CONSTRAINT IF EXISTS %I; - ALTER TABLE %I - ADD CONSTRAINT %I - CHECK (((datetime IS NULL) AND (end_datetime IS NULL))) NOT VALID - ; - ALTER TABLE %I - VALIDATE CONSTRAINT %I - ; - - - - EXCEPTION WHEN others THEN - RAISE WARNING '%%, Issue Altering Constraints. Please run update_partition_stats(%I)', SQLERRM USING ERRCODE = SQLSTATE; - END; - $block$; - $q$, - t, - format('%s_dt', t), - t, - format('%s_dt', t), - t, - format('%s_dt', t), - t - ); + + -- Clamp the widen to the partition's structural bound (_constraint_dtrange). A sub-partitioned + -- (month/year) partition's bound is finite: cover the whole partition so dtrange never misses and + -- never spills into a neighbour. An unbounded partition (NULL-partition_trunc or NULL bound) has + -- nothing to clamp to, so pad the batch range by the widen buffer each side. + is_unbounded := _constraint_dtrange IS NULL + OR (lower(_constraint_dtrange) = '-infinity'::timestamptz + AND upper(_constraint_dtrange) = 'infinity'::timestamptz); + tail := GREATEST(upper(_edtrange) - upper(_dtrange), '0'::interval); + IF is_unbounded THEN + target_dtrange := tstzrange(lower(_dtrange) - buf, upper(_dtrange) + buf, '[]'); + target_edtrange := tstzrange(lower(_dtrange) - buf, upper(_edtrange) + buf, '[]'); ELSE - q :=format( - $q$ - DO $block$ - BEGIN - - ALTER TABLE %I DROP CONSTRAINT IF EXISTS %I; - ALTER TABLE %I - ADD CONSTRAINT %I - CHECK ( - (datetime >= %L) - AND (datetime <= %L) - AND (end_datetime >= %L) - AND (end_datetime <= %L) - ) NOT VALID - ; - ALTER TABLE %I - VALIDATE CONSTRAINT %I - ; - - - - EXCEPTION WHEN others THEN - RAISE WARNING '%%, Issue Altering Constraints. Please run update_partition_stats(%I)', SQLERRM USING ERRCODE = SQLSTATE; - END; - $block$; - $q$, - t, - format('%s_dt', t), - t, - format('%s_dt', t), - lower(_dtrange), - upper(_dtrange), - lower(_edtrange), - upper(_edtrange), - t, - format('%s_dt', t), - t - ); + target_dtrange := _constraint_dtrange; + target_edtrange := tstzrange(lower(_constraint_dtrange), upper(_constraint_dtrange) + tail, '[]'); END IF; - PERFORM run_or_queue(q); - RETURN t; + + new_dtrange := range_merge(COALESCE(cur.dtrange, target_dtrange), target_dtrange); + new_edtrange := range_merge(COALESCE(cur.edtrange, target_edtrange), target_edtrange); + new_spatial := CASE WHEN spatial_covered THEN cur.spatial ELSE NULL END; + + UPDATE partition_stats + SET dtrange = new_dtrange, + edtrange = new_edtrange, + spatial = new_spatial, + dirty = true, + last_updated = now() + WHERE partition = _partition; END; $$ LANGUAGE PLPGSQL SECURITY DEFINER; @@ -410,18 +290,14 @@ $$ LANGUAGE PLPGSQL SECURITY DEFINER; CREATE OR REPLACE FUNCTION check_partition( _collection text, _dtrange tstzrange, - _edtrange tstzrange + _edtrange tstzrange, + _spatial geometry DEFAULT NULL ) RETURNS text AS $$ DECLARE c RECORD; - pm RECORD; _partition_name text; - _partition_dtrange tstzrange; - _constraint_dtrange tstzrange; - _constraint_edtrange tstzrange; - q text; - deferrable_q text; - err_context text; + _parent_name text; + _partition_range tstzrange; BEGIN SELECT * INTO c FROM pgstac.collections WHERE id=_collection; IF NOT FOUND THEN @@ -429,115 +305,99 @@ BEGIN END IF; IF c.partition_trunc IS NOT NULL THEN - _partition_dtrange := tstzrange( + _partition_range := tstzrange( date_trunc(c.partition_trunc, lower(_dtrange)), date_trunc(c.partition_trunc, lower(_dtrange)) + (concat('1 ', c.partition_trunc))::interval, '[)' ); ELSE - _partition_dtrange := '[-infinity, infinity]'::tstzrange; + _partition_range := '[-infinity, infinity]'::tstzrange; END IF; - IF NOT _partition_dtrange @> _dtrange THEN - RAISE EXCEPTION 'dtrange % is greater than the partition size % for collection %', _dtrange, c.partition_trunc, _collection; + IF NOT _partition_range @> _dtrange THEN + RAISE EXCEPTION 'dtrange % spans more than the % partition window for collection %', _dtrange, c.partition_trunc, _collection; END IF; - + _parent_name := format('_items_%s', c.key); IF c.partition_trunc = 'year' THEN - _partition_name := format('_items_%s_%s', c.key, to_char(lower(_partition_dtrange),'YYYY')); + _partition_name := format('%s_%s', _parent_name, to_char(lower(_partition_range),'YYYY')); ELSIF c.partition_trunc = 'month' THEN - _partition_name := format('_items_%s_%s', c.key, to_char(lower(_partition_dtrange),'YYYYMM')); + _partition_name := format('%s_%s', _parent_name, to_char(lower(_partition_range),'YYYYMM')); ELSE - _partition_name := format('_items_%s', c.key); + _partition_name := _parent_name; END IF; - SELECT * INTO pm FROM partition_sys_meta WHERE collection=_collection AND partition_dtrange @> _dtrange; - IF FOUND THEN - RAISE NOTICE '% % %', _edtrange, _dtrange, pm; - _constraint_edtrange := - tstzrange( - least( - lower(_edtrange), - nullif(lower(pm.constraint_edtrange), '-infinity') - ), - greatest( - upper(_edtrange), - nullif(upper(pm.constraint_edtrange), 'infinity') - ), - '[]' - ); - _constraint_dtrange := - tstzrange( - least( - lower(_dtrange), - nullif(lower(pm.constraint_dtrange), '-infinity') - ), - greatest( - upper(_dtrange), - nullif(upper(pm.constraint_dtrange), 'infinity') - ), - '[]' + -- Create the collection-level PARENT partition (_items_) first, for sub-partitioned collections. + -- It is shared across every child window, so concurrent setup of different children would race on its + -- CREATE TABLE. Guard it with a parent-scoped advisory lock, taken before any child lock and only when + -- the parent is missing: parent-before-child is one lock order (parent < child) that can't deadlock + -- with ensure_partitions' sorted child locks, and skipping it once the parent exists keeps steady-state + -- ingest off the parent lock. + IF c.partition_trunc IS NOT NULL AND to_regclass(format('pgstac.%I', _parent_name)) IS NULL THEN + PERFORM pg_advisory_xact_lock(hashtext('pgstac.check_partition'), hashtext(_parent_name)); + IF to_regclass(format('pgstac.%I', _parent_name)) IS NULL THEN + EXECUTE format( + 'CREATE TABLE IF NOT EXISTS %I partition OF items FOR VALUES IN (%L) PARTITION BY RANGE (datetime)', + _parent_name, _collection ); - - IF pm.constraint_edtrange @> _edtrange AND pm.constraint_dtrange @> _dtrange THEN - RETURN pm.partition; - ELSE - PERFORM drop_table_constraints(_partition_name); END IF; - ELSE - _constraint_edtrange := _edtrange; - _constraint_dtrange := _dtrange; END IF; - RAISE NOTICE 'EXISTING CONSTRAINTS % %, NEW % %', pm.constraint_dtrange, pm.constraint_edtrange, _constraint_dtrange, _constraint_edtrange; - RAISE NOTICE 'Creating partition % %', _partition_name, _partition_dtrange; - IF c.partition_trunc IS NULL THEN - q := format( - $q$ - CREATE TABLE IF NOT EXISTS %I partition OF items FOR VALUES IN (%L); - CREATE UNIQUE INDEX IF NOT EXISTS %I ON %I (id); - GRANT ALL ON %I to pgstac_ingest; - $q$, - _partition_name, - _collection, - concat(_partition_name,'_pk'), - _partition_name, - _partition_name - ); - ELSE - q := format( - $q$ - CREATE TABLE IF NOT EXISTS %I partition OF items FOR VALUES IN (%L) PARTITION BY RANGE (datetime); - CREATE TABLE IF NOT EXISTS %I partition OF %I FOR VALUES FROM (%L) TO (%L); - CREATE UNIQUE INDEX IF NOT EXISTS %I ON %I (id); - GRANT ALL ON %I TO pgstac_ingest; - $q$, - format('_items_%s', c.key), - _collection, - _partition_name, - format('_items_%s', c.key), - lower(_partition_dtrange), - upper(_partition_dtrange), - format('%s_pk', _partition_name), - _partition_name, - _partition_name - ); + + -- Serialize concurrent setup of THIS (leaf) partition with a partition-scoped advisory lock. Different + -- partitions hash to different keys, so this never serializes unrelated ingest; ensure_partitions + -- acquires these in sorted order, so concurrent multi-partition batches cannot deadlock. The lock + -- releases at commit (setup is fast + idempotent). + PERFORM pg_advisory_xact_lock(hashtext('pgstac.check_partition'), hashtext(_partition_name)); + + -- Create the leaf partition if missing. Parent-inherited indexes only (id PK here; datetime/geometry + -- come from the items parent); no CHECK constraints; queryable indexes are deferred via + -- indexes_pending. A SELECT grant lets read/ingest query it; writes reach it only through the + -- SECURITY DEFINER write functions (the privilege wall in 998_idempotent_post). + -- Skip the DDL when it already exists: re-running CREATE/GRANT takes a relation lock that deadlocks + -- with concurrent INSERTs. The existence check is race-safe under the advisory lock. + IF to_regclass(format('pgstac.%I', _partition_name)) IS NULL THEN + BEGIN + IF c.partition_trunc IS NULL THEN + EXECUTE format( + $q$ + CREATE TABLE IF NOT EXISTS %I partition OF items FOR VALUES IN (%L); + CREATE UNIQUE INDEX IF NOT EXISTS %I ON %I (id); + GRANT SELECT ON %I TO pgstac_read, pgstac_ingest; + $q$, + _partition_name, _collection, + concat(_partition_name, '_pk'), _partition_name, + _partition_name + ); + ELSE + EXECUTE format( + $q$ + CREATE TABLE IF NOT EXISTS %I partition OF %I FOR VALUES FROM (%L) TO (%L); + CREATE UNIQUE INDEX IF NOT EXISTS %I ON %I (id); + GRANT SELECT ON %I TO pgstac_read, pgstac_ingest; + $q$, + _partition_name, _parent_name, lower(_partition_range), upper(_partition_range), + concat(_partition_name, '_pk'), _partition_name, + _partition_name + ); + END IF; + EXCEPTION + -- A concurrent creator that finished between our checks: benign, it exists now. + WHEN duplicate_table THEN + RAISE DEBUG 'Partition % already exists.', _partition_name; + -- Do NOT swallow other errors: a failed creation must propagate so the caller never goes on to + -- write a partition that does not exist (invariant: check_partition succeeds before use). + END; END IF; - BEGIN - EXECUTE q; - EXCEPTION - WHEN duplicate_table THEN - RAISE NOTICE 'Partition % already exists.', _partition_name; - WHEN others THEN - GET STACKED DIAGNOSTICS err_context = PG_EXCEPTION_CONTEXT; - RAISE INFO 'Error Name:%',SQLERRM; - RAISE INFO 'Error State:%', SQLSTATE; - RAISE INFO 'Error Context:%', err_context; - END; - PERFORM maintain_partitions(_partition_name); - PERFORM update_partition_stats_q(_partition_name, true); - REFRESH MATERIALIZED VIEW partitions; - REFRESH MATERIALIZED VIEW partition_steps; + -- Seed the partition_stats row (collection only — the data ranges start NULL), then cover this batch + -- via the shared widen guard, which fills dtrange/edtrange. ON CONFLICT keeps an existing row (so a + -- sweep that cleared indexes_pending/dirty is not reset on a later check_partition for the same + -- partition). + INSERT INTO partition_stats (partition, collection, dirty, indexes_pending) + VALUES (_partition_name, _collection, true, true) + ON CONFLICT (partition) DO NOTHING; + PERFORM widen_partition_stats(_partition_name, _dtrange, _edtrange, _spatial, _partition_range); + RETURN _partition_name; END; $$ LANGUAGE PLPGSQL SECURITY DEFINER; diff --git a/src/pgstac/sql/003c_ingest.sql b/src/pgstac/sql/003c_ingest.sql new file mode 100644 index 00000000..c8769dba --- /dev/null +++ b/src/pgstac/sql/003c_ingest.sql @@ -0,0 +1,210 @@ +-- --------------------------------------------------------------------------- +-- Rust-first ingest: SECURITY DEFINER staging + flush + fragment helpers. +-- +-- These are the server-side seam for the Rust loader's binary-COPY path. The connecting role never +-- writes the real tables directly; it binary-COPYs fully-dehydrated rows into a session-local TEMP +-- table (make_binary_staging) and calls flush_items_staging_binary, which is the only thing that writes +-- `items`. Partition existence + stats (extent + n) and registry/fragment coverage are established BEFORE +-- the load in their own transactions (prepare_partition_for_load / ensure_fragments); the flush writes only +-- `items` and never touches partition_stats. +-- --------------------------------------------------------------------------- + +-- ensure_fragments: upsert per-collection fragment payloads and return, for each input position, the +-- item_fragments id the Rust loader stamps into items.fragment_id. +-- +-- The caller (the Rust loader) deduplicates fragments locally and sends only the DISTINCT set (one input +-- element per unique fragment), then maps each item -> its fragment via the returned `ord`. Sending one +-- fragment per item would ship millions of duplicates when a collection has only a handful of distinct +-- fragments. The function is nonetheless dup-safe (a DISTINCT ON guards the insert), so a caller that +-- passes duplicates still gets the correct id for every position. +-- +-- Each input element is {"content": , "links_template": }; the canonical +-- hash matches pgstac_hash_fragment (so it dedups identically to the SQL ingest path). ON CONFLICT keeps +-- existing rows; pre-existing ids come from item_fragments (snapshot), new ids from the INSERT's RETURNING. +CREATE OR REPLACE FUNCTION ensure_fragments( + _collection text, + _fragments jsonb[] +) RETURNS TABLE (ord int, frag_id bigint) AS $$ + WITH input AS ( + SELECT + o::int AS ord, + pgstac_hash_fragment( + jsonb_strip_nulls(jsonb_build_object( + 'content', NULLIF(f->'content', '{}'::jsonb), + 'links_template', f->'links_template' + )) + ) AS hash, + COALESCE(f->'content', '{}'::jsonb) AS content, + f->'links_template' AS links_template + FROM unnest(_fragments) WITH ORDINALITY AS u(f, o) + ), + distinct_hashes AS ( + SELECT DISTINCT ON (hash) hash, content, links_template + FROM input + ORDER BY hash + ), + upserted AS ( + INSERT INTO item_fragments (collection, hash, content, links_template) + SELECT _collection, hash, content, links_template FROM distinct_hashes + -- DO UPDATE (a no-op write) rather than DO NOTHING: it locks + RETURNS the conflicting row even when + -- a CONCURRENT transaction committed the same (collection, hash) mid-statement. DO NOTHING returns + -- nothing on conflict, and the old "UNION the existing rows via a separate SELECT" ran on the + -- statement-start snapshot, so it could MISS such a concurrent insert -> that hash dropped out of the + -- final join and the item was stamped with NO fragment_id. DO UPDATE returns every distinct hash + -- exactly once (newly inserted or pre-existing), making the stamp concurrency-safe. + ON CONFLICT (collection, hash) DO UPDATE SET content = item_fragments.content + RETURNING id, hash + ) + SELECT input.ord, upserted.id + FROM input JOIN upserted ON input.hash = upserted.hash + ORDER BY input.ord; +$$ LANGUAGE SQL SECURITY DEFINER; + + +-- ensure_partitions: create + widen every partition a batch will land in, BEFORE the load (widen-now). +-- The Rust loader passes parallel arrays of (collection, datetime, end_datetime) — one element per item. +-- Grouping happens HERE, server-side, so the month/year truncation uses the server's date_trunc (correct +-- timezone) rather than re-deriving partition boundaries in the client. One call per batch replaces the +-- loader's per-item check_partition round trips: it groups by (collection, partition window) and calls +-- check_partition once per distinct partition with that group's datetime range. Runs as its own committed +-- statement before the load transaction, so partitions exist + their stats cover the batch by the time +-- the binary COPY + flush run. +CREATE OR REPLACE FUNCTION ensure_partitions( + _collections text[], + _datetimes timestamptz[], + _end_datetimes timestamptz[] +) RETURNS void AS $$ +DECLARE + r RECORD; +BEGIN + -- Call check_partition per distinct (collection, partition window) in a DETERMINISTIC order: each + -- check_partition takes that partition's advisory lock, so locking in a consistent order prevents two + -- concurrent multi-partition batches from advisory-deadlocking on the locks. (flush locks in the same + -- sorted order.) A PLPGSQL FOR loop guarantees the order; a set-returning SELECT would not. + FOR r IN + WITH t AS ( + SELECT + unnest(_collections) AS collection, + unnest(_datetimes) AS dt, + unnest(_end_datetimes) AS edt + ), + j AS ( + SELECT t.collection, t.dt, t.edt, c.partition_trunc + FROM t JOIN pgstac.collections c ON t.collection = c.id + ) + SELECT + collection, + tstzrange(min(dt), max(dt), '[]') AS dtrange, + tstzrange(min(edt), max(edt), '[]') AS edtrange + FROM j + GROUP BY collection, COALESCE(date_trunc(partition_trunc::text, dt), '-infinity'::timestamptz) + ORDER BY collection, COALESCE(date_trunc(partition_trunc::text, dt), '-infinity'::timestamptz) + LOOP + PERFORM pgstac.check_partition(r.collection, r.dtrange, r.edtrange); + END LOOP; +END; +$$ LANGUAGE PLPGSQL SECURITY DEFINER; + + +-- prepare_partition_for_load: per-partition metadata for the Rust direct / precheck load paths. ONE small +-- self-contained txn per partition, run BEFORE any COPY: create + widen the partition to cover this batch +-- (dt + edt + the real SPATIAL envelope, unlike ensure_partitions which passes NULL) and bump n, so +-- partition_stats is at least as wide as the data (golden rule) and search treats the partition as non-empty +-- before the data lands. Returns the partition name + the pre-load row count (n BEFORE the bump) so the +-- loader can choose its adaptive precheck path: empty -> skip the precheck; batch > n -> pull the partition's +-- (id,item_hash) to the client; n >= batch -> COPY the batch (id,hash) to a temp table + JOIN this partition. +CREATE OR REPLACE FUNCTION prepare_partition_for_load( + _collection text, + _dt_lo timestamptz, _dt_hi timestamptz, + _edt_lo timestamptz, _edt_hi timestamptz, + _xmin float8, _ymin float8, _xmax float8, _ymax float8, + _n_add bigint, + OUT partition_name text, + OUT pre_load_n bigint +) AS $$ +BEGIN + partition_name := pgstac.check_partition( + _collection, + tstzrange(_dt_lo, _dt_hi, '[]'), + tstzrange(_edt_lo, _edt_hi, '[]'), + st_setsrid(st_makeenvelope(_xmin, _ymin, _xmax, _ymax), 4326) + ); + SELECT COALESCE(n, 0) INTO pre_load_n + FROM pgstac.partition_stats WHERE partition = partition_name; + -- Over-estimating n is the safe direction; the async tightener computes the exact count + extent off the + -- hot path. Single-row atomic UPDATE, so concurrent loads into the same partition serialize on the row. + UPDATE pgstac.partition_stats + SET n = COALESCE(n, 0) + _n_add, dirty = true, last_updated = now() + WHERE partition = partition_name; +END; +$$ LANGUAGE PLPGSQL SECURITY DEFINER; + + +-- make_binary_staging: create a session-local TEMP staging table shaped exactly like `items` +-- (ON COMMIT DROP) and grant INSERT to pgstac_ingest so the connecting role can binary-COPY into it. +-- Created inside this SECURITY DEFINER function, so the temp table is owned by pgstac_admin (which also +-- owns `items`), letting flush_items_staging_binary read it. Returns the generated table name. +CREATE OR REPLACE FUNCTION make_binary_staging() RETURNS text AS $$ +DECLARE + _name text := format('_staging_%s', replace(gen_random_uuid()::text, '-', '')); +BEGIN + EXECUTE format('CREATE TEMP TABLE %I (LIKE items) ON COMMIT DROP', _name); + EXECUTE format('GRANT INSERT ON pg_temp.%I TO pgstac_ingest', _name); + RETURN _name; +END; +$$ LANGUAGE PLPGSQL SECURITY DEFINER; + + +-- flush_items_staging_binary: move fully-dehydrated rows from a TEMP staging table into `items` with the +-- conflict policy. partition_stats (extent + n + dirty) is set entirely by prepare_partition_for_load in +-- the preflight, so this writes only `items`: +-- ignore -> ON CONFLICT DO NOTHING (idempotent; orphans the old row on a cross-partition move) +-- upsert -> delete a changed row IN THE PARTITION IT ROUTES TO (window-pruned), then insert. A +-- datetime change that stays in the partition is applied; one that moves the item to a +-- different partition orphans the old row (use 'delsert'). +-- delsert -> delete a changed row CROSS-partition (by collection+id, move-safe), then insert +-- error -> plain INSERT (raises on any duplicate) +-- Returns the number of rows inserted. +CREATE OR REPLACE FUNCTION flush_items_staging_binary( + _staging text, + _policy text DEFAULT 'ignore' +) RETURNS bigint AS $$ +DECLARE + nrows bigint; +BEGIN + IF _policy = 'ignore' THEN + EXECUTE format('INSERT INTO items SELECT * FROM %1$I ON CONFLICT DO NOTHING', _staging); + ELSIF _policy = 'error' THEN + EXECUTE format('INSERT INTO items SELECT * FROM %1$I', _staging); + ELSIF _policy = 'upsert' THEN + -- Fast SAME-partition upsert: delete a changed row only in the partition the incoming row routes to. + -- The window range [date_trunc(trunc, s.datetime), + 1 trunc) is derived per staged row, so the + -- planner runtime-prunes `items` to that one partition (no cross-partition scan). A datetime change + -- within the partition is applied; one that MOVES the item to another partition isn't seen here, so + -- the old row orphans (use 'delsert'). NULL partition_trunc => the collection's single partition. + EXECUTE format($q$ + DELETE FROM items i USING %1$I s JOIN collections c ON c.id = s.collection + WHERE i.collection = s.collection AND i.id = s.id + AND i.datetime >= (CASE WHEN c.partition_trunc IS NULL THEN '-infinity'::timestamptz + ELSE date_trunc(c.partition_trunc::text, s.datetime) END) + AND i.datetime < (CASE WHEN c.partition_trunc IS NULL THEN 'infinity'::timestamptz + ELSE date_trunc(c.partition_trunc::text, s.datetime) + + ('1 ' || c.partition_trunc::text)::interval END) + AND ( %2$s ) + $q$, _staging, items_content_distinct_sql('i', 's')); + EXECUTE format('INSERT INTO items SELECT * FROM %1$I ON CONFLICT DO NOTHING', _staging); + ELSIF _policy = 'delsert' THEN + -- Move-safe CROSS-partition upsert: delete the old row wherever it lives (by collection+id, no + -- datetime bound, so it probes every partition), then insert. + EXECUTE format($q$ + DELETE FROM items i USING %1$I s + WHERE i.id = s.id AND i.collection = s.collection AND ( %2$s ) + $q$, _staging, items_content_distinct_sql('i', 's')); + EXECUTE format('INSERT INTO items SELECT * FROM %1$I ON CONFLICT DO NOTHING', _staging); + ELSE + RAISE EXCEPTION 'unknown conflict policy % (expected ignore | upsert | delsert | error)', _policy; + END IF; + GET DIAGNOSTICS nrows = ROW_COUNT; + RETURN nrows; +END; +$$ LANGUAGE PLPGSQL SECURITY DEFINER; diff --git a/src/pgstac/sql/004_search.sql b/src/pgstac/sql/004_search.sql index ff173bf0..3a86520f 100644 --- a/src/pgstac/sql/004_search.sql +++ b/src/pgstac/sql/004_search.sql @@ -222,8 +222,16 @@ $$ LANGUAGE PLPGSQL STABLE PARALLEL SAFE; -- fields_to_itemcols: produce a SELECT list of item columns in attnum order. -- Heavy columns (geometry, bbox, assets, links, properties, extra) are emitted as --- NULL::type when their controlling field is excluded/not-included. -CREATE OR REPLACE FUNCTION fields_to_itemcols(fields jsonb DEFAULT '{}'::jsonb) RETURNS text AS $$ +-- NULL::type AS when their controlling field is excluded/not-included: the +-- column is kept (named) so the result schema is stable across any `fields` value +-- -- a client preparing this query once can step bands without re-binding columns -- +-- while the heavy value itself is never transferred. +-- When _skip_fragment is true (the caller determined the requested fields are +-- satisfiable from item columns alone, via needs_fragment), fragment_id is emitted +-- as NULL too: a client driving this query then sees no fragment_id and skips the +-- per-row item_fragments lookup entirely -- the fragment-skip is carried in the +-- projection itself, not a separate flag. +CREATE OR REPLACE FUNCTION fields_to_itemcols(fields jsonb DEFAULT '{}'::jsonb, _skip_fragment boolean DEFAULT false) RETURNS text AS $$ DECLARE includes text[] := ARRAY(SELECT jsonb_array_elements_text(fields->'include')); excludes text[] := ARRAY(SELECT jsonb_array_elements_text(fields->'exclude')); @@ -231,11 +239,13 @@ DECLARE BEGIN SELECT string_agg( CASE + WHEN a.attname = 'fragment_id' AND _skip_fragment + THEN format('NULL::%s AS %I', format_type(a.atttypid, a.atttypmod), a.attname) WHEN a.attname = ANY (ARRAY['geometry','bbox','assets','links','link_hrefs', 'extra','properties','stac_version','stac_extensions']) AND NOT field_included(CASE WHEN a.attname = 'link_hrefs' THEN 'links' ELSE a.attname END, includes, excludes) - THEN format('NULL::%s', format_type(a.atttypid, a.atttypmod)) + THEN format('NULL::%s AS %I', format_type(a.atttypid, a.atttypmod), a.attname) ELSE format('i.%I', a.attname) END, ', ' ORDER BY a.attnum) INTO cols @@ -276,7 +286,7 @@ DECLARE band_cap_months CONSTANT int := 18; page_rows items[] := '{}'::items[]; chunk_rows items[]; target int := _limit + 1; got int := 0; got_band int; - is_asc boolean; cursor_ts timestamptz; band record; + is_asc boolean; band record; band_target numeric; obs_sel numeric; band_where text; guard int := 0; cum_scanned bigint := 0; proj_expr text; mo interval := interval '1 month'; @@ -332,12 +342,14 @@ BEGIN END IF; IF datetime_leading AND array_length(bnds.months, 1) IS NOT NULL THEN - cursor_ts := CASE WHEN is_asc THEN bnds.months[1] ELSE bnds.months[array_length(bnds.months, 1)] + mo END; - cursor_idx := 1; + -- Start the band walk at the leading edge for the sort direction: oldest month for ascending, + -- most recent month for descending. (A descending walk that started at the oldest band would + -- collect the oldest items and never reach the recent months.) + cursor_idx := CASE WHEN is_asc THEN 1 ELSE array_length(bnds.months, 1) END; band_target := target * band_margin; WHILE got < target AND guard < 80 LOOP guard := guard + 1; - SELECT * INTO band FROM next_band(bnds.counts, cursor_idx, band_target, band_cap_months); + SELECT * INTO band FROM next_band(bnds.counts, cursor_idx, band_target, band_cap_months, NOT is_asc); -- a valid band must be processed even when next_band also flags done (it consumed the -- last bucket); only stop when there is no band at all. EXIT WHEN band.band_start_idx IS NULL; @@ -471,7 +483,6 @@ DECLARE bnds record; coll_clamp text := ''; clamp text; BEGIN IF _where IS NULL OR btrim(_where) = '' THEN _where := ' TRUE '; END IF; - collist := fields_to_itemcols(coalesce(_search->'fields', '{}'::jsonb)); orderby_str := keyset_orderby(_search, is_prev); SELECT (array_agg(field ORDER BY ord))[1], (array_agg(CASE WHEN is_prev THEN (CASE dir WHEN 'ASC' THEN 'DESC' ELSE 'ASC' END) ELSE dir END ORDER BY ord))[1] @@ -497,10 +508,20 @@ BEGIN coll_clamp := format('i.collection = ANY(%L::text[]) AND ', bnds.collections); END IF; + -- Build the projection now that the resolved collections are known: when the requested fields are + -- satisfiable from item columns alone (needs_fragment is false), fragment_id is nulled in the + -- projection so a client driving this query never looks up item_fragments -- the fragment-skip + -- rides in the returned query's SELECT list, not a separate flag. + collist := fields_to_itemcols( + coalesce(_search->'fields', '{}'::jsonb), + NOT needs_fragment(coalesce(_search->'fields', '{}'::jsonb), bnds.collections)); + IF datetime_leading AND array_length(bnds.months, 1) IS NOT NULL THEN + -- collections baked in as a literal so the query is parameterized only by $1 (band low), + -- $2 (band high), $3 (limit): the client prepares it once and binds per histogram band. query := format( - 'SELECT %s FROM items i WHERE i.collection = ANY($4) AND i.datetime >= $1 AND i.datetime < $2 AND (%s) ORDER BY %s LIMIT $3', - collist, full_where, orderby_str); + 'SELECT %s FROM items i WHERE %si.datetime >= $1 AND i.datetime < $2 AND (%s) ORDER BY %s LIMIT $3', + collist, coll_clamp, full_where, orderby_str); -- serialize the per-month histogram (months[]/counts[]) to jsonb for the streaming client histogram := ( SELECT jsonb_agg(jsonb_build_object('m', m, 'n', n) ORDER BY ord) @@ -509,7 +530,10 @@ BEGIN ELSE DECLARE dt_clamp text := ''; BEGIN IF array_length(bnds.months, 1) IS NOT NULL THEN - dt_clamp := format('i.datetime >= %L AND i.datetime <= %L AND ', bnds.months[1], bnds.months[array_length(bnds.months, 1)]); + -- months[] holds month *starts*; the last bucket spans [start, start + 1 month), so the + -- upper bound must be exclusive of the next month, not <= the last month's start (which + -- would drop items dated after the start of the final month). + dt_clamp := format('i.datetime >= %L AND i.datetime < %L AND ', bnds.months[1], bnds.months[array_length(bnds.months, 1)] + interval '1 month'); END IF; query := format( 'SELECT %s FROM items i WHERE %s%s(%s) ORDER BY %s LIMIT $1', diff --git a/src/pgstac/sql/006_tilesearch.sql b/src/pgstac/sql/006_tilesearch.sql index d33ec264..6aa557aa 100644 --- a/src/pgstac/sql/006_tilesearch.sql +++ b/src/pgstac/sql/006_tilesearch.sql @@ -78,12 +78,17 @@ BEGIN END IF; IF array_length(bnds.months, 1) IS NOT NULL THEN - cursor_idx := 1; + -- Walk bands in the SEARCH's datetime direction: newest month first for a descending search, + -- oldest first when sorting datetime ascending. Mirrors search_page (004_search). Walking the + -- wrong direction returns older items before newer ones for a descending limit/early-exit search, + -- and (because band grouping shifts with the partition_stats histogram) makes the result depend on + -- stats — a bug, since stats must only affect performance, never which rows come back. + cursor_idx := CASE WHEN is_asc THEN 1 ELSE array_length(bnds.months, 1) END; band_target := (_limit + 1) * band_margin; <> WHILE NOT exit_flag AND guard < 80 LOOP guard := guard + 1; - SELECT * INTO band FROM next_band(bnds.counts, cursor_idx, band_target, band_cap_months); + SELECT * INTO band FROM next_band(bnds.counts, cursor_idx, band_target, band_cap_months, NOT is_asc); -- process a valid band even when next_band also flags done; stop only on no band. EXIT bands WHEN band.band_start_idx IS NULL; query := format( diff --git a/src/pgstac/sql/997_maintenance.sql b/src/pgstac/sql/997_maintenance.sql index 3c3e3b6e..6d912ac0 100644 --- a/src/pgstac/sql/997_maintenance.sql +++ b/src/pgstac/sql/997_maintenance.sql @@ -1,4 +1,64 @@ +-- tighten_dirty_partition_stats: recompute the exact envelope + row count for dirty partitions (oldest +-- first), clearing dirty. Run off-hours (pg_cron or the maintenance CLI). Optional: a wide envelope only +-- over-includes a partition in search, so skipping it never loses rows. `_limit` caps the batch (NULL = +-- all dirty); returns the number of partitions tightened. +-- +-- pg_cron example (operators install this themselves): +-- SELECT cron.schedule('pgstac-tighten', '*/15 * * * *', +-- $$SELECT pgstac.tighten_dirty_partition_stats(200)$$); +CREATE OR REPLACE FUNCTION tighten_dirty_partition_stats(_limit int DEFAULT NULL) +RETURNS int AS $$ +DECLARE + _part text; + _count int := 0; +BEGIN + FOR _part IN + SELECT partition FROM pgstac.partition_stats + WHERE dirty + ORDER BY last_updated NULLS FIRST + LIMIT _limit + LOOP + PERFORM pgstac.tighten_partition_stats(_part); + _count := _count + 1; + END LOOP; + RETURN _count; +END; +$$ LANGUAGE PLPGSQL SECURITY DEFINER; + + +-- build_pending_indexes: build the queryable indexes for partitions flagged `indexes_pending` (new +-- partitions are created index-light for fast ingest; changing a queryable flags every partition via the +-- queryables trigger), then clear the flag. Builds the DDL directly (not via the queue), so schedule it +-- off-hours like the tighten sweep. `_limit` caps the batch (NULL = all pending); returns the count built. +-- +-- pg_cron example (operators install this themselves): +-- SELECT cron.schedule('pgstac-build-indexes', '*/30 * * * *', +-- $$SELECT pgstac.build_pending_indexes(50)$$); +CREATE OR REPLACE FUNCTION build_pending_indexes(_limit int DEFAULT NULL) +RETURNS int AS $$ +DECLARE + _part text; + _q text; + _count int := 0; +BEGIN + FOR _part IN + SELECT partition FROM pgstac.partition_stats + WHERE indexes_pending + ORDER BY last_updated NULLS FIRST + LIMIT _limit + LOOP + FOR _q IN SELECT * FROM pgstac.maintain_partition_queries(_part) LOOP + EXECUTE _q; + END LOOP; + UPDATE pgstac.partition_stats SET indexes_pending = false WHERE partition = _part; + _count := _count + 1; + END LOOP; + RETURN _count; +END; +$$ LANGUAGE PLPGSQL SECURITY DEFINER; + + CREATE OR REPLACE PROCEDURE analyze_items() AS $$ DECLARE q text; @@ -57,7 +117,7 @@ DECLARE extent jsonb; BEGIN IF runupdate THEN - PERFORM update_partition_stats_q(partition) + PERFORM tighten_partition_stats(partition) FROM partitions_view WHERE collection=_collection; END IF; SELECT @@ -149,3 +209,237 @@ BEGIN END; $$ LANGUAGE PLPGSQL; + +-- update_collection_extents: recompute every collection's extent from partitions_view. +CREATE OR REPLACE FUNCTION update_collection_extents() RETURNS VOID AS $$ +UPDATE collections + SET content = jsonb_set_lax( + content, + '{extent}'::text[], + collection_extent(id, FALSE), + true, + 'use_json_null' + ) +; +$$ LANGUAGE SQL; + + +-- --------------------------------------------------------------------------- +-- Field registry maintenance: track which paths (and value types) exist per +-- collection. Used for schema inference (e.g. the geoparquet export schema). +-- jsonb_field_rows is in 001a_jsonutils.sql. +-- --------------------------------------------------------------------------- + +-- update_field_registry_from_sample: UPSERT registry rows from an array of raw item content JSONBs (the +-- caller picks the sample); value_kinds accumulate observed types over time. +CREATE OR REPLACE FUNCTION update_field_registry_from_sample( + _collection text, + item_contents jsonb[] +) RETURNS void AS $$ + INSERT INTO item_field_registry (collection, path, is_leaf, value_kinds, first_seen, last_seen) + SELECT + _collection, + r.path, + bool_and(r.is_leaf) AS is_leaf, + array_agg(DISTINCT r.value_kind) FILTER (WHERE r.value_kind IS NOT NULL) AS value_kinds, + now(), + now() + FROM unnest(item_contents) AS item(content) + CROSS JOIN LATERAL jsonb_field_rows(item.content) AS r(path, is_leaf, value_kind) + GROUP BY r.path + ON CONFLICT (collection, path) DO UPDATE SET + is_leaf = EXCLUDED.is_leaf, + value_kinds = ( + SELECT array_agg(DISTINCT v) + FROM unnest(item_field_registry.value_kinds || EXCLUDED.value_kinds) t(v) + ), + last_seen = now() + ; +$$ LANGUAGE SQL VOLATILE SECURITY DEFINER; + + +-- update_field_registry_from_items: sample a live collection and UPSERT registry rows (TABLESAMPLE +-- BERNOULLI(5) above ~10k rows by pg_class estimate, else LIMIT 1000). Returns (registered_paths, +-- rows_processed). +CREATE OR REPLACE FUNCTION update_field_registry_from_items( + _collection text +) RETURNS TABLE (registered_paths int, rows_processed int) AS $$ +DECLARE + est_rows bigint; + nrows int; + npaths int; +BEGIN + -- Sum reltuples across the registered item partitions for this collection. + -- reltuples can be -1 (never analyzed); treat negative values as zero. + SELECT COALESCE(sum(GREATEST(c.reltuples::bigint, 0)), 0) INTO est_rows + FROM partitions_view p + JOIN pg_class c ON c.relname = p.partition + JOIN pg_namespace n ON n.oid = c.relnamespace + WHERE p.collection = _collection + AND n.nspname = 'pgstac' + AND c.relkind = 'r'; + + IF est_rows > 10000 THEN + -- Large collection: use statistical sampling to avoid full seq-scan. + WITH sampled AS ( + SELECT content_hydrate(i) AS content FROM items i TABLESAMPLE BERNOULLI(5) WHERE i.collection = _collection + ), + upserted AS ( + INSERT INTO item_field_registry (collection, path, is_leaf, value_kinds, first_seen, last_seen) + SELECT + _collection, + r.path, + bool_and(r.is_leaf) AS is_leaf, + array_agg(DISTINCT r.value_kind) FILTER (WHERE r.value_kind IS NOT NULL) AS value_kinds, + now(), now() + FROM sampled + CROSS JOIN LATERAL jsonb_field_rows(content) AS r(path, is_leaf, value_kind) + GROUP BY r.path + ON CONFLICT (collection, path) DO UPDATE SET + is_leaf = EXCLUDED.is_leaf, + value_kinds = ( + SELECT array_agg(DISTINCT v) + FROM unnest(item_field_registry.value_kinds || EXCLUDED.value_kinds) t(v) + ), + last_seen = now() + RETURNING 1 + ) + SELECT + (SELECT count(*)::int FROM upserted), + (SELECT count(*)::int FROM sampled) + INTO npaths, nrows; + ELSE + -- Small collection: process up to 1000 rows to avoid BERNOULLI returning 0 rows. + WITH sampled AS ( + SELECT content_hydrate(i) AS content FROM items i WHERE i.collection = _collection LIMIT 1000 + ), + upserted AS ( + INSERT INTO item_field_registry (collection, path, is_leaf, value_kinds, first_seen, last_seen) + SELECT + _collection, + r.path, + bool_and(r.is_leaf) AS is_leaf, + array_agg(DISTINCT r.value_kind) FILTER (WHERE r.value_kind IS NOT NULL) AS value_kinds, + now(), now() + FROM sampled + CROSS JOIN LATERAL jsonb_field_rows(content) AS r(path, is_leaf, value_kind) + GROUP BY r.path + ON CONFLICT (collection, path) DO UPDATE SET + is_leaf = EXCLUDED.is_leaf, + value_kinds = ( + SELECT array_agg(DISTINCT v) + FROM unnest(item_field_registry.value_kinds || EXCLUDED.value_kinds) t(v) + ), + last_seen = now() + RETURNING 1 + ) + SELECT + (SELECT count(*)::int FROM upserted), + (SELECT count(*)::int FROM sampled) + INTO npaths, nrows; + END IF; + + RETURN QUERY SELECT npaths, nrows; +END; +$$ LANGUAGE PLPGSQL VOLATILE SECURITY DEFINER; + + +-- refresh_field_registry: expire registry entries not seen within retention_interval (scheduled +-- maintenance). Returns (collection, expired_paths) per affected collection. +CREATE OR REPLACE FUNCTION refresh_field_registry( + _collection text DEFAULT NULL, + retention_interval interval DEFAULT '90 days' +) RETURNS TABLE (collection_id text, expired_paths int) AS $$ + WITH deleted AS ( + DELETE FROM item_field_registry + WHERE (_collection IS NULL OR collection = _collection) + AND last_seen < now() - retention_interval + RETURNING collection + ) + SELECT collection, count(*)::int + FROM deleted + GROUP BY collection; +$$ LANGUAGE SQL VOLATILE SECURITY DEFINER; + + +-- gc_fragments: garbage-collect orphaned item_fragments with a single set-based DELETE (NOT EXISTS +-- anti-join against items.fragment_id). items.fragment_id has no FK (partitioned-items incremental +-- NOT VALID FKs aren't supported), so a fragment unreferenced at the DELETE snapshot but referenced by a +-- concurrent insert can be removed; the retention_interval guard makes this unlikely. Run during +-- low-ingest periods. +CREATE OR REPLACE FUNCTION gc_fragments( + _collection text DEFAULT NULL, + retention_interval interval DEFAULT '90 days' +) RETURNS TABLE ( + collection_id text, + fragments_removed int +) AS $$ + WITH deleted AS ( + DELETE FROM item_fragments f + WHERE + (_collection IS NULL OR f.collection = _collection) + AND f.created_at < now() - retention_interval + AND NOT EXISTS (SELECT 1 FROM items i WHERE i.fragment_id = f.id) + RETURNING f.collection + ) + SELECT collection, count(*)::int + FROM deleted + GROUP BY collection; +$$ LANGUAGE SQL VOLATILE PARALLEL UNSAFE SECURITY DEFINER; + + +-- tighten_partition_stats: recompute the exact envelope (datetime/end_datetime min-max + spatial extent) +-- and row count for one partition, then clear `dirty`. The only function that narrows a partition_stats +-- envelope; the maintenance sweeps drive it over partition_stats WHERE dirty. An empty partition tightens +-- to an 'empty' range so search prunes it out. +CREATE OR REPLACE FUNCTION tighten_partition_stats(_partition text) RETURNS void AS $$ +DECLARE + _n bigint; + _dtmin timestamptz; _dtmax timestamptz; + _edtmin timestamptz; _edtmax timestamptz; + _spatial geometry; + _dtrange tstzrange; + _edtrange tstzrange; + _collection text; +BEGIN + -- Hold the partition's advisory lock across the scan + write: tighten narrows dtrange to the scanned + -- extent and clears dirty, so without the lock a concurrent ingest committing a row outside that + -- extent would be left uncovered (search would prune + miss it). Same lock check_partition uses, so + -- tighten serializes with ingest into this partition only. + PERFORM pg_advisory_xact_lock(hashtext('pgstac.check_partition'), hashtext(_partition)); + + EXECUTE format( + $q$ + SELECT count(*), min(datetime), max(datetime), min(end_datetime), max(end_datetime), + st_setsrid(st_extent(geometry)::geometry, 4326) + FROM %I + $q$, + _partition + ) INTO _n, _dtmin, _dtmax, _edtmin, _edtmax, _spatial; + + IF _n = 0 THEN + _dtrange := 'empty'::tstzrange; + _edtrange := 'empty'::tstzrange; + _spatial := NULL; + ELSE + _dtrange := tstzrange(_dtmin, _dtmax, '[]'); + _edtrange := tstzrange(_edtmin, _edtmax, '[]'); + END IF; + + SELECT pv.collection + INTO _collection + FROM partitions_view pv WHERE partition = _partition; + + INSERT INTO partition_stats + (partition, collection, dtrange, edtrange, spatial, n, dirty, last_updated) + VALUES (_partition, _collection, _dtrange, _edtrange, _spatial, _n, false, now()) + ON CONFLICT (partition) DO UPDATE SET + collection = EXCLUDED.collection, + dtrange = EXCLUDED.dtrange, + edtrange = EXCLUDED.edtrange, + spatial = EXCLUDED.spatial, + n = EXCLUDED.n, + dirty = false, + last_updated = now(); +END; +$$ LANGUAGE PLPGSQL SECURITY DEFINER; diff --git a/src/pgstac/sql/998_idempotent_post.sql b/src/pgstac/sql/998_idempotent_post.sql index 65e8d4fc..8b68605c 100644 --- a/src/pgstac/sql/998_idempotent_post.sql +++ b/src/pgstac/sql/998_idempotent_post.sql @@ -150,6 +150,21 @@ GRANT EXECUTE ON ALL FUNCTIONS IN SCHEMA pgstac to pgstac_ingest; GRANT ALL ON ALL TABLES IN SCHEMA pgstac to pgstac_ingest; GRANT USAGE ON ALL SEQUENCES IN SCHEMA pgstac to pgstac_ingest; +-- Privilege wall (INV-1): the connecting role (inheriting pgstac_ingest) may READ these tables but may +-- only MUTATE them through the SECURITY DEFINER write functions (check_partition, widen/tighten_partition_stats, +-- ensure_fragments, make_binary_staging, flush_items_staging_binary, the items_staging trigger, create/ +-- update/upsert/delete_item, the field-registry/fragment helpers). This makes an un-widened or +-- envelope-narrowing direct write structurally impossible. New partitions are created SELECT-only for +-- pgstac_ingest by check_partition; the items parent is revoked here. Staging tables (items_staging*) stay +-- writable so the SQL-only ingest path can COPY into them. +REVOKE INSERT, UPDATE, DELETE, TRUNCATE ON items, partition_stats, item_fragments FROM pgstac_ingest; +-- item_field_registry is the one exception to the wall: the loader WIDENS it directly with an add-only +-- INSERT ... ON CONFLICT DO UPDATE (no SD function), so INSERT + UPDATE stay granted. DELETE/TRUNCATE remain +-- revoked, so even a direct write cannot NARROW the registry — INV-1 (registry is a superset of the data) +-- holds structurally against narrowing, while the cheap widen stays on the hot load path. +REVOKE DELETE, TRUNCATE ON item_field_registry FROM pgstac_ingest; + + REVOKE ALL PRIVILEGES ON PROCEDURE run_queued_queries FROM public; GRANT ALL ON PROCEDURE run_queued_queries TO pgstac_admin; @@ -161,6 +176,6 @@ GRANT ALL ON PROCEDURE gc_deleted_items_log_committed(interval, integer) TO pgst RESET ROLE; -SET ROLE pgstac_ingest; -SELECT update_partition_stats_q(partition) FROM partitions_view; -RESET ROLE; +-- (No install-time stats seeding: a fresh install has no partitions, and partition_stats rows are now +-- seeded by check_partition as partitions are created. Exact stats come from tighten_partition_stats via +-- the maintenance sweeps.) diff --git a/src/pgstac/tests/basic/crud_functions.sql b/src/pgstac/tests/basic/crud_functions.sql index 5444346a..8657fdc2 100644 --- a/src/pgstac/tests/basic/crud_functions.sql +++ b/src/pgstac/tests/basic/crud_functions.sql @@ -2,8 +2,9 @@ SET ROLE pgstac_ingest; SET pgstac.use_queue=FALSE; SELECT get_setting_bool('use_queue'); -SET pgstac.update_collection_extent=TRUE; -SELECT get_setting_bool('update_collection_extent'); +-- NOTE: collection extent is no longer auto-updated on ingest. It is refreshed EXPLICITLY via +-- update_collection_extents() (exercised at the end of this test), so the inline checks below show an +-- unset extent right after ingest. --create base data to use with tests CREATE TEMP TABLE test_items AS SELECT jsonb_build_object( @@ -20,14 +21,14 @@ INSERT INTO collections (content, partition_trunc) VALUES ('{"id":"pgstactest-cr SELECT create_item((SELECT content FROM test_items LIMIT 1)); SELECT id, geometry, collection, datetime, end_datetime, properties, extra FROM items WHERE collection='pgstactest-crudtest' ORDER BY id; --- Check to see if extent got updated +-- Extent is NOT auto-updated on ingest (now explicit/async) -- expect it unset here SELECT content->'extent' FROM collections WHERE id='pgstactest-crudtest'; -- Update item with new datetime that is in a different partition SELECT update_item((SELECT content || '{"properties":{"datetime":"2023-01-01 00:00:00Z"}}'::jsonb FROM test_items LIMIT 1)); SELECT id, geometry, collection, datetime, end_datetime, properties, extra FROM items WHERE collection='pgstactest-crudtest' ORDER BY id; --- Check to see if extent got updated +-- Extent is NOT auto-updated on ingest (now explicit/async) -- expect it unset here SELECT content->'extent' FROM collections WHERE id='pgstactest-crudtest'; -- Update item with new datetime that is in a different partition @@ -43,7 +44,9 @@ aggregated AS (SELECT jsonb_agg(content) as items FROM c) SELECT create_items(items) FROM aggregated; SELECT id, geometry, collection, datetime, end_datetime, properties, extra FROM items WHERE collection='pgstactest-crudtest' ORDER BY id; -DELETE FROM items WHERE collection='pgstactest-crudtest'; +-- clear the collection via delete_item (direct DELETE on items is blocked by the privilege wall) +WITH ids AS MATERIALIZED (SELECT id FROM items WHERE collection='pgstactest-crudtest') +SELECT count(*) AS cleared FROM (SELECT delete_item(id, 'pgstactest-crudtest') FROM ids) d; -- upsert items that do not exist yet WITH c AS (SELECT content FROM test_items LIMIT 2), @@ -57,11 +60,10 @@ aggregated AS (SELECT jsonb_agg(content) as items FROM c) SELECT upsert_items(items) FROM aggregated; SELECT id, geometry, collection, datetime, end_datetime, properties, extra FROM items WHERE collection='pgstactest-crudtest' ORDER BY id; --- turn off update_collection_extent then add an item and verify that the extent did not get updated automatically -SET pgstac.update_collection_extent=FALSE; -SELECT get_setting_bool('update_collection_extent'); -SELECT update_item((SELECT content || '{"properties":{"datetime":"2024-01-01 00:00:00Z"}}'::jsonb FROM test_items LIMIT 1)); - +-- Collection extent is refreshed EXPLICITLY (maintenance), not automatically on ingest. Tighten the +-- partitions to exact stats, then refresh the stored extent, and verify it is now populated. +SELECT tighten_partition_stats(partition) FROM partitions_view WHERE collection='pgstactest-crudtest'; +SELECT update_collection_extents(); SELECT content->'extent' FROM collections WHERE id='pgstactest-crudtest'; -- check formatting of temporal extent diff --git a/src/pgstac/tests/basic/crud_functions.sql.out b/src/pgstac/tests/basic/crud_functions.sql.out index 0e4de1b6..1a556317 100644 --- a/src/pgstac/tests/basic/crud_functions.sql.out +++ b/src/pgstac/tests/basic/crud_functions.sql.out @@ -6,11 +6,9 @@ SET SELECT get_setting_bool('use_queue'); f -SET pgstac.update_collection_extent=TRUE; -SET -SELECT get_setting_bool('update_collection_extent'); - t - +-- NOTE: collection extent is no longer auto-updated on ingest. It is refreshed EXPLICITLY via +-- update_collection_extents() (exercised at the end of this test), so the inline checks below show an +-- unset extent right after ingest. --create base data to use with tests CREATE TEMP TABLE test_items AS SELECT jsonb_build_object( @@ -25,55 +23,58 @@ INSERT INTO collections (content, partition_trunc) VALUES ('{"id":"pgstactest-cr INSERT 0 1 -- Create an item SELECT create_item((SELECT content FROM test_items LIMIT 1)); - + SELECT id, geometry, collection, datetime, end_datetime, properties, extra FROM items WHERE collection='pgstactest-crudtest' ORDER BY id; pgstactest-crudtest-1 | 0103000020E610000001000000050000005B3FFD67CD5355C0C4211B4817EF3E400CE6AF90B95355C0A112D731AE003F4004C93B87325855C0BEBC00FBE8003F40FA0AD28C455855C000E5EFDE51EF3E405B3FFD67CD5355C0C4211B4817EF3E40 | pgstactest-crudtest | 2020-01-01 00:00:00+00 | 2020-01-01 00:00:00+00 | {} | {} --- Check to see if extent got updated +-- Extent is NOT auto-updated on ingest (now explicit/async) -- expect it unset here SELECT content->'extent' FROM collections WHERE id='pgstactest-crudtest'; - {"spatial": {"bbox": [[-85.3792495727539, 30.933948516845703, -85.30819702148438, 31.003555297851562]]}, "temporal": {"interval": [["2020-01-01T00:00:00+00:00", "2020-01-01T00:00:00+00:00"]]}} + -- Update item with new datetime that is in a different partition SELECT update_item((SELECT content || '{"properties":{"datetime":"2023-01-01 00:00:00Z"}}'::jsonb FROM test_items LIMIT 1)); - + SELECT id, geometry, collection, datetime, end_datetime, properties, extra FROM items WHERE collection='pgstactest-crudtest' ORDER BY id; pgstactest-crudtest-1 | 0103000020E610000001000000050000005B3FFD67CD5355C0C4211B4817EF3E400CE6AF90B95355C0A112D731AE003F4004C93B87325855C0BEBC00FBE8003F40FA0AD28C455855C000E5EFDE51EF3E405B3FFD67CD5355C0C4211B4817EF3E40 | pgstactest-crudtest | 2023-01-01 00:00:00+00 | 2023-01-01 00:00:00+00 | {} | {} --- Check to see if extent got updated +-- Extent is NOT auto-updated on ingest (now explicit/async) -- expect it unset here SELECT content->'extent' FROM collections WHERE id='pgstactest-crudtest'; - {"spatial": {"bbox": [[-85.3792495727539, 30.933948516845703, -85.30819702148438, 31.003555297851562]]}, "temporal": {"interval": [["2023-01-01T00:00:00+00:00", "2023-01-01T00:00:00+00:00"]]}} + -- Update item with new datetime that is in a different partition SELECT upsert_item((SELECT content || '{"properties":{"datetime":"2023-02-01 00:00:00Z"}}'::jsonb FROM test_items LIMIT 1)); - + SELECT id, geometry, collection, datetime, end_datetime, properties, extra FROM items WHERE collection='pgstactest-crudtest' ORDER BY id; pgstactest-crudtest-1 | 0103000020E610000001000000050000005B3FFD67CD5355C0C4211B4817EF3E400CE6AF90B95355C0A112D731AE003F4004C93B87325855C0BEBC00FBE8003F40FA0AD28C455855C000E5EFDE51EF3E405B3FFD67CD5355C0C4211B4817EF3E40 | pgstactest-crudtest | 2023-02-01 00:00:00+00 | 2023-02-01 00:00:00+00 | {} | {} -- Delete an item SELECT delete_item('pgstactest-crudtest-1', 'pgstactest-crudtest'); - + SELECT id, geometry, collection, datetime, end_datetime, properties, extra FROM items WHERE collection='pgstactest-crudtest' ORDER BY id; WITH c AS (SELECT content FROM test_items LIMIT 2), aggregated AS (SELECT jsonb_agg(content) as items FROM c) SELECT create_items(items) FROM aggregated; - + SELECT id, geometry, collection, datetime, end_datetime, properties, extra FROM items WHERE collection='pgstactest-crudtest' ORDER BY id; pgstactest-crudtest-1 | 0103000020E610000001000000050000005B3FFD67CD5355C0C4211B4817EF3E400CE6AF90B95355C0A112D731AE003F4004C93B87325855C0BEBC00FBE8003F40FA0AD28C455855C000E5EFDE51EF3E405B3FFD67CD5355C0C4211B4817EF3E40 | pgstactest-crudtest | 2020-01-01 00:00:00+00 | 2020-01-01 00:00:00+00 | {} | {} pgstactest-crudtest-2 | 0103000020E610000001000000050000005B3FFD67CD5355C0C4211B4817EF3E400CE6AF90B95355C0A112D731AE003F4004C93B87325855C0BEBC00FBE8003F40FA0AD28C455855C000E5EFDE51EF3E405B3FFD67CD5355C0C4211B4817EF3E40 | pgstactest-crudtest | 2020-02-01 00:00:00+00 | 2020-02-01 00:00:00+00 | {} | {} -DELETE FROM items WHERE collection='pgstactest-crudtest'; -DELETE 2 +-- clear the collection via delete_item (direct DELETE on items is blocked by the privilege wall) +WITH ids AS MATERIALIZED (SELECT id FROM items WHERE collection='pgstactest-crudtest') +SELECT count(*) AS cleared FROM (SELECT delete_item(id, 'pgstactest-crudtest') FROM ids) d; + 2 + -- upsert items that do not exist yet WITH c AS (SELECT content FROM test_items LIMIT 2), aggregated AS (SELECT jsonb_agg(content) as items FROM c) SELECT upsert_items(items) FROM aggregated; - + SELECT id, geometry, collection, datetime, end_datetime, properties, extra FROM items WHERE collection='pgstactest-crudtest' ORDER BY id; pgstactest-crudtest-1 | 0103000020E610000001000000050000005B3FFD67CD5355C0C4211B4817EF3E400CE6AF90B95355C0A112D731AE003F4004C93B87325855C0BEBC00FBE8003F40FA0AD28C455855C000E5EFDE51EF3E405B3FFD67CD5355C0C4211B4817EF3E40 | pgstactest-crudtest | 2020-01-01 00:00:00+00 | 2020-01-01 00:00:00+00 | {} | {} @@ -83,25 +84,25 @@ SELECT id, geometry, collection, datetime, end_datetime, properties, extra FROM WITH c AS (SELECT content || '{"properties":{"datetime":"2023-02-01 00:00:00Z"}}'::jsonb as content FROM test_items LIMIT 2), aggregated AS (SELECT jsonb_agg(content) as items FROM c) SELECT upsert_items(items) FROM aggregated; - + SELECT id, geometry, collection, datetime, end_datetime, properties, extra FROM items WHERE collection='pgstactest-crudtest' ORDER BY id; pgstactest-crudtest-1 | 0103000020E610000001000000050000005B3FFD67CD5355C0C4211B4817EF3E400CE6AF90B95355C0A112D731AE003F4004C93B87325855C0BEBC00FBE8003F40FA0AD28C455855C000E5EFDE51EF3E405B3FFD67CD5355C0C4211B4817EF3E40 | pgstactest-crudtest | 2023-02-01 00:00:00+00 | 2023-02-01 00:00:00+00 | {} | {} pgstactest-crudtest-2 | 0103000020E610000001000000050000005B3FFD67CD5355C0C4211B4817EF3E400CE6AF90B95355C0A112D731AE003F4004C93B87325855C0BEBC00FBE8003F40FA0AD28C455855C000E5EFDE51EF3E405B3FFD67CD5355C0C4211B4817EF3E40 | pgstactest-crudtest | 2023-02-01 00:00:00+00 | 2023-02-01 00:00:00+00 | {} | {} --- turn off update_collection_extent then add an item and verify that the extent did not get updated automatically -SET pgstac.update_collection_extent=FALSE; -SET -SELECT get_setting_bool('update_collection_extent'); - f +-- Collection extent is refreshed EXPLICITLY (maintenance), not automatically on ingest. Tighten the +-- partitions to exact stats, then refresh the stored extent, and verify it is now populated. +SELECT tighten_partition_stats(partition) FROM partitions_view WHERE collection='pgstactest-crudtest'; + + + +SELECT update_collection_extents(); -SELECT update_item((SELECT content || '{"properties":{"datetime":"2024-01-01 00:00:00Z"}}'::jsonb FROM test_items LIMIT 1)); - SELECT content->'extent' FROM collections WHERE id='pgstactest-crudtest'; - {"spatial": {"bbox": [[-85.3792495727539, 30.933948516845703, -85.30819702148438, 31.003555297851562]]}, "temporal": {"interval": [["2023-02-01T00:00:00+00:00", "2023-02-01T00:00:00+00:00"]]}} + {"spatial": {"bbox": [[-85.379245, 30.933949, -85.308201, 31.003555]]}, "temporal": {"interval": [["2023-02-01T00:00:00+00:00", "2023-02-01T00:00:00+00:00"]]}} -- check formatting of temporal extent SELECT collection_temporal_extent('pgstactest-crudtest'); - [["2023-02-01T00:00:00+00:00", "2024-01-01T00:00:00+00:00"]] + [["2023-02-01T00:00:00+00:00", "2023-02-01T00:00:00+00:00"]] diff --git a/src/pgstac/tests/basic/partitions.sql b/src/pgstac/tests/basic/partitions.sql index 1442c7b8..267ee5a6 100644 --- a/src/pgstac/tests/basic/partitions.sql +++ b/src/pgstac/tests/basic/partitions.sql @@ -17,36 +17,36 @@ SELECT jsonb_build_object( INSERT INTO collections (content) VALUES ('{"id":"pgstactest-partitioned"}'); INSERT INTO items_staging(content) SELECT content FROM test_items; -SELECT count(*) FROM partitions WHERE collection='pgstactest-partitioned'; +SELECT count(*) FROM partitions_view WHERE collection='pgstactest-partitioned'; --test collection partioned by year INSERT INTO collections (content, partition_trunc) VALUES ('{"id":"pgstactest-partitioned-year"}', 'year'); INSERT INTO items_staging(content) SELECT content || '{"collection":"pgstactest-partitioned-year"}'::jsonb FROM test_items; -SELECT count(*) FROM partitions WHERE collection='pgstactest-partitioned-year'; +SELECT count(*) FROM partitions_view WHERE collection='pgstactest-partitioned-year'; --test collection partioned by month INSERT INTO collections (content, partition_trunc) VALUES ('{"id":"pgstactest-partitioned-month"}', 'month'); INSERT INTO items_staging(content) SELECT content || '{"collection":"pgstactest-partitioned-month"}'::jsonb FROM test_items; -SELECT count(*) FROM partitions WHERE collection='pgstactest-partitioned-month'; +SELECT count(*) FROM partitions_view WHERE collection='pgstactest-partitioned-month'; --test repartitioning from year to non partitioned UPDATE collections SET partition_trunc=NULL WHERE id='pgstactest-partitioned-year'; -SELECT count(*) FROM partitions WHERE collection='pgstactest-partitioned-year'; +SELECT count(*) FROM partitions_view WHERE collection='pgstactest-partitioned-year'; SELECT count(*) FROM items WHERE collection='pgstactest-partitioned-year'; --test repartitioning from non-partitioned to year UPDATE collections SET partition_trunc='year' WHERE id='pgstactest-partitioned'; -SELECT count(*) FROM partitions WHERE collection='pgstactest-partitioned'; +SELECT count(*) FROM partitions_view WHERE collection='pgstactest-partitioned'; SELECT count(*) FROM items WHERE collection='pgstactest-partitioned'; --check that partition stats have been updated -SELECT count(*) FROM partitions WHERE collection='pgstactest-partitioned' and spatial IS NULL; +SELECT count(*) FROM partitions_view WHERE collection='pgstactest-partitioned' and spatial IS NULL; --test noop for repartitioning UPDATE collections SET content=content || '{"foo":"bar"}'::jsonb WHERE id='pgstactest-partitioned-month'; -SELECT count(*) FROM partitions WHERE collection='pgstactest-partitioned-month'; +SELECT count(*) FROM partitions_view WHERE collection='pgstactest-partitioned-month'; SELECT count(*) FROM items WHERE collection='pgstactest-partitioned-month'; --test using query queue @@ -56,25 +56,24 @@ SELECT get_setting_bool('use_queue'); INSERT INTO collections (content, partition_trunc) VALUES ('{"id":"pgstactest-partitioned-q"}', 'month'); INSERT INTO items_staging(content) SELECT content || '{"collection":"pgstactest-partitioned-q"}'::jsonb FROM test_items; -SELECT count(*) FROM partitions WHERE collection='pgstactest-partitioned-q'; +SELECT count(*) FROM partitions_view WHERE collection='pgstactest-partitioned-q'; ---check that partition stats haven't been updated -SELECT count(*) FROM partitions WHERE collection='pgstactest-partitioned-q' and spatial IS NULL; - ---check that queue has items +--Under the widen-now / tighten-async model, ingest widens partition stats inline with spatial left NULL +--(always a search candidate) until an async tighten, and queues NO stats work, so the query queue stays +--empty. +SELECT count(*) FROM partitions_view WHERE collection='pgstactest-partitioned-q' and spatial IS NULL; SELECT count(*)>0 FROM query_queue; ---run queue items to update partition stats -SELECT run_queued_queries_intransaction()>0; - ---check that queue has been emptied -SELECT count(*) FROM query_queue; -SELECT run_queued_queries_intransaction(); +--tighten the partitions explicitly (the async maintenance step) to compute exact stats +SELECT count(*)>0 AS tightened FROM ( + SELECT tighten_partition_stats(partition) FROM partitions_view WHERE collection='pgstactest-partitioned-q' +) t; ---check that partition stats have been updated -SELECT count(*) FROM partitions WHERE collection='pgstactest-partitioned-q' and spatial IS NULL; +--after tighten, non-empty partitions have an exact spatial extent (NULL remains only for empty partitions) +SELECT count(*) FROM partitions_view WHERE collection='pgstactest-partitioned-q' and spatial IS NULL; ---check that collection extents have been updated +--collection extents are refreshed explicitly (maintenance), not on ingest +SELECT update_collection_extents(); SELECT id, content->'extent' FROM collections WHERE id LIKE 'pgstactest-partitioned%' ORDER BY id; --check that values for datetimes that are non 4 digit or that have very high precision are ingesting correctly and that partitioning is working for them @@ -107,9 +106,9 @@ INSERT INTO collections (content, partition_trunc) VALUES ('{"id":"pgstactest-pa INSERT INTO items_staging(content) SELECT content || '{"collection":"pgstactest-partitioned-oddballs"}'::jsonb FROM test_items; -SELECT count(*) FROM partitions WHERE collection='pgstactest-partitioned-oddballs'; +SELECT count(*) FROM partitions_view WHERE collection='pgstactest-partitioned-oddballs'; -SELECT collection, partition_dtrange, constraint_dtrange, constraint_edtrange, dtrange, edtrange -FROM partitions +SELECT collection, constraint_dtrange, constraint_edtrange, dtrange, edtrange +FROM partitions_view WHERE collection='pgstactest-partitioned-oddballs' -ORDER BY partition_dtrange; +ORDER BY constraint_dtrange; diff --git a/src/pgstac/tests/basic/partitions.sql.out b/src/pgstac/tests/basic/partitions.sql.out index 580b49c6..c0b6fefd 100644 --- a/src/pgstac/tests/basic/partitions.sql.out +++ b/src/pgstac/tests/basic/partitions.sql.out @@ -26,7 +26,7 @@ INSERT 0 1 INSERT INTO items_staging(content) SELECT content FROM test_items; INSERT 0 105 -SELECT count(*) FROM partitions WHERE collection='pgstactest-partitioned'; +SELECT count(*) FROM partitions_view WHERE collection='pgstactest-partitioned'; 1 --test collection partioned by year @@ -35,7 +35,7 @@ INSERT 0 1 INSERT INTO items_staging(content) SELECT content || '{"collection":"pgstactest-partitioned-year"}'::jsonb FROM test_items; INSERT 0 105 -SELECT count(*) FROM partitions WHERE collection='pgstactest-partitioned-year'; +SELECT count(*) FROM partitions_view WHERE collection='pgstactest-partitioned-year'; 2 --test collection partioned by month @@ -44,13 +44,13 @@ INSERT 0 1 INSERT INTO items_staging(content) SELECT content || '{"collection":"pgstactest-partitioned-month"}'::jsonb FROM test_items; INSERT 0 105 -SELECT count(*) FROM partitions WHERE collection='pgstactest-partitioned-month'; +SELECT count(*) FROM partitions_view WHERE collection='pgstactest-partitioned-month'; 24 --test repartitioning from year to non partitioned UPDATE collections SET partition_trunc=NULL WHERE id='pgstactest-partitioned-year'; UPDATE 1 -SELECT count(*) FROM partitions WHERE collection='pgstactest-partitioned-year'; +SELECT count(*) FROM partitions_view WHERE collection='pgstactest-partitioned-year'; 1 SELECT count(*) FROM items WHERE collection='pgstactest-partitioned-year'; @@ -59,20 +59,20 @@ SELECT count(*) FROM items WHERE collection='pgstactest-partitioned-year'; --test repartitioning from non-partitioned to year UPDATE collections SET partition_trunc='year' WHERE id='pgstactest-partitioned'; UPDATE 1 -SELECT count(*) FROM partitions WHERE collection='pgstactest-partitioned'; +SELECT count(*) FROM partitions_view WHERE collection='pgstactest-partitioned'; 2 SELECT count(*) FROM items WHERE collection='pgstactest-partitioned'; 105 --check that partition stats have been updated -SELECT count(*) FROM partitions WHERE collection='pgstactest-partitioned' and spatial IS NULL; - 0 +SELECT count(*) FROM partitions_view WHERE collection='pgstactest-partitioned' and spatial IS NULL; + 2 --test noop for repartitioning UPDATE collections SET content=content || '{"foo":"bar"}'::jsonb WHERE id='pgstactest-partitioned-month'; UPDATE 1 -SELECT count(*) FROM partitions WHERE collection='pgstactest-partitioned-month'; +SELECT count(*) FROM partitions_view WHERE collection='pgstactest-partitioned-month'; 24 SELECT count(*) FROM items WHERE collection='pgstactest-partitioned-month'; @@ -89,38 +89,37 @@ INSERT 0 1 INSERT INTO items_staging(content) SELECT content || '{"collection":"pgstactest-partitioned-q"}'::jsonb FROM test_items; INSERT 0 105 -SELECT count(*) FROM partitions WHERE collection='pgstactest-partitioned-q'; +SELECT count(*) FROM partitions_view WHERE collection='pgstactest-partitioned-q'; 24 ---check that partition stats haven't been updated -SELECT count(*) FROM partitions WHERE collection='pgstactest-partitioned-q' and spatial IS NULL; +--Under the widen-now / tighten-async model, ingest widens partition stats inline with spatial left NULL +--(always a search candidate) until an async tighten, and queues NO stats work, so the query queue stays +--empty. +SELECT count(*) FROM partitions_view WHERE collection='pgstactest-partitioned-q' and spatial IS NULL; 24 ---check that queue has items SELECT count(*)>0 FROM query_queue; - t + f ---run queue items to update partition stats -SELECT run_queued_queries_intransaction()>0; +--tighten the partitions explicitly (the async maintenance step) to compute exact stats +SELECT count(*)>0 AS tightened FROM ( + SELECT tighten_partition_stats(partition) FROM partitions_view WHERE collection='pgstactest-partitioned-q' +) t; t ---check that queue has been emptied -SELECT count(*) FROM query_queue; +--after tighten, non-empty partitions have an exact spatial extent (NULL remains only for empty partitions) +SELECT count(*) FROM partitions_view WHERE collection='pgstactest-partitioned-q' and spatial IS NULL; 0 -SELECT run_queued_queries_intransaction(); - 0 +--collection extents are refreshed explicitly (maintenance), not on ingest +SELECT update_collection_extents(); ---check that partition stats have been updated -SELECT count(*) FROM partitions WHERE collection='pgstactest-partitioned-q' and spatial IS NULL; - 0 ---check that collection extents have been updated SELECT id, content->'extent' FROM collections WHERE id LIKE 'pgstactest-partitioned%' ORDER BY id; - pgstactest-partitioned | {"spatial": {"bbox": [[-85.3792495727539, 30.933948516845703, -85.30819702148438, 31.003555297851562]]}, "temporal": {"interval": [["2020-01-01T00:00:00+00:00", "2021-12-29T00:00:00+00:00"]]}} - pgstactest-partitioned-month | {"spatial": {"bbox": [[-85.3792495727539, 30.933948516845703, -85.30819702148438, 31.003555297851562]]}, "temporal": {"interval": [["2020-01-01T00:00:00+00:00", "2021-12-29T00:00:00+00:00"]]}} - pgstactest-partitioned-q | {"spatial": {"bbox": [[-85.3792495727539, 30.933948516845703, -85.30819702148438, 31.003555297851562]]}, "temporal": {"interval": [["2020-01-01T00:00:00+00:00", "2021-12-29T00:00:00+00:00"]]}} - pgstactest-partitioned-year | {"spatial": {"bbox": [[-85.3792495727539, 30.933948516845703, -85.30819702148438, 31.003555297851562]]}, "temporal": {"interval": [["2020-01-01T00:00:00+00:00", "2021-12-29T00:00:00+00:00"]]}} + pgstactest-partitioned | null + pgstactest-partitioned-month | null + pgstactest-partitioned-q | {"spatial": {"bbox": [[-85.379245, 30.933949, -85.308201, 31.003555]]}, "temporal": {"interval": [["2020-01-01T00:00:00+00:00", "2021-12-29T00:00:00+00:00"]]}} + pgstactest-partitioned-year | null --check that values for datetimes that are non 4 digit or that have very high precision are ingesting correctly and that partitioning is working for them SET pgstac.use_queue=FALSE; @@ -157,36 +156,37 @@ INSERT 0 1 INSERT INTO items_staging(content) SELECT content || '{"collection":"pgstactest-partitioned-oddballs"}'::jsonb FROM test_items; INSERT 0 108 -SELECT count(*) FROM partitions WHERE collection='pgstactest-partitioned-oddballs'; +SELECT count(*) FROM partitions_view WHERE collection='pgstactest-partitioned-oddballs'; 26 -SELECT collection, partition_dtrange, constraint_dtrange, constraint_edtrange, dtrange, edtrange -FROM partitions +SELECT collection, constraint_dtrange, constraint_edtrange, dtrange, edtrange +FROM partitions_view WHERE collection='pgstactest-partitioned-oddballs' -ORDER BY partition_dtrange; - pgstactest-partitioned-oddballs | ["2000-01-01 00:00:00+00","2000-02-01 00:00:00+00") | ["2000-01-01 00:00:00+00","2000-02-01 00:00:00+00") | [-infinity,infinity] | ["2000-01-01 00:00:00.123899+00","2000-01-01 00:00:00.123899+00"] | ["2000-01-01 00:00:00.123899+00","99999-01-01 00:00:00+00"] - pgstactest-partitioned-oddballs | ["2020-01-01 00:00:00+00","2020-02-01 00:00:00+00") | ["2020-01-01 00:00:00+00","2020-02-01 00:00:00+00") | [-infinity,infinity] | ["2020-01-01 00:00:00+00","2020-01-29 00:00:00+00"] | ["2020-01-01 00:00:00+00","2020-01-29 00:00:00+00"] - pgstactest-partitioned-oddballs | ["2020-02-01 00:00:00+00","2020-03-01 00:00:00+00") | ["2020-02-01 00:00:00+00","2020-03-01 00:00:00+00") | [-infinity,infinity] | ["2020-02-05 00:00:00+00","2020-02-26 00:00:00+00"] | ["2020-02-05 00:00:00+00","2020-02-26 00:00:00+00"] - pgstactest-partitioned-oddballs | ["2020-03-01 00:00:00+00","2020-04-01 00:00:00+00") | ["2020-03-01 00:00:00+00","2020-04-01 00:00:00+00") | [-infinity,infinity] | ["2020-03-04 00:00:00+00","2020-03-25 00:00:00+00"] | ["2020-03-04 00:00:00+00","2020-03-25 00:00:00+00"] - pgstactest-partitioned-oddballs | ["2020-04-01 00:00:00+00","2020-05-01 00:00:00+00") | ["2020-04-01 00:00:00+00","2020-05-01 00:00:00+00") | [-infinity,infinity] | ["2020-04-01 00:00:00+00","2020-04-29 00:00:00+00"] | ["2020-04-01 00:00:00+00","2020-04-29 00:00:00+00"] - pgstactest-partitioned-oddballs | ["2020-05-01 00:00:00+00","2020-06-01 00:00:00+00") | ["2020-05-01 00:00:00+00","2020-06-01 00:00:00+00") | [-infinity,infinity] | ["2020-05-06 00:00:00+00","2020-05-27 00:00:00+00"] | ["2020-05-06 00:00:00+00","2020-05-27 00:00:00+00"] - pgstactest-partitioned-oddballs | ["2020-06-01 00:00:00+00","2020-07-01 00:00:00+00") | ["2020-06-01 00:00:00+00","2020-07-01 00:00:00+00") | [-infinity,infinity] | ["2020-06-03 00:00:00+00","2020-06-24 00:00:00+00"] | ["2020-06-03 00:00:00+00","2020-06-24 00:00:00+00"] - pgstactest-partitioned-oddballs | ["2020-07-01 00:00:00+00","2020-08-01 00:00:00+00") | ["2020-07-01 00:00:00+00","2020-08-01 00:00:00+00") | [-infinity,infinity] | ["2020-07-01 00:00:00+00","2020-07-29 00:00:00+00"] | ["2020-07-01 00:00:00+00","2020-07-29 00:00:00+00"] - pgstactest-partitioned-oddballs | ["2020-08-01 00:00:00+00","2020-09-01 00:00:00+00") | ["2020-08-01 00:00:00+00","2020-09-01 00:00:00+00") | [-infinity,infinity] | ["2020-08-05 00:00:00+00","2020-08-26 00:00:00+00"] | ["2020-08-05 00:00:00+00","2020-08-26 00:00:00+00"] - pgstactest-partitioned-oddballs | ["2020-09-01 00:00:00+00","2020-10-01 00:00:00+00") | ["2020-09-01 00:00:00+00","2020-10-01 00:00:00+00") | [-infinity,infinity] | ["2020-09-02 00:00:00+00","2020-09-30 00:00:00+00"] | ["2020-09-02 00:00:00+00","2020-09-30 00:00:00+00"] - pgstactest-partitioned-oddballs | ["2020-10-01 00:00:00+00","2020-11-01 00:00:00+00") | ["2020-10-01 00:00:00+00","2020-11-01 00:00:00+00") | [-infinity,infinity] | ["2020-10-07 00:00:00+00","2020-10-28 00:00:00+00"] | ["2020-10-07 00:00:00+00","2020-10-28 00:00:00+00"] - pgstactest-partitioned-oddballs | ["2020-11-01 00:00:00+00","2020-12-01 00:00:00+00") | ["2020-11-01 00:00:00+00","2020-12-01 00:00:00+00") | [-infinity,infinity] | ["2020-11-04 00:00:00+00","2020-11-25 00:00:00+00"] | ["2020-11-04 00:00:00+00","2020-11-25 00:00:00+00"] - pgstactest-partitioned-oddballs | ["2020-12-01 00:00:00+00","2021-01-01 00:00:00+00") | ["2020-12-01 00:00:00+00","2021-01-01 00:00:00+00") | [-infinity,infinity] | ["2020-12-02 00:00:00+00","2020-12-30 00:00:00+00"] | ["2020-12-02 00:00:00+00","2020-12-30 00:00:00+00"] - pgstactest-partitioned-oddballs | ["2021-01-01 00:00:00+00","2021-02-01 00:00:00+00") | ["2021-01-01 00:00:00+00","2021-02-01 00:00:00+00") | [-infinity,infinity] | ["2021-01-06 00:00:00+00","2021-01-27 00:00:00+00"] | ["2021-01-06 00:00:00+00","2021-01-27 00:00:00+00"] - pgstactest-partitioned-oddballs | ["2021-02-01 00:00:00+00","2021-03-01 00:00:00+00") | ["2021-02-01 00:00:00+00","2021-03-01 00:00:00+00") | [-infinity,infinity] | ["2021-02-03 00:00:00+00","2021-02-24 00:00:00+00"] | ["2021-02-03 00:00:00+00","2021-02-24 00:00:00+00"] - pgstactest-partitioned-oddballs | ["2021-03-01 00:00:00+00","2021-04-01 00:00:00+00") | ["2021-03-01 00:00:00+00","2021-04-01 00:00:00+00") | [-infinity,infinity] | ["2021-03-03 00:00:00+00","2021-03-31 00:00:00+00"] | ["2021-03-03 00:00:00+00","2021-03-31 00:00:00+00"] - pgstactest-partitioned-oddballs | ["2021-04-01 00:00:00+00","2021-05-01 00:00:00+00") | ["2021-04-01 00:00:00+00","2021-05-01 00:00:00+00") | [-infinity,infinity] | ["2021-04-07 00:00:00+00","2021-04-28 00:00:00+00"] | ["2021-04-07 00:00:00+00","2021-04-28 00:00:00+00"] - pgstactest-partitioned-oddballs | ["2021-05-01 00:00:00+00","2021-06-01 00:00:00+00") | ["2021-05-01 00:00:00+00","2021-06-01 00:00:00+00") | [-infinity,infinity] | ["2021-05-05 00:00:00+00","2021-05-26 00:00:00+00"] | ["2021-05-05 00:00:00+00","2021-05-26 00:00:00+00"] - pgstactest-partitioned-oddballs | ["2021-06-01 00:00:00+00","2021-07-01 00:00:00+00") | ["2021-06-01 00:00:00+00","2021-07-01 00:00:00+00") | [-infinity,infinity] | ["2021-06-02 00:00:00+00","2021-06-30 00:00:00+00"] | ["2021-06-02 00:00:00+00","2021-06-30 00:00:00+00"] - pgstactest-partitioned-oddballs | ["2021-07-01 00:00:00+00","2021-08-01 00:00:00+00") | ["2021-07-01 00:00:00+00","2021-08-01 00:00:00+00") | [-infinity,infinity] | ["2021-07-07 00:00:00+00","2021-07-28 00:00:00+00"] | ["2021-07-07 00:00:00+00","2021-07-28 00:00:00+00"] - pgstactest-partitioned-oddballs | ["2021-08-01 00:00:00+00","2021-09-01 00:00:00+00") | ["2021-08-01 00:00:00+00","2021-09-01 00:00:00+00") | [-infinity,infinity] | ["2021-08-04 00:00:00+00","2021-08-25 00:00:00+00"] | ["2021-08-04 00:00:00+00","2021-08-25 00:00:00+00"] - pgstactest-partitioned-oddballs | ["2021-09-01 00:00:00+00","2021-10-01 00:00:00+00") | ["2021-09-01 00:00:00+00","2021-10-01 00:00:00+00") | [-infinity,infinity] | ["2021-09-01 00:00:00+00","2021-09-29 00:00:00+00"] | ["2021-09-01 00:00:00+00","2021-09-29 00:00:00+00"] - pgstactest-partitioned-oddballs | ["2021-10-01 00:00:00+00","2021-11-01 00:00:00+00") | ["2021-10-01 00:00:00+00","2021-11-01 00:00:00+00") | [-infinity,infinity] | ["2021-10-06 00:00:00+00","2021-10-27 00:00:00+00"] | ["2021-10-06 00:00:00+00","2021-10-27 00:00:00+00"] - pgstactest-partitioned-oddballs | ["2021-11-01 00:00:00+00","2021-12-01 00:00:00+00") | ["2021-11-01 00:00:00+00","2021-12-01 00:00:00+00") | [-infinity,infinity] | ["2021-11-03 00:00:00+00","2021-11-24 00:00:00+00"] | ["2021-11-03 00:00:00+00","2021-11-24 00:00:00+00"] - pgstactest-partitioned-oddballs | ["2021-12-01 00:00:00+00","2022-01-01 00:00:00+00") | ["2021-12-01 00:00:00+00","2022-01-01 00:00:00+00") | [-infinity,infinity] | ["2021-12-01 00:00:00+00","2021-12-29 00:00:00+00"] | ["2021-12-01 00:00:00+00","2021-12-29 00:00:00+00"] - pgstactest-partitioned-oddballs | ["10000-01-01 00:00:00+00","10000-02-01 00:00:00+00") | ["10000-01-01 00:00:00+00","10000-02-01 00:00:00+00") | [-infinity,infinity] | ["10000-01-01 00:00:00+00","10000-01-01 00:00:00+00"] | ["10000-01-01 00:00:00+00","10000-01-01 00:00:00+00"] +ORDER BY constraint_dtrange; + pgstactest-partitioned-oddballs | ["2000-01-01 00:00:00+00","2000-02-01 00:00:00+00") | [-infinity,infinity] | ["2000-01-01 00:00:00+00","2000-02-01 00:00:00+00") | ["2000-01-01 00:00:00+00","99999-01-31 23:59:59.876101+00"] + pgstactest-partitioned-oddballs | ["2020-01-01 00:00:00+00","2020-02-01 00:00:00+00") | [-infinity,infinity] | ["2020-01-01 00:00:00+00","2020-02-01 00:00:00+00") | ["2020-01-01 00:00:00+00","2020-02-01 00:00:00+00"] + pgstactest-partitioned-oddballs | ["2020-02-01 00:00:00+00","2020-03-01 00:00:00+00") | [-infinity,infinity] | ["2020-02-01 00:00:00+00","2020-03-01 00:00:00+00") | ["2020-02-01 00:00:00+00","2020-03-01 00:00:00+00"] + pgstactest-partitioned-oddballs | ["2020-03-01 00:00:00+00","2020-04-01 00:00:00+00") | [-infinity,infinity] | ["2020-03-01 00:00:00+00","2020-04-01 00:00:00+00") | ["2020-03-01 00:00:00+00","2020-04-01 00:00:00+00"] + pgstactest-partitioned-oddballs | ["2020-04-01 00:00:00+00","2020-05-01 00:00:00+00") | [-infinity,infinity] | ["2020-04-01 00:00:00+00","2020-05-01 00:00:00+00") | ["2020-04-01 00:00:00+00","2020-05-01 00:00:00+00"] + pgstactest-partitioned-oddballs | ["2020-05-01 00:00:00+00","2020-06-01 00:00:00+00") | [-infinity,infinity] | ["2020-05-01 00:00:00+00","2020-06-01 00:00:00+00") | ["2020-05-01 00:00:00+00","2020-06-01 00:00:00+00"] + pgstactest-partitioned-oddballs | ["2020-06-01 00:00:00+00","2020-07-01 00:00:00+00") | [-infinity,infinity] | ["2020-06-01 00:00:00+00","2020-07-01 00:00:00+00") | ["2020-06-01 00:00:00+00","2020-07-01 00:00:00+00"] + pgstactest-partitioned-oddballs | ["2020-07-01 00:00:00+00","2020-08-01 00:00:00+00") | [-infinity,infinity] | ["2020-07-01 00:00:00+00","2020-08-01 00:00:00+00") | ["2020-07-01 00:00:00+00","2020-08-01 00:00:00+00"] + pgstactest-partitioned-oddballs | ["2020-08-01 00:00:00+00","2020-09-01 00:00:00+00") | [-infinity,infinity] | ["2020-08-01 00:00:00+00","2020-09-01 00:00:00+00") | ["2020-08-01 00:00:00+00","2020-09-01 00:00:00+00"] + pgstactest-partitioned-oddballs | ["2020-09-01 00:00:00+00","2020-10-01 00:00:00+00") | [-infinity,infinity] | ["2020-09-01 00:00:00+00","2020-10-01 00:00:00+00") | ["2020-09-01 00:00:00+00","2020-10-01 00:00:00+00"] + pgstactest-partitioned-oddballs | ["2020-10-01 00:00:00+00","2020-11-01 00:00:00+00") | [-infinity,infinity] | ["2020-10-01 00:00:00+00","2020-11-01 00:00:00+00") | ["2020-10-01 00:00:00+00","2020-11-01 00:00:00+00"] + pgstactest-partitioned-oddballs | ["2020-11-01 00:00:00+00","2020-12-01 00:00:00+00") | [-infinity,infinity] | ["2020-11-01 00:00:00+00","2020-12-01 00:00:00+00") | ["2020-11-01 00:00:00+00","2020-12-01 00:00:00+00"] + pgstactest-partitioned-oddballs | ["2020-12-01 00:00:00+00","2021-01-01 00:00:00+00") | [-infinity,infinity] | ["2020-12-01 00:00:00+00","2021-01-01 00:00:00+00") | ["2020-12-01 00:00:00+00","2021-01-01 00:00:00+00"] + pgstactest-partitioned-oddballs | ["2021-01-01 00:00:00+00","2021-02-01 00:00:00+00") | [-infinity,infinity] | ["2021-01-01 00:00:00+00","2021-02-01 00:00:00+00") | ["2021-01-01 00:00:00+00","2021-02-01 00:00:00+00"] + pgstactest-partitioned-oddballs | ["2021-02-01 00:00:00+00","2021-03-01 00:00:00+00") | [-infinity,infinity] | ["2021-02-01 00:00:00+00","2021-03-01 00:00:00+00") | ["2021-02-01 00:00:00+00","2021-03-01 00:00:00+00"] + pgstactest-partitioned-oddballs | ["2021-03-01 00:00:00+00","2021-04-01 00:00:00+00") | [-infinity,infinity] | ["2021-03-01 00:00:00+00","2021-04-01 00:00:00+00") | ["2021-03-01 00:00:00+00","2021-04-01 00:00:00+00"] + pgstactest-partitioned-oddballs | ["2021-04-01 00:00:00+00","2021-05-01 00:00:00+00") | [-infinity,infinity] | ["2021-04-01 00:00:00+00","2021-05-01 00:00:00+00") | ["2021-04-01 00:00:00+00","2021-05-01 00:00:00+00"] + pgstactest-partitioned-oddballs | ["2021-05-01 00:00:00+00","2021-06-01 00:00:00+00") | [-infinity,infinity] | ["2021-05-01 00:00:00+00","2021-06-01 00:00:00+00") | ["2021-05-01 00:00:00+00","2021-06-01 00:00:00+00"] + pgstactest-partitioned-oddballs | ["2021-06-01 00:00:00+00","2021-07-01 00:00:00+00") | [-infinity,infinity] | ["2021-06-01 00:00:00+00","2021-07-01 00:00:00+00") | ["2021-06-01 00:00:00+00","2021-07-01 00:00:00+00"] + pgstactest-partitioned-oddballs | ["2021-07-01 00:00:00+00","2021-08-01 00:00:00+00") | [-infinity,infinity] | ["2021-07-01 00:00:00+00","2021-08-01 00:00:00+00") | ["2021-07-01 00:00:00+00","2021-08-01 00:00:00+00"] + pgstactest-partitioned-oddballs | ["2021-08-01 00:00:00+00","2021-09-01 00:00:00+00") | [-infinity,infinity] | ["2021-08-01 00:00:00+00","2021-09-01 00:00:00+00") | ["2021-08-01 00:00:00+00","2021-09-01 00:00:00+00"] + pgstactest-partitioned-oddballs | ["2021-09-01 00:00:00+00","2021-10-01 00:00:00+00") | [-infinity,infinity] | ["2021-09-01 00:00:00+00","2021-10-01 00:00:00+00") | ["2021-09-01 00:00:00+00","2021-10-01 00:00:00+00"] + pgstactest-partitioned-oddballs | ["2021-10-01 00:00:00+00","2021-11-01 00:00:00+00") | [-infinity,infinity] | ["2021-10-01 00:00:00+00","2021-11-01 00:00:00+00") | ["2021-10-01 00:00:00+00","2021-11-01 00:00:00+00"] + pgstactest-partitioned-oddballs | ["2021-11-01 00:00:00+00","2021-12-01 00:00:00+00") | [-infinity,infinity] | ["2021-11-01 00:00:00+00","2021-12-01 00:00:00+00") | ["2021-11-01 00:00:00+00","2021-12-01 00:00:00+00"] + pgstactest-partitioned-oddballs | ["2021-12-01 00:00:00+00","2022-01-01 00:00:00+00") | [-infinity,infinity] | ["2021-12-01 00:00:00+00","2022-01-01 00:00:00+00") | ["2021-12-01 00:00:00+00","2022-01-01 00:00:00+00"] + pgstactest-partitioned-oddballs | ["10000-01-01 00:00:00+00","10000-02-01 00:00:00+00") | [-infinity,infinity] | ["10000-01-01 00:00:00+00","10000-02-01 00:00:00+00") | ["10000-01-01 00:00:00+00","10000-02-01 00:00:00+00"] + diff --git a/src/pgstac/tests/basic/search_path.sql b/src/pgstac/tests/basic/search_path.sql index 74259a06..60087a37 100644 --- a/src/pgstac/tests/basic/search_path.sql +++ b/src/pgstac/tests/basic/search_path.sql @@ -21,20 +21,16 @@ CREATE TEMP TABLE sp_partition_stats AS SELECT * FROM partition_stats WHERE partition IN (SELECT partition FROM sp_partition_sys_meta) ORDER BY partition; CREATE TEMP TABLE sp_partitions AS -SELECT * FROM partitions WHERE collection='pgstactest-searchpath' ORDER BY partition; +SELECT * FROM partitions_view WHERE collection='pgstactest-searchpath' ORDER BY partition; CREATE TEMP TABLE sp_partitions_view AS SELECT * FROM partitions_view WHERE collection='pgstactest-searchpath' ORDER BY partition; -CREATE TEMP TABLE sp_partition_steps AS -SELECT * FROM partition_steps WHERE name IN (SELECT partition FROM sp_partition_sys_meta) ORDER BY name; - -- Verify we have data SELECT count(*) > 0 AS has_partition_sys_meta FROM sp_partition_sys_meta; SELECT count(*) > 0 AS has_partition_stats FROM sp_partition_stats; SELECT count(*) > 0 AS has_partitions FROM sp_partitions; SELECT count(*) > 0 AS has_partitions_view FROM sp_partitions_view; -SELECT count(*) > 0 AS has_partition_steps FROM sp_partition_steps; -- Now remove pgstac from search_path SET search_path TO public; @@ -73,22 +69,8 @@ SELECT count(*) = 0 AS partition_stats_data_match FROM ( SELECT partition, dtrange, edtrange FROM pgstac.partition_stats WHERE partition IN (SELECT partition FROM sp_partition_sys_meta)) ) diff; --- partitions (materialized view): compare counts and key columns -SELECT ( - SELECT count(*) FROM pgstac.partitions WHERE collection='pgstactest-searchpath' -) = ( - SELECT count(*) FROM sp_partitions -) AS partitions_count_match; - -SELECT count(*) = 0 AS partitions_data_match FROM ( - (SELECT partition, collection, partition_dtrange, constraint_dtrange, constraint_edtrange, dtrange, edtrange FROM pgstac.partitions WHERE collection='pgstactest-searchpath' - EXCEPT - SELECT partition, collection, partition_dtrange, constraint_dtrange, constraint_edtrange, dtrange, edtrange FROM sp_partitions) - UNION ALL - (SELECT partition, collection, partition_dtrange, constraint_dtrange, constraint_edtrange, dtrange, edtrange FROM sp_partitions - EXCEPT - SELECT partition, collection, partition_dtrange, constraint_dtrange, constraint_edtrange, dtrange, edtrange FROM pgstac.partitions WHERE collection='pgstactest-searchpath') -) diff; +-- (the `partitions` materialized view was dropped under the widen-now / tighten-async model; the live +-- partitions_view below is the equivalent and is the relation search-path behavior is verified against.) -- partitions_view: compare counts and key columns SELECT ( @@ -98,30 +80,13 @@ SELECT ( ) AS partitions_view_count_match; SELECT count(*) = 0 AS partitions_view_data_match FROM ( - (SELECT partition, collection, partition_dtrange, constraint_dtrange, constraint_edtrange, dtrange, edtrange FROM pgstac.partitions_view WHERE collection='pgstactest-searchpath' - EXCEPT - SELECT partition, collection, partition_dtrange, constraint_dtrange, constraint_edtrange, dtrange, edtrange FROM sp_partitions_view) - UNION ALL - (SELECT partition, collection, partition_dtrange, constraint_dtrange, constraint_edtrange, dtrange, edtrange FROM sp_partitions_view - EXCEPT - SELECT partition, collection, partition_dtrange, constraint_dtrange, constraint_edtrange, dtrange, edtrange FROM pgstac.partitions_view WHERE collection='pgstactest-searchpath') -) diff; - --- partition_steps: compare counts and key columns -SELECT ( - SELECT count(*) FROM pgstac.partition_steps WHERE name IN (SELECT partition FROM sp_partition_sys_meta) -) = ( - SELECT count(*) FROM sp_partition_steps -) AS partition_steps_count_match; - -SELECT count(*) = 0 AS partition_steps_data_match FROM ( - (SELECT name, sdate, edate FROM pgstac.partition_steps WHERE name IN (SELECT partition FROM sp_partition_sys_meta) + (SELECT partition, collection, constraint_dtrange, constraint_edtrange, dtrange, edtrange FROM pgstac.partitions_view WHERE collection='pgstactest-searchpath' EXCEPT - SELECT name, sdate, edate FROM sp_partition_steps) + SELECT partition, collection, constraint_dtrange, constraint_edtrange, dtrange, edtrange FROM sp_partitions_view) UNION ALL - (SELECT name, sdate, edate FROM sp_partition_steps + (SELECT partition, collection, constraint_dtrange, constraint_edtrange, dtrange, edtrange FROM sp_partitions_view EXCEPT - SELECT name, sdate, edate FROM pgstac.partition_steps WHERE name IN (SELECT partition FROM sp_partition_sys_meta)) + SELECT partition, collection, constraint_dtrange, constraint_edtrange, dtrange, edtrange FROM pgstac.partitions_view WHERE collection='pgstactest-searchpath') ) diff; -- Restore search_path diff --git a/src/pgstac/tests/basic/search_path.sql.out b/src/pgstac/tests/basic/search_path.sql.out index 30718c0d..c5597fec 100644 --- a/src/pgstac/tests/basic/search_path.sql.out +++ b/src/pgstac/tests/basic/search_path.sql.out @@ -24,14 +24,11 @@ CREATE TEMP TABLE sp_partition_stats AS SELECT * FROM partition_stats WHERE partition IN (SELECT partition FROM sp_partition_sys_meta) ORDER BY partition; SELECT 5 CREATE TEMP TABLE sp_partitions AS -SELECT * FROM partitions WHERE collection='pgstactest-searchpath' ORDER BY partition; +SELECT * FROM partitions_view WHERE collection='pgstactest-searchpath' ORDER BY partition; SELECT 5 CREATE TEMP TABLE sp_partitions_view AS SELECT * FROM partitions_view WHERE collection='pgstactest-searchpath' ORDER BY partition; SELECT 5 -CREATE TEMP TABLE sp_partition_steps AS -SELECT * FROM partition_steps WHERE name IN (SELECT partition FROM sp_partition_sys_meta) ORDER BY name; -SELECT 5 -- Verify we have data SELECT count(*) > 0 AS has_partition_sys_meta FROM sp_partition_sys_meta; t @@ -45,9 +42,6 @@ SELECT count(*) > 0 AS has_partitions FROM sp_partitions; SELECT count(*) > 0 AS has_partitions_view FROM sp_partitions_view; t -SELECT count(*) > 0 AS has_partition_steps FROM sp_partition_steps; - t - -- Now remove pgstac from search_path SET search_path TO public; SET @@ -89,25 +83,8 @@ SELECT count(*) = 0 AS partition_stats_data_match FROM ( ) diff; t --- partitions (materialized view): compare counts and key columns -SELECT ( - SELECT count(*) FROM pgstac.partitions WHERE collection='pgstactest-searchpath' -) = ( - SELECT count(*) FROM sp_partitions -) AS partitions_count_match; - t - -SELECT count(*) = 0 AS partitions_data_match FROM ( - (SELECT partition, collection, partition_dtrange, constraint_dtrange, constraint_edtrange, dtrange, edtrange FROM pgstac.partitions WHERE collection='pgstactest-searchpath' - EXCEPT - SELECT partition, collection, partition_dtrange, constraint_dtrange, constraint_edtrange, dtrange, edtrange FROM sp_partitions) - UNION ALL - (SELECT partition, collection, partition_dtrange, constraint_dtrange, constraint_edtrange, dtrange, edtrange FROM sp_partitions - EXCEPT - SELECT partition, collection, partition_dtrange, constraint_dtrange, constraint_edtrange, dtrange, edtrange FROM pgstac.partitions WHERE collection='pgstactest-searchpath') -) diff; - t - +-- (the `partitions` materialized view was dropped under the widen-now / tighten-async model; the live +-- partitions_view below is the equivalent and is the relation search-path behavior is verified against.) -- partitions_view: compare counts and key columns SELECT ( SELECT count(*) FROM pgstac.partitions_view WHERE collection='pgstactest-searchpath' @@ -117,32 +94,13 @@ SELECT ( t SELECT count(*) = 0 AS partitions_view_data_match FROM ( - (SELECT partition, collection, partition_dtrange, constraint_dtrange, constraint_edtrange, dtrange, edtrange FROM pgstac.partitions_view WHERE collection='pgstactest-searchpath' - EXCEPT - SELECT partition, collection, partition_dtrange, constraint_dtrange, constraint_edtrange, dtrange, edtrange FROM sp_partitions_view) - UNION ALL - (SELECT partition, collection, partition_dtrange, constraint_dtrange, constraint_edtrange, dtrange, edtrange FROM sp_partitions_view - EXCEPT - SELECT partition, collection, partition_dtrange, constraint_dtrange, constraint_edtrange, dtrange, edtrange FROM pgstac.partitions_view WHERE collection='pgstactest-searchpath') -) diff; - t - --- partition_steps: compare counts and key columns -SELECT ( - SELECT count(*) FROM pgstac.partition_steps WHERE name IN (SELECT partition FROM sp_partition_sys_meta) -) = ( - SELECT count(*) FROM sp_partition_steps -) AS partition_steps_count_match; - t - -SELECT count(*) = 0 AS partition_steps_data_match FROM ( - (SELECT name, sdate, edate FROM pgstac.partition_steps WHERE name IN (SELECT partition FROM sp_partition_sys_meta) + (SELECT partition, collection, constraint_dtrange, constraint_edtrange, dtrange, edtrange FROM pgstac.partitions_view WHERE collection='pgstactest-searchpath' EXCEPT - SELECT name, sdate, edate FROM sp_partition_steps) + SELECT partition, collection, constraint_dtrange, constraint_edtrange, dtrange, edtrange FROM sp_partitions_view) UNION ALL - (SELECT name, sdate, edate FROM sp_partition_steps + (SELECT partition, collection, constraint_dtrange, constraint_edtrange, dtrange, edtrange FROM sp_partitions_view EXCEPT - SELECT name, sdate, edate FROM pgstac.partition_steps WHERE name IN (SELECT partition FROM sp_partition_sys_meta)) + SELECT partition, collection, constraint_dtrange, constraint_edtrange, dtrange, edtrange FROM pgstac.partitions_view WHERE collection='pgstactest-searchpath') ) diff; t diff --git a/src/pgstac/tests/pgtap.sql b/src/pgstac/tests/pgtap.sql index 78e1d8d4..4d75eb89 100644 --- a/src/pgstac/tests/pgtap.sql +++ b/src/pgstac/tests/pgtap.sql @@ -17,12 +17,13 @@ CREATE EXTENSION IF NOT EXISTS pgtap; SET SEARCH_PATH TO pgstac, pgtap, public; -- Plan the tests. -SELECT plan(357); +SELECT plan(371); --SELECT * FROM no_plan(); -- Run the tests. -- Core +\i tests/pgtap/000_privileges.sql \i tests/pgtap/001_core.sql \i tests/pgtap/001a_jsonutils.sql \i tests/pgtap/001b_cursorutils.sql diff --git a/src/pgstac/tests/pgtap/000_privileges.sql b/src/pgstac/tests/pgtap/000_privileges.sql new file mode 100644 index 00000000..3441579a --- /dev/null +++ b/src/pgstac/tests/pgtap/000_privileges.sql @@ -0,0 +1,22 @@ +-- Privilege wall (INV-1): pgstac_ingest may READ the wall-protected tables but may MUTATE them only through +-- the SECURITY DEFINER write functions — no inherited grant lets it write directly. + +SELECT ok( + has_table_privilege('pgstac_ingest', 'pgstac.items', 'SELECT'), + 'pgstac_ingest may SELECT items' +); + +SELECT ok( + NOT has_table_privilege('pgstac_ingest', 'pgstac.items', 'INSERT'), + 'pgstac_ingest may not directly INSERT items (privilege wall)' +); + +SELECT ok( + NOT has_table_privilege('pgstac_ingest', 'pgstac.partition_stats', 'UPDATE'), + 'pgstac_ingest may not directly UPDATE partition_stats (privilege wall)' +); + +SELECT ok( + NOT has_table_privilege('pgstac_ingest', 'pgstac.item_fragments', 'DELETE'), + 'pgstac_ingest may not directly DELETE item_fragments (privilege wall)' +); diff --git a/src/pgstac/tests/pgtap/001_core.sql b/src/pgstac/tests/pgtap/001_core.sql index 37ea1ebf..4eb4b32d 100644 --- a/src/pgstac/tests/pgtap/001_core.sql +++ b/src/pgstac/tests/pgtap/001_core.sql @@ -40,10 +40,16 @@ SELECT lives_ok( ); RESET pgstac.context; -SELECT is_definer('update_partition_stats'); -SELECT is_definer('partition_after_triggerfunc'); -SELECT is_definer('drop_table_constraints'); -SELECT is_definer('create_table_constraints'); +SELECT is_definer('widen_partition_stats'); +SELECT is_definer('tighten_partition_stats'); +SELECT is_definer('ensure_fragments'); +SELECT is_definer('make_binary_staging'); +SELECT is_definer('flush_items_staging_binary'); +SELECT is_definer('delete_item'); +SELECT is_definer('items_staging_triggerfunc'); +SELECT is_definer('update_field_registry_from_sample'); +SELECT is_definer('refresh_field_registry'); +SELECT is_definer('gc_fragments'); SELECT is_definer('check_partition'); SELECT is_definer('repartition'); SELECT is_definer('where_stats'); diff --git a/src/pgstac/tests/pgtap/001a_jsonutils.sql b/src/pgstac/tests/pgtap/001a_jsonutils.sql index 95e951b9..5f5830b1 100644 --- a/src/pgstac/tests/pgtap/001a_jsonutils.sql +++ b/src/pgstac/tests/pgtap/001a_jsonutils.sql @@ -106,26 +106,3 @@ SELECT results_eq( $$ SELECT '{"a":1}'::jsonb $$, 'jsonb_merge_recursive returns the fragment when the item is empty' ); - --- jsonb_canonical / jsonb_hash / jsonb_field_rows: RFC 8785-aligned, externally reproducible. -SELECT has_function('pgstac'::name, 'jsonb_canonical', ARRAY['jsonb']); -SELECT has_function('pgstac'::name, 'jsonb_hash', ARRAY['jsonb']); -SELECT has_function('pgstac'::name, 'jsonb_field_rows', ARRAY['jsonb', 'text', 'integer']); - -SELECT results_eq( - $$ SELECT pgstac.jsonb_canonical('{"id":1,"bbox":2}'::jsonb) $$, - $$ SELECT '{"bbox":2,"id":1}'::text $$, - 'jsonb_canonical sorts object keys by code point (alphabetical), not jsonb length-then-byte order' -); - -SELECT results_eq( - $$ SELECT pgstac.jsonb_canonical('{"b":1.0,"a":[3,2,1],"c":{"y":0.10,"x":true},"d":null,"id":"x"}'::jsonb) $$, - $$ SELECT '{"a":[3,2,1],"b":1,"c":{"x":true,"y":0.1},"d":null,"id":"x"}'::text $$, - 'jsonb_canonical: nested key sort, array order preserved, numbers as shortest round-trip doubles' -); - -SELECT results_eq( - $$ SELECT pgstac.jsonb_hash('{"b":1.0,"a":[3,2,1],"c":{"y":0.10,"x":true},"d":null,"id":"x"}'::jsonb) $$, - $$ SELECT decode('77f18c0a2c2c9f9e4836045bae644ba3d00c0308c9d2c0bd024624c22d532bf7', 'hex') $$, - 'jsonb_hash returns a 32-byte bytea matching an externally-computed sha256 of the canonical form' -); diff --git a/src/pgstac/tests/pgtap/002c_envelope.sql b/src/pgstac/tests/pgtap/002c_envelope.sql index 19d86646..dc4f0394 100644 --- a/src/pgstac/tests/pgtap/002c_envelope.sql +++ b/src/pgstac/tests/pgtap/002c_envelope.sql @@ -3,7 +3,7 @@ SELECT has_function('pgstac'::name, 'env_and', ARRAY['pred_envelope','pred_envel SELECT has_function('pgstac'::name, 'env_or', ARRAY['pred_envelope','pred_envelope']); SELECT has_function('pgstac'::name, 'search_envelope', ARRAY['jsonb']); SELECT has_function('pgstac'::name, 'partition_bounds', ARRAY['pred_envelope']); -SELECT has_function('pgstac'::name, 'next_band', ARRAY['bigint[]','integer','numeric','integer']); +SELECT has_function('pgstac'::name, 'next_band', ARRAY['bigint[]','integer','numeric','integer','boolean']); SELECT results_eq($$ SELECT (env_full()).colls IS NULL $$, $$ SELECT true $$, 'env_full: colls NULL'); SELECT results_eq($$ SELECT (env_full()).dt IS NOT NULL $$, $$ SELECT true $$, 'env_full: dt not NULL'); @@ -24,6 +24,12 @@ SELECT results_eq($$ SELECT (nb).done FROM next_band(ARRAY[10,20,30,40]::bigint[ SELECT results_eq($$ SELECT (nb).done FROM next_band(ARRAY[10,20]::bigint[], 1, 100, 4) nb $$, $$ SELECT true $$, 'next_band: done when exceeds'); SELECT results_eq($$ SELECT (nb).band_end_idx FROM next_band(ARRAY[10,10,10,10,10,10]::bigint[], 1, 100, 2) nb $$, $$ SELECT 2 $$, 'next_band: cap'); SELECT results_eq($$ SELECT (nb).done FROM next_band(NULL::bigint[], 1, 10, 4) nb $$, $$ SELECT true $$, 'next_band: NULL done'); +-- Descending walk (the desc band-order fix): starting at the most recent band and walking toward +-- older months. Cursor at index 4, target 50 over counts 40+30 reaches the target at index 3. +SELECT results_eq($$ SELECT (nb).band_start_idx FROM next_band(ARRAY[10,20,30,40]::bigint[], 4, 50, 4, true) nb $$, $$ SELECT 3 $$, 'next_band desc: start (older bound)'); +SELECT results_eq($$ SELECT (nb).band_end_idx FROM next_band(ARRAY[10,20,30,40]::bigint[], 4, 50, 4, true) nb $$, $$ SELECT 4 $$, 'next_band desc: end (newer bound)'); +SELECT results_eq($$ SELECT (nb).next_cursor_idx FROM next_band(ARRAY[10,20,30,40]::bigint[], 4, 50, 4, true) nb $$, $$ SELECT 2 $$, 'next_band desc: cursor moves toward older'); +SELECT results_eq($$ SELECT (nb).done FROM next_band(ARRAY[10,20,30,40]::bigint[], 2, 1000, 4, true) nb $$, $$ SELECT true $$, 'next_band desc: done at oldest band'); SELECT throws_ok($$ SELECT cql2_envelope('{"op":"s_intersects","args":[{"property":"geometry"},{"type":"Invalid"}]}') $$, '22P02', NULL, 'cql2_envelope: bad GeoJSON raises'); diff --git a/src/pypgstac/src/pypgstac/load.py b/src/pypgstac/src/pypgstac/load.py index 657580cb..377a8b1f 100644 --- a/src/pypgstac/src/pypgstac/load.py +++ b/src/pypgstac/src/pypgstac/load.py @@ -11,12 +11,11 @@ from enum import Enum from pathlib import Path from typing import ( + IO, Any, - BinaryIO, Generator, Iterable, Iterator, - TextIO, ) import orjson @@ -98,7 +97,7 @@ def open_std( **kwargs: Any, ) -> Generator[Any, None, None]: """Open files and i/o streams transparently.""" - fh: TextIO | BinaryIO + fh: IO[Any] if ( filename is None or filename == "-"