pgstac-rs: Rust engine for pgstac (search, ingest, dump/export)

.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,14 @@ 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 all features
+        run: cargo check --all-features --all-targets --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

CHANGELOG.md

@@ -15,9 +15,15 @@ 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 NDJSON / a FeatureCollection page / stac-geoparquet off the keyset engine.
+- `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.
 
 ### 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 <col>` (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.

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"

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 <COMMAND>
+  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 <DSN>` | `$PGSTAC_DSN` / local dev | Postgres connection string. |
+| `-o, --out <OUT>` | *(required)* | Destination: a directory, `s3://bucket/prefix` (or other object-store URL), a `*.tar` / `*.tar.zst` file, or `-` for stdout. |
+| `-c, --collection <ID>` | *(all)* | Restrict to these collection ids (repeatable). |
+| `--datetime-start <RFC3339>` | — | Inclusive datetime prefilter start (requires `--datetime-end`). |
+| `--datetime-end <RFC3339>` | — | Exclusive datetime prefilter end (requires `--datetime-start`). |
+| `--bbox <W> <S> <E> <N>` | — | Bbox prefilter (EPSG:4326). |
+| `--compression <CODEC>` | `zstd` | Parquet codec: `zstd`, `snappy`, `uncompressed`. |
+| `--skip-errors` | off | Continue past per-item errors, recording skips in the report. |
+| `--memory-budget <BYTES>` | ~25% RAM | Memory budget for the buffered geoparquet path. |
+| `--concurrency <N>` | `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 <DSN>` | `$PGSTAC_DSN` / local dev | Postgres connection string. |
+| `-o, --out <OUT>` | `-` (stdout) | A file path or `-` for stdout. |
+| `--format <FORMAT>` | `ndjson` | `ndjson` (streams all pages), `itemcollection` (one SQL-faithful page with next/prev + context), or `geoparquet`. |
+| `-c, --collection <ID>` | *(all)* | Restrict to these collection ids (repeatable). |
+| `--id <ID>` | — | Restrict to these item ids (repeatable). |
+| `--bbox <W> <S> <E> <N>` | — | Bbox (EPSG:4326). |
+| `--datetime <DT>` | — | Datetime / interval (STAC datetime syntax). |
+| `--filter <CQL2>` | — | CQL2-JSON filter (a JSON object). |
+| `--limit <N>` | server default | Page size (the keyset limit). |
+| `--token <TOKEN>` | — | Continuation token (`next:…` / `prev:…`). With `itemcollection`, returns exactly that one page. |
+| `--max-items <N>` | *(unbounded)* | Cap total items streamed (`ndjson` / `geoparquet`). |
+| `--compression <CODEC>` | `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)`.

scripts/pgstac-rs-test-db

@@ -0,0 +1,43 @@
+#!/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)
+[[ -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 <<SQL
+SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname = '$TEMPLATE';
+DROP DATABASE IF EXISTS $TEMPLATE;
+CREATE DATABASE $TEMPLATE;
+ALTER DATABASE $TEMPLATE SET search_path TO pgstac, public;
+ALTER DATABASE $TEMPLATE SET client_min_messages TO WARNING;
+SQL
+
+psql -h "$PGHOST" -p "$PGPORT" -U "$PGUSER" -d "$TEMPLATE" -v ON_ERROR_STOP=1 -q -f "$PGSTAC_SQL"
+
+# Verify a pristine install.
+psql -h "$PGHOST" -p "$PGPORT" -U "$PGUSER" -d "$TEMPLATE" -tAc \
+    "SET search_path=pgstac,public; SELECT format('pgstac %s ready: %s collections, %s items', get_version(), (SELECT count(*) FROM collections), (SELECT count(*) FROM items));"

scripts/pgstac-rs-test-rich-db

@@ -0,0 +1,75 @@
+#!/bin/bash
+# Build the RICH pgstac test template used by the pgstac-rs hydration/search tests: real
+# landsat-c2-l2 / sentinel-2-l2a / naip collections (many assets + properties) with `fragment_config`
+# set, plus 1000 items each dehydrated into `item_fragments`. This exercises the fragment-hydration
+# path, which is one of the largest bottlenecks. Re-run after schema changes.
+#
+# Usage: scripts/pgstac-rs-test-rich-db [TEMPLATE_NAME]   (default pgstac_rs_test_rich)
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
+source "$SCRIPT_DIR/pgstacenv"   # cd to repo root, ensure .env
+
+[[ -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 <<SQL
+SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname = '$TEMPLATE';
+DROP DATABASE IF EXISTS $TEMPLATE;
+CREATE DATABASE $TEMPLATE;
+ALTER DATABASE $TEMPLATE SET search_path TO pgstac, public;
+ALTER DATABASE $TEMPLATE SET client_min_messages TO WARNING;
+SQL
+
+psql -d "$TEMPLATE" -v ON_ERROR_STOP=1 -q -f "$PGSTAC_SQL"
+
+# Load raw NDJSON into jsonb columns via CSV format with control-char quote/delimiter (chosen so they
+# never appear in the data): this reads each line verbatim as one field, avoiding COPY text-format
+# backslash de-escaping that would corrupt JSON escapes like \n.
+JSON_COPY="WITH (FORMAT csv, DELIMITER E'\\x02', QUOTE E'\\x01')"
+
+# Load collections, then BEFORE loading items: derive fragment_config from each collection's
+# item_assets (so items dehydrate into shared fragments) and set partition_trunc=month (so items land
+# in datetime sub-partitions and the search() band engine has a real histogram — matching how the
+# synth/landsat test collections are configured).
+psql -d "$TEMPLATE" -v ON_ERROR_STOP=1 -q <<SQL
+\copy collections (content) FROM '$FIXROOT/collections.ndjson' $JSON_COPY
+UPDATE collections
+   SET fragment_config = collection_fragment_config_default(content),
+       partition_trunc = 'month'
+ WHERE fragment_config IS NULL;
+SQL
+
+for fixture in "${FIXTURES[@]}"; do
+    items="$FIXROOT/data/$fixture/items.ndjson"
+    [[ -f "$items" ]] || { echo "ERROR: $items not found." >&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));"