diff --git a/.github/dependabot.yml b/.github/dependabot.yml new file mode 100644 index 0000000..a4f8b0f --- /dev/null +++ b/.github/dependabot.yml @@ -0,0 +1,24 @@ +version: 2 + +updates: + - package-ecosystem: "gomod" + directories: + - "/" + - "/extra/*" + - "/examples/*" + schedule: + interval: "weekly" + ignore: + # Local modules are resolved through go.work during repository builds. + - dependency-name: "github.com/mkbeh/xpg" + - dependency-name: "github.com/mkbeh/xpg/*" + groups: + go-dependencies: + patterns: + - "*" + group-by: dependency-name + + - package-ecosystem: "github-actions" + directory: "/" + schedule: + interval: "weekly" diff --git a/.github/scripts/create-go-workspace.sh b/.github/scripts/create-go-workspace.sh new file mode 100755 index 0000000..950ceee --- /dev/null +++ b/.github/scripts/create-go-workspace.sh @@ -0,0 +1,95 @@ +#!/usr/bin/env bash +set -euo pipefail + +if [[ -n "${GITHUB_WORKSPACE:-}" ]]; then + root="${GITHUB_WORKSPACE}" +else + root="$(git rev-parse --show-toplevel 2>/dev/null || pwd)" +fi + +include_examples=false + +for arg in "$@"; do + case "${arg}" in + --with-examples) + include_examples=true + ;; + *) + echo "unknown argument: ${arg}" >&2 + exit 2 + ;; + esac +done + +cd "${root}" + +rm -f go.work go.work.sum + +modules=(.) + +while IFS= read -r -d '' mod; do + modules+=("$(dirname "${mod}")") +done < <(find ./extra -name go.mod -type f -print0 2>/dev/null | sort -z) + +if [[ "${include_examples}" == "true" ]]; then + while IFS= read -r -d '' mod; do + modules+=("$(dirname "${mod}")") + done < <(find ./examples -name go.mod -type f -print0 2>/dev/null | sort -z) +fi + +# The workspace is ephemeral in CI and local-only for development. +GOWORK=off go work init "${modules[@]}" + +# Local modules may already require release versions such as v0.1.0 while the +# corresponding tags do not exist yet. Bind those exact requirements to the +# current checkout so all Go tooling, including go/packages-based linters, +# resolves the unpublished modules locally. +while read -r module version; do + [[ -n "${module}" && -n "${version}" ]] || continue + + case "${module}" in + github.com/mkbeh/xpg) + local_dir="." + ;; + github.com/mkbeh/xpg/*) + local_dir="./${module#github.com/mkbeh/xpg/}" + ;; + *) + continue + ;; + esac + + if [[ -f "${root}/${local_dir#./}/go.mod" ]]; then + GOWORK="${root}/go.work" go work edit \ + -replace="${module}@${version}=${local_dir}" + fi +done < <( + { + printf '%s\0' ./go.mod + + find ./extra ./examples \ + -name go.mod -type f -print0 2>/dev/null || true + } | + while IFS= read -r -d '' mod; do + awk ' + $1 == "require" && + $2 ~ /^github\.com\/mkbeh\/xpg(\/.*)?$/ && + $3 ~ /^v[0-9]/ { + print $2, $3 + next + } + + $1 ~ /^github\.com\/mkbeh\/xpg(\/.*)?$/ && + $2 ~ /^v[0-9]/ { + print $1, $2 + } + ' "${mod}" + done | + sort -u +) + +echo "Created temporary Go workspace at ${root}/go.work" + +if [[ "${WORKSPACE_DEBUG:-false}" == "true" ]]; then + GOWORK="${root}/go.work" go work edit -json +fi diff --git a/.github/workflows/codecov.yml b/.github/workflows/codecov.yml new file mode 100644 index 0000000..d0c3d86 --- /dev/null +++ b/.github/workflows/codecov.yml @@ -0,0 +1,60 @@ +name: Coverage + +on: + push: + branches: + - main + pull_request: + branches: + - main + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +env: + GOTOOLCHAIN: local + CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }} + +jobs: + coverage: + name: Coverage + runs-on: ubuntu-latest + timeout-minutes: 15 + + steps: + - name: Check out repository + uses: actions/checkout@v7 + + - name: Set up Go + uses: actions/setup-go@v7 + with: + go-version-file: go.mod + check-latest: true + cache-dependency-path: | + go.sum + extra/**/go.mod + + - name: Set up Task + uses: go-task/setup-task@v2 + + - name: Create temporary Go workspace + run: bash .github/scripts/create-go-workspace.sh + + - name: Run coverage + env: + GOWORK: ${{ github.workspace }}/go.work + run: task coverage + + - name: Upload coverage to Codecov + if: env.CODECOV_TOKEN != '' + uses: codecov/codecov-action@v7 + with: + token: ${{ env.CODECOV_TOKEN }} + files: ./coverage.out + disable_search: true + fail_ci_if_error: true diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml new file mode 100644 index 0000000..22648e6 --- /dev/null +++ b/.github/workflows/lint.yml @@ -0,0 +1,123 @@ +name: Lint + +on: + push: + branches: + - main + pull_request: + branches: + - main + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +env: + GOTOOLCHAIN: local + +jobs: + modules: + name: Discover modules + runs-on: ubuntu-latest + + outputs: + modules: ${{ steps.modules.outputs.modules }} + + steps: + - name: Check out repository + uses: actions/checkout@v7 + + - name: Discover Go modules + id: modules + shell: bash + run: | + set -euo pipefail + + modules="$( + { + printf '.\n' + + find ./extra ./examples \ + -name go.mod \ + -type f \ + -print | + while IFS= read -r mod; do + dirname "${mod#./}" + done + } | + sort | + jq -R -s -c 'split("\n")[:-1]' + )" + + echo "modules=${modules}" >> "${GITHUB_OUTPUT}" + + lint: + name: Lint (${{ matrix.module }}) + needs: + - modules + + runs-on: ubuntu-latest + timeout-minutes: 15 + + strategy: + fail-fast: false + matrix: + module: ${{ fromJSON(needs.modules.outputs.modules) }} + + steps: + - name: Check out repository + uses: actions/checkout@v7 + + - name: Set up Go + uses: actions/setup-go@v7 + with: + go-version-file: go.mod + check-latest: true + cache-dependency-path: | + go.sum + extra/**/go.mod + examples/**/go.mod + + - name: Create temporary Go workspace + run: bash .github/scripts/create-go-workspace.sh --with-examples + + - name: Verify workspace resolution + env: + GOWORK: ${{ github.workspace }}/go.work + run: | + set -euo pipefail + + test "$(go env GOWORK)" = "${GITHUB_WORKSPACE}/go.work" + + go list -m \ + -f '{{if .Main}}{{.Path}} => {{.Dir}}{{end}}' \ + all | + sed '/^$/d' + + - name: Run golangci-lint + uses: golangci/golangci-lint-action@v9 + env: + GOWORK: ${{ github.workspace }}/go.work + with: + version: v2.13.1 + working-directory: ${{ matrix.module }} + args: >- + --config=${{ github.workspace }}/.golangci.yml + --timeout=5m + --modules-download-mode=readonly + + result: + name: Lint + if: always() + needs: + - lint + runs-on: ubuntu-latest + timeout-minutes: 5 + + steps: + - name: Check lint jobs + run: test "${{ needs.lint.result }}" = "success" diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml new file mode 100644 index 0000000..70af09a --- /dev/null +++ b/.github/workflows/test.yml @@ -0,0 +1,107 @@ +name: Test + +on: + push: + branches: + - main + pull_request: + branches: + - main + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +env: + GOTOOLCHAIN: local + +jobs: + test: + name: Tests + runs-on: ubuntu-latest + timeout-minutes: 15 + + steps: + - name: Check out repository + uses: actions/checkout@v7 + + - name: Set up Go + uses: actions/setup-go@v7 + with: + go-version-file: go.mod + check-latest: true + cache-dependency-path: | + go.sum + extra/**/go.mod + examples/**/go.mod + + - name: Set up Task + uses: go-task/setup-task@v2 + + - name: Create temporary Go workspace + run: bash .github/scripts/create-go-workspace.sh --with-examples + + - name: Verify workspace resolution + env: + GOWORK: ${{ github.workspace }}/go.work + run: | + set -euo pipefail + + test "$(go env GOWORK)" = "${GITHUB_WORKSPACE}/go.work" + + go list -m \ + -f '{{if .Main}}{{.Path}} => {{.Dir}}{{end}}' \ + all | + sed '/^$/d' + + - name: Run tests + env: + GOWORK: ${{ github.workspace }}/go.work + run: task test + + race: + name: Race + runs-on: ubuntu-latest + timeout-minutes: 20 + + steps: + - name: Check out repository + uses: actions/checkout@v7 + + - name: Set up Go + uses: actions/setup-go@v7 + with: + go-version-file: go.mod + check-latest: true + cache-dependency-path: | + go.sum + extra/**/go.mod + examples/**/go.mod + + - name: Set up Task + uses: go-task/setup-task@v2 + + - name: Create temporary Go workspace + run: bash .github/scripts/create-go-workspace.sh --with-examples + + - name: Verify workspace resolution + env: + GOWORK: ${{ github.workspace }}/go.work + run: | + set -euo pipefail + + test "$(go env GOWORK)" = "${GITHUB_WORKSPACE}/go.work" + + go list -m \ + -f '{{if .Main}}{{.Path}} => {{.Dir}}{{end}}' \ + all | + sed '/^$/d' + + - name: Run race tests + env: + GOWORK: ${{ github.workspace }}/go.work + run: task test-race diff --git a/.gitignore b/.gitignore index 2149a16..15e302a 100644 --- a/.gitignore +++ b/.gitignore @@ -28,3 +28,5 @@ go.work.sum .vscode .idea +# other +/tmp/ diff --git a/.golangci.yml b/.golangci.yml index e419d53..9c32ca8 100755 --- a/.golangci.yml +++ b/.golangci.yml @@ -1,7 +1,10 @@ version: "2" + run: - go: "1.26" + go: "1.27" + timeout: 5m allow-parallel-runners: true + linters: default: none enable: @@ -10,6 +13,8 @@ linters: - bodyclose - copyloopvar - durationcheck + - errcheck + - errorlint - exhaustive - gocritic - goprintffuncname @@ -24,17 +29,24 @@ linters: - rowserrcheck - sloglint - sqlclosecheck + - staticcheck - unconvert - unparam - unused - usetesting - wastedassign - whitespace + settings: + errcheck: + check-type-assertions: true + errorlint: errorf: false + exhaustive: default-signifies-exhaustive: true + gocritic: disabled-checks: - appendAssign @@ -61,15 +73,18 @@ linters: sizeThreshold: 256 rangeValCopy: sizeThreshold: 256 + gosec: excludes: - G104 - - G404 - G115 + - G404 + nolintlint: require-explanation: true require-specific: true allow-unused: false + revive: confidence: 0.8 severity: warning @@ -105,7 +120,6 @@ linters: - name: superfluous-else - name: time-equal - name: time-naming - - name: var-declaration - name: unconditional-recursion - name: unexported-naming - name: unexported-return @@ -113,7 +127,9 @@ linters: - name: unreachable-code - name: unused-parameter - name: useless-break + - name: var-declaration - name: waitgroup-by-value + sloglint: no-mixed-args: true no-global: default @@ -122,30 +138,42 @@ linters: no-raw-keys: false key-naming-case: snake args-on-sep-lines: true + staticcheck: checks: - all - -SA1012 - -SA1019 + - -ST1000 + exclusions: generated: lax paths: - - third_party$ - - builtin$ - - examples + - third_party/ + - builtin/ + - examples/ + - _tmp/ + rules: + - path: _test\.go + linters: + - revive + text: "dot-imports" + issues: max-same-issues: 0 + formatters: enable: - - gofmt - gofumpt - goimports settings: gofumpt: - extra-rules: true + extra: + group-params: true exclusions: generated: lax paths: - - third_party$ - - builtin$ - - examples$ + - third_party/ + - builtin/ + - examples/ + - _tmp/ \ No newline at end of file diff --git a/CHANGELOG.md b/CHANGELOG.md index b4bf2c2..621ec21 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,51 @@ -# 0.1.13 (Jun 22, 2026) +# Changelog -* build: bump dependencies +All notable changes to this project will be documented in this file. + +## v0.2.0 + +Initial production release of `xpg`, built around `pgx` with PostgreSQL transaction helpers, primary/replica clustering, +application-level sharding, and observability integrations. + +### Added + +* **Pool Lifecycle Management:** Added a thin pool layer over `pgxpool` with pgx query operations, pool metadata, + runtime statistics, configuration options, and direct access to the underlying pool. +* **Transactions and Savepoints:** Added managed transactions and savepoints for isolating optional work within a + transaction. +* **Advisory Locks:** Added transaction-level PostgreSQL advisory locks for coordinating concurrent work. +* **PostgreSQL Error Classification:** Added SQLSTATE inspection and helpers for constraint violations, serialization + failures, deadlocks, lock errors, query cancellation, and connection failures. +* **Primary/Replica Clustering:** Added logical clusters with explicit read policies, round-robin and custom replica + selection, replica-preferred fallback when no replica can be selected, and enforced read-only transactions. +* **Application-Level Sharding:** Added immutable shard topologies with rendezvous hashing, ordered ranges, time ranges, + and custom resolvers, together with colocation checks, grouping by shard, and bounded parallel operations. +* **Logging and Tracing:** Added pgx-compatible logging and tracing hooks with support for combining multiple query + tracers. +* **Examples:** Added runnable examples covering pool usage, transactions, advisory locks, observability, clustering, + range-based sharding, and custom geographic routing. + +--- + +## extra/otelxpg/v0.1.0 + +Initial release of the `otelxpg` integration module. + +### Added + +* **OpenTelemetry Pool Metrics:** Added metrics for PostgreSQL connection-pool state, usage, acquisition behavior, and + lifecycle counters. +* **Metric Attributes:** Added pool names and custom labels as metric attributes for filtering and aggregation across + standalone pools, cluster nodes, and shard pools. +* **Meter Provider Integration:** Added support for an application-provided `metric.MeterProvider`. + +--- + +## extra/slogxpg/v0.1.0 + +Initial release of the `slogxpg` integration module. + +### Added + +* **slog Adapter:** Added an adapter from the standard library `log/slog` logger to `pgx/tracelog.Logger`. +* **Log Level Mapping:** Added deterministic mapping from pgx trace log levels to `slog` levels. \ No newline at end of file diff --git a/LICENSE b/LICENSE index f7372df..d22da70 100644 --- a/LICENSE +++ b/LICENSE @@ -1,6 +1,6 @@ MIT License -Copyright (c) 2024 Nia +Copyright (c) 2024 mkbeh Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal diff --git a/README.md b/README.md index a0d2be9..3c83c29 100644 --- a/README.md +++ b/README.md @@ -1,332 +1,237 @@
-# xpg +# Postgres toolkit for Go **Lightweight PostgreSQL wrapper for Go, built on top of [pgx](https://github.com/jackc/pgx).** -![Go Version](https://img.shields.io/badge/go-1.26%2B-blue) -[![License: MIT](https://img.shields.io/badge/license-MIT-green.svg)](LICENSE) +[![Go Reference](https://pkg.go.dev/badge/github.com/mkbeh/xpg.svg)](https://pkg.go.dev/github.com/mkbeh/xpg) +[![Test](https://github.com/mkbeh/xpg/actions/workflows/test.yml/badge.svg)](https://github.com/mkbeh/xpg/actions/workflows/test.yml) +[![Coverage](https://codecov.io/gh/mkbeh/xpg/graph/badge.svg)](https://codecov.io/gh/mkbeh/xpg)
-`xpg` wraps the excellent [`pgx`](https://github.com/jackc/pgx) PostgreSQL driver with a compact API for common -PostgreSQL workflows: read/write connection pool splitting, transaction helpers, embedded SQL migrations, query -building, normalized errors, and exposing PostgreSQL observability with OpenTelemetry and Prometheus. +`xpg` builds on the `pgx` client with a compact API for common PostgreSQL infrastructure patterns. It adds support for +pool lifecycle, transactions and savepoints, PostgreSQL error classification, advisory locks, primary/replica routing, +application-level sharding, and observability. + +The library uses `pgx` types and query model directly while keeping its core behavior and reducing boilerplate around +connection management, routing, and common production workflows. ## Features -* **Pools**: Separate writer and reader connection pools. -* **Queries**: PostgreSQL-friendly query builder support. -* **Transactions**: Transaction helpers with automatic rollback and panic recovery. -* **Migrations**: Embedded SQL migrations via [golang-migrate](https://github.com/golang-migrate/migrate). -* **Errors**: Normalized PostgreSQL error codes for common failure cases. -* **Observability**: OpenTelemetry tracing and Prometheus metrics out of the box. -* **Configuration**: Configure via Go structs or environment variables. +* **Pool Lifecycle Management:** Thin pool management on top of `pgx` with direct access to the underlying PostgreSQL + client. +* **Transactions and Savepoints:** Managed transactions, savepoints, and helpers for common multi-step transactional + workflows. +* **Error Classification:** Classification of PostgreSQL constraint, transaction, cancellation, connection, and other + common database errors. +* **Advisory Locking:** Transaction-level advisory locks for coordinating concurrent database operations. +* **Primary/Replica Routing:** Explicit read policies, replica selection, primary fallback, and read-only transactions + across PostgreSQL nodes. +* **Application-Level Sharding:** Hash, range, time-based, and custom routing with colocation checks, key grouping, and + bounded parallel operations across shards. +* **Observability:** Structured logging, tracing, pool statistics, and optional OpenTelemetry metrics. ## Installation +This repository contains the core `xpg` module. The core module is released from the repository root: + ```bash go get github.com/mkbeh/xpg ``` -## Quick start - -The example below creates separate writer and reader pools and runs a simple query. - -```go -package main - -import ( - "context" - "fmt" - "log" - - postgres "github.com/mkbeh/xpg" -) - -func main() { - ctx := context.Background() - - cfg := &postgres.Config{ - ClusterHost: "127.0.0.1", - ClusterPort: "5432", // writer/master port - ClusterReplicaPort: "5432", // reader/replica port; can be different in production - User: "user", - Password: "pass", - DB: "postgres", - } - - writer, err := postgres.NewWriter( - postgres.WithConfig(cfg), - postgres.WithClientID("my-service"), // appended to application_name and metrics labels - ) - if err != nil { - log.Fatal("failed to init writer pool:", err) - } - defer writer.Close() - - reader, err := postgres.NewReader( - postgres.WithConfig(cfg), - postgres.WithClientID("my-service"), - ) - if err != nil { - log.Fatal("failed to init reader pool:", err) - } - defer reader.Close() - - // Use the writer pool for write-side operations. - if _, err := writer.Exec(ctx, "select 1"); err != nil { - log.Fatal("writer query failed:", err) - } - - var greeting string - - // Use the reader pool for read-only queries. - if err := reader.QueryRow(ctx, "select 'Hello, world!'").Scan(&greeting); err != nil { - log.Fatal("query failed:", err) - } +Optional integrations are released independently under `extra`: - fmt.Println(greeting) -} +```bash +go get github.com/mkbeh/xpg/extra/otelxpg ``` -More examples: [examples/](https://github.com/mkbeh/xpg/tree/main/examples) - -## Query Builder +## Usage -Each pool includes a preconfigured [squirrel](https://github.com/Masterminds/squirrel) statement builder with PostgreSQL -dollar placeholders out of the box. +Open an `xpg` pool and execute a PostgreSQL query: ```go -sql, args, err := writer.QueryBuilder(). - Insert("orders"). - Columns("id", "status"). - Values(orderID, "created"). - ToSql() +// urlExample := "postgres://username:password@localhost:5432/database_name" +pool, err := xpg.Open( + context.Background(), + os.Getenv("DATABASE_URL"), + xpg.WithName("example-pool"), +) if err != nil { - log.Fatalf("failed to build query: %v", err) + log.Fatalf("failed to open pool: %v", err) } +defer pool.Close() -if _, err := writer.Exec(ctx, sql, args...); err != nil { - log.Fatalf("failed to execute insert: %v", err) +var message string +err = pool.QueryRow(context.Background(), "SELECT 'hello from xpg'").Scan(&message) +if err != nil { + log.Fatalf("query failed: %v", err) } + +fmt.Println(message) // Outputs: hello from xpg ``` -## Transactions +### Transactions -Use `RunInTxx` for transactions with default options. It acts as an alias for `RunInTx` using default `pgx.TxOptions`. +`xpg` provides managed transactions using the native `pgx` transaction API. Returning `nil` commits the transaction; +returning an error rolls it back. ```go -err := writer.RunInTxx(ctx, func(ctx context.Context) error { - _, err := writer.Exec(ctx, "INSERT INTO orders (id) VALUES (\$1)", orderID) +err := pool.InTx(ctx, pgx.TxOptions{}, func(ctx context.Context, tx pgx.Tx) error { + _, err := tx.Exec(ctx, "UPDATE users SET active = true WHERE id = $1", userID) return err }) -if err != nil { - log.Fatalf("transaction failed: %v", err) -} ``` -For a custom isolation level or access mode, use `RunInTx`: +Savepoints can isolate optional work without aborting the outer transaction. + +### Advisory Locks + +`xpg` provides transaction-level PostgreSQL advisory locks for coordinating concurrent work. ```go -err := writer.RunInTx(ctx, func(ctx context.Context) error { - _, err := writer.Exec(ctx, "INSERT INTO orders (id) VALUES (\$1)", orderID) +err := pool.InTx(ctx, pgx.TxOptions{}, func(ctx context.Context, tx pgx.Tx) error { + if err := xpg.AdvisoryXactLock(ctx, tx, lockID); err != nil { + return err + } + + _, err := tx.Exec(ctx, "UPDATE jobs SET status = 'running' WHERE id = $1", jobID) return err -}, pgx.TxOptions{ - IsoLevel: pgx.Serializable, }) -if err != nil { - log.Fatalf("serializable transaction failed: %v", err) -} ``` -Rollback is handled automatically if an error occurs. The transaction is committed only if the function returns `nil`. -Any panics inside the block are recovered and returned as standard Go errors. - -## Migrations +The lock is held for the duration of the transaction and released automatically on commit or rollback. -`xpg` supports embedded SQL migrations out of the box using [golang-migrate](https://github.com/golang-migrate/migrate). +### Error Handling -First, create an `embed.go` file inside your migrations directory: +`xpg` provides helpers for classifying PostgreSQL errors and inspecting SQLSTATE codes. ```go -package migrations - -import "embed" - -//go:embed *.sql -var FS embed.FS +_, err := pool.Exec(ctx, "INSERT INTO users (id, email) VALUES ($1, $2)", userID, email) + +switch { +case xpg.IsUniqueViolation(err): + // Handle duplicate data. +case xpg.IsRetryableTransaction(err): + // Retry the transaction when the operation is safe to replay. +case err != nil: + return err +} ``` -Then, pass the embedded filesystem using `WithMigrations`: +The underlying SQLSTATE code is also available through `xpg.SQLState`. Helpers cover constraint violations, +serialization failures, deadlocks, lock errors, query cancellation, and connection failures. + +## Clustering + +`xpg` groups primary and replica pools into a logical cluster with explicit read routing. ```go -writer, err := postgres.NewWriter( - postgres.WithConfig(&postgres.Config{ - ClusterHost: "127.0.0.1", - ClusterPort: "5432", - ClusterReplicaPort: "5432", - User: "user", - Password: "pass", - DB: "postgres", - MigrateEnabled: true, - }), - postgres.WithMigrations(migrations.FS), -) +orders, err := cluster.New(cluster.Config{ + ID: "orders", + Primary: primary, + Replicas: []*xpg.Pool{replicaA, replicaB}, +}) if err != nil { - log.Fatalf("failed to initialize writer and run migrations: %v", err) + panic(err) } -defer writer.Close() -``` - +defer orders.Close() -Migrations will run automatically during `NewWriter` initialization if `MigrateEnabled` is set to `true`. +// Route writes explicitly to the primary. +primaryPool := orders.Primary() -Your SQL migration files must follow the standard `golang-migrate` naming convention: - -```text -000001_create_users.up.sql -000001_create_users.down.sql -``` - -## Observability +_, err = primaryPool.Exec(ctx, "UPDATE orders SET status = 'processed' WHERE id = $1", orderID) +if err != nil { + panic(err) +} -`xpg` instruments PostgreSQL queries through native `pgx` tracing hooks and exposes pool metrics for Prometheus. +// Route reads according to the selected policy. +readPool, err := orders.ReadPool(ctx, cluster.ReadReplicaPreferred) +if err != nil { + panic(err) +} - -```go -writer, err := postgres.NewWriter( - postgres.WithConfig(cfg), - postgres.WithClientID("orders-service"), - postgres.WithTraceProvider(tracerProvider), - postgres.WithMetricsNamespace("orders"), -) +var status string +err = readPool.QueryRow(ctx, "SELECT status FROM orders WHERE id = $1", orderID).Scan(&status) if err != nil { - log.Fatalf("failed to initialize observed writer pool: %v", err) + panic(err) } -defer writer.Close() ``` +Read policies support primary-only, replica-required, and replica-preferred routing with primary fallback when no +replica is available. Replica selection is round-robin by default and can be customized. -The following Prometheus metric labels are added automatically: +## Sharding -| Label | Description | -| :--- | :--- | -| `client_id` | Generated client identifier or configured ID with a unique suffix. | -| `client_kind` | `writer` for writer pools, `reader` for reader pools. | -| `db` | Database name from the configuration. | -| `shard_id` | Shard ID from the configuration. | - -## Error Handling - -`xpg` provides normalized PostgreSQL error codes through `ConvertError`, so application code does not need to deal with -raw `pgx` and `pgconn` error types directly. +`xpg` provides application-level sharding with explicit key routing across an immutable shard topology. ```go -err := writer.QueryRow(ctx, "SELECT id FROM users WHERE id = \$1", userID).Scan(&id) +topology, err := shard.NewTopology([]shard.Config{ + {Cluster: shardA}, + {Cluster: shardB}, +}) if err != nil { - pgErr := postgres.ConvertError(err) - - if pgErr.Code() == postgres.ErrNoRows { - // handle missing row - return nil - } - - if pgErr.Code() == postgres.ErrSerializable { - // retry transaction - return nil - } - - return pgErr + panic(err) +} +defer topology.Close() + +// Partition user IDs into shard ranges. +users, err := resolver.NewRange( + topology, + []resolver.Range[uint64]{ + {Start: 0, End: 100, ShardID: "shard-a"}, + {Start: 100, End: 200, ShardID: "shard-b"}, + }, +) +if err != nil { + panic(err) } -``` - -Common PostgreSQL errors such as `ErrNoRows`, `ErrUniqViolation`, `ErrForeignKeyViolation`, and `ErrSerializable` are -mapped to stable `xpg` error codes. +// Resolve the target shard. +shard, err := users.Resolve(userID) +if err != nil { + panic(err) +} -## Configuration +// Write to the shard primary. +primaryPool := shard.Primary() -The `Config` struct can be initialized directly in Go. It also includes `envconfig` tags, allowing you to seamlessly -populate it from environment variables using your preferred configuration library. +_, err = primaryPool.Exec(ctx, "UPDATE users SET active = true WHERE id = $1", userID) +if err != nil { + panic(err) +} -### Config Struct +// Read from the same shard using the selected read policy. +readPool, err := shard.ReadPool(ctx, cluster.ReadReplicaPreferred) +if err != nil { + panic(err) +} - -```go -cfg := &postgres.Config{ - ClusterHost: "127.0.0.1", // required - ClusterPort: "5432", // required, writer port - ClusterReplicaPort: "5433", // required, reader port - User: "user", // required - Password: "pass", // required - DB: "mydb", // required - - MaxRWConn: 16, - MaxROConn: 16, - MaxConnLifetime: 5 * time.Minute, - MaxConnIdleTime: 30 * time.Second, - - MigrateEnabled: true, +var active bool +err = readPool.QueryRow(ctx, "SELECT active FROM users WHERE id = $1", userID).Scan(&active) +if err != nil { + panic(err) } ``` -The connection DSN is dynamically built from the `Config` fields using the following format: +Built-in routing strategies include rendezvous hashing, ordered ranges, time ranges, and custom resolvers. Sharding +utilities cover key colocation, grouping by shard, and parallel operations across shards. -```text -postgres://user:pass@host:port/db?sslmode=disable&application_name=& -``` +## Examples -### Environment Variables - -| Variable | Required | Default | Description | -| :--- | :---: | :--- | :--- | -| `POSTGRES_CLUSTER_HOST` | ✓ | — | Database host. | -| `POSTGRES_CLUSTER_PORT` | ✓ | — | Writer pool port. | -| `POSTGRES_CLUSTER_REPLICA_PORT` | ✓ | — | Reader pool port. | -| `POSTGRES_USER` | ✓ | — | Database user. | -| `POSTGRES_PASSWORD` | ✓ | — | Database password. | -| `POSTGRES_DB` | ✓ | — | Database name. | -| `POSTGRES_SHARD_ID` | | `0` | Shard ID exposed in metrics. | -| `POSTGRES_MIN_RW_CONN` | | `1` | Minimum connections in the writer pool. | -| `POSTGRES_MIN_RO_CONN` | | `1` | Minimum connections in the reader pool. | -| `POSTGRES_MAX_RW_CONN` | | `max(4, NumCPU)`| Maximum connections in the writer pool. | -| `POSTGRES_MAX_RO_CONN` | | `max(4, NumCPU)`| Maximum connections in the reader pool. | -| `POSTGRES_MAX_CONN_LIFETIME` | | `1m` | Maximum connection lifetime. | -| `POSTGRES_MAX_CONN_IDLE_TIME` | | `30s` | Maximum idle connection lifetime. | -| `POSTGRES_QUERY_EXEC_MODE` | | `cache_statement` | Query execution mode. | -| `POSTGRES_STATEMENT_CACHE_CAPACITY` | | `128` | Statement cache size. | -| `POSTGRES_DESCRIPTION_CACHE_CAPACITY`| | `512` | Description cache size. | -| `POSTGRES_WRITER_ARGS` | | — | Extra DSN args for the writer connection. | -| `POSTGRES_REPLICA_ARGS` | | — | Extra DSN args for the reader connection. | -| `POSTGRES_MIGRATE_ENABLED` | | `false` | Run migrations on writer startup. | -| `POSTGRES_MIGRATE_PORT` | | `POSTGRES_CLUSTER_PORT` | Port used for migrations. | -| `POSTGRES_MIGRATE_ARGS` | | — | Extra DSN args for the migration connection. | - -### Query Execution Modes - -| Value | Protocol | Round Trips | Description | -| :--- | :--- | :--- | :--- | -| `cache_statement` | Extended | 1 after warm-up | Automatically prepares and caches statements. **Default.** May fail on first execution after schema changes. | -| `cache_describe` | Extended | 1 after warm-up | Caches argument and result type descriptions instead of prepared statements. Has the same schema-change caveat. | -| `describe_exec` | Extended | 2 | Fetches description on every execution, then executes. Safer with concurrent schema changes, but may break with connection poolers that switch connections between round trips. | -| `exec` | Extended | 1 | No prepare and no describe. Infers PostgreSQL types from Go types using text format. Register custom types with `pgtype.Map.RegisterDefaultPgType`. | -| `simple_protocol` | Simple | 1 | Uses the simple protocol. Useful with PgBouncer or proxies that do not support the extended protocol. Prefer `exec` when possible. | - -> 💡 **Tip for connection pooling:** For PgBouncer **transaction** pooling, use `simple_protocol`. For **session** -> pooling, `cache_statement` usually works fine. +See the [examples](examples) directory for runnable examples covering the main `xpg` usage patterns. ## License diff --git a/Taskfile.yml b/Taskfile.yml index bc5cf57..2f18413 100755 --- a/Taskfile.yml +++ b/Taskfile.yml @@ -1,28 +1,171 @@ -version: 3 +version: "3" vars: - LINTER_VER: "v2.12.2" + LINTER_VER: "v2.13.1" + + GOMODCACHE: + sh: go env GOMODCACHE + + GOCACHE: + sh: go env GOCACHE + + GOSUMCACHE: + sh: echo "$(go env GOPATH)/pkg/sumdb" + + # Root module + all local submodules that are part of the project. + # Examples are included because they depend on unpublished local modules. + MODULE_DIRS: + sh: | + { + printf '.\n' + + find ./extra ./examples -name go.mod -print | + while IFS= read -r mod; do + dirname "$mod" + done + } | sort -u + + # Coverage is aggregated for production modules only. Example modules are + # still linted and tested, but do not affect the library coverage metric. + COVERAGE_MODULE_DIRS: + sh: | + { + printf '.\n' + + find ./extra -name go.mod -print | + while IFS= read -r mod; do + dirname "$mod" + done + } | sort -u | tr '\n' ' ' + +env: + # The workspace resolves unpublished local modules, including examples, + # against the current checkout instead of released module versions. + GOWORK: "{{.ROOT_DIR}}/go.work" tasks: + init-workspace: + desc: "Initialize or reset the local Go workspace" + cmds: + - bash .github/scripts/create-go-workspace.sh --with-examples + run: - desc: "Full cycle" - deps: - - lint + desc: "Run lint and race tests for all modules" + cmds: + - task: lint + - task: test-race lint: - desc: "lint" + desc: "Lint all modules" + cmds: + - for: { var: MODULE_DIRS } + task: _run-linter + vars: + WORK_DIR: "/app/{{.ITEM}}" + + test: + desc: "Run tests for all modules" + cmds: + - for: { var: MODULE_DIRS } + task: _run-tests + vars: + WORK_DIR: "{{.ITEM}}" + + test-race: + desc: "Run tests with the race detector for all modules" + cmds: + - for: { var: MODULE_DIRS } + task: _run-race-tests + vars: + WORK_DIR: "{{.ITEM}}" + + coverage: + desc: "Run tests and print total coverage for production modules" + cmds: + - | + set -eu + + root="{{.ROOT_DIR}}" + output="$root/coverage.out" + profiles="$(mktemp -d)" + + cleanup() { + rm -rf "$profiles" + } + trap cleanup EXIT + + printf 'mode: atomic\n' > "$output" + + profile_index=0 + + for module in {{.COVERAGE_MODULE_DIRS}}; do + echo "==> $module" + + profile="$profiles/$profile_index.out" + profile_index=$((profile_index + 1)) + + ( + cd "$root/$module" + + GOWORK="$root/go.work" go test \ + -mod=readonly \ + -count=1 \ + -covermode=atomic \ + -coverprofile="$profile" \ + ./... + ) + + tail -n +2 "$profile" >> "$output" + done + + GOWORK="$root/go.work" go tool cover -func="$output" | grep total + + coverage:html: + desc: "Generate HTML coverage report" + deps: + - coverage + cmds: + - go tool cover -html=coverage.out -o coverage.html + + _run-tests: + internal: true + requires: + vars: + - WORK_DIR + dir: "{{.ROOT_DIR}}/{{.WORK_DIR}}" + cmds: + - echo "==> {{.WORK_DIR}}" + - go test -mod=readonly -v -count=1 ./... + + _run-race-tests: + internal: true + requires: + vars: + - WORK_DIR + dir: "{{.ROOT_DIR}}/{{.WORK_DIR}}" + cmds: + - echo "==> {{.WORK_DIR}}" + - go test -mod=readonly -race -v -count=1 ./... + + _run-linter: + internal: true + requires: + vars: + - WORK_DIR cmds: - - docker run --rm + - echo "==> {{.WORK_DIR}}" + - mkdir -p "{{.GOSUMCACHE}}" + - > + docker run --rm -u $(id -u):$(id -g) - -v {{.PWD}}:/app - -v {{.GOCACHE_ENV}}:/go/pkg/mod - -e GOCACHE=/go/pkg/mod - -e GOLANGCI_LINT_CACHE=/go/pkg/mod - -w /app + -v {{.ROOT_DIR}}:/app + -v {{.GOMODCACHE}}:/go/pkg/mod + -v {{.GOSUMCACHE}}:/go/pkg/sumdb + -v {{.GOCACHE}}:/go/build-cache + -e GOMODCACHE=/go/pkg/mod + -e GOCACHE=/go/build-cache + -e GOLANGCI_LINT_CACHE=/go/build-cache/golangci-lint + -e GOWORK=/app/go.work + -w {{.WORK_DIR}} golangci/golangci-lint:{{.LINTER_VER}} - golangci-lint run --timeout 5m --fix --modules-download-mode vendor - vars: - GOCACHE_ENV: - sh: go env GOCACHE - PWD: - sh: pwd + golangci-lint run --timeout 5m --modules-download-mode readonly diff --git a/advisory.go b/advisory.go new file mode 100644 index 0000000..369c236 --- /dev/null +++ b/advisory.go @@ -0,0 +1,49 @@ +package xpg + +import ( + "context" + "errors" + "fmt" + + "github.com/jackc/pgx/v5" +) + +const ( + advisoryXactLockSQL = "SELECT pg_advisory_xact_lock($1)" + tryAdvisoryXactLockSQL = "SELECT pg_try_advisory_xact_lock($1)" +) + +// AdvisoryXactLock acquires an exclusive transaction-level advisory lock. +// +// The call blocks until the lock is acquired or ctx is canceled. PostgreSQL +// releases the lock automatically when tx is committed or rolled back. +func AdvisoryXactLock(ctx context.Context, tx pgx.Tx, key int64) error { + if tx == nil { + return errors.New("xpg: transaction is nil") + } + + if _, err := tx.Exec(ctx, advisoryXactLockSQL, key); err != nil { + return fmt.Errorf("xpg: acquire transaction advisory lock: %w", err) + } + + return nil +} + +// TryAdvisoryXactLock attempts to acquire an exclusive transaction-level +// advisory lock without waiting. +// +// PostgreSQL releases an acquired lock automatically when tx is committed or +// rolled back. +func TryAdvisoryXactLock(ctx context.Context, tx pgx.Tx, key int64) (bool, error) { + if tx == nil { + return false, errors.New("xpg: transaction is nil") + } + + var acquired bool + + if err := tx.QueryRow(ctx, tryAdvisoryXactLockSQL, key).Scan(&acquired); err != nil { + return false, fmt.Errorf("xpg: try transaction advisory lock: %w", err) + } + + return acquired, nil +} diff --git a/advisory_test.go b/advisory_test.go new file mode 100644 index 0000000..24cbcb8 --- /dev/null +++ b/advisory_test.go @@ -0,0 +1,232 @@ +package xpg + +import ( + "context" + "errors" + "testing" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgconn" +) + +func TestAdvisoryXactLockNilTx(t *testing.T) { + t.Parallel() + + err := AdvisoryXactLock( + t.Context(), + nil, + 1, + ) + + assertErrorMessage(t, err, "xpg: transaction is nil") +} + +func TestAdvisoryXactLock(t *testing.T) { + t.Parallel() + + tx := &advisoryTestTx{} + + err := AdvisoryXactLock( + t.Context(), + tx, + 42, + ) + if err != nil { + t.Fatalf("AdvisoryXactLock returned an error: %v", err) + } + + if tx.execSQL != advisoryXactLockSQL { + t.Fatalf( + "SQL = %q, want %q", + tx.execSQL, + advisoryXactLockSQL, + ) + } + + assertAdvisoryKey(t, tx.execArgs, 42) +} + +func TestAdvisoryXactLockPreservesError(t *testing.T) { + t.Parallel() + + expectedErr := errors.New("lock failed") + + err := AdvisoryXactLock( + t.Context(), + &advisoryTestTx{ + execErr: expectedErr, + }, + 1, + ) + if !errors.Is(err, expectedErr) { + t.Fatalf("original error was not preserved: %v", err) + } +} + +func TestTryAdvisoryXactLockNilTx(t *testing.T) { + t.Parallel() + + _, err := TryAdvisoryXactLock( + t.Context(), + nil, + 1, + ) + + assertErrorMessage(t, err, "xpg: transaction is nil") +} + +func TestTryAdvisoryXactLock(t *testing.T) { + t.Parallel() + + for _, acquired := range []bool{false, true} { + t.Run( + map[bool]string{ + false: "not acquired", + true: "acquired", + }[acquired], + func(t *testing.T) { + t.Parallel() + + tx := &advisoryTestTx{ + row: advisoryTestRow{ + acquired: acquired, + }, + } + + got, err := TryAdvisoryXactLock( + t.Context(), + tx, + 42, + ) + if err != nil { + t.Fatalf( + "TryAdvisoryXactLock returned an error: %v", + err, + ) + } + + if got != acquired { + t.Fatalf( + "acquired = %v, want %v", + got, + acquired, + ) + } + + if tx.querySQL != tryAdvisoryXactLockSQL { + t.Fatalf( + "SQL = %q, want %q", + tx.querySQL, + tryAdvisoryXactLockSQL, + ) + } + + assertAdvisoryKey(t, tx.queryArgs, 42) + }, + ) + } +} + +func TestTryAdvisoryXactLockPreservesError(t *testing.T) { + t.Parallel() + + expectedErr := errors.New("try lock failed") + + _, err := TryAdvisoryXactLock( + t.Context(), + &advisoryTestTx{ + row: advisoryTestRow{ + err: expectedErr, + }, + }, + 1, + ) + if !errors.Is(err, expectedErr) { + t.Fatalf("original error was not preserved: %v", err) + } +} + +type advisoryTestTx struct { + pgx.Tx + + execSQL string + execArgs []any + execErr error + + querySQL string + queryArgs []any + row pgx.Row +} + +func (tx *advisoryTestTx) Exec( + _ context.Context, + sql string, + args ...any, +) (pgconn.CommandTag, error) { + tx.execSQL = sql + tx.execArgs = args + + return pgconn.CommandTag{}, tx.execErr +} + +func (tx *advisoryTestTx) QueryRow( + _ context.Context, + sql string, + args ...any, +) pgx.Row { + tx.querySQL = sql + tx.queryArgs = args + + return tx.row +} + +type advisoryTestRow struct { + acquired bool + err error +} + +func (row advisoryTestRow) Scan(dest ...any) error { + if row.err != nil { + return row.err + } + + acquired, ok := dest[0].(*bool) + if !ok { + return errors.New("unexpected destination type") + } + + *acquired = row.acquired + + return nil +} + +func assertAdvisoryKey( + t *testing.T, + args []any, + want int64, +) { + t.Helper() + + if len(args) != 1 { + t.Fatalf( + "argument count = %d, want 1", + len(args), + ) + } + + key, ok := args[0].(int64) + if !ok { + t.Fatalf( + "argument type = %T, want int64", + args[0], + ) + } + + if key != want { + t.Fatalf( + "key = %d, want %d", + key, + want, + ) + } +} diff --git a/cluster/cluster.go b/cluster/cluster.go new file mode 100644 index 0000000..dbe925d --- /dev/null +++ b/cluster/cluster.go @@ -0,0 +1,191 @@ +package cluster + +import ( + "errors" + "fmt" + "maps" + "slices" + "sync" + + "github.com/mkbeh/xpg" +) + +// ID identifies one logical PostgreSQL cluster. +type ID string + +// Config configures a Cluster from independently created pools. +// +// ID and Labels are optional metadata. New takes ownership of Primary and +// Replicas only after it returns successfully. Cluster.Close closes all owned +// pools. +type Config struct { + ID ID + Labels map[string]string + + Primary *xpg.Pool + Replicas []*xpg.Pool + Selector ReplicaSelector +} + +// Cluster represents a logical PostgreSQL cluster composed of an optional +// primary pool and zero or more replica pools. +// +// A deployment with one PostgreSQL endpoint is represented by a Cluster with +// one Primary and no Replicas. A read-only deployment may omit Primary and +// configure only Replicas. +// +// Cluster does not inspect SQL, retry failed queries, promote replicas, or +// discover PostgreSQL nodes. Those responsibilities remain with the caller or +// the surrounding high-availability infrastructure. +type Cluster struct { + id ID + labels map[string]string + + primary *xpg.Pool + replicas []*xpg.Pool + + metadata replicaMetadata + selector ReplicaSelector + + closeOnce sync.Once +} + +// New creates a Cluster from independently configured pools. +// +// At least one pool is required. When Selector is nil, replicas are selected +// using round-robin. +func New(config Config) (*Cluster, error) { + if config.Primary != nil && config.Primary.Raw() == nil { + return nil, errors.New("xpg/cluster: primary pool is invalid") + } + + if config.Primary == nil && len(config.Replicas) == 0 { + return nil, errors.New("xpg/cluster: at least one pool is required") + } + + if err := validateLabels(config.Labels); err != nil { + return nil, fmt.Errorf("xpg/cluster: %w", err) + } + + replicas := slices.Clone(config.Replicas) + metadata := make(replicaMetadata, len(replicas)) + + for index, replica := range replicas { + if replica == nil || replica.Raw() == nil { + return nil, fmt.Errorf("xpg/cluster: replica %d is invalid", index) + } + + metadata[index] = ReplicaInfo{ + name: replica.Name(), + labels: replica.Labels(), + } + } + + selector := config.Selector + if selector == nil { + selector = RoundRobinSelector() + } + + return &Cluster{ + id: config.ID, + labels: cloneLabels(config.Labels), + primary: config.Primary, + replicas: replicas, + metadata: metadata, + selector: selector, + }, nil +} + +// ID returns the logical cluster ID. +func (c *Cluster) ID() ID { + if c == nil { + return "" + } + + return c.id +} + +// Label returns one cluster label without allocating a copy of all labels. +func (c *Cluster) Label(key string) (string, bool) { + if c == nil { + return "", false + } + + value, ok := c.labels[key] + + return value, ok +} + +// Labels returns a defensive copy of cluster labels. +func (c *Cluster) Labels() map[string]string { + if c == nil { + return nil + } + + return cloneLabels(c.labels) +} + +// Primary returns the primary pool owned by the cluster. +// +// Primary returns nil when no primary is configured. The returned pool is +// borrowed and must not be closed separately. +func (c *Cluster) Primary() *xpg.Pool { + if c == nil { + return nil + } + + return c.primary +} + +// ReplicaCount returns the number of replicas registered in the cluster. +func (c *Cluster) ReplicaCount() int { + if c == nil { + return 0 + } + + return len(c.replicas) +} + +// ReplicaAt returns the replica at index in registration order. +// +// The returned pool is borrowed and must not be closed separately. ReplicaAt +// panics when c is nil or index is out of range. +func (c *Cluster) ReplicaAt(index int) *xpg.Pool { + return c.replicas[index] +} + +// Close closes replicas in reverse registration order, then closes the primary +// when one is configured. Close is safe to call multiple times. +func (c *Cluster) Close() { + if c == nil { + return + } + + c.closeOnce.Do(func() { + for _, replica := range slices.Backward(c.replicas) { + replica.Close() + } + + if c.primary != nil { + c.primary.Close() + } + }) +} + +func cloneLabels(labels map[string]string) map[string]string { + if len(labels) == 0 { + return nil + } + + return maps.Clone(labels) +} + +func validateLabels(labels map[string]string) error { + for key := range labels { + if key == "" { + return errors.New("label key must not be empty") + } + } + + return nil +} diff --git a/cluster/cluster_test.go b/cluster/cluster_test.go new file mode 100644 index 0000000..403d64b --- /dev/null +++ b/cluster/cluster_test.go @@ -0,0 +1,320 @@ +package cluster + +import ( + "context" + "testing" + + "github.com/mkbeh/xpg" +) + +func TestNewRequiresPool(t *testing.T) { + t.Parallel() + + cluster, err := New(Config{}) + if err == nil { + cluster.Close() + t.Fatal("expected error") + } + + if got, want := err.Error(), "xpg/cluster: at least one pool is required"; got != want { + t.Fatalf("error = %q, want %q", got, want) + } +} + +func TestNewRejectsInvalidPrimary(t *testing.T) { + t.Parallel() + + cluster, err := New(Config{ + Primary: &xpg.Pool{}, + }) + if err == nil { + cluster.Close() + t.Fatal("expected error") + } + + if got, want := err.Error(), "xpg/cluster: primary pool is invalid"; got != want { + t.Fatalf("error = %q, want %q", got, want) + } +} + +func TestNewRejectsInvalidReplica(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + replica *xpg.Pool + }{ + { + name: "nil", + replica: nil, + }, + { + name: "zero value", + replica: &xpg.Pool{}, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + + cluster, err := New(Config{ + Replicas: []*xpg.Pool{test.replica}, + }) + if err == nil { + cluster.Close() + t.Fatal("expected error") + } + + if got, want := err.Error(), "xpg/cluster: replica 0 is invalid"; got != want { + t.Fatalf("error = %q, want %q", got, want) + } + }) + } +} + +func TestNewLabels(t *testing.T) { + t.Parallel() + + primary := newTestPool(t, "primary", nil) + labels := map[string]string{ + "environment": "", + " ": "whitespace-key", + } + + cluster, err := New(Config{ + ID: "orders", + Labels: labels, + Primary: primary, + }) + if err != nil { + t.Fatalf("New() error = %v", err) + } + t.Cleanup(cluster.Close) + + labels["environment"] = "changed" + labels["new"] = "value" + + if got, want := cluster.ID(), ID("orders"); got != want { + t.Fatalf("ID() = %q, want %q", got, want) + } + + if got, ok := cluster.Label("environment"); !ok || got != "" { + t.Fatalf("Label(environment) = %q, %v; want empty value, true", got, ok) + } + + if got, ok := cluster.Label(" "); !ok || got != "whitespace-key" { + t.Fatalf("Label(whitespace) = %q, %v; want %q, true", got, ok, "whitespace-key") + } + + if _, ok := cluster.Label("new"); ok { + t.Fatal("cluster labels changed after input map mutation") + } + + cloned := cluster.Labels() + cloned["environment"] = "mutated" + delete(cloned, " ") + + if got, _ := cluster.Label("environment"); got != "" { + t.Fatalf("Label(environment) after Labels mutation = %q, want empty value", got) + } + + if got, ok := cluster.Label(" "); !ok || got != "whitespace-key" { + t.Fatalf("Label(whitespace) after Labels mutation = %q, %v; want %q, true", got, ok, "whitespace-key") + } +} + +func TestNewRejectsEmptyLabelKey(t *testing.T) { + t.Parallel() + + primary := newTestPool(t, "primary", nil) + + cluster, err := New(Config{ + Labels: map[string]string{ + "": "value", + }, + Primary: primary, + }) + if err == nil { + cluster.Close() + t.Fatal("expected error") + } + + if got, want := err.Error(), "xpg/cluster: label key must not be empty"; got != want { + t.Fatalf("error = %q, want %q", got, want) + } +} + +func TestNewClonesReplicaSlice(t *testing.T) { + t.Parallel() + + replicaA := newTestPool(t, "replica-a", nil) + replicaB := newTestPool(t, "replica-b", nil) + replicas := []*xpg.Pool{replicaA} + + cluster, err := New(Config{ + Replicas: replicas, + }) + if err != nil { + t.Fatalf("New() error = %v", err) + } + t.Cleanup(cluster.Close) + + replicas[0] = replicaB + + if got := cluster.ReplicaAt(0); got != replicaA { + t.Fatalf("ReplicaAt(0) = %p, want original replica %p", got, replicaA) + } +} + +func TestNewAllowsDuplicatePools(t *testing.T) { + t.Parallel() + + pool := newTestPool(t, "shared", nil) + + cluster, err := New(Config{ + Primary: pool, + Replicas: []*xpg.Pool{ + pool, + pool, + }, + }) + if err != nil { + t.Fatalf("New() error = %v", err) + } + t.Cleanup(cluster.Close) + + if cluster.Primary() != pool { + t.Fatal("Primary() did not return the configured pool") + } + + if got, want := cluster.ReplicaCount(), 2; got != want { + t.Fatalf("ReplicaCount() = %d, want %d", got, want) + } + + if cluster.ReplicaAt(0) != pool || cluster.ReplicaAt(1) != pool { + t.Fatal("duplicate topology entries were not preserved") + } +} + +func TestNewCapturesReplicaMetadata(t *testing.T) { + t.Parallel() + + replica := newTestPool(t, "replica-a", map[string]string{ + "region": "eu", + "role": "read", + }) + + selectorCalls := 0 + selector := ReplicaSelectorFunc(func(_ context.Context, replicas ReplicaSet) (int, error) { + selectorCalls++ + + if got, want := replicas.Len(), 1; got != want { + t.Fatalf("replicas.Len() = %d, want %d", got, want) + } + + info := replicas.At(0) + + if got, want := info.Name(), "replica-a"; got != want { + t.Fatalf("replica name = %q, want %q", got, want) + } + + if got, ok := info.Label("region"); !ok || got != "eu" { + t.Fatalf("region label = %q, %v; want %q, true", got, ok, "eu") + } + + labels := info.Labels() + labels["region"] = "mutated" + + if got, _ := info.Label("region"); got != "eu" { + t.Fatalf("replica metadata mutated through Labels(): got %q", got) + } + + return 0, nil + }) + + cluster, err := New(Config{ + Replicas: []*xpg.Pool{replica}, + Selector: selector, + }) + if err != nil { + t.Fatalf("New() error = %v", err) + } + t.Cleanup(cluster.Close) + + resolved, err := cluster.ReadPool(t.Context(), ReadReplicaRequired) + if err != nil { + t.Fatalf("ReadPool() error = %v", err) + } + + if resolved != replica { + t.Fatalf("ReadPool() = %p, want %p", resolved, replica) + } + + if selectorCalls != 1 { + t.Fatalf("selector calls = %d, want 1", selectorCalls) + } +} + +func TestClusterNilReceiverMetadata(t *testing.T) { + t.Parallel() + + var cluster *Cluster + + if got := cluster.ID(); got != "" { + t.Fatalf("ID() = %q, want empty", got) + } + + if value, ok := cluster.Label("key"); ok || value != "" { + t.Fatalf("Label() = %q, %v; want empty, false", value, ok) + } + + if labels := cluster.Labels(); labels != nil { + t.Fatalf("Labels() = %v, want nil", labels) + } + + if primary := cluster.Primary(); primary != nil { + t.Fatalf("Primary() = %p, want nil", primary) + } + + if got := cluster.ReplicaCount(); got != 0 { + t.Fatalf("ReplicaCount() = %d, want 0", got) + } + + cluster.Close() +} + +func TestReplicaAtPanicsOutOfRange(t *testing.T) { + t.Parallel() + + replica := newTestPool(t, "replica", nil) + cluster := newTestCluster(t, Config{ + Replicas: []*xpg.Pool{replica}, + }) + + defer func() { + if recover() == nil { + t.Fatal("ReplicaAt() did not panic") + } + }() + + _ = cluster.ReplicaAt(1) +} + +func TestCloseIsIdempotent(t *testing.T) { + t.Parallel() + + primary := newTestPool(t, "primary", nil) + replica := newTestPool(t, "replica", nil) + + cluster, err := New(Config{ + Primary: primary, + Replicas: []*xpg.Pool{replica}, + }) + if err != nil { + t.Fatalf("New() error = %v", err) + } + + cluster.Close() + cluster.Close() +} diff --git a/cluster/doc.go b/cluster/doc.go new file mode 100644 index 0000000..d47cd57 --- /dev/null +++ b/cluster/doc.go @@ -0,0 +1,3 @@ +// Package cluster provides primary/replica routing for PostgreSQL connection +// pools. +package cluster diff --git a/cluster/errors.go b/cluster/errors.go new file mode 100644 index 0000000..03adc17 --- /dev/null +++ b/cluster/errors.go @@ -0,0 +1,13 @@ +package cluster + +import "errors" + +var ( + // ErrNoPrimary is returned when an operation requires a primary but none + // is configured. + ErrNoPrimary = errors.New("xpg/cluster: no primary available") + + // ErrNoReplica is returned when an operation requires a replica but none + // can be selected. + ErrNoReplica = errors.New("xpg/cluster: no replica available") +) diff --git a/cluster/helpers_test.go b/cluster/helpers_test.go new file mode 100644 index 0000000..092a235 --- /dev/null +++ b/cluster/helpers_test.go @@ -0,0 +1,50 @@ +package cluster + +import ( + "testing" + + "github.com/jackc/pgx/v5/pgxpool" + "github.com/mkbeh/xpg" +) + +const testDatabaseURL = "postgres://postgres:postgres@127.0.0.1:1/postgres?sslmode=disable" //nolint:gosec // Test-only DSN with non-production credentials. + +func newTestPool(t *testing.T, name string, labels map[string]string) *xpg.Pool { + t.Helper() + + config, err := pgxpool.ParseConfig(testDatabaseURL) + if err != nil { + t.Fatalf("pgxpool.ParseConfig() error = %v", err) + } + + config.MinConns = 0 + config.MaxConns = 1 + + options := []xpg.Option{ + xpg.WithName(name), + } + + if len(labels) != 0 { + options = append(options, xpg.WithLabels(labels)) + } + + pool, err := xpg.New(t.Context(), config, options...) + if err != nil { + t.Fatalf("xpg.New() error = %v", err) + } + t.Cleanup(pool.Close) + + return pool +} + +func newTestCluster(t *testing.T, config Config) *Cluster { + t.Helper() + + cluster, err := New(config) + if err != nil { + t.Fatalf("New() error = %v", err) + } + t.Cleanup(cluster.Close) + + return cluster +} diff --git a/cluster/resolver.go b/cluster/resolver.go new file mode 100644 index 0000000..0d0a902 --- /dev/null +++ b/cluster/resolver.go @@ -0,0 +1,124 @@ +package cluster + +import ( + "context" + "errors" + "fmt" + + "github.com/mkbeh/xpg" +) + +// ReadPolicy defines where a Cluster resolves a read operation. +type ReadPolicy uint8 + +const ( + // ReadPrimary always resolves reads to the primary pool. + ReadPrimary ReadPolicy = iota + + // ReadReplicaPreferred prefers a replica and falls back to the primary when + // no replica can be selected. + ReadReplicaPreferred + + // ReadReplicaRequired requires a replica and returns ErrNoReplica when none + // can be selected. + ReadReplicaRequired +) + +const ( + readPolicyPrimary = "primary" + readPolicyReplicaPreferred = "replica_preferred" + readPolicyReplicaRequired = "replica_required" +) + +// ParseReadPolicy parses a ReadPolicy from its string representation. +func ParseReadPolicy(value string) (ReadPolicy, error) { + switch value { + case readPolicyPrimary: + return ReadPrimary, nil + case readPolicyReplicaPreferred: + return ReadReplicaPreferred, nil + case readPolicyReplicaRequired: + return ReadReplicaRequired, nil + default: + return 0, fmt.Errorf( + "xpg/cluster: unknown read policy %q", + value, + ) + } +} + +// String returns the string representation of the read policy. +func (policy ReadPolicy) String() string { + switch policy { + case ReadPrimary: + return readPolicyPrimary + case ReadReplicaPreferred: + return readPolicyReplicaPreferred + case ReadReplicaRequired: + return readPolicyReplicaRequired + default: + return "unknown" + } +} + +// ReadPool returns a pool according to policy. +// +// ReadReplicaPreferred falls back to the primary only when no replica can be +// selected. Other selector errors are returned to the caller. +func (c *Cluster) ReadPool(ctx context.Context, policy ReadPolicy) (*xpg.Pool, error) { + if c == nil { + return nil, errors.New("xpg/cluster: cluster is nil") + } + + switch policy { + case ReadPrimary: + return c.resolvePrimary() + + case ReadReplicaPreferred: + replica, err := c.resolveReplica(ctx) + if err == nil { + return replica, nil + } + + if !errors.Is(err, ErrNoReplica) { + return nil, err + } + + return c.resolvePrimary() + + case ReadReplicaRequired: + return c.resolveReplica(ctx) + + default: + return nil, fmt.Errorf("xpg/cluster: unsupported read policy %d", policy) + } +} + +func (c *Cluster) resolvePrimary() (*xpg.Pool, error) { + if c.primary == nil { + return nil, ErrNoPrimary + } + + return c.primary, nil +} + +func (c *Cluster) resolveReplica(ctx context.Context) (*xpg.Pool, error) { + if len(c.replicas) == 0 { + return nil, ErrNoReplica + } + + index, err := c.selector.Select(ctx, c.metadata) + if err != nil { + return nil, fmt.Errorf("xpg/cluster: select replica: %w", err) + } + + if index < 0 || index >= len(c.replicas) { + return nil, fmt.Errorf( + "xpg/cluster: replica selector returned invalid index %d for %d replicas", + index, + len(c.replicas), + ) + } + + return c.replicas[index], nil +} diff --git a/cluster/routing_test.go b/cluster/routing_test.go new file mode 100644 index 0000000..bd08329 --- /dev/null +++ b/cluster/routing_test.go @@ -0,0 +1,390 @@ +package cluster + +import ( + "context" + "errors" + "fmt" + "testing" + + "github.com/mkbeh/xpg" +) + +func TestParseReadPolicy(t *testing.T) { + t.Parallel() + + tests := []struct { + value string + want ReadPolicy + }{ + {value: "primary", want: ReadPrimary}, + {value: "replica_preferred", want: ReadReplicaPreferred}, + {value: "replica_required", want: ReadReplicaRequired}, + } + + for _, test := range tests { + t.Run(test.value, func(t *testing.T) { + t.Parallel() + + got, err := ParseReadPolicy(test.value) + if err != nil { + t.Fatalf("ParseReadPolicy() error = %v", err) + } + + if got != test.want { + t.Fatalf("ParseReadPolicy(%q) = %v, want %v", test.value, got, test.want) + } + }) + } +} + +func TestParseReadPolicyRejectsUnknown(t *testing.T) { + t.Parallel() + + _, err := ParseReadPolicy("nearest") + if err == nil { + t.Fatal("expected error") + } + + if got, want := err.Error(), `xpg/cluster: unknown read policy "nearest"`; got != want { + t.Fatalf("error = %q, want %q", got, want) + } +} + +func TestReadPolicyString(t *testing.T) { + t.Parallel() + + tests := []struct { + policy ReadPolicy + want string + }{ + {policy: ReadPrimary, want: "primary"}, + {policy: ReadReplicaPreferred, want: "replica_preferred"}, + {policy: ReadReplicaRequired, want: "replica_required"}, + {policy: ReadPolicy(255), want: "unknown"}, + } + + for _, test := range tests { + if got := test.policy.String(); got != test.want { + t.Fatalf("ReadPolicy(%d).String() = %q, want %q", test.policy, got, test.want) + } + } +} + +func TestReadPoolNilCluster(t *testing.T) { + t.Parallel() + + var cluster *Cluster + + pool, err := cluster.ReadPool(t.Context(), ReadPrimary) + if pool != nil { + t.Fatalf("ReadPool() pool = %p, want nil", pool) + } + + if err == nil { + t.Fatal("expected error") + } + + if got, want := err.Error(), "xpg/cluster: cluster is nil"; got != want { + t.Fatalf("error = %q, want %q", got, want) + } +} + +func TestReadPoolPrimary(t *testing.T) { + t.Parallel() + + primary := newTestPool(t, "primary", nil) + cluster := newTestCluster(t, Config{Primary: primary}) + + got, err := cluster.ReadPool(t.Context(), ReadPrimary) + if err != nil { + t.Fatalf("ReadPool() error = %v", err) + } + + if got != primary { + t.Fatalf("ReadPool() = %p, want %p", got, primary) + } +} + +func TestReadPoolPrimaryWithoutPrimary(t *testing.T) { + t.Parallel() + + replica := newTestPool(t, "replica", nil) + cluster := newTestCluster(t, Config{ + Replicas: []*xpg.Pool{replica}, + }) + + pool, err := cluster.ReadPool(t.Context(), ReadPrimary) + if pool != nil { + t.Fatalf("ReadPool() pool = %p, want nil", pool) + } + + if !errors.Is(err, ErrNoPrimary) { + t.Fatalf("ReadPool() error = %v, want ErrNoPrimary", err) + } +} + +func TestReadPoolReplicaPreferredUsesReplica(t *testing.T) { + t.Parallel() + + primary := newTestPool(t, "primary", nil) + replica := newTestPool(t, "replica", nil) + cluster := newTestCluster(t, Config{ + Primary: primary, + Replicas: []*xpg.Pool{replica}, + }) + + got, err := cluster.ReadPool(t.Context(), ReadReplicaPreferred) + if err != nil { + t.Fatalf("ReadPool() error = %v", err) + } + + if got != replica { + t.Fatalf("ReadPool() = %p, want replica %p", got, replica) + } +} + +func TestReadPoolReplicaPreferredFallsBackWithoutReplicas(t *testing.T) { + t.Parallel() + + primary := newTestPool(t, "primary", nil) + cluster := newTestCluster(t, Config{Primary: primary}) + + got, err := cluster.ReadPool(t.Context(), ReadReplicaPreferred) + if err != nil { + t.Fatalf("ReadPool() error = %v", err) + } + + if got != primary { + t.Fatalf("ReadPool() = %p, want primary %p", got, primary) + } +} + +func TestReadPoolReplicaPreferredFallsBackOnErrNoReplica(t *testing.T) { + t.Parallel() + + primary := newTestPool(t, "primary", nil) + replica := newTestPool(t, "replica", nil) + selector := ReplicaSelectorFunc(func(context.Context, ReplicaSet) (int, error) { + return -1, errors.Join(errors.New("temporarily unavailable"), ErrNoReplica) + }) + + cluster := newTestCluster(t, Config{ + Primary: primary, + Replicas: []*xpg.Pool{replica}, + Selector: selector, + }) + + got, err := cluster.ReadPool(t.Context(), ReadReplicaPreferred) + if err != nil { + t.Fatalf("ReadPool() error = %v", err) + } + + if got != primary { + t.Fatalf("ReadPool() = %p, want primary %p", got, primary) + } +} + +func TestReadPoolReplicaPreferredDoesNotFallbackOnSelectorError(t *testing.T) { + t.Parallel() + + errSelector := errors.New("selector failed") + primary := newTestPool(t, "primary", nil) + replica := newTestPool(t, "replica", nil) + selector := ReplicaSelectorFunc(func(context.Context, ReplicaSet) (int, error) { + return -1, errSelector + }) + + cluster := newTestCluster(t, Config{ + Primary: primary, + Replicas: []*xpg.Pool{replica}, + Selector: selector, + }) + + pool, err := cluster.ReadPool(t.Context(), ReadReplicaPreferred) + if pool != nil { + t.Fatalf("ReadPool() pool = %p, want nil", pool) + } + + if !errors.Is(err, errSelector) { + t.Fatalf("ReadPool() error = %v, want selector error", err) + } +} + +func TestReadPoolReplicaPreferredWithoutPrimary(t *testing.T) { + t.Parallel() + + replica := newTestPool(t, "replica", nil) + selector := ReplicaSelectorFunc(func(context.Context, ReplicaSet) (int, error) { + return -1, ErrNoReplica + }) + + cluster := newTestCluster(t, Config{ + Replicas: []*xpg.Pool{replica}, + Selector: selector, + }) + + pool, err := cluster.ReadPool(t.Context(), ReadReplicaPreferred) + if pool != nil { + t.Fatalf("ReadPool() pool = %p, want nil", pool) + } + + if !errors.Is(err, ErrNoPrimary) { + t.Fatalf("ReadPool() error = %v, want ErrNoPrimary", err) + } +} + +func TestReadPoolReplicaRequired(t *testing.T) { + t.Parallel() + + replica := newTestPool(t, "replica", nil) + cluster := newTestCluster(t, Config{ + Replicas: []*xpg.Pool{replica}, + }) + + got, err := cluster.ReadPool(t.Context(), ReadReplicaRequired) + if err != nil { + t.Fatalf("ReadPool() error = %v", err) + } + + if got != replica { + t.Fatalf("ReadPool() = %p, want %p", got, replica) + } +} + +func TestReadPoolReplicaRequiredWithoutReplicas(t *testing.T) { + t.Parallel() + + primary := newTestPool(t, "primary", nil) + cluster := newTestCluster(t, Config{Primary: primary}) + + pool, err := cluster.ReadPool(t.Context(), ReadReplicaRequired) + if pool != nil { + t.Fatalf("ReadPool() pool = %p, want nil", pool) + } + + if !errors.Is(err, ErrNoReplica) { + t.Fatalf("ReadPool() error = %v, want ErrNoReplica", err) + } +} + +func TestReadPoolDefaultSelectorRoundRobin(t *testing.T) { + t.Parallel() + + replicaA := newTestPool(t, "replica-a", nil) + replicaB := newTestPool(t, "replica-b", nil) + cluster := newTestCluster(t, Config{ + Replicas: []*xpg.Pool{replicaA, replicaB}, + }) + + want := []*xpg.Pool{ + replicaA, + replicaB, + replicaA, + replicaB, + } + + for call, wantPool := range want { + got, err := cluster.ReadPool(t.Context(), ReadReplicaRequired) + if err != nil { + t.Fatalf("ReadPool() call %d error = %v", call, err) + } + + if got != wantPool { + t.Fatalf("ReadPool() call %d = %p, want %p", call, got, wantPool) + } + } +} + +func TestReadPoolRejectsSelectorIndex(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + index int + }{ + {name: "negative", index: -1}, + {name: "past end", index: 1}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + + replica := newTestPool(t, "replica", nil) + selector := ReplicaSelectorFunc(func(context.Context, ReplicaSet) (int, error) { + return test.index, nil + }) + + cluster := newTestCluster(t, Config{ + Replicas: []*xpg.Pool{replica}, + Selector: selector, + }) + + pool, err := cluster.ReadPool(t.Context(), ReadReplicaRequired) + if pool != nil { + t.Fatalf("ReadPool() pool = %p, want nil", pool) + } + + if err == nil { + t.Fatal("expected error") + } + + want := fmt.Sprintf( + "xpg/cluster: replica selector returned invalid index %d for 1 replicas", + test.index, + ) + + if got := err.Error(); got != want { + t.Fatalf("error = %q, want %q", got, want) + } + }) + } +} + +func TestReadPoolPreservesSelectorError(t *testing.T) { + t.Parallel() + + errSelector := errors.New("boom") + replica := newTestPool(t, "replica", nil) + selector := ReplicaSelectorFunc(func(context.Context, ReplicaSet) (int, error) { + return -1, errSelector + }) + + cluster := newTestCluster(t, Config{ + Replicas: []*xpg.Pool{replica}, + Selector: selector, + }) + + pool, err := cluster.ReadPool(t.Context(), ReadReplicaRequired) + if pool != nil { + t.Fatalf("ReadPool() pool = %p, want nil", pool) + } + + if !errors.Is(err, errSelector) { + t.Fatalf("ReadPool() error = %v, want wrapped selector error", err) + } + + if got, want := err.Error(), "xpg/cluster: select replica: boom"; got != want { + t.Fatalf("error = %q, want %q", got, want) + } +} + +func TestReadPoolRejectsUnsupportedPolicy(t *testing.T) { + t.Parallel() + + primary := newTestPool(t, "primary", nil) + cluster := newTestCluster(t, Config{Primary: primary}) + + pool, err := cluster.ReadPool(t.Context(), ReadPolicy(255)) + if pool != nil { + t.Fatalf("ReadPool() pool = %p, want nil", pool) + } + + if err == nil { + t.Fatal("expected error") + } + + if got, want := err.Error(), "xpg/cluster: unsupported read policy 255"; got != want { + t.Fatalf("error = %q, want %q", got, want) + } +} diff --git a/cluster/selector.go b/cluster/selector.go new file mode 100644 index 0000000..ab43dd6 --- /dev/null +++ b/cluster/selector.go @@ -0,0 +1,94 @@ +package cluster + +import ( + "context" + "errors" + "sync/atomic" +) + +// ReplicaInfo contains immutable metadata captured from a replica pool when the +// Cluster is created. +type ReplicaInfo struct { + name string + labels map[string]string +} + +// Name returns the logical replica pool name. +func (info ReplicaInfo) Name() string { + return info.name +} + +// Label returns one replica label without allocating a copy of all labels. +func (info ReplicaInfo) Label(key string) (string, bool) { + value, ok := info.labels[key] + + return value, ok +} + +// Labels returns a defensive copy of replica labels. +func (info ReplicaInfo) Labels() map[string]string { + return cloneLabels(info.labels) +} + +// ReplicaSet provides read-only access to replica metadata. +type ReplicaSet interface { + Len() int + At(index int) ReplicaInfo +} + +type replicaMetadata []ReplicaInfo + +func (replicas replicaMetadata) Len() int { + return len(replicas) +} + +func (replicas replicaMetadata) At(index int) ReplicaInfo { + return replicas[index] +} + +// ReplicaSelector selects one replica from the supplied metadata. +// +// Implementations must be safe for concurrent use. Select must return +// ErrNoReplica when no replica is eligible for selection. Other errors are +// propagated to the caller. A successful call must return a valid replica +// index. +type ReplicaSelector interface { + Select(ctx context.Context, replicas ReplicaSet) (index int, err error) +} + +// ReplicaSelectorFunc adapts a function to ReplicaSelector. +type ReplicaSelectorFunc func(context.Context, ReplicaSet) (int, error) + +// Select calls the wrapped selector function. +func (selector ReplicaSelectorFunc) Select(ctx context.Context, replicas ReplicaSet) (int, error) { + if selector == nil { + return -1, errors.New("xpg/cluster: replica selector function is nil") + } + + return selector(ctx, replicas) +} + +// RoundRobinSelector returns a concurrency-safe selector that distributes +// selections across replicas in registration order. +func RoundRobinSelector() ReplicaSelector { + return &roundRobinSelector{} +} + +type roundRobinSelector struct { + next atomic.Uint64 +} + +func (selector *roundRobinSelector) Select(_ context.Context, replicas ReplicaSet) (int, error) { + length := replicas.Len() + + switch length { + case 0: + return -1, ErrNoReplica + case 1: + return 0, nil + } + + next := selector.next.Add(1) - 1 + + return int(next % uint64(length)), nil +} diff --git a/cluster/selector_test.go b/cluster/selector_test.go new file mode 100644 index 0000000..348eefa --- /dev/null +++ b/cluster/selector_test.go @@ -0,0 +1,203 @@ +package cluster + +import ( + "context" + "errors" + "sync" + "testing" +) + +func TestReplicaInfoLabelsAreDefensive(t *testing.T) { + t.Parallel() + + info := ReplicaInfo{ + name: "replica", + labels: map[string]string{ + "region": "eu", + }, + } + + labels := info.Labels() + labels["region"] = "us" + labels["new"] = "value" + + if got, want := info.Name(), "replica"; got != want { + t.Fatalf("Name() = %q, want %q", got, want) + } + + if got, ok := info.Label("region"); !ok || got != "eu" { + t.Fatalf("Label(region) = %q, %v; want %q, true", got, ok, "eu") + } + + if _, ok := info.Label("new"); ok { + t.Fatal("Labels() exposed internal map") + } +} + +func TestReplicaMetadataAtPanicsOutOfRange(t *testing.T) { + t.Parallel() + + replicas := replicaMetadata{{name: "replica"}} + + defer func() { + if recover() == nil { + t.Fatal("At() did not panic") + } + }() + + _ = replicas.At(1) +} + +func TestReplicaSelectorFunc(t *testing.T) { + t.Parallel() + + type contextKey struct{} + ctx := context.WithValue(t.Context(), contextKey{}, "value") + replicas := replicaMetadata{{name: "replica"}} + + selector := ReplicaSelectorFunc(func(gotCtx context.Context, gotReplicas ReplicaSet) (int, error) { + if got := gotCtx.Value(contextKey{}); got != "value" { + t.Fatalf("context value = %v, want %q", got, "value") + } + + if got, want := gotReplicas.Len(), 1; got != want { + t.Fatalf("replicas.Len() = %d, want %d", got, want) + } + + return 0, nil + }) + + index, err := selector.Select(ctx, replicas) + if err != nil { + t.Fatalf("Select() error = %v", err) + } + + if index != 0 { + t.Fatalf("Select() index = %d, want 0", index) + } +} + +func TestReplicaSelectorFuncNil(t *testing.T) { + t.Parallel() + + var selector ReplicaSelectorFunc + + index, err := selector.Select(t.Context(), nil) + if index != -1 { + t.Fatalf("Select() index = %d, want -1", index) + } + + if err == nil { + t.Fatal("expected error") + } + + if got, want := err.Error(), "xpg/cluster: replica selector function is nil"; got != want { + t.Fatalf("error = %q, want %q", got, want) + } +} + +func TestRoundRobinSelectorEmpty(t *testing.T) { + t.Parallel() + + selector := RoundRobinSelector() + + index, err := selector.Select(t.Context(), replicaMetadata(nil)) + if index != -1 { + t.Fatalf("Select() index = %d, want -1", index) + } + + if !errors.Is(err, ErrNoReplica) { + t.Fatalf("Select() error = %v, want ErrNoReplica", err) + } +} + +func TestRoundRobinSelectorSingleReplica(t *testing.T) { + t.Parallel() + + selector := RoundRobinSelector() + replicas := replicaMetadata{{name: "replica"}} + + for range 10 { + index, err := selector.Select(t.Context(), replicas) + if err != nil { + t.Fatalf("Select() error = %v", err) + } + + if index != 0 { + t.Fatalf("Select() index = %d, want 0", index) + } + } +} + +func TestRoundRobinSelectorSequence(t *testing.T) { + t.Parallel() + + selector := RoundRobinSelector() + replicas := replicaMetadata{ + {name: "replica-a"}, + {name: "replica-b"}, + {name: "replica-c"}, + } + want := []int{0, 1, 2, 0, 1, 2, 0} + + for call, wantIndex := range want { + index, err := selector.Select(t.Context(), replicas) + if err != nil { + t.Fatalf("Select() call %d error = %v", call, err) + } + + if index != wantIndex { + t.Fatalf("Select() call %d index = %d, want %d", call, index, wantIndex) + } + } +} + +func TestRoundRobinSelectorConcurrent(t *testing.T) { + t.Parallel() + + const ( + replicaCount = 3 + callCount = 600 + ) + + selector := RoundRobinSelector() + replicas := make(replicaMetadata, replicaCount) + results := make(chan int, callCount) + ctx := t.Context() + + var waitGroup sync.WaitGroup + waitGroup.Add(callCount) + + for range callCount { + go func() { + defer waitGroup.Done() + + index, err := selector.Select(ctx, replicas) + if err != nil { + results <- -1 + return + } + + results <- index + }() + } + + waitGroup.Wait() + close(results) + + counts := make([]int, replicaCount) + + for index := range results { + if index < 0 || index >= replicaCount { + t.Fatalf("Select() returned invalid index %d", index) + } + + counts[index]++ + } + + for index, count := range counts { + if count != callCount/replicaCount { + t.Fatalf("replica %d selections = %d, want %d", index, count, callCount/replicaCount) + } + } +} diff --git a/cluster/tx.go b/cluster/tx.go new file mode 100644 index 0000000..f2d5818 --- /dev/null +++ b/cluster/tx.go @@ -0,0 +1,61 @@ +package cluster + +import ( + "context" + "errors" + + "github.com/jackc/pgx/v5" +) + +// ReadTxOptions configures a read-only transaction. +// +// AccessMode is always pgx.ReadOnly. BeginQuery and CommitQuery are not exposed +// so callers cannot override the read-only transaction semantics. +type ReadTxOptions struct { + IsoLevel pgx.TxIsoLevel + DeferrableMode pgx.TxDeferrableMode +} + +// InPrimaryTx executes fn in a transaction on the primary pool. +func (c *Cluster) InPrimaryTx( + ctx context.Context, + options pgx.TxOptions, + fn func(context.Context, pgx.Tx) error, +) error { + if c == nil { + return errors.New("xpg/cluster: cluster is nil") + } + + pool, err := c.resolvePrimary() + if err != nil { + return err + } + + return pool.InTx(ctx, options, fn) +} + +// InReadTx selects a pool according to policy and executes fn in a read-only +// transaction. +// +// The transaction remains read-only when policy resolves to the primary. +func (c *Cluster) InReadTx( + ctx context.Context, + policy ReadPolicy, + options ReadTxOptions, + fn func(context.Context, pgx.Tx) error, +) error { + pool, err := c.ReadPool(ctx, policy) + if err != nil { + return err + } + + return pool.InTx( + ctx, + pgx.TxOptions{ + IsoLevel: options.IsoLevel, + AccessMode: pgx.ReadOnly, + DeferrableMode: options.DeferrableMode, + }, + fn, + ) +} diff --git a/cluster/tx_test.go b/cluster/tx_test.go new file mode 100644 index 0000000..8478e22 --- /dev/null +++ b/cluster/tx_test.go @@ -0,0 +1,152 @@ +package cluster + +import ( + "context" + "errors" + "testing" + + "github.com/jackc/pgx/v5" + "github.com/mkbeh/xpg" +) + +func TestInPrimaryTxNilCluster(t *testing.T) { + t.Parallel() + + var cluster *Cluster + called := false + + err := cluster.InPrimaryTx( + t.Context(), + pgx.TxOptions{}, + func(context.Context, pgx.Tx) error { + called = true + return nil + }, + ) + if err == nil { + t.Fatal("expected error") + } + + if got, want := err.Error(), "xpg/cluster: cluster is nil"; got != want { + t.Fatalf("error = %q, want %q", got, want) + } + + if called { + t.Fatal("transaction callback was called") + } +} + +func TestInPrimaryTxWithoutPrimary(t *testing.T) { + t.Parallel() + + replica := newTestPool(t, "replica", nil) + cluster := newTestCluster(t, Config{ + Replicas: []*xpg.Pool{replica}, + }) + + called := false + err := cluster.InPrimaryTx( + t.Context(), + pgx.TxOptions{}, + func(context.Context, pgx.Tx) error { + called = true + return nil + }, + ) + + if !errors.Is(err, ErrNoPrimary) { + t.Fatalf("InPrimaryTx() error = %v, want ErrNoPrimary", err) + } + + if called { + t.Fatal("transaction callback was called") + } +} + +func TestInPrimaryTxContextCancellation(t *testing.T) { + t.Parallel() + + primary := newTestPool(t, "primary", nil) + cluster := newTestCluster(t, Config{Primary: primary}) + + ctx, cancel := context.WithCancel(t.Context()) + cancel() + + called := false + err := cluster.InPrimaryTx( + ctx, + pgx.TxOptions{}, + func(context.Context, pgx.Tx) error { + called = true + return nil + }, + ) + + if !errors.Is(err, context.Canceled) { + t.Fatalf("InPrimaryTx() error = %v, want context.Canceled", err) + } + + if called { + t.Fatal("transaction callback was called") + } +} + +func TestInReadTxRoutingError(t *testing.T) { + t.Parallel() + + primary := newTestPool(t, "primary", nil) + cluster := newTestCluster(t, Config{Primary: primary}) + + called := false + err := cluster.InReadTx( + t.Context(), + ReadReplicaRequired, + ReadTxOptions{}, + func(context.Context, pgx.Tx) error { + called = true + return nil + }, + ) + + if !errors.Is(err, ErrNoReplica) { + t.Fatalf("InReadTx() error = %v, want ErrNoReplica", err) + } + + if called { + t.Fatal("transaction callback was called") + } +} + +func TestInReadTxContextCancellation(t *testing.T) { + t.Parallel() + + replica := newTestPool(t, "replica", nil) + cluster := newTestCluster(t, Config{ + Replicas: []*xpg.Pool{replica}, + }) + + ctx, cancel := context.WithCancel(t.Context()) + cancel() + + called := false + err := cluster.InReadTx( + ctx, + ReadReplicaRequired, + ReadTxOptions{ + IsoLevel: pgx.Serializable, + DeferrableMode: pgx.Deferrable, + }, + func(context.Context, pgx.Tx) error { + called = true + return nil + }, + ) + + if !errors.Is(err, context.Canceled) { + t.Fatalf("InReadTx() error = %v, want context.Canceled", err) + } + + if called { + t.Fatal("transaction callback was called") + } +} diff --git a/doc.go b/doc.go new file mode 100644 index 0000000..0dec990 --- /dev/null +++ b/doc.go @@ -0,0 +1,2 @@ +// Package xpg provides PostgreSQL infrastructure utilities built on pgx. +package xpg diff --git a/errors.go b/errors.go index b923dfc..66c963d 100644 --- a/errors.go +++ b/errors.go @@ -1,82 +1,121 @@ -package postgres +package xpg import ( - "context" "errors" + "io" "net" - "github.com/jackc/pgerrcode" "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgconn" ) -type PgErrorCode int - const ( - ErrContextDeadline PgErrorCode = iota - ErrNoRows - ErrUniqViolation - ErrForeignKeyViolation - ErrSerializable - ErrOther - ErrBeginTransaction - ErrCommitTransaction - ErrNoConnection + sqlStateUniqueViolation = "23505" + sqlStateForeignKeyViolation = "23503" + sqlStateNotNullViolation = "23502" + sqlStateCheckViolation = "23514" + sqlStateSerializationFailure = "40001" + sqlStateDeadlockDetected = "40P01" + sqlStateLockNotAvailable = "55P03" + sqlStateQueryCanceled = "57014" + sqlStateConnectionExceptionClass = "08" ) -type PgError struct { - code PgErrorCode - msg string +// SQLState returns the PostgreSQL SQLSTATE code carried by err. +// It returns an empty string when the error tree does not contain a +// pgconn.PgError. +func SQLState(err error) string { + pgErr, ok := errors.AsType[*pgconn.PgError](err) + if !ok || pgErr == nil { + return "" + } + + return pgErr.SQLState() +} + +// IsNoRows reports whether err indicates that a query returned no rows. +func IsNoRows(err error) bool { + return errors.Is(err, pgx.ErrNoRows) +} + +// IsUniqueViolation reports whether err is a PostgreSQL unique_violation. +func IsUniqueViolation(err error) bool { + return SQLState(err) == sqlStateUniqueViolation +} + +// IsForeignKeyViolation reports whether err is a PostgreSQL +// foreign_key_violation. +func IsForeignKeyViolation(err error) bool { + return SQLState(err) == sqlStateForeignKeyViolation } -func (e PgError) Error() string { - return e.msg +// IsNotNullViolation reports whether err is a PostgreSQL not_null_violation. +func IsNotNullViolation(err error) bool { + return SQLState(err) == sqlStateNotNullViolation } -func (e PgError) Code() PgErrorCode { - return e.code +// IsCheckViolation reports whether err is a PostgreSQL check_violation. +func IsCheckViolation(err error) bool { + return SQLState(err) == sqlStateCheckViolation } -func NewPgError(code PgErrorCode, err error) *PgError { - return &PgError{code, err.Error()} +// IsSerializationFailure reports whether err is a PostgreSQL +// serialization_failure. +func IsSerializationFailure(err error) bool { + return SQLState(err) == sqlStateSerializationFailure } -func ConvertError(err error) *PgError { +// IsDeadlock reports whether err is a PostgreSQL deadlock_detected error. +func IsDeadlock(err error) bool { + return SQLState(err) == sqlStateDeadlockDetected +} + +// IsLockNotAvailable reports whether err is a PostgreSQL lock_not_available +// error. +func IsLockNotAvailable(err error) bool { + return SQLState(err) == sqlStateLockNotAvailable +} + +// IsQueryCanceled reports whether PostgreSQL canceled the query. +// +// Client-side context cancellation remains available through errors.Is with +// context.Canceled or context.DeadlineExceeded. +func IsQueryCanceled(err error) bool { + return SQLState(err) == sqlStateQueryCanceled +} + +// IsConnectionError reports whether err represents a PostgreSQL connection +// failure known to pgx or the Go networking stack. +func IsConnectionError(err error) bool { if err == nil { - return nil + return false } - if pgErr, ok := err.(*PgError); ok { - return pgErr + if errors.Is(err, pgconn.ErrConnClosed) || + errors.Is(err, io.EOF) || + errors.Is(err, io.ErrUnexpectedEOF) { + return true } - if ne, ok := err.(net.Error); ok { - return NewPgError(ErrNoConnection, ne) + if _, ok := errors.AsType[*pgconn.ConnectError](err); ok { + return true } - switch { - case errors.Is(err, context.Canceled), pgconn.Timeout(err): - return NewPgError(ErrContextDeadline, err) - case errors.Is(err, pgx.ErrNoRows): - return NewPgError(ErrNoRows, err) - default: - var pgErr *pgconn.PgError - if !errors.As(err, &pgErr) { - return NewPgError(ErrOther, err) - } - return NewPgError(pgCodeToError(pgErr.Code), err) + if _, ok := errors.AsType[*net.OpError](err); ok { + return true } -} -var pgCodeMap = map[string]PgErrorCode{ - pgerrcode.UniqueViolation: ErrUniqViolation, - pgerrcode.ForeignKeyViolation: ErrForeignKeyViolation, - pgerrcode.SerializationFailure: ErrSerializable, + state := SQLState(err) + + return len(state) >= 2 && + state[:2] == sqlStateConnectionExceptionClass } -func pgCodeToError(code string) PgErrorCode { - if c, ok := pgCodeMap[code]; ok { - return c - } - return ErrOther +// IsRetryableTransaction reports whether PostgreSQL aborted the transaction +// because of a serialization failure or a deadlock. +// +// The entire transaction callback must still be safe to replay. Connection +// failures are deliberately not classified as transaction-retryable. +func IsRetryableTransaction(err error) bool { + return IsSerializationFailure(err) || IsDeadlock(err) } diff --git a/errors_test.go b/errors_test.go new file mode 100644 index 0000000..a21c04e --- /dev/null +++ b/errors_test.go @@ -0,0 +1,268 @@ +package xpg + +import ( + "context" + "errors" + "fmt" + "io" + "net" + "testing" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgconn" +) + +func TestSQLState(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + err error + want string + }{ + { + name: "PostgreSQL error", + err: fmt.Errorf( + "wrapped: %w", + &pgconn.PgError{Code: sqlStateUniqueViolation}, + ), + want: sqlStateUniqueViolation, + }, + { + name: "generic error", + err: errors.New("generic"), + want: "", + }, + { + name: "nil", + err: nil, + want: "", + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + + if got := SQLState(test.err); got != test.want { + t.Fatalf( + "SQLState() = %q, want %q", + got, + test.want, + ) + } + }) + } +} + +func TestIsNoRows(t *testing.T) { + t.Parallel() + + if !IsNoRows(fmt.Errorf("wrapped: %w", pgx.ErrNoRows)) { + t.Fatal("IsNoRows returned false for pgx.ErrNoRows") + } + + if IsNoRows(errors.New("generic")) { + t.Fatal("IsNoRows returned true for generic error") + } +} + +func TestErrorClassifiers(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + state string + classifier func(error) bool + }{ + { + name: "unique violation", + state: sqlStateUniqueViolation, + classifier: IsUniqueViolation, + }, + { + name: "foreign key violation", + state: sqlStateForeignKeyViolation, + classifier: IsForeignKeyViolation, + }, + { + name: "not null violation", + state: sqlStateNotNullViolation, + classifier: IsNotNullViolation, + }, + { + name: "check violation", + state: sqlStateCheckViolation, + classifier: IsCheckViolation, + }, + { + name: "serialization failure", + state: sqlStateSerializationFailure, + classifier: IsSerializationFailure, + }, + { + name: "deadlock", + state: sqlStateDeadlockDetected, + classifier: IsDeadlock, + }, + { + name: "lock not available", + state: sqlStateLockNotAvailable, + classifier: IsLockNotAvailable, + }, + { + name: "query canceled", + state: sqlStateQueryCanceled, + classifier: IsQueryCanceled, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + + err := fmt.Errorf( + "wrapped: %w", + &pgconn.PgError{Code: test.state}, + ) + + if !test.classifier(err) { + t.Fatalf( + "classifier returned false for SQLSTATE %q", + test.state, + ) + } + + if test.classifier( + &pgconn.PgError{Code: "00000"}, + ) { + t.Fatalf( + "classifier returned true for unrelated SQLSTATE", + ) + } + }) + } +} + +func TestIsQueryCanceledDoesNotClassifyContextCancellation(t *testing.T) { + t.Parallel() + + if IsQueryCanceled(context.Canceled) { + t.Fatal("IsQueryCanceled returned true for context.Canceled") + } + + if IsQueryCanceled(context.DeadlineExceeded) { + t.Fatal("IsQueryCanceled returned true for context.DeadlineExceeded") + } +} + +func TestIsConnectionError(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + err error + want bool + }{ + { + name: "SQLSTATE connection exception", + err: &pgconn.PgError{Code: "08006"}, + want: true, + }, + { + name: "closed connection", + err: fmt.Errorf("wrapped: %w", pgconn.ErrConnClosed), + want: true, + }, + { + name: "EOF", + err: io.EOF, + want: true, + }, + { + name: "unexpected EOF", + err: io.ErrUnexpectedEOF, + want: true, + }, + { + name: "network operation", + err: &net.OpError{ + Op: "read", + Net: "tcp", + Err: errors.New("connection reset"), + }, + want: true, + }, + { + name: "generic error", + err: errors.New("generic"), + want: false, + }, + { + name: "nil", + err: nil, + want: false, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + + if got := IsConnectionError(test.err); got != test.want { + t.Fatalf( + "IsConnectionError() = %v, want %v", + got, + test.want, + ) + } + }) + } +} + +func TestIsRetryableTransaction(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + state string + want bool + }{ + { + name: "serialization failure", + state: sqlStateSerializationFailure, + want: true, + }, + { + name: "deadlock", + state: sqlStateDeadlockDetected, + want: true, + }, + { + name: "unique violation", + state: sqlStateUniqueViolation, + want: false, + }, + { + name: "connection exception", + state: "08006", + want: false, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + + err := &pgconn.PgError{Code: test.state} + + if got := IsRetryableTransaction(err); got != test.want { + t.Fatalf( + "IsRetryableTransaction() = %v, want %v", + got, + test.want, + ) + } + }) + } +} diff --git a/examples/README.md b/examples/README.md new file mode 100644 index 0000000..d84b6ba --- /dev/null +++ b/examples/README.md @@ -0,0 +1,19 @@ +# Examples + +This directory contains runnable examples covering the main `xpg` usage patterns. + +| Example | Covers | +|:---------------------------------|:----------------------------------------------------------------------------| +| [`basic`](basic) | Core `xpg.Pool` usage for common PostgreSQL operations | +| [`transactions`](transactions) | Transactions, savepoints, and recovering from an optional operation failure | +| [`advisory`](advisory) | Coordinating concurrent work with transaction-level advisory locks | +| [`observability`](observability) | `slog` logging, OpenTelemetry tracing, and Prometheus pool metrics | +| [`cluster`](cluster) | Primary and replica routing, round-robin reads, and read-only transactions | +| [`shard`](shard) | Range-based shard routing and grouping keys by shard | +| [`shard_geo`](shard_geo) | Custom geographic routing built from shard metadata | + +## Running the examples + +Each example is self-contained and includes its own setup and run instructions. + +Open the corresponding directory and follow its README. diff --git a/examples/advisory/README.md b/examples/advisory/README.md new file mode 100644 index 0000000..dcfbaf8 --- /dev/null +++ b/examples/advisory/README.md @@ -0,0 +1,111 @@ +# Transaction advisory locks + +This example shows how transaction-level advisory locks protect a shared operation across concurrent workers: + +* Acquire a lock for the lifetime of a transaction +* Check the same lock without blocking +* Release the lock automatically when the transaction completes + +Worker A holds the lock while its transaction is open. Worker B cannot enter the protected section, and worker C +acquires the lock after worker A commits. + +## Local setup + +From this directory, start PostgreSQL and Adminer: + +```shell +docker compose up -d +``` + +Apply the example schema: + +```shell +psql 'postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable' \ + < sql/schema.sql +``` + +The services are available at: + +```text +PostgreSQL: localhost:5432 +Adminer: http://localhost:8080 +``` + +To inspect the example data in Adminer, sign in with: + +```text +System: PostgreSQL +Server: postgres +Username: postgres +Password: postgres +Database: postgres +``` + +## Configuration + +By default, the example connects to: + +```text +postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable +``` + +To use another PostgreSQL instance, set `XPG_DATABASE_URL`: + +```shell +export XPG_DATABASE_URL='postgres://user:password@localhost:5432/database?sslmode=disable' +``` + +The target database must contain the schema from `sql/schema.sql`. + +> [!NOTE] +> The example runs two transactions concurrently, so the pool must allow at least two connections. + +## Run + +From this directory: + +```shell +go run . +``` + +Or from the repository root: + +```shell +go run ./examples/advisory +``` + +## Expected output + +```text +worker-a acquired the lock +worker-b acquired the lock: false +worker-a committed and released the lock +worker-c acquired the lock: true +recorded job runs: +- worker-a (lock key: 2026) +- worker-c (lock key: 2026) +``` + +Transaction-level advisory locks are released automatically on commit or rollback. Keep the protected transaction short +because it holds both the lock and a pool connection. + +## Cleanup + +To remove the example schema and data: + +```shell +psql 'postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable' \ + -c 'DROP SCHEMA IF EXISTS xpg_advisory_example CASCADE;' +``` + +Stop the local services: + +```shell +docker compose down +``` + +To also remove the PostgreSQL data volume: + +```shell +docker compose down -v +``` diff --git a/examples/advisory/docker-compose.yml b/examples/advisory/docker-compose.yml new file mode 100644 index 0000000..c7ff2bf --- /dev/null +++ b/examples/advisory/docker-compose.yml @@ -0,0 +1,30 @@ +services: + postgres: + image: postgres:18-alpine + environment: + POSTGRES_DB: postgres + POSTGRES_USER: postgres + POSTGRES_PASSWORD: postgres + ports: + - "5432:5432" + volumes: + - postgres-data:/var/lib/postgresql + healthcheck: + test: "pg_isready -U $${POSTGRES_USER} -d $${POSTGRES_DB}" + interval: 2s + timeout: 5s + retries: 10 + start_period: 5s + + adminer: + image: adminer:standalone + environment: + ADMINER_DEFAULT_SERVER: postgres + ports: + - "8080:8080" + depends_on: + postgres: + condition: service_healthy + +volumes: + postgres-data: diff --git a/examples/advisory/go.mod b/examples/advisory/go.mod new file mode 100644 index 0000000..da98d91 --- /dev/null +++ b/examples/advisory/go.mod @@ -0,0 +1,8 @@ +module advisory + +go 1.27 + +require ( + github.com/jackc/pgx/v5 v5.10.0 + github.com/mkbeh/xpg v0.2.0 +) diff --git a/examples/advisory/main.go b/examples/advisory/main.go new file mode 100644 index 0000000..25fd7c4 --- /dev/null +++ b/examples/advisory/main.go @@ -0,0 +1,239 @@ +package main + +import ( + "context" + "errors" + "fmt" + "log" + "os" + + "github.com/jackc/pgx/v5" + "github.com/mkbeh/xpg" +) + +const ( + defaultDatabaseURL = "postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable" + jobLockKey = int64(2026) +) + +type jobRun struct { + Worker string + LockKey int64 +} + +func main() { + if err := run(context.Background()); err != nil { + log.Fatal(err) + } +} + +func run(ctx context.Context) error { + pool, err := xpg.Open( + ctx, + databaseURL(), + xpg.WithName("advisory-example"), + ) + if err != nil { + return fmt.Errorf("open pool: %w", err) + } + defer pool.Close() + + if err := pool.Ping(ctx); err != nil { + return fmt.Errorf("ping PostgreSQL: %w", err) + } + + lockAcquired := make(chan struct{}) + releaseLock := make(chan struct{}) + holderDone := make(chan error, 1) + + go func() { + holderDone <- holdJobLock( + ctx, + pool, + "worker-a", + jobLockKey, + lockAcquired, + releaseLock, + ) + }() + + select { + case <-lockAcquired: + fmt.Println("worker-a acquired the lock") + case err := <-holderDone: + if err == nil { + return errors.New("worker-a exited before acquiring the lock") + } + + return fmt.Errorf("worker-a: %w", err) + case <-ctx.Done(): + return ctx.Err() + } + + workerBAcquired, workerBErr := tryRunJob(ctx, pool, "worker-b", jobLockKey) + + close(releaseLock) + holderErr := <-holderDone + + if workerBErr != nil { + return fmt.Errorf("worker-b: %w", workerBErr) + } + + if holderErr != nil { + return fmt.Errorf("worker-a: %w", holderErr) + } + + if workerBAcquired { + return errors.New("worker-b acquired a lock that should still be held") + } + + fmt.Printf("worker-b acquired the lock: %t\n", workerBAcquired) + fmt.Println("worker-a committed and released the lock") + + workerCAcquired, err := tryRunJob(ctx, pool, "worker-c", jobLockKey) + if err != nil { + return fmt.Errorf("worker-c: %w", err) + } + + if !workerCAcquired { + return errors.New("worker-c did not acquire the released lock") + } + + fmt.Printf("worker-c acquired the lock: %t\n", workerCAcquired) + + runs, err := loadJobRuns(ctx, pool) + if err != nil { + return fmt.Errorf("load job runs: %w", err) + } + + fmt.Println("recorded job runs:") + for _, current := range runs { + fmt.Printf("- %s (lock key: %d)\n", current.Worker, current.LockKey) + } + + return nil +} + +func holdJobLock( + ctx context.Context, + pool *xpg.Pool, + worker string, + lockKey int64, + lockAcquired chan<- struct{}, + releaseLock <-chan struct{}, +) error { + return pool.InTx( + ctx, + pgx.TxOptions{}, + func(ctx context.Context, tx pgx.Tx) error { + if err := xpg.AdvisoryXactLock(ctx, tx, lockKey); err != nil { + return err + } + + if err := recordJobRun(ctx, tx, worker, lockKey); err != nil { + return fmt.Errorf("record job run: %w", err) + } + + close(lockAcquired) + + select { + case <-releaseLock: + return nil + case <-ctx.Done(): + return ctx.Err() + } + }, + ) +} + +func tryRunJob( + ctx context.Context, + pool *xpg.Pool, + worker string, + lockKey int64, +) (bool, error) { + var acquired bool + + err := pool.InTx( + ctx, + pgx.TxOptions{}, + func(ctx context.Context, tx pgx.Tx) error { + var err error + + acquired, err = xpg.TryAdvisoryXactLock(ctx, tx, lockKey) + if err != nil || !acquired { + return err + } + + if err := recordJobRun(ctx, tx, worker, lockKey); err != nil { + return fmt.Errorf("record job run: %w", err) + } + + return nil + }, + ) + if err != nil { + return false, err + } + + return acquired, nil +} + +func recordJobRun( + ctx context.Context, + tx pgx.Tx, + worker string, + lockKey int64, +) error { + _, err := tx.Exec( + ctx, + `INSERT INTO xpg_advisory_example.job_runs (worker, lock_key) + VALUES ($1, $2)`, + worker, + lockKey, + ) + + return err +} + +func loadJobRuns(ctx context.Context, pool *xpg.Pool) ([]jobRun, error) { + rows, err := pool.Query( + ctx, + `SELECT worker, lock_key + FROM xpg_advisory_example.job_runs + ORDER BY id`, + ) + if err != nil { + return nil, err + } + defer rows.Close() + + runs := make([]jobRun, 0, 2) + + for rows.Next() { + var current jobRun + + if err := rows.Scan( + ¤t.Worker, + ¤t.LockKey, + ); err != nil { + return nil, err + } + + runs = append(runs, current) + } + + if err := rows.Err(); err != nil { + return nil, err + } + + return runs, nil +} + +func databaseURL() string { + if value := os.Getenv("XPG_DATABASE_URL"); value != "" { + return value + } + + return defaultDatabaseURL +} diff --git a/examples/advisory/sql/schema.sql b/examples/advisory/sql/schema.sql new file mode 100644 index 0000000..bab9a39 --- /dev/null +++ b/examples/advisory/sql/schema.sql @@ -0,0 +1,14 @@ +BEGIN; + +DROP SCHEMA IF EXISTS xpg_advisory_example CASCADE; + +CREATE SCHEMA xpg_advisory_example; + +CREATE TABLE xpg_advisory_example.job_runs ( + id bigint GENERATED ALWAYS AS IDENTITY PRIMARY KEY, + worker text NOT NULL, + lock_key bigint NOT NULL, + created_at timestamptz NOT NULL DEFAULT clock_timestamp() +); + +COMMIT; diff --git a/examples/basic/README.md b/examples/basic/README.md new file mode 100644 index 0000000..c971969 --- /dev/null +++ b/examples/basic/README.md @@ -0,0 +1,98 @@ +# Basic pool usage + +This example shows how to use `xpg.Pool` for common PostgreSQL operations: + +* Create a named PostgreSQL connection pool and verify connectivity +* Execute a write and read a single record back +* Query and iterate over multiple rows + +## Local setup + +From this directory, start PostgreSQL and Adminer: + +```shell +docker compose up -d +``` + +Apply the example schema: + +```shell +psql 'postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable' \ + < sql/schema.sql +``` + +The services are available at: + +```text +PostgreSQL: localhost:5432 +Adminer: http://localhost:8080 +``` + +To inspect the example data in Adminer, sign in with: + +```text +System: PostgreSQL +Server: postgres +Username: postgres +Password: postgres +Database: postgres +``` + +## Configuration + +By default, the example connects to: + +```text +postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable +``` + +To use another PostgreSQL instance, set `XPG_DATABASE_URL`: + +```shell +export XPG_DATABASE_URL='postgres://user:password@localhost:5432/database?sslmode=disable' +``` + +## Run + +From this directory: + +```shell +go run . +``` + +Or from the repository root: + +```shell +go run ./examples/basic +``` + +## Expected output + +```text +pool: basic-example +upserted users: 2 +selected user: 1 Alice active=true +active users: +- 1 Alice +``` + +## Cleanup + +To remove the example schema and data: + +```shell +psql 'postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable' \ + -c 'DROP SCHEMA IF EXISTS xpg_basic_example CASCADE;' +``` + +Stop the local services: + +```shell +docker compose down +``` + +To also remove the PostgreSQL data volume: + +```shell +docker compose down -v +``` diff --git a/examples/basic/docker-compose.yml b/examples/basic/docker-compose.yml new file mode 100644 index 0000000..046626a --- /dev/null +++ b/examples/basic/docker-compose.yml @@ -0,0 +1,30 @@ +services: + postgres: + image: postgres:18-alpine + environment: + POSTGRES_DB: postgres + POSTGRES_USER: postgres + POSTGRES_PASSWORD: postgres + ports: + - "5432:5432" + volumes: + - postgres-data:/var/lib/postgresql + healthcheck: + test: "pg_isready -U $${POSTGRES_USER} -d $${POSTGRES_DB}" + interval: 2s + timeout: 5s + retries: 10 + start_period: 5s + + adminer: + image: adminer:standalone + environment: + ADMINER_DEFAULT_SERVER: postgres + ports: + - "8080:8080" + depends_on: + postgres: + condition: service_healthy + +volumes: + postgres-data: \ No newline at end of file diff --git a/examples/basic/go.mod b/examples/basic/go.mod new file mode 100644 index 0000000..7e5031b --- /dev/null +++ b/examples/basic/go.mod @@ -0,0 +1,5 @@ +module basic + +go 1.27 + +require github.com/mkbeh/xpg v0.2.0 diff --git a/examples/basic/main.go b/examples/basic/main.go new file mode 100644 index 0000000..b3c1244 --- /dev/null +++ b/examples/basic/main.go @@ -0,0 +1,187 @@ +package main + +import ( + "context" + "fmt" + "log" + "os" + + "github.com/mkbeh/xpg" +) + +const defaultDatabaseURL = "postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable" + +type user struct { + ID int64 + Name string + Email string + Active bool +} + +func main() { + if err := run(context.Background()); err != nil { + log.Fatal(err) + } +} + +func run(ctx context.Context) error { + pool, err := xpg.Open( + ctx, + databaseURL(), + xpg.WithName("basic-example"), + ) + if err != nil { + return fmt.Errorf("open pool: %w", err) + } + defer pool.Close() + + if err := pool.Ping(ctx); err != nil { + return fmt.Errorf("ping PostgreSQL: %w", err) + } + + inserted, err := upsertUsers(ctx, pool) + if err != nil { + return fmt.Errorf("upsert users: %w", err) + } + + selected, err := loadUser(ctx, pool, 1) + if err != nil { + return fmt.Errorf("load user: %w", err) + } + + active, err := listActiveUsers(ctx, pool) + if err != nil { + return fmt.Errorf("list active users: %w", err) + } + + fmt.Printf("pool: %s\n", pool.Name()) + fmt.Printf("upserted users: %d\n", inserted) + fmt.Printf( + "selected user: %d %s <%s> active=%t\n", + selected.ID, + selected.Name, + selected.Email, + selected.Active, + ) + + fmt.Println("active users:") + for _, current := range active { + fmt.Printf("- %d %s <%s>\n", current.ID, current.Name, current.Email) + } + + return nil +} + +func upsertUsers(ctx context.Context, pool *xpg.Pool) (int64, error) { + tag, err := pool.Exec( + ctx, + `INSERT INTO xpg_basic_example.users ( + id, + name, + email, + active + ) + VALUES + ($1, $2, $3, $4), + ($5, $6, $7, $8) + ON CONFLICT (id) DO UPDATE SET + name = EXCLUDED.name, + email = EXCLUDED.email, + active = EXCLUDED.active`, + int64(1), + "Alice", + "alice@example.com", + true, + int64(2), + "Bob", + "bob@example.com", + false, + ) + if err != nil { + return 0, err + } + + return tag.RowsAffected(), nil +} + +func loadUser( + ctx context.Context, + pool *xpg.Pool, + userID int64, +) (user, error) { + var selected user + + err := pool.QueryRow( + ctx, + `SELECT + id, + name, + email, + active + FROM xpg_basic_example.users + WHERE id = $1`, + userID, + ).Scan( + &selected.ID, + &selected.Name, + &selected.Email, + &selected.Active, + ) + if err != nil { + return user{}, err + } + + return selected, nil +} + +func listActiveUsers( + ctx context.Context, + pool *xpg.Pool, +) ([]user, error) { + rows, err := pool.Query( + ctx, + `SELECT + id, + name, + email, + active + FROM xpg_basic_example.users + WHERE active + ORDER BY id`, + ) + if err != nil { + return nil, err + } + defer rows.Close() + + var users []user + + for rows.Next() { + var current user + + if err := rows.Scan( + ¤t.ID, + ¤t.Name, + ¤t.Email, + ¤t.Active, + ); err != nil { + return nil, err + } + + users = append(users, current) + } + + if err := rows.Err(); err != nil { + return nil, err + } + + return users, nil +} + +func databaseURL() string { + if value := os.Getenv("XPG_DATABASE_URL"); value != "" { + return value + } + + return defaultDatabaseURL +} diff --git a/examples/basic/sql/schema.sql b/examples/basic/sql/schema.sql new file mode 100644 index 0000000..846e4b1 --- /dev/null +++ b/examples/basic/sql/schema.sql @@ -0,0 +1,8 @@ +CREATE SCHEMA IF NOT EXISTS xpg_basic_example; + +CREATE TABLE IF NOT EXISTS xpg_basic_example.users ( + id bigint PRIMARY KEY, + name text NOT NULL, + email text NOT NULL UNIQUE, + active boolean NOT NULL +); diff --git a/examples/cluster/README.md b/examples/cluster/README.md new file mode 100644 index 0000000..114f36d --- /dev/null +++ b/examples/cluster/README.md @@ -0,0 +1,111 @@ +# Cluster routing + +This example shows how `cluster.Cluster` coordinates primary and replica access: + +* Route reads explicitly to the primary or replicas +* Distribute replica reads with round-robin selection +* Run write transactions on the primary +* Run read-only transactions on replicas + +> [!NOTE] +> The local services are independent PostgreSQL instances used only to demonstrate routing. They do not configure +> streaming replication. + +## Local setup + +From this directory, start the three PostgreSQL nodes and Adminer: + +```shell +docker compose up -d +``` + +Apply the node-specific setup: + +```shell +psql 'postgres://postgres:postgres@localhost:55432/postgres?sslmode=disable' \ + < sql/primary.sql + +psql 'postgres://postgres:postgres@localhost:55433/postgres?sslmode=disable' \ + < sql/replica-one.sql + +psql 'postgres://postgres:postgres@localhost:55434/postgres?sslmode=disable' \ + < sql/replica-two.sql +``` + +The services are available at: + +```text +Primary: localhost:55432 +Replica 1: localhost:55433 +Replica 2: localhost:55434 +Adminer: http://localhost:8080 +``` + +To inspect the nodes in Adminer, sign in with: + +```text +System: PostgreSQL +Server: postgres-primary +Username: postgres +Password: postgres +Database: postgres +``` + +Use `postgres-replica-one` or `postgres-replica-two` in the **Server** field to inspect a replica. + +## Configuration + +By default, the example connects to: + +```text +Primary: postgres://postgres:postgres@localhost:55432/postgres?sslmode=disable&target_session_attrs=read-write +Replica 1: postgres://postgres:postgres@localhost:55433/postgres?sslmode=disable&target_session_attrs=read-only +Replica 2: postgres://postgres:postgres@localhost:55434/postgres?sslmode=disable&target_session_attrs=read-only +``` + +To use other PostgreSQL endpoints, set `XPG_PRIMARY_DATABASE_URL`, `XPG_REPLICA_ONE_DATABASE_URL`, and +`XPG_REPLICA_TWO_DATABASE_URL`. + +## Run + +From this directory: + +```shell +go run . +``` + +Or from the repository root: + +```shell +go run ./examples/cluster +``` + +## Expected output + +```text +primary read: +- pool=cluster.primary node=primary role=primary +replica reads: +- pool=cluster.replica-one node=replica-one role=replica +- pool=cluster.replica-two node=replica-two role=replica +transactions: +- primary node=primary read_only=off +- replica node=replica-one read_only=on +``` + +The two replica reads show one complete round-robin cycle. The following read transaction continues from the same +selector and therefore resolves the first replica again. + +## Cleanup + +Stop the local services: + +```shell +docker compose down +``` + +To also remove all PostgreSQL data volumes: + +```shell +docker compose down -v +``` diff --git a/examples/cluster/docker-compose.yml b/examples/cluster/docker-compose.yml new file mode 100644 index 0000000..44e7765 --- /dev/null +++ b/examples/cluster/docker-compose.yml @@ -0,0 +1,70 @@ +services: + postgres-primary: + image: postgres:18-alpine + environment: + POSTGRES_DB: postgres + POSTGRES_USER: postgres + POSTGRES_PASSWORD: postgres + ports: + - "55432:5432" + volumes: + - primary-data:/var/lib/postgresql + healthcheck: + test: "pg_isready -U $${POSTGRES_USER} -d $${POSTGRES_DB}" + interval: 2s + timeout: 5s + retries: 10 + start_period: 5s + + postgres-replica-one: + image: postgres:18-alpine + environment: + POSTGRES_DB: postgres + POSTGRES_USER: postgres + POSTGRES_PASSWORD: postgres + ports: + - "55433:5432" + volumes: + - replica-one-data:/var/lib/postgresql + healthcheck: + test: "pg_isready -U $${POSTGRES_USER} -d $${POSTGRES_DB}" + interval: 2s + timeout: 5s + retries: 10 + start_period: 5s + + postgres-replica-two: + image: postgres:18-alpine + environment: + POSTGRES_DB: postgres + POSTGRES_USER: postgres + POSTGRES_PASSWORD: postgres + ports: + - "55434:5432" + volumes: + - replica-two-data:/var/lib/postgresql + healthcheck: + test: "pg_isready -U $${POSTGRES_USER} -d $${POSTGRES_DB}" + interval: 2s + timeout: 5s + retries: 10 + start_period: 5s + + adminer: + image: adminer:standalone + environment: + ADMINER_DEFAULT_SERVER: postgres-primary + ports: + - "8080:8080" + depends_on: + postgres-primary: + condition: service_healthy + postgres-replica-one: + condition: service_healthy + postgres-replica-two: + condition: service_healthy + +volumes: + primary-data: + replica-one-data: + replica-two-data: diff --git a/examples/cluster/go.mod b/examples/cluster/go.mod new file mode 100644 index 0000000..ba9090d --- /dev/null +++ b/examples/cluster/go.mod @@ -0,0 +1,8 @@ +module cluster + +go 1.27 + +require ( + github.com/jackc/pgx/v5 v5.10.0 + github.com/mkbeh/xpg v0.2.0 +) diff --git a/examples/cluster/main.go b/examples/cluster/main.go new file mode 100644 index 0000000..0d614cc --- /dev/null +++ b/examples/cluster/main.go @@ -0,0 +1,178 @@ +package main + +import ( + "context" + "fmt" + "log" + + "github.com/jackc/pgx/v5" + "github.com/mkbeh/xpg/cluster" +) + +type nodeInfo struct { + Name string + Role string +} + +type rowQuerier interface { + QueryRow(context.Context, string, ...any) pgx.Row +} + +func main() { + if err := run(context.Background()); err != nil { + log.Fatal(err) + } +} + +func run(ctx context.Context) error { + dbCluster, err := openCluster(ctx) + if err != nil { + return err + } + defer dbCluster.Close() + + if err := showRouting(ctx, dbCluster); err != nil { + return err + } + + if err := showTransactions(ctx, dbCluster); err != nil { + return err + } + + return nil +} + +func showRouting(ctx context.Context, dbCluster *cluster.Cluster) error { + primary, err := dbCluster.ReadPool(ctx, cluster.ReadPrimary) + if err != nil { + return fmt.Errorf("resolve primary read: %w", err) + } + + node, err := loadNode(ctx, primary) + if err != nil { + return fmt.Errorf("read primary node: %w", err) + } + + fmt.Println("primary read:") + fmt.Printf( + "- pool=%s node=%s role=%s\n", + primary.Name(), + node.Name, + node.Role, + ) + + fmt.Println("replica reads:") + + for range dbCluster.ReplicaCount() { + replica, err := dbCluster.ReadPool( + ctx, + cluster.ReadReplicaRequired, + ) + if err != nil { + return fmt.Errorf("resolve replica read: %w", err) + } + + node, err := loadNode(ctx, replica) + if err != nil { + return fmt.Errorf("read replica node: %w", err) + } + + fmt.Printf( + "- pool=%s node=%s role=%s\n", + replica.Name(), + node.Name, + node.Role, + ) + } + + return nil +} + +func showTransactions(ctx context.Context, dbCluster *cluster.Cluster) error { + var ( + primaryNode nodeInfo + primaryReadOnly string + ) + + err := dbCluster.InPrimaryTx( + ctx, + pgx.TxOptions{}, + func(ctx context.Context, tx pgx.Tx) error { + var err error + + primaryNode, err = loadNode(ctx, tx) + if err != nil { + return err + } + + return tx.QueryRow( + ctx, + "SHOW transaction_read_only", + ).Scan(&primaryReadOnly) + }, + ) + if err != nil { + return fmt.Errorf("run primary transaction: %w", err) + } + + var ( + replicaNode nodeInfo + replicaReadOnly string + ) + + err = dbCluster.InReadTx( + ctx, + cluster.ReadReplicaRequired, + cluster.ReadTxOptions{ + IsoLevel: pgx.RepeatableRead, + }, + func(ctx context.Context, tx pgx.Tx) error { + var err error + + replicaNode, err = loadNode(ctx, tx) + if err != nil { + return err + } + + return tx.QueryRow( + ctx, + "SHOW transaction_read_only", + ).Scan(&replicaReadOnly) + }, + ) + if err != nil { + return fmt.Errorf("run replica transaction: %w", err) + } + + fmt.Println("transactions:") + fmt.Printf( + "- primary node=%s read_only=%s\n", + primaryNode.Name, + primaryReadOnly, + ) + fmt.Printf( + "- replica node=%s read_only=%s\n", + replicaNode.Name, + replicaReadOnly, + ) + + return nil +} + +func loadNode(ctx context.Context, db rowQuerier) (nodeInfo, error) { + var node nodeInfo + + err := db.QueryRow( + ctx, + `SELECT node_name, node_role + FROM xpg_cluster_example.node_info`, + ).Scan( + &node.Name, + &node.Role, + ) + if err != nil { + return nodeInfo{}, err + } + + return node, nil +} diff --git a/examples/cluster/setup.go b/examples/cluster/setup.go new file mode 100644 index 0000000..fe62855 --- /dev/null +++ b/examples/cluster/setup.go @@ -0,0 +1,96 @@ +package main + +import ( + "context" + "fmt" + "os" + + "github.com/mkbeh/xpg" + "github.com/mkbeh/xpg/cluster" +) + +const ( + defaultPrimaryDatabaseURL = "postgres://postgres:postgres@localhost:55432/postgres?sslmode=disable&target_session_attrs=read-write" + defaultReplicaOneURL = "postgres://postgres:postgres@localhost:55433/postgres?sslmode=disable&target_session_attrs=read-only" + defaultReplicaTwoURL = "postgres://postgres:postgres@localhost:55434/postgres?sslmode=disable&target_session_attrs=read-only" +) + +func openCluster(ctx context.Context) (*cluster.Cluster, error) { + primary, err := openPool( + ctx, + environment("XPG_PRIMARY_DATABASE_URL", defaultPrimaryDatabaseURL), + "cluster.primary", + "primary", + ) + if err != nil { + return nil, fmt.Errorf("open primary pool: %w", err) + } + + replicaOne, err := openPool( + ctx, + environment("XPG_REPLICA_ONE_DATABASE_URL", defaultReplicaOneURL), + "cluster.replica-one", + "replica", + ) + if err != nil { + primary.Close() + + return nil, fmt.Errorf("open replica-one pool: %w", err) + } + + replicaTwo, err := openPool( + ctx, + environment("XPG_REPLICA_TWO_DATABASE_URL", defaultReplicaTwoURL), + "cluster.replica-two", + "replica", + ) + if err != nil { + replicaOne.Close() + primary.Close() + + return nil, fmt.Errorf("open replica-two pool: %w", err) + } + + dbCluster, err := cluster.New(cluster.Config{ + ID: "cluster-example", + Primary: primary, + Replicas: []*xpg.Pool{replicaOne, replicaTwo}, + }) + if err != nil { + replicaTwo.Close() + replicaOne.Close() + primary.Close() + + return nil, fmt.Errorf("create cluster: %w", err) + } + + return dbCluster, nil +} + +func openPool(ctx context.Context, databaseURL, name, role string) (*xpg.Pool, error) { + pool, err := xpg.Open( + ctx, + databaseURL, + xpg.WithName(name), + xpg.WithLabel("xpg.pool.role", role), + ) + if err != nil { + return nil, err + } + + if err := pool.Ping(ctx); err != nil { + pool.Close() + + return nil, fmt.Errorf("ping pool: %w", err) + } + + return pool, nil +} + +func environment(key, fallback string) string { + if value := os.Getenv(key); value != "" { + return value + } + + return fallback +} diff --git a/examples/cluster/sql/primary.sql b/examples/cluster/sql/primary.sql new file mode 100644 index 0000000..d265322 --- /dev/null +++ b/examples/cluster/sql/primary.sql @@ -0,0 +1,20 @@ +ALTER DATABASE postgres +SET default_transaction_read_only = off; + +DROP SCHEMA IF EXISTS xpg_cluster_example CASCADE; + +CREATE SCHEMA xpg_cluster_example; + +CREATE TABLE xpg_cluster_example.node_info ( + node_name text PRIMARY KEY, + node_role text NOT NULL +); + +INSERT INTO xpg_cluster_example.node_info ( + node_name, + node_role +) +VALUES ( + 'primary', + 'primary' +); diff --git a/examples/cluster/sql/replica-one.sql b/examples/cluster/sql/replica-one.sql new file mode 100644 index 0000000..c502e73 --- /dev/null +++ b/examples/cluster/sql/replica-one.sql @@ -0,0 +1,22 @@ +SET default_transaction_read_only = off; + +DROP SCHEMA IF EXISTS xpg_cluster_example CASCADE; + +CREATE SCHEMA xpg_cluster_example; + +CREATE TABLE xpg_cluster_example.node_info ( + node_name text PRIMARY KEY, + node_role text NOT NULL +); + +INSERT INTO xpg_cluster_example.node_info ( + node_name, + node_role +) +VALUES ( + 'replica-one', + 'replica' +); + +ALTER DATABASE postgres +SET default_transaction_read_only = on; diff --git a/examples/cluster/sql/replica-two.sql b/examples/cluster/sql/replica-two.sql new file mode 100644 index 0000000..d60ed58 --- /dev/null +++ b/examples/cluster/sql/replica-two.sql @@ -0,0 +1,22 @@ +SET default_transaction_read_only = off; + +DROP SCHEMA IF EXISTS xpg_cluster_example CASCADE; + +CREATE SCHEMA xpg_cluster_example; + +CREATE TABLE xpg_cluster_example.node_info ( + node_name text PRIMARY KEY, + node_role text NOT NULL +); + +INSERT INTO xpg_cluster_example.node_info ( + node_name, + node_role +) +VALUES ( + 'replica-two', + 'replica' +); + +ALTER DATABASE postgres +SET default_transaction_read_only = on; diff --git a/examples/observability/README.md b/examples/observability/README.md new file mode 100644 index 0000000..9f6fff1 --- /dev/null +++ b/examples/observability/README.md @@ -0,0 +1,101 @@ +# Observability + +This example shows how to add observability to `xpg`: + +* Log PostgreSQL activity with `slog` +* Trace database operations with OpenTelemetry +* Export connection-pool metrics to Prometheus +* Generate concurrent load to observe pool behavior + +## Configuration + +By default, the example connects to: + +```text +postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable +``` + +To use another PostgreSQL instance, set `XPG_DATABASE_URL`: + +```shell +export XPG_DATABASE_URL='postgres://user:password@localhost:5432/database?sslmode=disable' +``` + +The HTTP server listens on `localhost:9464`. To use another address, set `HTTP_ADDR`: + +```shell +export HTTP_ADDR='localhost:9464' +``` + +## Local setup + +From this directory, start PostgreSQL and Adminer: + +```shell +docker compose up -d +``` + +The services are available at: + +```text +PostgreSQL: localhost:5432 +Adminer: http://localhost:8080 +``` + +## Run + +From this directory: + +```shell +go run . +``` + +Or from the repository root: + +```shell +go run ./examples/observability +``` + +The HTTP server starts on: + +```shell +localhost:9464 +``` + +## Generate load + +Run six concurrent one-second queries against a pool limited to two connections: + +```shell +curl -X POST 'http://localhost:9464/load' +``` + +While the request is running, inspect pool contention from another terminal: + +```shell +curl -s 'http://localhost:9464/metrics' \ + | grep -E 'db_client_connection_count|xpg_pool_connection_acquire_' +``` + +To inspect all exported metrics: + +```shell +curl 'http://localhost:9464/metrics' +``` + +The request produces pgx logs and OpenTelemetry spans in the application output, while pool metrics remain available +from the `/metrics` endpoint. + +## Cleanup + +Stop the local services: + +```shell +docker compose down +``` + +To also remove the PostgreSQL data volume: + +```shell +docker compose down -v +``` \ No newline at end of file diff --git a/examples/observability/docker-compose.yml b/examples/observability/docker-compose.yml new file mode 100644 index 0000000..c7ff2bf --- /dev/null +++ b/examples/observability/docker-compose.yml @@ -0,0 +1,30 @@ +services: + postgres: + image: postgres:18-alpine + environment: + POSTGRES_DB: postgres + POSTGRES_USER: postgres + POSTGRES_PASSWORD: postgres + ports: + - "5432:5432" + volumes: + - postgres-data:/var/lib/postgresql + healthcheck: + test: "pg_isready -U $${POSTGRES_USER} -d $${POSTGRES_DB}" + interval: 2s + timeout: 5s + retries: 10 + start_period: 5s + + adminer: + image: adminer:standalone + environment: + ADMINER_DEFAULT_SERVER: postgres + ports: + - "8080:8080" + depends_on: + postgres: + condition: service_healthy + +volumes: + postgres-data: diff --git a/examples/observability/go.mod b/examples/observability/go.mod new file mode 100644 index 0000000..0176b64 --- /dev/null +++ b/examples/observability/go.mod @@ -0,0 +1,18 @@ +module observability + +go 1.27 + +require ( + github.com/exaring/otelpgx v0.11.1 + github.com/jackc/pgx/v5 v5.10.0 + github.com/mkbeh/xpg v0.2.0 + github.com/mkbeh/xpg/extra/otelxpg v0.1.0 + github.com/mkbeh/xpg/extra/slogxpg v0.1.0 + github.com/prometheus/client_golang v1.24.1 + go.opentelemetry.io/otel v1.45.0 + go.opentelemetry.io/otel/exporters/prometheus v0.68.0 + go.opentelemetry.io/otel/exporters/stdout/stdouttrace v1.46.0 + go.opentelemetry.io/otel/sdk v1.46.0 + go.opentelemetry.io/otel/sdk/metric v1.46.0 + go.opentelemetry.io/otel/trace v1.45.0 +) diff --git a/examples/observability/main.go b/examples/observability/main.go new file mode 100644 index 0000000..6b8f920 --- /dev/null +++ b/examples/observability/main.go @@ -0,0 +1,273 @@ +package main + +import ( + "context" + "errors" + "fmt" + "log/slog" + "net/http" + "os" + "os/signal" + "syscall" + "time" + + "github.com/exaring/otelpgx" + "github.com/jackc/pgx/v5/pgxpool" + "github.com/jackc/pgx/v5/tracelog" + "github.com/mkbeh/xpg" + "github.com/mkbeh/xpg/extra/otelxpg" + "github.com/mkbeh/xpg/extra/slogxpg" + "go.opentelemetry.io/otel/codes" + "go.opentelemetry.io/otel/trace" +) + +const ( + defaultDatabaseURL = "postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable" + defaultHTTPAddress = "localhost:9464" +) + +func main() { + logger := slog.New( + slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{ + Level: slog.LevelDebug, + }), + ) + + ctx, stop := signal.NotifyContext( + context.Background(), + os.Interrupt, + syscall.SIGTERM, + ) + defer stop() + + if err := run(ctx, logger); err != nil { + logger.Error( + "observability example failed", + slog.Any("error", err), + ) + os.Exit(1) + } +} + +func run(ctx context.Context, logger *slog.Logger) (runErr error) { + resource, err := newOTelResource(ctx) + if err != nil { + return err + } + + meterProvider, metricsHandler, err := newMeterProvider(resource) + if err != nil { + return fmt.Errorf("initialize metrics: %w", err) + } + defer func() { + shutdownCtx, cancel := context.WithTimeout( + context.Background(), + 5*time.Second, + ) + defer cancel() + + runErr = errors.Join( + runErr, + meterProvider.Shutdown(shutdownCtx), + ) + }() + + tracerProvider, err := newTracerProvider(resource) + if err != nil { + return fmt.Errorf("initialize tracing: %w", err) + } + defer func() { + shutdownCtx, cancel := context.WithTimeout( + context.Background(), + 5*time.Second, + ) + defer cancel() + + runErr = errors.Join( + runErr, + tracerProvider.Shutdown(shutdownCtx), + ) + }() + + config, err := pgxpool.ParseConfig(databaseURL()) + if err != nil { + return fmt.Errorf("parse PostgreSQL config: %w", err) + } + + // A small pool makes contention visible when POST /load runs. + config.MaxConns = 2 + + pgxLogger := logger.With( + slog.String("component", "pgx"), + ) + + pool, err := xpg.New( + ctx, + config, + xpg.WithName("observability-example"), + xpg.WithLabel("xpg.pool.role", "primary"), + xpg.WithLogger( + slogxpg.New(pgxLogger), + tracelog.LogLevelInfo, + ), + xpg.WithTracer( + otelpgx.NewTracer( + otelpgx.WithTracerProvider(tracerProvider), + otelpgx.WithTrimSQLInSpanName(), + ), + ), + xpg.WithMetrics( + otelxpg.NewMetrics( + otelxpg.WithMeterProvider(meterProvider), + ), + ), + ) + if err != nil { + return fmt.Errorf("create pool: %w", err) + } + defer pool.Close() + + if err := pool.Ping(ctx); err != nil { + return fmt.Errorf("ping PostgreSQL: %w", err) + } + + mux := http.NewServeMux() + mux.Handle("GET /metrics", metricsHandler) + mux.HandleFunc( + "POST /load", + loadHandler( + pool, + tracerProvider.Tracer(tracingInstrumentationName), + ), + ) + + server := &http.Server{ + Addr: httpAddress(), + Handler: mux, + ReadHeaderTimeout: 5 * time.Second, + } + + logger.Info( + "observability example listening", + slog.String("address", "http://"+server.Addr), + ) + + if err := serveHTTP(ctx, server); err != nil { + return fmt.Errorf("serve HTTP: %w", err) + } + + return nil +} + +func serveHTTP(ctx context.Context, server *http.Server) error { + serverErr := make(chan error, 1) + + go func() { + serverErr <- server.ListenAndServe() + }() + + select { + case err := <-serverErr: + if errors.Is(err, http.ErrServerClosed) { + return nil + } + + return err + case <-ctx.Done(): + } + + shutdownCtx, cancel := context.WithTimeout( + context.Background(), + 5*time.Second, + ) + defer cancel() + + if err := server.Shutdown(shutdownCtx); err != nil { + return fmt.Errorf("shutdown: %w", err) + } + + if err := <-serverErr; !errors.Is(err, http.ErrServerClosed) { + return err + } + + return nil +} + +func loadHandler(pool *xpg.Pool, tracer trace.Tracer) http.HandlerFunc { + return func(w http.ResponseWriter, request *http.Request) { + ctx, span := tracer.Start( + request.Context(), + "run-load", + ) + defer span.End() + + startedAt := time.Now() + + if err := runWorkload(ctx, pool); err != nil { + span.RecordError(err) + span.SetStatus(codes.Error, "workload failed") + + http.Error( + w, + fmt.Sprintf("run workload: %v", err), + http.StatusInternalServerError, + ) + + return + } + + _, _ = fmt.Fprintf( + w, + "workload completed in %s\n", + time.Since(startedAt), + ) + } +} + +func runWorkload(ctx context.Context, pool *xpg.Pool) error { + const workerCount = 6 + + start := make(chan struct{}) + results := make(chan error, workerCount) + + for range workerCount { + go func() { + <-start + + _, err := pool.Exec( + ctx, + "SELECT pg_sleep(1)", + ) + results <- err + }() + } + + close(start) + + var workloadErr error + + for range workerCount { + workloadErr = errors.Join( + workloadErr, + <-results, + ) + } + + return workloadErr +} + +func databaseURL() string { + if value := os.Getenv("XPG_DATABASE_URL"); value != "" { + return value + } + + return defaultDatabaseURL +} + +func httpAddress() string { + if value := os.Getenv("HTTP_ADDR"); value != "" { + return value + } + + return defaultHTTPAddress +} diff --git a/examples/observability/metrics.go b/examples/observability/metrics.go new file mode 100644 index 0000000..c49115f --- /dev/null +++ b/examples/observability/metrics.go @@ -0,0 +1,36 @@ +package main + +import ( + "fmt" + "net/http" + + promclient "github.com/prometheus/client_golang/prometheus" + "github.com/prometheus/client_golang/prometheus/promhttp" + otelprom "go.opentelemetry.io/otel/exporters/prometheus" + sdkmetric "go.opentelemetry.io/otel/sdk/metric" + "go.opentelemetry.io/otel/sdk/resource" +) + +func newMeterProvider(resource *resource.Resource) (*sdkmetric.MeterProvider, http.Handler, error) { + registry := promclient.NewRegistry() + + exporter, err := otelprom.New( + otelprom.WithRegisterer(registry), + otelprom.WithoutScopeInfo(), + ) + if err != nil { + return nil, nil, fmt.Errorf("create Prometheus exporter: %w", err) + } + + meterProvider := sdkmetric.NewMeterProvider( + sdkmetric.WithResource(resource), + sdkmetric.WithReader(exporter), + ) + + handler := promhttp.HandlerFor( + registry, + promhttp.HandlerOpts{}, + ) + + return meterProvider, handler, nil +} diff --git a/examples/observability/otel.go b/examples/observability/otel.go new file mode 100644 index 0000000..b0d319d --- /dev/null +++ b/examples/observability/otel.go @@ -0,0 +1,49 @@ +package main + +import ( + "context" + "fmt" + + "go.opentelemetry.io/otel/exporters/stdout/stdouttrace" + "go.opentelemetry.io/otel/sdk/resource" + sdktrace "go.opentelemetry.io/otel/sdk/trace" + semconv "go.opentelemetry.io/otel/semconv/v1.43.0" +) + +const tracingInstrumentationName = "github.com/mkbeh/xpg/examples/observability" + +func newOTelResource(ctx context.Context) (*resource.Resource, error) { + res, err := resource.New( + ctx, + resource.WithFromEnv(), + resource.WithTelemetrySDK(), + resource.WithAttributes( + semconv.ServiceName("xpg-observability-example"), + semconv.ServiceVersion("dev"), + ), + ) + if err != nil { + return nil, fmt.Errorf("create OpenTelemetry resource: %w", err) + } + + return res, nil +} + +func newTracerProvider(resource *resource.Resource) (*sdktrace.TracerProvider, error) { + exporter, err := stdouttrace.New( + stdouttrace.WithPrettyPrint(), + ) + if err != nil { + return nil, fmt.Errorf("create stdout trace exporter: %w", err) + } + + tracerProvider := sdktrace.NewTracerProvider( + sdktrace.WithResource(resource), + // The synchronous processor keeps the example easy to inspect. + // Production applications should normally use a batch processor with + // an OTLP exporter. + sdktrace.WithSyncer(exporter), + ) + + return tracerProvider, nil +} diff --git a/examples/sample/README.md b/examples/sample/README.md deleted file mode 100644 index 8b21258..0000000 --- a/examples/sample/README.md +++ /dev/null @@ -1,46 +0,0 @@ -# Description - -This is a sample REST API service implemented using pgx as the connector to a PostgreSQL data store. - -# Usage - -Create a PostgreSQL database. - -Configure the database connection with environment variables: - -```text -POSTGRES_CLUSTER_HOST=127.0.0.1 -POSTGRES_CLUSTER_PORT=54320 -POSTGRES_DB=db -POSTGRES_PASSWORD=pass -POSTGRES_USER=user -``` - -Set up docker-compose: -```shell -docker-compose up --build -d -``` - -Run main.go: - -``` -go run main.go -``` - -## Create tasks - -```shell -curl '127.0.0.1:8080/create' -``` - -## Get tasks - -```shell -curl '127.0.0.1:8080/get' -``` - -## Metrics - -```shell -curl 'http://localhost:8080/metrics' -``` \ No newline at end of file diff --git a/examples/sample/docker-compose.yml b/examples/sample/docker-compose.yml deleted file mode 100644 index be9059b..0000000 --- a/examples/sample/docker-compose.yml +++ /dev/null @@ -1,30 +0,0 @@ -version: '3.9' -services: - postgres: - image: postgres:14.2-alpine - container_name: sample-postgres - environment: - POSTGRES_DB: "db" - POSTGRES_PASSWORD: "pass" - POSTGRES_USER: "user" - ports: - - "54320:5432" - volumes: - - ./scripts/init_postgres.sh:/docker-entrypoint-initdb.d/init_postgres.sh - - sample-pgdata:/var/lib/postgresql/data - networks: - - sample-network - healthcheck: - test: [ "CMD-SHELL", "pg_isready -d $${POSTGRES_DB} -U $${POSTGRES_USER}" ] - interval: 10s - timeout: 5s - retries: 5 - start_period: 10s - - -networks: - sample-network: - driver: bridge - -volumes: - sample-pgdata: \ No newline at end of file diff --git a/examples/sample/main.go b/examples/sample/main.go deleted file mode 100644 index 2528e12..0000000 --- a/examples/sample/main.go +++ /dev/null @@ -1,114 +0,0 @@ -package main - -import ( - "encoding/json" - "log" - "net/http" - "os" - - postgres "github.com/mkbeh/xpg" - "github.com/mkbeh/xpg/examples/sample/migrations" - "github.com/prometheus/client_golang/prometheus/promhttp" -) - -var ( - writer *postgres.Pool - reader *postgres.Pool -) - -var ( - host string - port string - user string - pass string - db string -) - -func init() { - host = os.Getenv("POSTGRES_CLUSTER_HOST") - port = os.Getenv("POSTGRES_CLUSTER_PORT") - user = os.Getenv("POSTGRES_USER") - pass = os.Getenv("POSTGRES_PASSWORD") - db = os.Getenv("POSTGRES_DB") -} - -func getTasksHandler(w http.ResponseWriter, req *http.Request) { - type task struct { - Id int `json:"id"` - Description string `json:"description"` - } - - rows, err := reader.Query(req.Context(), "SELECT id, description FROM tasks;") - if err != nil { - http.Error(w, err.Error(), http.StatusInternalServerError) - return - } - defer rows.Close() - - tasks := make([]task, 0) - for rows.Next() { - var v task - if err := rows.Scan(&v.Id, &v.Description); err != nil { - http.Error(w, err.Error(), http.StatusInternalServerError) - return - } - tasks = append(tasks, v) - } - - data, _ := json.Marshal(tasks) - - w.Header().Set("Content-Type", "application/json") - w.Write(data) -} - -func createTasksHandler(w http.ResponseWriter, req *http.Request) { - _, err := writer.Exec(req.Context(), "INSERT INTO tasks VALUES (1, 'test1'), (2, 'test-2') ON CONFLICT DO NOTHING;") - if err != nil { - http.Error(w, err.Error(), http.StatusInternalServerError) - } else { - w.Header().Set("Content-Type", "text/plain") - w.Write([]byte("OK")) - } -} - -func main() { - var err error - - cfg := &postgres.Config{ - ClusterHost: host, - ClusterPort: port, - ClusterReplicaPort: port, - User: user, - Password: pass, - DB: db, - MigrateEnabled: true, - } - - writer, err = postgres.NewWriter( - postgres.WithConfig(cfg), - postgres.WithClientID("test-client"), - postgres.WithMigrations(migrations.FS), - ) - if err != nil { - log.Fatalln("failed init master pool", err) - } - defer writer.Close() - - reader, err = postgres.NewReader( - postgres.WithConfig(cfg), - postgres.WithClientID("test-client"), - ) - if err != nil { - log.Fatalln("failed init reader pool", err) - } - defer reader.Close() - - http.HandleFunc("/get", getTasksHandler) - http.HandleFunc("/create", createTasksHandler) - http.Handle("/metrics", promhttp.Handler()) - - err = http.ListenAndServe("localhost:8080", nil) - if err != nil { - log.Fatalln("Unable to start web server:", err) - } -} diff --git a/examples/sample/migrations/000001_init.down.sql b/examples/sample/migrations/000001_init.down.sql deleted file mode 100644 index 22c971f..0000000 --- a/examples/sample/migrations/000001_init.down.sql +++ /dev/null @@ -1 +0,0 @@ -drop table tasks; \ No newline at end of file diff --git a/examples/sample/migrations/000001_init.up.sql b/examples/sample/migrations/000001_init.up.sql deleted file mode 100644 index d7798b1..0000000 --- a/examples/sample/migrations/000001_init.up.sql +++ /dev/null @@ -1,5 +0,0 @@ -create table tasks -( - id serial primary key, - description text not null -); \ No newline at end of file diff --git a/examples/sample/migrations/embed.go b/examples/sample/migrations/embed.go deleted file mode 100644 index 91cca1c..0000000 --- a/examples/sample/migrations/embed.go +++ /dev/null @@ -1,6 +0,0 @@ -package migrations - -import "embed" - -//go:embed *.sql -var FS embed.FS diff --git a/examples/shard/README.md b/examples/shard/README.md new file mode 100644 index 0000000..4a890d1 --- /dev/null +++ b/examples/shard/README.md @@ -0,0 +1,116 @@ +# Sharding + +This example shows how to distribute application data across PostgreSQL shards: + +* Route user IDs with a range-based shard resolver +* Write records to the resolved shard +* Group keys by shard for efficient batch processing + +The example uses two primary-only shards: + +```text +[0, 100) -> shard-a +[100, 200) -> shard-b +``` + +## Local setup + +From this directory, start both PostgreSQL shards and Adminer: + +```shell +docker compose up -d +``` + +Apply the example schema to both shards: + +```shell +psql 'postgres://postgres:postgres@localhost:56431/postgres?sslmode=disable' \ + < sql/schema.sql + +psql 'postgres://postgres:postgres@localhost:56432/postgres?sslmode=disable' \ + < sql/schema.sql +``` + +The services are available at: + +```text +Shard A: localhost:56431 +Shard B: localhost:56432 +Adminer: http://localhost:8080 +``` + +To inspect a shard in Adminer, sign in with: + +```text +System: PostgreSQL +Server: postgres-shard-a +Username: postgres +Password: postgres +Database: postgres +``` + +Use `postgres-shard-b` in the **Server** field to inspect the second shard. + +## Configuration + +By default, the example connects to: + +```text +Shard A: postgres://postgres:postgres@localhost:56431/postgres?sslmode=disable +Shard B: postgres://postgres:postgres@localhost:56432/postgres?sslmode=disable +``` + +To use other PostgreSQL endpoints, set `XPG_SHARD_A_DATABASE_URL` and `XPG_SHARD_B_DATABASE_URL`. + +## Run + +From this directory: + +```shell +go run . +``` + +Or from the repository root: + +```shell +go run ./examples/shard +``` + +## Expected output + +```text +range routing: +- user_id=42 shard=shard-a pool=shard.shard-a.primary +- user_id=142 shard=shard-b pool=shard.shard-b.primary + +grouping: +- shard=shard-b user_ids=[142 143] +- shard=shard-a user_ids=[42 43] +``` + +`GroupByShard` preserves the order in which shards first appear in the input and the relative order of keys within each +group. + +## Cleanup + +To remove the example schema and data: + +```shell +psql 'postgres://postgres:postgres@localhost:56431/postgres?sslmode=disable' \ + -c 'DROP SCHEMA IF EXISTS xpg_shard_example CASCADE;' + +psql 'postgres://postgres:postgres@localhost:56432/postgres?sslmode=disable' \ + -c 'DROP SCHEMA IF EXISTS xpg_shard_example CASCADE;' +``` + +Stop the local services: + +```shell +docker compose down +``` + +To also remove both PostgreSQL data volumes: + +```shell +docker compose down -v +``` diff --git a/examples/shard/docker-compose.yml b/examples/shard/docker-compose.yml new file mode 100644 index 0000000..ae82f0a --- /dev/null +++ b/examples/shard/docker-compose.yml @@ -0,0 +1,50 @@ +services: + postgres-shard-a: + image: postgres:18-alpine + environment: + POSTGRES_DB: postgres + POSTGRES_USER: postgres + POSTGRES_PASSWORD: postgres + ports: + - "56431:5432" + volumes: + - shard-a-data:/var/lib/postgresql + healthcheck: + test: "pg_isready -U $${POSTGRES_USER} -d $${POSTGRES_DB}" + interval: 2s + timeout: 5s + retries: 10 + start_period: 5s + + postgres-shard-b: + image: postgres:18-alpine + environment: + POSTGRES_DB: postgres + POSTGRES_USER: postgres + POSTGRES_PASSWORD: postgres + ports: + - "56432:5432" + volumes: + - shard-b-data:/var/lib/postgresql + healthcheck: + test: "pg_isready -U $${POSTGRES_USER} -d $${POSTGRES_DB}" + interval: 2s + timeout: 5s + retries: 10 + start_period: 5s + + adminer: + image: adminer:standalone + environment: + ADMINER_DEFAULT_SERVER: postgres-shard-a + ports: + - "8080:8080" + depends_on: + postgres-shard-a: + condition: service_healthy + postgres-shard-b: + condition: service_healthy + +volumes: + shard-a-data: + shard-b-data: diff --git a/examples/shard/go.mod b/examples/shard/go.mod new file mode 100644 index 0000000..81abd8b --- /dev/null +++ b/examples/shard/go.mod @@ -0,0 +1,5 @@ +module shard + +go 1.27 + +require github.com/mkbeh/xpg v0.2.0 diff --git a/examples/shard/main.go b/examples/shard/main.go new file mode 100644 index 0000000..9f45891 --- /dev/null +++ b/examples/shard/main.go @@ -0,0 +1,113 @@ +package main + +import ( + "context" + "fmt" + "log" + + "github.com/mkbeh/xpg/shard" + "github.com/mkbeh/xpg/shard/resolver" +) + +const ( + shardARangeStart uint64 = 0 + shardBoundary uint64 = 100 + shardBRangeEnd uint64 = 200 +) + +type user struct { + ID uint64 + Name string +} + +func main() { + if err := run(context.Background()); err != nil { + log.Fatal(err) + } +} + +func run(ctx context.Context) error { + topology, err := openTopology(ctx) + if err != nil { + return err + } + defer topology.Close() + + userResolver, err := resolver.NewRange( + topology, + []resolver.Range[uint64]{ + { + Start: shardARangeStart, + End: shardBoundary, + ShardID: shardAID, + }, + { + Start: shardBoundary, + End: shardBRangeEnd, + ShardID: shardBID, + }, + }, + ) + if err != nil { + return fmt.Errorf("create user resolver: %w", err) + } + + users := []user{ + {ID: 42, Name: "alice"}, + {ID: 142, Name: "bob"}, + } + + fmt.Println("range routing:") + + for _, current := range users { + resolved, err := userResolver.Resolve(current.ID) + if err != nil { + return fmt.Errorf("resolve user %d: %w", current.ID, err) + } + + primary := resolved.Primary() + if primary == nil { + return fmt.Errorf("shard %q has no primary", resolved.ID()) + } + + if _, err := primary.Exec( + ctx, + `INSERT INTO xpg_shard_example.users (id, name) + VALUES ($1, $2) + ON CONFLICT (id) DO UPDATE + SET name = EXCLUDED.name`, + current.ID, + current.Name, + ); err != nil { + return fmt.Errorf("write user %d: %w", current.ID, err) + } + + fmt.Printf( + "- user_id=%d shard=%s pool=%s\n", + current.ID, + resolved.ID(), + primary.Name(), + ) + } + + groups, err := shard.GroupByShard( + userResolver, + []uint64{142, 42, 143, 43}, + ) + if err != nil { + return fmt.Errorf("group user IDs: %w", err) + } + + fmt.Println() + fmt.Println("grouping:") + + for _, group := range groups { + fmt.Printf( + "- shard=%s user_ids=%v\n", + group.Shard.ID(), + group.Keys, + ) + } + + return nil +} diff --git a/examples/shard/setup.go b/examples/shard/setup.go new file mode 100644 index 0000000..2a43d9d --- /dev/null +++ b/examples/shard/setup.go @@ -0,0 +1,101 @@ +package main + +import ( + "context" + "fmt" + "os" + + "github.com/mkbeh/xpg" + "github.com/mkbeh/xpg/cluster" + "github.com/mkbeh/xpg/shard" +) + +const ( + defaultShardADatabaseURL = "postgres://postgres:postgres@localhost:56431/postgres?sslmode=disable" + defaultShardBDatabaseURL = "postgres://postgres:postgres@localhost:56432/postgres?sslmode=disable" + + shardAID shard.ID = "shard-a" + shardBID shard.ID = "shard-b" +) + +func openTopology(ctx context.Context) (*shard.Topology, error) { + shardA, err := openShard( + ctx, + shardAID, + "shard.shard-a.primary", + environment( + "XPG_SHARD_A_DATABASE_URL", + defaultShardADatabaseURL, + ), + ) + if err != nil { + return nil, fmt.Errorf("open shard-a: %w", err) + } + + shardB, err := openShard( + ctx, + shardBID, + "shard.shard-b.primary", + environment( + "XPG_SHARD_B_DATABASE_URL", + defaultShardBDatabaseURL, + ), + ) + if err != nil { + shardA.Close() + + return nil, fmt.Errorf("open shard-b: %w", err) + } + + topology, err := shard.NewTopology([]shard.Config{ + {Cluster: shardA}, + {Cluster: shardB}, + }) + if err != nil { + shardB.Close() + shardA.Close() + + return nil, fmt.Errorf("create topology: %w", err) + } + + return topology, nil +} + +func openShard(ctx context.Context, id shard.ID, name, databaseURL string) (*cluster.Cluster, error) { + pool, err := xpg.Open( + ctx, + databaseURL, + xpg.WithName(name), + xpg.WithLabel("xpg.shard.id", string(id)), + xpg.WithLabel("xpg.pool.role", "primary"), + ) + if err != nil { + return nil, fmt.Errorf("open pool: %w", err) + } + + if err := pool.Ping(ctx); err != nil { + pool.Close() + + return nil, fmt.Errorf("ping pool: %w", err) + } + + shardCluster, err := cluster.New(cluster.Config{ + ID: id, + Primary: pool, + }) + if err != nil { + pool.Close() + + return nil, fmt.Errorf("create cluster: %w", err) + } + + return shardCluster, nil +} + +func environment(name, fallback string) string { + if value := os.Getenv(name); value != "" { + return value + } + + return fallback +} diff --git a/examples/shard/sql/schema.sql b/examples/shard/sql/schema.sql new file mode 100644 index 0000000..27c3abd --- /dev/null +++ b/examples/shard/sql/schema.sql @@ -0,0 +1,8 @@ +DROP SCHEMA IF EXISTS xpg_shard_example CASCADE; + +CREATE SCHEMA xpg_shard_example; + +CREATE TABLE xpg_shard_example.users ( + id bigint PRIMARY KEY, + name text NOT NULL +); diff --git a/examples/shard_geo/README.md b/examples/shard_geo/README.md new file mode 100644 index 0000000..f013fb4 --- /dev/null +++ b/examples/shard_geo/README.md @@ -0,0 +1,114 @@ +# Geographic sharding + +This example shows how to route tenants by region across PostgreSQL shards: + +* Associate each shard with a region +* Build a custom resolver from shard metadata +* Route tenant operations to the matching shard +* Handle unsupported regions with `shard.ErrNoShard` + +The topology contains two primary-only shards: + +```text +eu -> shard-eu +us -> shard-us +``` + +## Local setup + +From this directory, start both PostgreSQL shards and Adminer: + +```shell +docker compose up -d +``` + +Apply the example schema to both shards: + +```shell +psql 'postgres://postgres:postgres@localhost:57431/postgres?sslmode=disable' \ + < sql/schema.sql + +psql 'postgres://postgres:postgres@localhost:57432/postgres?sslmode=disable' \ + < sql/schema.sql +``` + +The services are available at: + +```text +EU shard: localhost:57431 +US shard: localhost:57432 +Adminer: http://localhost:8080 +``` + +To inspect a shard in Adminer, sign in with: + +```text +System: PostgreSQL +Server: postgres-shard-eu +Username: postgres +Password: postgres +Database: postgres +``` + +Use `postgres-shard-us` in the **Server** field to inspect the US shard. + +## Configuration + +By default, the example connects to: + +```text +EU shard: postgres://postgres:postgres@localhost:57431/postgres?sslmode=disable +US shard: postgres://postgres:postgres@localhost:57432/postgres?sslmode=disable +``` + +To use other PostgreSQL endpoints, set `XPG_SHARD_EU_DATABASE_URL` and `XPG_SHARD_US_DATABASE_URL`. + +## Run + +From this directory: + +```shell +go run . +``` + +Or from the repository root: + +```shell +go run ./examples/shard_geo +``` + +## Expected output + +```text +geo routing: +- tenant=tenant-42 region=eu shard=shard-eu name=Alice +- tenant=tenant-77 region=us shard=shard-us name=Bob +unsupported region: no shard +``` + +The region index is built once from immutable shard metadata. Each subsequent route is a local map lookup followed by +the normal `shard.Shard` execution API. + +## Cleanup + +To remove the example schema and data: + +```shell +psql 'postgres://postgres:postgres@localhost:57431/postgres?sslmode=disable' \ + -c 'DROP SCHEMA IF EXISTS xpg_shard_geo_example CASCADE;' + +psql 'postgres://postgres:postgres@localhost:57432/postgres?sslmode=disable' \ + -c 'DROP SCHEMA IF EXISTS xpg_shard_geo_example CASCADE;' +``` + +Stop the local services: + +```shell +docker compose down +``` + +To also remove both PostgreSQL data volumes: + +```shell +docker compose down -v +``` diff --git a/examples/shard_geo/docker-compose.yml b/examples/shard_geo/docker-compose.yml new file mode 100644 index 0000000..10c7cb0 --- /dev/null +++ b/examples/shard_geo/docker-compose.yml @@ -0,0 +1,50 @@ +services: + postgres-shard-eu: + image: postgres:18-alpine + environment: + POSTGRES_DB: postgres + POSTGRES_USER: postgres + POSTGRES_PASSWORD: postgres + ports: + - "57431:5432" + volumes: + - shard-eu-data:/var/lib/postgresql + healthcheck: + test: "pg_isready -U $${POSTGRES_USER} -d $${POSTGRES_DB}" + interval: 2s + timeout: 5s + retries: 10 + start_period: 5s + + postgres-shard-us: + image: postgres:18-alpine + environment: + POSTGRES_DB: postgres + POSTGRES_USER: postgres + POSTGRES_PASSWORD: postgres + ports: + - "57432:5432" + volumes: + - shard-us-data:/var/lib/postgresql + healthcheck: + test: "pg_isready -U $${POSTGRES_USER} -d $${POSTGRES_DB}" + interval: 2s + timeout: 5s + retries: 10 + start_period: 5s + + adminer: + image: adminer:standalone + environment: + ADMINER_DEFAULT_SERVER: postgres-shard-eu + ports: + - "8080:8080" + depends_on: + postgres-shard-eu: + condition: service_healthy + postgres-shard-us: + condition: service_healthy + +volumes: + shard-eu-data: + shard-us-data: diff --git a/examples/shard_geo/go.mod b/examples/shard_geo/go.mod new file mode 100644 index 0000000..7f51df5 --- /dev/null +++ b/examples/shard_geo/go.mod @@ -0,0 +1,8 @@ +module shard_geo + +go 1.27 + +require ( + github.com/jackc/pgx/v5 v5.10.0 + github.com/mkbeh/xpg v0.2.0 +) diff --git a/examples/shard_geo/main.go b/examples/shard_geo/main.go new file mode 100644 index 0000000..f2d75fd --- /dev/null +++ b/examples/shard_geo/main.go @@ -0,0 +1,151 @@ +package main + +import ( + "context" + "errors" + "fmt" + "log" + + "github.com/jackc/pgx/v5" + "github.com/mkbeh/xpg/shard" +) + +type tenantKey struct { + TenantID string + Region string +} + +type tenant struct { + ID string + Name string + Region string +} + +func main() { + if err := run(context.Background()); err != nil { + log.Fatal(err) + } +} + +func run(ctx context.Context) error { + topology, err := openTopology(ctx) + if err != nil { + return err + } + defer topology.Close() + + tenantResolver, err := newTenantResolver(topology) + if err != nil { + return fmt.Errorf("create tenant resolver: %w", err) + } + + tenants := []tenant{ + { + ID: "tenant-42", + Name: "Alice", + Region: "eu", + }, + { + ID: "tenant-77", + Name: "Bob", + Region: "us", + }, + } + + fmt.Println("geo routing:") + + for _, current := range tenants { + key := tenantKey{ + TenantID: current.ID, + Region: current.Region, + } + + resolved, err := tenantResolver.Resolve(key) + if err != nil { + return fmt.Errorf("resolve tenant %q: %w", current.ID, err) + } + + if err := upsertTenant(ctx, resolved, current); err != nil { + return fmt.Errorf("upsert tenant %q: %w", current.ID, err) + } + + stored, err := loadTenant(ctx, resolved, current.ID) + if err != nil { + return fmt.Errorf("load tenant %q: %w", current.ID, err) + } + + fmt.Printf( + "- tenant=%s region=%s shard=%s name=%s\n", + stored.ID, + stored.Region, + resolved.ID(), + stored.Name, + ) + } + + _, err = tenantResolver.Resolve(tenantKey{ + TenantID: "tenant-99", + Region: "apac", + }) + if !errors.Is(err, shard.ErrNoShard) { + return fmt.Errorf("resolve unsupported region: got %v, want shard.ErrNoShard", err) + } + + fmt.Println("unsupported region: no shard") + + return nil +} + +func upsertTenant(ctx context.Context, resolved shard.Shard, current tenant) error { + return resolved.InPrimaryTx( + ctx, + pgx.TxOptions{}, + func(ctx context.Context, tx pgx.Tx) error { + _, err := tx.Exec( + ctx, + `INSERT INTO xpg_shard_geo_example.tenants ( + id, + name, + region, + last_seen_at + ) + VALUES ($1, $2, $3, now()) + ON CONFLICT (id) DO UPDATE SET + name = EXCLUDED.name, + region = EXCLUDED.region, + last_seen_at = EXCLUDED.last_seen_at`, + current.ID, + current.Name, + current.Region, + ) + + return err + }, + ) +} + +func loadTenant(ctx context.Context, resolved shard.Shard, tenantID string) (tenant, error) { + primary := resolved.Primary() + if primary == nil { + return tenant{}, fmt.Errorf("shard %q has no primary", resolved.ID()) + } + + var stored tenant + + err := primary.QueryRow( + ctx, + `SELECT id, name, region + FROM xpg_shard_geo_example.tenants + WHERE id = $1`, + tenantID, + ).Scan( + &stored.ID, + &stored.Name, + &stored.Region, + ) + if err != nil { + return tenant{}, err + } + + return stored, nil +} diff --git a/examples/shard_geo/resolver.go b/examples/shard_geo/resolver.go new file mode 100644 index 0000000..dca5e14 --- /dev/null +++ b/examples/shard_geo/resolver.go @@ -0,0 +1,50 @@ +package main + +import ( + "fmt" + + "github.com/mkbeh/xpg/shard" + "github.com/mkbeh/xpg/shard/resolver" +) + +func newTenantResolver(topology *shard.Topology) (shard.Resolver[tenantKey], error) { + shardByRegion := make(map[string]shard.ID, topology.Len()) + + // Build the routing index once instead of scanning the topology on every + // tenant lookup. + for _, candidate := range topology.Shards() { + region, ok := candidate.Label("region") + if !ok || region == "" { + return nil, fmt.Errorf("shard %q has no region label", candidate.ID()) + } + + if existing, exists := shardByRegion[region]; exists { + return nil, fmt.Errorf( + "region %q is assigned to both shard %q and shard %q", + region, + existing, + candidate.ID(), + ) + } + + shardByRegion[region] = candidate.ID() + } + + resolve := resolver.ResolveFunc[tenantKey]( + func(key tenantKey, _ *shard.Topology) (shard.ID, error) { + id, ok := shardByRegion[key.Region] + if !ok { + return "", fmt.Errorf( + "tenant %q has unsupported region %q: %w", + key.TenantID, + key.Region, + shard.ErrNoShard, + ) + } + + return id, nil + }, + ) + + return resolver.NewCustom(topology, resolve) +} diff --git a/examples/shard_geo/setup.go b/examples/shard_geo/setup.go new file mode 100644 index 0000000..e189457 --- /dev/null +++ b/examples/shard_geo/setup.go @@ -0,0 +1,110 @@ +package main + +import ( + "context" + "fmt" + "os" + + "github.com/mkbeh/xpg" + "github.com/mkbeh/xpg/cluster" + "github.com/mkbeh/xpg/shard" +) + +const ( + defaultShardEUDatabaseURL = "postgres://postgres:postgres@localhost:57431/postgres?sslmode=disable" + defaultShardUSDatabaseURL = "postgres://postgres:postgres@localhost:57432/postgres?sslmode=disable" + + shardEUID shard.ID = "shard-eu" + shardUSID shard.ID = "shard-us" +) + +func openTopology(ctx context.Context) (*shard.Topology, error) { + shardEU, err := openShard( + ctx, + shardEUID, + "eu", + "geo.shard-eu.primary", + environment( + "XPG_SHARD_EU_DATABASE_URL", + defaultShardEUDatabaseURL, + ), + ) + if err != nil { + return nil, fmt.Errorf("open shard-eu: %w", err) + } + + shardUS, err := openShard( + ctx, + shardUSID, + "us", + "geo.shard-us.primary", + environment( + "XPG_SHARD_US_DATABASE_URL", + defaultShardUSDatabaseURL, + ), + ) + if err != nil { + shardEU.Close() + + return nil, fmt.Errorf("open shard-us: %w", err) + } + + topology, err := shard.NewTopology([]shard.Config{ + {Cluster: shardEU}, + {Cluster: shardUS}, + }) + if err != nil { + shardUS.Close() + shardEU.Close() + + return nil, fmt.Errorf("create topology: %w", err) + } + + return topology, nil +} + +func openShard( + ctx context.Context, + id shard.ID, + region string, + name string, + databaseURL string, +) (*cluster.Cluster, error) { + pool, err := xpg.Open( + ctx, + databaseURL, + xpg.WithName(name), + xpg.WithLabel("xpg.shard.id", string(id)), + xpg.WithLabel("xpg.pool.role", "primary"), + ) + if err != nil { + return nil, fmt.Errorf("open pool: %w", err) + } + + if err := pool.Ping(ctx); err != nil { + pool.Close() + return nil, fmt.Errorf("ping pool: %w", err) + } + + shardCluster, err := cluster.New(cluster.Config{ + ID: id, + Labels: map[string]string{ + "region": region, + }, + Primary: pool, + }) + if err != nil { + pool.Close() + return nil, fmt.Errorf("create cluster: %w", err) + } + + return shardCluster, nil +} + +func environment(name, fallback string) string { + if value := os.Getenv(name); value != "" { + return value + } + + return fallback +} diff --git a/examples/shard_geo/sql/schema.sql b/examples/shard_geo/sql/schema.sql new file mode 100644 index 0000000..1b30d39 --- /dev/null +++ b/examples/shard_geo/sql/schema.sql @@ -0,0 +1,10 @@ +DROP SCHEMA IF EXISTS xpg_shard_geo_example CASCADE; + +CREATE SCHEMA xpg_shard_geo_example; + +CREATE TABLE xpg_shard_geo_example.tenants ( + id text PRIMARY KEY, + name text NOT NULL, + region text NOT NULL, + last_seen_at timestamptz NOT NULL +); diff --git a/examples/transactions/README.md b/examples/transactions/README.md new file mode 100644 index 0000000..d5ad7d8 --- /dev/null +++ b/examples/transactions/README.md @@ -0,0 +1,106 @@ +# Transactions and savepoints + +This example shows how to keep an outer transaction commit-able when an optional operation fails: + +* Execute the main operation inside a transaction +* Isolate optional work with a savepoint +* Detect an expected PostgreSQL constraint violation +* Roll back only the savepoint while allowing the outer transaction to commit + +The example attempts to redeem `PROMO2026`, which is already in use. The promo write fails and is rolled back to the +savepoint, while the order is committed successfully. + +## Local setup + +From this directory, start PostgreSQL and Adminer: + +```shell +docker compose up -d +``` + +Apply the example schema: + +```shell +psql 'postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable' \ + < sql/schema.sql +``` + +The services are available at: + +```text +PostgreSQL: localhost:5432 +Adminer: http://localhost:8080 +``` + +To inspect the example data in Adminer, sign in with: + +```text +System: PostgreSQL +Server: postgres +Username: postgres +Password: postgres +Database: postgres +``` + +## Configuration + +By default, the example connects to: + +```text +postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable +``` + +To use another PostgreSQL instance, set `XPG_DATABASE_URL`: + +```shell +export XPG_DATABASE_URL='postgres://user:password@localhost:5432/database?sslmode=disable' +``` + +The target database must contain the schema from `sql/schema.sql`. + +## Run + +From this directory: + +```shell +go run . +``` + +Or from the repository root: + +```shell +go run ./examples/transactions +``` + +## Expected output + +```text +order ID: 1 +order status: new +promo code: PROMO2026 +promo applied: false +``` + +The failed promo insert is rolled back to the savepoint. The outer transaction then returns `nil`, so the order is +committed. + +## Cleanup + +To remove the example schema and data: + +```shell +psql 'postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable' \ + -c 'DROP SCHEMA IF EXISTS xpg_transactions_example CASCADE;' +``` + +Stop the local services: + +```shell +docker compose down +``` + +To also remove the PostgreSQL data volume: + +```shell +docker compose down -v +``` diff --git a/examples/transactions/docker-compose.yml b/examples/transactions/docker-compose.yml new file mode 100644 index 0000000..c7ff2bf --- /dev/null +++ b/examples/transactions/docker-compose.yml @@ -0,0 +1,30 @@ +services: + postgres: + image: postgres:18-alpine + environment: + POSTGRES_DB: postgres + POSTGRES_USER: postgres + POSTGRES_PASSWORD: postgres + ports: + - "5432:5432" + volumes: + - postgres-data:/var/lib/postgresql + healthcheck: + test: "pg_isready -U $${POSTGRES_USER} -d $${POSTGRES_DB}" + interval: 2s + timeout: 5s + retries: 10 + start_period: 5s + + adminer: + image: adminer:standalone + environment: + ADMINER_DEFAULT_SERVER: postgres + ports: + - "8080:8080" + depends_on: + postgres: + condition: service_healthy + +volumes: + postgres-data: diff --git a/examples/transactions/go.mod b/examples/transactions/go.mod new file mode 100644 index 0000000..ed21479 --- /dev/null +++ b/examples/transactions/go.mod @@ -0,0 +1,8 @@ +module transactions + +go 1.27 + +require ( + github.com/jackc/pgx/v5 v5.10.0 + github.com/mkbeh/xpg v0.2.0 +) diff --git a/examples/transactions/main.go b/examples/transactions/main.go new file mode 100644 index 0000000..782dfd6 --- /dev/null +++ b/examples/transactions/main.go @@ -0,0 +1,154 @@ +package main + +import ( + "context" + "fmt" + "log" + "os" + + "github.com/jackc/pgx/v5" + "github.com/mkbeh/xpg" +) + +const defaultDatabaseURL = "postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable" + +func main() { + if err := run(context.Background()); err != nil { + log.Fatal(err) + } +} + +func run(ctx context.Context) error { + pool, err := xpg.Open( + ctx, + databaseURL(), + xpg.WithName("transactions-example"), + ) + if err != nil { + return fmt.Errorf("open pool: %w", err) + } + defer pool.Close() + + if err := pool.Ping(ctx); err != nil { + return fmt.Errorf("ping PostgreSQL: %w", err) + } + + const ( + orderID = int64(1) + promoCode = "PROMO2026" + ) + + if err := processOrder(ctx, pool, orderID, promoCode); err != nil { + return fmt.Errorf("process order: %w", err) + } + + status, promoApplied, err := loadOrder(ctx, pool, orderID, promoCode) + if err != nil { + return fmt.Errorf("load order: %w", err) + } + + fmt.Printf("order ID: %d\n", orderID) + fmt.Printf("order status: %s\n", status) + fmt.Printf("promo code: %s\n", promoCode) + fmt.Printf("promo applied: %t\n", promoApplied) + + return nil +} + +func processOrder( + ctx context.Context, + pool *xpg.Pool, + orderID int64, + promoCode string, +) error { + return pool.InTx( + ctx, + pgx.TxOptions{}, + func(ctx context.Context, tx pgx.Tx) error { + if _, err := tx.Exec( + ctx, + `INSERT INTO xpg_transactions_example.orders (id, status) + VALUES ($1, $2) + ON CONFLICT (id) DO UPDATE + SET status = EXCLUDED.status`, + orderID, + "new", + ); err != nil { + return fmt.Errorf("upsert order: %w", err) + } + + err := xpg.InSavepoint( + ctx, + tx, + func(ctx context.Context, savepoint pgx.Tx) error { + _, err := savepoint.Exec( + ctx, + `INSERT INTO xpg_transactions_example.promo_redemptions ( + code, + order_id + ) + VALUES ($1, $2)`, + promoCode, + orderID, + ) + + return err + }, + ) + if err == nil { + return nil + } + + if xpg.IsUniqueViolation(err) { + // InSavepoint already rolled back the failed promo insert. + return nil + } + + return fmt.Errorf("apply promo: %w", err) + }, + ) +} + +func loadOrder( + ctx context.Context, + pool *xpg.Pool, + orderID int64, + promoCode string, +) (string, bool, error) { + var ( + status string + promoApplied bool + ) + + err := pool.QueryRow( + ctx, + `SELECT + o.status, + EXISTS ( + SELECT 1 + FROM xpg_transactions_example.promo_redemptions AS p + WHERE p.order_id = o.id + AND p.code = $2 + ) + FROM xpg_transactions_example.orders AS o + WHERE o.id = $1`, + orderID, + promoCode, + ).Scan( + &status, + &promoApplied, + ) + if err != nil { + return "", false, err + } + + return status, promoApplied, nil +} + +func databaseURL() string { + if value := os.Getenv("XPG_DATABASE_URL"); value != "" { + return value + } + + return defaultDatabaseURL +} diff --git a/examples/transactions/sql/schema.sql b/examples/transactions/sql/schema.sql new file mode 100644 index 0000000..94e4185 --- /dev/null +++ b/examples/transactions/sql/schema.sql @@ -0,0 +1,16 @@ +CREATE SCHEMA IF NOT EXISTS xpg_transactions_example; + +CREATE TABLE IF NOT EXISTS xpg_transactions_example.orders ( + id bigint PRIMARY KEY, + status text NOT NULL +); + +CREATE TABLE IF NOT EXISTS xpg_transactions_example.promo_redemptions ( + code text PRIMARY KEY, + order_id bigint NOT NULL +); + +INSERT INTO xpg_transactions_example.promo_redemptions (code, order_id) +VALUES ('PROMO2026', 100) +ON CONFLICT (code) DO UPDATE +SET order_id = EXCLUDED.order_id; diff --git a/extra/otelxpg/README.md b/extra/otelxpg/README.md new file mode 100644 index 0000000..492f65c --- /dev/null +++ b/extra/otelxpg/README.md @@ -0,0 +1,46 @@ +# OpenTelemetry Metrics for xpg + +`otelxpg` provides optional OpenTelemetry metrics integration for `xpg`. + +The package is exporter-agnostic: applications own the OpenTelemetry SDK lifecycle and exporter configuration, while +`otelxpg` uses the configured `MeterProvider` to expose connection pool metrics. + +## Installation + +```bash +go get github.com/mkbeh/xpg/extra/otelxpg +``` + +## Usage + + +```go +import ( + "context" + + "github.com/mkbeh/xpg" + "github.com/mkbeh/xpg/extra/otelxpg" + sdkmetric "go.opentelemetry.io/otel/sdk/metric" +) + +meterProvider := sdkmetric.NewMeterProvider() +defer meterProvider.Shutdown(context.Background()) + +// Create a reusable OpenTelemetry metrics integration. +metrics := otelxpg.NewMetrics( + otelxpg.WithMeterProvider(meterProvider), +) + +// Attach metrics when creating the pool. +pool, _ := xpg.Open( + context.Background(), + databaseURL, + xpg.WithName("example-pool"), + xpg.WithMetrics(metrics), +) +defer pool.Close() +``` + + +For a complete runnable setup using OpenTelemetry and Prometheus, see the +[observability example](../../examples/observability). \ No newline at end of file diff --git a/extra/otelxpg/go.mod b/extra/otelxpg/go.mod new file mode 100644 index 0000000..38c5329 --- /dev/null +++ b/extra/otelxpg/go.mod @@ -0,0 +1,9 @@ +module github.com/mkbeh/xpg/extra/otelxpg + +go 1.27 + +require ( + github.com/mkbeh/xpg v0.2.0 + go.opentelemetry.io/otel v1.46.0 + go.opentelemetry.io/otel/metric v1.46.0 +) diff --git a/extra/otelxpg/metrics.go b/extra/otelxpg/metrics.go new file mode 100644 index 0000000..7427388 --- /dev/null +++ b/extra/otelxpg/metrics.go @@ -0,0 +1,36 @@ +package otelxpg + +import ( + "fmt" + "sync" + + "go.opentelemetry.io/otel" + "go.opentelemetry.io/otel/metric" +) + +const instrumentationName = "github.com/mkbeh/xpg/extra/otelxpg" + +// Metrics exports xpg pool statistics through OpenTelemetry. +// +// Metrics is safe for concurrent use and reuse across multiple pools. +type Metrics struct { + meterProvider metric.MeterProvider +} + +// metricsRegistration represents one OpenTelemetry callback registration. +type metricsRegistration struct { + registration metric.Registration + closeOnce sync.Once +} + +func (m *metricsRegistration) Close() { + if m == nil || m.registration == nil { + return + } + + m.closeOnce.Do(func() { + if err := m.registration.Unregister(); err != nil { + otel.Handle(fmt.Errorf("otelxpg: unregister metrics: %w", err)) + } + }) +} diff --git a/extra/otelxpg/metrics_test.go b/extra/otelxpg/metrics_test.go new file mode 100644 index 0000000..4a50451 --- /dev/null +++ b/extra/otelxpg/metrics_test.go @@ -0,0 +1,508 @@ +package otelxpg + +import ( + "slices" + "testing" + "time" + + "github.com/jackc/pgx/v5/pgxpool" + "github.com/mkbeh/xpg" + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/metric" + "go.opentelemetry.io/otel/metric/noop" +) + +func TestMetricsRegistration(t *testing.T) { + t.Parallel() + + metrics := NewMetrics( + WithMeterProvider(noop.NewMeterProvider()), + ) + + pool := newTestPool(t, metrics) + + pool.Close() + pool.Close() +} + +func TestMetricsRegisterNilReceiver(t *testing.T) { + t.Parallel() + + var metrics *Metrics + + registration, err := metrics.Register(nil) + if err == nil { + t.Fatal("expected error") + } + + if registration != nil { + t.Fatal("expected nil registration") + } +} + +func TestMetricsRegisterNilPool(t *testing.T) { + t.Parallel() + + metrics := NewMetrics( + WithMeterProvider(noop.NewMeterProvider()), + ) + + registration, err := metrics.Register(nil) + if err == nil { + t.Fatal("expected error") + } + + if registration != nil { + t.Fatal("expected nil registration") + } +} + +func TestWithMeterProviderNilUsesGlobalProvider(t *testing.T) { + metrics := NewMetrics( + WithMeterProvider(nil), + ) + + pool := newTestPool(t, metrics) + pool.Close() +} + +func TestPoolMetricInstruments(t *testing.T) { + t.Parallel() + + meter := &recordingMeter{} + + instruments, err := newPoolMetricInstruments(meter) + if err != nil { + t.Fatalf("create instruments: %v", err) + } + + want := []instrumentSpec{ + { + name: connectionCountMetricName, + kind: "int64_observable_up_down_counter", + description: "The number of connections currently in the state described by db.client.connection.state.", + unit: "{connection}", + }, + { + name: connectionMaxMetricName, + kind: "int64_observable_up_down_counter", + description: "The maximum number of open connections allowed by the pool.", + unit: "{connection}", + }, + { + name: connectionConstructingMetricName, + kind: "int64_observable_gauge", + description: "The number of connections currently being created by the pool.", + unit: "{connection}", + }, + { + name: connectionAcquireCountMetricName, + kind: "int64_observable_counter", + description: "The cumulative number of successful connection acquires.", + unit: "{request}", + }, + { + name: connectionAcquireTimeMetricName, + kind: "float64_observable_counter", + description: "The cumulative time spent acquiring connections.", + unit: "s", + }, + { + name: connectionAcquireCanceledCountMetricName, + kind: "int64_observable_counter", + description: "The cumulative number of connection acquires canceled by context.", + unit: "{request}", + }, + { + name: connectionAcquireEmptyCountMetricName, + kind: "int64_observable_counter", + description: "The cumulative number of successful acquires that waited because the pool was empty.", + unit: "{request}", + }, + { + name: connectionAcquireEmptyWaitTimeMetricName, + kind: "float64_observable_counter", + description: "The cumulative time spent waiting for a connection while the pool was empty.", + unit: "s", + }, + { + name: connectionCreateCountMetricName, + kind: "int64_observable_counter", + description: "The cumulative number of connections created by the pool.", + unit: "{connection}", + }, + { + name: connectionDestroyCountMetricName, + kind: "int64_observable_counter", + description: "The cumulative number of connections destroyed by pool lifecycle limits.", + unit: "{connection}", + }, + } + + if !slices.Equal(meter.instruments, want) { + t.Fatalf( + "instruments = %#v, want %#v", + meter.instruments, + want, + ) + } + + if got, want := len(instruments.observables()), len(want); got != want { + t.Fatalf("observable count = %d, want %d", got, want) + } +} + +func TestPoolMetricAttributes(t *testing.T) { + t.Parallel() + + attributes := newPoolMetricAttributes( + "test-pool", + map[string]string{ + "region": "eu", + dbSystemNameAttribute: "custom-system", + poolNameAttribute: "custom-pool", + connectionStateAttribute: "custom-state", + destroyReasonAttribute: "custom-reason", + }, + ) + + base := observeAttributes(attributes.base) + + requireStringAttribute(t, base, dbSystemNameAttribute, dbSystemPostgreSQL) + requireStringAttribute(t, base, poolNameAttribute, "test-pool") + requireStringAttribute(t, base, "region", "eu") + requireStringAttribute(t, base, connectionStateAttribute, "custom-state") + requireStringAttribute(t, base, destroyReasonAttribute, "custom-reason") + + idle := observeAttributes(attributes.idle) + requireStringAttribute(t, idle, connectionStateAttribute, connectionStateIdle) + + used := observeAttributes(attributes.used) + requireStringAttribute(t, used, connectionStateAttribute, connectionStateUsed) + + destroyedIdle := observeAttributes(attributes.destroyedIdle) + requireStringAttribute(t, destroyedIdle, destroyReasonAttribute, destroyReasonIdleTimeout) + + destroyedLifetime := observeAttributes(attributes.destroyedLifetime) + requireStringAttribute(t, destroyedLifetime, destroyReasonAttribute, destroyReasonLifetime) +} + +func TestPoolMetricInstrumentsObserve(t *testing.T) { + t.Parallel() + + stats := xpg.PoolStats{ + AcquiredConns: 1, + ConstructingConns: 2, + IdleConns: 3, + MaxConns: 4, + AcquireCount: 5, + AcquireDuration: 6 * time.Second, + CanceledAcquireCount: 7, + EmptyAcquireCount: 8, + EmptyAcquireWaitTime: 9 * time.Second, + NewConnsCount: 10, + MaxIdleDestroyCount: 11, + MaxLifetimeDestroyCount: 12, + } + + attributes := newPoolMetricAttributes( + "test-pool", + map[string]string{ + "region": "eu", + }, + ) + + observer := &recordingObserver{} + + var instruments poolMetricInstruments + + instruments.observe( + observer, + stats, + attributes, + ) + + wantInt64 := []int64{ + 3, + 1, + 4, + 2, + 5, + 7, + 8, + 10, + 11, + 12, + } + + if !slices.Equal(observer.int64Values, wantInt64) { + t.Fatalf( + "int64 observations = %v, want %v", + observer.int64Values, + wantInt64, + ) + } + + wantFloat64 := []float64{ + 6, + 9, + } + + if !slices.Equal(observer.float64Values, wantFloat64) { + t.Fatalf( + "float64 observations = %v, want %v", + observer.float64Values, + wantFloat64, + ) + } + + requireStringAttribute( + t, + observer.int64Attributes[0], + connectionStateAttribute, + connectionStateIdle, + ) + requireStringAttribute( + t, + observer.int64Attributes[1], + connectionStateAttribute, + connectionStateUsed, + ) + requireStringAttribute( + t, + observer.int64Attributes[8], + destroyReasonAttribute, + destroyReasonIdleTimeout, + ) + requireStringAttribute( + t, + observer.int64Attributes[9], + destroyReasonAttribute, + destroyReasonLifetime, + ) + + for _, attributes := range observer.int64Attributes { + requireStringAttribute(t, attributes, poolNameAttribute, "test-pool") + requireStringAttribute(t, attributes, "region", "eu") + } + + for _, attributes := range observer.float64Attributes { + requireStringAttribute(t, attributes, poolNameAttribute, "test-pool") + requireStringAttribute(t, attributes, "region", "eu") + } +} + +func TestMetricsRegistrationCloseOnce(t *testing.T) { + t.Parallel() + + registration := ®istrationStub{} + + metricsRegistration := &metricsRegistration{ + registration: registration, + } + + metricsRegistration.Close() + metricsRegistration.Close() + + if registration.calls != 1 { + t.Fatalf( + "unregister calls = %d, want 1", + registration.calls, + ) + } +} + +func newTestPool( + t *testing.T, + metrics xpg.Metrics, +) *xpg.Pool { + t.Helper() + + poolConfig, err := pgxpool.ParseConfig("") + if err != nil { + t.Fatalf("parse pool config: %v", err) + } + + pool, err := xpg.New( + t.Context(), + poolConfig, + xpg.WithName("test-pool"), + xpg.WithMetrics(metrics), + ) + if err != nil { + t.Fatalf("create pool: %v", err) + } + + return pool +} + +func observeAttributes(option metric.ObserveOption) attribute.Set { + return metric.NewObserveConfig( + []metric.ObserveOption{option}, + ).Attributes() +} + +func requireStringAttribute( + t *testing.T, + attributes attribute.Set, + key string, + want string, +) { + t.Helper() + + value, ok := attributes.Value(attribute.Key(key)) + if !ok { + t.Fatalf("attribute %q is missing", key) + } + + if got := value.AsString(); got != want { + t.Fatalf( + "attribute %q = %q, want %q", + key, + got, + want, + ) + } +} + +type instrumentSpec struct { + name string + kind string + description string + unit string +} + +type recordingMeter struct { + noop.Meter + + instruments []instrumentSpec +} + +func (m *recordingMeter) Int64ObservableUpDownCounter( + name string, + options ...metric.Int64ObservableUpDownCounterOption, +) (metric.Int64ObservableUpDownCounter, error) { + config := metric.NewInt64ObservableUpDownCounterConfig(options...) + + m.instruments = append( + m.instruments, + instrumentSpec{ + name: name, + kind: "int64_observable_up_down_counter", + description: config.Description(), + unit: config.Unit(), + }, + ) + + return m.Meter.Int64ObservableUpDownCounter(name, options...) +} + +func (m *recordingMeter) Int64ObservableGauge( + name string, + options ...metric.Int64ObservableGaugeOption, +) (metric.Int64ObservableGauge, error) { + config := metric.NewInt64ObservableGaugeConfig(options...) + + m.instruments = append( + m.instruments, + instrumentSpec{ + name: name, + kind: "int64_observable_gauge", + description: config.Description(), + unit: config.Unit(), + }, + ) + + return m.Meter.Int64ObservableGauge(name, options...) +} + +func (m *recordingMeter) Int64ObservableCounter( + name string, + options ...metric.Int64ObservableCounterOption, +) (metric.Int64ObservableCounter, error) { + config := metric.NewInt64ObservableCounterConfig(options...) + + m.instruments = append( + m.instruments, + instrumentSpec{ + name: name, + kind: "int64_observable_counter", + description: config.Description(), + unit: config.Unit(), + }, + ) + + return m.Meter.Int64ObservableCounter(name, options...) +} + +func (m *recordingMeter) Float64ObservableCounter( + name string, + options ...metric.Float64ObservableCounterOption, +) (metric.Float64ObservableCounter, error) { + config := metric.NewFloat64ObservableCounterConfig(options...) + + m.instruments = append( + m.instruments, + instrumentSpec{ + name: name, + kind: "float64_observable_counter", + description: config.Description(), + unit: config.Unit(), + }, + ) + + return m.Meter.Float64ObservableCounter(name, options...) +} + +type recordingObserver struct { + noop.Observer + + int64Values []int64 + int64Attributes []attribute.Set + float64Values []float64 + float64Attributes []attribute.Set +} + +func (o *recordingObserver) ObserveInt64( + _ metric.Int64Observable, + value int64, + options ...metric.ObserveOption, +) { + o.int64Values = append( + o.int64Values, + value, + ) + + o.int64Attributes = append( + o.int64Attributes, + metric.NewObserveConfig(options).Attributes(), + ) +} + +func (o *recordingObserver) ObserveFloat64( + _ metric.Float64Observable, + value float64, + options ...metric.ObserveOption, +) { + o.float64Values = append( + o.float64Values, + value, + ) + + o.float64Attributes = append( + o.float64Attributes, + metric.NewObserveConfig(options).Attributes(), + ) +} + +type registrationStub struct { + noop.Registration + + calls int +} + +func (r *registrationStub) Unregister() error { + r.calls++ + + return nil +} diff --git a/extra/otelxpg/options.go b/extra/otelxpg/options.go new file mode 100644 index 0000000..786960f --- /dev/null +++ b/extra/otelxpg/options.go @@ -0,0 +1,56 @@ +package otelxpg + +import ( + "go.opentelemetry.io/otel/metric" +) + +// MetricsOption configures OpenTelemetry metrics. +// +// The interface is sealed so options can only be created by this package. +type MetricsOption interface { + apply(*metricsSettings) +} + +// NewMetrics creates an OpenTelemetry metrics integration. +// +// When no MeterProvider is configured, the global OpenTelemetry MeterProvider +// is used. The returned value is safe for concurrent use and reuse across +// multiple pools. +func NewMetrics(options ...MetricsOption) *Metrics { + settings := metricsSettings{} + + for _, option := range options { + if option == nil { + continue + } + + option.apply(&settings) + } + + return &Metrics{ + meterProvider: settings.meterProvider, + } +} + +// WithMeterProvider configures the MeterProvider used for metrics. +// +// A nil provider leaves the global OpenTelemetry MeterProvider in use. The +// caller retains ownership of a non-nil provider and is responsible for +// shutting it down after all instrumented pools have been closed. +func WithMeterProvider(provider metric.MeterProvider) MetricsOption { + return metricsOptionFunc(func(settings *metricsSettings) { + if provider != nil { + settings.meterProvider = provider + } + }) +} + +type metricsOptionFunc func(*metricsSettings) + +func (option metricsOptionFunc) apply(settings *metricsSettings) { + option(settings) +} + +type metricsSettings struct { + meterProvider metric.MeterProvider +} diff --git a/extra/otelxpg/options_test.go b/extra/otelxpg/options_test.go new file mode 100644 index 0000000..df39469 --- /dev/null +++ b/extra/otelxpg/options_test.go @@ -0,0 +1,66 @@ +package otelxpg + +import ( + "testing" + + "go.opentelemetry.io/otel/metric/noop" +) + +func TestNewMetrics(t *testing.T) { + t.Parallel() + + metrics := NewMetrics(nil) + + if metrics == nil { + t.Fatal("expected metrics") + } + + if metrics.meterProvider != nil { + t.Fatal("expected nil meter provider") + } +} + +func TestWithMeterProvider(t *testing.T) { + t.Parallel() + + provider := noop.NewMeterProvider() + + metrics := NewMetrics( + WithMeterProvider(provider), + ) + + if metrics.meterProvider != provider { + t.Fatal("unexpected meter provider") + } +} + +func TestWithMeterProviderLastWins(t *testing.T) { + t.Parallel() + + first := noop.NewMeterProvider() + second := noop.NewMeterProvider() + + metrics := NewMetrics( + WithMeterProvider(first), + WithMeterProvider(second), + ) + + if metrics.meterProvider != second { + t.Fatal("expected last meter provider to win") + } +} + +func TestWithMeterProviderNilIgnored(t *testing.T) { + t.Parallel() + + provider := noop.NewMeterProvider() + + metrics := NewMetrics( + WithMeterProvider(provider), + WithMeterProvider(nil), + ) + + if metrics.meterProvider != provider { + t.Fatal("expected nil meter provider option to be ignored") + } +} diff --git a/extra/otelxpg/pool.go b/extra/otelxpg/pool.go new file mode 100644 index 0000000..634de96 --- /dev/null +++ b/extra/otelxpg/pool.go @@ -0,0 +1,425 @@ +package otelxpg + +import ( + "context" + "errors" + "fmt" + "slices" + + "github.com/mkbeh/xpg" + "go.opentelemetry.io/otel" + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/metric" +) + +const ( + connectionCountMetricName = "db.client.connection.count" + connectionMaxMetricName = "db.client.connection.max" + connectionConstructingMetricName = "xpg.pool.connection.constructing" + connectionAcquireCountMetricName = "xpg.pool.connection.acquire.count" + connectionAcquireTimeMetricName = "xpg.pool.connection.acquire.time" + connectionAcquireCanceledCountMetricName = "xpg.pool.connection.acquire.canceled.count" + connectionAcquireEmptyCountMetricName = "xpg.pool.connection.acquire.empty.count" + connectionAcquireEmptyWaitTimeMetricName = "xpg.pool.connection.acquire.empty.wait_time" + connectionCreateCountMetricName = "xpg.pool.connection.create.count" + connectionDestroyCountMetricName = "xpg.pool.connection.destroy.count" +) + +const ( + dbSystemNameAttribute = "db.system.name" + poolNameAttribute = "db.client.connection.pool.name" + connectionStateAttribute = "db.client.connection.state" + destroyReasonAttribute = "xpg.pool.connection.destroy.reason" +) + +const ( + dbSystemPostgreSQL = "postgresql" + + connectionStateIdle = "idle" + connectionStateUsed = "used" + + destroyReasonIdleTimeout = "idle_timeout" + destroyReasonLifetime = "lifetime" +) + +var _ xpg.Metrics = (*Metrics)(nil) + +type poolMetricInstruments struct { + connectionCount metric.Int64ObservableUpDownCounter + connectionMax metric.Int64ObservableUpDownCounter + constructingConnections metric.Int64ObservableGauge + acquireCount metric.Int64ObservableCounter + acquireTime metric.Float64ObservableCounter + canceledAcquireCount metric.Int64ObservableCounter + emptyAcquireCount metric.Int64ObservableCounter + emptyAcquireWaitTime metric.Float64ObservableCounter + createdConnections metric.Int64ObservableCounter + destroyedConnections metric.Int64ObservableCounter +} + +type poolMetricAttributes struct { + base metric.ObserveOption + idle metric.ObserveOption + used metric.ObserveOption + destroyedIdle metric.ObserveOption + destroyedLifetime metric.ObserveOption +} + +// Register registers metrics for one xpg Pool. +func (m *Metrics) Register(pool *xpg.Pool) (xpg.MetricsRegistration, error) { + if m == nil { + return nil, errors.New("otelxpg: metrics is nil") + } + + if pool == nil { + return nil, errors.New("otelxpg: pool is nil") + } + + meterProvider := m.meterProvider + if meterProvider == nil { + meterProvider = otel.GetMeterProvider() + } + + return registerPoolMetrics(pool, meterProvider) +} + +func registerPoolMetrics(pool *xpg.Pool, provider metric.MeterProvider) (xpg.MetricsRegistration, error) { + meter := provider.Meter(instrumentationName) + + instruments, err := newPoolMetricInstruments(meter) + if err != nil { + return nil, err + } + + attributes := newPoolMetricAttributes(pool.Name(), pool.Labels()) + + registration, err := meter.RegisterCallback( + func(_ context.Context, observer metric.Observer) error { + instruments.observe(observer, pool.Stats(), attributes) + return nil + }, + instruments.observables()..., + ) + if err != nil { + return nil, fmt.Errorf("otelxpg: register pool metrics callback: %w", err) + } + + return &metricsRegistration{ + registration: registration, + }, nil +} + +func newPoolMetricInstruments( + meter metric.Meter, +) (poolMetricInstruments, error) { + var instruments poolMetricInstruments + + var err error + + instruments.connectionCount, err = meter.Int64ObservableUpDownCounter( + connectionCountMetricName, + metric.WithDescription( + "The number of connections currently in the state described by db.client.connection.state.", + ), + metric.WithUnit("{connection}"), + ) + if err != nil { + return poolMetricInstruments{}, fmt.Errorf( + "otelxpg: create %s: %w", + connectionCountMetricName, + err, + ) + } + + instruments.connectionMax, err = meter.Int64ObservableUpDownCounter( + connectionMaxMetricName, + metric.WithDescription( + "The maximum number of open connections allowed by the pool.", + ), + metric.WithUnit("{connection}"), + ) + if err != nil { + return poolMetricInstruments{}, fmt.Errorf( + "otelxpg: create %s: %w", + connectionMaxMetricName, + err, + ) + } + + instruments.constructingConnections, err = meter.Int64ObservableGauge( + connectionConstructingMetricName, + metric.WithDescription( + "The number of connections currently being created by the pool.", + ), + metric.WithUnit("{connection}"), + ) + if err != nil { + return poolMetricInstruments{}, fmt.Errorf( + "otelxpg: create %s: %w", + connectionConstructingMetricName, + err, + ) + } + + instruments.acquireCount, err = meter.Int64ObservableCounter( + connectionAcquireCountMetricName, + metric.WithDescription( + "The cumulative number of successful connection acquires.", + ), + metric.WithUnit("{request}"), + ) + if err != nil { + return poolMetricInstruments{}, fmt.Errorf( + "otelxpg: create %s: %w", + connectionAcquireCountMetricName, + err, + ) + } + + instruments.acquireTime, err = meter.Float64ObservableCounter( + connectionAcquireTimeMetricName, + metric.WithDescription( + "The cumulative time spent acquiring connections.", + ), + metric.WithUnit("s"), + ) + if err != nil { + return poolMetricInstruments{}, fmt.Errorf( + "otelxpg: create %s: %w", + connectionAcquireTimeMetricName, + err, + ) + } + + instruments.canceledAcquireCount, err = meter.Int64ObservableCounter( + connectionAcquireCanceledCountMetricName, + metric.WithDescription( + "The cumulative number of connection acquires canceled by context.", + ), + metric.WithUnit("{request}"), + ) + if err != nil { + return poolMetricInstruments{}, fmt.Errorf( + "otelxpg: create %s: %w", + connectionAcquireCanceledCountMetricName, + err, + ) + } + + instruments.emptyAcquireCount, err = meter.Int64ObservableCounter( + connectionAcquireEmptyCountMetricName, + metric.WithDescription( + "The cumulative number of successful acquires that waited because the pool was empty.", + ), + metric.WithUnit("{request}"), + ) + if err != nil { + return poolMetricInstruments{}, fmt.Errorf( + "otelxpg: create %s: %w", + connectionAcquireEmptyCountMetricName, + err, + ) + } + + instruments.emptyAcquireWaitTime, err = meter.Float64ObservableCounter( + connectionAcquireEmptyWaitTimeMetricName, + metric.WithDescription( + "The cumulative time spent waiting for a connection while the pool was empty.", + ), + metric.WithUnit("s"), + ) + if err != nil { + return poolMetricInstruments{}, fmt.Errorf( + "otelxpg: create %s: %w", + connectionAcquireEmptyWaitTimeMetricName, + err, + ) + } + + instruments.createdConnections, err = meter.Int64ObservableCounter( + connectionCreateCountMetricName, + metric.WithDescription( + "The cumulative number of connections created by the pool.", + ), + metric.WithUnit("{connection}"), + ) + if err != nil { + return poolMetricInstruments{}, fmt.Errorf( + "otelxpg: create %s: %w", + connectionCreateCountMetricName, + err, + ) + } + + instruments.destroyedConnections, err = meter.Int64ObservableCounter( + connectionDestroyCountMetricName, + metric.WithDescription( + "The cumulative number of connections destroyed by pool lifecycle limits.", + ), + metric.WithUnit("{connection}"), + ) + if err != nil { + return poolMetricInstruments{}, fmt.Errorf( + "otelxpg: create %s: %w", + connectionDestroyCountMetricName, + err, + ) + } + + return instruments, nil +} + +func newPoolMetricAttributes(name string, labels map[string]string) poolMetricAttributes { + var base []attribute.KeyValue + + for key, value := range labels { + base = append(base, attribute.String(key, value)) + } + + // System attributes are appended last, so xpg-controlled values win when + // user labels contain duplicate keys. + base = append( + base, + attribute.String( + dbSystemNameAttribute, + dbSystemPostgreSQL, + ), + attribute.String( + poolNameAttribute, + name, + ), + ) + + option := func(extra ...attribute.KeyValue) metric.ObserveOption { + return metric.WithAttributeSet( + attribute.NewSet( + slices.Concat(base, extra)..., + ), + ) + } + + return poolMetricAttributes{ + base: option(), + + idle: option( + attribute.String( + connectionStateAttribute, + connectionStateIdle, + ), + ), + + used: option( + attribute.String( + connectionStateAttribute, + connectionStateUsed, + ), + ), + + destroyedIdle: option( + attribute.String( + destroyReasonAttribute, + destroyReasonIdleTimeout, + ), + ), + + destroyedLifetime: option( + attribute.String( + destroyReasonAttribute, + destroyReasonLifetime, + ), + ), + } +} + +func (i poolMetricInstruments) observe( + observer metric.Observer, + stats xpg.PoolStats, + attributes poolMetricAttributes, +) { + observer.ObserveInt64( + i.connectionCount, + int64(stats.IdleConns), + attributes.idle, + ) + + observer.ObserveInt64( + i.connectionCount, + int64(stats.AcquiredConns), + attributes.used, + ) + + observer.ObserveInt64( + i.connectionMax, + int64(stats.MaxConns), + attributes.base, + ) + + observer.ObserveInt64( + i.constructingConnections, + int64(stats.ConstructingConns), + attributes.base, + ) + + observer.ObserveInt64( + i.acquireCount, + stats.AcquireCount, + attributes.base, + ) + + observer.ObserveFloat64( + i.acquireTime, + stats.AcquireDuration.Seconds(), + attributes.base, + ) + + observer.ObserveInt64( + i.canceledAcquireCount, + stats.CanceledAcquireCount, + attributes.base, + ) + + observer.ObserveInt64( + i.emptyAcquireCount, + stats.EmptyAcquireCount, + attributes.base, + ) + + observer.ObserveFloat64( + i.emptyAcquireWaitTime, + stats.EmptyAcquireWaitTime.Seconds(), + attributes.base, + ) + + observer.ObserveInt64( + i.createdConnections, + stats.NewConnsCount, + attributes.base, + ) + + observer.ObserveInt64( + i.destroyedConnections, + stats.MaxIdleDestroyCount, + attributes.destroyedIdle, + ) + + observer.ObserveInt64( + i.destroyedConnections, + stats.MaxLifetimeDestroyCount, + attributes.destroyedLifetime, + ) +} + +func (i poolMetricInstruments) observables() []metric.Observable { + return []metric.Observable{ + i.connectionCount, + i.connectionMax, + i.constructingConnections, + i.acquireCount, + i.acquireTime, + i.canceledAcquireCount, + i.emptyAcquireCount, + i.emptyAcquireWaitTime, + i.createdConnections, + i.destroyedConnections, + } +} diff --git a/extra/slogxpg/go.mod b/extra/slogxpg/go.mod new file mode 100644 index 0000000..3fe3a33 --- /dev/null +++ b/extra/slogxpg/go.mod @@ -0,0 +1,13 @@ +module github.com/mkbeh/xpg/extra/slogxpg + +go 1.27 + +require github.com/jackc/pgx/v5 v5.10.0 + +require ( + github.com/jackc/pgpassfile v1.0.0 // indirect + github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect + github.com/jackc/puddle/v2 v2.2.2 // indirect + golang.org/x/sync v0.17.0 // indirect + golang.org/x/text v0.29.0 // indirect +) diff --git a/extra/slogxpg/go.sum b/extra/slogxpg/go.sum new file mode 100644 index 0000000..c0e505b --- /dev/null +++ b/extra/slogxpg/go.sum @@ -0,0 +1,26 @@ +github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= +github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM= +github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg= +github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 h1:iCEnooe7UlwOQYpKFhBabPMi4aNAfoODPEFNiAnClxo= +github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761/go.mod h1:5TJZWKEWniPve33vlWYSoGYefn3gLQRzjfDlhSJ9ZKM= +github.com/jackc/pgx/v5 v5.10.0 h1:VhSvgU2jSli8o3AqIEOTJr7rZwAEUVo4E4XhR94Zfr0= +github.com/jackc/pgx/v5 v5.10.0/go.mod h1:mal1tBGAFfLHvZzaYh77YS/eC6IX9OWbRV1QIIM0Jn4= +github.com/jackc/puddle/v2 v2.2.2 h1:PR8nw+E/1w0GLuRFSmiioY6UooMp6KJv0/61nB7icHo= +github.com/jackc/puddle/v2 v2.2.2/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= +github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= +github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= +github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U= +github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U= +golang.org/x/sync v0.17.0 h1:l60nONMj9l5drqw6jlhIELNv9I0A4OFgRsG9k2oT9Ug= +golang.org/x/sync v0.17.0/go.mod h1:9KTHXmSnoGruLpwFjVSX0lNNA75CykiMECbovNTZqGI= +golang.org/x/text v0.29.0 h1:1neNs90w9YzJ9BocxfsQNHKuAT4pkghyXc4nhZ6sJvk= +golang.org/x/text v0.29.0/go.mod h1:7MhJOA9CD2qZyOKYazxdYMF85OwPdEr9jTtBpO7ydH4= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/extra/slogxpg/logger.go b/extra/slogxpg/logger.go new file mode 100644 index 0000000..de5ec2c --- /dev/null +++ b/extra/slogxpg/logger.go @@ -0,0 +1,64 @@ +package slogxpg + +import ( + "context" + "log/slog" + + "github.com/jackc/pgx/v5/tracelog" +) + +type adapter struct { + logger *slog.Logger +} + +var _ tracelog.Logger = (*adapter)(nil) + +// New adapts logger to pgx tracelog.Logger. +// +// New returns nil when logger is nil. +func New(logger *slog.Logger) tracelog.Logger { + if logger == nil { + return nil + } + + return &adapter{ + logger: logger, + } +} + +func (a *adapter) Log( + ctx context.Context, + level tracelog.LogLevel, + msg string, + data map[string]any, +) { + attrs := make([]slog.Attr, 0, len(data)) + + for key, value := range data { + attrs = append(attrs, slog.Any(key, value)) + } + + a.logger.LogAttrs( + ctx, + slogLevel(level), + msg, //nolint:sloglint // pgx provides the log message dynamically. + attrs..., + ) +} + +func slogLevel(level tracelog.LogLevel) slog.Level { + switch level { + case tracelog.LogLevelTrace: + return slog.LevelDebug - 1 + case tracelog.LogLevelDebug: + return slog.LevelDebug + case tracelog.LogLevelInfo: + return slog.LevelInfo + case tracelog.LogLevelWarn: + return slog.LevelWarn + case tracelog.LogLevelError: + return slog.LevelError + default: + return slog.LevelError + } +} diff --git a/extra/slogxpg/logger_test.go b/extra/slogxpg/logger_test.go new file mode 100644 index 0000000..ccc8350 --- /dev/null +++ b/extra/slogxpg/logger_test.go @@ -0,0 +1,166 @@ +package slogxpg + +import ( + "context" + "log/slog" + "testing" + + "github.com/jackc/pgx/v5/tracelog" +) + +func TestNewNil(t *testing.T) { + t.Parallel() + + if logger := New(nil); logger != nil { + t.Fatal("expected nil logger") + } +} + +func TestSlogLevel(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + pgx tracelog.LogLevel + want slog.Level + }{ + { + name: "trace", + pgx: tracelog.LogLevelTrace, + want: slog.LevelDebug - 1, + }, + { + name: "debug", + pgx: tracelog.LogLevelDebug, + want: slog.LevelDebug, + }, + { + name: "info", + pgx: tracelog.LogLevelInfo, + want: slog.LevelInfo, + }, + { + name: "warn", + pgx: tracelog.LogLevelWarn, + want: slog.LevelWarn, + }, + { + name: "error", + pgx: tracelog.LogLevelError, + want: slog.LevelError, + }, + { + name: "unknown", + pgx: tracelog.LogLevel(255), + want: slog.LevelError, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + + if got := slogLevel(test.pgx); got != test.want { + t.Fatalf("slogLevel(%v) = %v, want %v", test.pgx, got, test.want) + } + }) + } +} + +func TestAdapterLog(t *testing.T) { + t.Parallel() + + handler := &captureHandler{} + logger := New(slog.New(handler)) + + logger.Log( + context.Background(), + tracelog.LogLevelInfo, + "query", + map[string]any{ + "sql": "select 1", + "args": 1, + }, + ) + + record := handler.record + + if record.Message != "query" { + t.Fatalf("message = %q, want %q", record.Message, "query") + } + + if record.Level != slog.LevelInfo { + t.Fatalf("level = %v, want %v", record.Level, slog.LevelInfo) + } + + attrs := recordAttrs(&record) + + if attrs["sql"] != "select 1" { + t.Fatalf("sql = %v, want %q", attrs["sql"], "select 1") + } + + if attrs["args"] != int64(1) { + t.Fatalf("args = %v, want %v", attrs["args"], int64(1)) + } +} + +func TestAdapterUnknownLevel(t *testing.T) { + t.Parallel() + + handler := &captureHandler{} + logger := New(slog.New(handler)) + + logger.Log( + context.Background(), + tracelog.LogLevel(255), + "unknown", + nil, + ) + + record := handler.record + + if record.Message != "unknown" { + t.Fatalf("message = %q, want %q", record.Message, "unknown") + } + + if record.Level != slog.LevelError { + t.Fatalf("level = %v, want %v", record.Level, slog.LevelError) + } +} + +type captureHandler struct { + record slog.Record +} + +func (h *captureHandler) Enabled(context.Context, slog.Level) bool { + return true +} + +func (h *captureHandler) Handle( + _ context.Context, + record slog.Record, //nolint:gocritic // slog.Handler requires slog.Record by value. +) error { + h.record = record.Clone() + + return nil +} + +func (h *captureHandler) WithAttrs([]slog.Attr) slog.Handler { + return h +} + +func (h *captureHandler) WithGroup(string) slog.Handler { + return h +} + +func recordAttrs(record *slog.Record) map[string]any { + attrs := make(map[string]any, record.NumAttrs()) + + record.Attrs(func(attr slog.Attr) bool { + attrs[attr.Key] = attr.Value.Any() + + return true + }) + + return attrs +} diff --git a/go.mod b/go.mod index f94ec39..e7fcb52 100644 --- a/go.mod +++ b/go.mod @@ -1,38 +1,13 @@ module github.com/mkbeh/xpg -go 1.26 +go 1.27 -require ( - github.com/Masterminds/squirrel v1.5.4 - github.com/exaring/otelpgx v0.11.1 - github.com/golang-migrate/migrate/v4 v4.19.1 - github.com/google/uuid v1.6.0 - github.com/jackc/pgerrcode v0.0.0-20250907135507-afb5586c32a6 - github.com/jackc/pgx/v5 v5.10.0 - github.com/prometheus/client_golang v1.23.2 - go.opentelemetry.io/otel v1.44.0 - go.opentelemetry.io/otel/trace v1.44.0 -) +require github.com/jackc/pgx/v5 v5.10.0 require ( - github.com/beorn7/perks v1.0.1 // indirect - github.com/cespare/xxhash/v2 v2.3.0 // indirect - github.com/go-logr/logr v1.4.3 // indirect - github.com/go-logr/stdr v1.2.2 // indirect github.com/jackc/pgpassfile v1.0.0 // indirect github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect github.com/jackc/puddle/v2 v2.2.2 // indirect - github.com/lann/builder v0.0.0-20180802200727-47ae307949d0 // indirect - github.com/lann/ps v0.0.0-20150810152359-62de8c46ede0 // indirect - github.com/lib/pq v1.12.3 // indirect - github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect - github.com/prometheus/client_model v0.6.2 // indirect - github.com/prometheus/common v0.69.0 // indirect - github.com/prometheus/procfs v0.20.1 // indirect - go.opentelemetry.io/auto/sdk v1.2.1 // indirect - go.opentelemetry.io/otel/metric v1.44.0 // indirect - golang.org/x/sync v0.21.0 // indirect - golang.org/x/sys v0.46.0 // indirect - golang.org/x/text v0.38.0 // indirect - google.golang.org/protobuf v1.36.11 // indirect + golang.org/x/sync v0.22.0 // indirect + golang.org/x/text v0.41.0 // indirect ) diff --git a/go.sum b/go.sum index 6b66baa..a277497 100644 --- a/go.sum +++ b/go.sum @@ -1,136 +1,21 @@ -github.com/Azure/go-ansiterm v0.0.0-20230124172434-306776ec8161 h1:L/gRVlceqvL25UVaW/CKtUDjefjrs0SPonmDGUVOYP0= -github.com/Azure/go-ansiterm v0.0.0-20230124172434-306776ec8161/go.mod h1:xomTg63KZ2rFqZQzSB4Vz2SUXa1BpHTVz9L5PTmPC4E= -github.com/Masterminds/squirrel v1.5.4 h1:uUcX/aBc8O7Fg9kaISIUsHXdKuqehiXAMQTYX8afzqM= -github.com/Masterminds/squirrel v1.5.4/go.mod h1:NNaOrjSoIDfDA40n7sr2tPNZRfjzjA400rg+riTZj10= -github.com/Microsoft/go-winio v0.6.2 h1:F2VQgta7ecxGYO8k3ZZz3RS8fVIXVxONVUPlNERoyfY= -github.com/Microsoft/go-winio v0.6.2/go.mod h1:yd8OoFMLzJbo9gZq8j5qaps8bJ9aShtEA8Ipt1oGCvU= -github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM= -github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw= -github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= -github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= -github.com/containerd/errdefs v1.0.0 h1:tg5yIfIlQIrxYtu9ajqY42W3lpS19XqdxRQeEwYG8PI= -github.com/containerd/errdefs v1.0.0/go.mod h1:+YBYIdtsnF4Iw6nWZhJcqGSg/dwvV7tyJ/kCkyJ2k+M= -github.com/containerd/errdefs/pkg v0.3.0 h1:9IKJ06FvyNlexW690DXuQNx2KA2cUJXx151Xdx3ZPPE= -github.com/containerd/errdefs/pkg v0.3.0/go.mod h1:NJw6s9HwNuRhnjJhM7pylWwMyAkmCQvQ4GpJHEqRLVk= github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= -github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= -github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM= -github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= -github.com/dhui/dktest v0.4.6 h1:+DPKyScKSEp3VLtbMDHcUq6V5Lm5zfZZVb0Sk7Ahom4= -github.com/dhui/dktest v0.4.6/go.mod h1:JHTSYDtKkvFNFHJKqCzVzqXecyv+tKt8EzceOmQOgbU= -github.com/distribution/reference v0.6.0 h1:0IXCQ5g4/QMHHkarYzh5l+u8T3t73zM5QvfrDyIgxBk= -github.com/distribution/reference v0.6.0/go.mod h1:BbU0aIcezP1/5jX/8MP0YiH4SdvB5Y4f/wlDRiLyi3E= -github.com/docker/docker v28.3.3+incompatible h1:Dypm25kh4rmk49v1eiVbsAtpAsYURjYkaKubwuBdxEI= -github.com/docker/docker v28.3.3+incompatible/go.mod h1:eEKB0N0r5NX/I1kEveEz05bcu8tLC/8azJZsviup8Sk= -github.com/docker/go-connections v0.5.0 h1:USnMq7hx7gwdVZq1L49hLXaFtUdTADjXGp+uj1Br63c= -github.com/docker/go-connections v0.5.0/go.mod h1:ov60Kzw0kKElRwhNs9UlUHAE/F9Fe6GLaXnqyDdmEXc= -github.com/docker/go-units v0.5.0 h1:69rxXcBk27SvSaaxTtLh/8llcHD8vYHT7WSdRZ/jvr4= -github.com/docker/go-units v0.5.0/go.mod h1:fgPhTUdO+D/Jk86RDLlptpiXQzgHJF7gydDDbaIK4Dk= -github.com/exaring/otelpgx v0.11.1 h1:pE79fIg/qh/Lpu00kvswFC5dKfqyJJhMJ4Y4N3w5Lj4= -github.com/exaring/otelpgx v0.11.1/go.mod h1:3OojrUKhhy3lTbYIMBijP3YjMey/jo14eHAW5cXcUdk= -github.com/felixge/httpsnoop v1.0.4 h1:NFTV2Zj1bL4mc9sqWACXbQFVBBg2W3GPvqp8/ESS2Wg= -github.com/felixge/httpsnoop v1.0.4/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U= -github.com/go-logr/logr v1.2.2/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A= -github.com/go-logr/logr v1.4.3 h1:CjnDlHq8ikf6E492q6eKboGOC0T8CDaOvkHCIg8idEI= -github.com/go-logr/logr v1.4.3/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY= -github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag= -github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE= -github.com/gogo/protobuf v1.3.2 h1:Ov1cvc58UF3b5XjBnZv7+opcTcQFZebYjWzi34vdm4Q= -github.com/gogo/protobuf v1.3.2/go.mod h1:P1XiOD3dCwIKUDQYPy72D8LYyHL2YPYrpS2s69NZV8Q= -github.com/golang-migrate/migrate/v4 v4.19.1 h1:OCyb44lFuQfYXYLx1SCxPZQGU7mcaZ7gH9yH4jSFbBA= -github.com/golang-migrate/migrate/v4 v4.19.1/go.mod h1:CTcgfjxhaUtsLipnLoQRWCrjYXycRz/g5+RWDuYgPrE= -github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= -github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= -github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= -github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= -github.com/jackc/pgerrcode v0.0.0-20250907135507-afb5586c32a6 h1:D/V0gu4zQ3cL2WKeVNVM4r2gLxGGf6McLwgXzRTo2RQ= -github.com/jackc/pgerrcode v0.0.0-20250907135507-afb5586c32a6/go.mod h1:a/s9Lp5W7n/DD0VrVoyJ00FbP2ytTPDVOivvn2bMlds= github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM= github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg= github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 h1:iCEnooe7UlwOQYpKFhBabPMi4aNAfoODPEFNiAnClxo= github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761/go.mod h1:5TJZWKEWniPve33vlWYSoGYefn3gLQRzjfDlhSJ9ZKM= -github.com/jackc/pgx/v5 v5.9.2 h1:3ZhOzMWnR4yJ+RW1XImIPsD1aNSz4T4fyP7zlQb56hw= -github.com/jackc/pgx/v5 v5.9.2/go.mod h1:mal1tBGAFfLHvZzaYh77YS/eC6IX9OWbRV1QIIM0Jn4= github.com/jackc/pgx/v5 v5.10.0 h1:VhSvgU2jSli8o3AqIEOTJr7rZwAEUVo4E4XhR94Zfr0= github.com/jackc/pgx/v5 v5.10.0/go.mod h1:mal1tBGAFfLHvZzaYh77YS/eC6IX9OWbRV1QIIM0Jn4= github.com/jackc/puddle/v2 v2.2.2 h1:PR8nw+E/1w0GLuRFSmiioY6UooMp6KJv0/61nB7icHo= github.com/jackc/puddle/v2 v2.2.2/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4= -github.com/klauspost/compress v1.18.0 h1:c/Cqfb0r+Yi+JtIEq73FWXVkRonBlf0CRNYc8Zttxdo= -github.com/klauspost/compress v1.18.0/go.mod h1:2Pp+KzxcywXVXMr50+X0Q/Lsb43OQHYWRCY2AiWywWQ= -github.com/kylelemons/godebug v1.1.0 h1:RPNrshWIDI6G2gRW9EHilWtl7Z6Sb1BR0xunSBf0SNc= -github.com/kylelemons/godebug v1.1.0/go.mod h1:9/0rRGxNHcop5bhtWyNeEfOS8JIWk580+fNqagV/RAw= -github.com/lann/builder v0.0.0-20180802200727-47ae307949d0 h1:SOEGU9fKiNWd/HOJuq6+3iTQz8KNCLtVX6idSoTLdUw= -github.com/lann/builder v0.0.0-20180802200727-47ae307949d0/go.mod h1:dXGbAdH5GtBTC4WfIxhKZfyBF/HBFgRZSWwZ9g/He9o= -github.com/lann/ps v0.0.0-20150810152359-62de8c46ede0 h1:P6pPBnrTSX3DEVR4fDembhRWSsG5rVo6hYhAB/ADZrk= -github.com/lann/ps v0.0.0-20150810152359-62de8c46ede0/go.mod h1:vmVJ0l/dxyfGW6FmdpVm2joNMFikkuWg0EoCKLGUMNw= -github.com/lib/pq v1.12.3 h1:tTWxr2YLKwIvK90ZXEw8GP7UFHtcbTtty8zsI+YjrfQ= -github.com/lib/pq v1.12.3/go.mod h1:/p+8NSbOcwzAEI7wiMXFlgydTwcgTr3OSKMsD2BitpA= -github.com/moby/docker-image-spec v1.3.1 h1:jMKff3w6PgbfSa69GfNg+zN/XLhfXJGnEx3Nl2EsFP0= -github.com/moby/docker-image-spec v1.3.1/go.mod h1:eKmb5VW8vQEh/BAr2yvVNvuiJuY6UIocYsFu/DxxRpo= -github.com/moby/term v0.5.0 h1:xt8Q1nalod/v7BqbG21f8mQPqH+xAaC9C3N3wfWbVP0= -github.com/moby/term v0.5.0/go.mod h1:8FzsFHVUBGZdbDsJw/ot+X+d5HLUbvklYLJ9uGfcI3Y= -github.com/morikuni/aec v1.0.0 h1:nP9CBfwrvYnBRgY6qfDQkygYDmYwOilePFkwzv4dU8A= -github.com/morikuni/aec v1.0.0/go.mod h1:BbKIizmSmc5MMPqRYbxO4ZU0S0+P200+tUnFx7PXmsc= -github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 h1:C3w9PqII01/Oq1c1nUAm88MOHcQC9l5mIlSMApZMrHA= -github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822/go.mod h1:+n7T8mK8HuQTcFwEeznm/DIxMOiR9yIdICNftLE1DvQ= -github.com/opencontainers/go-digest v1.0.0 h1:apOUWs51W5PlhuyGyz9FCeeBIOUDA/6nW8Oi/yOhh5U= -github.com/opencontainers/go-digest v1.0.0/go.mod h1:0JzlMkj0TRzQZfJkVvzbP0HBR3IKzErnv2BNG4W4MAM= -github.com/opencontainers/image-spec v1.1.0 h1:8SG7/vwALn54lVB/0yZ/MMwhFrPYtpEHQb2IpWsCzug= -github.com/opencontainers/image-spec v1.1.0/go.mod h1:W4s4sFTMaBeK1BQLXbG4AdM2szdn85PY75RI83NrTrM= -github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4= -github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= -github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRIccs7FGNTlIRMkT8wgtp5eCXdBlqhYGL6U= -github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= -github.com/prometheus/client_golang v1.23.2 h1:Je96obch5RDVy3FDMndoUsjAhG5Edi49h0RJWRi/o0o= -github.com/prometheus/client_golang v1.23.2/go.mod h1:Tb1a6LWHB3/SPIzCoaDXI4I8UHKeFTEQ1YCr+0Gyqmg= -github.com/prometheus/client_model v0.6.2 h1:oBsgwpGs7iVziMvrGhE53c/GrLUsZdHnqNwqPLxwZyk= -github.com/prometheus/client_model v0.6.2/go.mod h1:y3m2F6Gdpfy6Ut/GBsUqTWZqCUvMVzSfMLjcu6wAwpE= -github.com/prometheus/common v0.68.0 h1:8rQJvQmYltsR2L7h8Zw0Iyj8WYNNmpwikoQTZXwfVeA= -github.com/prometheus/common v0.68.0/go.mod h1:4soH+U8yJSROk7OJ//hmTiWKsxapv6zRGgTt3keN8gQ= -github.com/prometheus/common v0.69.0 h1:OA85nJQS/T/MaYh/Q2CcgDKSGWqNIgrBDvDH85CuiNk= -github.com/prometheus/common v0.69.0/go.mod h1:ZzL3f6u94qUxh9p+tJTrF+FvBS1XXbbRAZCQkytAL0Y= -github.com/prometheus/procfs v0.20.1 h1:XwbrGOIplXW/AU3YhIhLODXMJYyC1isLFfYCsTEycfc= -github.com/prometheus/procfs v0.20.1/go.mod h1:o9EMBZGRyvDrSPH1RqdxhojkuXstoe4UlK79eF5TGGo= github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= -github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs= github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= -github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U= -github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U= -go.opentelemetry.io/auto/sdk v1.2.1 h1:jXsnJ4Lmnqd11kwkBV2LgLoFMZKizbCi5fNZ/ipaZ64= -go.opentelemetry.io/auto/sdk v1.2.1/go.mod h1:KRTj+aOaElaLi+wW1kO/DZRXwkF4C5xPbEe3ZiIhN7Y= -go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.61.0 h1:F7Jx+6hwnZ41NSFTO5q4LYDtJRXBf2PD0rNBkeB/lus= -go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.61.0/go.mod h1:UHB22Z8QsdRDrnAtX4PntOl36ajSxcdUMt1sF7Y6E7Q= -go.opentelemetry.io/otel v1.44.0 h1:JjwHmHpA4iZ3wBxluu2fbbE7j4kqlE8jXyAyPXH7HqU= -go.opentelemetry.io/otel v1.44.0/go.mod h1:BMgjTHL9WPRlRjL2oZCBTL4whCGtXch2H4BhOPIAyYc= -go.opentelemetry.io/otel/metric v1.44.0 h1:1w0gILTcHdr3YI+ixLyjemwrVnsMURbTZFrSYCdDdmc= -go.opentelemetry.io/otel/metric v1.44.0/go.mod h1:8O7hanEPBNgEMmybD3s2VBKcgWOCsA6tzHBPODAiquo= -go.opentelemetry.io/otel/sdk v1.43.0 h1:pi5mE86i5rTeLXqoF/hhiBtUNcrAGHLKQdhg4h4V9Dg= -go.opentelemetry.io/otel/sdk v1.43.0/go.mod h1:P+IkVU3iWukmiit/Yf9AWvpyRDlUeBaRg6Y+C58QHzg= -go.opentelemetry.io/otel/sdk/metric v1.43.0 h1:S88dyqXjJkuBNLeMcVPRFXpRw2fuwdvfCGLEo89fDkw= -go.opentelemetry.io/otel/sdk/metric v1.43.0/go.mod h1:C/RJtwSEJ5hzTiUz5pXF1kILHStzb9zFlIEe85bhj6A= -go.opentelemetry.io/otel/trace v1.44.0 h1:jxF5CsGYCe74MCRx2X4g7WsY/VBKRqqpNvXlX/6gtIk= -go.opentelemetry.io/otel/trace v1.44.0/go.mod h1:oLl1jrMQAVo6v3GAggN+1VH9VIz9iUSvW53sW1Q8PIE= -go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto= -go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE= -go.yaml.in/yaml/v2 v2.4.4 h1:tuyd0P+2Ont/d6e2rl3be67goVK4R6deVxCUX5vyPaQ= -go.yaml.in/yaml/v2 v2.4.4/go.mod h1:gMZqIpDtDqOfM0uNfy0SkpRhvUryYH0Z6wdMYcacYXQ= -golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4= -golang.org/x/sync v0.20.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= -golang.org/x/sync v0.21.0 h1:HLII4xRRTtCRkxYp4HNFF0Js/Og6q2i++KXbg0gHCwM= -golang.org/x/sync v0.21.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= -golang.org/x/sys v0.45.0 h1:dO4czNzziLiiXplLQgBCEpCvXQ3dnkn0SdaZSYdQ+FY= -golang.org/x/sys v0.45.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= -golang.org/x/sys v0.46.0 h1:noSf2Fq6F8DBgS+LysIkx7rIExoNHJsxOAtPp4rthXw= -golang.org/x/sys v0.46.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= -golang.org/x/text v0.37.0 h1:Cqjiwd9eSg8e0QAkyCaQTNHFIIzWtidPahFWR83rTrc= -golang.org/x/text v0.37.0/go.mod h1:a5sjxXGs9hsn/AJVwuElvCAo9v8QYLzvavO5z2PiM38= -golang.org/x/text v0.38.0 h1:sXmwo9DwP3OK9EZ7PqAdaooSGozfl/3a6/xJcbzPRhE= -golang.org/x/text v0.38.0/go.mod h1:YXZt3QhHUKYT53r2lLKFIVi6Ao1jdzrTR/KQ09qyxF4= -google.golang.org/protobuf v1.36.11 h1:fV6ZwhNocDyBLK0dj+fg8ektcVegBBuEolpbTQyBNVE= -google.golang.org/protobuf v1.36.11/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco= +github.com/stretchr/testify v1.12.1 h1:EuwCh5fleGS7H32xRwO3wRGT7DxrDhLAT6FF8MpWDWE= +go.yaml.in/yaml/v3 v3.0.5 h1:N6y/pJk8buWs9NY5ERU2HSMfm+IuD/OtfdAnq6kESPw= +golang.org/x/sync v0.22.0 h1:SZjpbeLmrCk4xhRSZFNZW5gFUeCeFgjekvI/+gfScek= +golang.org/x/sync v0.22.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= +golang.org/x/text v0.41.0 h1:vz/seA0lnX87Othu2f/0L24RcgrXD9/YFTSuGjj3rH8= +golang.org/x/text v0.41.0/go.mod h1:jvf1O8ajNzZqhSrQBPbutR/EB83Cc0CFrezNQIwbb5M= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= -gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= -gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/internal/pkg/pgxpoolcollector/v5/collector.go b/internal/pkg/pgxpoolcollector/v5/collector.go deleted file mode 100644 index f9cd7ed..0000000 --- a/internal/pkg/pgxpoolcollector/v5/collector.go +++ /dev/null @@ -1,181 +0,0 @@ -// Package v5 provides prometheus plug-in metrics for a pgx client. -// -// This package tracks the following metrics under the following names: -// -// #ns_postgres_acquire_count{} -// #ns_postgres_acquire_duration{} -// #ns_postgres_acquired_conns{} -// #ns_postgres_canceled_acquire_count{} -// #ns_postgres_constructing_conns{} -// #ns_postgres_empty_acquire_count{} -// #ns_postgres_idle_conns{} -// #ns_postgres_max_conns{} -// #ns_postgres_total_conns{} -// -// Labels list: -// client_id=#{client_id} -// client_kind=#{master/replica} -// db=#{db} -// shard_id=#{shard_id} - -package v5 - -import ( - "github.com/jackc/pgx/v5/pgxpool" - "github.com/prometheus/client_golang/prometheus" -) - -// StatsGetter is an interface that gets sql.DBStats. -// It's implemented by e.g. *sql.DB or *sqlx.DB. -type StatsGetter interface { - Stat() *pgxpool.Stat -} - -// StatsCollector implements the prometheus.Collector interface. -type StatsCollector struct { - sg StatsGetter - - // descriptions of exported metrics - acquireCount *prometheus.Desc - acquireDuration *prometheus.Desc - acquiredConns *prometheus.Desc - canceledAcquireCount *prometheus.Desc - constructingConns *prometheus.Desc - emptyAcquireCount *prometheus.Desc - idleConns *prometheus.Desc - maxConns *prometheus.Desc - totalConns *prometheus.Desc -} - -// NewStatsCollector creates a new StatsCollector. -func NewStatsCollector(namespace, subsystem string, constLabels prometheus.Labels, sg StatsGetter) *StatsCollector { - return &StatsCollector{ - sg: sg, - acquireCount: prometheus.NewDesc( - prometheus.BuildFQName(namespace, subsystem, "acquire_count"), - "Cumulative count of successful acquires from the pool.", - nil, - constLabels, - ), - acquireDuration: prometheus.NewDesc( - prometheus.BuildFQName(namespace, subsystem, "acquire_duration"), - "Total duration of all successful acquires from the pool.", - nil, - constLabels, - ), - acquiredConns: prometheus.NewDesc( - prometheus.BuildFQName(namespace, subsystem, "acquired_conns"), - "Number of currently acquired connections in the pool.", - nil, - constLabels, - ), - canceledAcquireCount: prometheus.NewDesc( - prometheus.BuildFQName(namespace, subsystem, "canceled_acquire_count"), - "Cumulative count of acquires from the pool that were canceled by a context.", - nil, - constLabels, - ), - constructingConns: prometheus.NewDesc( - prometheus.BuildFQName(namespace, subsystem, "constructing_conns"), - "Number of conns with construction in progress in the pool.", - nil, - constLabels, - ), - emptyAcquireCount: prometheus.NewDesc( - prometheus.BuildFQName(namespace, subsystem, "empty_acquire_count"), - "Cumulative count of successful acquires from the pool that waited for a resource to be released or constructed because the pool was empty.", - nil, - constLabels, - ), - idleConns: prometheus.NewDesc( - prometheus.BuildFQName(namespace, subsystem, "idle_conns"), - "Number of currently idle conns in the pool.", - nil, - constLabels, - ), - maxConns: prometheus.NewDesc( - prometheus.BuildFQName(namespace, subsystem, "max_conns"), - "Maximum size of the pool.", - nil, - constLabels, - ), - totalConns: prometheus.NewDesc( - prometheus.BuildFQName(namespace, subsystem, "total_conns"), - "Total number of resources currently in the pool. The value is the sum of ConstructingConns, AcquiredConns, and IdleConns.", - nil, - constLabels, - ), - } -} - -// Describe implements the prometheus.Collector interface. -func (c StatsCollector) Describe(ch chan<- *prometheus.Desc) { - ch <- c.acquireCount - ch <- c.acquireDuration - ch <- c.acquiredConns - ch <- c.canceledAcquireCount - ch <- c.constructingConns - ch <- c.emptyAcquireCount - ch <- c.idleConns - ch <- c.maxConns - ch <- c.totalConns -} - -// Collect implements the prometheus.Collector interface. -func (c StatsCollector) Collect(ch chan<- prometheus.Metric) { - stats := c.sg.Stat() - - ch <- prometheus.MustNewConstMetric( - c.acquireCount, - prometheus.GaugeValue, - float64(stats.AcquireCount()), - ) - - ch <- prometheus.MustNewConstMetric( - c.acquireDuration, - prometheus.GaugeValue, - float64(stats.AcquireDuration()), - ) - - ch <- prometheus.MustNewConstMetric( - c.acquiredConns, - prometheus.GaugeValue, - float64(stats.AcquiredConns()), - ) - - ch <- prometheus.MustNewConstMetric( - c.canceledAcquireCount, - prometheus.GaugeValue, - float64(stats.CanceledAcquireCount()), - ) - - ch <- prometheus.MustNewConstMetric( - c.constructingConns, - prometheus.GaugeValue, - float64(stats.ConstructingConns()), - ) - - ch <- prometheus.MustNewConstMetric( - c.emptyAcquireCount, - prometheus.GaugeValue, - float64(stats.EmptyAcquireCount()), - ) - - ch <- prometheus.MustNewConstMetric( - c.idleConns, - prometheus.GaugeValue, - float64(stats.IdleConns()), - ) - - ch <- prometheus.MustNewConstMetric( - c.maxConns, - prometheus.GaugeValue, - float64(stats.MaxConns()), - ) - - ch <- prometheus.MustNewConstMetric( - c.totalConns, - prometheus.GaugeValue, - float64(stats.TotalConns()), - ) -} diff --git a/internal/pkg/pgxslog/adapter.go b/internal/pkg/pgxslog/adapter.go deleted file mode 100644 index 9651cf4..0000000 --- a/internal/pkg/pgxslog/adapter.go +++ /dev/null @@ -1,63 +0,0 @@ -package pgxslog - -import ( - "context" - "log/slog" - - "github.com/jackc/pgx/v5/tracelog" -) - -type Logger struct { - l *slog.Logger -} - -func NewLogger(l *slog.Logger) *Logger { - return &Logger{l} -} - -func (l *Logger) Log(ctx context.Context, level tracelog.LogLevel, msg string, data map[string]interface{}) { - attrs := make([]slog.Attr, 0, len(data)) - for k, v := range data { - attrs = append(attrs, slog.Any(k, v)) - } - - var lvl slog.Level - switch level { - case tracelog.LogLevelTrace: - lvl = slog.LevelDebug - 1 - attrs = append(attrs, slog.Any("pgx_log_level", level)) - case tracelog.LogLevelDebug: - lvl = slog.LevelDebug - case tracelog.LogLevelInfo: - lvl = slog.LevelInfo - case tracelog.LogLevelWarn: - lvl = slog.LevelWarn - case tracelog.LogLevelError: - lvl = slog.LevelError - default: - lvl = slog.LevelError - attrs = append(attrs, slog.Any("invalid_pgx_log_level", level)) - } - //nolint:sloglint // there is no other option to pass the message - l.l.LogAttrs(ctx, lvl, msg, attrs...) -} - -// all key constants must be defined with the "Key" suffix. - -// Constants and constructors to standardize the keys used in logs. - -const ( - // errorKey - to pass error to log. - errorKey string = "error" - // componentKey - to identify the component that is logging (e.g.: "kafka-consumer") - componentKey = "component" -) - -// Error for passing error to log. -func Error(err error) slog.Attr { return slog.Any(errorKey, err) } - -// Component to identify the component that is logging (e.g.: "kafka-consumer"). -// here any distinct level of your application abstraction can be used. -// it's not mandatory associated with kind of an external connection, -// e.g.: "location-event-processor" is also applicable for the field. -func Component(component string) slog.Attr { return slog.String(componentKey, component) } diff --git a/internal/pkg/pgxtracer/tracer.go b/internal/pkg/pgxtracer/tracer.go deleted file mode 100644 index b92fec0..0000000 --- a/internal/pkg/pgxtracer/tracer.go +++ /dev/null @@ -1,40 +0,0 @@ -// Package pgxtracer is a chaining QueryTracer for pgx. -package pgxtracer - -import ( - "context" - - "github.com/jackc/pgx/v5" -) - -// QueryTracer traces Query, QueryRow, and Exec. -type QueryTracer interface { - // TraceQueryStart is called at the beginning of Query, QueryRow, and Exec calls. The returned context is used for the - // rest of the call and will be passed to TraceQueryEnd. - TraceQueryStart(ctx context.Context, conn *pgx.Conn, data pgx.TraceQueryStartData) context.Context - - TraceQueryEnd(ctx context.Context, conn *pgx.Conn, data pgx.TraceQueryEndData) -} - -type tracer struct { - tracers []QueryTracer -} - -func (t *tracer) TraceQueryStart(ctx context.Context, conn *pgx.Conn, data pgx.TraceQueryStartData) context.Context { - for i := range t.tracers { - ctx = t.tracers[i].TraceQueryStart(ctx, conn, data) - } - return ctx -} - -func (t *tracer) TraceQueryEnd(ctx context.Context, conn *pgx.Conn, data pgx.TraceQueryEndData) { - for i := range t.tracers { - t.tracers[i].TraceQueryEnd(ctx, conn, data) - } -} - -func New(tracers ...QueryTracer) QueryTracer { - return &tracer{ - tracers: tracers, - } -} diff --git a/metrics.go b/metrics.go new file mode 100644 index 0000000..3d08aa8 --- /dev/null +++ b/metrics.go @@ -0,0 +1,16 @@ +package xpg + +// Metrics registers metrics for a Pool. +// +// Implementations must be safe to reuse across multiple pools. Register is +// called after the underlying pgxpool.Pool has been created. +type Metrics interface { + Register(pool *Pool) (MetricsRegistration, error) +} + +// MetricsRegistration represents a metrics registration for one Pool. +// +// Close is called once before the underlying pgxpool.Pool is closed. +type MetricsRegistration interface { + Close() +} diff --git a/migrate.go b/migrate.go deleted file mode 100644 index a23e882..0000000 --- a/migrate.go +++ /dev/null @@ -1,52 +0,0 @@ -package postgres - -import ( - "context" - "embed" - "errors" - "log/slog" - - "github.com/golang-migrate/migrate/v4" - _ "github.com/golang-migrate/migrate/v4/database/postgres" // init - "github.com/golang-migrate/migrate/v4/source" - "github.com/golang-migrate/migrate/v4/source/iofs" -) - -func applyMigrations(fs embed.FS, dsn string, l *slog.Logger) (err error) { - var src source.Driver - ctx := context.Background() - - src, err = iofs.New(fs, ".") - if err != nil { - err = errors.Join(errors.New("embed.FS init failed"), err) - return err - } - defer func() { - if ce := src.Close(); ce != nil { - err = errors.Join(err, ce) - } - }() - - var instance *migrate.Migrate - instance, err = migrate.NewWithSourceInstance("iofs", src, dsn) - if err != nil { - err = errors.Join(errors.New("db instance init failed"), err) - return err - } - defer func() { - if sErr, ie := instance.Close(); sErr != nil || ie != nil { - err = errors.Join(err, sErr, ie) - } - }() - - if mErr := instance.Up(); mErr != nil && !errors.Is(mErr, migrate.ErrNoChange) { - err = errors.Join(errors.New("migrate-up failed"), mErr) - } else { - ver, dirty, _ := instance.Version() - l.InfoContext(ctx, "migrate-up done", - slog.Any("version", ver), - slog.Any("dirty", dirty)) - } - - return err -} diff --git a/options.go b/options.go index 2ba07d3..c41dbd9 100644 --- a/options.go +++ b/options.go @@ -1,239 +1,204 @@ -package postgres +package xpg import ( - "embed" + "errors" "fmt" - "log/slog" - "runtime" - "time" + "maps" + "net" + "strconv" + "strings" "github.com/jackc/pgx/v5" - "go.opentelemetry.io/otel/trace" + "github.com/jackc/pgx/v5/multitracer" + "github.com/jackc/pgx/v5/tracelog" ) -// An Option lets you add opts to pberrors interceptors using With* funcs. +// Option configures a Pool. +// +// The interface is sealed so options can only be created by this package. type Option interface { - apply(p *Pool) + apply(*settings) error } -type optionFunc func(p *Pool) +// WithName assigns a stable logical name to the pool. +// +// Name is metadata for diagnostics and observability. It does not change the +// PostgreSQL application_name runtime parameter. +func WithName(name string) Option { + name = strings.TrimSpace(name) -func (f optionFunc) apply(p *Pool) { - f(p) -} - -func WithLogger(l *slog.Logger) Option { - return optionFunc(func(p *Pool) { - if l != nil { - p.logger = l + return optionFunc(func(settings *settings) error { + if name == "" { + return errors.New("pool name must not be blank") } - }) -} -func WithConfig(config *Config) Option { - return optionFunc(func(p *Pool) { - if config != nil { - p.cfg = config - } + settings.name = name + + return nil }) } -func WithClientID(id string) Option { - return optionFunc(func(p *Pool) { - if id != "" { - p.id = fmt.Sprintf("%s-%s", id, GenerateUUID()) +// WithLabel adds or replaces one pool label. +func WithLabel(key, value string) Option { + return optionFunc(func(settings *settings) error { + if key == "" { + return errors.New("label key must not be empty") } - }) -} -func WithTraceProvider(provider trace.TracerProvider) Option { - return optionFunc(func(p *Pool) { - p.traceProvider = provider + settings.labels[key] = value + + return nil }) } -func WithMigrations(migrations ...embed.FS) Option { - return optionFunc(func(p *Pool) { - if len(migrations) > 0 { - p.migrations = migrations +// WithLabels merges labels into the pool metadata. +// +// Labels are defensively copied. When the same key is configured more than +// once, the last value wins. +func WithLabels(labels map[string]string) Option { + labels = cloneLabels(labels) + + return optionFunc(func(settings *settings) error { + for key, value := range labels { + if key == "" { + return errors.New("label key must not be empty") + } + + settings.labels[key] = value } + + return nil }) } -func WithMetricsNamespace(ns string) Option { - return optionFunc(func(p *Pool) { - if ns != "" { - p.namespace = ns +// WithLogger attaches a pgx-compatible logger to the pool. +// +// Logging uses pgx tracelog and participates in the same tracing pipeline as +// tracers configured through xpg. pgx tracelog may include SQL text and query +// arguments in log records; applications are responsible for choosing an +// appropriate level and handling sensitive values. +func WithLogger(logger tracelog.Logger, level tracelog.LogLevel) Option { + return optionFunc(func(settings *settings) error { + if logger == nil { + return errors.New("pool logger is nil") } + + settings.tracers = append(settings.tracers, &tracelog.TraceLog{ + Logger: logger, + LogLevel: level, + }) + + return nil }) } -type Config struct { - ShardID int `envconfig:"POSTGRES_SHARD_ID"` - ClusterHost string `envconfig:"POSTGRES_CLUSTER_HOST" required:"true"` - ClusterPort string `envconfig:"POSTGRES_CLUSTER_PORT" required:"true"` - ClusterReplicaPort string `envconfig:"POSTGRES_CLUSTER_REPLICA_PORT" required:"true"` - User string `envconfig:"POSTGRES_USER" required:"true"` - Password string `envconfig:"POSTGRES_PASSWORD" required:"true"` - DB string `envconfig:"POSTGRES_DB" required:"true"` - - MinRWConn int32 `envconfig:"POSTGRES_MIN_RW_CONN"` - MinROConn int32 `envconfig:"POSTGRES_MIN_RO_CONN"` - MaxRWConn int32 `envconfig:"POSTGRES_MAX_RW_CONN"` - MaxROConn int32 `envconfig:"POSTGRES_MAX_RO_CONN"` - MaxConnLifetime time.Duration `envconfig:"POSTGRES_MAX_CONN_LIFETIME"` - MaxConnIdleTime time.Duration `envconfig:"POSTGRES_MAX_CONN_IDLE_TIME"` - QueryExecMode string `envconfig:"POSTGRES_QUERY_EXEC_MODE"` - StatementCacheCapacity int `envconfig:"POSTGRES_STATEMENT_CACHE_CAPACITY"` - DescriptionCacheCapacity int `envconfig:"POSTGRES_DESCRIPTION_CACHE_CAPACITY"` - - MasterArgs string `envconfig:"POSTGRES_MASTER_ARGS"` - ReplicaArgs string `envconfig:"POSTGRES_REPLICA_ARGS"` - - MigrateEnabled bool `envconfig:"POSTGRES_MIGRATE_ENABLED"` - MigrateArgs string `envconfig:"POSTGRES_MIGRATE_ARGS"` - MigratePort string `envconfig:"POSTGRES_MIGRATE_PORT"` - - writer bool - appName string -} +// WithTracer attaches one pgx query tracer to the pool. +// +// The option may be specified multiple times. Configured loggers and tracers +// are combined through pgx multitracer. When xpg logging or tracing options are +// configured, the resulting tracing pipeline replaces any tracer already +// configured on the pgx connection config. +func WithTracer(tracer pgx.QueryTracer) Option { + return WithTracers(tracer) +} + +// WithTracers attaches multiple pgx query tracers to the pool. +// +// Configured loggers and tracers are invoked in configuration order and +// combined through pgx multitracer. When xpg logging or tracing options are +// configured, the resulting tracing pipeline replaces any tracer already +// configured on the pgx connection config. +func WithTracers(tracers ...pgx.QueryTracer) Option { + return optionFunc(func(settings *settings) error { + for _, tracer := range tracers { + if tracer == nil { + return errors.New("pool tracer is nil") + } + } -// DSN postgres://username:password@host:port/db?sslmode=disable& -func (c *Config) getDSN() string { - return formatDSN(c.User, c.Password, c.ClusterHost, c.getPort(), c.DB, c.appName, c.getArgs()) -} + settings.tracers = append(settings.tracers, tracers...) -func (c *Config) getMigrateDSN() string { - return formatDSN(c.User, c.Password, c.ClusterHost, c.getMigratePort(), c.DB, c.appName, c.MigrateArgs) + return nil + }) } -func (c *Config) getPort() string { - if c.writer { - return c.ClusterPort - } - return c.ClusterReplicaPort -} +// WithMetrics attaches one metrics implementation to the pool. +// +// Metrics are registered when the pool is created and unregistered +// automatically when the Pool is closed. +func WithMetrics(metrics Metrics) Option { + return optionFunc(func(settings *settings) error { + if metrics == nil { + return errors.New("pool metrics is nil") + } -func (c *Config) getMigratePort() string { - if c.MigratePort != "" { - return c.MigratePort - } - return c.ClusterPort -} + settings.metrics = metrics -func (c *Config) getArgs() string { - if c.writer { - return c.MasterArgs - } - return c.ReplicaArgs + return nil + }) } -func (c *Config) getMinConns() int32 { - if c.writer { - return c.MinRWConn - } - return c.MinROConn -} +type optionFunc func(*settings) error -func (c *Config) getMaxConns() int32 { - if c.writer { - return c.MaxRWConn - } - return c.MaxROConn +func (option optionFunc) apply(settings *settings) error { + return option(settings) } -func (c *Config) resolvedMinConns() int32 { - if v := c.getMinConns(); v > 0 { - return v - } - return 1 +type settings struct { + name string + labels map[string]string + metrics Metrics + tracers []pgx.QueryTracer } -func (c *Config) resolvedMaxConns() int32 { - v := c.getMaxConns() - if v <= 0 { - v = 4 +func defaultSettings() *settings { + return &settings{ + labels: make(map[string]string), } - if numCPU := int32(runtime.NumCPU()); numCPU > v { - return numCPU - } - return v } -func formatDSN(user, pass, host, port, db, appName, args string) string { - return fmt.Sprintf("postgres://%s:%s@%s:%s/%s?sslmode=disable&application_name=%s&%s", - user, pass, host, port, db, appName, args, - ) -} - -const ( - QueryExecModeCacheStatement = "cache_statement" - QueryExecModeCacheDescribe = "cache_describe" - QueryExecModeDescribeExec = "describe_exec" - QueryExecModeExec = "exec" - QueryExecModeSimpleProtocol = "simple_protocol" -) - -func getQueryExecMode(mode string) pgx.QueryExecMode { - switch mode { - case QueryExecModeCacheStatement: - return pgx.QueryExecModeCacheStatement - - case QueryExecModeCacheDescribe: - return pgx.QueryExecModeCacheDescribe - - case QueryExecModeDescribeExec: - return pgx.QueryExecModeDescribeExec - - case QueryExecModeExec: - return pgx.QueryExecModeExec - - case QueryExecModeSimpleProtocol: - return pgx.QueryExecModeSimpleProtocol +func applyOptions(settings *settings, opts ...Option) error { + for _, opt := range opts { + if opt == nil { + return errors.New("xpg: option is nil") + } - default: - return pgx.QueryExecModeCacheStatement + if err := opt.apply(settings); err != nil { + return fmt.Errorf("xpg: apply option: %w", err) + } } -} -func parseConfig(cfg *Config) *options { - o := &options{ - dsn: cfg.getDSN(), - minConns: cfg.resolvedMinConns(), - maxConns: cfg.resolvedMaxConns(), - maxConnLifetime: time.Minute * 1, - maxConnIdleTime: time.Second * 30, - defaultQueryExecMode: getQueryExecMode(cfg.QueryExecMode), - statementCacheCapacity: 128, - descriptionCacheCapacity: 512, - } + return nil +} - if cfg.getMinConns() > 0 { - o.minConns = cfg.getMinConns() +func (s *settings) poolName(host string, port uint16, database string) string { + if s.name != "" { + return s.name } - if cfg.getMaxConns() > 0 { - o.maxConns = cfg.getMaxConns() - if numCPU := int32(runtime.NumCPU()); numCPU > cfg.getMaxConns() { - o.maxConns = numCPU - } + address := net.JoinHostPort( + host, + strconv.Itoa(int(port)), + ) + if database == "" { + return address } - if cfg.MaxConnLifetime > 0 { - o.maxConnLifetime = cfg.MaxConnLifetime - } + return address + "/" + database +} - if cfg.MaxConnIdleTime > 0 { - o.maxConnIdleTime = cfg.MaxConnIdleTime +func (s *settings) buildTracer() pgx.QueryTracer { + if len(s.tracers) == 0 { + return nil } - if cfg.StatementCacheCapacity > 0 { - o.statementCacheCapacity = cfg.StatementCacheCapacity - } + return multitracer.New(s.tracers...) +} - if cfg.DescriptionCacheCapacity > 0 { - o.descriptionCacheCapacity = cfg.DescriptionCacheCapacity +func cloneLabels(labels map[string]string) map[string]string { + if len(labels) == 0 { + return nil } - return o + return maps.Clone(labels) } diff --git a/pool.go b/pool.go index 6ad6756..72d5a29 100644 --- a/pool.go +++ b/pool.go @@ -1,260 +1,159 @@ -package postgres +package xpg import ( "context" - "embed" "errors" "fmt" - "log/slog" - "strconv" - "time" + "sync" - "github.com/Masterminds/squirrel" - "github.com/exaring/otelpgx" "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgconn" "github.com/jackc/pgx/v5/pgxpool" - "github.com/jackc/pgx/v5/tracelog" - poolcollector "github.com/mkbeh/xpg/internal/pkg/pgxpoolcollector/v5" - "github.com/mkbeh/xpg/internal/pkg/pgxslog" - "github.com/mkbeh/xpg/internal/pkg/pgxtracer" - "github.com/prometheus/client_golang/prometheus" - "go.opentelemetry.io/otel" - "go.opentelemetry.io/otel/trace" ) +// Pool is a concurrency-safe PostgreSQL connection pool backed by pgxpool. type Pool struct { - *pgxpool.Pool - - id string - cfg *Config - logger *slog.Logger - traceProvider trace.TracerProvider - qBuilder squirrel.StatementBuilderType - migrations []embed.FS - namespace string - labels prometheus.Labels -} - -type options struct { - dsn string - minConns int32 - maxConns int32 - maxConnLifetime time.Duration - maxConnIdleTime time.Duration - statementCacheCapacity int - descriptionCacheCapacity int - defaultQueryExecMode pgx.QueryExecMode - logger *slog.Logger - traceProvider trace.TracerProvider - tracers []pgxtracer.QueryTracer -} + pool *pgxpool.Pool + metrics MetricsRegistration -func NewWriter(opts ...Option) (*Pool, error) { - return newPool(true, opts) -} + name string + labels map[string]string -func NewReader(opts ...Option) (*Pool, error) { - return newPool(false, opts) + closeOnce sync.Once } -func newPool(writer bool, opts []Option) (*Pool, error) { - p := &Pool{ - cfg: &Config{}, - logger: slog.Default(), - qBuilder: squirrel.StatementBuilder.PlaceholderFormat(squirrel.Dollar), +// Open parses a PostgreSQL connection string and creates a Pool. +func Open(ctx context.Context, connString string, options ...Option) (*Pool, error) { + config, err := pgxpool.ParseConfig(connString) + if err != nil { + return nil, fmt.Errorf("xpg: parse pool config: %w", err) } - for _, opt := range opts { - opt.apply(p) + return New(ctx, config, options...) +} + +// New creates a Pool from config. +// +// Config must have been created by pgxpool.ParseConfig. New passes a defensive +// copy to pgxpool, so subsequent changes to config do not affect the Pool. +// +// As with pgxpool.Config.Copy, the referenced tls.Config remains shared and +// must not be modified after it has been used to create connections. +func New(ctx context.Context, config *pgxpool.Config, options ...Option) (*Pool, error) { + if config == nil || config.ConnConfig == nil { + return nil, errors.New("xpg: pool config is nil") } - p.cfg.writer = writer - p.cfg.appName = p.getID() + settings := defaultSettings() - if p.traceProvider == nil { - p.traceProvider = otel.GetTracerProvider() + if err := applyOptions(settings, options...); err != nil { + return nil, err } - if writer { - p.logger = p.logger.With(pgxslog.Component("postgres_master")) - } else { - p.logger = p.logger.With(pgxslog.Component("postgres_replica")) - } + poolConfig := config.Copy() + connConfig := poolConfig.ConnConfig - connOpts := parseConfig(p.cfg) - connOpts.logger = p.logger - connOpts.traceProvider = p.traceProvider + if tracer := settings.buildTracer(); tracer != nil { + connConfig.Tracer = tracer + } - conn, err := connect(connOpts) + pgxPool, err := pgxpool.NewWithConfig(ctx, poolConfig) if err != nil { - return nil, err + return nil, fmt.Errorf("xpg: create pool: %w", err) } - p.Pool = conn - p.exposeMetrics(writer) + pool := &Pool{ + pool: pgxPool, + name: settings.poolName( + connConfig.Host, + connConfig.Port, + connConfig.Database, + ), + labels: cloneLabels(settings.labels), + } - collector := poolcollector.NewStatsCollector(p.namespace, "postgres", p.labels, p.Pool) - prometheus.MustRegister(collector) + if err := pool.registerMetrics(settings.metrics); err != nil { + pool.Close() - if p.cfg.writer && p.cfg.MigrateEnabled { - for _, fs := range p.migrations { - if err := applyMigrations(fs, p.cfg.getMigrateDSN(), p.logger); err != nil { - return nil, err - } - } + return nil, fmt.Errorf("xpg: register pool metrics: %w", err) } - return p, err -} - -func (p *Pool) QueryBuilder() squirrel.StatementBuilderType { - return p.qBuilder + return pool, nil } -func (p *Pool) Logger() *slog.Logger { - return p.logger +// Name returns the logical pool name. +// +// If WithName is not configured, the name is derived from the connection host, +// port, and database. +func (p *Pool) Name() string { + return p.name } -func (p *Pool) Close() error { - p.Pool.Close() - return nil +// Labels returns a copy of the pool labels. +func (p *Pool) Labels() map[string]string { + return cloneLabels(p.labels) } -func (p *Pool) SendBatch(ctx context.Context, b *pgx.Batch) pgx.BatchResults { - if tx := extractTx(ctx); tx != nil { - return tx.SendBatch(ctx, b) - } - return p.Pool.SendBatch(ctx, b) +// Raw returns the underlying pgxpool.Pool. +// +// The returned pool is owned by Pool and must not be closed directly. +func (p *Pool) Raw() *pgxpool.Pool { + return p.pool } -func (p *Pool) Exec(ctx context.Context, sql string, arguments ...any) (pgconn.CommandTag, error) { - if tx := extractTx(ctx); tx != nil { - return tx.Exec(ctx, sql, arguments...) - } - return p.Pool.Exec(ctx, sql, arguments...) +// Ping verifies connectivity to PostgreSQL. +func (p *Pool) Ping(ctx context.Context) error { + return p.pool.Ping(ctx) } -func (p *Pool) Query(ctx context.Context, sql string, args ...any) (pgx.Rows, error) { - if tx := extractTx(ctx); tx != nil { - return tx.Query(ctx, sql, args...) - } - return p.Pool.Query(ctx, sql, args...) -} +// Close closes the pool and waits for acquired connections to be returned. +// Close is safe to call multiple times. +func (p *Pool) Close() { + p.closeOnce.Do(func() { + if p.metrics != nil { + p.metrics.Close() + } -func (p *Pool) QueryRow(ctx context.Context, sql string, args ...any) pgx.Row { - if tx := extractTx(ctx); tx != nil { - return tx.QueryRow(ctx, sql, args...) - } - return p.Pool.QueryRow(ctx, sql, args...) + p.pool.Close() + }) } -// RunInTxx alias for RunInTx. -func (p *Pool) RunInTxx(ctx context.Context, fn func(ctx context.Context) error) error { - return p.RunInTx(ctx, fn, pgx.TxOptions{}) +// Exec executes SQL against the pool. +func (p *Pool) Exec(ctx context.Context, sql string, arguments ...any) (pgconn.CommandTag, error) { + return p.pool.Exec(ctx, sql, arguments...) } -func (p *Pool) RunInTx(ctx context.Context, fn func(ctx context.Context) error, txOptions pgx.TxOptions) (err error) { - tx, err := p.Pool.BeginTx(ctx, txOptions) - if err != nil { - p.Logger().ErrorContext(ctx, "failed to begin transaction", pgxslog.Error(err)) - return NewPgError(ErrBeginTransaction, err) - } - - defer func() { - if r := recover(); r != nil { - p.Logger().ErrorContext(ctx, "panic recovered", slog.Any("error", r)) - err = NewPgError(ErrOther, fmt.Errorf("%v", r)) - } - - if rErr := tx.Rollback(ctx); rErr != nil { - if !errors.Is(rErr, pgx.ErrTxClosed) { - p.Logger().ErrorContext(ctx, "failed to rollback transaction", pgxslog.Error(rErr)) - } - } - }() - - if err = fn(injectTx(ctx, tx)); err != nil { - return err - } - - if err = tx.Commit(ctx); err != nil { - p.Logger().ErrorContext(ctx, "failed to commit transaction", pgxslog.Error(err)) - return NewPgError(ErrCommitTransaction, err) - } - - return nil +// Query executes SQL and returns the resulting rows. +func (p *Pool) Query(ctx context.Context, sql string, args ...any) (pgx.Rows, error) { + return p.pool.Query(ctx, sql, args...) } -func (p *Pool) AcquireTxLock(ctx context.Context, key string, durationSeconds float64) (isLocked bool, err error) { - row := p.QueryRow(ctx, ` - SELECT CASE - WHEN pg_try_advisory_xact_lock($1) THEN (SELECT concat(pg_sleep($2), 'false'))::bool - ELSE true - END AS is_locked;`, - int64(StringAsHash64(key)), - durationSeconds) - err = row.Scan(&isLocked) - return isLocked, err +// QueryRow executes SQL that is expected to return at most one row. +func (p *Pool) QueryRow(ctx context.Context, sql string, args ...any) pgx.Row { + return p.pool.QueryRow(ctx, sql, args...) } -func (p *Pool) getID() string { - if p.id == "" { - return GenerateUUID() - } - return p.id +// SendBatch sends a batch of queries through the pool. +func (p *Pool) SendBatch(ctx context.Context, batch *pgx.Batch) pgx.BatchResults { + return p.pool.SendBatch(ctx, batch) } -func (p *Pool) exposeMetrics(writer bool) { - if p.labels == nil { - p.labels = make(prometheus.Labels) - } - - p.labels["client_id"] = p.getID() - p.labels["db"] = p.cfg.DB - p.labels["shard_id"] = strconv.Itoa(p.cfg.ShardID) - - if writer { - p.labels["client_kind"] = "master" - } else { - p.labels["client_kind"] = "replica" - } +// CopyFrom copies rows into the specified table. +func (p *Pool) CopyFrom(ctx context.Context, tableName pgx.Identifier, columnNames []string, rowSrc pgx.CopyFromSource) (int64, error) { + return p.pool.CopyFrom(ctx, tableName, columnNames, rowSrc) } -func connect(opts *options) (*pgxpool.Pool, error) { - poolCfg, err := pgxpool.ParseConfig(opts.dsn) - if err != nil { - return nil, err +func (p *Pool) registerMetrics(metrics Metrics) error { + if metrics == nil { + return nil } - opts.tracers = append(opts.tracers, - &tracelog.TraceLog{Logger: pgxslog.NewLogger(opts.logger), LogLevel: tracelog.LogLevelError}, - otelpgx.NewTracer( - otelpgx.WithTrimSQLInSpanName(), - otelpgx.WithTracerProvider(opts.traceProvider), - ), - ) - - poolCfg.MinConns = opts.minConns - poolCfg.MaxConns = opts.maxConns - poolCfg.MaxConnLifetime = opts.maxConnLifetime - poolCfg.MaxConnIdleTime = opts.maxConnIdleTime - - poolCfg.ConnConfig.StatementCacheCapacity = opts.statementCacheCapacity - poolCfg.ConnConfig.DescriptionCacheCapacity = opts.descriptionCacheCapacity - poolCfg.ConnConfig.DefaultQueryExecMode = opts.defaultQueryExecMode - poolCfg.ConnConfig.Tracer = pgxtracer.New(opts.tracers...) - - ctx := context.Background() - - pool, err := pgxpool.NewWithConfig(ctx, poolCfg) + registration, err := metrics.Register(p) if err != nil { - return nil, err + return err } - if err := pool.Ping(ctx); err != nil { - return nil, err - } + p.metrics = registration - return pool, nil + return nil } diff --git a/pool_test.go b/pool_test.go new file mode 100644 index 0000000..b68d91b --- /dev/null +++ b/pool_test.go @@ -0,0 +1,274 @@ +package xpg + +import ( + "errors" + "testing" + + "github.com/jackc/pgx/v5/pgxpool" +) + +func TestNewRejectsNilConfig(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + config *pgxpool.Config + }{ + { + name: "nil config", + config: nil, + }, + { + name: "nil connection config", + config: &pgxpool.Config{}, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + + _, err := New( + t.Context(), + test.config, + ) + + assertErrorMessage( + t, + err, + "xpg: pool config is nil", + ) + }) + } +} + +func TestPoolMetadata(t *testing.T) { + t.Parallel() + + config := newPoolTestConfig(t) + + labels := map[string]string{ + "region": "eu", + "role": "reader", + } + + labelsOption := WithLabels(labels) + labels["region"] = "changed" + + pool, err := New( + t.Context(), + config, + WithName(" orders "), + labelsOption, + WithLabel("role", "primary"), + ) + if err != nil { + t.Fatalf("create pool: %v", err) + } + defer pool.Close() + + if got, want := pool.Name(), "orders"; got != want { + t.Fatalf( + "Name() = %q, want %q", + got, + want, + ) + } + + gotLabels := pool.Labels() + + if got, want := gotLabels["region"], "eu"; got != want { + t.Fatalf( + "region label = %q, want %q", + got, + want, + ) + } + + if got, want := gotLabels["role"], "primary"; got != want { + t.Fatalf( + "role label = %q, want %q", + got, + want, + ) + } + + gotLabels["region"] = "mutated" + + if got, want := pool.Labels()["region"], "eu"; got != want { + t.Fatalf( + "region label after mutation = %q, want %q", + got, + want, + ) + } + + if pool.Raw() == nil { + t.Fatal("Raw returned nil") + } +} + +func TestPoolDerivedName(t *testing.T) { + t.Parallel() + + pool, err := New( + t.Context(), + newPoolTestConfig(t), + ) + if err != nil { + t.Fatalf("create pool: %v", err) + } + defer pool.Close() + + if got, want := pool.Name(), "localhost:5432/testdb"; got != want { + t.Fatalf( + "Name() = %q, want %q", + got, + want, + ) + } +} + +func TestPoolStats(t *testing.T) { + t.Parallel() + + config := newPoolTestConfig(t) + config.MaxConns = 17 + + pool, err := New( + t.Context(), + config, + ) + if err != nil { + t.Fatalf("create pool: %v", err) + } + defer pool.Close() + + stats := pool.Stats() + + if got, want := stats.MaxConns, int32(17); got != want { + t.Fatalf( + "MaxConns = %d, want %d", + got, + want, + ) + } + + if stats.TotalConns != 0 { + t.Fatalf( + "TotalConns = %d, want 0", + stats.TotalConns, + ) + } +} + +func TestPoolMetricsLifecycle(t *testing.T) { + t.Parallel() + + registration := &poolTestMetricsRegistration{} + metrics := &poolTestMetrics{ + registration: registration, + } + + pool, err := New( + t.Context(), + newPoolTestConfig(t), + WithMetrics(metrics), + ) + if err != nil { + t.Fatalf("create pool: %v", err) + } + + if metrics.registerCalls != 1 { + t.Fatalf( + "register calls = %d, want 1", + metrics.registerCalls, + ) + } + + if metrics.pool != pool { + t.Fatal("metrics registered with unexpected pool") + } + + pool.Close() + pool.Close() + + if registration.closeCalls != 1 { + t.Fatalf( + "registration close calls = %d, want 1", + registration.closeCalls, + ) + } +} + +func TestNewPreservesMetricsRegistrationError(t *testing.T) { + t.Parallel() + + expectedErr := errors.New("register failed") + metrics := &poolTestMetrics{ + err: expectedErr, + } + + pool, err := New( + t.Context(), + newPoolTestConfig(t), + WithMetrics(metrics), + ) + + if pool != nil { + t.Fatal("expected nil pool") + } + + if !errors.Is(err, expectedErr) { + t.Fatalf("original error was not preserved: %v", err) + } + + if metrics.registerCalls != 1 { + t.Fatalf( + "register calls = %d, want 1", + metrics.registerCalls, + ) + } +} + +func newPoolTestConfig(t *testing.T) *pgxpool.Config { + t.Helper() + + config, err := pgxpool.ParseConfig( + "postgres://user:password@localhost:5432/testdb?sslmode=disable", + ) + if err != nil { + t.Fatalf("parse pool config: %v", err) + } + + return config +} + +type poolTestMetrics struct { + registration MetricsRegistration + err error + + pool *Pool + registerCalls int +} + +func (m *poolTestMetrics) Register( + pool *Pool, +) (MetricsRegistration, error) { + m.registerCalls++ + m.pool = pool + + if m.err != nil { + return nil, m.err + } + + return m.registration, nil +} + +type poolTestMetricsRegistration struct { + closeCalls int +} + +func (r *poolTestMetricsRegistration) Close() { + r.closeCalls++ +} diff --git a/shard/doc.go b/shard/doc.go new file mode 100644 index 0000000..aa07ba8 --- /dev/null +++ b/shard/doc.go @@ -0,0 +1,9 @@ +// Package shard provides application-level routing across PostgreSQL clusters. +// +// A Topology owns an immutable set of logical shards backed by cluster.Cluster +// values. Resolvers map application keys to shards using rendezvous hashing, +// ordered ranges, time ranges, or custom routing logic. +// +// The package also provides shard grouping, colocation checks, and bounded +// parallel operations across shards. +package shard diff --git a/shard/errors.go b/shard/errors.go new file mode 100644 index 0000000..908af96 --- /dev/null +++ b/shard/errors.go @@ -0,0 +1,54 @@ +package shard + +import ( + "errors" + "fmt" +) + +var ( + // ErrNoShard is returned when a resolver cannot map a key to any shard. + ErrNoShard = errors.New("xpg/shard: no shard resolved") + + // ErrUnknownShard is returned when routing references a shard that does not + // exist in the topology. + ErrUnknownShard = errors.New("xpg/shard: unknown shard") + + // ErrShardMismatch is returned when keys expected to be colocated resolve to + // different shards. + ErrShardMismatch = errors.New("xpg/shard: keys resolve to different shards") +) + +// UnknownShardError identifies a shard referenced by routing that does not +// exist in the topology. +type UnknownShardError struct { + ShardID ID +} + +func (e *UnknownShardError) Error() string { + return fmt.Sprintf("xpg/shard: unknown shard %q", e.ShardID) +} + +func (e *UnknownShardError) Unwrap() error { + return ErrUnknownShard +} + +// MismatchError describes the first key whose resolved shard differs from the +// shard of the first key. +type MismatchError struct { + Expected ID + Actual ID + Index int +} + +func (e *MismatchError) Error() string { + return fmt.Sprintf( + "xpg/shard: key %d resolved to shard %q instead of %q", + e.Index, + e.Actual, + e.Expected, + ) +} + +func (e *MismatchError) Unwrap() error { + return ErrShardMismatch +} diff --git a/shard/errors_test.go b/shard/errors_test.go new file mode 100644 index 0000000..5523408 --- /dev/null +++ b/shard/errors_test.go @@ -0,0 +1,38 @@ +package shard + +import ( + "errors" + "testing" +) + +func TestUnknownShardError(t *testing.T) { + t.Parallel() + + err := &UnknownShardError{ShardID: "missing"} + + if got, want := err.Error(), `xpg/shard: unknown shard "missing"`; got != want { + t.Fatalf("Error() = %q, want %q", got, want) + } + + if !errors.Is(err, ErrUnknownShard) { + t.Fatal("errors.Is() = false, want ErrUnknownShard") + } +} + +func TestMismatchError(t *testing.T) { + t.Parallel() + + err := &MismatchError{ + Expected: "shard-a", + Actual: "shard-b", + Index: 2, + } + + if got, want := err.Error(), `xpg/shard: key 2 resolved to shard "shard-b" instead of "shard-a"`; got != want { + t.Fatalf("Error() = %q, want %q", got, want) + } + + if !errors.Is(err, ErrShardMismatch) { + t.Fatal("errors.Is() = false, want ErrShardMismatch") + } +} diff --git a/shard/foreach.go b/shard/foreach.go new file mode 100644 index 0000000..9effadc --- /dev/null +++ b/shard/foreach.go @@ -0,0 +1,127 @@ +package shard + +import ( + "context" + "errors" + "fmt" + "sync" +) + +// ForEachShardResult contains the result associated with one shard. +type ForEachShardResult struct { + ShardID ID + Err error +} + +// ForEachShardResults contains results in topology registration order. +type ForEachShardResults []ForEachShardResult + +// Err returns all shard failures joined in registration order. +func (results ForEachShardResults) Err() error { + var errs []error + + for _, result := range results { + if result.Err == nil { + continue + } + + errs = append( + errs, + fmt.Errorf( + "xpg/shard: shard %q: %w", + result.ShardID, + result.Err, + ), + ) + } + + return errors.Join(errs...) +} + +// ForEachShard invokes fn across the topology with at most concurrency +// callbacks running at once. Results are returned in topology registration +// order; callback execution order is not guaranteed. +// +// Callback failures and context cancellation are stored in the corresponding +// results and can be joined with ForEachShardResults.Err. The returned error is +// reserved for invalid invocation arguments. +// +// Once context cancellation is observed, callbacks that have not started are +// skipped and their results contain ctx.Err(). Callbacks already running are +// responsible for observing ctx. ForEachShard waits for all started callbacks +// to finish before returning. +func (t *Topology) ForEachShard( + ctx context.Context, + concurrency int, + fn func(context.Context, Shard) error, +) (ForEachShardResults, error) { + if t == nil { + return nil, errors.New("xpg/shard: topology is nil") + } + + if len(t.shards) == 0 { + return nil, errors.New("xpg/shard: topology is empty") + } + + if concurrency <= 0 { + return nil, errors.New("xpg/shard: concurrency must be positive") + } + + if fn == nil { + return nil, errors.New("xpg/shard: callback is nil") + } + + results := make(ForEachShardResults, len(t.shards)) + for index, shard := range t.shards { + results[index].ShardID = shard.ID() + } + + workerCount := min(concurrency, len(t.shards)) + jobs := make(chan int) + + var workers sync.WaitGroup + workers.Add(workerCount) + + for range workerCount { + go func() { + defer workers.Done() + + for index := range jobs { + // An index may have been scheduled immediately before context + // cancellation. Skip callbacks that have not started yet. + if err := ctx.Err(); err != nil { + results[index].Err = err + continue + } + + results[index].Err = fn(ctx, t.shards[index]) + } + }() + } + + nextIndex := 0 + + for nextIndex < len(t.shards) && ctx.Err() == nil { + // The explicit context check above prevents scheduling new work after + // cancellation has already been observed. The select still handles + // cancellation that happens while waiting for a worker. + select { + case jobs <- nextIndex: + nextIndex++ + case <-ctx.Done(): + } + } + + close(jobs) + workers.Wait() + + // Workers own results for scheduled indexes [0, nextIndex). After all + // workers finish, remaining indexes can be marked canceled without races. + if err := ctx.Err(); err != nil { + for index := nextIndex; index < len(results); index++ { + results[index].Err = err + } + } + + return results, nil +} diff --git a/shard/foreach_test.go b/shard/foreach_test.go new file mode 100644 index 0000000..ceb8060 --- /dev/null +++ b/shard/foreach_test.go @@ -0,0 +1,356 @@ +package shard + +import ( + "context" + "errors" + "sync/atomic" + "testing" + "time" +) + +func TestForEachShardValidatesArguments(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a") + + tests := []struct { + name string + topology *Topology + concurrency int + fn func(context.Context, Shard) error + wantError string + }{ + { + name: "nil topology", + topology: nil, + concurrency: 1, + fn: func(context.Context, Shard) error { return nil }, + wantError: "xpg/shard: topology is nil", + }, + { + name: "empty topology", + topology: &Topology{}, + concurrency: 1, + fn: func(context.Context, Shard) error { return nil }, + wantError: "xpg/shard: topology is empty", + }, + { + name: "zero concurrency", + topology: topology, + concurrency: 0, + fn: func(context.Context, Shard) error { return nil }, + wantError: "xpg/shard: concurrency must be positive", + }, + { + name: "nil callback", + topology: topology, + concurrency: 1, + fn: nil, + wantError: "xpg/shard: callback is nil", + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + _, err := test.topology.ForEachShard( + t.Context(), + test.concurrency, + test.fn, + ) + if err == nil { + t.Fatal("expected error") + } + + if got := err.Error(); got != test.wantError { + t.Fatalf("error = %q, want %q", got, test.wantError) + } + }) + } +} + +func TestForEachShardPreservesRegistrationOrder(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-c", "shard-a", "shard-b") + + results, err := topology.ForEachShard( + t.Context(), + 2, + func(context.Context, Shard) error { return nil }, + ) + if err != nil { + t.Fatalf("ForEachShard() error = %v", err) + } + + want := []ID{"shard-c", "shard-a", "shard-b"} + for index, wantID := range want { + if got := results[index].ShardID; got != wantID { + t.Fatalf("results[%d].ShardID = %q, want %q", index, got, wantID) + } + + if results[index].Err != nil { + t.Fatalf("results[%d].Err = %v", index, results[index].Err) + } + } +} + +func TestForEachShardHonorsConcurrencyLimit(t *testing.T) { + t.Parallel() + + topology := newTestTopology( + t, + "shard-a", + "shard-b", + "shard-c", + "shard-d", + "shard-e", + "shard-f", + ) + + ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second) + defer cancel() + + started := make(chan struct{}, topology.Len()) + release := make(chan struct{}) + + var active atomic.Int32 + var maximum atomic.Int32 + var calls atomic.Int32 + + done := make(chan struct { + results ForEachShardResults + err error + }, 1) + + go func() { + results, err := topology.ForEachShard( + ctx, + 2, + func(ctx context.Context, _ Shard) error { + current := active.Add(1) + defer active.Add(-1) + + calls.Add(1) + + for { + observed := maximum.Load() + if current <= observed || maximum.CompareAndSwap(observed, current) { + break + } + } + + started <- struct{}{} + + select { + case <-release: + return nil + case <-ctx.Done(): + return ctx.Err() + } + }, + ) + + done <- struct { + results ForEachShardResults + err error + }{ + results: results, + err: err, + } + }() + + for range 2 { + select { + case <-started: + case <-ctx.Done(): + close(release) + t.Fatal("two callbacks did not start concurrently") + } + } + + close(release) + + var outcome struct { + results ForEachShardResults + err error + } + + select { + case outcome = <-done: + case <-ctx.Done(): + t.Fatal("ForEachShard() did not finish") + } + + if outcome.err != nil { + t.Fatalf("ForEachShard() error = %v", outcome.err) + } + + if got, want := calls.Load(), int32(topology.Len()); got != want { + t.Fatalf("callback calls = %d, want %d", got, want) + } + + if got, want := maximum.Load(), int32(2); got != want { + t.Fatalf("maximum concurrent callbacks = %d, want %d", got, want) + } + + if err := outcome.results.Err(); err != nil { + t.Fatalf("results.Err() = %v", err) + } +} + +func TestForEachShardCallbackErrorsDoNotStopOtherShards(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a", "shard-b", "shard-c") + sentinel := errors.New("callback failed") + var calls atomic.Int32 + + results, err := topology.ForEachShard( + t.Context(), + 2, + func(_ context.Context, current Shard) error { + calls.Add(1) + if current.ID() == "shard-b" { + return sentinel + } + + return nil + }, + ) + if err != nil { + t.Fatalf("ForEachShard() error = %v", err) + } + + if got, want := calls.Load(), int32(3); got != want { + t.Fatalf("callback calls = %d, want %d", got, want) + } + + if results[0].Err != nil || !errors.Is(results[1].Err, sentinel) || results[2].Err != nil { + t.Fatalf("results = %+v", results) + } + + joined := results.Err() + if !errors.Is(joined, sentinel) { + t.Fatalf("results.Err() = %v, want wrapped sentinel", joined) + } + + if got, want := joined.Error(), + `xpg/shard: shard "shard-b": callback failed`; got != want { + t.Fatalf("results.Err() = %q, want %q", got, want) + } +} + +func TestForEachShardCanceledBeforeScheduling(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a", "shard-b", "shard-c") + ctx, cancel := context.WithCancel(t.Context()) + cancel() + + var calls atomic.Int32 + + results, err := topology.ForEachShard( + ctx, + 2, + func(context.Context, Shard) error { + calls.Add(1) + return nil + }, + ) + if err != nil { + t.Fatalf("ForEachShard() error = %v", err) + } + + if got := calls.Load(); got != 0 { + t.Fatalf("callback calls = %d, want 0", got) + } + + for index, result := range results { + if !errors.Is(result.Err, context.Canceled) { + t.Fatalf("results[%d].Err = %v, want context.Canceled", index, result.Err) + } + } +} + +func TestForEachShardCancellationSkipsCallbacksNotStarted(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a", "shard-b", "shard-c", "shard-d") + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + + started := make(chan struct{}) + var calls atomic.Int32 + + done := make(chan struct { + results ForEachShardResults + err error + }, 1) + + go func() { + results, err := topology.ForEachShard( + ctx, + 1, + func(ctx context.Context, _ Shard) error { + if calls.Add(1) == 1 { + close(started) + } + + <-ctx.Done() + return ctx.Err() + }, + ) + + done <- struct { + results ForEachShardResults + err error + }{results: results, err: err} + }() + + <-started + cancel() + + outcome := <-done + if outcome.err != nil { + t.Fatalf("ForEachShard() error = %v", outcome.err) + } + + if got := calls.Load(); got != 1 { + t.Fatalf("callback calls = %d, want 1", got) + } + + for index, result := range outcome.results { + if !errors.Is(result.Err, context.Canceled) { + t.Fatalf("results[%d].Err = %v, want context.Canceled", index, result.Err) + } + } +} + +func TestForEachShardResultsErr(t *testing.T) { + t.Parallel() + + first := errors.New("first") + second := errors.New("second") + + results := ForEachShardResults{ + {ShardID: "shard-a", Err: first}, + {ShardID: "shard-b"}, + {ShardID: "shard-c", Err: second}, + } + + err := results.Err() + if !errors.Is(err, first) || !errors.Is(err, second) { + t.Fatalf("Err() = %v, want both failures", err) + } + + want := "xpg/shard: shard \"shard-a\": first\n" + + "xpg/shard: shard \"shard-c\": second" + + if got := err.Error(); got != want { + t.Fatalf("Err() = %q, want %q", got, want) + } + + if err := (ForEachShardResults{{ShardID: "shard-a"}}).Err(); err != nil { + t.Fatalf("Err() = %v, want nil", err) + } +} diff --git a/shard/group.go b/shard/group.go new file mode 100644 index 0000000..5d04c0d --- /dev/null +++ b/shard/group.go @@ -0,0 +1,90 @@ +package shard + +import ( + "errors" + "fmt" +) + +// SameShard resolves keys and verifies that they all belong to the same shard. +// +// It returns ErrNoShard when no keys are provided and MismatchError when a key +// resolves to a different shard. +func SameShard[K any](resolver Resolver[K], keys ...K) (Shard, error) { + if resolver == nil { + return Shard{}, errors.New("xpg/shard: resolver is nil") + } + + if len(keys) == 0 { + return Shard{}, ErrNoShard + } + + expected, err := resolver.Resolve(keys[0]) + if err != nil { + return Shard{}, fmt.Errorf("xpg/shard: resolve key 0: %w", err) + } + + expectedID := expected.ID() + + for index := 1; index < len(keys); index++ { + actual, err := resolver.Resolve(keys[index]) + if err != nil { + return Shard{}, fmt.Errorf("xpg/shard: resolve key %d: %w", index, err) + } + + actualID := actual.ID() + if actualID == expectedID { + continue + } + + return Shard{}, &MismatchError{ + Expected: expectedID, + Actual: actualID, + Index: index, + } + } + + return expected, nil +} + +// Group contains keys that resolve to the same shard. +// Keys preserve their original relative order. +type Group[K any] struct { + Shard Shard + Keys []K +} + +// GroupByShard resolves each key once and groups keys by shard. +// +// Groups are returned in order of each shard's first appearance in keys. +// Keys within each group preserve their original relative order. +func GroupByShard[K any](resolver Resolver[K], keys []K) ([]Group[K], error) { + if resolver == nil { + return nil, errors.New("xpg/shard: resolver is nil") + } + + groups := make([]Group[K], 0) + indexByID := make(map[ID]int) + + for keyIndex, key := range keys { + resolved, err := resolver.Resolve(key) + if err != nil { + return nil, fmt.Errorf("xpg/shard: resolve key %d: %w", keyIndex, err) + } + + id := resolved.ID() + + groupIndex, exists := indexByID[id] + if !exists { + groupIndex = len(groups) + indexByID[id] = groupIndex + + groups = append(groups, Group[K]{ + Shard: resolved, + }) + } + + groups[groupIndex].Keys = append(groups[groupIndex].Keys, key) + } + + return groups, nil +} diff --git a/shard/group_test.go b/shard/group_test.go new file mode 100644 index 0000000..5d548f5 --- /dev/null +++ b/shard/group_test.go @@ -0,0 +1,248 @@ +package shard + +import ( + "errors" + "slices" + "testing" +) + +func TestSameShard(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a", "shard-b") + shardA := topology.At(0) + shardB := topology.At(1) + + resolver := testResolverFunc[int](func(key int) (Shard, error) { + if key < 100 { + return shardA, nil + } + + return shardB, nil + }) + + resolved, err := SameShard(resolver, 1, 2, 3) + if err != nil { + t.Fatalf("SameShard() error = %v", err) + } + + if got, want := resolved.ID(), ID("shard-a"); got != want { + t.Fatalf("SameShard().ID() = %q, want %q", got, want) + } +} + +func TestSameShardRejectsNilResolver(t *testing.T) { + t.Parallel() + + _, err := SameShard[int](nil, 1) + if err == nil { + t.Fatal("expected error") + } + + if got, want := err.Error(), "xpg/shard: resolver is nil"; got != want { + t.Fatalf("error = %q, want %q", got, want) + } +} + +func TestSameShardRequiresKey(t *testing.T) { + t.Parallel() + + resolver := testResolverFunc[int](func(int) (Shard, error) { + t.Fatal("resolver should not be called") + return Shard{}, nil + }) + + _, err := SameShard(resolver) + if !errors.Is(err, ErrNoShard) { + t.Fatalf("error = %v, want ErrNoShard", err) + } +} + +func TestSameShardWrapsFirstResolveError(t *testing.T) { + t.Parallel() + + sentinel := errors.New("resolve failed") + resolver := testResolverFunc[int](func(int) (Shard, error) { + return Shard{}, sentinel + }) + + _, err := SameShard(resolver, 1) + if !errors.Is(err, sentinel) { + t.Fatalf("error = %v, want wrapped sentinel", err) + } + + if got, want := err.Error(), "xpg/shard: resolve key 0: resolve failed"; got != want { + t.Fatalf("error = %q, want %q", got, want) + } +} + +func TestSameShardWrapsResolveErrorWithIndex(t *testing.T) { + t.Parallel() + + sentinel := errors.New("resolve failed") + resolver := testResolverFunc[int](func(key int) (Shard, error) { + if key == 2 { + return Shard{}, sentinel + } + + return Shard{}, nil + }) + + _, err := SameShard(resolver, 1, 2) + if !errors.Is(err, sentinel) { + t.Fatalf("error = %v, want wrapped sentinel", err) + } + + if got, want := err.Error(), "xpg/shard: resolve key 1: resolve failed"; got != want { + t.Fatalf("error = %q, want %q", got, want) + } +} + +func TestSameShardReturnsMismatchDetails(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a", "shard-b") + shardA := topology.At(0) + shardB := topology.At(1) + + resolver := testResolverFunc[int](func(key int) (Shard, error) { + if key == 3 { + return shardB, nil + } + + return shardA, nil + }) + + _, err := SameShard(resolver, 1, 2, 3) + if !errors.Is(err, ErrShardMismatch) { + t.Fatalf("error = %v, want ErrShardMismatch", err) + } + + var mismatch *MismatchError + if !errors.As(err, &mismatch) { + t.Fatalf("error = %T, want *MismatchError", err) + } + + if mismatch.Expected != "shard-a" || mismatch.Actual != "shard-b" || mismatch.Index != 2 { + t.Fatalf("mismatch = %+v", mismatch) + } +} + +func TestGroupByShardPreservesGroupAndKeyOrder(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a", "shard-b") + shardA := topology.At(0) + shardB := topology.At(1) + + resolver := testResolverFunc[int](func(key int) (Shard, error) { + if key < 100 { + return shardA, nil + } + + return shardB, nil + }) + + groups, err := GroupByShard(resolver, []int{142, 42, 143, 43}) + if err != nil { + t.Fatalf("GroupByShard() error = %v", err) + } + + if got, want := len(groups), 2; got != want { + t.Fatalf("len(groups) = %d, want %d", got, want) + } + + if got, want := groups[0].Shard.ID(), ID("shard-b"); got != want { + t.Fatalf("groups[0].Shard.ID() = %q, want %q", got, want) + } + if got, want := groups[0].Keys, []int{142, 143}; !slices.Equal(got, want) { + t.Fatalf("groups[0].Keys = %v, want %v", got, want) + } + + if got, want := groups[1].Shard.ID(), ID("shard-a"); got != want { + t.Fatalf("groups[1].Shard.ID() = %q, want %q", got, want) + } + if got, want := groups[1].Keys, []int{42, 43}; !slices.Equal(got, want) { + t.Fatalf("groups[1].Keys = %v, want %v", got, want) + } +} + +func TestGroupByShardResolvesEachKeyOnce(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a") + resolved := topology.At(0) + calls := 0 + + resolver := testResolverFunc[int](func(int) (Shard, error) { + calls++ + return resolved, nil + }) + + keys := []int{1, 2, 3, 4} + groups, err := GroupByShard(resolver, keys) + if err != nil { + t.Fatalf("GroupByShard() error = %v", err) + } + + if got, want := calls, len(keys); got != want { + t.Fatalf("resolve calls = %d, want %d", got, want) + } + + if len(groups) != 1 { + t.Fatalf("len(groups) = %d, want 1", len(groups)) + } +} + +func TestGroupByShardRejectsNilResolver(t *testing.T) { + t.Parallel() + + _, err := GroupByShard[int](nil, []int{1}) + if err == nil { + t.Fatal("expected error") + } + + if got, want := err.Error(), "xpg/shard: resolver is nil"; got != want { + t.Fatalf("error = %q, want %q", got, want) + } +} + +func TestGroupByShardEmptyKeys(t *testing.T) { + t.Parallel() + + resolver := testResolverFunc[int](func(int) (Shard, error) { + t.Fatal("resolver should not be called") + return Shard{}, nil + }) + + groups, err := GroupByShard(resolver, nil) + if err != nil { + t.Fatalf("GroupByShard() error = %v", err) + } + + if len(groups) != 0 { + t.Fatalf("len(groups) = %d, want 0", len(groups)) + } +} + +func TestGroupByShardWrapsResolveErrorWithIndex(t *testing.T) { + t.Parallel() + + sentinel := errors.New("resolve failed") + resolver := testResolverFunc[int](func(key int) (Shard, error) { + if key == 3 { + return Shard{}, sentinel + } + + return Shard{}, nil + }) + + _, err := GroupByShard(resolver, []int{1, 2, 3}) + if !errors.Is(err, sentinel) { + t.Fatalf("error = %v, want wrapped sentinel", err) + } + + if got, want := err.Error(), "xpg/shard: resolve key 2: resolve failed"; got != want { + t.Fatalf("error = %q, want %q", got, want) + } +} diff --git a/shard/helpers_test.go b/shard/helpers_test.go new file mode 100644 index 0000000..17aab99 --- /dev/null +++ b/shard/helpers_test.go @@ -0,0 +1,73 @@ +package shard + +import ( + "testing" + + "github.com/jackc/pgx/v5/pgxpool" + "github.com/mkbeh/xpg" + "github.com/mkbeh/xpg/cluster" +) + +const testDatabaseURL = "postgres://postgres@127.0.0.1:1/postgres?sslmode=disable" + +func newTestCluster(t *testing.T, id ID, labels map[string]string) *cluster.Cluster { + t.Helper() + + config, err := pgxpool.ParseConfig(testDatabaseURL) + if err != nil { + t.Fatalf("pgxpool.ParseConfig() error = %v", err) + } + + config.MinConns = 0 + config.MaxConns = 1 + + pool, err := xpg.New( + t.Context(), + config, + xpg.WithName("shard."+string(id)+".primary"), + ) + if err != nil { + t.Fatalf("xpg.New() error = %v", err) + } + + shardCluster, err := cluster.New(cluster.Config{ + ID: id, + Labels: labels, + Primary: pool, + }) + if err != nil { + pool.Close() + t.Fatalf("cluster.New() error = %v", err) + } + + t.Cleanup(shardCluster.Close) + + return shardCluster +} + +func newTestTopology(t *testing.T, ids ...ID) *Topology { + t.Helper() + + configs := make([]Config, len(ids)) + + for index, id := range ids { + configs[index] = Config{ + Cluster: newTestCluster(t, id, nil), + } + } + + topology, err := NewTopology(configs) + if err != nil { + t.Fatalf("NewTopology() error = %v", err) + } + + t.Cleanup(topology.Close) + + return topology +} + +type testResolverFunc[K any] func(K) (Shard, error) + +func (resolve testResolverFunc[K]) Resolve(key K) (Shard, error) { + return resolve(key) +} diff --git a/shard/resolver.go b/shard/resolver.go new file mode 100644 index 0000000..500fcb8 --- /dev/null +++ b/shard/resolver.go @@ -0,0 +1,9 @@ +package shard + +// Resolver maps a typed application key to one shard. +// +// Resolve should return ErrNoShard when the key cannot be mapped to a shard. +// Implementations must be safe for concurrent use. +type Resolver[K any] interface { + Resolve(key K) (Shard, error) +} diff --git a/shard/resolver/custom.go b/shard/resolver/custom.go new file mode 100644 index 0000000..cae2096 --- /dev/null +++ b/shard/resolver/custom.go @@ -0,0 +1,59 @@ +package resolver + +import ( + "errors" + + "github.com/mkbeh/xpg/shard" +) + +// ResolveFunc maps an application key to a shard ID within topology. +// +// Resolve functions must be deterministic and safe for concurrent use. They +// should return shard.ErrNoShard when a key cannot be mapped and should not +// perform hidden I/O or modify topology. +type ResolveFunc[K any] func(key K, topology *shard.Topology) (shard.ID, error) + +// CustomResolver adapts ResolveFunc to shard.Resolver. +// +// CustomResolver borrows its topology and must not outlive it. +type CustomResolver[K any] struct { + topology *shard.Topology + resolve ResolveFunc[K] +} + +// NewCustom binds custom routing logic to an immutable topology. +func NewCustom[K any](topology *shard.Topology, resolve ResolveFunc[K]) (*CustomResolver[K], error) { + if err := requireTopology(topology); err != nil { + return nil, err + } + + if resolve == nil { + return nil, errors.New("xpg/shard/resolver: custom resolve function is nil") + } + + return &CustomResolver[K]{ + topology: topology, + resolve: resolve, + }, nil +} + +// Resolve maps key to a shard in the bound topology. +func (resolver *CustomResolver[K]) Resolve(key K) (shard.Shard, error) { + if resolver == nil || resolver.topology == nil || resolver.resolve == nil { + return shard.Shard{}, errors.New("xpg/shard/resolver: custom resolver is not initialized") + } + + id, err := resolver.resolve(key, resolver.topology) + if err != nil { + return shard.Shard{}, err + } + + resolved, ok := resolver.topology.Shard(id) + if !ok { + return shard.Shard{}, &shard.UnknownShardError{ + ShardID: id, + } + } + + return resolved, nil +} diff --git a/shard/resolver/custom_test.go b/shard/resolver/custom_test.go new file mode 100644 index 0000000..764c4f4 --- /dev/null +++ b/shard/resolver/custom_test.go @@ -0,0 +1,177 @@ +package resolver + +import ( + "errors" + "testing" + + "github.com/mkbeh/xpg/shard" +) + +func TestNewCustomValidatesArguments(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a") + + validResolve := ResolveFunc[int]( + func(int, *shard.Topology) (shard.ID, error) { + return "shard-a", nil + }, + ) + + tests := []struct { + name string + topology *shard.Topology + resolve ResolveFunc[int] + wantError string + }{ + { + name: "nil topology", + resolve: validResolve, + wantError: "xpg/shard/resolver: topology is nil or empty", + }, + { + name: "nil resolve function", + topology: topology, + wantError: "xpg/shard/resolver: custom resolve function is nil", + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + + _, err := NewCustom( + test.topology, + test.resolve, + ) + if err == nil { + t.Fatal("expected error") + } + + if got := err.Error(); got != test.wantError { + t.Fatalf("error = %q, want %q", got, test.wantError) + } + }) + } +} + +func TestCustomResolverResolve(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a", "shard-b") + + resolver, err := NewCustom( + topology, + func(key int, gotTopology *shard.Topology) (shard.ID, error) { + if gotTopology != topology { + t.Fatal("resolve function received a different topology") + } + + if key < 100 { + return "shard-a", nil + } + + return "shard-b", nil + }, + ) + if err != nil { + t.Fatalf("NewCustom() error = %v", err) + } + + resolved, err := resolver.Resolve(142) + if err != nil { + t.Fatalf("Resolve() error = %v", err) + } + + if got, want := resolved.ID(), shard.ID("shard-b"); got != want { + t.Fatalf("Resolve().ID() = %q, want %q", got, want) + } +} + +func TestCustomResolverPropagatesResolveError(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a") + sentinel := errors.New("resolve failed") + + resolver, err := NewCustom( + topology, + func(int, *shard.Topology) (shard.ID, error) { + return "", sentinel + }, + ) + if err != nil { + t.Fatalf("NewCustom() error = %v", err) + } + + _, err = resolver.Resolve(1) + if !errors.Is(err, sentinel) { + t.Fatalf("Resolve() error = %v, want sentinel", err) + } +} + +func TestCustomResolverPropagatesErrNoShard(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a") + + resolver, err := NewCustom( + topology, + func(int, *shard.Topology) (shard.ID, error) { + return "", shard.ErrNoShard + }, + ) + if err != nil { + t.Fatalf("NewCustom() error = %v", err) + } + + _, err = resolver.Resolve(1) + if !errors.Is(err, shard.ErrNoShard) { + t.Fatalf("Resolve() error = %v, want ErrNoShard", err) + } +} + +func TestCustomResolverRejectsUnknownShard(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a") + + resolver, err := NewCustom( + topology, + func(int, *shard.Topology) (shard.ID, error) { + return "missing", nil + }, + ) + if err != nil { + t.Fatalf("NewCustom() error = %v", err) + } + + _, err = resolver.Resolve(1) + if !errors.Is(err, shard.ErrUnknownShard) { + t.Fatalf("Resolve() error = %v, want ErrUnknownShard", err) + } + + var unknown *shard.UnknownShardError + if !errors.As(err, &unknown) { + t.Fatalf("Resolve() error = %T, want *shard.UnknownShardError", err) + } + + if got, want := unknown.ShardID, shard.ID("missing"); got != want { + t.Fatalf("ShardID = %q, want %q", got, want) + } +} + +func TestCustomResolverUninitialized(t *testing.T) { + t.Parallel() + + var resolver *CustomResolver[int] + + _, err := resolver.Resolve(1) + if err == nil { + t.Fatal("expected error") + } + + if got, want := err.Error(), "xpg/shard/resolver: custom resolver is not initialized"; got != want { + t.Fatalf("error = %q, want %q", got, want) + } +} diff --git a/shard/resolver/doc.go b/shard/resolver/doc.go new file mode 100644 index 0000000..a565f8f --- /dev/null +++ b/shard/resolver/doc.go @@ -0,0 +1,7 @@ +// Package resolver provides routing strategies for shard.Topology. +// +// Resolvers map application keys to shards using rendezvous hashing, ordered +// ranges, time ranges, or custom routing logic. +// +// Resolvers borrow their topology and must not outlive it. +package resolver diff --git a/shard/resolver/encoder.go b/shard/resolver/encoder.go new file mode 100644 index 0000000..995357e --- /dev/null +++ b/shard/resolver/encoder.go @@ -0,0 +1,90 @@ +package resolver + +import ( + "bytes" + "encoding/binary" + "errors" +) + +// KeyEncoder converts a typed key into stable canonical bytes. +// +// Implementations must be deterministic and safe for concurrent use. Changing +// an encoder changes persistent hash placement and may require data migration. +type KeyEncoder[K any] interface { + Encode(K) ([]byte, error) +} + +// KeyEncoderFunc adapts a function to KeyEncoder. +type KeyEncoderFunc[K any] func(K) ([]byte, error) + +// Encode calls the wrapped encoder function. +func (encoder KeyEncoderFunc[K]) Encode(key K) ([]byte, error) { + if encoder == nil { + return nil, errors.New("xpg/shard/resolver: key encoder function is nil") + } + + return encoder(key) +} + +// StringKeyEncoder encodes the exact bytes stored in a string without +// normalization. +func StringKeyEncoder() KeyEncoder[string] { + return KeyEncoderFunc[string]( + func(key string) ([]byte, error) { + return []byte(key), nil + }, + ) +} + +// BytesKeyEncoder encodes exact bytes and returns a defensive copy. +func BytesKeyEncoder() KeyEncoder[[]byte] { + return KeyEncoderFunc[[]byte]( + func(key []byte) ([]byte, error) { + return bytes.Clone(key), nil + }, + ) +} + +// Bytes16KeyEncoder encodes a 16-byte key exactly. +func Bytes16KeyEncoder() KeyEncoder[[16]byte] { + return KeyEncoderFunc[[16]byte]( + func(key [16]byte) ([]byte, error) { + return bytes.Clone(key[:]), nil + }, + ) +} + +// Bytes32KeyEncoder encodes a 32-byte key exactly. +func Bytes32KeyEncoder() KeyEncoder[[32]byte] { + return KeyEncoderFunc[[32]byte]( + func(key [32]byte) ([]byte, error) { + return bytes.Clone(key[:]), nil + }, + ) +} + +// Int64KeyEncoder encodes a signed integer as big-endian two's-complement. +func Int64KeyEncoder() KeyEncoder[int64] { + return KeyEncoderFunc[int64]( + func(key int64) ([]byte, error) { + encoded := make([]byte, 8) + + binary.BigEndian.PutUint64(encoded, uint64(key)) + + return encoded, nil + }, + ) +} + +// Uint64KeyEncoder encodes an unsigned integer as big-endian bytes. +func Uint64KeyEncoder() KeyEncoder[uint64] { + return KeyEncoderFunc[uint64]( + func(key uint64) ([]byte, error) { + encoded := make([]byte, 8) + + binary.BigEndian.PutUint64(encoded, key) + + return encoded, nil + }, + ) +} diff --git a/shard/resolver/encoder_test.go b/shard/resolver/encoder_test.go new file mode 100644 index 0000000..067187b --- /dev/null +++ b/shard/resolver/encoder_test.go @@ -0,0 +1,148 @@ +package resolver + +import ( + "bytes" + "encoding/hex" + "errors" + "testing" +) + +func TestKeyEncoderFuncNil(t *testing.T) { + t.Parallel() + + var encoder KeyEncoderFunc[int] + + _, err := encoder.Encode(1) + if err == nil { + t.Fatal("expected error") + } + + if got, want := err.Error(), "xpg/shard/resolver: key encoder function is nil"; got != want { + t.Fatalf("error = %q, want %q", got, want) + } +} + +func TestKeyEncoderFunc(t *testing.T) { + t.Parallel() + + sentinel := errors.New("encode failed") + encoder := KeyEncoderFunc[int](func(key int) ([]byte, error) { + if key < 0 { + return nil, sentinel + } + + return []byte{byte(key)}, nil + }) + + encoded, err := encoder.Encode(7) + if err != nil { + t.Fatalf("Encode() error = %v", err) + } + if got, want := encoded, []byte{7}; !bytes.Equal(got, want) { + t.Fatalf("Encode() = %v, want %v", got, want) + } + + if _, err := encoder.Encode(-1); !errors.Is(err, sentinel) { + t.Fatalf("Encode() error = %v, want sentinel", err) + } +} + +func TestStringKeyEncoder(t *testing.T) { + t.Parallel() + + encoded, err := StringKeyEncoder().Encode("a\x00b") + if err != nil { + t.Fatalf("Encode() error = %v", err) + } + + if got, want := string(encoded), "a\x00b"; got != want { + t.Fatalf("Encode() = %q, want %q", got, want) + } +} + +func TestBytesKeyEncoderReturnsDefensiveCopy(t *testing.T) { + t.Parallel() + + key := []byte{1, 2, 3} + encoded, err := BytesKeyEncoder().Encode(key) + if err != nil { + t.Fatalf("Encode() error = %v", err) + } + + encoded[0] = 9 + + if got, want := key[0], byte(1); got != want { + t.Fatalf("input key changed to %d, want %d", got, want) + } +} + +func TestIntegerKeyEncoders(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + got func() ([]byte, error) + want string + }{ + { + name: "int64 positive", + got: func() ([]byte, error) { return Int64KeyEncoder().Encode(1) }, + want: "0000000000000001", + }, + { + name: "int64 negative", + got: func() ([]byte, error) { return Int64KeyEncoder().Encode(-1) }, + want: "ffffffffffffffff", + }, + { + name: "uint64", + got: func() ([]byte, error) { return Uint64KeyEncoder().Encode(0x0102030405060708) }, + want: "0102030405060708", + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + + encoded, err := test.got() + if err != nil { + t.Fatalf("Encode() error = %v", err) + } + + if got := hex.EncodeToString(encoded); got != test.want { + t.Fatalf("Encode() = %s, want %s", got, test.want) + } + }) + } +} + +func TestFixedSizeKeyEncoders(t *testing.T) { + t.Parallel() + + var key16 [16]byte + for index := range key16 { + key16[index] = byte(index) + } + + encoded16, err := Bytes16KeyEncoder().Encode(key16) + if err != nil { + t.Fatalf("Bytes16KeyEncoder.Encode() error = %v", err) + } + if got, want := encoded16, key16[:]; !bytes.Equal(got, want) { + t.Fatalf("Bytes16KeyEncoder.Encode() = %v, want %v", got, want) + } + + var key32 [32]byte + for index := range key32 { + key32[index] = byte(31 - index) + } + + encoded32, err := Bytes32KeyEncoder().Encode(key32) + if err != nil { + t.Fatalf("Bytes32KeyEncoder.Encode() error = %v", err) + } + if got, want := encoded32, key32[:]; !bytes.Equal(got, want) { + t.Fatalf("Bytes32KeyEncoder.Encode() = %v, want %v", got, want) + } +} diff --git a/shard/resolver/hash.go b/shard/resolver/hash.go new file mode 100644 index 0000000..a631afe --- /dev/null +++ b/shard/resolver/hash.go @@ -0,0 +1,175 @@ +package resolver + +import ( + "bytes" + "crypto/sha256" + "encoding/binary" + "errors" + "fmt" + "math" + + "github.com/mkbeh/xpg/shard" +) + +const ( + // rendezvousDomain identifies the persistent hash-placement algorithm. + // Changing it changes shard placement and therefore requires data migration. + rendezvousDomain = "xpg.shard.rendezvous.v1" + + // rendezvousLengthSize is the size of each uint32 length prefix in bytes. + rendezvousLengthSize = 4 +) + +// HashResolver routes keys using rendezvous/HRW hashing with SHA-256. +// +// HashResolver captures the shard set when it is created. The shards remain +// borrowed from the topology, so the resolver must not outlive it. +type HashResolver[K any] struct { + shards []shard.Shard + prefix []byte + encoder KeyEncoder[K] + maxShardIDLength int +} + +// NewHash creates a rendezvous hash resolver bound to topology. +// +// Namespace is an opaque non-empty string and part of the persistent placement +// contract. Changing the namespace, key encoder, shard IDs, or placement format +// changes shard placement and may require data migration. +func NewHash[K any]( + topology *shard.Topology, + namespace string, + encoder KeyEncoder[K], +) (*HashResolver[K], error) { + if err := requireTopology(topology); err != nil { + return nil, err + } + + if encoder == nil { + return nil, errors.New("xpg/shard/resolver: key encoder is nil") + } + + if namespace == "" { + return nil, errors.New("xpg/shard/resolver: hash namespace must not be empty") + } + + if len(namespace) > math.MaxUint32 { + return nil, errors.New("xpg/shard/resolver: hash namespace is too large") + } + + shards := topology.Shards() + maxShardIDLength := 0 + + for _, candidate := range shards { + id := candidate.ID() + + if len(id) > math.MaxUint32 { + return nil, errors.New("xpg/shard/resolver: shard ID is too large") + } + + maxShardIDLength = max(maxShardIDLength, len(id)) + } + + prefixSize := len(rendezvousDomain) + rendezvousLengthSize + len(namespace) + prefix := make([]byte, prefixSize) + + lengthOffset := len(rendezvousDomain) + namespaceOffset := lengthOffset + rendezvousLengthSize + + copy(prefix, rendezvousDomain) + + binary.BigEndian.PutUint32( + prefix[lengthOffset:namespaceOffset], + uint32(len(namespace)), + ) + + copy(prefix[namespaceOffset:], namespace) + + return &HashResolver[K]{ + shards: shards, + prefix: prefix, + encoder: encoder, + maxShardIDLength: maxShardIDLength, + }, nil +} + +// Resolve maps key to a shard using rendezvous hashing. +func (resolver *HashResolver[K]) Resolve(key K) (shard.Shard, error) { + if resolver == nil || len(resolver.shards) == 0 || resolver.encoder == nil { + return shard.Shard{}, errors.New("xpg/shard/resolver: hash resolver is not initialized") + } + + encoded, err := resolver.encoder.Encode(key) + if err != nil { + return shard.Shard{}, fmt.Errorf("xpg/shard/resolver: encode hash key: %w", err) + } + + if len(encoded) > math.MaxUint32 { + return shard.Shard{}, errors.New("xpg/shard/resolver: encoded key is too large") + } + + keyLengthOffset := len(resolver.prefix) + keyOffset := keyLengthOffset + rendezvousLengthSize + idLengthOffset := keyOffset + len(encoded) + idOffset := idLengthOffset + rendezvousLengthSize + + // Persistent placement format: + // + // domain || namespace_length || namespace || + // key_length || key || shard_id_length || shard_id + // + // The candidate-independent prefix and key are written once. Only the shard + // ID suffix is overwritten while evaluating candidates. + scoreInput := make( + []byte, + idOffset+resolver.maxShardIDLength, + ) + + copy(scoreInput, resolver.prefix) + + binary.BigEndian.PutUint32( + scoreInput[keyLengthOffset:keyOffset], + uint32(len(encoded)), + ) + + copy(scoreInput[keyOffset:idLengthOffset], encoded) + + var ( + selected shard.Shard + best [sha256.Size]byte + bestID shard.ID + hasBest bool + ) + + for _, candidate := range resolver.shards { + candidateID := candidate.ID() + inputEnd := idOffset + len(candidateID) + + binary.BigEndian.PutUint32( + scoreInput[idLengthOffset:idOffset], + uint32(len(candidateID)), + ) + + copy(scoreInput[idOffset:inputEnd], candidateID) + + score := sha256.Sum256(scoreInput[:inputEnd]) + comparison := bytes.Compare(score[:], best[:]) + + // Shard ID is the deterministic tie-breaker, so placement does not + // depend on topology registration order when scores are equal. + if !hasBest || + comparison > 0 || + (comparison == 0 && candidateID < bestID) { + selected = candidate + best = score + bestID = candidateID + hasBest = true + } + } + + if !hasBest { + return shard.Shard{}, shard.ErrNoShard + } + + return selected, nil +} diff --git a/shard/resolver/hash_test.go b/shard/resolver/hash_test.go new file mode 100644 index 0000000..f29c5c4 --- /dev/null +++ b/shard/resolver/hash_test.go @@ -0,0 +1,255 @@ +package resolver + +import ( + "errors" + "fmt" + "testing" + + "github.com/mkbeh/xpg/shard" +) + +func TestNewHashValidatesArguments(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a") + + tests := []struct { + name string + topology *shard.Topology + namespace string + encoder KeyEncoder[string] + wantError string + }{ + { + name: "nil topology", + namespace: "users", + encoder: StringKeyEncoder(), + wantError: "xpg/shard/resolver: topology is nil or empty", + }, + { + name: "nil encoder", + topology: topology, + namespace: "users", + wantError: "xpg/shard/resolver: key encoder is nil", + }, + { + name: "empty namespace", + topology: topology, + encoder: StringKeyEncoder(), + wantError: "xpg/shard/resolver: hash namespace must not be empty", + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + + _, err := NewHash( + test.topology, + test.namespace, + test.encoder, + ) + if err == nil { + t.Fatal("expected error") + } + + if got := err.Error(); got != test.wantError { + t.Fatalf("error = %q, want %q", got, test.wantError) + } + }) + } +} + +func TestNewHashTreatsNamespaceAsOpaqueNonEmptyString(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a") + + for _, namespace := range []string{"users", " users ", " "} { + resolver, err := NewHash(topology, namespace, StringKeyEncoder()) + if err != nil { + t.Fatalf("NewHash(%q) error = %v", namespace, err) + } + + resolved, err := resolver.Resolve("alice") + if err != nil { + t.Fatalf("Resolve() error = %v", err) + } + + if got, want := resolved.ID(), shard.ID("shard-a"); got != want { + t.Fatalf("Resolve().ID() = %q, want %q", got, want) + } + } +} + +func TestHashResolverStablePlacementVectors(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a", "shard-b", "shard-c") + resolver, err := NewHash(topology, "users", StringKeyEncoder()) + if err != nil { + t.Fatalf("NewHash() error = %v", err) + } + + tests := []struct { + key string + want shard.ID + }{ + {key: "alice", want: "shard-a"}, + {key: "bob", want: "shard-b"}, + {key: "carol", want: "shard-b"}, + {key: "dave", want: "shard-b"}, + {key: "eve", want: "shard-c"}, + {key: "0", want: "shard-a"}, + {key: "1", want: "shard-b"}, + {key: "2", want: "shard-c"}, + } + + for _, test := range tests { + t.Run(test.key, func(t *testing.T) { + t.Parallel() + + resolved, err := resolver.Resolve(test.key) + if err != nil { + t.Fatalf("Resolve() error = %v", err) + } + + if got := resolved.ID(); got != test.want { + t.Fatalf( + "Resolve(%q).ID() = %q, want %q", + test.key, + got, + test.want, + ) + } + }) + } +} + +func TestHashResolverPlacementDoesNotDependOnTopologyOrder(t *testing.T) { + t.Parallel() + + first := newTestTopology(t, "shard-a", "shard-b", "shard-c") + second := newTestTopology(t, "shard-c", "shard-a", "shard-b") + + firstResolver, err := NewHash(first, "users", StringKeyEncoder()) + if err != nil { + t.Fatalf("NewHash(first) error = %v", err) + } + + secondResolver, err := NewHash(second, "users", StringKeyEncoder()) + if err != nil { + t.Fatalf("NewHash(second) error = %v", err) + } + + for _, key := range []string{"alice", "bob", "carol", "dave", "eve", "user-123"} { + firstShard, err := firstResolver.Resolve(key) + if err != nil { + t.Fatalf("first Resolve(%q) error = %v", key, err) + } + + secondShard, err := secondResolver.Resolve(key) + if err != nil { + t.Fatalf("second Resolve(%q) error = %v", key, err) + } + + if firstShard.ID() != secondShard.ID() { + t.Fatalf( + "Resolve(%q) = %q and %q for different topology orders", + key, + firstShard.ID(), + secondShard.ID(), + ) + } + } +} + +func TestHashResolverAddingShardOnlyMovesKeysToNewShard(t *testing.T) { + t.Parallel() + + before := newTestTopology(t, "shard-a", "shard-b") + after := newTestTopology(t, "shard-a", "shard-b", "shard-c") + + beforeResolver, err := NewHash(before, "users", StringKeyEncoder()) + if err != nil { + t.Fatalf("NewHash(before) error = %v", err) + } + + afterResolver, err := NewHash(after, "users", StringKeyEncoder()) + if err != nil { + t.Fatalf("NewHash(after) error = %v", err) + } + + moved := 0 + + for index := range 256 { + key := fmt.Sprintf("user-%d", index) + + previous, err := beforeResolver.Resolve(key) + if err != nil { + t.Fatalf("before Resolve(%q) error = %v", key, err) + } + + current, err := afterResolver.Resolve(key) + if err != nil { + t.Fatalf("after Resolve(%q) error = %v", key, err) + } + + if previous.ID() == current.ID() { + continue + } + + moved++ + + if current.ID() != "shard-c" { + t.Fatalf( + "Resolve(%q) moved from %q to existing shard %q", + key, + previous.ID(), + current.ID(), + ) + } + } + + if moved == 0 { + t.Fatal("expected at least one key to move to the new shard") + } +} + +func TestHashResolverWrapsEncoderError(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a") + sentinel := errors.New("encode failed") + + resolver, err := NewHash( + topology, + "users", + KeyEncoderFunc[string](func(string) ([]byte, error) { + return nil, sentinel + }), + ) + if err != nil { + t.Fatalf("NewHash() error = %v", err) + } + + _, err = resolver.Resolve("alice") + if !errors.Is(err, sentinel) { + t.Fatalf("Resolve() error = %v, want wrapped sentinel", err) + } +} + +func TestHashResolverUninitialized(t *testing.T) { + t.Parallel() + + var resolver *HashResolver[string] + + _, err := resolver.Resolve("alice") + if err == nil { + t.Fatal("expected error") + } + + if got, want := err.Error(), "xpg/shard/resolver: hash resolver is not initialized"; got != want { + t.Fatalf("error = %q, want %q", got, want) + } +} diff --git a/shard/resolver/helpers_test.go b/shard/resolver/helpers_test.go new file mode 100644 index 0000000..a12991b --- /dev/null +++ b/shard/resolver/helpers_test.go @@ -0,0 +1,59 @@ +package resolver + +import ( + "testing" + + "github.com/jackc/pgx/v5/pgxpool" + "github.com/mkbeh/xpg" + "github.com/mkbeh/xpg/cluster" + "github.com/mkbeh/xpg/shard" +) + +const testDatabaseURL = "postgres://postgres@127.0.0.1:1/postgres?sslmode=disable" + +func newTestTopology(t *testing.T, ids ...shard.ID) *shard.Topology { + t.Helper() + + configs := make([]shard.Config, len(ids)) + + for index, id := range ids { + poolConfig, err := pgxpool.ParseConfig(testDatabaseURL) + if err != nil { + t.Fatalf("pgxpool.ParseConfig() error = %v", err) + } + + poolConfig.MinConns = 0 + poolConfig.MaxConns = 1 + + pool, err := xpg.New( + t.Context(), + poolConfig, + xpg.WithName("shard."+string(id)+".primary"), + ) + if err != nil { + t.Fatalf("xpg.New() error = %v", err) + } + + shardCluster, err := cluster.New(cluster.Config{ + ID: id, + Primary: pool, + }) + if err != nil { + pool.Close() + t.Fatalf("cluster.New() error = %v", err) + } + + t.Cleanup(shardCluster.Close) + + configs[index] = shard.Config{Cluster: shardCluster} + } + + topology, err := shard.NewTopology(configs) + if err != nil { + t.Fatalf("shard.NewTopology() error = %v", err) + } + + t.Cleanup(topology.Close) + + return topology +} diff --git a/shard/resolver/range.go b/shard/resolver/range.go new file mode 100644 index 0000000..9714777 --- /dev/null +++ b/shard/resolver/range.go @@ -0,0 +1,141 @@ +package resolver + +import ( + "cmp" + "errors" + "fmt" + "slices" + "sort" + + "github.com/mkbeh/xpg/shard" +) + +// Range maps the bounded half-open interval [Start, End) to one shard. +// +// Start must be less than End. Ranges may be supplied in any order. +// Gaps are allowed and resolve to shard.ErrNoShard. +type Range[K cmp.Ordered] struct { + Start K + End K + ShardID shard.ID +} + +// RangeResolver routes ordered keys through non-overlapping ranges. +type RangeResolver[K cmp.Ordered] struct { + ranges []rangeEntry[K] +} + +// NewRange creates a resolver from non-overlapping half-open ranges. +// +// The supplied ranges may be unordered. NewRange copies and sorts them by Start, +// validates their boundaries and overlap, and leaves the caller's slice unchanged. +func NewRange[K cmp.Ordered](topology *shard.Topology, ranges []Range[K]) (*RangeResolver[K], error) { + if err := requireTopology(topology); err != nil { + return nil, err + } + + if len(ranges) == 0 { + return nil, errors.New("xpg/shard/resolver: range resolver requires at least one range") + } + + entries := make([]rangeEntry[K], len(ranges)) + + for index, valueRange := range ranges { + if err := requireShardID(valueRange.ShardID); err != nil { + return nil, fmt.Errorf("xpg/shard/resolver: range %d: %w", index, err) + } + + // Using < intentionally rejects empty and reversed ranges as well as + // ranges with NaN boundaries for floating-point key types. + if !(valueRange.Start < valueRange.End) { //nolint:staticcheck // Negated comparison intentionally rejects NaN boundaries. + return nil, fmt.Errorf("xpg/shard/resolver: range %d must satisfy start < end", index) + } + + resolved, ok := topology.Shard(valueRange.ShardID) + if !ok { + return nil, fmt.Errorf( + "xpg/shard/resolver: range %d: %w", + index, + &shard.UnknownShardError{ + ShardID: valueRange.ShardID, + }, + ) + } + + entries[index] = rangeEntry[K]{ + start: valueRange.Start, + end: valueRange.End, + shard: resolved, + sourceIndex: index, + } + } + + slices.SortStableFunc( + entries, + func(left, right rangeEntry[K]) int { + return cmp.Compare(left.start, right.start) + }, + ) + + // Once ranges are sorted by Start, checking adjacent entries is sufficient + // to detect every overlap. + for index := 1; index < len(entries); index++ { + previous := entries[index-1] + current := entries[index] + + // Adjacent half-open ranges are valid: + // + // [0, 100) and [100, 200) + if previous.end <= current.start { + continue + } + + return nil, fmt.Errorf( + "xpg/shard/resolver: ranges %d and %d overlap", + previous.sourceIndex, + current.sourceIndex, + ) + } + + return &RangeResolver[K]{ + ranges: entries, + }, nil +} + +// Resolve returns the shard whose configured range contains key. +func (resolver *RangeResolver[K]) Resolve(key K) (shard.Shard, error) { + if resolver == nil || len(resolver.ranges) == 0 { + return shard.Shard{}, errors.New("xpg/shard/resolver: range resolver is not initialized") + } + + // Non-overlap validation guarantees strictly increasing upper boundaries, + // making this search predicate monotonic. + index := sort.Search( + len(resolver.ranges), + func(index int) bool { + return key < resolver.ranges[index].end + }, + ) + + if index == len(resolver.ranges) { + return shard.Shard{}, shard.ErrNoShard + } + + entry := resolver.ranges[index] + + // The first range ending after key may still start after it when there is a + // gap between configured ranges. + if key < entry.start { + return shard.Shard{}, shard.ErrNoShard + } + + return entry.shard, nil +} + +type rangeEntry[K cmp.Ordered] struct { + start K + end K + shard shard.Shard + + sourceIndex int +} diff --git a/shard/resolver/range_test.go b/shard/resolver/range_test.go new file mode 100644 index 0000000..d001386 --- /dev/null +++ b/shard/resolver/range_test.go @@ -0,0 +1,275 @@ +package resolver + +import ( + "errors" + "math" + "slices" + "testing" + + "github.com/mkbeh/xpg/shard" +) + +func TestNewRangeValidatesArguments(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a") + + tests := []struct { + name string + topology *shard.Topology + ranges []Range[int] + wantError string + }{ + { + name: "nil topology", + ranges: []Range[int]{ + {Start: 0, End: 10, ShardID: "shard-a"}, + }, + wantError: "xpg/shard/resolver: topology is nil or empty", + }, + { + name: "empty ranges", + topology: topology, + wantError: "xpg/shard/resolver: range resolver requires at least one range", + }, + { + name: "empty shard ID", + topology: topology, + ranges: []Range[int]{ + {Start: 0, End: 10}, + }, + wantError: "xpg/shard/resolver: range 0: shard ID must not be empty", + }, + { + name: "empty interval", + topology: topology, + ranges: []Range[int]{ + {Start: 10, End: 10, ShardID: "shard-a"}, + }, + wantError: "xpg/shard/resolver: range 0 must satisfy start < end", + }, + { + name: "reversed interval", + topology: topology, + ranges: []Range[int]{ + {Start: 20, End: 10, ShardID: "shard-a"}, + }, + wantError: "xpg/shard/resolver: range 0 must satisfy start < end", + }, + { + name: "unknown shard", + topology: topology, + ranges: []Range[int]{ + {Start: 0, End: 10, ShardID: "missing"}, + }, + wantError: `xpg/shard/resolver: range 0: xpg/shard: unknown shard "missing"`, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + + _, err := NewRange( + test.topology, + test.ranges, + ) + if err == nil { + t.Fatal("expected error") + } + + if got := err.Error(); got != test.wantError { + t.Fatalf("error = %q, want %q", got, test.wantError) + } + }) + } +} + +func TestNewRangeRejectsNaNBoundaries(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a") + + tests := []struct { + name string + valueRange Range[float64] + }{ + { + name: "NaN start", + valueRange: Range[float64]{ + Start: math.NaN(), + End: 10, + ShardID: "shard-a", + }, + }, + { + name: "NaN end", + valueRange: Range[float64]{ + Start: 0, + End: math.NaN(), + ShardID: "shard-a", + }, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + + _, err := NewRange( + topology, + []Range[float64]{test.valueRange}, + ) + if err == nil { + t.Fatal("expected error") + } + + if got, want := err.Error(), + "xpg/shard/resolver: range 0 must satisfy start < end"; got != want { + t.Fatalf("error = %q, want %q", got, want) + } + }) + } +} + +func TestNewRangeRejectsOverlapUsingSourceIndexes(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a", "shard-b") + + _, err := NewRange(topology, []Range[int]{ + {Start: 100, End: 200, ShardID: "shard-b"}, + {Start: 50, End: 150, ShardID: "shard-a"}, + }) + if err == nil { + t.Fatal("expected overlap error") + } + + if got, want := err.Error(), "xpg/shard/resolver: ranges 1 and 0 overlap"; got != want { + t.Fatalf("error = %q, want %q", got, want) + } +} + +func TestNewRangeDoesNotModifyInputOrder(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a", "shard-b") + ranges := []Range[int]{ + {Start: 100, End: 200, ShardID: "shard-b"}, + {Start: 0, End: 100, ShardID: "shard-a"}, + } + want := slices.Clone(ranges) + + if _, err := NewRange(topology, ranges); err != nil { + t.Fatalf("NewRange() error = %v", err) + } + + if !slices.Equal(ranges, want) { + t.Fatalf("ranges = %+v, want unchanged %+v", ranges, want) + } +} + +func TestRangeResolverHalfOpenBoundariesAndGaps(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a", "shard-b") + resolver, err := NewRange(topology, []Range[int]{ + {Start: 20, End: 30, ShardID: "shard-b"}, + {Start: 0, End: 10, ShardID: "shard-a"}, + }) + if err != nil { + t.Fatalf("NewRange() error = %v", err) + } + + tests := []struct { + key int + wantID shard.ID + wantErr error + }{ + {key: 0, wantID: "shard-a"}, + {key: 9, wantID: "shard-a"}, + {key: 10, wantErr: shard.ErrNoShard}, + {key: 19, wantErr: shard.ErrNoShard}, + {key: 20, wantID: "shard-b"}, + {key: 29, wantID: "shard-b"}, + {key: 30, wantErr: shard.ErrNoShard}, + } + + for _, test := range tests { + resolved, err := resolver.Resolve(test.key) + if test.wantErr != nil { + if !errors.Is(err, test.wantErr) { + t.Fatalf("Resolve(%d) error = %v, want %v", test.key, err, test.wantErr) + } + + continue + } + + if err != nil { + t.Fatalf("Resolve(%d) error = %v", test.key, err) + } + + if got := resolved.ID(); got != test.wantID { + t.Fatalf("Resolve(%d).ID() = %q, want %q", test.key, got, test.wantID) + } + } +} + +func TestRangeResolverAllowsAdjacentRanges(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a", "shard-b") + resolver, err := NewRange(topology, []Range[int]{ + {Start: 0, End: 100, ShardID: "shard-a"}, + {Start: 100, End: 200, ShardID: "shard-b"}, + }) + if err != nil { + t.Fatalf("NewRange() error = %v", err) + } + + resolved, err := resolver.Resolve(100) + if err != nil { + t.Fatalf("Resolve() error = %v", err) + } + + if got, want := resolved.ID(), shard.ID("shard-b"); got != want { + t.Fatalf("Resolve().ID() = %q, want %q", got, want) + } +} + +func TestRangeResolverSupportsStrings(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a", "shard-b") + resolver, err := NewRange(topology, []Range[string]{ + {Start: "a", End: "m", ShardID: "shard-a"}, + {Start: "m", End: "z", ShardID: "shard-b"}, + }) + if err != nil { + t.Fatalf("NewRange() error = %v", err) + } + + resolved, err := resolver.Resolve("m") + if err != nil { + t.Fatalf("Resolve() error = %v", err) + } + + if got, want := resolved.ID(), shard.ID("shard-b"); got != want { + t.Fatalf("Resolve().ID() = %q, want %q", got, want) + } +} + +func TestRangeResolverUninitialized(t *testing.T) { + t.Parallel() + + var resolver *RangeResolver[int] + + _, err := resolver.Resolve(1) + if err == nil { + t.Fatal("expected error") + } + + if got, want := err.Error(), "xpg/shard/resolver: range resolver is not initialized"; got != want { + t.Fatalf("error = %q, want %q", got, want) + } +} diff --git a/shard/resolver/time_range.go b/shard/resolver/time_range.go new file mode 100644 index 0000000..84d329b --- /dev/null +++ b/shard/resolver/time_range.go @@ -0,0 +1,147 @@ +package resolver + +import ( + "errors" + "fmt" + "slices" + "sort" + "time" + + "github.com/mkbeh/xpg/shard" +) + +// TimeRange maps the bounded half-open interval [Start, End) to one shard. +// +// Start must be before End. Ranges may be supplied in any order. +// Gaps are allowed and resolve to shard.ErrNoShard. +type TimeRange struct { + Start time.Time + End time.Time + ShardID shard.ID +} + +// TimeRangeResolver routes time instants through non-overlapping ranges. +type TimeRangeResolver struct { + ranges []timeRangeEntry +} + +// NewTimeRange creates a resolver from non-overlapping half-open time ranges. +// +// The supplied ranges may be unordered. NewTimeRange normalizes boundaries to +// UTC, sorts ranges by Start, validates their overlap, and leaves the caller's +// slice unchanged. +func NewTimeRange(topology *shard.Topology, ranges []TimeRange) (*TimeRangeResolver, error) { + if err := requireTopology(topology); err != nil { + return nil, err + } + + if len(ranges) == 0 { + return nil, errors.New("xpg/shard/resolver: time range resolver requires at least one range") + } + + entries := make([]timeRangeEntry, len(ranges)) + + for index, valueRange := range ranges { + if err := requireShardID(valueRange.ShardID); err != nil { + return nil, fmt.Errorf("xpg/shard/resolver: time range %d: %w", index, err) + } + + start := timeToUTC(valueRange.Start) + end := timeToUTC(valueRange.End) + + if !start.Before(end) { + return nil, fmt.Errorf("xpg/shard/resolver: time range %d must satisfy start < end", index) + } + + resolved, ok := topology.Shard(valueRange.ShardID) + if !ok { + return nil, fmt.Errorf( + "xpg/shard/resolver: time range %d: %w", + index, + &shard.UnknownShardError{ + ShardID: valueRange.ShardID, + }, + ) + } + + entries[index] = timeRangeEntry{ + start: start, + end: end, + shard: resolved, + sourceIndex: index, + } + } + + slices.SortStableFunc( + entries, + func(left, right timeRangeEntry) int { + return left.start.Compare(right.start) + }, + ) + + // Once ranges are sorted by Start, checking adjacent entries is sufficient + // to detect every overlap. + for index := 1; index < len(entries); index++ { + previous := entries[index-1] + current := entries[index] + + // Adjacent half-open ranges are valid: + // + // [00:00, 01:00) and [01:00, 02:00) + if previous.end.After(current.start) { + return nil, fmt.Errorf( + "xpg/shard/resolver: time ranges %d and %d overlap", + previous.sourceIndex, + current.sourceIndex, + ) + } + } + + return &TimeRangeResolver{ + ranges: entries, + }, nil +} + +// Resolve returns the shard whose configured time range contains key. +func (resolver *TimeRangeResolver) Resolve(key time.Time) (shard.Shard, error) { + if resolver == nil || len(resolver.ranges) == 0 { + return shard.Shard{}, errors.New("xpg/shard/resolver: time range resolver is not initialized") + } + + key = timeToUTC(key) + + // Non-overlap validation guarantees strictly increasing upper boundaries, + // making this search predicate monotonic. + index := sort.Search( + len(resolver.ranges), + func(index int) bool { + return key.Before(resolver.ranges[index].end) + }, + ) + + if index == len(resolver.ranges) { + return shard.Shard{}, shard.ErrNoShard + } + + entry := resolver.ranges[index] + + // The first range ending after key may still start after it when there is a + // gap between configured ranges. + if key.Before(entry.start) { + return shard.Shard{}, shard.ErrNoShard + } + + return entry.shard, nil +} + +type timeRangeEntry struct { + start time.Time + end time.Time + shard shard.Shard + + sourceIndex int +} + +func timeToUTC(value time.Time) time.Time { + return value.UTC() +} diff --git a/shard/resolver/time_range_test.go b/shard/resolver/time_range_test.go new file mode 100644 index 0000000..246c34d --- /dev/null +++ b/shard/resolver/time_range_test.go @@ -0,0 +1,237 @@ +package resolver + +import ( + "errors" + "slices" + "testing" + "time" + + "github.com/mkbeh/xpg/shard" +) + +func TestNewTimeRangeValidatesArguments(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a") + start := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC) + end := start.Add(time.Hour) + + tests := []struct { + name string + topology *shard.Topology + ranges []TimeRange + wantError string + }{ + { + name: "nil topology", + ranges: []TimeRange{ + {Start: start, End: end, ShardID: "shard-a"}, + }, + wantError: "xpg/shard/resolver: topology is nil or empty", + }, + { + name: "empty ranges", + topology: topology, + wantError: "xpg/shard/resolver: time range resolver requires at least one range", + }, + { + name: "empty shard ID", + topology: topology, + ranges: []TimeRange{ + {Start: start, End: end}, + }, + wantError: "xpg/shard/resolver: time range 0: shard ID must not be empty", + }, + { + name: "empty interval", + topology: topology, + ranges: []TimeRange{ + {Start: start, End: start, ShardID: "shard-a"}, + }, + wantError: "xpg/shard/resolver: time range 0 must satisfy start < end", + }, + { + name: "reversed interval", + topology: topology, + ranges: []TimeRange{ + {Start: end, End: start, ShardID: "shard-a"}, + }, + wantError: "xpg/shard/resolver: time range 0 must satisfy start < end", + }, + { + name: "unknown shard", + topology: topology, + ranges: []TimeRange{ + {Start: start, End: end, ShardID: "missing"}, + }, + wantError: `xpg/shard/resolver: time range 0: xpg/shard: unknown shard "missing"`, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + + _, err := NewTimeRange( + test.topology, + test.ranges, + ) + if err == nil { + t.Fatal("expected error") + } + + if got := err.Error(); got != test.wantError { + t.Fatalf("error = %q, want %q", got, test.wantError) + } + }) + } +} + +func TestNewTimeRangeRejectsOverlapUsingSourceIndexes(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a", "shard-b") + base := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC) + + _, err := NewTimeRange(topology, []TimeRange{ + {Start: base.Add(2 * time.Hour), End: base.Add(4 * time.Hour), ShardID: "shard-b"}, + {Start: base.Add(time.Hour), End: base.Add(3 * time.Hour), ShardID: "shard-a"}, + }) + if err == nil { + t.Fatal("expected overlap error") + } + + if got, want := err.Error(), "xpg/shard/resolver: time ranges 1 and 0 overlap"; got != want { + t.Fatalf("error = %q, want %q", got, want) + } +} + +func TestNewTimeRangeDoesNotModifyInput(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a", "shard-b") + location := time.FixedZone("UTC+3", 3*60*60) + base := time.Date(2026, 1, 1, 0, 0, 0, 0, location) + ranges := []TimeRange{ + {Start: base.Add(time.Hour), End: base.Add(2 * time.Hour), ShardID: "shard-b"}, + {Start: base, End: base.Add(time.Hour), ShardID: "shard-a"}, + } + want := slices.Clone(ranges) + + if _, err := NewTimeRange(topology, ranges); err != nil { + t.Fatalf("NewTimeRange() error = %v", err) + } + + if !slices.Equal(ranges, want) { + t.Fatalf("ranges = %+v, want unchanged %+v", ranges, want) + } +} + +func TestTimeRangeResolverHalfOpenBoundariesAndGaps(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a", "shard-b") + base := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC) + resolver, err := NewTimeRange(topology, []TimeRange{ + {Start: base.Add(2 * time.Hour), End: base.Add(3 * time.Hour), ShardID: "shard-b"}, + {Start: base, End: base.Add(time.Hour), ShardID: "shard-a"}, + }) + if err != nil { + t.Fatalf("NewTimeRange() error = %v", err) + } + + tests := []struct { + key time.Time + wantID shard.ID + wantErr error + }{ + {key: base, wantID: "shard-a"}, + {key: base.Add(time.Hour - time.Nanosecond), wantID: "shard-a"}, + {key: base.Add(time.Hour), wantErr: shard.ErrNoShard}, + {key: base.Add(2 * time.Hour), wantID: "shard-b"}, + {key: base.Add(3 * time.Hour), wantErr: shard.ErrNoShard}, + } + + for _, test := range tests { + resolved, err := resolver.Resolve(test.key) + if test.wantErr != nil { + if !errors.Is(err, test.wantErr) { + t.Fatalf("Resolve(%v) error = %v, want %v", test.key, err, test.wantErr) + } + + continue + } + + if err != nil { + t.Fatalf("Resolve(%v) error = %v", test.key, err) + } + + if got := resolved.ID(); got != test.wantID { + t.Fatalf("Resolve(%v).ID() = %q, want %q", test.key, got, test.wantID) + } + } +} + +func TestTimeRangeResolverNormalizesToUTC(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a") + start := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC) + end := start.Add(time.Hour) + resolver, err := NewTimeRange(topology, []TimeRange{ + {Start: start, End: end, ShardID: "shard-a"}, + }) + if err != nil { + t.Fatalf("NewTimeRange() error = %v", err) + } + + location := time.FixedZone("UTC+3", 3*60*60) + key := start.Add(30 * time.Minute).In(location) + + resolved, err := resolver.Resolve(key) + if err != nil { + t.Fatalf("Resolve() error = %v", err) + } + + if got, want := resolved.ID(), shard.ID("shard-a"); got != want { + t.Fatalf("Resolve().ID() = %q, want %q", got, want) + } +} + +func TestTimeRangeResolverAllowsAdjacentRanges(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a", "shard-b") + base := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC) + resolver, err := NewTimeRange(topology, []TimeRange{ + {Start: base, End: base.Add(time.Hour), ShardID: "shard-a"}, + {Start: base.Add(time.Hour), End: base.Add(2 * time.Hour), ShardID: "shard-b"}, + }) + if err != nil { + t.Fatalf("NewTimeRange() error = %v", err) + } + + resolved, err := resolver.Resolve(base.Add(time.Hour)) + if err != nil { + t.Fatalf("Resolve() error = %v", err) + } + + if got, want := resolved.ID(), shard.ID("shard-b"); got != want { + t.Fatalf("Resolve().ID() = %q, want %q", got, want) + } +} + +func TestTimeRangeResolverUninitialized(t *testing.T) { + t.Parallel() + + var resolver *TimeRangeResolver + + _, err := resolver.Resolve(time.Now()) + if err == nil { + t.Fatal("expected error") + } + + if got, want := err.Error(), "xpg/shard/resolver: time range resolver is not initialized"; got != want { + t.Fatalf("error = %q, want %q", got, want) + } +} diff --git a/shard/resolver/validation.go b/shard/resolver/validation.go new file mode 100644 index 0000000..c70966f --- /dev/null +++ b/shard/resolver/validation.go @@ -0,0 +1,23 @@ +package resolver + +import ( + "errors" + + "github.com/mkbeh/xpg/shard" +) + +func requireTopology(topology *shard.Topology) error { + if topology == nil || topology.Len() == 0 { + return errors.New("xpg/shard/resolver: topology is nil or empty") + } + + return nil +} + +func requireShardID(id shard.ID) error { + if id == "" { + return errors.New("shard ID must not be empty") + } + + return nil +} diff --git a/shard/shard.go b/shard/shard.go new file mode 100644 index 0000000..c3afb65 --- /dev/null +++ b/shard/shard.go @@ -0,0 +1,99 @@ +package shard + +import ( + "context" + + "github.com/jackc/pgx/v5" + "github.com/mkbeh/xpg" + "github.com/mkbeh/xpg/cluster" +) + +// ID identifies one logical shard. +type ID = cluster.ID + +// Shard is a borrowed handle to one cluster registered in a Topology. +// +// Shard exposes shard-local operations without exposing cluster lifecycle or +// replica-set management. +type Shard struct { + cluster *cluster.Cluster +} + +// ID returns the logical shard ID. +func (s Shard) ID() ID { + if s.cluster == nil { + return "" + } + + return s.cluster.ID() +} + +// Label returns one shard label without allocating a copy of all labels. +func (s Shard) Label(key string) (string, bool) { + if s.cluster == nil { + return "", false + } + + return s.cluster.Label(key) +} + +// Labels returns a defensive copy of shard labels. +func (s Shard) Labels() map[string]string { + if s.cluster == nil { + return nil + } + + return s.cluster.Labels() +} + +// Primary returns the shard primary pool. +// +// Primary returns nil when no primary is configured. The returned pool is +// borrowed and must not be closed separately. +func (s Shard) Primary() *xpg.Pool { + if s.cluster == nil { + return nil + } + + return s.cluster.Primary() +} + +// ReadPool returns a borrowed pool according to policy. +func (s Shard) ReadPool( + ctx context.Context, + policy cluster.ReadPolicy, +) (*xpg.Pool, error) { + if s.cluster == nil { + return nil, ErrNoShard + } + + return s.cluster.ReadPool(ctx, policy) +} + +// InPrimaryTx executes fn in a transaction on the shard primary. +func (s Shard) InPrimaryTx( + ctx context.Context, + options pgx.TxOptions, + fn func(context.Context, pgx.Tx) error, +) error { + if s.cluster == nil { + return ErrNoShard + } + + return s.cluster.InPrimaryTx(ctx, options, fn) +} + +// InReadTx executes fn in a read-only transaction on a pool selected according +// to policy within the shard. +func (s Shard) InReadTx( + ctx context.Context, + policy cluster.ReadPolicy, + options cluster.ReadTxOptions, + fn func(context.Context, pgx.Tx) error, +) error { + if s.cluster == nil { + return ErrNoShard + } + + return s.cluster.InReadTx(ctx, policy, options, fn) +} diff --git a/shard/shard_test.go b/shard/shard_test.go new file mode 100644 index 0000000..400822b --- /dev/null +++ b/shard/shard_test.go @@ -0,0 +1,97 @@ +package shard + +import ( + "errors" + "testing" + + "github.com/jackc/pgx/v5" + "github.com/mkbeh/xpg/cluster" +) + +func TestShardZeroValue(t *testing.T) { + t.Parallel() + + var shard Shard + + if got := shard.ID(); got != "" { + t.Fatalf("ID() = %q, want empty", got) + } + + if value, ok := shard.Label("region"); ok || value != "" { + t.Fatalf("Label() = %q, %v; want empty, false", value, ok) + } + + if labels := shard.Labels(); labels != nil { + t.Fatalf("Labels() = %#v, want nil", labels) + } + + if primary := shard.Primary(); primary != nil { + t.Fatalf("Primary() = %p, want nil", primary) + } + + if _, err := shard.ReadPool(t.Context(), cluster.ReadPrimary); !errors.Is(err, ErrNoShard) { + t.Fatalf("ReadPool() error = %v, want ErrNoShard", err) + } + + if err := shard.InPrimaryTx(t.Context(), pgx.TxOptions{}, nil); !errors.Is(err, ErrNoShard) { + t.Fatalf("InPrimaryTx() error = %v, want ErrNoShard", err) + } + + if err := shard.InReadTx( + t.Context(), + cluster.ReadPrimary, + cluster.ReadTxOptions{}, + nil, + ); !errors.Is(err, ErrNoShard) { + t.Fatalf("InReadTx() error = %v, want ErrNoShard", err) + } +} + +func TestShardDelegatesClusterMetadataAndRouting(t *testing.T) { + t.Parallel() + + shardCluster := newTestCluster(t, "shard-a", map[string]string{ + "region": "eu-west", + "role": "", + }) + + topology, err := NewTopology([]Config{{Cluster: shardCluster}}) + if err != nil { + t.Fatalf("NewTopology() error = %v", err) + } + t.Cleanup(topology.Close) + + resolved := topology.At(0) + + if got, want := resolved.ID(), ID("shard-a"); got != want { + t.Fatalf("ID() = %q, want %q", got, want) + } + + if got, ok := resolved.Label("region"); !ok || got != "eu-west" { + t.Fatalf("Label(region) = %q, %v", got, ok) + } + + if got, ok := resolved.Label("role"); !ok || got != "" { + t.Fatalf("Label(role) = %q, %v", got, ok) + } + + labels := resolved.Labels() + labels["region"] = "changed" + + if got, _ := resolved.Label("region"); got != "eu-west" { + t.Fatalf("Label(region) after mutation = %q, want eu-west", got) + } + + if resolved.Primary() != shardCluster.Primary() { + t.Fatal("Primary() did not return cluster primary") + } + + pool, err := resolved.ReadPool(t.Context(), cluster.ReadPrimary) + if err != nil { + t.Fatalf("ReadPool() error = %v", err) + } + + if pool != shardCluster.Primary() { + t.Fatal("ReadPool() did not return cluster primary") + } +} diff --git a/shard/topology.go b/shard/topology.go new file mode 100644 index 0000000..db7468a --- /dev/null +++ b/shard/topology.go @@ -0,0 +1,133 @@ +package shard + +import ( + "errors" + "fmt" + "slices" + "sync" + + "github.com/mkbeh/xpg/cluster" +) + +// Config registers one Cluster as a logical shard. +// +// The shard ID and labels are provided by Cluster. +type Config struct { + Cluster *cluster.Cluster +} + +// Topology is an immutable ordered set of logical shards and their PostgreSQL +// clusters. +// +// NewTopology takes ownership of all configured clusters only after it returns +// successfully. Close closes every owned cluster exactly once. +type Topology struct { + shards []Shard + indexByID map[ID]int + + closeOnce sync.Once +} + +// NewTopology validates and creates an immutable topology. +// +// Shards retain their registration order. Every cluster must have a non-empty +// and unique ID. +func NewTopology(configs []Config) (*Topology, error) { + if len(configs) == 0 { + return nil, errors.New( + "xpg/shard: topology must contain at least one shard", + ) + } + + shards := make([]Shard, len(configs)) + indexByID := make(map[ID]int, len(configs)) + + for index, config := range configs { + if config.Cluster == nil { + return nil, fmt.Errorf( + "xpg/shard: shard %d: cluster is nil", + index, + ) + } + + id := config.Cluster.ID() + if id == "" { + return nil, fmt.Errorf( + "xpg/shard: shard %d: cluster ID must not be empty", + index, + ) + } + + if previousIndex, exists := indexByID[id]; exists { + return nil, fmt.Errorf( + "xpg/shard: duplicate shard ID %q at indexes %d and %d", + id, + previousIndex, + index, + ) + } + + shards[index] = Shard{ + cluster: config.Cluster, + } + indexByID[id] = index + } + + return &Topology{ + shards: shards, + indexByID: indexByID, + }, nil +} + +// Len returns the number of registered shards. +func (t *Topology) Len() int { + if t == nil { + return 0 + } + + return len(t.shards) +} + +// At returns the shard at index in registration order. +// +// At panics when t is nil or index is out of range. +func (t *Topology) At(index int) Shard { + return t.shards[index] +} + +// Shard returns the shard with id. +func (t *Topology) Shard(id ID) (Shard, bool) { + if t == nil { + return Shard{}, false + } + + index, ok := t.indexByID[id] + if !ok { + return Shard{}, false + } + + return t.shards[index], true +} + +// Shards returns a defensive copy of shards in registration order. +func (t *Topology) Shards() []Shard { + if t == nil { + return nil + } + + return slices.Clone(t.shards) +} + +// Close closes owned clusters in reverse registration order. Close is safe to +// call multiple times. +func (t *Topology) Close() { + if t == nil { + return + } + + t.closeOnce.Do(func() { + for _, shard := range slices.Backward(t.shards) { + shard.cluster.Close() + } + }) +} diff --git a/shard/topology_test.go b/shard/topology_test.go new file mode 100644 index 0000000..81712e5 --- /dev/null +++ b/shard/topology_test.go @@ -0,0 +1,162 @@ +package shard + +import ( + "testing" +) + +func TestNewTopologyRequiresShard(t *testing.T) { + t.Parallel() + + topology, err := NewTopology(nil) + if err == nil { + topology.Close() + t.Fatal("expected error") + } + + if got, want := err.Error(), "xpg/shard: topology must contain at least one shard"; got != want { + t.Fatalf("error = %q, want %q", got, want) + } +} + +func TestNewTopologyRejectsNilCluster(t *testing.T) { + t.Parallel() + + topology, err := NewTopology([]Config{{}}) + if err == nil { + topology.Close() + t.Fatal("expected error") + } + + if got, want := err.Error(), "xpg/shard: shard 0: cluster is nil"; got != want { + t.Fatalf("error = %q, want %q", got, want) + } +} + +func TestNewTopologyRejectsEmptyClusterID(t *testing.T) { + t.Parallel() + + shardCluster := newTestCluster(t, "", nil) + topology, err := NewTopology([]Config{{Cluster: shardCluster}}) + if err == nil { + topology.Close() + t.Fatal("expected error") + } + + if got, want := err.Error(), "xpg/shard: shard 0: cluster ID must not be empty"; got != want { + t.Fatalf("error = %q, want %q", got, want) + } +} + +func TestNewTopologyRejectsDuplicateIDs(t *testing.T) { + t.Parallel() + + first := newTestCluster(t, "shard-a", nil) + second := newTestCluster(t, "shard-a", nil) + topology, err := NewTopology([]Config{ + {Cluster: first}, + {Cluster: second}, + }) + if err == nil { + topology.Close() + t.Fatal("expected error") + } + + if got, want := err.Error(), + `xpg/shard: duplicate shard ID "shard-a" at indexes 0 and 1`; got != want { + t.Fatalf("error = %q, want %q", got, want) + } +} + +func TestTopologyPreservesRegistrationOrder(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-b", "shard-a", "shard-c") + + if got, want := topology.Len(), 3; got != want { + t.Fatalf("Len() = %d, want %d", got, want) + } + + want := []ID{"shard-b", "shard-a", "shard-c"} + + for index, wantID := range want { + if got := topology.At(index).ID(); got != wantID { + t.Fatalf("At(%d).ID() = %q, want %q", index, got, wantID) + } + + resolved, ok := topology.Shard(wantID) + if !ok { + t.Fatalf("Shard(%q) not found", wantID) + } + + if got := resolved.ID(); got != wantID { + t.Fatalf("Shard(%q).ID() = %q", wantID, got) + } + } +} + +func TestTopologyShardsReturnsDefensiveCopy(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a", "shard-b") + + shards := topology.Shards() + shards[0] = Shard{} + + if got, want := topology.At(0).ID(), ID("shard-a"); got != want { + t.Fatalf("At(0).ID() = %q, want %q", got, want) + } +} + +func TestTopologyShardUnknownID(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a") + + resolved, ok := topology.Shard("missing") + if ok { + t.Fatalf("Shard() = %+v, true; want false", resolved) + } +} + +func TestTopologyAtPanicsOutOfRange(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a") + + defer func() { + if recover() == nil { + t.Fatal("expected panic") + } + }() + + _ = topology.At(1) +} + +func TestTopologyNilReceiver(t *testing.T) { + t.Parallel() + + var topology *Topology + + if got := topology.Len(); got != 0 { + t.Fatalf("Len() = %d, want 0", got) + } + + if shards := topology.Shards(); shards != nil { + t.Fatalf("Shards() = %#v, want nil", shards) + } + + if resolved, ok := topology.Shard("shard-a"); ok || resolved.ID() != "" { + t.Fatalf("Shard() = %+v, %v; want zero, false", resolved, ok) + } + + topology.Close() +} + +func TestTopologyCloseIsIdempotent(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a", "shard-b") + + topology.Close() + topology.Close() +} diff --git a/stats.go b/stats.go new file mode 100644 index 0000000..8d1aafb --- /dev/null +++ b/stats.go @@ -0,0 +1,81 @@ +package xpg + +import "time" + +// PoolStats is a detached point-in-time snapshot of connection pool statistics. +// +// Counter fields are cumulative for the lifetime of the pool. +type PoolStats struct { + // Current state. + + // AcquiredConns is the number of connections currently checked out from the + // pool. + AcquiredConns int32 + + // ConstructingConns is the number of connections currently being created. + ConstructingConns int32 + + // IdleConns is the number of currently idle connections. + IdleConns int32 + + // MaxConns is the maximum number of connections allowed by the pool. + MaxConns int32 + + // TotalConns is the number of acquired, idle, and constructing connections. + TotalConns int32 + + // Acquire lifecycle. + + // AcquireCount is the cumulative number of successful connection acquires. + AcquireCount int64 + + // AcquireDuration is the cumulative duration of successful connection + // acquires. + AcquireDuration time.Duration + + // CanceledAcquireCount is the cumulative number of connection acquires + // canceled by context cancellation. + CanceledAcquireCount int64 + + // EmptyAcquireCount is the cumulative number of successful acquires that + // waited because the pool was empty. + EmptyAcquireCount int64 + + // EmptyAcquireWaitTime is the cumulative time spent waiting on successful + // acquires while the pool was empty. + EmptyAcquireWaitTime time.Duration + + // Connection lifecycle. + + // NewConnsCount is the cumulative number of connections created by the pool. + NewConnsCount int64 + + // MaxIdleDestroyCount is the cumulative number of connections closed because + // they exceeded MaxConnIdleTime. + MaxIdleDestroyCount int64 + + // MaxLifetimeDestroyCount is the cumulative number of connections closed + // because they exceeded MaxConnLifetime. + MaxLifetimeDestroyCount int64 +} + +// Stats returns a detached snapshot of the current pool statistics. +func (p *Pool) Stats() PoolStats { + stats := p.pool.Stat() + + return PoolStats{ + AcquiredConns: stats.AcquiredConns(), + ConstructingConns: stats.ConstructingConns(), + IdleConns: stats.IdleConns(), + MaxConns: stats.MaxConns(), + TotalConns: stats.TotalConns(), + AcquireCount: stats.AcquireCount(), + AcquireDuration: stats.AcquireDuration(), + CanceledAcquireCount: stats.CanceledAcquireCount(), + EmptyAcquireCount: stats.EmptyAcquireCount(), + EmptyAcquireWaitTime: stats.EmptyAcquireWaitTime(), + NewConnsCount: stats.NewConnsCount(), + MaxIdleDestroyCount: stats.MaxIdleDestroyCount(), + MaxLifetimeDestroyCount: stats.MaxLifetimeDestroyCount(), + } +} diff --git a/tx.go b/tx.go index ed977b4..b33ebfc 100644 --- a/tx.go +++ b/tx.go @@ -1,33 +1,78 @@ -package postgres +package xpg import ( "context" + "errors" + "fmt" "github.com/jackc/pgx/v5" ) -type ctxTx struct { - tx pgx.Tx -} +// InTx executes fn in a transaction configured by txOptions. +// +// If fn returns nil, the transaction is committed; otherwise it is rolled back. +// If fn panics, rollback is attempted before the panic is propagated. The +// callback must not call Commit or Rollback; InTx owns transaction +// finalization. +// +// The callback receives ctx unchanged. Context cancellation does not +// automatically finalize the transaction while fn is running; fn should +// observe ctx and return promptly. +func (p *Pool) InTx( + ctx context.Context, + txOptions pgx.TxOptions, + fn func(context.Context, pgx.Tx) error, +) error { + if fn == nil { + return errors.New("xpg: transaction function is nil") + } -type txKey struct{} + err := pgx.BeginTxFunc( + ctx, + p.pool, + txOptions, + func(tx pgx.Tx) error { + return fn(ctx, tx) + }, + ) + if err != nil { + return fmt.Errorf("xpg: transaction: %w", err) + } -var ( - txMarkerKey = &txKey{} - nullTx = &ctxTx{} -) + return nil +} -func injectTx(ctx context.Context, tx pgx.Tx) context.Context { - t := &ctxTx{ - tx: tx, +// InSavepoint executes fn within a PostgreSQL savepoint. +// +// If fn returns nil, the savepoint is released; otherwise it is rolled back. +// If fn panics, rollback is attempted before the panic is propagated. The +// callback must not call Commit or Rollback; InSavepoint owns savepoint +// finalization. +// +// The callback receives ctx unchanged and should observe its cancellation. +func InSavepoint( + ctx context.Context, + tx pgx.Tx, + fn func(context.Context, pgx.Tx) error, +) error { + if tx == nil { + return errors.New("xpg: transaction is nil") + } + + if fn == nil { + return errors.New("xpg: savepoint function is nil") } - return context.WithValue(ctx, txMarkerKey, t) -} -func extractTx(ctx context.Context) pgx.Tx { - t, ok := ctx.Value(txMarkerKey).(*ctxTx) - if !ok || t == nil { - return nullTx.tx + err := pgx.BeginFunc( + ctx, + tx, + func(savepoint pgx.Tx) error { + return fn(ctx, savepoint) + }, + ) + if err != nil { + return fmt.Errorf("xpg: savepoint: %w", err) } - return t.tx + + return nil } diff --git a/tx_test.go b/tx_test.go new file mode 100644 index 0000000..f44d20e --- /dev/null +++ b/tx_test.go @@ -0,0 +1,233 @@ +package xpg + +import ( + "context" + "errors" + "testing" + + "github.com/jackc/pgx/v5" +) + +func TestPoolInTxNilFunc(t *testing.T) { + t.Parallel() + + err := (&Pool{}).InTx( + t.Context(), + pgx.TxOptions{}, + nil, + ) + + assertErrorMessage( + t, + err, + "xpg: transaction function is nil", + ) +} + +func TestInSavepointNilTx(t *testing.T) { + t.Parallel() + + called := false + + err := InSavepoint( + t.Context(), + nil, + func(context.Context, pgx.Tx) error { + called = true + + return nil + }, + ) + + assertErrorMessage( + t, + err, + "xpg: transaction is nil", + ) + + if called { + t.Fatal("savepoint function was called with a nil transaction") + } +} + +func TestInSavepointNilFunc(t *testing.T) { + t.Parallel() + + err := InSavepoint( + t.Context(), + &savepointParentTx{}, + nil, + ) + + assertErrorMessage( + t, + err, + "xpg: savepoint function is nil", + ) +} + +func TestInSavepoint(t *testing.T) { + t.Parallel() + + savepoint := &savepointTestTx{} + parent := &savepointParentTx{ + savepoint: savepoint, + } + + called := false + + err := InSavepoint( + t.Context(), + parent, + func(_ context.Context, tx pgx.Tx) error { + called = true + + if tx != savepoint { + t.Fatal("unexpected savepoint transaction") + } + + return nil + }, + ) + if err != nil { + t.Fatalf("InSavepoint returned an error: %v", err) + } + + if !called { + t.Fatal("savepoint function was not called") + } + + if savepoint.commitCalls != 1 { + t.Fatalf( + "commit calls = %d, want 1", + savepoint.commitCalls, + ) + } +} + +func TestInSavepointPreservesCallbackError(t *testing.T) { + t.Parallel() + + expectedErr := errors.New("callback failed") + savepoint := &savepointTestTx{} + parent := &savepointParentTx{ + savepoint: savepoint, + } + + err := InSavepoint( + t.Context(), + parent, + func(context.Context, pgx.Tx) error { + return expectedErr + }, + ) + if !errors.Is(err, expectedErr) { + t.Fatalf("original error was not preserved: %v", err) + } + + if savepoint.commitCalls != 0 { + t.Fatalf( + "commit calls = %d, want 0", + savepoint.commitCalls, + ) + } + + if savepoint.rollbackCalls == 0 { + t.Fatal("savepoint was not rolled back") + } +} + +func TestInSavepointPreservesBeginError(t *testing.T) { + t.Parallel() + + expectedErr := errors.New("begin failed") + parent := &savepointParentTx{ + beginErr: expectedErr, + } + + called := false + + err := InSavepoint( + t.Context(), + parent, + func(context.Context, pgx.Tx) error { + called = true + + return nil + }, + ) + if !errors.Is(err, expectedErr) { + t.Fatalf("original error was not preserved: %v", err) + } + + if called { + t.Fatal("savepoint function was called after begin failure") + } +} + +type savepointParentTx struct { + pgx.Tx + + savepoint pgx.Tx + beginErr error +} + +func (tx *savepointParentTx) Begin(context.Context) (pgx.Tx, error) { + if tx.beginErr != nil { + return nil, tx.beginErr + } + + return tx.savepoint, nil +} + +type savepointTestTx struct { + pgx.Tx + + closed bool + commitCalls int + rollbackCalls int +} + +func (tx *savepointTestTx) Commit(context.Context) error { + tx.commitCalls++ + + if tx.closed { + return pgx.ErrTxClosed + } + + tx.closed = true + + return nil +} + +func (tx *savepointTestTx) Rollback(context.Context) error { + tx.rollbackCalls++ + + if tx.closed { + return pgx.ErrTxClosed + } + + tx.closed = true + + return nil +} + +func assertErrorMessage( + t *testing.T, + err error, + want string, +) { + t.Helper() + + if err == nil { + t.Fatalf("expected error %q, got nil", want) + } + + if err.Error() != want { + t.Fatalf( + "error = %q, want %q", + err, + want, + ) + } +} diff --git a/utils.go b/utils.go deleted file mode 100644 index 455cb02..0000000 --- a/utils.go +++ /dev/null @@ -1,17 +0,0 @@ -package postgres - -import ( - "hash/fnv" - - "github.com/google/uuid" -) - -func GenerateUUID() string { - return uuid.New().String() -} - -func StringAsHash64(s string) uint64 { - hash := fnv.New64() - _, _ = hash.Write([]byte(s)) - return hash.Sum64() -}