diff --git a/.claude/settings.json b/.claude/settings.json index 4140b11098..da49281a4c 100644 --- a/.claude/settings.json +++ b/.claude/settings.json @@ -4,7 +4,8 @@ "golang@ai-helpers": true, "github@ai-helpers": true, "gopls-lsp@claude-plugins-official": true, - "typescript-lsp@claude-plugins-official": true + "typescript-lsp@claude-plugins-official": true, + "gopls-lsp@ai-helpers": true }, "extraKnownMarketplaces": { "ai-helpers": { diff --git a/cmd/sippy-daemon/main.go b/cmd/sippy-daemon/main.go index 6b11c665f1..edec5cb3f5 100644 --- a/cmd/sippy-daemon/main.go +++ b/cmd/sippy-daemon/main.go @@ -7,9 +7,12 @@ import ( "os" "time" + "github.com/openshift/sippy/pkg/api/jobartifacts" + "github.com/openshift/sippy/pkg/api/jobrunscan" "github.com/openshift/sippy/pkg/bigquery/bqlabel" "github.com/pkg/errors" "github.com/prometheus/client_golang/prometheus/promhttp" + "github.com/riverqueue/river" log "github.com/sirupsen/logrus" "github.com/spf13/cobra" "github.com/spf13/pflag" @@ -20,6 +23,7 @@ import ( "github.com/openshift/sippy/pkg/flags" "github.com/openshift/sippy/pkg/github/commenter" "github.com/openshift/sippy/pkg/sippyserver" + "github.com/openshift/sippy/pkg/sippyserver/workqueue" "github.com/openshift/sippy/pkg/version" ) @@ -122,6 +126,14 @@ func NewSippyDaemonCommand() *cobra.Command { processes = append(processes, sippyserver.NewWorkProcessor(dbc, bigQueryClient, gcsClient.Bucket(f.GoogleCloudFlags.StorageBucket), cacheClient, ghCommenter, 10, 5*time.Minute, 5*time.Second, f.GithubCommenterFlags.CommentProcessingDryRun)) } + // Set up River work queue for async job processing + riverProcess, err := setupRiverProcess(cmd.Context(), f) + if err != nil { + log.WithError(err).Error("failed to set up River work queue, async re-evaluation will not be available") + } else if riverProcess != nil { + processes = append(processes, riverProcess) + } + daemonServer := sippyserver.NewDaemonServer(processes) // Serve our metrics endpoint for prometheus to scrape @@ -146,6 +158,68 @@ func NewSippyDaemonCommand() *cobra.Command { return cmd } +const ( + reevaluateWorkerCount = 8 + jobRetentionPeriod = 8 * 24 * time.Hour // 8 days +) + +func setupRiverProcess(ctx context.Context, f *SippyDaemonFlags) (sippyserver.DaemonProcess, error) { + dsn := f.DBFlags.DSN + if dsn == "" { + log.Info("no database DSN configured, skipping River work queue setup") + return nil, nil + } + + dbc, err := f.DBFlags.GetDBClient() + if err != nil { + return nil, fmt.Errorf("getting DB client for River: %w", err) + } + + cacheClient, err := f.CacheFlags.GetCacheClient() + if err != nil { + return nil, fmt.Errorf("getting cache client for River: %w", err) + } + + opCtx := bqlabel.OperationalContext{ + App: bqlabel.AppSippy, + Command: "sippy-daemon", + Environment: bqlabel.EnvDaemon, + } + bqClient, err := f.BigQueryFlags.GetBigQueryClient(ctx, opCtx, cacheClient, f.GoogleCloudFlags.ServiceAccountCredentialFile) + if err != nil { + return nil, fmt.Errorf("getting BigQuery client for River: %w", err) + } + + gcsClient, err := gcs.NewGCSClient(ctx, + f.GoogleCloudFlags.ServiceAccountCredentialFile, + f.GoogleCloudFlags.OAuthClientCredentialFile, + ) + if err != nil { + return nil, fmt.Errorf("getting GCS client for River: %w", err) + } + + artifactMgr := jobartifacts.NewManager(ctx) + evaluator := jobrunscan.NewReEvaluator(bqClient, gcsClient, f.GoogleCloudFlags.StorageBucket, dbc, cacheClient, artifactMgr, false) + + workers := river.NewWorkers() + river.AddWorker(workers, jobrunscan.NewReevaluateWorker(evaluator)) + + riverSetup, err := workqueue.Setup(ctx, workqueue.SetupConfig{ + DatabaseDSN: dsn, + Queues: map[string]river.QueueConfig{ + jobrunscan.ReevaluateQueue: {MaxWorkers: reevaluateWorkerCount}, + }, + Workers: workers, + CompletedJobRetention: jobRetentionPeriod, + DiscardedJobRetention: jobRetentionPeriod, + }) + if err != nil { + return nil, fmt.Errorf("setting up River: %w", err) + } + + return workqueue.NewRiverProcess(riverSetup.Client), nil +} + func main() { // Set log level level, err := log.ParseLevel(logLevel) diff --git a/cmd/sippy/serve.go b/cmd/sippy/serve.go index a1d6e665f5..8d08f8af78 100644 --- a/cmd/sippy/serve.go +++ b/cmd/sippy/serve.go @@ -31,6 +31,7 @@ import ( "github.com/openshift/sippy/pkg/flags/configflags" "github.com/openshift/sippy/pkg/sippyserver" "github.com/openshift/sippy/pkg/sippyserver/metrics" + "github.com/openshift/sippy/pkg/sippyserver/workqueue" "github.com/openshift/sippy/pkg/testidentification" "github.com/openshift/sippy/pkg/util" ) @@ -202,6 +203,22 @@ func NewServeCommand() *cobra.Command { jiraClient, ) + // Set up River insert-only client for async batch submission + if dbc != nil && f.APIFlags.EnableWriteEndpoints { + riverSetup, err := workqueue.Setup(cmd.Context(), workqueue.SetupConfig{ + DatabaseDSN: f.DBFlags.DSN, + }) + if err != nil { + log.WithError(err).Warn("failed to set up River work queue, async re-evaluation will not be available") + } else { + server.SetWorkqueue( + workqueue.NewSubmitter(dbc.DB, riverSetup.Client), + workqueue.NewStatusQuerier(dbc.DB), + ) + log.Info("River work queue configured for insert-only mode") + } + } + if f.APIFlags.MetricsAddr != "" { // Do an immediate metrics update err = metrics.RefreshMetricsDB( diff --git a/docs/features/job-analysis-symptoms.md b/docs/features/job-analysis-symptoms.md index 55ef878338..1268ed00a0 100644 --- a/docs/features/job-analysis-symptoms.md +++ b/docs/features/job-analysis-symptoms.md @@ -121,7 +121,28 @@ The `POST /api/jobs/runs/reevaluate` endpoint re-runs symptom detection for spec Unlike the cloud function (which processes files as they arrive), the re-evaluator scans all artifacts at once for completed job runs. -Flow: +**Async batch mode** (default, `dry_run: false` or omitted): + +The endpoint creates an async batch and returns `202 Accepted` with a batch ID. Individual job +runs are enqueued as River jobs and processed by workers in the sippy-daemon. The UI polls +`GET /api/jobs/runs/reevaluate/{batch_id}` for progress. + +- Up to 10,000 job runs per batch. +- Deduplication: the same job run is not re-evaluated within 90 minutes (BigQuery streaming + buffer window), unless the symptom definitions have changed. A hash of the active symptom + set is included in the deduplication key, so submitting a batch after editing symptoms + always triggers fresh evaluation. +- Retry: transient failures (e.g. GCS timeouts) are retried up to 3 times with exponential + backoff. +- Multiple batches can reference the same underlying River job. The batch tracks which items + belong to it, while River handles execution and uniqueness. + +**Dry-run mode** (`dry_run: true`): + +Processes synchronously with a `200 OK` response. Limited to 50 job runs. Useful for testing +symptom changes before committing to a full re-evaluation. + +**Per-item flow** (same for both modes): 1. Load all active symptom definitions from PostgreSQL (excluding unimplemented matcher types). 2. For each job run, run one `JobArtifactQuery` per symptom against GCS artifacts. @@ -133,6 +154,12 @@ The delete-then-insert strategy makes re-evaluation idempotent: if a symptom is removed, re-evaluating produces the correct result. Manually-applied labels (those with empty `symptom_id`) are preserved through re-evaluation. +**Known limitations:** + +- No batch cancellation. Once enqueued, River jobs run to completion or exhaust retries. +- Batch completion is detected lazily when the status endpoint is polled. Unpollable batches + remain in "running" status until cleaned up. + ## Key Code Locations ### Sippy (`openshift/sippy`) @@ -146,6 +173,8 @@ removed, re-evaluating produces the correct result. Manually-applied labels (tho | `pkg/api/jobrunscan/` | API handlers for symptom/label CRUD and re-evaluation, with validation logic. | | `pkg/api/jobrunscan/reevaluate.go` | Re-evaluation service: symptom scanning, BQ/GCS/PostgreSQL write logic. | | `pkg/sippyserver/job_run_scan.go` | HTTP route handlers delegating to the jobrunscan API package. | +| `pkg/sippyserver/workqueue/` | Generic async batch processing via River: batch models, submitter, status querier, River process adapter. | +| `pkg/api/jobrunscan/reevaluate_worker.go` | River worker definition and job args for async re-evaluation. | | `pkg/sippyclient/jobrunscan/` | Go client library for symptom/label APIs (used by cloud function). | | `pkg/componentreadiness/jobrunannotator/jobrunannotator.go` | `JobRunAnnotator` - the `annotate-job-runs` tool which can add labels but doesn't (yet) know about symptoms. | | `pkg/componentreadiness/jobrunannotator/prow_bucket.go` | `JobRunBucketLabel`, `WriteHTMLSummaryToBucket` - writes label files and HTML summaries to GCS. Shared with cloud function. | @@ -171,6 +200,8 @@ All endpoints are under `/api/jobs/` and support standard CRUD: - `GET/POST /api/jobs/symptoms` - list / create symptoms - `GET/PUT/DELETE /api/jobs/symptoms/{id}` - read / update / delete - `POST /api/jobs/runs/reevaluate` - re-evaluate symptoms for specified job runs + (returns `202 Accepted` for async, `200 OK` for dry-run) +- `GET /api/jobs/runs/reevaluate/{batch_id}` - poll async re-evaluation batch status See `pkg/api/jobrunscan/` for validation rules and `pkg/api/README.md` for broader API documentation. @@ -183,7 +214,10 @@ documentation. | PostgreSQL `job_run_labels` | Label definitions | Authoritative source for label metadata. | | PostgreSQL `prow_job_runs.labels` | Applied label IDs per job run | Sippy queries and UI display. | | PostgreSQL `release_job_runs.labels` | Applied label IDs per payload job run | Sippy queries and UI display. | -| PostgreSQL `triage_symptoms` | Symptom↔triage associations | Triage UI symptom summaries. | +| PostgreSQL `triage_symptoms` | Symptom-triage associations | Triage UI symptom summaries. | +| PostgreSQL `workqueue_batches` | Async batch metadata | Tracks batch lifecycle and progress. | +| PostgreSQL `workqueue_batch_items` | Batch-to-River-job associations | Links batches to individual work items. | +| PostgreSQL `river_job` (managed by River) | Job queue rows | Execution state, retry tracking, deduplication. | | BigQuery `ci_analysis_us.job_labels` | Applied labels with provenance | Warehouse for analytics; source of truth during fetchdata. | | GCS `artifacts/job_labels/*.json` | Per-match label files | Provenance and Spyglass display. | | GCS `artifacts/job_labels/label-summary.html` | HTML summary | Rendered by Spyglass html lens. | @@ -197,3 +231,5 @@ Active/planned work includes: - richer CEL-based label composition. - **Full management UI** ([TRT-2479](https://redhat.atlassian.net/browse/TRT-2479)) - dedicated UI for label/symptom lifecycle. +- **Async re-evaluation** ([TRT-2867](https://redhat.atlassian.net/browse/TRT-2867)) + - batch processing via River job queue (implemented, pending UI polling integration). diff --git a/docs/plans/trt-2867-async-reevaluation-plan.md b/docs/plans/trt-2867-async-reevaluation-plan.md new file mode 100644 index 0000000000..5b38518fef --- /dev/null +++ b/docs/plans/trt-2867-async-reevaluation-plan.md @@ -0,0 +1,459 @@ +# TRT-2867: Async Symptom Re-evaluation via River Job Queue + +## Overview + +Move the synchronous symptom re-evaluation API (`POST /api/jobs/runs/reevaluate`) to an +asynchronous batch model. The API handler creates a batch request and enqueues individual +re-evaluation work items via [River](https://github.com/riverqueue/river), a PostgreSQL-backed +job queue. The sippy-daemon processes work items using the existing `ReEvaluator` logic. The UI +polls a status endpoint for progress until all items complete. + +**Jira:** [TRT-2867](https://redhat.atlassian.net/browse/TRT-2867) +**Predecessor:** [TRT-2695](https://redhat.atlassian.net/browse/TRT-2695) — synchronous +re-evaluation API (already implemented) + +## Motivation + +The current `POST /api/jobs/runs/reevaluate` endpoint processes all requested job runs +synchronously in the HTTP request. For large batches this causes: + +- HTTP timeouts when evaluating many job runs against many symptoms +- No progress visibility — the UI blocks until the entire batch completes +- No deduplication — concurrent requests can re-evaluate the same job run, causing BigQuery + streaming buffer conflicts (rows are not deletable within 90 minutes of insertion) +- No retry of individual failures — one GCS timeout fails the entire request + +## Design Principles + +1. **Reuse, don't rewrite** — the existing `ReEvaluator.reEvaluateOne()` becomes the work unit + inside each River job worker, unchanged. +2. **Appropriate abstraction** — the new package (`pkg/sippyserver/workqueue`) is generic enough + to support future async workloads (report generation, bulk annotation, etc.) without being + tied to symptom re-evaluation specifics. +3. **River owns execution and dedup; application tables own batch semantics** — River handles job + scheduling, concurrency, retries, and uniqueness. Custom tables handle batch-to-item + associations and user-facing status. + +## Prerequisites: Orientation + +Before writing code, read and understand these files in addition to those listed in the +TRT-2695 plan: + +| File | What to learn | +|------|---------------| +| `pkg/sippyserver/daemon_server.go` | `DaemonProcess` interface and `DaemonServer.Serve()` goroutine lifecycle. Each process runs independently with its own context. | +| `pkg/sippyserver/pr_commenting_processor.go` | `WorkProcessor` — existing daemon process example. Ticker-based polling, bounded worker goroutines, channel-based dispatch. | +| `cmd/sippy-daemon/main.go` | Daemon process registration and startup. New processes added via `processes = append(processes, ...)`. | +| `pkg/api/jobrunscan/reevaluate.go` | `ReEvaluator` struct, `ReEvaluateJobRuns()`, `reEvaluateOne()`. The synchronous implementation to be reused as the per-item worker. | +| `pkg/sippyserver/job_run_scan.go` | Current HTTP handler `jsonReEvaluateJobRunSymptoms` — the handler to be replaced with async dispatch. | + +## Dependency: River and pgx/v5 + +River requires `jackc/pgx/v5`. Sippy currently uses `jackc/pgx/v4` (v4.18.2). + +**Approach:** Add `pgx/v5` alongside the existing `pgx/v4` dependency. Create a separate +`*pgxpool.Pool` (v5) for River's use. The rest of Sippy continues using v4 + gorm until a +broader migration is undertaken. Both driver versions can coexist in a Go module; they have +different import paths (`github.com/jackc/pgx/v4` vs `github.com/jackc/pgx/v5`). + +## Step 1: Add the `pkg/sippyserver/workqueue` Package + +Create a generic work queue abstraction in `pkg/sippyserver/workqueue/`. This package wraps +River and provides the batch submission and status query patterns used by the re-evaluation +API, but is designed for reuse by future async workloads. + +### 1.1: Database models (`pkg/sippyserver/workqueue/models.go`) + +```go +type BatchStatus string + +const ( + BatchStatusPending BatchStatus = "pending" + BatchStatusRunning BatchStatus = "running" + BatchStatusComplete BatchStatus = "complete" + BatchStatusFailed BatchStatus = "failed" +) + +// Batch represents a user-initiated batch of work items. +// A batch groups related work items for status tracking and progress reporting. +type Batch struct { + ID uuid.UUID `gorm:"type:uuid;primaryKey" json:"id"` + Kind string `gorm:"not null;index" json:"kind"` + RequestedCount int `gorm:"not null" json:"requested_count"` + EnqueuedCount int `gorm:"not null" json:"enqueued_count"` + DedupedCount int `gorm:"not null" json:"deduped_count"` + Status BatchStatus `gorm:"not null;default:'pending'" json:"status"` + CreatedAt time.Time `gorm:"autoCreateTime" json:"created_at"` + CompletedAt *time.Time ` json:"completed_at,omitempty"` +} + +// BatchItem associates a batch with a River job for many-to-many status tracking. +// Multiple batches can reference the same River job (when deduplication occurs). +type BatchItem struct { + ID uint64 `gorm:"primaryKey;autoIncrement" json:"id"` + BatchID uuid.UUID `gorm:"type:uuid;not null;index:idx_batch_items" json:"batch_id"` + RiverJobID int64 `gorm:"not null;index:idx_batch_items" json:"river_job_id"` + ItemKey string `gorm:"not null" json:"item_key"` +} +``` + +- `Batch.Kind` identifies the type of work (e.g. `"reevaluate_symptoms"`) so the table can + serve multiple async workflows. +- `BatchItem.ItemKey` stores the human-readable work item identifier (e.g. the + `prow_job_build_id`) for display in status responses without needing to decode River job args. +- Add a unique constraint on `(batch_id, item_key)` to prevent a batch from listing the same + item twice. + +### 1.2: Batch submission (`pkg/sippyserver/workqueue/submitter.go`) + +Define a `Submitter` struct that accepts a gorm DB and a River client. Provide a `Submit` +method that: + +1. Creates a `Batch` row in the database. +2. Calls River's `InsertManyTx` to insert all work items as individual River jobs with + deduplication (see Step 2 for `UniqueOpts`). +3. For each insert result, records a `BatchItem` row linking the batch to the River job ID. + River returns the existing job row for duplicates, so `BatchItem` rows are created + regardless of whether the job was new or deduped. +4. Counts how many jobs were newly enqueued vs. deduplicated (River's insert result indicates + this via `UniqueSkippedAsDuplicate`) and updates the `Batch` row with those counts and + sets status to `running`. +5. Returns a `SubmitResult` with the batch ID and new/dedup counts. + +**Transaction boundaries:** The batch/batch-item writes use gorm (pgx/v4) while River's +`InsertManyTx` requires a pgx/v5 transaction. Because these are different driver versions, +they cannot share a single database transaction. Use two separate transactions: one gorm +transaction for the batch and batch-item rows, and one pgx/v5 transaction for River job +insertion. If the gorm transaction succeeds but the River insert fails, the batch row will +exist with no corresponding River jobs. If the River insert succeeds but the gorm transaction +fails, the River jobs will run without a tracking batch. In either failure case, the user can +submit a new batch with the same entries (deduplication prevents duplicate work) and poll that +batch for progress. + +### 1.3: Batch status query (`pkg/sippyserver/workqueue/status.go`) + +Define a `StatusQuerier` struct with a `Query` method that: + +1. Loads the `Batch` row by ID. +2. Joins `workqueue_batch_items` with `river_job` to get the current River state for each item. +3. Aggregates counts by state category: completed, failed (discarded/cancelled), running, + and pending (available/scheduled/retryable). +4. Returns a `BatchStatusResponse` containing the batch ID, overall status, aggregate counts, + and per-item status list. + +**Lazy completion detection:** When all items have reached a terminal state +(`completed + failed >= total`), the query marks the batch as complete (or failed if all items +failed) and sets `completed_at`. This is idempotent — no asynchronous bookkeeping is required. +The batch status is simply computed fresh on each poll and the batch row is updated as a +side effect when the terminal condition is met. + +### 1.4: Database migration + +Add a migration using `golang-migrate/v4` (following existing patterns in `pkg/db/migrations/`): + +```sql +-- Up +CREATE TABLE workqueue_batches ( + id UUID PRIMARY KEY, + kind TEXT NOT NULL, + requested_count INT NOT NULL, + enqueued_count INT NOT NULL DEFAULT 0, + deduped_count INT NOT NULL DEFAULT 0, + status TEXT NOT NULL DEFAULT 'pending', + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + completed_at TIMESTAMPTZ +); + +CREATE INDEX idx_workqueue_batches_kind_status ON workqueue_batches (kind, status); + +CREATE TABLE workqueue_batch_items ( + id BIGSERIAL PRIMARY KEY, + batch_id UUID NOT NULL REFERENCES workqueue_batches(id) ON DELETE CASCADE, + river_job_id BIGINT NOT NULL, + item_key TEXT NOT NULL, + UNIQUE(batch_id, item_key) +); + +CREATE INDEX idx_workqueue_batch_items_batch ON workqueue_batch_items (batch_id, river_job_id); + +-- River's own migrations are handled by river.Migrator at startup (see Step 3). + +-- Down +DROP TABLE IF EXISTS workqueue_batch_items; +DROP TABLE IF EXISTS workqueue_batches; +``` + +## Step 2: Define the Re-evaluation River Job + +Create `pkg/api/jobrunscan/reevaluate_worker.go` alongside the existing re-evaluation code. + +### 2.1: Job args and constants + +```go +const ( + ReevaluateJobKind = "reevaluate_job_run" + ReevaluateQueue = "reevaluate" + ReevaluateDedupPeriod = 90 * time.Minute + ReevaluateMaxAttempts = 3 + MaxJobRunsPerBatch = 10000 +) + +type ReevaluateJobRunArgs struct { + ProwJobBuildID string `json:"prow_job_build_id" river:"unique"` + SymptomHash string `json:"symptom_hash" river:"unique"` +} +``` + +The `InsertOpts` method on `ReevaluateJobRunArgs` should return: + +- `Queue`: `ReevaluateQueue` +- `MaxAttempts`: `ReevaluateMaxAttempts` +- `UniqueOpts`: `ByArgs: true`, `ByPeriod: ReevaluateDedupPeriod` + +Key design decisions: + +- Both `ProwJobBuildID` and `SymptomHash` are tagged `river:"unique"`, so deduplication is + scoped to the combination of job run identity and symptom state. The same job run will be + re-evaluated if symptoms have changed since the last evaluation. +- `ByPeriod: 90 * time.Minute` prevents re-evaluating the same job run within the BigQuery + streaming buffer window, avoiding the "rows not deletable within 90m" constraint. +- No `BatchID` in the args — batch association is tracked via `workqueue_batch_items`, not + River job args. This keeps the unique hash clean and avoids the subtlety where different + batch IDs in the args would defeat deduplication. + +### 2.2: Job worker + +Define a `ReevaluateWorker` struct that embeds `river.WorkerDefaults[ReevaluateJobRunArgs]` +and holds a reference to a `ReEvaluator`. The `Work` method calls the existing +`reEvaluateOne()` for the given `ProwJobBuildID` against the evaluator's cached symptoms, +returning an error if the evaluation fails (which triggers River's retry logic). + +### 2.3: Symptom caching and hash-based uniqueness + +The `ReEvaluator` should maintain a concurrency-safe cache of active symptoms to avoid +reloading them from the database for every individual work item. + +Add a `RefreshSymptomCache()` method to `ReEvaluator` that loads all active symptoms via the +existing `loadActiveSymptoms()` and stores them in a field guarded by a `sync.RWMutex`. The +method also computes a hash (e.g. SHA-256) of the sorted symptom IDs/versions and stores it +alongside the cached symptoms. The `reEvaluateOne()` method reads from the cache under an +`RLock`. + +**Refresh trigger:** The symptom cache is refreshed when a new batch is created. The API +handler (or `Submitter`) calls `RefreshSymptomCache()` during batch submission. The resulting +symptom hash is included in the `ReevaluateJobRunArgs` (see below) so that it participates +in River's deduplication. This means that if a user modifies symptoms and submits a new batch +for the same job runs, the changed hash defeats deduplication and the jobs run again with the +updated symptoms. + +We are not interested in tracking symptom changes that occur while a batch is processing. +The expectation is that the user submits a batch after making the symptom changes they care +about. If they make further changes, they submit another batch, and the new symptom hash +ensures those jobs are not deduplicated against the earlier run. + +### 2.4: Retry policy + +River's default retry policy uses exponential backoff with jitter. Configure +`MaxAttempts: 3` — a transient GCS timeout retries twice, then marks the job `discarded` +(permanent failure). This is surfaced in the batch status response as a failed item. + +## Step 3: Integrate River into sippy-daemon + +### 3.1: River client setup + +In `cmd/sippy-daemon/main.go`: + +1. Create a `pgx/v5` connection pool using the existing database DSN. This pool is used + exclusively by River and coexists with the existing pgx/v4 pool. +2. Run River's built-in migrations on startup via `rivermigrate.Migrator`. +3. Construct a `ReEvaluator` with the same dependencies used by the API server (BigQuery + client, GCS client, GCS bucket, gorm DB, cache, job artifacts manager). If + `cmd/sippy-daemon/main.go` does not currently initialize these clients, add the + initialization following the patterns in `cmd/sippy/main.go`. +4. Register the `ReevaluateWorker` with River's worker registry. +5. Create and start the River client with the `reevaluate` queue configured (suggested + starting concurrency: 8 workers, tunable based on load testing). + +### 3.2: Daemon lifecycle integration + +River manages its own goroutines internally. Wrap the River client in a thin adapter struct +implementing the `DaemonProcess` interface so that it participates in the existing +`DaemonServer.Serve()` lifecycle. The adapter's `Run` method starts the River client, blocks +until the context is cancelled, then calls `Stop` with a graceful shutdown timeout. + +Register this adapter as a daemon process in `cmd/sippy-daemon/main.go` alongside the +existing `WorkProcessor`. + +### 3.3: Job artifact query worker pool + +The existing `ReEvaluator.reEvaluateOne()` uses `pkg/api/jobartifacts.Manager` for concurrent +GCS artifact scanning. Ensure the daemon's `ReEvaluator` receives a properly initialized +`jobartifacts.Manager` instance — this requires a GCS client, bucket name, and configuration +matching the API server's setup. + +## Step 4: Modify the API + +### 4.1: Change the existing endpoint to async + +Modify the handler for `POST /api/jobs/runs/reevaluate` in `pkg/sippyserver/job_run_scan.go`: + +- Validate the request body (same fields as today: `prow_job_build_ids` and `dry_run`). +- Enforce a maximum of 10,000 job runs per request. +- If `dry_run: true`, process synchronously using the existing `ReEvaluator` and return + results immediately with `200 OK` (preserving current behavior). +- If `dry_run: false` (or omitted), use the `Submitter` to create a batch and enqueue River + jobs, then return `202 Accepted` with the batch ID and dedup counts. +- Trigger a symptom cache refresh on the `ReEvaluator` during batch submission (see Step 2.3). + +### 4.2: Add the batch status endpoint + +Register `GET /api/jobs/runs/reevaluate/{batch_id}` with the `LocalDBCapability` requirement. +The handler parses the batch UUID from the path, calls `StatusQuerier.Query()`, and returns +the result. Return `404` if the batch ID is not found. + +### 4.3: API contract + +**Submit batch (POST):** + +``` +POST /api/jobs/runs/reevaluate +Content-Type: application/json + +{ + "prow_job_build_ids": ["1234567890", "1234567891", ...], + "dry_run": false +} +``` + +Response (`202 Accepted`): + +```json +{ + "batch_id": "a1b2c3d4-...", + "requested": 50, + "enqueued": 42, + "deduped": 8, + "links": { + "status": "/api/jobs/runs/reevaluate/a1b2c3d4-..." + } +} +``` + +**Poll status (GET):** + +``` +GET /api/jobs/runs/reevaluate/{batch_id} +``` + +Response (`200 OK`): + +```json +{ + "batch_id": "a1b2c3d4-...", + "status": "running", + "total": 50, + "completed": 35, + "failed": 1, + "running": 6, + "pending": 8, + "items": [ + {"item_key": "1234567890", "state": "completed"}, + {"item_key": "1234567891", "state": "running"}, + {"item_key": "1234567892", "state": "available"}, + {"item_key": "1234567893", "state": "discarded"} + ] +} +``` + +When `dry_run: true`, the existing synchronous behavior is preserved — results are returned +immediately with a `200 OK` response. This keeps the dry-run path simple and useful for testing. + +## Step 5: Wire Up the API Server + +The API server (sippy) needs a `Submitter` and `StatusQuerier` from +`pkg/sippyserver/workqueue` but does **not** run River workers — those run in sippy-daemon only. + +Create a River client in insert-only mode in `cmd/sippy/main.go`: configure the client without +registering any workers or queues and without calling `Start()`. The client can still insert +jobs via `InsertMany()`. + +Add `workqueueSubmitter` and `workqueueStatusQuerier` fields to the `Server` struct. Initialize +them during server construction, gated on `LocalDBCapability` (required for both endpoints). +The existing `POST /api/jobs/runs/reevaluate` already requires both `LocalDBCapability` and +`WriteEndpointsCapability`; the new `GET /api/jobs/runs/reevaluate/{batch_id}` status endpoint +requires only `LocalDBCapability` (it is read-only). + +## Step 6: Batch Lifecycle and Cleanup + +### 6.1: Completed batch retention + +Add a periodic cleanup job (either a River periodic job or a simple cron-like `DaemonProcess`) +that deletes `workqueue_batches` rows where `completed_at` is older than 7 days. The +`ON DELETE CASCADE` foreign key handles the `workqueue_batch_items` rows automatically. + +### 6.2: River job retention + +Configure River's `CompletedJobRetention` and `DiscardedJobRetention` to 8 days, matching +the batch retention period. This ensures `river_job` rows don't grow unbounded while keeping +enough history for status queries on recent batches. + +## Step 7: Testing + +Avoid mocking any substantial part of these packages (River client, ReEvaluator, storage +clients). Instead, follow the project's functional test pattern (see +`pkg/api/jobrunscan/reevaluate_functional_test.go`): tests that require external dependencies +(PostgreSQL, GCS, BigQuery) are gated behind environment variables and skipped when those +variables are not set. A human runs them by providing the necessary credentials. + +- **Unit tests** for pure logic functions only: dedup counting, symptom hash computation, + batch status aggregation from pre-populated row structs. No mocking of River or database + clients. +- **Functional tests** for `pkg/sippyserver/workqueue/` — test `Submitter.Submit()` and + `StatusQuerier.Query()` against a real PostgreSQL instance with River migrations applied. + Verify dedup counting, batch item creation, completion detection, and status transitions. + Skip unless `SIPPY_FUNCTIONAL_TEST_DSN` (or similar) is set. +- **Functional tests** for `ReevaluateWorker.Work()` using River's `rivertest` package + against a real PostgreSQL instance. Verify success and error/retry paths. +- **End-to-end functional test** of the full flow: POST to the API, verify 202, poll status, + verify items transition through `available` to `running` to `completed`. Requires + PostgreSQL with River migrations, and optionally GCS/BigQuery credentials for full coverage. + +## Step 8: Implementation Order + +1. Add `pgx/v5` and River dependencies to `go.mod`. +2. Create `pkg/sippyserver/workqueue/` — models, submitter, status querier. +3. Create `ReevaluateJobRunArgs` and `ReevaluateWorker` in `pkg/api/jobrunscan/`. +4. Add symptom cache with `sync.RWMutex` and symptom hash computation to `ReEvaluator`. +5. Write database migration for `workqueue_batches` and `workqueue_batch_items`. +6. Integrate River into `cmd/sippy-daemon/main.go` — pgx/v5 pool, migrations, worker + registration, daemon process adapter, GCS/BQ client initialization. +7. Modify API handler — async dispatch for non-dry-run, new status endpoint. +8. Wire up API server — insert-only River client, submitter, status querier on Server struct. +9. Tests (unit, integration, end-to-end). +10. Update `docs/features/job-analysis-symptoms.md` — document the async API behavior, + new status endpoint, deduplication semantics, and polling pattern. + +## Known Limitations + +1. **No batch cancellation** — There is no API to cancel an in-flight batch. Once River jobs + are enqueued, they will run to completion (or exhaust retries). A future enhancement could + add `DELETE /api/jobs/runs/reevaluate/{batch_id}` to cancel pending River jobs associated + with a batch. + +2. **Lazy batch completion** — The batch status is only updated to "complete" or "failed" when + the status endpoint is polled and detects that all items have reached a terminal state. If + no one polls after the last item completes, the batch row stays in "running" status + indefinitely. This is cosmetic; the cleanup job (Step 6.1) will eventually delete it + regardless of status. + +## Open Questions + +1. **pgx/v5 coexistence** — Are there any known issues running pgx/v4 and v5 pools against + the same PostgreSQL instance? Initial research suggests this is safe (separate import paths, + separate connection pools), but worth a quick verification. +2. **MaxWorkers tuning** — 8 concurrent re-evaluation workers is a starting point. The right + number depends on GCS rate limits and the job artifact query concurrency within each worker. + May need adjustment based on load testing. +3. **UI polling interval** — Suggest 2-3 second polling interval from the UI. Consider + implementing exponential backoff (e.g. 1s → 2s → 5s → 10s) or Server-Sent Events (SSE) + in a future iteration to reduce polling overhead. diff --git a/go.mod b/go.mod index 953c5949b0..6bf2057ddb 100644 --- a/go.mod +++ b/go.mod @@ -37,7 +37,7 @@ require ( github.com/tcnksm/go-gitconfig v0.1.2 github.com/testcontainers/testcontainers-go v0.43.0 github.com/testcontainers/testcontainers-go/modules/postgres v0.43.0 - github.com/tidwall/gjson v1.14.4 + github.com/tidwall/gjson v1.19.0 github.com/trivago/tgo v1.0.7 github.com/yuin/goldmark v1.7.16 golang.org/x/exp v0.0.0-20260410095643-746e56fc9e2f @@ -110,6 +110,7 @@ require ( github.com/jackc/pgpassfile v1.0.0 // indirect github.com/jackc/pgproto3/v2 v2.3.3 // indirect github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect + github.com/jackc/pgx/v5 v5.10.0 // indirect github.com/jinzhu/inflection v1.0.0 // indirect github.com/jinzhu/now v1.1.5 // indirect github.com/jtolds/gls v4.20.0+incompatible // indirect @@ -143,7 +144,7 @@ require ( github.com/skelterjohn/go.matrix v0.0.0-20130517144113-daa59528eefd // indirect github.com/spf13/cast v1.7.1 // indirect github.com/spiffe/go-spiffe/v2 v2.6.0 // indirect - github.com/tidwall/match v1.1.1 // indirect + github.com/tidwall/match v1.2.0 // indirect github.com/tidwall/pretty v1.2.1 // indirect github.com/tklauser/go-sysconf v0.3.16 // indirect github.com/tklauser/numcpus v0.11.0 // indirect @@ -161,7 +162,7 @@ require ( go.opentelemetry.io/otel/sdk/metric v1.43.0 // indirect go.opentelemetry.io/otel/trace v1.43.0 // indirect golang.org/x/crypto v0.54.0 // indirect - golang.org/x/mod v0.37.0 // indirect + golang.org/x/mod v0.38.0 // indirect golang.org/x/net v0.56.0 // indirect golang.org/x/sys v0.47.0 // indirect golang.org/x/telemetry v0.0.0-20260625142307-59b4966ccb57 // indirect diff --git a/go.sum b/go.sum index e019d6ebfb..98124f161c 100644 --- a/go.sum +++ b/go.sum @@ -246,8 +246,8 @@ github.com/jackc/pgx/v4 v4.12.1-0.20210724153913-640aa07df17c/go.mod h1:1QD0+tgS github.com/jackc/pgx/v4 v4.13.0/go.mod h1:9P4X524sErlaxj0XSGZk7s+LD0eOyu1ZDUrrpznYDF0= github.com/jackc/pgx/v4 v4.18.2 h1:xVpYkNR5pk5bMCZGfClbO962UIqVABcAGt7ha1s/FeU= github.com/jackc/pgx/v4 v4.18.2/go.mod h1:Ey4Oru5tH5sB6tV7hDmfWFahwF15Eb7DNXlRKx2CkVw= -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 v0.0.0-20190413234325-e4ced69a3a2b/go.mod h1:m4B5Dj62Y0fbyuIc15OsIqK0+JU8nkqQjsgx7dvjSWk= github.com/jackc/puddle v0.0.0-20190608224051-11cab39313c9/go.mod h1:m4B5Dj62Y0fbyuIc15OsIqK0+JU8nkqQjsgx7dvjSWk= github.com/jackc/puddle v1.1.3/go.mod h1:m4B5Dj62Y0fbyuIc15OsIqK0+JU8nkqQjsgx7dvjSWk= @@ -416,11 +416,10 @@ github.com/testcontainers/testcontainers-go v0.43.0 h1:oEQx5MW2DGd9z3AeEQfB2lPM0 github.com/testcontainers/testcontainers-go v0.43.0/go.mod h1:+VxkT2NQnKOZPKi6praMuMKYHYyOGXr0XSBSlSMCzFo= github.com/testcontainers/testcontainers-go/modules/postgres v0.43.0 h1:ShNOFYAF4lKHvdIG258hi69bSxC88uXnxJkJvNs/IVs= github.com/testcontainers/testcontainers-go/modules/postgres v0.43.0/go.mod h1:vdq5/RqmGfWeefzyfcVI/pID1rzmc1TDvqXa15bPJks= -github.com/tidwall/gjson v1.14.4 h1:uo0p8EbA09J7RQaflQ1aBRffTR7xedD2bcIVSYxLnkM= -github.com/tidwall/gjson v1.14.4/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk= -github.com/tidwall/match v1.1.1 h1:+Ho715JplO36QYgwN9PGYNhgZvoUSc9X2c80KVTi+GA= -github.com/tidwall/match v1.1.1/go.mod h1:eRSPERbgtNPcGhD8UCthc6PmLEQXEWd3PRB5JTxsfmM= -github.com/tidwall/pretty v1.2.0/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU= +github.com/tidwall/gjson v1.19.0 h1:xwxm7n691Uf3u5OFjzngavjGTh55KX5q/9w9xHW88JU= +github.com/tidwall/gjson v1.19.0/go.mod h1:V37/opeE/JbLUOfH0QTXiNez2l0RUjYUhpT4szFQAfc= +github.com/tidwall/match v1.2.0 h1:0pt8FlkOwjN2fPt4bIl4BoNxb98gGHN2ObFEDkrfZnM= +github.com/tidwall/match v1.2.0/go.mod h1:eRSPERbgtNPcGhD8UCthc6PmLEQXEWd3PRB5JTxsfmM= github.com/tidwall/pretty v1.2.1 h1:qjsOFOWWQl+N3RsoF5/ssm1pHmJJwhjlSbZ51I6wMl4= github.com/tidwall/pretty v1.2.1/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU= github.com/tklauser/go-sysconf v0.3.16 h1:frioLaCQSsF5Cy1jgRBrzr6t502KIIwQ0MArYICU0nA= @@ -492,8 +491,8 @@ golang.org/x/lint v0.0.0-20190930215403-16217165b5de/go.mod h1:6SW0HCj/g11FgYtHl golang.org/x/mod v0.0.0-20190513183733-4bf6d317e70e/go.mod h1:mXi4GBBbnImb6dmsKGUJ2LatrhH/nqhxcFungHvyanc= golang.org/x/mod v0.1.1-0.20191105210325-c90efee705ee/go.mod h1:QqPTAvyqsEbceGzBzNggFXnrqF1CaUcvgkdR5Ot7KZg= golang.org/x/mod v0.3.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA= -golang.org/x/mod v0.37.0 h1:vF1DjpVEshcIqoEaauuHebaLk1O1forxjxBaVn884JQ= -golang.org/x/mod v0.37.0/go.mod h1:m8S8VeM9r4dzDwjrKO0a1sZP3YjeMamRRlD+fmR2Q/0= +golang.org/x/mod v0.38.0 h1:MECBjubtXD7yj4HrhIUcywNaGeNVUdfVnxmPajOk4yk= +golang.org/x/mod v0.38.0/go.mod h1:V6Xz0pq8TQ3dGqVQ1FVHuelZpAL0uNhSkk9ogYP3c40= golang.org/x/net v0.0.0-20180906233101-161cd47e91fd/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= golang.org/x/net v0.0.0-20190311183353-d8887717615a/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg= golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg= diff --git a/pkg/api/jobrunscan/reevaluate.go b/pkg/api/jobrunscan/reevaluate.go index d33910b1a3..108d77096c 100644 --- a/pkg/api/jobrunscan/reevaluate.go +++ b/pkg/api/jobrunscan/reevaluate.go @@ -7,6 +7,7 @@ import ( "regexp" "strconv" "strings" + "sync" "time" "github.com/lib/pq" @@ -95,6 +96,10 @@ type ReEvaluator struct { cache cache.Cache artifactMgr *jobartifacts.Manager dryRun bool + + symptomMu sync.RWMutex + cachedHash string + cachedSymps []jobrunscan.Symptom } // NewReEvaluator creates a ReEvaluator with the given clients. @@ -110,6 +115,33 @@ func NewReEvaluator(bqClient *bqclient.Client, gcsClient *storage.Client, gcsBuc } } +// RefreshSymptomCache loads active symptoms from the database and stores them +// along with their hash. Returns the hash for use in job deduplication. +func (r *ReEvaluator) RefreshSymptomCache() (string, error) { + symptoms, err := r.loadActiveSymptoms() + if err != nil { + return "", fmt.Errorf("loading symptoms for cache: %w", err) + } + hash := computeSymptomHash(symptoms) + + r.symptomMu.Lock() + r.cachedSymps = symptoms + r.cachedHash = hash + r.symptomMu.Unlock() + + log.WithFields(log.Fields{"count": len(symptoms), "hash": hash}). + Info("symptom cache refreshed") + return hash, nil +} + +// CachedSymptoms returns the cached symptom list, or nil if the cache has not +// been populated via RefreshSymptomCache. +func (r *ReEvaluator) CachedSymptoms() []jobrunscan.Symptom { + r.symptomMu.RLock() + defer r.symptomMu.RUnlock() + return r.cachedSymps +} + // symptomMatch records that a symptom matched a file/text in a job run. type symptomMatch struct { symptom jobrunscan.Symptom diff --git a/pkg/api/jobrunscan/reevaluate_worker.go b/pkg/api/jobrunscan/reevaluate_worker.go new file mode 100644 index 0000000000..a4324a021d --- /dev/null +++ b/pkg/api/jobrunscan/reevaluate_worker.go @@ -0,0 +1,86 @@ +package jobrunscan + +import ( + "context" + "fmt" + "time" + + "github.com/riverqueue/river" + + "github.com/openshift/sippy/pkg/db/models/jobrunscan" +) + +const ( + ReevaluateJobKind = "reevaluate_job_run" + ReevaluateQueue = "reevaluate" + ReevaluateDedupPeriod = 90 * time.Minute + ReevaluateMaxAttempts = 3 + MaxJobRunsPerBatch = 10000 +) + +// ReevaluateJobRunArgs are the arguments for a single re-evaluation River job. +// Both ProwJobBuildID and SymptomHash participate in uniqueness, so the same +// job run will be re-evaluated if the symptom set changes. +type ReevaluateJobRunArgs struct { + ProwJobBuildID string `json:"prow_job_build_id" river:"unique"` + SymptomHash string `json:"symptom_hash" river:"unique"` +} + +func (ReevaluateJobRunArgs) Kind() string { return ReevaluateJobKind } + +func (ReevaluateJobRunArgs) InsertOpts() river.InsertOpts { + return river.InsertOpts{ + Queue: ReevaluateQueue, + MaxAttempts: ReevaluateMaxAttempts, + UniqueOpts: river.UniqueOpts{ + ByArgs: true, + ByPeriod: ReevaluateDedupPeriod, + }, + } +} + +// ReevaluateWorker processes a single re-evaluation job using the cached +// symptoms from the ReEvaluator. +type ReevaluateWorker struct { + river.WorkerDefaults[ReevaluateJobRunArgs] + evaluator *ReEvaluator +} + +// NewReevaluateWorker creates a worker that delegates to the given ReEvaluator. +func NewReevaluateWorker(evaluator *ReEvaluator) *ReevaluateWorker { + return &ReevaluateWorker{evaluator: evaluator} +} + +func (w *ReevaluateWorker) Work(ctx context.Context, job *river.Job[ReevaluateJobRunArgs]) error { + symptoms := w.evaluator.CachedSymptoms() + if symptoms == nil { + return fmt.Errorf("symptom cache not initialized") + } + + result := w.evaluator.reEvaluateOne(ctx, job.Args.ProwJobBuildID, symptoms) + if result.Status != ReEvalSuccess { + return fmt.Errorf("re-evaluation failed for %s: %s", job.Args.ProwJobBuildID, result.Error) + } + return nil +} + +// BuildInsertParams creates River insert parameters for a batch of build IDs. +func BuildInsertParams(buildIDs []string, symptomHash string) ([]river.InsertManyParams, []string) { + params := make([]river.InsertManyParams, len(buildIDs)) + keys := make([]string, len(buildIDs)) + for i, id := range buildIDs { + params[i] = river.InsertManyParams{ + Args: ReevaluateJobRunArgs{ + ProwJobBuildID: id, + SymptomHash: symptomHash, + }, + } + keys[i] = id + } + return params, keys +} + +// SymptomHash computes a stable hash of symptom definitions for dedup purposes. +func SymptomHash(symptoms []jobrunscan.Symptom) string { + return computeSymptomHash(symptoms) +} diff --git a/pkg/api/jobrunscan/reevaluate_worker_test.go b/pkg/api/jobrunscan/reevaluate_worker_test.go new file mode 100644 index 0000000000..2d3d34086a --- /dev/null +++ b/pkg/api/jobrunscan/reevaluate_worker_test.go @@ -0,0 +1,60 @@ +package jobrunscan + +import ( + "testing" +) + +func TestBuildInsertParams(t *testing.T) { + buildIDs := []string{"111", "222", "333"} + hash := "abc123" + + params, keys := BuildInsertParams(buildIDs, hash) + + if len(params) != 3 { + t.Fatalf("expected 3 params, got %d", len(params)) + } + if len(keys) != 3 { + t.Fatalf("expected 3 keys, got %d", len(keys)) + } + + for i, id := range buildIDs { + args, ok := params[i].Args.(ReevaluateJobRunArgs) + if !ok { + t.Fatalf("params[%d].Args is not ReevaluateJobRunArgs", i) + } + if args.ProwJobBuildID != id { + t.Errorf("params[%d].ProwJobBuildID = %q, want %q", i, args.ProwJobBuildID, id) + } + if args.SymptomHash != hash { + t.Errorf("params[%d].SymptomHash = %q, want %q", i, args.SymptomHash, hash) + } + if keys[i] != id { + t.Errorf("keys[%d] = %q, want %q", i, keys[i], id) + } + } +} + +func TestReevaluateJobRunArgs_Kind(t *testing.T) { + args := ReevaluateJobRunArgs{} + if args.Kind() != ReevaluateJobKind { + t.Errorf("Kind() = %q, want %q", args.Kind(), ReevaluateJobKind) + } +} + +func TestReevaluateJobRunArgs_InsertOpts(t *testing.T) { + args := ReevaluateJobRunArgs{} + opts := args.InsertOpts() + + if opts.Queue != ReevaluateQueue { + t.Errorf("Queue = %q, want %q", opts.Queue, ReevaluateQueue) + } + if opts.MaxAttempts != ReevaluateMaxAttempts { + t.Errorf("MaxAttempts = %d, want %d", opts.MaxAttempts, ReevaluateMaxAttempts) + } + if !opts.UniqueOpts.ByArgs { + t.Error("UniqueOpts.ByArgs should be true") + } + if opts.UniqueOpts.ByPeriod != ReevaluateDedupPeriod { + t.Errorf("UniqueOpts.ByPeriod = %v, want %v", opts.UniqueOpts.ByPeriod, ReevaluateDedupPeriod) + } +} diff --git a/pkg/api/jobrunscan/symptom_hash.go b/pkg/api/jobrunscan/symptom_hash.go new file mode 100644 index 0000000000..2f7e735917 --- /dev/null +++ b/pkg/api/jobrunscan/symptom_hash.go @@ -0,0 +1,34 @@ +package jobrunscan + +import ( + "crypto/sha256" + "fmt" + "sort" + "strings" + + "github.com/openshift/sippy/pkg/db/models/jobrunscan" +) + +// computeSymptomHash produces a deterministic hash over the sorted symptom +// definitions. Two calls with the same set of symptoms (regardless of order) +// will return the same hash. The hash includes each symptom's ID, matcher +// type, match string, file pattern, and label IDs so that any change to +// symptom configuration produces a different hash. +func computeSymptomHash(symptoms []jobrunscan.Symptom) string { + entries := make([]string, len(symptoms)) + for i, s := range symptoms { + labels := make([]string, len(s.LabelIDs)) + copy(labels, s.LabelIDs) + sort.Strings(labels) + entries[i] = fmt.Sprintf("%s|%s|%s|%s|%s", + s.ID, s.MatcherType, s.MatchString, s.FilePattern, strings.Join(labels, ",")) + } + sort.Strings(entries) + + h := sha256.New() + for _, e := range entries { + h.Write([]byte(e)) + h.Write([]byte{0}) + } + return fmt.Sprintf("%x", h.Sum(nil))[:16] +} diff --git a/pkg/api/jobrunscan/symptom_hash_test.go b/pkg/api/jobrunscan/symptom_hash_test.go new file mode 100644 index 0000000000..d79ac94cbf --- /dev/null +++ b/pkg/api/jobrunscan/symptom_hash_test.go @@ -0,0 +1,101 @@ +package jobrunscan + +import ( + "testing" + + "github.com/lib/pq" + "github.com/openshift/sippy/pkg/db/models/jobrunscan" +) + +func symptom(id, matcherType, matchString, filePattern string, labels ...string) jobrunscan.Symptom { + return jobrunscan.Symptom{ + SymptomContent: jobrunscan.SymptomContent{ + ID: id, + MatcherType: matcherType, + MatchString: matchString, + FilePattern: filePattern, + LabelIDs: pq.StringArray(labels), + }, + } +} + +func TestComputeSymptomHash(t *testing.T) { + tests := []struct { + name string + a, b []jobrunscan.Symptom + wantSame bool + }{ + { + name: "empty slices are equal", + a: nil, + b: []jobrunscan.Symptom{}, + wantSame: true, + }, + { + name: "same symptoms in same order", + a: []jobrunscan.Symptom{ + symptom("s1", "string", "foo", "*.log", "l1"), + symptom("s2", "regex", "bar.*", "*.txt", "l2"), + }, + b: []jobrunscan.Symptom{ + symptom("s1", "string", "foo", "*.log", "l1"), + symptom("s2", "regex", "bar.*", "*.txt", "l2"), + }, + wantSame: true, + }, + { + name: "same symptoms in different order", + a: []jobrunscan.Symptom{ + symptom("s2", "regex", "bar.*", "", "l2"), + symptom("s1", "string", "foo", "", "l1"), + }, + b: []jobrunscan.Symptom{ + symptom("s1", "string", "foo", "", "l1"), + symptom("s2", "regex", "bar.*", "", "l2"), + }, + wantSame: true, + }, + { + name: "different match string", + a: []jobrunscan.Symptom{symptom("s1", "string", "foo", "", "l1")}, + b: []jobrunscan.Symptom{symptom("s1", "string", "bar", "", "l1")}, + wantSame: false, + }, + { + name: "different labels", + a: []jobrunscan.Symptom{symptom("s1", "string", "foo", "", "l1")}, + b: []jobrunscan.Symptom{symptom("s1", "string", "foo", "", "l1", "l2")}, + wantSame: false, + }, + { + name: "label order does not matter", + a: []jobrunscan.Symptom{symptom("s1", "", "", "", "l2", "l1")}, + b: []jobrunscan.Symptom{symptom("s1", "", "", "", "l1", "l2")}, + wantSame: true, + }, + { + name: "extra symptom differs", + a: []jobrunscan.Symptom{symptom("s1", "", "", "", "l1")}, + b: []jobrunscan.Symptom{ + symptom("s1", "", "", "", "l1"), + symptom("s2", "", "", "", "l2"), + }, + wantSame: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + hashA := computeSymptomHash(tt.a) + hashB := computeSymptomHash(tt.b) + + if len(hashA) != 16 { + t.Errorf("hash length = %d, want 16", len(hashA)) + } + + if (hashA == hashB) != tt.wantSame { + t.Errorf("hashA=%s hashB=%s, wantSame=%v", hashA, hashB, tt.wantSame) + } + }) + } +} diff --git a/pkg/db/migrations/000013_create_workqueue_tables.down.sql b/pkg/db/migrations/000013_create_workqueue_tables.down.sql new file mode 100644 index 0000000000..190b7c0218 --- /dev/null +++ b/pkg/db/migrations/000013_create_workqueue_tables.down.sql @@ -0,0 +1,2 @@ +DROP TABLE IF EXISTS workqueue_batch_items; +DROP TABLE IF EXISTS workqueue_batches; diff --git a/pkg/db/migrations/000013_create_workqueue_tables.up.sql b/pkg/db/migrations/000013_create_workqueue_tables.up.sql new file mode 100644 index 0000000000..f8b6fa03e7 --- /dev/null +++ b/pkg/db/migrations/000013_create_workqueue_tables.up.sql @@ -0,0 +1,22 @@ +CREATE TABLE workqueue_batches ( + id UUID PRIMARY KEY, + kind TEXT NOT NULL, + requested_count INT NOT NULL, + enqueued_count INT NOT NULL DEFAULT 0, + deduped_count INT NOT NULL DEFAULT 0, + status TEXT NOT NULL DEFAULT 'pending', + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + completed_at TIMESTAMPTZ +); + +CREATE INDEX idx_workqueue_batches_kind_status ON workqueue_batches (kind, status); + +CREATE TABLE workqueue_batch_items ( + id BIGSERIAL PRIMARY KEY, + batch_id UUID NOT NULL REFERENCES workqueue_batches(id) ON DELETE CASCADE, + river_job_id BIGINT NOT NULL, + item_key TEXT NOT NULL, + UNIQUE(batch_id, item_key) +); + +CREATE INDEX idx_workqueue_batch_items_batch ON workqueue_batch_items (batch_id, river_job_id); diff --git a/pkg/db/migrations/MANIFEST b/pkg/db/migrations/MANIFEST index 0c6d1c812e..4246798b0a 100644 --- a/pkg/db/migrations/MANIFEST +++ b/pkg/db/migrations/MANIFEST @@ -19,3 +19,4 @@ 000010_drop_test_analysis_by_job_by_dates 000011_add_lifecycle_to_summaries 000012_drop_test_daily_totals_date_index +000013_create_workqueue_tables diff --git a/pkg/sippyserver/job_run_scan.go b/pkg/sippyserver/job_run_scan.go index 5b78ec2fab..a199ad4c42 100644 --- a/pkg/sippyserver/job_run_scan.go +++ b/pkg/sippyserver/job_run_scan.go @@ -2,8 +2,11 @@ package sippyserver import ( "encoding/json" + "fmt" "net/http" + "strings" + "github.com/google/uuid" "github.com/gorilla/mux" "github.com/openshift/sippy/pkg/api" apijobrunscan "github.com/openshift/sippy/pkg/api/jobrunscan" @@ -182,24 +185,32 @@ func (s *Server) jsonReEvaluateJobRunSymptoms(w http.ResponseWriter, req *http.R } req.Body = http.MaxBytesReader(w, req.Body, 1<<20) // 1 MiB limit to prevent DoS dec := json.NewDecoder(req.Body) - dec.DisallowUnknownFields() // catch client errors faster + dec.DisallowUnknownFields() if err := dec.Decode(&body); err != nil { failureResponse(w, http.StatusBadRequest, "invalid request body: "+err.Error()) return } - if err := apijobrunscan.ValidateReEvalRequest(body.ProwJobBuildIDs); err != nil { - failureResponse(w, http.StatusBadRequest, err.Error()) + if s.bigQueryClient == nil || s.gcsClient == nil || s.gcsBucket == "" { + failureResponse(w, http.StatusServiceUnavailable, "symptom re-evaluation requires BigQuery and GCS configuration") return } - if s.bigQueryClient == nil || s.gcsClient == nil || s.gcsBucket == "" { - failureResponse(w, http.StatusServiceUnavailable, "symptom re-evaluation requires BigQuery and GCS configuration") + if body.DryRun { + s.reEvaluateSynchronous(w, req, body.ProwJobBuildIDs) return } + s.reEvaluateAsync(w, req, body.ProwJobBuildIDs) +} - re := apijobrunscan.NewReEvaluator(s.bigQueryClient, s.gcsClient, s.gcsBucket, s.db, s.cache, s.jobartifactsManager, body.DryRun) - results, err := re.ReEvaluateJobRuns(req.Context(), body.ProwJobBuildIDs) +func (s *Server) reEvaluateSynchronous(w http.ResponseWriter, req *http.Request, buildIDs []string) { + if err := apijobrunscan.ValidateReEvalRequest(buildIDs); err != nil { + failureResponse(w, http.StatusBadRequest, err.Error()) + return + } + + re := apijobrunscan.NewReEvaluator(s.bigQueryClient, s.gcsClient, s.gcsBucket, s.db, s.cache, s.jobartifactsManager, true) + results, err := re.ReEvaluateJobRuns(req.Context(), buildIDs) if err != nil { failureResponse(w, http.StatusInternalServerError, err.Error()) return @@ -208,3 +219,70 @@ func (s *Server) jsonReEvaluateJobRunSymptoms(w http.ResponseWriter, req *http.R apijobrunscan.InjectReEvalHATEOASLinks(&resp, api.GetBaseURL(req)) api.RespondWithJSON(http.StatusOK, w, resp) } + +func (s *Server) reEvaluateAsync(w http.ResponseWriter, req *http.Request, buildIDs []string) { + if s.workqueueSubmitter == nil { + failureResponse(w, http.StatusServiceUnavailable, "async re-evaluation is not configured") + return + } + + if len(buildIDs) == 0 { + failureResponse(w, http.StatusBadRequest, "prow_job_build_ids is required") + return + } + if len(buildIDs) > apijobrunscan.MaxJobRunsPerBatch { + failureResponse(w, http.StatusBadRequest, fmt.Sprintf("maximum %d job runs per batch", apijobrunscan.MaxJobRunsPerBatch)) + return + } + + re := apijobrunscan.NewReEvaluator(s.bigQueryClient, s.gcsClient, s.gcsBucket, s.db, s.cache, s.jobartifactsManager, false) + symptomHash, err := re.RefreshSymptomCache() + if err != nil { + failureResponse(w, http.StatusInternalServerError, "failed to load symptoms: "+err.Error()) + return + } + + insertParams, itemKeys := apijobrunscan.BuildInsertParams(buildIDs, symptomHash) + result, err := s.workqueueSubmitter.Submit(req.Context(), apijobrunscan.ReevaluateJobKind, insertParams, itemKeys) + if err != nil { + failureResponse(w, http.StatusInternalServerError, "failed to submit batch: "+err.Error()) + return + } + + baseURL := api.GetBaseURL(req) + api.RespondWithJSON(http.StatusAccepted, w, map[string]interface{}{ + "batch_id": result.BatchID, + "requested": result.Requested, + "enqueued": result.Enqueued, + "deduped": result.Deduped, + "links": map[string]string{ + "status": baseURL + "/api/jobs/runs/reevaluate/" + result.BatchID.String(), + }, + }) +} + +func (s *Server) jsonReEvaluateBatchStatus(w http.ResponseWriter, req *http.Request) { + if s.workqueueStatusQuerier == nil { + failureResponse(w, http.StatusServiceUnavailable, "batch status is not configured") + return + } + + batchIDStr := mux.Vars(req)["batch_id"] + batchID, err := uuid.Parse(batchIDStr) + if err != nil { + failureResponse(w, http.StatusBadRequest, "invalid batch_id: "+err.Error()) + return + } + + status, err := s.workqueueStatusQuerier.Query(req.Context(), batchID) + if err != nil { + if strings.Contains(err.Error(), "record not found") { + failureResponse(w, http.StatusNotFound, "batch not found") + return + } + failureResponse(w, http.StatusInternalServerError, "failed to query batch status: "+err.Error()) + return + } + + api.RespondWithJSON(http.StatusOK, w, status) +} diff --git a/pkg/sippyserver/server.go b/pkg/sippyserver/server.go index 2b04a63f35..42e4519d71 100644 --- a/pkg/sippyserver/server.go +++ b/pkg/sippyserver/server.go @@ -62,6 +62,7 @@ import ( "github.com/openshift/sippy/pkg/db/models" "github.com/openshift/sippy/pkg/db/query" "github.com/openshift/sippy/pkg/filter" + "github.com/openshift/sippy/pkg/sippyserver/workqueue" "github.com/openshift/sippy/pkg/synthetictests" "github.com/openshift/sippy/pkg/testidentification" "github.com/openshift/sippy/pkg/util" @@ -137,6 +138,12 @@ func NewServer( return server } +// SetWorkqueue configures the server for async batch processing via River. +func (s *Server) SetWorkqueue(submitter *workqueue.Submitter, statusQuerier *workqueue.StatusQuerier) { + s.workqueueSubmitter = submitter + s.workqueueStatusQuerier = statusQuerier +} + var matViewRefreshMetric = promauto.NewHistogramVec(prometheus.HistogramOpts{ Name: "sippy_matview_refresh_millis", Help: "Milliseconds to refresh our postgresql materialized views", @@ -160,31 +167,33 @@ var matViewUniqueNumberOfJobRuns = promauto.NewGaugeVec(prometheus.GaugeOpts{ }, []string{"lookback_days"}) type Server struct { - mode Mode - listenAddr string - corsAllowedOrigin string - syntheticTestManager synthetictests.SyntheticTestManager - variantManager testidentification.VariantManager - jobartifactsManager *jobartifacts.Manager - sippyNG fs.FS - static fs.FS - httpServer *http.Server - db *db.DB - bigQueryClient *sippybq.Client - crDataProvider dataprovider.DataProvider - pinnedDateTime *time.Time - gcsClient *storage.Client - gcsBucket string - cache cache.Cache - crTimeRoundingFactor time.Duration - crTimeRoundingOffset time.Duration - capabilities []string - views *apitype.SippyViews - config *v1.SippyConfig - enableWriteAPIs bool - chatAPIURL string - jiraClient *jira.Client - rateLimiters map[string]*rateLimiter + mode Mode + listenAddr string + corsAllowedOrigin string + syntheticTestManager synthetictests.SyntheticTestManager + variantManager testidentification.VariantManager + jobartifactsManager *jobartifacts.Manager + sippyNG fs.FS + static fs.FS + httpServer *http.Server + db *db.DB + bigQueryClient *sippybq.Client + crDataProvider dataprovider.DataProvider + pinnedDateTime *time.Time + gcsClient *storage.Client + gcsBucket string + cache cache.Cache + crTimeRoundingFactor time.Duration + crTimeRoundingOffset time.Duration + capabilities []string + views *apitype.SippyViews + config *v1.SippyConfig + enableWriteAPIs bool + workqueueSubmitter *workqueue.Submitter + workqueueStatusQuerier *workqueue.StatusQuerier + chatAPIURL string + jiraClient *jira.Client + rateLimiters map[string]*rateLimiter } // getReleases returns release data via the configured data provider. @@ -2641,6 +2650,13 @@ func (s *Server) Serve() { Capabilities: []string{LocalDBCapability, WriteEndpointsCapability}, HandlerFunc: s.jsonReEvaluateJobRunSymptoms, }, + { + EndpointPath: "/api/jobs/runs/reevaluate/{batch_id}", + Description: "Get status of an async re-evaluation batch", + Methods: []string{http.MethodGet}, + Capabilities: []string{LocalDBCapability}, + HandlerFunc: s.jsonReEvaluateBatchStatus, + }, { EndpointPath: "/api/job_variants", Description: "Reports all job variants", diff --git a/pkg/sippyserver/workqueue/models.go b/pkg/sippyserver/workqueue/models.go new file mode 100644 index 0000000000..5c816e611c --- /dev/null +++ b/pkg/sippyserver/workqueue/models.go @@ -0,0 +1,47 @@ +package workqueue + +import ( + "time" + + "github.com/google/uuid" +) + +// BatchStatus represents the lifecycle state of a batch. +type BatchStatus string + +const ( + BatchStatusPending BatchStatus = "pending" + BatchStatusRunning BatchStatus = "running" + BatchStatusComplete BatchStatus = "complete" + BatchStatusFailed BatchStatus = "failed" +) + +// Batch represents a user-initiated batch of work items. +// A batch groups related work items for status tracking and progress reporting. +type Batch struct { + ID uuid.UUID `gorm:"type:uuid;primaryKey" json:"id"` + Kind string `gorm:"not null;index" json:"kind"` + RequestedCount int `gorm:"not null" json:"requested_count"` + EnqueuedCount int `gorm:"not null" json:"enqueued_count"` + DedupedCount int `gorm:"not null" json:"deduped_count"` + Status BatchStatus `gorm:"not null;default:'pending'" json:"status"` + CreatedAt time.Time `gorm:"autoCreateTime" json:"created_at"` + CompletedAt *time.Time ` json:"completed_at,omitempty"` +} + +func (Batch) TableName() string { + return "workqueue_batches" +} + +// BatchItem associates a batch with a River job for many-to-many status tracking. +// Multiple batches can reference the same River job (when deduplication occurs). +type BatchItem struct { + ID uint64 `gorm:"primaryKey;autoIncrement" json:"id"` + BatchID uuid.UUID `gorm:"type:uuid;not null;index:idx_batch_items" json:"batch_id"` + RiverJobID int64 `gorm:"not null;index:idx_batch_items" json:"river_job_id"` + ItemKey string `gorm:"not null" json:"item_key"` +} + +func (BatchItem) TableName() string { + return "workqueue_batch_items" +} diff --git a/pkg/sippyserver/workqueue/river_process.go b/pkg/sippyserver/workqueue/river_process.go new file mode 100644 index 0000000000..8c2608fe32 --- /dev/null +++ b/pkg/sippyserver/workqueue/river_process.go @@ -0,0 +1,42 @@ +package workqueue + +import ( + "context" + "time" + + "github.com/jackc/pgx/v5" + "github.com/riverqueue/river" + log "github.com/sirupsen/logrus" +) + +const shutdownTimeout = 30 * time.Second + +// RiverProcess adapts a River client to the DaemonProcess interface. +type RiverProcess struct { + client *river.Client[pgx.Tx] +} + +// NewRiverProcess wraps a River client as a DaemonProcess. +func NewRiverProcess(client *river.Client[pgx.Tx]) *RiverProcess { + return &RiverProcess{client: client} +} + +// Run starts the River client and blocks until the context is cancelled. +func (p *RiverProcess) Run(ctx context.Context) { + if err := p.client.Start(ctx); err != nil { + log.WithError(err).Error("failed to start River client") + return + } + log.Info("River work queue started") + + <-ctx.Done() + + shutdownCtx, cancel := context.WithTimeout(context.Background(), shutdownTimeout) + defer cancel() + + if err := p.client.Stop(shutdownCtx); err != nil { + log.WithError(err).Error("error stopping River client") + } else { + log.Info("River work queue stopped gracefully") + } +} diff --git a/pkg/sippyserver/workqueue/setup.go b/pkg/sippyserver/workqueue/setup.go new file mode 100644 index 0000000000..171437c1e5 --- /dev/null +++ b/pkg/sippyserver/workqueue/setup.go @@ -0,0 +1,77 @@ +package workqueue + +import ( + "context" + "fmt" + "time" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgxpool" + "github.com/riverqueue/river" + "github.com/riverqueue/river/riverdriver/riverpgxv5" + "github.com/riverqueue/river/rivermigrate" + log "github.com/sirupsen/logrus" +) + +// SetupConfig holds the configuration for setting up River. +type SetupConfig struct { + DatabaseDSN string + Queues map[string]river.QueueConfig + Workers *river.Workers + CompletedJobRetention time.Duration + DiscardedJobRetention time.Duration +} + +// SetupResult holds the outputs of a successful River setup. +type SetupResult struct { + Pool *pgxpool.Pool + Client *river.Client[pgx.Tx] +} + +// Setup creates a pgx/v5 pool, runs River migrations, and returns a configured +// River client. If Queues is empty, the client operates in insert-only mode +// (no workers started). +func Setup(ctx context.Context, cfg SetupConfig) (*SetupResult, error) { + pool, err := pgxpool.New(ctx, cfg.DatabaseDSN) + if err != nil { + return nil, fmt.Errorf("creating pgx/v5 pool: %w", err) + } + + driver := riverpgxv5.New(pool) + + migrator, err := rivermigrate.New(driver, nil) + if err != nil { + pool.Close() + return nil, fmt.Errorf("creating River migrator: %w", err) + } + + res, err := migrator.Migrate(ctx, rivermigrate.DirectionUp, nil) + if err != nil { + pool.Close() + return nil, fmt.Errorf("running River migrations: %w", err) + } + if len(res.Versions) > 0 { + log.WithField("versions", len(res.Versions)).Info("River migrations applied") + } + + riverConfig := &river.Config{ + Workers: cfg.Workers, + } + if len(cfg.Queues) > 0 { + riverConfig.Queues = cfg.Queues + } + if cfg.CompletedJobRetention > 0 { + riverConfig.CompletedJobRetentionPeriod = cfg.CompletedJobRetention + } + if cfg.DiscardedJobRetention > 0 { + riverConfig.DiscardedJobRetentionPeriod = cfg.DiscardedJobRetention + } + + client, err := river.NewClient(driver, riverConfig) + if err != nil { + pool.Close() + return nil, fmt.Errorf("creating River client: %w", err) + } + + return &SetupResult{Pool: pool, Client: client}, nil +} diff --git a/pkg/sippyserver/workqueue/status.go b/pkg/sippyserver/workqueue/status.go new file mode 100644 index 0000000000..41e71b401f --- /dev/null +++ b/pkg/sippyserver/workqueue/status.go @@ -0,0 +1,152 @@ +package workqueue + +import ( + "context" + "fmt" + "time" + + "github.com/google/uuid" + "github.com/riverqueue/river/rivertype" + "gorm.io/gorm" +) + +// ItemStatus represents a single work item's current state. +type ItemStatus struct { + ItemKey string `json:"item_key"` + State rivertype.JobState `json:"state"` + Errors []map[string]string `json:"errors,omitempty"` +} + +// BatchStatusResponse contains the full status of a batch. +type BatchStatusResponse struct { + BatchID uuid.UUID `json:"batch_id"` + Status BatchStatus `json:"status"` + Total int `json:"total"` + Completed int `json:"completed"` + Failed int `json:"failed"` + Running int `json:"running"` + Pending int `json:"pending"` + Items []ItemStatus `json:"items"` +} + +// batchItemRow is the result of joining batch_items with river_job. +type batchItemRow struct { + ItemKey string + JobState string +} + +// StatusQuerier retrieves batch status by joining batch items with River jobs. +type StatusQuerier struct { + db *gorm.DB +} + +// NewStatusQuerier creates a StatusQuerier. +func NewStatusQuerier(db *gorm.DB) *StatusQuerier { + return &StatusQuerier{db: db} +} + +// Query loads a batch and its items' current River job states. +func (q *StatusQuerier) Query(ctx context.Context, batchID uuid.UUID) (*BatchStatusResponse, error) { + var batch Batch + if err := q.db.WithContext(ctx).First(&batch, "id = ?", batchID).Error; err != nil { + return nil, fmt.Errorf("loading batch: %w", err) + } + + var rows []batchItemRow + err := q.db.WithContext(ctx). + Table("workqueue_batch_items bi"). + Select("bi.item_key, rj.state as job_state"). + Joins("JOIN river_job rj ON rj.id = bi.river_job_id"). + Where("bi.batch_id = ?", batchID). + Scan(&rows).Error + if err != nil { + return nil, fmt.Errorf("querying batch items: %w", err) + } + + resp := &BatchStatusResponse{ + BatchID: batchID, + Total: len(rows), + } + + resp.Items = make([]ItemStatus, len(rows)) + for i, row := range rows { + state := rivertype.JobState(row.JobState) + resp.Items[i] = ItemStatus{ + ItemKey: row.ItemKey, + State: state, + } + + switch categorizeState(state) { + case stateCompleted: + resp.Completed++ + case stateFailed: + resp.Failed++ + case stateRunning: + resp.Running++ + case statePending: + resp.Pending++ + } + } + + resp.Status = computeBatchStatus(&batch, resp) + q.maybeFinalizeBatch(ctx, &batch, resp) + + return resp, nil +} + +type stateCategory int + +const ( + statePending stateCategory = iota + stateRunning + stateCompleted + stateFailed +) + +func categorizeState(s rivertype.JobState) stateCategory { + switch s { + case rivertype.JobStateCompleted: + return stateCompleted + case rivertype.JobStateDiscarded, rivertype.JobStateCancelled: + return stateFailed + case rivertype.JobStateRunning: + return stateRunning + default: + return statePending + } +} + +func computeBatchStatus(batch *Batch, resp *BatchStatusResponse) BatchStatus { + if batch.CompletedAt != nil { + return batch.Status + } + if resp.Total == 0 { + return BatchStatusComplete + } + terminal := resp.Completed + resp.Failed + if terminal < resp.Total { + return BatchStatusRunning + } + if resp.Completed == 0 { + return BatchStatusFailed + } + return BatchStatusComplete +} + +// maybeFinalizeBatch updates the batch row if all items have reached a terminal +// state. This is idempotent — repeated polls are harmless. +func (q *StatusQuerier) maybeFinalizeBatch(ctx context.Context, batch *Batch, resp *BatchStatusResponse) { + if batch.CompletedAt != nil { + return + } + terminal := resp.Completed + resp.Failed + if resp.Total == 0 || terminal < resp.Total { + return + } + + now := time.Now() + q.db.WithContext(ctx).Model(batch).Updates(map[string]interface{}{ + "status": resp.Status, + "completed_at": now, + }) +} diff --git a/pkg/sippyserver/workqueue/status_test.go b/pkg/sippyserver/workqueue/status_test.go new file mode 100644 index 0000000000..cec17e68aa --- /dev/null +++ b/pkg/sippyserver/workqueue/status_test.go @@ -0,0 +1,92 @@ +package workqueue + +import ( + "testing" + + "github.com/google/uuid" + "github.com/riverqueue/river/rivertype" +) + +func TestCategorizeState(t *testing.T) { + tests := []struct { + state rivertype.JobState + want stateCategory + }{ + {rivertype.JobStateCompleted, stateCompleted}, + {rivertype.JobStateDiscarded, stateFailed}, + {rivertype.JobStateCancelled, stateFailed}, + {rivertype.JobStateRunning, stateRunning}, + {rivertype.JobStateAvailable, statePending}, + {rivertype.JobStateScheduled, statePending}, + {rivertype.JobStateRetryable, statePending}, + {rivertype.JobStatePending, statePending}, + } + for _, tt := range tests { + t.Run(string(tt.state), func(t *testing.T) { + if got := categorizeState(tt.state); got != tt.want { + t.Errorf("categorizeState(%q) = %d, want %d", tt.state, got, tt.want) + } + }) + } +} + +func TestComputeBatchStatus(t *testing.T) { + tests := []struct { + name string + resp *BatchStatusResponse + want BatchStatus + }{ + { + name: "empty batch", + resp: &BatchStatusResponse{Total: 0}, + want: BatchStatusComplete, + }, + { + name: "all completed", + resp: &BatchStatusResponse{Total: 5, Completed: 5}, + want: BatchStatusComplete, + }, + { + name: "some still running", + resp: &BatchStatusResponse{Total: 5, Completed: 3, Running: 2}, + want: BatchStatusRunning, + }, + { + name: "all failed", + resp: &BatchStatusResponse{Total: 3, Failed: 3}, + want: BatchStatusFailed, + }, + { + name: "mixed completed and failed", + resp: &BatchStatusResponse{Total: 5, Completed: 3, Failed: 2}, + want: BatchStatusComplete, + }, + { + name: "some pending", + resp: &BatchStatusResponse{Total: 5, Completed: 2, Pending: 3}, + want: BatchStatusRunning, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + batch := &Batch{ID: uuid.New()} + got := computeBatchStatus(batch, tt.resp) + if got != tt.want { + t.Errorf("computeBatchStatus() = %q, want %q", got, tt.want) + } + }) + } +} + +func TestComputeBatchStatus_AlreadyFinalized(t *testing.T) { + batch := &Batch{ID: uuid.New(), Status: BatchStatusFailed} + now := batch.CreatedAt + batch.CompletedAt = &now + + resp := &BatchStatusResponse{Total: 3, Completed: 3} + got := computeBatchStatus(batch, resp) + if got != BatchStatusFailed { + t.Errorf("expected finalized status to be preserved, got %q", got) + } +} diff --git a/pkg/sippyserver/workqueue/submitter.go b/pkg/sippyserver/workqueue/submitter.go new file mode 100644 index 0000000000..f9e61dd069 --- /dev/null +++ b/pkg/sippyserver/workqueue/submitter.go @@ -0,0 +1,115 @@ +package workqueue + +import ( + "context" + "fmt" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + "github.com/riverqueue/river" + "github.com/riverqueue/river/rivertype" + log "github.com/sirupsen/logrus" + "gorm.io/gorm" +) + +// SubmitResult contains the outcome of a batch submission. +type SubmitResult struct { + BatchID uuid.UUID `json:"batch_id"` + Requested int `json:"requested"` + Enqueued int `json:"enqueued"` + Deduped int `json:"deduped"` +} + +// Submitter creates batches and enqueues River jobs for async processing. +type Submitter struct { + db *gorm.DB + riverClient *river.Client[pgx.Tx] +} + +// NewSubmitter creates a Submitter. +func NewSubmitter(db *gorm.DB, riverClient *river.Client[pgx.Tx]) *Submitter { + return &Submitter{ + db: db, + riverClient: riverClient, + } +} + +// Submit creates a batch, enqueues River jobs for each item, and records +// batch-item associations. Uses two separate transactions (gorm for batch +// tracking, pgx/v5 for River jobs) because the two driver versions cannot +// share a transaction. +func (s *Submitter) Submit(ctx context.Context, kind string, items []river.InsertManyParams, itemKeys []string) (*SubmitResult, error) { + if len(items) != len(itemKeys) { + return nil, fmt.Errorf("items and itemKeys must have the same length") + } + + batchID := uuid.New() + batch := Batch{ + ID: batchID, + Kind: kind, + RequestedCount: len(items), + Status: BatchStatusPending, + } + + // Insert River jobs first. If this succeeds but the batch tracking fails, + // the jobs still run (harmless) and the user can submit another batch. + results, err := s.riverClient.InsertMany(ctx, items) + if err != nil { + return nil, fmt.Errorf("inserting River jobs: %w", err) + } + + enqueued, deduped := countInsertResults(results) + + // Now create the batch and batch-item rows in gorm. + err = s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { + batch.EnqueuedCount = enqueued + batch.DedupedCount = deduped + batch.Status = BatchStatusRunning + + if err := tx.Create(&batch).Error; err != nil { + return fmt.Errorf("creating batch: %w", err) + } + + batchItems := buildBatchItems(batchID, results, itemKeys) + if err := tx.CreateInBatches(batchItems, 500).Error; err != nil { + return fmt.Errorf("creating batch items: %w", err) + } + + return nil + }) + if err != nil { + log.WithError(err).WithField("batch_id", batchID). + Warn("River jobs were enqueued but batch tracking failed") + return nil, fmt.Errorf("recording batch: %w", err) + } + + return &SubmitResult{ + BatchID: batchID, + Requested: len(items), + Enqueued: enqueued, + Deduped: deduped, + }, nil +} + +func countInsertResults(results []*rivertype.JobInsertResult) (enqueued, deduped int) { + for _, r := range results { + if r.UniqueSkippedAsDuplicate { + deduped++ + } else { + enqueued++ + } + } + return +} + +func buildBatchItems(batchID uuid.UUID, results []*rivertype.JobInsertResult, itemKeys []string) []BatchItem { + items := make([]BatchItem, len(results)) + for i, r := range results { + items[i] = BatchItem{ + BatchID: batchID, + RiverJobID: r.Job.ID, + ItemKey: itemKeys[i], + } + } + return items +} diff --git a/pkg/sippyserver/workqueue/submitter_test.go b/pkg/sippyserver/workqueue/submitter_test.go new file mode 100644 index 0000000000..246784b071 --- /dev/null +++ b/pkg/sippyserver/workqueue/submitter_test.go @@ -0,0 +1,87 @@ +package workqueue + +import ( + "testing" + + "github.com/google/uuid" + "github.com/riverqueue/river/rivertype" +) + +func TestCountInsertResults(t *testing.T) { + tests := []struct { + name string + results []*rivertype.JobInsertResult + wantNew int + wantDeduped int + }{ + { + name: "empty", + results: nil, + }, + { + name: "all new", + results: []*rivertype.JobInsertResult{ + {Job: &rivertype.JobRow{ID: 1}}, + {Job: &rivertype.JobRow{ID: 2}}, + }, + wantNew: 2, + }, + { + name: "all deduped", + results: []*rivertype.JobInsertResult{ + {Job: &rivertype.JobRow{ID: 1}, UniqueSkippedAsDuplicate: true}, + {Job: &rivertype.JobRow{ID: 2}, UniqueSkippedAsDuplicate: true}, + }, + wantDeduped: 2, + }, + { + name: "mixed", + results: []*rivertype.JobInsertResult{ + {Job: &rivertype.JobRow{ID: 1}}, + {Job: &rivertype.JobRow{ID: 2}, UniqueSkippedAsDuplicate: true}, + {Job: &rivertype.JobRow{ID: 3}}, + }, + wantNew: 2, + wantDeduped: 1, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + gotNew, gotDeduped := countInsertResults(tt.results) + if gotNew != tt.wantNew { + t.Errorf("enqueued = %d, want %d", gotNew, tt.wantNew) + } + if gotDeduped != tt.wantDeduped { + t.Errorf("deduped = %d, want %d", gotDeduped, tt.wantDeduped) + } + }) + } +} + +func TestBuildBatchItems(t *testing.T) { + batchID := uuid.New() + results := []*rivertype.JobInsertResult{ + {Job: &rivertype.JobRow{ID: 100}}, + {Job: &rivertype.JobRow{ID: 200}, UniqueSkippedAsDuplicate: true}, + } + keys := []string{"build-1", "build-2"} + + items := buildBatchItems(batchID, results, keys) + + if len(items) != 2 { + t.Fatalf("expected 2 items, got %d", len(items)) + } + + if items[0].BatchID != batchID { + t.Errorf("items[0].BatchID = %v, want %v", items[0].BatchID, batchID) + } + if items[0].RiverJobID != 100 { + t.Errorf("items[0].RiverJobID = %d, want 100", items[0].RiverJobID) + } + if items[0].ItemKey != "build-1" { + t.Errorf("items[0].ItemKey = %q, want %q", items[0].ItemKey, "build-1") + } + if items[1].RiverJobID != 200 { + t.Errorf("items[1].RiverJobID = %d, want 200", items[1].RiverJobID) + } +} diff --git a/vendor/github.com/jackc/pgx/v5/.gitignore b/vendor/github.com/jackc/pgx/v5/.gitignore new file mode 100644 index 0000000000..a2ebbe9c60 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/.gitignore @@ -0,0 +1,27 @@ +# Compiled Object files, Static and Dynamic libs (Shared Objects) +*.o +*.a +*.so + +# Folders +_obj +_test + +# Architecture specific extensions/prefixes +*.[568vq] +[568vq].out + +*.cgo1.go +*.cgo2.c +_cgo_defun.c +_cgo_gotypes.go +_cgo_export.* + +_testmain.go + +*.exe + +.envrc +/.testdb + +.DS_Store diff --git a/vendor/github.com/jackc/pgx/v5/.golangci.yml b/vendor/github.com/jackc/pgx/v5/.golangci.yml new file mode 100644 index 0000000000..ec98840f55 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/.golangci.yml @@ -0,0 +1,28 @@ +# See for configurations: https://golangci-lint.run/usage/configuration/ +version: "2" + +linters: + default: none + enable: + - govet + - ineffassign + - unconvert + - gocritic + +# See: https://golangci-lint.run/usage/formatters/ +formatters: + enable: + - gofmt # https://pkg.go.dev/cmd/gofmt + - gofumpt # https://github.com/mvdan/gofumpt + + settings: + gofmt: + simplify: true # Simplify code: gofmt with `-s` option. + + gofumpt: + # Module path which contains the source code being formatted. + # Default: "" + module-path: github.com/jackc/pgx/v5 # Should match with module in go.mod + # Choose whether to use the extra rules. + # Default: false + extra-rules: true diff --git a/vendor/github.com/jackc/pgx/v5/CHANGELOG.md b/vendor/github.com/jackc/pgx/v5/CHANGELOG.md new file mode 100644 index 0000000000..12e633a4bb --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/CHANGELOG.md @@ -0,0 +1,605 @@ +# 5.10.0 (June 3, 2026) + +This release includes a significant amount of hardening against malicious or compromised PostgreSQL servers, +contributed by Sean Chittenden at CrowdStrike, Inc. This work bounds binary decoders against attacker-controlled +message sizes, caps server-supplied SCRAM iteration counts, adds `require_auth` to restrict which authentication +methods a server may use (mitigating downgrade attacks under `sslmode=prefer`), and ensures cancellation requests are +sent over TLS when the original connection used TLS. + +## Features + +* Add `require_auth` to restrict accepted server authentication methods (Sean Chittenden at CrowdStrike, Inc.) +* Add `ParseConfigOptions.ConnStringAllowedKeys` to restrict allowed connection string keys (Sean Chittenden at CrowdStrike, Inc.) +* Add `StructArgs` and `StrictStructArgs` for `@`-named queries (Tubelight30) +* Add `ErrConnClosed` sentinel error and unwrap it from `connLockError` (Charlie Tonneslan) +* pgxpool: check if connection is expired before acquire (arthurdotwork) + +## Security Hardening + +* Encrypt `CancelRequest` connection when the primary connection used TLS (Sean Chittenden at CrowdStrike, Inc.) +* Cap server-supplied SCRAM iteration count (Sean Chittenden at CrowdStrike, Inc.) +* Default Frontend max message body length to ~1 GiB (Sean Chittenden at CrowdStrike, Inc.) +* Bound hstore binary decode against malicious server input (Sean Chittenden at CrowdStrike, Inc.) +* Bound array binary decode element length against remaining message bytes (Sean Chittenden at CrowdStrike, Inc.) +* Bound array element count against remaining message bytes (Sean Chittenden at CrowdStrike, Inc.) +* Bound range, multirange, and tsvector binary decoders (Sean Chittenden at CrowdStrike, Inc.) +* Document secure connection configuration (Sean Chittenden at CrowdStrike, Inc.) +* Fix panic on malformed geometric text; return an error instead (MaIII) + +## Fixes + +* Fix scanning `"char"` (OID 18) into `*string` in binary format (luongs3) +* Fix handling of typed-nil `driver.Valuer` in array and composite codecs (Donncha Fahy) +* Fix `CopyData.Data` hex decoding in `UnmarshalJSON` (Charlie Tonneslan) +* Fix data race when context is cancelled during connect +* Fix `parseKeywordValueSettings` rejecting trailing whitespace (alliasgher) +* pgconn: preserve full error chain in `normalizeTimeoutError` (Charlie Tonneslan) +* pgconn: use a fresh context for the fallback connection in `connectPreferred` (Charlie Tonneslan) +* pgxpool: fix `MaxLifetimeDestroyCount` and ping order for acquire-time expiry check +* Add missing error check of `rows.Err` to load types (Jen Altavilla) + +# 5.9.2 (April 18, 2026) + +Fix SQL Injection via placeholder confusion with dollar quoted string literals (GHSA-j88v-2chj-qfwx) + +SQL injection can occur when: + +1. The non-default simple protocol is used. +2. A dollar quoted string literal is used in the SQL query. +3. That query contains text that would be would be interpreted outside as a placeholder outside of a string literal. +4. The value of that placeholder is controllable by the attacker. + +e.g. + +```go +attackValue := `$tag$; drop table canary; --` +_, err = tx.Exec(ctx, `select $tag$ $1 $tag$, $1`, pgx.QueryExecModeSimpleProtocol, attackValue) +``` + +This is unlikely to occur outside of a contrived scenario. + +# 5.9.1 (March 22, 2026) + +* Fix: batch result format corruption when using cached prepared statements (reported by Dirkjan Bussink) + +# 5.9.0 (March 21, 2026) + +This release includes a number of new features such as SCRAM-SHA-256-PLUS support, OAuth authentication support, and +PostgreSQL protocol 3.2 support. + +It significantly reduces the amount of network traffic when using prepared statements (which are used automatically by +default) by avoiding unnecessary Describe Portal messages. This also reduces local memory usage. + +It also includes multiple fixes for potential DoS due to panic or OOM if connected to a malicious server that sends +deliberately malformed messages. + +* Require Go 1.25+ +* Add SCRAM-SHA-256-PLUS support (Adam Brightwell) +* Add OAuth authentication support for PostgreSQL 18 (David Schneider) +* Add PostgreSQL protocol 3.2 support (Dirkjan Bussink) +* Add tsvector type support (Adam Brightwell) +* Skip Describe Portal for cached prepared statements reducing network round trips +* Make LoadTypes query easier to support on "postgres-like" servers (Jelte Fennema-Nio) +* Default empty user to current OS user matching libpq behavior (ShivangSrivastava) +* Optimize LRU statement cache with custom linked list and node pooling (Mathias Bogaert) +* Optimize date scanning by replacing regex with manual parsing (Mathias Bogaert) +* Optimize pgio append/set functions with direct byte shifts (Mathias Bogaert) +* Make RowsAffected faster (Abhishek Chanda) +* Fix: Pipeline.Close panic when server sends multiple FATAL errors (Varun Chawla) +* Fix: ContextWatcher goroutine leak (Hank Donnay) +* Fix: stdlib discard connections with open transactions in ResetSession (Jeremy Schneider) +* Fix: pipelineBatchResults.Exec silently swallowing lastRows error +* Fix: ColumnTypeLength using BPCharArrayOID instead of BPCharOID +* Fix: TSVector text encoding returning nil for valid empty tsvector +* Fix: wrong error messages for Int2 and Int4 underflow +* Fix: Numeric nil Int pointer dereference with Valid: true +* Fix: reversed strings.ContainsAny arguments in Numeric.ScanScientific +* Fix: message length parsing on 32-bit platforms +* Fix: FunctionCallResponse.Decode mishandling of signed result size +* Fix: returning wrong error in configTLS when DecryptPEMBlock fails (Maxim Motyshen) +* Fix: misleading ParseConfig error when default_query_exec_mode is invalid (Skarm) +* Fix: missed Unwatch in Pipeline error paths +* Clarify too many failed acquire attempts error message +* Better error wrapping with context and SQL statement (Aneesh Makala) +* Enable govet and ineffassign linters (Federico Guerinoni) +* Guard against various malformed binary messages (arrays, hstore, multirange, protocol messages) +* Fix various godoc comments (ferhat elmas) +* Fix typos in comments (Oleksandr Redko) + +# 5.8.0 (December 26, 2025) + +* Require Go 1.24+ +* Remove golang.org/x/crypto dependency +* Add OptionShouldPing to control ResetSession ping behavior (ilyam8) +* Fix: Avoid overflow when MaxConns is set to MaxInt32 +* Fix: Close batch pipeline after a query error (Anthonin Bonnefoy) +* Faster shutdown of pgxpool.Pool background goroutines (Blake Gentry) +* Add pgxpool ping timeout (Amirsalar Safaei) +* Fix: Rows.FieldDescriptions for empty query +* Scan unknown types into *any as string or []byte based on format code +* Optimize pgtype.Numeric (Philip Dubé) +* Add AfterNetConnect hook to pgconn.Config +* Fix: Handle for preparing statements that fail during the Describe phase +* Fix overflow in numeric scanning (Ilia Demianenko) +* Fix: json/jsonb sql.Scanner source type is []byte +* Migrate from math/rand to math/rand/v2 (Mathias Bogaert) +* Optimize internal iobufpool (Mathias Bogaert) +* Optimize stmtcache invalidation (Mathias Bogaert) +* Fix: missing error case in interval parsing (Maxime Soulé) +* Fix: invalidate statement/description cache in Exec (James Hartig) +* ColumnTypeLength method return the type length for varbit type (DengChan) +* Array and Composite codecs handle typed nils + +# 5.7.6 (September 8, 2025) + +* Use ParseConfigError in pgx.ParseConfig and pgxpool.ParseConfig (Yurasov Ilia) +* Add PrepareConn hook to pgxpool (Jonathan Hall) +* Reduce allocations in QueryContext (Dominique Lefevre) +* Add MarshalJSON and UnmarshalJSON for pgtype.Uint32 (Panos Koutsovasilis) +* Configure ping behavior on pgxpool with ShouldPing (Christian Kiely) +* zeronull int types implement Int64Valuer and Int64Scanner (Li Zeghong) +* Fix panic when receiving terminate connection message during CopyFrom (Michal Drausowski) +* Fix statement cache not being invalidated on error during batch (Muhammadali Nazarov) + +# 5.7.5 (May 17, 2025) + +* Support sslnegotiation connection option (divyam234) +* Update golang.org/x/crypto to v0.37.0. This placates security scanners that were unable to see that pgx did not use the behavior affected by https://pkg.go.dev/vuln/GO-2025-3487. +* TraceLog now logs Acquire and Release at the debug level (dave sinclair) +* Add support for PGTZ environment variable +* Add support for PGOPTIONS environment variable +* Unpin memory used by Rows quicker +* Remove PlanScan memoization. This resolves a rare issue where scanning could be broken for one type by first scanning another. The problem was in the memoization system and benchmarking revealed that memoization was not providing any meaningful benefit. + +# 5.7.4 (March 24, 2025) + +* Fix / revert change to scanning JSON `null` (Felix Röhrich) + +# 5.7.3 (March 21, 2025) + +* Expose EmptyAcquireWaitTime in pgxpool.Stat (vamshiaruru32) +* Improve SQL sanitizer performance (ninedraft) +* Fix Scan confusion with json(b), sql.Scanner, and automatic dereferencing (moukoublen, felix-roehrich) +* Fix Values() for xml type always returning nil instead of []byte +* Add ability to send Flush message in pipeline mode (zenkovev) +* Fix pgtype.Timestamp's JSON behavior to match PostgreSQL (pconstantinou) +* Better error messages when scanning structs (logicbomb) +* Fix handling of error on batch write (bonnefoa) +* Match libpq's connection fallback behavior more closely (felix-roehrich) +* Add MinIdleConns to pgxpool (djahandarie) + +# 5.7.2 (December 21, 2024) + +* Fix prepared statement already exists on batch prepare failure +* Add commit query to tx options (Lucas Hild) +* Fix pgtype.Timestamp json unmarshal (Shean de Montigny-Desautels) +* Add message body size limits in frontend and backend (zene) +* Add xid8 type +* Ensure planning encodes and scans cannot infinitely recurse +* Implement pgtype.UUID.String() (Konstantin Grachev) +* Switch from ExecParams to Exec in ValidateConnectTargetSessionAttrs functions (Alexander Rumyantsev) +* Update golang.org/x/crypto +* Fix json(b) columns prefer sql.Scanner interface like database/sql (Ludovico Russo) + +# 5.7.1 (September 10, 2024) + +* Fix data race in tracelog.TraceLog +* Update puddle to v2.2.2. This removes the import of nanotime via linkname. +* Update golang.org/x/crypto and golang.org/x/text + +# 5.7.0 (September 7, 2024) + +* Add support for sslrootcert=system (Yann Soubeyrand) +* Add LoadTypes to load multiple types in a single SQL query (Nick Farrell) +* Add XMLCodec supports encoding + scanning XML column type like json (nickcruess-soda) +* Add MultiTrace (Stepan Rabotkin) +* Add TraceLogConfig with customizable TimeKey (stringintech) +* pgx.ErrNoRows wraps sql.ErrNoRows to aid in database/sql compatibility with native pgx functions (merlin) +* Support scanning binary formatted uint32 into string / TextScanner (jennifersp) +* Fix interval encoding to allow 0s and avoid extra spaces (Carlos Pérez-Aradros Herce) +* Update pgservicefile - fixes panic when parsing invalid file +* Better error message when reading past end of batch +* Don't print url when url.Parse returns an error (Kevin Biju) +* Fix snake case name normalization collision in RowToStructByName with db tag (nolandseigler) +* Fix: Scan and encode types with underlying types of arrays + +# 5.6.0 (May 25, 2024) + +* Add StrictNamedArgs (Tomas Zahradnicek) +* Add support for macaddr8 type (Carlos Pérez-Aradros Herce) +* Add SeverityUnlocalized field to PgError / Notice +* Performance optimization of RowToStructByPos/Name (Zach Olstein) +* Allow customizing context canceled behavior for pgconn +* Add ScanLocation to pgtype.Timestamp[tz]Codec +* Add custom data to pgconn.PgConn +* Fix ResultReader.Read() to handle nil values +* Do not encode interval microseconds when they are 0 (Carlos Pérez-Aradros Herce) +* pgconn.SafeToRetry checks for wrapped errors (tjasko) +* Failed connection attempts include all errors +* Optimize LargeObject.Read (Mitar) +* Add tracing for connection acquire and release from pool (ngavinsir) +* Fix encode driver.Valuer not called when nil +* Add support for custom JSON marshal and unmarshal (Mitar) +* Use Go default keepalive for TCP connections (Hans-Joachim Kliemeck) + +# 5.5.5 (March 9, 2024) + +Use spaces instead of parentheses for SQL sanitization. + +This still solves the problem of negative numbers creating a line comment, but this avoids breaking edge cases such as +`set foo to $1` where the substitution is taking place in a location where an arbitrary expression is not allowed. + +# 5.5.4 (March 4, 2024) + +Fix CVE-2024-27304 + +SQL injection can occur if an attacker can cause a single query or bind message to exceed 4 GB in size. An integer +overflow in the calculated message size can cause the one large message to be sent as multiple messages under the +attacker's control. + +Thanks to Paul Gerste for reporting this issue. + +* Fix behavior of CollectRows to return empty slice if Rows are empty (Felix) +* Fix simple protocol encoding of json.RawMessage +* Fix *Pipeline.getResults should close pipeline on error +* Fix panic in TryFindUnderlyingTypeScanPlan (David Kurman) +* Fix deallocation of invalidated cached statements in a transaction +* Handle invalid sslkey file +* Fix scan float4 into sql.Scanner +* Fix pgtype.Bits not making copy of data from read buffer. This would cause the data to be corrupted by future reads. + +# 5.5.3 (February 3, 2024) + +* Fix: prepared statement already exists +* Improve CopyFrom auto-conversion of text-ish values +* Add ltree type support (Florent Viel) +* Make some properties of Batch and QueuedQuery public (Pavlo Golub) +* Add AppendRows function (Edoardo Spadolini) +* Optimize convert UUID [16]byte to string (Kirill Malikov) +* Fix: LargeObject Read and Write of more than ~1GB at a time (Mitar) + +# 5.5.2 (January 13, 2024) + +* Allow NamedArgs to start with underscore +* pgproto3: Maximum message body length support (jeremy.spriet) +* Upgrade golang.org/x/crypto to v0.17.0 +* Add snake_case support to RowToStructByName (Tikhon Fedulov) +* Fix: update description cache after exec prepare (James Hartig) +* Fix: pipeline checks if it is closed (James Hartig and Ryan Fowler) +* Fix: normalize timeout / context errors during TLS startup (Samuel Stauffer) +* Add OnPgError for easier centralized error handling (James Hartig) + +# 5.5.1 (December 9, 2023) + +* Add CopyFromFunc helper function. (robford) +* Add PgConn.Deallocate method that uses PostgreSQL protocol Close message. +* pgx uses new PgConn.Deallocate method. This allows deallocating statements to work in a failed transaction. This fixes a case where the prepared statement map could become invalid. +* Fix: Prefer driver.Valuer over json.Marshaler for json fields. (Jacopo) +* Fix: simple protocol SQL sanitizer previously panicked if an invalid $0 placeholder was used. This now returns an error instead. (maksymnevajdev) +* Add pgtype.Numeric.ScanScientific (Eshton Robateau) + +# 5.5.0 (November 4, 2023) + +* Add CollectExactlyOneRow. (Julien GOTTELAND) +* Add OpenDBFromPool to create *database/sql.DB from *pgxpool.Pool. (Lev Zakharov) +* Prepare can automatically choose statement name based on sql. This makes it easier to explicitly manage prepared statements. +* Statement cache now uses deterministic, stable statement names. +* database/sql prepared statement names are deterministically generated. +* Fix: SendBatch wasn't respecting context cancellation. +* Fix: Timeout error from pipeline is now normalized. +* Fix: database/sql encoding json.RawMessage to []byte. +* CancelRequest: Wait for the cancel request to be acknowledged by the server. This should improve PgBouncer compatibility. (Anton Levakin) +* stdlib: Use Ping instead of CheckConn in ResetSession +* Add json.Marshaler and json.Unmarshaler for Float4, Float8 (Kirill Mironov) + +# 5.4.3 (August 5, 2023) + +* Fix: QCharArrayOID was defined with the wrong OID (Christoph Engelbert) +* Fix: connect_timeout for sslmode=allow|prefer (smaher-edb) +* Fix: pgxpool: background health check cannot overflow pool +* Fix: Check for nil in defer when sending batch (recover properly from panic) +* Fix: json scan of non-string pointer to pointer +* Fix: zeronull.Timestamptz should use pgtype.Timestamptz +* Fix: NewConnsCount was not correctly counting connections created by Acquire directly. (James Hartig) +* RowTo(AddrOf)StructByPos ignores fields with "-" db tag +* Optimization: improve text format numeric parsing (horpto) + +# 5.4.2 (July 11, 2023) + +* Fix: RowScanner errors are fatal to Rows +* Fix: Enable failover efforts when pg_hba.conf disallows non-ssl connections (Brandon Kauffman) +* Hstore text codec internal improvements (Evan Jones) +* Fix: Stop timers for background reader when not in use. Fixes memory leak when closing connections (Adrian-Stefan Mares) +* Fix: Stop background reader as soon as possible. +* Add PgConn.SyncConn(). This combined with the above fix makes it safe to directly use the underlying net.Conn. + +# 5.4.1 (June 18, 2023) + +* Fix: concurrency bug with pgtypeDefaultMap and simple protocol (Lev Zakharov) +* Add TxOptions.BeginQuery to allow overriding the default BEGIN query + +# 5.4.0 (June 14, 2023) + +* Replace platform specific syscalls for non-blocking IO with more traditional goroutines and deadlines. This returns to the v4 approach with some additional improvements and fixes. This restores the ability to use a pgx.Conn over an ssh.Conn as well as other non-TCP or Unix socket connections. In addition, it is a significantly simpler implementation that is less likely to have cross platform issues. +* Optimization: The default type registrations are now shared among all connections. This saves about 100KB of memory per connection. `pgtype.Type` and `pgtype.Codec` values are now required to be immutable after registration. This was already necessary in most cases but wasn't documented until now. (Lev Zakharov) +* Fix: Ensure pgxpool.Pool.QueryRow.Scan releases connection on panic +* CancelRequest: don't try to read the reply (Nicola Murino) +* Fix: correctly handle bool type aliases (Wichert Akkerman) +* Fix: pgconn.CancelRequest: Fix unix sockets: don't use RemoteAddr() +* Fix: pgx.Conn memory leak with prepared statement caching (Evan Jones) +* Add BeforeClose to pgxpool.Pool (Evan Cordell) +* Fix: various hstore fixes and optimizations (Evan Jones) +* Fix: RowToStructByPos with embedded unexported struct +* Support different bool string representations (Lev Zakharov) +* Fix: error when using BatchResults.Exec on a select that returns an error after some rows. +* Fix: pipelineBatchResults.Exec() not returning error from ResultReader +* Fix: pipeline batch results not closing pipeline when error occurs while reading directly from results instead of using + a callback. +* Fix: scanning a table type into a struct +* Fix: scan array of record to pointer to slice of struct +* Fix: handle null for json (Cemre Mengu) +* Batch Query callback is called even when there is an error +* Add RowTo(AddrOf)StructByNameLax (Audi P. Risa P) + +# 5.3.1 (February 27, 2023) + +* Fix: Support v4 and v5 stdlib in same program (Tomáš Procházka) +* Fix: sql.Scanner not being used in certain cases +* Add text format jsonpath support +* Fix: fake non-blocking read adaptive wait time + +# 5.3.0 (February 11, 2023) + +* Fix: json values work with sql.Scanner +* Fixed / improved error messages (Mark Chambers and Yevgeny Pats) +* Fix: support scan into single dimensional arrays +* Fix: MaxConnLifetimeJitter setting actually jitter (Ben Weintraub) +* Fix: driver.Value representation of bytea should be []byte not string +* Fix: better handling of unregistered OIDs +* CopyFrom can use query cache to avoid extra round trip to get OIDs (Alejandro Do Nascimento Mora) +* Fix: encode to json ignoring driver.Valuer +* Support sql.Scanner on renamed base type +* Fix: pgtype.Numeric text encoding of negative numbers (Mark Chambers) +* Fix: connect with multiple hostnames when one can't be resolved +* Upgrade puddle to remove dependency on uber/atomic and fix alignment issue on 32-bit platform +* Fix: scanning json column into **string +* Multiple reductions in memory allocations +* Fake non-blocking read adapts its max wait time +* Improve CopyFrom performance and reduce memory usage +* Fix: encode []any to array +* Fix: LoadType for composite with dropped attributes (Felix Röhrich) +* Support v4 and v5 stdlib in same program +* Fix: text format array decoding with string of "NULL" +* Prefer binary format for arrays + +# 5.2.0 (December 5, 2022) + +* `tracelog.TraceLog` implements the pgx.PrepareTracer interface. (Vitalii Solodilov) +* Optimize creating begin transaction SQL string (Petr Evdokimov and ksco) +* `Conn.LoadType` supports range and multirange types (Vitalii Solodilov) +* Fix scan `uint` and `uint64` `ScanNumeric`. This resolves a PostgreSQL `numeric` being incorrectly scanned into `uint` and `uint64`. + +# 5.1.1 (November 17, 2022) + +* Fix simple query sanitizer where query text contains a Unicode replacement character. +* Remove erroneous `name` argument from `DeallocateAll()`. Technically, this is a breaking change, but given that method was only added 5 days ago this change was accepted. (Bodo Kaiser) + +# 5.1.0 (November 12, 2022) + +* Update puddle to v2.1.2. This resolves a race condition and a deadlock in pgxpool. +* `QueryRewriter.RewriteQuery` now returns an error. Technically, this is a breaking change for any external implementers, but given the minimal likelihood that there are actually any external implementers this change was accepted. +* Expose `GetSSLPassword` support to pgx. +* Fix encode `ErrorResponse` unknown field handling. This would only affect pgproto3 being used directly as a proxy with a non-PostgreSQL server that included additional error fields. +* Fix date text format encoding with 5 digit years. +* Fix date values passed to a `sql.Scanner` as `string` instead of `time.Time`. +* DateCodec.DecodeValue can return `pgtype.InfinityModifier` instead of `string` for infinite values. This now matches the behavior of the timestamp types. +* Add domain type support to `Conn.LoadType()`. +* Add `RowToStructByName` and `RowToAddrOfStructByName`. (Pavlo Golub) +* Add `Conn.DeallocateAll()` to clear all prepared statements including the statement cache. (Bodo Kaiser) + +# 5.0.4 (October 24, 2022) + +* Fix: CollectOneRow prefers PostgreSQL error over pgx.ErrorNoRows +* Fix: some reflect Kind checks to first check for nil +* Bump golang.org/x/text dependency to placate snyk +* Fix: RowToStructByPos on structs with multiple anonymous sub-structs (Baptiste Fontaine) +* Fix: Exec checks if tx is closed + +# 5.0.3 (October 14, 2022) + +* Fix `driver.Valuer` handling edge cases that could cause infinite loop or crash + +# v5.0.2 (October 8, 2022) + +* Fix date encoding in text format to always use 2 digits for month and day +* Prefer driver.Valuer over wrap plans when encoding +* Fix scan to pointer to pointer to renamed type +* Allow scanning NULL even if PG and Go types are incompatible + +# v5.0.1 (September 24, 2022) + +* Fix 32-bit atomic usage +* Add MarshalJSON for Float8 (yogipristiawan) +* Add `[` and `]` to text encoding of `Lseg` +* Fix sqlScannerWrapper NULL handling + +# v5.0.0 (September 17, 2022) + +## Merged Packages + +`github.com/jackc/pgtype`, `github.com/jackc/pgconn`, and `github.com/jackc/pgproto3` are now included in the main +`github.com/jackc/pgx` repository. Previously there was confusion as to where issues should be reported, additional +release work due to releasing multiple packages, and less clear changelogs. + +## pgconn + +`CommandTag` is now an opaque type instead of directly exposing an underlying `[]byte`. + +The return value `ResultReader.Values()` is no longer safe to retain a reference to after a subsequent call to `NextRow()` or `Close()`. + +`Trace()` method adds low level message tracing similar to the `PQtrace` function in `libpq`. + +pgconn now uses non-blocking IO. This is a significant internal restructuring, but it should not cause any visible changes on its own. However, it is important in implementing other new features. + +`CheckConn()` checks a connection's liveness by doing a non-blocking read. This can be used to detect database restarts or network interruptions without executing a query or a ping. + +pgconn now supports pipeline mode. + +`*PgConn.ReceiveResults` removed. Use pipeline mode instead. + +`Timeout()` no longer considers `context.Canceled` as a timeout error. `context.DeadlineExceeded` still is considered a timeout error. + +## pgxpool + +`Connect` and `ConnectConfig` have been renamed to `New` and `NewWithConfig` respectively. The `LazyConnect` option has been removed. Pools always lazily connect. + +## pgtype + +The `pgtype` package has been significantly changed. + +### NULL Representation + +Previously, types had a `Status` field that could be `Undefined`, `Null`, or `Present`. This has been changed to a +`Valid` `bool` field to harmonize with how `database/sql` represents `NULL` and to make the zero value useable. + +Previously, a type that implemented `driver.Valuer` would have the `Value` method called even on a nil pointer. All nils +whether typed or untyped now represent `NULL`. + +### Codec and Value Split + +Previously, the type system combined decoding and encoding values with the value types. e.g. Type `Int8` both handled +encoding and decoding the PostgreSQL representation and acted as a value object. This caused some difficulties when +there was not an exact 1 to 1 relationship between the Go types and the PostgreSQL types For example, scanning a +PostgreSQL binary `numeric` into a Go `float64` was awkward (see https://github.com/jackc/pgtype/issues/147). This +concepts have been separated. A `Codec` only has responsibility for encoding and decoding values. Value types are +generally defined by implementing an interface that a particular `Codec` understands (e.g. `PointScanner` and +`PointValuer` for the PostgreSQL `point` type). + +### Array Types + +All array types are now handled by `ArrayCodec` instead of using code generation for each new array type. This also +means that less common array types such as `point[]` are now supported. `Array[T]` supports PostgreSQL multi-dimensional +arrays. + +### Composite Types + +Composite types must be registered before use. `CompositeFields` may still be used to construct and destruct composite +values, but any type may now implement `CompositeIndexGetter` and `CompositeIndexScanner` to be used as a composite. + +### Range Types + +Range types are now handled with types `RangeCodec` and `Range[T]`. This allows additional user defined range types to +easily be handled. Multirange types are handled similarly with `MultirangeCodec` and `Multirange[T]`. + +### pgxtype + +`LoadDataType` moved to `*Conn` as `LoadType`. + +### Bytea + +The `Bytea` and `GenericBinary` types have been replaced. Use the following instead: + +* `[]byte` - For normal usage directly use `[]byte`. +* `DriverBytes` - Uses driver memory only available until next database method call. Avoids a copy and an allocation. +* `PreallocBytes` - Uses preallocated byte slice to avoid an allocation. +* `UndecodedBytes` - Avoids any decoding. Allows working with raw bytes. + +### Dropped lib/pq Support + +`pgtype` previously supported and was tested against [lib/pq](https://github.com/lib/pq). While it will continue to work +in most cases this is no longer supported. + +### database/sql Scan + +Previously, most `Scan` implementations would convert `[]byte` to `string` automatically to decode a text value. Now +only `string` is handled. This is to allow the possibility of future binary support in `database/sql` mode by +considering `[]byte` to be binary format and `string` text format. This change should have no effect for any use with +`pgx`. The previous behavior was only necessary for `lib/pq` compatibility. + +Added `*Map.SQLScanner` to create a `sql.Scanner` for types such as `[]int32` and `Range[T]` that do not implement +`sql.Scanner` directly. + +### Number Type Fields Include Bit size + +`Int2`, `Int4`, `Int8`, `Float4`, `Float8`, and `Uint32` fields now include bit size. e.g. `Int` is renamed to `Int64`. +This matches the convention set by `database/sql`. In addition, for comparable types like `pgtype.Int8` and +`sql.NullInt64` the structures are identical. This means they can be directly converted one to another. + +### 3rd Party Type Integrations + +* Extracted integrations with https://github.com/shopspring/decimal and https://github.com/gofrs/uuid to + https://github.com/jackc/pgx-shopspring-decimal and https://github.com/jackc/pgx-gofrs-uuid respectively. This trims + the pgx dependency tree. + +### Other Changes + +* `Bit` and `Varbit` are both replaced by the `Bits` type. +* `CID`, `OID`, `OIDValue`, and `XID` are replaced by the `Uint32` type. +* `Hstore` is now defined as `map[string]*string`. +* `JSON` and `JSONB` types removed. Use `[]byte` or `string` directly. +* `QChar` type removed. Use `rune` or `byte` directly. +* `Inet` and `Cidr` types removed. Use `netip.Addr` and `netip.Prefix` directly. These types are more memory efficient than the previous `net.IPNet`. +* `Macaddr` type removed. Use `net.HardwareAddr` directly. +* Renamed `pgtype.ConnInfo` to `pgtype.Map`. +* Renamed `pgtype.DataType` to `pgtype.Type`. +* Renamed `pgtype.None` to `pgtype.Finite`. +* `RegisterType` now accepts a `*Type` instead of `Type`. +* Assorted array helper methods and types made private. + +## stdlib + +* Removed `AcquireConn` and `ReleaseConn` as that functionality has been built in since Go 1.13. + +## Reduced Memory Usage by Reusing Read Buffers + +Previously, the connection read buffer would allocate large chunks of memory and never reuse them. This allowed +transferring ownership to anything such as scanned values without incurring an additional allocation and memory copy. +However, this came at the cost of overall increased memory allocation size. But worse it was also possible to pin large +chunks of memory by retaining a reference to a small value that originally came directly from the read buffer. Now +ownership remains with the read buffer and anything needing to retain a value must make a copy. + +## Query Execution Modes + +Control over automatic prepared statement caching and simple protocol use are now combined into query execution mode. +See documentation for `QueryExecMode`. + +## QueryRewriter Interface and NamedArgs + +pgx now supports named arguments with the `NamedArgs` type. This is implemented via the new `QueryRewriter` interface which +allows arbitrary rewriting of query SQL and arguments. + +## RowScanner Interface + +The `RowScanner` interface allows a single argument to Rows.Scan to scan the entire row. + +## Rows Result Helpers + +* `CollectRows` and `RowTo*` functions simplify collecting results into a slice. +* `CollectOneRow` collects one row using `RowTo*` functions. +* `ForEachRow` simplifies scanning each row and executing code using the scanned values. `ForEachRow` replaces `QueryFunc`. + +## Tx Helpers + +Rather than every type that implemented `Begin` or `BeginTx` methods also needing to implement `BeginFunc` and +`BeginTxFunc` these methods have been converted to functions that take a db that implements `Begin` or `BeginTx`. + +## Improved Batch Query Ergonomics + +Previously, the code for building a batch went in one place before the call to `SendBatch`, and the code for reading the +results went in one place after the call to `SendBatch`. This could make it difficult to match up the query and the code +to handle the results. Now `Queue` returns a `QueuedQuery` which has methods `Query`, `QueryRow`, and `Exec` which can +be used to register a callback function that will handle the result. Callback functions are called automatically when +`BatchResults.Close` is called. + +## SendBatch Uses Pipeline Mode When Appropriate + +Previously, a batch with 10 unique parameterized statements executed 100 times would entail 11 network round trips. 1 +for each prepare / describe and 1 for executing them all. Now pipeline mode is used to prepare / describe all statements +in a single network round trip. So it would only take 2 round trips. + +## Tracing and Logging + +Internal logging support has been replaced with tracing hooks. This allows custom tracing integration with tools like OpenTelemetry. Package tracelog provides an adapter for pgx v4 loggers to act as a tracer. + +All integrations with 3rd party loggers have been extracted to separate repositories. This trims the pgx dependency +tree. diff --git a/vendor/github.com/jackc/pgx/v5/CLAUDE.md b/vendor/github.com/jackc/pgx/v5/CLAUDE.md new file mode 100644 index 0000000000..71a8fc164d --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/CLAUDE.md @@ -0,0 +1,73 @@ +# CLAUDE.md + +This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository. + +## Project Overview + +pgx is a PostgreSQL driver and toolkit for Go (`github.com/jackc/pgx/v5`). It provides both a native PostgreSQL interface and a `database/sql` compatible driver. Requires Go 1.25+ and supports PostgreSQL 14+ and CockroachDB. + +## Build & Test Commands + +```bash +# Run all tests (requires PGX_TEST_DATABASE to be set) +go test ./... + +# Run a specific test +go test -run TestFunctionName ./... + +# Run tests for a specific package +go test ./pgconn/... + +# Run tests with race detector +go test -race ./... + +# DevContainer: run tests against specific PostgreSQL versions +./test.sh pg18 # Default: PostgreSQL 18 +./test.sh pg16 -run TestConnect # Specific test against PG16 +./test.sh crdb # CockroachDB +./test.sh all # All targets (pg14-18 + crdb) + +# Format (always run after making changes) +goimports -w . + +# Lint +golangci-lint run ./... +``` + +## Test Database Setup + +Tests require `PGX_TEST_DATABASE` environment variable. In the devcontainer, `test.sh` handles this. For local development: + +```bash +export PGX_TEST_DATABASE="host=localhost user=postgres password=postgres dbname=pgx_test" +``` + +The test database needs extensions: `hstore`, `ltree`, and a `uint64` domain. See `testsetup/postgresql_setup.sql` for full setup. Many tests are skipped unless additional `PGX_TEST_*` env vars are set (for TLS, SCRAM, MD5, unix socket, PgBouncer testing). + +## Architecture + +The codebase is a layered architecture, bottom-up: + +- **pgproto3/** — PostgreSQL wire protocol v3 encoder/decoder. Defines `FrontendMessage` and `BackendMessage` types for every protocol message. +- **pgconn/** — Low-level connection layer (roughly libpq-equivalent). Handles authentication, TLS, query execution, COPY protocol, and notifications. `PgConn` is the core type. +- **pgx** (root package) — High-level query interface built on `pgconn`. Provides `Conn`, `Rows`, `Tx`, `Batch`, `CopyFrom`, and generic helpers like `CollectRows`/`ForEachRow`. Includes automatic statement caching (LRU). +- **pgtype/** — Type system mapping between Go and PostgreSQL types (70+ types). Key interfaces: `Codec`, `Type`, `TypeMap`. Custom types (enums, composites, domains) are registered through `TypeMap`. +- **pgxpool/** — Concurrency-safe connection pool built on `puddle/v2`. `Pool` is the main type; wraps `pgx.Conn`. +- **stdlib/** — `database/sql` compatibility adapter. + +Supporting packages: +- **internal/stmtcache/** — Prepared statement cache with LRU eviction +- **internal/sanitize/** — SQL query sanitization +- **tracelog/** — Logging adapter that implements tracer interfaces +- **multitracer/** — Composes multiple tracers into one +- **pgxtest/** — Test helpers for running tests across connection types + +## Key Design Conventions + +- **Semantic versioning** — strictly followed. Do not break the public API (no removing or renaming exported types, functions, methods, or fields; no changing function signatures). +- **Minimal dependencies** — adding new dependencies is strongly discouraged (see CONTRIBUTING.md). +- **Context-based** — all blocking operations take `context.Context`. +- **Tracer interfaces** — observability via `QueryTracer`, `BatchTracer`, `CopyFromTracer`, `PrepareTracer` on `ConnConfig.Tracer`. +- **Formatting** — always run `goimports -w .` after making changes to ensure code is properly formatted. CI checks formatting via `gofmt -l -s -w . && git diff --exit-code`. `gofumpt` with extra rules is also enforced via `golangci-lint`. +- **Linters** — `govet`, `ineffassign`, and `unconvert` only (configured in `.golangci.yml`). +- **CI matrix** — tests run against Go 1.25/1.26 × PostgreSQL 14-18 + CockroachDB, on Linux and Windows. Race detector enabled on Linux only. diff --git a/vendor/github.com/jackc/pgx/v5/CONTRIBUTING.md b/vendor/github.com/jackc/pgx/v5/CONTRIBUTING.md new file mode 100644 index 0000000000..2283ae6700 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/CONTRIBUTING.md @@ -0,0 +1,139 @@ +# Contributing + +## Discuss Significant Changes + +Before you invest a significant amount of time on a change, please create a discussion or issue describing your +proposal. This will help to ensure your proposed change has a reasonable chance of being merged. + +## Avoid Dependencies + +Adding a dependency is a big deal. While on occasion a new dependency may be accepted, the default answer to any change +that adds a dependency is no. + +## AI + +Using AI is acceptable (not that it can really be stopped) under one the following conditions. + +* AI was used, but you deeply understand the code and you can answer questions regarding your change. You are not going + to answer questions with "I don't know", AI did it. You are not going to "answer" questions by relaying them to your + agent. This is wasteful of the code reviewer's time. +* AI was used to solve a problem without your deep understanding. This can still be a good starting point for a fix or + feature. But you need to clearly state that this is an AI proposal. You should include additional information such as + the AI used and what prompts were used. You should also be aware that large, complicated, or subtle changes may be + rejected simply because the reviewer is not confident in a change that no human understands. + +## Development Environment Setup + +pgx tests naturally require a PostgreSQL database. It will connect to the database specified in the `PGX_TEST_DATABASE` +environment variable. The `PGX_TEST_DATABASE` environment variable can either be a URL or key-value pairs. In addition, +the standard `PG*` environment variables will be respected. Consider using [direnv](https://github.com/direnv/direnv) to +simplify environment variable handling. + +### Devcontainer + +The easiest way to start development is with the included devcontainer. It includes containers for each supported +PostgreSQL version as well as CockroachDB. `./test.sh all` will run the tests against all database types. + +### Using an Existing PostgreSQL Cluster Outside of a Devcontainer + +If you already have a PostgreSQL development server this is the quickest way to start and run the majority of the pgx +test suite. Some tests will be skipped that require server configuration changes (e.g. those testing different +authentication methods). + +Create and setup a test database: + +``` +export PGDATABASE=pgx_test +createdb +psql -c 'create extension hstore;' +psql -c 'create extension ltree;' +psql -c 'create domain uint64 as numeric(20,0);' +``` + +Ensure a `postgres` user exists. This happens by default in normal PostgreSQL installs, but some installation methods +such as Homebrew do not. + +``` +createuser -s postgres +``` + +Ensure your `PGX_TEST_DATABASE` environment variable points to the database you just created and run the tests. + +``` +export PGX_TEST_DATABASE="host=/private/tmp database=pgx_test" +go test ./... +``` + +This will run the vast majority of the tests, but some tests will be skipped (e.g. those testing different connection methods). + +### Creating a New PostgreSQL Cluster Exclusively for Testing Outside of a Devcontainer + +The following environment variables need to be set both for initial setup and whenever the tests are run. (direnv is +highly recommended). Depending on your platform, you may need to change the host for `PGX_TEST_UNIX_SOCKET_CONN_STRING`. + +``` +export PGPORT=5015 +export PGUSER=postgres +export PGDATABASE=pgx_test +export POSTGRESQL_DATA_DIR=postgresql + +export PGX_TEST_DATABASE="host=127.0.0.1 database=pgx_test user=pgx_md5 password=secret" +export PGX_TEST_UNIX_SOCKET_CONN_STRING="host=/private/tmp database=pgx_test" +export PGX_TEST_TCP_CONN_STRING="host=127.0.0.1 database=pgx_test user=pgx_md5 password=secret" +export PGX_TEST_SCRAM_PASSWORD_CONN_STRING="host=127.0.0.1 user=pgx_scram password=secret database=pgx_test channel_binding=disable" +export PGX_TEST_SCRAM_PLUS_CONN_STRING="host=localhost user=pgx_ssl password=secret sslmode=verify-full sslrootcert=`pwd`/.testdb/ca.pem database=pgx_test channel_binding=require" +export PGX_TEST_MD5_PASSWORD_CONN_STRING="host=127.0.0.1 database=pgx_test user=pgx_md5 password=secret" +export PGX_TEST_PLAIN_PASSWORD_CONN_STRING="host=127.0.0.1 user=pgx_pw password=secret" +export PGX_TEST_TLS_CONN_STRING="host=localhost user=pgx_ssl password=secret sslmode=verify-full sslrootcert=`pwd`/.testdb/ca.pem channel_binding=disable" +export PGX_SSL_PASSWORD=certpw +export PGX_TEST_TLS_CLIENT_CONN_STRING="host=localhost user=pgx_sslcert sslmode=verify-full sslrootcert=`pwd`/.testdb/ca.pem database=pgx_test sslcert=`pwd`/.testdb/pgx_sslcert.crt sslkey=`pwd`/.testdb/pgx_sslcert.key" +``` + +Create a new database cluster. + +``` +initdb --locale=en_US -E UTF-8 --username=postgres .testdb/$POSTGRESQL_DATA_DIR + +echo "listen_addresses = '127.0.0.1'" >> .testdb/$POSTGRESQL_DATA_DIR/postgresql.conf +echo "port = $PGPORT" >> .testdb/$POSTGRESQL_DATA_DIR/postgresql.conf +cat testsetup/postgresql_ssl.conf >> .testdb/$POSTGRESQL_DATA_DIR/postgresql.conf +cp testsetup/pg_hba.conf .testdb/$POSTGRESQL_DATA_DIR/pg_hba.conf + +cd .testdb + +# Generate CA, server, and encrypted client certificates. +go run ../testsetup/generate_certs.go + +# Copy certificates to server directory and set permissions. +cp ca.pem $POSTGRESQL_DATA_DIR/root.crt +cp localhost.key $POSTGRESQL_DATA_DIR/server.key +chmod 600 $POSTGRESQL_DATA_DIR/server.key +cp localhost.crt $POSTGRESQL_DATA_DIR/server.crt + +cd .. +``` + + +Start the new cluster. This will be necessary whenever you are running pgx tests. + +``` +postgres -D .testdb/$POSTGRESQL_DATA_DIR +``` + +Setup the test database in the new cluster. + +``` +createdb +psql --no-psqlrc -f testsetup/postgresql_setup.sql +``` + +### PgBouncer + +There are tests specific for PgBouncer that will be executed if `PGX_TEST_PGBOUNCER_CONN_STRING` is set. + +### Optional Tests + +pgx supports multiple connection types and means of authentication. These tests are optional. They will only run if the +appropriate environment variables are set. In addition, there may be tests specific to particular PostgreSQL versions, +non-PostgreSQL servers (e.g. CockroachDB), or connection poolers (e.g. PgBouncer). `go test ./... -v | grep SKIP` to see +if any tests are being skipped. diff --git a/vendor/github.com/jackc/pgx/v5/LICENSE b/vendor/github.com/jackc/pgx/v5/LICENSE new file mode 100644 index 0000000000..5c486c39a2 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/LICENSE @@ -0,0 +1,22 @@ +Copyright (c) 2013-2021 Jack Christensen + +MIT License + +Permission is hereby granted, free of charge, to any person obtaining +a copy of this software and associated documentation files (the +"Software"), to deal in the Software without restriction, including +without limitation the rights to use, copy, modify, merge, publish, +distribute, sublicense, and/or sell copies of the Software, and to +permit persons to whom the Software is furnished to do so, subject to +the following conditions: + +The above copyright notice and this permission notice shall be +included in all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, +EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF +MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND +NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE +LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION +OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION +WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. diff --git a/vendor/github.com/jackc/pgx/v5/README.md b/vendor/github.com/jackc/pgx/v5/README.md new file mode 100644 index 0000000000..aa35e4a3d6 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/README.md @@ -0,0 +1,192 @@ +[![Go Reference](https://pkg.go.dev/badge/github.com/jackc/pgx/v5.svg)](https://pkg.go.dev/github.com/jackc/pgx/v5) +[![Build Status](https://github.com/jackc/pgx/actions/workflows/ci.yml/badge.svg)](https://github.com/jackc/pgx/actions/workflows/ci.yml) + +# pgx - PostgreSQL Driver and Toolkit + +pgx is a pure Go driver and toolkit for PostgreSQL. + +The pgx driver is a low-level, high performance interface that exposes PostgreSQL-specific features such as `LISTEN` / +`NOTIFY` and `COPY`. It also includes an adapter for the standard `database/sql` interface. + +The toolkit component is a related set of packages that implement PostgreSQL functionality such as parsing the wire protocol +and type mapping between PostgreSQL and Go. These underlying packages can be used to implement alternative drivers, +proxies, load balancers, logical replication clients, etc. + +## Example Usage + +```go +package main + +import ( + "context" + "fmt" + "os" + + "github.com/jackc/pgx/v5" +) + +func main() { + // urlExample := "postgres://username:password@localhost:5432/database_name" + conn, err := pgx.Connect(context.Background(), os.Getenv("DATABASE_URL")) + if err != nil { + fmt.Fprintf(os.Stderr, "Unable to connect to database: %v\n", err) + os.Exit(1) + } + defer conn.Close(context.Background()) + + var name string + var weight int64 + err = conn.QueryRow(context.Background(), "select name, weight from widgets where id=$1", 42).Scan(&name, &weight) + if err != nil { + fmt.Fprintf(os.Stderr, "QueryRow failed: %v\n", err) + os.Exit(1) + } + + fmt.Println(name, weight) +} +``` + +See the [getting started guide](https://github.com/jackc/pgx/wiki/Getting-started-with-pgx) for more information. + +## Features + +* Support for approximately 70 different PostgreSQL types +* Automatic statement preparation and caching +* Batch queries +* Single-round trip query mode +* Full TLS connection control +* Binary format support for custom types (allows for much quicker encoding/decoding) +* `COPY` protocol support for faster bulk data loads +* Tracing and logging support +* Connection pool with after-connect hook for arbitrary connection setup +* `LISTEN` / `NOTIFY` +* Conversion of PostgreSQL arrays to Go slice mappings for integers, floats, and strings +* `hstore` support +* `json` and `jsonb` support +* Maps `inet` and `cidr` PostgreSQL types to `netip.Addr` and `netip.Prefix` +* Large object support +* NULL mapping to pointer to pointer +* Supports `database/sql.Scanner` and `database/sql/driver.Valuer` interfaces for custom types +* Notice response handling +* Simulated nested transactions with savepoints + +## Choosing Between the pgx and database/sql Interfaces + +The pgx interface is faster. Many PostgreSQL specific features such as `LISTEN` / `NOTIFY` and `COPY` are not available +through the `database/sql` interface. + +The pgx interface is recommended when: + +1. The application only targets PostgreSQL. +2. No other libraries that require `database/sql` are in use. + +It is also possible to use the `database/sql` interface and convert a connection to the lower-level pgx interface as needed. + +## Testing + +See [CONTRIBUTING.md](./CONTRIBUTING.md) for setup instructions. + +## Architecture + +See the presentation at Golang Estonia, [PGX Top to Bottom](https://www.youtube.com/watch?v=sXMSWhcHCf8) for a description of pgx architecture. + +## Supported Go and PostgreSQL Versions + +pgx supports the same versions of Go and PostgreSQL that are supported by their respective teams. For [Go](https://golang.org/doc/devel/release.html#policy) that is the two most recent major releases and for [PostgreSQL](https://www.postgresql.org/support/versioning/) the major releases in the last 5 years. This means pgx supports Go 1.25 and higher and PostgreSQL 14 and higher. pgx also is tested against the latest version of [CockroachDB](https://www.cockroachlabs.com/product/). + +## Version Policy + +pgx follows semantic versioning for the documented public API on stable releases. `v5` is the latest stable major version. + +## PGX Family Libraries + +### [github.com/jackc/pglogrepl](https://github.com/jackc/pglogrepl) + +pglogrepl provides functionality to act as a client for PostgreSQL logical replication. + +### [github.com/jackc/pgmock](https://github.com/jackc/pgmock) + +pgmock offers the ability to create a server that mocks the PostgreSQL wire protocol. This is used internally to test pgx by purposely inducing unusual errors. pgproto3 and pgmock together provide most of the foundational tooling required to implement a PostgreSQL proxy or MitM (such as for a custom connection pooler). + +### [github.com/jackc/tern](https://github.com/jackc/tern) + +tern is a stand-alone SQL migration system. + +### [github.com/jackc/pgerrcode](https://github.com/jackc/pgerrcode) + +pgerrcode contains constants for the PostgreSQL error codes. + +## Adapters for 3rd Party Types + +* [github.com/jackc/pgx-gofrs-uuid](https://github.com/jackc/pgx-gofrs-uuid) +* [github.com/jackc/pgx-shopspring-decimal](https://github.com/jackc/pgx-shopspring-decimal) +* [github.com/ColeBurch/pgx-govalues-decimal](https://github.com/ColeBurch/pgx-govalues-decimal) +* [github.com/twpayne/pgx-geos](https://github.com/twpayne/pgx-geos) ([PostGIS](https://postgis.net/) and [GEOS](https://libgeos.org/) via [go-geos](https://github.com/twpayne/go-geos)) +* [github.com/vgarvardt/pgx-google-uuid](https://github.com/vgarvardt/pgx-google-uuid) + + +## Adapters for 3rd Party Tracers + +* [github.com/jackhopner/pgx-xray-tracer](https://github.com/jackhopner/pgx-xray-tracer) +* [github.com/exaring/otelpgx](https://github.com/exaring/otelpgx) + +## Adapters for 3rd Party Loggers + +These adapters can be used with the tracelog package. + +* [github.com/jackc/pgx-go-kit-log](https://github.com/jackc/pgx-go-kit-log) +* [github.com/jackc/pgx-log15](https://github.com/jackc/pgx-log15) +* [github.com/jackc/pgx-logrus](https://github.com/jackc/pgx-logrus) +* [github.com/jackc/pgx-zap](https://github.com/jackc/pgx-zap) +* [github.com/jackc/pgx-zerolog](https://github.com/jackc/pgx-zerolog) +* [github.com/mcosta74/pgx-slog](https://github.com/mcosta74/pgx-slog) +* [github.com/kataras/pgx-golog](https://github.com/kataras/pgx-golog) + +## 3rd Party Libraries with PGX Support + +### [github.com/pashagolub/pgxmock](https://github.com/pashagolub/pgxmock) + +pgxmock is a mock library implementing pgx interfaces. +pgxmock has one and only purpose - to simulate pgx behavior in tests, without needing a real database connection. + +### [github.com/georgysavva/scany](https://github.com/georgysavva/scany) + +Library for scanning data from a database into Go structs and more. + +### [github.com/vingarcia/ksql](https://github.com/vingarcia/ksql) + +A carefully designed SQL client for making using SQL easier, +more productive, and less error-prone on Golang. + +### [github.com/otan/gopgkrb5](https://github.com/otan/gopgkrb5) + +Adds GSSAPI / Kerberos authentication support. + +### [github.com/wcamarao/pmx](https://github.com/wcamarao/pmx) + +Explicit data mapping and scanning library for Go structs and slices. + +### [github.com/stephenafamo/scan](https://github.com/stephenafamo/scan) + +Type safe and flexible package for scanning database data into Go types. +Supports, structs, maps, slices and custom mapping functions. + +### [github.com/z0ne-dev/mgx](https://github.com/z0ne-dev/mgx) + +Code first migration library for native pgx (no database/sql abstraction). + +### [github.com/amirsalarsafaei/sqlc-pgx-monitoring](https://github.com/amirsalarsafaei/sqlc-pgx-monitoring) + +A database monitoring/metrics library for pgx and sqlc. Trace, log and monitor your sqlc query performance using OpenTelemetry. + +### [https://github.com/nikolayk812/pgx-outbox](https://github.com/nikolayk812/pgx-outbox) + +Simple Golang implementation for transactional outbox pattern for PostgreSQL using jackc/pgx driver. + +### [https://github.com/Arlandaren/pgxWrappy](https://github.com/Arlandaren/pgxWrappy) + +Simplifies working with the pgx library, providing convenient scanning of nested structures. + +### [https://github.com/KoNekoD/pgx-colon-query-rewriter](https://github.com/KoNekoD/pgx-colon-query-rewriter) + +Implementation of the pgx query rewriter to use ':' instead of '@' in named query parameters. diff --git a/vendor/github.com/jackc/pgx/v5/Rakefile b/vendor/github.com/jackc/pgx/v5/Rakefile new file mode 100644 index 0000000000..3e3aa5030d --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/Rakefile @@ -0,0 +1,18 @@ +require "erb" + +rule '.go' => '.go.erb' do |task| + erb = ERB.new(File.read(task.source)) + File.write(task.name, "// Code generated from #{task.source}. DO NOT EDIT.\n\n" + erb.result(binding)) + sh "goimports", "-w", task.name +end + +generated_code_files = [ + "pgtype/int.go", + "pgtype/int_test.go", + "pgtype/integration_benchmark_test.go", + "pgtype/zeronull/int.go", + "pgtype/zeronull/int_test.go" +] + +desc "Generate code" +task generate: generated_code_files diff --git a/vendor/github.com/jackc/pgx/v5/batch.go b/vendor/github.com/jackc/pgx/v5/batch.go new file mode 100644 index 0000000000..805cc39efa --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/batch.go @@ -0,0 +1,537 @@ +package pgx + +import ( + "context" + "errors" + "fmt" + + "github.com/jackc/pgx/v5/pgconn" +) + +// QueuedQuery is a query that has been queued for execution via a [Batch]. +type QueuedQuery struct { + SQL string + Arguments []any + Fn batchItemFunc + sd *pgconn.StatementDescription +} + +type batchItemFunc func(br BatchResults) error + +// Query sets fn to be called when the response to qq is received. +func (qq *QueuedQuery) Query(fn func(rows Rows) error) { + qq.Fn = func(br BatchResults) error { + rows, _ := br.Query() + defer rows.Close() + + err := fn(rows) + if err != nil { + return err + } + rows.Close() + + return rows.Err() + } +} + +// Query sets fn to be called when the response to qq is received. +func (qq *QueuedQuery) QueryRow(fn func(row Row) error) { + qq.Fn = func(br BatchResults) error { + row := br.QueryRow() + return fn(row) + } +} + +// Exec sets fn to be called when the response to qq is received. +// +// Note: for simple batch insert uses where it is not required to handle +// each potential error individually, it's sufficient to not set any callbacks, +// and just handle the return value of [BatchResults.Close]. +func (qq *QueuedQuery) Exec(fn func(ct pgconn.CommandTag) error) { + qq.Fn = func(br BatchResults) error { + ct, err := br.Exec() + if err != nil { + return err + } + + return fn(ct) + } +} + +// Batch queries are a way of bundling multiple queries together to avoid +// unnecessary network round trips. A Batch must only be sent once. +type Batch struct { + QueuedQueries []*QueuedQuery +} + +// Queue queues a query to batch b. query can be an SQL query or the name of a prepared statement. The only pgx option +// argument that is supported is [QueryRewriter]. Queries are executed using the connection's DefaultQueryExecMode +// (see [ConnConfig.DefaultQueryExecMode]). +// +// While query can contain multiple statements if the connection's DefaultQueryExecMode is [QueryExecModeSimpleProtocol], +// this should be avoided. QueuedQuery.Fn must not be set as it will only be called for the first query. That is, +// [QueuedQuery.Query], [QueuedQuery.QueryRow], and [QueuedQuery.Exec] must not be called. In addition, any error +// messages or tracing that include the current query may reference the wrong query. +func (b *Batch) Queue(query string, arguments ...any) *QueuedQuery { + qq := &QueuedQuery{ + SQL: query, + Arguments: arguments, + } + b.QueuedQueries = append(b.QueuedQueries, qq) + return qq +} + +// Len returns number of queries that have been queued so far. +func (b *Batch) Len() int { + return len(b.QueuedQueries) +} + +type BatchResults interface { + // Exec reads the results from the next query in the batch as if the query has been sent with [Conn.Exec]. Prefer + // calling Exec on the QueuedQuery, or just calling Close. + Exec() (pgconn.CommandTag, error) + + // Query reads the results from the next query in the batch as if the query has been sent with [Conn.Query]. Prefer + // calling [QueuedQuery.Query]. + Query() (Rows, error) + + // QueryRow reads the results from the next query in the batch as if the query has been sent with [Conn.QueryRow]. + // Prefer calling [QueuedQuery.QueryRow]. + QueryRow() Row + + // Close closes the batch operation. All unread results are read and any callback functions registered with + // [QueuedQuery.Query], [QueuedQuery.QueryRow], or [QueuedQuery.Exec] will be called. If a callback function returns an + // error or the batch encounters an error subsequent callback functions will not be called. + // + // For simple batch inserts inside a transaction or similar queries, it's sufficient to not set any callbacks, + // and just handle the return value of Close. + // + // Close must be called before the underlying connection can be used again. Any error that occurred during a batch + // operation may have made it impossible to resyncronize the connection with the server. In this case the underlying + // connection will have been closed. + // + // Close is safe to call multiple times. If it returns an error subsequent calls will return the same error. Callback + // functions will not be rerun. + Close() error +} + +type batchResults struct { + ctx context.Context + conn *Conn + mrr *pgconn.MultiResultReader + err error + b *Batch + qqIdx int + closed bool + endTraced bool +} + +// Exec reads the results from the next query in the batch as if the query has been sent with Exec. +func (br *batchResults) Exec() (pgconn.CommandTag, error) { + if br.err != nil { + return pgconn.CommandTag{}, br.err + } + if br.closed { + return pgconn.CommandTag{}, fmt.Errorf("batch already closed") + } + + query, arguments, _ := br.nextQueryAndArgs() + + if !br.mrr.NextResult() { + err := br.mrr.Close() + if err == nil { + err = errors.New("no more results in batch") + } + if br.conn.batchTracer != nil { + br.conn.batchTracer.TraceBatchQuery(br.ctx, br.conn, TraceBatchQueryData{ + SQL: query, + Args: arguments, + Err: err, + }) + } + return pgconn.CommandTag{}, err + } + + commandTag, err := br.mrr.ResultReader().Close() + if err != nil { + br.err = err + br.mrr.Close() + } + + if br.conn.batchTracer != nil { + br.conn.batchTracer.TraceBatchQuery(br.ctx, br.conn, TraceBatchQueryData{ + SQL: query, + Args: arguments, + CommandTag: commandTag, + Err: br.err, + }) + } + + return commandTag, br.err +} + +// Query reads the results from the next query in the batch as if the query has been sent with Query. +func (br *batchResults) Query() (Rows, error) { + query, arguments, ok := br.nextQueryAndArgs() + if !ok { + query = "batch query" + } + + if br.err != nil { + return &baseRows{err: br.err, closed: true}, br.err + } + + if br.closed { + alreadyClosedErr := fmt.Errorf("batch already closed") + return &baseRows{err: alreadyClosedErr, closed: true}, alreadyClosedErr + } + + rows := br.conn.getRows(br.ctx, query, arguments) + rows.batchTracer = br.conn.batchTracer + + if !br.mrr.NextResult() { + rows.err = br.mrr.Close() + if rows.err == nil { + rows.err = errors.New("no more results in batch") + } + rows.closed = true + + if br.conn.batchTracer != nil { + br.conn.batchTracer.TraceBatchQuery(br.ctx, br.conn, TraceBatchQueryData{ + SQL: query, + Args: arguments, + Err: rows.err, + }) + } + + return rows, rows.err + } + + rows.resultReader = br.mrr.ResultReader() + return rows, nil +} + +// QueryRow reads the results from the next query in the batch as if the query has been sent with QueryRow. +func (br *batchResults) QueryRow() Row { + rows, _ := br.Query() + return (*connRow)(rows.(*baseRows)) +} + +// Close closes the batch operation. Any error that occurred during a batch operation may have made it impossible to +// resyncronize the connection with the server. In this case the underlying connection will have been closed. +func (br *batchResults) Close() error { + defer func() { + if !br.endTraced { + if br.conn != nil && br.conn.batchTracer != nil { + br.conn.batchTracer.TraceBatchEnd(br.ctx, br.conn, TraceBatchEndData{Err: br.err}) + } + br.endTraced = true + } + + invalidateCachesOnBatchResultsError(br.conn, br.b, br.err) + }() + + if br.err != nil { + return br.err + } + + if br.closed { + return nil + } + + // Read and run fn for all remaining items + for br.err == nil && !br.closed && br.b != nil && br.qqIdx < len(br.b.QueuedQueries) { + if br.b.QueuedQueries[br.qqIdx].Fn != nil { + err := br.b.QueuedQueries[br.qqIdx].Fn(br) + if err != nil { + br.err = err + } + } else { + br.Exec() + } + } + + br.closed = true + + err := br.mrr.Close() + if br.err == nil { + br.err = err + } + + return br.err +} + +func (br *batchResults) earlyError() error { + return br.err +} + +func (br *batchResults) nextQueryAndArgs() (query string, args []any, ok bool) { + if br.b != nil && br.qqIdx < len(br.b.QueuedQueries) { + bi := br.b.QueuedQueries[br.qqIdx] + query = bi.SQL + args = bi.Arguments + ok = true + br.qqIdx++ + } + return query, args, ok +} + +type pipelineBatchResults struct { + ctx context.Context + conn *Conn + pipeline *pgconn.Pipeline + lastRows *baseRows + err error + b *Batch + qqIdx int + closed bool + endTraced bool +} + +// Exec reads the results from the next query in the batch as if the query has been sent with Exec. +func (br *pipelineBatchResults) Exec() (pgconn.CommandTag, error) { + if br.err != nil { + return pgconn.CommandTag{}, br.err + } + if br.closed { + return pgconn.CommandTag{}, fmt.Errorf("batch already closed") + } + if br.lastRows != nil && br.lastRows.err != nil { + br.err = br.lastRows.err + return pgconn.CommandTag{}, br.err + } + + query, arguments, err := br.nextQueryAndArgs() + if err != nil { + return pgconn.CommandTag{}, err + } + + results, err := br.pipeline.GetResults() + if err != nil { + br.err = err + return pgconn.CommandTag{}, br.err + } + var commandTag pgconn.CommandTag + switch results := results.(type) { + case *pgconn.ResultReader: + commandTag, br.err = results.Close() + default: + return pgconn.CommandTag{}, fmt.Errorf("unexpected pipeline result: %T", results) + } + + if br.conn.batchTracer != nil { + br.conn.batchTracer.TraceBatchQuery(br.ctx, br.conn, TraceBatchQueryData{ + SQL: query, + Args: arguments, + CommandTag: commandTag, + Err: br.err, + }) + } + + return commandTag, br.err +} + +// Query reads the results from the next query in the batch as if the query has been sent with Query. +func (br *pipelineBatchResults) Query() (Rows, error) { + if br.err != nil { + return &baseRows{err: br.err, closed: true}, br.err + } + + if br.closed { + alreadyClosedErr := fmt.Errorf("batch already closed") + return &baseRows{err: alreadyClosedErr, closed: true}, alreadyClosedErr + } + + if br.lastRows != nil && br.lastRows.err != nil { + br.err = br.lastRows.err + return &baseRows{err: br.err, closed: true}, br.err + } + + query, arguments, err := br.nextQueryAndArgs() + if err != nil { + return &baseRows{err: err, closed: true}, err + } + + rows := br.conn.getRows(br.ctx, query, arguments) + rows.batchTracer = br.conn.batchTracer + br.lastRows = rows + + results, err := br.pipeline.GetResults() + if err != nil { + br.err = err + rows.err = err + rows.closed = true + + if br.conn.batchTracer != nil { + br.conn.batchTracer.TraceBatchQuery(br.ctx, br.conn, TraceBatchQueryData{ + SQL: query, + Args: arguments, + Err: err, + }) + } + } else { + switch results := results.(type) { + case *pgconn.ResultReader: + rows.resultReader = results + default: + err = fmt.Errorf("unexpected pipeline result: %T", results) + br.err = err + rows.err = err + rows.closed = true + } + } + + return rows, rows.err +} + +// QueryRow reads the results from the next query in the batch as if the query has been sent with QueryRow. +func (br *pipelineBatchResults) QueryRow() Row { + rows, _ := br.Query() + return (*connRow)(rows.(*baseRows)) +} + +// Close closes the batch operation. Any error that occurred during a batch operation may have made it impossible to +// resyncronize the connection with the server. In this case the underlying connection will have been closed. +func (br *pipelineBatchResults) Close() error { + defer func() { + if !br.endTraced { + if br.conn.batchTracer != nil { + br.conn.batchTracer.TraceBatchEnd(br.ctx, br.conn, TraceBatchEndData{Err: br.err}) + } + br.endTraced = true + } + + invalidateCachesOnBatchResultsError(br.conn, br.b, br.err) + }() + + if br.err == nil && br.lastRows != nil && br.lastRows.err != nil { + br.err = br.lastRows.err + } + + if br.closed { + return br.err + } + + // Read and run fn for all remaining items + for br.err == nil && !br.closed && br.b != nil && br.qqIdx < len(br.b.QueuedQueries) { + if br.b.QueuedQueries[br.qqIdx].Fn != nil { + err := br.b.QueuedQueries[br.qqIdx].Fn(br) + if err != nil { + br.err = err + } + } else { + br.Exec() + } + } + + br.closed = true + + err := br.pipeline.Close() + if br.err == nil { + br.err = err + } + + return br.err +} + +func (br *pipelineBatchResults) earlyError() error { + return br.err +} + +func (br *pipelineBatchResults) nextQueryAndArgs() (query string, args []any, err error) { + if br.b == nil { + return "", nil, errors.New("no reference to batch") + } + + if br.qqIdx >= len(br.b.QueuedQueries) { + return "", nil, errors.New("no more results in batch") + } + + bi := br.b.QueuedQueries[br.qqIdx] + br.qqIdx++ + return bi.SQL, bi.Arguments, nil +} + +type emptyBatchResults struct { + conn *Conn + closed bool +} + +// Exec reads the results from the next query in the batch as if the query has been sent with Exec. +func (br *emptyBatchResults) Exec() (pgconn.CommandTag, error) { + if br.closed { + return pgconn.CommandTag{}, fmt.Errorf("batch already closed") + } + return pgconn.CommandTag{}, errors.New("no more results in batch") +} + +// Query reads the results from the next query in the batch as if the query has been sent with Query. +func (br *emptyBatchResults) Query() (Rows, error) { + if br.closed { + alreadyClosedErr := fmt.Errorf("batch already closed") + return &baseRows{err: alreadyClosedErr, closed: true}, alreadyClosedErr + } + + rows := br.conn.getRows(context.Background(), "", nil) + rows.err = errors.New("no more results in batch") + rows.closed = true + return rows, rows.err +} + +// QueryRow reads the results from the next query in the batch as if the query has been sent with QueryRow. +func (br *emptyBatchResults) QueryRow() Row { + rows, _ := br.Query() + return (*connRow)(rows.(*baseRows)) +} + +// Close closes the batch operation. Any error that occurred during a batch operation may have made it impossible to +// resyncronize the connection with the server. In this case the underlying connection will have been closed. +func (br *emptyBatchResults) Close() error { + br.closed = true + return nil +} + +// invalidates statement and description caches on batch results error +func invalidateCachesOnBatchResultsError(conn *Conn, b *Batch, err error) { + if err != nil && conn != nil && b != nil { + if sc := conn.statementCache; sc != nil { + for _, bi := range b.QueuedQueries { + sc.Invalidate(bi.SQL) + } + } + + if sc := conn.descriptionCache; sc != nil { + for _, bi := range b.QueuedQueries { + sc.Invalidate(bi.SQL) + } + } + } +} + +// ErrPreprocessingBatch occurs when an error is encountered while preprocessing a batch. +// The two preprocessing steps are "prepare" (server-side SQL parse/plan) and +// "build" (client-side argument encoding). +type ErrPreprocessingBatch struct { + step string // "prepare" or "build" + sql string + err error +} + +func newErrPreprocessingBatch(step, sql string, err error) ErrPreprocessingBatch { + return ErrPreprocessingBatch{step: step, sql: sql, err: err} +} + +func (e ErrPreprocessingBatch) Error() string { + // intentionally not including the SQL query in the error message + // to avoid leaking potentially sensitive information into logs. + // If the user wants the SQL, they can call SQL(). + return fmt.Sprintf("error preprocessing batch (%s): %v", e.step, e.err) +} + +func (e ErrPreprocessingBatch) Unwrap() error { + return e.err +} + +func (e ErrPreprocessingBatch) SQL() string { + return e.sql +} diff --git a/vendor/github.com/jackc/pgx/v5/conn.go b/vendor/github.com/jackc/pgx/v5/conn.go new file mode 100644 index 0000000000..bc5d064939 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/conn.go @@ -0,0 +1,1472 @@ +package pgx + +import ( + "context" + "crypto/sha256" + "database/sql" + "encoding/hex" + "errors" + "fmt" + "strconv" + "strings" + "time" + + "github.com/jackc/pgx/v5/internal/sanitize" + "github.com/jackc/pgx/v5/internal/stmtcache" + "github.com/jackc/pgx/v5/pgconn" + "github.com/jackc/pgx/v5/pgtype" +) + +// ConnConfig contains all the options used to establish a connection. It must be created by [ParseConfig] and +// then it can be modified. A manually initialized ConnConfig will cause [ConnectConfig] to panic. +type ConnConfig struct { + pgconn.Config + + Tracer QueryTracer + + // Original connection string that was parsed into config. + connString string + + // StatementCacheCapacity is maximum size of the statement cache used when executing a query with "cache_statement" + // query exec mode. + StatementCacheCapacity int + + // DescriptionCacheCapacity is the maximum size of the description cache used when executing a query with + // "cache_describe" query exec mode. + DescriptionCacheCapacity int + + // DefaultQueryExecMode controls the default mode for executing queries. By default pgx uses the extended protocol + // and automatically prepares and caches prepared statements. However, this may be incompatible with proxies such as + // PGBouncer. In this case it may be preferable to use [QueryExecModeExec] or [QueryExecModeSimpleProtocol]. The same + // functionality can be controlled on a per query basis by passing a [QueryExecMode] as the first query argument. + DefaultQueryExecMode QueryExecMode + + createdByParseConfig bool // Used to enforce created by ParseConfig rule. +} + +// ParseConfigOptions contains options that control how a config is built such as getsslpassword. +type ParseConfigOptions struct { + pgconn.ParseConfigOptions +} + +// Copy returns a deep copy of the config that is safe to use and modify. +// The only exception is the tls.Config: +// according to the tls.Config docs it must not be modified after creation. +func (cc *ConnConfig) Copy() *ConnConfig { + newConfig := new(ConnConfig) + *newConfig = *cc + newConfig.Config = *newConfig.Config.Copy() + return newConfig +} + +// ConnString returns the connection string as parsed by pgx.ParseConfig into pgx.ConnConfig. +func (cc *ConnConfig) ConnString() string { return cc.connString } + +// Conn is a PostgreSQL connection handle. It is not safe for concurrent usage. Use a connection pool to manage access +// to multiple database connections from multiple goroutines. +type Conn struct { + pgConn *pgconn.PgConn + config *ConnConfig // config used when establishing this connection + preparedStatements map[string]*pgconn.StatementDescription + failedDescribeStatement string + statementCache stmtcache.Cache + descriptionCache stmtcache.Cache + + queryTracer QueryTracer + batchTracer BatchTracer + copyFromTracer CopyFromTracer + prepareTracer PrepareTracer + + notifications []*pgconn.Notification + + doneChan chan struct{} + closedChan chan error + + typeMap *pgtype.Map + + wbuf []byte + eqb ExtendedQueryBuilder +} + +// Identifier a PostgreSQL identifier or name. Identifiers can be composed of +// multiple parts such as ["schema", "table"] or ["table", "column"]. +type Identifier []string + +// Sanitize returns a sanitized string safe for SQL interpolation. +func (ident Identifier) Sanitize() string { + parts := make([]string, len(ident)) + for i := range ident { + s := strings.ReplaceAll(ident[i], string([]byte{0}), "") + parts[i] = `"` + strings.ReplaceAll(s, `"`, `""`) + `"` + } + return strings.Join(parts, ".") +} + +var ( + // ErrNoRows occurs when rows are expected but none are returned. + ErrNoRows = newProxyErr(sql.ErrNoRows, "no rows in result set") + // ErrTooManyRows occurs when more rows than expected are returned. + ErrTooManyRows = errors.New("too many rows in result set") +) + +func newProxyErr(background error, msg string) error { + return &proxyError{ + msg: msg, + background: background, + } +} + +type proxyError struct { + msg string + background error +} + +func (err *proxyError) Error() string { return err.msg } + +func (err *proxyError) Unwrap() error { return err.background } + +var ( + errDisabledStatementCache = fmt.Errorf("cannot use QueryExecModeCacheStatement with disabled statement cache") + errDisabledDescriptionCache = fmt.Errorf("cannot use QueryExecModeCacheDescribe with disabled description cache") +) + +// Connect establishes a connection with a PostgreSQL server with a connection string. See +// [pgconn.Connect] for details. +func Connect(ctx context.Context, connString string) (*Conn, error) { + connConfig, err := ParseConfig(connString) + if err != nil { + return nil, err + } + return connect(ctx, connConfig) +} + +// ConnectWithOptions behaves exactly like Connect with the addition of options. At the present options is only used to +// provide a [pgconn.GetSSLPasswordFunc] function. +func ConnectWithOptions(ctx context.Context, connString string, options ParseConfigOptions) (*Conn, error) { + connConfig, err := ParseConfigWithOptions(connString, options) + if err != nil { + return nil, err + } + return connect(ctx, connConfig) +} + +// ConnectConfig establishes a connection with a PostgreSQL server with a configuration struct. +// connConfig must have been created by [ParseConfig]. +func ConnectConfig(ctx context.Context, connConfig *ConnConfig) (*Conn, error) { + // In general this improves safety. In particular avoid the config.Config.OnNotification mutation from affecting other + // connections with the same config. See https://github.com/jackc/pgx/issues/618. + connConfig = connConfig.Copy() + + return connect(ctx, connConfig) +} + +// ParseConfigWithOptions behaves exactly as [ParseConfig] does with the addition of options. At the present options is +// only used to provide a [pgconn.GetSSLPasswordFunc] function. +func ParseConfigWithOptions(connString string, options ParseConfigOptions) (*ConnConfig, error) { + config, err := pgconn.ParseConfigWithOptions(connString, options.ParseConfigOptions) + if err != nil { + return nil, err + } + + statementCacheCapacity := 512 + if s, ok := config.RuntimeParams["statement_cache_capacity"]; ok { + delete(config.RuntimeParams, "statement_cache_capacity") + n, err := strconv.ParseInt(s, 10, 32) + if err != nil { + return nil, pgconn.NewParseConfigError(connString, "cannot parse statement_cache_capacity", err) + } + statementCacheCapacity = int(n) + } + + descriptionCacheCapacity := 512 + if s, ok := config.RuntimeParams["description_cache_capacity"]; ok { + delete(config.RuntimeParams, "description_cache_capacity") + n, err := strconv.ParseInt(s, 10, 32) + if err != nil { + return nil, pgconn.NewParseConfigError(connString, "cannot parse description_cache_capacity", err) + } + descriptionCacheCapacity = int(n) + } + + defaultQueryExecMode := QueryExecModeCacheStatement + if s, ok := config.RuntimeParams["default_query_exec_mode"]; ok { + delete(config.RuntimeParams, "default_query_exec_mode") + switch s { + case "cache_statement": + defaultQueryExecMode = QueryExecModeCacheStatement + case "cache_describe": + defaultQueryExecMode = QueryExecModeCacheDescribe + case "describe_exec": + defaultQueryExecMode = QueryExecModeDescribeExec + case "exec": + defaultQueryExecMode = QueryExecModeExec + case "simple_protocol": + defaultQueryExecMode = QueryExecModeSimpleProtocol + default: + return nil, pgconn.NewParseConfigError( + connString, "invalid default_query_exec_mode", fmt.Errorf("unknown value %q", s), + ) + } + } + + connConfig := &ConnConfig{ + Config: *config, + createdByParseConfig: true, + StatementCacheCapacity: statementCacheCapacity, + DescriptionCacheCapacity: descriptionCacheCapacity, + DefaultQueryExecMode: defaultQueryExecMode, + connString: connString, + } + + return connConfig, nil +} + +// ParseConfig creates a ConnConfig from a connection string. ParseConfig handles all options that [pgconn.ParseConfig] +// does. In addition, it accepts the following options: +// +// - default_query_exec_mode. +// Possible values: "cache_statement", "cache_describe", "describe_exec", "exec", and "simple_protocol". See +// QueryExecMode constant documentation for the meaning of these values. Default: "cache_statement". +// +// - statement_cache_capacity. +// The maximum size of the statement cache used when executing a query with "cache_statement" query exec mode. +// Default: 512. +// +// - description_cache_capacity. +// The maximum size of the description cache used when executing a query with "cache_describe" query exec mode. +// Default: 512. +func ParseConfig(connString string) (*ConnConfig, error) { + return ParseConfigWithOptions(connString, ParseConfigOptions{}) +} + +// connect connects to a database. connect takes ownership of config. The caller must not use or access it again. +func connect(ctx context.Context, config *ConnConfig) (c *Conn, err error) { + if connectTracer, ok := config.Tracer.(ConnectTracer); ok { + ctx = connectTracer.TraceConnectStart(ctx, TraceConnectStartData{ConnConfig: config}) + defer func() { + connectTracer.TraceConnectEnd(ctx, TraceConnectEndData{Conn: c, Err: err}) + }() + } + + // Default values are set in ParseConfig. Enforce initial creation by ParseConfig rather than setting defaults from + // zero values. + if !config.createdByParseConfig { + panic("config must be created by ParseConfig") + } + + c = &Conn{ + config: config, + typeMap: pgtype.NewMap(), + queryTracer: config.Tracer, + } + + if t, ok := c.queryTracer.(BatchTracer); ok { + c.batchTracer = t + } + if t, ok := c.queryTracer.(CopyFromTracer); ok { + c.copyFromTracer = t + } + if t, ok := c.queryTracer.(PrepareTracer); ok { + c.prepareTracer = t + } + + // Only install pgx notification system if no other callback handler is present. + if config.Config.OnNotification == nil { + config.Config.OnNotification = c.bufferNotifications + } + + c.pgConn, err = pgconn.ConnectConfig(ctx, &config.Config) + if err != nil { + return nil, err + } + + c.preparedStatements = make(map[string]*pgconn.StatementDescription) + c.doneChan = make(chan struct{}) + c.closedChan = make(chan error) + c.wbuf = make([]byte, 0, 1024) + + if c.config.StatementCacheCapacity > 0 { + c.statementCache = stmtcache.NewLRUCache(c.config.StatementCacheCapacity) + } + + if c.config.DescriptionCacheCapacity > 0 { + c.descriptionCache = stmtcache.NewLRUCache(c.config.DescriptionCacheCapacity) + } + + return c, nil +} + +// Close closes a connection. It is safe to call Close on an already closed +// connection. +func (c *Conn) Close(ctx context.Context) error { + if c.IsClosed() { + return nil + } + + err := c.pgConn.Close(ctx) + return err +} + +// Prepare creates a prepared statement with name and sql. sql can contain placeholders for bound parameters. These +// placeholders are referenced positionally as $1, $2, etc. name can be used instead of sql with [Conn.Query], +// [Conn.QueryRow], and [Conn.Exec] to execute the statement. It can also be used with [Batch.Queue]. +// +// The underlying PostgreSQL identifier for the prepared statement will be name if name != sql or a digest of sql if +// name == sql. +// +// Prepare is idempotent; i.e. it is safe to call Prepare multiple times with the same name and sql arguments. This +// allows a code path to Prepare and Query/Exec without concern for if the statement has already been prepared. +func (c *Conn) Prepare(ctx context.Context, name, sql string) (sd *pgconn.StatementDescription, err error) { + if c.failedDescribeStatement != "" { + err = c.Deallocate(ctx, c.failedDescribeStatement) + if err != nil { + return nil, fmt.Errorf("failed to deallocate previously failed statement %q: %w", c.failedDescribeStatement, err) + } + c.failedDescribeStatement = "" + } + + if c.prepareTracer != nil { + ctx = c.prepareTracer.TracePrepareStart(ctx, c, TracePrepareStartData{Name: name, SQL: sql}) + } + + if name != "" { + var ok bool + if sd, ok = c.preparedStatements[name]; ok && sd.SQL == sql { + if c.prepareTracer != nil { + c.prepareTracer.TracePrepareEnd(ctx, c, TracePrepareEndData{AlreadyPrepared: true}) + } + return sd, nil + } + } + + if c.prepareTracer != nil { + defer func() { + c.prepareTracer.TracePrepareEnd(ctx, c, TracePrepareEndData{Err: err}) + }() + } + + var psName, psKey string + if name == sql { + digest := sha256.Sum256([]byte(sql)) + psName = "stmt_" + hex.EncodeToString(digest[0:24]) + psKey = sql + } else { + psName = name + psKey = name + } + + sd, err = c.pgConn.Prepare(ctx, psName, sql, nil) + if err != nil { + var pErr *pgconn.PrepareError + if errors.As(err, &pErr) { + c.failedDescribeStatement = psKey + } + return nil, err + } + + if psKey != "" { + c.preparedStatements[psKey] = sd + } + + return sd, nil +} + +// Deallocate releases a prepared statement. Calling Deallocate on a non-existent prepared statement will succeed. +func (c *Conn) Deallocate(ctx context.Context, name string) error { + var psName string + sd := c.preparedStatements[name] + if sd != nil { + psName = sd.Name + } else { + psName = name + } + + err := c.pgConn.Deallocate(ctx, psName) + if err != nil { + return err + } + + if sd != nil { + delete(c.preparedStatements, name) + } + + return nil +} + +// DeallocateAll releases all previously prepared statements from the server and client, where it also resets the statement and description cache. +func (c *Conn) DeallocateAll(ctx context.Context) error { + c.preparedStatements = map[string]*pgconn.StatementDescription{} + if c.config.StatementCacheCapacity > 0 { + c.statementCache = stmtcache.NewLRUCache(c.config.StatementCacheCapacity) + } + if c.config.DescriptionCacheCapacity > 0 { + c.descriptionCache = stmtcache.NewLRUCache(c.config.DescriptionCacheCapacity) + } + _, err := c.pgConn.Exec(ctx, "deallocate all").ReadAll() + return err +} + +func (c *Conn) bufferNotifications(_ *pgconn.PgConn, n *pgconn.Notification) { + c.notifications = append(c.notifications, n) +} + +// WaitForNotification waits for a PostgreSQL notification. It wraps the underlying pgconn notification system in a +// slightly more convenient form. +func (c *Conn) WaitForNotification(ctx context.Context) (*pgconn.Notification, error) { + var n *pgconn.Notification + + // Return already received notification immediately + if len(c.notifications) > 0 { + n = c.notifications[0] + c.notifications = c.notifications[1:] + return n, nil + } + + err := c.pgConn.WaitForNotification(ctx) + if len(c.notifications) > 0 { + n = c.notifications[0] + c.notifications = c.notifications[1:] + } + return n, err +} + +// IsClosed reports if the connection has been closed. +func (c *Conn) IsClosed() bool { + return c.pgConn.IsClosed() +} + +func (c *Conn) die() { + if c.IsClosed() { + return + } + + ctx, cancel := context.WithCancel(context.Background()) + cancel() // force immediate hard cancel + c.pgConn.Close(ctx) +} + +func quoteIdentifier(s string) string { + return `"` + strings.ReplaceAll(s, `"`, `""`) + `"` +} + +// Ping delegates to the underlying *pgconn.PgConn.Ping. +func (c *Conn) Ping(ctx context.Context) error { + return c.pgConn.Ping(ctx) +} + +// PgConn returns the underlying *pgconn.PgConn. This is an escape hatch method that allows lower level access to the +// PostgreSQL connection than pgx exposes. +// +// It is strongly recommended that the connection be idle (no in-progress queries) before the underlying *pgconn.PgConn +// is used and the connection must be returned to the same state before any *pgx.Conn methods are again used. +func (c *Conn) PgConn() *pgconn.PgConn { return c.pgConn } + +// TypeMap returns the connection info used for this connection. +func (c *Conn) TypeMap() *pgtype.Map { return c.typeMap } + +// Config returns a copy of config that was used to establish this connection. +func (c *Conn) Config() *ConnConfig { return c.config.Copy() } + +// Exec executes sql. sql can be either a prepared statement name or an SQL string. arguments should be referenced +// positionally from the sql string as $1, $2, etc. +func (c *Conn) Exec(ctx context.Context, sql string, arguments ...any) (pgconn.CommandTag, error) { + if c.queryTracer != nil { + ctx = c.queryTracer.TraceQueryStart(ctx, c, TraceQueryStartData{SQL: sql, Args: arguments}) + } + + if err := c.deallocateInvalidatedCachedStatements(ctx); err != nil { + return pgconn.CommandTag{}, err + } + + commandTag, err := c.exec(ctx, sql, arguments...) + + if c.queryTracer != nil { + c.queryTracer.TraceQueryEnd(ctx, c, TraceQueryEndData{CommandTag: commandTag, Err: err}) + } + + return commandTag, err +} + +func (c *Conn) exec(ctx context.Context, sql string, arguments ...any) (commandTag pgconn.CommandTag, err error) { + mode := c.config.DefaultQueryExecMode + var queryRewriter QueryRewriter + +optionLoop: + for len(arguments) > 0 { + switch arg := arguments[0].(type) { + case QueryExecMode: + mode = arg + arguments = arguments[1:] + case QueryRewriter: + queryRewriter = arg + arguments = arguments[1:] + default: + break optionLoop + } + } + + if queryRewriter != nil { + sql, arguments, err = queryRewriter.RewriteQuery(ctx, c, sql, arguments) + if err != nil { + return pgconn.CommandTag{}, fmt.Errorf("rewrite query failed: %w", err) + } + } + + // Always use simple protocol when there are no arguments. + if len(arguments) == 0 { + mode = QueryExecModeSimpleProtocol + } + + defer func() { + if err != nil { + if sc := c.statementCache; sc != nil { + sc.Invalidate(sql) + } + + if sc := c.descriptionCache; sc != nil { + sc.Invalidate(sql) + } + } + }() + + if sd, ok := c.preparedStatements[sql]; ok { + return c.execPrepared(ctx, sd, arguments) + } + + switch mode { + case QueryExecModeCacheStatement: + if c.statementCache == nil { + return pgconn.CommandTag{}, errDisabledStatementCache + } + sd := c.statementCache.Get(sql) + if sd == nil { + sd, err = c.Prepare(ctx, stmtcache.StatementName(sql), sql) + if err != nil { + return pgconn.CommandTag{}, err + } + c.statementCache.Put(sd) + } + + return c.execPrepared(ctx, sd, arguments) + case QueryExecModeCacheDescribe: + if c.descriptionCache == nil { + return pgconn.CommandTag{}, errDisabledDescriptionCache + } + sd := c.descriptionCache.Get(sql) + if sd == nil { + sd, err = c.Prepare(ctx, "", sql) + if err != nil { + return pgconn.CommandTag{}, err + } + c.descriptionCache.Put(sd) + } + + return c.execParams(ctx, sd, arguments) + case QueryExecModeDescribeExec: + sd, err := c.Prepare(ctx, "", sql) + if err != nil { + return pgconn.CommandTag{}, err + } + return c.execPrepared(ctx, sd, arguments) + case QueryExecModeExec: + return c.execSQLParams(ctx, sql, arguments) + case QueryExecModeSimpleProtocol: + return c.execSimpleProtocol(ctx, sql, arguments) + default: + return pgconn.CommandTag{}, fmt.Errorf("unknown QueryExecMode: %v", mode) + } +} + +func (c *Conn) execSimpleProtocol(ctx context.Context, sql string, arguments []any) (commandTag pgconn.CommandTag, err error) { + if len(arguments) > 0 { + sql, err = c.sanitizeForSimpleQuery(sql, arguments...) + if err != nil { + return pgconn.CommandTag{}, err + } + } + + mrr := c.pgConn.Exec(ctx, sql) + for mrr.NextResult() { + commandTag, _ = mrr.ResultReader().Close() + } + err = mrr.Close() + return commandTag, err +} + +func (c *Conn) execParams(ctx context.Context, sd *pgconn.StatementDescription, arguments []any) (pgconn.CommandTag, error) { + err := c.eqb.Build(c.typeMap, sd, arguments) + if err != nil { + return pgconn.CommandTag{}, err + } + + result := c.pgConn.ExecParams(ctx, sd.SQL, c.eqb.ParamValues, sd.ParamOIDs, c.eqb.ParamFormats, c.eqb.ResultFormats).Read() + c.eqb.reset() // Allow c.eqb internal memory to be GC'ed as soon as possible. + return result.CommandTag, result.Err +} + +func (c *Conn) execPrepared(ctx context.Context, sd *pgconn.StatementDescription, arguments []any) (pgconn.CommandTag, error) { + err := c.eqb.Build(c.typeMap, sd, arguments) + if err != nil { + return pgconn.CommandTag{}, err + } + + result := c.pgConn.ExecStatement(ctx, sd, c.eqb.ParamValues, c.eqb.ParamFormats, c.eqb.ResultFormats).Read() + c.eqb.reset() // Allow c.eqb internal memory to be GC'ed as soon as possible. + return result.CommandTag, result.Err +} + +func (c *Conn) execSQLParams(ctx context.Context, sql string, args []any) (pgconn.CommandTag, error) { + err := c.eqb.Build(c.typeMap, nil, args) + if err != nil { + return pgconn.CommandTag{}, err + } + + result := c.pgConn.ExecParams(ctx, sql, c.eqb.ParamValues, nil, c.eqb.ParamFormats, c.eqb.ResultFormats).Read() + c.eqb.reset() // Allow c.eqb internal memory to be GC'ed as soon as possible. + return result.CommandTag, result.Err +} + +func (c *Conn) getRows(ctx context.Context, sql string, args []any) *baseRows { + r := &baseRows{} + + r.ctx = ctx + r.queryTracer = c.queryTracer + r.typeMap = c.typeMap + r.startTime = time.Now() + r.sql = sql + r.args = args + r.conn = c + + return r +} + +type QueryExecMode int32 + +const ( + _ QueryExecMode = iota + + // Automatically prepare and cache statements. This uses the extended protocol. Queries are executed in a single round + // trip after the statement is cached. This is the default. If the database schema is modified or the search_path is + // changed after a statement is cached then the first execution of a previously cached query may fail. e.g. If the + // number of columns returned by a "SELECT *" changes or the type of a column is changed. + QueryExecModeCacheStatement + + // Cache statement descriptions (i.e. argument and result types) and assume they do not change. This uses the extended + // protocol. Queries are executed in a single round trip after the description is cached. If the database schema is + // modified or the search_path is changed after a statement is cached then the first execution of a previously cached + // query may fail. e.g. If the number of columns returned by a "SELECT *" changes or the type of a column is changed. + QueryExecModeCacheDescribe + + // Get the statement description on every execution. This uses the extended protocol. Queries require two round trips + // to execute. It does not use named prepared statements. But it does use the unnamed prepared statement to get the + // statement description on the first round trip and then uses it to execute the query on the second round trip. This + // may cause problems with connection poolers that switch the underlying connection between round trips. It is safe + // even when the database schema is modified concurrently. + QueryExecModeDescribeExec + + // Assume the PostgreSQL query parameter types based on the Go type of the arguments. This uses the extended protocol + // with text formatted parameters and results. Queries are executed in a single round trip. Type mappings can be + // registered with pgtype.Map.RegisterDefaultPgType. Queries will be rejected that have arguments that are + // unregistered or ambiguous. e.g. A map[string]string may have the PostgreSQL type json or hstore. Modes that know + // the PostgreSQL type can use a map[string]string directly as an argument. This mode cannot. + // + // On rare occasions user defined types may behave differently when encoded in the text format instead of the binary + // format. For example, this could happen if a "type RomanNumeral int32" implements fmt.Stringer to format integers as + // Roman numerals (e.g. 7 is VII). The binary format would properly encode the integer 7 as the binary value for 7. + // But the text format would encode the integer 7 as the string "VII". As QueryExecModeExec uses the text format, it + // is possible that changing query mode from another mode to QueryExecModeExec could change the behavior of the query. + // This should not occur with types pgx supports directly and can be avoided by registering the types with + // pgtype.Map.RegisterDefaultPgType and implementing the appropriate type interfaces. In the cas of RomanNumeral, it + // should implement pgtype.Int64Valuer. + QueryExecModeExec + + // Use the simple protocol. Assume the PostgreSQL query parameter types based on the Go type of the arguments. This is + // especially significant for []byte values. []byte values are encoded as PostgreSQL bytea. string must be used + // instead for text type values including json and jsonb. Type mappings can be registered with + // pgtype.Map.RegisterDefaultPgType. Queries will be rejected that have arguments that are unregistered or ambiguous. + // e.g. A map[string]string may have the PostgreSQL type json or hstore. Modes that know the PostgreSQL type can use a + // map[string]string directly as an argument. This mode cannot. Queries are executed in a single round trip. + // + // QueryExecModeSimpleProtocol should have the user application visible behavior as QueryExecModeExec. This includes + // the warning regarding differences in text format and binary format encoding with user defined types. There may be + // other minor exceptions such as behavior when multiple result returning queries are erroneously sent in a single + // string. + // + // QueryExecModeSimpleProtocol uses client side parameter interpolation. All values are quoted and escaped. Prefer + // QueryExecModeExec over QueryExecModeSimpleProtocol whenever possible. In general QueryExecModeSimpleProtocol should + // only be used if connecting to a proxy server, connection pool server, or non-PostgreSQL server that does not + // support the extended protocol. + QueryExecModeSimpleProtocol +) + +func (m QueryExecMode) String() string { + switch m { + case QueryExecModeCacheStatement: + return "cache statement" + case QueryExecModeCacheDescribe: + return "cache describe" + case QueryExecModeDescribeExec: + return "describe exec" + case QueryExecModeExec: + return "exec" + case QueryExecModeSimpleProtocol: + return "simple protocol" + default: + return "invalid" + } +} + +// QueryResultFormats controls the result format (text=0, binary=1) of a query by result column position. +type QueryResultFormats []int16 + +// QueryResultFormatsByOID controls the result format (text=0, binary=1) of a query by the result column OID. +type QueryResultFormatsByOID map[uint32]int16 + +// QueryRewriter rewrites a query when used as the first arguments to a query method. +type QueryRewriter interface { + RewriteQuery(ctx context.Context, conn *Conn, sql string, args []any) (newSQL string, newArgs []any, err error) +} + +// Query sends a query to the server and returns a Rows to read the results. Only errors encountered sending the query +// and initializing Rows will be returned. Err() on the returned Rows must be checked after the Rows is closed to +// determine if the query executed successfully. +// +// The returned Rows must be closed before the connection can be used again. It is safe to attempt to read from the +// returned Rows even if an error is returned. The error will be the available in rows.Err() after rows are closed. It +// is allowed to ignore the error returned from Query and handle it in Rows. +// +// It is possible for a call of FieldDescriptions on the returned Rows to return nil even if the Query call did not +// return an error. +// +// It is possible for a query to return one or more rows before encountering an error. In most cases the rows should be +// collected before processing rather than processed while receiving each row. This avoids the possibility of the +// application processing rows from a query that the server rejected. The CollectRows function is useful here. +// +// An implementor of QueryRewriter may be passed as the first element of args. It can rewrite the sql and change or +// replace args. For example, NamedArgs is QueryRewriter that implements named arguments. +// +// For extra control over how the query is executed, the types QueryExecMode, QueryResultFormats, and +// QueryResultFormatsByOID may be used as the first args to control exactly how the query is executed. This is rarely +// needed. See the documentation for those types for details. +func (c *Conn) Query(ctx context.Context, sql string, args ...any) (Rows, error) { + if c.queryTracer != nil { + ctx = c.queryTracer.TraceQueryStart(ctx, c, TraceQueryStartData{SQL: sql, Args: args}) + } + + if err := c.deallocateInvalidatedCachedStatements(ctx); err != nil { + if c.queryTracer != nil { + c.queryTracer.TraceQueryEnd(ctx, c, TraceQueryEndData{Err: err}) + } + return &baseRows{err: err, closed: true}, err + } + + var resultFormats QueryResultFormats + var resultFormatsByOID QueryResultFormatsByOID + mode := c.config.DefaultQueryExecMode + var queryRewriter QueryRewriter + +optionLoop: + for len(args) > 0 { + switch arg := args[0].(type) { + case QueryResultFormats: + resultFormats = arg + args = args[1:] + case QueryResultFormatsByOID: + resultFormatsByOID = arg + args = args[1:] + case QueryExecMode: + mode = arg + args = args[1:] + case QueryRewriter: + queryRewriter = arg + args = args[1:] + default: + break optionLoop + } + } + + if queryRewriter != nil { + var err error + originalSQL := sql + originalArgs := args + sql, args, err = queryRewriter.RewriteQuery(ctx, c, sql, args) + if err != nil { + rows := c.getRows(ctx, originalSQL, originalArgs) + err = fmt.Errorf("rewrite query failed: %w", err) + rows.fatal(err) + return rows, err + } + } + + // Bypass any statement caching. + if sql == "" { + mode = QueryExecModeSimpleProtocol + } + + c.eqb.reset() + rows := c.getRows(ctx, sql, args) + + var err error + sd, explicitPreparedStatement := c.preparedStatements[sql] + switch { + case sd != nil || mode == QueryExecModeCacheStatement || mode == QueryExecModeCacheDescribe || mode == QueryExecModeDescribeExec: + if sd == nil { + sd, err = c.getStatementDescription(ctx, mode, sql) + if err != nil { + rows.fatal(err) + return rows, err + } + } + + if len(sd.ParamOIDs) != len(args) { + rows.fatal(fmt.Errorf("expected %d arguments, got %d", len(sd.ParamOIDs), len(args))) + return rows, rows.err + } + + rows.sql = sd.SQL + + err = c.eqb.Build(c.typeMap, sd, args) + if err != nil { + rows.fatal(err) + return rows, rows.err + } + + if resultFormatsByOID != nil { + resultFormats = make([]int16, len(sd.Fields)) + for i := range resultFormats { + resultFormats[i] = resultFormatsByOID[sd.Fields[i].DataTypeOID] + } + } + + if resultFormats == nil { + resultFormats = c.eqb.ResultFormats + } + + if !explicitPreparedStatement && mode == QueryExecModeCacheDescribe { + rows.resultReader = c.pgConn.ExecParams(ctx, sql, c.eqb.ParamValues, sd.ParamOIDs, c.eqb.ParamFormats, resultFormats) + } else { + rows.resultReader = c.pgConn.ExecStatement(ctx, sd, c.eqb.ParamValues, c.eqb.ParamFormats, resultFormats) + } + case mode == QueryExecModeExec: + err := c.eqb.Build(c.typeMap, nil, args) + if err != nil { + rows.fatal(err) + return rows, rows.err + } + + rows.resultReader = c.pgConn.ExecParams(ctx, sql, c.eqb.ParamValues, nil, c.eqb.ParamFormats, c.eqb.ResultFormats) + case mode == QueryExecModeSimpleProtocol: + sql, err = c.sanitizeForSimpleQuery(sql, args...) + if err != nil { + rows.fatal(err) + return rows, err + } + + mrr := c.pgConn.Exec(ctx, sql) + if mrr.NextResult() { + rows.resultReader = mrr.ResultReader() + rows.multiResultReader = mrr + } else { + err = mrr.Close() + rows.fatal(err) + return rows, err + } + + return rows, nil + default: + err = fmt.Errorf("unknown QueryExecMode: %v", mode) + rows.fatal(err) + return rows, rows.err + } + + c.eqb.reset() // Allow c.eqb internal memory to be GC'ed as soon as possible. + + return rows, rows.err +} + +// getStatementDescription returns the statement description of the sql query +// according to the given mode. +// +// If the mode is one that doesn't require to know the param and result OIDs +// then nil is returned without error. +func (c *Conn) getStatementDescription( + ctx context.Context, + mode QueryExecMode, + sql string, +) (sd *pgconn.StatementDescription, err error) { + switch mode { + case QueryExecModeCacheStatement: + if c.statementCache == nil { + return nil, errDisabledStatementCache + } + sd = c.statementCache.Get(sql) + if sd == nil { + sd, err = c.Prepare(ctx, stmtcache.StatementName(sql), sql) + if err != nil { + return nil, err + } + c.statementCache.Put(sd) + } + case QueryExecModeCacheDescribe: + if c.descriptionCache == nil { + return nil, errDisabledDescriptionCache + } + sd = c.descriptionCache.Get(sql) + if sd == nil { + sd, err = c.Prepare(ctx, "", sql) + if err != nil { + return nil, err + } + c.descriptionCache.Put(sd) + } + case QueryExecModeDescribeExec: + return c.Prepare(ctx, "", sql) + } + return sd, err +} + +// QueryRow is a convenience wrapper over Query. Any error that occurs while +// querying is deferred until calling Scan on the returned Row. That Row will +// error with ErrNoRows if no rows are returned. +func (c *Conn) QueryRow(ctx context.Context, sql string, args ...any) Row { + rows, _ := c.Query(ctx, sql, args...) + return (*connRow)(rows.(*baseRows)) +} + +// SendBatch sends all queued queries to the server at once. All queries are run in an implicit transaction unless +// explicit transaction control statements are executed. The returned [BatchResults] must be closed before the connection +// is used again. +// +// Depending on the QueryExecMode, all queries may be prepared before any are executed. This means that creating a table +// and using it in a subsequent query in the same batch can fail. +func (c *Conn) SendBatch(ctx context.Context, b *Batch) (br BatchResults) { + if len(b.QueuedQueries) == 0 { + return &emptyBatchResults{conn: c} + } + + if c.batchTracer != nil { + ctx = c.batchTracer.TraceBatchStart(ctx, c, TraceBatchStartData{Batch: b}) + defer func() { + err := br.(interface{ earlyError() error }).earlyError() + if err != nil { + c.batchTracer.TraceBatchEnd(ctx, c, TraceBatchEndData{Err: err}) + } + }() + } + + if err := c.deallocateInvalidatedCachedStatements(ctx); err != nil { + return &batchResults{ctx: ctx, conn: c, err: err} + } + + for _, bi := range b.QueuedQueries { + var queryRewriter QueryRewriter + sql := bi.SQL + arguments := bi.Arguments + + optionLoop: + for len(arguments) > 0 { + // Update Batch.Queue function comment when additional options are implemented + switch arg := arguments[0].(type) { + case QueryRewriter: + queryRewriter = arg + arguments = arguments[1:] + default: + break optionLoop + } + } + + if queryRewriter != nil { + var err error + sql, arguments, err = queryRewriter.RewriteQuery(ctx, c, sql, arguments) + if err != nil { + return &batchResults{ctx: ctx, conn: c, err: fmt.Errorf("rewrite query failed: %w", err)} + } + } + + bi.SQL = sql + bi.Arguments = arguments + } + + // TODO: changing mode per batch? Update Batch.Queue function comment when implemented + mode := c.config.DefaultQueryExecMode + if mode == QueryExecModeSimpleProtocol { + return c.sendBatchQueryExecModeSimpleProtocol(ctx, b) + } + + // All other modes use extended protocol and thus can use prepared statements. + for _, bi := range b.QueuedQueries { + if sd, ok := c.preparedStatements[bi.SQL]; ok { + bi.sd = sd + } + } + + switch mode { + case QueryExecModeExec: + return c.sendBatchQueryExecModeExec(ctx, b) + case QueryExecModeCacheStatement: + return c.sendBatchQueryExecModeCacheStatement(ctx, b) + case QueryExecModeCacheDescribe: + return c.sendBatchQueryExecModeCacheDescribe(ctx, b) + case QueryExecModeDescribeExec: + return c.sendBatchQueryExecModeDescribeExec(ctx, b) + default: + panic("unknown QueryExecMode") + } +} + +func (c *Conn) sendBatchQueryExecModeSimpleProtocol(ctx context.Context, b *Batch) *batchResults { + var sb strings.Builder + for i, bi := range b.QueuedQueries { + if i > 0 { + sb.WriteByte(';') + } + sql, err := c.sanitizeForSimpleQuery(bi.SQL, bi.Arguments...) + if err != nil { + return &batchResults{ctx: ctx, conn: c, err: err} + } + sb.WriteString(sql) + } + mrr := c.pgConn.Exec(ctx, sb.String()) + return &batchResults{ + ctx: ctx, + conn: c, + mrr: mrr, + b: b, + qqIdx: 0, + } +} + +func (c *Conn) sendBatchQueryExecModeExec(ctx context.Context, b *Batch) *batchResults { + batch := &pgconn.Batch{} + + for _, bi := range b.QueuedQueries { + sd := bi.sd + if sd != nil { + err := c.eqb.Build(c.typeMap, sd, bi.Arguments) + if err != nil { + return &batchResults{ctx: ctx, conn: c, err: err} + } + + batch.ExecPrepared(sd.Name, c.eqb.ParamValues, c.eqb.ParamFormats, c.eqb.ResultFormats) + } else { + err := c.eqb.Build(c.typeMap, nil, bi.Arguments) + if err != nil { + return &batchResults{ctx: ctx, conn: c, err: err} + } + batch.ExecParams(bi.SQL, c.eqb.ParamValues, nil, c.eqb.ParamFormats, c.eqb.ResultFormats) + } + } + + c.eqb.reset() // Allow c.eqb internal memory to be GC'ed as soon as possible. + + mrr := c.pgConn.ExecBatch(ctx, batch) + + return &batchResults{ + ctx: ctx, + conn: c, + mrr: mrr, + b: b, + qqIdx: 0, + } +} + +func (c *Conn) sendBatchQueryExecModeCacheStatement(ctx context.Context, b *Batch) (pbr *pipelineBatchResults) { + if c.statementCache == nil { + return &pipelineBatchResults{ctx: ctx, conn: c, err: errDisabledStatementCache, closed: true} + } + + distinctNewQueries := []*pgconn.StatementDescription{} + distinctNewQueriesIdxMap := make(map[string]int) + + for _, bi := range b.QueuedQueries { + if bi.sd == nil { + sd := c.statementCache.Get(bi.SQL) + if sd != nil { + bi.sd = sd + } else { + if idx, present := distinctNewQueriesIdxMap[bi.SQL]; present { + bi.sd = distinctNewQueries[idx] + } else { + sd = &pgconn.StatementDescription{ + Name: stmtcache.StatementName(bi.SQL), + SQL: bi.SQL, + } + distinctNewQueriesIdxMap[sd.SQL] = len(distinctNewQueries) + distinctNewQueries = append(distinctNewQueries, sd) + bi.sd = sd + } + } + } + } + + return c.sendBatchExtendedWithDescription(ctx, b, distinctNewQueries, c.statementCache) +} + +func (c *Conn) sendBatchQueryExecModeCacheDescribe(ctx context.Context, b *Batch) (pbr *pipelineBatchResults) { + if c.descriptionCache == nil { + return &pipelineBatchResults{ctx: ctx, conn: c, err: errDisabledDescriptionCache, closed: true} + } + + distinctNewQueries := []*pgconn.StatementDescription{} + distinctNewQueriesIdxMap := make(map[string]int) + + for _, bi := range b.QueuedQueries { + if bi.sd == nil { + sd := c.descriptionCache.Get(bi.SQL) + if sd != nil { + bi.sd = sd + } else { + if idx, present := distinctNewQueriesIdxMap[bi.SQL]; present { + bi.sd = distinctNewQueries[idx] + } else { + sd = &pgconn.StatementDescription{ + SQL: bi.SQL, + } + distinctNewQueriesIdxMap[sd.SQL] = len(distinctNewQueries) + distinctNewQueries = append(distinctNewQueries, sd) + bi.sd = sd + } + } + } + } + + return c.sendBatchExtendedWithDescription(ctx, b, distinctNewQueries, c.descriptionCache) +} + +func (c *Conn) sendBatchQueryExecModeDescribeExec(ctx context.Context, b *Batch) (pbr *pipelineBatchResults) { + distinctNewQueries := []*pgconn.StatementDescription{} + distinctNewQueriesIdxMap := make(map[string]int) + + for _, bi := range b.QueuedQueries { + if bi.sd == nil { + if idx, present := distinctNewQueriesIdxMap[bi.SQL]; present { + bi.sd = distinctNewQueries[idx] + } else { + sd := &pgconn.StatementDescription{ + SQL: bi.SQL, + } + distinctNewQueriesIdxMap[sd.SQL] = len(distinctNewQueries) + distinctNewQueries = append(distinctNewQueries, sd) + bi.sd = sd + } + } + } + + return c.sendBatchExtendedWithDescription(ctx, b, distinctNewQueries, nil) +} + +func (c *Conn) sendBatchExtendedWithDescription(ctx context.Context, b *Batch, distinctNewQueries []*pgconn.StatementDescription, sdCache stmtcache.Cache) (pbr *pipelineBatchResults) { + pipeline := c.pgConn.StartPipeline(ctx) + defer func() { + if pbr != nil && pbr.err != nil { + pipeline.Close() + } + }() + + // Prepare any needed queries + if len(distinctNewQueries) > 0 { + err := func() (err error) { + for _, sd := range distinctNewQueries { + pipeline.SendPrepare(sd.Name, sd.SQL, nil) + } + + // Store all statements we are preparing into the cache. It's fine if it overflows because HandleInvalidated will + // clean them up later. + if sdCache != nil { + for _, sd := range distinctNewQueries { + sdCache.Put(sd) + } + } + + // If something goes wrong preparing the statements, we need to invalidate the cache entries we just added. + defer func() { + if err != nil && sdCache != nil { + for _, sd := range distinctNewQueries { + sdCache.Invalidate(sd.SQL) + } + } + }() + + err = pipeline.Sync() + if err != nil { + return err + } + + for _, sd := range distinctNewQueries { + results, err := pipeline.GetResults() + if err != nil { + return newErrPreprocessingBatch("prepare", sd.SQL, err) + } + + resultSD, ok := results.(*pgconn.StatementDescription) + if !ok { + return fmt.Errorf("expected statement description, got %T", results) + } + + // Fill in the previously empty / pending statement descriptions. + sd.ParamOIDs = resultSD.ParamOIDs + sd.Fields = resultSD.Fields + } + + results, err := pipeline.GetResults() + if err != nil { + return err + } + + _, ok := results.(*pgconn.PipelineSync) + if !ok { + return fmt.Errorf("expected sync, got %T", results) + } + + return nil + }() + if err != nil { + return &pipelineBatchResults{ctx: ctx, conn: c, err: err, closed: true} + } + } + + // Queue the queries. + for _, bi := range b.QueuedQueries { + err := c.eqb.Build(c.typeMap, bi.sd, bi.Arguments) + if err != nil { + err = newErrPreprocessingBatch("build", bi.SQL, err) + return &pipelineBatchResults{ctx: ctx, conn: c, err: err, closed: true} + } + + if bi.sd.Name == "" { + pipeline.SendQueryParams(bi.sd.SQL, c.eqb.ParamValues, bi.sd.ParamOIDs, c.eqb.ParamFormats, c.eqb.ResultFormats) + } else { + // Copy ResultFormats because SendQueryStatement stores the slice for later use, and eqb.Build reuses the + // backing array on the next iteration. + resultFormats := make([]int16, len(c.eqb.ResultFormats)) + copy(resultFormats, c.eqb.ResultFormats) + pipeline.SendQueryStatement(bi.sd, c.eqb.ParamValues, c.eqb.ParamFormats, resultFormats) + } + } + + err := pipeline.Sync() + if err != nil { + return &pipelineBatchResults{ctx: ctx, conn: c, err: err, closed: true} + } + + return &pipelineBatchResults{ + ctx: ctx, + conn: c, + pipeline: pipeline, + b: b, + } +} + +func (c *Conn) sanitizeForSimpleQuery(sql string, args ...any) (string, error) { + if c.pgConn.ParameterStatus("standard_conforming_strings") != "on" { + return "", errors.New("simple protocol queries must be run with standard_conforming_strings=on") + } + + if c.pgConn.ParameterStatus("client_encoding") != "UTF8" { + return "", errors.New("simple protocol queries must be run with client_encoding=UTF8") + } + + var err error + valueArgs := make([]any, len(args)) + for i, a := range args { + valueArgs[i], err = convertSimpleArgument(c.typeMap, a) + if err != nil { + return "", err + } + } + + return sanitize.SanitizeSQL(sql, valueArgs...) +} + +// LoadType inspects the database for typeName and produces a [pgtype.Type] suitable for registration. typeName must be +// the name of a type where the underlying type(s) is already understood by pgx. It is for derived types. In particular, +// typeName must be one of the following: +// - An array type name of a type that is already registered. e.g. "_foo" when "foo" is registered. +// - A composite type name where all field types are already registered. +// - A domain type name where the base type is already registered. +// - An enum type name. +// - A range type name where the element type is already registered. +// - A multirange type name where the element type is already registered. +func (c *Conn) LoadType(ctx context.Context, typeName string) (*pgtype.Type, error) { + var oid uint32 + + err := c.QueryRow(ctx, "select $1::text::regtype::oid;", typeName).Scan(&oid) + if err != nil { + return nil, err + } + + var typtype string + var typbasetype uint32 + + err = c.QueryRow(ctx, "select typtype::text, typbasetype from pg_type where oid=$1", oid).Scan(&typtype, &typbasetype) + if err != nil { + return nil, err + } + + switch typtype { + case "b": // array + elementOID, err := c.getArrayElementOID(ctx, oid) + if err != nil { + return nil, err + } + + dt, ok := c.TypeMap().TypeForOID(elementOID) + if !ok { + return nil, errors.New("array element OID not registered") + } + + return &pgtype.Type{Name: typeName, OID: oid, Codec: &pgtype.ArrayCodec{ElementType: dt}}, nil + case "c": // composite + fields, err := c.getCompositeFields(ctx, oid) + if err != nil { + return nil, err + } + + return &pgtype.Type{Name: typeName, OID: oid, Codec: &pgtype.CompositeCodec{Fields: fields}}, nil + case "d": // domain + dt, ok := c.TypeMap().TypeForOID(typbasetype) + if !ok { + return nil, errors.New("domain base type OID not registered") + } + + return &pgtype.Type{Name: typeName, OID: oid, Codec: dt.Codec}, nil + case "e": // enum + return &pgtype.Type{Name: typeName, OID: oid, Codec: &pgtype.EnumCodec{}}, nil + case "r": // range + elementOID, err := c.getRangeElementOID(ctx, oid) + if err != nil { + return nil, err + } + + dt, ok := c.TypeMap().TypeForOID(elementOID) + if !ok { + return nil, errors.New("range element OID not registered") + } + + return &pgtype.Type{Name: typeName, OID: oid, Codec: &pgtype.RangeCodec{ElementType: dt}}, nil + case "m": // multirange + elementOID, err := c.getMultiRangeElementOID(ctx, oid) + if err != nil { + return nil, err + } + + dt, ok := c.TypeMap().TypeForOID(elementOID) + if !ok { + return nil, errors.New("multirange element OID not registered") + } + + return &pgtype.Type{Name: typeName, OID: oid, Codec: &pgtype.MultirangeCodec{ElementType: dt}}, nil + default: + return &pgtype.Type{}, errors.New("unknown typtype") + } +} + +func (c *Conn) getArrayElementOID(ctx context.Context, oid uint32) (uint32, error) { + var typelem uint32 + + err := c.QueryRow(ctx, "select typelem from pg_type where oid=$1", oid).Scan(&typelem) + if err != nil { + return 0, err + } + + return typelem, nil +} + +func (c *Conn) getRangeElementOID(ctx context.Context, oid uint32) (uint32, error) { + var typelem uint32 + + err := c.QueryRow(ctx, "select rngsubtype from pg_range where rngtypid=$1", oid).Scan(&typelem) + if err != nil { + return 0, err + } + + return typelem, nil +} + +func (c *Conn) getMultiRangeElementOID(ctx context.Context, oid uint32) (uint32, error) { + var typelem uint32 + + err := c.QueryRow(ctx, "select rngtypid from pg_range where rngmultitypid=$1", oid).Scan(&typelem) + if err != nil { + return 0, err + } + + return typelem, nil +} + +func (c *Conn) getCompositeFields(ctx context.Context, oid uint32) ([]pgtype.CompositeCodecField, error) { + var typrelid uint32 + + err := c.QueryRow(ctx, "select typrelid from pg_type where oid=$1", oid).Scan(&typrelid) + if err != nil { + return nil, err + } + + var fields []pgtype.CompositeCodecField + var fieldName string + var fieldOID uint32 + rows, _ := c.Query(ctx, `select attname, atttypid +from pg_attribute +where attrelid=$1 + and not attisdropped + and attnum > 0 +order by attnum`, + typrelid, + ) + _, err = ForEachRow(rows, []any{&fieldName, &fieldOID}, func() error { + dt, ok := c.TypeMap().TypeForOID(fieldOID) + if !ok { + return fmt.Errorf("unknown composite type field OID: %v", fieldOID) + } + fields = append(fields, pgtype.CompositeCodecField{Name: fieldName, Type: dt}) + return nil + }) + if err != nil { + return nil, err + } + + return fields, nil +} + +func (c *Conn) deallocateInvalidatedCachedStatements(ctx context.Context) error { + if txStatus := c.pgConn.TxStatus(); txStatus != 'I' && txStatus != 'T' { + return nil + } + + if c.descriptionCache != nil { + c.descriptionCache.RemoveInvalidated() + } + + var invalidatedStatements []*pgconn.StatementDescription + if c.statementCache != nil { + invalidatedStatements = c.statementCache.GetInvalidated() + } + + if len(invalidatedStatements) == 0 { + return nil + } + + pipeline := c.pgConn.StartPipeline(ctx) + defer pipeline.Close() + + for _, sd := range invalidatedStatements { + pipeline.SendDeallocate(sd.Name) + } + + err := pipeline.Sync() + if err != nil { + return fmt.Errorf("failed to deallocate cached statement(s): %w", err) + } + + err = pipeline.Close() + if err != nil { + return fmt.Errorf("failed to deallocate cached statement(s): %w", err) + } + + c.statementCache.RemoveInvalidated() + for _, sd := range invalidatedStatements { + delete(c.preparedStatements, sd.Name) + } + + return nil +} diff --git a/vendor/github.com/jackc/pgx/v5/copy_from.go b/vendor/github.com/jackc/pgx/v5/copy_from.go new file mode 100644 index 0000000000..038c568cfe --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/copy_from.go @@ -0,0 +1,276 @@ +package pgx + +import ( + "bytes" + "context" + "fmt" + "io" + + "github.com/jackc/pgx/v5/internal/pgio" + "github.com/jackc/pgx/v5/pgconn" +) + +// CopyFromRows returns a [CopyFromSource] interface over the provided rows slice +// making it usable by [Conn.CopyFrom]. +func CopyFromRows(rows [][]any) CopyFromSource { + return ©FromRows{rows: rows, idx: -1} +} + +type copyFromRows struct { + rows [][]any + idx int +} + +func (ctr *copyFromRows) Next() bool { + ctr.idx++ + return ctr.idx < len(ctr.rows) +} + +func (ctr *copyFromRows) Values() ([]any, error) { + return ctr.rows[ctr.idx], nil +} + +func (ctr *copyFromRows) Err() error { + return nil +} + +// CopyFromSlice returns a [CopyFromSource] interface over a dynamic func +// making it usable by [Conn.CopyFrom]. +func CopyFromSlice(length int, next func(int) ([]any, error)) CopyFromSource { + return ©FromSlice{next: next, idx: -1, len: length} +} + +type copyFromSlice struct { + next func(int) ([]any, error) + idx int + len int + err error +} + +func (cts *copyFromSlice) Next() bool { + cts.idx++ + return cts.idx < cts.len +} + +func (cts *copyFromSlice) Values() ([]any, error) { + values, err := cts.next(cts.idx) + if err != nil { + cts.err = err + } + return values, err +} + +func (cts *copyFromSlice) Err() error { + return cts.err +} + +// CopyFromFunc returns a [CopyFromSource] interface that relies on nxtf for values. +// nxtf returns rows until it either signals an 'end of data' by returning row=nil and err=nil, +// or it returns an error. If nxtf returns an error, the copy is aborted. +func CopyFromFunc(nxtf func() (row []any, err error)) CopyFromSource { + return ©FromFunc{next: nxtf} +} + +type copyFromFunc struct { + next func() ([]any, error) + valueRow []any + err error +} + +func (g *copyFromFunc) Next() bool { + g.valueRow, g.err = g.next() + // only return true if valueRow exists and no error + return g.valueRow != nil && g.err == nil +} + +func (g *copyFromFunc) Values() ([]any, error) { + return g.valueRow, g.err +} + +func (g *copyFromFunc) Err() error { + return g.err +} + +// CopyFromSource is the interface used by [Conn.CopyFrom] as the source for copy data. +type CopyFromSource interface { + // Next returns true if there is another row and makes the next row data + // available to Values(). When there are no more rows available or an error + // has occurred it returns false. + Next() bool + + // Values returns the values for the current row. + Values() ([]any, error) + + // Err returns any error that has been encountered by the CopyFromSource. If + // this is not nil *Conn.CopyFrom will abort the copy. + Err() error +} + +type copyFrom struct { + conn *Conn + tableName Identifier + columnNames []string + rowSrc CopyFromSource + readerErrChan chan error + mode QueryExecMode +} + +func (ct *copyFrom) run(ctx context.Context) (int64, error) { + if ct.conn.copyFromTracer != nil { + ctx = ct.conn.copyFromTracer.TraceCopyFromStart(ctx, ct.conn, TraceCopyFromStartData{ + TableName: ct.tableName, + ColumnNames: ct.columnNames, + }) + } + + quotedTableName := ct.tableName.Sanitize() + cbuf := &bytes.Buffer{} + for i, cn := range ct.columnNames { + if i != 0 { + cbuf.WriteString(", ") + } + cbuf.WriteString(quoteIdentifier(cn)) + } + quotedColumnNames := cbuf.String() + + var sd *pgconn.StatementDescription + switch ct.mode { + case QueryExecModeExec, QueryExecModeSimpleProtocol: + // These modes don't support the binary format. Before the inclusion of the + // QueryExecModes, Conn.Prepare was called on every COPY operation to get + // the OIDs. These prepared statements were not cached. + // + // Since that's the same behavior provided by QueryExecModeDescribeExec, + // we'll default to that mode. + ct.mode = QueryExecModeDescribeExec + fallthrough + case QueryExecModeCacheStatement, QueryExecModeCacheDescribe, QueryExecModeDescribeExec: + var err error + sd, err = ct.conn.getStatementDescription( + ctx, + ct.mode, + fmt.Sprintf("select %s from %s", quotedColumnNames, quotedTableName), + ) + if err != nil { + return 0, fmt.Errorf("statement description failed: %w", err) + } + default: + return 0, fmt.Errorf("unknown QueryExecMode: %v", ct.mode) + } + + r, w := io.Pipe() + doneChan := make(chan struct{}) + + go func() { + defer close(doneChan) + + // Purposely NOT using defer w.Close(). See https://github.com/golang/go/issues/24283. + buf := ct.conn.wbuf + + buf = append(buf, "PGCOPY\n\377\r\n\000"...) + buf = pgio.AppendInt32(buf, 0) + buf = pgio.AppendInt32(buf, 0) + + moreRows := true + for moreRows { + var err error + moreRows, buf, err = ct.buildCopyBuf(buf, sd) + if err != nil { + w.CloseWithError(err) + return + } + + if ct.rowSrc.Err() != nil { + w.CloseWithError(ct.rowSrc.Err()) + return + } + + if len(buf) > 0 { + _, err = w.Write(buf) + if err != nil { + w.Close() + return + } + } + + buf = buf[:0] + } + + w.Close() + }() + + commandTag, err := ct.conn.pgConn.CopyFrom(ctx, r, fmt.Sprintf("copy %s ( %s ) from stdin binary;", quotedTableName, quotedColumnNames)) + + r.Close() + <-doneChan + + if ct.conn.copyFromTracer != nil { + ct.conn.copyFromTracer.TraceCopyFromEnd(ctx, ct.conn, TraceCopyFromEndData{ + CommandTag: commandTag, + Err: err, + }) + } + + return commandTag.RowsAffected(), err +} + +func (ct *copyFrom) buildCopyBuf(buf []byte, sd *pgconn.StatementDescription) (bool, []byte, error) { + const sendBufSize = 65536 - 5 // The packet has a 5-byte header + lastBufLen := 0 + largestRowLen := 0 + + for ct.rowSrc.Next() { + lastBufLen = len(buf) + + values, err := ct.rowSrc.Values() + if err != nil { + return false, nil, err + } + if len(values) != len(ct.columnNames) { + return false, nil, fmt.Errorf("expected %d values, got %d values", len(ct.columnNames), len(values)) + } + + buf = pgio.AppendInt16(buf, int16(len(ct.columnNames))) + for i, val := range values { + buf, err = encodeCopyValue(ct.conn.typeMap, buf, sd.Fields[i].DataTypeOID, val) + if err != nil { + return false, nil, err + } + } + + rowLen := len(buf) - lastBufLen + if rowLen > largestRowLen { + largestRowLen = rowLen + } + + // Try not to overflow size of the buffer PgConn.CopyFrom will be reading into. If that happens then the nature of + // io.Pipe means that the next Read will be short. This can lead to pathological send sizes such as 65531, 13, 65531 + // 13, 65531, 13, 65531, 13. + if len(buf) > sendBufSize-largestRowLen { + return true, buf, nil + } + } + + return false, buf, nil +} + +// CopyFrom uses the PostgreSQL copy protocol to perform bulk data insertion. It returns the number of rows copied and +// an error. +// +// CopyFrom requires all values use the binary format. A pgtype.Type that supports the binary format must be registered +// for the type of each column. Almost all types implemented by pgx support the binary format. +// +// Even though enum types appear to be strings they still must be registered to use with [Conn.CopyFrom]. This can be done with +// [Conn.LoadType] and [pgtype.Map.RegisterType]. +func (c *Conn) CopyFrom(ctx context.Context, tableName Identifier, columnNames []string, rowSrc CopyFromSource) (int64, error) { + ct := ©From{ + conn: c, + tableName: tableName, + columnNames: columnNames, + rowSrc: rowSrc, + readerErrChan: make(chan error), + mode: c.config.DefaultQueryExecMode, + } + + return ct.run(ctx) +} diff --git a/vendor/github.com/jackc/pgx/v5/derived_types.go b/vendor/github.com/jackc/pgx/v5/derived_types.go new file mode 100644 index 0000000000..3916006bee --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/derived_types.go @@ -0,0 +1,261 @@ +package pgx + +import ( + "context" + "fmt" + "regexp" + "strconv" + "strings" + + "github.com/jackc/pgx/v5/pgtype" +) + +/* +buildLoadDerivedTypesSQL generates the correct query for retrieving type information. + + pgVersion: the major version of the PostgreSQL server + typeNames: the names of the types to load. If nil, load all types. +*/ +func buildLoadDerivedTypesSQL(pgVersion int64, typeNames []string) string { + supportsMultirange := (pgVersion >= 14) + var typeNamesClause string + + if typeNames == nil { + // This should not occur; this will not return any types + typeNamesClause = "= ''" + } else { + typeNamesClause = "= ANY($1::text[])" + } + parts := make([]string, 0, 10) + + // Each of the type names provided might be found in pg_class or pg_type. + // Additionally, it may or may not include a schema portion. + parts = append(parts, ` +WITH RECURSIVE +-- find the OIDs in pg_class which match one of the provided type names +selected_classes(oid,reltype) AS ( + -- this query uses the namespace search path, so will match type names without a schema prefix + SELECT pg_class.oid, pg_class.reltype + FROM pg_catalog.pg_class + LEFT JOIN pg_catalog.pg_namespace n ON n.oid = pg_class.relnamespace + WHERE pg_catalog.pg_table_is_visible(pg_class.oid) + AND relname `, typeNamesClause, ` +UNION ALL + -- this query will only match type names which include the schema prefix + SELECT pg_class.oid, pg_class.reltype + FROM pg_class + INNER JOIN pg_namespace ON (pg_class.relnamespace = pg_namespace.oid) + WHERE nspname || '.' || relname `, typeNamesClause, ` +), +selected_types(oid) AS ( + -- collect the OIDs from pg_types which correspond to the selected classes + SELECT reltype AS oid + FROM selected_classes +UNION ALL + -- as well as any other type names which match our criteria + SELECT pg_type.oid + FROM pg_type + LEFT OUTER JOIN pg_namespace ON (pg_type.typnamespace = pg_namespace.oid) + WHERE typname `, typeNamesClause, ` + OR nspname || '.' || typname `, typeNamesClause, ` +), +-- this builds a parent/child mapping of objects, allowing us to know +-- all the child (ie: dependent) types that a parent (type) requires +-- As can be seen, there are 3 ways this can occur (the last of which +-- is due to being a composite class, where the composite fields are children) +pc(parent, child) AS ( + SELECT parent.oid, parent.typelem + FROM pg_type parent + WHERE parent.typtype = 'b' AND parent.typelem != 0 +UNION ALL + SELECT parent.oid, parent.typbasetype + FROM pg_type parent + WHERE parent.typtypmod = -1 AND parent.typbasetype != 0 +UNION ALL + SELECT pg_type.oid, atttypid + FROM pg_attribute + INNER JOIN pg_class ON (pg_class.oid = pg_attribute.attrelid) + INNER JOIN pg_type ON (pg_type.oid = pg_class.reltype) + WHERE NOT attisdropped + AND attnum > 0 +), +-- Now construct a recursive query which includes a 'depth' element. +-- This is used to ensure that the "youngest" children are registered before +-- their parents. +relationships(parent, child, depth) AS ( + SELECT DISTINCT 0::OID, selected_types.oid, 0 + FROM selected_types +UNION ALL + SELECT pg_type.oid AS parent, pg_attribute.atttypid AS child, 1 + FROM selected_classes c + inner join pg_type ON (c.reltype = pg_type.oid) + inner join pg_attribute on (c.oid = pg_attribute.attrelid) +UNION ALL + SELECT pc.parent, pc.child, relationships.depth + 1 + FROM pc + INNER JOIN relationships ON (pc.parent = relationships.child) +), +-- composite fields need to be encapsulated as a couple of arrays to provide the required information for registration +composite AS ( + SELECT pg_type.oid, ARRAY_AGG(attname ORDER BY attnum) AS attnames, ARRAY_AGG(atttypid ORDER BY ATTNUM) AS atttypids + FROM pg_attribute + INNER JOIN pg_class ON (pg_class.oid = pg_attribute.attrelid) + INNER JOIN pg_type ON (pg_type.oid = pg_class.reltype) + WHERE NOT attisdropped + AND attnum > 0 + GROUP BY pg_type.oid +) +-- Bring together this information, showing all the information which might possibly be required +-- to complete the registration, applying filters to only show the items which relate to the selected +-- types/classes. +SELECT typname, + pg_namespace.nspname, + typtype, + typbasetype, + typelem, + pg_type.oid,`) + if supportsMultirange { + parts = append(parts, ` + COALESCE(multirange.rngtypid, 0) AS rngtypid,`) + } else { + parts = append(parts, ` + 0 AS rngtypid,`) + } + parts = append(parts, ` + COALESCE(pg_range.rngsubtype, 0) AS rngsubtype, + attnames, atttypids + FROM relationships + INNER JOIN pg_type ON (pg_type.oid = relationships.child) + LEFT OUTER JOIN pg_range ON (pg_type.oid = pg_range.rngtypid)`) + if supportsMultirange { + parts = append(parts, ` + LEFT OUTER JOIN pg_range multirange ON (pg_type.oid = multirange.rngmultitypid)`) + } + + parts = append(parts, ` + LEFT OUTER JOIN composite USING (oid) + LEFT OUTER JOIN pg_namespace ON (pg_type.typnamespace = pg_namespace.oid) + WHERE NOT (typtype = 'b' AND typelem = 0)`) + parts = append(parts, ` + GROUP BY typname, pg_namespace.nspname, typtype, typbasetype, typelem, pg_type.oid, pg_range.rngsubtype,`) + if supportsMultirange { + parts = append(parts, ` + multirange.rngtypid,`) + } + parts = append(parts, ` + attnames, atttypids + ORDER BY MAX(depth) desc, typname;`) + return strings.Join(parts, "") +} + +type derivedTypeInfo struct { + Oid, Typbasetype, Typelem, Rngsubtype, Rngtypid uint32 + TypeName, Typtype, NspName string + Attnames []string + Atttypids []uint32 +} + +// LoadTypes performs a single (complex) query, returning all the required +// information to register the named types, as well as any other types directly +// or indirectly required to complete the registration. +// The result of this call can be passed into RegisterTypes to complete the process. +func (c *Conn) LoadTypes(ctx context.Context, typeNames []string) ([]*pgtype.Type, error) { + m := c.TypeMap() + if len(typeNames) == 0 { + return nil, fmt.Errorf("No type names were supplied.") + } + + // Disregard server version errors. This will result in + // the SQL not support recent structures such as multirange + serverVersion, _ := serverVersion(c) + sql := buildLoadDerivedTypesSQL(serverVersion, typeNames) + rows, err := c.Query(ctx, sql, QueryResultFormats{TextFormatCode}, typeNames) + if err != nil { + return nil, fmt.Errorf("While generating load types query: %w", err) + } + defer rows.Close() + result := make([]*pgtype.Type, 0, 100) + for rows.Next() { + ti := derivedTypeInfo{} + err = rows.Scan(&ti.TypeName, &ti.NspName, &ti.Typtype, &ti.Typbasetype, &ti.Typelem, &ti.Oid, &ti.Rngtypid, &ti.Rngsubtype, &ti.Attnames, &ti.Atttypids) + if err != nil { + return nil, fmt.Errorf("While scanning type information: %w", err) + } + var type_ *pgtype.Type + switch ti.Typtype { + case "b": // array + dt, ok := m.TypeForOID(ti.Typelem) + if !ok { + return nil, fmt.Errorf("Array element OID %v not registered while loading pgtype %q", ti.Typelem, ti.TypeName) + } + type_ = &pgtype.Type{Name: ti.TypeName, OID: ti.Oid, Codec: &pgtype.ArrayCodec{ElementType: dt}} + case "c": // composite + var fields []pgtype.CompositeCodecField + for i, fieldName := range ti.Attnames { + dt, ok := m.TypeForOID(ti.Atttypids[i]) + if !ok { + return nil, fmt.Errorf("Unknown field for composite type %q: field %q (OID %v) is not already registered.", ti.TypeName, fieldName, ti.Atttypids[i]) + } + fields = append(fields, pgtype.CompositeCodecField{Name: fieldName, Type: dt}) + } + + type_ = &pgtype.Type{Name: ti.TypeName, OID: ti.Oid, Codec: &pgtype.CompositeCodec{Fields: fields}} + case "d": // domain + dt, ok := m.TypeForOID(ti.Typbasetype) + if !ok { + return nil, fmt.Errorf("Domain base type OID %v was not already registered, needed for %q", ti.Typbasetype, ti.TypeName) + } + + type_ = &pgtype.Type{Name: ti.TypeName, OID: ti.Oid, Codec: dt.Codec} + case "e": // enum + type_ = &pgtype.Type{Name: ti.TypeName, OID: ti.Oid, Codec: &pgtype.EnumCodec{}} + case "r": // range + dt, ok := m.TypeForOID(ti.Rngsubtype) + if !ok { + return nil, fmt.Errorf("Range element OID %v was not already registered, needed for %q", ti.Rngsubtype, ti.TypeName) + } + + type_ = &pgtype.Type{Name: ti.TypeName, OID: ti.Oid, Codec: &pgtype.RangeCodec{ElementType: dt}} + case "m": // multirange + dt, ok := m.TypeForOID(ti.Rngtypid) + if !ok { + return nil, fmt.Errorf("Multirange element OID %v was not already registered, needed for %q", ti.Rngtypid, ti.TypeName) + } + + type_ = &pgtype.Type{Name: ti.TypeName, OID: ti.Oid, Codec: &pgtype.MultirangeCodec{ElementType: dt}} + default: + return nil, fmt.Errorf("Unknown typtype %q was found while registering %q", ti.Typtype, ti.TypeName) + } + + // the type_ is impossible to be null + m.RegisterType(type_) + if ti.NspName != "" { + nspType := &pgtype.Type{Name: ti.NspName + "." + type_.Name, OID: type_.OID, Codec: type_.Codec} + m.RegisterType(nspType) + result = append(result, nspType) + } + result = append(result, type_) + } + + if err := rows.Err(); err != nil { + return nil, fmt.Errorf("While processing rows: %w", err) + } + + return result, nil +} + +// serverVersion returns the postgresql server version. +func serverVersion(c *Conn) (int64, error) { + serverVersionStr := c.PgConn().ParameterStatus("server_version") + serverVersionStr = regexp.MustCompile(`^[0-9]+`).FindString(serverVersionStr) + // if not PostgreSQL do nothing + if serverVersionStr == "" { + return 0, fmt.Errorf("Cannot identify server version in %q", serverVersionStr) + } + + version, err := strconv.ParseInt(serverVersionStr, 10, 64) + if err != nil { + return 0, fmt.Errorf("postgres version parsing failed: %w", err) + } + return version, nil +} diff --git a/vendor/github.com/jackc/pgx/v5/doc.go b/vendor/github.com/jackc/pgx/v5/doc.go new file mode 100644 index 0000000000..5e4870191a --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/doc.go @@ -0,0 +1,220 @@ +// Package pgx is a PostgreSQL database driver. +/* +pgx provides a native PostgreSQL driver and can act as a [database/sql/driver]. The native PostgreSQL interface is similar +to the [database/sql] interface while providing better speed and access to PostgreSQL specific features. Use +[github.com/jackc/pgx/v5/stdlib] to use pgx as a database/sql compatible driver. See that package's documentation for +details. + +Establishing a Connection + +The primary way of establishing a connection is with [pgx.Connect]: + + conn, err := pgx.Connect(context.Background(), os.Getenv("DATABASE_URL")) + +The database connection string can be in URL or key/value format. Both PostgreSQL settings and pgx settings can be +specified here. In addition, a config struct can be created by [ParseConfig] and modified before establishing the +connection with [ConnectConfig] to configure settings such as tracing that cannot be configured with a connection +string. + +Connection Pool + +[*pgx.Conn] represents a single connection to the database and is not concurrency safe. Use package +[github.com/jackc/pgx/v5/pgxpool] for a concurrency safe connection pool. + +Query Interface + +pgx implements [Conn.Query] in the familiar database/sql style. However, pgx provides generic functions such as [CollectRows] and +[ForEachRow] that are a simpler and safer way of processing rows than manually calling defer [Rows.Close], [Rows.Next], +[Rows.Scan], and [Rows.Err]. + +[CollectRows] can be used collect all returned rows into a slice. + + rows, _ := conn.Query(context.Background(), "select generate_series(1,$1)", 5) + numbers, err := pgx.CollectRows(rows, pgx.RowTo[int32]) + if err != nil { + return err + } + // numbers => [1 2 3 4 5] + +[ForEachRow] can be used to execute a callback function for every row. This is often easier than iterating over rows +directly. + + var sum, n int32 + rows, _ := conn.Query(context.Background(), "select generate_series(1,$1)", 10) + _, err := pgx.ForEachRow(rows, []any{&n}, func() error { + sum += n + return nil + }) + if err != nil { + return err + } + +pgx also implements [Conn.QueryRow] in the same style as database/sql. + + var name string + var weight int64 + err := conn.QueryRow(context.Background(), "select name, weight from widgets where id=$1", 42).Scan(&name, &weight) + if err != nil { + return err + } + +Use [Conn.Exec] to execute a query that does not return a result set. + + commandTag, err := conn.Exec(context.Background(), "delete from widgets where id=$1", 42) + if err != nil { + return err + } + if commandTag.RowsAffected() != 1 { + return errors.New("No row found to delete") + } + +PostgreSQL Data Types + +pgx uses the [pgtype] package to converting Go values to and from PostgreSQL values. It supports many PostgreSQL types +directly and is customizable and extendable. User defined data types such as enums, domains, and composite types may +require type registration. See that package's documentation for details. + +PostgreSQL arrays (including results from set-returning aggregates such as array_agg) +can be scanned directly into a matching Go slice. For scalar columns, pass the slice +as the scan destination: + + var ids []int64 + err := conn.QueryRow(ctx, "select array_agg(id) from things").Scan(&ids) + +For a column that is part of a row, combine the slice with the usual row-to-struct +helpers. A struct field of slice type will pick up the array_agg column when +collected via [CollectRows] and [RowToStructByName] (or +[RowToAddrOfStructByPos]): + + type ThingEntry struct { + GroupID int64 + ThingIDs []int64 `db:"thing_ids"` + } + + rows, _ := conn.Query(ctx, + "select group_id, array_agg(thing_id) as thing_ids from things group by group_id") + entries, err := pgx.CollectRows(rows, pgx.RowToStructByName[ThingEntry]) + +Transactions + +Transactions are started by calling [Conn.Begin]. + + tx, err := conn.Begin(context.Background()) + if err != nil { + return err + } + // Rollback is safe to call even if the tx is already closed, so if + // the tx commits successfully, this is a no-op + defer tx.Rollback(context.Background()) + + _, err = tx.Exec(context.Background(), "insert into foo(id) values (1)") + if err != nil { + return err + } + + err = tx.Commit(context.Background()) + if err != nil { + return err + } + +The [Tx] returned from [Conn.Begin] also implements the [Tx.Begin] method. This can be used to implement pseudo nested transactions. +These are internally implemented with savepoints. + +Use [Conn.BeginTx] to control the transaction mode. [Conn.BeginTx] also can be used to ensure a new transaction is created instead of +a pseudo nested transaction. + +[BeginFunc] and [BeginTxFunc] are functions that begin a transaction, execute a function, and commit or rollback the +transaction depending on the return value of the function. These can be simpler and less error prone to use. + + err = pgx.BeginFunc(context.Background(), conn, func(tx pgx.Tx) error { + _, err := tx.Exec(context.Background(), "insert into foo(id) values (1)") + return err + }) + if err != nil { + return err + } + +Prepared Statements + +Prepared statements can be manually created with the [Conn.Prepare] method. However, this is rarely necessary because pgx +includes an automatic statement cache by default. Queries run through the normal [Conn.Query], [Conn.QueryRow], and [Conn.Exec] +functions are automatically prepared on first execution and the prepared statement is reused on subsequent executions. +See [ParseConfig] for information on how to customize or disable the statement cache. + +Copy Protocol + +Use [Conn.CopyFrom] to efficiently insert multiple rows at a time using the PostgreSQL copy protocol. [Conn.CopyFrom] accepts a +[CopyFromSource] interface. If the data is already in a [][]any use [CopyFromRows] to wrap it in a [CopyFromSource] interface. +Or implement [CopyFromSource] to avoid buffering the entire data set in memory. + + rows := [][]any{ + {"John", "Smith", int32(36)}, + {"Jane", "Doe", int32(29)}, + } + + copyCount, err := conn.CopyFrom( + context.Background(), + pgx.Identifier{"people"}, + []string{"first_name", "last_name", "age"}, + pgx.CopyFromRows(rows), + ) + +When you already have a typed array using [CopyFromSlice] can be more convenient. + + rows := []User{ + {"John", "Smith", 36}, + {"Jane", "Doe", 29}, + } + + copyCount, err := conn.CopyFrom( + context.Background(), + pgx.Identifier{"people"}, + []string{"first_name", "last_name", "age"}, + pgx.CopyFromSlice(len(rows), func(i int) ([]any, error) { + return []any{rows[i].FirstName, rows[i].LastName, rows[i].Age}, nil + }), + ) + +CopyFrom can be faster than an insert with as few as 5 rows. + +Listen and Notify + +pgx can listen to the PostgreSQL notification system with the [Conn.WaitForNotification] method. It blocks until a +notification is received or the context is canceled. + + _, err := conn.Exec(context.Background(), "listen channelname") + if err != nil { + return err + } + + notification, err := conn.WaitForNotification(context.Background()) + if err != nil { + return err + } + // do something with notification + + +Tracing and Logging + +pgx supports tracing by setting [ConnConfig.Tracer]. To combine several tracers you can use the [github.com/jackc/pgx/v5/multitracer.Tracer]. + +In addition, the [github.com/jackc/pgx/v5/tracelog] package provides the [github.com/jackc/pgx/v5/tracelog.TraceLog] type which lets a +traditional logger act as a [QueryTracer]. + +For debug tracing of the actual PostgreSQL wire protocol messages see [github.com/jackc/pgx/v5/pgproto3]. + +Lower Level PostgreSQL Functionality + +[github.com/jackc/pgx/v5/pgconn] contains a lower level PostgreSQL driver roughly at the level of libpq. [Conn] is +implemented on top of [pgconn.PgConn]. The [Conn.PgConn] method can be used to access this lower layer. + +PgBouncer + +By default pgx automatically uses prepared statements. Prepared statements are incompatible with PgBouncer. This can be +disabled by setting a different [QueryExecMode] in [ConnConfig.DefaultQueryExecMode]. +*/ +package pgx + +import ( + _ "github.com/jackc/pgx/v5/pgconn" // Just for allowing godoc to resolve "pgconn" +) diff --git a/vendor/github.com/jackc/pgx/v5/extended_query_builder.go b/vendor/github.com/jackc/pgx/v5/extended_query_builder.go new file mode 100644 index 0000000000..526b0e953b --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/extended_query_builder.go @@ -0,0 +1,146 @@ +package pgx + +import ( + "fmt" + + "github.com/jackc/pgx/v5/pgconn" + "github.com/jackc/pgx/v5/pgtype" +) + +// ExtendedQueryBuilder is used to choose the parameter formats, to format the parameters and to choose the result +// formats for an extended query. +type ExtendedQueryBuilder struct { + ParamValues [][]byte + paramValueBytes []byte + ParamFormats []int16 + ResultFormats []int16 +} + +// Build sets ParamValues, ParamFormats, and ResultFormats for use with *PgConn.ExecParams or *PgConn.ExecPrepared. If +// sd is nil then QueryExecModeExec behavior will be used. +func (eqb *ExtendedQueryBuilder) Build(m *pgtype.Map, sd *pgconn.StatementDescription, args []any) error { + eqb.reset() + + if sd == nil { + for i := range args { + err := eqb.appendParam(m, 0, pgtype.TextFormatCode, args[i]) + if err != nil { + err = fmt.Errorf("failed to encode args[%d]: %w", i, err) + return err + } + } + return nil + } + + if len(sd.ParamOIDs) != len(args) { + return fmt.Errorf("mismatched param and argument count") + } + + for i := range args { + err := eqb.appendParam(m, sd.ParamOIDs[i], -1, args[i]) + if err != nil { + err = fmt.Errorf("failed to encode args[%d]: %w", i, err) + return err + } + } + + for i := range sd.Fields { + eqb.appendResultFormat(m.FormatCodeForOID(sd.Fields[i].DataTypeOID)) + } + + return nil +} + +// appendParam appends a parameter to the query. format may be -1 to automatically choose the format. If arg is nil it +// must be an untyped nil. +func (eqb *ExtendedQueryBuilder) appendParam(m *pgtype.Map, oid uint32, format int16, arg any) error { + if format == -1 { + preferredFormat := eqb.chooseParameterFormatCode(m, oid, arg) + preferredErr := eqb.appendParam(m, oid, preferredFormat, arg) + if preferredErr == nil { + return nil + } + + var otherFormat int16 + if preferredFormat == TextFormatCode { + otherFormat = BinaryFormatCode + } else { + otherFormat = TextFormatCode + } + + otherErr := eqb.appendParam(m, oid, otherFormat, arg) + if otherErr == nil { + return nil + } + + return preferredErr // return the error from the preferred format + } + + v, err := eqb.encodeExtendedParamValue(m, oid, format, arg) + if err != nil { + return err + } + + eqb.ParamFormats = append(eqb.ParamFormats, format) + eqb.ParamValues = append(eqb.ParamValues, v) + + return nil +} + +// appendResultFormat appends a result format to the query. +func (eqb *ExtendedQueryBuilder) appendResultFormat(format int16) { + eqb.ResultFormats = append(eqb.ResultFormats, format) +} + +// reset readies eqb to build another query. +func (eqb *ExtendedQueryBuilder) reset() { + eqb.ParamValues = eqb.ParamValues[0:0] + eqb.paramValueBytes = eqb.paramValueBytes[0:0] + eqb.ParamFormats = eqb.ParamFormats[0:0] + eqb.ResultFormats = eqb.ResultFormats[0:0] + + if cap(eqb.ParamValues) > 64 { + eqb.ParamValues = make([][]byte, 0, 64) + } + + if cap(eqb.paramValueBytes) > 256 { + eqb.paramValueBytes = make([]byte, 0, 256) + } + + if cap(eqb.ParamFormats) > 64 { + eqb.ParamFormats = make([]int16, 0, 64) + } + if cap(eqb.ResultFormats) > 64 { + eqb.ResultFormats = make([]int16, 0, 64) + } +} + +func (eqb *ExtendedQueryBuilder) encodeExtendedParamValue(m *pgtype.Map, oid uint32, formatCode int16, arg any) ([]byte, error) { + if eqb.paramValueBytes == nil { + eqb.paramValueBytes = make([]byte, 0, 128) + } + + pos := len(eqb.paramValueBytes) + + buf, err := m.Encode(oid, formatCode, arg, eqb.paramValueBytes) + if err != nil { + return nil, err + } + if buf == nil { + return nil, nil + } + eqb.paramValueBytes = buf + return eqb.paramValueBytes[pos:], nil +} + +// chooseParameterFormatCode determines the correct format code for an +// argument to a prepared statement. It defaults to TextFormatCode if no +// determination can be made. +func (eqb *ExtendedQueryBuilder) chooseParameterFormatCode(m *pgtype.Map, oid uint32, arg any) int16 { + switch arg.(type) { + case string, *string: + return TextFormatCode + } + + return m.FormatCodeForOID(oid) +} diff --git a/vendor/github.com/jackc/pgx/v5/internal/iobufpool/iobufpool.go b/vendor/github.com/jackc/pgx/v5/internal/iobufpool/iobufpool.go new file mode 100644 index 0000000000..abc41f657f --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/internal/iobufpool/iobufpool.go @@ -0,0 +1,78 @@ +// Package iobufpool implements a global segregated-fit pool of buffers for IO. +// +// It uses *[]byte instead of []byte to avoid the sync.Pool allocation with Put. Unfortunately, using a pointer to avoid +// an allocation is purposely not documented. https://github.com/golang/go/issues/16323 +package iobufpool + +import ( + "math/bits" + "sync" +) + +const minPoolExpOf2 = 8 + +var pools [18]*sync.Pool + +func init() { + for i := range pools { + bufLen := 1 << (minPoolExpOf2 + i) + pools[i] = &sync.Pool{ + New: func() any { + buf := make([]byte, bufLen) + return &buf + }, + } + } +} + +// Get gets a []byte of len size with cap <= size*2. +func Get(size int) *[]byte { + i := getPoolIdx(size) + if i >= len(pools) { + buf := make([]byte, size) + return &buf + } + + ptrBuf := (pools[i].Get().(*[]byte)) + *ptrBuf = (*ptrBuf)[:size] + + return ptrBuf +} + +func getPoolIdx(size int) int { + if size < 2 { + return 0 + } + idx := bits.Len(uint(size-1)) - minPoolExpOf2 + if idx < 0 { + return 0 + } + return idx +} + +// Put returns buf to the pool. +func Put(buf *[]byte) { + i := putPoolIdx(cap(*buf)) + if i < 0 { + return + } + + pools[i].Put(buf) +} + +func putPoolIdx(size int) int { + // Only exact power-of-2 sizes match pool buckets + if size&(size-1) != 0 { + return -1 + } + + // Calculate log2(size) using trailing zeros count + exp := bits.TrailingZeros(uint(size)) + idx := exp - minPoolExpOf2 + + if idx < 0 || idx >= len(pools) { + return -1 + } + + return idx +} diff --git a/vendor/github.com/jackc/pgx/v5/internal/pgio/README.md b/vendor/github.com/jackc/pgx/v5/internal/pgio/README.md new file mode 100644 index 0000000000..b2fc58014a --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/internal/pgio/README.md @@ -0,0 +1,6 @@ +# pgio + +Package pgio is a low-level toolkit building messages in the PostgreSQL wire protocol. + +pgio provides functions for appending integers to a []byte while doing byte +order conversion. diff --git a/vendor/github.com/jackc/pgx/v5/internal/pgio/doc.go b/vendor/github.com/jackc/pgx/v5/internal/pgio/doc.go new file mode 100644 index 0000000000..ef2dcc7f72 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/internal/pgio/doc.go @@ -0,0 +1,6 @@ +// Package pgio is a low-level toolkit building messages in the PostgreSQL wire protocol. +/* +pgio provides functions for appending integers to a []byte while doing byte +order conversion. +*/ +package pgio diff --git a/vendor/github.com/jackc/pgx/v5/internal/pgio/write.go b/vendor/github.com/jackc/pgx/v5/internal/pgio/write.go new file mode 100644 index 0000000000..3a6700dc4c --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/internal/pgio/write.go @@ -0,0 +1,32 @@ +package pgio + +func AppendUint16(buf []byte, n uint16) []byte { + return append(buf, byte(n>>8), byte(n)) +} + +func AppendUint32(buf []byte, n uint32) []byte { + return append(buf, byte(n>>24), byte(n>>16), byte(n>>8), byte(n)) +} + +func AppendUint64(buf []byte, n uint64) []byte { + return append(buf, + byte(n>>56), byte(n>>48), byte(n>>40), byte(n>>32), + byte(n>>24), byte(n>>16), byte(n>>8), byte(n), + ) +} + +func AppendInt16(buf []byte, n int16) []byte { + return AppendUint16(buf, uint16(n)) +} + +func AppendInt32(buf []byte, n int32) []byte { + return AppendUint32(buf, uint32(n)) +} + +func AppendInt64(buf []byte, n int64) []byte { + return AppendUint64(buf, uint64(n)) +} + +func SetInt32(buf []byte, n int32) { + *(*[4]byte)(buf) = [4]byte{byte(n >> 24), byte(n >> 16), byte(n >> 8), byte(n)} +} diff --git a/vendor/github.com/jackc/pgx/v5/internal/sanitize/benchmark.sh b/vendor/github.com/jackc/pgx/v5/internal/sanitize/benchmark.sh new file mode 100644 index 0000000000..b4ee3fe744 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/internal/sanitize/benchmark.sh @@ -0,0 +1,60 @@ +#!/usr/bin/env bash + +current_branch=$(git rev-parse --abbrev-ref HEAD) +if [ "$current_branch" == "HEAD" ]; then + current_branch=$(git rev-parse HEAD) +fi + +restore_branch() { + echo "Restoring original branch/commit: $current_branch" + git checkout "$current_branch" +} +trap restore_branch EXIT + +# Check if there are uncommitted changes +if ! git diff --quiet || ! git diff --cached --quiet; then + echo "There are uncommitted changes. Please commit or stash them before running this script." + exit 1 +fi + +# Ensure that at least one commit argument is passed +if [ "$#" -lt 1 ]; then + echo "Usage: $0 ... " + exit 1 +fi + +commits=("$@") +benchmarks_dir=benchmarks + +if ! mkdir -p "${benchmarks_dir}"; then + echo "Unable to create dir for benchmarks data" + exit 1 +fi + +# Benchmark results +bench_files=() + +# Run benchmark for each listed commit +for i in "${!commits[@]}"; do + commit="${commits[i]}" + git checkout "$commit" || { + echo "Failed to checkout $commit" + exit 1 + } + + # Sanitized commit message + commit_message=$(git log -1 --pretty=format:"%s" | tr -c '[:alnum:]-_' '_') + + # Benchmark data will go there + bench_file="${benchmarks_dir}/${i}_${commit_message}.bench" + + if ! go test -bench=. -count=10 >"$bench_file"; then + echo "Benchmarking failed for commit $commit" + exit 1 + fi + + bench_files+=("$bench_file") +done + +# go install golang.org/x/perf/cmd/benchstat[@latest] +benchstat "${bench_files[@]}" diff --git a/vendor/github.com/jackc/pgx/v5/internal/sanitize/sanitize.go b/vendor/github.com/jackc/pgx/v5/internal/sanitize/sanitize.go new file mode 100644 index 0000000000..033a4143b8 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/internal/sanitize/sanitize.go @@ -0,0 +1,541 @@ +package sanitize + +import ( + "bytes" + "encoding/hex" + "fmt" + "math" + "slices" + "strconv" + "strings" + "sync" + "time" + "unicode/utf8" +) + +// Part is either a string or an int. A string is raw SQL. An int is a +// argument placeholder. +type Part any + +type Query struct { + Parts []Part +} + +// utf.DecodeRune returns the utf8.RuneError for errors. But that is actually rune U+FFFD -- the unicode replacement +// character. utf8.RuneError is not an error if it is also width 3. +// +// https://github.com/jackc/pgx/issues/1380 +const replacementcharacterwidth = 3 + +const maxBufSize = 16384 // 16 Ki + +var bufPool = &pool[*bytes.Buffer]{ + new: func() *bytes.Buffer { + return &bytes.Buffer{} + }, + reset: func(b *bytes.Buffer) bool { + n := b.Len() + b.Reset() + return n < maxBufSize + }, +} + +var null = []byte("null") + +func (q *Query) Sanitize(args ...any) (string, error) { + argUse := make([]bool, len(args)) + buf := bufPool.get() + defer bufPool.put(buf) + + for _, part := range q.Parts { + switch part := part.(type) { + case string: + buf.WriteString(part) + case int: + argIdx := part - 1 + var p []byte + if argIdx < 0 { + return "", fmt.Errorf("first sql argument must be > 0") + } + + if argIdx >= len(args) { + return "", fmt.Errorf("insufficient arguments") + } + + // Prevent SQL injection via Line Comment Creation + // https://github.com/jackc/pgx/security/advisories/GHSA-m7wr-2xf7-cm9p + buf.WriteByte(' ') + + arg := args[argIdx] + switch arg := arg.(type) { + case nil: + p = null + case int64: + p = strconv.AppendInt(buf.AvailableBuffer(), arg, 10) + case float64: + p = strconv.AppendFloat(buf.AvailableBuffer(), arg, 'f', -1, 64) + case bool: + p = strconv.AppendBool(buf.AvailableBuffer(), arg) + case []byte: + p = QuoteBytes(buf.AvailableBuffer(), arg) + case string: + p = QuoteString(buf.AvailableBuffer(), arg) + case time.Time: + p = arg.Truncate(time.Microsecond). + AppendFormat(buf.AvailableBuffer(), "'2006-01-02 15:04:05.999999999Z07:00:00'") + default: + return "", fmt.Errorf("invalid arg type: %T", arg) + } + argUse[argIdx] = true + + buf.Write(p) + + // Prevent SQL injection via Line Comment Creation + // https://github.com/jackc/pgx/security/advisories/GHSA-m7wr-2xf7-cm9p + buf.WriteByte(' ') + default: + return "", fmt.Errorf("invalid Part type: %T", part) + } + } + + for i, used := range argUse { + if !used { + return "", fmt.Errorf("unused argument: %d", i) + } + } + return buf.String(), nil +} + +func NewQuery(sql string) (*Query, error) { + query := &Query{} + query.init(sql) + + return query, nil +} + +var sqlLexerPool = &pool[*sqlLexer]{ + new: func() *sqlLexer { + return &sqlLexer{} + }, + reset: func(sl *sqlLexer) bool { + *sl = sqlLexer{} + return true + }, +} + +func (q *Query) init(sql string) { + parts := q.Parts[:0] + if parts == nil { + // dirty, but fast heuristic to preallocate for ~90% usecases + n := strings.Count(sql, "$") + strings.Count(sql, "--") + 1 + parts = make([]Part, 0, n) + } + + l := sqlLexerPool.get() + defer sqlLexerPool.put(l) + + l.src = sql + l.stateFn = rawState + l.parts = parts + + for l.stateFn != nil { + l.stateFn = l.stateFn(l) + } + + q.Parts = l.parts +} + +func QuoteString(dst []byte, str string) []byte { + const quote = '\'' + + // Preallocate space for the worst case scenario + dst = slices.Grow(dst, len(str)*2+2) + + // Add opening quote + dst = append(dst, quote) + + // Iterate through the string without allocating + for i := 0; i < len(str); i++ { + if str[i] == quote { + dst = append(dst, quote, quote) + } else { + dst = append(dst, str[i]) + } + } + + // Add closing quote + dst = append(dst, quote) + + return dst +} + +func QuoteBytes(dst, buf []byte) []byte { + if len(buf) == 0 { + return append(dst, `'\x'`...) + } + + // Calculate required length + requiredLen := 3 + hex.EncodedLen(len(buf)) + 1 + + // Ensure dst has enough capacity + if cap(dst)-len(dst) < requiredLen { + newDst := make([]byte, len(dst), len(dst)+requiredLen) + copy(newDst, dst) + dst = newDst + } + + // Record original length and extend slice + origLen := len(dst) + dst = dst[:origLen+requiredLen] + + // Add prefix + dst[origLen] = '\'' + dst[origLen+1] = '\\' + dst[origLen+2] = 'x' + + // Encode bytes directly into dst + hex.Encode(dst[origLen+3:len(dst)-1], buf) + + // Add suffix + dst[len(dst)-1] = '\'' + + return dst +} + +type sqlLexer struct { + src string + start int + pos int + nested int // multiline comment nesting level. + dollarTag string // active tag while inside a dollar-quoted string (may be empty for $$). + stateFn stateFn + parts []Part +} + +type stateFn func(*sqlLexer) stateFn + +func rawState(l *sqlLexer) stateFn { + for { + r, width := utf8.DecodeRuneInString(l.src[l.pos:]) + l.pos += width + + switch r { + case 'e', 'E': + nextRune, width := utf8.DecodeRuneInString(l.src[l.pos:]) + if nextRune == '\'' { + l.pos += width + return escapeStringState + } + case '\'': + return singleQuoteState + case '"': + return doubleQuoteState + case '$': + nextRune, _ := utf8.DecodeRuneInString(l.src[l.pos:]) + if '0' <= nextRune && nextRune <= '9' { + if l.pos-l.start > 0 { + l.parts = append(l.parts, l.src[l.start:l.pos-width]) + } + l.start = l.pos + return placeholderState + } + // PostgreSQL dollar-quoted string: $[tag]$...$[tag]$. The $ was + // just consumed; try to match the rest of the opening tag. + // Without this, placeholders embedded inside dollar-quoted + // literals would be incorrectly substituted. + if tagLen, ok := scanDollarQuoteTag(l.src[l.pos:]); ok { + l.dollarTag = l.src[l.pos : l.pos+tagLen] + l.pos += tagLen + 1 // advance past tag and closing '$' + return dollarQuoteState + } + case '-': + nextRune, width := utf8.DecodeRuneInString(l.src[l.pos:]) + if nextRune == '-' { + l.pos += width + return oneLineCommentState + } + case '/': + nextRune, width := utf8.DecodeRuneInString(l.src[l.pos:]) + if nextRune == '*' { + l.pos += width + return multilineCommentState + } + case utf8.RuneError: + if width != replacementcharacterwidth { + if l.pos-l.start > 0 { + l.parts = append(l.parts, l.src[l.start:l.pos]) + l.start = l.pos + } + return nil + } + } + } +} + +func singleQuoteState(l *sqlLexer) stateFn { + for { + r, width := utf8.DecodeRuneInString(l.src[l.pos:]) + l.pos += width + + switch r { + case '\'': + nextRune, width := utf8.DecodeRuneInString(l.src[l.pos:]) + if nextRune != '\'' { + return rawState + } + l.pos += width + case utf8.RuneError: + if width != replacementcharacterwidth { + if l.pos-l.start > 0 { + l.parts = append(l.parts, l.src[l.start:l.pos]) + l.start = l.pos + } + return nil + } + } + } +} + +func doubleQuoteState(l *sqlLexer) stateFn { + for { + r, width := utf8.DecodeRuneInString(l.src[l.pos:]) + l.pos += width + + switch r { + case '"': + nextRune, width := utf8.DecodeRuneInString(l.src[l.pos:]) + if nextRune != '"' { + return rawState + } + l.pos += width + case utf8.RuneError: + if width != replacementcharacterwidth { + if l.pos-l.start > 0 { + l.parts = append(l.parts, l.src[l.start:l.pos]) + l.start = l.pos + } + return nil + } + } + } +} + +// placeholderState consumes a placeholder value. The $ must have already has +// already been consumed. The first rune must be a digit. +func placeholderState(l *sqlLexer) stateFn { + num := 0 + + for { + r, width := utf8.DecodeRuneInString(l.src[l.pos:]) + l.pos += width + + if '0' <= r && r <= '9' { + // Clamp rather than silently wrap on pathological input like + // "$92233720368547758070" which would otherwise overflow int and + // could land on a valid args index. Any value above MaxInt32 far + // exceeds any plausible args length, so Sanitize will correctly + // return "insufficient arguments". + if num > (math.MaxInt32-9)/10 { + num = math.MaxInt32 + } else { + num = num*10 + int(r-'0') + } + } else { + l.parts = append(l.parts, num) + l.pos -= width + l.start = l.pos + return rawState + } + } +} + +// dollarQuoteState consumes the body of a PostgreSQL dollar-quoted string +// ($[tag]$...$[tag]$). The opening tag (including its terminating '$') has +// already been consumed. +func dollarQuoteState(l *sqlLexer) stateFn { + closer := "$" + l.dollarTag + "$" + idx := strings.Index(l.src[l.pos:], closer) + if idx < 0 { + // Unterminated — mirror the behavior of other quoted-string states by + // consuming the remaining input into the current part and stopping. + if len(l.src)-l.start > 0 { + l.parts = append(l.parts, l.src[l.start:]) + l.start = len(l.src) + } + l.pos = len(l.src) + return nil + } + l.pos += idx + len(closer) + l.dollarTag = "" + return rawState +} + +// scanDollarQuoteTag checks whether src begins with an optional dollar-quoted +// string tag followed by a closing '$'. src must point just past the opening +// '$'. Returns the byte length of the tag (zero for an anonymous $$) and +// whether a valid tag was found. +// +// Tag grammar matches the PostgreSQL lexer (scan.l): +// +// dolq_start: [A-Za-z_\x80-\xff] +// dolq_cont: [A-Za-z0-9_\x80-\xff] +func scanDollarQuoteTag(src string) (int, bool) { + first := true + for i := 0; i < len(src); { + r, w := utf8.DecodeRuneInString(src[i:]) + if r == '$' { + return i, true + } + if !isDollarTagRune(r, first) { + return 0, false + } + first = false + i += w + } + return 0, false +} + +func isDollarTagRune(r rune, first bool) bool { + switch { + case r == '_': + return true + case 'a' <= r && r <= 'z': + return true + case 'A' <= r && r <= 'Z': + return true + case !first && '0' <= r && r <= '9': + return true + case r >= 0x80 && r != utf8.RuneError: + return true + } + return false +} + +func escapeStringState(l *sqlLexer) stateFn { + for { + r, width := utf8.DecodeRuneInString(l.src[l.pos:]) + l.pos += width + + switch r { + case '\\': + _, width = utf8.DecodeRuneInString(l.src[l.pos:]) + l.pos += width + case '\'': + nextRune, width := utf8.DecodeRuneInString(l.src[l.pos:]) + if nextRune != '\'' { + return rawState + } + l.pos += width + case utf8.RuneError: + if width != replacementcharacterwidth { + if l.pos-l.start > 0 { + l.parts = append(l.parts, l.src[l.start:l.pos]) + l.start = l.pos + } + return nil + } + } + } +} + +func oneLineCommentState(l *sqlLexer) stateFn { + for { + r, width := utf8.DecodeRuneInString(l.src[l.pos:]) + l.pos += width + + switch r { + case '\\': + _, width = utf8.DecodeRuneInString(l.src[l.pos:]) + l.pos += width + case '\n', '\r': + return rawState + case utf8.RuneError: + if width != replacementcharacterwidth { + if l.pos-l.start > 0 { + l.parts = append(l.parts, l.src[l.start:l.pos]) + l.start = l.pos + } + return nil + } + } + } +} + +func multilineCommentState(l *sqlLexer) stateFn { + for { + r, width := utf8.DecodeRuneInString(l.src[l.pos:]) + l.pos += width + + switch r { + case '/': + nextRune, width := utf8.DecodeRuneInString(l.src[l.pos:]) + if nextRune == '*' { + l.pos += width + l.nested++ + } + case '*': + nextRune, width := utf8.DecodeRuneInString(l.src[l.pos:]) + if nextRune != '/' { + continue + } + + l.pos += width + if l.nested == 0 { + return rawState + } + l.nested-- + + case utf8.RuneError: + if width != replacementcharacterwidth { + if l.pos-l.start > 0 { + l.parts = append(l.parts, l.src[l.start:l.pos]) + l.start = l.pos + } + return nil + } + } + } +} + +var queryPool = &pool[*Query]{ + new: func() *Query { + return &Query{} + }, + reset: func(q *Query) bool { + n := len(q.Parts) + q.Parts = q.Parts[:0] + return n < 64 // drop too large queries + }, +} + +// SanitizeSQL replaces placeholder values with args. It quotes and escapes args +// as necessary. This function is only safe when standard_conforming_strings is +// on. +func SanitizeSQL(sql string, args ...any) (string, error) { + query := queryPool.get() + query.init(sql) + defer queryPool.put(query) + + return query.Sanitize(args...) +} + +type pool[E any] struct { + p sync.Pool + new func() E + reset func(E) bool +} + +func (pool *pool[E]) get() E { + v, ok := pool.p.Get().(E) + if !ok { + v = pool.new() + } + + return v +} + +func (p *pool[E]) put(v E) { + if p.reset(v) { + p.p.Put(v) + } +} diff --git a/vendor/github.com/jackc/pgx/v5/internal/stmtcache/lru_cache.go b/vendor/github.com/jackc/pgx/v5/internal/stmtcache/lru_cache.go new file mode 100644 index 0000000000..b677d29cb9 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/internal/stmtcache/lru_cache.go @@ -0,0 +1,187 @@ +package stmtcache + +import ( + "github.com/jackc/pgx/v5/pgconn" +) + +// lruNode is a typed doubly-linked list node with freelist support. +type lruNode struct { + sd *pgconn.StatementDescription + prev *lruNode + next *lruNode +} + +// LRUCache implements Cache with a Least Recently Used (LRU) cache. +type LRUCache struct { + m map[string]*lruNode + head *lruNode + + tail *lruNode + len int + cap int + freelist *lruNode + + invalidStmts []*pgconn.StatementDescription + invalidSet map[string]struct{} +} + +// NewLRUCache creates a new LRUCache. cap is the maximum size of the cache. +func NewLRUCache(cap int) *LRUCache { + head := &lruNode{} + tail := &lruNode{} + head.next = tail + tail.prev = head + + return &LRUCache{ + cap: cap, + m: make(map[string]*lruNode, cap), + head: head, + tail: tail, + invalidSet: make(map[string]struct{}), + } +} + +// Get returns the statement description for sql. Returns nil if not found. +func (c *LRUCache) Get(key string) *pgconn.StatementDescription { + node, ok := c.m[key] + if !ok { + return nil + } + c.moveToFront(node) + return node.sd +} + +// Put stores sd in the cache. Put panics if sd.SQL is "". Put does nothing if sd.SQL already exists in the cache or +// sd.SQL has been invalidated and HandleInvalidated has not been called yet. +func (c *LRUCache) Put(sd *pgconn.StatementDescription) { + if sd.SQL == "" { + panic("cannot store statement description with empty SQL") + } + + if _, present := c.m[sd.SQL]; present { + return + } + + // The statement may have been invalidated but not yet handled. Do not readd it to the cache. + if _, invalidated := c.invalidSet[sd.SQL]; invalidated { + return + } + + if c.len == c.cap { + c.invalidateOldest() + } + + node := c.allocNode() + node.sd = sd + c.insertAfter(c.head, node) + c.m[sd.SQL] = node + c.len++ +} + +// Invalidate invalidates statement description identified by sql. Does nothing if not found. +func (c *LRUCache) Invalidate(sql string) { + node, ok := c.m[sql] + if !ok { + return + } + delete(c.m, sql) + c.invalidStmts = append(c.invalidStmts, node.sd) + c.invalidSet[sql] = struct{}{} + c.unlink(node) + c.len-- + c.freeNode(node) +} + +// InvalidateAll invalidates all statement descriptions. +func (c *LRUCache) InvalidateAll() { + for node := c.head.next; node != c.tail; { + next := node.next + c.invalidStmts = append(c.invalidStmts, node.sd) + c.invalidSet[node.sd.SQL] = struct{}{} + c.freeNode(node) + node = next + } + + clear(c.m) + c.head.next = c.tail + c.tail.prev = c.head + c.len = 0 +} + +// GetInvalidated returns a slice of all statement descriptions invalidated since the last call to RemoveInvalidated. +func (c *LRUCache) GetInvalidated() []*pgconn.StatementDescription { + return c.invalidStmts +} + +// RemoveInvalidated removes all invalidated statement descriptions. No other calls to Cache must be made between a +// call to GetInvalidated and RemoveInvalidated or RemoveInvalidated may remove statement descriptions that were +// never seen by the call to GetInvalidated. +func (c *LRUCache) RemoveInvalidated() { + c.invalidStmts = c.invalidStmts[:0] + clear(c.invalidSet) +} + +// Len returns the number of cached prepared statement descriptions. +func (c *LRUCache) Len() int { + return c.len +} + +// Cap returns the maximum number of cached prepared statement descriptions. +func (c *LRUCache) Cap() int { + return c.cap +} + +func (c *LRUCache) invalidateOldest() { + node := c.tail.prev + if node == c.head { + return + } + c.invalidStmts = append(c.invalidStmts, node.sd) + c.invalidSet[node.sd.SQL] = struct{}{} + delete(c.m, node.sd.SQL) + c.unlink(node) + c.len-- + c.freeNode(node) +} + +// List operations - sentinel nodes eliminate nil checks + +func (c *LRUCache) insertAfter(at, node *lruNode) { + node.prev = at + node.next = at.next + at.next.prev = node + at.next = node +} + +func (c *LRUCache) unlink(node *lruNode) { + node.prev.next = node.next + node.next.prev = node.prev +} + +func (c *LRUCache) moveToFront(node *lruNode) { + if node.prev == c.head { + return + } + c.unlink(node) + c.insertAfter(c.head, node) +} + +// Node pool operations - reuse evicted nodes to avoid allocations + +func (c *LRUCache) allocNode() *lruNode { + if c.freelist != nil { + node := c.freelist + c.freelist = node.next + node.next = nil + node.prev = nil + return node + } + return &lruNode{} +} + +func (c *LRUCache) freeNode(node *lruNode) { + node.sd = nil + node.prev = nil + node.next = c.freelist + c.freelist = node +} diff --git a/vendor/github.com/jackc/pgx/v5/internal/stmtcache/stmtcache.go b/vendor/github.com/jackc/pgx/v5/internal/stmtcache/stmtcache.go new file mode 100644 index 0000000000..d57bdd29e6 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/internal/stmtcache/stmtcache.go @@ -0,0 +1,45 @@ +// Package stmtcache is a cache for statement descriptions. +package stmtcache + +import ( + "crypto/sha256" + "encoding/hex" + + "github.com/jackc/pgx/v5/pgconn" +) + +// StatementName returns a statement name that will be stable for sql across multiple connections and program +// executions. +func StatementName(sql string) string { + digest := sha256.Sum256([]byte(sql)) + return "stmtcache_" + hex.EncodeToString(digest[0:24]) +} + +// Cache caches statement descriptions. +type Cache interface { + // Get returns the statement description for sql. Returns nil if not found. + Get(sql string) *pgconn.StatementDescription + + // Put stores sd in the cache. Put panics if sd.SQL is "". Put does nothing if sd.SQL already exists in the cache. + Put(sd *pgconn.StatementDescription) + + // Invalidate invalidates statement description identified by sql. Does nothing if not found. + Invalidate(sql string) + + // InvalidateAll invalidates all statement descriptions. + InvalidateAll() + + // GetInvalidated returns a slice of all statement descriptions invalidated since the last call to RemoveInvalidated. + GetInvalidated() []*pgconn.StatementDescription + + // RemoveInvalidated removes all invalidated statement descriptions. No other calls to Cache must be made between a + // call to GetInvalidated and RemoveInvalidated or RemoveInvalidated may remove statement descriptions that were + // never seen by the call to GetInvalidated. + RemoveInvalidated() + + // Len returns the number of cached prepared statement descriptions. + Len() int + + // Cap returns the maximum number of cached prepared statement descriptions. + Cap() int +} diff --git a/vendor/github.com/jackc/pgx/v5/large_objects.go b/vendor/github.com/jackc/pgx/v5/large_objects.go new file mode 100644 index 0000000000..9d21afdce9 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/large_objects.go @@ -0,0 +1,161 @@ +package pgx + +import ( + "context" + "errors" + "io" + + "github.com/jackc/pgx/v5/pgtype" +) + +// The PostgreSQL wire protocol has a limit of 1 GB - 1 per message. See definition of +// PQ_LARGE_MESSAGE_LIMIT in the PostgreSQL source code. To allow for the other data +// in the message,maxLargeObjectMessageLength should be no larger than 1 GB - 1 KB. +var maxLargeObjectMessageLength = 1024*1024*1024 - 1024 + +// LargeObjects is a structure used to access the large objects API. It is only valid within the transaction where it +// was created. +// +// For more details see: http://www.postgresql.org/docs/current/static/largeobjects.html +type LargeObjects struct { + tx Tx +} + +type LargeObjectMode int32 + +const ( + LargeObjectModeWrite LargeObjectMode = 0x20000 + LargeObjectModeRead LargeObjectMode = 0x40000 +) + +// Create creates a new large object. If oid is zero, the server assigns an unused OID. +func (o *LargeObjects) Create(ctx context.Context, oid uint32) (uint32, error) { + err := o.tx.QueryRow(ctx, "select lo_create($1)", oid).Scan(&oid) + return oid, err +} + +// Open opens an existing large object with the given mode. ctx will also be used for all operations on the opened large +// object. +func (o *LargeObjects) Open(ctx context.Context, oid uint32, mode LargeObjectMode) (*LargeObject, error) { + var fd int32 + err := o.tx.QueryRow(ctx, "select lo_open($1, $2)", oid, mode).Scan(&fd) + if err != nil { + return nil, err + } + return &LargeObject{fd: fd, tx: o.tx, ctx: ctx}, nil +} + +// Unlink removes a large object from the database. +func (o *LargeObjects) Unlink(ctx context.Context, oid uint32) error { + var result int32 + err := o.tx.QueryRow(ctx, "select lo_unlink($1)", oid).Scan(&result) + if err != nil { + return err + } + + if result != 1 { + return errors.New("failed to remove large object") + } + + return nil +} + +// A LargeObject is a large object stored on the server. It is only valid within the transaction that it was initialized +// in. It uses the context it was initialized with for all operations. It implements these interfaces: +// +// io.Writer +// io.Reader +// io.Seeker +// io.Closer +type LargeObject struct { + ctx context.Context + tx Tx + fd int32 +} + +// Write writes p to the large object and returns the number of bytes written and an error if not all of p was written. +func (o *LargeObject) Write(p []byte) (int, error) { + nTotal := 0 + for { + expected := len(p) - nTotal + if expected == 0 { + break + } else if expected > maxLargeObjectMessageLength { + expected = maxLargeObjectMessageLength + } + + var n int + err := o.tx.QueryRow(o.ctx, "select lowrite($1, $2)", o.fd, p[nTotal:nTotal+expected]).Scan(&n) + if err != nil { + return nTotal, err + } + + if n < 0 { + return nTotal, errors.New("failed to write to large object") + } + + nTotal += n + + if n < expected { + return nTotal, errors.New("short write to large object") + } else if n > expected { + return nTotal, errors.New("invalid write to large object") + } + } + + return nTotal, nil +} + +// Read reads up to len(p) bytes into p returning the number of bytes read. +func (o *LargeObject) Read(p []byte) (int, error) { + nTotal := 0 + for { + expected := len(p) - nTotal + if expected == 0 { + break + } else if expected > maxLargeObjectMessageLength { + expected = maxLargeObjectMessageLength + } + + res := pgtype.PreallocBytes(p[nTotal:]) + err := o.tx.QueryRow(o.ctx, "select loread($1, $2)", o.fd, expected).Scan(&res) + // We compute expected so that it always fits into p, so it should never happen + // that PreallocBytes's ScanBytes had to allocate a new slice. + nTotal += len(res) + if err != nil { + return nTotal, err + } + + if len(res) < expected { + return nTotal, io.EOF + } else if len(res) > expected { + return nTotal, errors.New("invalid read of large object") + } + } + + return nTotal, nil +} + +// Seek moves the current location pointer to the new location specified by offset. +func (o *LargeObject) Seek(offset int64, whence int) (n int64, err error) { + err = o.tx.QueryRow(o.ctx, "select lo_lseek64($1, $2, $3)", o.fd, offset, whence).Scan(&n) + return n, err +} + +// Tell returns the current read or write location of the large object descriptor. +func (o *LargeObject) Tell() (n int64, err error) { + err = o.tx.QueryRow(o.ctx, "select lo_tell64($1)", o.fd).Scan(&n) + return n, err +} + +// Truncate the large object to size. +func (o *LargeObject) Truncate(size int64) (err error) { + _, err = o.tx.Exec(o.ctx, "select lo_truncate64($1, $2)", o.fd, size) + return err +} + +// Close the large object descriptor. +func (o *LargeObject) Close() error { + _, err := o.tx.Exec(o.ctx, "select lo_close($1)", o.fd) + return err +} diff --git a/vendor/github.com/jackc/pgx/v5/mise.toml b/vendor/github.com/jackc/pgx/v5/mise.toml new file mode 100644 index 0000000000..78610d4ccc --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/mise.toml @@ -0,0 +1,8 @@ +[tools] +go = '1.26.3' +"go:github.com/go-critic/go-critic/cmd/gocritic" = "latest" +"go:github.com/gordonklaus/ineffassign" = "latest" +"go:github.com/mdempsky/unconvert" = "latest" +"go:golang.org/x/tools/cmd/goimports" = "latest" +"go:mvdan.cc/gofumpt" = "latest" +ruby = '4.0.4' diff --git a/vendor/github.com/jackc/pgx/v5/named_args.go b/vendor/github.com/jackc/pgx/v5/named_args.go new file mode 100644 index 0000000000..1300c91c0f --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/named_args.go @@ -0,0 +1,410 @@ +package pgx + +import ( + "context" + "fmt" + "reflect" + "strconv" + "strings" + "unicode/utf8" +) + +// NamedArgs can be used as the first argument to a query method. It will replace every '@' named placeholder with a '$' +// ordinal placeholder and construct the appropriate arguments. +// +// For example, the following two queries are equivalent: +// +// conn.Query(ctx, "select * from widgets where foo = @foo and bar = @bar", pgx.NamedArgs{"foo": 1, "bar": 2}) +// conn.Query(ctx, "select * from widgets where foo = $1 and bar = $2", 1, 2) +// +// Named placeholders are case sensitive and must start with a letter or underscore. Subsequent characters can be +// letters, numbers, or underscores. + +type NamedArgs map[string]any + +// RewriteQuery implements the QueryRewriter interface. +func (na NamedArgs) RewriteQuery(ctx context.Context, conn *Conn, sql string, args []any) (newSQL string, newArgs []any, err error) { + return rewriteQuery(na, sql, false) +} + +// StrictNamedArgs can be used in the same way as NamedArgs, but provided arguments are also checked to include all +// named arguments that the sql query uses, and no extra arguments. +type StrictNamedArgs map[string]any + +// RewriteQuery implements the QueryRewriter interface. +func (sna StrictNamedArgs) RewriteQuery(ctx context.Context, conn *Conn, sql string, args []any) (newSQL string, newArgs []any, err error) { + return rewriteQuery(sna, sql, true) +} + +type errorQueryRewriter struct { + err error +} + +func (r errorQueryRewriter) RewriteQuery(ctx context.Context, conn *Conn, sql string, args []any) (newSQL string, newArgs []any, err error) { + return "", nil, r.err +} + +// StructArgs converts exported fields of a struct into a QueryRewriter so it can +// be used as the first argument to a query method (e.g. "where id=@id"). +// +// Field names are taken from the `db` struct tag if present. Tag values may +// include comma-separated options (e.g. `db:"id,omitempty"`). A `db:"-"` field is +// ignored. If no `db` tag is present, the Go field name is used. +// +// sa may be a struct or a pointer to a struct. +func StructArgs(sa any) QueryRewriter { + args, err := structArgs(sa) + if err != nil { + return errorQueryRewriter{err: err} + } + return NamedArgs(args) +} + +// StrictStructArgs is like StructArgs but uses StrictNamedArgs rewriting +// semantics (i.e. errors if the SQL query references missing arguments or if +// extra arguments are provided). +func StrictStructArgs(sa any) QueryRewriter { + args, err := structArgs(sa) + if err != nil { + return errorQueryRewriter{err: err} + } + return StrictNamedArgs(args) +} + +func structArgs(sa any) (map[string]any, error) { + if sa == nil { + return nil, fmt.Errorf("StructArgs requires a struct or pointer to struct, got nil") + } + + v := reflect.ValueOf(sa) + t := v.Type() + + if t.Kind() == reflect.Pointer { + if v.IsNil() { + return nil, fmt.Errorf("StructArgs requires a non-nil pointer to struct") + } + v = v.Elem() + t = v.Type() + } + + if t.Kind() != reflect.Struct { + return nil, fmt.Errorf("StructArgs requires a struct or pointer to struct, got %s", t) + } + + out := make(map[string]any, t.NumField()) + for i := 0; i < t.NumField(); i++ { + sf := t.Field(i) + + // Ignore unexported fields. + if sf.PkgPath != "" { + continue + } + + key, ok, err := dbTagKey(sf) + if err != nil { + return nil, err + } + if !ok { + continue + } + + if _, exists := out[key]; exists { + return nil, fmt.Errorf("duplicate StructArgs key %q", key) + } + + out[key] = v.Field(i).Interface() + } + + return out, nil +} + +// dbTagKey derives the named-argument key for a struct field. Tag parsing matches +// RowToStructByName* in rows.go (structTagKey, Lookup, comma options, db:"-"). +// Anonymous embedded structs are skipped without flattening (unlike row scanning). +func dbTagKey(sf reflect.StructField) (key string, ok bool, err error) { + if sf.Anonymous { + ft := sf.Type + if ft.Kind() == reflect.Pointer { + ft = ft.Elem() + } + if ft.Kind() == reflect.Struct { + return "", false, nil + } + } + + dbTag, dbTagPresent := sf.Tag.Lookup(structTagKey) + if dbTagPresent { + dbTag, _, _ = strings.Cut(dbTag, ",") + } + if dbTag == "-" { + return "", false, nil + } + if dbTagPresent { + if dbTag == "" { + return "", false, fmt.Errorf("field %s has empty `%s` tag", sf.Name, structTagKey) + } + return dbTag, true, nil + } + + return sf.Name, true, nil +} + +type namedArg string + +type sqlLexer struct { + src string + start int + pos int + nested int // multiline comment nesting level. + stateFn stateFn + parts []any + + nameToOrdinal map[namedArg]int +} + +type stateFn func(*sqlLexer) stateFn + +func rewriteQuery(na map[string]any, sql string, isStrict bool) (newSQL string, newArgs []any, err error) { + l := &sqlLexer{ + src: sql, + stateFn: rawState, + nameToOrdinal: make(map[namedArg]int, len(na)), + } + + for l.stateFn != nil { + l.stateFn = l.stateFn(l) + } + + sb := strings.Builder{} + for _, p := range l.parts { + switch p := p.(type) { + case string: + sb.WriteString(p) + case namedArg: + sb.WriteRune('$') + sb.WriteString(strconv.Itoa(l.nameToOrdinal[p])) + } + } + + newArgs = make([]any, len(l.nameToOrdinal)) + for name, ordinal := range l.nameToOrdinal { + var found bool + newArgs[ordinal-1], found = na[string(name)] + if isStrict && !found { + return "", nil, fmt.Errorf("argument %s found in sql query but not present in StrictNamedArgs", name) + } + } + + if isStrict { + for name := range na { + if _, found := l.nameToOrdinal[namedArg(name)]; !found { + return "", nil, fmt.Errorf("argument %s of StrictNamedArgs not found in sql query", name) + } + } + } + + return sb.String(), newArgs, nil +} + +func rawState(l *sqlLexer) stateFn { + for { + r, width := utf8.DecodeRuneInString(l.src[l.pos:]) + l.pos += width + + switch r { + case 'e', 'E': + nextRune, width := utf8.DecodeRuneInString(l.src[l.pos:]) + if nextRune == '\'' { + l.pos += width + return escapeStringState + } + case '\'': + return singleQuoteState + case '"': + return doubleQuoteState + case '@': + nextRune, _ := utf8.DecodeRuneInString(l.src[l.pos:]) + if isLetter(nextRune) || nextRune == '_' { + if l.pos-l.start > 0 { + l.parts = append(l.parts, l.src[l.start:l.pos-width]) + } + l.start = l.pos + return namedArgState + } + case '-': + nextRune, width := utf8.DecodeRuneInString(l.src[l.pos:]) + if nextRune == '-' { + l.pos += width + return oneLineCommentState + } + case '/': + nextRune, width := utf8.DecodeRuneInString(l.src[l.pos:]) + if nextRune == '*' { + l.pos += width + return multilineCommentState + } + case utf8.RuneError: + if l.pos-l.start > 0 { + l.parts = append(l.parts, l.src[l.start:l.pos]) + l.start = l.pos + } + return nil + } + } +} + +func isLetter(r rune) bool { + return (r >= 'a' && r <= 'z') || (r >= 'A' && r <= 'Z') +} + +func namedArgState(l *sqlLexer) stateFn { + for { + r, width := utf8.DecodeRuneInString(l.src[l.pos:]) + l.pos += width + + if r == utf8.RuneError { + if l.pos-l.start > 0 { + na := namedArg(l.src[l.start:l.pos]) + if _, found := l.nameToOrdinal[na]; !found { + l.nameToOrdinal[na] = len(l.nameToOrdinal) + 1 + } + l.parts = append(l.parts, na) + l.start = l.pos + } + return nil + } else if !(isLetter(r) || (r >= '0' && r <= '9') || r == '_') { + l.pos -= width + na := namedArg(l.src[l.start:l.pos]) + if _, found := l.nameToOrdinal[na]; !found { + l.nameToOrdinal[na] = len(l.nameToOrdinal) + 1 + } + l.parts = append(l.parts, na) + l.start = l.pos + return rawState + } + } +} + +func singleQuoteState(l *sqlLexer) stateFn { + for { + r, width := utf8.DecodeRuneInString(l.src[l.pos:]) + l.pos += width + + switch r { + case '\'': + nextRune, width := utf8.DecodeRuneInString(l.src[l.pos:]) + if nextRune != '\'' { + return rawState + } + l.pos += width + case utf8.RuneError: + if l.pos-l.start > 0 { + l.parts = append(l.parts, l.src[l.start:l.pos]) + l.start = l.pos + } + return nil + } + } +} + +func doubleQuoteState(l *sqlLexer) stateFn { + for { + r, width := utf8.DecodeRuneInString(l.src[l.pos:]) + l.pos += width + + switch r { + case '"': + nextRune, width := utf8.DecodeRuneInString(l.src[l.pos:]) + if nextRune != '"' { + return rawState + } + l.pos += width + case utf8.RuneError: + if l.pos-l.start > 0 { + l.parts = append(l.parts, l.src[l.start:l.pos]) + l.start = l.pos + } + return nil + } + } +} + +func escapeStringState(l *sqlLexer) stateFn { + for { + r, width := utf8.DecodeRuneInString(l.src[l.pos:]) + l.pos += width + + switch r { + case '\\': + _, width = utf8.DecodeRuneInString(l.src[l.pos:]) + l.pos += width + case '\'': + nextRune, width := utf8.DecodeRuneInString(l.src[l.pos:]) + if nextRune != '\'' { + return rawState + } + l.pos += width + case utf8.RuneError: + if l.pos-l.start > 0 { + l.parts = append(l.parts, l.src[l.start:l.pos]) + l.start = l.pos + } + return nil + } + } +} + +func oneLineCommentState(l *sqlLexer) stateFn { + for { + r, width := utf8.DecodeRuneInString(l.src[l.pos:]) + l.pos += width + + switch r { + case '\\': + _, width = utf8.DecodeRuneInString(l.src[l.pos:]) + l.pos += width + case '\n', '\r': + return rawState + case utf8.RuneError: + if l.pos-l.start > 0 { + l.parts = append(l.parts, l.src[l.start:l.pos]) + l.start = l.pos + } + return nil + } + } +} + +func multilineCommentState(l *sqlLexer) stateFn { + for { + r, width := utf8.DecodeRuneInString(l.src[l.pos:]) + l.pos += width + + switch r { + case '/': + nextRune, width := utf8.DecodeRuneInString(l.src[l.pos:]) + if nextRune == '*' { + l.pos += width + l.nested++ + } + case '*': + nextRune, width := utf8.DecodeRuneInString(l.src[l.pos:]) + if nextRune != '/' { + continue + } + + l.pos += width + if l.nested == 0 { + return rawState + } + l.nested-- + + case utf8.RuneError: + if l.pos-l.start > 0 { + l.parts = append(l.parts, l.src[l.start:l.pos]) + l.start = l.pos + } + return nil + } + } +} diff --git a/vendor/github.com/jackc/pgx/v5/pgconn/README.md b/vendor/github.com/jackc/pgx/v5/pgconn/README.md new file mode 100644 index 0000000000..1fe15c2686 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgconn/README.md @@ -0,0 +1,29 @@ +# pgconn + +Package pgconn is a low-level PostgreSQL database driver. It operates at nearly the same level as the C library libpq. +It is primarily intended to serve as the foundation for higher level libraries such as https://github.com/jackc/pgx. +Applications should handle normal queries with a higher level library and only use pgconn directly when required for +low-level access to PostgreSQL functionality. + +## Example Usage + +```go +pgConn, err := pgconn.Connect(context.Background(), os.Getenv("DATABASE_URL")) +if err != nil { + log.Fatalln("pgconn failed to connect:", err) +} +defer pgConn.Close(context.Background()) + +result := pgConn.ExecParams(context.Background(), "SELECT email FROM users WHERE id=$1", [][]byte{[]byte("123")}, nil, nil, nil) +for result.NextRow() { + fmt.Println("User 123 has email:", string(result.Values()[0])) +} +_, err = result.Close() +if err != nil { + log.Fatalln("failed reading result:", err) +} +``` + +## Testing + +See CONTRIBUTING.md for setup instructions. diff --git a/vendor/github.com/jackc/pgx/v5/pgconn/auth_oauth.go b/vendor/github.com/jackc/pgx/v5/pgconn/auth_oauth.go new file mode 100644 index 0000000000..991f6585db --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgconn/auth_oauth.go @@ -0,0 +1,67 @@ +package pgconn + +import ( + "context" + "encoding/json" + "errors" + "fmt" + + "github.com/jackc/pgx/v5/pgproto3" +) + +func (c *PgConn) oauthAuth(ctx context.Context) error { + if c.config.OAuthTokenProvider == nil { + return errors.New("OAuth authentication required but no token provider configured") + } + + token, err := c.config.OAuthTokenProvider(ctx) + if err != nil { + return fmt.Errorf("failed to obtain OAuth token: %w", err) + } + + // https://www.rfc-editor.org/rfc/rfc7628.html#section-3.1 + initialResponse := []byte("n,,\x01auth=Bearer " + token + "\x01\x01") + + saslInitialResponse := &pgproto3.SASLInitialResponse{ + AuthMechanism: "OAUTHBEARER", + Data: initialResponse, + } + c.frontend.Send(saslInitialResponse) + err = c.flushWithPotentialWriteReadDeadlock() + if err != nil { + return err + } + + msg, err := c.receiveMessage() + if err != nil { + return err + } + + switch m := msg.(type) { + case *pgproto3.AuthenticationOk: + return nil + case *pgproto3.AuthenticationSASLContinue: + // Server sent error response in SASL continue + // https://www.rfc-editor.org/rfc/rfc7628.html#section-3.2.2 + // https://www.rfc-editor.org/rfc/rfc7628.html#section-3.2.3 + errResponse := struct { + Status string `json:"status"` + Scope string `json:"scope"` + OpenIDConfiguration string `json:"openid-configuration"` + }{} + err := json.Unmarshal(m.Data, &errResponse) + if err != nil { + return fmt.Errorf("invalid OAuth error response from server: %w", err) + } + + // Per RFC 7628 section 3.2.3, we should send a SASLResponse which only contains \x01. + // However, since the connection will be closed anyway, we can skip this + return fmt.Errorf("OAuth authentication failed: %s", errResponse.Status) + + case *pgproto3.ErrorResponse: + return ErrorResponseToPgError(m) + + default: + return fmt.Errorf("unexpected message type during OAuth auth: %T", msg) + } +} diff --git a/vendor/github.com/jackc/pgx/v5/pgconn/auth_scram.go b/vendor/github.com/jackc/pgx/v5/pgconn/auth_scram.go new file mode 100644 index 0000000000..aeaf8cb1c5 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgconn/auth_scram.go @@ -0,0 +1,407 @@ +// SCRAM-SHA-256 and SCRAM-SHA-256-PLUS authentication +// +// Resources: +// https://tools.ietf.org/html/rfc5802 +// https://tools.ietf.org/html/rfc5929 +// https://tools.ietf.org/html/rfc8265 +// https://www.postgresql.org/docs/current/sasl-authentication.html +// +// Inspiration drawn from other implementations: +// https://github.com/lib/pq/pull/608 +// https://github.com/lib/pq/pull/788 +// https://github.com/lib/pq/pull/833 + +package pgconn + +import ( + "bytes" + "crypto/hmac" + "crypto/pbkdf2" + "crypto/rand" + "crypto/sha256" + "crypto/sha512" + "crypto/tls" + "crypto/x509" + "encoding/base64" + "errors" + "fmt" + "hash" + "slices" + "strconv" + + "github.com/jackc/pgx/v5/pgproto3" + "golang.org/x/text/secure/precis" +) + +const ( + clientNonceLen = 18 + scramSHA256Name = "SCRAM-SHA-256" + scramSHA256PlusName = "SCRAM-SHA-256-PLUS" +) + +// Perform SCRAM authentication. +func (c *PgConn) scramAuth(serverAuthMechanisms []string) error { + sc, err := newScramClient(serverAuthMechanisms, c.config.Password) + if err != nil { + return err + } + + serverHasPlus := slices.Contains(sc.serverAuthMechanisms, scramSHA256PlusName) + if c.config.ChannelBinding == "require" && !serverHasPlus { + return errors.New("channel binding required but server does not support SCRAM-SHA-256-PLUS") + } + + // If we have a TLS connection and channel binding is not disabled, attempt to + // extract the server certificate hash for tls-server-end-point channel binding. + if tlsConn, ok := c.conn.(*tls.Conn); ok && c.config.ChannelBinding != "disable" { + certHash, err := getTLSCertificateHash(tlsConn) + if err != nil && c.config.ChannelBinding == "require" { + return fmt.Errorf("channel binding required but failed to get server certificate hash: %w", err) + } + + // Upgrade to SCRAM-SHA-256-PLUS if we have binding data and the server supports it. + if certHash != nil && serverHasPlus { + sc.authMechanism = scramSHA256PlusName + } + + sc.channelBindingData = certHash + sc.hasTLS = true + } + + if c.config.ChannelBinding == "require" && sc.channelBindingData == nil { + return errors.New("channel binding required but channel binding data is not available") + } + + // Send client-first-message in a SASLInitialResponse + saslInitialResponse := &pgproto3.SASLInitialResponse{ + AuthMechanism: sc.authMechanism, + Data: sc.clientFirstMessage(), + } + c.frontend.Send(saslInitialResponse) + err = c.flushWithPotentialWriteReadDeadlock() + if err != nil { + return err + } + + // Receive server-first-message payload in an AuthenticationSASLContinue. + saslContinue, err := c.rxSASLContinue() + if err != nil { + return err + } + err = sc.recvServerFirstMessage(saslContinue.Data) + if err != nil { + return err + } + + // Send client-final-message in a SASLResponse + saslResponse := &pgproto3.SASLResponse{ + Data: []byte(sc.clientFinalMessage()), + } + c.frontend.Send(saslResponse) + err = c.flushWithPotentialWriteReadDeadlock() + if err != nil { + return err + } + + // Receive server-final-message payload in an AuthenticationSASLFinal. + saslFinal, err := c.rxSASLFinal() + if err != nil { + return err + } + return sc.recvServerFinalMessage(saslFinal.Data) +} + +func (c *PgConn) rxSASLContinue() (*pgproto3.AuthenticationSASLContinue, error) { + msg, err := c.receiveMessage() + if err != nil { + return nil, err + } + switch m := msg.(type) { + case *pgproto3.AuthenticationSASLContinue: + return m, nil + case *pgproto3.ErrorResponse: + return nil, ErrorResponseToPgError(m) + } + + return nil, fmt.Errorf("expected AuthenticationSASLContinue message but received unexpected message %T", msg) +} + +func (c *PgConn) rxSASLFinal() (*pgproto3.AuthenticationSASLFinal, error) { + msg, err := c.receiveMessage() + if err != nil { + return nil, err + } + switch m := msg.(type) { + case *pgproto3.AuthenticationSASLFinal: + return m, nil + case *pgproto3.ErrorResponse: + return nil, ErrorResponseToPgError(m) + } + + return nil, fmt.Errorf("expected AuthenticationSASLFinal message but received unexpected message %T", msg) +} + +type scramClient struct { + serverAuthMechanisms []string + password string + clientNonce []byte + + // authMechanism is the selected SASL mechanism for the client. Must be + // either SCRAM-SHA-256 (default) or SCRAM-SHA-256-PLUS. + // + // Upgraded to SCRAM-SHA-256-PLUS during authentication when channel binding + // is not disabled, channel binding data is available (TLS connection with + // an obtainable server certificate hash) and the server advertises + // SCRAM-SHA-256-PLUS. + authMechanism string + + // hasTLS indicates whether the connection is using TLS. This is + // needed because the GS2 header must distinguish between a client that + // supports channel binding but the server does not ("y,,") versus one + // that does not support it at all ("n,,"). + hasTLS bool + + // channelBindingData is the hash of the server's TLS certificate, computed + // per the tls-server-end-point channel binding type (RFC 5929). Used as + // the binding input in SCRAM-SHA-256-PLUS. nil when not in use. + channelBindingData []byte + + clientFirstMessageBare []byte + clientGS2Header []byte + + serverFirstMessage []byte + clientAndServerNonce []byte + salt []byte + iterations int + + saltedPassword []byte + authMessage []byte +} + +func newScramClient(serverAuthMechanisms []string, password string) (*scramClient, error) { + sc := &scramClient{ + serverAuthMechanisms: serverAuthMechanisms, + authMechanism: scramSHA256Name, + } + + // Ensure the server supports SCRAM-SHA-256. SCRAM-SHA-256-PLUS is the + // channel binding variant and is only advertised when the server supports + // SSL. PostgreSQL always advertises the base SCRAM-SHA-256 mechanism + // regardless of SSL. + if !slices.Contains(sc.serverAuthMechanisms, scramSHA256Name) { + return nil, errors.New("server does not support SCRAM-SHA-256") + } + + // precis.OpaqueString is equivalent to SASLprep for password. + var err error + sc.password, err = precis.OpaqueString.String(password) + if err != nil { + // PostgreSQL allows passwords invalid according to SCRAM / SASLprep. + sc.password = password + } + + buf := make([]byte, clientNonceLen) + _, err = rand.Read(buf) + if err != nil { + return nil, err + } + sc.clientNonce = make([]byte, base64.RawStdEncoding.EncodedLen(len(buf))) + base64.RawStdEncoding.Encode(sc.clientNonce, buf) + + return sc, nil +} + +func (sc *scramClient) clientFirstMessage() []byte { + // The client-first-message is the GS2 header concatenated with the bare + // message (username + client nonce). The GS2 header communicates the + // client's channel binding capability to the server: + // + // "n,," - client is not using TLS (channel binding not possible) + // "y,," - client is using TLS but channel binding is not + // in use (e.g., server did not advertise SCRAM-SHA-256-PLUS + // or the server certificate hash was not obtainable) + // "p=tls-server-end-point,," - channel binding is active via SCRAM-SHA-256-PLUS + // + // See: + // https://www.rfc-editor.org/rfc/rfc5802#section-6 + // https://www.rfc-editor.org/rfc/rfc5929#section-4 + // https://www.postgresql.org/docs/current/sasl-authentication.html#SASL-SCRAM-SHA-256 + + sc.clientFirstMessageBare = fmt.Appendf(nil, "n=,r=%s", sc.clientNonce) + + switch { + case sc.authMechanism == scramSHA256PlusName: + sc.clientGS2Header = []byte("p=tls-server-end-point,,") + case sc.hasTLS: + sc.clientGS2Header = []byte("y,,") + default: + sc.clientGS2Header = []byte("n,,") + } + + return append(sc.clientGS2Header, sc.clientFirstMessageBare...) +} + +func (sc *scramClient) recvServerFirstMessage(serverFirstMessage []byte) error { + sc.serverFirstMessage = serverFirstMessage + buf := serverFirstMessage + if !bytes.HasPrefix(buf, []byte("r=")) { + return errors.New("invalid SCRAM server-first-message received from server: did not include r=") + } + buf = buf[2:] + + idx := bytes.IndexByte(buf, ',') + if idx == -1 { + return errors.New("invalid SCRAM server-first-message received from server: did not include s=") + } + sc.clientAndServerNonce = buf[:idx] + buf = buf[idx+1:] + + if !bytes.HasPrefix(buf, []byte("s=")) { + return errors.New("invalid SCRAM server-first-message received from server: did not include s=") + } + buf = buf[2:] + + idx = bytes.IndexByte(buf, ',') + if idx == -1 { + return errors.New("invalid SCRAM server-first-message received from server: did not include i=") + } + saltStr := buf[:idx] + buf = buf[idx+1:] + + if !bytes.HasPrefix(buf, []byte("i=")) { + return errors.New("invalid SCRAM server-first-message received from server: did not include i=") + } + buf = buf[2:] + iterationsStr := buf + + var err error + sc.salt, err = base64.StdEncoding.DecodeString(string(saltStr)) + if err != nil { + return fmt.Errorf("invalid SCRAM salt received from server: %w", err) + } + + sc.iterations, err = strconv.Atoi(string(iterationsStr)) + if err != nil || sc.iterations <= 0 { + return fmt.Errorf("invalid SCRAM iteration count received from server: %w", err) + } + // Bound server-supplied iteration count to prevent a malicious server from forcing the client + // to spend unbounded CPU in PBKDF2. PostgreSQL's scram_iterations defaults to 4096; this ceiling + // is ~2500x that. + const maxScramIterations = 10_000_000 + if sc.iterations > maxScramIterations { + return fmt.Errorf("SCRAM iteration count from server too high: %d (max %d)", sc.iterations, maxScramIterations) + } + + if !bytes.HasPrefix(sc.clientAndServerNonce, sc.clientNonce) { + return errors.New("invalid SCRAM nonce: did not start with client nonce") + } + + if len(sc.clientAndServerNonce) <= len(sc.clientNonce) { + return errors.New("invalid SCRAM nonce: did not include server nonce") + } + + return nil +} + +func (sc *scramClient) clientFinalMessage() string { + // The c= attribute carries the base64-encoded channel binding input. + // + // Without channel binding this is just the GS2 header alone ("biws" for + // "n,," or "eSws" for "y,,"). + // + // With channel binding, this is the GS2 header with the channel binding data + // (certificate hash) appended. + channelBindInput := sc.clientGS2Header + if sc.authMechanism == scramSHA256PlusName { + channelBindInput = slices.Concat(sc.clientGS2Header, sc.channelBindingData) + } + channelBindingEncoded := base64.StdEncoding.EncodeToString(channelBindInput) + clientFinalMessageWithoutProof := fmt.Appendf(nil, "c=%s,r=%s", channelBindingEncoded, sc.clientAndServerNonce) + + var err error + sc.saltedPassword, err = pbkdf2.Key(sha256.New, sc.password, sc.salt, sc.iterations, 32) + if err != nil { + panic(err) // This should never happen. + } + sc.authMessage = bytes.Join([][]byte{sc.clientFirstMessageBare, sc.serverFirstMessage, clientFinalMessageWithoutProof}, []byte(",")) + + clientProof := computeClientProof(sc.saltedPassword, sc.authMessage) + + return fmt.Sprintf("%s,p=%s", clientFinalMessageWithoutProof, clientProof) +} + +func (sc *scramClient) recvServerFinalMessage(serverFinalMessage []byte) error { + if !bytes.HasPrefix(serverFinalMessage, []byte("v=")) { + return errors.New("invalid SCRAM server-final-message received from server") + } + + serverSignature := serverFinalMessage[2:] + + if !hmac.Equal(serverSignature, computeServerSignature(sc.saltedPassword, sc.authMessage)) { + return errors.New("invalid SCRAM ServerSignature received from server") + } + + return nil +} + +func computeHMAC(key, msg []byte) []byte { + mac := hmac.New(sha256.New, key) + mac.Write(msg) + return mac.Sum(nil) +} + +func computeClientProof(saltedPassword, authMessage []byte) []byte { + clientKey := computeHMAC(saltedPassword, []byte("Client Key")) + storedKey := sha256.Sum256(clientKey) + clientSignature := computeHMAC(storedKey[:], authMessage) + + clientProof := make([]byte, len(clientSignature)) + for i := range clientSignature { + clientProof[i] = clientKey[i] ^ clientSignature[i] + } + + buf := make([]byte, base64.StdEncoding.EncodedLen(len(clientProof))) + base64.StdEncoding.Encode(buf, clientProof) + return buf +} + +func computeServerSignature(saltedPassword, authMessage []byte) []byte { + serverKey := computeHMAC(saltedPassword, []byte("Server Key")) + serverSignature := computeHMAC(serverKey, authMessage) + buf := make([]byte, base64.StdEncoding.EncodedLen(len(serverSignature))) + base64.StdEncoding.Encode(buf, serverSignature) + return buf +} + +// Get the server certificate hash for SCRAM channel binding type +// tls-server-end-point. +func getTLSCertificateHash(conn *tls.Conn) ([]byte, error) { + state := conn.ConnectionState() + if len(state.PeerCertificates) == 0 { + return nil, errors.New("no peer certificates for channel binding") + } + + cert := state.PeerCertificates[0] + + // Per RFC 5929 section 4.1: If the certificate's signatureAlgorithm uses + // MD5 or SHA-1, use SHA-256. Otherwise use the hash from the signature + // algorithm. + // + // See: https://www.rfc-editor.org/rfc/rfc5929.html#section-4.1 + var h hash.Hash + switch cert.SignatureAlgorithm { + case x509.MD5WithRSA, x509.SHA1WithRSA, x509.ECDSAWithSHA1: + h = sha256.New() + case x509.SHA256WithRSA, x509.SHA256WithRSAPSS, x509.ECDSAWithSHA256: + h = sha256.New() + case x509.SHA384WithRSA, x509.SHA384WithRSAPSS, x509.ECDSAWithSHA384: + h = sha512.New384() + case x509.SHA512WithRSA, x509.SHA512WithRSAPSS, x509.ECDSAWithSHA512: + h = sha512.New() + default: + return nil, fmt.Errorf("tls-server-end-point channel binding is undefined for certificate signature algorithm %v", cert.SignatureAlgorithm) + } + + h.Write(cert.Raw) + return h.Sum(nil), nil +} diff --git a/vendor/github.com/jackc/pgx/v5/pgconn/config.go b/vendor/github.com/jackc/pgx/v5/pgconn/config.go new file mode 100644 index 0000000000..eec7de6394 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgconn/config.go @@ -0,0 +1,1088 @@ +package pgconn + +import ( + "context" + "crypto/tls" + "crypto/x509" + "encoding/pem" + "errors" + "fmt" + "maps" + "math" + "net" + "net/url" + "os" + "path/filepath" + "strconv" + "strings" + "time" + + "github.com/jackc/pgpassfile" + "github.com/jackc/pgservicefile" + "github.com/jackc/pgx/v5/pgconn/ctxwatch" + "github.com/jackc/pgx/v5/pgproto3" +) + +type ( + AfterConnectFunc func(ctx context.Context, pgconn *PgConn) error + ValidateConnectFunc func(ctx context.Context, pgconn *PgConn) error + GetSSLPasswordFunc func(ctx context.Context) string +) + +// Config is the settings used to establish a connection to a PostgreSQL server. It must be created by [ParseConfig]. A +// manually initialized Config will cause ConnectConfig to panic. +type Config struct { + Host string // host (e.g. localhost) or absolute path to unix domain socket directory (e.g. /private/tmp) + Port uint16 + Database string + User string + Password string + TLSConfig *tls.Config // nil disables TLS + ConnectTimeout time.Duration + DialFunc DialFunc // e.g. net.Dialer.DialContext + LookupFunc LookupFunc // e.g. net.Resolver.LookupHost + BuildFrontend BuildFrontendFunc + + // BuildContextWatcherHandler is called to create a ContextWatcherHandler for a connection. The handler is called + // when a context passed to a PgConn method is canceled. + BuildContextWatcherHandler func(*PgConn) ctxwatch.Handler + + RuntimeParams map[string]string // Run-time parameters to set on connection as session default values (e.g. search_path or application_name) + + KerberosSrvName string + KerberosSpn string + Fallbacks []*FallbackConfig + + SSLNegotiation string // sslnegotiation=postgres or sslnegotiation=direct + + // AfterNetConnect is called after the network connection, including TLS if applicable, is established but before any + // PostgreSQL protocol communication. It takes the established net.Conn and returns a net.Conn that will be used in + // its place. It can be used to wrap the net.Conn (e.g. for logging, diagnostics, or testing). Its functionality has + // some overlap with DialFunc. However, DialFunc takes place before TLS is established and cannot be used to control + // the final net.Conn used for PostgreSQL protocol communication while AfterNetConnect can. + AfterNetConnect func(ctx context.Context, config *Config, conn net.Conn) (net.Conn, error) + + // ValidateConnect is called during a connection attempt after a successful authentication with the PostgreSQL server. + // It can be used to validate that the server is acceptable. If this returns an error the connection is closed and the next + // fallback config is tried. This allows implementing high availability behavior such as libpq does with target_session_attrs. + ValidateConnect ValidateConnectFunc + + // AfterConnect is called after ValidateConnect. It can be used to set up the connection (e.g. Set session variables + // or prepare statements). If this returns an error the connection attempt fails. + AfterConnect AfterConnectFunc + + // OnNotice is a callback function called when a notice response is received. + OnNotice NoticeHandler + + // OnNotification is a callback function called when a notification from the LISTEN/NOTIFY system is received. + OnNotification NotificationHandler + + // OnPgError is a callback function called when a Postgres error is received by the server. The default handler will close + // the connection on any FATAL errors. If you override this handler you should call the previously set handler or ensure + // that you close on FATAL errors by returning false. + OnPgError PgErrorHandler + + // OAuthTokenProvider is a function that returns an OAuth token for authentication. If set, it will be used for + // OAUTHBEARER SASL authentication when the server requests it. + OAuthTokenProvider func(context.Context) (string, error) + + // MinProtocolVersion is the minimum acceptable PostgreSQL protocol version. + // If the server does not support at least this version, the connection will fail. + // Valid values: "3.0", "3.2", "latest". Defaults to "3.0". + MinProtocolVersion string + + // MaxProtocolVersion is the maximum PostgreSQL protocol version to request from the server. + // Valid values: "3.0", "3.2", "latest". Defaults to "3.0" for compatibility. + MaxProtocolVersion string + + // ChannelBinding is the channel_binding parameter for SCRAM-SHA-256-PLUS authentication. + // Valid values: "disable", "prefer", "require". Defaults to "prefer". + ChannelBinding string + + // RequireAuth restricts which authentication methods the client will accept from the server, + // matching libpq's require_auth parameter. It is a comma-separated list of method names + // (password, md5, gss, sspi, scram-sha-256, oauth, none). A leading "!" on every entry negates + // the list (forbid these methods, allow all others). Empty (the default) means all methods are + // accepted. + RequireAuth string + + createdByParseConfig bool // Used to enforce created by ParseConfig rule. +} + +// connStringKeyAliases maps libpq parameter keywords to the canonical key names this package +// uses internally in the parsed-settings map. Most keywords are already canonical; this map +// holds only those whose pgx-internal name differs from the libpq spelling. +var connStringKeyAliases = map[string]string{ + "dbname": "database", +} + +// canonicalConnStringKey returns the canonical settings-map key for a libpq parameter keyword. +func canonicalConnStringKey(k string) string { + if c, ok := connStringKeyAliases[k]; ok { + return c + } + return k +} + +// ParseConfigOptions contains options that control how a config is built such as GetSSLPassword. +type ParseConfigOptions struct { + // GetSSLPassword gets the password to decrypt a SSL client certificate. This is analogous to the libpq function + // PQsetSSLKeyPassHook_OpenSSL. + GetSSLPassword GetSSLPasswordFunc + + // ConnStringAllowedKeys, if non-nil, restricts which parameter keys may appear in connString + // itself. Any other key (whether connString is in keyword/value or URL form) causes + // ParseConfigWithOptions to return an error before any filesystem access or network + // resolution is attempted. Environment variables (PGHOST, PGSERVICEFILE, ...) and built-in + // defaults are not checked: only keys that originate from the connString argument. + // + // Keys may be given in either their libpq spelling ("dbname") or pgx-internal spelling + // ("database"); both are accepted. + // + // A nil slice (the default) applies no restriction and matches libpq behaviour. An empty + // non-nil slice rejects every key, i.e. connString must be empty. + // + // Use this when any part of connString is built from input the application does not fully + // control (tenant configuration, RPC parameters, admin UI fields). List only the keys that + // input is expected to supply. This fails closed: a future libpq parameter that pgconn learns + // to parse will be rejected unless the application has explicitly allowed it, rather than + // silently passing through. + ConnStringAllowedKeys []string +} + +// Copy returns a deep copy of the config that is safe to use and modify. +// The only exception is the TLSConfig field: +// according to the tls.Config docs it must not be modified after creation. +func (c *Config) Copy() *Config { + newConf := new(Config) + *newConf = *c + if newConf.TLSConfig != nil { + newConf.TLSConfig = c.TLSConfig.Clone() + } + if newConf.RuntimeParams != nil { + newConf.RuntimeParams = make(map[string]string, len(c.RuntimeParams)) + maps.Copy(newConf.RuntimeParams, c.RuntimeParams) + } + if newConf.Fallbacks != nil { + newConf.Fallbacks = make([]*FallbackConfig, len(c.Fallbacks)) + for i, fallback := range c.Fallbacks { + newFallback := new(FallbackConfig) + *newFallback = *fallback + if newFallback.TLSConfig != nil { + newFallback.TLSConfig = fallback.TLSConfig.Clone() + } + newConf.Fallbacks[i] = newFallback + } + } + return newConf +} + +// FallbackConfig is additional settings to attempt a connection with when the primary Config fails to establish a +// network connection. It is used for TLS fallback such as sslmode=prefer and high availability (HA) connections. +type FallbackConfig struct { + Host string // host (e.g. localhost) or path to unix domain socket directory (e.g. /private/tmp) + Port uint16 + TLSConfig *tls.Config // nil disables TLS +} + +// connectOneConfig is the configuration for a single attempt to connect to a single host. +type connectOneConfig struct { + network string + address string + originalHostname string // original hostname before resolving + tlsConfig *tls.Config // nil disables TLS +} + +// isAbsolutePath checks if the provided value is an absolute path either +// beginning with a forward slash (as on Linux-based systems) or with a capital +// letter A-Z followed by a colon and a backslash, e.g., "C:\", (as on Windows). +func isAbsolutePath(path string) bool { + isWindowsPath := func(p string) bool { + if len(p) < 3 { + return false + } + drive := p[0] + colon := p[1] + backslash := p[2] + if drive >= 'A' && drive <= 'Z' && colon == ':' && backslash == '\\' { + return true + } + return false + } + return strings.HasPrefix(path, "/") || isWindowsPath(path) +} + +// NetworkAddress converts a PostgreSQL host and port into network and address suitable for use with +// net.Dial. +func NetworkAddress(host string, port uint16) (network, address string) { + if isAbsolutePath(host) { + network = "unix" + address = filepath.Join(host, ".s.PGSQL.") + strconv.FormatInt(int64(port), 10) + } else { + network = "tcp" + address = net.JoinHostPort(host, strconv.Itoa(int(port))) + } + return network, address +} + +// ParseConfig builds a *Config from connString with similar behavior to the PostgreSQL standard C library libpq. It +// uses the same defaults as libpq (e.g. port=5432) and understands most PG* environment variables. ParseConfig closely +// matches the parsing behavior of libpq. connString may either be in URL format or keyword = value format. See +// https://www.postgresql.org/docs/current/libpq-connect.html#LIBPQ-CONNSTRING for details. connString also may be empty +// to only read from the environment. If a password is not supplied it will attempt to read the .pgpass file. +// +// # Example Keyword/Value +// user=jack password=secret host=pg.example.com port=5432 dbname=mydb sslmode=verify-ca +// +// # Example URL +// postgres://jack:secret@pg.example.com:5432/mydb?sslmode=verify-ca +// +// The returned *Config may be modified. However, it is strongly recommended that any configuration that can be done +// through the connection string be done there. In particular the fields Host, Port, TLSConfig, and Fallbacks can be +// interdependent (e.g. TLSConfig needs knowledge of the host to validate the server certificate). These fields should +// not be modified individually. They should all be modified or all left unchanged. +// +// ParseConfig supports specifying multiple hosts in similar manner to libpq. Host and port may include comma separated +// values that will be tried in order. This can be used as part of a high availability system. See +// https://www.postgresql.org/docs/current/libpq-connect.html#LIBPQ-MULTIPLE-HOSTS for more information. +// +// # Example URL +// postgres://jack:secret@foo.example.com:5432,bar.example.com:5432/mydb +// +// ParseConfig currently recognizes the following environment variable and their parameter key word equivalents passed +// via database URL or keyword/value: +// +// PGHOST +// PGPORT +// PGDATABASE +// PGUSER +// PGPASSWORD +// PGPASSFILE +// PGSERVICE +// PGSERVICEFILE +// PGSSLMODE +// PGSSLCERT +// PGSSLKEY +// PGSSLROOTCERT +// PGSSLPASSWORD +// PGOPTIONS +// PGAPPNAME +// PGCONNECT_TIMEOUT +// PGTARGETSESSIONATTRS +// PGTZ +// PGMINPROTOCOLVERSION +// PGMAXPROTOCOLVERSION +// +// See http://www.postgresql.org/docs/current/static/libpq-envars.html for details on the meaning of environment variables. +// +// See https://www.postgresql.org/docs/current/libpq-connect.html#LIBPQ-PARAMKEYWORDS for parameter key word names. They are +// usually but not always the environment variable name downcased and without the "PG" prefix. +// +// Important Security Notes: +// +// ParseConfig tries to match libpq behavior with regard to PGSSLMODE. This includes defaulting to "prefer" behavior if +// not set. +// +// See http://www.postgresql.org/docs/current/static/libpq-ssl.html#LIBPQ-SSL-PROTECTION for details on what level of +// security each sslmode provides. +// +// The sslmode "prefer" (the default), sslmode "allow", and multiple hosts are implemented via the Fallbacks field of +// the Config struct. If TLSConfig is manually changed it will not affect the fallbacks. For example, in the case of +// sslmode "prefer" this means it will first try the main Config settings which use TLS, then it will try the fallback +// which does not use TLS. This can lead to an unexpected unencrypted connection if the main TLS config is manually +// changed later but the unencrypted fallback is present. Ensure there are no stale fallbacks when manually setting +// TLSConfig. +// +// Several connection parameters cause ParseConfig to read files from the local filesystem: servicefile, passfile, +// sslkey, sslcert, and sslrootcert. Applications that build connection strings from untrusted input must not allow +// these keys to be set by that input. In particular, servicefile (which pgconn accepts in the connection string; +// libpq does not) is read as an INI file whose entries override other connection settings including host, port, and +// sslmode, so an attacker who controls servicefile and service can redirect the connection. If any portion of the +// connection string is externally supplied, use ParseConfigWithOptions and set ParseConfigOptions.ConnStringAllowedKeys +// to an allow-list of the keys that input is expected to supply; any other key in the connection string is then +// rejected before any filesystem access occurs. +// +// Other known differences with libpq: +// +// When multiple hosts are specified, libpq allows them to have different passwords set via the .pgpass file. pgconn +// does not. +// +// In addition, ParseConfig accepts the following options: +// +// - servicefile. +// libpq only reads servicefile from the PGSERVICEFILE environment variable. ParseConfig accepts servicefile as a +// part of the connection string. +func ParseConfig(connString string) (*Config, error) { + var parseConfigOptions ParseConfigOptions + return ParseConfigWithOptions(connString, parseConfigOptions) +} + +// ParseConfigWithOptions builds a *Config from connString and options with similar behavior to the PostgreSQL standard +// C library libpq. options contains settings that cannot be specified in a connString such as providing a function to +// get the SSL password. +func ParseConfigWithOptions(connString string, options ParseConfigOptions) (*Config, error) { + defaultSettings := defaultSettings() + envSettings := parseEnvSettings() + + connStringSettings := make(map[string]string) + if connString != "" { + var err error + // connString may be a database URL or in PostgreSQL keyword/value format + if strings.HasPrefix(connString, "postgres://") || strings.HasPrefix(connString, "postgresql://") { + connStringSettings, err = parseURLSettings(connString) + if err != nil { + return nil, &ParseConfigError{ConnString: connString, msg: "failed to parse as URL", err: err} + } + } else { + connStringSettings, err = parseKeywordValueSettings(connString) + if err != nil { + return nil, &ParseConfigError{ConnString: connString, msg: "failed to parse as keyword/value", err: err} + } + } + } + + if options.ConnStringAllowedKeys != nil { + allowed := make(map[string]struct{}, len(options.ConnStringAllowedKeys)) + for _, k := range options.ConnStringAllowedKeys { + allowed[canonicalConnStringKey(k)] = struct{}{} + } + for k := range connStringSettings { + if _, ok := allowed[k]; !ok { + return nil, &ParseConfigError{ConnString: connString, msg: fmt.Sprintf("connection string key %q is not in ConnStringAllowedKeys", k)} + } + } + } + + settings := mergeSettings(defaultSettings, envSettings, connStringSettings) + if service, present := settings["service"]; present { + serviceSettings, err := parseServiceSettings(settings["servicefile"], service) + if err != nil { + return nil, &ParseConfigError{ConnString: connString, msg: "failed to read service", err: err} + } + + settings = mergeSettings(defaultSettings, envSettings, serviceSettings, connStringSettings) + } + + config := &Config{ + createdByParseConfig: true, + Database: settings["database"], + User: settings["user"], + Password: settings["password"], + RuntimeParams: make(map[string]string), + BuildFrontend: pgproto3.NewFrontend, + BuildContextWatcherHandler: func(pgConn *PgConn) ctxwatch.Handler { + return &DeadlineContextWatcherHandler{Conn: pgConn.conn} + }, + OnPgError: func(_ *PgConn, pgErr *PgError) bool { + // we want to automatically close any fatal errors + if strings.EqualFold(pgErr.Severity, "FATAL") { + return false + } + return true + }, + } + + if connectTimeoutSetting, present := settings["connect_timeout"]; present { + connectTimeout, err := parseConnectTimeoutSetting(connectTimeoutSetting) + if err != nil { + return nil, &ParseConfigError{ConnString: connString, msg: "invalid connect_timeout", err: err} + } + config.ConnectTimeout = connectTimeout + config.DialFunc = makeConnectTimeoutDialFunc(connectTimeout) + } else { + defaultDialer := makeDefaultDialer() + config.DialFunc = defaultDialer.DialContext + } + + config.LookupFunc = makeDefaultResolver().LookupHost + + notRuntimeParams := map[string]struct{}{ + "host": {}, + "port": {}, + "database": {}, + "user": {}, + "password": {}, + "passfile": {}, + "connect_timeout": {}, + "sslmode": {}, + "sslkey": {}, + "sslcert": {}, + "sslrootcert": {}, + "sslnegotiation": {}, + "sslpassword": {}, + "sslsni": {}, + "krbspn": {}, + "krbsrvname": {}, + "target_session_attrs": {}, + "service": {}, + "servicefile": {}, + "min_protocol_version": {}, + "max_protocol_version": {}, + "channel_binding": {}, + "require_auth": {}, + } + + // Adding kerberos configuration + if _, present := settings["krbsrvname"]; present { + config.KerberosSrvName = settings["krbsrvname"] + } + if _, present := settings["krbspn"]; present { + config.KerberosSpn = settings["krbspn"] + } + + for k, v := range settings { + if _, present := notRuntimeParams[k]; present { + continue + } + config.RuntimeParams[k] = v + } + + fallbacks := []*FallbackConfig{} + + hosts := strings.Split(settings["host"], ",") + ports := strings.Split(settings["port"], ",") + + for i, host := range hosts { + var portStr string + if i < len(ports) { + portStr = ports[i] + } else { + portStr = ports[0] + } + + port, err := parsePort(portStr) + if err != nil { + return nil, &ParseConfigError{ConnString: connString, msg: "invalid port", err: err} + } + + var tlsConfigs []*tls.Config + + // Ignore TLS settings if Unix domain socket like libpq + if network, _ := NetworkAddress(host, port); network == "unix" { + tlsConfigs = append(tlsConfigs, nil) + } else { + var err error + tlsConfigs, err = configTLS(settings, host, options) + if err != nil { + return nil, &ParseConfigError{ConnString: connString, msg: "failed to configure TLS", err: err} + } + } + + for _, tlsConfig := range tlsConfigs { + fallbacks = append(fallbacks, &FallbackConfig{ + Host: host, + Port: port, + TLSConfig: tlsConfig, + }) + } + } + + config.Host = fallbacks[0].Host + config.Port = fallbacks[0].Port + config.TLSConfig = fallbacks[0].TLSConfig + config.Fallbacks = fallbacks[1:] + config.SSLNegotiation = settings["sslnegotiation"] + + passfile, err := pgpassfile.ReadPassfile(settings["passfile"]) + if err == nil { + if config.Password == "" { + host := config.Host + if network, _ := NetworkAddress(config.Host, config.Port); network == "unix" { + host = "localhost" + } + + config.Password = passfile.FindPassword(host, strconv.Itoa(int(config.Port)), config.Database, config.User) + } + } + + switch tsa := settings["target_session_attrs"]; tsa { + case "read-write": + config.ValidateConnect = ValidateConnectTargetSessionAttrsReadWrite + case "read-only": + config.ValidateConnect = ValidateConnectTargetSessionAttrsReadOnly + case "primary": + config.ValidateConnect = ValidateConnectTargetSessionAttrsPrimary + case "standby": + config.ValidateConnect = ValidateConnectTargetSessionAttrsStandby + case "prefer-standby": + config.ValidateConnect = ValidateConnectTargetSessionAttrsPreferStandby + case "any": + // do nothing + default: + return nil, &ParseConfigError{ConnString: connString, msg: fmt.Sprintf("unknown target_session_attrs value: %v", tsa)} + } + + minProto, err := parseProtocolVersion(settings["min_protocol_version"]) + if err != nil { + return nil, &ParseConfigError{ConnString: connString, msg: fmt.Sprintf("invalid min_protocol_version: %q", settings["min_protocol_version"]), err: err} + } + maxProto, err := parseProtocolVersion(settings["max_protocol_version"]) + if err != nil { + return nil, &ParseConfigError{ConnString: connString, msg: fmt.Sprintf("invalid max_protocol_version: %q", settings["max_protocol_version"]), err: err} + } + + config.MinProtocolVersion = settings["min_protocol_version"] + config.MaxProtocolVersion = settings["max_protocol_version"] + + if config.MinProtocolVersion == "" { + config.MinProtocolVersion = "3.0" + } + + // When max_protocol_version is not explicitly set, default based on + // min_protocol_version. This matches libpq behavior: if min > 3.0, + // default max to latest; otherwise default to 3.0 for compatibility + // with older servers/poolers that don't support NegotiateProtocolVersion. + if config.MaxProtocolVersion == "" { + if minProto > pgproto3.ProtocolVersion30 { + config.MaxProtocolVersion = "latest" + } else { + config.MaxProtocolVersion = "3.0" + } + } + + // Only error when max_protocol_version was explicitly set and conflicts + // with min_protocol_version. When max_protocol_version is not explicitly + // set, the auto-raise logic above already ensures a valid default. + if minProto > maxProto && settings["max_protocol_version"] != "" { + return nil, &ParseConfigError{ConnString: connString, msg: "min_protocol_version cannot be greater than max_protocol_version"} + } + + switch channelBinding := settings["channel_binding"]; channelBinding { + case "", "prefer": + config.ChannelBinding = "prefer" + case "disable": + config.ChannelBinding = "disable" + case "require": + config.ChannelBinding = "require" + default: + return nil, &ParseConfigError{ConnString: connString, msg: fmt.Sprintf("unknown channel_binding value: %v", channelBinding)} + } + + config.RequireAuth = settings["require_auth"] + if _, err := parseRequireAuth(config.RequireAuth); err != nil { + return nil, &ParseConfigError{ConnString: connString, msg: "invalid require_auth", err: err} + } + + return config, nil +} + +func mergeSettings(settingSets ...map[string]string) map[string]string { + settings := make(map[string]string) + + for _, s2 := range settingSets { + maps.Copy(settings, s2) + } + + return settings +} + +func parseEnvSettings() map[string]string { + settings := make(map[string]string) + + nameMap := map[string]string{ + "PGHOST": "host", + "PGPORT": "port", + "PGDATABASE": "database", + "PGUSER": "user", + "PGPASSWORD": "password", + "PGPASSFILE": "passfile", + "PGAPPNAME": "application_name", + "PGCONNECT_TIMEOUT": "connect_timeout", + "PGSSLMODE": "sslmode", + "PGSSLKEY": "sslkey", + "PGSSLCERT": "sslcert", + "PGSSLSNI": "sslsni", + "PGSSLROOTCERT": "sslrootcert", + "PGSSLPASSWORD": "sslpassword", + "PGSSLNEGOTIATION": "sslnegotiation", + "PGTARGETSESSIONATTRS": "target_session_attrs", + "PGSERVICE": "service", + "PGSERVICEFILE": "servicefile", + "PGTZ": "timezone", + "PGOPTIONS": "options", + "PGMINPROTOCOLVERSION": "min_protocol_version", + "PGMAXPROTOCOLVERSION": "max_protocol_version", + "PGCHANNELBINDING": "channel_binding", + "PGREQUIREAUTH": "require_auth", + } + + for envname, realname := range nameMap { + value := os.Getenv(envname) + if value != "" { + settings[realname] = value + } + } + + return settings +} + +func parseURLSettings(connString string) (map[string]string, error) { + settings := make(map[string]string) + + parsedURL, err := url.Parse(connString) + if err != nil { + if urlErr := new(url.Error); errors.As(err, &urlErr) { + return nil, urlErr.Err + } + return nil, err + } + + if parsedURL.User != nil { + if u := parsedURL.User.Username(); u != "" { + settings["user"] = u + } + if password, present := parsedURL.User.Password(); present { + settings["password"] = password + } + } + + // Handle multiple host:port's in url.Host by splitting them into host,host,host and port,port,port. + var hosts []string + var ports []string + for host := range strings.SplitSeq(parsedURL.Host, ",") { + if host == "" { + continue + } + if isIPOnly(host) { + hosts = append(hosts, strings.Trim(host, "[]")) + continue + } + h, p, err := net.SplitHostPort(host) + if err != nil { + return nil, fmt.Errorf("failed to split host:port in '%s', err: %w", host, err) + } + if h != "" { + hosts = append(hosts, h) + } + if p != "" { + ports = append(ports, p) + } + } + if len(hosts) > 0 { + settings["host"] = strings.Join(hosts, ",") + } + if len(ports) > 0 { + settings["port"] = strings.Join(ports, ",") + } + + database := strings.TrimLeft(parsedURL.Path, "/") + if database != "" { + settings["database"] = database + } + + for k, v := range parsedURL.Query() { + settings[canonicalConnStringKey(k)] = v[0] + } + + return settings, nil +} + +func isIPOnly(host string) bool { + return net.ParseIP(strings.Trim(host, "[]")) != nil || !strings.Contains(host, ":") +} + +var asciiSpace = [256]uint8{'\t': 1, '\n': 1, '\v': 1, '\f': 1, '\r': 1, ' ': 1} + +func parseKeywordValueSettings(s string) (map[string]string, error) { + settings := make(map[string]string) + + // Trim any leading whitespace so that the loop exits cleanly when only + // spaces remain (e.g. trailing spaces after the last value). + s = strings.TrimLeft(s, " \t\n\r\v\f") + for len(s) > 0 { + var key, val string + eqIdx := strings.IndexRune(s, '=') + if eqIdx < 0 { + return nil, errors.New("invalid keyword/value") + } + + key = strings.Trim(s[:eqIdx], " \t\n\r\v\f") + s = strings.TrimLeft(s[eqIdx+1:], " \t\n\r\v\f") + switch { + case len(s) == 0: + case s[0] != '\'': + end := 0 + for ; end < len(s); end++ { + if asciiSpace[s[end]] == 1 { + break + } + if s[end] == '\\' { + end++ + if end == len(s) { + return nil, errors.New("invalid backslash") + } + } + } + val = strings.ReplaceAll(strings.ReplaceAll(s[:end], "\\\\", "\\"), "\\'", "'") + // Consume the value and trim any subsequent whitespace so that + // multiple trailing spaces don't cause a spurious parse failure. + s = strings.TrimLeft(s[end:], " \t\n\r\v\f") + default: // quoted string + s = s[1:] + end := 0 + for ; end < len(s); end++ { + if s[end] == '\'' { + break + } + if s[end] == '\\' { + end++ + } + } + if end == len(s) { + return nil, errors.New("unterminated quoted string in connection info string") + } + val = strings.ReplaceAll(strings.ReplaceAll(s[:end], "\\\\", "\\"), "\\'", "'") + // Consume the closing quote and any subsequent whitespace. + s = strings.TrimLeft(s[end+1:], " \t\n\r\v\f") + } + + key = canonicalConnStringKey(key) + + if key == "" { + return nil, errors.New("invalid keyword/value") + } + + if key == "user" && val == "" { + continue + } + settings[key] = val + } + + return settings, nil +} + +func parseServiceSettings(servicefilePath, serviceName string) (map[string]string, error) { + servicefile, err := pgservicefile.ReadServicefile(servicefilePath) + if err != nil { + return nil, fmt.Errorf("failed to read service file: %v", servicefilePath) + } + + service, err := servicefile.GetService(serviceName) + if err != nil { + return nil, fmt.Errorf("unable to find service: %v", serviceName) + } + + settings := make(map[string]string, len(service.Settings)) + for k, v := range service.Settings { + settings[canonicalConnStringKey(k)] = v + } + + return settings, nil +} + +// configTLS uses libpq's TLS parameters to construct []*tls.Config. It is +// necessary to allow returning multiple TLS configs as sslmode "allow" and +// "prefer" allow fallback. +func configTLS(settings map[string]string, thisHost string, parseConfigOptions ParseConfigOptions) ([]*tls.Config, error) { + host := thisHost + sslmode := settings["sslmode"] + sslrootcert := settings["sslrootcert"] + sslcert := settings["sslcert"] + sslkey := settings["sslkey"] + sslpassword := settings["sslpassword"] + sslsni := settings["sslsni"] + sslnegotiation := settings["sslnegotiation"] + + // Match libpq default behavior + if sslmode == "" { + sslmode = "prefer" + } + if sslsni == "" { + sslsni = "1" + } + + tlsConfig := &tls.Config{} + + if sslnegotiation == "direct" { + tlsConfig.NextProtos = []string{"postgresql"} + if sslmode == "prefer" { + sslmode = "require" + } + } + + if sslrootcert != "" { + var caCertPool *x509.CertPool + + if sslrootcert == "system" { + var err error + + caCertPool, err = x509.SystemCertPool() + if err != nil { + return nil, fmt.Errorf("unable to load system certificate pool: %w", err) + } + + sslmode = "verify-full" + } else { + caCertPool = x509.NewCertPool() + + caPath := sslrootcert + caCert, err := os.ReadFile(caPath) + if err != nil { + return nil, fmt.Errorf("unable to read CA file: %w", err) + } + + if !caCertPool.AppendCertsFromPEM(caCert) { + return nil, errors.New("unable to add CA to cert pool") + } + } + + tlsConfig.RootCAs = caCertPool + tlsConfig.ClientCAs = caCertPool + } + + switch sslmode { + case "disable": + return []*tls.Config{nil}, nil + case "allow", "prefer": + tlsConfig.InsecureSkipVerify = true + case "require": + // According to PostgreSQL documentation, if a root CA file exists, + // the behavior of sslmode=require should be the same as that of verify-ca + // + // See https://www.postgresql.org/docs/current/libpq-ssl.html + if sslrootcert != "" { + goto nextCase + } + tlsConfig.InsecureSkipVerify = true + break + nextCase: + fallthrough + case "verify-ca": + // Don't perform the default certificate verification because it + // will verify the hostname. Instead, verify the server's + // certificate chain ourselves in VerifyPeerCertificate and + // ignore the server name. This emulates libpq's verify-ca + // behavior. + // + // See https://github.com/golang/go/issues/21971#issuecomment-332693931 + // and https://pkg.go.dev/crypto/tls?tab=doc#example-Config-VerifyPeerCertificate + // for more info. + tlsConfig.InsecureSkipVerify = true + tlsConfig.VerifyPeerCertificate = func(certificates [][]byte, _ [][]*x509.Certificate) error { + certs := make([]*x509.Certificate, len(certificates)) + for i, asn1Data := range certificates { + cert, err := x509.ParseCertificate(asn1Data) + if err != nil { + return errors.New("failed to parse certificate from server: " + err.Error()) + } + certs[i] = cert + } + + // Leave DNSName empty to skip hostname verification. + opts := x509.VerifyOptions{ + Roots: tlsConfig.RootCAs, + Intermediates: x509.NewCertPool(), + } + // Skip the first cert because it's the leaf. All others + // are intermediates. + for _, cert := range certs[1:] { + opts.Intermediates.AddCert(cert) + } + _, err := certs[0].Verify(opts) + return err + } + case "verify-full": + tlsConfig.ServerName = host + default: + return nil, errors.New("sslmode is invalid") + } + + if (sslcert != "" && sslkey == "") || (sslcert == "" && sslkey != "") { + return nil, errors.New(`both "sslcert" and "sslkey" are required`) + } + + if sslcert != "" && sslkey != "" { + buf, err := os.ReadFile(sslkey) + if err != nil { + return nil, fmt.Errorf("unable to read sslkey: %w", err) + } + block, _ := pem.Decode(buf) + if block == nil { + return nil, errors.New("failed to decode sslkey") + } + var pemKey []byte + var decryptedKey []byte + var decryptedError error + // If PEM is encrypted, attempt to decrypt using pass phrase + if x509.IsEncryptedPEMBlock(block) { + // Attempt decryption with pass phrase + // NOTE: only supports RSA (PKCS#1) + if sslpassword != "" { + decryptedKey, decryptedError = x509.DecryptPEMBlock(block, []byte(sslpassword)) //nolint:ineffassign + } + // if sslpassword not provided or has decryption error when use it + // try to find sslpassword with callback function + if sslpassword == "" || decryptedError != nil { + if parseConfigOptions.GetSSLPassword != nil { + sslpassword = parseConfigOptions.GetSSLPassword(context.Background()) + } + if sslpassword == "" { + return nil, fmt.Errorf("unable to find sslpassword") + } + } + decryptedKey, decryptedError = x509.DecryptPEMBlock(block, []byte(sslpassword)) + // Should we also provide warning for PKCS#1 needed? + if decryptedError != nil { + return nil, fmt.Errorf("unable to decrypt key: %w", decryptedError) + } + + pemBytes := pem.Block{ + Type: "RSA PRIVATE KEY", + Bytes: decryptedKey, + } + pemKey = pem.EncodeToMemory(&pemBytes) + } else { + pemKey = pem.EncodeToMemory(block) + } + certfile, err := os.ReadFile(sslcert) + if err != nil { + return nil, fmt.Errorf("unable to read cert: %w", err) + } + cert, err := tls.X509KeyPair(certfile, pemKey) + if err != nil { + return nil, fmt.Errorf("unable to load cert: %w", err) + } + tlsConfig.Certificates = []tls.Certificate{cert} + } + + // Set Server Name Indication (SNI), if enabled by connection parameters. + // Per RFC 6066, do not set it if the host is a literal IP address (IPv4 + // or IPv6). + if sslsni == "1" && net.ParseIP(host) == nil { + tlsConfig.ServerName = host + } + + switch sslmode { + case "allow": + return []*tls.Config{nil, tlsConfig}, nil + case "prefer": + return []*tls.Config{tlsConfig, nil}, nil + case "require", "verify-ca", "verify-full": + return []*tls.Config{tlsConfig}, nil + default: + panic("BUG: bad sslmode should already have been caught") + } +} + +func parsePort(s string) (uint16, error) { + port, err := strconv.ParseUint(s, 10, 16) + if err != nil { + return 0, err + } + if port < 1 || port > math.MaxUint16 { + return 0, errors.New("outside range") + } + return uint16(port), nil +} + +func makeDefaultDialer() *net.Dialer { + // rely on GOLANG KeepAlive settings + return &net.Dialer{} +} + +func makeDefaultResolver() *net.Resolver { + return net.DefaultResolver +} + +func parseConnectTimeoutSetting(s string) (time.Duration, error) { + timeout, err := strconv.ParseInt(s, 10, 64) + if err != nil { + return 0, err + } + if timeout < 0 { + return 0, errors.New("negative timeout") + } + return time.Duration(timeout) * time.Second, nil +} + +func makeConnectTimeoutDialFunc(timeout time.Duration) DialFunc { + d := makeDefaultDialer() + d.Timeout = timeout + return d.DialContext +} + +// ValidateConnectTargetSessionAttrsReadWrite is a ValidateConnectFunc that implements libpq compatible +// target_session_attrs=read-write. +func ValidateConnectTargetSessionAttrsReadWrite(ctx context.Context, pgConn *PgConn) error { + result, err := pgConn.Exec(ctx, "show transaction_read_only").ReadAll() + if err != nil { + return err + } + + if string(result[0].Rows[0][0]) == "on" { + return errors.New("read only connection") + } + + return nil +} + +// ValidateConnectTargetSessionAttrsReadOnly is a ValidateConnectFunc that implements libpq compatible +// target_session_attrs=read-only. +func ValidateConnectTargetSessionAttrsReadOnly(ctx context.Context, pgConn *PgConn) error { + result, err := pgConn.Exec(ctx, "show transaction_read_only").ReadAll() + if err != nil { + return err + } + + if string(result[0].Rows[0][0]) != "on" { + return errors.New("connection is not read only") + } + + return nil +} + +// ValidateConnectTargetSessionAttrsStandby is a ValidateConnectFunc that implements libpq compatible +// target_session_attrs=standby. +func ValidateConnectTargetSessionAttrsStandby(ctx context.Context, pgConn *PgConn) error { + result, err := pgConn.Exec(ctx, "select pg_is_in_recovery()").ReadAll() + if err != nil { + return err + } + + if string(result[0].Rows[0][0]) != "t" { + return errors.New("server is not in hot standby mode") + } + + return nil +} + +// ValidateConnectTargetSessionAttrsPrimary is a ValidateConnectFunc that implements libpq compatible +// target_session_attrs=primary. +func ValidateConnectTargetSessionAttrsPrimary(ctx context.Context, pgConn *PgConn) error { + result, err := pgConn.Exec(ctx, "select pg_is_in_recovery()").ReadAll() + if err != nil { + return err + } + + if string(result[0].Rows[0][0]) == "t" { + return errors.New("server is in standby mode") + } + + return nil +} + +// ValidateConnectTargetSessionAttrsPreferStandby is a ValidateConnectFunc that implements libpq compatible +// target_session_attrs=prefer-standby. +func ValidateConnectTargetSessionAttrsPreferStandby(ctx context.Context, pgConn *PgConn) error { + result, err := pgConn.Exec(ctx, "select pg_is_in_recovery()").ReadAll() + if err != nil { + return err + } + + if string(result[0].Rows[0][0]) != "t" { + return &NotPreferredError{err: errors.New("server is not in hot standby mode")} + } + + return nil +} + +func parseProtocolVersion(s string) (uint32, error) { + switch s { + case "", "3.0": + return pgproto3.ProtocolVersion30, nil + case "3.2", "latest": + return pgproto3.ProtocolVersion32, nil + default: + return 0, fmt.Errorf("invalid protocol version: %q", s) + } +} diff --git a/vendor/github.com/jackc/pgx/v5/pgconn/ctxwatch/context_watcher.go b/vendor/github.com/jackc/pgx/v5/pgconn/ctxwatch/context_watcher.go new file mode 100644 index 0000000000..b8892e68b1 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgconn/ctxwatch/context_watcher.go @@ -0,0 +1,72 @@ +package ctxwatch + +import ( + "context" + "sync" +) + +// ContextWatcher watches a context and performs an action when the context is canceled. It can watch one context at a +// time. +type ContextWatcher struct { + handler Handler + + // Lock protects the members below. + lock sync.Mutex + // Stop is the handle for an "after func". See [context.AfterFunc]. + stop func() bool + done chan struct{} +} + +// NewContextWatcher returns a ContextWatcher. onCancel will be called when a watched context is canceled. +// OnUnwatchAfterCancel will be called when Unwatch is called and the watched context had already been canceled and +// onCancel called. +func NewContextWatcher(handler Handler) *ContextWatcher { + cw := &ContextWatcher{ + handler: handler, + } + + return cw +} + +// Watch starts watching ctx. If ctx is canceled then the onCancel function passed to NewContextWatcher will be called. +func (cw *ContextWatcher) Watch(ctx context.Context) { + cw.lock.Lock() + defer cw.lock.Unlock() + + if cw.stop != nil { + panic("watch already in progress") + } + + if ctx.Done() != nil { + cw.done = make(chan struct{}) + cw.stop = context.AfterFunc(ctx, func() { + cw.handler.HandleCancel(ctx) + close(cw.done) + }) + } +} + +// Unwatch stops watching the previously watched context. If the onCancel function passed to NewContextWatcher was +// called then onUnwatchAfterCancel will also be called. +func (cw *ContextWatcher) Unwatch() { + cw.lock.Lock() + defer cw.lock.Unlock() + + if cw.stop != nil { + if !cw.stop() { + <-cw.done + cw.handler.HandleUnwatchAfterCancel() + } + cw.stop = nil + cw.done = nil + } +} + +type Handler interface { + // HandleCancel is called when the context that a ContextWatcher is currently watching is canceled. canceledCtx is the + // context that was canceled. + HandleCancel(canceledCtx context.Context) + + // HandleUnwatchAfterCancel is called when a ContextWatcher that called HandleCancel on this Handler is unwatched. + HandleUnwatchAfterCancel() +} diff --git a/vendor/github.com/jackc/pgx/v5/pgconn/defaults.go b/vendor/github.com/jackc/pgx/v5/pgconn/defaults.go new file mode 100644 index 0000000000..1dd514ff44 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgconn/defaults.go @@ -0,0 +1,63 @@ +//go:build !windows +// +build !windows + +package pgconn + +import ( + "os" + "os/user" + "path/filepath" +) + +func defaultSettings() map[string]string { + settings := make(map[string]string) + + settings["host"] = defaultHost() + settings["port"] = "5432" + + // Default to the OS user name. Purposely ignoring err getting user name from + // OS. The client application will simply have to specify the user in that + // case (which they typically will be doing anyway). + user, err := user.Current() + if err == nil { + settings["user"] = user.Username + settings["passfile"] = filepath.Join(user.HomeDir, ".pgpass") + settings["servicefile"] = filepath.Join(user.HomeDir, ".pg_service.conf") + sslcert := filepath.Join(user.HomeDir, ".postgresql", "postgresql.crt") + sslkey := filepath.Join(user.HomeDir, ".postgresql", "postgresql.key") + if _, err := os.Stat(sslcert); err == nil { + if _, err := os.Stat(sslkey); err == nil { + // Both the cert and key must be present to use them, or do not use either + settings["sslcert"] = sslcert + settings["sslkey"] = sslkey + } + } + sslrootcert := filepath.Join(user.HomeDir, ".postgresql", "root.crt") + if _, err := os.Stat(sslrootcert); err == nil { + settings["sslrootcert"] = sslrootcert + } + } + + settings["target_session_attrs"] = "any" + + return settings +} + +// defaultHost attempts to mimic libpq's default host. libpq uses the default unix socket location on *nix and localhost +// on Windows. The default socket location is compiled into libpq. Since pgx does not have access to that default it +// checks the existence of common locations. +func defaultHost() string { + candidatePaths := []string{ + "/var/run/postgresql", // Debian + "/private/tmp", // OSX - homebrew + "/tmp", // standard PostgreSQL + } + + for _, path := range candidatePaths { + if _, err := os.Stat(path); err == nil { + return path + } + } + + return "localhost" +} diff --git a/vendor/github.com/jackc/pgx/v5/pgconn/defaults_windows.go b/vendor/github.com/jackc/pgx/v5/pgconn/defaults_windows.go new file mode 100644 index 0000000000..33b4a1ff85 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgconn/defaults_windows.go @@ -0,0 +1,57 @@ +package pgconn + +import ( + "os" + "os/user" + "path/filepath" + "strings" +) + +func defaultSettings() map[string]string { + settings := make(map[string]string) + + settings["host"] = defaultHost() + settings["port"] = "5432" + + // Default to the OS user name. Purposely ignoring err getting user name from + // OS. The client application will simply have to specify the user in that + // case (which they typically will be doing anyway). + user, err := user.Current() + appData := os.Getenv("APPDATA") + if err == nil { + // Windows gives us the username here as `DOMAIN\user` or `LOCALPCNAME\user`, + // but the libpq default is just the `user` portion, so we strip off the first part. + username := user.Username + if strings.Contains(username, "\\") { + username = username[strings.LastIndex(username, "\\")+1:] + } + + settings["user"] = username + settings["passfile"] = filepath.Join(appData, "postgresql", "pgpass.conf") + settings["servicefile"] = filepath.Join(user.HomeDir, ".pg_service.conf") + sslcert := filepath.Join(appData, "postgresql", "postgresql.crt") + sslkey := filepath.Join(appData, "postgresql", "postgresql.key") + if _, err := os.Stat(sslcert); err == nil { + if _, err := os.Stat(sslkey); err == nil { + // Both the cert and key must be present to use them, or do not use either + settings["sslcert"] = sslcert + settings["sslkey"] = sslkey + } + } + sslrootcert := filepath.Join(appData, "postgresql", "root.crt") + if _, err := os.Stat(sslrootcert); err == nil { + settings["sslrootcert"] = sslrootcert + } + } + + settings["target_session_attrs"] = "any" + + return settings +} + +// defaultHost attempts to mimic libpq's default host. libpq uses the default unix socket location on *nix and localhost +// on Windows. The default socket location is compiled into libpq. Since pgx does not have access to that default it +// checks the existence of common locations. +func defaultHost() string { + return "localhost" +} diff --git a/vendor/github.com/jackc/pgx/v5/pgconn/doc.go b/vendor/github.com/jackc/pgx/v5/pgconn/doc.go new file mode 100644 index 0000000000..f9707262e7 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgconn/doc.go @@ -0,0 +1,64 @@ +// Package pgconn is a low-level PostgreSQL database driver. +/* +pgconn provides lower level access to a PostgreSQL connection than a database/sql or pgx connection. It operates at +nearly the same level is the C library libpq. + +Establishing a Connection + +Use Connect to establish a connection. It accepts a connection string in URL or keyword/value format and will read the +environment for libpq style environment variables. + +Connecting Securely + +By default ParseConfig matches libpq and uses sslmode=prefer, which silently falls back to an unencrypted connection +if the server does not offer TLS. For connections that traverse an untrusted network, set the following parameters +explicitly: + + # URL form + postgres://user@db.example.com/mydb?sslmode=verify-full&sslrootcert=/path/to/root.crt&channel_binding=require&require_auth=scram-sha-256 + + # keyword/value form + host=db.example.com user=user dbname=mydb sslmode=verify-full sslrootcert=/path/to/root.crt channel_binding=require require_auth=scram-sha-256 + + sslmode=verify-full Require TLS, verify the server certificate against sslrootcert, and verify that the + certificate's host name matches the host being connected to. Weaker modes (disable, + allow, prefer, require) either permit plaintext fallback or skip certificate + verification, allowing a network attacker to impersonate the server. + channel_binding=require Require SCRAM-SHA-256-PLUS, which binds the authentication exchange to the TLS channel + so that a TLS-terminating intermediary cannot relay credentials to the real server. + require_auth=scram-sha-256 + Refuse to respond to AuthenticationCleartextPassword or AuthenticationMD5Password + requests from the server. Without this, a server (or interceptor) can request the + password in cleartext and the client will send it. + +These parameters may also be set via the PGSSLMODE, PGSSLROOTCERT, PGCHANNELBINDING, and PGREQUIREAUTH environment +variables. + +Executing a Query + +ExecParams and ExecPrepared execute a single query. They return readers that iterate over each row. The Read method +reads all rows into memory. + +Executing Multiple Queries in a Single Round Trip + +Exec and ExecBatch can execute multiple queries in a single round trip. They return readers that iterate over each query +result. The ReadAll method reads all query results into memory. + +Pipeline Mode + +Pipeline mode allows sending queries without having read the results of previously sent queries. It allows control of +exactly how many and when network round trips occur. + +Context Support + +All potentially blocking operations take a context.Context. The default behavior when a context is canceled is for the +method to immediately return. In most circumstances, this will also close the underlying connection. This behavior can +be customized by using BuildContextWatcherHandler on the Config to create a ctxwatch.Handler with different behavior. +This can be especially useful when queries that are frequently canceled and the overhead of creating new connections is +a problem. DeadlineContextWatcherHandler and CancelRequestContextWatcherHandler can be used to introduce a delay before +interrupting the query in such a way as to close the connection. + +The CancelRequest method may be used to request the PostgreSQL server cancel an in-progress query without forcing the +client to abort. +*/ +package pgconn diff --git a/vendor/github.com/jackc/pgx/v5/pgconn/errors.go b/vendor/github.com/jackc/pgx/v5/pgconn/errors.go new file mode 100644 index 0000000000..9fbe68cb65 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgconn/errors.go @@ -0,0 +1,287 @@ +package pgconn + +import ( + "context" + "errors" + "fmt" + "net" + "net/url" + "regexp" + "strings" +) + +// SafeToRetry checks if the err is guaranteed to have occurred before sending any data to the server. +func SafeToRetry(err error) bool { + var retryableErr interface{ SafeToRetry() bool } + if errors.As(err, &retryableErr) { + return retryableErr.SafeToRetry() + } + return false +} + +// Timeout checks if err was caused by a timeout. To be specific, it is true if err was caused within pgconn by a +// context.DeadlineExceeded or an implementer of net.Error where Timeout() is true. +func Timeout(err error) bool { + var timeoutErr *errTimeout + return errors.As(err, &timeoutErr) +} + +// PgError represents an error reported by the PostgreSQL server. See +// http://www.postgresql.org/docs/current/static/protocol-error-fields.html for +// detailed field description. +type PgError struct { + Severity string + SeverityUnlocalized string + Code string + Message string + Detail string + Hint string + Position int32 + InternalPosition int32 + InternalQuery string + Where string + SchemaName string + TableName string + ColumnName string + DataTypeName string + ConstraintName string + File string + Line int32 + Routine string +} + +func (pe *PgError) Error() string { + return pe.Severity + ": " + pe.Message + " (SQLSTATE " + pe.Code + ")" +} + +// SQLState returns the SQLState of the error. +func (pe *PgError) SQLState() string { + return pe.Code +} + +// ConnectError is the error returned when a connection attempt fails. +type ConnectError struct { + Config *Config // The configuration that was used in the connection attempt. + err error +} + +func (e *ConnectError) Error() string { + prefix := fmt.Sprintf("failed to connect to `user=%s database=%s`:", e.Config.User, e.Config.Database) + details := e.err.Error() + if strings.Contains(details, "\n") { + return prefix + "\n\t" + strings.ReplaceAll(details, "\n", "\n\t") + } else { + return prefix + " " + details + } +} + +func (e *ConnectError) Unwrap() error { + return e.err +} + +type perDialConnectError struct { + address string + originalHostname string + err error +} + +func (e *perDialConnectError) Error() string { + return fmt.Sprintf("%s (%s): %s", e.address, e.originalHostname, e.err.Error()) +} + +func (e *perDialConnectError) Unwrap() error { + return e.err +} + +// ErrConnClosed is returned (possibly wrapped) when an operation is attempted +// on a connection that the driver has already closed, e.g. because a prior +// query was cancelled mid-flight or the underlying socket went away. Use +// errors.Is to test for it, since it shows up wrapped inside connLockError. +var ErrConnClosed = errors.New("conn closed") + +type connLockError struct { + status string +} + +func (e *connLockError) SafeToRetry() bool { + return true // a lock failure by definition happens before the connection is used. +} + +func (e *connLockError) Error() string { + return e.status +} + +func (e *connLockError) Unwrap() error { + if e.status == "conn closed" { + return ErrConnClosed + } + return nil +} + +// ParseConfigError is the error returned when a connection string cannot be parsed. +type ParseConfigError struct { + ConnString string // The connection string that could not be parsed. + msg string + err error +} + +func NewParseConfigError(conn, msg string, err error) error { + return &ParseConfigError{ + ConnString: conn, + msg: msg, + err: err, + } +} + +func (e *ParseConfigError) Error() string { + // Now that ParseConfigError is public and ConnString is available to the developer, perhaps it would be better only + // return a static string. That would ensure that the error message cannot leak a password. The ConnString field would + // allow access to the original string if desired and Unwrap would allow access to the underlying error. + connString := redactPW(e.ConnString) + if e.err == nil { + return fmt.Sprintf("cannot parse `%s`: %s", connString, e.msg) + } + return fmt.Sprintf("cannot parse `%s`: %s (%s)", connString, e.msg, e.err.Error()) +} + +func (e *ParseConfigError) Unwrap() error { + return e.err +} + +func normalizeTimeoutError(ctx context.Context, err error) error { + var netErr net.Error + if errors.As(err, &netErr) && netErr.Timeout() { + switch ctx.Err() { + case context.Canceled: + // Since the timeout was caused by a context cancellation, the actual error is context.Canceled not the timeout error. + return context.Canceled + case context.DeadlineExceeded: + return &errTimeout{err: ctx.Err()} + default: + return &errTimeout{err: err} + } + } + return err +} + +type pgconnError struct { + msg string + err error + safeToRetry bool +} + +func (e *pgconnError) Error() string { + if e.msg == "" { + return e.err.Error() + } + if e.err == nil { + return e.msg + } + return fmt.Sprintf("%s: %s", e.msg, e.err.Error()) +} + +func (e *pgconnError) SafeToRetry() bool { + return e.safeToRetry +} + +func (e *pgconnError) Unwrap() error { + return e.err +} + +// errTimeout occurs when an error was caused by a timeout. Specifically, it wraps an error which is +// context.Canceled, context.DeadlineExceeded, or an implementer of net.Error where Timeout() is true. +type errTimeout struct { + err error +} + +func (e *errTimeout) Error() string { + return fmt.Sprintf("timeout: %s", e.err.Error()) +} + +func (e *errTimeout) SafeToRetry() bool { + return SafeToRetry(e.err) +} + +func (e *errTimeout) Unwrap() error { + return e.err +} + +type contextAlreadyDoneError struct { + err error +} + +func (e *contextAlreadyDoneError) Error() string { + return fmt.Sprintf("context already done: %s", e.err.Error()) +} + +func (e *contextAlreadyDoneError) SafeToRetry() bool { + return true +} + +func (e *contextAlreadyDoneError) Unwrap() error { + return e.err +} + +// newContextAlreadyDoneError double-wraps a context error in `contextAlreadyDoneError` and `errTimeout`. +func newContextAlreadyDoneError(ctx context.Context) (err error) { + return &errTimeout{&contextAlreadyDoneError{err: ctx.Err()}} +} + +func redactPW(connString string) string { + if strings.HasPrefix(connString, "postgres://") || strings.HasPrefix(connString, "postgresql://") { + if u, err := url.Parse(connString); err == nil { + return redactURL(u) + } + } + quotedKV := regexp.MustCompile(`password='[^']*'`) + connString = quotedKV.ReplaceAllLiteralString(connString, "password=xxxxx") + plainKV := regexp.MustCompile(`password=[^ ]*`) + connString = plainKV.ReplaceAllLiteralString(connString, "password=xxxxx") + brokenURL := regexp.MustCompile(`:[^:@]+?@`) + connString = brokenURL.ReplaceAllLiteralString(connString, ":xxxxxx@") + return connString +} + +func redactURL(u *url.URL) string { + if u == nil { + return "" + } + if _, pwSet := u.User.Password(); pwSet { + u.User = url.UserPassword(u.User.Username(), "xxxxx") + } + return u.String() +} + +type NotPreferredError struct { + err error + safeToRetry bool +} + +func (e *NotPreferredError) Error() string { + return fmt.Sprintf("standby server not found: %s", e.err.Error()) +} + +func (e *NotPreferredError) SafeToRetry() bool { + return e.safeToRetry +} + +func (e *NotPreferredError) Unwrap() error { + return e.err +} + +type PrepareError struct { + err error + + ParseComplete bool // Indicates whether the error occurred after a ParseComplete message was received. +} + +func (e *PrepareError) Error() string { + if e.ParseComplete { + return fmt.Sprintf("prepare failed after ParseComplete: %s", e.err.Error()) + } + return e.err.Error() +} + +func (e *PrepareError) Unwrap() error { + return e.err +} diff --git a/vendor/github.com/jackc/pgx/v5/pgconn/internal/bgreader/bgreader.go b/vendor/github.com/jackc/pgx/v5/pgconn/internal/bgreader/bgreader.go new file mode 100644 index 0000000000..e65c2c2bf2 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgconn/internal/bgreader/bgreader.go @@ -0,0 +1,139 @@ +// Package bgreader provides a io.Reader that can optionally buffer reads in the background. +package bgreader + +import ( + "io" + "sync" + + "github.com/jackc/pgx/v5/internal/iobufpool" +) + +const ( + StatusStopped = iota + StatusRunning + StatusStopping +) + +// BGReader is an io.Reader that can optionally buffer reads in the background. It is safe for concurrent use. +type BGReader struct { + r io.Reader + + cond *sync.Cond + status int32 + readResults []readResult +} + +type readResult struct { + buf *[]byte + err error +} + +// Start starts the backgrounder reader. If the background reader is already running this is a no-op. The background +// reader will stop automatically when the underlying reader returns an error. +func (r *BGReader) Start() { + r.cond.L.Lock() + defer r.cond.L.Unlock() + + switch r.status { + case StatusStopped: + r.status = StatusRunning + go r.bgRead() + case StatusRunning: + // no-op + case StatusStopping: + r.status = StatusRunning + } +} + +// Stop tells the background reader to stop after the in progress Read returns. It is safe to call Stop when the +// background reader is not running. +func (r *BGReader) Stop() { + r.cond.L.Lock() + defer r.cond.L.Unlock() + + switch r.status { + case StatusStopped: + // no-op + case StatusRunning: + r.status = StatusStopping + case StatusStopping: + // no-op + } +} + +// Status returns the current status of the background reader. +func (r *BGReader) Status() int32 { + r.cond.L.Lock() + defer r.cond.L.Unlock() + return r.status +} + +func (r *BGReader) bgRead() { + keepReading := true + for keepReading { + buf := iobufpool.Get(8192) + n, err := r.r.Read(*buf) + *buf = (*buf)[:n] + + r.cond.L.Lock() + r.readResults = append(r.readResults, readResult{buf: buf, err: err}) + if r.status == StatusStopping || err != nil { + r.status = StatusStopped + keepReading = false + } + r.cond.L.Unlock() + r.cond.Broadcast() + } +} + +// Read implements the io.Reader interface. +func (r *BGReader) Read(p []byte) (int, error) { + r.cond.L.Lock() + defer r.cond.L.Unlock() + + if len(r.readResults) > 0 { + return r.readFromReadResults(p) + } + + // There are no unread background read results and the background reader is stopped. + if r.status == StatusStopped { + return r.r.Read(p) + } + + // Wait for results from the background reader + for len(r.readResults) == 0 { + r.cond.Wait() + } + return r.readFromReadResults(p) +} + +// readBackgroundResults reads a result previously read by the background reader. r.cond.L must be held. +func (r *BGReader) readFromReadResults(p []byte) (int, error) { + buf := r.readResults[0].buf + var err error + + n := copy(p, *buf) + if n == len(*buf) { + err = r.readResults[0].err + iobufpool.Put(buf) + if len(r.readResults) == 1 { + r.readResults = nil + } else { + r.readResults = r.readResults[1:] + } + } else { + *buf = (*buf)[n:] + r.readResults[0].buf = buf + } + + return n, err +} + +func New(r io.Reader) *BGReader { + return &BGReader{ + r: r, + cond: &sync.Cond{ + L: &sync.Mutex{}, + }, + } +} diff --git a/vendor/github.com/jackc/pgx/v5/pgconn/krb5.go b/vendor/github.com/jackc/pgx/v5/pgconn/krb5.go new file mode 100644 index 0000000000..efb0d61b87 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgconn/krb5.go @@ -0,0 +1,100 @@ +package pgconn + +import ( + "errors" + "fmt" + + "github.com/jackc/pgx/v5/pgproto3" +) + +// NewGSSFunc creates a GSS authentication provider, for use with +// RegisterGSSProvider. +type NewGSSFunc func() (GSS, error) + +var newGSS NewGSSFunc + +// RegisterGSSProvider registers a GSS authentication provider. For example, if +// you need to use Kerberos to authenticate with your server, add this to your +// main package: +// +// import "github.com/otan/gopgkrb5" +// +// func init() { +// pgconn.RegisterGSSProvider(func() (pgconn.GSS, error) { return gopgkrb5.NewGSS() }) +// } +func RegisterGSSProvider(newGSSArg NewGSSFunc) { + newGSS = newGSSArg +} + +// GSS provides GSSAPI authentication (e.g., Kerberos). +type GSS interface { + GetInitToken(host, service string) ([]byte, error) + GetInitTokenFromSPN(spn string) ([]byte, error) + Continue(inToken []byte) (done bool, outToken []byte, err error) +} + +func (c *PgConn) gssAuth() error { + if newGSS == nil { + return errors.New("kerberos error: no GSSAPI provider registered, see https://github.com/otan/gopgkrb5") + } + cli, err := newGSS() + if err != nil { + return err + } + + var nextData []byte + if c.config.KerberosSpn != "" { + // Use the supplied SPN if provided. + nextData, err = cli.GetInitTokenFromSPN(c.config.KerberosSpn) + } else { + // Allow the kerberos service name to be overridden + service := "postgres" + if c.config.KerberosSrvName != "" { + service = c.config.KerberosSrvName + } + nextData, err = cli.GetInitToken(c.config.Host, service) + } + if err != nil { + return err + } + + for { + gssResponse := &pgproto3.GSSResponse{ + Data: nextData, + } + c.frontend.Send(gssResponse) + err = c.flushWithPotentialWriteReadDeadlock() + if err != nil { + return err + } + resp, err := c.rxGSSContinue() + if err != nil { + return err + } + var done bool + done, nextData, err = cli.Continue(resp.Data) + if err != nil { + return err + } + if done { + break + } + } + return nil +} + +func (c *PgConn) rxGSSContinue() (*pgproto3.AuthenticationGSSContinue, error) { + msg, err := c.receiveMessage() + if err != nil { + return nil, err + } + + switch m := msg.(type) { + case *pgproto3.AuthenticationGSSContinue: + return m, nil + case *pgproto3.ErrorResponse: + return nil, ErrorResponseToPgError(m) + } + + return nil, fmt.Errorf("expected AuthenticationGSSContinue message but received unexpected message %T", msg) +} diff --git a/vendor/github.com/jackc/pgx/v5/pgconn/pgconn.go b/vendor/github.com/jackc/pgx/v5/pgconn/pgconn.go new file mode 100644 index 0000000000..d181f7f5f2 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgconn/pgconn.go @@ -0,0 +1,3042 @@ +package pgconn + +import ( + "container/list" + "context" + "crypto/md5" + "crypto/tls" + "encoding/binary" + "encoding/hex" + "errors" + "fmt" + "io" + "maps" + "math" + "net" + "strconv" + "strings" + "sync" + "time" + + "github.com/jackc/pgx/v5/internal/iobufpool" + "github.com/jackc/pgx/v5/internal/pgio" + "github.com/jackc/pgx/v5/pgconn/ctxwatch" + "github.com/jackc/pgx/v5/pgconn/internal/bgreader" + "github.com/jackc/pgx/v5/pgproto3" + "github.com/jackc/pgx/v5/pgtype" +) + +const ( + connStatusUninitialized = iota + connStatusConnecting + connStatusClosed + connStatusIdle + connStatusBusy +) + +// Notice represents a notice response message reported by the PostgreSQL server. Be aware that this is distinct from +// LISTEN/NOTIFY notification. +type Notice PgError + +// Notification is a message received from the PostgreSQL LISTEN/NOTIFY system +type Notification struct { + PID uint32 // backend pid that sent the notification + Channel string // channel from which notification was received + Payload string +} + +// DialFunc is a function that can be used to connect to a PostgreSQL server. +type DialFunc func(ctx context.Context, network, addr string) (net.Conn, error) + +// LookupFunc is a function that can be used to lookup IPs addrs from host. Optionally an ip:port combination can be +// returned in order to override the connection string's port. +type LookupFunc func(ctx context.Context, host string) (addrs []string, err error) + +// BuildFrontendFunc is a function that can be used to create Frontend implementation for connection. +type BuildFrontendFunc func(r io.Reader, w io.Writer) *pgproto3.Frontend + +// PgErrorHandler is a function that handles errors returned from Postgres. This function must return true to keep +// the connection open. Returning false will cause the connection to be closed immediately. You should return +// false on any FATAL-severity errors. This will not receive network errors. The *PgConn is provided so the handler is +// aware of the origin of the error, but it must not invoke any query method. +type PgErrorHandler func(*PgConn, *PgError) bool + +// NoticeHandler is a function that can handle notices received from the PostgreSQL server. Notices can be received at +// any time, usually during handling of a query response. The *PgConn is provided so the handler is aware of the origin +// of the notice, but it must not invoke any query method. Be aware that this is distinct from LISTEN/NOTIFY +// notification. +type NoticeHandler func(*PgConn, *Notice) + +// NotificationHandler is a function that can handle notifications received from the PostgreSQL server. Notifications +// can be received at any time, usually during handling of a query response. The *PgConn is provided so the handler is +// aware of the origin of the notice, but it must not invoke any query method. Be aware that this is distinct from a +// notice event. +type NotificationHandler func(*PgConn, *Notification) + +// PgConn is a low-level PostgreSQL connection handle. It is not safe for concurrent usage. +type PgConn struct { + conn net.Conn + tlsConfig *tls.Config // tls.Config that conn was negotiated with; nil if conn is not TLS + pid uint32 // backend pid + secretKey []byte // key to use to send a cancel query message to the server + parameterStatuses map[string]string // parameters that have been reported by the server + txStatus byte + frontend *pgproto3.Frontend + bgReader *bgreader.BGReader + slowWriteTimer *time.Timer + bgReaderStarted chan struct{} + + customData map[string]any + + config *Config + + status byte // One of connStatus* constants + + bufferingReceive bool + bufferingReceiveMux sync.Mutex + bufferingReceiveMsg pgproto3.BackendMessage + bufferingReceiveErr error + + peekedMsg pgproto3.BackendMessage + + // Reusable / preallocated resources + resultReader ResultReader + multiResultReader MultiResultReader + pipeline Pipeline + contextWatcher *ctxwatch.ContextWatcher + fieldDescriptions [16]FieldDescription + + cleanupDone chan struct{} +} + +// Connect establishes a connection to a PostgreSQL server using the environment and connString (in URL or keyword/value +// format) to provide configuration. See documentation for [ParseConfig] for details. ctx can be used to cancel a +// connect attempt. +func Connect(ctx context.Context, connString string) (*PgConn, error) { + config, err := ParseConfig(connString) + if err != nil { + return nil, err + } + + return ConnectConfig(ctx, config) +} + +// Connect establishes a connection to a PostgreSQL server using the environment and connString (in URL or keyword/value +// format) and ParseConfigOptions to provide additional configuration. See documentation for [ParseConfig] for details. +// ctx can be used to cancel a connect attempt. +func ConnectWithOptions(ctx context.Context, connString string, parseConfigOptions ParseConfigOptions) (*PgConn, error) { + config, err := ParseConfigWithOptions(connString, parseConfigOptions) + if err != nil { + return nil, err + } + + return ConnectConfig(ctx, config) +} + +// Connect establishes a connection to a PostgreSQL server using config. config must have been constructed with +// [ParseConfig]. ctx can be used to cancel a connect attempt. +// +// If config.Fallbacks are present they will sequentially be tried in case of error establishing network connection. An +// authentication error will terminate the chain of attempts (like libpq: +// https://www.postgresql.org/docs/current/libpq-connect.html#LIBPQ-MULTIPLE-HOSTS) and be returned as the error. +func ConnectConfig(ctx context.Context, config *Config) (*PgConn, error) { + // Default values are set in ParseConfig. Enforce initial creation by ParseConfig rather than setting defaults from + // zero values. + if !config.createdByParseConfig { + panic("config must be created by ParseConfig") + } + + var allErrors []error + + connectConfigs, errs := buildConnectOneConfigs(ctx, config) + if len(errs) > 0 { + allErrors = append(allErrors, errs...) + } + + if len(connectConfigs) == 0 { + return nil, &ConnectError{Config: config, err: fmt.Errorf("hostname resolving error: %w", errors.Join(allErrors...))} + } + + pgConn, errs := connectPreferred(ctx, config, connectConfigs) + if len(errs) > 0 { + allErrors = append(allErrors, errs...) + return nil, &ConnectError{Config: config, err: errors.Join(allErrors...)} + } + + if config.AfterConnect != nil { + err := config.AfterConnect(ctx, pgConn) + if err != nil { + pgConn.conn.Close() + return nil, &ConnectError{Config: config, err: fmt.Errorf("AfterConnect error: %w", err)} + } + } + + return pgConn, nil +} + +// buildConnectOneConfigs resolves hostnames and builds a list of connectOneConfigs to try connecting to. It returns a +// slice of successfully resolved connectOneConfigs and a slice of errors. It is possible for both slices to contain +// values if some hosts were successfully resolved and others were not. +func buildConnectOneConfigs(ctx context.Context, config *Config) ([]*connectOneConfig, []error) { + // Simplify usage by treating primary config and fallbacks the same. + fallbackConfigs := []*FallbackConfig{ + { + Host: config.Host, + Port: config.Port, + TLSConfig: config.TLSConfig, + }, + } + fallbackConfigs = append(fallbackConfigs, config.Fallbacks...) + + var configs []*connectOneConfig + + var allErrors []error + + for _, fb := range fallbackConfigs { + // skip resolve for unix sockets + if isAbsolutePath(fb.Host) { + network, address := NetworkAddress(fb.Host, fb.Port) + configs = append(configs, &connectOneConfig{ + network: network, + address: address, + originalHostname: fb.Host, + tlsConfig: fb.TLSConfig, + }) + + continue + } + + ips, err := config.LookupFunc(ctx, fb.Host) + if err != nil { + allErrors = append(allErrors, err) + continue + } + + for _, ip := range ips { + splitIP, splitPort, err := net.SplitHostPort(ip) + if err == nil { + port, err := strconv.ParseUint(splitPort, 10, 16) + if err != nil { + return nil, []error{fmt.Errorf("error parsing port (%s) from lookup: %w", splitPort, err)} + } + network, address := NetworkAddress(splitIP, uint16(port)) + configs = append(configs, &connectOneConfig{ + network: network, + address: address, + originalHostname: fb.Host, + tlsConfig: fb.TLSConfig, + }) + } else { + network, address := NetworkAddress(ip, fb.Port) + configs = append(configs, &connectOneConfig{ + network: network, + address: address, + originalHostname: fb.Host, + tlsConfig: fb.TLSConfig, + }) + } + } + } + + return configs, allErrors +} + +// connectPreferred attempts to connect to the preferred host from connectOneConfigs. The connections are attempted in +// order. If a connection is successful it is returned. If no connection is successful then all errors are returned. If +// a connection attempt returns a [NotPreferredError], then that host will be used if no other hosts are successful. +func connectPreferred(ctx context.Context, config *Config, connectOneConfigs []*connectOneConfig) (*PgConn, []error) { + octx := ctx + var allErrors []error + + var fallbackConnectOneConfig *connectOneConfig + for i, c := range connectOneConfigs { + // ConnectTimeout restricts the whole connection process. + if config.ConnectTimeout != 0 { + // create new context first time or when previous host was different + if i == 0 || (connectOneConfigs[i].address != connectOneConfigs[i-1].address) { + var cancel context.CancelFunc + ctx, cancel = context.WithTimeout(octx, config.ConnectTimeout) + defer cancel() + } + } else { + ctx = octx + } + + pgConn, err := connectOne(ctx, config, c, false) + if pgConn != nil { + return pgConn, nil + } + + allErrors = append(allErrors, err) + + var pgErr *PgError + if errors.As(err, &pgErr) { + // pgx will try next host even if libpq does not in certain cases (see #2246) + // consider change for the next major version + + const ERRCODE_INVALID_PASSWORD = "28P01" + const ERRCODE_INVALID_CATALOG_NAME = "3D000" // db does not exist + const ERRCODE_INSUFFICIENT_PRIVILEGE = "42501" // missing connect privilege + + // auth failed due to invalid password, db does not exist or user has no permission + if pgErr.Code == ERRCODE_INVALID_PASSWORD || + pgErr.Code == ERRCODE_INVALID_CATALOG_NAME || + pgErr.Code == ERRCODE_INSUFFICIENT_PRIVILEGE { + return nil, allErrors + } + } + + var npErr *NotPreferredError + if errors.As(err, &npErr) { + fallbackConnectOneConfig = c + } + } + + if fallbackConnectOneConfig != nil { + fallbackCtx := octx + if config.ConnectTimeout != 0 { + var cancel context.CancelFunc + fallbackCtx, cancel = context.WithTimeout(octx, config.ConnectTimeout) + defer cancel() + } + pgConn, err := connectOne(fallbackCtx, config, fallbackConnectOneConfig, true) + if err == nil { + return pgConn, nil + } + allErrors = append(allErrors, err) + } + + return nil, allErrors +} + +// connectOne makes one connection attempt to a single host. +func connectOne(ctx context.Context, config *Config, connectConfig *connectOneConfig, + ignoreNotPreferredErr bool, +) (*PgConn, error) { + pgConn := new(PgConn) + pgConn.config = config + pgConn.cleanupDone = make(chan struct{}) + pgConn.customData = make(map[string]any) + + var err error + + newPerDialConnectError := func(msg string, err error) *perDialConnectError { + err = normalizeTimeoutError(ctx, err) + e := &perDialConnectError{address: connectConfig.address, originalHostname: connectConfig.originalHostname, err: fmt.Errorf("%s: %w", msg, err)} + return e + } + + maxProtocolVersion, err := parseProtocolVersion(config.MaxProtocolVersion) + if err != nil { + return nil, newPerDialConnectError("invalid max_protocol_version", err) + } + minProtocolVersion, err := parseProtocolVersion(config.MinProtocolVersion) + if err != nil { + return nil, newPerDialConnectError("invalid min_protocol_version", err) + } + + pgConn.conn, err = config.DialFunc(ctx, connectConfig.network, connectConfig.address) + if err != nil { + return nil, newPerDialConnectError("dial error", err) + } + + if connectConfig.tlsConfig != nil { + pgConn.contextWatcher = ctxwatch.NewContextWatcher(&DeadlineContextWatcherHandler{Conn: pgConn.conn}) + pgConn.contextWatcher.Watch(ctx) + var ( + tlsConn net.Conn + err error + ) + if config.SSLNegotiation == "direct" { + tlsConn = tls.Client(pgConn.conn, connectConfig.tlsConfig) + } else { + tlsConn, err = startTLS(pgConn.conn, connectConfig.tlsConfig) + } + pgConn.contextWatcher.Unwatch() // Always unwatch `netConn` after TLS. + if err != nil { + pgConn.conn.Close() + return nil, newPerDialConnectError("tls error", err) + } + + pgConn.conn = tlsConn + pgConn.tlsConfig = connectConfig.tlsConfig + } + + if config.AfterNetConnect != nil { + pgConn.conn, err = config.AfterNetConnect(ctx, config, pgConn.conn) + if err != nil { + pgConn.conn.Close() + return nil, newPerDialConnectError("AfterNetConnect failed", err) + } + } + + // Use a deadline-only watcher during connect. The application-supplied + // BuildContextWatcherHandler may read *PgConn fields (e.g. + // CancelRequestContextWatcherHandler reads pgConn.pid and + // pgConn.secretKey), which would race with this function's writes to + // those fields when handling BackendKeyData. The application handler is + // installed below, after the connection reaches connStatusIdle. + pgConn.contextWatcher = ctxwatch.NewContextWatcher(&DeadlineContextWatcherHandler{Conn: pgConn.conn}) + pgConn.contextWatcher.Watch(ctx) + defer pgConn.contextWatcher.Unwatch() + + pgConn.parameterStatuses = make(map[string]string) + pgConn.status = connStatusConnecting + pgConn.bgReader = bgreader.New(pgConn.conn) + pgConn.slowWriteTimer = time.AfterFunc(time.Duration(math.MaxInt64), + func() { + pgConn.bgReader.Start() + pgConn.bgReaderStarted <- struct{}{} + }, + ) + pgConn.slowWriteTimer.Stop() + pgConn.bgReaderStarted = make(chan struct{}) + pgConn.frontend = config.BuildFrontend(pgConn.bgReader, pgConn.conn) + + startupMsg := pgproto3.StartupMessage{ + ProtocolVersion: maxProtocolVersion, + Parameters: make(map[string]string), + } + + // Copy default run-time params + maps.Copy(startupMsg.Parameters, config.RuntimeParams) + + startupMsg.Parameters["user"] = config.User + if config.Database != "" { + startupMsg.Parameters["database"] = config.Database + } + + pgConn.frontend.Send(&startupMsg) + if err := pgConn.flushWithPotentialWriteReadDeadlock(); err != nil { + pgConn.conn.Close() + return nil, newPerDialConnectError("failed to write startup message", err) + } + + // Parse require_auth on each connect so that callers who mutate + // Config.RequireAuth after ParseConfig see their change take effect. + // The parser is pure and cheap; ParseConfigWithOptions validates the + // value up front so any parse error here indicates post-parse mutation. + requireAuthPolicy, err := parseRequireAuth(config.RequireAuth) + if err != nil { + pgConn.conn.Close() + return nil, newPerDialConnectError("invalid require_auth", err) + } + requireAuthFail := func(err error) (*PgConn, error) { + pgConn.conn.Close() + return nil, newPerDialConnectError("require_auth check failed", err) + } + clientFinishedAuth := false + + for { + msg, err := pgConn.receiveMessage() + if err != nil { + pgConn.conn.Close() + if err, ok := err.(*PgError); ok { + return nil, newPerDialConnectError("server error", err) + } + return nil, newPerDialConnectError("failed to receive message", err) + } + + switch msg := msg.(type) { + case *pgproto3.BackendKeyData: + pgConn.pid = msg.ProcessID + pgConn.secretKey = msg.SecretKey + + case *pgproto3.AuthenticationOk: + if requireAuthPolicy.authRequired && !clientFinishedAuth { + return requireAuthFail(requireAuthPolicy.check(authMethodNone)) + } + case *pgproto3.AuthenticationCleartextPassword: + if err := requireAuthPolicy.check(authMethodPassword); err != nil { + return requireAuthFail(err) + } + err = pgConn.txPasswordMessage(pgConn.config.Password) + if err != nil { + pgConn.conn.Close() + return nil, newPerDialConnectError("failed to write password message", err) + } + clientFinishedAuth = true + case *pgproto3.AuthenticationMD5Password: + if err := requireAuthPolicy.check(authMethodMD5); err != nil { + return requireAuthFail(err) + } + digestedPassword := "md5" + hexMD5(hexMD5(pgConn.config.Password+pgConn.config.User)+string(msg.Salt[:])) + err = pgConn.txPasswordMessage(digestedPassword) + if err != nil { + pgConn.conn.Close() + return nil, newPerDialConnectError("failed to write password message", err) + } + clientFinishedAuth = true + case *pgproto3.AuthenticationSASL: + // Check if OAUTHBEARER is supported + serverSupportsOAuthBearer := false + for _, mech := range msg.AuthMechanisms { + if mech == "OAUTHBEARER" { + serverSupportsOAuthBearer = true + break + } + } + + if serverSupportsOAuthBearer && pgConn.config.OAuthTokenProvider != nil { + if err := requireAuthPolicy.check(authMethodOAuth); err != nil { + return requireAuthFail(err) + } + err = pgConn.oauthAuth(ctx) + } else { + if err := requireAuthPolicy.check(authMethodSCRAMSHA256); err != nil { + return requireAuthFail(err) + } + err = pgConn.scramAuth(msg.AuthMechanisms) + } + if err != nil { + pgConn.conn.Close() + return nil, newPerDialConnectError("failed SASL auth", err) + } + clientFinishedAuth = true + case *pgproto3.AuthenticationGSS: + if err := requireAuthPolicy.check(authMethodGSS); err != nil { + return requireAuthFail(err) + } + err = pgConn.gssAuth() + if err != nil { + pgConn.conn.Close() + return nil, newPerDialConnectError("failed GSS auth", err) + } + clientFinishedAuth = true + case *pgproto3.ReadyForQuery: + pgConn.status = connStatusIdle + // The connect-phase deadline-only watcher is no longer needed; replace + // it with the application-supplied watcher so subsequent operations + // (including any queries run by ValidateConnect) use it. + pgConn.contextWatcher.Unwatch() + pgConn.contextWatcher = ctxwatch.NewContextWatcher(config.BuildContextWatcherHandler(pgConn)) + + if config.ValidateConnect != nil { + err := config.ValidateConnect(ctx, pgConn) + if err != nil { + if _, ok := err.(*NotPreferredError); ignoreNotPreferredErr && ok { + return pgConn, nil + } + pgConn.conn.Close() + return nil, newPerDialConnectError("ValidateConnect failed", err) + } + } + return pgConn, nil + case *pgproto3.ParameterStatus, *pgproto3.NoticeResponse: + // handled by ReceiveMessage + case *pgproto3.NegotiateProtocolVersion: + serverVersion := pgproto3.ProtocolVersion30&0xFFFF0000 | msg.NewestMinorProtocol + if serverVersion < minProtocolVersion { + pgConn.conn.Close() + return nil, newPerDialConnectError("server protocol version too low", nil) + } + case *pgproto3.ErrorResponse: + pgConn.conn.Close() + return nil, newPerDialConnectError("server error", ErrorResponseToPgError(msg)) + default: + pgConn.conn.Close() + return nil, newPerDialConnectError("received unexpected message", err) + } + } +} + +func startTLS(conn net.Conn, tlsConfig *tls.Config) (net.Conn, error) { + err := binary.Write(conn, binary.BigEndian, []int32{8, 80877103}) + if err != nil { + return nil, err + } + + response := make([]byte, 1) + if _, err = io.ReadFull(conn, response); err != nil { + return nil, err + } + + if response[0] != 'S' { + return nil, errors.New("server refused TLS connection") + } + + return tls.Client(conn, tlsConfig), nil +} + +func (pgConn *PgConn) txPasswordMessage(password string) (err error) { + pgConn.frontend.Send(&pgproto3.PasswordMessage{Password: password}) + return pgConn.flushWithPotentialWriteReadDeadlock() +} + +func hexMD5(s string) string { + hash := md5.New() + io.WriteString(hash, s) + return hex.EncodeToString(hash.Sum(nil)) +} + +func (pgConn *PgConn) signalMessage() chan struct{} { + if pgConn.bufferingReceive { + panic("BUG: signalMessage when already in progress") + } + + pgConn.bufferingReceive = true + pgConn.bufferingReceiveMux.Lock() + + ch := make(chan struct{}) + go func() { + pgConn.bufferingReceiveMsg, pgConn.bufferingReceiveErr = pgConn.frontend.Receive() + pgConn.bufferingReceiveMux.Unlock() + close(ch) + }() + + return ch +} + +// ReceiveMessage receives one wire protocol message from the PostgreSQL server. It must only be used when the +// connection is not busy. e.g. It is an error to call [PgConn.ReceiveMessage] while reading the result of a query. The messages +// are still handled by the core pgconn message handling system so receiving a NotificationResponse will still trigger +// the OnNotification callback. +// +// This is a very low level method that requires deep understanding of the PostgreSQL wire protocol to use correctly. +// See https://www.postgresql.org/docs/current/protocol.html. +func (pgConn *PgConn) ReceiveMessage(ctx context.Context) (pgproto3.BackendMessage, error) { + if err := pgConn.lock(); err != nil { + return nil, err + } + defer pgConn.unlock() + + if ctx != context.Background() { + select { + case <-ctx.Done(): + return nil, newContextAlreadyDoneError(ctx) + default: + } + pgConn.contextWatcher.Watch(ctx) + defer pgConn.contextWatcher.Unwatch() + } + + msg, err := pgConn.receiveMessage() + if err != nil { + err = &pgconnError{ + msg: "receive message failed", + err: normalizeTimeoutError(ctx, err), + safeToRetry: true, + } + } + return msg, err +} + +// peekMessage peeks at the next message without setting up context cancellation. +func (pgConn *PgConn) peekMessage() (pgproto3.BackendMessage, error) { + if pgConn.peekedMsg != nil { + return pgConn.peekedMsg, nil + } + + var msg pgproto3.BackendMessage + var err error + if pgConn.bufferingReceive { + pgConn.bufferingReceiveMux.Lock() + msg = pgConn.bufferingReceiveMsg + err = pgConn.bufferingReceiveErr + pgConn.bufferingReceiveMux.Unlock() + pgConn.bufferingReceive = false + + // If a timeout error happened in the background try the read again. + var netErr net.Error + if errors.As(err, &netErr) && netErr.Timeout() { + msg, err = pgConn.frontend.Receive() + } + } else { + msg, err = pgConn.frontend.Receive() + } + + if err != nil { + // Close on anything other than timeout error - everything else is fatal + var netErr net.Error + isNetErr := errors.As(err, &netErr) + if !(isNetErr && netErr.Timeout()) { + pgConn.asyncClose() + } + + return nil, err + } + + pgConn.peekedMsg = msg + return msg, nil +} + +// receiveMessage receives a message without setting up context cancellation +func (pgConn *PgConn) receiveMessage() (pgproto3.BackendMessage, error) { + if pgConn.status == connStatusClosed { + return nil, &connLockError{status: "conn closed"} + } + + msg, err := pgConn.peekMessage() + if err != nil { + return nil, err + } + pgConn.peekedMsg = nil + + switch msg := msg.(type) { + case *pgproto3.ReadyForQuery: + pgConn.txStatus = msg.TxStatus + case *pgproto3.ParameterStatus: + pgConn.parameterStatuses[msg.Name] = msg.Value + case *pgproto3.ErrorResponse: + err := ErrorResponseToPgError(msg) + if pgConn.config.OnPgError != nil && !pgConn.config.OnPgError(pgConn, err) { + pgConn.status = connStatusClosed + pgConn.conn.Close() // Ignore error as the connection is already broken and there is already an error to return. + close(pgConn.cleanupDone) + return nil, err + } + case *pgproto3.NoticeResponse: + if pgConn.config.OnNotice != nil { + pgConn.config.OnNotice(pgConn, noticeResponseToNotice(msg)) + } + case *pgproto3.NotificationResponse: + if pgConn.config.OnNotification != nil { + pgConn.config.OnNotification(pgConn, &Notification{PID: msg.PID, Channel: msg.Channel, Payload: msg.Payload}) + } + } + + return msg, nil +} + +// Conn returns the underlying net.Conn. This rarely necessary. If the connection will be directly used for reading or +// writing then SyncConn should usually be called before Conn. +func (pgConn *PgConn) Conn() net.Conn { + return pgConn.conn +} + +// PID returns the backend PID. +func (pgConn *PgConn) PID() uint32 { + return pgConn.pid +} + +// TxStatus returns the current TxStatus as reported by the server in the ReadyForQuery message. +// +// Possible return values: +// +// 'I' - idle / not in transaction +// 'T' - in a transaction +// 'E' - in a failed transaction +// +// See https://www.postgresql.org/docs/current/protocol-message-formats.html. +func (pgConn *PgConn) TxStatus() byte { + return pgConn.txStatus +} + +// SecretKey returns the backend secret key used to send a cancel query message to the server. +func (pgConn *PgConn) SecretKey() []byte { + return pgConn.secretKey +} + +// Frontend returns the underlying *pgproto3.Frontend. This rarely necessary. +func (pgConn *PgConn) Frontend() *pgproto3.Frontend { + return pgConn.frontend +} + +// Close closes a connection. It is safe to call Close on an already closed connection. Close attempts a clean close by +// sending the exit message to PostgreSQL. However, this could block so ctx is available to limit the time to wait. The +// underlying net.Conn.Close() will always be called regardless of any other errors. +func (pgConn *PgConn) Close(ctx context.Context) error { + if pgConn.status == connStatusClosed { + return nil + } + pgConn.status = connStatusClosed + + defer close(pgConn.cleanupDone) + defer pgConn.conn.Close() + + if ctx != context.Background() { + // Close may be called while a cancellable query is in progress. This will most often be triggered by panic when + // a defer closes the connection (possibly indirectly via a transaction or a connection pool). Unwatch to end any + // previous watch. It is safe to Unwatch regardless of whether a watch is already is progress. + // + // See https://github.com/jackc/pgconn/issues/29 + pgConn.contextWatcher.Unwatch() + + pgConn.contextWatcher.Watch(ctx) + defer pgConn.contextWatcher.Unwatch() + } + + // Ignore any errors sending Terminate message and waiting for server to close connection. + // This mimics the behavior of libpq PQfinish. It calls closePGconn which calls sendTerminateConn which purposefully + // ignores errors. + // + // See https://github.com/jackc/pgx/issues/637 + pgConn.frontend.Send(&pgproto3.Terminate{}) + pgConn.flushWithPotentialWriteReadDeadlock() + + return pgConn.conn.Close() +} + +// asyncClose marks the connection as closed and asynchronously sends a cancel query message and closes the underlying +// connection. +func (pgConn *PgConn) asyncClose() { + if pgConn.status == connStatusClosed { + return + } + pgConn.status = connStatusClosed + + go func() { + defer close(pgConn.cleanupDone) + defer pgConn.conn.Close() + + deadline := time.Now().Add(time.Second * 15) + + ctx, cancel := context.WithDeadline(context.Background(), deadline) + defer cancel() + + pgConn.CancelRequest(ctx) + + pgConn.conn.SetDeadline(deadline) + + pgConn.frontend.Send(&pgproto3.Terminate{}) + pgConn.flushWithPotentialWriteReadDeadlock() + }() +} + +// CleanupDone returns a channel that will be closed after all underlying resources have been cleaned up. A closed +// connection is no longer usable, but underlying resources, in particular the net.Conn, may not have finished closing +// yet. This is because certain errors such as a context cancellation require that the interrupted function call return +// immediately, but the error may also cause the connection to be closed. In these cases the underlying resources are +// closed asynchronously. +// +// This is only likely to be useful to connection pools. It gives them a way avoid establishing a new connection while +// an old connection is still being cleaned up and thereby exceeding the maximum pool size. +func (pgConn *PgConn) CleanupDone() chan (struct{}) { + return pgConn.cleanupDone +} + +// IsClosed reports if the connection has been closed. +// +// CleanupDone() can be used to determine if all cleanup has been completed. +func (pgConn *PgConn) IsClosed() bool { + return pgConn.status < connStatusIdle +} + +// IsBusy reports if the connection is busy. +func (pgConn *PgConn) IsBusy() bool { + return pgConn.status == connStatusBusy +} + +// lock locks the connection. +func (pgConn *PgConn) lock() error { + switch pgConn.status { + case connStatusBusy: + return &connLockError{status: "conn busy"} // This only should be possible in case of an application bug. + case connStatusClosed: + return &connLockError{status: "conn closed"} + case connStatusUninitialized: + return &connLockError{status: "conn uninitialized"} + } + pgConn.status = connStatusBusy + return nil +} + +func (pgConn *PgConn) unlock() { + switch pgConn.status { + case connStatusBusy: + pgConn.status = connStatusIdle + case connStatusClosed: + default: + panic("BUG: cannot unlock unlocked connection") // This should only be possible if there is a bug in this package. + } +} + +// ParameterStatus returns the value of a parameter reported by the server (e.g. +// server_version). Returns an empty string for unknown parameters. +func (pgConn *PgConn) ParameterStatus(key string) string { + return pgConn.parameterStatuses[key] +} + +// CommandTag is the status text returned by PostgreSQL for a query. +type CommandTag struct { + s string +} + +// NewCommandTag makes a CommandTag from s. +func NewCommandTag(s string) CommandTag { + return CommandTag{s: s} +} + +// RowsAffected returns the number of rows affected. If the CommandTag was not +// for a row affecting command (e.g. "CREATE TABLE") then it returns 0. +func (ct CommandTag) RowsAffected() int64 { + // Parse the number from the end in a single pass. + var n int64 + var mult int64 = 1 + + for i := len(ct.s) - 1; i >= 0; i-- { + c := ct.s[i] + if c >= '0' && c <= '9' { + n += int64(c-'0') * mult + mult *= 10 + } else { + break + } + } + + return n +} + +func (ct CommandTag) String() string { + return ct.s +} + +// Insert is true if the command tag starts with "INSERT". +func (ct CommandTag) Insert() bool { + return strings.HasPrefix(ct.s, "INSERT") +} + +// Update is true if the command tag starts with "UPDATE". +func (ct CommandTag) Update() bool { + return strings.HasPrefix(ct.s, "UPDATE") +} + +// Delete is true if the command tag starts with "DELETE". +func (ct CommandTag) Delete() bool { + return strings.HasPrefix(ct.s, "DELETE") +} + +// Select is true if the command tag starts with "SELECT". +func (ct CommandTag) Select() bool { + return strings.HasPrefix(ct.s, "SELECT") +} + +type FieldDescription struct { + Name string + TableOID uint32 + TableAttributeNumber uint16 + DataTypeOID uint32 + DataTypeSize int16 + TypeModifier int32 + Format int16 +} + +func (pgConn *PgConn) getFieldDescriptionSlice(n int) []FieldDescription { + if cap(pgConn.fieldDescriptions) >= n { + return pgConn.fieldDescriptions[:n:n] + } else { + return make([]FieldDescription, n) + } +} + +func convertRowDescription(dst []FieldDescription, rd *pgproto3.RowDescription) { + for i := range rd.Fields { + dst[i].Name = string(rd.Fields[i].Name) + dst[i].TableOID = rd.Fields[i].TableOID + dst[i].TableAttributeNumber = rd.Fields[i].TableAttributeNumber + dst[i].DataTypeOID = rd.Fields[i].DataTypeOID + dst[i].DataTypeSize = rd.Fields[i].DataTypeSize + dst[i].TypeModifier = rd.Fields[i].TypeModifier + dst[i].Format = rd.Fields[i].Format + } +} + +type StatementDescription struct { + Name string + SQL string + ParamOIDs []uint32 + Fields []FieldDescription +} + +// Prepare creates a prepared statement. If the name is empty, the anonymous prepared statement will be used. This +// allows Prepare to also to describe statements without creating a server-side prepared statement. +// +// Prepare does not send a PREPARE statement to the server. It uses the PostgreSQL Parse and Describe protocol messages +// directly. +// +// In extremely rare cases, Prepare may fail after the Parse is successful, but before the Describe is complete. In this +// case, the returned error will be an error where errors.As with a *PrepareError succeeds and the *PrepareError has +// ParseComplete set to true. +func (pgConn *PgConn) Prepare(ctx context.Context, name, sql string, paramOIDs []uint32) (*StatementDescription, error) { + if err := pgConn.lock(); err != nil { + return nil, err + } + defer pgConn.unlock() + + if ctx != context.Background() { + select { + case <-ctx.Done(): + return nil, newContextAlreadyDoneError(ctx) + default: + } + pgConn.contextWatcher.Watch(ctx) + defer pgConn.contextWatcher.Unwatch() + } + + pgConn.frontend.SendParse(&pgproto3.Parse{Name: name, Query: sql, ParameterOIDs: paramOIDs}) + pgConn.frontend.SendDescribe(&pgproto3.Describe{ObjectType: 'S', Name: name}) + pgConn.frontend.SendSync(&pgproto3.Sync{}) + err := pgConn.flushWithPotentialWriteReadDeadlock() + if err != nil { + pgConn.asyncClose() + return nil, err + } + + psd := &StatementDescription{Name: name, SQL: sql} + + var ParseComplete bool + var pgErr *PgError + +readloop: + for { + msg, err := pgConn.receiveMessage() + if err != nil { + pgConn.asyncClose() + return nil, normalizeTimeoutError(ctx, err) + } + + switch msg := msg.(type) { + case *pgproto3.ParseComplete: + ParseComplete = true + case *pgproto3.ParameterDescription: + psd.ParamOIDs = make([]uint32, len(msg.ParameterOIDs)) + copy(psd.ParamOIDs, msg.ParameterOIDs) + case *pgproto3.RowDescription: + psd.Fields = make([]FieldDescription, len(msg.Fields)) + convertRowDescription(psd.Fields, msg) + case *pgproto3.ErrorResponse: + pgErr = ErrorResponseToPgError(msg) + case *pgproto3.ReadyForQuery: + break readloop + } + } + + if pgErr != nil { + return nil, &PrepareError{err: pgErr, ParseComplete: ParseComplete} + } + return psd, nil +} + +// Deallocate deallocates a prepared statement. +// +// Deallocate does not send a DEALLOCATE statement to the server. It uses the PostgreSQL Close protocol message +// directly. This has slightly different behavior than executing DEALLOCATE statement. +// - Deallocate can succeed in an aborted transaction. +// - Deallocating a non-existent prepared statement is not an error. +func (pgConn *PgConn) Deallocate(ctx context.Context, name string) error { + if err := pgConn.lock(); err != nil { + return err + } + defer pgConn.unlock() + + if ctx != context.Background() { + select { + case <-ctx.Done(): + return newContextAlreadyDoneError(ctx) + default: + } + pgConn.contextWatcher.Watch(ctx) + defer pgConn.contextWatcher.Unwatch() + } + + pgConn.frontend.SendClose(&pgproto3.Close{ObjectType: 'S', Name: name}) + pgConn.frontend.SendSync(&pgproto3.Sync{}) + err := pgConn.flushWithPotentialWriteReadDeadlock() + if err != nil { + pgConn.asyncClose() + return err + } + + for { + msg, err := pgConn.receiveMessage() + if err != nil { + pgConn.asyncClose() + return normalizeTimeoutError(ctx, err) + } + + switch msg := msg.(type) { + case *pgproto3.ErrorResponse: + return ErrorResponseToPgError(msg) + case *pgproto3.ReadyForQuery: + return nil + } + } +} + +// ErrorResponseToPgError converts a wire protocol error message to a *PgError. +func ErrorResponseToPgError(msg *pgproto3.ErrorResponse) *PgError { + return &PgError{ + Severity: msg.Severity, + SeverityUnlocalized: msg.SeverityUnlocalized, + Code: msg.Code, + Message: msg.Message, + Detail: msg.Detail, + Hint: msg.Hint, + Position: msg.Position, + InternalPosition: msg.InternalPosition, + InternalQuery: msg.InternalQuery, + Where: msg.Where, + SchemaName: msg.SchemaName, + TableName: msg.TableName, + ColumnName: msg.ColumnName, + DataTypeName: msg.DataTypeName, + ConstraintName: msg.ConstraintName, + File: msg.File, + Line: msg.Line, + Routine: msg.Routine, + } +} + +func noticeResponseToNotice(msg *pgproto3.NoticeResponse) *Notice { + pgerr := ErrorResponseToPgError((*pgproto3.ErrorResponse)(msg)) + return (*Notice)(pgerr) +} + +// CancelRequest sends a cancel request to the PostgreSQL server. It returns an error if unable to deliver the cancel +// request, but lack of an error does not ensure that the query was canceled. As specified in the documentation, there +// is no way to be sure a query was canceled. +// See https://www.postgresql.org/docs/current/protocol-flow.html#PROTOCOL-FLOW-CANCELING-REQUESTS +func (pgConn *PgConn) CancelRequest(ctx context.Context) error { + // Open a cancellation request to the same server. The address is taken from the net.Conn directly instead of reusing + // the connection config. This is important in high availability configurations where fallback connections may be + // specified or DNS may be used to load balance. + serverAddr := pgConn.conn.RemoteAddr() + var serverNetwork string + var serverAddress string + if serverAddr.Network() == "unix" { + // for unix sockets, RemoteAddr() calls getpeername() which returns the name the + // server passed to bind(). For Postgres, this is always a relative path "./.s.PGSQL.5432" + // so connecting to it will fail. Fall back to the config's value + serverNetwork, serverAddress = NetworkAddress(pgConn.config.Host, pgConn.config.Port) + } else { + serverNetwork, serverAddress = serverAddr.Network(), serverAddr.String() + } + cancelConn, err := pgConn.config.DialFunc(ctx, serverNetwork, serverAddress) + if err != nil { + // In case of unix sockets, RemoteAddr() returns only the file part of the path. If the + // first connect failed, try the config. + if serverAddr.Network() != "unix" { + return err + } + serverNetwork, serverAddr := NetworkAddress(pgConn.config.Host, pgConn.config.Port) + cancelConn, err = pgConn.config.DialFunc(ctx, serverNetwork, serverAddr) + if err != nil { + return err + } + } + defer cancelConn.Close() + + if ctx != context.Background() { + contextWatcher := ctxwatch.NewContextWatcher(&DeadlineContextWatcherHandler{Conn: cancelConn}) + contextWatcher.Watch(ctx) + defer contextWatcher.Unwatch() + } + + // If the primary connection is encrypted, encrypt the cancel connection the same way so the + // backend pid and secret key are not exposed to a passive network observer. This mirrors libpq's + // PQcancelCreate (PG17+), which reuses the original connection's sslmode/gssencmode for the + // cancel connection. The legacy unencrypted path is still used when the primary connection is + // plaintext (e.g. unix sockets or sslmode=disable). + if pgConn.tlsConfig != nil { + var tlsCancelConn net.Conn + if pgConn.config.SSLNegotiation == "direct" { + tlsCancelConn = tls.Client(cancelConn, pgConn.tlsConfig) + } else { + tlsCancelConn, err = startTLS(cancelConn, pgConn.tlsConfig) + if err != nil { + return fmt.Errorf("tls error on cancel connection: %w", err) + } + } + cancelConn = tlsCancelConn + defer cancelConn.Close() + } + + buf := make([]byte, 12+len(pgConn.secretKey)) + binary.BigEndian.PutUint32(buf[0:4], uint32(len(buf))) + binary.BigEndian.PutUint32(buf[4:8], 80877102) + binary.BigEndian.PutUint32(buf[8:12], pgConn.pid) + copy(buf[12:], pgConn.secretKey) + + if _, err := cancelConn.Write(buf); err != nil { + return fmt.Errorf("write to connection for cancellation: %w", err) + } + + // Wait for the cancel request to be acknowledged by the server. + // It copies the behavior of the libpq: https://github.com/postgres/postgres/blob/REL_16_0/src/interfaces/libpq/fe-connect.c#L4946-L4960 + _, _ = cancelConn.Read(buf) + + return nil +} + +// WaitForNotification waits for a LISTEN/NOTIFY message to be received. It returns an error if a notification was not +// received. +func (pgConn *PgConn) WaitForNotification(ctx context.Context) error { + if err := pgConn.lock(); err != nil { + return err + } + defer pgConn.unlock() + + if ctx != context.Background() { + select { + case <-ctx.Done(): + return newContextAlreadyDoneError(ctx) + default: + } + + pgConn.contextWatcher.Watch(ctx) + defer pgConn.contextWatcher.Unwatch() + } + + for { + msg, err := pgConn.receiveMessage() + if err != nil { + return normalizeTimeoutError(ctx, err) + } + + if _, ok := msg.(*pgproto3.NotificationResponse); ok { + return nil + } + } +} + +// Exec executes SQL via the PostgreSQL simple query protocol. SQL may contain multiple queries. Execution is +// implicitly wrapped in a transaction unless a transaction is already in progress or SQL contains transaction control +// statements. +// +// Prefer [PgConn.ExecParams] unless executing arbitrary SQL that may contain multiple queries. +func (pgConn *PgConn) Exec(ctx context.Context, sql string) *MultiResultReader { + if err := pgConn.lock(); err != nil { + return &MultiResultReader{ + closed: true, + err: err, + } + } + + pgConn.multiResultReader = MultiResultReader{ + pgConn: pgConn, + ctx: ctx, + } + multiResult := &pgConn.multiResultReader + if ctx != context.Background() { + select { + case <-ctx.Done(): + multiResult.closed = true + multiResult.err = newContextAlreadyDoneError(ctx) + pgConn.unlock() + return multiResult + default: + } + pgConn.contextWatcher.Watch(ctx) + } + + pgConn.frontend.SendQuery(&pgproto3.Query{String: sql}) + err := pgConn.flushWithPotentialWriteReadDeadlock() + if err != nil { + pgConn.asyncClose() + pgConn.contextWatcher.Unwatch() + multiResult.closed = true + multiResult.err = err + pgConn.unlock() + return multiResult + } + + return multiResult +} + +// ExecParams executes a command via the PostgreSQL extended query protocol. +// +// sql is a SQL command string. It may only contain one query. Parameter substitution is positional using $1, $2, $3, +// etc. +// +// paramValues are the parameter values. It must be encoded in the format given by paramFormats. +// +// paramOIDs is a slice of data type OIDs for paramValues. If paramOIDs is nil, the server will infer the data type for +// all parameters. Any paramOID element that is 0 that will cause the server to infer the data type for that parameter. +// ExecParams will panic if len(paramOIDs) is not 0, 1, or len(paramValues). +// +// paramFormats is a slice of format codes determining for each paramValue column whether it is encoded in text or +// binary format. If paramFormats is nil all params are text format. ExecParams will panic if +// len(paramFormats) is not 0, 1, or len(paramValues). +// +// resultFormats is a slice of format codes determining for each result column whether it is encoded in text or +// binary format. If resultFormats is nil all results will be in text format. +// +// [ResultReader] must be closed before [PgConn] can be used again. +func (pgConn *PgConn) ExecParams(ctx context.Context, sql string, paramValues [][]byte, paramOIDs []uint32, paramFormats, resultFormats []int16) *ResultReader { + result := pgConn.execExtendedPrefix(ctx, paramValues) + if result.closed { + return result + } + + pgConn.frontend.SendParse(&pgproto3.Parse{Query: sql, ParameterOIDs: paramOIDs}) + pgConn.frontend.SendBind(&pgproto3.Bind{ParameterFormatCodes: paramFormats, Parameters: paramValues, ResultFormatCodes: resultFormats}) + + pgConn.execExtendedSuffix(result, nil, nil) + + return result +} + +// ExecPrepared enqueues the execution of a prepared statement via the PostgreSQL extended query protocol. +// +// paramValues are the parameter values. It must be encoded in the format given by paramFormats. +// +// paramFormats is a slice of format codes determining for each paramValue column whether it is encoded in text or +// binary format. If paramFormats is nil all params are text format. ExecPrepared will panic if +// len(paramFormats) is not 0, 1, or len(paramValues). +// +// resultFormats is a slice of format codes determining for each result column whether it is encoded in text or +// binary format. If resultFormats is nil all results will be in text format. +// +// [ResultReader] must be closed before [PgConn] can be used again. +func (pgConn *PgConn) ExecPrepared(ctx context.Context, stmtName string, paramValues [][]byte, paramFormats, resultFormats []int16) *ResultReader { + result := pgConn.execExtendedPrefix(ctx, paramValues) + if result.closed { + return result + } + + pgConn.frontend.SendBind(&pgproto3.Bind{PreparedStatement: stmtName, ParameterFormatCodes: paramFormats, Parameters: paramValues, ResultFormatCodes: resultFormats}) + + pgConn.execExtendedSuffix(result, nil, nil) + + return result +} + +// ExecStatement enqueues the execution of a prepared statement via the PostgreSQL extended query protocol. +// +// This differs from [PgConn.ExecPrepared] in that it takes a [*StatementDescription] instead of the prepared statement name. +// Because it has the [*StatementDescription] it can avoid the Describe Portal message that [PgConn.ExecPrepared] must send to get +// the result column descriptions. +// +// paramValues are the parameter values. It must be encoded in the format given by paramFormats. +// +// paramFormats is a slice of format codes determining for each paramValue column whether it is encoded in text or +// binary format. If paramFormats is nil all params are text format. ExecStatement will panic if len(paramFormats) is not +// 0, 1, or len(paramValues). +// +// resultFormats is a slice of format codes determining for each result column whether it is encoded in text or binary +// format. If resultFormats is nil all results will be in text format. +// +// [ResultReader] must be closed before [PgConn] can be used again. +func (pgConn *PgConn) ExecStatement(ctx context.Context, statementDescription *StatementDescription, paramValues [][]byte, paramFormats, resultFormats []int16) *ResultReader { + result := pgConn.execExtendedPrefix(ctx, paramValues) + if result.closed { + return result + } + + pgConn.frontend.SendBind(&pgproto3.Bind{PreparedStatement: statementDescription.Name, ParameterFormatCodes: paramFormats, Parameters: paramValues, ResultFormatCodes: resultFormats}) + + pgConn.execExtendedSuffix(result, statementDescription, resultFormats) + + return result +} + +func (pgConn *PgConn) execExtendedPrefix(ctx context.Context, paramValues [][]byte) *ResultReader { + pgConn.resultReader = ResultReader{ + pgConn: pgConn, + ctx: ctx, + } + result := &pgConn.resultReader + + if err := pgConn.lock(); err != nil { + result.concludeCommand(CommandTag{}, err) + result.closed = true + return result + } + + if len(paramValues) > math.MaxUint16 { + result.concludeCommand(CommandTag{}, fmt.Errorf("extended protocol limited to %v parameters", math.MaxUint16)) + result.closed = true + pgConn.unlock() + return result + } + + if ctx != context.Background() { + select { + case <-ctx.Done(): + result.concludeCommand(CommandTag{}, newContextAlreadyDoneError(ctx)) + result.closed = true + pgConn.unlock() + return result + default: + } + pgConn.contextWatcher.Watch(ctx) + } + + return result +} + +func (pgConn *PgConn) execExtendedSuffix(result *ResultReader, statementDescription *StatementDescription, resultFormats []int16) { + if statementDescription == nil { + pgConn.frontend.SendDescribe(&pgproto3.Describe{ObjectType: 'P'}) + } + pgConn.frontend.SendExecute(&pgproto3.Execute{}) + pgConn.frontend.SendSync(&pgproto3.Sync{}) + + err := pgConn.flushWithPotentialWriteReadDeadlock() + if err != nil { + pgConn.asyncClose() + result.concludeCommand(CommandTag{}, err) + pgConn.contextWatcher.Unwatch() + result.closed = true + pgConn.unlock() + return + } + + result.readUntilRowDescription(statementDescription, resultFormats) +} + +// CopyTo executes the copy command sql and copies the results to w. +func (pgConn *PgConn) CopyTo(ctx context.Context, w io.Writer, sql string) (CommandTag, error) { + if err := pgConn.lock(); err != nil { + return CommandTag{}, err + } + + if ctx != context.Background() { + select { + case <-ctx.Done(): + pgConn.unlock() + return CommandTag{}, newContextAlreadyDoneError(ctx) + default: + } + pgConn.contextWatcher.Watch(ctx) + defer pgConn.contextWatcher.Unwatch() + } + + // Send copy to command + pgConn.frontend.SendQuery(&pgproto3.Query{String: sql}) + + err := pgConn.flushWithPotentialWriteReadDeadlock() + if err != nil { + pgConn.asyncClose() + pgConn.unlock() + return CommandTag{}, err + } + + // Read results + var commandTag CommandTag + var pgErr error + for { + msg, err := pgConn.receiveMessage() + if err != nil { + pgConn.asyncClose() + return CommandTag{}, normalizeTimeoutError(ctx, err) + } + + switch msg := msg.(type) { + case *pgproto3.CopyDone: + case *pgproto3.CopyData: + _, err := w.Write(msg.Data) + if err != nil { + pgConn.asyncClose() + return CommandTag{}, err + } + case *pgproto3.ReadyForQuery: + pgConn.unlock() + return commandTag, pgErr + case *pgproto3.CommandComplete: + commandTag = pgConn.makeCommandTag(msg.CommandTag) + case *pgproto3.ErrorResponse: + pgErr = ErrorResponseToPgError(msg) + } + } +} + +// CopyFrom executes the copy command sql and copies all of r to the PostgreSQL server. +// +// Note: context cancellation will only interrupt operations on the underlying PostgreSQL network connection. Reads on r +// could still block. +func (pgConn *PgConn) CopyFrom(ctx context.Context, r io.Reader, sql string) (CommandTag, error) { + if err := pgConn.lock(); err != nil { + return CommandTag{}, err + } + defer pgConn.unlock() + + if ctx != context.Background() { + select { + case <-ctx.Done(): + return CommandTag{}, newContextAlreadyDoneError(ctx) + default: + } + pgConn.contextWatcher.Watch(ctx) + defer pgConn.contextWatcher.Unwatch() + } + + // Send copy from query + pgConn.frontend.SendQuery(&pgproto3.Query{String: sql}) + err := pgConn.flushWithPotentialWriteReadDeadlock() + if err != nil { + pgConn.asyncClose() + return CommandTag{}, err + } + + // Send copy data + abortCopyChan := make(chan struct{}) + copyErrChan := make(chan error, 1) + signalMessageChan := pgConn.signalMessage() + var wg sync.WaitGroup + wg.Go(func() { + buf := iobufpool.Get(65536) + defer iobufpool.Put(buf) + (*buf)[0] = 'd' + + for { + n, readErr := r.Read((*buf)[5:cap(*buf)]) + if n > 0 { + *buf = (*buf)[0 : n+5] + pgio.SetInt32((*buf)[1:], int32(n+4)) + + writeErr := pgConn.frontend.SendUnbufferedEncodedCopyData(*buf) + if writeErr != nil { + // Write errors are always fatal, but we can't use asyncClose because we are in a different goroutine. Not + // setting pgConn.status or closing pgConn.cleanupDone for the same reason. + pgConn.conn.Close() + + copyErrChan <- writeErr + return + } + } + if readErr != nil { + copyErrChan <- readErr + return + } + + select { + case <-abortCopyChan: + return + default: + } + } + }) + + var pgErr error + var copyErr error + for copyErr == nil && pgErr == nil { + select { + case copyErr = <-copyErrChan: + case <-signalMessageChan: + // If pgConn.receiveMessage encounters an error it will call pgConn.asyncClose. But that is a race condition with + // the goroutine. So instead check pgConn.bufferingReceiveErr which will have been set by the signalMessage. If an + // error is found then forcibly close the connection without sending the Terminate message. + if err := pgConn.bufferingReceiveErr; err != nil { + pgConn.status = connStatusClosed + pgConn.conn.Close() + close(pgConn.cleanupDone) + return CommandTag{}, normalizeTimeoutError(ctx, err) + } + // peekMessage never returns err in the bufferingReceive mode - it only forwards the bufferingReceive variables. + // Therefore, the only case for receiveMessage to return err is during handling of the ErrorResponse message type + // and using pgOnError handler to determine the connection is no longer valid (and thus closing the conn). + msg, serverError := pgConn.receiveMessage() + if serverError != nil { + close(abortCopyChan) + return CommandTag{}, serverError + } + + switch msg := msg.(type) { + case *pgproto3.ErrorResponse: + pgErr = ErrorResponseToPgError(msg) + default: + signalMessageChan = pgConn.signalMessage() + } + } + } + close(abortCopyChan) + // Make sure io goroutine finishes before writing. + wg.Wait() + + if copyErr == io.EOF || pgErr != nil { + pgConn.frontend.Send(&pgproto3.CopyDone{}) + } else { + pgConn.frontend.Send(&pgproto3.CopyFail{Message: copyErr.Error()}) + } + err = pgConn.flushWithPotentialWriteReadDeadlock() + if err != nil { + pgConn.asyncClose() + return CommandTag{}, err + } + + // Read results + var commandTag CommandTag + for { + msg, err := pgConn.receiveMessage() + if err != nil { + pgConn.asyncClose() + return CommandTag{}, normalizeTimeoutError(ctx, err) + } + + switch msg := msg.(type) { + case *pgproto3.ReadyForQuery: + return commandTag, pgErr + case *pgproto3.CommandComplete: + commandTag = pgConn.makeCommandTag(msg.CommandTag) + case *pgproto3.ErrorResponse: + pgErr = ErrorResponseToPgError(msg) + } + } +} + +// MultiResultReader is a reader for a command that could return multiple results such as Exec or ExecBatch. +type MultiResultReader struct { + pgConn *PgConn + ctx context.Context + + rr *ResultReader + + // Data from when the batch was queued. + statementDescriptions []*StatementDescription + resultFormats [][]int16 + + closed bool + err error +} + +// ReadAll reads all available results. Calling ReadAll is mutually exclusive with all other MultiResultReader methods. +func (mrr *MultiResultReader) ReadAll() ([]*Result, error) { + var results []*Result + + for mrr.NextResult() { + results = append(results, mrr.ResultReader().Read()) + } + err := mrr.Close() + + return results, err +} + +func (mrr *MultiResultReader) receiveMessage() (pgproto3.BackendMessage, error) { + msg, err := mrr.pgConn.receiveMessage() + if err != nil { + mrr.pgConn.contextWatcher.Unwatch() + mrr.err = normalizeTimeoutError(mrr.ctx, err) + mrr.closed = true + mrr.pgConn.asyncClose() + return nil, mrr.err + } + + switch msg := msg.(type) { + case *pgproto3.ReadyForQuery: + mrr.closed = true + mrr.pgConn.contextWatcher.Unwatch() + mrr.pgConn.unlock() + case *pgproto3.ErrorResponse: + mrr.err = ErrorResponseToPgError(msg) + } + + return msg, nil +} + +// NextResult returns advances the MultiResultReader to the next result and returns true if a result is available. +func (mrr *MultiResultReader) NextResult() bool { + for !mrr.closed && mrr.err == nil { + msg, _ := mrr.pgConn.peekMessage() + if _, ok := msg.(*pgproto3.DataRow); ok { + if len(mrr.statementDescriptions) > 0 { + rr := ResultReader{ + pgConn: mrr.pgConn, + multiResultReader: mrr, + ctx: mrr.ctx, + } + + // This result corresponds to a prepared statement description that was provided when queuing the batch. + sd := mrr.statementDescriptions[0] + mrr.statementDescriptions = mrr.statementDescriptions[1:] + + resultFormats := mrr.resultFormats[0] + mrr.resultFormats = mrr.resultFormats[1:] + + sdFields := sd.Fields + rr.fieldDescriptions = rr.pgConn.getFieldDescriptionSlice(len(sdFields)) + + err := combineFieldDescriptionsAndResultFormats(rr.fieldDescriptions, sdFields, resultFormats) + if err != nil { + rr.concludeCommand(CommandTag{}, err) + } + + mrr.pgConn.resultReader = rr + mrr.rr = &mrr.pgConn.resultReader + return true + } + + mrr.err = fmt.Errorf("unexpected DataRow message without preceding RowDescription") + return false + } + + msg, err := mrr.receiveMessage() + if err != nil { + return false + } + + switch msg := msg.(type) { + case *pgproto3.RowDescription: + mrr.pgConn.resultReader = ResultReader{ + pgConn: mrr.pgConn, + multiResultReader: mrr, + ctx: mrr.ctx, + fieldDescriptions: mrr.pgConn.getFieldDescriptionSlice(len(msg.Fields)), + } + convertRowDescription(mrr.pgConn.resultReader.fieldDescriptions, msg) + + mrr.rr = &mrr.pgConn.resultReader + return true + case *pgproto3.CommandComplete: + mrr.pgConn.resultReader = ResultReader{ + commandTag: mrr.pgConn.makeCommandTag(msg.CommandTag), + commandConcluded: true, + closed: true, + } + mrr.rr = &mrr.pgConn.resultReader + return true + case *pgproto3.EmptyQueryResponse: + mrr.pgConn.resultReader = ResultReader{ + commandConcluded: true, + closed: true, + } + mrr.rr = &mrr.pgConn.resultReader + return true + } + } + + return false +} + +// ResultReader returns the current ResultReader. +func (mrr *MultiResultReader) ResultReader() *ResultReader { + return mrr.rr +} + +// Close closes the MultiResultReader and returns the first error that occurred during the MultiResultReader's use. +func (mrr *MultiResultReader) Close() error { + for !mrr.closed { + _, err := mrr.receiveMessage() + if err != nil { + return mrr.err + } + } + + return mrr.err +} + +// ResultReader is a reader for the result of a single query. +type ResultReader struct { + pgConn *PgConn + multiResultReader *MultiResultReader + pipeline *Pipeline + ctx context.Context + + fieldDescriptions []FieldDescription + rowValues [][]byte + commandTag CommandTag + preloaded bool + commandConcluded bool + closed bool + err error +} + +// Result is the saved query response that is returned by calling Read on a ResultReader. +type Result struct { + FieldDescriptions []FieldDescription + Rows [][][]byte + CommandTag CommandTag + Err error +} + +// Read saves the query response to a Result. +func (rr *ResultReader) Read() *Result { + br := &Result{} + + for rr.NextRow() { + if br.FieldDescriptions == nil { + br.FieldDescriptions = make([]FieldDescription, len(rr.FieldDescriptions())) + copy(br.FieldDescriptions, rr.FieldDescriptions()) + } + + values := rr.Values() + row := make([][]byte, len(values)) + for i := range row { + if values[i] != nil { + row[i] = make([]byte, len(values[i])) + copy(row[i], values[i]) + } + } + br.Rows = append(br.Rows, row) + } + + br.CommandTag, br.Err = rr.Close() + + return br +} + +// NextRow advances the ResultReader to the next row and returns true if a row is available. +func (rr *ResultReader) NextRow() bool { + if rr.preloaded { + rr.preloaded = false + return true + } + + for !rr.commandConcluded { + msg, err := rr.receiveMessage() + if err != nil { + return false + } + + if msg, ok := msg.(*pgproto3.DataRow); ok { + rr.rowValues = msg.Values + return true + } + } + + return false +} + +func (rr *ResultReader) preloadRowValues(values [][]byte) { + rr.rowValues = values + rr.preloaded = true +} + +// FieldDescriptions returns the field descriptions for the current result set. The returned slice is only valid until +// the ResultReader is closed. It may return nil (for example, if the query did not return a result set or an error was +// encountered.) +func (rr *ResultReader) FieldDescriptions() []FieldDescription { + return rr.fieldDescriptions +} + +// Values returns the current row data. NextRow must have been previously been called. The returned [][]byte is only +// valid until the next NextRow call or the ResultReader is closed. +func (rr *ResultReader) Values() [][]byte { + return rr.rowValues +} + +// Close consumes any remaining result data and returns the command tag or +// error. +func (rr *ResultReader) Close() (CommandTag, error) { + if rr.closed { + return rr.commandTag, rr.err + } + rr.closed = true + + for !rr.commandConcluded { + _, err := rr.receiveMessage() + if err != nil { + return CommandTag{}, rr.err + } + } + + if rr.multiResultReader == nil && rr.pipeline == nil { + for { + msg, err := rr.receiveMessage() + if err != nil { + return CommandTag{}, rr.err + } + + switch msg := msg.(type) { + // Detect a deferred constraint violation where the ErrorResponse is sent after CommandComplete. + case *pgproto3.ErrorResponse: + rr.err = ErrorResponseToPgError(msg) + case *pgproto3.ReadyForQuery: + rr.pgConn.contextWatcher.Unwatch() + rr.pgConn.unlock() + return rr.commandTag, rr.err + } + } + } + + return rr.commandTag, rr.err +} + +// readUntilRowDescription ensures the ResultReader's fieldDescriptions are loaded. It does not return an error as any +// error will be stored in the ResultReader. +func (rr *ResultReader) readUntilRowDescription(statementDescription *StatementDescription, resultFormats []int16) { + for !rr.commandConcluded { + msg, _ := rr.receiveMessage() + switch msg := msg.(type) { + case *pgproto3.RowDescription: + return + case *pgproto3.DataRow: + rr.preloadRowValues(msg.Values) + if statementDescription != nil { + sdFields := statementDescription.Fields + rr.fieldDescriptions = rr.pgConn.getFieldDescriptionSlice(len(sdFields)) + + err := combineFieldDescriptionsAndResultFormats(rr.fieldDescriptions, sdFields, resultFormats) + if err != nil { + rr.concludeCommand(CommandTag{}, err) + } + } + return + case *pgproto3.CommandComplete: + if statementDescription != nil { + sdFields := statementDescription.Fields + rr.fieldDescriptions = rr.pgConn.getFieldDescriptionSlice(len(sdFields)) + + err := combineFieldDescriptionsAndResultFormats(rr.fieldDescriptions, sdFields, resultFormats) + if err != nil { + rr.concludeCommand(CommandTag{}, err) + } + } + return + } + } +} + +func (rr *ResultReader) receiveMessage() (msg pgproto3.BackendMessage, err error) { + if rr.multiResultReader == nil { + msg, err = rr.pgConn.receiveMessage() + } else { + msg, err = rr.multiResultReader.receiveMessage() + } + + if err != nil { + err = normalizeTimeoutError(rr.ctx, err) + rr.concludeCommand(CommandTag{}, err) + rr.pgConn.contextWatcher.Unwatch() + rr.closed = true + if rr.multiResultReader == nil { + rr.pgConn.asyncClose() + } + + return nil, rr.err + } + + switch msg := msg.(type) { + case *pgproto3.RowDescription: + rr.fieldDescriptions = rr.pgConn.getFieldDescriptionSlice(len(msg.Fields)) + convertRowDescription(rr.fieldDescriptions, msg) + case *pgproto3.CommandComplete: + rr.concludeCommand(rr.pgConn.makeCommandTag(msg.CommandTag), nil) + case *pgproto3.EmptyQueryResponse: + rr.concludeCommand(CommandTag{}, nil) + case *pgproto3.ErrorResponse: + pgErr := ErrorResponseToPgError(msg) + if rr.pipeline != nil { + rr.pipeline.state.HandleError(pgErr) + } + rr.concludeCommand(CommandTag{}, pgErr) + } + + return msg, nil +} + +func (rr *ResultReader) concludeCommand(commandTag CommandTag, err error) { + // Keep the first error that is recorded. Store the error before checking if the command is already concluded to + // allow for receiving an error after CommandComplete but before ReadyForQuery. + if err != nil && rr.err == nil { + rr.err = err + } + + if rr.commandConcluded { + return + } + + rr.commandTag = commandTag + rr.rowValues = nil + rr.commandConcluded = true +} + +// Batch is a collection of queries that can be sent to the PostgreSQL server in a single round-trip. +type Batch struct { + buf []byte + statementDescriptions []*StatementDescription + resultFormats [][]int16 + err error +} + +// ExecParams appends an ExecParams command to the batch. See PgConn.ExecParams for parameter descriptions. +func (batch *Batch) ExecParams(sql string, paramValues [][]byte, paramOIDs []uint32, paramFormats, resultFormats []int16) { + if batch.err != nil { + return + } + + batch.buf, batch.err = (&pgproto3.Parse{Query: sql, ParameterOIDs: paramOIDs}).Encode(batch.buf) + if batch.err != nil { + return + } + batch.ExecPrepared("", paramValues, paramFormats, resultFormats) +} + +// ExecPrepared appends an ExecPrepared e command to the batch. See PgConn.ExecPrepared for parameter descriptions. +func (batch *Batch) ExecPrepared(stmtName string, paramValues [][]byte, paramFormats, resultFormats []int16) { + if batch.err != nil { + return + } + + batch.buf, batch.err = (&pgproto3.Bind{PreparedStatement: stmtName, ParameterFormatCodes: paramFormats, Parameters: paramValues, ResultFormatCodes: resultFormats}).Encode(batch.buf) + if batch.err != nil { + return + } + + batch.buf, batch.err = (&pgproto3.Describe{ObjectType: 'P'}).Encode(batch.buf) + if batch.err != nil { + return + } + + batch.buf, batch.err = (&pgproto3.Execute{}).Encode(batch.buf) + if batch.err != nil { + return + } +} + +// ExecStatement appends an ExecStatement command to the batch. See PgConn.ExecPrepared for parameter descriptions. +// +// This differs from ExecPrepared in that it takes a *StatementDescription instead of just the prepared statement name. +// Because it has the *StatementDescription it can avoid the Describe Portal message that ExecPrepared must send to get +// the result column descriptions. +func (batch *Batch) ExecStatement(statementDescription *StatementDescription, paramValues [][]byte, paramFormats, resultFormats []int16) { + if batch.err != nil { + return + } + + batch.buf, batch.err = (&pgproto3.Bind{PreparedStatement: statementDescription.Name, ParameterFormatCodes: paramFormats, Parameters: paramValues, ResultFormatCodes: resultFormats}).Encode(batch.buf) + if batch.err != nil { + return + } + + batch.statementDescriptions = append(batch.statementDescriptions, statementDescription) + batch.resultFormats = append(batch.resultFormats, resultFormats) + + batch.buf, batch.err = (&pgproto3.Execute{}).Encode(batch.buf) + if batch.err != nil { + return + } +} + +// ExecBatch executes all the queries in batch in a single round-trip. Execution is implicitly transactional unless a +// transaction is already in progress or SQL contains transaction control statements. This is a simpler way of executing +// multiple queries in a single round trip than using pipeline mode. +func (pgConn *PgConn) ExecBatch(ctx context.Context, batch *Batch) *MultiResultReader { + if batch.err != nil { + return &MultiResultReader{ + closed: true, + err: batch.err, + } + } + + if err := pgConn.lock(); err != nil { + return &MultiResultReader{ + closed: true, + err: err, + } + } + + pgConn.multiResultReader = MultiResultReader{ + pgConn: pgConn, + ctx: ctx, + statementDescriptions: batch.statementDescriptions, + resultFormats: batch.resultFormats, + } + multiResult := &pgConn.multiResultReader + + if ctx != context.Background() { + select { + case <-ctx.Done(): + multiResult.closed = true + multiResult.err = newContextAlreadyDoneError(ctx) + pgConn.unlock() + return multiResult + default: + } + pgConn.contextWatcher.Watch(ctx) + } + + batch.buf, batch.err = (&pgproto3.Sync{}).Encode(batch.buf) + if batch.err != nil { + pgConn.contextWatcher.Unwatch() + multiResult.err = normalizeTimeoutError(multiResult.ctx, batch.err) + multiResult.closed = true + pgConn.asyncClose() + return multiResult + } + + _, err := func(buf []byte) (int, error) { + pgConn.enterPotentialWriteReadDeadlock() + defer pgConn.exitPotentialWriteReadDeadlock() + return pgConn.conn.Write(buf) + }(batch.buf) + if err != nil { + pgConn.contextWatcher.Unwatch() + multiResult.err = normalizeTimeoutError(multiResult.ctx, err) + multiResult.closed = true + pgConn.asyncClose() + return multiResult + } + + return multiResult +} + +// EscapeString escapes a string such that it can safely be interpolated into a SQL command string. It does not include +// the surrounding single quotes. +// +// The current implementation requires that standard_conforming_strings=on and client_encoding="UTF8". If these +// conditions are not met an error will be returned. It is possible these restrictions will be lifted in the future. +func (pgConn *PgConn) EscapeString(s string) (string, error) { + if pgConn.ParameterStatus("standard_conforming_strings") != "on" { + return "", errors.New("EscapeString must be run with standard_conforming_strings=on") + } + + if pgConn.ParameterStatus("client_encoding") != "UTF8" { + return "", errors.New("EscapeString must be run with client_encoding=UTF8") + } + + return strings.ReplaceAll(s, "'", "''"), nil +} + +// CheckConn checks the underlying connection without writing any bytes. This is currently implemented by doing a read +// with a very short deadline. This can be useful because a TCP connection can be broken such that a write will appear +// to succeed even though it will never actually reach the server. Reading immediately before a write will detect this +// condition. If this is done immediately before sending a query it reduces the chances a query will be sent that fails +// without the client knowing whether the server received it or not. +// +// Deprecated: CheckConn is deprecated in favor of Ping. CheckConn cannot detect all types of broken connections where +// the write would still appear to succeed. Prefer Ping unless on a high latency connection. +func (pgConn *PgConn) CheckConn() error { + ctx, cancel := context.WithTimeout(context.Background(), 1*time.Millisecond) + defer cancel() + + _, err := pgConn.ReceiveMessage(ctx) + if err != nil { + if !Timeout(err) { + return err + } + } + + return nil +} + +// Ping pings the server. This can be useful because a TCP connection can be broken such that a write will appear to +// succeed even though it will never actually reach the server. Pinging immediately before sending a query reduces the +// chances a query will be sent that fails without the client knowing whether the server received it or not. +func (pgConn *PgConn) Ping(ctx context.Context) error { + return pgConn.Exec(ctx, "-- ping").Close() +} + +// makeCommandTag makes a CommandTag. It does not retain a reference to buf or buf's underlying memory. +func (pgConn *PgConn) makeCommandTag(buf []byte) CommandTag { + return CommandTag{s: string(buf)} +} + +// enterPotentialWriteReadDeadlock must be called before a write that could deadlock if the server is simultaneously +// blocked writing to us. +func (pgConn *PgConn) enterPotentialWriteReadDeadlock() { + // The time to wait is somewhat arbitrary. A Write should only take as long as the syscall and memcpy to the OS + // outbound network buffer unless the buffer is full (which potentially is a block). It needs to be long enough for + // the normal case, but short enough not to kill performance if a block occurs. + // + // In addition, on Windows the default timer resolution is 15.6ms. So setting the timer to less than that is + // ineffective. + if pgConn.slowWriteTimer.Reset(15 * time.Millisecond) { + panic("BUG: slow write timer already active") + } +} + +// exitPotentialWriteReadDeadlock must be called after a call to enterPotentialWriteReadDeadlock. +func (pgConn *PgConn) exitPotentialWriteReadDeadlock() { + if !pgConn.slowWriteTimer.Stop() { + // The timer starts its function in a separate goroutine. It is necessary to ensure the background reader has + // started before calling Stop. Otherwise, the background reader may not be stopped. That on its own is not a + // serious problem. But what is a serious problem is that the background reader may start at an inopportune time in + // a subsequent query. For example, if a subsequent query was canceled then a deadline may be set on the net.Conn to + // interrupt an in-progress read. After the read is interrupted, but before the deadline is cleared, the background + // reader could start and read a deadline error. Then the next query would receive the an unexpected deadline error. + <-pgConn.bgReaderStarted + pgConn.bgReader.Stop() + } +} + +func (pgConn *PgConn) flushWithPotentialWriteReadDeadlock() error { + pgConn.enterPotentialWriteReadDeadlock() + defer pgConn.exitPotentialWriteReadDeadlock() + err := pgConn.frontend.Flush() + return err +} + +// SyncConn prepares the underlying net.Conn for direct use. PgConn may internally buffer reads or use goroutines for +// background IO. This means that any direct use of the underlying net.Conn may be corrupted if a read is already +// buffered or a read is in progress. SyncConn drains read buffers and stops background IO. In some cases this may +// require sending a ping to the server. ctx can be used to cancel this operation. This should be called before any +// operation that will use the underlying net.Conn directly. e.g. Before Conn() or Hijack(). +// +// This should not be confused with the PostgreSQL protocol Sync message. +func (pgConn *PgConn) SyncConn(ctx context.Context) error { + for range 10 { + if pgConn.bgReader.Status() == bgreader.StatusStopped && pgConn.frontend.ReadBufferLen() == 0 { + return nil + } + + err := pgConn.Ping(ctx) + if err != nil { + return fmt.Errorf("SyncConn: Ping failed while syncing conn: %w", err) + } + } + + // This should never happen. Only way I can imagine this occurring is if the server is constantly sending data such as + // LISTEN/NOTIFY or log notifications such that we never can get an empty buffer. + return errors.New("SyncConn: conn never synchronized") +} + +// CustomData returns a map that can be used to associate custom data with the connection. +func (pgConn *PgConn) CustomData() map[string]any { + return pgConn.customData +} + +// HijackedConn is the result of hijacking a connection. +// +// Due to the necessary exposure of internal implementation details, it is not covered by the semantic versioning +// compatibility. +type HijackedConn struct { + Conn net.Conn + TLSConfig *tls.Config // tls.Config that Conn was negotiated with; nil if Conn is not TLS + PID uint32 // backend pid + SecretKey []byte // key to use to send a cancel query message to the server + ParameterStatuses map[string]string // parameters that have been reported by the server + TxStatus byte + Frontend *pgproto3.Frontend + Config *Config + CustomData map[string]any +} + +// Hijack extracts the internal connection data. pgConn must be in an idle state. SyncConn should be called immediately +// before Hijack. pgConn is unusable after hijacking. Hijacking is typically only useful when using pgconn to establish +// a connection, but taking complete control of the raw connection after that (e.g. a load balancer or proxy). +// +// Due to the necessary exposure of internal implementation details, it is not covered by the semantic versioning +// compatibility. +func (pgConn *PgConn) Hijack() (*HijackedConn, error) { + if err := pgConn.lock(); err != nil { + return nil, err + } + pgConn.status = connStatusClosed + + return &HijackedConn{ + Conn: pgConn.conn, + TLSConfig: pgConn.tlsConfig, + PID: pgConn.pid, + SecretKey: pgConn.secretKey, + ParameterStatuses: pgConn.parameterStatuses, + TxStatus: pgConn.txStatus, + Frontend: pgConn.frontend, + Config: pgConn.config, + CustomData: pgConn.customData, + }, nil +} + +// Construct created a PgConn from an already established connection to a PostgreSQL server. This is the inverse of +// PgConn.Hijack. The connection must be in an idle state. +// +// hc.Frontend is replaced by a new pgproto3.Frontend built by hc.Config.BuildFrontend. +// +// Due to the necessary exposure of internal implementation details, it is not covered by the semantic versioning +// compatibility. +func Construct(hc *HijackedConn) (*PgConn, error) { + pgConn := &PgConn{ + conn: hc.Conn, + tlsConfig: hc.TLSConfig, + pid: hc.PID, + secretKey: hc.SecretKey, + parameterStatuses: hc.ParameterStatuses, + txStatus: hc.TxStatus, + frontend: hc.Frontend, + config: hc.Config, + customData: hc.CustomData, + + status: connStatusIdle, + + cleanupDone: make(chan struct{}), + } + + pgConn.contextWatcher = ctxwatch.NewContextWatcher(hc.Config.BuildContextWatcherHandler(pgConn)) + pgConn.bgReader = bgreader.New(pgConn.conn) + pgConn.slowWriteTimer = time.AfterFunc(time.Duration(math.MaxInt64), + func() { + pgConn.bgReader.Start() + pgConn.bgReaderStarted <- struct{}{} + }, + ) + pgConn.slowWriteTimer.Stop() + pgConn.bgReaderStarted = make(chan struct{}) + pgConn.frontend = hc.Config.BuildFrontend(pgConn.bgReader, pgConn.conn) + + return pgConn, nil +} + +// Pipeline represents a connection in pipeline mode. +// +// SendPrepare, SendQueryParams, SendQueryPrepared, and SendQueryStatement queue requests to the server. These requests +// are not written until pipeline is flushed by Flush or Sync. Sync must be called after the last request is queued. +// Requests between synchronization points are implicitly transactional unless explicit transaction control statements +// have been issued. +// +// The context the pipeline was started with is in effect for the entire life of the Pipeline. +// +// For a deeper understanding of pipeline mode see the PostgreSQL documentation for the extended query protocol +// (https://www.postgresql.org/docs/current/protocol-flow.html#PROTOCOL-FLOW-EXT-QUERY) and the libpq pipeline mode +// (https://www.postgresql.org/docs/current/libpq-pipeline-mode.html). +type Pipeline struct { + conn *PgConn + ctx context.Context + + state pipelineState + err error + closed bool +} + +// PipelineSync is returned by GetResults when a ReadyForQuery message is received. +type PipelineSync struct{} + +// CloseComplete is returned by GetResults when a CloseComplete message is received. +type CloseComplete struct{} + +type pipelineRequestType int + +const ( + pipelineNil pipelineRequestType = iota + pipelinePrepare + pipelineQueryParams + pipelineQueryPrepared + pipelineQueryStatement + pipelineDeallocate + pipelineSyncRequest + pipelineFlushRequest +) + +type pipelineRequestEvent struct { + RequestType pipelineRequestType + WasSentToServer bool + BeforeFlushOrSync bool +} + +type pipelineState struct { + requestEventQueue list.List + statementDescriptionsQueue list.List + resultFormatsQueue list.List + lastRequestType pipelineRequestType + pgErr *PgError + expectedReadyForQueryCount int +} + +func (s *pipelineState) Init() { + s.requestEventQueue.Init() + s.statementDescriptionsQueue.Init() + s.resultFormatsQueue.Init() + s.lastRequestType = pipelineNil +} + +func (s *pipelineState) RegisterSendingToServer() { + for elem := s.requestEventQueue.Back(); elem != nil; elem = elem.Prev() { + val := elem.Value.(pipelineRequestEvent) + if val.WasSentToServer { + return + } + val.WasSentToServer = true + elem.Value = val + } +} + +func (s *pipelineState) registerFlushingBufferOnServer() { + for elem := s.requestEventQueue.Back(); elem != nil; elem = elem.Prev() { + val := elem.Value.(pipelineRequestEvent) + if val.BeforeFlushOrSync { + return + } + val.BeforeFlushOrSync = true + elem.Value = val + } +} + +func (s *pipelineState) PushBackRequestType(req pipelineRequestType) { + if req == pipelineNil { + return + } + + if req != pipelineFlushRequest { + s.requestEventQueue.PushBack(pipelineRequestEvent{RequestType: req}) + } + if req == pipelineFlushRequest || req == pipelineSyncRequest { + s.registerFlushingBufferOnServer() + } + s.lastRequestType = req + + if req == pipelineSyncRequest { + s.expectedReadyForQueryCount++ + } +} + +func (s *pipelineState) ExtractFrontRequestType() pipelineRequestType { + for { + elem := s.requestEventQueue.Front() + if elem == nil { + return pipelineNil + } + val := elem.Value.(pipelineRequestEvent) + if !(val.WasSentToServer && val.BeforeFlushOrSync) { + return pipelineNil + } + + s.requestEventQueue.Remove(elem) + if val.RequestType == pipelineSyncRequest { + s.pgErr = nil + } + if s.pgErr == nil { + return val.RequestType + } + } +} + +func (s *pipelineState) PushBackStatementData(sd *StatementDescription, resultFormats []int16) { + s.statementDescriptionsQueue.PushBack(sd) + s.resultFormatsQueue.PushBack(resultFormats) +} + +func (s *pipelineState) ExtractFrontStatementData() (*StatementDescription, []int16) { + sdElem := s.statementDescriptionsQueue.Front() + var sd *StatementDescription + if sdElem != nil { + s.statementDescriptionsQueue.Remove(sdElem) + sd = sdElem.Value.(*StatementDescription) + } + + rfElem := s.resultFormatsQueue.Front() + var resultFormats []int16 + if rfElem != nil { + s.resultFormatsQueue.Remove(rfElem) + resultFormats = rfElem.Value.([]int16) + } + + return sd, resultFormats +} + +func (s *pipelineState) HandleError(err *PgError) { + s.pgErr = err +} + +func (s *pipelineState) HandleReadyForQuery() { + s.expectedReadyForQueryCount-- +} + +func (s *pipelineState) PendingSync() bool { + var notPendingSync bool + + if elem := s.requestEventQueue.Back(); elem != nil { + val := elem.Value.(pipelineRequestEvent) + notPendingSync = (val.RequestType == pipelineSyncRequest) && val.WasSentToServer + } else { + notPendingSync = (s.lastRequestType == pipelineSyncRequest) || (s.lastRequestType == pipelineNil) + } + + return !notPendingSync +} + +func (s *pipelineState) ExpectedReadyForQuery() int { + return s.expectedReadyForQueryCount +} + +// StartPipeline switches the connection to pipeline mode and returns a *Pipeline. In pipeline mode requests can be sent +// to the server without waiting for a response. Close must be called on the returned *Pipeline to return the connection +// to normal mode. While in pipeline mode, no methods that communicate with the server may be called except +// CancelRequest and Close. ctx is in effect for entire life of the *Pipeline. +// +// Prefer ExecBatch when only sending one group of queries at once. +func (pgConn *PgConn) StartPipeline(ctx context.Context) *Pipeline { + if err := pgConn.lock(); err != nil { + pipeline := &Pipeline{ + closed: true, + err: err, + } + pipeline.state.Init() + + return pipeline + } + + pgConn.resultReader = ResultReader{closed: true} + + pgConn.pipeline = Pipeline{ + conn: pgConn, + ctx: ctx, + } + pgConn.pipeline.state.Init() + + pipeline := &pgConn.pipeline + + if ctx != context.Background() { + select { + case <-ctx.Done(): + pipeline.closed = true + pipeline.err = newContextAlreadyDoneError(ctx) + pgConn.unlock() + return pipeline + default: + } + pgConn.contextWatcher.Watch(ctx) + } + + return pipeline +} + +// SendPrepare is the pipeline version of *PgConn.Prepare. +func (p *Pipeline) SendPrepare(name, sql string, paramOIDs []uint32) { + if p.closed { + return + } + + p.conn.frontend.SendParse(&pgproto3.Parse{Name: name, Query: sql, ParameterOIDs: paramOIDs}) + p.conn.frontend.SendDescribe(&pgproto3.Describe{ObjectType: 'S', Name: name}) + p.state.PushBackRequestType(pipelinePrepare) +} + +// SendDeallocate deallocates a prepared statement. +func (p *Pipeline) SendDeallocate(name string) { + if p.closed { + return + } + + p.conn.frontend.SendClose(&pgproto3.Close{ObjectType: 'S', Name: name}) + p.state.PushBackRequestType(pipelineDeallocate) +} + +// SendQueryParams is the pipeline version of *PgConn.ExecParams. +func (p *Pipeline) SendQueryParams(sql string, paramValues [][]byte, paramOIDs []uint32, paramFormats, resultFormats []int16) { + if p.closed { + return + } + + p.conn.frontend.SendParse(&pgproto3.Parse{Query: sql, ParameterOIDs: paramOIDs}) + p.conn.frontend.SendBind(&pgproto3.Bind{ParameterFormatCodes: paramFormats, Parameters: paramValues, ResultFormatCodes: resultFormats}) + p.conn.frontend.SendDescribe(&pgproto3.Describe{ObjectType: 'P'}) + p.conn.frontend.SendExecute(&pgproto3.Execute{}) + p.state.PushBackRequestType(pipelineQueryParams) +} + +// SendQueryPrepared is the pipeline version of *PgConn.ExecPrepared. +func (p *Pipeline) SendQueryPrepared(stmtName string, paramValues [][]byte, paramFormats, resultFormats []int16) { + if p.closed { + return + } + + p.conn.frontend.SendBind(&pgproto3.Bind{PreparedStatement: stmtName, ParameterFormatCodes: paramFormats, Parameters: paramValues, ResultFormatCodes: resultFormats}) + p.conn.frontend.SendDescribe(&pgproto3.Describe{ObjectType: 'P'}) + p.conn.frontend.SendExecute(&pgproto3.Execute{}) + p.state.PushBackRequestType(pipelineQueryPrepared) +} + +// SendQueryStatement is the pipeline version of *PgConn.ExecStatement. +func (p *Pipeline) SendQueryStatement(statementDescription *StatementDescription, paramValues [][]byte, paramFormats, resultFormats []int16) { + if p.closed { + return + } + + p.conn.frontend.SendBind(&pgproto3.Bind{PreparedStatement: statementDescription.Name, ParameterFormatCodes: paramFormats, Parameters: paramValues, ResultFormatCodes: resultFormats}) + p.conn.frontend.SendExecute(&pgproto3.Execute{}) + p.state.PushBackRequestType(pipelineQueryStatement) + p.state.PushBackStatementData(statementDescription, resultFormats) +} + +// SendFlushRequest sends a request for the server to flush its output buffer. +// +// The server flushes its output buffer automatically as a result of Sync being called, +// or on any request when not in pipeline mode; this function is useful to cause the server +// to flush its output buffer in pipeline mode without establishing a synchronization point. +// Note that the request is not itself flushed to the server automatically; use Flush if +// necessary. This copies the behavior of libpq PQsendFlushRequest. +func (p *Pipeline) SendFlushRequest() { + if p.closed { + return + } + + p.conn.frontend.Send(&pgproto3.Flush{}) + p.state.PushBackRequestType(pipelineFlushRequest) +} + +// SendPipelineSync marks a synchronization point in a pipeline by sending a sync message +// without flushing the send buffer. This serves as the delimiter of an implicit +// transaction and an error recovery point. +// +// Note that the request is not itself flushed to the server automatically; use Flush if +// necessary. This copies the behavior of libpq PQsendPipelineSync. +func (p *Pipeline) SendPipelineSync() { + if p.closed { + return + } + + p.conn.frontend.SendSync(&pgproto3.Sync{}) + p.state.PushBackRequestType(pipelineSyncRequest) +} + +// Flush flushes the queued requests without establishing a synchronization point. +func (p *Pipeline) Flush() error { + if p.closed { + if p.err != nil { + return p.err + } + return errors.New("pipeline closed") + } + + err := p.conn.flushWithPotentialWriteReadDeadlock() + if err != nil { + err = normalizeTimeoutError(p.ctx, err) + + p.conn.asyncClose() + + p.conn.contextWatcher.Unwatch() + p.conn.unlock() + p.closed = true + p.err = err + return err + } + + p.state.RegisterSendingToServer() + return nil +} + +// Sync establishes a synchronization point and flushes the queued requests. +func (p *Pipeline) Sync() error { + p.SendPipelineSync() + return p.Flush() +} + +// GetResults gets the next results. If results are present, results may be a *ResultReader, *StatementDescription, or +// *PipelineSync. If an ErrorResponse is received from the server, results will be nil and err will be a *PgError. If no +// results are available, results and err will both be nil. +func (p *Pipeline) GetResults() (results any, err error) { + if p.closed { + if p.err != nil { + return nil, p.err + } + return nil, errors.New("pipeline closed") + } + + return p.getResults() +} + +func (p *Pipeline) getResults() (results any, err error) { + if !p.conn.resultReader.closed { + _, err := p.conn.resultReader.Close() + if err != nil { + return nil, err + } + } + + currentRequestType := p.state.ExtractFrontRequestType() + switch currentRequestType { + case pipelineNil: + return nil, nil + case pipelinePrepare: + return p.getResultsPrepare() + case pipelineQueryParams: + return p.getResultsQueryParams() + case pipelineQueryPrepared: + return p.getResultsQueryPrepared() + case pipelineQueryStatement: + return p.getResultsQueryStatement() + case pipelineDeallocate: + return p.getResultsDeallocate() + case pipelineSyncRequest: + return p.getResultsSync() + case pipelineFlushRequest: + return nil, errors.New("BUG: pipelineFlushRequest should not be in request queue") + default: + return nil, errors.New("BUG: unknown pipeline request type") + } +} + +func (p *Pipeline) getResultsPrepare() (*StatementDescription, error) { + err := p.receiveParseComplete("Prepare") + if err != nil { + return nil, err + } + + psd := &StatementDescription{} + + msg, err := p.receiveMessage() + if err != nil { + return nil, err + } + + switch msg := msg.(type) { + case *pgproto3.ParameterDescription: + psd.ParamOIDs = make([]uint32, len(msg.ParameterOIDs)) + copy(psd.ParamOIDs, msg.ParameterOIDs) + case *pgproto3.ErrorResponse: + pgErr := ErrorResponseToPgError(msg) + p.state.HandleError(pgErr) + return nil, pgErr + default: + return nil, p.handleUnexpectedMessage("Prepare ParameterDescription", msg) + } + + msg, err = p.receiveMessage() + if err != nil { + return nil, err + } + + switch msg := msg.(type) { + case *pgproto3.RowDescription: + psd.Fields = make([]FieldDescription, len(msg.Fields)) + convertRowDescription(psd.Fields, msg) + return psd, nil + + // NoData is returned instead of RowDescription when there is no expected result. e.g. An INSERT without a RETURNING + // clause. + case *pgproto3.NoData: + return psd, nil + + case *pgproto3.ErrorResponse: + pgErr := ErrorResponseToPgError(msg) + p.state.HandleError(pgErr) + return nil, pgErr + default: + return nil, p.handleUnexpectedMessage("Prepare RowDescription", msg) + } +} + +func (p *Pipeline) getResultsQueryParams() (*ResultReader, error) { + err := p.receiveParseComplete("QueryParams") + if err != nil { + return nil, err + } + + err = p.receiveBindComplete("QueryParams") + if err != nil { + return nil, err + } + + return p.receiveDescribedResultReader("QueryParams") +} + +func (p *Pipeline) getResultsQueryPrepared() (*ResultReader, error) { + err := p.receiveBindComplete("QueryPrepared") + if err != nil { + return nil, err + } + + return p.receiveDescribedResultReader("QueryPrepared") +} + +func (p *Pipeline) getResultsQueryStatement() (*ResultReader, error) { + err := p.receiveBindComplete("QueryStatement") + if err != nil { + return nil, err + } + + msg, err := p.receiveMessage() + if err != nil { + return nil, err + } + + sd, resultFormats := p.state.ExtractFrontStatementData() + if sd == nil { + return nil, errors.New("BUG: missing statement description or result formats for QueryStatement") + } + sdFields := sd.Fields + fieldDescriptions := p.conn.getFieldDescriptionSlice(len(sdFields)) + err = combineFieldDescriptionsAndResultFormats(fieldDescriptions, sdFields, resultFormats) + if err != nil { + return nil, err + } + + switch msg := msg.(type) { + case *pgproto3.DataRow: + rr := ResultReader{ + pgConn: p.conn, + pipeline: p, + ctx: p.ctx, + fieldDescriptions: fieldDescriptions, + } + rr.preloadRowValues(msg.Values) + p.conn.resultReader = rr + return &p.conn.resultReader, nil + case *pgproto3.CommandComplete: + p.conn.resultReader = ResultReader{ + commandTag: p.conn.makeCommandTag(msg.CommandTag), + commandConcluded: true, + closed: true, + fieldDescriptions: fieldDescriptions, + } + return &p.conn.resultReader, nil + case *pgproto3.ErrorResponse: + pgErr := ErrorResponseToPgError(msg) + p.state.HandleError(pgErr) + p.conn.resultReader.closed = true + return nil, pgErr + default: + return nil, p.handleUnexpectedMessage("QueryStatement", msg) + } +} + +func (p *Pipeline) getResultsDeallocate() (*CloseComplete, error) { + msg, err := p.receiveMessage() + if err != nil { + return nil, err + } + + switch msg := msg.(type) { + case *pgproto3.CloseComplete: + return &CloseComplete{}, nil + case *pgproto3.ErrorResponse: + pgErr := ErrorResponseToPgError(msg) + p.state.HandleError(pgErr) + p.conn.resultReader.closed = true + return nil, pgErr + default: + return nil, p.handleUnexpectedMessage("Deallocate", msg) + } +} + +func (p *Pipeline) getResultsSync() (*PipelineSync, error) { + msg, err := p.receiveMessage() + if err != nil { + return nil, err + } + + switch msg := msg.(type) { + case *pgproto3.ReadyForQuery: + p.state.HandleReadyForQuery() + return &PipelineSync{}, nil + case *pgproto3.ErrorResponse: + // Error message that is received while expecting a Sync message still consumes the expected Sync. Put it back. + p.state.requestEventQueue.PushFront(pipelineRequestEvent{RequestType: pipelineSyncRequest, WasSentToServer: true, BeforeFlushOrSync: true}) + + pgErr := ErrorResponseToPgError(msg) + p.state.HandleError(pgErr) + p.conn.resultReader.closed = true + return nil, pgErr + default: + return nil, p.handleUnexpectedMessage("Sync", msg) + } +} + +func (p *Pipeline) receiveParseComplete(errStr string) error { + msg, err := p.receiveMessage() + if err != nil { + return err + } + + switch msg := msg.(type) { + case *pgproto3.ParseComplete: + return nil + case *pgproto3.ErrorResponse: + pgErr := ErrorResponseToPgError(msg) + p.state.HandleError(pgErr) + return pgErr + default: + return p.handleUnexpectedMessage(fmt.Sprintf("%s Parse", errStr), msg) + } +} + +func (p *Pipeline) receiveBindComplete(errStr string) error { + msg, err := p.receiveMessage() + if err != nil { + return err + } + + switch msg := msg.(type) { + case *pgproto3.BindComplete: + return nil + case *pgproto3.ErrorResponse: + pgErr := ErrorResponseToPgError(msg) + p.state.HandleError(pgErr) + return pgErr + default: + return p.handleUnexpectedMessage(fmt.Sprintf("%s Bind", errStr), msg) + } +} + +func (p *Pipeline) receiveDescribedResultReader(errStr string) (*ResultReader, error) { + msg, err := p.receiveMessage() + if err != nil { + return nil, err + } + + switch msg := msg.(type) { + case *pgproto3.RowDescription: + p.conn.resultReader = ResultReader{ + pgConn: p.conn, + pipeline: p, + ctx: p.ctx, + fieldDescriptions: p.conn.getFieldDescriptionSlice(len(msg.Fields)), + } + convertRowDescription(p.conn.resultReader.fieldDescriptions, msg) + return &p.conn.resultReader, nil + case *pgproto3.NoData: + case *pgproto3.ErrorResponse: + pgErr := ErrorResponseToPgError(msg) + p.state.HandleError(pgErr) + p.conn.resultReader.closed = true + return nil, pgErr + default: + return nil, p.handleUnexpectedMessage(fmt.Sprintf("%s RowDescription or NoData", errStr), msg) + } + + msg, err = p.receiveMessage() + if err != nil { + return nil, err + } + + switch msg := msg.(type) { + case *pgproto3.CommandComplete: + p.conn.resultReader = ResultReader{ + commandTag: p.conn.makeCommandTag(msg.CommandTag), + commandConcluded: true, + closed: true, + } + return &p.conn.resultReader, nil + case *pgproto3.ErrorResponse: + pgErr := ErrorResponseToPgError(msg) + p.state.HandleError(pgErr) + p.conn.resultReader.closed = true + return nil, pgErr + default: + return nil, p.handleUnexpectedMessage(fmt.Sprintf("%s CommandComplete", errStr), msg) + } +} + +func (p *Pipeline) receiveMessage() (pgproto3.BackendMessage, error) { + for { + msg, err := p.conn.receiveMessage() + if err != nil { + p.err = err + p.conn.asyncClose() + return nil, normalizeTimeoutError(p.ctx, err) + } + + switch msg := msg.(type) { + case *pgproto3.ParameterStatus, *pgproto3.NoticeResponse, *pgproto3.NotificationResponse: + // Filter these message types out in pipeline mode. The normal processing is handled by PgConn.receiveMessage. + default: + return msg, nil + } + } +} + +func (p *Pipeline) handleUnexpectedMessage(errStr string, msg pgproto3.BackendMessage) error { + p.err = fmt.Errorf("pipeline: %s: received unexpected message type %T", errStr, msg) + p.conn.asyncClose() + return p.err +} + +// Close closes the pipeline and returns the connection to normal mode. +func (p *Pipeline) Close() error { + if p.closed { + return p.err + } + + p.closed = true + + if p.state.PendingSync() { + p.conn.asyncClose() + p.err = errors.New("pipeline has unsynced requests") + p.conn.contextWatcher.Unwatch() + p.conn.unlock() + + return p.err + } + + for p.state.ExpectedReadyForQuery() > 0 { + results, err := p.getResults() + if err != nil { + p.err = err + var pgErr *PgError + if !errors.As(err, &pgErr) { + p.conn.asyncClose() + break + } + } else if results == nil { + // getResults returns (nil, nil) when the request queue is exhausted but + // ExpectedReadyForQuery is still > 0. This can happen when FATAL errors consume + // queued request slots without the server ever sending ReadyForQuery. + p.conn.asyncClose() + if p.err == nil { + p.err = errors.New("pipeline: no more results but expected ReadyForQuery") + } + break + } + } + + p.conn.contextWatcher.Unwatch() + p.conn.unlock() + + return p.err +} + +// DeadlineContextWatcherHandler handles canceled contexts by setting a deadline on a net.Conn. +type DeadlineContextWatcherHandler struct { + Conn net.Conn + + // DeadlineDelay is the delay to set on the deadline set on net.Conn when the context is canceled. + DeadlineDelay time.Duration +} + +func (h *DeadlineContextWatcherHandler) HandleCancel(ctx context.Context) { + h.Conn.SetDeadline(time.Now().Add(h.DeadlineDelay)) +} + +func (h *DeadlineContextWatcherHandler) HandleUnwatchAfterCancel() { + h.Conn.SetDeadline(time.Time{}) +} + +// CancelRequestContextWatcherHandler handles canceled contexts by sending a cancel request to the server. It also sets +// a deadline on a net.Conn as a fallback. +type CancelRequestContextWatcherHandler struct { + Conn *PgConn + + // CancelRequestDelay is the delay before sending the cancel request to the server. + CancelRequestDelay time.Duration + + // DeadlineDelay is the delay to set on the deadline set on net.Conn when the context is canceled. + DeadlineDelay time.Duration + + cancelFinishedChan chan struct{} + handleUnwatchAfterCancelCalled func() +} + +func (h *CancelRequestContextWatcherHandler) HandleCancel(context.Context) { + h.cancelFinishedChan = make(chan struct{}) + var handleUnwatchedAfterCancelCalledCtx context.Context + handleUnwatchedAfterCancelCalledCtx, h.handleUnwatchAfterCancelCalled = context.WithCancel(context.Background()) + + deadline := time.Now().Add(h.DeadlineDelay) + h.Conn.conn.SetDeadline(deadline) + + go func() { + defer close(h.cancelFinishedChan) + + select { + case <-handleUnwatchedAfterCancelCalledCtx.Done(): + return + case <-time.After(h.CancelRequestDelay): + } + + cancelRequestCtx, cancel := context.WithDeadline(handleUnwatchedAfterCancelCalledCtx, deadline) + defer cancel() + h.Conn.CancelRequest(cancelRequestCtx) + + // CancelRequest is inherently racy. Even though the cancel request has been received by the server at this point, + // it hasn't necessarily been delivered to the other connection. If we immediately return and the connection is + // immediately used then it is possible the CancelRequest will actually cancel our next query. The + // TestCancelRequestContextWatcherHandler Stress test can produce this error without the sleep below. The sleep time + // is arbitrary, but should be sufficient to prevent this error case. + time.Sleep(100 * time.Millisecond) + }() +} + +func (h *CancelRequestContextWatcherHandler) HandleUnwatchAfterCancel() { + h.handleUnwatchAfterCancelCalled() + <-h.cancelFinishedChan + + h.Conn.conn.SetDeadline(time.Time{}) +} + +func combineFieldDescriptionsAndResultFormats(outputFields, inputFields []FieldDescription, resultFormats []int16) error { + switch { + case len(resultFormats) == 0: + // No format codes provided means text format for all columns. + for i := range inputFields { + outputFields[i] = inputFields[i] + outputFields[i].Format = pgtype.TextFormatCode + } + case len(resultFormats) == 1: + // Single format code applies to all columns. + format := resultFormats[0] + for i := range inputFields { + outputFields[i] = inputFields[i] + outputFields[i].Format = format + } + case len(resultFormats) == len(inputFields): + // One format code per column. + for i := range inputFields { + outputFields[i] = inputFields[i] + outputFields[i].Format = resultFormats[i] + } + default: + // This should not occur if Bind validation is correct, but handle gracefully + return fmt.Errorf("result format codes length %d does not match field count %d", len(resultFormats), len(inputFields)) + } + + return nil +} diff --git a/vendor/github.com/jackc/pgx/v5/pgconn/require_auth.go b/vendor/github.com/jackc/pgx/v5/pgconn/require_auth.go new file mode 100644 index 0000000000..3449a2fcc8 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgconn/require_auth.go @@ -0,0 +1,146 @@ +package pgconn + +import ( + "fmt" + "strings" +) + +// authMethod is one of the method keywords accepted by libpq's require_auth. +type authMethod uint8 + +const ( + authMethodPassword authMethod = iota + authMethodMD5 + authMethodGSS + authMethodSSPI + authMethodSCRAMSHA256 + authMethodOAuth + authMethodNone + authMethodCount +) + +var authMethodNames = [authMethodCount]string{ + authMethodPassword: "password", + authMethodMD5: "md5", + authMethodGSS: "gss", + authMethodSSPI: "sspi", + authMethodSCRAMSHA256: "scram-sha-256", + authMethodOAuth: "oauth", + authMethodNone: "none", +} + +// requireAuth is the parsed form of the require_auth connection parameter. It mirrors libpq's +// auth_required / allowed_auth_methods bookkeeping (see fe-connect.c, conn->allowed_auth_methods). +type requireAuth struct { + // raw is the original parameter value, used in error messages. + raw string + + // authRequired is true when the server must complete an authentication exchange before sending + // AuthenticationOk. It is false when the parameter is unset, fully negated, or "none" is in the + // allowed set. + authRequired bool + + // allowed is a bitmask of permitted authMethod values. + allowed uint8 +} + +func (ra requireAuth) allows(m authMethod) bool { + return ra.allowed&(1< 1 { + idx := bytes.IndexByte(authMechanisms, 0) + if idx == -1 { + return &invalidMessageFormatErr{messageType: "AuthenticationSASL", details: "unterminated string"} + } + dst.AuthMechanisms = append(dst.AuthMechanisms, string(authMechanisms[:idx])) + authMechanisms = authMechanisms[idx+1:] + } + + return nil +} + +// Encode encodes src into dst. dst will include the 1 byte message type identifier and the 4 byte message length. +func (src *AuthenticationSASL) Encode(dst []byte) ([]byte, error) { + dst, sp := beginMessage(dst, 'R') + dst = pgio.AppendUint32(dst, AuthTypeSASL) + + for _, s := range src.AuthMechanisms { + dst = append(dst, []byte(s)...) + dst = append(dst, 0) + } + dst = append(dst, 0) + + return finishMessage(dst, sp) +} + +// MarshalJSON implements encoding/json.Marshaler. +func (src AuthenticationSASL) MarshalJSON() ([]byte, error) { + return json.Marshal(struct { + Type string + AuthMechanisms []string + }{ + Type: "AuthenticationSASL", + AuthMechanisms: src.AuthMechanisms, + }) +} diff --git a/vendor/github.com/jackc/pgx/v5/pgproto3/authentication_sasl_continue.go b/vendor/github.com/jackc/pgx/v5/pgproto3/authentication_sasl_continue.go new file mode 100644 index 0000000000..70fba4a67f --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgproto3/authentication_sasl_continue.go @@ -0,0 +1,75 @@ +package pgproto3 + +import ( + "encoding/binary" + "encoding/json" + "errors" + + "github.com/jackc/pgx/v5/internal/pgio" +) + +// AuthenticationSASLContinue is a message sent from the backend containing a SASL challenge. +type AuthenticationSASLContinue struct { + Data []byte +} + +// Backend identifies this message as sendable by the PostgreSQL backend. +func (*AuthenticationSASLContinue) Backend() {} + +// Backend identifies this message as an authentication response. +func (*AuthenticationSASLContinue) AuthenticationResponse() {} + +// Decode decodes src into dst. src must contain the complete message with the exception of the initial 1 byte message +// type identifier and 4 byte message length. +func (dst *AuthenticationSASLContinue) Decode(src []byte) error { + if len(src) < 4 { + return errors.New("authentication message too short") + } + + authType := binary.BigEndian.Uint32(src) + + if authType != AuthTypeSASLContinue { + return errors.New("bad auth type") + } + + dst.Data = src[4:] + + return nil +} + +// Encode encodes src into dst. dst will include the 1 byte message type identifier and the 4 byte message length. +func (src *AuthenticationSASLContinue) Encode(dst []byte) ([]byte, error) { + dst, sp := beginMessage(dst, 'R') + dst = pgio.AppendUint32(dst, AuthTypeSASLContinue) + dst = append(dst, src.Data...) + return finishMessage(dst, sp) +} + +// MarshalJSON implements encoding/json.Marshaler. +func (src AuthenticationSASLContinue) MarshalJSON() ([]byte, error) { + return json.Marshal(struct { + Type string + Data string + }{ + Type: "AuthenticationSASLContinue", + Data: string(src.Data), + }) +} + +// UnmarshalJSON implements encoding/json.Unmarshaler. +func (dst *AuthenticationSASLContinue) UnmarshalJSON(data []byte) error { + // Ignore null, like in the main JSON package. + if string(data) == "null" { + return nil + } + + var msg struct { + Data string + } + if err := json.Unmarshal(data, &msg); err != nil { + return err + } + + dst.Data = []byte(msg.Data) + return nil +} diff --git a/vendor/github.com/jackc/pgx/v5/pgproto3/authentication_sasl_final.go b/vendor/github.com/jackc/pgx/v5/pgproto3/authentication_sasl_final.go new file mode 100644 index 0000000000..84976c2a31 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgproto3/authentication_sasl_final.go @@ -0,0 +1,75 @@ +package pgproto3 + +import ( + "encoding/binary" + "encoding/json" + "errors" + + "github.com/jackc/pgx/v5/internal/pgio" +) + +// AuthenticationSASLFinal is a message sent from the backend indicating a SASL authentication has completed. +type AuthenticationSASLFinal struct { + Data []byte +} + +// Backend identifies this message as sendable by the PostgreSQL backend. +func (*AuthenticationSASLFinal) Backend() {} + +// Backend identifies this message as an authentication response. +func (*AuthenticationSASLFinal) AuthenticationResponse() {} + +// Decode decodes src into dst. src must contain the complete message with the exception of the initial 1 byte message +// type identifier and 4 byte message length. +func (dst *AuthenticationSASLFinal) Decode(src []byte) error { + if len(src) < 4 { + return errors.New("authentication message too short") + } + + authType := binary.BigEndian.Uint32(src) + + if authType != AuthTypeSASLFinal { + return errors.New("bad auth type") + } + + dst.Data = src[4:] + + return nil +} + +// Encode encodes src into dst. dst will include the 1 byte message type identifier and the 4 byte message length. +func (src *AuthenticationSASLFinal) Encode(dst []byte) ([]byte, error) { + dst, sp := beginMessage(dst, 'R') + dst = pgio.AppendUint32(dst, AuthTypeSASLFinal) + dst = append(dst, src.Data...) + return finishMessage(dst, sp) +} + +// MarshalJSON implements encoding/json.Unmarshaler. +func (src AuthenticationSASLFinal) MarshalJSON() ([]byte, error) { + return json.Marshal(struct { + Type string + Data string + }{ + Type: "AuthenticationSASLFinal", + Data: string(src.Data), + }) +} + +// UnmarshalJSON implements encoding/json.Unmarshaler. +func (dst *AuthenticationSASLFinal) UnmarshalJSON(data []byte) error { + // Ignore null, like in the main JSON package. + if string(data) == "null" { + return nil + } + + var msg struct { + Data string + } + if err := json.Unmarshal(data, &msg); err != nil { + return err + } + + dst.Data = []byte(msg.Data) + return nil +} diff --git a/vendor/github.com/jackc/pgx/v5/pgproto3/backend.go b/vendor/github.com/jackc/pgx/v5/pgproto3/backend.go new file mode 100644 index 0000000000..65388ad49b --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgproto3/backend.go @@ -0,0 +1,299 @@ +package pgproto3 + +import ( + "bytes" + "encoding/binary" + "fmt" + "io" +) + +// Backend acts as a server for the PostgreSQL wire protocol version 3. +type Backend struct { + cr *chunkReader + w io.Writer + + // tracer is used to trace messages when Send or Receive is called. This means an outbound message is traced + // before it is actually transmitted (i.e. before Flush). + tracer *tracer + + wbuf []byte + encodeError error + + // Frontend message flyweights + bind Bind + cancelRequest CancelRequest + _close Close + copyFail CopyFail + copyData CopyData + copyDone CopyDone + describe Describe + execute Execute + flush Flush + functionCall FunctionCall + gssEncRequest GSSEncRequest + parse Parse + query Query + sslRequest SSLRequest + startupMessage StartupMessage + sync Sync + terminate Terminate + + bodyLen int + maxBodyLen int // maxBodyLen is the maximum length of a message body in octets. If a message body exceeds this length, Receive will return an error. + msgType byte + partialMsg bool + authType uint32 +} + +const ( + minStartupPacketLen = 4 // minStartupPacketLen is a single 32-bit int version or code. + maxStartupPacketLen = 10_000 // maxStartupPacketLen is MAX_STARTUP_PACKET_LENGTH from PG source. +) + +// NewBackend creates a new Backend. +func NewBackend(r io.Reader, w io.Writer) *Backend { + cr := newChunkReader(r, 0) + return &Backend{cr: cr, w: w} +} + +// Send sends a message to the frontend (i.e. the client). The message is buffered until Flush is called. Any error +// encountered will be returned from Flush. +func (b *Backend) Send(msg BackendMessage) { + if b.encodeError != nil { + return + } + + prevLen := len(b.wbuf) + newBuf, err := msg.Encode(b.wbuf) + if err != nil { + b.encodeError = err + return + } + b.wbuf = newBuf + + if b.tracer != nil { + b.tracer.traceMessage('B', int32(len(b.wbuf)-prevLen), msg) + } +} + +// Flush writes any pending messages to the frontend (i.e. the client). +func (b *Backend) Flush() error { + if err := b.encodeError; err != nil { + b.encodeError = nil + b.wbuf = b.wbuf[:0] + return &writeError{err: err, safeToRetry: true} + } + + n, err := b.w.Write(b.wbuf) + + const maxLen = 1024 + if len(b.wbuf) > maxLen { + b.wbuf = make([]byte, 0, maxLen) + } else { + b.wbuf = b.wbuf[:0] + } + + if err != nil { + return &writeError{err: err, safeToRetry: n == 0} + } + + return nil +} + +// Trace starts tracing the message traffic to w. It writes in a similar format to that produced by the libpq function +// PQtrace. +func (b *Backend) Trace(w io.Writer, options TracerOptions) { + b.tracer = &tracer{ + w: w, + buf: &bytes.Buffer{}, + TracerOptions: options, + } +} + +// Untrace stops tracing. +func (b *Backend) Untrace() { + b.tracer = nil +} + +// ReceiveStartupMessage receives the initial connection message. This method is used of the normal Receive method +// because the initial connection message is "special" and does not include the message type as the first byte. This +// will return either a StartupMessage, SSLRequest, GSSEncRequest, or CancelRequest. +func (b *Backend) ReceiveStartupMessage() (FrontendMessage, error) { + buf, err := b.cr.Next(4) + if err != nil { + return nil, err + } + msgSize := int(int32(binary.BigEndian.Uint32(buf)) - 4) + + if msgSize < minStartupPacketLen || msgSize > maxStartupPacketLen { + return nil, fmt.Errorf("invalid length of startup packet: %d", msgSize) + } + + buf, err = b.cr.Next(msgSize) + if err != nil { + return nil, translateEOFtoErrUnexpectedEOF(err) + } + + code := binary.BigEndian.Uint32(buf) + + switch code { + case ProtocolVersion30, ProtocolVersion32: + err = b.startupMessage.Decode(buf) + if err != nil { + return nil, err + } + return &b.startupMessage, nil + case sslRequestNumber: + err = b.sslRequest.Decode(buf) + if err != nil { + return nil, err + } + return &b.sslRequest, nil + case cancelRequestCode: + err = b.cancelRequest.Decode(buf) + if err != nil { + return nil, err + } + return &b.cancelRequest, nil + case gssEncReqNumber: + err = b.gssEncRequest.Decode(buf) + if err != nil { + return nil, err + } + return &b.gssEncRequest, nil + default: + return nil, fmt.Errorf("unknown startup message code: %d", code) + } +} + +// Receive receives a message from the frontend. The returned message is only valid until the next call to Receive. +func (b *Backend) Receive() (FrontendMessage, error) { + if !b.partialMsg { + header, err := b.cr.Next(5) + if err != nil { + return nil, translateEOFtoErrUnexpectedEOF(err) + } + + b.msgType = header[0] + + msgLength := int(int32(binary.BigEndian.Uint32(header[1:]))) + if msgLength < 4 { + return nil, fmt.Errorf("invalid message length: %d", msgLength) + } + + b.bodyLen = msgLength - 4 + if b.maxBodyLen > 0 && b.bodyLen > b.maxBodyLen { + return nil, &ExceededMaxBodyLenErr{b.maxBodyLen, b.bodyLen} + } + b.partialMsg = true + } + + var msg FrontendMessage + switch b.msgType { + case 'B': + msg = &b.bind + case 'C': + msg = &b._close + case 'D': + msg = &b.describe + case 'E': + msg = &b.execute + case 'F': + msg = &b.functionCall + case 'f': + msg = &b.copyFail + case 'd': + msg = &b.copyData + case 'c': + msg = &b.copyDone + case 'H': + msg = &b.flush + case 'P': + msg = &b.parse + case 'p': + switch b.authType { + case AuthTypeSASL: + msg = &SASLInitialResponse{} + case AuthTypeSASLContinue: + msg = &SASLResponse{} + case AuthTypeSASLFinal: + msg = &SASLResponse{} + case AuthTypeGSS, AuthTypeGSSCont: + msg = &GSSResponse{} + case AuthTypeCleartextPassword, AuthTypeMD5Password: + fallthrough + default: + // to maintain backwards compatibility + msg = &PasswordMessage{} + } + case 'Q': + msg = &b.query + case 'S': + msg = &b.sync + case 'X': + msg = &b.terminate + default: + return nil, fmt.Errorf("unknown message type: %c", b.msgType) + } + + msgBody, err := b.cr.Next(b.bodyLen) + if err != nil { + return nil, translateEOFtoErrUnexpectedEOF(err) + } + + b.partialMsg = false + + err = msg.Decode(msgBody) + if err != nil { + return nil, err + } + + if b.tracer != nil { + b.tracer.traceMessage('F', int32(5+len(msgBody)), msg) + } + + return msg, nil +} + +// SetAuthType sets the authentication type in the backend. +// Since multiple message types can start with 'p', SetAuthType allows +// contextual identification of FrontendMessages. For example, in the +// PG message flow documentation for PasswordMessage: +// +// Byte1('p') +// +// Identifies the message as a password response. Note that this is also used for +// GSSAPI, SSPI and SASL response messages. The exact message type can be deduced from +// the context. +// +// Since the Frontend does not know about the state of a backend, it is important +// to call SetAuthType() after an authentication request is received by the Frontend. +func (b *Backend) SetAuthType(authType uint32) error { + switch authType { + case AuthTypeOk, + AuthTypeCleartextPassword, + AuthTypeMD5Password, + AuthTypeSCMCreds, + AuthTypeGSS, + AuthTypeGSSCont, + AuthTypeSSPI, + AuthTypeSASL, + AuthTypeSASLContinue, + AuthTypeSASLFinal: + b.authType = authType + default: + return fmt.Errorf("authType not recognized: %d", authType) + } + + return nil +} + +// SetMaxBodyLen sets the maximum length of a message body in octets. +// If a message body exceeds this length, Receive will return an error. +// This is useful for protecting against malicious clients that send +// large messages with the intent of causing memory exhaustion. +// The default value is 0. +// If maxBodyLen is 0, then no maximum is enforced. +func (b *Backend) SetMaxBodyLen(maxBodyLen int) { + b.maxBodyLen = maxBodyLen +} diff --git a/vendor/github.com/jackc/pgx/v5/pgproto3/backend_key_data.go b/vendor/github.com/jackc/pgx/v5/pgproto3/backend_key_data.go new file mode 100644 index 0000000000..c73b2da0cc --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgproto3/backend_key_data.go @@ -0,0 +1,71 @@ +package pgproto3 + +import ( + "encoding/binary" + "encoding/hex" + "encoding/json" + + "github.com/jackc/pgx/v5/internal/pgio" +) + +type BackendKeyData struct { + ProcessID uint32 + SecretKey []byte +} + +// Backend identifies this message as sendable by the PostgreSQL backend. +func (*BackendKeyData) Backend() {} + +// Decode decodes src into dst. src must contain the complete message with the exception of the initial 1 byte message +// type identifier and 4 byte message length. +func (dst *BackendKeyData) Decode(src []byte) error { + if len(src) < 8 { + return &invalidMessageLenErr{messageType: "BackendKeyData", expectedLen: 8, actualLen: len(src)} + } + + dst.ProcessID = binary.BigEndian.Uint32(src[:4]) + dst.SecretKey = make([]byte, len(src)-4) + copy(dst.SecretKey, src[4:]) + + return nil +} + +// Encode encodes src into dst. dst will include the 1 byte message type identifier and the 4 byte message length. +func (src *BackendKeyData) Encode(dst []byte) ([]byte, error) { + dst, sp := beginMessage(dst, 'K') + dst = pgio.AppendUint32(dst, src.ProcessID) + dst = append(dst, src.SecretKey...) + return finishMessage(dst, sp) +} + +// MarshalJSON implements encoding/json.Marshaler. +func (src BackendKeyData) MarshalJSON() ([]byte, error) { + return json.Marshal(struct { + Type string + ProcessID uint32 + SecretKey string + }{ + Type: "BackendKeyData", + ProcessID: src.ProcessID, + SecretKey: hex.EncodeToString(src.SecretKey), + }) +} + +// UnmarshalJSON implements encoding/json.Unmarshaler. +func (dst *BackendKeyData) UnmarshalJSON(data []byte) error { + var msg struct { + ProcessID uint32 + SecretKey string + } + if err := json.Unmarshal(data, &msg); err != nil { + return err + } + + dst.ProcessID = msg.ProcessID + secretKey, err := hex.DecodeString(msg.SecretKey) + if err != nil { + return err + } + dst.SecretKey = secretKey + return nil +} diff --git a/vendor/github.com/jackc/pgx/v5/pgproto3/big_endian.go b/vendor/github.com/jackc/pgx/v5/pgproto3/big_endian.go new file mode 100644 index 0000000000..f7bdb97eb7 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgproto3/big_endian.go @@ -0,0 +1,37 @@ +package pgproto3 + +import ( + "encoding/binary" +) + +type BigEndianBuf [8]byte + +func (b BigEndianBuf) Int16(n int16) []byte { + buf := b[0:2] + binary.BigEndian.PutUint16(buf, uint16(n)) + return buf +} + +func (b BigEndianBuf) Uint16(n uint16) []byte { + buf := b[0:2] + binary.BigEndian.PutUint16(buf, n) + return buf +} + +func (b BigEndianBuf) Int32(n int32) []byte { + buf := b[0:4] + binary.BigEndian.PutUint32(buf, uint32(n)) + return buf +} + +func (b BigEndianBuf) Uint32(n uint32) []byte { + buf := b[0:4] + binary.BigEndian.PutUint32(buf, n) + return buf +} + +func (b BigEndianBuf) Int64(n int64) []byte { + buf := b[0:8] + binary.BigEndian.PutUint64(buf, uint64(n)) + return buf +} diff --git a/vendor/github.com/jackc/pgx/v5/pgproto3/bind.go b/vendor/github.com/jackc/pgx/v5/pgproto3/bind.go new file mode 100644 index 0000000000..fb56e4dca7 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgproto3/bind.go @@ -0,0 +1,223 @@ +package pgproto3 + +import ( + "bytes" + "encoding/binary" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "math" + + "github.com/jackc/pgx/v5/internal/pgio" +) + +type Bind struct { + DestinationPortal string + PreparedStatement string + ParameterFormatCodes []int16 + Parameters [][]byte + ResultFormatCodes []int16 +} + +// Frontend identifies this message as sendable by a PostgreSQL frontend. +func (*Bind) Frontend() {} + +// Decode decodes src into dst. src must contain the complete message with the exception of the initial 1 byte message +// type identifier and 4 byte message length. +func (dst *Bind) Decode(src []byte) error { + *dst = Bind{} + + idx := bytes.IndexByte(src, 0) + if idx < 0 { + return &invalidMessageFormatErr{messageType: "Bind"} + } + dst.DestinationPortal = string(src[:idx]) + rp := idx + 1 + + idx = bytes.IndexByte(src[rp:], 0) + if idx < 0 { + return &invalidMessageFormatErr{messageType: "Bind"} + } + dst.PreparedStatement = string(src[rp : rp+idx]) + rp += idx + 1 + + if len(src[rp:]) < 2 { + return &invalidMessageFormatErr{messageType: "Bind"} + } + parameterFormatCodeCount := int(binary.BigEndian.Uint16(src[rp:])) + rp += 2 + + if parameterFormatCodeCount > 0 { + dst.ParameterFormatCodes = make([]int16, parameterFormatCodeCount) + + if len(src[rp:]) < len(dst.ParameterFormatCodes)*2 { + return &invalidMessageFormatErr{messageType: "Bind"} + } + for i := range parameterFormatCodeCount { + dst.ParameterFormatCodes[i] = int16(binary.BigEndian.Uint16(src[rp:])) + rp += 2 + } + } + + if len(src[rp:]) < 2 { + return &invalidMessageFormatErr{messageType: "Bind"} + } + parameterCount := int(binary.BigEndian.Uint16(src[rp:])) + rp += 2 + + if parameterCount > 0 { + dst.Parameters = make([][]byte, parameterCount) + + for i := range parameterCount { + if len(src[rp:]) < 4 { + return &invalidMessageFormatErr{messageType: "Bind"} + } + + msgSize := int(int32(binary.BigEndian.Uint32(src[rp:]))) + rp += 4 + + // null + if msgSize == -1 { + continue + } + + if msgSize < 0 || len(src[rp:]) < msgSize { + return &invalidMessageFormatErr{messageType: "Bind"} + } + + dst.Parameters[i] = src[rp : rp+msgSize] + rp += msgSize + } + } + + if len(src[rp:]) < 2 { + return &invalidMessageFormatErr{messageType: "Bind"} + } + resultFormatCodeCount := int(binary.BigEndian.Uint16(src[rp:])) + rp += 2 + + dst.ResultFormatCodes = make([]int16, resultFormatCodeCount) + if len(src[rp:]) < len(dst.ResultFormatCodes)*2 { + return &invalidMessageFormatErr{messageType: "Bind"} + } + for i := range resultFormatCodeCount { + dst.ResultFormatCodes[i] = int16(binary.BigEndian.Uint16(src[rp:])) + rp += 2 + } + + return nil +} + +// Encode encodes src into dst. dst will include the 1 byte message type identifier and the 4 byte message length. +func (src *Bind) Encode(dst []byte) ([]byte, error) { + dst, sp := beginMessage(dst, 'B') + + dst = append(dst, src.DestinationPortal...) + dst = append(dst, 0) + dst = append(dst, src.PreparedStatement...) + dst = append(dst, 0) + + if len(src.ParameterFormatCodes) > math.MaxUint16 { + return nil, errors.New("too many parameter format codes") + } + dst = pgio.AppendUint16(dst, uint16(len(src.ParameterFormatCodes))) + for _, fc := range src.ParameterFormatCodes { + dst = pgio.AppendInt16(dst, fc) + } + + if len(src.Parameters) > math.MaxUint16 { + return nil, errors.New("too many parameters") + } + dst = pgio.AppendUint16(dst, uint16(len(src.Parameters))) + for _, p := range src.Parameters { + if p == nil { + dst = pgio.AppendInt32(dst, -1) + continue + } + + dst = pgio.AppendInt32(dst, int32(len(p))) + dst = append(dst, p...) + } + + if len(src.ResultFormatCodes) > math.MaxUint16 { + return nil, errors.New("too many result format codes") + } + dst = pgio.AppendUint16(dst, uint16(len(src.ResultFormatCodes))) + for _, fc := range src.ResultFormatCodes { + dst = pgio.AppendInt16(dst, fc) + } + + return finishMessage(dst, sp) +} + +// MarshalJSON implements encoding/json.Marshaler. +func (src Bind) MarshalJSON() ([]byte, error) { + formattedParameters := make([]map[string]string, len(src.Parameters)) + for i, p := range src.Parameters { + if p == nil { + continue + } + + textFormat := true + if len(src.ParameterFormatCodes) == 1 { + textFormat = src.ParameterFormatCodes[0] == 0 + } else if len(src.ParameterFormatCodes) > 1 { + textFormat = src.ParameterFormatCodes[i] == 0 + } + + if textFormat { + formattedParameters[i] = map[string]string{"text": string(p)} + } else { + formattedParameters[i] = map[string]string{"binary": hex.EncodeToString(p)} + } + } + + return json.Marshal(struct { + Type string + DestinationPortal string + PreparedStatement string + ParameterFormatCodes []int16 + Parameters []map[string]string + ResultFormatCodes []int16 + }{ + Type: "Bind", + DestinationPortal: src.DestinationPortal, + PreparedStatement: src.PreparedStatement, + ParameterFormatCodes: src.ParameterFormatCodes, + Parameters: formattedParameters, + ResultFormatCodes: src.ResultFormatCodes, + }) +} + +// UnmarshalJSON implements encoding/json.Unmarshaler. +func (dst *Bind) UnmarshalJSON(data []byte) error { + // Ignore null, like in the main JSON package. + if string(data) == "null" { + return nil + } + + var msg struct { + DestinationPortal string + PreparedStatement string + ParameterFormatCodes []int16 + Parameters []map[string]string + ResultFormatCodes []int16 + } + err := json.Unmarshal(data, &msg) + if err != nil { + return err + } + dst.DestinationPortal = msg.DestinationPortal + dst.PreparedStatement = msg.PreparedStatement + dst.ParameterFormatCodes = msg.ParameterFormatCodes + dst.Parameters = make([][]byte, len(msg.Parameters)) + dst.ResultFormatCodes = msg.ResultFormatCodes + for n, parameter := range msg.Parameters { + dst.Parameters[n], err = getValueFromJSON(parameter) + if err != nil { + return fmt.Errorf("cannot get param %d: %w", n, err) + } + } + return nil +} diff --git a/vendor/github.com/jackc/pgx/v5/pgproto3/bind_complete.go b/vendor/github.com/jackc/pgx/v5/pgproto3/bind_complete.go new file mode 100644 index 0000000000..bacf30d88a --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgproto3/bind_complete.go @@ -0,0 +1,34 @@ +package pgproto3 + +import ( + "encoding/json" +) + +type BindComplete struct{} + +// Backend identifies this message as sendable by the PostgreSQL backend. +func (*BindComplete) Backend() {} + +// Decode decodes src into dst. src must contain the complete message with the exception of the initial 1 byte message +// type identifier and 4 byte message length. +func (dst *BindComplete) Decode(src []byte) error { + if len(src) != 0 { + return &invalidMessageLenErr{messageType: "BindComplete", expectedLen: 0, actualLen: len(src)} + } + + return nil +} + +// Encode encodes src into dst. dst will include the 1 byte message type identifier and the 4 byte message length. +func (src *BindComplete) Encode(dst []byte) ([]byte, error) { + return append(dst, '2', 0, 0, 0, 4), nil +} + +// MarshalJSON implements encoding/json.Marshaler. +func (src BindComplete) MarshalJSON() ([]byte, error) { + return json.Marshal(struct { + Type string + }{ + Type: "BindComplete", + }) +} diff --git a/vendor/github.com/jackc/pgx/v5/pgproto3/cancel_request.go b/vendor/github.com/jackc/pgx/v5/pgproto3/cancel_request.go new file mode 100644 index 0000000000..63ebe5c47f --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgproto3/cancel_request.go @@ -0,0 +1,85 @@ +package pgproto3 + +import ( + "encoding/binary" + "encoding/hex" + "encoding/json" + "errors" + + "github.com/jackc/pgx/v5/internal/pgio" +) + +const cancelRequestCode = 80877102 + +type CancelRequest struct { + ProcessID uint32 + SecretKey []byte +} + +// Frontend identifies this message as sendable by a PostgreSQL frontend. +func (*CancelRequest) Frontend() {} + +func (dst *CancelRequest) Decode(src []byte) error { + if len(src) < 12 { + return errors.New("cancel request too short") + } + if len(src) > 264 { + return errors.New("cancel request too long") + } + + requestCode := binary.BigEndian.Uint32(src) + if requestCode != cancelRequestCode { + return errors.New("bad cancel request code") + } + + dst.ProcessID = binary.BigEndian.Uint32(src[4:]) + dst.SecretKey = make([]byte, len(src)-8) + copy(dst.SecretKey, src[8:]) + + return nil +} + +// Encode encodes src into dst. dst will include the 4 byte message length. +func (src *CancelRequest) Encode(dst []byte) ([]byte, error) { + if len(src.SecretKey) > 256 { + return nil, errors.New("secret key too long") + } + msgLen := int32(12 + len(src.SecretKey)) + dst = pgio.AppendInt32(dst, msgLen) + dst = pgio.AppendInt32(dst, cancelRequestCode) + dst = pgio.AppendUint32(dst, src.ProcessID) + dst = append(dst, src.SecretKey...) + return dst, nil +} + +// MarshalJSON implements encoding/json.Marshaler. +func (src CancelRequest) MarshalJSON() ([]byte, error) { + return json.Marshal(struct { + Type string + ProcessID uint32 + SecretKey string + }{ + Type: "CancelRequest", + ProcessID: src.ProcessID, + SecretKey: hex.EncodeToString(src.SecretKey), + }) +} + +// UnmarshalJSON implements encoding/json.Unmarshaler. +func (dst *CancelRequest) UnmarshalJSON(data []byte) error { + var msg struct { + ProcessID uint32 + SecretKey string + } + if err := json.Unmarshal(data, &msg); err != nil { + return err + } + + dst.ProcessID = msg.ProcessID + secretKey, err := hex.DecodeString(msg.SecretKey) + if err != nil { + return err + } + dst.SecretKey = secretKey + return nil +} diff --git a/vendor/github.com/jackc/pgx/v5/pgproto3/chunkreader.go b/vendor/github.com/jackc/pgx/v5/pgproto3/chunkreader.go new file mode 100644 index 0000000000..fc0fa61e9c --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgproto3/chunkreader.go @@ -0,0 +1,90 @@ +package pgproto3 + +import ( + "io" + + "github.com/jackc/pgx/v5/internal/iobufpool" +) + +// chunkReader is a io.Reader wrapper that minimizes IO reads and memory allocations. It allocates memory in chunks and +// will read as much as will fit in the current buffer in a single call regardless of how large a read is actually +// requested. The memory returned via Next is only valid until the next call to Next. +// +// This is roughly equivalent to a bufio.Reader that only uses Peek and Discard to never copy bytes. +type chunkReader struct { + r io.Reader + + buf *[]byte + rp, wp int // buf read position and write position + + minBufSize int +} + +// newChunkReader creates and returns a new chunkReader for r with default configuration. If minBufSize is <= 0 it uses +// a default value. +func newChunkReader(r io.Reader, minBufSize int) *chunkReader { + if minBufSize <= 0 { + // By historical reasons Postgres currently has 8KB send buffer inside, + // so here we want to have at least the same size buffer. + // @see https://github.com/postgres/postgres/blob/249d64999615802752940e017ee5166e726bc7cd/src/backend/libpq/pqcomm.c#L134 + // @see https://www.postgresql.org/message-id/0cdc5485-cb3c-5e16-4a46-e3b2f7a41322%40ya.ru + // + // In addition, testing has found no benefit of any larger buffer. + minBufSize = 8192 + } + + return &chunkReader{ + r: r, + minBufSize: minBufSize, + buf: iobufpool.Get(minBufSize), + } +} + +// Next returns buf filled with the next n bytes. buf is only valid until next call of Next. If an error occurs, buf +// will be nil. +func (r *chunkReader) Next(n int) (buf []byte, err error) { + // Reset the buffer if it is empty + if r.rp == r.wp { + if len(*r.buf) != r.minBufSize { + iobufpool.Put(r.buf) + r.buf = iobufpool.Get(r.minBufSize) + } + r.rp = 0 + r.wp = 0 + } + + // n bytes already in buf + if (r.wp - r.rp) >= n { + buf = (*r.buf)[r.rp : r.rp+n : r.rp+n] + r.rp += n + return buf, err + } + + // buf is smaller than requested number of bytes + if len(*r.buf) < n { + bigBuf := iobufpool.Get(n) + r.wp = copy((*bigBuf), (*r.buf)[r.rp:r.wp]) + r.rp = 0 + iobufpool.Put(r.buf) + r.buf = bigBuf + } + + // buf is large enough, but need to shift filled area to start to make enough contiguous space + minReadCount := n - (r.wp - r.rp) + if (len(*r.buf) - r.wp) < minReadCount { + r.wp = copy((*r.buf), (*r.buf)[r.rp:r.wp]) + r.rp = 0 + } + + // Read at least the required number of bytes from the underlying io.Reader + readBytesCount, err := io.ReadAtLeast(r.r, (*r.buf)[r.wp:], minReadCount) + r.wp += readBytesCount + // fmt.Println("read", n) + if err != nil { + return nil, err + } + + buf = (*r.buf)[r.rp : r.rp+n : r.rp+n] + r.rp += n + return buf, nil +} diff --git a/vendor/github.com/jackc/pgx/v5/pgproto3/close.go b/vendor/github.com/jackc/pgx/v5/pgproto3/close.go new file mode 100644 index 0000000000..0e7e049522 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgproto3/close.go @@ -0,0 +1,81 @@ +package pgproto3 + +import ( + "bytes" + "encoding/json" + "errors" +) + +type Close struct { + ObjectType byte // 'S' = prepared statement, 'P' = portal + Name string +} + +// Frontend identifies this message as sendable by a PostgreSQL frontend. +func (*Close) Frontend() {} + +// Decode decodes src into dst. src must contain the complete message with the exception of the initial 1 byte message +// type identifier and 4 byte message length. +func (dst *Close) Decode(src []byte) error { + if len(src) < 2 { + return &invalidMessageFormatErr{messageType: "Close"} + } + + dst.ObjectType = src[0] + rp := 1 + + idx := bytes.IndexByte(src[rp:], 0) + if idx != len(src[rp:])-1 { + return &invalidMessageFormatErr{messageType: "Close"} + } + + dst.Name = string(src[rp : len(src)-1]) + + return nil +} + +// Encode encodes src into dst. dst will include the 1 byte message type identifier and the 4 byte message length. +func (src *Close) Encode(dst []byte) ([]byte, error) { + dst, sp := beginMessage(dst, 'C') + dst = append(dst, src.ObjectType) + dst = append(dst, src.Name...) + dst = append(dst, 0) + return finishMessage(dst, sp) +} + +// MarshalJSON implements encoding/json.Marshaler. +func (src Close) MarshalJSON() ([]byte, error) { + return json.Marshal(struct { + Type string + ObjectType string + Name string + }{ + Type: "Close", + ObjectType: string(src.ObjectType), + Name: src.Name, + }) +} + +// UnmarshalJSON implements encoding/json.Unmarshaler. +func (dst *Close) UnmarshalJSON(data []byte) error { + // Ignore null, like in the main JSON package. + if string(data) == "null" { + return nil + } + + var msg struct { + ObjectType string + Name string + } + if err := json.Unmarshal(data, &msg); err != nil { + return err + } + + if len(msg.ObjectType) != 1 { + return errors.New("invalid length for Close.ObjectType") + } + + dst.ObjectType = msg.ObjectType[0] + dst.Name = msg.Name + return nil +} diff --git a/vendor/github.com/jackc/pgx/v5/pgproto3/close_complete.go b/vendor/github.com/jackc/pgx/v5/pgproto3/close_complete.go new file mode 100644 index 0000000000..833f7a12c8 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgproto3/close_complete.go @@ -0,0 +1,34 @@ +package pgproto3 + +import ( + "encoding/json" +) + +type CloseComplete struct{} + +// Backend identifies this message as sendable by the PostgreSQL backend. +func (*CloseComplete) Backend() {} + +// Decode decodes src into dst. src must contain the complete message with the exception of the initial 1 byte message +// type identifier and 4 byte message length. +func (dst *CloseComplete) Decode(src []byte) error { + if len(src) != 0 { + return &invalidMessageLenErr{messageType: "CloseComplete", expectedLen: 0, actualLen: len(src)} + } + + return nil +} + +// Encode encodes src into dst. dst will include the 1 byte message type identifier and the 4 byte message length. +func (src *CloseComplete) Encode(dst []byte) ([]byte, error) { + return append(dst, '3', 0, 0, 0, 4), nil +} + +// MarshalJSON implements encoding/json.Marshaler. +func (src CloseComplete) MarshalJSON() ([]byte, error) { + return json.Marshal(struct { + Type string + }{ + Type: "CloseComplete", + }) +} diff --git a/vendor/github.com/jackc/pgx/v5/pgproto3/command_complete.go b/vendor/github.com/jackc/pgx/v5/pgproto3/command_complete.go new file mode 100644 index 0000000000..eba70947d9 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgproto3/command_complete.go @@ -0,0 +1,66 @@ +package pgproto3 + +import ( + "bytes" + "encoding/json" +) + +type CommandComplete struct { + CommandTag []byte +} + +// Backend identifies this message as sendable by the PostgreSQL backend. +func (*CommandComplete) Backend() {} + +// Decode decodes src into dst. src must contain the complete message with the exception of the initial 1 byte message +// type identifier and 4 byte message length. +func (dst *CommandComplete) Decode(src []byte) error { + idx := bytes.IndexByte(src, 0) + if idx == -1 { + return &invalidMessageFormatErr{messageType: "CommandComplete", details: "unterminated string"} + } + if idx != len(src)-1 { + return &invalidMessageFormatErr{messageType: "CommandComplete", details: "string terminated too early"} + } + + dst.CommandTag = src[:idx] + + return nil +} + +// Encode encodes src into dst. dst will include the 1 byte message type identifier and the 4 byte message length. +func (src *CommandComplete) Encode(dst []byte) ([]byte, error) { + dst, sp := beginMessage(dst, 'C') + dst = append(dst, src.CommandTag...) + dst = append(dst, 0) + return finishMessage(dst, sp) +} + +// MarshalJSON implements encoding/json.Marshaler. +func (src CommandComplete) MarshalJSON() ([]byte, error) { + return json.Marshal(struct { + Type string + CommandTag string + }{ + Type: "CommandComplete", + CommandTag: string(src.CommandTag), + }) +} + +// UnmarshalJSON implements encoding/json.Unmarshaler. +func (dst *CommandComplete) UnmarshalJSON(data []byte) error { + // Ignore null, like in the main JSON package. + if string(data) == "null" { + return nil + } + + var msg struct { + CommandTag string + } + if err := json.Unmarshal(data, &msg); err != nil { + return err + } + + dst.CommandTag = []byte(msg.CommandTag) + return nil +} diff --git a/vendor/github.com/jackc/pgx/v5/pgproto3/copy_both_response.go b/vendor/github.com/jackc/pgx/v5/pgproto3/copy_both_response.go new file mode 100644 index 0000000000..e2a402f9a8 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgproto3/copy_both_response.go @@ -0,0 +1,95 @@ +package pgproto3 + +import ( + "bytes" + "encoding/binary" + "encoding/json" + "errors" + "math" + + "github.com/jackc/pgx/v5/internal/pgio" +) + +type CopyBothResponse struct { + OverallFormat byte + ColumnFormatCodes []uint16 +} + +// Backend identifies this message as sendable by the PostgreSQL backend. +func (*CopyBothResponse) Backend() {} + +// Decode decodes src into dst. src must contain the complete message with the exception of the initial 1 byte message +// type identifier and 4 byte message length. +func (dst *CopyBothResponse) Decode(src []byte) error { + buf := bytes.NewBuffer(src) + + if buf.Len() < 3 { + return &invalidMessageFormatErr{messageType: "CopyBothResponse"} + } + + overallFormat := buf.Next(1)[0] + + columnCount := int(binary.BigEndian.Uint16(buf.Next(2))) + if buf.Len() != columnCount*2 { + return &invalidMessageFormatErr{messageType: "CopyBothResponse"} + } + + columnFormatCodes := make([]uint16, columnCount) + for i := range columnCount { + columnFormatCodes[i] = binary.BigEndian.Uint16(buf.Next(2)) + } + + *dst = CopyBothResponse{OverallFormat: overallFormat, ColumnFormatCodes: columnFormatCodes} + + return nil +} + +// Encode encodes src into dst. dst will include the 1 byte message type identifier and the 4 byte message length. +func (src *CopyBothResponse) Encode(dst []byte) ([]byte, error) { + dst, sp := beginMessage(dst, 'W') + dst = append(dst, src.OverallFormat) + if len(src.ColumnFormatCodes) > math.MaxUint16 { + return nil, errors.New("too many column format codes") + } + dst = pgio.AppendUint16(dst, uint16(len(src.ColumnFormatCodes))) + for _, fc := range src.ColumnFormatCodes { + dst = pgio.AppendUint16(dst, fc) + } + + return finishMessage(dst, sp) +} + +// MarshalJSON implements encoding/json.Marshaler. +func (src CopyBothResponse) MarshalJSON() ([]byte, error) { + return json.Marshal(struct { + Type string + ColumnFormatCodes []uint16 + }{ + Type: "CopyBothResponse", + ColumnFormatCodes: src.ColumnFormatCodes, + }) +} + +// UnmarshalJSON implements encoding/json.Unmarshaler. +func (dst *CopyBothResponse) UnmarshalJSON(data []byte) error { + // Ignore null, like in the main JSON package. + if string(data) == "null" { + return nil + } + + var msg struct { + OverallFormat string + ColumnFormatCodes []uint16 + } + if err := json.Unmarshal(data, &msg); err != nil { + return err + } + + if len(msg.OverallFormat) != 1 { + return errors.New("invalid length for CopyBothResponse.OverallFormat") + } + + dst.OverallFormat = msg.OverallFormat[0] + dst.ColumnFormatCodes = msg.ColumnFormatCodes + return nil +} diff --git a/vendor/github.com/jackc/pgx/v5/pgproto3/copy_data.go b/vendor/github.com/jackc/pgx/v5/pgproto3/copy_data.go new file mode 100644 index 0000000000..d72e1f353b --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgproto3/copy_data.go @@ -0,0 +1,67 @@ +package pgproto3 + +import ( + "encoding/hex" + "encoding/json" +) + +type CopyData struct { + Data []byte +} + +// Backend identifies this message as sendable by the PostgreSQL backend. +func (*CopyData) Backend() {} + +// Frontend identifies this message as sendable by a PostgreSQL frontend. +func (*CopyData) Frontend() {} + +// Decode decodes src into dst. src must contain the complete message with the exception of the initial 1 byte message +// type identifier and 4 byte message length. +func (dst *CopyData) Decode(src []byte) error { + dst.Data = src + return nil +} + +// Encode encodes src into dst. dst will include the 1 byte message type identifier and the 4 byte message length. +func (src *CopyData) Encode(dst []byte) ([]byte, error) { + dst, sp := beginMessage(dst, 'd') + dst = append(dst, src.Data...) + return finishMessage(dst, sp) +} + +// MarshalJSON implements encoding/json.Marshaler. +func (src CopyData) MarshalJSON() ([]byte, error) { + return json.Marshal(struct { + Type string + Data string + }{ + Type: "CopyData", + Data: hex.EncodeToString(src.Data), + }) +} + +// UnmarshalJSON implements encoding/json.Unmarshaler. +func (dst *CopyData) UnmarshalJSON(data []byte) error { + // Ignore null, like in the main JSON package. + if string(data) == "null" { + return nil + } + + var msg struct { + Data string + } + if err := json.Unmarshal(data, &msg); err != nil { + return err + } + + if msg.Data == "" { + dst.Data = []byte{} + return nil + } + b, err := hex.DecodeString(msg.Data) + if err != nil { + return err + } + dst.Data = b + return nil +} diff --git a/vendor/github.com/jackc/pgx/v5/pgproto3/copy_done.go b/vendor/github.com/jackc/pgx/v5/pgproto3/copy_done.go new file mode 100644 index 0000000000..c3421a9b5b --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgproto3/copy_done.go @@ -0,0 +1,37 @@ +package pgproto3 + +import ( + "encoding/json" +) + +type CopyDone struct{} + +// Backend identifies this message as sendable by the PostgreSQL backend. +func (*CopyDone) Backend() {} + +// Frontend identifies this message as sendable by a PostgreSQL frontend. +func (*CopyDone) Frontend() {} + +// Decode decodes src into dst. src must contain the complete message with the exception of the initial 1 byte message +// type identifier and 4 byte message length. +func (dst *CopyDone) Decode(src []byte) error { + if len(src) != 0 { + return &invalidMessageLenErr{messageType: "CopyDone", expectedLen: 0, actualLen: len(src)} + } + + return nil +} + +// Encode encodes src into dst. dst will include the 1 byte message type identifier and the 4 byte message length. +func (src *CopyDone) Encode(dst []byte) ([]byte, error) { + return append(dst, 'c', 0, 0, 0, 4), nil +} + +// MarshalJSON implements encoding/json.Marshaler. +func (src CopyDone) MarshalJSON() ([]byte, error) { + return json.Marshal(struct { + Type string + }{ + Type: "CopyDone", + }) +} diff --git a/vendor/github.com/jackc/pgx/v5/pgproto3/copy_fail.go b/vendor/github.com/jackc/pgx/v5/pgproto3/copy_fail.go new file mode 100644 index 0000000000..f8a00b8b74 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgproto3/copy_fail.go @@ -0,0 +1,49 @@ +package pgproto3 + +import ( + "bytes" + "encoding/json" +) + +type CopyFail struct { + Message string +} + +// Frontend identifies this message as sendable by a PostgreSQL frontend. +func (*CopyFail) Frontend() {} + +// Decode decodes src into dst. src must contain the complete message with the exception of the initial 1 byte message +// type identifier and 4 byte message length. +func (dst *CopyFail) Decode(src []byte) error { + if len(src) == 0 { + return &invalidMessageFormatErr{messageType: "CopyFail"} + } + + idx := bytes.IndexByte(src, 0) + if idx != len(src)-1 { + return &invalidMessageFormatErr{messageType: "CopyFail"} + } + + dst.Message = string(src[:idx]) + + return nil +} + +// Encode encodes src into dst. dst will include the 1 byte message type identifier and the 4 byte message length. +func (src *CopyFail) Encode(dst []byte) ([]byte, error) { + dst, sp := beginMessage(dst, 'f') + dst = append(dst, src.Message...) + dst = append(dst, 0) + return finishMessage(dst, sp) +} + +// MarshalJSON implements encoding/json.Marshaler. +func (src CopyFail) MarshalJSON() ([]byte, error) { + return json.Marshal(struct { + Type string + Message string + }{ + Type: "CopyFail", + Message: src.Message, + }) +} diff --git a/vendor/github.com/jackc/pgx/v5/pgproto3/copy_in_response.go b/vendor/github.com/jackc/pgx/v5/pgproto3/copy_in_response.go new file mode 100644 index 0000000000..0633935b96 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgproto3/copy_in_response.go @@ -0,0 +1,96 @@ +package pgproto3 + +import ( + "bytes" + "encoding/binary" + "encoding/json" + "errors" + "math" + + "github.com/jackc/pgx/v5/internal/pgio" +) + +type CopyInResponse struct { + OverallFormat byte + ColumnFormatCodes []uint16 +} + +// Backend identifies this message as sendable by the PostgreSQL backend. +func (*CopyInResponse) Backend() {} + +// Decode decodes src into dst. src must contain the complete message with the exception of the initial 1 byte message +// type identifier and 4 byte message length. +func (dst *CopyInResponse) Decode(src []byte) error { + buf := bytes.NewBuffer(src) + + if buf.Len() < 3 { + return &invalidMessageFormatErr{messageType: "CopyInResponse"} + } + + overallFormat := buf.Next(1)[0] + + columnCount := int(binary.BigEndian.Uint16(buf.Next(2))) + if buf.Len() != columnCount*2 { + return &invalidMessageFormatErr{messageType: "CopyInResponse"} + } + + columnFormatCodes := make([]uint16, columnCount) + for i := range columnCount { + columnFormatCodes[i] = binary.BigEndian.Uint16(buf.Next(2)) + } + + *dst = CopyInResponse{OverallFormat: overallFormat, ColumnFormatCodes: columnFormatCodes} + + return nil +} + +// Encode encodes src into dst. dst will include the 1 byte message type identifier and the 4 byte message length. +func (src *CopyInResponse) Encode(dst []byte) ([]byte, error) { + dst, sp := beginMessage(dst, 'G') + + dst = append(dst, src.OverallFormat) + if len(src.ColumnFormatCodes) > math.MaxUint16 { + return nil, errors.New("too many column format codes") + } + dst = pgio.AppendUint16(dst, uint16(len(src.ColumnFormatCodes))) + for _, fc := range src.ColumnFormatCodes { + dst = pgio.AppendUint16(dst, fc) + } + + return finishMessage(dst, sp) +} + +// MarshalJSON implements encoding/json.Marshaler. +func (src CopyInResponse) MarshalJSON() ([]byte, error) { + return json.Marshal(struct { + Type string + ColumnFormatCodes []uint16 + }{ + Type: "CopyInResponse", + ColumnFormatCodes: src.ColumnFormatCodes, + }) +} + +// UnmarshalJSON implements encoding/json.Unmarshaler. +func (dst *CopyInResponse) UnmarshalJSON(data []byte) error { + // Ignore null, like in the main JSON package. + if string(data) == "null" { + return nil + } + + var msg struct { + OverallFormat string + ColumnFormatCodes []uint16 + } + if err := json.Unmarshal(data, &msg); err != nil { + return err + } + + if len(msg.OverallFormat) != 1 { + return errors.New("invalid length for CopyInResponse.OverallFormat") + } + + dst.OverallFormat = msg.OverallFormat[0] + dst.ColumnFormatCodes = msg.ColumnFormatCodes + return nil +} diff --git a/vendor/github.com/jackc/pgx/v5/pgproto3/copy_out_response.go b/vendor/github.com/jackc/pgx/v5/pgproto3/copy_out_response.go new file mode 100644 index 0000000000..006864ac8c --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgproto3/copy_out_response.go @@ -0,0 +1,96 @@ +package pgproto3 + +import ( + "bytes" + "encoding/binary" + "encoding/json" + "errors" + "math" + + "github.com/jackc/pgx/v5/internal/pgio" +) + +type CopyOutResponse struct { + OverallFormat byte + ColumnFormatCodes []uint16 +} + +func (*CopyOutResponse) Backend() {} + +// Decode decodes src into dst. src must contain the complete message with the exception of the initial 1 byte message +// type identifier and 4 byte message length. +func (dst *CopyOutResponse) Decode(src []byte) error { + buf := bytes.NewBuffer(src) + + if buf.Len() < 3 { + return &invalidMessageFormatErr{messageType: "CopyOutResponse"} + } + + overallFormat := buf.Next(1)[0] + + columnCount := int(binary.BigEndian.Uint16(buf.Next(2))) + if buf.Len() != columnCount*2 { + return &invalidMessageFormatErr{messageType: "CopyOutResponse"} + } + + columnFormatCodes := make([]uint16, columnCount) + for i := range columnCount { + columnFormatCodes[i] = binary.BigEndian.Uint16(buf.Next(2)) + } + + *dst = CopyOutResponse{OverallFormat: overallFormat, ColumnFormatCodes: columnFormatCodes} + + return nil +} + +// Encode encodes src into dst. dst will include the 1 byte message type identifier and the 4 byte message length. +func (src *CopyOutResponse) Encode(dst []byte) ([]byte, error) { + dst, sp := beginMessage(dst, 'H') + + dst = append(dst, src.OverallFormat) + + if len(src.ColumnFormatCodes) > math.MaxUint16 { + return nil, errors.New("too many column format codes") + } + dst = pgio.AppendUint16(dst, uint16(len(src.ColumnFormatCodes))) + for _, fc := range src.ColumnFormatCodes { + dst = pgio.AppendUint16(dst, fc) + } + + return finishMessage(dst, sp) +} + +// MarshalJSON implements encoding/json.Marshaler. +func (src CopyOutResponse) MarshalJSON() ([]byte, error) { + return json.Marshal(struct { + Type string + ColumnFormatCodes []uint16 + }{ + Type: "CopyOutResponse", + ColumnFormatCodes: src.ColumnFormatCodes, + }) +} + +// UnmarshalJSON implements encoding/json.Unmarshaler. +func (dst *CopyOutResponse) UnmarshalJSON(data []byte) error { + // Ignore null, like in the main JSON package. + if string(data) == "null" { + return nil + } + + var msg struct { + OverallFormat string + ColumnFormatCodes []uint16 + } + if err := json.Unmarshal(data, &msg); err != nil { + return err + } + + if len(msg.OverallFormat) != 1 { + return errors.New("invalid length for CopyOutResponse.OverallFormat") + } + + dst.OverallFormat = msg.OverallFormat[0] + dst.ColumnFormatCodes = msg.ColumnFormatCodes + return nil +} diff --git a/vendor/github.com/jackc/pgx/v5/pgproto3/data_row.go b/vendor/github.com/jackc/pgx/v5/pgproto3/data_row.go new file mode 100644 index 0000000000..54418d58ca --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgproto3/data_row.go @@ -0,0 +1,140 @@ +package pgproto3 + +import ( + "encoding/binary" + "encoding/hex" + "encoding/json" + "errors" + "math" + + "github.com/jackc/pgx/v5/internal/pgio" +) + +type DataRow struct { + Values [][]byte +} + +// Backend identifies this message as sendable by the PostgreSQL backend. +func (*DataRow) Backend() {} + +// Decode decodes src into dst. src must contain the complete message with the exception of the initial 1 byte message +// type identifier and 4 byte message length. +func (dst *DataRow) Decode(src []byte) error { + if len(src) < 2 { + return &invalidMessageFormatErr{messageType: "DataRow"} + } + rp := 0 + fieldCount := int(binary.BigEndian.Uint16(src[rp:])) + rp += 2 + + // If the capacity of the values slice is too small OR substantially too + // large reallocate. This is too avoid one row with many columns from + // permanently allocating memory. + if cap(dst.Values) < fieldCount || cap(dst.Values)-fieldCount > 32 { + newCap := max(32, fieldCount) + dst.Values = make([][]byte, fieldCount, newCap) + } else { + dst.Values = dst.Values[:fieldCount] + } + + for i := range fieldCount { + if len(src[rp:]) < 4 { + return &invalidMessageFormatErr{messageType: "DataRow"} + } + + valueLen := int(int32(binary.BigEndian.Uint32(src[rp:]))) + rp += 4 + + // null + if valueLen == -1 { + dst.Values[i] = nil + } else { + if len(src[rp:]) < valueLen || valueLen < 0 { + return &invalidMessageFormatErr{messageType: "DataRow"} + } + + dst.Values[i] = src[rp : rp+valueLen : rp+valueLen] + rp += valueLen + } + } + + return nil +} + +// Encode encodes src into dst. dst will include the 1 byte message type identifier and the 4 byte message length. +func (src *DataRow) Encode(dst []byte) ([]byte, error) { + dst, sp := beginMessage(dst, 'D') + + if len(src.Values) > math.MaxUint16 { + return nil, errors.New("too many values") + } + dst = pgio.AppendUint16(dst, uint16(len(src.Values))) + for _, v := range src.Values { + if v == nil { + dst = pgio.AppendInt32(dst, -1) + continue + } + + dst = pgio.AppendInt32(dst, int32(len(v))) + dst = append(dst, v...) + } + + return finishMessage(dst, sp) +} + +// MarshalJSON implements encoding/json.Marshaler. +func (src DataRow) MarshalJSON() ([]byte, error) { + formattedValues := make([]map[string]string, len(src.Values)) + for i, v := range src.Values { + if v == nil { + continue + } + + var hasNonPrintable bool + for _, b := range v { + if b < 32 { + hasNonPrintable = true + break + } + } + + if hasNonPrintable { + formattedValues[i] = map[string]string{"binary": hex.EncodeToString(v)} + } else { + formattedValues[i] = map[string]string{"text": string(v)} + } + } + + return json.Marshal(struct { + Type string + Values []map[string]string + }{ + Type: "DataRow", + Values: formattedValues, + }) +} + +// UnmarshalJSON implements encoding/json.Unmarshaler. +func (dst *DataRow) UnmarshalJSON(data []byte) error { + // Ignore null, like in the main JSON package. + if string(data) == "null" { + return nil + } + + var msg struct { + Values []map[string]string + } + if err := json.Unmarshal(data, &msg); err != nil { + return err + } + + dst.Values = make([][]byte, len(msg.Values)) + for n, parameter := range msg.Values { + var err error + dst.Values[n], err = getValueFromJSON(parameter) + if err != nil { + return err + } + } + return nil +} diff --git a/vendor/github.com/jackc/pgx/v5/pgproto3/describe.go b/vendor/github.com/jackc/pgx/v5/pgproto3/describe.go new file mode 100644 index 0000000000..0c396f1ba8 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgproto3/describe.go @@ -0,0 +1,80 @@ +package pgproto3 + +import ( + "bytes" + "encoding/json" + "errors" +) + +type Describe struct { + ObjectType byte // 'S' = prepared statement, 'P' = portal + Name string +} + +// Frontend identifies this message as sendable by a PostgreSQL frontend. +func (*Describe) Frontend() {} + +// Decode decodes src into dst. src must contain the complete message with the exception of the initial 1 byte message +// type identifier and 4 byte message length. +func (dst *Describe) Decode(src []byte) error { + if len(src) < 2 { + return &invalidMessageFormatErr{messageType: "Describe"} + } + + dst.ObjectType = src[0] + rp := 1 + + idx := bytes.IndexByte(src[rp:], 0) + if idx != len(src[rp:])-1 { + return &invalidMessageFormatErr{messageType: "Describe"} + } + + dst.Name = string(src[rp : len(src)-1]) + + return nil +} + +// Encode encodes src into dst. dst will include the 1 byte message type identifier and the 4 byte message length. +func (src *Describe) Encode(dst []byte) ([]byte, error) { + dst, sp := beginMessage(dst, 'D') + dst = append(dst, src.ObjectType) + dst = append(dst, src.Name...) + dst = append(dst, 0) + return finishMessage(dst, sp) +} + +// MarshalJSON implements encoding/json.Marshaler. +func (src Describe) MarshalJSON() ([]byte, error) { + return json.Marshal(struct { + Type string + ObjectType string + Name string + }{ + Type: "Describe", + ObjectType: string(src.ObjectType), + Name: src.Name, + }) +} + +// UnmarshalJSON implements encoding/json.Unmarshaler. +func (dst *Describe) UnmarshalJSON(data []byte) error { + // Ignore null, like in the main JSON package. + if string(data) == "null" { + return nil + } + + var msg struct { + ObjectType string + Name string + } + if err := json.Unmarshal(data, &msg); err != nil { + return err + } + if len(msg.ObjectType) != 1 { + return errors.New("invalid length for Describe.ObjectType") + } + + dst.ObjectType = msg.ObjectType[0] + dst.Name = msg.Name + return nil +} diff --git a/vendor/github.com/jackc/pgx/v5/pgproto3/doc.go b/vendor/github.com/jackc/pgx/v5/pgproto3/doc.go new file mode 100644 index 0000000000..0afd18e294 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgproto3/doc.go @@ -0,0 +1,11 @@ +// Package pgproto3 is an encoder and decoder of the PostgreSQL wire protocol version 3. +// +// The primary interfaces are Frontend and Backend. They correspond to a client and server respectively. Messages are +// sent with Send (or a specialized Send variant). Messages are automatically buffered to minimize small writes. Call +// Flush to ensure a message has actually been sent. +// +// The Trace method of Frontend and Backend can be used to examine the wire-level message traffic. It outputs in a +// similar format to the PQtrace function in libpq. +// +// See https://www.postgresql.org/docs/current/protocol-message-formats.html for meanings of the different messages. +package pgproto3 diff --git a/vendor/github.com/jackc/pgx/v5/pgproto3/empty_query_response.go b/vendor/github.com/jackc/pgx/v5/pgproto3/empty_query_response.go new file mode 100644 index 0000000000..cb6cca0735 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgproto3/empty_query_response.go @@ -0,0 +1,34 @@ +package pgproto3 + +import ( + "encoding/json" +) + +type EmptyQueryResponse struct{} + +// Backend identifies this message as sendable by the PostgreSQL backend. +func (*EmptyQueryResponse) Backend() {} + +// Decode decodes src into dst. src must contain the complete message with the exception of the initial 1 byte message +// type identifier and 4 byte message length. +func (dst *EmptyQueryResponse) Decode(src []byte) error { + if len(src) != 0 { + return &invalidMessageLenErr{messageType: "EmptyQueryResponse", expectedLen: 0, actualLen: len(src)} + } + + return nil +} + +// Encode encodes src into dst. dst will include the 1 byte message type identifier and the 4 byte message length. +func (src *EmptyQueryResponse) Encode(dst []byte) ([]byte, error) { + return append(dst, 'I', 0, 0, 0, 4), nil +} + +// MarshalJSON implements encoding/json.Marshaler. +func (src EmptyQueryResponse) MarshalJSON() ([]byte, error) { + return json.Marshal(struct { + Type string + }{ + Type: "EmptyQueryResponse", + }) +} diff --git a/vendor/github.com/jackc/pgx/v5/pgproto3/error_response.go b/vendor/github.com/jackc/pgx/v5/pgproto3/error_response.go new file mode 100644 index 0000000000..6ef9bd0614 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgproto3/error_response.go @@ -0,0 +1,326 @@ +package pgproto3 + +import ( + "bytes" + "encoding/json" + "strconv" +) + +type ErrorResponse struct { + Severity string + SeverityUnlocalized string // only in 9.6 and greater + Code string + Message string + Detail string + Hint string + Position int32 + InternalPosition int32 + InternalQuery string + Where string + SchemaName string + TableName string + ColumnName string + DataTypeName string + ConstraintName string + File string + Line int32 + Routine string + + UnknownFields map[byte]string +} + +// Backend identifies this message as sendable by the PostgreSQL backend. +func (*ErrorResponse) Backend() {} + +// Decode decodes src into dst. src must contain the complete message with the exception of the initial 1 byte message +// type identifier and 4 byte message length. +func (dst *ErrorResponse) Decode(src []byte) error { + *dst = ErrorResponse{} + + buf := bytes.NewBuffer(src) + + for { + k, err := buf.ReadByte() + if err != nil { + return err + } + if k == 0 { + break + } + + vb, err := buf.ReadBytes(0) + if err != nil { + return err + } + v := string(vb[:len(vb)-1]) + + switch k { + case 'S': + dst.Severity = v + case 'V': + dst.SeverityUnlocalized = v + case 'C': + dst.Code = v + case 'M': + dst.Message = v + case 'D': + dst.Detail = v + case 'H': + dst.Hint = v + case 'P': + s := v + n, _ := strconv.ParseInt(s, 10, 32) + dst.Position = int32(n) + case 'p': + s := v + n, _ := strconv.ParseInt(s, 10, 32) + dst.InternalPosition = int32(n) + case 'q': + dst.InternalQuery = v + case 'W': + dst.Where = v + case 's': + dst.SchemaName = v + case 't': + dst.TableName = v + case 'c': + dst.ColumnName = v + case 'd': + dst.DataTypeName = v + case 'n': + dst.ConstraintName = v + case 'F': + dst.File = v + case 'L': + s := v + n, _ := strconv.ParseInt(s, 10, 32) + dst.Line = int32(n) + case 'R': + dst.Routine = v + + default: + if dst.UnknownFields == nil { + dst.UnknownFields = make(map[byte]string) + } + dst.UnknownFields[k] = v + } + } + + return nil +} + +// Encode encodes src into dst. dst will include the 1 byte message type identifier and the 4 byte message length. +func (src *ErrorResponse) Encode(dst []byte) ([]byte, error) { + dst, sp := beginMessage(dst, 'E') + dst = src.appendFields(dst) + return finishMessage(dst, sp) +} + +func (src *ErrorResponse) appendFields(dst []byte) []byte { + if src.Severity != "" { + dst = append(dst, 'S') + dst = append(dst, src.Severity...) + dst = append(dst, 0) + } + if src.SeverityUnlocalized != "" { + dst = append(dst, 'V') + dst = append(dst, src.SeverityUnlocalized...) + dst = append(dst, 0) + } + if src.Code != "" { + dst = append(dst, 'C') + dst = append(dst, src.Code...) + dst = append(dst, 0) + } + if src.Message != "" { + dst = append(dst, 'M') + dst = append(dst, src.Message...) + dst = append(dst, 0) + } + if src.Detail != "" { + dst = append(dst, 'D') + dst = append(dst, src.Detail...) + dst = append(dst, 0) + } + if src.Hint != "" { + dst = append(dst, 'H') + dst = append(dst, src.Hint...) + dst = append(dst, 0) + } + if src.Position != 0 { + dst = append(dst, 'P') + dst = append(dst, strconv.Itoa(int(src.Position))...) + dst = append(dst, 0) + } + if src.InternalPosition != 0 { + dst = append(dst, 'p') + dst = append(dst, strconv.Itoa(int(src.InternalPosition))...) + dst = append(dst, 0) + } + if src.InternalQuery != "" { + dst = append(dst, 'q') + dst = append(dst, src.InternalQuery...) + dst = append(dst, 0) + } + if src.Where != "" { + dst = append(dst, 'W') + dst = append(dst, src.Where...) + dst = append(dst, 0) + } + if src.SchemaName != "" { + dst = append(dst, 's') + dst = append(dst, src.SchemaName...) + dst = append(dst, 0) + } + if src.TableName != "" { + dst = append(dst, 't') + dst = append(dst, src.TableName...) + dst = append(dst, 0) + } + if src.ColumnName != "" { + dst = append(dst, 'c') + dst = append(dst, src.ColumnName...) + dst = append(dst, 0) + } + if src.DataTypeName != "" { + dst = append(dst, 'd') + dst = append(dst, src.DataTypeName...) + dst = append(dst, 0) + } + if src.ConstraintName != "" { + dst = append(dst, 'n') + dst = append(dst, src.ConstraintName...) + dst = append(dst, 0) + } + if src.File != "" { + dst = append(dst, 'F') + dst = append(dst, src.File...) + dst = append(dst, 0) + } + if src.Line != 0 { + dst = append(dst, 'L') + dst = append(dst, strconv.Itoa(int(src.Line))...) + dst = append(dst, 0) + } + if src.Routine != "" { + dst = append(dst, 'R') + dst = append(dst, src.Routine...) + dst = append(dst, 0) + } + + for k, v := range src.UnknownFields { + dst = append(dst, k) + dst = append(dst, v...) + dst = append(dst, 0) + } + + dst = append(dst, 0) + + return dst +} + +// MarshalJSON implements encoding/json.Marshaler. +func (src ErrorResponse) MarshalJSON() ([]byte, error) { + return json.Marshal(struct { + Type string + Severity string + SeverityUnlocalized string // only in 9.6 and greater + Code string + Message string + Detail string + Hint string + Position int32 + InternalPosition int32 + InternalQuery string + Where string + SchemaName string + TableName string + ColumnName string + DataTypeName string + ConstraintName string + File string + Line int32 + Routine string + + UnknownFields map[byte]string + }{ + Type: "ErrorResponse", + Severity: src.Severity, + SeverityUnlocalized: src.SeverityUnlocalized, + Code: src.Code, + Message: src.Message, + Detail: src.Detail, + Hint: src.Hint, + Position: src.Position, + InternalPosition: src.InternalPosition, + InternalQuery: src.InternalQuery, + Where: src.Where, + SchemaName: src.SchemaName, + TableName: src.TableName, + ColumnName: src.ColumnName, + DataTypeName: src.DataTypeName, + ConstraintName: src.ConstraintName, + File: src.File, + Line: src.Line, + Routine: src.Routine, + UnknownFields: src.UnknownFields, + }) +} + +// UnmarshalJSON implements encoding/json.Unmarshaler. +func (dst *ErrorResponse) UnmarshalJSON(data []byte) error { + // Ignore null, like in the main JSON package. + if string(data) == "null" { + return nil + } + + var msg struct { + Type string + Severity string + SeverityUnlocalized string // only in 9.6 and greater + Code string + Message string + Detail string + Hint string + Position int32 + InternalPosition int32 + InternalQuery string + Where string + SchemaName string + TableName string + ColumnName string + DataTypeName string + ConstraintName string + File string + Line int32 + Routine string + + UnknownFields map[byte]string + } + if err := json.Unmarshal(data, &msg); err != nil { + return err + } + + dst.Severity = msg.Severity + dst.SeverityUnlocalized = msg.SeverityUnlocalized + dst.Code = msg.Code + dst.Message = msg.Message + dst.Detail = msg.Detail + dst.Hint = msg.Hint + dst.Position = msg.Position + dst.InternalPosition = msg.InternalPosition + dst.InternalQuery = msg.InternalQuery + dst.Where = msg.Where + dst.SchemaName = msg.SchemaName + dst.TableName = msg.TableName + dst.ColumnName = msg.ColumnName + dst.DataTypeName = msg.DataTypeName + dst.ConstraintName = msg.ConstraintName + dst.File = msg.File + dst.Line = msg.Line + dst.Routine = msg.Routine + + dst.UnknownFields = msg.UnknownFields + + return nil +} diff --git a/vendor/github.com/jackc/pgx/v5/pgproto3/execute.go b/vendor/github.com/jackc/pgx/v5/pgproto3/execute.go new file mode 100644 index 0000000000..31bc714d1a --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgproto3/execute.go @@ -0,0 +1,58 @@ +package pgproto3 + +import ( + "bytes" + "encoding/binary" + "encoding/json" + + "github.com/jackc/pgx/v5/internal/pgio" +) + +type Execute struct { + Portal string + MaxRows uint32 +} + +// Frontend identifies this message as sendable by a PostgreSQL frontend. +func (*Execute) Frontend() {} + +// Decode decodes src into dst. src must contain the complete message with the exception of the initial 1 byte message +// type identifier and 4 byte message length. +func (dst *Execute) Decode(src []byte) error { + buf := bytes.NewBuffer(src) + + b, err := buf.ReadBytes(0) + if err != nil { + return err + } + dst.Portal = string(b[:len(b)-1]) + + if buf.Len() < 4 { + return &invalidMessageFormatErr{messageType: "Execute"} + } + dst.MaxRows = binary.BigEndian.Uint32(buf.Next(4)) + + return nil +} + +// Encode encodes src into dst. dst will include the 1 byte message type identifier and the 4 byte message length. +func (src *Execute) Encode(dst []byte) ([]byte, error) { + dst, sp := beginMessage(dst, 'E') + dst = append(dst, src.Portal...) + dst = append(dst, 0) + dst = pgio.AppendUint32(dst, src.MaxRows) + return finishMessage(dst, sp) +} + +// MarshalJSON implements encoding/json.Marshaler. +func (src Execute) MarshalJSON() ([]byte, error) { + return json.Marshal(struct { + Type string + Portal string + MaxRows uint32 + }{ + Type: "Execute", + Portal: src.Portal, + MaxRows: src.MaxRows, + }) +} diff --git a/vendor/github.com/jackc/pgx/v5/pgproto3/flush.go b/vendor/github.com/jackc/pgx/v5/pgproto3/flush.go new file mode 100644 index 0000000000..e5dc1fbbd3 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgproto3/flush.go @@ -0,0 +1,34 @@ +package pgproto3 + +import ( + "encoding/json" +) + +type Flush struct{} + +// Frontend identifies this message as sendable by a PostgreSQL frontend. +func (*Flush) Frontend() {} + +// Decode decodes src into dst. src must contain the complete message with the exception of the initial 1 byte message +// type identifier and 4 byte message length. +func (dst *Flush) Decode(src []byte) error { + if len(src) != 0 { + return &invalidMessageLenErr{messageType: "Flush", expectedLen: 0, actualLen: len(src)} + } + + return nil +} + +// Encode encodes src into dst. dst will include the 1 byte message type identifier and the 4 byte message length. +func (src *Flush) Encode(dst []byte) ([]byte, error) { + return append(dst, 'H', 0, 0, 0, 4), nil +} + +// MarshalJSON implements encoding/json.Marshaler. +func (src Flush) MarshalJSON() ([]byte, error) { + return json.Marshal(struct { + Type string + }{ + Type: "Flush", + }) +} diff --git a/vendor/github.com/jackc/pgx/v5/pgproto3/frontend.go b/vendor/github.com/jackc/pgx/v5/pgproto3/frontend.go new file mode 100644 index 0000000000..9fc85f2c74 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgproto3/frontend.go @@ -0,0 +1,475 @@ +package pgproto3 + +import ( + "bytes" + "encoding/binary" + "errors" + "fmt" + "io" +) + +// Frontend acts as a client for the PostgreSQL wire protocol version 3. +type Frontend struct { + cr *chunkReader + w io.Writer + + // tracer is used to trace messages when Send or Receive is called. This means an outbound message is traced + // before it is actually transmitted (i.e. before Flush). It is safe to change this variable when the Frontend is + // idle. Setting and unsetting tracer provides equivalent functionality to PQtrace and PQuntrace in libpq. + tracer *tracer + + wbuf []byte + encodeError error + + // Backend message flyweights + authenticationOk AuthenticationOk + authenticationCleartextPassword AuthenticationCleartextPassword + authenticationMD5Password AuthenticationMD5Password + authenticationGSS AuthenticationGSS + authenticationGSSContinue AuthenticationGSSContinue + authenticationSASL AuthenticationSASL + authenticationSASLContinue AuthenticationSASLContinue + authenticationSASLFinal AuthenticationSASLFinal + backendKeyData BackendKeyData + bindComplete BindComplete + closeComplete CloseComplete + commandComplete CommandComplete + copyBothResponse CopyBothResponse + copyData CopyData + copyInResponse CopyInResponse + copyOutResponse CopyOutResponse + copyDone CopyDone + dataRow DataRow + emptyQueryResponse EmptyQueryResponse + errorResponse ErrorResponse + functionCallResponse FunctionCallResponse + noData NoData + noticeResponse NoticeResponse + notificationResponse NotificationResponse + parameterDescription ParameterDescription + parameterStatus ParameterStatus + parseComplete ParseComplete + readyForQuery ReadyForQuery + rowDescription RowDescription + portalSuspended PortalSuspended + negotiateProtocolVersion NegotiateProtocolVersion + + bodyLen int + maxBodyLen int // maxBodyLen is the maximum length of a message body in octets. If a message body exceeds this length, Receive will return an error. + msgType byte + partialMsg bool + authType uint32 +} + +// NewFrontend creates a new Frontend. +// +// The maximum accepted message body length defaults to the same ~1 GiB limit the PostgreSQL +// server enforces on inbound messages (PQ_LARGE_MESSAGE_LIMIT in src/include/libpq/libpq.h). +// Use [Frontend.SetMaxBodyLen] to change or remove the limit. +func NewFrontend(r io.Reader, w io.Writer) *Frontend { + cr := newChunkReader(r, 0) + return &Frontend{cr: cr, w: w, maxBodyLen: maxMessageBodyLen} +} + +// Send sends a message to the backend (i.e. the server). The message is buffered until Flush is called. Any error +// encountered will be returned from Flush. +// +// Send can work with any FrontendMessage. Some commonly used message types such as Bind have specialized send methods +// such as SendBind. These methods should be preferred when the type of message is known up front (e.g. when building an +// extended query protocol query) as they may be faster due to knowing the type of msg rather than it being hidden +// behind an interface. +func (f *Frontend) Send(msg FrontendMessage) { + if f.encodeError != nil { + return + } + + prevLen := len(f.wbuf) + newBuf, err := msg.Encode(f.wbuf) + if err != nil { + f.encodeError = err + return + } + f.wbuf = newBuf + + if f.tracer != nil { + f.tracer.traceMessage('F', int32(len(f.wbuf)-prevLen), msg) + } +} + +// Flush writes any pending messages to the backend (i.e. the server). +func (f *Frontend) Flush() error { + if err := f.encodeError; err != nil { + f.encodeError = nil + f.wbuf = f.wbuf[:0] + return &writeError{err: err, safeToRetry: true} + } + + if len(f.wbuf) == 0 { + return nil + } + + n, err := f.w.Write(f.wbuf) + + const maxLen = 1024 + if len(f.wbuf) > maxLen { + f.wbuf = make([]byte, 0, maxLen) + } else { + f.wbuf = f.wbuf[:0] + } + + if err != nil { + return &writeError{err: err, safeToRetry: n == 0} + } + + return nil +} + +// Trace starts tracing the message traffic to w. It writes in a similar format to that produced by the libpq function +// PQtrace. +func (f *Frontend) Trace(w io.Writer, options TracerOptions) { + f.tracer = &tracer{ + w: w, + buf: &bytes.Buffer{}, + TracerOptions: options, + } +} + +// Untrace stops tracing. +func (f *Frontend) Untrace() { + f.tracer = nil +} + +// SendBind sends a Bind message to the backend (i.e. the server). The message is buffered until Flush is called. Any +// error encountered will be returned from Flush. +func (f *Frontend) SendBind(msg *Bind) { + if f.encodeError != nil { + return + } + + prevLen := len(f.wbuf) + newBuf, err := msg.Encode(f.wbuf) + if err != nil { + f.encodeError = err + return + } + f.wbuf = newBuf + + if f.tracer != nil { + f.tracer.traceBind('F', int32(len(f.wbuf)-prevLen), msg) + } +} + +// SendParse sends a Parse message to the backend (i.e. the server). The message is buffered until Flush is called. Any +// error encountered will be returned from Flush. +func (f *Frontend) SendParse(msg *Parse) { + if f.encodeError != nil { + return + } + + prevLen := len(f.wbuf) + newBuf, err := msg.Encode(f.wbuf) + if err != nil { + f.encodeError = err + return + } + f.wbuf = newBuf + + if f.tracer != nil { + f.tracer.traceParse('F', int32(len(f.wbuf)-prevLen), msg) + } +} + +// SendClose sends a Close message to the backend (i.e. the server). The message is buffered until Flush is called. Any +// error encountered will be returned from Flush. +func (f *Frontend) SendClose(msg *Close) { + if f.encodeError != nil { + return + } + + prevLen := len(f.wbuf) + newBuf, err := msg.Encode(f.wbuf) + if err != nil { + f.encodeError = err + return + } + f.wbuf = newBuf + + if f.tracer != nil { + f.tracer.traceClose('F', int32(len(f.wbuf)-prevLen), msg) + } +} + +// SendDescribe sends a Describe message to the backend (i.e. the server). The message is buffered until Flush is +// called. Any error encountered will be returned from Flush. +func (f *Frontend) SendDescribe(msg *Describe) { + if f.encodeError != nil { + return + } + + prevLen := len(f.wbuf) + newBuf, err := msg.Encode(f.wbuf) + if err != nil { + f.encodeError = err + return + } + f.wbuf = newBuf + + if f.tracer != nil { + f.tracer.traceDescribe('F', int32(len(f.wbuf)-prevLen), msg) + } +} + +// SendExecute sends an Execute message to the backend (i.e. the server). The message is buffered until Flush is called. +// Any error encountered will be returned from Flush. +func (f *Frontend) SendExecute(msg *Execute) { + if f.encodeError != nil { + return + } + + prevLen := len(f.wbuf) + newBuf, err := msg.Encode(f.wbuf) + if err != nil { + f.encodeError = err + return + } + f.wbuf = newBuf + + if f.tracer != nil { + f.tracer.traceExecute('F', int32(len(f.wbuf)-prevLen), msg) + } +} + +// SendSync sends a Sync message to the backend (i.e. the server). The message is buffered until Flush is called. Any +// error encountered will be returned from Flush. +func (f *Frontend) SendSync(msg *Sync) { + if f.encodeError != nil { + return + } + + prevLen := len(f.wbuf) + newBuf, err := msg.Encode(f.wbuf) + if err != nil { + f.encodeError = err + return + } + f.wbuf = newBuf + + if f.tracer != nil { + f.tracer.traceSync('F', int32(len(f.wbuf)-prevLen), msg) + } +} + +// SendQuery sends a Query message to the backend (i.e. the server). The message is buffered until Flush is called. Any +// error encountered will be returned from Flush. +func (f *Frontend) SendQuery(msg *Query) { + if f.encodeError != nil { + return + } + + prevLen := len(f.wbuf) + newBuf, err := msg.Encode(f.wbuf) + if err != nil { + f.encodeError = err + return + } + f.wbuf = newBuf + + if f.tracer != nil { + f.tracer.traceQuery('F', int32(len(f.wbuf)-prevLen), msg) + } +} + +// SendUnbufferedEncodedCopyData immediately sends an encoded CopyData message to the backend (i.e. the server). This method +// is more efficient than sending a CopyData message with Send as the message data is not copied to the internal buffer +// before being written out. The internal buffer is flushed before the message is sent. +func (f *Frontend) SendUnbufferedEncodedCopyData(msg []byte) error { + err := f.Flush() + if err != nil { + return err + } + + n, err := f.w.Write(msg) + if err != nil { + return &writeError{err: err, safeToRetry: n == 0} + } + + if f.tracer != nil { + f.tracer.traceCopyData('F', int32(len(msg)-1), &CopyData{}) + } + + return nil +} + +func translateEOFtoErrUnexpectedEOF(err error) error { + if err == io.EOF { + return io.ErrUnexpectedEOF + } + return err +} + +// Receive receives a message from the backend. The returned message is only valid until the next call to Receive. +func (f *Frontend) Receive() (BackendMessage, error) { + if !f.partialMsg { + header, err := f.cr.Next(5) + if err != nil { + return nil, translateEOFtoErrUnexpectedEOF(err) + } + + f.msgType = header[0] + + msgLength := int(int32(binary.BigEndian.Uint32(header[1:]))) + if msgLength < 4 { + return nil, fmt.Errorf("invalid message length: %d", msgLength) + } + + f.bodyLen = msgLength - 4 + if f.maxBodyLen > 0 && f.bodyLen > f.maxBodyLen { + return nil, &ExceededMaxBodyLenErr{f.maxBodyLen, f.bodyLen} + } + f.partialMsg = true + } + + msgBody, err := f.cr.Next(f.bodyLen) + if err != nil { + return nil, translateEOFtoErrUnexpectedEOF(err) + } + + f.partialMsg = false + + var msg BackendMessage + switch f.msgType { + case '1': + msg = &f.parseComplete + case '2': + msg = &f.bindComplete + case '3': + msg = &f.closeComplete + case 'A': + msg = &f.notificationResponse + case 'c': + msg = &f.copyDone + case 'C': + msg = &f.commandComplete + case 'd': + msg = &f.copyData + case 'D': + msg = &f.dataRow + case 'E': + msg = &f.errorResponse + case 'G': + msg = &f.copyInResponse + case 'H': + msg = &f.copyOutResponse + case 'I': + msg = &f.emptyQueryResponse + case 'K': + msg = &f.backendKeyData + case 'n': + msg = &f.noData + case 'N': + msg = &f.noticeResponse + case 'R': + var err error + msg, err = f.findAuthenticationMessageType(msgBody) + if err != nil { + return nil, err + } + case 's': + msg = &f.portalSuspended + case 'S': + msg = &f.parameterStatus + case 't': + msg = &f.parameterDescription + case 'T': + msg = &f.rowDescription + case 'V': + msg = &f.functionCallResponse + case 'W': + msg = &f.copyBothResponse + case 'Z': + msg = &f.readyForQuery + case 'v': + msg = &f.negotiateProtocolVersion + default: + return nil, fmt.Errorf("unknown message type: %c", f.msgType) + } + + err = msg.Decode(msgBody) + if err != nil { + return nil, err + } + + if f.tracer != nil { + f.tracer.traceMessage('B', int32(5+len(msgBody)), msg) + } + + return msg, nil +} + +// Authentication message type constants. +// See src/include/libpq/pqcomm.h for all +// constants. +const ( + AuthTypeOk = 0 + AuthTypeCleartextPassword = 3 + AuthTypeMD5Password = 5 + AuthTypeSCMCreds = 6 + AuthTypeGSS = 7 + AuthTypeGSSCont = 8 + AuthTypeSSPI = 9 + AuthTypeSASL = 10 + AuthTypeSASLContinue = 11 + AuthTypeSASLFinal = 12 +) + +func (f *Frontend) findAuthenticationMessageType(src []byte) (BackendMessage, error) { + if len(src) < 4 { + return nil, errors.New("authentication message too short") + } + f.authType = binary.BigEndian.Uint32(src[:4]) + + switch f.authType { + case AuthTypeOk: + return &f.authenticationOk, nil + case AuthTypeCleartextPassword: + return &f.authenticationCleartextPassword, nil + case AuthTypeMD5Password: + return &f.authenticationMD5Password, nil + case AuthTypeSCMCreds: + return nil, errors.New("AuthTypeSCMCreds is unimplemented") + case AuthTypeGSS: + return &f.authenticationGSS, nil + case AuthTypeGSSCont: + return &f.authenticationGSSContinue, nil + case AuthTypeSSPI: + return nil, errors.New("AuthTypeSSPI is unimplemented") + case AuthTypeSASL: + return &f.authenticationSASL, nil + case AuthTypeSASLContinue: + return &f.authenticationSASLContinue, nil + case AuthTypeSASLFinal: + return &f.authenticationSASLFinal, nil + default: + return nil, fmt.Errorf("unknown authentication type: %d", f.authType) + } +} + +// GetAuthType returns the authType used in the current state of the frontend. +// See SetAuthType for more information. +func (f *Frontend) GetAuthType() uint32 { + return f.authType +} + +func (f *Frontend) ReadBufferLen() int { + return f.cr.wp - f.cr.rp +} + +// SetMaxBodyLen sets the maximum length of a message body in octets. +// If a message body exceeds this length, Receive will return an error. +// This is useful for protecting against a corrupted server that sends +// messages with incorrect length, which can cause memory exhaustion. +// The default value is 0. +// If maxBodyLen is 0, then no maximum is enforced. +func (f *Frontend) SetMaxBodyLen(maxBodyLen int) { + f.maxBodyLen = maxBodyLen +} diff --git a/vendor/github.com/jackc/pgx/v5/pgproto3/function_call.go b/vendor/github.com/jackc/pgx/v5/pgproto3/function_call.go new file mode 100644 index 0000000000..ef3cfd3b80 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgproto3/function_call.go @@ -0,0 +1,124 @@ +package pgproto3 + +import ( + "encoding/binary" + "errors" + "math" + + "github.com/jackc/pgx/v5/internal/pgio" +) + +type FunctionCall struct { + Function uint32 + ArgFormatCodes []uint16 + Arguments [][]byte + ResultFormatCode uint16 +} + +// Frontend identifies this message as sendable by a PostgreSQL frontend. +func (*FunctionCall) Frontend() {} + +// Decode decodes src into dst. src must contain the complete message with the exception of the initial 1 byte message +// type identifier and 4 byte message length. +func (dst *FunctionCall) Decode(src []byte) error { + *dst = FunctionCall{} + rp := 0 + + if len(src) < 8 { + return &invalidMessageFormatErr{messageType: "FunctionCall"} + } + + // Specifies the object ID of the function to call. + dst.Function = binary.BigEndian.Uint32(src[rp:]) + rp += 4 + // The number of argument format codes that follow (denoted C below). + // This can be zero to indicate that there are no arguments or that the arguments all use the default format (text); + // or one, in which case the specified format code is applied to all arguments; + // or it can equal the actual number of arguments. + nArgumentCodes := int(binary.BigEndian.Uint16(src[rp:])) + rp += 2 + + if len(src[rp:]) < nArgumentCodes*2+2 { + return &invalidMessageFormatErr{messageType: "FunctionCall"} + } + + argumentCodes := make([]uint16, nArgumentCodes) + for i := range nArgumentCodes { + // The argument format codes. Each must presently be zero (text) or one (binary). + ac := binary.BigEndian.Uint16(src[rp:]) + if ac != 0 && ac != 1 { + return &invalidMessageFormatErr{messageType: "FunctionCall"} + } + argumentCodes[i] = ac + rp += 2 + } + dst.ArgFormatCodes = argumentCodes + + // Specifies the number of arguments being supplied to the function. + nArguments := int(binary.BigEndian.Uint16(src[rp:])) + rp += 2 + arguments := make([][]byte, nArguments) + for i := range nArguments { + if len(src[rp:]) < 4 { + return &invalidMessageFormatErr{messageType: "FunctionCall"} + } + // The length of the argument value, in bytes (this count does not include itself). Can be zero. + // As a special case, -1 indicates a NULL argument value. No value bytes follow in the NULL case. + argumentLength := int(int32(binary.BigEndian.Uint32(src[rp:]))) + rp += 4 + switch { + case argumentLength == -1: + arguments[i] = nil + case argumentLength < 0: + return &invalidMessageFormatErr{messageType: "FunctionCall"} + default: + if len(src[rp:]) < argumentLength { + return &invalidMessageFormatErr{messageType: "FunctionCall"} + } + // The value of the argument, in the format indicated by the associated format code. n is the above length. + argumentValue := src[rp : rp+argumentLength] + rp += argumentLength + arguments[i] = argumentValue + } + } + dst.Arguments = arguments + // The format code for the function result. Must presently be zero (text) or one (binary). + if len(src[rp:]) < 2 { + return &invalidMessageFormatErr{messageType: "FunctionCall"} + } + resultFormatCode := binary.BigEndian.Uint16(src[rp:]) + if resultFormatCode != 0 && resultFormatCode != 1 { + return &invalidMessageFormatErr{messageType: "FunctionCall"} + } + dst.ResultFormatCode = resultFormatCode + return nil +} + +// Encode encodes src into dst. dst will include the 1 byte message type identifier and the 4 byte message length. +func (src *FunctionCall) Encode(dst []byte) ([]byte, error) { + dst, sp := beginMessage(dst, 'F') + dst = pgio.AppendUint32(dst, src.Function) + + if len(src.ArgFormatCodes) > math.MaxUint16 { + return nil, errors.New("too many arg format codes") + } + dst = pgio.AppendUint16(dst, uint16(len(src.ArgFormatCodes))) + for _, argFormatCode := range src.ArgFormatCodes { + dst = pgio.AppendUint16(dst, argFormatCode) + } + + if len(src.Arguments) > math.MaxUint16 { + return nil, errors.New("too many arguments") + } + dst = pgio.AppendUint16(dst, uint16(len(src.Arguments))) + for _, argument := range src.Arguments { + if argument == nil { + dst = pgio.AppendInt32(dst, -1) + } else { + dst = pgio.AppendInt32(dst, int32(len(argument))) + dst = append(dst, argument...) + } + } + dst = pgio.AppendUint16(dst, src.ResultFormatCode) + return finishMessage(dst, sp) +} diff --git a/vendor/github.com/jackc/pgx/v5/pgproto3/function_call_response.go b/vendor/github.com/jackc/pgx/v5/pgproto3/function_call_response.go new file mode 100644 index 0000000000..6b6ed8b929 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgproto3/function_call_response.go @@ -0,0 +1,97 @@ +package pgproto3 + +import ( + "encoding/binary" + "encoding/hex" + "encoding/json" + + "github.com/jackc/pgx/v5/internal/pgio" +) + +type FunctionCallResponse struct { + Result []byte +} + +// Backend identifies this message as sendable by the PostgreSQL backend. +func (*FunctionCallResponse) Backend() {} + +// Decode decodes src into dst. src must contain the complete message with the exception of the initial 1 byte message +// type identifier and 4 byte message length. +func (dst *FunctionCallResponse) Decode(src []byte) error { + if len(src) < 4 { + return &invalidMessageFormatErr{messageType: "FunctionCallResponse"} + } + rp := 0 + resultSize := int(int32(binary.BigEndian.Uint32(src[rp:]))) + rp += 4 + + if resultSize == -1 { + dst.Result = nil + return nil + } + + if resultSize < 0 || len(src[rp:]) != resultSize { + return &invalidMessageFormatErr{messageType: "FunctionCallResponse"} + } + + dst.Result = src[rp:] + return nil +} + +// Encode encodes src into dst. dst will include the 1 byte message type identifier and the 4 byte message length. +func (src *FunctionCallResponse) Encode(dst []byte) ([]byte, error) { + dst, sp := beginMessage(dst, 'V') + + if src.Result == nil { + dst = pgio.AppendInt32(dst, -1) + } else { + dst = pgio.AppendInt32(dst, int32(len(src.Result))) + dst = append(dst, src.Result...) + } + + return finishMessage(dst, sp) +} + +// MarshalJSON implements encoding/json.Marshaler. +func (src FunctionCallResponse) MarshalJSON() ([]byte, error) { + var formattedValue map[string]string + var hasNonPrintable bool + for _, b := range src.Result { + if b < 32 { + hasNonPrintable = true + break + } + } + + if hasNonPrintable { + formattedValue = map[string]string{"binary": hex.EncodeToString(src.Result)} + } else { + formattedValue = map[string]string{"text": string(src.Result)} + } + + return json.Marshal(struct { + Type string + Result map[string]string + }{ + Type: "FunctionCallResponse", + Result: formattedValue, + }) +} + +// UnmarshalJSON implements encoding/json.Unmarshaler. +func (dst *FunctionCallResponse) UnmarshalJSON(data []byte) error { + // Ignore null, like in the main JSON package. + if string(data) == "null" { + return nil + } + + var msg struct { + Result map[string]string + } + err := json.Unmarshal(data, &msg) + if err != nil { + return err + } + dst.Result, err = getValueFromJSON(msg.Result) + return err +} diff --git a/vendor/github.com/jackc/pgx/v5/pgproto3/gss_enc_request.go b/vendor/github.com/jackc/pgx/v5/pgproto3/gss_enc_request.go new file mode 100644 index 0000000000..122d1341c6 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgproto3/gss_enc_request.go @@ -0,0 +1,48 @@ +package pgproto3 + +import ( + "encoding/binary" + "encoding/json" + "errors" + + "github.com/jackc/pgx/v5/internal/pgio" +) + +const gssEncReqNumber = 80877104 + +type GSSEncRequest struct{} + +// Frontend identifies this message as sendable by a PostgreSQL frontend. +func (*GSSEncRequest) Frontend() {} + +func (dst *GSSEncRequest) Decode(src []byte) error { + if len(src) < 4 { + return errors.New("gss encoding request too short") + } + + requestCode := binary.BigEndian.Uint32(src) + + if requestCode != gssEncReqNumber { + return errors.New("bad gss encoding request code") + } + + return nil +} + +// Encode encodes src into dst. dst will include the 4 byte message length. +func (src *GSSEncRequest) Encode(dst []byte) ([]byte, error) { + dst = pgio.AppendInt32(dst, 8) + dst = pgio.AppendInt32(dst, gssEncReqNumber) + return dst, nil +} + +// MarshalJSON implements encoding/json.Marshaler. +func (src GSSEncRequest) MarshalJSON() ([]byte, error) { + return json.Marshal(struct { + Type string + ProtocolVersion uint32 + Parameters map[string]string + }{ + Type: "GSSEncRequest", + }) +} diff --git a/vendor/github.com/jackc/pgx/v5/pgproto3/gss_response.go b/vendor/github.com/jackc/pgx/v5/pgproto3/gss_response.go new file mode 100644 index 0000000000..10d9377593 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgproto3/gss_response.go @@ -0,0 +1,46 @@ +package pgproto3 + +import ( + "encoding/json" +) + +type GSSResponse struct { + Data []byte +} + +// Frontend identifies this message as sendable by a PostgreSQL frontend. +func (g *GSSResponse) Frontend() {} + +func (g *GSSResponse) Decode(data []byte) error { + g.Data = data + return nil +} + +func (g *GSSResponse) Encode(dst []byte) ([]byte, error) { + dst, sp := beginMessage(dst, 'p') + dst = append(dst, g.Data...) + return finishMessage(dst, sp) +} + +// MarshalJSON implements encoding/json.Marshaler. +func (g *GSSResponse) MarshalJSON() ([]byte, error) { + return json.Marshal(struct { + Type string + Data []byte + }{ + Type: "GSSResponse", + Data: g.Data, + }) +} + +// UnmarshalJSON implements encoding/json.Unmarshaler. +func (g *GSSResponse) UnmarshalJSON(data []byte) error { + var msg struct { + Data []byte + } + if err := json.Unmarshal(data, &msg); err != nil { + return err + } + g.Data = msg.Data + return nil +} diff --git a/vendor/github.com/jackc/pgx/v5/pgproto3/negotiate_protocol_version.go b/vendor/github.com/jackc/pgx/v5/pgproto3/negotiate_protocol_version.go new file mode 100644 index 0000000000..43bd7ec636 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgproto3/negotiate_protocol_version.go @@ -0,0 +1,93 @@ +package pgproto3 + +import ( + "encoding/binary" + "encoding/json" + + "github.com/jackc/pgx/v5/internal/pgio" +) + +type NegotiateProtocolVersion struct { + NewestMinorProtocol uint32 + UnrecognizedOptions []string +} + +// Backend identifies this message as sendable by the PostgreSQL backend. +func (*NegotiateProtocolVersion) Backend() {} + +// Decode decodes src into dst. src must contain the complete message with the exception of the initial 1 byte message +// type identifier and 4 byte message length. +func (dst *NegotiateProtocolVersion) Decode(src []byte) error { + if len(src) < 8 { + return &invalidMessageLenErr{messageType: "NegotiateProtocolVersion", expectedLen: 8, actualLen: len(src)} + } + + dst.NewestMinorProtocol = binary.BigEndian.Uint32(src[:4]) + optionCount := int(binary.BigEndian.Uint32(src[4:8])) + + rp := 8 + + // Use the remaining message size as an upper bound for capacity to prevent + // malicious optionCount values from causing excessive memory allocation. + capHint := optionCount + if remaining := len(src) - rp; capHint > remaining { + capHint = remaining + } + dst.UnrecognizedOptions = make([]string, 0, capHint) + for i := 0; i < optionCount; i++ { + if rp >= len(src) { + return &invalidMessageFormatErr{messageType: "NegotiateProtocolVersion"} + } + end := rp + for end < len(src) && src[end] != 0 { + end++ + } + if end >= len(src) { + return &invalidMessageFormatErr{messageType: "NegotiateProtocolVersion"} + } + dst.UnrecognizedOptions = append(dst.UnrecognizedOptions, string(src[rp:end])) + rp = end + 1 + } + + return nil +} + +// Encode encodes src into dst. dst will include the 1 byte message type identifier and the 4 byte message length. +func (src *NegotiateProtocolVersion) Encode(dst []byte) ([]byte, error) { + dst, sp := beginMessage(dst, 'v') + dst = pgio.AppendUint32(dst, src.NewestMinorProtocol) + dst = pgio.AppendUint32(dst, uint32(len(src.UnrecognizedOptions))) + for _, option := range src.UnrecognizedOptions { + dst = append(dst, option...) + dst = append(dst, 0) + } + return finishMessage(dst, sp) +} + +// MarshalJSON implements encoding/json.Marshaler. +func (src NegotiateProtocolVersion) MarshalJSON() ([]byte, error) { + return json.Marshal(struct { + Type string + NewestMinorProtocol uint32 + UnrecognizedOptions []string + }{ + Type: "NegotiateProtocolVersion", + NewestMinorProtocol: src.NewestMinorProtocol, + UnrecognizedOptions: src.UnrecognizedOptions, + }) +} + +// UnmarshalJSON implements encoding/json.Unmarshaler. +func (dst *NegotiateProtocolVersion) UnmarshalJSON(data []byte) error { + var msg struct { + NewestMinorProtocol uint32 + UnrecognizedOptions []string + } + if err := json.Unmarshal(data, &msg); err != nil { + return err + } + + dst.NewestMinorProtocol = msg.NewestMinorProtocol + dst.UnrecognizedOptions = msg.UnrecognizedOptions + return nil +} diff --git a/vendor/github.com/jackc/pgx/v5/pgproto3/no_data.go b/vendor/github.com/jackc/pgx/v5/pgproto3/no_data.go new file mode 100644 index 0000000000..cbcaad40c4 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgproto3/no_data.go @@ -0,0 +1,34 @@ +package pgproto3 + +import ( + "encoding/json" +) + +type NoData struct{} + +// Backend identifies this message as sendable by the PostgreSQL backend. +func (*NoData) Backend() {} + +// Decode decodes src into dst. src must contain the complete message with the exception of the initial 1 byte message +// type identifier and 4 byte message length. +func (dst *NoData) Decode(src []byte) error { + if len(src) != 0 { + return &invalidMessageLenErr{messageType: "NoData", expectedLen: 0, actualLen: len(src)} + } + + return nil +} + +// Encode encodes src into dst. dst will include the 1 byte message type identifier and the 4 byte message length. +func (src *NoData) Encode(dst []byte) ([]byte, error) { + return append(dst, 'n', 0, 0, 0, 4), nil +} + +// MarshalJSON implements encoding/json.Marshaler. +func (src NoData) MarshalJSON() ([]byte, error) { + return json.Marshal(struct { + Type string + }{ + Type: "NoData", + }) +} diff --git a/vendor/github.com/jackc/pgx/v5/pgproto3/notice_response.go b/vendor/github.com/jackc/pgx/v5/pgproto3/notice_response.go new file mode 100644 index 0000000000..497aba6dd5 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgproto3/notice_response.go @@ -0,0 +1,19 @@ +package pgproto3 + +type NoticeResponse ErrorResponse + +// Backend identifies this message as sendable by the PostgreSQL backend. +func (*NoticeResponse) Backend() {} + +// Decode decodes src into dst. src must contain the complete message with the exception of the initial 1 byte message +// type identifier and 4 byte message length. +func (dst *NoticeResponse) Decode(src []byte) error { + return (*ErrorResponse)(dst).Decode(src) +} + +// Encode encodes src into dst. dst will include the 1 byte message type identifier and the 4 byte message length. +func (src *NoticeResponse) Encode(dst []byte) ([]byte, error) { + dst, sp := beginMessage(dst, 'N') + dst = (*ErrorResponse)(src).appendFields(dst) + return finishMessage(dst, sp) +} diff --git a/vendor/github.com/jackc/pgx/v5/pgproto3/notification_response.go b/vendor/github.com/jackc/pgx/v5/pgproto3/notification_response.go new file mode 100644 index 0000000000..243b6bf7c6 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgproto3/notification_response.go @@ -0,0 +1,71 @@ +package pgproto3 + +import ( + "bytes" + "encoding/binary" + "encoding/json" + + "github.com/jackc/pgx/v5/internal/pgio" +) + +type NotificationResponse struct { + PID uint32 + Channel string + Payload string +} + +// Backend identifies this message as sendable by the PostgreSQL backend. +func (*NotificationResponse) Backend() {} + +// Decode decodes src into dst. src must contain the complete message with the exception of the initial 1 byte message +// type identifier and 4 byte message length. +func (dst *NotificationResponse) Decode(src []byte) error { + buf := bytes.NewBuffer(src) + + if buf.Len() < 4 { + return &invalidMessageFormatErr{messageType: "NotificationResponse", details: "too short"} + } + + pid := binary.BigEndian.Uint32(buf.Next(4)) + + b, err := buf.ReadBytes(0) + if err != nil { + return err + } + channel := string(b[:len(b)-1]) + + b, err = buf.ReadBytes(0) + if err != nil { + return err + } + payload := string(b[:len(b)-1]) + + *dst = NotificationResponse{PID: pid, Channel: channel, Payload: payload} + return nil +} + +// Encode encodes src into dst. dst will include the 1 byte message type identifier and the 4 byte message length. +func (src *NotificationResponse) Encode(dst []byte) ([]byte, error) { + dst, sp := beginMessage(dst, 'A') + dst = pgio.AppendUint32(dst, src.PID) + dst = append(dst, src.Channel...) + dst = append(dst, 0) + dst = append(dst, src.Payload...) + dst = append(dst, 0) + return finishMessage(dst, sp) +} + +// MarshalJSON implements encoding/json.Marshaler. +func (src NotificationResponse) MarshalJSON() ([]byte, error) { + return json.Marshal(struct { + Type string + PID uint32 + Channel string + Payload string + }{ + Type: "NotificationResponse", + PID: src.PID, + Channel: src.Channel, + Payload: src.Payload, + }) +} diff --git a/vendor/github.com/jackc/pgx/v5/pgproto3/parameter_description.go b/vendor/github.com/jackc/pgx/v5/pgproto3/parameter_description.go new file mode 100644 index 0000000000..58eb26ef0f --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgproto3/parameter_description.go @@ -0,0 +1,67 @@ +package pgproto3 + +import ( + "bytes" + "encoding/binary" + "encoding/json" + "errors" + "math" + + "github.com/jackc/pgx/v5/internal/pgio" +) + +type ParameterDescription struct { + ParameterOIDs []uint32 +} + +// Backend identifies this message as sendable by the PostgreSQL backend. +func (*ParameterDescription) Backend() {} + +// Decode decodes src into dst. src must contain the complete message with the exception of the initial 1 byte message +// type identifier and 4 byte message length. +func (dst *ParameterDescription) Decode(src []byte) error { + buf := bytes.NewBuffer(src) + + if buf.Len() < 2 { + return &invalidMessageFormatErr{messageType: "ParameterDescription"} + } + + // Reported parameter count will be incorrect when number of args is greater than uint16 + buf.Next(2) + // Instead infer parameter count by remaining size of message + parameterCount := buf.Len() / 4 + + *dst = ParameterDescription{ParameterOIDs: make([]uint32, parameterCount)} + + for i := range parameterCount { + dst.ParameterOIDs[i] = binary.BigEndian.Uint32(buf.Next(4)) + } + + return nil +} + +// Encode encodes src into dst. dst will include the 1 byte message type identifier and the 4 byte message length. +func (src *ParameterDescription) Encode(dst []byte) ([]byte, error) { + dst, sp := beginMessage(dst, 't') + + if len(src.ParameterOIDs) > math.MaxUint16 { + return nil, errors.New("too many parameter oids") + } + dst = pgio.AppendUint16(dst, uint16(len(src.ParameterOIDs))) + for _, oid := range src.ParameterOIDs { + dst = pgio.AppendUint32(dst, oid) + } + + return finishMessage(dst, sp) +} + +// MarshalJSON implements encoding/json.Marshaler. +func (src ParameterDescription) MarshalJSON() ([]byte, error) { + return json.Marshal(struct { + Type string + ParameterOIDs []uint32 + }{ + Type: "ParameterDescription", + ParameterOIDs: src.ParameterOIDs, + }) +} diff --git a/vendor/github.com/jackc/pgx/v5/pgproto3/parameter_status.go b/vendor/github.com/jackc/pgx/v5/pgproto3/parameter_status.go new file mode 100644 index 0000000000..9ee0720b54 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgproto3/parameter_status.go @@ -0,0 +1,58 @@ +package pgproto3 + +import ( + "bytes" + "encoding/json" +) + +type ParameterStatus struct { + Name string + Value string +} + +// Backend identifies this message as sendable by the PostgreSQL backend. +func (*ParameterStatus) Backend() {} + +// Decode decodes src into dst. src must contain the complete message with the exception of the initial 1 byte message +// type identifier and 4 byte message length. +func (dst *ParameterStatus) Decode(src []byte) error { + buf := bytes.NewBuffer(src) + + b, err := buf.ReadBytes(0) + if err != nil { + return err + } + name := string(b[:len(b)-1]) + + b, err = buf.ReadBytes(0) + if err != nil { + return err + } + value := string(b[:len(b)-1]) + + *dst = ParameterStatus{Name: name, Value: value} + return nil +} + +// Encode encodes src into dst. dst will include the 1 byte message type identifier and the 4 byte message length. +func (src *ParameterStatus) Encode(dst []byte) ([]byte, error) { + dst, sp := beginMessage(dst, 'S') + dst = append(dst, src.Name...) + dst = append(dst, 0) + dst = append(dst, src.Value...) + dst = append(dst, 0) + return finishMessage(dst, sp) +} + +// MarshalJSON implements encoding/json.Marshaler. +func (ps ParameterStatus) MarshalJSON() ([]byte, error) { + return json.Marshal(struct { + Type string + Name string + Value string + }{ + Type: "ParameterStatus", + Name: ps.Name, + Value: ps.Value, + }) +} diff --git a/vendor/github.com/jackc/pgx/v5/pgproto3/parse.go b/vendor/github.com/jackc/pgx/v5/pgproto3/parse.go new file mode 100644 index 0000000000..8fb8de5d47 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgproto3/parse.go @@ -0,0 +1,89 @@ +package pgproto3 + +import ( + "bytes" + "encoding/binary" + "encoding/json" + "errors" + "math" + + "github.com/jackc/pgx/v5/internal/pgio" +) + +type Parse struct { + Name string + Query string + ParameterOIDs []uint32 +} + +// Frontend identifies this message as sendable by a PostgreSQL frontend. +func (*Parse) Frontend() {} + +// Decode decodes src into dst. src must contain the complete message with the exception of the initial 1 byte message +// type identifier and 4 byte message length. +func (dst *Parse) Decode(src []byte) error { + *dst = Parse{} + + buf := bytes.NewBuffer(src) + + b, err := buf.ReadBytes(0) + if err != nil { + return err + } + dst.Name = string(b[:len(b)-1]) + + b, err = buf.ReadBytes(0) + if err != nil { + return err + } + dst.Query = string(b[:len(b)-1]) + + if buf.Len() < 2 { + return &invalidMessageFormatErr{messageType: "Parse"} + } + parameterOIDCount := int(binary.BigEndian.Uint16(buf.Next(2))) + + for range parameterOIDCount { + if buf.Len() < 4 { + return &invalidMessageFormatErr{messageType: "Parse"} + } + dst.ParameterOIDs = append(dst.ParameterOIDs, binary.BigEndian.Uint32(buf.Next(4))) + } + + return nil +} + +// Encode encodes src into dst. dst will include the 1 byte message type identifier and the 4 byte message length. +func (src *Parse) Encode(dst []byte) ([]byte, error) { + dst, sp := beginMessage(dst, 'P') + + dst = append(dst, src.Name...) + dst = append(dst, 0) + dst = append(dst, src.Query...) + dst = append(dst, 0) + + if len(src.ParameterOIDs) > math.MaxUint16 { + return nil, errors.New("too many parameter oids") + } + dst = pgio.AppendUint16(dst, uint16(len(src.ParameterOIDs))) + for _, oid := range src.ParameterOIDs { + dst = pgio.AppendUint32(dst, oid) + } + + return finishMessage(dst, sp) +} + +// MarshalJSON implements encoding/json.Marshaler. +func (src Parse) MarshalJSON() ([]byte, error) { + return json.Marshal(struct { + Type string + Name string + Query string + ParameterOIDs []uint32 + }{ + Type: "Parse", + Name: src.Name, + Query: src.Query, + ParameterOIDs: src.ParameterOIDs, + }) +} diff --git a/vendor/github.com/jackc/pgx/v5/pgproto3/parse_complete.go b/vendor/github.com/jackc/pgx/v5/pgproto3/parse_complete.go new file mode 100644 index 0000000000..cff9e27d06 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgproto3/parse_complete.go @@ -0,0 +1,34 @@ +package pgproto3 + +import ( + "encoding/json" +) + +type ParseComplete struct{} + +// Backend identifies this message as sendable by the PostgreSQL backend. +func (*ParseComplete) Backend() {} + +// Decode decodes src into dst. src must contain the complete message with the exception of the initial 1 byte message +// type identifier and 4 byte message length. +func (dst *ParseComplete) Decode(src []byte) error { + if len(src) != 0 { + return &invalidMessageLenErr{messageType: "ParseComplete", expectedLen: 0, actualLen: len(src)} + } + + return nil +} + +// Encode encodes src into dst. dst will include the 1 byte message type identifier and the 4 byte message length. +func (src *ParseComplete) Encode(dst []byte) ([]byte, error) { + return append(dst, '1', 0, 0, 0, 4), nil +} + +// MarshalJSON implements encoding/json.Marshaler. +func (src ParseComplete) MarshalJSON() ([]byte, error) { + return json.Marshal(struct { + Type string + }{ + Type: "ParseComplete", + }) +} diff --git a/vendor/github.com/jackc/pgx/v5/pgproto3/password_message.go b/vendor/github.com/jackc/pgx/v5/pgproto3/password_message.go new file mode 100644 index 0000000000..67b78515d0 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgproto3/password_message.go @@ -0,0 +1,49 @@ +package pgproto3 + +import ( + "bytes" + "encoding/json" +) + +type PasswordMessage struct { + Password string +} + +// Frontend identifies this message as sendable by a PostgreSQL frontend. +func (*PasswordMessage) Frontend() {} + +// InitialResponse identifies this message as an authentication response. +func (*PasswordMessage) InitialResponse() {} + +// Decode decodes src into dst. src must contain the complete message with the exception of the initial 1 byte message +// type identifier and 4 byte message length. +func (dst *PasswordMessage) Decode(src []byte) error { + buf := bytes.NewBuffer(src) + + b, err := buf.ReadBytes(0) + if err != nil { + return err + } + dst.Password = string(b[:len(b)-1]) + + return nil +} + +// Encode encodes src into dst. dst will include the 1 byte message type identifier and the 4 byte message length. +func (src *PasswordMessage) Encode(dst []byte) ([]byte, error) { + dst, sp := beginMessage(dst, 'p') + dst = append(dst, src.Password...) + dst = append(dst, 0) + return finishMessage(dst, sp) +} + +// MarshalJSON implements encoding/json.Marshaler. +func (src PasswordMessage) MarshalJSON() ([]byte, error) { + return json.Marshal(struct { + Type string + Password string + }{ + Type: "PasswordMessage", + Password: src.Password, + }) +} diff --git a/vendor/github.com/jackc/pgx/v5/pgproto3/pgproto3.go b/vendor/github.com/jackc/pgx/v5/pgproto3/pgproto3.go new file mode 100644 index 0000000000..128f97f871 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgproto3/pgproto3.go @@ -0,0 +1,120 @@ +package pgproto3 + +import ( + "encoding/hex" + "errors" + "fmt" + + "github.com/jackc/pgx/v5/internal/pgio" +) + +// maxMessageBodyLen is the maximum length of a message body in bytes. See PG_LARGE_MESSAGE_LIMIT in the PostgreSQL +// source. It is defined as (MaxAllocSize - 1). MaxAllocSize is defined as 0x3fffffff. +const maxMessageBodyLen = (0x3fffffff - 1) + +// Message is the interface implemented by an object that can decode and encode +// a particular PostgreSQL message. +type Message interface { + // Decode is allowed and expected to retain a reference to data after + // returning (unlike encoding.BinaryUnmarshaler). + Decode(data []byte) error + + // Encode appends itself to dst and returns the new buffer. + Encode(dst []byte) ([]byte, error) +} + +// FrontendMessage is a message sent by the frontend (i.e. the client). +type FrontendMessage interface { + Message + Frontend() // no-op method to distinguish frontend from backend methods +} + +// BackendMessage is a message sent by the backend (i.e. the server). +type BackendMessage interface { + Message + Backend() // no-op method to distinguish frontend from backend methods +} + +type AuthenticationResponseMessage interface { + BackendMessage + AuthenticationResponse() // no-op method to distinguish authentication responses +} + +type invalidMessageLenErr struct { + messageType string + expectedLen int + actualLen int +} + +func (e *invalidMessageLenErr) Error() string { + return fmt.Sprintf("%s body must have length of %d, but it is %d", e.messageType, e.expectedLen, e.actualLen) +} + +type invalidMessageFormatErr struct { + messageType string + details string +} + +func (e *invalidMessageFormatErr) Error() string { + return fmt.Sprintf("%s body is invalid %s", e.messageType, e.details) +} + +type writeError struct { + err error + safeToRetry bool +} + +func (e *writeError) Error() string { + return fmt.Sprintf("write failed: %s", e.err.Error()) +} + +func (e *writeError) SafeToRetry() bool { + return e.safeToRetry +} + +func (e *writeError) Unwrap() error { + return e.err +} + +type ExceededMaxBodyLenErr struct { + MaxExpectedBodyLen int + ActualBodyLen int +} + +func (e *ExceededMaxBodyLenErr) Error() string { + return fmt.Sprintf("invalid body length: expected at most %d, but got %d", e.MaxExpectedBodyLen, e.ActualBodyLen) +} + +// getValueFromJSON gets the value from a protocol message representation in JSON. +func getValueFromJSON(v map[string]string) ([]byte, error) { + if v == nil { + return nil, nil + } + if text, ok := v["text"]; ok { + return []byte(text), nil + } + if binary, ok := v["binary"]; ok { + return hex.DecodeString(binary) + } + return nil, errors.New("unknown protocol representation") +} + +// beginMessage begins a new message of type t. It appends the message type and a placeholder for the message length to +// dst. It returns the new buffer and the position of the message length placeholder. +func beginMessage(dst []byte, t byte) ([]byte, int) { + dst = append(dst, t) + sp := len(dst) + dst = pgio.AppendInt32(dst, -1) + return dst, sp +} + +// finishMessage finishes a message that was started with beginMessage. It computes the message length and writes it to +// dst[sp]. If the message length is too large it returns an error. Otherwise it returns the final message buffer. +func finishMessage(dst []byte, sp int) ([]byte, error) { + messageBodyLen := len(dst[sp:]) + if messageBodyLen > maxMessageBodyLen { + return nil, errors.New("message body too large") + } + pgio.SetInt32(dst[sp:], int32(messageBodyLen)) + return dst, nil +} diff --git a/vendor/github.com/jackc/pgx/v5/pgproto3/portal_suspended.go b/vendor/github.com/jackc/pgx/v5/pgproto3/portal_suspended.go new file mode 100644 index 0000000000..9e2f8cbc41 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgproto3/portal_suspended.go @@ -0,0 +1,34 @@ +package pgproto3 + +import ( + "encoding/json" +) + +type PortalSuspended struct{} + +// Backend identifies this message as sendable by the PostgreSQL backend. +func (*PortalSuspended) Backend() {} + +// Decode decodes src into dst. src must contain the complete message with the exception of the initial 1 byte message +// type identifier and 4 byte message length. +func (dst *PortalSuspended) Decode(src []byte) error { + if len(src) != 0 { + return &invalidMessageLenErr{messageType: "PortalSuspended", expectedLen: 0, actualLen: len(src)} + } + + return nil +} + +// Encode encodes src into dst. dst will include the 1 byte message type identifier and the 4 byte message length. +func (src *PortalSuspended) Encode(dst []byte) ([]byte, error) { + return append(dst, 's', 0, 0, 0, 4), nil +} + +// MarshalJSON implements encoding/json.Marshaler. +func (src PortalSuspended) MarshalJSON() ([]byte, error) { + return json.Marshal(struct { + Type string + }{ + Type: "PortalSuspended", + }) +} diff --git a/vendor/github.com/jackc/pgx/v5/pgproto3/query.go b/vendor/github.com/jackc/pgx/v5/pgproto3/query.go new file mode 100644 index 0000000000..9e16465c25 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgproto3/query.go @@ -0,0 +1,49 @@ +package pgproto3 + +import ( + "bytes" + "encoding/json" +) + +type Query struct { + String string +} + +// Frontend identifies this message as sendable by a PostgreSQL frontend. +func (*Query) Frontend() {} + +// Decode decodes src into dst. src must contain the complete message with the exception of the initial 1 byte message +// type identifier and 4 byte message length. +func (dst *Query) Decode(src []byte) error { + if len(src) == 0 { + return &invalidMessageFormatErr{messageType: "Query"} + } + + i := bytes.IndexByte(src, 0) + if i != len(src)-1 { + return &invalidMessageFormatErr{messageType: "Query"} + } + + dst.String = string(src[:i]) + + return nil +} + +// Encode encodes src into dst. dst will include the 1 byte message type identifier and the 4 byte message length. +func (src *Query) Encode(dst []byte) ([]byte, error) { + dst, sp := beginMessage(dst, 'Q') + dst = append(dst, src.String...) + dst = append(dst, 0) + return finishMessage(dst, sp) +} + +// MarshalJSON implements encoding/json.Marshaler. +func (src Query) MarshalJSON() ([]byte, error) { + return json.Marshal(struct { + Type string + String string + }{ + Type: "Query", + String: src.String, + }) +} diff --git a/vendor/github.com/jackc/pgx/v5/pgproto3/ready_for_query.go b/vendor/github.com/jackc/pgx/v5/pgproto3/ready_for_query.go new file mode 100644 index 0000000000..a56af9fb24 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgproto3/ready_for_query.go @@ -0,0 +1,61 @@ +package pgproto3 + +import ( + "encoding/json" + "errors" +) + +type ReadyForQuery struct { + TxStatus byte +} + +// Backend identifies this message as sendable by the PostgreSQL backend. +func (*ReadyForQuery) Backend() {} + +// Decode decodes src into dst. src must contain the complete message with the exception of the initial 1 byte message +// type identifier and 4 byte message length. +func (dst *ReadyForQuery) Decode(src []byte) error { + if len(src) != 1 { + return &invalidMessageLenErr{messageType: "ReadyForQuery", expectedLen: 1, actualLen: len(src)} + } + + dst.TxStatus = src[0] + + return nil +} + +// Encode encodes src into dst. dst will include the 1 byte message type identifier and the 4 byte message length. +func (src *ReadyForQuery) Encode(dst []byte) ([]byte, error) { + return append(dst, 'Z', 0, 0, 0, 5, src.TxStatus), nil +} + +// MarshalJSON implements encoding/json.Marshaler. +func (src ReadyForQuery) MarshalJSON() ([]byte, error) { + return json.Marshal(struct { + Type string + TxStatus string + }{ + Type: "ReadyForQuery", + TxStatus: string(src.TxStatus), + }) +} + +// UnmarshalJSON implements encoding/json.Unmarshaler. +func (dst *ReadyForQuery) UnmarshalJSON(data []byte) error { + // Ignore null, like in the main JSON package. + if string(data) == "null" { + return nil + } + + var msg struct { + TxStatus string + } + if err := json.Unmarshal(data, &msg); err != nil { + return err + } + if len(msg.TxStatus) != 1 { + return errors.New("invalid length for ReadyForQuery.TxStatus") + } + dst.TxStatus = msg.TxStatus[0] + return nil +} diff --git a/vendor/github.com/jackc/pgx/v5/pgproto3/row_description.go b/vendor/github.com/jackc/pgx/v5/pgproto3/row_description.go new file mode 100644 index 0000000000..b46f510dc1 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgproto3/row_description.go @@ -0,0 +1,165 @@ +package pgproto3 + +import ( + "bytes" + "encoding/binary" + "encoding/json" + "errors" + "math" + + "github.com/jackc/pgx/v5/internal/pgio" +) + +const ( + TextFormat = 0 + BinaryFormat = 1 +) + +type FieldDescription struct { + Name []byte + TableOID uint32 + TableAttributeNumber uint16 + DataTypeOID uint32 + DataTypeSize int16 + TypeModifier int32 + Format int16 +} + +// MarshalJSON implements encoding/json.Marshaler. +func (fd FieldDescription) MarshalJSON() ([]byte, error) { + return json.Marshal(struct { + Name string + TableOID uint32 + TableAttributeNumber uint16 + DataTypeOID uint32 + DataTypeSize int16 + TypeModifier int32 + Format int16 + }{ + Name: string(fd.Name), + TableOID: fd.TableOID, + TableAttributeNumber: fd.TableAttributeNumber, + DataTypeOID: fd.DataTypeOID, + DataTypeSize: fd.DataTypeSize, + TypeModifier: fd.TypeModifier, + Format: fd.Format, + }) +} + +type RowDescription struct { + Fields []FieldDescription +} + +// Backend identifies this message as sendable by the PostgreSQL backend. +func (*RowDescription) Backend() {} + +// Decode decodes src into dst. src must contain the complete message with the exception of the initial 1 byte message +// type identifier and 4 byte message length. +func (dst *RowDescription) Decode(src []byte) error { + if len(src) < 2 { + return &invalidMessageFormatErr{messageType: "RowDescription"} + } + fieldCount := int(binary.BigEndian.Uint16(src)) + rp := 2 + + dst.Fields = dst.Fields[0:0] + + for range fieldCount { + var fd FieldDescription + + idx := bytes.IndexByte(src[rp:], 0) + if idx < 0 { + return &invalidMessageFormatErr{messageType: "RowDescription"} + } + fd.Name = src[rp : rp+idx] + rp += idx + 1 + + // Since buf.Next() doesn't return an error if we hit the end of the buffer + // check Len ahead of time + if len(src[rp:]) < 18 { + return &invalidMessageFormatErr{messageType: "RowDescription"} + } + + fd.TableOID = binary.BigEndian.Uint32(src[rp:]) + rp += 4 + fd.TableAttributeNumber = binary.BigEndian.Uint16(src[rp:]) + rp += 2 + fd.DataTypeOID = binary.BigEndian.Uint32(src[rp:]) + rp += 4 + fd.DataTypeSize = int16(binary.BigEndian.Uint16(src[rp:])) + rp += 2 + fd.TypeModifier = int32(binary.BigEndian.Uint32(src[rp:])) + rp += 4 + fd.Format = int16(binary.BigEndian.Uint16(src[rp:])) + rp += 2 + + dst.Fields = append(dst.Fields, fd) + } + + return nil +} + +// Encode encodes src into dst. dst will include the 1 byte message type identifier and the 4 byte message length. +func (src *RowDescription) Encode(dst []byte) ([]byte, error) { + dst, sp := beginMessage(dst, 'T') + + if len(src.Fields) > math.MaxUint16 { + return nil, errors.New("too many fields") + } + dst = pgio.AppendUint16(dst, uint16(len(src.Fields))) + for _, fd := range src.Fields { + dst = append(dst, fd.Name...) + dst = append(dst, 0) + + dst = pgio.AppendUint32(dst, fd.TableOID) + dst = pgio.AppendUint16(dst, fd.TableAttributeNumber) + dst = pgio.AppendUint32(dst, fd.DataTypeOID) + dst = pgio.AppendInt16(dst, fd.DataTypeSize) + dst = pgio.AppendInt32(dst, fd.TypeModifier) + dst = pgio.AppendInt16(dst, fd.Format) + } + + return finishMessage(dst, sp) +} + +// MarshalJSON implements encoding/json.Marshaler. +func (src RowDescription) MarshalJSON() ([]byte, error) { + return json.Marshal(struct { + Type string + Fields []FieldDescription + }{ + Type: "RowDescription", + Fields: src.Fields, + }) +} + +// UnmarshalJSON implements encoding/json.Unmarshaler. +func (dst *RowDescription) UnmarshalJSON(data []byte) error { + var msg struct { + Fields []struct { + Name string + TableOID uint32 + TableAttributeNumber uint16 + DataTypeOID uint32 + DataTypeSize int16 + TypeModifier int32 + Format int16 + } + } + if err := json.Unmarshal(data, &msg); err != nil { + return err + } + dst.Fields = make([]FieldDescription, len(msg.Fields)) + for n, field := range msg.Fields { + dst.Fields[n] = FieldDescription{ + Name: []byte(field.Name), + TableOID: field.TableOID, + TableAttributeNumber: field.TableAttributeNumber, + DataTypeOID: field.DataTypeOID, + DataTypeSize: field.DataTypeSize, + TypeModifier: field.TypeModifier, + Format: field.Format, + } + } + return nil +} diff --git a/vendor/github.com/jackc/pgx/v5/pgproto3/sasl_initial_response.go b/vendor/github.com/jackc/pgx/v5/pgproto3/sasl_initial_response.go new file mode 100644 index 0000000000..123f3cd66a --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgproto3/sasl_initial_response.go @@ -0,0 +1,93 @@ +package pgproto3 + +import ( + "bytes" + "encoding/hex" + "encoding/json" + "errors" + + "github.com/jackc/pgx/v5/internal/pgio" +) + +type SASLInitialResponse struct { + AuthMechanism string + Data []byte +} + +// Frontend identifies this message as sendable by a PostgreSQL frontend. +func (*SASLInitialResponse) Frontend() {} + +// Decode decodes src into dst. src must contain the complete message with the exception of the initial 1 byte message +// type identifier and 4 byte message length. +func (dst *SASLInitialResponse) Decode(src []byte) error { + *dst = SASLInitialResponse{} + + rp := 0 + + idx := bytes.IndexByte(src, 0) + if idx < 0 { + return errors.New("invalid SASLInitialResponse") + } + + dst.AuthMechanism = string(src[rp:idx]) + rp = idx + 1 + + if len(src[rp:]) < 4 { + return errors.New("invalid SASLInitialResponse") + } + rp += 4 // The rest of the message is data so we can just skip the size + dst.Data = src[rp:] + + return nil +} + +// Encode encodes src into dst. dst will include the 1 byte message type identifier and the 4 byte message length. +func (src *SASLInitialResponse) Encode(dst []byte) ([]byte, error) { + dst, sp := beginMessage(dst, 'p') + + dst = append(dst, []byte(src.AuthMechanism)...) + dst = append(dst, 0) + + dst = pgio.AppendInt32(dst, int32(len(src.Data))) + dst = append(dst, src.Data...) + + return finishMessage(dst, sp) +} + +// MarshalJSON implements encoding/json.Marshaler. +func (src SASLInitialResponse) MarshalJSON() ([]byte, error) { + return json.Marshal(struct { + Type string + AuthMechanism string + Data string + }{ + Type: "SASLInitialResponse", + AuthMechanism: src.AuthMechanism, + Data: string(src.Data), + }) +} + +// UnmarshalJSON implements encoding/json.Unmarshaler. +func (dst *SASLInitialResponse) UnmarshalJSON(data []byte) error { + // Ignore null, like in the main JSON package. + if string(data) == "null" { + return nil + } + + var msg struct { + AuthMechanism string + Data string + } + if err := json.Unmarshal(data, &msg); err != nil { + return err + } + dst.AuthMechanism = msg.AuthMechanism + if msg.Data != "" { + decoded, err := hex.DecodeString(msg.Data) + if err != nil { + return err + } + dst.Data = decoded + } + return nil +} diff --git a/vendor/github.com/jackc/pgx/v5/pgproto3/sasl_response.go b/vendor/github.com/jackc/pgx/v5/pgproto3/sasl_response.go new file mode 100644 index 0000000000..1b604c2542 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgproto3/sasl_response.go @@ -0,0 +1,56 @@ +package pgproto3 + +import ( + "encoding/hex" + "encoding/json" +) + +type SASLResponse struct { + Data []byte +} + +// Frontend identifies this message as sendable by a PostgreSQL frontend. +func (*SASLResponse) Frontend() {} + +// Decode decodes src into dst. src must contain the complete message with the exception of the initial 1 byte message +// type identifier and 4 byte message length. +func (dst *SASLResponse) Decode(src []byte) error { + *dst = SASLResponse{Data: src} + return nil +} + +// Encode encodes src into dst. dst will include the 1 byte message type identifier and the 4 byte message length. +func (src *SASLResponse) Encode(dst []byte) ([]byte, error) { + dst, sp := beginMessage(dst, 'p') + dst = append(dst, src.Data...) + return finishMessage(dst, sp) +} + +// MarshalJSON implements encoding/json.Marshaler. +func (src SASLResponse) MarshalJSON() ([]byte, error) { + return json.Marshal(struct { + Type string + Data string + }{ + Type: "SASLResponse", + Data: string(src.Data), + }) +} + +// UnmarshalJSON implements encoding/json.Unmarshaler. +func (dst *SASLResponse) UnmarshalJSON(data []byte) error { + var msg struct { + Data string + } + if err := json.Unmarshal(data, &msg); err != nil { + return err + } + if msg.Data != "" { + decoded, err := hex.DecodeString(msg.Data) + if err != nil { + return err + } + dst.Data = decoded + } + return nil +} diff --git a/vendor/github.com/jackc/pgx/v5/pgproto3/ssl_request.go b/vendor/github.com/jackc/pgx/v5/pgproto3/ssl_request.go new file mode 100644 index 0000000000..bdfc7c427a --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgproto3/ssl_request.go @@ -0,0 +1,48 @@ +package pgproto3 + +import ( + "encoding/binary" + "encoding/json" + "errors" + + "github.com/jackc/pgx/v5/internal/pgio" +) + +const sslRequestNumber = 80877103 + +type SSLRequest struct{} + +// Frontend identifies this message as sendable by a PostgreSQL frontend. +func (*SSLRequest) Frontend() {} + +func (dst *SSLRequest) Decode(src []byte) error { + if len(src) < 4 { + return errors.New("ssl request too short") + } + + requestCode := binary.BigEndian.Uint32(src) + + if requestCode != sslRequestNumber { + return errors.New("bad ssl request code") + } + + return nil +} + +// Encode encodes src into dst. dst will include the 4 byte message length. +func (src *SSLRequest) Encode(dst []byte) ([]byte, error) { + dst = pgio.AppendInt32(dst, 8) + dst = pgio.AppendInt32(dst, sslRequestNumber) + return dst, nil +} + +// MarshalJSON implements encoding/json.Marshaler. +func (src SSLRequest) MarshalJSON() ([]byte, error) { + return json.Marshal(struct { + Type string + ProtocolVersion uint32 + Parameters map[string]string + }{ + Type: "SSLRequest", + }) +} diff --git a/vendor/github.com/jackc/pgx/v5/pgproto3/startup_message.go b/vendor/github.com/jackc/pgx/v5/pgproto3/startup_message.go new file mode 100644 index 0000000000..eb48f72bf0 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgproto3/startup_message.go @@ -0,0 +1,99 @@ +package pgproto3 + +import ( + "bytes" + "encoding/binary" + "encoding/json" + "errors" + "fmt" + + "github.com/jackc/pgx/v5/internal/pgio" +) + +const ( + ProtocolVersion30 = 196608 // 3.0 + ProtocolVersion32 = 196610 // 3.2 + ProtocolVersionLatest = ProtocolVersion32 // Latest is 3.2 + ProtocolVersionNumber = ProtocolVersion30 // Default is still 3.0 +) + +type StartupMessage struct { + ProtocolVersion uint32 + Parameters map[string]string +} + +// Frontend identifies this message as sendable by a PostgreSQL frontend. +func (*StartupMessage) Frontend() {} + +// Decode decodes src into dst. src must contain the complete message with the exception of the initial 1 byte message +// type identifier and 4 byte message length. +func (dst *StartupMessage) Decode(src []byte) error { + if len(src) < 4 { + return errors.New("startup message too short") + } + + dst.ProtocolVersion = binary.BigEndian.Uint32(src) + rp := 4 + + if dst.ProtocolVersion != ProtocolVersion30 && dst.ProtocolVersion != ProtocolVersion32 { + return fmt.Errorf("Bad startup message version number. Expected %d or %d, got %d", ProtocolVersion30, ProtocolVersion32, dst.ProtocolVersion) + } + + dst.Parameters = make(map[string]string) + for { + idx := bytes.IndexByte(src[rp:], 0) + if idx < 0 { + return &invalidMessageFormatErr{messageType: "StartupMessage"} + } + key := string(src[rp : rp+idx]) + rp += idx + 1 + + idx = bytes.IndexByte(src[rp:], 0) + if idx < 0 { + return &invalidMessageFormatErr{messageType: "StartupMessage"} + } + value := string(src[rp : rp+idx]) + rp += idx + 1 + + dst.Parameters[key] = value + + if len(src[rp:]) == 1 { + if src[rp] != 0 { + return fmt.Errorf("Bad startup message last byte. Expected 0, got %d", src[rp]) + } + break + } + } + + return nil +} + +// Encode encodes src into dst. dst will include the 1 byte message type identifier and the 4 byte message length. +func (src *StartupMessage) Encode(dst []byte) ([]byte, error) { + sp := len(dst) + dst = pgio.AppendInt32(dst, -1) + + dst = pgio.AppendUint32(dst, src.ProtocolVersion) + for k, v := range src.Parameters { + dst = append(dst, k...) + dst = append(dst, 0) + dst = append(dst, v...) + dst = append(dst, 0) + } + dst = append(dst, 0) + + return finishMessage(dst, sp) +} + +// MarshalJSON implements encoding/json.Marshaler. +func (src StartupMessage) MarshalJSON() ([]byte, error) { + return json.Marshal(struct { + Type string + ProtocolVersion uint32 + Parameters map[string]string + }{ + Type: "StartupMessage", + ProtocolVersion: src.ProtocolVersion, + Parameters: src.Parameters, + }) +} diff --git a/vendor/github.com/jackc/pgx/v5/pgproto3/sync.go b/vendor/github.com/jackc/pgx/v5/pgproto3/sync.go new file mode 100644 index 0000000000..ea4fc9594c --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgproto3/sync.go @@ -0,0 +1,34 @@ +package pgproto3 + +import ( + "encoding/json" +) + +type Sync struct{} + +// Frontend identifies this message as sendable by a PostgreSQL frontend. +func (*Sync) Frontend() {} + +// Decode decodes src into dst. src must contain the complete message with the exception of the initial 1 byte message +// type identifier and 4 byte message length. +func (dst *Sync) Decode(src []byte) error { + if len(src) != 0 { + return &invalidMessageLenErr{messageType: "Sync", expectedLen: 0, actualLen: len(src)} + } + + return nil +} + +// Encode encodes src into dst. dst will include the 1 byte message type identifier and the 4 byte message length. +func (src *Sync) Encode(dst []byte) ([]byte, error) { + return append(dst, 'S', 0, 0, 0, 4), nil +} + +// MarshalJSON implements encoding/json.Marshaler. +func (src Sync) MarshalJSON() ([]byte, error) { + return json.Marshal(struct { + Type string + }{ + Type: "Sync", + }) +} diff --git a/vendor/github.com/jackc/pgx/v5/pgproto3/terminate.go b/vendor/github.com/jackc/pgx/v5/pgproto3/terminate.go new file mode 100644 index 0000000000..35a9dc837d --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgproto3/terminate.go @@ -0,0 +1,34 @@ +package pgproto3 + +import ( + "encoding/json" +) + +type Terminate struct{} + +// Frontend identifies this message as sendable by a PostgreSQL frontend. +func (*Terminate) Frontend() {} + +// Decode decodes src into dst. src must contain the complete message with the exception of the initial 1 byte message +// type identifier and 4 byte message length. +func (dst *Terminate) Decode(src []byte) error { + if len(src) != 0 { + return &invalidMessageLenErr{messageType: "Terminate", expectedLen: 0, actualLen: len(src)} + } + + return nil +} + +// Encode encodes src into dst. dst will include the 1 byte message type identifier and the 4 byte message length. +func (src *Terminate) Encode(dst []byte) ([]byte, error) { + return append(dst, 'X', 0, 0, 0, 4), nil +} + +// MarshalJSON implements encoding/json.Marshaler. +func (src Terminate) MarshalJSON() ([]byte, error) { + return json.Marshal(struct { + Type string + }{ + Type: "Terminate", + }) +} diff --git a/vendor/github.com/jackc/pgx/v5/pgproto3/trace.go b/vendor/github.com/jackc/pgx/v5/pgproto3/trace.go new file mode 100644 index 0000000000..2f9da6289c --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgproto3/trace.go @@ -0,0 +1,416 @@ +package pgproto3 + +import ( + "bytes" + "fmt" + "io" + "strconv" + "strings" + "sync" + "time" +) + +// tracer traces the messages send to and from a Backend or Frontend. The format it produces roughly mimics the +// format produced by the libpq C function PQtrace. +type tracer struct { + TracerOptions + + mux sync.Mutex + w io.Writer + buf *bytes.Buffer +} + +// TracerOptions controls tracing behavior. It is roughly equivalent to the libpq function PQsetTraceFlags. +type TracerOptions struct { + // SuppressTimestamps prevents printing of timestamps. + SuppressTimestamps bool + + // RegressMode redacts fields that may be vary between executions. + RegressMode bool +} + +func (t *tracer) traceMessage(sender byte, encodedLen int32, msg Message) { + switch msg := msg.(type) { + case *AuthenticationCleartextPassword: + t.traceAuthenticationCleartextPassword(sender, encodedLen, msg) + case *AuthenticationGSS: + t.traceAuthenticationGSS(sender, encodedLen, msg) + case *AuthenticationGSSContinue: + t.traceAuthenticationGSSContinue(sender, encodedLen, msg) + case *AuthenticationMD5Password: + t.traceAuthenticationMD5Password(sender, encodedLen, msg) + case *AuthenticationOk: + t.traceAuthenticationOk(sender, encodedLen, msg) + case *AuthenticationSASL: + t.traceAuthenticationSASL(sender, encodedLen, msg) + case *AuthenticationSASLContinue: + t.traceAuthenticationSASLContinue(sender, encodedLen, msg) + case *AuthenticationSASLFinal: + t.traceAuthenticationSASLFinal(sender, encodedLen, msg) + case *BackendKeyData: + t.traceBackendKeyData(sender, encodedLen, msg) + case *Bind: + t.traceBind(sender, encodedLen, msg) + case *BindComplete: + t.traceBindComplete(sender, encodedLen, msg) + case *CancelRequest: + t.traceCancelRequest(sender, encodedLen, msg) + case *Close: + t.traceClose(sender, encodedLen, msg) + case *CloseComplete: + t.traceCloseComplete(sender, encodedLen, msg) + case *CommandComplete: + t.traceCommandComplete(sender, encodedLen, msg) + case *CopyBothResponse: + t.traceCopyBothResponse(sender, encodedLen, msg) + case *CopyData: + t.traceCopyData(sender, encodedLen, msg) + case *CopyDone: + t.traceCopyDone(sender, encodedLen, msg) + case *CopyFail: + t.traceCopyFail(sender, encodedLen, msg) + case *CopyInResponse: + t.traceCopyInResponse(sender, encodedLen, msg) + case *CopyOutResponse: + t.traceCopyOutResponse(sender, encodedLen, msg) + case *DataRow: + t.traceDataRow(sender, encodedLen, msg) + case *Describe: + t.traceDescribe(sender, encodedLen, msg) + case *EmptyQueryResponse: + t.traceEmptyQueryResponse(sender, encodedLen, msg) + case *ErrorResponse: + t.traceErrorResponse(sender, encodedLen, msg) + case *Execute: + t.traceExecute(sender, encodedLen, msg) + case *Flush: + t.traceFlush(sender, encodedLen, msg) + case *FunctionCall: + t.traceFunctionCall(sender, encodedLen, msg) + case *FunctionCallResponse: + t.traceFunctionCallResponse(sender, encodedLen, msg) + case *GSSEncRequest: + t.traceGSSEncRequest(sender, encodedLen, msg) + case *NoData: + t.traceNoData(sender, encodedLen, msg) + case *NoticeResponse: + t.traceNoticeResponse(sender, encodedLen, msg) + case *NotificationResponse: + t.traceNotificationResponse(sender, encodedLen, msg) + case *ParameterDescription: + t.traceParameterDescription(sender, encodedLen, msg) + case *ParameterStatus: + t.traceParameterStatus(sender, encodedLen, msg) + case *Parse: + t.traceParse(sender, encodedLen, msg) + case *ParseComplete: + t.traceParseComplete(sender, encodedLen, msg) + case *PortalSuspended: + t.tracePortalSuspended(sender, encodedLen, msg) + case *Query: + t.traceQuery(sender, encodedLen, msg) + case *ReadyForQuery: + t.traceReadyForQuery(sender, encodedLen, msg) + case *RowDescription: + t.traceRowDescription(sender, encodedLen, msg) + case *SSLRequest: + t.traceSSLRequest(sender, encodedLen, msg) + case *StartupMessage: + t.traceStartupMessage(sender, encodedLen, msg) + case *Sync: + t.traceSync(sender, encodedLen, msg) + case *Terminate: + t.traceTerminate(sender, encodedLen, msg) + default: + t.writeTrace(sender, encodedLen, "Unknown", nil) + } +} + +func (t *tracer) traceAuthenticationCleartextPassword(sender byte, encodedLen int32, msg *AuthenticationCleartextPassword) { + t.writeTrace(sender, encodedLen, "AuthenticationCleartextPassword", nil) +} + +func (t *tracer) traceAuthenticationGSS(sender byte, encodedLen int32, msg *AuthenticationGSS) { + t.writeTrace(sender, encodedLen, "AuthenticationGSS", nil) +} + +func (t *tracer) traceAuthenticationGSSContinue(sender byte, encodedLen int32, msg *AuthenticationGSSContinue) { + t.writeTrace(sender, encodedLen, "AuthenticationGSSContinue", nil) +} + +func (t *tracer) traceAuthenticationMD5Password(sender byte, encodedLen int32, msg *AuthenticationMD5Password) { + t.writeTrace(sender, encodedLen, "AuthenticationMD5Password", nil) +} + +func (t *tracer) traceAuthenticationOk(sender byte, encodedLen int32, msg *AuthenticationOk) { + t.writeTrace(sender, encodedLen, "AuthenticationOk", nil) +} + +func (t *tracer) traceAuthenticationSASL(sender byte, encodedLen int32, msg *AuthenticationSASL) { + t.writeTrace(sender, encodedLen, "AuthenticationSASL", nil) +} + +func (t *tracer) traceAuthenticationSASLContinue(sender byte, encodedLen int32, msg *AuthenticationSASLContinue) { + t.writeTrace(sender, encodedLen, "AuthenticationSASLContinue", nil) +} + +func (t *tracer) traceAuthenticationSASLFinal(sender byte, encodedLen int32, msg *AuthenticationSASLFinal) { + t.writeTrace(sender, encodedLen, "AuthenticationSASLFinal", nil) +} + +func (t *tracer) traceBackendKeyData(sender byte, encodedLen int32, msg *BackendKeyData) { + t.writeTrace(sender, encodedLen, "BackendKeyData", func() { + if t.RegressMode { + t.buf.WriteString("\t NNNN NNNN") + } else { + fmt.Fprintf(t.buf, "\t %d %d", msg.ProcessID, msg.SecretKey) + } + }) +} + +func (t *tracer) traceBind(sender byte, encodedLen int32, msg *Bind) { + t.writeTrace(sender, encodedLen, "Bind", func() { + fmt.Fprintf(t.buf, "\t %s %s %d", traceDoubleQuotedString([]byte(msg.DestinationPortal)), traceDoubleQuotedString([]byte(msg.PreparedStatement)), len(msg.ParameterFormatCodes)) + for _, fc := range msg.ParameterFormatCodes { + fmt.Fprintf(t.buf, " %d", fc) + } + fmt.Fprintf(t.buf, " %d", len(msg.Parameters)) + for _, p := range msg.Parameters { + fmt.Fprintf(t.buf, " %s", traceSingleQuotedString(p)) + } + fmt.Fprintf(t.buf, " %d", len(msg.ResultFormatCodes)) + for _, fc := range msg.ResultFormatCodes { + fmt.Fprintf(t.buf, " %d", fc) + } + }) +} + +func (t *tracer) traceBindComplete(sender byte, encodedLen int32, msg *BindComplete) { + t.writeTrace(sender, encodedLen, "BindComplete", nil) +} + +func (t *tracer) traceCancelRequest(sender byte, encodedLen int32, msg *CancelRequest) { + t.writeTrace(sender, encodedLen, "CancelRequest", nil) +} + +func (t *tracer) traceClose(sender byte, encodedLen int32, msg *Close) { + t.writeTrace(sender, encodedLen, "Close", nil) +} + +func (t *tracer) traceCloseComplete(sender byte, encodedLen int32, msg *CloseComplete) { + t.writeTrace(sender, encodedLen, "CloseComplete", nil) +} + +func (t *tracer) traceCommandComplete(sender byte, encodedLen int32, msg *CommandComplete) { + t.writeTrace(sender, encodedLen, "CommandComplete", func() { + fmt.Fprintf(t.buf, "\t %s", traceDoubleQuotedString(msg.CommandTag)) + }) +} + +func (t *tracer) traceCopyBothResponse(sender byte, encodedLen int32, msg *CopyBothResponse) { + t.writeTrace(sender, encodedLen, "CopyBothResponse", nil) +} + +func (t *tracer) traceCopyData(sender byte, encodedLen int32, msg *CopyData) { + t.writeTrace(sender, encodedLen, "CopyData", nil) +} + +func (t *tracer) traceCopyDone(sender byte, encodedLen int32, msg *CopyDone) { + t.writeTrace(sender, encodedLen, "CopyDone", nil) +} + +func (t *tracer) traceCopyFail(sender byte, encodedLen int32, msg *CopyFail) { + t.writeTrace(sender, encodedLen, "CopyFail", func() { + fmt.Fprintf(t.buf, "\t %s", traceDoubleQuotedString([]byte(msg.Message))) + }) +} + +func (t *tracer) traceCopyInResponse(sender byte, encodedLen int32, msg *CopyInResponse) { + t.writeTrace(sender, encodedLen, "CopyInResponse", nil) +} + +func (t *tracer) traceCopyOutResponse(sender byte, encodedLen int32, msg *CopyOutResponse) { + t.writeTrace(sender, encodedLen, "CopyOutResponse", nil) +} + +func (t *tracer) traceDataRow(sender byte, encodedLen int32, msg *DataRow) { + t.writeTrace(sender, encodedLen, "DataRow", func() { + fmt.Fprintf(t.buf, "\t %d", len(msg.Values)) + for _, v := range msg.Values { + if v == nil { + t.buf.WriteString(" -1") + } else { + fmt.Fprintf(t.buf, " %d %s", len(v), traceSingleQuotedString(v)) + } + } + }) +} + +func (t *tracer) traceDescribe(sender byte, encodedLen int32, msg *Describe) { + t.writeTrace(sender, encodedLen, "Describe", func() { + fmt.Fprintf(t.buf, "\t %c %s", msg.ObjectType, traceDoubleQuotedString([]byte(msg.Name))) + }) +} + +func (t *tracer) traceEmptyQueryResponse(sender byte, encodedLen int32, msg *EmptyQueryResponse) { + t.writeTrace(sender, encodedLen, "EmptyQueryResponse", nil) +} + +func (t *tracer) traceErrorResponse(sender byte, encodedLen int32, msg *ErrorResponse) { + t.writeTrace(sender, encodedLen, "ErrorResponse", nil) +} + +func (t *tracer) traceExecute(sender byte, encodedLen int32, msg *Execute) { + t.writeTrace(sender, encodedLen, "Execute", func() { + fmt.Fprintf(t.buf, "\t %s %d", traceDoubleQuotedString([]byte(msg.Portal)), msg.MaxRows) + }) +} + +func (t *tracer) traceFlush(sender byte, encodedLen int32, msg *Flush) { + t.writeTrace(sender, encodedLen, "Flush", nil) +} + +func (t *tracer) traceFunctionCall(sender byte, encodedLen int32, msg *FunctionCall) { + t.writeTrace(sender, encodedLen, "FunctionCall", nil) +} + +func (t *tracer) traceFunctionCallResponse(sender byte, encodedLen int32, msg *FunctionCallResponse) { + t.writeTrace(sender, encodedLen, "FunctionCallResponse", nil) +} + +func (t *tracer) traceGSSEncRequest(sender byte, encodedLen int32, msg *GSSEncRequest) { + t.writeTrace(sender, encodedLen, "GSSEncRequest", nil) +} + +func (t *tracer) traceNoData(sender byte, encodedLen int32, msg *NoData) { + t.writeTrace(sender, encodedLen, "NoData", nil) +} + +func (t *tracer) traceNoticeResponse(sender byte, encodedLen int32, msg *NoticeResponse) { + t.writeTrace(sender, encodedLen, "NoticeResponse", nil) +} + +func (t *tracer) traceNotificationResponse(sender byte, encodedLen int32, msg *NotificationResponse) { + t.writeTrace(sender, encodedLen, "NotificationResponse", func() { + fmt.Fprintf(t.buf, "\t %d %s %s", msg.PID, traceDoubleQuotedString([]byte(msg.Channel)), traceDoubleQuotedString([]byte(msg.Payload))) + }) +} + +func (t *tracer) traceParameterDescription(sender byte, encodedLen int32, msg *ParameterDescription) { + t.writeTrace(sender, encodedLen, "ParameterDescription", nil) +} + +func (t *tracer) traceParameterStatus(sender byte, encodedLen int32, msg *ParameterStatus) { + t.writeTrace(sender, encodedLen, "ParameterStatus", func() { + fmt.Fprintf(t.buf, "\t %s %s", traceDoubleQuotedString([]byte(msg.Name)), traceDoubleQuotedString([]byte(msg.Value))) + }) +} + +func (t *tracer) traceParse(sender byte, encodedLen int32, msg *Parse) { + t.writeTrace(sender, encodedLen, "Parse", func() { + fmt.Fprintf(t.buf, "\t %s %s %d", traceDoubleQuotedString([]byte(msg.Name)), traceDoubleQuotedString([]byte(msg.Query)), len(msg.ParameterOIDs)) + for _, oid := range msg.ParameterOIDs { + fmt.Fprintf(t.buf, " %d", oid) + } + }) +} + +func (t *tracer) traceParseComplete(sender byte, encodedLen int32, msg *ParseComplete) { + t.writeTrace(sender, encodedLen, "ParseComplete", nil) +} + +func (t *tracer) tracePortalSuspended(sender byte, encodedLen int32, msg *PortalSuspended) { + t.writeTrace(sender, encodedLen, "PortalSuspended", nil) +} + +func (t *tracer) traceQuery(sender byte, encodedLen int32, msg *Query) { + t.writeTrace(sender, encodedLen, "Query", func() { + fmt.Fprintf(t.buf, "\t %s", traceDoubleQuotedString([]byte(msg.String))) + }) +} + +func (t *tracer) traceReadyForQuery(sender byte, encodedLen int32, msg *ReadyForQuery) { + t.writeTrace(sender, encodedLen, "ReadyForQuery", func() { + fmt.Fprintf(t.buf, "\t %c", msg.TxStatus) + }) +} + +func (t *tracer) traceRowDescription(sender byte, encodedLen int32, msg *RowDescription) { + t.writeTrace(sender, encodedLen, "RowDescription", func() { + fmt.Fprintf(t.buf, "\t %d", len(msg.Fields)) + for _, fd := range msg.Fields { + fmt.Fprintf(t.buf, ` %s %d %d %d %d %d %d`, traceDoubleQuotedString(fd.Name), fd.TableOID, fd.TableAttributeNumber, fd.DataTypeOID, fd.DataTypeSize, fd.TypeModifier, fd.Format) + } + }) +} + +func (t *tracer) traceSSLRequest(sender byte, encodedLen int32, msg *SSLRequest) { + t.writeTrace(sender, encodedLen, "SSLRequest", nil) +} + +func (t *tracer) traceStartupMessage(sender byte, encodedLen int32, msg *StartupMessage) { + t.writeTrace(sender, encodedLen, "StartupMessage", nil) +} + +func (t *tracer) traceSync(sender byte, encodedLen int32, msg *Sync) { + t.writeTrace(sender, encodedLen, "Sync", nil) +} + +func (t *tracer) traceTerminate(sender byte, encodedLen int32, msg *Terminate) { + t.writeTrace(sender, encodedLen, "Terminate", nil) +} + +func (t *tracer) writeTrace(sender byte, encodedLen int32, msgType string, writeDetails func()) { + t.mux.Lock() + defer t.mux.Unlock() + defer func() { + if t.buf.Cap() > 1024 { + t.buf = &bytes.Buffer{} + } else { + t.buf.Reset() + } + }() + + if !t.SuppressTimestamps { + now := time.Now() + t.buf.WriteString(now.Format("2006-01-02 15:04:05.000000")) + t.buf.WriteByte('\t') + } + + t.buf.WriteByte(sender) + t.buf.WriteByte('\t') + t.buf.WriteString(msgType) + t.buf.WriteByte('\t') + t.buf.WriteString(strconv.FormatInt(int64(encodedLen), 10)) + + if writeDetails != nil { + writeDetails() + } + + t.buf.WriteByte('\n') + t.buf.WriteTo(t.w) +} + +// traceDoubleQuotedString returns t.buf as a double-quoted string without any escaping. It is roughly equivalent to +// pqTraceOutputString in libpq. +func traceDoubleQuotedString(buf []byte) string { + return `"` + string(buf) + `"` +} + +// traceSingleQuotedString returns buf as a single-quoted string with non-printable characters hex-escaped. It is +// roughly equivalent to pqTraceOutputNchar in libpq. +func traceSingleQuotedString(buf []byte) string { + sb := &strings.Builder{} + + sb.WriteByte('\'') + for _, b := range buf { + if b < 32 || b > 126 { + fmt.Fprintf(sb, `\x%x`, b) + } else { + sb.WriteByte(b) + } + } + sb.WriteByte('\'') + + return sb.String() +} diff --git a/vendor/github.com/jackc/pgx/v5/pgtype/array.go b/vendor/github.com/jackc/pgx/v5/pgtype/array.go new file mode 100644 index 0000000000..26505fb841 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgtype/array.go @@ -0,0 +1,469 @@ +package pgtype + +import ( + "bytes" + "encoding/binary" + "fmt" + "io" + "strconv" + "strings" + "unicode" + + "github.com/jackc/pgx/v5/internal/pgio" +) + +// Information on the internals of PostgreSQL arrays can be found in +// src/include/utils/array.h and src/backend/utils/adt/arrayfuncs.c. Of +// particular interest is the array_send function. + +type arrayHeader struct { + ContainsNull bool + ElementOID uint32 + Dimensions []ArrayDimension +} + +type ArrayDimension struct { + Length int32 + LowerBound int32 +} + +// cardinality returns the number of elements in an array of dimensions size. +func cardinality(dimensions []ArrayDimension) int { + if len(dimensions) == 0 { + return 0 + } + + elementCount := int(dimensions[0].Length) + for _, d := range dimensions[1:] { + elementCount *= int(d.Length) + } + + if elementCount < 0 { + return 0 + } + + return elementCount +} + +func (dst *arrayHeader) DecodeBinary(m *Map, src []byte) (int, error) { + if len(src) < 12 { + return 0, fmt.Errorf("array header too short: %d", len(src)) + } + + rp := 0 + + numDims := int(binary.BigEndian.Uint32(src[rp:])) + rp += 4 + + if numDims > 6 { + return 0, fmt.Errorf("array has too many dimensions: %d", numDims) + } + + dst.ContainsNull = binary.BigEndian.Uint32(src[rp:]) == 1 + rp += 4 + + dst.ElementOID = binary.BigEndian.Uint32(src[rp:]) + rp += 4 + + if len(src) < 12+numDims*8 { + return 0, fmt.Errorf("array header too short for %d dimensions: %d", numDims, len(src)) + } + dst.Dimensions = make([]ArrayDimension, numDims) + for i := range dst.Dimensions { + dst.Dimensions[i].Length = int32(binary.BigEndian.Uint32(src[rp:])) + rp += 4 + + dst.Dimensions[i].LowerBound = int32(binary.BigEndian.Uint32(src[rp:])) + rp += 4 + } + + return rp, nil +} + +func (src arrayHeader) EncodeBinary(buf []byte) []byte { + buf = pgio.AppendInt32(buf, int32(len(src.Dimensions))) + + var containsNull int32 + if src.ContainsNull { + containsNull = 1 + } + buf = pgio.AppendInt32(buf, containsNull) + + buf = pgio.AppendUint32(buf, src.ElementOID) + + for i := range src.Dimensions { + buf = pgio.AppendInt32(buf, src.Dimensions[i].Length) + buf = pgio.AppendInt32(buf, src.Dimensions[i].LowerBound) + } + + return buf +} + +type untypedTextArray struct { + Elements []string + Quoted []bool + Dimensions []ArrayDimension +} + +func parseUntypedTextArray(src string) (*untypedTextArray, error) { + dst := &untypedTextArray{ + Elements: []string{}, + Quoted: []bool{}, + Dimensions: []ArrayDimension{}, + } + + buf := bytes.NewBufferString(src) + + skipWhitespace(buf) + + r, _, err := buf.ReadRune() + if err != nil { + return nil, fmt.Errorf("invalid array: %w", err) + } + + var explicitDimensions []ArrayDimension + + // Array has explicit dimensions + if r == '[' { + buf.UnreadRune() + + for { + r, _, err = buf.ReadRune() + if err != nil { + return nil, fmt.Errorf("invalid array: %w", err) + } + + if r == '=' { + break + } else if r != '[' { + return nil, fmt.Errorf("invalid array, expected '[' or '=' got %v", r) + } + + lower, err := arrayParseInteger(buf) + if err != nil { + return nil, fmt.Errorf("invalid array: %w", err) + } + + r, _, err = buf.ReadRune() + if err != nil { + return nil, fmt.Errorf("invalid array: %w", err) + } + + if r != ':' { + return nil, fmt.Errorf("invalid array, expected ':' got %v", r) + } + + upper, err := arrayParseInteger(buf) + if err != nil { + return nil, fmt.Errorf("invalid array: %w", err) + } + + r, _, err = buf.ReadRune() + if err != nil { + return nil, fmt.Errorf("invalid array: %w", err) + } + + if r != ']' { + return nil, fmt.Errorf("invalid array, expected ']' got %v", r) + } + + explicitDimensions = append(explicitDimensions, ArrayDimension{LowerBound: lower, Length: upper - lower + 1}) + } + + r, _, err = buf.ReadRune() + if err != nil { + return nil, fmt.Errorf("invalid array: %w", err) + } + } + + if r != '{' { + return nil, fmt.Errorf("invalid array, expected '{' got %v", r) + } + + implicitDimensions := []ArrayDimension{{LowerBound: 1, Length: 0}} + + // Consume all initial opening brackets. This provides number of dimensions. + for { + r, _, err = buf.ReadRune() + if err != nil { + return nil, fmt.Errorf("invalid array: %w", err) + } + + if r == '{' { + implicitDimensions[len(implicitDimensions)-1].Length = 1 + implicitDimensions = append(implicitDimensions, ArrayDimension{LowerBound: 1}) + } else { + buf.UnreadRune() + break + } + } + currentDim := len(implicitDimensions) - 1 + counterDim := currentDim + + for { + r, _, err = buf.ReadRune() + if err != nil { + return nil, fmt.Errorf("invalid array: %w", err) + } + + switch r { + case '{': + if currentDim == counterDim { + implicitDimensions[currentDim].Length++ + } + currentDim++ + case ',': + case '}': + currentDim-- + if currentDim < counterDim { + counterDim = currentDim + } + default: + buf.UnreadRune() + value, quoted, err := arrayParseValue(buf) + if err != nil { + return nil, fmt.Errorf("invalid array value: %w", err) + } + if currentDim == counterDim { + implicitDimensions[currentDim].Length++ + } + dst.Quoted = append(dst.Quoted, quoted) + dst.Elements = append(dst.Elements, value) + } + + if currentDim < 0 { + break + } + } + + skipWhitespace(buf) + + if buf.Len() > 0 { + return nil, fmt.Errorf("unexpected trailing data: %v", buf.String()) + } + + switch { + case len(dst.Elements) == 0: + case len(explicitDimensions) > 0: + dst.Dimensions = explicitDimensions + default: + dst.Dimensions = implicitDimensions + } + + return dst, nil +} + +func skipWhitespace(buf *bytes.Buffer) { + var r rune + var err error + for r, _, _ = buf.ReadRune(); unicode.IsSpace(r); r, _, _ = buf.ReadRune() { + } + + if err != io.EOF { + buf.UnreadRune() + } +} + +func arrayParseValue(buf *bytes.Buffer) (string, bool, error) { + r, _, err := buf.ReadRune() + if err != nil { + return "", false, err + } + if r == '"' { + return arrayParseQuotedValue(buf) + } + buf.UnreadRune() + + s := &bytes.Buffer{} + + for { + r, _, err := buf.ReadRune() + if err != nil { + return "", false, err + } + + switch r { + case ',', '}': + buf.UnreadRune() + return s.String(), false, nil + } + + s.WriteRune(r) + } +} + +func arrayParseQuotedValue(buf *bytes.Buffer) (string, bool, error) { + s := &bytes.Buffer{} + + for { + r, _, err := buf.ReadRune() + if err != nil { + return "", false, err + } + + switch r { + case '\\': + r, _, err = buf.ReadRune() + if err != nil { + return "", false, err + } + case '"': + _, _, err = buf.ReadRune() + if err != nil { + return "", false, err + } + buf.UnreadRune() + return s.String(), true, nil + } + s.WriteRune(r) + } +} + +func arrayParseInteger(buf *bytes.Buffer) (int32, error) { + s := &bytes.Buffer{} + + for { + r, _, err := buf.ReadRune() + if err != nil { + return 0, err + } + + if ('0' <= r && r <= '9') || r == '-' { + s.WriteRune(r) + } else { + buf.UnreadRune() + n, err := strconv.ParseInt(s.String(), 10, 32) + if err != nil { + return 0, err + } + return int32(n), nil + } + } +} + +func encodeTextArrayDimensions(buf []byte, dimensions []ArrayDimension) []byte { + var customDimensions bool + for _, dim := range dimensions { + if dim.LowerBound != 1 { + customDimensions = true + } + } + + if !customDimensions { + return buf + } + + for _, dim := range dimensions { + buf = append(buf, '[') + buf = append(buf, strconv.FormatInt(int64(dim.LowerBound), 10)...) + buf = append(buf, ':') + buf = append(buf, strconv.FormatInt(int64(dim.LowerBound+dim.Length-1), 10)...) + buf = append(buf, ']') + } + + return append(buf, '=') +} + +var quoteArrayReplacer = strings.NewReplacer(`\`, `\\`, `"`, `\"`) + +func quoteArrayElement(src string) string { + return `"` + quoteArrayReplacer.Replace(src) + `"` +} + +func isSpace(ch byte) bool { + // see array_isspace: + // https://github.com/postgres/postgres/blob/master/src/backend/utils/adt/arrayfuncs.c + return ch == ' ' || ch == '\t' || ch == '\n' || ch == '\r' || ch == '\v' || ch == '\f' +} + +func quoteArrayElementIfNeeded(src string) string { + if src == "" || (len(src) == 4 && strings.EqualFold(src, "null")) || isSpace(src[0]) || isSpace(src[len(src)-1]) || strings.ContainsAny(src, `{},"\`) { + return quoteArrayElement(src) + } + return src +} + +// Array represents a PostgreSQL array for T. It implements the [ArrayGetter] and [ArraySetter] interfaces. It preserves +// PostgreSQL dimensions and custom lower bounds. Use [FlatArray] if these are not needed. +type Array[T any] struct { + Elements []T + Dims []ArrayDimension + Valid bool +} + +func (a Array[T]) Dimensions() []ArrayDimension { + return a.Dims +} + +func (a Array[T]) Index(i int) any { + return a.Elements[i] +} + +func (a Array[T]) IndexType() any { + var el T + return el +} + +func (a *Array[T]) SetDimensions(dimensions []ArrayDimension) error { + if dimensions == nil { + *a = Array[T]{} + return nil + } + + elementCount := cardinality(dimensions) + *a = Array[T]{ + Elements: make([]T, elementCount), + Dims: dimensions, + Valid: true, + } + + return nil +} + +func (a Array[T]) ScanIndex(i int) any { + return &a.Elements[i] +} + +func (a Array[T]) ScanIndexType() any { + return new(T) +} + +// FlatArray implements the [ArrayGetter] and [ArraySetter] interfaces for any slice of T. It ignores PostgreSQL dimensions +// and custom lower bounds. Use [Array] to preserve these. +type FlatArray[T any] []T + +func (a FlatArray[T]) Dimensions() []ArrayDimension { + if a == nil { + return nil + } + + return []ArrayDimension{{Length: int32(len(a)), LowerBound: 1}} +} + +func (a FlatArray[T]) Index(i int) any { + return a[i] +} + +func (a FlatArray[T]) IndexType() any { + var el T + return el +} + +func (a *FlatArray[T]) SetDimensions(dimensions []ArrayDimension) error { + if dimensions == nil { + *a = nil + return nil + } + + elementCount := cardinality(dimensions) + *a = make(FlatArray[T], elementCount) + return nil +} + +func (a FlatArray[T]) ScanIndex(i int) any { + return &a[i] +} + +func (a FlatArray[T]) ScanIndexType() any { + return new(T) +} diff --git a/vendor/github.com/jackc/pgx/v5/pgtype/array_codec.go b/vendor/github.com/jackc/pgx/v5/pgtype/array_codec.go new file mode 100644 index 0000000000..ac01496ad1 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgtype/array_codec.go @@ -0,0 +1,428 @@ +package pgtype + +import ( + "database/sql/driver" + "encoding/binary" + "fmt" + "reflect" + + "github.com/jackc/pgx/v5/internal/pgio" +) + +// ArrayGetter is a type that can be converted into a PostgreSQL array. +type ArrayGetter interface { + // Dimensions returns the array dimensions. If array is nil then nil is returned. + Dimensions() []ArrayDimension + + // Index returns the element at i. + Index(i int) any + + // IndexType returns a non-nil scan target of the type Index will return. This is used by ArrayCodec.PlanEncode. + IndexType() any +} + +// ArraySetter is a type can be set from a PostgreSQL array. +type ArraySetter interface { + // SetDimensions prepares the value such that ScanIndex can be called for each element. This will remove any existing + // elements. dimensions may be nil to indicate a NULL array. If unable to exactly preserve dimensions SetDimensions + // may return an error or silently flatten the array dimensions. + SetDimensions(dimensions []ArrayDimension) error + + // ScanIndex returns a value usable as a scan target for i. SetDimensions must be called before ScanIndex. + ScanIndex(i int) any + + // ScanIndexType returns a non-nil scan target of the type ScanIndex will return. This is used by + // ArrayCodec.PlanScan. + ScanIndexType() any +} + +// ArrayCodec is a codec for any array type. +type ArrayCodec struct { + ElementType *Type +} + +func (c *ArrayCodec) FormatSupported(format int16) bool { + return c.ElementType.Codec.FormatSupported(format) +} + +func (c *ArrayCodec) PreferredFormat() int16 { + // The binary format should always be preferred for arrays if it is supported. Usually, this will happen automatically + // because most types that support binary prefer it. However, text, json, and jsonb support binary but prefer the text + // format. This is because it is simpler for jsonb and PostgreSQL can be significantly faster using the text format + // for text-like data types than binary. However, arrays appear to always be faster in binary. + // + // https://www.postgresql.org/message-id/CAMovtNoHFod2jMAKQjjxv209PCTJx5Kc66anwWvX0mEiaXwgmA%40mail.gmail.com + if c.ElementType.Codec.FormatSupported(BinaryFormatCode) { + return BinaryFormatCode + } + return TextFormatCode +} + +func (c *ArrayCodec) PlanEncode(m *Map, oid uint32, format int16, value any) EncodePlan { + arrayValuer, ok := value.(ArrayGetter) + if !ok { + return nil + } + + elementType := arrayValuer.IndexType() + + elementEncodePlan := m.PlanEncode(c.ElementType.OID, format, elementType) + if elementEncodePlan == nil { + if reflect.TypeOf(elementType) != nil { + return nil + } + } + + switch format { + case BinaryFormatCode: + return &encodePlanArrayCodecBinary{ac: c, m: m, oid: oid} + case TextFormatCode: + return &encodePlanArrayCodecText{ac: c, m: m, oid: oid} + } + + return nil +} + +type encodePlanArrayCodecText struct { + ac *ArrayCodec + m *Map + oid uint32 +} + +func (p *encodePlanArrayCodecText) Encode(value any, buf []byte) (newBuf []byte, err error) { + array := value.(ArrayGetter) + + dimensions := array.Dimensions() + if dimensions == nil { + return nil, nil + } + + elementCount := cardinality(dimensions) + if elementCount == 0 { + return append(buf, '{', '}'), nil + } + + buf = encodeTextArrayDimensions(buf, dimensions) + + // dimElemCounts is the multiples of elements that each array lies on. For + // example, a single dimension array of length 4 would have a dimElemCounts of + // [4]. A multi-dimensional array of lengths [3,5,2] would have a + // dimElemCounts of [30,10,2]. This is used to simplify when to render a '{' + // or '}'. + dimElemCounts := make([]int, len(dimensions)) + dimElemCounts[len(dimensions)-1] = int(dimensions[len(dimensions)-1].Length) + for i := len(dimensions) - 2; i > -1; i-- { + dimElemCounts[i] = int(dimensions[i].Length) * dimElemCounts[i+1] + } + + var encodePlan EncodePlan + var lastElemType reflect.Type + inElemBuf := make([]byte, 0, 32) + for i := range elementCount { + if i > 0 { + buf = append(buf, ',') + } + + for _, dec := range dimElemCounts { + if i%dec == 0 { + buf = append(buf, '{') + } + } + + elem := array.Index(i) + var elemBuf []byte + isNil, callNilDriverValuer := isNilDriverValuer(elem) + if !isNil { + elemType := reflect.TypeOf(elem) + if lastElemType != elemType { + lastElemType = elemType + encodePlan = p.m.PlanEncode(p.ac.ElementType.OID, TextFormatCode, elem) + if encodePlan == nil { + return nil, fmt.Errorf("unable to encode %v", array.Index(i)) + } + } + elemBuf, err = encodePlan.Encode(elem, inElemBuf) + if err != nil { + return nil, err + } + } else if callNilDriverValuer { + elemBuf, err = (&encodePlanDriverValuer{m: p.m, oid: p.ac.ElementType.OID, formatCode: TextFormatCode}).Encode(elem, inElemBuf) + if err != nil { + return nil, err + } + } + + if elemBuf == nil { + buf = append(buf, `NULL`...) + } else { + buf = append(buf, quoteArrayElementIfNeeded(string(elemBuf))...) + } + + for _, dec := range dimElemCounts { + if (i+1)%dec == 0 { + buf = append(buf, '}') + } + } + } + + return buf, nil +} + +type encodePlanArrayCodecBinary struct { + ac *ArrayCodec + m *Map + oid uint32 +} + +func (p *encodePlanArrayCodecBinary) Encode(value any, buf []byte) (newBuf []byte, err error) { + array := value.(ArrayGetter) + + dimensions := array.Dimensions() + if dimensions == nil { + return nil, nil + } + + arrayHeader := arrayHeader{ + Dimensions: dimensions, + ElementOID: p.ac.ElementType.OID, + } + + containsNullIndex := len(buf) + 4 + + buf = arrayHeader.EncodeBinary(buf) + + elementCount := cardinality(dimensions) + + var encodePlan EncodePlan + var lastElemType reflect.Type + for i := range elementCount { + sp := len(buf) + buf = pgio.AppendInt32(buf, -1) + + elem := array.Index(i) + var elemBuf []byte + isNil, callNilDriverValuer := isNilDriverValuer(elem) + if !isNil { + elemType := reflect.TypeOf(elem) + if lastElemType != elemType { + lastElemType = elemType + encodePlan = p.m.PlanEncode(p.ac.ElementType.OID, BinaryFormatCode, elem) + if encodePlan == nil { + return nil, fmt.Errorf("unable to encode %v", array.Index(i)) + } + } + elemBuf, err = encodePlan.Encode(elem, buf) + if err != nil { + return nil, err + } + } else if callNilDriverValuer { + elemBuf, err = (&encodePlanDriverValuer{m: p.m, oid: p.ac.ElementType.OID, formatCode: BinaryFormatCode}).Encode(elem, buf) + if err != nil { + return nil, err + } + } + + if elemBuf == nil { + pgio.SetInt32(buf[containsNullIndex:], 1) + } else { + buf = elemBuf + pgio.SetInt32(buf[sp:], int32(len(buf[sp:])-4)) + } + } + + return buf, nil +} + +func (c *ArrayCodec) PlanScan(m *Map, oid uint32, format int16, target any) ScanPlan { + arrayScanner, ok := target.(ArraySetter) + if !ok { + return nil + } + + // target / arrayScanner might be a pointer to a nil. If it is create one so we can call ScanIndexType to plan the + // scan of the elements. + if isNil, _ := isNilDriverValuer(target); isNil { + arrayScanner = reflect.New(reflect.TypeOf(target).Elem()).Interface().(ArraySetter) + } + + elementType := arrayScanner.ScanIndexType() + + elementScanPlan := m.PlanScan(c.ElementType.OID, format, elementType) + if _, ok := elementScanPlan.(*scanPlanFail); ok { + return nil + } + + return &scanPlanArrayCodec{ + arrayCodec: c, + m: m, + oid: oid, + formatCode: format, + } +} + +func (c *ArrayCodec) decodeBinary(m *Map, arrayOID uint32, src []byte, array ArraySetter) error { + var arrayHeader arrayHeader + rp, err := arrayHeader.DecodeBinary(m, src) + if err != nil { + return err + } + + elementCount := cardinality(arrayHeader.Dimensions) + // Each element carries at minimum a 4-byte length header, so elementCount cannot exceed the + // remaining bytes / 4. This bounds the allocation in SetDimensions and the loop below against a + // malicious server claiming huge dimensions in a small message. + if maxElements := len(src[rp:]) / 4; elementCount > maxElements { + return fmt.Errorf("array claims %d elements but only %d bytes remain", elementCount, len(src[rp:])) + } + + err = array.SetDimensions(arrayHeader.Dimensions) + if err != nil { + return err + } + + if elementCount == 0 { + return nil + } + + elementScanPlan := c.ElementType.Codec.PlanScan(m, c.ElementType.OID, BinaryFormatCode, array.ScanIndex(0)) + if elementScanPlan == nil { + elementScanPlan = m.PlanScan(c.ElementType.OID, BinaryFormatCode, array.ScanIndex(0)) + } + + for i := range elementCount { + if len(src[rp:]) < 4 { + return fmt.Errorf("array body truncated at element %d", i) + } + elem := array.ScanIndex(i) + elemLen := int(int32(binary.BigEndian.Uint32(src[rp:]))) + rp += 4 + var elemSrc []byte + if elemLen >= 0 { + if len(src[rp:]) < elemLen { + return fmt.Errorf("array element %d length %d exceeds remaining %d bytes", i, elemLen, len(src[rp:])) + } + elemSrc = src[rp : rp+elemLen] + rp += elemLen + } + err = elementScanPlan.Scan(elemSrc, elem) + if err != nil { + return fmt.Errorf("failed to scan array element %d: %w", i, err) + } + } + + return nil +} + +func (c *ArrayCodec) decodeText(m *Map, arrayOID uint32, src []byte, array ArraySetter) error { + uta, err := parseUntypedTextArray(string(src)) + if err != nil { + return err + } + + err = array.SetDimensions(uta.Dimensions) + if err != nil { + return err + } + + if len(uta.Elements) == 0 { + return nil + } + + elementScanPlan := c.ElementType.Codec.PlanScan(m, c.ElementType.OID, TextFormatCode, array.ScanIndex(0)) + if elementScanPlan == nil { + elementScanPlan = m.PlanScan(c.ElementType.OID, TextFormatCode, array.ScanIndex(0)) + } + + for i, s := range uta.Elements { + elem := array.ScanIndex(i) + var elemSrc []byte + if s != "NULL" || uta.Quoted[i] { + elemSrc = []byte(s) + } + + err = elementScanPlan.Scan(elemSrc, elem) + if err != nil { + return err + } + } + + return nil +} + +type scanPlanArrayCodec struct { + arrayCodec *ArrayCodec + m *Map + oid uint32 + formatCode int16 + elementScanPlan ScanPlan +} + +func (spac *scanPlanArrayCodec) Scan(src []byte, dst any) error { + c := spac.arrayCodec + m := spac.m + oid := spac.oid + formatCode := spac.formatCode + + array := dst.(ArraySetter) + + if src == nil { + return array.SetDimensions(nil) + } + + switch formatCode { + case BinaryFormatCode: + return c.decodeBinary(m, oid, src, array) + case TextFormatCode: + return c.decodeText(m, oid, src, array) + default: + return fmt.Errorf("unknown format code %d", formatCode) + } +} + +func (c *ArrayCodec) DecodeDatabaseSQLValue(m *Map, oid uint32, format int16, src []byte) (driver.Value, error) { + if src == nil { + return nil, nil + } + + switch format { + case TextFormatCode: + return string(src), nil + case BinaryFormatCode: + buf := make([]byte, len(src)) + copy(buf, src) + return buf, nil + default: + return nil, fmt.Errorf("unknown format code %d", format) + } +} + +func (c *ArrayCodec) DecodeValue(m *Map, oid uint32, format int16, src []byte) (any, error) { + if src == nil { + return nil, nil + } + + var slice []any + err := m.PlanScan(oid, format, &slice).Scan(src, &slice) + return slice, err +} + +func isRagged(slice reflect.Value) bool { + if slice.Type().Elem().Kind() != reflect.Slice { + return false + } + + sliceLen := slice.Len() + innerLen := 0 + for i := range sliceLen { + if i == 0 { + innerLen = slice.Index(i).Len() + } else if slice.Index(i).Len() != innerLen { + return true + } + if isRagged(slice.Index(i)) { + return true + } + } + + return false +} diff --git a/vendor/github.com/jackc/pgx/v5/pgtype/bits.go b/vendor/github.com/jackc/pgx/v5/pgtype/bits.go new file mode 100644 index 0000000000..986fe2311a --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgtype/bits.go @@ -0,0 +1,208 @@ +package pgtype + +import ( + "database/sql/driver" + "encoding/binary" + "fmt" + + "github.com/jackc/pgx/v5/internal/pgio" +) + +type BitsScanner interface { + ScanBits(v Bits) error +} + +type BitsValuer interface { + BitsValue() (Bits, error) +} + +// Bits represents the PostgreSQL bit and varbit types. +type Bits struct { + Bytes []byte + Len int32 // Number of bits + Valid bool +} + +// ScanBits implements the [BitsScanner] interface. +func (b *Bits) ScanBits(v Bits) error { + *b = v + return nil +} + +// BitsValue implements the [BitsValuer] interface. +func (b Bits) BitsValue() (Bits, error) { + return b, nil +} + +// Scan implements the [database/sql.Scanner] interface. +func (dst *Bits) Scan(src any) error { + if src == nil { + *dst = Bits{} + return nil + } + + if src, ok := src.(string); ok { + return scanPlanTextAnyToBitsScanner{}.Scan([]byte(src), dst) + } + + return fmt.Errorf("cannot scan %T", src) +} + +// Value implements the [database/sql/driver.Valuer] interface. +func (src Bits) Value() (driver.Value, error) { + if !src.Valid { + return nil, nil + } + + buf, err := BitsCodec{}.PlanEncode(nil, 0, TextFormatCode, src).Encode(src, nil) + if err != nil { + return nil, err + } + return string(buf), err +} + +type BitsCodec struct{} + +func (BitsCodec) FormatSupported(format int16) bool { + return format == TextFormatCode || format == BinaryFormatCode +} + +func (BitsCodec) PreferredFormat() int16 { + return BinaryFormatCode +} + +func (BitsCodec) PlanEncode(m *Map, oid uint32, format int16, value any) EncodePlan { + if _, ok := value.(BitsValuer); !ok { + return nil + } + + switch format { + case BinaryFormatCode: + return encodePlanBitsCodecBinary{} + case TextFormatCode: + return encodePlanBitsCodecText{} + } + + return nil +} + +type encodePlanBitsCodecBinary struct{} + +func (encodePlanBitsCodecBinary) Encode(value any, buf []byte) (newBuf []byte, err error) { + bits, err := value.(BitsValuer).BitsValue() + if err != nil { + return nil, err + } + + if !bits.Valid { + return nil, nil + } + + buf = pgio.AppendInt32(buf, bits.Len) + return append(buf, bits.Bytes...), nil +} + +type encodePlanBitsCodecText struct{} + +func (encodePlanBitsCodecText) Encode(value any, buf []byte) (newBuf []byte, err error) { + bits, err := value.(BitsValuer).BitsValue() + if err != nil { + return nil, err + } + + if !bits.Valid { + return nil, nil + } + + for i := int32(0); i < bits.Len; i++ { + byteIdx := i / 8 + bitMask := byte(128 >> byte(i%8)) + char := byte('0') + if bits.Bytes[byteIdx]&bitMask > 0 { + char = '1' + } + buf = append(buf, char) + } + + return buf, nil +} + +func (BitsCodec) PlanScan(m *Map, oid uint32, format int16, target any) ScanPlan { + switch format { + case BinaryFormatCode: + if _, ok := target.(BitsScanner); ok { + return scanPlanBinaryBitsToBitsScanner{} + } + case TextFormatCode: + if _, ok := target.(BitsScanner); ok { + return scanPlanTextAnyToBitsScanner{} + } + } + + return nil +} + +func (c BitsCodec) DecodeDatabaseSQLValue(m *Map, oid uint32, format int16, src []byte) (driver.Value, error) { + return codecDecodeToTextFormat(c, m, oid, format, src) +} + +func (c BitsCodec) DecodeValue(m *Map, oid uint32, format int16, src []byte) (any, error) { + if src == nil { + return nil, nil + } + + var box Bits + err := codecScan(c, m, oid, format, src, &box) + if err != nil { + return nil, err + } + return box, nil +} + +type scanPlanBinaryBitsToBitsScanner struct{} + +func (scanPlanBinaryBitsToBitsScanner) Scan(src []byte, dst any) error { + scanner := (dst).(BitsScanner) + + if src == nil { + return scanner.ScanBits(Bits{}) + } + + if len(src) < 4 { + return fmt.Errorf("invalid length for bit/varbit: %v", len(src)) + } + + bitLen := int32(binary.BigEndian.Uint32(src)) + rp := 4 + buf := make([]byte, len(src[rp:])) + copy(buf, src[rp:]) + + return scanner.ScanBits(Bits{Bytes: buf, Len: bitLen, Valid: true}) +} + +type scanPlanTextAnyToBitsScanner struct{} + +func (scanPlanTextAnyToBitsScanner) Scan(src []byte, dst any) error { + scanner := (dst).(BitsScanner) + + if src == nil { + return scanner.ScanBits(Bits{}) + } + + bitLen := len(src) + byteLen := bitLen / 8 + if bitLen%8 > 0 { + byteLen++ + } + buf := make([]byte, byteLen) + + for i, b := range src { + if b == '1' { + byteIdx := i / 8 + bitIdx := uint(i % 8) + buf[byteIdx] |= 128 >> bitIdx + } + } + + return scanner.ScanBits(Bits{Bytes: buf, Len: int32(bitLen), Valid: true}) +} diff --git a/vendor/github.com/jackc/pgx/v5/pgtype/bool.go b/vendor/github.com/jackc/pgx/v5/pgtype/bool.go new file mode 100644 index 0000000000..077668e79d --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgtype/bool.go @@ -0,0 +1,346 @@ +package pgtype + +import ( + "bytes" + "database/sql/driver" + "encoding/json" + "fmt" + "strconv" + "strings" +) + +type BoolScanner interface { + ScanBool(v Bool) error +} + +type BoolValuer interface { + BoolValue() (Bool, error) +} + +type Bool struct { + Bool bool + Valid bool +} + +// ScanBool implements the [BoolScanner] interface. +func (b *Bool) ScanBool(v Bool) error { + *b = v + return nil +} + +// BoolValue implements the [BoolValuer] interface. +func (b Bool) BoolValue() (Bool, error) { + return b, nil +} + +// Scan implements the [database/sql.Scanner] interface. +func (dst *Bool) Scan(src any) error { + if src == nil { + *dst = Bool{} + return nil + } + + switch src := src.(type) { + case bool: + *dst = Bool{Bool: src, Valid: true} + return nil + case string: + b, err := strconv.ParseBool(src) + if err != nil { + return err + } + *dst = Bool{Bool: b, Valid: true} + return nil + case []byte: + b, err := strconv.ParseBool(string(src)) + if err != nil { + return err + } + *dst = Bool{Bool: b, Valid: true} + return nil + } + + return fmt.Errorf("cannot scan %T", src) +} + +// Value implements the [database/sql/driver.Valuer] interface. +func (src Bool) Value() (driver.Value, error) { + if !src.Valid { + return nil, nil + } + + return src.Bool, nil +} + +// MarshalJSON implements the [encoding/json.Marshaler] interface. +func (src Bool) MarshalJSON() ([]byte, error) { + if !src.Valid { + return []byte("null"), nil + } + + if src.Bool { + return []byte("true"), nil + } else { + return []byte("false"), nil + } +} + +// UnmarshalJSON implements the [encoding/json.Unmarshaler] interface. +func (dst *Bool) UnmarshalJSON(b []byte) error { + var v *bool + err := json.Unmarshal(b, &v) + if err != nil { + return err + } + + if v == nil { + *dst = Bool{} + } else { + *dst = Bool{Bool: *v, Valid: true} + } + + return nil +} + +type BoolCodec struct{} + +func (BoolCodec) FormatSupported(format int16) bool { + return format == TextFormatCode || format == BinaryFormatCode +} + +func (BoolCodec) PreferredFormat() int16 { + return BinaryFormatCode +} + +func (BoolCodec) PlanEncode(m *Map, oid uint32, format int16, value any) EncodePlan { + switch format { + case BinaryFormatCode: + switch value.(type) { + case bool: + return encodePlanBoolCodecBinaryBool{} + case BoolValuer: + return encodePlanBoolCodecBinaryBoolValuer{} + } + case TextFormatCode: + switch value.(type) { + case bool: + return encodePlanBoolCodecTextBool{} + case BoolValuer: + return encodePlanBoolCodecTextBoolValuer{} + } + } + + return nil +} + +type encodePlanBoolCodecBinaryBool struct{} + +func (encodePlanBoolCodecBinaryBool) Encode(value any, buf []byte) (newBuf []byte, err error) { + v := value.(bool) + + if v { + buf = append(buf, 1) + } else { + buf = append(buf, 0) + } + + return buf, nil +} + +type encodePlanBoolCodecTextBoolValuer struct{} + +func (encodePlanBoolCodecTextBoolValuer) Encode(value any, buf []byte) (newBuf []byte, err error) { + b, err := value.(BoolValuer).BoolValue() + if err != nil { + return nil, err + } + + if !b.Valid { + return nil, nil + } + + if b.Bool { + buf = append(buf, 't') + } else { + buf = append(buf, 'f') + } + + return buf, nil +} + +type encodePlanBoolCodecBinaryBoolValuer struct{} + +func (encodePlanBoolCodecBinaryBoolValuer) Encode(value any, buf []byte) (newBuf []byte, err error) { + b, err := value.(BoolValuer).BoolValue() + if err != nil { + return nil, err + } + + if !b.Valid { + return nil, nil + } + + if b.Bool { + buf = append(buf, 1) + } else { + buf = append(buf, 0) + } + + return buf, nil +} + +type encodePlanBoolCodecTextBool struct{} + +func (encodePlanBoolCodecTextBool) Encode(value any, buf []byte) (newBuf []byte, err error) { + v := value.(bool) + + if v { + buf = append(buf, 't') + } else { + buf = append(buf, 'f') + } + + return buf, nil +} + +func (BoolCodec) PlanScan(m *Map, oid uint32, format int16, target any) ScanPlan { + switch format { + case BinaryFormatCode: + switch target.(type) { + case *bool: + return scanPlanBinaryBoolToBool{} + case BoolScanner: + return scanPlanBinaryBoolToBoolScanner{} + } + case TextFormatCode: + switch target.(type) { + case *bool: + return scanPlanTextAnyToBool{} + case BoolScanner: + return scanPlanTextAnyToBoolScanner{} + } + } + + return nil +} + +func (c BoolCodec) DecodeDatabaseSQLValue(m *Map, oid uint32, format int16, src []byte) (driver.Value, error) { + return c.DecodeValue(m, oid, format, src) +} + +func (c BoolCodec) DecodeValue(m *Map, oid uint32, format int16, src []byte) (any, error) { + if src == nil { + return nil, nil + } + + var b bool + err := codecScan(c, m, oid, format, src, &b) + if err != nil { + return nil, err + } + return b, nil +} + +type scanPlanBinaryBoolToBool struct{} + +func (scanPlanBinaryBoolToBool) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + if len(src) != 1 { + return fmt.Errorf("invalid length for bool: %v", len(src)) + } + + p, ok := (dst).(*bool) + if !ok { + return ErrScanTargetTypeChanged + } + + *p = src[0] == 1 + + return nil +} + +type scanPlanTextAnyToBool struct{} + +func (scanPlanTextAnyToBool) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + if len(src) == 0 { + return fmt.Errorf("cannot scan empty string into %T", dst) + } + + p, ok := (dst).(*bool) + if !ok { + return ErrScanTargetTypeChanged + } + + v, err := planTextToBool(src) + if err != nil { + return err + } + + *p = v + + return nil +} + +type scanPlanBinaryBoolToBoolScanner struct{} + +func (scanPlanBinaryBoolToBoolScanner) Scan(src []byte, dst any) error { + s, ok := (dst).(BoolScanner) + if !ok { + return ErrScanTargetTypeChanged + } + + if src == nil { + return s.ScanBool(Bool{}) + } + + if len(src) != 1 { + return fmt.Errorf("invalid length for bool: %v", len(src)) + } + + return s.ScanBool(Bool{Bool: src[0] == 1, Valid: true}) +} + +type scanPlanTextAnyToBoolScanner struct{} + +func (scanPlanTextAnyToBoolScanner) Scan(src []byte, dst any) error { + s, ok := (dst).(BoolScanner) + if !ok { + return ErrScanTargetTypeChanged + } + + if src == nil { + return s.ScanBool(Bool{}) + } + + if len(src) == 0 { + return fmt.Errorf("cannot scan empty string into %T", dst) + } + + v, err := planTextToBool(src) + if err != nil { + return err + } + + return s.ScanBool(Bool{Bool: v, Valid: true}) +} + +// https://www.postgresql.org/docs/current/datatype-boolean.html +func planTextToBool(src []byte) (bool, error) { + s := string(bytes.ToLower(bytes.TrimSpace(src))) + + switch { + case strings.HasPrefix("true", s), strings.HasPrefix("yes", s), s == "on", s == "1": //nolint:gocritic // s is intentionally the prefix argument so partial inputs (t, tr, tru) also match. + return true, nil + case strings.HasPrefix("false", s), strings.HasPrefix("no", s), strings.HasPrefix("off", s), s == "0": //nolint:gocritic // s is intentionally the prefix argument so partial inputs (f, fa, fal) also match. + return false, nil + default: + return false, fmt.Errorf("unknown boolean string representation %q", src) + } +} diff --git a/vendor/github.com/jackc/pgx/v5/pgtype/box.go b/vendor/github.com/jackc/pgx/v5/pgtype/box.go new file mode 100644 index 0000000000..8270aafd57 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgtype/box.go @@ -0,0 +1,235 @@ +package pgtype + +import ( + "database/sql/driver" + "encoding/binary" + "fmt" + "math" + "strconv" + "strings" + + "github.com/jackc/pgx/v5/internal/pgio" +) + +type BoxScanner interface { + ScanBox(v Box) error +} + +type BoxValuer interface { + BoxValue() (Box, error) +} + +type Box struct { + P [2]Vec2 + Valid bool +} + +// ScanBox implements the [BoxScanner] interface. +func (b *Box) ScanBox(v Box) error { + *b = v + return nil +} + +// BoxValue implements the [BoxValuer] interface. +func (b Box) BoxValue() (Box, error) { + return b, nil +} + +// Scan implements the [database/sql.Scanner] interface. +func (dst *Box) Scan(src any) error { + if src == nil { + *dst = Box{} + return nil + } + + if src, ok := src.(string); ok { + return scanPlanTextAnyToBoxScanner{}.Scan([]byte(src), dst) + } + + return fmt.Errorf("cannot scan %T", src) +} + +// Value implements the [database/sql/driver.Valuer] interface. +func (src Box) Value() (driver.Value, error) { + if !src.Valid { + return nil, nil + } + + buf, err := BoxCodec{}.PlanEncode(nil, 0, TextFormatCode, src).Encode(src, nil) + if err != nil { + return nil, err + } + return string(buf), err +} + +type BoxCodec struct{} + +func (BoxCodec) FormatSupported(format int16) bool { + return format == TextFormatCode || format == BinaryFormatCode +} + +func (BoxCodec) PreferredFormat() int16 { + return BinaryFormatCode +} + +func (BoxCodec) PlanEncode(m *Map, oid uint32, format int16, value any) EncodePlan { + if _, ok := value.(BoxValuer); !ok { + return nil + } + + switch format { + case BinaryFormatCode: + return encodePlanBoxCodecBinary{} + case TextFormatCode: + return encodePlanBoxCodecText{} + } + + return nil +} + +type encodePlanBoxCodecBinary struct{} + +func (encodePlanBoxCodecBinary) Encode(value any, buf []byte) (newBuf []byte, err error) { + box, err := value.(BoxValuer).BoxValue() + if err != nil { + return nil, err + } + + if !box.Valid { + return nil, nil + } + + buf = pgio.AppendUint64(buf, math.Float64bits(box.P[0].X)) + buf = pgio.AppendUint64(buf, math.Float64bits(box.P[0].Y)) + buf = pgio.AppendUint64(buf, math.Float64bits(box.P[1].X)) + buf = pgio.AppendUint64(buf, math.Float64bits(box.P[1].Y)) + return buf, nil +} + +type encodePlanBoxCodecText struct{} + +func (encodePlanBoxCodecText) Encode(value any, buf []byte) (newBuf []byte, err error) { + box, err := value.(BoxValuer).BoxValue() + if err != nil { + return nil, err + } + + if !box.Valid { + return nil, nil + } + + buf = append(buf, fmt.Sprintf(`(%s,%s),(%s,%s)`, + strconv.FormatFloat(box.P[0].X, 'f', -1, 64), + strconv.FormatFloat(box.P[0].Y, 'f', -1, 64), + strconv.FormatFloat(box.P[1].X, 'f', -1, 64), + strconv.FormatFloat(box.P[1].Y, 'f', -1, 64), + )...) + return buf, nil +} + +func (BoxCodec) PlanScan(m *Map, oid uint32, format int16, target any) ScanPlan { + switch format { + case BinaryFormatCode: + if _, ok := target.(BoxScanner); ok { + return scanPlanBinaryBoxToBoxScanner{} + } + case TextFormatCode: + if _, ok := target.(BoxScanner); ok { + return scanPlanTextAnyToBoxScanner{} + } + } + + return nil +} + +type scanPlanBinaryBoxToBoxScanner struct{} + +func (scanPlanBinaryBoxToBoxScanner) Scan(src []byte, dst any) error { + scanner := (dst).(BoxScanner) + + if src == nil { + return scanner.ScanBox(Box{}) + } + + if len(src) != 32 { + return fmt.Errorf("invalid length for Box: %v", len(src)) + } + + x1 := binary.BigEndian.Uint64(src) + y1 := binary.BigEndian.Uint64(src[8:]) + x2 := binary.BigEndian.Uint64(src[16:]) + y2 := binary.BigEndian.Uint64(src[24:]) + + return scanner.ScanBox(Box{ + P: [2]Vec2{ + {math.Float64frombits(x1), math.Float64frombits(y1)}, + {math.Float64frombits(x2), math.Float64frombits(y2)}, + }, + Valid: true, + }) +} + +type scanPlanTextAnyToBoxScanner struct{} + +func (scanPlanTextAnyToBoxScanner) Scan(src []byte, dst any) error { + scanner := (dst).(BoxScanner) + + if src == nil { + return scanner.ScanBox(Box{}) + } + + if len(src) < 11 { + return fmt.Errorf("invalid length for Box: %v", len(src)) + } + + // Expected format: (x1,y1),(x2,y2) + sp1, sp2, found := strings.Cut(string(src[1:len(src)-1]), "),(") + if !found { + return fmt.Errorf("invalid format for Box") + } + + sx1, sy1, found := strings.Cut(sp1, ",") + if !found { + return fmt.Errorf("invalid format for Box") + } + sx2, sy2, found := strings.Cut(sp2, ",") + if !found { + return fmt.Errorf("invalid format for Box") + } + + x1, err := strconv.ParseFloat(sx1, 64) + if err != nil { + return err + } + y1, err := strconv.ParseFloat(sy1, 64) + if err != nil { + return err + } + x2, err := strconv.ParseFloat(sx2, 64) + if err != nil { + return err + } + y2, err := strconv.ParseFloat(sy2, 64) + if err != nil { + return err + } + + return scanner.ScanBox(Box{P: [2]Vec2{{x1, y1}, {x2, y2}}, Valid: true}) +} + +func (c BoxCodec) DecodeDatabaseSQLValue(m *Map, oid uint32, format int16, src []byte) (driver.Value, error) { + return codecDecodeToTextFormat(c, m, oid, format, src) +} + +func (c BoxCodec) DecodeValue(m *Map, oid uint32, format int16, src []byte) (any, error) { + if src == nil { + return nil, nil + } + + var box Box + err := codecScan(c, m, oid, format, src, &box) + if err != nil { + return nil, err + } + return box, nil +} diff --git a/vendor/github.com/jackc/pgx/v5/pgtype/builtin_wrappers.go b/vendor/github.com/jackc/pgx/v5/pgtype/builtin_wrappers.go new file mode 100644 index 0000000000..a412763a9b --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgtype/builtin_wrappers.go @@ -0,0 +1,952 @@ +package pgtype + +import ( + "errors" + "fmt" + "math" + "math/big" + "net" + "net/netip" + "reflect" + "time" +) + +type int8Wrapper int8 + +func (w int8Wrapper) SkipUnderlyingTypePlan() {} + +func (w *int8Wrapper) ScanInt64(v Int8) error { + if !v.Valid { + return fmt.Errorf("cannot scan NULL into *int8") + } + + if v.Int64 < math.MinInt8 { + return fmt.Errorf("%d is less than minimum value for int8", v.Int64) + } + if v.Int64 > math.MaxInt8 { + return fmt.Errorf("%d is greater than maximum value for int8", v.Int64) + } + *w = int8Wrapper(v.Int64) + + return nil +} + +func (w int8Wrapper) Int64Value() (Int8, error) { + return Int8{Int64: int64(w), Valid: true}, nil +} + +type int16Wrapper int16 + +func (w int16Wrapper) SkipUnderlyingTypePlan() {} + +func (w *int16Wrapper) ScanInt64(v Int8) error { + if !v.Valid { + return fmt.Errorf("cannot scan NULL into *int16") + } + + if v.Int64 < math.MinInt16 { + return fmt.Errorf("%d is less than minimum value for int16", v.Int64) + } + if v.Int64 > math.MaxInt16 { + return fmt.Errorf("%d is greater than maximum value for int16", v.Int64) + } + *w = int16Wrapper(v.Int64) + + return nil +} + +func (w int16Wrapper) Int64Value() (Int8, error) { + return Int8{Int64: int64(w), Valid: true}, nil +} + +type int32Wrapper int32 + +func (w int32Wrapper) SkipUnderlyingTypePlan() {} + +func (w *int32Wrapper) ScanInt64(v Int8) error { + if !v.Valid { + return fmt.Errorf("cannot scan NULL into *int32") + } + + if v.Int64 < math.MinInt32 { + return fmt.Errorf("%d is less than minimum value for int32", v.Int64) + } + if v.Int64 > math.MaxInt32 { + return fmt.Errorf("%d is greater than maximum value for int32", v.Int64) + } + *w = int32Wrapper(v.Int64) + + return nil +} + +func (w int32Wrapper) Int64Value() (Int8, error) { + return Int8{Int64: int64(w), Valid: true}, nil +} + +type int64Wrapper int64 + +func (w int64Wrapper) SkipUnderlyingTypePlan() {} + +func (w *int64Wrapper) ScanInt64(v Int8) error { + if !v.Valid { + return fmt.Errorf("cannot scan NULL into *int64") + } + + *w = int64Wrapper(v.Int64) + + return nil +} + +func (w int64Wrapper) Int64Value() (Int8, error) { + return Int8{Int64: int64(w), Valid: true}, nil +} + +type intWrapper int + +func (w intWrapper) SkipUnderlyingTypePlan() {} + +func (w *intWrapper) ScanInt64(v Int8) error { + if !v.Valid { + return fmt.Errorf("cannot scan NULL into *int") + } + + if v.Int64 < math.MinInt { + return fmt.Errorf("%d is less than minimum value for int", v.Int64) + } + if v.Int64 > math.MaxInt { + return fmt.Errorf("%d is greater than maximum value for int", v.Int64) + } + + *w = intWrapper(v.Int64) + + return nil +} + +func (w intWrapper) Int64Value() (Int8, error) { + return Int8{Int64: int64(w), Valid: true}, nil +} + +type uint8Wrapper uint8 + +func (w uint8Wrapper) SkipUnderlyingTypePlan() {} + +func (w *uint8Wrapper) ScanInt64(v Int8) error { + if !v.Valid { + return fmt.Errorf("cannot scan NULL into *uint8") + } + + if v.Int64 < 0 { + return fmt.Errorf("%d is less than minimum value for uint8", v.Int64) + } + if v.Int64 > math.MaxUint8 { + return fmt.Errorf("%d is greater than maximum value for uint8", v.Int64) + } + *w = uint8Wrapper(v.Int64) + + return nil +} + +func (w uint8Wrapper) Int64Value() (Int8, error) { + return Int8{Int64: int64(w), Valid: true}, nil +} + +type uint16Wrapper uint16 + +func (w uint16Wrapper) SkipUnderlyingTypePlan() {} + +func (w *uint16Wrapper) ScanInt64(v Int8) error { + if !v.Valid { + return fmt.Errorf("cannot scan NULL into *uint16") + } + + if v.Int64 < 0 { + return fmt.Errorf("%d is less than minimum value for uint16", v.Int64) + } + if v.Int64 > math.MaxUint16 { + return fmt.Errorf("%d is greater than maximum value for uint16", v.Int64) + } + *w = uint16Wrapper(v.Int64) + + return nil +} + +func (w uint16Wrapper) Int64Value() (Int8, error) { + return Int8{Int64: int64(w), Valid: true}, nil +} + +type uint32Wrapper uint32 + +func (w uint32Wrapper) SkipUnderlyingTypePlan() {} + +func (w *uint32Wrapper) ScanInt64(v Int8) error { + if !v.Valid { + return fmt.Errorf("cannot scan NULL into *uint32") + } + + if v.Int64 < 0 { + return fmt.Errorf("%d is less than minimum value for uint32", v.Int64) + } + if v.Int64 > math.MaxUint32 { + return fmt.Errorf("%d is greater than maximum value for uint32", v.Int64) + } + *w = uint32Wrapper(v.Int64) + + return nil +} + +func (w uint32Wrapper) Int64Value() (Int8, error) { + return Int8{Int64: int64(w), Valid: true}, nil +} + +type uint64Wrapper uint64 + +func (w uint64Wrapper) SkipUnderlyingTypePlan() {} + +func (w *uint64Wrapper) ScanInt64(v Int8) error { + if !v.Valid { + return fmt.Errorf("cannot scan NULL into *uint64") + } + + if v.Int64 < 0 { + return fmt.Errorf("%d is less than minimum value for uint64", v.Int64) + } + + *w = uint64Wrapper(v.Int64) + + return nil +} + +func (w uint64Wrapper) Int64Value() (Int8, error) { + if uint64(w) > uint64(math.MaxInt64) { + return Int8{}, fmt.Errorf("%d is greater than maximum value for int64", w) + } + + return Int8{Int64: int64(w), Valid: true}, nil +} + +func (w *uint64Wrapper) ScanNumeric(v Numeric) error { + if !v.Valid { + return fmt.Errorf("cannot scan NULL into *uint64") + } + + bi, err := v.toBigInt() + if err != nil { + return fmt.Errorf("cannot scan into *uint64: %w", err) + } + + if !bi.IsUint64() { + return fmt.Errorf("cannot scan %v into *uint64", bi.String()) + } + + *w = uint64Wrapper(bi.Uint64()) + + return nil +} + +func (w uint64Wrapper) NumericValue() (Numeric, error) { + return Numeric{Int: new(big.Int).SetUint64(uint64(w)), Valid: true}, nil +} + +type uintWrapper uint + +func (w uintWrapper) SkipUnderlyingTypePlan() {} + +func (w *uintWrapper) ScanInt64(v Int8) error { + if !v.Valid { + return fmt.Errorf("cannot scan NULL into *uint64") + } + + if v.Int64 < 0 { + return fmt.Errorf("%d is less than minimum value for uint64", v.Int64) + } + + if uint64(v.Int64) > math.MaxUint { + return fmt.Errorf("%d is greater than maximum value for uint", v.Int64) + } + + *w = uintWrapper(v.Int64) + + return nil +} + +func (w uintWrapper) Int64Value() (Int8, error) { + if uint64(w) > uint64(math.MaxInt64) { + return Int8{}, fmt.Errorf("%d is greater than maximum value for int64", w) + } + + return Int8{Int64: int64(w), Valid: true}, nil +} + +func (w *uintWrapper) ScanNumeric(v Numeric) error { + if !v.Valid { + return fmt.Errorf("cannot scan NULL into *uint") + } + + bi, err := v.toBigInt() + if err != nil { + return fmt.Errorf("cannot scan into *uint: %w", err) + } + + if !bi.IsUint64() { + return fmt.Errorf("cannot scan %v into *uint", bi.String()) + } + + ui := bi.Uint64() + + if math.MaxUint < ui { + return fmt.Errorf("cannot scan %v into *uint", ui) + } + + *w = uintWrapper(ui) + + return nil +} + +func (w uintWrapper) NumericValue() (Numeric, error) { + return Numeric{Int: new(big.Int).SetUint64(uint64(w)), Valid: true}, nil +} + +type float32Wrapper float32 + +func (w float32Wrapper) SkipUnderlyingTypePlan() {} + +func (w *float32Wrapper) ScanInt64(v Int8) error { + if !v.Valid { + return fmt.Errorf("cannot scan NULL into *float32") + } + + *w = float32Wrapper(v.Int64) + + return nil +} + +func (w float32Wrapper) Int64Value() (Int8, error) { + if w > math.MaxInt64 { + return Int8{}, fmt.Errorf("%f is greater than maximum value for int64", w) + } + + return Int8{Int64: int64(w), Valid: true}, nil +} + +func (w *float32Wrapper) ScanFloat64(v Float8) error { + if !v.Valid { + return fmt.Errorf("cannot scan NULL into *float32") + } + + *w = float32Wrapper(v.Float64) + + return nil +} + +func (w float32Wrapper) Float64Value() (Float8, error) { + return Float8{Float64: float64(w), Valid: true}, nil +} + +type float64Wrapper float64 + +func (w float64Wrapper) SkipUnderlyingTypePlan() {} + +func (w *float64Wrapper) ScanInt64(v Int8) error { + if !v.Valid { + return fmt.Errorf("cannot scan NULL into *float64") + } + + *w = float64Wrapper(v.Int64) + + return nil +} + +func (w float64Wrapper) Int64Value() (Int8, error) { + if w > math.MaxInt64 { + return Int8{}, fmt.Errorf("%f is greater than maximum value for int64", w) + } + + return Int8{Int64: int64(w), Valid: true}, nil +} + +func (w *float64Wrapper) ScanFloat64(v Float8) error { + if !v.Valid { + return fmt.Errorf("cannot scan NULL into *float64") + } + + *w = float64Wrapper(v.Float64) + + return nil +} + +func (w float64Wrapper) Float64Value() (Float8, error) { + return Float8{Float64: float64(w), Valid: true}, nil +} + +type stringWrapper string + +func (w stringWrapper) SkipUnderlyingTypePlan() {} + +func (w *stringWrapper) ScanText(v Text) error { + if !v.Valid { + return fmt.Errorf("cannot scan NULL into *string") + } + + *w = stringWrapper(v.String) + return nil +} + +func (w stringWrapper) TextValue() (Text, error) { + return Text{String: string(w), Valid: true}, nil +} + +type timeWrapper time.Time + +func (w *timeWrapper) ScanDate(v Date) error { + if !v.Valid { + return fmt.Errorf("cannot scan NULL into *time.Time") + } + + switch v.InfinityModifier { + case Finite: + *w = timeWrapper(v.Time) + return nil + case Infinity: + return fmt.Errorf("cannot scan Infinity into *time.Time") + case NegativeInfinity: + return fmt.Errorf("cannot scan -Infinity into *time.Time") + default: + return fmt.Errorf("invalid InfinityModifier: %v", v.InfinityModifier) + } +} + +func (w timeWrapper) DateValue() (Date, error) { + return Date{Time: time.Time(w), Valid: true}, nil +} + +func (w *timeWrapper) ScanTimestamp(v Timestamp) error { + if !v.Valid { + return fmt.Errorf("cannot scan NULL into *time.Time") + } + + switch v.InfinityModifier { + case Finite: + *w = timeWrapper(v.Time) + return nil + case Infinity: + return fmt.Errorf("cannot scan Infinity into *time.Time") + case NegativeInfinity: + return fmt.Errorf("cannot scan -Infinity into *time.Time") + default: + return fmt.Errorf("invalid InfinityModifier: %v", v.InfinityModifier) + } +} + +func (w timeWrapper) TimestampValue() (Timestamp, error) { + return Timestamp{Time: time.Time(w), Valid: true}, nil +} + +func (w *timeWrapper) ScanTimestamptz(v Timestamptz) error { + if !v.Valid { + return fmt.Errorf("cannot scan NULL into *time.Time") + } + + switch v.InfinityModifier { + case Finite: + *w = timeWrapper(v.Time) + return nil + case Infinity: + return fmt.Errorf("cannot scan Infinity into *time.Time") + case NegativeInfinity: + return fmt.Errorf("cannot scan -Infinity into *time.Time") + default: + return fmt.Errorf("invalid InfinityModifier: %v", v.InfinityModifier) + } +} + +func (w timeWrapper) TimestamptzValue() (Timestamptz, error) { + return Timestamptz{Time: time.Time(w), Valid: true}, nil +} + +func (w *timeWrapper) ScanTime(v Time) error { + if !v.Valid { + return fmt.Errorf("cannot scan NULL into *time.Time") + } + + // 24:00:00 is max allowed time in PostgreSQL, but time.Time will normalize that to 00:00:00 the next day. + var maxRepresentableByTime int64 = 24*60*60*1000000 - 1 + if v.Microseconds > maxRepresentableByTime { + return fmt.Errorf("%d microseconds cannot be represented as time.Time", v.Microseconds) + } + + usec := v.Microseconds + hours := usec / microsecondsPerHour + usec -= hours * microsecondsPerHour + minutes := usec / microsecondsPerMinute + usec -= minutes * microsecondsPerMinute + seconds := usec / microsecondsPerSecond + usec -= seconds * microsecondsPerSecond + ns := usec * 1000 + *w = timeWrapper(time.Date(2000, 1, 1, int(hours), int(minutes), int(seconds), int(ns), time.UTC)) + return nil +} + +func (w timeWrapper) TimeValue() (Time, error) { + t := time.Time(w) + usec := int64(t.Hour())*microsecondsPerHour + + int64(t.Minute())*microsecondsPerMinute + + int64(t.Second())*microsecondsPerSecond + + int64(t.Nanosecond())/1000 + return Time{Microseconds: usec, Valid: true}, nil +} + +type durationWrapper time.Duration + +func (w durationWrapper) SkipUnderlyingTypePlan() {} + +func (w *durationWrapper) ScanInterval(v Interval) error { + if !v.Valid { + return fmt.Errorf("cannot scan NULL into *time.Interval") + } + + us := int64(v.Months)*microsecondsPerMonth + int64(v.Days)*microsecondsPerDay + v.Microseconds + *w = durationWrapper(time.Duration(us) * time.Microsecond) + return nil +} + +func (w durationWrapper) IntervalValue() (Interval, error) { + return Interval{Microseconds: int64(w) / 1000, Valid: true}, nil +} + +type netIPNetWrapper net.IPNet + +func (w *netIPNetWrapper) ScanNetipPrefix(v netip.Prefix) error { + if !v.IsValid() { + return fmt.Errorf("cannot scan NULL into *net.IPNet") + } + + *w = netIPNetWrapper{ + IP: v.Addr().AsSlice(), + Mask: net.CIDRMask(v.Bits(), v.Addr().BitLen()), + } + + return nil +} + +func (w netIPNetWrapper) NetipPrefixValue() (netip.Prefix, error) { + ip, ok := netip.AddrFromSlice(w.IP) + if !ok { + return netip.Prefix{}, errors.New("invalid net.IPNet") + } + + ones, _ := w.Mask.Size() + + return netip.PrefixFrom(ip, ones), nil +} + +type netIPWrapper net.IP + +func (w netIPWrapper) SkipUnderlyingTypePlan() {} + +func (w *netIPWrapper) ScanNetipPrefix(v netip.Prefix) error { + if !v.IsValid() { + *w = nil + return nil + } + + if v.Addr().BitLen() != v.Bits() { + return fmt.Errorf("cannot scan %v to *net.IP", v) + } + + *w = netIPWrapper(v.Addr().AsSlice()) + return nil +} + +func (w netIPWrapper) NetipPrefixValue() (netip.Prefix, error) { + if w == nil { + return netip.Prefix{}, nil + } + + addr, ok := netip.AddrFromSlice([]byte(w)) + if !ok { + return netip.Prefix{}, errors.New("invalid net.IP") + } + + return netip.PrefixFrom(addr, addr.BitLen()), nil +} + +type netipPrefixWrapper netip.Prefix + +func (w *netipPrefixWrapper) ScanNetipPrefix(v netip.Prefix) error { + *w = netipPrefixWrapper(v) + return nil +} + +func (w netipPrefixWrapper) NetipPrefixValue() (netip.Prefix, error) { + return netip.Prefix(w), nil +} + +type netipAddrWrapper netip.Addr + +func (w *netipAddrWrapper) ScanNetipPrefix(v netip.Prefix) error { + if !v.IsValid() { + *w = netipAddrWrapper(netip.Addr{}) + return nil + } + + if v.Addr().BitLen() != v.Bits() { + return fmt.Errorf("cannot scan %v to netip.Addr", v) + } + + *w = netipAddrWrapper(v.Addr()) + + return nil +} + +func (w netipAddrWrapper) NetipPrefixValue() (netip.Prefix, error) { + addr := (netip.Addr)(w) + if !addr.IsValid() { + return netip.Prefix{}, nil + } + + return netip.PrefixFrom(addr, addr.BitLen()), nil +} + +type mapStringToPointerStringWrapper map[string]*string + +func (w *mapStringToPointerStringWrapper) ScanHstore(v Hstore) error { + *w = mapStringToPointerStringWrapper(v) + return nil +} + +func (w mapStringToPointerStringWrapper) HstoreValue() (Hstore, error) { + return Hstore(w), nil +} + +type mapStringToStringWrapper map[string]string + +func (w *mapStringToStringWrapper) ScanHstore(v Hstore) error { + *w = make(mapStringToStringWrapper, len(v)) + for k, v := range v { + if v == nil { + return fmt.Errorf("cannot scan NULL to string") + } + (*w)[k] = *v + } + return nil +} + +func (w mapStringToStringWrapper) HstoreValue() (Hstore, error) { + if w == nil { + return nil, nil + } + + hstore := make(Hstore, len(w)) + for k, v := range w { + s := v + hstore[k] = &s + } + return hstore, nil +} + +type fmtStringerWrapper struct { + s fmt.Stringer +} + +func (w fmtStringerWrapper) TextValue() (Text, error) { + return Text{String: w.s.String(), Valid: true}, nil +} + +type byte16Wrapper [16]byte + +func (w *byte16Wrapper) ScanUUID(v UUID) error { + if !v.Valid { + return fmt.Errorf("cannot scan NULL into *[16]byte") + } + *w = byte16Wrapper(v.Bytes) + return nil +} + +func (w byte16Wrapper) UUIDValue() (UUID, error) { + return UUID{Bytes: [16]byte(w), Valid: true}, nil +} + +type byteSliceWrapper []byte + +func (w byteSliceWrapper) SkipUnderlyingTypePlan() {} + +func (w *byteSliceWrapper) ScanText(v Text) error { + if !v.Valid { + *w = nil + return nil + } + + *w = byteSliceWrapper(v.String) + return nil +} + +func (w byteSliceWrapper) TextValue() (Text, error) { + if w == nil { + return Text{}, nil + } + + return Text{String: string(w), Valid: true}, nil +} + +func (w *byteSliceWrapper) ScanUUID(v UUID) error { + if !v.Valid { + *w = nil + return nil + } + *w = make(byteSliceWrapper, 16) + copy(*w, v.Bytes[:]) + return nil +} + +func (w byteSliceWrapper) UUIDValue() (UUID, error) { + if w == nil { + return UUID{}, nil + } + + uuid := UUID{Valid: true} + copy(uuid.Bytes[:], w) + return uuid, nil +} + +// structWrapper implements CompositeIndexGetter for a struct. +type structWrapper struct { + s any + exportedFields []reflect.Value +} + +func (w structWrapper) IsNull() bool { + return w.s == nil +} + +func (w structWrapper) Index(i int) any { + if i >= len(w.exportedFields) { + return fmt.Errorf("%#v only has %d public fields - %d is out of bounds", w.s, len(w.exportedFields), i) + } + + return w.exportedFields[i].Interface() +} + +// ptrStructWrapper implements CompositeIndexScanner for a pointer to a struct. +type ptrStructWrapper struct { + s any + exportedFields []reflect.Value +} + +func (w *ptrStructWrapper) ScanNull() error { + return fmt.Errorf("cannot scan NULL into %#v", w.s) +} + +func (w *ptrStructWrapper) ScanIndex(i int) any { + if i >= len(w.exportedFields) { + return fmt.Errorf("%#v only has %d public fields - %d is out of bounds", w.s, len(w.exportedFields), i) + } + + return w.exportedFields[i].Addr().Interface() +} + +type anySliceArrayReflect struct { + slice reflect.Value +} + +func (a anySliceArrayReflect) Dimensions() []ArrayDimension { + if a.slice.IsNil() { + return nil + } + + return []ArrayDimension{{Length: int32(a.slice.Len()), LowerBound: 1}} +} + +func (a anySliceArrayReflect) Index(i int) any { + return a.slice.Index(i).Interface() +} + +func (a anySliceArrayReflect) IndexType() any { + return reflect.New(a.slice.Type().Elem()).Elem().Interface() +} + +func (a *anySliceArrayReflect) SetDimensions(dimensions []ArrayDimension) error { + sliceType := a.slice.Type() + + if dimensions == nil { + a.slice.Set(reflect.Zero(sliceType)) + return nil + } + + elementCount := cardinality(dimensions) + slice := reflect.MakeSlice(sliceType, elementCount, elementCount) + a.slice.Set(slice) + return nil +} + +func (a *anySliceArrayReflect) ScanIndex(i int) any { + return a.slice.Index(i).Addr().Interface() +} + +func (a *anySliceArrayReflect) ScanIndexType() any { + return reflect.New(a.slice.Type().Elem()).Interface() +} + +type anyMultiDimSliceArray struct { + slice reflect.Value + dims []ArrayDimension +} + +func (a *anyMultiDimSliceArray) Dimensions() []ArrayDimension { + if a.slice.IsNil() { + return nil + } + + s := a.slice + for { + a.dims = append(a.dims, ArrayDimension{Length: int32(s.Len()), LowerBound: 1}) + if s.Len() > 0 { + s = s.Index(0) + } else { + break + } + if s.Type().Kind() == reflect.Slice { + } else { + break + } + } + + return a.dims +} + +func (a *anyMultiDimSliceArray) Index(i int) any { + if len(a.dims) == 1 { + return a.slice.Index(i).Interface() + } + + indexes := make([]int, len(a.dims)) + for j := len(a.dims) - 1; j >= 0; j-- { + dimLen := int(a.dims[j].Length) + indexes[j] = i % dimLen + i /= dimLen + } + + v := a.slice + for _, si := range indexes { + v = v.Index(si) + } + + return v.Interface() +} + +func (a *anyMultiDimSliceArray) IndexType() any { + lowestSliceType := a.slice.Type() + for ; lowestSliceType.Elem().Kind() == reflect.Slice; lowestSliceType = lowestSliceType.Elem() { + } + return reflect.New(lowestSliceType.Elem()).Elem().Interface() +} + +func (a *anyMultiDimSliceArray) SetDimensions(dimensions []ArrayDimension) error { + sliceType := a.slice.Type() + + if dimensions == nil { + a.slice.Set(reflect.Zero(sliceType)) + return nil + } + + switch len(dimensions) { + case 0: + // Empty, but non-nil array + slice := reflect.MakeSlice(sliceType, 0, 0) + a.slice.Set(slice) + return nil + case 1: + elementCount := cardinality(dimensions) + slice := reflect.MakeSlice(sliceType, elementCount, elementCount) + a.slice.Set(slice) + return nil + default: + sliceDimensionCount := 1 + lowestSliceType := sliceType + for ; lowestSliceType.Elem().Kind() == reflect.Slice; lowestSliceType = lowestSliceType.Elem() { + sliceDimensionCount++ + } + + if sliceDimensionCount != len(dimensions) { + return fmt.Errorf("PostgreSQL array has %d dimensions but slice has %d dimensions", len(dimensions), sliceDimensionCount) + } + + elementCount := cardinality(dimensions) + flatSlice := reflect.MakeSlice(lowestSliceType, elementCount, elementCount) + + multiDimSlice := a.makeMultidimensionalSlice(sliceType, dimensions, flatSlice, 0) + a.slice.Set(multiDimSlice) + + // Now that a.slice is a multi-dimensional slice with the underlying data pointed at flatSlice change a.slice to + // flatSlice so ScanIndex only has to handle simple one dimensional slices. + a.slice = flatSlice + + return nil + } +} + +func (a *anyMultiDimSliceArray) makeMultidimensionalSlice(sliceType reflect.Type, dimensions []ArrayDimension, flatSlice reflect.Value, flatSliceIdx int) reflect.Value { + if len(dimensions) == 1 { + endIdx := flatSliceIdx + int(dimensions[0].Length) + return flatSlice.Slice3(flatSliceIdx, endIdx, endIdx) + } + + sliceLen := int(dimensions[0].Length) + slice := reflect.MakeSlice(sliceType, sliceLen, sliceLen) + for i := range sliceLen { + subSlice := a.makeMultidimensionalSlice(sliceType.Elem(), dimensions[1:], flatSlice, flatSliceIdx+(i*int(dimensions[1].Length))) + slice.Index(i).Set(subSlice) + } + + return slice +} + +func (a *anyMultiDimSliceArray) ScanIndex(i int) any { + return a.slice.Index(i).Addr().Interface() +} + +func (a *anyMultiDimSliceArray) ScanIndexType() any { + lowestSliceType := a.slice.Type() + for ; lowestSliceType.Elem().Kind() == reflect.Slice; lowestSliceType = lowestSliceType.Elem() { + } + return reflect.New(lowestSliceType.Elem()).Interface() +} + +type anyArrayArrayReflect struct { + array reflect.Value +} + +func (a anyArrayArrayReflect) Dimensions() []ArrayDimension { + return []ArrayDimension{{Length: int32(a.array.Len()), LowerBound: 1}} +} + +func (a anyArrayArrayReflect) Index(i int) any { + return a.array.Index(i).Interface() +} + +func (a anyArrayArrayReflect) IndexType() any { + return reflect.New(a.array.Type().Elem()).Elem().Interface() +} + +func (a *anyArrayArrayReflect) SetDimensions(dimensions []ArrayDimension) error { + if dimensions == nil { + return fmt.Errorf("anyArrayArrayReflect: cannot scan NULL into %v", a.array.Type().String()) + } + + if len(dimensions) != 1 { + return fmt.Errorf("anyArrayArrayReflect: cannot scan multi-dimensional array into %v", a.array.Type().String()) + } + + if int(dimensions[0].Length) != a.array.Len() { + return fmt.Errorf("anyArrayArrayReflect: cannot scan array with length %v into %v", dimensions[0].Length, a.array.Type().String()) + } + + return nil +} + +func (a *anyArrayArrayReflect) ScanIndex(i int) any { + return a.array.Index(i).Addr().Interface() +} + +func (a *anyArrayArrayReflect) ScanIndexType() any { + return reflect.New(a.array.Type().Elem()).Interface() +} diff --git a/vendor/github.com/jackc/pgx/v5/pgtype/bytea.go b/vendor/github.com/jackc/pgx/v5/pgtype/bytea.go new file mode 100644 index 0000000000..6c4f0c5eaf --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgtype/bytea.go @@ -0,0 +1,254 @@ +package pgtype + +import ( + "database/sql/driver" + "encoding/hex" + "fmt" +) + +type BytesScanner interface { + // ScanBytes receives a byte slice of driver memory that is only valid until the next database method call. + ScanBytes(v []byte) error +} + +type BytesValuer interface { + // BytesValue returns a byte slice of the byte data. The caller must not change the returned slice. + BytesValue() ([]byte, error) +} + +// DriverBytes is a byte slice that holds a reference to memory owned by the driver. It is only valid from the time it +// is scanned until Rows.Next or Rows.Close is called. It is never safe to use DriverBytes with QueryRow as Row.Scan +// internally calls Rows.Close before returning. +type DriverBytes []byte + +func (b *DriverBytes) ScanBytes(v []byte) error { + *b = v + return nil +} + +// PreallocBytes is a byte slice of preallocated memory that scanned bytes will be copied to. If it is too small a new +// slice will be allocated. +type PreallocBytes []byte + +func (b *PreallocBytes) ScanBytes(v []byte) error { + if v == nil { + *b = nil + return nil + } + + if len(v) <= len(*b) { + *b = (*b)[:len(v)] + } else { + *b = make(PreallocBytes, len(v)) + } + copy(*b, v) + return nil +} + +// UndecodedBytes can be used as a scan target to get the raw bytes from PostgreSQL without any decoding. +type UndecodedBytes []byte + +type scanPlanAnyToUndecodedBytes struct{} + +func (scanPlanAnyToUndecodedBytes) Scan(src []byte, dst any) error { + dstBuf := dst.(*UndecodedBytes) + if src == nil { + *dstBuf = nil + return nil + } + + *dstBuf = make([]byte, len(src)) + copy(*dstBuf, src) + return nil +} + +type ByteaCodec struct{} + +func (ByteaCodec) FormatSupported(format int16) bool { + return format == TextFormatCode || format == BinaryFormatCode +} + +func (ByteaCodec) PreferredFormat() int16 { + return BinaryFormatCode +} + +func (ByteaCodec) PlanEncode(m *Map, oid uint32, format int16, value any) EncodePlan { + switch format { + case BinaryFormatCode: + switch value.(type) { + case []byte: + return encodePlanBytesCodecBinaryBytes{} + case BytesValuer: + return encodePlanBytesCodecBinaryBytesValuer{} + } + case TextFormatCode: + switch value.(type) { + case []byte: + return encodePlanBytesCodecTextBytes{} + case BytesValuer: + return encodePlanBytesCodecTextBytesValuer{} + } + } + + return nil +} + +type encodePlanBytesCodecBinaryBytes struct{} + +func (encodePlanBytesCodecBinaryBytes) Encode(value any, buf []byte) (newBuf []byte, err error) { + b := value.([]byte) + if b == nil { + return nil, nil + } + + return append(buf, b...), nil +} + +type encodePlanBytesCodecBinaryBytesValuer struct{} + +func (encodePlanBytesCodecBinaryBytesValuer) Encode(value any, buf []byte) (newBuf []byte, err error) { + b, err := value.(BytesValuer).BytesValue() + if err != nil { + return nil, err + } + if b == nil { + return nil, nil + } + + return append(buf, b...), nil +} + +type encodePlanBytesCodecTextBytes struct{} + +func (encodePlanBytesCodecTextBytes) Encode(value any, buf []byte) (newBuf []byte, err error) { + b := value.([]byte) + if b == nil { + return nil, nil + } + + buf = append(buf, `\x`...) + buf = append(buf, hex.EncodeToString(b)...) + return buf, nil +} + +type encodePlanBytesCodecTextBytesValuer struct{} + +func (encodePlanBytesCodecTextBytesValuer) Encode(value any, buf []byte) (newBuf []byte, err error) { + b, err := value.(BytesValuer).BytesValue() + if err != nil { + return nil, err + } + if b == nil { + return nil, nil + } + + buf = append(buf, `\x`...) + buf = append(buf, hex.EncodeToString(b)...) + return buf, nil +} + +func (ByteaCodec) PlanScan(m *Map, oid uint32, format int16, target any) ScanPlan { + switch format { + case BinaryFormatCode: + switch target.(type) { + case *[]byte: + return scanPlanBinaryBytesToBytes{} + case BytesScanner: + return scanPlanBinaryBytesToBytesScanner{} + } + case TextFormatCode: + switch target.(type) { + case *[]byte: + return scanPlanTextByteaToBytes{} + case BytesScanner: + return scanPlanTextByteaToBytesScanner{} + } + } + + return nil +} + +type scanPlanBinaryBytesToBytes struct{} + +func (scanPlanBinaryBytesToBytes) Scan(src []byte, dst any) error { + dstBuf := dst.(*[]byte) + if src == nil { + *dstBuf = nil + return nil + } + + *dstBuf = make([]byte, len(src)) + copy(*dstBuf, src) + return nil +} + +type scanPlanBinaryBytesToBytesScanner struct{} + +func (scanPlanBinaryBytesToBytesScanner) Scan(src []byte, dst any) error { + scanner := (dst).(BytesScanner) + return scanner.ScanBytes(src) +} + +type scanPlanTextByteaToBytes struct{} + +func (scanPlanTextByteaToBytes) Scan(src []byte, dst any) error { + dstBuf := dst.(*[]byte) + if src == nil { + *dstBuf = nil + return nil + } + + buf, err := decodeHexBytea(src) + if err != nil { + return err + } + *dstBuf = buf + + return nil +} + +type scanPlanTextByteaToBytesScanner struct{} + +func (scanPlanTextByteaToBytesScanner) Scan(src []byte, dst any) error { + scanner := (dst).(BytesScanner) + buf, err := decodeHexBytea(src) + if err != nil { + return err + } + return scanner.ScanBytes(buf) +} + +func decodeHexBytea(src []byte) ([]byte, error) { + if src == nil { + return nil, nil + } + + if len(src) < 2 || src[0] != '\\' || src[1] != 'x' { + return nil, fmt.Errorf("invalid hex format") + } + + buf := make([]byte, (len(src)-2)/2) + _, err := hex.Decode(buf, src[2:]) + if err != nil { + return nil, err + } + + return buf, nil +} + +func (c ByteaCodec) DecodeDatabaseSQLValue(m *Map, oid uint32, format int16, src []byte) (driver.Value, error) { + return c.DecodeValue(m, oid, format, src) +} + +func (c ByteaCodec) DecodeValue(m *Map, oid uint32, format int16, src []byte) (any, error) { + if src == nil { + return nil, nil + } + + var buf []byte + err := codecScan(c, m, oid, format, src, &buf) + if err != nil { + return nil, err + } + return buf, nil +} diff --git a/vendor/github.com/jackc/pgx/v5/pgtype/circle.go b/vendor/github.com/jackc/pgx/v5/pgtype/circle.go new file mode 100644 index 0000000000..ea1d629958 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgtype/circle.go @@ -0,0 +1,231 @@ +package pgtype + +import ( + "database/sql/driver" + "encoding/binary" + "fmt" + "math" + "strconv" + "strings" + + "github.com/jackc/pgx/v5/internal/pgio" +) + +type CircleScanner interface { + ScanCircle(v Circle) error +} + +type CircleValuer interface { + CircleValue() (Circle, error) +} + +type Circle struct { + P Vec2 + R float64 + Valid bool +} + +// ScanCircle implements the [CircleScanner] interface. +func (c *Circle) ScanCircle(v Circle) error { + *c = v + return nil +} + +// CircleValue implements the [CircleValuer] interface. +func (c Circle) CircleValue() (Circle, error) { + return c, nil +} + +// Scan implements the [database/sql.Scanner] interface. +func (dst *Circle) Scan(src any) error { + if src == nil { + *dst = Circle{} + return nil + } + + if src, ok := src.(string); ok { + return scanPlanTextAnyToCircleScanner{}.Scan([]byte(src), dst) + } + + return fmt.Errorf("cannot scan %T", src) +} + +// Value implements the [database/sql/driver.Valuer] interface. +func (src Circle) Value() (driver.Value, error) { + if !src.Valid { + return nil, nil + } + + buf, err := CircleCodec{}.PlanEncode(nil, 0, TextFormatCode, src).Encode(src, nil) + if err != nil { + return nil, err + } + return string(buf), err +} + +type CircleCodec struct{} + +func (CircleCodec) FormatSupported(format int16) bool { + return format == TextFormatCode || format == BinaryFormatCode +} + +func (CircleCodec) PreferredFormat() int16 { + return BinaryFormatCode +} + +func (CircleCodec) PlanEncode(m *Map, oid uint32, format int16, value any) EncodePlan { + if _, ok := value.(CircleValuer); !ok { + return nil + } + + switch format { + case BinaryFormatCode: + return encodePlanCircleCodecBinary{} + case TextFormatCode: + return encodePlanCircleCodecText{} + } + + return nil +} + +type encodePlanCircleCodecBinary struct{} + +func (encodePlanCircleCodecBinary) Encode(value any, buf []byte) (newBuf []byte, err error) { + circle, err := value.(CircleValuer).CircleValue() + if err != nil { + return nil, err + } + + if !circle.Valid { + return nil, nil + } + + buf = pgio.AppendUint64(buf, math.Float64bits(circle.P.X)) + buf = pgio.AppendUint64(buf, math.Float64bits(circle.P.Y)) + buf = pgio.AppendUint64(buf, math.Float64bits(circle.R)) + return buf, nil +} + +type encodePlanCircleCodecText struct{} + +func (encodePlanCircleCodecText) Encode(value any, buf []byte) (newBuf []byte, err error) { + circle, err := value.(CircleValuer).CircleValue() + if err != nil { + return nil, err + } + + if !circle.Valid { + return nil, nil + } + + buf = append(buf, fmt.Sprintf(`<(%s,%s),%s>`, + strconv.FormatFloat(circle.P.X, 'f', -1, 64), + strconv.FormatFloat(circle.P.Y, 'f', -1, 64), + strconv.FormatFloat(circle.R, 'f', -1, 64), + )...) + return buf, nil +} + +func (CircleCodec) PlanScan(m *Map, oid uint32, format int16, target any) ScanPlan { + switch format { + case BinaryFormatCode: + if _, ok := target.(CircleScanner); ok { + return scanPlanBinaryCircleToCircleScanner{} + } + case TextFormatCode: + if _, ok := target.(CircleScanner); ok { + return scanPlanTextAnyToCircleScanner{} + } + } + + return nil +} + +func (c CircleCodec) DecodeDatabaseSQLValue(m *Map, oid uint32, format int16, src []byte) (driver.Value, error) { + return codecDecodeToTextFormat(c, m, oid, format, src) +} + +func (c CircleCodec) DecodeValue(m *Map, oid uint32, format int16, src []byte) (any, error) { + if src == nil { + return nil, nil + } + + var circle Circle + err := codecScan(c, m, oid, format, src, &circle) + if err != nil { + return nil, err + } + return circle, nil +} + +type scanPlanBinaryCircleToCircleScanner struct{} + +func (scanPlanBinaryCircleToCircleScanner) Scan(src []byte, dst any) error { + scanner := (dst).(CircleScanner) + + if src == nil { + return scanner.ScanCircle(Circle{}) + } + + if len(src) != 24 { + return fmt.Errorf("invalid length for Circle: %v", len(src)) + } + + x := binary.BigEndian.Uint64(src) + y := binary.BigEndian.Uint64(src[8:]) + r := binary.BigEndian.Uint64(src[16:]) + + return scanner.ScanCircle(Circle{ + P: Vec2{math.Float64frombits(x), math.Float64frombits(y)}, + R: math.Float64frombits(r), + Valid: true, + }) +} + +type scanPlanTextAnyToCircleScanner struct{} + +func (scanPlanTextAnyToCircleScanner) Scan(src []byte, dst any) error { + scanner := (dst).(CircleScanner) + + if src == nil { + return scanner.ScanCircle(Circle{}) + } + + if len(src) < 9 { + return fmt.Errorf("invalid length for Circle: %v", len(src)) + } + + // Expected format: <(x,y),r> + str, ok := strings.CutPrefix(string(src), "<(") + if !ok { + return fmt.Errorf("invalid format for Circle") + } + str, ok = strings.CutSuffix(str, ">") + if !ok { + return fmt.Errorf("invalid format for Circle") + } + + sx, str, found := strings.Cut(str, ",") + if !found { + return fmt.Errorf("invalid format for Circle") + } + sy, sr, found := strings.Cut(str, "),") + if !found { + return fmt.Errorf("invalid format for Circle") + } + + x, err := strconv.ParseFloat(sx, 64) + if err != nil { + return err + } + y, err := strconv.ParseFloat(sy, 64) + if err != nil { + return err + } + r, err := strconv.ParseFloat(sr, 64) + if err != nil { + return err + } + + return scanner.ScanCircle(Circle{P: Vec2{x, y}, R: r, Valid: true}) +} diff --git a/vendor/github.com/jackc/pgx/v5/pgtype/composite.go b/vendor/github.com/jackc/pgx/v5/pgtype/composite.go new file mode 100644 index 0000000000..7f96ab4902 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgtype/composite.go @@ -0,0 +1,613 @@ +package pgtype + +import ( + "database/sql/driver" + "encoding/binary" + "errors" + "fmt" + "strings" + + "github.com/jackc/pgx/v5/internal/pgio" +) + +// CompositeIndexGetter is a type accessed by index that can be converted into a PostgreSQL composite. +type CompositeIndexGetter interface { + // IsNull returns true if the value is SQL NULL. + IsNull() bool + + // Index returns the element at i. + Index(i int) any +} + +// CompositeIndexScanner is a type accessed by index that can be scanned from a PostgreSQL composite. +type CompositeIndexScanner interface { + // ScanNull sets the value to SQL NULL. + ScanNull() error + + // ScanIndex returns a value usable as a scan target for i. + ScanIndex(i int) any +} + +type CompositeCodecField struct { + Name string + Type *Type +} + +type CompositeCodec struct { + Fields []CompositeCodecField +} + +func (c *CompositeCodec) FormatSupported(format int16) bool { + for _, f := range c.Fields { + if !f.Type.Codec.FormatSupported(format) { + return false + } + } + + return true +} + +func (c *CompositeCodec) PreferredFormat() int16 { + if c.FormatSupported(BinaryFormatCode) { + return BinaryFormatCode + } + return TextFormatCode +} + +func (c *CompositeCodec) PlanEncode(m *Map, oid uint32, format int16, value any) EncodePlan { + if _, ok := value.(CompositeIndexGetter); !ok { + return nil + } + + switch format { + case BinaryFormatCode: + return &encodePlanCompositeCodecCompositeIndexGetterToBinary{cc: c, m: m} + case TextFormatCode: + return &encodePlanCompositeCodecCompositeIndexGetterToText{cc: c, m: m} + } + + return nil +} + +type encodePlanCompositeCodecCompositeIndexGetterToBinary struct { + cc *CompositeCodec + m *Map +} + +func (plan *encodePlanCompositeCodecCompositeIndexGetterToBinary) Encode(value any, buf []byte) (newBuf []byte, err error) { + getter := value.(CompositeIndexGetter) + + if getter.IsNull() { + return nil, nil + } + + builder := NewCompositeBinaryBuilder(plan.m, buf) + for i, field := range plan.cc.Fields { + builder.AppendValue(field.Type.OID, getter.Index(i)) + } + + return builder.Finish() +} + +type encodePlanCompositeCodecCompositeIndexGetterToText struct { + cc *CompositeCodec + m *Map +} + +func (plan *encodePlanCompositeCodecCompositeIndexGetterToText) Encode(value any, buf []byte) (newBuf []byte, err error) { + getter := value.(CompositeIndexGetter) + + if getter.IsNull() { + return nil, nil + } + + b := NewCompositeTextBuilder(plan.m, buf) + for i, field := range plan.cc.Fields { + b.AppendValue(field.Type.OID, getter.Index(i)) + } + + return b.Finish() +} + +func (c *CompositeCodec) PlanScan(m *Map, oid uint32, format int16, target any) ScanPlan { + switch format { + case BinaryFormatCode: + if _, ok := target.(CompositeIndexScanner); ok { + return &scanPlanBinaryCompositeToCompositeIndexScanner{cc: c, m: m} + } + case TextFormatCode: + if _, ok := target.(CompositeIndexScanner); ok { + return &scanPlanTextCompositeToCompositeIndexScanner{cc: c, m: m} + } + } + + return nil +} + +type scanPlanBinaryCompositeToCompositeIndexScanner struct { + cc *CompositeCodec + m *Map +} + +func (plan *scanPlanBinaryCompositeToCompositeIndexScanner) Scan(src []byte, target any) error { + targetScanner := (target).(CompositeIndexScanner) + + if src == nil { + return targetScanner.ScanNull() + } + + scanner := NewCompositeBinaryScanner(plan.m, src) + for i, field := range plan.cc.Fields { + if scanner.Next() { + fieldTarget := targetScanner.ScanIndex(i) + if fieldTarget != nil { + fieldPlan := plan.m.PlanScan(field.Type.OID, BinaryFormatCode, fieldTarget) + if fieldPlan == nil { + return fmt.Errorf("unable to encode %v into OID %d in binary format", field, field.Type.OID) + } + + err := fieldPlan.Scan(scanner.Bytes(), fieldTarget) + if err != nil { + return err + } + } + } else { + return errors.New("read past end of composite") + } + } + + if err := scanner.Err(); err != nil { + return err + } + + return nil +} + +type scanPlanTextCompositeToCompositeIndexScanner struct { + cc *CompositeCodec + m *Map +} + +func (plan *scanPlanTextCompositeToCompositeIndexScanner) Scan(src []byte, target any) error { + targetScanner := (target).(CompositeIndexScanner) + + if src == nil { + return targetScanner.ScanNull() + } + + scanner := NewCompositeTextScanner(plan.m, src) + for i, field := range plan.cc.Fields { + if scanner.Next() { + fieldTarget := targetScanner.ScanIndex(i) + if fieldTarget != nil { + fieldPlan := plan.m.PlanScan(field.Type.OID, TextFormatCode, fieldTarget) + if fieldPlan == nil { + return fmt.Errorf("unable to encode %v into OID %d in text format", field, field.Type.OID) + } + + err := fieldPlan.Scan(scanner.Bytes(), fieldTarget) + if err != nil { + return err + } + } + } else { + return errors.New("read past end of composite") + } + } + + if err := scanner.Err(); err != nil { + return err + } + + return nil +} + +func (c *CompositeCodec) DecodeDatabaseSQLValue(m *Map, oid uint32, format int16, src []byte) (driver.Value, error) { + if src == nil { + return nil, nil + } + + switch format { + case TextFormatCode: + return string(src), nil + case BinaryFormatCode: + buf := make([]byte, len(src)) + copy(buf, src) + return buf, nil + default: + return nil, fmt.Errorf("unknown format code %d", format) + } +} + +func (c *CompositeCodec) DecodeValue(m *Map, oid uint32, format int16, src []byte) (any, error) { + if src == nil { + return nil, nil + } + + switch format { + case TextFormatCode: + scanner := NewCompositeTextScanner(m, src) + values := make(map[string]any, len(c.Fields)) + for i := 0; scanner.Next() && i < len(c.Fields); i++ { + var v any + fieldPlan := m.PlanScan(c.Fields[i].Type.OID, TextFormatCode, &v) + if fieldPlan == nil { + return nil, fmt.Errorf("unable to scan OID %d in text format into %v", c.Fields[i].Type.OID, v) + } + + err := fieldPlan.Scan(scanner.Bytes(), &v) + if err != nil { + return nil, err + } + + values[c.Fields[i].Name] = v + } + + if err := scanner.Err(); err != nil { + return nil, err + } + + return values, nil + case BinaryFormatCode: + scanner := NewCompositeBinaryScanner(m, src) + values := make(map[string]any, len(c.Fields)) + for i := 0; scanner.Next() && i < len(c.Fields); i++ { + var v any + fieldPlan := m.PlanScan(scanner.OID(), BinaryFormatCode, &v) + if fieldPlan == nil { + return nil, fmt.Errorf("unable to scan OID %d in binary format into %v", scanner.OID(), v) + } + + err := fieldPlan.Scan(scanner.Bytes(), &v) + if err != nil { + return nil, err + } + + values[c.Fields[i].Name] = v + } + + if err := scanner.Err(); err != nil { + return nil, err + } + + return values, nil + default: + return nil, fmt.Errorf("unknown format code %d", format) + } +} + +type CompositeBinaryScanner struct { + m *Map + rp int + src []byte + + fieldCount int32 + fieldBytes []byte + fieldOID uint32 + err error +} + +// NewCompositeBinaryScanner a scanner over a binary encoded composite value. +func NewCompositeBinaryScanner(m *Map, src []byte) *CompositeBinaryScanner { + rp := 0 + if len(src[rp:]) < 4 { + return &CompositeBinaryScanner{err: fmt.Errorf("Record incomplete %v", src)} + } + + fieldCount := int32(binary.BigEndian.Uint32(src[rp:])) + rp += 4 + + return &CompositeBinaryScanner{ + m: m, + rp: rp, + src: src, + fieldCount: fieldCount, + } +} + +// Next advances the scanner to the next field. It returns false after the last field is read or an error occurs. After +// Next returns false, the Err method can be called to check if any errors occurred. +func (cfs *CompositeBinaryScanner) Next() bool { + if cfs.err != nil { + return false + } + + if cfs.rp == len(cfs.src) { + return false + } + + if len(cfs.src[cfs.rp:]) < 8 { + cfs.err = fmt.Errorf("Record incomplete %v", cfs.src) + return false + } + cfs.fieldOID = binary.BigEndian.Uint32(cfs.src[cfs.rp:]) + cfs.rp += 4 + + fieldLen := int(int32(binary.BigEndian.Uint32(cfs.src[cfs.rp:]))) + cfs.rp += 4 + + if fieldLen >= 0 { + if len(cfs.src[cfs.rp:]) < fieldLen { + cfs.err = fmt.Errorf("Record incomplete rp=%d src=%v", cfs.rp, cfs.src) + return false + } + cfs.fieldBytes = cfs.src[cfs.rp : cfs.rp+fieldLen] + cfs.rp += fieldLen + } else { + cfs.fieldBytes = nil + } + + return true +} + +func (cfs *CompositeBinaryScanner) FieldCount() int { + return int(cfs.fieldCount) +} + +// Bytes returns the bytes of the field most recently read by Scan(). +func (cfs *CompositeBinaryScanner) Bytes() []byte { + return cfs.fieldBytes +} + +// OID returns the OID of the field most recently read by Scan(). +func (cfs *CompositeBinaryScanner) OID() uint32 { + return cfs.fieldOID +} + +// Err returns any error encountered by the scanner. +func (cfs *CompositeBinaryScanner) Err() error { + return cfs.err +} + +type CompositeTextScanner struct { + m *Map + rp int + src []byte + + fieldBytes []byte + err error +} + +// NewCompositeTextScanner a scanner over a text encoded composite value. +func NewCompositeTextScanner(m *Map, src []byte) *CompositeTextScanner { + if len(src) < 2 { + return &CompositeTextScanner{err: fmt.Errorf("Record incomplete %v", src)} + } + + if src[0] != '(' { + return &CompositeTextScanner{err: fmt.Errorf("composite text format must start with '('")} + } + + if src[len(src)-1] != ')' { + return &CompositeTextScanner{err: fmt.Errorf("composite text format must end with ')'")} + } + + return &CompositeTextScanner{ + m: m, + rp: 1, + src: src, + } +} + +// Next advances the scanner to the next field. It returns false after the last field is read or an error occurs. After +// Next returns false, the Err method can be called to check if any errors occurred. +func (cfs *CompositeTextScanner) Next() bool { + if cfs.err != nil { + return false + } + + if cfs.rp == len(cfs.src) { + return false + } + + switch cfs.src[cfs.rp] { + case ',', ')': // null + cfs.rp++ + cfs.fieldBytes = nil + return true + case '"': // quoted value + cfs.rp++ + cfs.fieldBytes = make([]byte, 0, 16) + quotedValue: + for { + ch := cfs.src[cfs.rp] + + switch ch { + case '"': + cfs.rp++ + if cfs.src[cfs.rp] == '"' { + cfs.fieldBytes = append(cfs.fieldBytes, '"') + cfs.rp++ + } else { + break quotedValue + } + case '\\': + cfs.rp++ + cfs.fieldBytes = append(cfs.fieldBytes, cfs.src[cfs.rp]) + cfs.rp++ + default: + cfs.fieldBytes = append(cfs.fieldBytes, ch) + cfs.rp++ + } + } + cfs.rp++ + return true + default: // unquoted value + start := cfs.rp + for { + ch := cfs.src[cfs.rp] + if ch == ',' || ch == ')' { + break + } + cfs.rp++ + } + cfs.fieldBytes = cfs.src[start:cfs.rp] + cfs.rp++ + return true + } +} + +// Bytes returns the bytes of the field most recently read by Scan(). +func (cfs *CompositeTextScanner) Bytes() []byte { + return cfs.fieldBytes +} + +// Err returns any error encountered by the scanner. +func (cfs *CompositeTextScanner) Err() error { + return cfs.err +} + +type CompositeBinaryBuilder struct { + m *Map + buf []byte + startIdx int + fieldCount uint32 + err error +} + +func NewCompositeBinaryBuilder(m *Map, buf []byte) *CompositeBinaryBuilder { + startIdx := len(buf) + buf = append(buf, 0, 0, 0, 0) // allocate room for number of fields + return &CompositeBinaryBuilder{m: m, buf: buf, startIdx: startIdx} +} + +func (b *CompositeBinaryBuilder) AppendValue(oid uint32, field any) { + if b.err != nil { + return + } + + isNil, callNilDriverValuer := isNilDriverValuer(field) + if isNil && !callNilDriverValuer { + b.buf = pgio.AppendUint32(b.buf, oid) + b.buf = pgio.AppendInt32(b.buf, -1) + b.fieldCount++ + return + } + + var plan EncodePlan + if isNil { + plan = &encodePlanDriverValuer{m: b.m, oid: oid, formatCode: BinaryFormatCode} + } else { + plan = b.m.PlanEncode(oid, BinaryFormatCode, field) + if plan == nil { + b.err = fmt.Errorf("unable to encode %v into OID %d in binary format", field, oid) + return + } + } + + b.buf = pgio.AppendUint32(b.buf, oid) + lengthPos := len(b.buf) + b.buf = pgio.AppendInt32(b.buf, -1) + fieldBuf, err := plan.Encode(field, b.buf) + if err != nil { + b.err = err + return + } + if fieldBuf != nil { + binary.BigEndian.PutUint32(fieldBuf[lengthPos:], uint32(len(fieldBuf)-len(b.buf))) + b.buf = fieldBuf + } + + b.fieldCount++ +} + +func (b *CompositeBinaryBuilder) Finish() ([]byte, error) { + if b.err != nil { + return nil, b.err + } + + binary.BigEndian.PutUint32(b.buf[b.startIdx:], b.fieldCount) + return b.buf, nil +} + +type CompositeTextBuilder struct { + m *Map + buf []byte + startIdx int + fieldCount uint32 + err error + fieldBuf [32]byte +} + +func NewCompositeTextBuilder(m *Map, buf []byte) *CompositeTextBuilder { + buf = append(buf, '(') // allocate room for number of fields + return &CompositeTextBuilder{m: m, buf: buf} +} + +func (b *CompositeTextBuilder) AppendValue(oid uint32, field any) { + if b.err != nil { + return + } + + isNil, callNilDriverValuer := isNilDriverValuer(field) + if isNil && !callNilDriverValuer { + b.buf = append(b.buf, ',') + return + } + + var plan EncodePlan + if isNil { + plan = &encodePlanDriverValuer{m: b.m, oid: oid, formatCode: TextFormatCode} + } else { + plan = b.m.PlanEncode(oid, TextFormatCode, field) + if plan == nil { + b.err = fmt.Errorf("unable to encode %v into OID %d in text format", field, oid) + return + } + } + + fieldBuf, err := plan.Encode(field, b.fieldBuf[0:0]) + if err != nil { + b.err = err + return + } + if fieldBuf != nil { + b.buf = append(b.buf, quoteCompositeFieldIfNeeded(string(fieldBuf))...) + } + + b.buf = append(b.buf, ',') +} + +func (b *CompositeTextBuilder) Finish() ([]byte, error) { + if b.err != nil { + return nil, b.err + } + + b.buf[len(b.buf)-1] = ')' + return b.buf, nil +} + +var quoteCompositeReplacer = strings.NewReplacer(`\`, `\\`, `"`, `\"`) + +func quoteCompositeField(src string) string { + return `"` + quoteCompositeReplacer.Replace(src) + `"` +} + +func quoteCompositeFieldIfNeeded(src string) string { + if src == "" || src[0] == ' ' || src[len(src)-1] == ' ' || strings.ContainsAny(src, `(),"\`) { + return quoteCompositeField(src) + } + return src +} + +// CompositeFields represents the values of a composite value. It can be used as an encoding source or as a scan target. +// It cannot scan a NULL, but the composite fields can be NULL. +type CompositeFields []any + +func (cf CompositeFields) SkipUnderlyingTypePlan() {} + +func (cf CompositeFields) IsNull() bool { + return cf == nil +} + +func (cf CompositeFields) Index(i int) any { + return cf[i] +} + +func (cf CompositeFields) ScanNull() error { + return fmt.Errorf("cannot scan NULL into CompositeFields") +} + +func (cf CompositeFields) ScanIndex(i int) any { + return cf[i] +} diff --git a/vendor/github.com/jackc/pgx/v5/pgtype/convert.go b/vendor/github.com/jackc/pgx/v5/pgtype/convert.go new file mode 100644 index 0000000000..6693338410 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgtype/convert.go @@ -0,0 +1,108 @@ +package pgtype + +import ( + "reflect" +) + +func NullAssignTo(dst any) error { + dstPtr := reflect.ValueOf(dst) + + // AssignTo dst must always be a pointer + if dstPtr.Kind() != reflect.Pointer { + return &nullAssignmentError{dst: dst} + } + + dstVal := dstPtr.Elem() + + switch dstVal.Kind() { + case reflect.Pointer, reflect.Slice, reflect.Map: + dstVal.Set(reflect.Zero(dstVal.Type())) + return nil + } + + return &nullAssignmentError{dst: dst} +} + +var kindTypes map[reflect.Kind]reflect.Type + +func toInterface(dst reflect.Value, t reflect.Type) (any, bool) { + nextDst := dst.Convert(t) + return nextDst.Interface(), dst.Type() != nextDst.Type() +} + +// GetAssignToDstType attempts to convert dst to something AssignTo can assign +// to. If dst is a pointer to pointer it allocates a value and returns the +// dereferences pointer. If dst is a named type such as *Foo where Foo is type +// Foo int16, it converts dst to *int16. +// +// GetAssignToDstType returns the converted dst and a bool representing if any +// change was made. +func GetAssignToDstType(dst any) (any, bool) { + dstPtr := reflect.ValueOf(dst) + + // AssignTo dst must always be a pointer + if dstPtr.Kind() != reflect.Pointer { + return nil, false + } + + dstVal := dstPtr.Elem() + + // if dst is a pointer to pointer, allocate space try again with the dereferenced pointer + if dstVal.Kind() == reflect.Pointer { + dstVal.Set(reflect.New(dstVal.Type().Elem())) + return dstVal.Interface(), true + } + + // if dst is pointer to a base type that has been renamed + if baseValType, ok := kindTypes[dstVal.Kind()]; ok { + return toInterface(dstPtr, reflect.PointerTo(baseValType)) + } + + if dstVal.Kind() == reflect.Slice { + if baseElemType, ok := kindTypes[dstVal.Type().Elem().Kind()]; ok { + return toInterface(dstPtr, reflect.PointerTo(reflect.SliceOf(baseElemType))) + } + } + + if dstVal.Kind() == reflect.Array { + if baseElemType, ok := kindTypes[dstVal.Type().Elem().Kind()]; ok { + return toInterface(dstPtr, reflect.PointerTo(reflect.ArrayOf(dstVal.Len(), baseElemType))) + } + } + + if dstVal.Kind() == reflect.Struct { + if dstVal.Type().NumField() == 1 && dstVal.Type().Field(0).Anonymous { + dstPtr = dstVal.Field(0).Addr() + nested := dstVal.Type().Field(0).Type + if nested.Kind() == reflect.Array { + if baseElemType, ok := kindTypes[nested.Elem().Kind()]; ok { + return toInterface(dstPtr, reflect.PointerTo(reflect.ArrayOf(nested.Len(), baseElemType))) + } + } + if _, ok := kindTypes[nested.Kind()]; ok && dstPtr.CanInterface() { + return dstPtr.Interface(), true + } + } + } + + return nil, false +} + +func init() { + kindTypes = map[reflect.Kind]reflect.Type{ + reflect.Bool: reflect.TypeFor[bool](), + reflect.Float32: reflect.TypeFor[float32](), + reflect.Float64: reflect.TypeFor[float64](), + reflect.Int: reflect.TypeFor[int](), + reflect.Int8: reflect.TypeFor[int8](), + reflect.Int16: reflect.TypeFor[int16](), + reflect.Int32: reflect.TypeFor[int32](), + reflect.Int64: reflect.TypeFor[int64](), + reflect.Uint: reflect.TypeFor[uint](), + reflect.Uint8: reflect.TypeFor[uint8](), + reflect.Uint16: reflect.TypeFor[uint16](), + reflect.Uint32: reflect.TypeFor[uint32](), + reflect.Uint64: reflect.TypeFor[uint64](), + reflect.String: reflect.TypeFor[string](), + } +} diff --git a/vendor/github.com/jackc/pgx/v5/pgtype/date.go b/vendor/github.com/jackc/pgx/v5/pgtype/date.go new file mode 100644 index 0000000000..305d83c31a --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgtype/date.go @@ -0,0 +1,412 @@ +package pgtype + +import ( + "database/sql/driver" + "encoding/binary" + "encoding/json" + "fmt" + "strconv" + "time" + + "github.com/jackc/pgx/v5/internal/pgio" +) + +type DateScanner interface { + ScanDate(v Date) error +} + +type DateValuer interface { + DateValue() (Date, error) +} + +type Date struct { + Time time.Time + InfinityModifier InfinityModifier + Valid bool +} + +// ScanDate implements the [DateScanner] interface. +func (d *Date) ScanDate(v Date) error { + *d = v + return nil +} + +// DateValue implements the [DateValuer] interface. +func (d Date) DateValue() (Date, error) { + return d, nil +} + +const ( + negativeInfinityDayOffset = -2147483648 + infinityDayOffset = 2147483647 +) + +// Scan implements the [database/sql.Scanner] interface. +func (dst *Date) Scan(src any) error { + if src == nil { + *dst = Date{} + return nil + } + + switch src := src.(type) { + case string: + return scanPlanTextAnyToDateScanner{}.Scan([]byte(src), dst) + case time.Time: + *dst = Date{Time: src, Valid: true} + return nil + } + + return fmt.Errorf("cannot scan %T", src) +} + +// Value implements the [database/sql/driver.Valuer] interface. +func (src Date) Value() (driver.Value, error) { + if !src.Valid { + return nil, nil + } + + if src.InfinityModifier != Finite { + return src.InfinityModifier.String(), nil + } + return src.Time, nil +} + +// MarshalJSON implements the [encoding/json.Marshaler] interface. +func (src Date) MarshalJSON() ([]byte, error) { + if !src.Valid { + return []byte("null"), nil + } + + var s string + + switch src.InfinityModifier { + case Finite: + s = src.Time.Format("2006-01-02") + case Infinity: + s = "infinity" + case NegativeInfinity: + s = "-infinity" + } + + return json.Marshal(s) +} + +// UnmarshalJSON implements the [encoding/json.Unmarshaler] interface. +func (dst *Date) UnmarshalJSON(b []byte) error { + var s *string + err := json.Unmarshal(b, &s) + if err != nil { + return err + } + + if s == nil { + *dst = Date{} + return nil + } + + switch *s { + case "infinity": + *dst = Date{Valid: true, InfinityModifier: Infinity} + case "-infinity": + *dst = Date{Valid: true, InfinityModifier: -Infinity} + default: + t, err := time.ParseInLocation("2006-01-02", *s, time.UTC) + if err != nil { + return err + } + + *dst = Date{Time: t, Valid: true} + } + + return nil +} + +type DateCodec struct{} + +func (DateCodec) FormatSupported(format int16) bool { + return format == TextFormatCode || format == BinaryFormatCode +} + +func (DateCodec) PreferredFormat() int16 { + return BinaryFormatCode +} + +func (DateCodec) PlanEncode(m *Map, oid uint32, format int16, value any) EncodePlan { + if _, ok := value.(DateValuer); !ok { + return nil + } + + switch format { + case BinaryFormatCode: + return encodePlanDateCodecBinary{} + case TextFormatCode: + return encodePlanDateCodecText{} + } + + return nil +} + +type encodePlanDateCodecBinary struct{} + +func (encodePlanDateCodecBinary) Encode(value any, buf []byte) (newBuf []byte, err error) { + date, err := value.(DateValuer).DateValue() + if err != nil { + return nil, err + } + + if !date.Valid { + return nil, nil + } + + var daysSinceDateEpoch int32 + switch date.InfinityModifier { + case Finite: + tUnix := time.Date(date.Time.Year(), date.Time.Month(), date.Time.Day(), 0, 0, 0, 0, time.UTC).Unix() + dateEpoch := time.Date(2000, 1, 1, 0, 0, 0, 0, time.UTC).Unix() + + secSinceDateEpoch := tUnix - dateEpoch + daysSinceDateEpoch = int32(secSinceDateEpoch / 86400) + case Infinity: + daysSinceDateEpoch = infinityDayOffset + case NegativeInfinity: + daysSinceDateEpoch = negativeInfinityDayOffset + } + + return pgio.AppendInt32(buf, daysSinceDateEpoch), nil +} + +type encodePlanDateCodecText struct{} + +func (encodePlanDateCodecText) Encode(value any, buf []byte) (newBuf []byte, err error) { + date, err := value.(DateValuer).DateValue() + if err != nil { + return nil, err + } + + if !date.Valid { + return nil, nil + } + + switch date.InfinityModifier { + case Finite: + // Year 0000 is 1 BC + bc := false + year := date.Time.Year() + if year <= 0 { + year = -year + 1 + bc = true + } + + yearBytes := strconv.AppendInt(make([]byte, 0, 6), int64(year), 10) + for i := len(yearBytes); i < 4; i++ { + buf = append(buf, '0') + } + buf = append(buf, yearBytes...) + buf = append(buf, '-') + if date.Time.Month() < 10 { + buf = append(buf, '0') + } + buf = strconv.AppendInt(buf, int64(date.Time.Month()), 10) + buf = append(buf, '-') + if date.Time.Day() < 10 { + buf = append(buf, '0') + } + buf = strconv.AppendInt(buf, int64(date.Time.Day()), 10) + + if bc { + buf = append(buf, " BC"...) + } + case Infinity: + buf = append(buf, "infinity"...) + case NegativeInfinity: + buf = append(buf, "-infinity"...) + } + + return buf, nil +} + +func (DateCodec) PlanScan(m *Map, oid uint32, format int16, target any) ScanPlan { + switch format { + case BinaryFormatCode: + if _, ok := target.(DateScanner); ok { + return scanPlanBinaryDateToDateScanner{} + } + case TextFormatCode: + if _, ok := target.(DateScanner); ok { + return scanPlanTextAnyToDateScanner{} + } + } + + return nil +} + +type scanPlanBinaryDateToDateScanner struct{} + +func (scanPlanBinaryDateToDateScanner) Scan(src []byte, dst any) error { + scanner := (dst).(DateScanner) + + if src == nil { + return scanner.ScanDate(Date{}) + } + + if len(src) != 4 { + return fmt.Errorf("invalid length for date: %v", len(src)) + } + + dayOffset := int32(binary.BigEndian.Uint32(src)) + + switch dayOffset { + case infinityDayOffset: + return scanner.ScanDate(Date{InfinityModifier: Infinity, Valid: true}) + case negativeInfinityDayOffset: + return scanner.ScanDate(Date{InfinityModifier: -Infinity, Valid: true}) + default: + t := time.Date(2000, 1, int(1+dayOffset), 0, 0, 0, 0, time.UTC) + return scanner.ScanDate(Date{Time: t, Valid: true}) + } +} + +type scanPlanTextAnyToDateScanner struct{} + +func (scanPlanTextAnyToDateScanner) Scan(src []byte, dst any) error { + scanner := (dst).(DateScanner) + + if src == nil { + return scanner.ScanDate(Date{}) + } + + // Check infinity cases first + if len(src) == 8 && string(src) == "infinity" { + return scanner.ScanDate(Date{InfinityModifier: Infinity, Valid: true}) + } + if len(src) == 9 && string(src) == "-infinity" { + return scanner.ScanDate(Date{InfinityModifier: -Infinity, Valid: true}) + } + + // Format: YYYY-MM-DD or YYYY...-MM-DD BC + // Minimum: 10 chars (2000-01-01), with BC: 13 chars + if len(src) < 10 { + return fmt.Errorf("invalid date format") + } + + // Check for BC suffix + bc := false + datePart := src + if len(src) >= 13 && string(src[len(src)-3:]) == " BC" { + bc = true + datePart = src[:len(src)-3] + } + + // Find year-month separator (first dash after at least 4 digits) + yearEnd := -1 + for i := 4; i < len(datePart); i++ { + if datePart[i] == '-' { + yearEnd = i + break + } + if datePart[i] < '0' || datePart[i] > '9' { + return fmt.Errorf("invalid date format") + } + } + if yearEnd == -1 || yearEnd+6 > len(datePart) { + return fmt.Errorf("invalid date format") + } + + // Validate: -MM-DD structure after year + if datePart[yearEnd+3] != '-' { + return fmt.Errorf("invalid date format") + } + + // Parse year + year, err := parseDigits(datePart[:yearEnd]) + if err != nil { + return fmt.Errorf("invalid date format") + } + + // Parse month (2 digits) + month, err := parse2Digits(datePart[yearEnd+1 : yearEnd+3]) + if err != nil { + return fmt.Errorf("invalid date format") + } + + // Parse day (2 digits) + day, err := parse2Digits(datePart[yearEnd+4 : yearEnd+6]) + if err != nil { + return fmt.Errorf("invalid date format") + } + + // Ensure nothing extra after day + if yearEnd+6 != len(datePart) { + return fmt.Errorf("invalid date format") + } + + if bc { + year = -year + 1 + } + + t := time.Date(int(year), time.Month(month), int(day), 0, 0, 0, 0, time.UTC) + return scanner.ScanDate(Date{Time: t, Valid: true}) +} + +// parse2Digits parses exactly 2 ASCII digits. +func parse2Digits(b []byte) (int64, error) { + if len(b) != 2 { + return 0, fmt.Errorf("expected 2 digits") + } + d1, d2 := b[0], b[1] + if d1 < '0' || d1 > '9' || d2 < '0' || d2 > '9' { + return 0, fmt.Errorf("expected digits") + } + return int64(d1-'0')*10 + int64(d2-'0'), nil +} + +// parseDigits parses a sequence of ASCII digits. +func parseDigits(b []byte) (int64, error) { + if len(b) == 0 { + return 0, fmt.Errorf("empty") + } + var n int64 + for _, c := range b { + if c < '0' || c > '9' { + return 0, fmt.Errorf("non-digit") + } + n = n*10 + int64(c-'0') + } + return n, nil +} + +func (c DateCodec) DecodeDatabaseSQLValue(m *Map, oid uint32, format int16, src []byte) (driver.Value, error) { + if src == nil { + return nil, nil + } + + var date Date + err := codecScan(c, m, oid, format, src, &date) + if err != nil { + return nil, err + } + + if date.InfinityModifier != Finite { + return date.InfinityModifier.String(), nil + } + + return date.Time, nil +} + +func (c DateCodec) DecodeValue(m *Map, oid uint32, format int16, src []byte) (any, error) { + if src == nil { + return nil, nil + } + + var date Date + err := codecScan(c, m, oid, format, src, &date) + if err != nil { + return nil, err + } + + if date.InfinityModifier != Finite { + return date.InfinityModifier, nil + } + + return date.Time, nil +} diff --git a/vendor/github.com/jackc/pgx/v5/pgtype/doc.go b/vendor/github.com/jackc/pgx/v5/pgtype/doc.go new file mode 100644 index 0000000000..dbcdf692f6 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgtype/doc.go @@ -0,0 +1,196 @@ +// Package pgtype converts between Go and PostgreSQL values. +/* +The primary type is the [Map] type. It is a map of PostgreSQL types identified by OID (object ID) to a [Codec]. A [Codec] is +responsible for converting between Go and PostgreSQL values. [NewMap] creates a [Map] with all supported standard PostgreSQL +types already registered. Additional types can be registered with [Map.RegisterType]. + +Use [Map.Scan] and [Map.Encode] to decode PostgreSQL values to Go and encode Go values to PostgreSQL respectively. + +Base Type Mapping + +pgtype maps between all common base types directly between Go and PostgreSQL. In particular: + + Go PostgreSQL + ----------------------- + string varchar + text + + // Integers are automatically be converted to any other integer type if + // it can be done without overflow or underflow. + int8 + int16 smallint + int32 int + int64 bigint + int + uint8 + uint16 + uint32 + uint64 + uint + + // Floats are strict and do not automatically convert like integers. + float32 float4 + float64 float8 + + time.Time date + timestamp + timestamptz + + netip.Addr inet + netip.Prefix cidr + + []byte bytea + +Null Values + +pgtype can map NULLs in two ways. The first is types that can directly represent NULL such as Int4. They work in a +similar fashion to database/sql. The second is to use a pointer to a pointer. + + var foo pgtype.Text + var bar *string + err := conn.QueryRow("select foo, bar from widgets where id=$1", 42).Scan(&foo, &bar) + if err != nil { + return err + } + +When using nullable pgtype types as parameters for queries, one has to remember to explicitly set their Valid field to +true, otherwise the parameter's value will be NULL. + +JSON Support + +pgtype automatically marshals and unmarshals data from json and jsonb PostgreSQL types. + +Extending Existing PostgreSQL Type Support + +Generally, all Codecs will support interfaces that can be implemented to enable scanning and encoding. For example, +[PointCodec] can use any Go type that implements the [PointScanner] and [PointValuer] interfaces. So rather than use +[Point] an application can directly use its own point type with pgtype as long as it implements those interfaces. + +See example_custom_type_test.go for an example of a custom type for the PostgreSQL point type. + +Sometimes pgx supports a PostgreSQL type such as numeric but the Go type is in an external package that does not have +pgx support such as github.com/shopspring/decimal. These types can be registered with pgtype with custom conversion +logic. See https://github.com/jackc/pgx-shopspring-decimal and https://github.com/jackc/pgx-gofrs-uuid for example +integrations. + +New PostgreSQL Type Support + +pgtype uses the PostgreSQL OID to determine how to encode or decode a value. pgtype supports array, composite, domain, +and enum types. However, any type created in PostgreSQL with CREATE TYPE will receive a new OID. This means that the OID +of each new PostgreSQL type must be registered for pgtype to handle values of that type with the correct [Codec]. + +The [github.com/jackc/pgx/v5.Conn.LoadType] method can return a [*Type] for array, composite, domain, and enum types by +inspecting the database metadata. This [*Type] can then be registered with [Map.RegisterType]. + +For example, the following function could be called after a connection is established: + + func RegisterDataTypes(ctx context.Context, conn *pgx.Conn) error { + dataTypeNames := []string{ + "foo", + "_foo", + "bar", + "_bar", + } + + for _, typeName := range dataTypeNames { + dataType, err := conn.LoadType(ctx, typeName) + if err != nil { + return err + } + conn.TypeMap().RegisterType(dataType) + } + + return nil + } + +A type cannot be registered unless all types it depends on are already registered. e.g. An array type cannot be +registered until its element type is registered. + +[ArrayCodec] implements support for arrays. If pgtype supports type T then it can easily support []T by registering an +[ArrayCodec] for the appropriate PostgreSQL OID. In addition, [Array] type can support multi-dimensional arrays. + +[CompositeCodec] implements support for PostgreSQL composite types. Go structs can be scanned into if the public fields of +the struct are in the exact order and type of the PostgreSQL type or by implementing [CompositeIndexScanner] and +[CompositeIndexGetter]. + +Domain types are treated as their underlying type if the underlying type and the domain type are registered. + +PostgreSQL enums can usually be treated as text. However, [EnumCodec] implements support for interning strings which can +reduce memory usage. + +While pgtype will often still work with unregistered types it is highly recommended that all types be registered due to +an improvement in performance and the elimination of certain edge cases. + +If an entirely new PostgreSQL type (e.g. PostGIS types) is used then the application or a library can create a new +[Codec]. Then the OID / [Codec] mapping can be registered with [Map.RegisterType]. There is no difference between a [Codec] +defined and registered by the application and a [Codec] built in to pgtype. See any of the [Codec]s in pgtype for [Codec] +examples and for examples of type registration. + +Encoding Unknown Types + +pgtype works best when the OID of the PostgreSQL type is known. But in some cases such as using the simple protocol the +OID is unknown. In this case [Map.RegisterDefaultPgType] can be used to register an assumed OID for a particular Go type. + +Renamed Types + +If pgtype does not recognize a type and that type is a renamed simple type simple (e.g. type MyInt32 int32) pgtype acts +as if it is the underlying type. It currently cannot automatically detect the underlying type of renamed structs (eg.g. +type MyTime time.Time). + +Compatibility with [database/sql] + +pgtype also includes support for custom types implementing the [database/sql.Scanner] and [database/sql/driver.Valuer] +interfaces. + +Encoding Typed Nils + +pgtype encodes untyped and typed nils (e.g. nil and []byte(nil)) to the SQL NULL value without going through the [Codec] +system. This means that [Codec]s and other encoding logic do not have to handle nil or *T(nil). + +However, [database/sql] compatibility requires Value to be called on T(nil) when T implements [database/sql/driver.Valuer]. Therefore, +[database/sql/driver.Valuer] values are only considered NULL when *T(nil) where [database/sql/driver.Valuer] is implemented on T not on *T. See +https://github.com/golang/go/issues/8415 and +https://github.com/golang/go/commit/0ce1d79a6a771f7449ec493b993ed2a720917870. + +Child Records + +pgtype's support for arrays and composite records can be used to load records and their children in a single query. See +example_child_records_test.go for an example. + +Overview of Scanning Implementation + +The first step is to use the OID to lookup the correct [Codec]. The [Map] will call the [Codec.PlanScan] method to get a +plan for scanning into the Go value. A [Codec] will support scanning into one or more Go types. Oftentime these Go types +are interfaces rather than explicit types. For example, [PointCodec] can use any Go type that implements the [PointScanner] +and [PointValuer] interfaces. + +If a Go value is not supported directly by a [Codec] then [Map] will try see if it is a [database/sql.Scanner]. If is then that +interface will be used to scan the value. Most [database/sql.Scanner]s require the input to be in the text format (e.g. UUIDs and +numeric). However, pgx will typically have received the value in the binary format. In this case the binary value will be +parsed, reencoded as text, and then passed to the [database/sql.Scanner]. This may incur additional overhead for query results with +a large number of affected values. + +If a Go value is not supported directly by a [Codec] then [Map] will try wrapping it with additional logic and try again. +For example, [Int8Codec] does not support scanning into a renamed type (e.g. type myInt64 int64). But [Map] will detect that +myInt64 is a renamed type and create a plan that converts the value to the underlying int64 type and then passes that to +the [Codec] (see [TryFindUnderlyingTypeScanPlan]). + +These plan wrappers are contained in [Map.TryWrapScanPlanFuncs]. By default these contain shared logic to handle renamed +types, pointers to pointers, slices, composite types, etc. Additional plan wrappers can be added to seamlessly integrate +types that do not support pgx directly. For example, the before mentioned +https://github.com/jackc/pgx-shopspring-decimal package detects decimal.Decimal values, wraps them in something +implementing [NumericScanner] and passes that to the [Codec]. + +[Map.Scan] and [Map.Encode] are convenience methods that wrap [Map.PlanScan] and [Map.PlanEncode]. Determining how to scan or +encode a particular type may be a time consuming operation. Hence the planning and execution steps of a conversion are +internally separated. + +Reducing Compiled Binary Size + +[github.com/jackc/pgx/v5.QueryExecModeExec] and [github.com/jackc/pgx/v5.QueryExecModeSimpleProtocol] require the default +PostgreSQL type to be registered for each Go type used as a query parameter. By default pgx does this for all supported +types and their array variants. If an application does not use those query execution modes or manually registers the default +PostgreSQL type for the types it uses as query parameters it can use the build tag nopgxregisterdefaulttypes. This omits +the default type registration and reduces the compiled binary size by ~2MB. +*/ +package pgtype diff --git a/vendor/github.com/jackc/pgx/v5/pgtype/enum_codec.go b/vendor/github.com/jackc/pgx/v5/pgtype/enum_codec.go new file mode 100644 index 0000000000..5e787c1e29 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgtype/enum_codec.go @@ -0,0 +1,109 @@ +package pgtype + +import ( + "database/sql/driver" + "fmt" +) + +// EnumCodec is a codec that caches the strings it decodes. If the same string is read multiple times only one copy is +// allocated. These strings are only garbage collected when the EnumCodec is garbage collected. EnumCodec can be used +// for any text type not only enums, but it should only be used when there are a small number of possible values. +type EnumCodec struct { + membersMap map[string]string // map to quickly lookup member and reuse string instead of allocating +} + +func (EnumCodec) FormatSupported(format int16) bool { + return format == TextFormatCode || format == BinaryFormatCode +} + +func (EnumCodec) PreferredFormat() int16 { + return TextFormatCode +} + +func (EnumCodec) PlanEncode(m *Map, oid uint32, format int16, value any) EncodePlan { + switch format { + case TextFormatCode, BinaryFormatCode: + switch value.(type) { + case string: + return encodePlanTextCodecString{} + case []byte: + return encodePlanTextCodecByteSlice{} + case TextValuer: + return encodePlanTextCodecTextValuer{} + } + } + + return nil +} + +func (c *EnumCodec) PlanScan(m *Map, oid uint32, format int16, target any) ScanPlan { + switch format { + case TextFormatCode, BinaryFormatCode: + switch target.(type) { + case *string: + return &scanPlanTextAnyToEnumString{codec: c} + case *[]byte: + return scanPlanAnyToNewByteSlice{} + case TextScanner: + return &scanPlanTextAnyToEnumTextScanner{codec: c} + } + } + + return nil +} + +func (c *EnumCodec) DecodeDatabaseSQLValue(m *Map, oid uint32, format int16, src []byte) (driver.Value, error) { + return c.DecodeValue(m, oid, format, src) +} + +func (c *EnumCodec) DecodeValue(m *Map, oid uint32, format int16, src []byte) (any, error) { + if src == nil { + return nil, nil + } + + return c.lookupAndCacheString(src), nil +} + +// lookupAndCacheString looks for src in the members map. If it is not found it is added to the map. +func (c *EnumCodec) lookupAndCacheString(src []byte) string { + if c.membersMap == nil { + c.membersMap = make(map[string]string) + } + + if s, found := c.membersMap[string(src)]; found { + return s + } + + s := string(src) + c.membersMap[s] = s + return s +} + +type scanPlanTextAnyToEnumString struct { + codec *EnumCodec +} + +func (plan *scanPlanTextAnyToEnumString) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + p := (dst).(*string) + *p = plan.codec.lookupAndCacheString(src) + + return nil +} + +type scanPlanTextAnyToEnumTextScanner struct { + codec *EnumCodec +} + +func (plan *scanPlanTextAnyToEnumTextScanner) Scan(src []byte, dst any) error { + scanner := (dst).(TextScanner) + + if src == nil { + return scanner.ScanText(Text{}) + } + + return scanner.ScanText(Text{String: plan.codec.lookupAndCacheString(src), Valid: true}) +} diff --git a/vendor/github.com/jackc/pgx/v5/pgtype/float4.go b/vendor/github.com/jackc/pgx/v5/pgtype/float4.go new file mode 100644 index 0000000000..a43553f672 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgtype/float4.go @@ -0,0 +1,323 @@ +package pgtype + +import ( + "database/sql/driver" + "encoding/binary" + "encoding/json" + "fmt" + "math" + "strconv" + + "github.com/jackc/pgx/v5/internal/pgio" +) + +type Float4 struct { + Float32 float32 + Valid bool +} + +// ScanFloat64 implements the [Float64Scanner] interface. +func (f *Float4) ScanFloat64(n Float8) error { + *f = Float4{Float32: float32(n.Float64), Valid: n.Valid} + return nil +} + +// Float64Value implements the [Float64Valuer] interface. +func (f Float4) Float64Value() (Float8, error) { + return Float8{Float64: float64(f.Float32), Valid: f.Valid}, nil +} + +// ScanInt64 implements the [Int64Scanner] interface. +func (f *Float4) ScanInt64(n Int8) error { + *f = Float4{Float32: float32(n.Int64), Valid: n.Valid} + return nil +} + +// Int64Value implements the [Int64Valuer] interface. +func (f Float4) Int64Value() (Int8, error) { + return Int8{Int64: int64(f.Float32), Valid: f.Valid}, nil +} + +// Scan implements the [database/sql.Scanner] interface. +func (f *Float4) Scan(src any) error { + if src == nil { + *f = Float4{} + return nil + } + + switch src := src.(type) { + case float64: + *f = Float4{Float32: float32(src), Valid: true} + return nil + case string: + n, err := strconv.ParseFloat(src, 32) + if err != nil { + return err + } + *f = Float4{Float32: float32(n), Valid: true} + return nil + } + + return fmt.Errorf("cannot scan %T", src) +} + +// Value implements the [database/sql/driver.Valuer] interface. +func (f Float4) Value() (driver.Value, error) { + if !f.Valid { + return nil, nil + } + return float64(f.Float32), nil +} + +// MarshalJSON implements the [encoding/json.Marshaler] interface. +func (f Float4) MarshalJSON() ([]byte, error) { + if !f.Valid { + return []byte("null"), nil + } + return json.Marshal(f.Float32) +} + +// UnmarshalJSON implements the [encoding/json.Unmarshaler] interface. +func (f *Float4) UnmarshalJSON(b []byte) error { + var n *float32 + err := json.Unmarshal(b, &n) + if err != nil { + return err + } + + if n == nil { + *f = Float4{} + } else { + *f = Float4{Float32: *n, Valid: true} + } + + return nil +} + +type Float4Codec struct{} + +func (Float4Codec) FormatSupported(format int16) bool { + return format == TextFormatCode || format == BinaryFormatCode +} + +func (Float4Codec) PreferredFormat() int16 { + return BinaryFormatCode +} + +func (Float4Codec) PlanEncode(m *Map, oid uint32, format int16, value any) EncodePlan { + switch format { + case BinaryFormatCode: + switch value.(type) { + case float32: + return encodePlanFloat4CodecBinaryFloat32{} + case Float64Valuer: + return encodePlanFloat4CodecBinaryFloat64Valuer{} + case Int64Valuer: + return encodePlanFloat4CodecBinaryInt64Valuer{} + } + case TextFormatCode: + switch value.(type) { + case float32: + return encodePlanTextFloat32{} + case Float64Valuer: + return encodePlanTextFloat64Valuer{} + case Int64Valuer: + return encodePlanTextInt64Valuer{} + } + } + + return nil +} + +type encodePlanFloat4CodecBinaryFloat32 struct{} + +func (encodePlanFloat4CodecBinaryFloat32) Encode(value any, buf []byte) (newBuf []byte, err error) { + n := value.(float32) + return pgio.AppendUint32(buf, math.Float32bits(n)), nil +} + +type encodePlanTextFloat32 struct{} + +func (encodePlanTextFloat32) Encode(value any, buf []byte) (newBuf []byte, err error) { + n := value.(float32) + return append(buf, strconv.FormatFloat(float64(n), 'f', -1, 32)...), nil +} + +type encodePlanFloat4CodecBinaryFloat64Valuer struct{} + +func (encodePlanFloat4CodecBinaryFloat64Valuer) Encode(value any, buf []byte) (newBuf []byte, err error) { + n, err := value.(Float64Valuer).Float64Value() + if err != nil { + return nil, err + } + + if !n.Valid { + return nil, nil + } + + return pgio.AppendUint32(buf, math.Float32bits(float32(n.Float64))), nil +} + +type encodePlanFloat4CodecBinaryInt64Valuer struct{} + +func (encodePlanFloat4CodecBinaryInt64Valuer) Encode(value any, buf []byte) (newBuf []byte, err error) { + n, err := value.(Int64Valuer).Int64Value() + if err != nil { + return nil, err + } + + if !n.Valid { + return nil, nil + } + + f := float32(n.Int64) + return pgio.AppendUint32(buf, math.Float32bits(f)), nil +} + +func (Float4Codec) PlanScan(m *Map, oid uint32, format int16, target any) ScanPlan { + switch format { + case BinaryFormatCode: + switch target.(type) { + case *float32: + return scanPlanBinaryFloat4ToFloat32{} + case Float64Scanner: + return scanPlanBinaryFloat4ToFloat64Scanner{} + case Int64Scanner: + return scanPlanBinaryFloat4ToInt64Scanner{} + case TextScanner: + return scanPlanBinaryFloat4ToTextScanner{} + } + case TextFormatCode: + switch target.(type) { + case *float32: + return scanPlanTextAnyToFloat32{} + case Float64Scanner: + return scanPlanTextAnyToFloat64Scanner{} + case Int64Scanner: + return scanPlanTextAnyToInt64Scanner{} + } + } + + return nil +} + +type scanPlanBinaryFloat4ToFloat32 struct{} + +func (scanPlanBinaryFloat4ToFloat32) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + if len(src) != 4 { + return fmt.Errorf("invalid length for float4: %v", len(src)) + } + + n := int32(binary.BigEndian.Uint32(src)) + f := (dst).(*float32) + *f = math.Float32frombits(uint32(n)) + + return nil +} + +type scanPlanBinaryFloat4ToFloat64Scanner struct{} + +func (scanPlanBinaryFloat4ToFloat64Scanner) Scan(src []byte, dst any) error { + s := (dst).(Float64Scanner) + + if src == nil { + return s.ScanFloat64(Float8{}) + } + + if len(src) != 4 { + return fmt.Errorf("invalid length for float4: %v", len(src)) + } + + n := int32(binary.BigEndian.Uint32(src)) + return s.ScanFloat64(Float8{Float64: float64(math.Float32frombits(uint32(n))), Valid: true}) +} + +type scanPlanBinaryFloat4ToInt64Scanner struct{} + +func (scanPlanBinaryFloat4ToInt64Scanner) Scan(src []byte, dst any) error { + s := (dst).(Int64Scanner) + + if src == nil { + return s.ScanInt64(Int8{}) + } + + if len(src) != 4 { + return fmt.Errorf("invalid length for float4: %v", len(src)) + } + + ui32 := int32(binary.BigEndian.Uint32(src)) + f32 := math.Float32frombits(uint32(ui32)) + i64 := int64(f32) + if f32 != float32(i64) { + return fmt.Errorf("cannot losslessly convert %v to int64", f32) + } + + return s.ScanInt64(Int8{Int64: i64, Valid: true}) +} + +type scanPlanBinaryFloat4ToTextScanner struct{} + +func (scanPlanBinaryFloat4ToTextScanner) Scan(src []byte, dst any) error { + s := (dst).(TextScanner) + + if src == nil { + return s.ScanText(Text{}) + } + + if len(src) != 4 { + return fmt.Errorf("invalid length for float4: %v", len(src)) + } + + ui32 := int32(binary.BigEndian.Uint32(src)) + f32 := math.Float32frombits(uint32(ui32)) + + return s.ScanText(Text{String: strconv.FormatFloat(float64(f32), 'f', -1, 32), Valid: true}) +} + +type scanPlanTextAnyToFloat32 struct{} + +func (scanPlanTextAnyToFloat32) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + n, err := strconv.ParseFloat(string(src), 32) + if err != nil { + return err + } + + f := (dst).(*float32) + *f = float32(n) + + return nil +} + +func (c Float4Codec) DecodeDatabaseSQLValue(m *Map, oid uint32, format int16, src []byte) (driver.Value, error) { + if src == nil { + return nil, nil + } + + var n float32 + err := codecScan(c, m, oid, format, src, &n) + if err != nil { + return nil, err + } + return float64(n), nil +} + +func (c Float4Codec) DecodeValue(m *Map, oid uint32, format int16, src []byte) (any, error) { + if src == nil { + return nil, nil + } + + var n float32 + err := codecScan(c, m, oid, format, src, &n) + if err != nil { + return nil, err + } + return n, nil +} diff --git a/vendor/github.com/jackc/pgx/v5/pgtype/float8.go b/vendor/github.com/jackc/pgx/v5/pgtype/float8.go new file mode 100644 index 0000000000..6234231d7a --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgtype/float8.go @@ -0,0 +1,369 @@ +package pgtype + +import ( + "database/sql/driver" + "encoding/binary" + "encoding/json" + "fmt" + "math" + "strconv" + + "github.com/jackc/pgx/v5/internal/pgio" +) + +type Float64Scanner interface { + ScanFloat64(Float8) error +} + +type Float64Valuer interface { + Float64Value() (Float8, error) +} + +type Float8 struct { + Float64 float64 + Valid bool +} + +// ScanFloat64 implements the [Float64Scanner] interface. +func (f *Float8) ScanFloat64(n Float8) error { + *f = n + return nil +} + +// Float64Value implements the [Float64Valuer] interface. +func (f Float8) Float64Value() (Float8, error) { + return f, nil +} + +// ScanInt64 implements the [Int64Scanner] interface. +func (f *Float8) ScanInt64(n Int8) error { + *f = Float8{Float64: float64(n.Int64), Valid: n.Valid} + return nil +} + +// Int64Value implements the [Int64Valuer] interface. +func (f Float8) Int64Value() (Int8, error) { + return Int8{Int64: int64(f.Float64), Valid: f.Valid}, nil +} + +// Scan implements the [database/sql.Scanner] interface. +func (f *Float8) Scan(src any) error { + if src == nil { + *f = Float8{} + return nil + } + + switch src := src.(type) { + case float64: + *f = Float8{Float64: src, Valid: true} + return nil + case string: + n, err := strconv.ParseFloat(src, 64) + if err != nil { + return err + } + *f = Float8{Float64: n, Valid: true} + return nil + } + + return fmt.Errorf("cannot scan %T", src) +} + +// Value implements the [database/sql/driver.Valuer] interface. +func (f Float8) Value() (driver.Value, error) { + if !f.Valid { + return nil, nil + } + return f.Float64, nil +} + +// MarshalJSON implements the [encoding/json.Marshaler] interface. +func (f Float8) MarshalJSON() ([]byte, error) { + if !f.Valid { + return []byte("null"), nil + } + return json.Marshal(f.Float64) +} + +// UnmarshalJSON implements the [encoding/json.Unmarshaler] interface. +func (f *Float8) UnmarshalJSON(b []byte) error { + var n *float64 + err := json.Unmarshal(b, &n) + if err != nil { + return err + } + + if n == nil { + *f = Float8{} + } else { + *f = Float8{Float64: *n, Valid: true} + } + + return nil +} + +type Float8Codec struct{} + +func (Float8Codec) FormatSupported(format int16) bool { + return format == TextFormatCode || format == BinaryFormatCode +} + +func (Float8Codec) PreferredFormat() int16 { + return BinaryFormatCode +} + +func (Float8Codec) PlanEncode(m *Map, oid uint32, format int16, value any) EncodePlan { + switch format { + case BinaryFormatCode: + switch value.(type) { + case float64: + return encodePlanFloat8CodecBinaryFloat64{} + case Float64Valuer: + return encodePlanFloat8CodecBinaryFloat64Valuer{} + case Int64Valuer: + return encodePlanFloat8CodecBinaryInt64Valuer{} + } + case TextFormatCode: + switch value.(type) { + case float64: + return encodePlanTextFloat64{} + case Float64Valuer: + return encodePlanTextFloat64Valuer{} + case Int64Valuer: + return encodePlanTextInt64Valuer{} + } + } + + return nil +} + +type encodePlanFloat8CodecBinaryFloat64 struct{} + +func (encodePlanFloat8CodecBinaryFloat64) Encode(value any, buf []byte) (newBuf []byte, err error) { + n := value.(float64) + return pgio.AppendUint64(buf, math.Float64bits(n)), nil +} + +type encodePlanTextFloat64 struct{} + +func (encodePlanTextFloat64) Encode(value any, buf []byte) (newBuf []byte, err error) { + n := value.(float64) + return append(buf, strconv.FormatFloat(n, 'f', -1, 64)...), nil +} + +type encodePlanFloat8CodecBinaryFloat64Valuer struct{} + +func (encodePlanFloat8CodecBinaryFloat64Valuer) Encode(value any, buf []byte) (newBuf []byte, err error) { + n, err := value.(Float64Valuer).Float64Value() + if err != nil { + return nil, err + } + + if !n.Valid { + return nil, nil + } + + return pgio.AppendUint64(buf, math.Float64bits(n.Float64)), nil +} + +type encodePlanTextFloat64Valuer struct{} + +func (encodePlanTextFloat64Valuer) Encode(value any, buf []byte) (newBuf []byte, err error) { + n, err := value.(Float64Valuer).Float64Value() + if err != nil { + return nil, err + } + + if !n.Valid { + return nil, nil + } + + return append(buf, strconv.FormatFloat(n.Float64, 'f', -1, 64)...), nil +} + +type encodePlanFloat8CodecBinaryInt64Valuer struct{} + +func (encodePlanFloat8CodecBinaryInt64Valuer) Encode(value any, buf []byte) (newBuf []byte, err error) { + n, err := value.(Int64Valuer).Int64Value() + if err != nil { + return nil, err + } + + if !n.Valid { + return nil, nil + } + + f := float64(n.Int64) + return pgio.AppendUint64(buf, math.Float64bits(f)), nil +} + +type encodePlanTextInt64Valuer struct{} + +func (encodePlanTextInt64Valuer) Encode(value any, buf []byte) (newBuf []byte, err error) { + n, err := value.(Int64Valuer).Int64Value() + if err != nil { + return nil, err + } + + if !n.Valid { + return nil, nil + } + + return append(buf, strconv.FormatInt(n.Int64, 10)...), nil +} + +func (Float8Codec) PlanScan(m *Map, oid uint32, format int16, target any) ScanPlan { + switch format { + case BinaryFormatCode: + switch target.(type) { + case *float64: + return scanPlanBinaryFloat8ToFloat64{} + case Float64Scanner: + return scanPlanBinaryFloat8ToFloat64Scanner{} + case Int64Scanner: + return scanPlanBinaryFloat8ToInt64Scanner{} + case TextScanner: + return scanPlanBinaryFloat8ToTextScanner{} + } + case TextFormatCode: + switch target.(type) { + case *float64: + return scanPlanTextAnyToFloat64{} + case Float64Scanner: + return scanPlanTextAnyToFloat64Scanner{} + case Int64Scanner: + return scanPlanTextAnyToInt64Scanner{} + } + } + + return nil +} + +type scanPlanBinaryFloat8ToFloat64 struct{} + +func (scanPlanBinaryFloat8ToFloat64) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + if len(src) != 8 { + return fmt.Errorf("invalid length for float8: %v", len(src)) + } + + n := int64(binary.BigEndian.Uint64(src)) + f := (dst).(*float64) + *f = math.Float64frombits(uint64(n)) + + return nil +} + +type scanPlanBinaryFloat8ToFloat64Scanner struct{} + +func (scanPlanBinaryFloat8ToFloat64Scanner) Scan(src []byte, dst any) error { + s := (dst).(Float64Scanner) + + if src == nil { + return s.ScanFloat64(Float8{}) + } + + if len(src) != 8 { + return fmt.Errorf("invalid length for float8: %v", len(src)) + } + + n := int64(binary.BigEndian.Uint64(src)) + return s.ScanFloat64(Float8{Float64: math.Float64frombits(uint64(n)), Valid: true}) +} + +type scanPlanBinaryFloat8ToInt64Scanner struct{} + +func (scanPlanBinaryFloat8ToInt64Scanner) Scan(src []byte, dst any) error { + s := (dst).(Int64Scanner) + + if src == nil { + return s.ScanInt64(Int8{}) + } + + if len(src) != 8 { + return fmt.Errorf("invalid length for float8: %v", len(src)) + } + + ui64 := int64(binary.BigEndian.Uint64(src)) + f64 := math.Float64frombits(uint64(ui64)) + i64 := int64(f64) + if f64 != float64(i64) { + return fmt.Errorf("cannot losslessly convert %v to int64", f64) + } + + return s.ScanInt64(Int8{Int64: i64, Valid: true}) +} + +type scanPlanBinaryFloat8ToTextScanner struct{} + +func (scanPlanBinaryFloat8ToTextScanner) Scan(src []byte, dst any) error { + s := (dst).(TextScanner) + + if src == nil { + return s.ScanText(Text{}) + } + + if len(src) != 8 { + return fmt.Errorf("invalid length for float8: %v", len(src)) + } + + ui64 := int64(binary.BigEndian.Uint64(src)) + f64 := math.Float64frombits(uint64(ui64)) + + return s.ScanText(Text{String: strconv.FormatFloat(f64, 'f', -1, 64), Valid: true}) +} + +type scanPlanTextAnyToFloat64 struct{} + +func (scanPlanTextAnyToFloat64) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + n, err := strconv.ParseFloat(string(src), 64) + if err != nil { + return err + } + + f := (dst).(*float64) + *f = n + + return nil +} + +type scanPlanTextAnyToFloat64Scanner struct{} + +func (scanPlanTextAnyToFloat64Scanner) Scan(src []byte, dst any) error { + s := (dst).(Float64Scanner) + + if src == nil { + return s.ScanFloat64(Float8{}) + } + + n, err := strconv.ParseFloat(string(src), 64) + if err != nil { + return err + } + + return s.ScanFloat64(Float8{Float64: n, Valid: true}) +} + +func (c Float8Codec) DecodeDatabaseSQLValue(m *Map, oid uint32, format int16, src []byte) (driver.Value, error) { + return c.DecodeValue(m, oid, format, src) +} + +func (c Float8Codec) DecodeValue(m *Map, oid uint32, format int16, src []byte) (any, error) { + if src == nil { + return nil, nil + } + + var n float64 + err := codecScan(c, m, oid, format, src, &n) + if err != nil { + return nil, err + } + return n, nil +} diff --git a/vendor/github.com/jackc/pgx/v5/pgtype/hstore.go b/vendor/github.com/jackc/pgx/v5/pgtype/hstore.go new file mode 100644 index 0000000000..4a2bb0a5bf --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgtype/hstore.go @@ -0,0 +1,502 @@ +package pgtype + +import ( + "database/sql/driver" + "encoding/binary" + "errors" + "fmt" + "strings" + + "github.com/jackc/pgx/v5/internal/pgio" +) + +type HstoreScanner interface { + ScanHstore(v Hstore) error +} + +type HstoreValuer interface { + HstoreValue() (Hstore, error) +} + +// Hstore represents an hstore column that can be null or have null values +// associated with its keys. +type Hstore map[string]*string + +// ScanHstore implements the [HstoreScanner] interface. +func (h *Hstore) ScanHstore(v Hstore) error { + *h = v + return nil +} + +// HstoreValue implements the [HstoreValuer] interface. +func (h Hstore) HstoreValue() (Hstore, error) { + return h, nil +} + +// Scan implements the [database/sql.Scanner] interface. +func (h *Hstore) Scan(src any) error { + if src == nil { + *h = nil + return nil + } + + if src, ok := src.(string); ok { + return scanPlanTextAnyToHstoreScanner{}.scanString(src, h) + } + + return fmt.Errorf("cannot scan %T", src) +} + +// Value implements the [database/sql/driver.Valuer] interface. +func (h Hstore) Value() (driver.Value, error) { + if h == nil { + return nil, nil + } + + buf, err := HstoreCodec{}.PlanEncode(nil, 0, TextFormatCode, h).Encode(h, nil) + if err != nil { + return nil, err + } + return string(buf), err +} + +type HstoreCodec struct{} + +func (HstoreCodec) FormatSupported(format int16) bool { + return format == TextFormatCode || format == BinaryFormatCode +} + +func (HstoreCodec) PreferredFormat() int16 { + return BinaryFormatCode +} + +func (HstoreCodec) PlanEncode(m *Map, oid uint32, format int16, value any) EncodePlan { + if _, ok := value.(HstoreValuer); !ok { + return nil + } + + switch format { + case BinaryFormatCode: + return encodePlanHstoreCodecBinary{} + case TextFormatCode: + return encodePlanHstoreCodecText{} + } + + return nil +} + +type encodePlanHstoreCodecBinary struct{} + +func (encodePlanHstoreCodecBinary) Encode(value any, buf []byte) (newBuf []byte, err error) { + hstore, err := value.(HstoreValuer).HstoreValue() + if err != nil { + return nil, err + } + + if hstore == nil { + return nil, nil + } + + buf = pgio.AppendInt32(buf, int32(len(hstore))) + + for k, v := range hstore { + buf = pgio.AppendInt32(buf, int32(len(k))) + buf = append(buf, k...) + + if v == nil { + buf = pgio.AppendInt32(buf, -1) + } else { + buf = pgio.AppendInt32(buf, int32(len(*v))) + buf = append(buf, (*v)...) + } + } + + return buf, nil +} + +type encodePlanHstoreCodecText struct{} + +func (encodePlanHstoreCodecText) Encode(value any, buf []byte) (newBuf []byte, err error) { + hstore, err := value.(HstoreValuer).HstoreValue() + if err != nil { + return nil, err + } + + if len(hstore) == 0 { + // distinguish between empty and nil: Not strictly required by Postgres, since its protocol + // explicitly marks NULL column values separately. However, the Binary codec does this, and + // this means we can "round trip" Encode and Scan without data loss. + // nil: []byte(nil); empty: []byte{} + if hstore == nil { + return nil, nil + } + return []byte{}, nil + } + + firstPair := true + + for k, v := range hstore { + if firstPair { + firstPair = false + } else { + buf = append(buf, ',', ' ') + } + + // unconditionally quote hstore keys/values like Postgres does + // this avoids a Mac OS X Postgres hstore parsing bug: + // https://www.postgresql.org/message-id/CA%2BHWA9awUW0%2BRV_gO9r1ABZwGoZxPztcJxPy8vMFSTbTfi4jig%40mail.gmail.com + buf = append(buf, '"') + buf = append(buf, quoteArrayReplacer.Replace(k)...) + buf = append(buf, '"') + buf = append(buf, "=>"...) + + if v == nil { + buf = append(buf, "NULL"...) + } else { + buf = append(buf, '"') + buf = append(buf, quoteArrayReplacer.Replace(*v)...) + buf = append(buf, '"') + } + } + + return buf, nil +} + +func (HstoreCodec) PlanScan(m *Map, oid uint32, format int16, target any) ScanPlan { + switch format { + case BinaryFormatCode: + if _, ok := target.(HstoreScanner); ok { + return scanPlanBinaryHstoreToHstoreScanner{} + } + case TextFormatCode: + if _, ok := target.(HstoreScanner); ok { + return scanPlanTextAnyToHstoreScanner{} + } + } + + return nil +} + +type scanPlanBinaryHstoreToHstoreScanner struct{} + +func (scanPlanBinaryHstoreToHstoreScanner) Scan(src []byte, dst any) error { + scanner := (dst).(HstoreScanner) + + if src == nil { + return scanner.ScanHstore(Hstore(nil)) + } + + rp := 0 + + const uint32Len = 4 + if len(src[rp:]) < uint32Len { + return fmt.Errorf("hstore incomplete %v", src) + } + pairCount := int(int32(binary.BigEndian.Uint32(src[rp:]))) + rp += uint32Len + + if pairCount < 0 { + return fmt.Errorf("hstore invalid pair count: %d", pairCount) + } + // Each pair carries at minimum two int32 length headers (key, value), so pairCount cannot + // exceed the remaining bytes / 8. This bounds the up-front make() against a malicious server + // claiming a huge pair count in a small message. + if maxPairs := len(src[rp:]) / (2 * uint32Len); pairCount > maxPairs { + return fmt.Errorf("hstore invalid pair count %d for %d remaining bytes", pairCount, len(src[rp:])) + } + + hstore := make(Hstore, pairCount) + // one allocation for all *string, rather than one per string, just like text parsing + valueStrings := make([]string, pairCount) + + for i := range pairCount { + if len(src[rp:]) < uint32Len { + return fmt.Errorf("hstore incomplete %v", src) + } + keyLen := int(int32(binary.BigEndian.Uint32(src[rp:]))) + rp += uint32Len + + if keyLen < 0 { + return fmt.Errorf("hstore invalid key length: %d", keyLen) + } + if len(src[rp:]) < keyLen { + return fmt.Errorf("hstore incomplete %v", src) + } + key := string(src[rp : rp+keyLen]) + rp += keyLen + + if len(src[rp:]) < uint32Len { + return fmt.Errorf("hstore incomplete %v", src) + } + valueLen := int(int32(binary.BigEndian.Uint32(src[rp:]))) + rp += 4 + + if valueLen >= 0 { + if len(src[rp:]) < valueLen { + return fmt.Errorf("hstore incomplete %v", src) + } + valueStrings[i] = string(src[rp : rp+valueLen]) + rp += valueLen + + hstore[key] = &valueStrings[i] + } else { + hstore[key] = nil + } + } + + return scanner.ScanHstore(hstore) +} + +type scanPlanTextAnyToHstoreScanner struct{} + +func (s scanPlanTextAnyToHstoreScanner) Scan(src []byte, dst any) error { + scanner := (dst).(HstoreScanner) + + if src == nil { + return scanner.ScanHstore(Hstore(nil)) + } + return s.scanString(string(src), scanner) +} + +// scanString does not return nil hstore values because string cannot be nil. +func (scanPlanTextAnyToHstoreScanner) scanString(src string, scanner HstoreScanner) error { + hstore, err := parseHstore(src) + if err != nil { + return err + } + return scanner.ScanHstore(hstore) +} + +func (c HstoreCodec) DecodeDatabaseSQLValue(m *Map, oid uint32, format int16, src []byte) (driver.Value, error) { + return codecDecodeToTextFormat(c, m, oid, format, src) +} + +func (c HstoreCodec) DecodeValue(m *Map, oid uint32, format int16, src []byte) (any, error) { + if src == nil { + return nil, nil + } + + var hstore Hstore + err := codecScan(c, m, oid, format, src, &hstore) + if err != nil { + return nil, err + } + return hstore, nil +} + +type hstoreParser struct { + str string + pos int + nextBackslash int +} + +func newHSP(in string) *hstoreParser { + return &hstoreParser{ + pos: 0, + str: in, + nextBackslash: strings.IndexByte(in, '\\'), + } +} + +func (p *hstoreParser) atEnd() bool { + return p.pos >= len(p.str) +} + +// consume returns the next byte of the string, or end if the string is done. +func (p *hstoreParser) consume() (b byte, end bool) { + if p.pos >= len(p.str) { + return 0, true + } + b = p.str[p.pos] + p.pos++ + return b, false +} + +func unexpectedByteErr(actualB, expectedB byte) error { + return fmt.Errorf("expected '%c' ('%#v'); found '%c' ('%#v')", expectedB, expectedB, actualB, actualB) +} + +// consumeExpectedByte consumes expectedB from the string, or returns an error. +func (p *hstoreParser) consumeExpectedByte(expectedB byte) error { + nextB, end := p.consume() + if end { + return fmt.Errorf("expected '%c' ('%#v'); found end", expectedB, expectedB) + } + if nextB != expectedB { + return unexpectedByteErr(nextB, expectedB) + } + return nil +} + +// consumeExpected2 consumes two expected bytes or returns an error. +// This was a bit faster than using a string argument (better inlining? Not sure). +func (p *hstoreParser) consumeExpected2(one, two byte) error { + if p.pos+2 > len(p.str) { + return errors.New("unexpected end of string") + } + if p.str[p.pos] != one { + return unexpectedByteErr(p.str[p.pos], one) + } + if p.str[p.pos+1] != two { + return unexpectedByteErr(p.str[p.pos+1], two) + } + p.pos += 2 + return nil +} + +var errEOSInQuoted = errors.New(`found end before closing double-quote ('"')`) + +// consumeDoubleQuoted consumes a double-quoted string from p. The double quote must have been +// parsed already. This copies the string from the backing string so it can be garbage collected. +func (p *hstoreParser) consumeDoubleQuoted() (string, error) { + // fast path: assume most keys/values do not contain escapes + nextDoubleQuote := strings.IndexByte(p.str[p.pos:], '"') + if nextDoubleQuote == -1 { + return "", errEOSInQuoted + } + nextDoubleQuote += p.pos + if p.nextBackslash == -1 || p.nextBackslash > nextDoubleQuote { + // clone the string from the source string to ensure it can be garbage collected separately + // TODO: use strings.Clone on Go 1.20; this could get optimized away + s := strings.Clone(p.str[p.pos:nextDoubleQuote]) + p.pos = nextDoubleQuote + 1 + return s, nil + } + + // slow path: string contains escapes + s, err := p.consumeDoubleQuotedWithEscapes(p.nextBackslash) + p.nextBackslash = strings.IndexByte(p.str[p.pos:], '\\') + if p.nextBackslash != -1 { + p.nextBackslash += p.pos + } + return s, err +} + +// consumeDoubleQuotedWithEscapes consumes a double-quoted string containing escapes, starting +// at p.pos, and with the first backslash at firstBackslash. This copies the string so it can be +// garbage collected separately. +func (p *hstoreParser) consumeDoubleQuotedWithEscapes(firstBackslash int) (string, error) { + // copy the prefix that does not contain backslashes + var builder strings.Builder + builder.WriteString(p.str[p.pos:firstBackslash]) + + // skip to the backslash + p.pos = firstBackslash + + // copy bytes until the end, unescaping backslashes +quotedString: + for { + nextB, end := p.consume() + switch { + case end: + return "", errEOSInQuoted + case nextB == '"': + break quotedString + case nextB == '\\': + // escape: skip the backslash and copy the char + nextB, end = p.consume() + if end { + return "", errEOSInQuoted + } + if !(nextB == '\\' || nextB == '"') { + return "", fmt.Errorf("unexpected escape in quoted string: found '%#v'", nextB) + } + builder.WriteByte(nextB) + default: + // normal byte: copy it + builder.WriteByte(nextB) + } + } + return builder.String(), nil +} + +// consumePairSeparator consumes the Hstore pair separator ", " or returns an error. +func (p *hstoreParser) consumePairSeparator() error { + return p.consumeExpected2(',', ' ') +} + +// consumeKVSeparator consumes the Hstore key/value separator "=>" or returns an error. +func (p *hstoreParser) consumeKVSeparator() error { + return p.consumeExpected2('=', '>') +} + +// consumeDoubleQuotedOrNull consumes the Hstore key/value separator "=>" or returns an error. +func (p *hstoreParser) consumeDoubleQuotedOrNull() (Text, error) { + // peek at the next byte + if p.atEnd() { + return Text{}, errors.New("found end instead of value") + } + next := p.str[p.pos] + if next == 'N' { + // must be the exact string NULL: use consumeExpected2 twice + err := p.consumeExpected2('N', 'U') + if err != nil { + return Text{}, err + } + err = p.consumeExpected2('L', 'L') + if err != nil { + return Text{}, err + } + return Text{String: "", Valid: false}, nil + } else if next != '"' { + return Text{}, unexpectedByteErr(next, '"') + } + + // skip the double quote + p.pos += 1 + s, err := p.consumeDoubleQuoted() + if err != nil { + return Text{}, err + } + return Text{String: s, Valid: true}, nil +} + +func parseHstore(s string) (Hstore, error) { + p := newHSP(s) + + // This is an over-estimate of the number of key/value pairs. Use '>' because I am guessing it + // is less likely to occur in keys/values than '=' or ','. + numPairsEstimate := strings.Count(s, ">") + // makes one allocation of strings for the entire Hstore, rather than one allocation per value. + valueStrings := make([]string, 0, numPairsEstimate) + result := make(Hstore, numPairsEstimate) + first := true + for !p.atEnd() { + if !first { + err := p.consumePairSeparator() + if err != nil { + return nil, err + } + } else { + first = false + } + + err := p.consumeExpectedByte('"') + if err != nil { + return nil, err + } + + key, err := p.consumeDoubleQuoted() + if err != nil { + return nil, err + } + + err = p.consumeKVSeparator() + if err != nil { + return nil, err + } + + value, err := p.consumeDoubleQuotedOrNull() + if err != nil { + return nil, err + } + if value.Valid { + valueStrings = append(valueStrings, value.String) + result[key] = &valueStrings[len(valueStrings)-1] + } else { + result[key] = nil + } + } + + return result, nil +} diff --git a/vendor/github.com/jackc/pgx/v5/pgtype/inet.go b/vendor/github.com/jackc/pgx/v5/pgtype/inet.go new file mode 100644 index 0000000000..2592a5b5b3 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgtype/inet.go @@ -0,0 +1,197 @@ +package pgtype + +import ( + "bytes" + "database/sql/driver" + "errors" + "fmt" + "net/netip" +) + +// Network address family is dependent on server socket.h value for AF_INET. +// In practice, all platforms appear to have the same value. See +// src/include/utils/inet.h for more information. +const ( + defaultAFInet = 2 + defaultAFInet6 = 3 +) + +type NetipPrefixScanner interface { + ScanNetipPrefix(v netip.Prefix) error +} + +type NetipPrefixValuer interface { + NetipPrefixValue() (netip.Prefix, error) +} + +// InetCodec handles both inet and cidr PostgreSQL types. The preferred Go types are [netip.Prefix] and [netip.Addr]. If +// IsValid() is false then they are treated as SQL NULL. +type InetCodec struct{} + +func (InetCodec) FormatSupported(format int16) bool { + return format == TextFormatCode || format == BinaryFormatCode +} + +func (InetCodec) PreferredFormat() int16 { + return BinaryFormatCode +} + +func (InetCodec) PlanEncode(m *Map, oid uint32, format int16, value any) EncodePlan { + if _, ok := value.(NetipPrefixValuer); !ok { + return nil + } + + switch format { + case BinaryFormatCode: + return encodePlanInetCodecBinary{} + case TextFormatCode: + return encodePlanInetCodecText{} + } + + return nil +} + +type encodePlanInetCodecBinary struct{} + +func (encodePlanInetCodecBinary) Encode(value any, buf []byte) (newBuf []byte, err error) { + prefix, err := value.(NetipPrefixValuer).NetipPrefixValue() + if err != nil { + return nil, err + } + + if !prefix.IsValid() { + return nil, nil + } + + var family byte + if prefix.Addr().Is4() { + family = defaultAFInet + } else { + family = defaultAFInet6 + } + + buf = append(buf, family) + + ones := prefix.Bits() + buf = append(buf, byte(ones)) + + // is_cidr is ignored on server + buf = append(buf, 0) + + if family == defaultAFInet { + buf = append(buf, byte(4)) + b := prefix.Addr().As4() + buf = append(buf, b[:]...) + } else { + buf = append(buf, byte(16)) + b := prefix.Addr().As16() + buf = append(buf, b[:]...) + } + + return buf, nil +} + +type encodePlanInetCodecText struct{} + +func (encodePlanInetCodecText) Encode(value any, buf []byte) (newBuf []byte, err error) { + prefix, err := value.(NetipPrefixValuer).NetipPrefixValue() + if err != nil { + return nil, err + } + + if !prefix.IsValid() { + return nil, nil + } + + return append(buf, prefix.String()...), nil +} + +func (InetCodec) PlanScan(m *Map, oid uint32, format int16, target any) ScanPlan { + switch format { + case BinaryFormatCode: + if _, ok := target.(NetipPrefixScanner); ok { + return scanPlanBinaryInetToNetipPrefixScanner{} + } + case TextFormatCode: + if _, ok := target.(NetipPrefixScanner); ok { + return scanPlanTextAnyToNetipPrefixScanner{} + } + } + + return nil +} + +func (c InetCodec) DecodeDatabaseSQLValue(m *Map, oid uint32, format int16, src []byte) (driver.Value, error) { + return codecDecodeToTextFormat(c, m, oid, format, src) +} + +func (c InetCodec) DecodeValue(m *Map, oid uint32, format int16, src []byte) (any, error) { + if src == nil { + return nil, nil + } + + var prefix netip.Prefix + err := codecScan(c, m, oid, format, src, (*netipPrefixWrapper)(&prefix)) + if err != nil { + return nil, err + } + + if !prefix.IsValid() { + return nil, nil + } + + return prefix, nil +} + +type scanPlanBinaryInetToNetipPrefixScanner struct{} + +func (scanPlanBinaryInetToNetipPrefixScanner) Scan(src []byte, dst any) error { + scanner := (dst).(NetipPrefixScanner) + + if src == nil { + return scanner.ScanNetipPrefix(netip.Prefix{}) + } + + if len(src) != 8 && len(src) != 20 { + return fmt.Errorf("Received an invalid size for an inet: %d", len(src)) + } + + // ignore family + bits := src[1] + // ignore is_cidr + // ignore addressLength - implicit in length of message + + addr, ok := netip.AddrFromSlice(src[4:]) + if !ok { + return errors.New("netip.AddrFromSlice failed") + } + + return scanner.ScanNetipPrefix(netip.PrefixFrom(addr, int(bits))) +} + +type scanPlanTextAnyToNetipPrefixScanner struct{} + +func (scanPlanTextAnyToNetipPrefixScanner) Scan(src []byte, dst any) error { + scanner := (dst).(NetipPrefixScanner) + + if src == nil { + return scanner.ScanNetipPrefix(netip.Prefix{}) + } + + var prefix netip.Prefix + if bytes.IndexByte(src, '/') == -1 { + addr, err := netip.ParseAddr(string(src)) + if err != nil { + return err + } + prefix = netip.PrefixFrom(addr, addr.BitLen()) + } else { + var err error + prefix, err = netip.ParsePrefix(string(src)) + if err != nil { + return err + } + } + + return scanner.ScanNetipPrefix(prefix) +} diff --git a/vendor/github.com/jackc/pgx/v5/pgtype/int.go b/vendor/github.com/jackc/pgx/v5/pgtype/int.go new file mode 100644 index 0000000000..95032e5a27 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgtype/int.go @@ -0,0 +1,1990 @@ +// Code generated from pgtype/int.go.erb. DO NOT EDIT. + +package pgtype + +import ( + "database/sql/driver" + "encoding/binary" + "encoding/json" + "fmt" + "math" + "strconv" + + "github.com/jackc/pgx/v5/internal/pgio" +) + +type Int64Scanner interface { + ScanInt64(Int8) error +} + +type Int64Valuer interface { + Int64Value() (Int8, error) +} + +type Int2 struct { + Int16 int16 + Valid bool +} + +// ScanInt64 implements the [Int64Scanner] interface. +func (dst *Int2) ScanInt64(n Int8) error { + if !n.Valid { + *dst = Int2{} + return nil + } + + if n.Int64 < math.MinInt16 { + return fmt.Errorf("%d is less than minimum value for Int2", n.Int64) + } + if n.Int64 > math.MaxInt16 { + return fmt.Errorf("%d is greater than maximum value for Int2", n.Int64) + } + *dst = Int2{Int16: int16(n.Int64), Valid: true} + + return nil +} + +// Int64Value implements the [Int64Valuer] interface. +func (n Int2) Int64Value() (Int8, error) { + return Int8{Int64: int64(n.Int16), Valid: n.Valid}, nil +} + +// Scan implements the [database/sql.Scanner] interface. +func (dst *Int2) Scan(src any) error { + if src == nil { + *dst = Int2{} + return nil + } + + var n int64 + + switch src := src.(type) { + case int64: + n = src + case string: + var err error + n, err = strconv.ParseInt(src, 10, 16) + if err != nil { + return err + } + case []byte: + var err error + n, err = strconv.ParseInt(string(src), 10, 16) + if err != nil { + return err + } + default: + return fmt.Errorf("cannot scan %T", src) + } + + if n < math.MinInt16 { + return fmt.Errorf("%d is less than minimum value for Int2", n) + } + if n > math.MaxInt16 { + return fmt.Errorf("%d is greater than maximum value for Int2", n) + } + *dst = Int2{Int16: int16(n), Valid: true} + + return nil +} + +// Value implements the [database/sql/driver.Valuer] interface. +func (src Int2) Value() (driver.Value, error) { + if !src.Valid { + return nil, nil + } + return int64(src.Int16), nil +} + +// MarshalJSON implements the [encoding/json.Marshaler] interface. +func (src Int2) MarshalJSON() ([]byte, error) { + if !src.Valid { + return []byte("null"), nil + } + return []byte(strconv.FormatInt(int64(src.Int16), 10)), nil +} + +// UnmarshalJSON implements the [encoding/json.Unmarshaler] interface. +func (dst *Int2) UnmarshalJSON(b []byte) error { + var n *int16 + err := json.Unmarshal(b, &n) + if err != nil { + return err + } + + if n == nil { + *dst = Int2{} + } else { + *dst = Int2{Int16: *n, Valid: true} + } + + return nil +} + +type Int2Codec struct{} + +func (Int2Codec) FormatSupported(format int16) bool { + return format == TextFormatCode || format == BinaryFormatCode +} + +func (Int2Codec) PreferredFormat() int16 { + return BinaryFormatCode +} + +func (Int2Codec) PlanEncode(m *Map, oid uint32, format int16, value any) EncodePlan { + switch format { + case BinaryFormatCode: + switch value.(type) { + case int16: + return encodePlanInt2CodecBinaryInt16{} + case Int64Valuer: + return encodePlanInt2CodecBinaryInt64Valuer{} + } + case TextFormatCode: + switch value.(type) { + case int16: + return encodePlanInt2CodecTextInt16{} + case Int64Valuer: + return encodePlanInt2CodecTextInt64Valuer{} + } + } + + return nil +} + +type encodePlanInt2CodecBinaryInt16 struct{} + +func (encodePlanInt2CodecBinaryInt16) Encode(value any, buf []byte) (newBuf []byte, err error) { + n := value.(int16) + return pgio.AppendInt16(buf, int16(n)), nil +} + +type encodePlanInt2CodecTextInt16 struct{} + +func (encodePlanInt2CodecTextInt16) Encode(value any, buf []byte) (newBuf []byte, err error) { + n := value.(int16) + return append(buf, strconv.FormatInt(int64(n), 10)...), nil +} + +type encodePlanInt2CodecBinaryInt64Valuer struct{} + +func (encodePlanInt2CodecBinaryInt64Valuer) Encode(value any, buf []byte) (newBuf []byte, err error) { + n, err := value.(Int64Valuer).Int64Value() + if err != nil { + return nil, err + } + + if !n.Valid { + return nil, nil + } + + if n.Int64 > math.MaxInt16 { + return nil, fmt.Errorf("%d is greater than maximum value for int2", n.Int64) + } + if n.Int64 < math.MinInt16 { + return nil, fmt.Errorf("%d is less than minimum value for int2", n.Int64) + } + + return pgio.AppendInt16(buf, int16(n.Int64)), nil +} + +type encodePlanInt2CodecTextInt64Valuer struct{} + +func (encodePlanInt2CodecTextInt64Valuer) Encode(value any, buf []byte) (newBuf []byte, err error) { + n, err := value.(Int64Valuer).Int64Value() + if err != nil { + return nil, err + } + + if !n.Valid { + return nil, nil + } + + if n.Int64 > math.MaxInt16 { + return nil, fmt.Errorf("%d is greater than maximum value for int2", n.Int64) + } + if n.Int64 < math.MinInt16 { + return nil, fmt.Errorf("%d is less than minimum value for int2", n.Int64) + } + + return append(buf, strconv.FormatInt(n.Int64, 10)...), nil +} + +func (Int2Codec) PlanScan(m *Map, oid uint32, format int16, target any) ScanPlan { + + switch format { + case BinaryFormatCode: + switch target.(type) { + case *int8: + return scanPlanBinaryInt2ToInt8{} + case *int16: + return scanPlanBinaryInt2ToInt16{} + case *int32: + return scanPlanBinaryInt2ToInt32{} + case *int64: + return scanPlanBinaryInt2ToInt64{} + case *int: + return scanPlanBinaryInt2ToInt{} + case *uint8: + return scanPlanBinaryInt2ToUint8{} + case *uint16: + return scanPlanBinaryInt2ToUint16{} + case *uint32: + return scanPlanBinaryInt2ToUint32{} + case *uint64: + return scanPlanBinaryInt2ToUint64{} + case *uint: + return scanPlanBinaryInt2ToUint{} + case Int64Scanner: + return scanPlanBinaryInt2ToInt64Scanner{} + case TextScanner: + return scanPlanBinaryInt2ToTextScanner{} + } + case TextFormatCode: + switch target.(type) { + case *int8: + return scanPlanTextAnyToInt8{} + case *int16: + return scanPlanTextAnyToInt16{} + case *int32: + return scanPlanTextAnyToInt32{} + case *int64: + return scanPlanTextAnyToInt64{} + case *int: + return scanPlanTextAnyToInt{} + case *uint8: + return scanPlanTextAnyToUint8{} + case *uint16: + return scanPlanTextAnyToUint16{} + case *uint32: + return scanPlanTextAnyToUint32{} + case *uint64: + return scanPlanTextAnyToUint64{} + case *uint: + return scanPlanTextAnyToUint{} + case Int64Scanner: + return scanPlanTextAnyToInt64Scanner{} + } + } + + return nil +} + +func (c Int2Codec) DecodeDatabaseSQLValue(m *Map, oid uint32, format int16, src []byte) (driver.Value, error) { + if src == nil { + return nil, nil + } + + var n int64 + err := codecScan(c, m, oid, format, src, &n) + if err != nil { + return nil, err + } + return n, nil +} + +func (c Int2Codec) DecodeValue(m *Map, oid uint32, format int16, src []byte) (any, error) { + if src == nil { + return nil, nil + } + + var n int16 + err := codecScan(c, m, oid, format, src, &n) + if err != nil { + return nil, err + } + return n, nil +} + +type scanPlanBinaryInt2ToInt8 struct{} + +func (scanPlanBinaryInt2ToInt8) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + if len(src) != 2 { + return fmt.Errorf("invalid length for int2: %v", len(src)) + } + + p, ok := (dst).(*int8) + if !ok { + return ErrScanTargetTypeChanged + } + + n := int16(binary.BigEndian.Uint16(src)) + if n < math.MinInt8 { + return fmt.Errorf("%d is less than minimum value for int8", n) + } else if n > math.MaxInt8 { + return fmt.Errorf("%d is greater than maximum value for int8", n) + } + + *p = int8(n) + + return nil +} + +type scanPlanBinaryInt2ToUint8 struct{} + +func (scanPlanBinaryInt2ToUint8) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + if len(src) != 2 { + return fmt.Errorf("invalid length for uint2: %v", len(src)) + } + + p, ok := (dst).(*uint8) + if !ok { + return ErrScanTargetTypeChanged + } + + n := int16(binary.BigEndian.Uint16(src)) + if n < 0 { + return fmt.Errorf("%d is less than minimum value for uint8", n) + } + + if n > math.MaxUint8 { + return fmt.Errorf("%d is greater than maximum value for uint8", n) + } + + *p = uint8(n) + + return nil +} + +type scanPlanBinaryInt2ToInt16 struct{} + +func (scanPlanBinaryInt2ToInt16) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + if len(src) != 2 { + return fmt.Errorf("invalid length for int2: %v", len(src)) + } + + p, ok := (dst).(*int16) + if !ok { + return ErrScanTargetTypeChanged + } + + *p = int16(binary.BigEndian.Uint16(src)) + + return nil +} + +type scanPlanBinaryInt2ToUint16 struct{} + +func (scanPlanBinaryInt2ToUint16) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + if len(src) != 2 { + return fmt.Errorf("invalid length for uint2: %v", len(src)) + } + + p, ok := (dst).(*uint16) + if !ok { + return ErrScanTargetTypeChanged + } + + n := int16(binary.BigEndian.Uint16(src)) + if n < 0 { + return fmt.Errorf("%d is less than minimum value for uint16", n) + } + + *p = uint16(n) + + return nil +} + +type scanPlanBinaryInt2ToInt32 struct{} + +func (scanPlanBinaryInt2ToInt32) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + if len(src) != 2 { + return fmt.Errorf("invalid length for int2: %v", len(src)) + } + + p, ok := (dst).(*int32) + if !ok { + return ErrScanTargetTypeChanged + } + + *p = int32(int16(binary.BigEndian.Uint16(src))) + + return nil +} + +type scanPlanBinaryInt2ToUint32 struct{} + +func (scanPlanBinaryInt2ToUint32) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + if len(src) != 2 { + return fmt.Errorf("invalid length for uint2: %v", len(src)) + } + + p, ok := (dst).(*uint32) + if !ok { + return ErrScanTargetTypeChanged + } + + n := int16(binary.BigEndian.Uint16(src)) + if n < 0 { + return fmt.Errorf("%d is less than minimum value for uint32", n) + } + + *p = uint32(n) + + return nil +} + +type scanPlanBinaryInt2ToInt64 struct{} + +func (scanPlanBinaryInt2ToInt64) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + if len(src) != 2 { + return fmt.Errorf("invalid length for int2: %v", len(src)) + } + + p, ok := (dst).(*int64) + if !ok { + return ErrScanTargetTypeChanged + } + + *p = int64(int16(binary.BigEndian.Uint16(src))) + + return nil +} + +type scanPlanBinaryInt2ToUint64 struct{} + +func (scanPlanBinaryInt2ToUint64) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + if len(src) != 2 { + return fmt.Errorf("invalid length for uint2: %v", len(src)) + } + + p, ok := (dst).(*uint64) + if !ok { + return ErrScanTargetTypeChanged + } + + n := int16(binary.BigEndian.Uint16(src)) + if n < 0 { + return fmt.Errorf("%d is less than minimum value for uint64", n) + } + + *p = uint64(n) + + return nil +} + +type scanPlanBinaryInt2ToInt struct{} + +func (scanPlanBinaryInt2ToInt) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + if len(src) != 2 { + return fmt.Errorf("invalid length for int2: %v", len(src)) + } + + p, ok := (dst).(*int) + if !ok { + return ErrScanTargetTypeChanged + } + + *p = int(int16(binary.BigEndian.Uint16(src))) + + return nil +} + +type scanPlanBinaryInt2ToUint struct{} + +func (scanPlanBinaryInt2ToUint) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + if len(src) != 2 { + return fmt.Errorf("invalid length for uint2: %v", len(src)) + } + + p, ok := (dst).(*uint) + if !ok { + return ErrScanTargetTypeChanged + } + + n := int64(int16(binary.BigEndian.Uint16(src))) + if n < 0 { + return fmt.Errorf("%d is less than minimum value for uint", n) + } + + *p = uint(n) + + return nil +} + +type scanPlanBinaryInt2ToInt64Scanner struct{} + +func (scanPlanBinaryInt2ToInt64Scanner) Scan(src []byte, dst any) error { + s, ok := (dst).(Int64Scanner) + if !ok { + return ErrScanTargetTypeChanged + } + + if src == nil { + return s.ScanInt64(Int8{}) + } + + if len(src) != 2 { + return fmt.Errorf("invalid length for int2: %v", len(src)) + } + + n := int64(int16(binary.BigEndian.Uint16(src))) + + return s.ScanInt64(Int8{Int64: n, Valid: true}) +} + +type scanPlanBinaryInt2ToTextScanner struct{} + +func (scanPlanBinaryInt2ToTextScanner) Scan(src []byte, dst any) error { + s, ok := (dst).(TextScanner) + if !ok { + return ErrScanTargetTypeChanged + } + + if src == nil { + return s.ScanText(Text{}) + } + + if len(src) != 2 { + return fmt.Errorf("invalid length for int2: %v", len(src)) + } + + n := int64(int16(binary.BigEndian.Uint16(src))) + + return s.ScanText(Text{String: strconv.FormatInt(n, 10), Valid: true}) +} + +type Int4 struct { + Int32 int32 + Valid bool +} + +// ScanInt64 implements the [Int64Scanner] interface. +func (dst *Int4) ScanInt64(n Int8) error { + if !n.Valid { + *dst = Int4{} + return nil + } + + if n.Int64 < math.MinInt32 { + return fmt.Errorf("%d is less than minimum value for Int4", n.Int64) + } + if n.Int64 > math.MaxInt32 { + return fmt.Errorf("%d is greater than maximum value for Int4", n.Int64) + } + *dst = Int4{Int32: int32(n.Int64), Valid: true} + + return nil +} + +// Int64Value implements the [Int64Valuer] interface. +func (n Int4) Int64Value() (Int8, error) { + return Int8{Int64: int64(n.Int32), Valid: n.Valid}, nil +} + +// Scan implements the [database/sql.Scanner] interface. +func (dst *Int4) Scan(src any) error { + if src == nil { + *dst = Int4{} + return nil + } + + var n int64 + + switch src := src.(type) { + case int64: + n = src + case string: + var err error + n, err = strconv.ParseInt(src, 10, 32) + if err != nil { + return err + } + case []byte: + var err error + n, err = strconv.ParseInt(string(src), 10, 32) + if err != nil { + return err + } + default: + return fmt.Errorf("cannot scan %T", src) + } + + if n < math.MinInt32 { + return fmt.Errorf("%d is less than minimum value for Int4", n) + } + if n > math.MaxInt32 { + return fmt.Errorf("%d is greater than maximum value for Int4", n) + } + *dst = Int4{Int32: int32(n), Valid: true} + + return nil +} + +// Value implements the [database/sql/driver.Valuer] interface. +func (src Int4) Value() (driver.Value, error) { + if !src.Valid { + return nil, nil + } + return int64(src.Int32), nil +} + +// MarshalJSON implements the [encoding/json.Marshaler] interface. +func (src Int4) MarshalJSON() ([]byte, error) { + if !src.Valid { + return []byte("null"), nil + } + return []byte(strconv.FormatInt(int64(src.Int32), 10)), nil +} + +// UnmarshalJSON implements the [encoding/json.Unmarshaler] interface. +func (dst *Int4) UnmarshalJSON(b []byte) error { + var n *int32 + err := json.Unmarshal(b, &n) + if err != nil { + return err + } + + if n == nil { + *dst = Int4{} + } else { + *dst = Int4{Int32: *n, Valid: true} + } + + return nil +} + +type Int4Codec struct{} + +func (Int4Codec) FormatSupported(format int16) bool { + return format == TextFormatCode || format == BinaryFormatCode +} + +func (Int4Codec) PreferredFormat() int16 { + return BinaryFormatCode +} + +func (Int4Codec) PlanEncode(m *Map, oid uint32, format int16, value any) EncodePlan { + switch format { + case BinaryFormatCode: + switch value.(type) { + case int32: + return encodePlanInt4CodecBinaryInt32{} + case Int64Valuer: + return encodePlanInt4CodecBinaryInt64Valuer{} + } + case TextFormatCode: + switch value.(type) { + case int32: + return encodePlanInt4CodecTextInt32{} + case Int64Valuer: + return encodePlanInt4CodecTextInt64Valuer{} + } + } + + return nil +} + +type encodePlanInt4CodecBinaryInt32 struct{} + +func (encodePlanInt4CodecBinaryInt32) Encode(value any, buf []byte) (newBuf []byte, err error) { + n := value.(int32) + return pgio.AppendInt32(buf, int32(n)), nil +} + +type encodePlanInt4CodecTextInt32 struct{} + +func (encodePlanInt4CodecTextInt32) Encode(value any, buf []byte) (newBuf []byte, err error) { + n := value.(int32) + return append(buf, strconv.FormatInt(int64(n), 10)...), nil +} + +type encodePlanInt4CodecBinaryInt64Valuer struct{} + +func (encodePlanInt4CodecBinaryInt64Valuer) Encode(value any, buf []byte) (newBuf []byte, err error) { + n, err := value.(Int64Valuer).Int64Value() + if err != nil { + return nil, err + } + + if !n.Valid { + return nil, nil + } + + if n.Int64 > math.MaxInt32 { + return nil, fmt.Errorf("%d is greater than maximum value for int4", n.Int64) + } + if n.Int64 < math.MinInt32 { + return nil, fmt.Errorf("%d is less than minimum value for int4", n.Int64) + } + + return pgio.AppendInt32(buf, int32(n.Int64)), nil +} + +type encodePlanInt4CodecTextInt64Valuer struct{} + +func (encodePlanInt4CodecTextInt64Valuer) Encode(value any, buf []byte) (newBuf []byte, err error) { + n, err := value.(Int64Valuer).Int64Value() + if err != nil { + return nil, err + } + + if !n.Valid { + return nil, nil + } + + if n.Int64 > math.MaxInt32 { + return nil, fmt.Errorf("%d is greater than maximum value for int4", n.Int64) + } + if n.Int64 < math.MinInt32 { + return nil, fmt.Errorf("%d is less than minimum value for int4", n.Int64) + } + + return append(buf, strconv.FormatInt(n.Int64, 10)...), nil +} + +func (Int4Codec) PlanScan(m *Map, oid uint32, format int16, target any) ScanPlan { + + switch format { + case BinaryFormatCode: + switch target.(type) { + case *int8: + return scanPlanBinaryInt4ToInt8{} + case *int16: + return scanPlanBinaryInt4ToInt16{} + case *int32: + return scanPlanBinaryInt4ToInt32{} + case *int64: + return scanPlanBinaryInt4ToInt64{} + case *int: + return scanPlanBinaryInt4ToInt{} + case *uint8: + return scanPlanBinaryInt4ToUint8{} + case *uint16: + return scanPlanBinaryInt4ToUint16{} + case *uint32: + return scanPlanBinaryInt4ToUint32{} + case *uint64: + return scanPlanBinaryInt4ToUint64{} + case *uint: + return scanPlanBinaryInt4ToUint{} + case Int64Scanner: + return scanPlanBinaryInt4ToInt64Scanner{} + case TextScanner: + return scanPlanBinaryInt4ToTextScanner{} + } + case TextFormatCode: + switch target.(type) { + case *int8: + return scanPlanTextAnyToInt8{} + case *int16: + return scanPlanTextAnyToInt16{} + case *int32: + return scanPlanTextAnyToInt32{} + case *int64: + return scanPlanTextAnyToInt64{} + case *int: + return scanPlanTextAnyToInt{} + case *uint8: + return scanPlanTextAnyToUint8{} + case *uint16: + return scanPlanTextAnyToUint16{} + case *uint32: + return scanPlanTextAnyToUint32{} + case *uint64: + return scanPlanTextAnyToUint64{} + case *uint: + return scanPlanTextAnyToUint{} + case Int64Scanner: + return scanPlanTextAnyToInt64Scanner{} + } + } + + return nil +} + +func (c Int4Codec) DecodeDatabaseSQLValue(m *Map, oid uint32, format int16, src []byte) (driver.Value, error) { + if src == nil { + return nil, nil + } + + var n int64 + err := codecScan(c, m, oid, format, src, &n) + if err != nil { + return nil, err + } + return n, nil +} + +func (c Int4Codec) DecodeValue(m *Map, oid uint32, format int16, src []byte) (any, error) { + if src == nil { + return nil, nil + } + + var n int32 + err := codecScan(c, m, oid, format, src, &n) + if err != nil { + return nil, err + } + return n, nil +} + +type scanPlanBinaryInt4ToInt8 struct{} + +func (scanPlanBinaryInt4ToInt8) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + if len(src) != 4 { + return fmt.Errorf("invalid length for int4: %v", len(src)) + } + + p, ok := (dst).(*int8) + if !ok { + return ErrScanTargetTypeChanged + } + + n := int32(binary.BigEndian.Uint32(src)) + if n < math.MinInt8 { + return fmt.Errorf("%d is less than minimum value for int8", n) + } else if n > math.MaxInt8 { + return fmt.Errorf("%d is greater than maximum value for int8", n) + } + + *p = int8(n) + + return nil +} + +type scanPlanBinaryInt4ToUint8 struct{} + +func (scanPlanBinaryInt4ToUint8) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + if len(src) != 4 { + return fmt.Errorf("invalid length for uint4: %v", len(src)) + } + + p, ok := (dst).(*uint8) + if !ok { + return ErrScanTargetTypeChanged + } + + n := int32(binary.BigEndian.Uint32(src)) + if n < 0 { + return fmt.Errorf("%d is less than minimum value for uint8", n) + } + + if n > math.MaxUint8 { + return fmt.Errorf("%d is greater than maximum value for uint8", n) + } + + *p = uint8(n) + + return nil +} + +type scanPlanBinaryInt4ToInt16 struct{} + +func (scanPlanBinaryInt4ToInt16) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + if len(src) != 4 { + return fmt.Errorf("invalid length for int4: %v", len(src)) + } + + p, ok := (dst).(*int16) + if !ok { + return ErrScanTargetTypeChanged + } + + n := int32(binary.BigEndian.Uint32(src)) + if n < math.MinInt16 { + return fmt.Errorf("%d is less than minimum value for int16", n) + } else if n > math.MaxInt16 { + return fmt.Errorf("%d is greater than maximum value for int16", n) + } + + *p = int16(n) + + return nil +} + +type scanPlanBinaryInt4ToUint16 struct{} + +func (scanPlanBinaryInt4ToUint16) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + if len(src) != 4 { + return fmt.Errorf("invalid length for uint4: %v", len(src)) + } + + p, ok := (dst).(*uint16) + if !ok { + return ErrScanTargetTypeChanged + } + + n := int32(binary.BigEndian.Uint32(src)) + if n < 0 { + return fmt.Errorf("%d is less than minimum value for uint16", n) + } + + if n > math.MaxUint16 { + return fmt.Errorf("%d is greater than maximum value for uint16", n) + } + + *p = uint16(n) + + return nil +} + +type scanPlanBinaryInt4ToInt32 struct{} + +func (scanPlanBinaryInt4ToInt32) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + if len(src) != 4 { + return fmt.Errorf("invalid length for int4: %v", len(src)) + } + + p, ok := (dst).(*int32) + if !ok { + return ErrScanTargetTypeChanged + } + + *p = int32(binary.BigEndian.Uint32(src)) + + return nil +} + +type scanPlanBinaryInt4ToUint32 struct{} + +func (scanPlanBinaryInt4ToUint32) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + if len(src) != 4 { + return fmt.Errorf("invalid length for uint4: %v", len(src)) + } + + p, ok := (dst).(*uint32) + if !ok { + return ErrScanTargetTypeChanged + } + + n := int32(binary.BigEndian.Uint32(src)) + if n < 0 { + return fmt.Errorf("%d is less than minimum value for uint32", n) + } + + *p = uint32(n) + + return nil +} + +type scanPlanBinaryInt4ToInt64 struct{} + +func (scanPlanBinaryInt4ToInt64) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + if len(src) != 4 { + return fmt.Errorf("invalid length for int4: %v", len(src)) + } + + p, ok := (dst).(*int64) + if !ok { + return ErrScanTargetTypeChanged + } + + *p = int64(int32(binary.BigEndian.Uint32(src))) + + return nil +} + +type scanPlanBinaryInt4ToUint64 struct{} + +func (scanPlanBinaryInt4ToUint64) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + if len(src) != 4 { + return fmt.Errorf("invalid length for uint4: %v", len(src)) + } + + p, ok := (dst).(*uint64) + if !ok { + return ErrScanTargetTypeChanged + } + + n := int32(binary.BigEndian.Uint32(src)) + if n < 0 { + return fmt.Errorf("%d is less than minimum value for uint64", n) + } + + *p = uint64(n) + + return nil +} + +type scanPlanBinaryInt4ToInt struct{} + +func (scanPlanBinaryInt4ToInt) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + if len(src) != 4 { + return fmt.Errorf("invalid length for int4: %v", len(src)) + } + + p, ok := (dst).(*int) + if !ok { + return ErrScanTargetTypeChanged + } + + *p = int(int32(binary.BigEndian.Uint32(src))) + + return nil +} + +type scanPlanBinaryInt4ToUint struct{} + +func (scanPlanBinaryInt4ToUint) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + if len(src) != 4 { + return fmt.Errorf("invalid length for uint4: %v", len(src)) + } + + p, ok := (dst).(*uint) + if !ok { + return ErrScanTargetTypeChanged + } + + n := int64(int32(binary.BigEndian.Uint32(src))) + if n < 0 { + return fmt.Errorf("%d is less than minimum value for uint", n) + } + + *p = uint(n) + + return nil +} + +type scanPlanBinaryInt4ToInt64Scanner struct{} + +func (scanPlanBinaryInt4ToInt64Scanner) Scan(src []byte, dst any) error { + s, ok := (dst).(Int64Scanner) + if !ok { + return ErrScanTargetTypeChanged + } + + if src == nil { + return s.ScanInt64(Int8{}) + } + + if len(src) != 4 { + return fmt.Errorf("invalid length for int4: %v", len(src)) + } + + n := int64(int32(binary.BigEndian.Uint32(src))) + + return s.ScanInt64(Int8{Int64: n, Valid: true}) +} + +type scanPlanBinaryInt4ToTextScanner struct{} + +func (scanPlanBinaryInt4ToTextScanner) Scan(src []byte, dst any) error { + s, ok := (dst).(TextScanner) + if !ok { + return ErrScanTargetTypeChanged + } + + if src == nil { + return s.ScanText(Text{}) + } + + if len(src) != 4 { + return fmt.Errorf("invalid length for int4: %v", len(src)) + } + + n := int64(int32(binary.BigEndian.Uint32(src))) + + return s.ScanText(Text{String: strconv.FormatInt(n, 10), Valid: true}) +} + +type Int8 struct { + Int64 int64 + Valid bool +} + +// ScanInt64 implements the [Int64Scanner] interface. +func (dst *Int8) ScanInt64(n Int8) error { + if !n.Valid { + *dst = Int8{} + return nil + } + + if n.Int64 < math.MinInt64 { + return fmt.Errorf("%d is less than minimum value for Int8", n.Int64) + } + if n.Int64 > math.MaxInt64 { + return fmt.Errorf("%d is greater than maximum value for Int8", n.Int64) + } + *dst = Int8{Int64: int64(n.Int64), Valid: true} + + return nil +} + +// Int64Value implements the [Int64Valuer] interface. +func (n Int8) Int64Value() (Int8, error) { + return Int8{Int64: int64(n.Int64), Valid: n.Valid}, nil +} + +// Scan implements the [database/sql.Scanner] interface. +func (dst *Int8) Scan(src any) error { + if src == nil { + *dst = Int8{} + return nil + } + + var n int64 + + switch src := src.(type) { + case int64: + n = src + case string: + var err error + n, err = strconv.ParseInt(src, 10, 64) + if err != nil { + return err + } + case []byte: + var err error + n, err = strconv.ParseInt(string(src), 10, 64) + if err != nil { + return err + } + default: + return fmt.Errorf("cannot scan %T", src) + } + + if n < math.MinInt64 { + return fmt.Errorf("%d is greater than maximum value for Int8", n) + } + if n > math.MaxInt64 { + return fmt.Errorf("%d is greater than maximum value for Int8", n) + } + *dst = Int8{Int64: int64(n), Valid: true} + + return nil +} + +// Value implements the [database/sql/driver.Valuer] interface. +func (src Int8) Value() (driver.Value, error) { + if !src.Valid { + return nil, nil + } + return int64(src.Int64), nil +} + +// MarshalJSON implements the [encoding/json.Marshaler] interface. +func (src Int8) MarshalJSON() ([]byte, error) { + if !src.Valid { + return []byte("null"), nil + } + return []byte(strconv.FormatInt(int64(src.Int64), 10)), nil +} + +// UnmarshalJSON implements the [encoding/json.Unmarshaler] interface. +func (dst *Int8) UnmarshalJSON(b []byte) error { + var n *int64 + err := json.Unmarshal(b, &n) + if err != nil { + return err + } + + if n == nil { + *dst = Int8{} + } else { + *dst = Int8{Int64: *n, Valid: true} + } + + return nil +} + +type Int8Codec struct{} + +func (Int8Codec) FormatSupported(format int16) bool { + return format == TextFormatCode || format == BinaryFormatCode +} + +func (Int8Codec) PreferredFormat() int16 { + return BinaryFormatCode +} + +func (Int8Codec) PlanEncode(m *Map, oid uint32, format int16, value any) EncodePlan { + switch format { + case BinaryFormatCode: + switch value.(type) { + case int64: + return encodePlanInt8CodecBinaryInt64{} + case Int64Valuer: + return encodePlanInt8CodecBinaryInt64Valuer{} + } + case TextFormatCode: + switch value.(type) { + case int64: + return encodePlanInt8CodecTextInt64{} + case Int64Valuer: + return encodePlanInt8CodecTextInt64Valuer{} + } + } + + return nil +} + +type encodePlanInt8CodecBinaryInt64 struct{} + +func (encodePlanInt8CodecBinaryInt64) Encode(value any, buf []byte) (newBuf []byte, err error) { + n := value.(int64) + return pgio.AppendInt64(buf, int64(n)), nil +} + +type encodePlanInt8CodecTextInt64 struct{} + +func (encodePlanInt8CodecTextInt64) Encode(value any, buf []byte) (newBuf []byte, err error) { + n := value.(int64) + return append(buf, strconv.FormatInt(int64(n), 10)...), nil +} + +type encodePlanInt8CodecBinaryInt64Valuer struct{} + +func (encodePlanInt8CodecBinaryInt64Valuer) Encode(value any, buf []byte) (newBuf []byte, err error) { + n, err := value.(Int64Valuer).Int64Value() + if err != nil { + return nil, err + } + + if !n.Valid { + return nil, nil + } + + if n.Int64 > math.MaxInt64 { + return nil, fmt.Errorf("%d is greater than maximum value for int8", n.Int64) + } + if n.Int64 < math.MinInt64 { + return nil, fmt.Errorf("%d is less than minimum value for int8", n.Int64) + } + + return pgio.AppendInt64(buf, int64(n.Int64)), nil +} + +type encodePlanInt8CodecTextInt64Valuer struct{} + +func (encodePlanInt8CodecTextInt64Valuer) Encode(value any, buf []byte) (newBuf []byte, err error) { + n, err := value.(Int64Valuer).Int64Value() + if err != nil { + return nil, err + } + + if !n.Valid { + return nil, nil + } + + if n.Int64 > math.MaxInt64 { + return nil, fmt.Errorf("%d is greater than maximum value for int8", n.Int64) + } + if n.Int64 < math.MinInt64 { + return nil, fmt.Errorf("%d is less than minimum value for int8", n.Int64) + } + + return append(buf, strconv.FormatInt(n.Int64, 10)...), nil +} + +func (Int8Codec) PlanScan(m *Map, oid uint32, format int16, target any) ScanPlan { + + switch format { + case BinaryFormatCode: + switch target.(type) { + case *int8: + return scanPlanBinaryInt8ToInt8{} + case *int16: + return scanPlanBinaryInt8ToInt16{} + case *int32: + return scanPlanBinaryInt8ToInt32{} + case *int64: + return scanPlanBinaryInt8ToInt64{} + case *int: + return scanPlanBinaryInt8ToInt{} + case *uint8: + return scanPlanBinaryInt8ToUint8{} + case *uint16: + return scanPlanBinaryInt8ToUint16{} + case *uint32: + return scanPlanBinaryInt8ToUint32{} + case *uint64: + return scanPlanBinaryInt8ToUint64{} + case *uint: + return scanPlanBinaryInt8ToUint{} + case Int64Scanner: + return scanPlanBinaryInt8ToInt64Scanner{} + case TextScanner: + return scanPlanBinaryInt8ToTextScanner{} + } + case TextFormatCode: + switch target.(type) { + case *int8: + return scanPlanTextAnyToInt8{} + case *int16: + return scanPlanTextAnyToInt16{} + case *int32: + return scanPlanTextAnyToInt32{} + case *int64: + return scanPlanTextAnyToInt64{} + case *int: + return scanPlanTextAnyToInt{} + case *uint8: + return scanPlanTextAnyToUint8{} + case *uint16: + return scanPlanTextAnyToUint16{} + case *uint32: + return scanPlanTextAnyToUint32{} + case *uint64: + return scanPlanTextAnyToUint64{} + case *uint: + return scanPlanTextAnyToUint{} + case Int64Scanner: + return scanPlanTextAnyToInt64Scanner{} + } + } + + return nil +} + +func (c Int8Codec) DecodeDatabaseSQLValue(m *Map, oid uint32, format int16, src []byte) (driver.Value, error) { + if src == nil { + return nil, nil + } + + var n int64 + err := codecScan(c, m, oid, format, src, &n) + if err != nil { + return nil, err + } + return n, nil +} + +func (c Int8Codec) DecodeValue(m *Map, oid uint32, format int16, src []byte) (any, error) { + if src == nil { + return nil, nil + } + + var n int64 + err := codecScan(c, m, oid, format, src, &n) + if err != nil { + return nil, err + } + return n, nil +} + +type scanPlanBinaryInt8ToInt8 struct{} + +func (scanPlanBinaryInt8ToInt8) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + if len(src) != 8 { + return fmt.Errorf("invalid length for int8: %v", len(src)) + } + + p, ok := (dst).(*int8) + if !ok { + return ErrScanTargetTypeChanged + } + + n := int64(binary.BigEndian.Uint64(src)) + if n < math.MinInt8 { + return fmt.Errorf("%d is less than minimum value for int8", n) + } else if n > math.MaxInt8 { + return fmt.Errorf("%d is greater than maximum value for int8", n) + } + + *p = int8(n) + + return nil +} + +type scanPlanBinaryInt8ToUint8 struct{} + +func (scanPlanBinaryInt8ToUint8) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + if len(src) != 8 { + return fmt.Errorf("invalid length for uint8: %v", len(src)) + } + + p, ok := (dst).(*uint8) + if !ok { + return ErrScanTargetTypeChanged + } + + n := int64(binary.BigEndian.Uint64(src)) + if n < 0 { + return fmt.Errorf("%d is less than minimum value for uint8", n) + } + + if n > math.MaxUint8 { + return fmt.Errorf("%d is greater than maximum value for uint8", n) + } + + *p = uint8(n) + + return nil +} + +type scanPlanBinaryInt8ToInt16 struct{} + +func (scanPlanBinaryInt8ToInt16) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + if len(src) != 8 { + return fmt.Errorf("invalid length for int8: %v", len(src)) + } + + p, ok := (dst).(*int16) + if !ok { + return ErrScanTargetTypeChanged + } + + n := int64(binary.BigEndian.Uint64(src)) + if n < math.MinInt16 { + return fmt.Errorf("%d is less than minimum value for int16", n) + } else if n > math.MaxInt16 { + return fmt.Errorf("%d is greater than maximum value for int16", n) + } + + *p = int16(n) + + return nil +} + +type scanPlanBinaryInt8ToUint16 struct{} + +func (scanPlanBinaryInt8ToUint16) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + if len(src) != 8 { + return fmt.Errorf("invalid length for uint8: %v", len(src)) + } + + p, ok := (dst).(*uint16) + if !ok { + return ErrScanTargetTypeChanged + } + + n := int64(binary.BigEndian.Uint64(src)) + if n < 0 { + return fmt.Errorf("%d is less than minimum value for uint16", n) + } + + if n > math.MaxUint16 { + return fmt.Errorf("%d is greater than maximum value for uint16", n) + } + + *p = uint16(n) + + return nil +} + +type scanPlanBinaryInt8ToInt32 struct{} + +func (scanPlanBinaryInt8ToInt32) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + if len(src) != 8 { + return fmt.Errorf("invalid length for int8: %v", len(src)) + } + + p, ok := (dst).(*int32) + if !ok { + return ErrScanTargetTypeChanged + } + + n := int64(binary.BigEndian.Uint64(src)) + if n < math.MinInt32 { + return fmt.Errorf("%d is less than minimum value for int32", n) + } else if n > math.MaxInt32 { + return fmt.Errorf("%d is greater than maximum value for int32", n) + } + + *p = int32(n) + + return nil +} + +type scanPlanBinaryInt8ToUint32 struct{} + +func (scanPlanBinaryInt8ToUint32) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + if len(src) != 8 { + return fmt.Errorf("invalid length for uint8: %v", len(src)) + } + + p, ok := (dst).(*uint32) + if !ok { + return ErrScanTargetTypeChanged + } + + n := int64(binary.BigEndian.Uint64(src)) + if n < 0 { + return fmt.Errorf("%d is less than minimum value for uint32", n) + } + + if n > math.MaxUint32 { + return fmt.Errorf("%d is greater than maximum value for uint32", n) + } + + *p = uint32(n) + + return nil +} + +type scanPlanBinaryInt8ToInt64 struct{} + +func (scanPlanBinaryInt8ToInt64) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + if len(src) != 8 { + return fmt.Errorf("invalid length for int8: %v", len(src)) + } + + p, ok := (dst).(*int64) + if !ok { + return ErrScanTargetTypeChanged + } + + *p = int64(binary.BigEndian.Uint64(src)) + + return nil +} + +type scanPlanBinaryInt8ToUint64 struct{} + +func (scanPlanBinaryInt8ToUint64) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + if len(src) != 8 { + return fmt.Errorf("invalid length for uint8: %v", len(src)) + } + + p, ok := (dst).(*uint64) + if !ok { + return ErrScanTargetTypeChanged + } + + n := int64(binary.BigEndian.Uint64(src)) + if n < 0 { + return fmt.Errorf("%d is less than minimum value for uint64", n) + } + + *p = uint64(n) + + return nil +} + +type scanPlanBinaryInt8ToInt struct{} + +func (scanPlanBinaryInt8ToInt) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + if len(src) != 8 { + return fmt.Errorf("invalid length for int8: %v", len(src)) + } + + p, ok := (dst).(*int) + if !ok { + return ErrScanTargetTypeChanged + } + + n := int64(binary.BigEndian.Uint64(src)) + if n < math.MinInt { + return fmt.Errorf("%d is less than minimum value for int", n) + } else if n > math.MaxInt { + return fmt.Errorf("%d is greater than maximum value for int", n) + } + + *p = int(n) + + return nil +} + +type scanPlanBinaryInt8ToUint struct{} + +func (scanPlanBinaryInt8ToUint) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + if len(src) != 8 { + return fmt.Errorf("invalid length for uint8: %v", len(src)) + } + + p, ok := (dst).(*uint) + if !ok { + return ErrScanTargetTypeChanged + } + + n := int64(int64(binary.BigEndian.Uint64(src))) + if n < 0 { + return fmt.Errorf("%d is less than minimum value for uint", n) + } + + if uint64(n) > math.MaxUint { + return fmt.Errorf("%d is greater than maximum value for uint", n) + } + + *p = uint(n) + + return nil +} + +type scanPlanBinaryInt8ToInt64Scanner struct{} + +func (scanPlanBinaryInt8ToInt64Scanner) Scan(src []byte, dst any) error { + s, ok := (dst).(Int64Scanner) + if !ok { + return ErrScanTargetTypeChanged + } + + if src == nil { + return s.ScanInt64(Int8{}) + } + + if len(src) != 8 { + return fmt.Errorf("invalid length for int8: %v", len(src)) + } + + n := int64(int64(binary.BigEndian.Uint64(src))) + + return s.ScanInt64(Int8{Int64: n, Valid: true}) +} + +type scanPlanBinaryInt8ToTextScanner struct{} + +func (scanPlanBinaryInt8ToTextScanner) Scan(src []byte, dst any) error { + s, ok := (dst).(TextScanner) + if !ok { + return ErrScanTargetTypeChanged + } + + if src == nil { + return s.ScanText(Text{}) + } + + if len(src) != 8 { + return fmt.Errorf("invalid length for int8: %v", len(src)) + } + + n := int64(int64(binary.BigEndian.Uint64(src))) + + return s.ScanText(Text{String: strconv.FormatInt(n, 10), Valid: true}) +} + +type scanPlanTextAnyToInt8 struct{} + +func (scanPlanTextAnyToInt8) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + p, ok := (dst).(*int8) + if !ok { + return ErrScanTargetTypeChanged + } + + n, err := strconv.ParseInt(string(src), 10, 8) + if err != nil { + return err + } + + *p = int8(n) + return nil +} + +type scanPlanTextAnyToUint8 struct{} + +func (scanPlanTextAnyToUint8) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + p, ok := (dst).(*uint8) + if !ok { + return ErrScanTargetTypeChanged + } + + n, err := strconv.ParseUint(string(src), 10, 8) + if err != nil { + return err + } + + *p = uint8(n) + return nil +} + +type scanPlanTextAnyToInt16 struct{} + +func (scanPlanTextAnyToInt16) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + p, ok := (dst).(*int16) + if !ok { + return ErrScanTargetTypeChanged + } + + n, err := strconv.ParseInt(string(src), 10, 16) + if err != nil { + return err + } + + *p = int16(n) + return nil +} + +type scanPlanTextAnyToUint16 struct{} + +func (scanPlanTextAnyToUint16) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + p, ok := (dst).(*uint16) + if !ok { + return ErrScanTargetTypeChanged + } + + n, err := strconv.ParseUint(string(src), 10, 16) + if err != nil { + return err + } + + *p = uint16(n) + return nil +} + +type scanPlanTextAnyToInt32 struct{} + +func (scanPlanTextAnyToInt32) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + p, ok := (dst).(*int32) + if !ok { + return ErrScanTargetTypeChanged + } + + n, err := strconv.ParseInt(string(src), 10, 32) + if err != nil { + return err + } + + *p = int32(n) + return nil +} + +type scanPlanTextAnyToUint32 struct{} + +func (scanPlanTextAnyToUint32) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + p, ok := (dst).(*uint32) + if !ok { + return ErrScanTargetTypeChanged + } + + n, err := strconv.ParseUint(string(src), 10, 32) + if err != nil { + return err + } + + *p = uint32(n) + return nil +} + +type scanPlanTextAnyToInt64 struct{} + +func (scanPlanTextAnyToInt64) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + p, ok := (dst).(*int64) + if !ok { + return ErrScanTargetTypeChanged + } + + n, err := strconv.ParseInt(string(src), 10, 64) + if err != nil { + return err + } + + *p = int64(n) + return nil +} + +type scanPlanTextAnyToUint64 struct{} + +func (scanPlanTextAnyToUint64) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + p, ok := (dst).(*uint64) + if !ok { + return ErrScanTargetTypeChanged + } + + n, err := strconv.ParseUint(string(src), 10, 64) + if err != nil { + return err + } + + *p = uint64(n) + return nil +} + +type scanPlanTextAnyToInt struct{} + +func (scanPlanTextAnyToInt) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + p, ok := (dst).(*int) + if !ok { + return ErrScanTargetTypeChanged + } + + n, err := strconv.ParseInt(string(src), 10, 0) + if err != nil { + return err + } + + *p = int(n) + return nil +} + +type scanPlanTextAnyToUint struct{} + +func (scanPlanTextAnyToUint) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + p, ok := (dst).(*uint) + if !ok { + return ErrScanTargetTypeChanged + } + + n, err := strconv.ParseUint(string(src), 10, 0) + if err != nil { + return err + } + + *p = uint(n) + return nil +} + +type scanPlanTextAnyToInt64Scanner struct{} + +func (scanPlanTextAnyToInt64Scanner) Scan(src []byte, dst any) error { + s, ok := (dst).(Int64Scanner) + if !ok { + return ErrScanTargetTypeChanged + } + + if src == nil { + return s.ScanInt64(Int8{}) + } + + n, err := strconv.ParseInt(string(src), 10, 64) + if err != nil { + return err + } + + err = s.ScanInt64(Int8{Int64: n, Valid: true}) + if err != nil { + return err + } + + return nil +} diff --git a/vendor/github.com/jackc/pgx/v5/pgtype/int.go.erb b/vendor/github.com/jackc/pgx/v5/pgtype/int.go.erb new file mode 100644 index 0000000000..c2d40f60b9 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgtype/int.go.erb @@ -0,0 +1,551 @@ +package pgtype + +import ( + "database/sql/driver" + "encoding/binary" + "encoding/json" + "fmt" + "math" + "strconv" + + "github.com/jackc/pgx/v5/internal/pgio" +) + +type Int64Scanner interface { + ScanInt64(Int8) error +} + +type Int64Valuer interface { + Int64Value() (Int8, error) +} + + +<% [2, 4, 8].each do |pg_byte_size| %> +<% pg_bit_size = pg_byte_size * 8 %> +type Int<%= pg_byte_size %> struct { + Int<%= pg_bit_size %> int<%= pg_bit_size %> + Valid bool +} + +// ScanInt64 implements the [Int64Scanner] interface. +func (dst *Int<%= pg_byte_size %>) ScanInt64(n Int8) error { + if !n.Valid { + *dst = Int<%= pg_byte_size %>{} + return nil + } + + if n.Int64 < math.MinInt<%= pg_bit_size %> { + return fmt.Errorf("%d is less than minimum value for Int<%= pg_byte_size %>", n.Int64) + } + if n.Int64 > math.MaxInt<%= pg_bit_size %> { + return fmt.Errorf("%d is greater than maximum value for Int<%= pg_byte_size %>", n.Int64) + } + *dst = Int<%= pg_byte_size %>{Int<%= pg_bit_size %>: int<%= pg_bit_size %>(n.Int64), Valid: true} + + return nil +} + +// Int64Value implements the [Int64Valuer] interface. +func (n Int<%= pg_byte_size %>) Int64Value() (Int8, error) { + return Int8{Int64: int64(n.Int<%= pg_bit_size %>), Valid: n.Valid}, nil +} + +// Scan implements the [database/sql.Scanner] interface. +func (dst *Int<%= pg_byte_size %>) Scan(src any) error { + if src == nil { + *dst = Int<%= pg_byte_size %>{} + return nil + } + + var n int64 + + switch src := src.(type) { + case int64: + n = src + case string: + var err error + n, err = strconv.ParseInt(src, 10, <%= pg_bit_size %>) + if err != nil { + return err + } + case []byte: + var err error + n, err = strconv.ParseInt(string(src), 10, <%= pg_bit_size %>) + if err != nil { + return err + } + default: + return fmt.Errorf("cannot scan %T", src) + } + + if n < math.MinInt<%= pg_bit_size %> { + return fmt.Errorf("%d is greater than maximum value for Int<%= pg_byte_size %>", n) + } + if n > math.MaxInt<%= pg_bit_size %> { + return fmt.Errorf("%d is greater than maximum value for Int<%= pg_byte_size %>", n) + } + *dst = Int<%= pg_byte_size %>{Int<%= pg_bit_size %>: int<%= pg_bit_size %>(n), Valid: true} + + return nil +} + +// Value implements the [database/sql/driver.Valuer] interface. +func (src Int<%= pg_byte_size %>) Value() (driver.Value, error) { + if !src.Valid { + return nil, nil + } + return int64(src.Int<%= pg_bit_size %>), nil +} + +// MarshalJSON implements the [encoding/json.Marshaler] interface. +func (src Int<%= pg_byte_size %>) MarshalJSON() ([]byte, error) { + if !src.Valid { + return []byte("null"), nil + } + return []byte(strconv.FormatInt(int64(src.Int<%= pg_bit_size %>), 10)), nil +} + +// UnmarshalJSON implements the [encoding/json.Unmarshaler] interface. +func (dst *Int<%= pg_byte_size %>) UnmarshalJSON(b []byte) error { + var n *int<%= pg_bit_size %> + err := json.Unmarshal(b, &n) + if err != nil { + return err + } + + if n == nil { + *dst = Int<%= pg_byte_size %>{} + } else { + *dst = Int<%= pg_byte_size %>{Int<%= pg_bit_size %>: *n, Valid: true} + } + + return nil +} + +type Int<%= pg_byte_size %>Codec struct{} + +func (Int<%= pg_byte_size %>Codec) FormatSupported(format int16) bool { + return format == TextFormatCode || format == BinaryFormatCode +} + +func (Int<%= pg_byte_size %>Codec) PreferredFormat() int16 { + return BinaryFormatCode +} + +func (Int<%= pg_byte_size %>Codec) PlanEncode(m *Map, oid uint32, format int16, value any) EncodePlan { + switch format { + case BinaryFormatCode: + switch value.(type) { + case int<%= pg_bit_size %>: + return encodePlanInt<%= pg_byte_size %>CodecBinaryInt<%= pg_bit_size %>{} + case Int64Valuer: + return encodePlanInt<%= pg_byte_size %>CodecBinaryInt64Valuer{} + } + case TextFormatCode: + switch value.(type) { + case int<%= pg_bit_size %>: + return encodePlanInt<%= pg_byte_size %>CodecTextInt<%= pg_bit_size %>{} + case Int64Valuer: + return encodePlanInt<%= pg_byte_size %>CodecTextInt64Valuer{} + } + } + + return nil +} + +type encodePlanInt<%= pg_byte_size %>CodecBinaryInt<%= pg_bit_size %> struct{} + +func (encodePlanInt<%= pg_byte_size %>CodecBinaryInt<%= pg_bit_size %>) Encode(value any, buf []byte) (newBuf []byte, err error) { + n := value.(int<%= pg_bit_size %>) + return pgio.AppendInt<%= pg_bit_size %>(buf, int<%= pg_bit_size %>(n)), nil +} + +type encodePlanInt<%= pg_byte_size %>CodecTextInt<%= pg_bit_size %> struct{} + +func (encodePlanInt<%= pg_byte_size %>CodecTextInt<%= pg_bit_size %>) Encode(value any, buf []byte) (newBuf []byte, err error) { + n := value.(int<%= pg_bit_size %>) + return append(buf, strconv.FormatInt(int64(n), 10)...), nil +} + +type encodePlanInt<%= pg_byte_size %>CodecBinaryInt64Valuer struct{} + +func (encodePlanInt<%= pg_byte_size %>CodecBinaryInt64Valuer) Encode(value any, buf []byte) (newBuf []byte, err error) { + n, err := value.(Int64Valuer).Int64Value() + if err != nil { + return nil, err + } + + if !n.Valid { + return nil, nil + } + + if n.Int64 > math.MaxInt<%= pg_bit_size %> { + return nil, fmt.Errorf("%d is greater than maximum value for int<%= pg_byte_size %>", n.Int64) + } + if n.Int64 < math.MinInt<%= pg_bit_size %> { + return nil, fmt.Errorf("%d is less than minimum value for int<%= pg_byte_size %>", n.Int64) + } + + return pgio.AppendInt<%= pg_bit_size %>(buf, int<%= pg_bit_size %>(n.Int64)), nil +} + +type encodePlanInt<%= pg_byte_size %>CodecTextInt64Valuer struct{} + +func (encodePlanInt<%= pg_byte_size %>CodecTextInt64Valuer) Encode(value any, buf []byte) (newBuf []byte, err error) { + n, err := value.(Int64Valuer).Int64Value() + if err != nil { + return nil, err + } + + if !n.Valid { + return nil, nil + } + + if n.Int64 > math.MaxInt<%= pg_bit_size %> { + return nil, fmt.Errorf("%d is greater than maximum value for int<%= pg_byte_size %>", n.Int64) + } + if n.Int64 < math.MinInt<%= pg_bit_size %> { + return nil, fmt.Errorf("%d is less than minimum value for int<%= pg_byte_size %>", n.Int64) + } + + return append(buf, strconv.FormatInt(n.Int64, 10)...), nil +} + +func (Int<%= pg_byte_size %>Codec) PlanScan(m *Map, oid uint32, format int16, target any) ScanPlan { + + switch format { + case BinaryFormatCode: + switch target.(type) { + case *int8: + return scanPlanBinaryInt<%= pg_byte_size %>ToInt8{} + case *int16: + return scanPlanBinaryInt<%= pg_byte_size %>ToInt16{} + case *int32: + return scanPlanBinaryInt<%= pg_byte_size %>ToInt32{} + case *int64: + return scanPlanBinaryInt<%= pg_byte_size %>ToInt64{} + case *int: + return scanPlanBinaryInt<%= pg_byte_size %>ToInt{} + case *uint8: + return scanPlanBinaryInt<%= pg_byte_size %>ToUint8{} + case *uint16: + return scanPlanBinaryInt<%= pg_byte_size %>ToUint16{} + case *uint32: + return scanPlanBinaryInt<%= pg_byte_size %>ToUint32{} + case *uint64: + return scanPlanBinaryInt<%= pg_byte_size %>ToUint64{} + case *uint: + return scanPlanBinaryInt<%= pg_byte_size %>ToUint{} + case Int64Scanner: + return scanPlanBinaryInt<%= pg_byte_size %>ToInt64Scanner{} + case TextScanner: + return scanPlanBinaryInt<%= pg_byte_size %>ToTextScanner{} + } + case TextFormatCode: + switch target.(type) { + case *int8: + return scanPlanTextAnyToInt8{} + case *int16: + return scanPlanTextAnyToInt16{} + case *int32: + return scanPlanTextAnyToInt32{} + case *int64: + return scanPlanTextAnyToInt64{} + case *int: + return scanPlanTextAnyToInt{} + case *uint8: + return scanPlanTextAnyToUint8{} + case *uint16: + return scanPlanTextAnyToUint16{} + case *uint32: + return scanPlanTextAnyToUint32{} + case *uint64: + return scanPlanTextAnyToUint64{} + case *uint: + return scanPlanTextAnyToUint{} + case Int64Scanner: + return scanPlanTextAnyToInt64Scanner{} + } + } + + return nil +} + +func (c Int<%= pg_byte_size %>Codec) DecodeDatabaseSQLValue(m *Map, oid uint32, format int16, src []byte) (driver.Value, error) { + if src == nil { + return nil, nil + } + + var n int64 + err := codecScan(c, m, oid, format, src, &n) + if err != nil { + return nil, err + } + return n, nil +} + +func (c Int<%= pg_byte_size %>Codec) DecodeValue(m *Map, oid uint32, format int16, src []byte) (any, error) { + if src == nil { + return nil, nil + } + + var n int<%= pg_bit_size %> + err := codecScan(c, m, oid, format, src, &n) + if err != nil { + return nil, err + } + return n, nil +} + +<%# PostgreSQL binary format integer to fixed size Go integers %> +<% [8, 16, 32, 64].each do |dst_bit_size| %> +type scanPlanBinaryInt<%= pg_byte_size %>ToInt<%= dst_bit_size %> struct{} + +func (scanPlanBinaryInt<%= pg_byte_size %>ToInt<%= dst_bit_size %>) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + if len(src) != <%= pg_byte_size %> { + return fmt.Errorf("invalid length for int<%= pg_byte_size %>: %v", len(src)) + } + + p, ok := (dst).(*int<%= dst_bit_size %>) + if !ok { + return ErrScanTargetTypeChanged + } + + <% if dst_bit_size < pg_bit_size %> + n := int<%= pg_bit_size %>(binary.BigEndian.Uint<%= pg_bit_size %>(src)) + if n < math.MinInt<%= dst_bit_size %> { + return fmt.Errorf("%d is less than minimum value for int<%= dst_bit_size %>", n) + } else if n > math.MaxInt<%= dst_bit_size %> { + return fmt.Errorf("%d is greater than maximum value for int<%= dst_bit_size %>", n) + } + + *p = int<%= dst_bit_size %>(n) + <% elsif dst_bit_size == pg_bit_size %> + *p = int<%= dst_bit_size %>(binary.BigEndian.Uint<%= pg_bit_size %>(src)) + <% else %> + *p = int<%= dst_bit_size %>(int<%= pg_bit_size %>(binary.BigEndian.Uint<%= pg_bit_size %>(src))) + <% end %> + + return nil +} + +type scanPlanBinaryInt<%= pg_byte_size %>ToUint<%= dst_bit_size %> struct{} + +func (scanPlanBinaryInt<%= pg_byte_size %>ToUint<%= dst_bit_size %>) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + if len(src) != <%= pg_byte_size %> { + return fmt.Errorf("invalid length for uint<%= pg_byte_size %>: %v", len(src)) + } + + p, ok := (dst).(*uint<%= dst_bit_size %>) + if !ok { + return ErrScanTargetTypeChanged + } + + n := int<%= pg_bit_size %>(binary.BigEndian.Uint<%= pg_bit_size %>(src)) + if n < 0 { + return fmt.Errorf("%d is less than minimum value for uint<%= dst_bit_size %>", n) + } + <% if dst_bit_size < pg_bit_size %> + if n > math.MaxUint<%= dst_bit_size %> { + return fmt.Errorf("%d is greater than maximum value for uint<%= dst_bit_size %>", n) + } + <% end %> + *p = uint<%= dst_bit_size %>(n) + + return nil +} +<% end %> + +<%# PostgreSQL binary format integer to Go machine integers %> +type scanPlanBinaryInt<%= pg_byte_size %>ToInt struct{} + +func (scanPlanBinaryInt<%= pg_byte_size %>ToInt) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + if len(src) != <%= pg_byte_size %> { + return fmt.Errorf("invalid length for int<%= pg_byte_size %>: %v", len(src)) + } + + p, ok := (dst).(*int) + if !ok { + return ErrScanTargetTypeChanged + } + + <% if 32 < pg_bit_size %> + n := int64(binary.BigEndian.Uint<%= pg_bit_size %>(src)) + if n < math.MinInt { + return fmt.Errorf("%d is less than minimum value for int", n) + } else if n > math.MaxInt { + return fmt.Errorf("%d is greater than maximum value for int", n) + } + + *p = int(n) + <% else %> + *p = int(int<%= pg_bit_size %>(binary.BigEndian.Uint<%= pg_bit_size %>(src))) + <% end %> + + return nil +} + +type scanPlanBinaryInt<%= pg_byte_size %>ToUint struct{} + +func (scanPlanBinaryInt<%= pg_byte_size %>ToUint) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + if len(src) != <%= pg_byte_size %> { + return fmt.Errorf("invalid length for uint<%= pg_byte_size %>: %v", len(src)) + } + + p, ok := (dst).(*uint) + if !ok { + return ErrScanTargetTypeChanged + } + + n := int64(int<%= pg_bit_size %>(binary.BigEndian.Uint<%= pg_bit_size %>(src))) + if n < 0 { + return fmt.Errorf("%d is less than minimum value for uint", n) + } + <% if 32 < pg_bit_size %> + if uint64(n) > math.MaxUint { + return fmt.Errorf("%d is greater than maximum value for uint", n) + } + <% end %> + *p = uint(n) + + return nil +} + +<%# PostgreSQL binary format integer to Go Int64Scanner %> +type scanPlanBinaryInt<%= pg_byte_size %>ToInt64Scanner struct{} + +func (scanPlanBinaryInt<%= pg_byte_size %>ToInt64Scanner) Scan(src []byte, dst any) error { + s, ok := (dst).(Int64Scanner) + if !ok { + return ErrScanTargetTypeChanged + } + + if src == nil { + return s.ScanInt64(Int8{}) + } + + if len(src) != <%= pg_byte_size %> { + return fmt.Errorf("invalid length for int<%= pg_byte_size %>: %v", len(src)) + } + + + n := int64(int<%= pg_bit_size %>(binary.BigEndian.Uint<%= pg_bit_size %>(src))) + + return s.ScanInt64(Int8{Int64: n, Valid: true}) +} + +<%# PostgreSQL binary format integer to Go TextScanner %> +type scanPlanBinaryInt<%= pg_byte_size %>ToTextScanner struct{} + +func (scanPlanBinaryInt<%= pg_byte_size %>ToTextScanner) Scan(src []byte, dst any) error { + s, ok := (dst).(TextScanner) + if !ok { + return ErrScanTargetTypeChanged + } + + if src == nil { + return s.ScanText(Text{}) + } + + if len(src) != <%= pg_byte_size %> { + return fmt.Errorf("invalid length for int<%= pg_byte_size %>: %v", len(src)) + } + + + n := int64(int<%= pg_bit_size %>(binary.BigEndian.Uint<%= pg_bit_size %>(src))) + + return s.ScanText(Text{String: strconv.FormatInt(n, 10), Valid: true}) +} +<% end %> + +<%# Any text to all integer types %> +<% [ + ["8", 8], + ["16", 16], + ["32", 32], + ["64", 64], + ["", 0] +].each do |type_suffix, bit_size| %> +type scanPlanTextAnyToInt<%= type_suffix %> struct{} + +func (scanPlanTextAnyToInt<%= type_suffix %>) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + p, ok := (dst).(*int<%= type_suffix %>) + if !ok { + return ErrScanTargetTypeChanged + } + + n, err := strconv.ParseInt(string(src), 10, <%= bit_size %>) + if err != nil { + return err + } + + *p = int<%= type_suffix %>(n) + return nil +} + +type scanPlanTextAnyToUint<%= type_suffix %> struct{} + +func (scanPlanTextAnyToUint<%= type_suffix %>) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + p, ok := (dst).(*uint<%= type_suffix %>) + if !ok { + return ErrScanTargetTypeChanged + } + + n, err := strconv.ParseUint(string(src), 10, <%= bit_size %>) + if err != nil { + return err + } + + *p = uint<%= type_suffix %>(n) + return nil +} +<% end %> + +type scanPlanTextAnyToInt64Scanner struct{} + +func (scanPlanTextAnyToInt64Scanner) Scan(src []byte, dst any) error { + s, ok := (dst).(Int64Scanner) + if !ok { + return ErrScanTargetTypeChanged + } + + if src == nil { + return s.ScanInt64(Int8{}) + } + + n, err := strconv.ParseInt(string(src), 10, 64) + if err != nil { + return err + } + + err = s.ScanInt64(Int8{Int64: n, Valid: true}) + if err != nil { + return err + } + + return nil +} diff --git a/vendor/github.com/jackc/pgx/v5/pgtype/int_test.go.erb b/vendor/github.com/jackc/pgx/v5/pgtype/int_test.go.erb new file mode 100644 index 0000000000..ac9a3f1430 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgtype/int_test.go.erb @@ -0,0 +1,93 @@ +package pgtype_test + +import ( + "math" + "testing" + + "github.com/jackc/pgx/v5/pgtype" +) + +<% [2, 4, 8].each do |pg_byte_size| %> +<% pg_bit_size = pg_byte_size * 8 %> +func TestInt<%= pg_byte_size %>Codec(t *testing.T) { + pgxtest.RunValueRoundTripTests(context.Background(), t, defaultConnTestRunner, nil, "int<%= pg_byte_size %>", []pgxtest.ValueRoundTripTest{ + {int8(1), new(int<%= pg_bit_size %>), isExpectedEq(int<%= pg_bit_size %>(1))}, + {int16(1), new(int<%= pg_bit_size %>), isExpectedEq(int<%= pg_bit_size %>(1))}, + {int32(1), new(int<%= pg_bit_size %>), isExpectedEq(int<%= pg_bit_size %>(1))}, + {int64(1), new(int<%= pg_bit_size %>), isExpectedEq(int<%= pg_bit_size %>(1))}, + {uint8(1), new(int<%= pg_bit_size %>), isExpectedEq(int<%= pg_bit_size %>(1))}, + {uint16(1), new(int<%= pg_bit_size %>), isExpectedEq(int<%= pg_bit_size %>(1))}, + {uint32(1), new(int<%= pg_bit_size %>), isExpectedEq(int<%= pg_bit_size %>(1))}, + {uint64(1), new(int<%= pg_bit_size %>), isExpectedEq(int<%= pg_bit_size %>(1))}, + {int(1), new(int<%= pg_bit_size %>), isExpectedEq(int<%= pg_bit_size %>(1))}, + {uint(1), new(int<%= pg_bit_size %>), isExpectedEq(int<%= pg_bit_size %>(1))}, + {pgtype.Int<%= pg_byte_size %>{Int<%= pg_bit_size %>: 1, Valid: true}, new(int<%= pg_bit_size %>), isExpectedEq(int<%= pg_bit_size %>(1))}, + {int32(-1), new(pgtype.Int<%= pg_byte_size %>), isExpectedEq(pgtype.Int<%= pg_byte_size %>{Int<%= pg_bit_size %>: -1, Valid: true})}, + {1, new(int8), isExpectedEq(int8(1))}, + {1, new(int16), isExpectedEq(int16(1))}, + {1, new(int32), isExpectedEq(int32(1))}, + {1, new(int64), isExpectedEq(int64(1))}, + {1, new(uint8), isExpectedEq(uint8(1))}, + {1, new(uint16), isExpectedEq(uint16(1))}, + {1, new(uint32), isExpectedEq(uint32(1))}, + {1, new(uint64), isExpectedEq(uint64(1))}, + {1, new(int), isExpectedEq(int(1))}, + {1, new(uint), isExpectedEq(uint(1))}, + {-1, new(int8), isExpectedEq(int8(-1))}, + {-1, new(int16), isExpectedEq(int16(-1))}, + {-1, new(int32), isExpectedEq(int32(-1))}, + {-1, new(int64), isExpectedEq(int64(-1))}, + {-1, new(int), isExpectedEq(int(-1))}, + {math.MinInt<%= pg_bit_size %>, new(int<%= pg_bit_size %>), isExpectedEq(int<%= pg_bit_size %>(math.MinInt<%= pg_bit_size %>))}, + {-1, new(int<%= pg_bit_size %>), isExpectedEq(int<%= pg_bit_size %>(-1))}, + {0, new(int<%= pg_bit_size %>), isExpectedEq(int<%= pg_bit_size %>(0))}, + {1, new(int<%= pg_bit_size %>), isExpectedEq(int<%= pg_bit_size %>(1))}, + {math.MaxInt<%= pg_bit_size %>, new(int<%= pg_bit_size %>), isExpectedEq(int<%= pg_bit_size %>(math.MaxInt<%= pg_bit_size %>))}, + {1, new(pgtype.Int<%= pg_byte_size %>), isExpectedEq(pgtype.Int<%= pg_byte_size %>{Int<%= pg_bit_size %>: 1, Valid: true})}, + {"1", new(string), isExpectedEq("1")}, + {pgtype.Int<%= pg_byte_size %>{}, new(pgtype.Int<%= pg_byte_size %>), isExpectedEq(pgtype.Int<%= pg_byte_size %>{})}, + {nil, new(*int<%= pg_bit_size %>), isExpectedEq((*int<%= pg_bit_size %>)(nil))}, + }) +} + +func TestInt<%= pg_byte_size %>MarshalJSON(t *testing.T) { + successfulTests := []struct { + source pgtype.Int<%= pg_byte_size %> + result string + }{ + {source: pgtype.Int<%= pg_byte_size %>{Int<%= pg_bit_size %>: 0}, result: "null"}, + {source: pgtype.Int<%= pg_byte_size %>{Int<%= pg_bit_size %>: 1, Valid: true}, result: "1"}, + } + for i, tt := range successfulTests { + r, err := tt.source.MarshalJSON() + if err != nil { + t.Errorf("%d: %v", i, err) + } + + if string(r) != tt.result { + t.Errorf("%d: expected %v to convert to %v, but it was %v", i, tt.source, tt.result, string(r)) + } + } +} + +func TestInt<%= pg_byte_size %>UnmarshalJSON(t *testing.T) { + successfulTests := []struct { + source string + result pgtype.Int<%= pg_byte_size %> + }{ + {source: "null", result: pgtype.Int<%= pg_byte_size %>{Int<%= pg_bit_size %>: 0}}, + {source: "1", result: pgtype.Int<%= pg_byte_size %>{Int<%= pg_bit_size %>: 1, Valid: true}}, + } + for i, tt := range successfulTests { + var r pgtype.Int<%= pg_byte_size %> + err := r.UnmarshalJSON([]byte(tt.source)) + if err != nil { + t.Errorf("%d: %v", i, err) + } + + if r != tt.result { + t.Errorf("%d: expected %v to convert to %v, but it was %v", i, tt.source, tt.result, r) + } + } +} +<% end %> diff --git a/vendor/github.com/jackc/pgx/v5/pgtype/integration_benchmark_test.go.erb b/vendor/github.com/jackc/pgx/v5/pgtype/integration_benchmark_test.go.erb new file mode 100644 index 0000000000..6f40115340 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgtype/integration_benchmark_test.go.erb @@ -0,0 +1,62 @@ +package pgtype_test + +import ( + "context" + "testing" + + "github.com/jackc/pgx/v5/pgtype/testutil" + "github.com/jackc/pgx/v5" +) + +<% + [ + ["int4", ["int16", "int32", "int64", "uint64", "pgtype.Int4"], [[1, 1], [1, 10], [10, 1], [100, 10]]], + ["numeric", ["int64", "float64", "pgtype.Numeric"], [[1, 1], [1, 10], [10, 1], [100, 10]]], + ].each do |pg_type, go_types, rows_columns| +%> +<% go_types.each do |go_type| %> +<% rows_columns.each do |rows, columns| %> +<% [["Text", "pgx.TextFormatCode"], ["Binary", "pgx.BinaryFormatCode"]].each do |format_name, format_code| %> +func BenchmarkQuery<%= format_name %>FormatDecode_PG_<%= pg_type %>_to_Go_<%= go_type.gsub(/\W/, "_") %>_<%= rows %>_rows_<%= columns %>_columns(b *testing.B) { + defaultConnTestRunner.RunTest(context.Background(), b, func(ctx context.Context, _ testing.TB, conn *pgx.Conn) { + b.ResetTimer() + var v [<%= columns %>]<%= go_type %> + for i := 0; i < b.N; i++ { + rows, _ := conn.Query( + ctx, + `select <% columns.times do |col_idx| %><% if col_idx != 0 %>, <% end %>n::<%= pg_type %> + <%= col_idx%><% end %> from generate_series(1, <%= rows %>) n`, + pgx.QueryResultFormats{<%= format_code %>}, + ) + _, err := pgx.ForEachRow(rows, []any{<% columns.times do |col_idx| %><% if col_idx != 0 %>, <% end %>&v[<%= col_idx%>]<% end %>}, func() error { return nil }) + if err != nil { + b.Fatal(err) + } + } + }) +} +<% end %> +<% end %> +<% end %> +<% end %> + +<% [10, 100, 1000].each do |array_size| %> +<% [["Text", "pgx.TextFormatCode"], ["Binary", "pgx.BinaryFormatCode"]].each do |format_name, format_code| %> +func BenchmarkQuery<%= format_name %>FormatDecode_PG_Int4Array_With_Go_Int4Array_<%= array_size %>(b *testing.B) { + defaultConnTestRunner.RunTest(context.Background(), b, func(ctx context.Context, _ testing.TB, conn *pgx.Conn) { + b.ResetTimer() + var v []int32 + for i := 0; i < b.N; i++ { + rows, _ := conn.Query( + ctx, + `select array_agg(n) from generate_series(1, <%= array_size %>) n`, + pgx.QueryResultFormats{<%= format_code %>}, + ) + _, err := pgx.ForEachRow(rows, []any{&v}, func() error { return nil }) + if err != nil { + b.Fatal(err) + } + } + }) +} +<% end %> +<% end %> diff --git a/vendor/github.com/jackc/pgx/v5/pgtype/integration_benchmark_test_gen.sh b/vendor/github.com/jackc/pgx/v5/pgtype/integration_benchmark_test_gen.sh new file mode 100644 index 0000000000..22ac01aaf4 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgtype/integration_benchmark_test_gen.sh @@ -0,0 +1,2 @@ +erb integration_benchmark_test.go.erb > integration_benchmark_test.go +goimports -w integration_benchmark_test.go diff --git a/vendor/github.com/jackc/pgx/v5/pgtype/interval.go b/vendor/github.com/jackc/pgx/v5/pgtype/interval.go new file mode 100644 index 0000000000..be8decdd12 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgtype/interval.go @@ -0,0 +1,297 @@ +package pgtype + +import ( + "database/sql/driver" + "encoding/binary" + "fmt" + "strconv" + "strings" + + "github.com/jackc/pgx/v5/internal/pgio" +) + +const ( + microsecondsPerSecond = 1_000_000 + microsecondsPerMinute = 60 * microsecondsPerSecond + microsecondsPerHour = 60 * microsecondsPerMinute + microsecondsPerDay = 24 * microsecondsPerHour + microsecondsPerMonth = 30 * microsecondsPerDay +) + +type IntervalScanner interface { + ScanInterval(v Interval) error +} + +type IntervalValuer interface { + IntervalValue() (Interval, error) +} + +type Interval struct { + Microseconds int64 + Days int32 + Months int32 + Valid bool +} + +// ScanInterval implements the [IntervalScanner] interface. +func (interval *Interval) ScanInterval(v Interval) error { + *interval = v + return nil +} + +// IntervalValue implements the [IntervalValuer] interface. +func (interval Interval) IntervalValue() (Interval, error) { + return interval, nil +} + +// Scan implements the [database/sql.Scanner] interface. +func (interval *Interval) Scan(src any) error { + if src == nil { + *interval = Interval{} + return nil + } + + if src, ok := src.(string); ok { + return scanPlanTextAnyToIntervalScanner{}.Scan([]byte(src), interval) + } + + return fmt.Errorf("cannot scan %T", src) +} + +// Value implements the [database/sql/driver.Valuer] interface. +func (interval Interval) Value() (driver.Value, error) { + if !interval.Valid { + return nil, nil + } + + buf, err := IntervalCodec{}.PlanEncode(nil, 0, TextFormatCode, interval).Encode(interval, nil) + if err != nil { + return nil, err + } + return string(buf), err +} + +type IntervalCodec struct{} + +func (IntervalCodec) FormatSupported(format int16) bool { + return format == TextFormatCode || format == BinaryFormatCode +} + +func (IntervalCodec) PreferredFormat() int16 { + return BinaryFormatCode +} + +func (IntervalCodec) PlanEncode(m *Map, oid uint32, format int16, value any) EncodePlan { + if _, ok := value.(IntervalValuer); !ok { + return nil + } + + switch format { + case BinaryFormatCode: + return encodePlanIntervalCodecBinary{} + case TextFormatCode: + return encodePlanIntervalCodecText{} + } + + return nil +} + +type encodePlanIntervalCodecBinary struct{} + +func (encodePlanIntervalCodecBinary) Encode(value any, buf []byte) (newBuf []byte, err error) { + interval, err := value.(IntervalValuer).IntervalValue() + if err != nil { + return nil, err + } + + if !interval.Valid { + return nil, nil + } + + buf = pgio.AppendInt64(buf, interval.Microseconds) + buf = pgio.AppendInt32(buf, interval.Days) + buf = pgio.AppendInt32(buf, interval.Months) + return buf, nil +} + +type encodePlanIntervalCodecText struct{} + +func (encodePlanIntervalCodecText) Encode(value any, buf []byte) (newBuf []byte, err error) { + interval, err := value.(IntervalValuer).IntervalValue() + if err != nil { + return nil, err + } + + if !interval.Valid { + return nil, nil + } + + if interval.Months != 0 { + buf = append(buf, strconv.FormatInt(int64(interval.Months), 10)...) + buf = append(buf, " mon "...) + } + + if interval.Days != 0 { + buf = append(buf, strconv.FormatInt(int64(interval.Days), 10)...) + buf = append(buf, " day "...) + } + + absMicroseconds := interval.Microseconds + if absMicroseconds < 0 { + absMicroseconds = -absMicroseconds + buf = append(buf, '-') + } + + hours := absMicroseconds / microsecondsPerHour + minutes := (absMicroseconds % microsecondsPerHour) / microsecondsPerMinute + seconds := (absMicroseconds % microsecondsPerMinute) / microsecondsPerSecond + + timeStr := fmt.Sprintf("%02d:%02d:%02d", hours, minutes, seconds) + buf = append(buf, timeStr...) + + microseconds := absMicroseconds % microsecondsPerSecond + if microseconds != 0 { + buf = append(buf, fmt.Sprintf(".%06d", microseconds)...) + } + + return buf, nil +} + +func (IntervalCodec) PlanScan(m *Map, oid uint32, format int16, target any) ScanPlan { + switch format { + case BinaryFormatCode: + if _, ok := target.(IntervalScanner); ok { + return scanPlanBinaryIntervalToIntervalScanner{} + } + case TextFormatCode: + if _, ok := target.(IntervalScanner); ok { + return scanPlanTextAnyToIntervalScanner{} + } + } + + return nil +} + +type scanPlanBinaryIntervalToIntervalScanner struct{} + +func (scanPlanBinaryIntervalToIntervalScanner) Scan(src []byte, dst any) error { + scanner := (dst).(IntervalScanner) + + if src == nil { + return scanner.ScanInterval(Interval{}) + } + + if len(src) != 16 { + return fmt.Errorf("Received an invalid size for an interval: %d", len(src)) + } + + microseconds := int64(binary.BigEndian.Uint64(src)) + days := int32(binary.BigEndian.Uint32(src[8:])) + months := int32(binary.BigEndian.Uint32(src[12:])) + + return scanner.ScanInterval(Interval{Microseconds: microseconds, Days: days, Months: months, Valid: true}) +} + +type scanPlanTextAnyToIntervalScanner struct{} + +func (scanPlanTextAnyToIntervalScanner) Scan(src []byte, dst any) error { + scanner := (dst).(IntervalScanner) + + if src == nil { + return scanner.ScanInterval(Interval{}) + } + + var microseconds int64 + var days int32 + var months int32 + + parts := strings.Split(string(src), " ") + + for i := 0; i < len(parts)-1; i += 2 { + scalar, err := strconv.ParseInt(parts[i], 10, 64) + if err != nil { + return fmt.Errorf("bad interval format") + } + + switch parts[i+1] { + case "year", "years": + months += int32(scalar * 12) + case "mon", "mons": + months += int32(scalar) + case "day", "days": + days = int32(scalar) + default: + return fmt.Errorf("bad interval format: %q", parts[i+1]) + } + } + + if len(parts)%2 == 1 { + timeParts := strings.SplitN(parts[len(parts)-1], ":", 3) + if len(timeParts) != 3 { + return fmt.Errorf("bad interval format") + } + + var negative bool + if timeParts[0][0] == '-' { + negative = true + timeParts[0] = timeParts[0][1:] + } + + hours, err := strconv.ParseInt(timeParts[0], 10, 64) + if err != nil { + return fmt.Errorf("bad interval hour format: %s", timeParts[0]) + } + + minutes, err := strconv.ParseInt(timeParts[1], 10, 64) + if err != nil { + return fmt.Errorf("bad interval minute format: %s", timeParts[1]) + } + + sec, secFrac, secFracFound := strings.Cut(timeParts[2], ".") + + seconds, err := strconv.ParseInt(sec, 10, 64) + if err != nil { + return fmt.Errorf("bad interval second format: %s", sec) + } + + var uSeconds int64 + if secFracFound { + uSeconds, err = strconv.ParseInt(secFrac, 10, 64) + if err != nil { + return fmt.Errorf("bad interval decimal format: %s", secFrac) + } + + for i := 0; i < 6-len(secFrac); i++ { + uSeconds *= 10 + } + } + + microseconds = hours * microsecondsPerHour + microseconds += minutes * microsecondsPerMinute + microseconds += seconds * microsecondsPerSecond + microseconds += uSeconds + + if negative { + microseconds = -microseconds + } + } + + return scanner.ScanInterval(Interval{Months: months, Days: days, Microseconds: microseconds, Valid: true}) +} + +func (c IntervalCodec) DecodeDatabaseSQLValue(m *Map, oid uint32, format int16, src []byte) (driver.Value, error) { + return codecDecodeToTextFormat(c, m, oid, format, src) +} + +func (c IntervalCodec) DecodeValue(m *Map, oid uint32, format int16, src []byte) (any, error) { + if src == nil { + return nil, nil + } + + var interval Interval + err := codecScan(c, m, oid, format, src, &interval) + if err != nil { + return nil, err + } + return interval, nil +} diff --git a/vendor/github.com/jackc/pgx/v5/pgtype/json.go b/vendor/github.com/jackc/pgx/v5/pgtype/json.go new file mode 100644 index 0000000000..a5f74eaed5 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgtype/json.go @@ -0,0 +1,243 @@ +package pgtype + +import ( + "database/sql" + "database/sql/driver" + "encoding/json" + "fmt" + "reflect" +) + +type JSONCodec struct { + Marshal func(v any) ([]byte, error) + Unmarshal func(data []byte, v any) error +} + +func (*JSONCodec) FormatSupported(format int16) bool { + return format == TextFormatCode || format == BinaryFormatCode +} + +func (*JSONCodec) PreferredFormat() int16 { + return TextFormatCode +} + +func (c *JSONCodec) PlanEncode(m *Map, oid uint32, format int16, value any) EncodePlan { + switch value.(type) { + case string: + return encodePlanJSONCodecEitherFormatString{} + case []byte: + return encodePlanJSONCodecEitherFormatByteSlice{} + + // Handle json.RawMessage specifically because if it is run through json.Marshal it may be mutated. + // e.g. `{"foo": "bar"}` -> `{"foo":"bar"}`. + case json.RawMessage: + return encodePlanJSONCodecEitherFormatJSONRawMessage{} + + // Cannot rely on driver.Valuer being handled later because anything can be marshalled. + // + // https://github.com/jackc/pgx/issues/1430 + // + // Check for driver.Valuer must come before json.Marshaler so that it is guaranteed to be used + // when both are implemented https://github.com/jackc/pgx/issues/1805 + case driver.Valuer: + return &encodePlanDriverValuer{m: m, oid: oid, formatCode: format} + + // Must come before trying wrap encode plans because a pointer to a struct may be unwrapped to a struct that can be + // marshalled. + // + // https://github.com/jackc/pgx/issues/1681 + case json.Marshaler: + return &encodePlanJSONCodecEitherFormatMarshal{ + marshal: c.Marshal, + } + } + + // Because anything can be marshalled the normal wrapping in Map.PlanScan doesn't get a chance to run. So try the + // appropriate wrappers here. + for _, f := range []TryWrapEncodePlanFunc{ + TryWrapDerefPointerEncodePlan, + TryWrapFindUnderlyingTypeEncodePlan, + } { + if wrapperPlan, nextValue, ok := f(value); ok { + if nextPlan := c.PlanEncode(m, oid, format, nextValue); nextPlan != nil { + wrapperPlan.SetNext(nextPlan) + return wrapperPlan + } + } + } + + return &encodePlanJSONCodecEitherFormatMarshal{ + marshal: c.Marshal, + } +} + +// JSON needs its on scan plan for pointers to handle 'null'::json(b). +// Consider making pointerPointerScanPlan more flexible in the future. +type jsonPointerScanPlan struct { + next ScanPlan +} + +func (p jsonPointerScanPlan) Scan(src []byte, dst any) error { + el := reflect.ValueOf(dst).Elem() + if src == nil || string(src) == "null" { + el.SetZero() + return nil + } + + el.Set(reflect.New(el.Type().Elem())) + if p.next != nil { + return p.next.Scan(src, el.Interface()) + } + + return nil +} + +type encodePlanJSONCodecEitherFormatString struct{} + +func (encodePlanJSONCodecEitherFormatString) Encode(value any, buf []byte) (newBuf []byte, err error) { + jsonString := value.(string) + buf = append(buf, jsonString...) + return buf, nil +} + +type encodePlanJSONCodecEitherFormatByteSlice struct{} + +func (encodePlanJSONCodecEitherFormatByteSlice) Encode(value any, buf []byte) (newBuf []byte, err error) { + jsonBytes := value.([]byte) + if jsonBytes == nil { + return nil, nil + } + + buf = append(buf, jsonBytes...) + return buf, nil +} + +type encodePlanJSONCodecEitherFormatJSONRawMessage struct{} + +func (encodePlanJSONCodecEitherFormatJSONRawMessage) Encode(value any, buf []byte) (newBuf []byte, err error) { + jsonBytes := value.(json.RawMessage) + if jsonBytes == nil { + return nil, nil + } + + buf = append(buf, jsonBytes...) + return buf, nil +} + +type encodePlanJSONCodecEitherFormatMarshal struct { + marshal func(v any) ([]byte, error) +} + +func (e *encodePlanJSONCodecEitherFormatMarshal) Encode(value any, buf []byte) (newBuf []byte, err error) { + jsonBytes, err := e.marshal(value) + if err != nil { + return nil, err + } + + buf = append(buf, jsonBytes...) + return buf, nil +} + +func (c *JSONCodec) PlanScan(m *Map, oid uint32, formatCode int16, target any) ScanPlan { + return c.planScan(m, oid, formatCode, target, 0) +} + +// JSON cannot fallback to pointerPointerScanPlan because of 'null'::json(b), +// so we need to duplicate the logic here. +func (c *JSONCodec) planScan(m *Map, oid uint32, formatCode int16, target any, depth int) ScanPlan { + if depth > 8 { + return &scanPlanFail{m: m, oid: oid, formatCode: formatCode} + } + + switch target.(type) { + case *string: + return &scanPlanAnyToString{} + case *[]byte: + return &scanPlanJSONToByteSlice{} + case BytesScanner: + return &scanPlanBinaryBytesToBytesScanner{} + case sql.Scanner: + return &scanPlanCodecSQLScanner{c: c, m: m, oid: oid, formatCode: formatCode} + } + + rv := reflect.ValueOf(target) + if rv.Kind() == reflect.Pointer && rv.Elem().Kind() == reflect.Pointer { + var plan jsonPointerScanPlan + plan.next = c.planScan(m, oid, formatCode, rv.Elem().Interface(), depth+1) + return plan + } else { + return &scanPlanJSONToJSONUnmarshal{unmarshal: c.Unmarshal} + } +} + +type scanPlanAnyToString struct{} + +func (scanPlanAnyToString) Scan(src []byte, dst any) error { + p := dst.(*string) + *p = string(src) + return nil +} + +type scanPlanJSONToByteSlice struct{} + +func (scanPlanJSONToByteSlice) Scan(src []byte, dst any) error { + dstBuf := dst.(*[]byte) + if src == nil { + *dstBuf = nil + return nil + } + + *dstBuf = make([]byte, len(src)) + copy(*dstBuf, src) + return nil +} + +type scanPlanJSONToJSONUnmarshal struct { + unmarshal func(data []byte, v any) error +} + +func (s *scanPlanJSONToJSONUnmarshal) Scan(src []byte, dst any) error { + if src == nil { + dstValue := reflect.ValueOf(dst) + if dstValue.Kind() == reflect.Pointer { + el := dstValue.Elem() + switch el.Kind() { + case reflect.Pointer, reflect.Slice, reflect.Map, reflect.Interface: + el.Set(reflect.Zero(el.Type())) + return nil + } + } + + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + v := reflect.ValueOf(dst) + if v.Kind() != reflect.Pointer || v.IsNil() { + return fmt.Errorf("cannot scan into non-pointer or nil destinations %T", dst) + } + + elem := v.Elem() + elem.Set(reflect.Zero(elem.Type())) + + return s.unmarshal(src, dst) +} + +func (c *JSONCodec) DecodeDatabaseSQLValue(m *Map, oid uint32, format int16, src []byte) (driver.Value, error) { + if src == nil { + return nil, nil + } + + dstBuf := make([]byte, len(src)) + copy(dstBuf, src) + return dstBuf, nil +} + +func (c *JSONCodec) DecodeValue(m *Map, oid uint32, format int16, src []byte) (any, error) { + if src == nil { + return nil, nil + } + + var dst any + err := c.Unmarshal(src, &dst) + return dst, err +} diff --git a/vendor/github.com/jackc/pgx/v5/pgtype/jsonb.go b/vendor/github.com/jackc/pgx/v5/pgtype/jsonb.go new file mode 100644 index 0000000000..4d4eb58e5b --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgtype/jsonb.go @@ -0,0 +1,129 @@ +package pgtype + +import ( + "database/sql/driver" + "fmt" +) + +type JSONBCodec struct { + Marshal func(v any) ([]byte, error) + Unmarshal func(data []byte, v any) error +} + +func (*JSONBCodec) FormatSupported(format int16) bool { + return format == TextFormatCode || format == BinaryFormatCode +} + +func (*JSONBCodec) PreferredFormat() int16 { + return TextFormatCode +} + +func (c *JSONBCodec) PlanEncode(m *Map, oid uint32, format int16, value any) EncodePlan { + switch format { + case BinaryFormatCode: + plan := (&JSONCodec{Marshal: c.Marshal, Unmarshal: c.Unmarshal}).PlanEncode(m, oid, TextFormatCode, value) + if plan != nil { + return &encodePlanJSONBCodecBinaryWrapper{textPlan: plan} + } + case TextFormatCode: + return (&JSONCodec{Marshal: c.Marshal, Unmarshal: c.Unmarshal}).PlanEncode(m, oid, format, value) + } + + return nil +} + +type encodePlanJSONBCodecBinaryWrapper struct { + textPlan EncodePlan +} + +func (plan *encodePlanJSONBCodecBinaryWrapper) Encode(value any, buf []byte) (newBuf []byte, err error) { + buf = append(buf, 1) + return plan.textPlan.Encode(value, buf) +} + +func (c *JSONBCodec) PlanScan(m *Map, oid uint32, format int16, target any) ScanPlan { + switch format { + case BinaryFormatCode: + plan := (&JSONCodec{Marshal: c.Marshal, Unmarshal: c.Unmarshal}).PlanScan(m, oid, TextFormatCode, target) + if plan != nil { + return &scanPlanJSONBCodecBinaryUnwrapper{textPlan: plan} + } + case TextFormatCode: + return (&JSONCodec{Marshal: c.Marshal, Unmarshal: c.Unmarshal}).PlanScan(m, oid, format, target) + } + + return nil +} + +type scanPlanJSONBCodecBinaryUnwrapper struct { + textPlan ScanPlan +} + +func (plan *scanPlanJSONBCodecBinaryUnwrapper) Scan(src []byte, dst any) error { + if src == nil { + return plan.textPlan.Scan(src, dst) + } + + if len(src) == 0 { + return fmt.Errorf("jsonb too short") + } + + if src[0] != 1 { + return fmt.Errorf("unknown jsonb version number %d", src[0]) + } + + return plan.textPlan.Scan(src[1:], dst) +} + +func (c *JSONBCodec) DecodeDatabaseSQLValue(m *Map, oid uint32, format int16, src []byte) (driver.Value, error) { + if src == nil { + return nil, nil + } + + switch format { + case BinaryFormatCode: + if len(src) == 0 { + return nil, fmt.Errorf("jsonb too short") + } + + if src[0] != 1 { + return nil, fmt.Errorf("unknown jsonb version number %d", src[0]) + } + + dstBuf := make([]byte, len(src)-1) + copy(dstBuf, src[1:]) + return dstBuf, nil + case TextFormatCode: + dstBuf := make([]byte, len(src)) + copy(dstBuf, src) + return dstBuf, nil + default: + return nil, fmt.Errorf("unknown format code: %v", format) + } +} + +func (c *JSONBCodec) DecodeValue(m *Map, oid uint32, format int16, src []byte) (any, error) { + if src == nil { + return nil, nil + } + + switch format { + case BinaryFormatCode: + if len(src) == 0 { + return nil, fmt.Errorf("jsonb too short") + } + + if src[0] != 1 { + return nil, fmt.Errorf("unknown jsonb version number %d", src[0]) + } + + src = src[1:] + case TextFormatCode: + default: + return nil, fmt.Errorf("unknown format code: %v", format) + } + + var dst any + err := c.Unmarshal(src, &dst) + return dst, err +} diff --git a/vendor/github.com/jackc/pgx/v5/pgtype/line.go b/vendor/github.com/jackc/pgx/v5/pgtype/line.go new file mode 100644 index 0000000000..73b0636494 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgtype/line.go @@ -0,0 +1,223 @@ +package pgtype + +import ( + "database/sql/driver" + "encoding/binary" + "fmt" + "math" + "strconv" + "strings" + + "github.com/jackc/pgx/v5/internal/pgio" +) + +type LineScanner interface { + ScanLine(v Line) error +} + +type LineValuer interface { + LineValue() (Line, error) +} + +type Line struct { + A, B, C float64 + Valid bool +} + +// ScanLine implements the [LineScanner] interface. +func (line *Line) ScanLine(v Line) error { + *line = v + return nil +} + +// LineValue implements the [LineValuer] interface. +func (line Line) LineValue() (Line, error) { + return line, nil +} + +func (line *Line) Set(src any) error { + return fmt.Errorf("cannot convert %v to Line", src) +} + +// Scan implements the [database/sql.Scanner] interface. +func (line *Line) Scan(src any) error { + if src == nil { + *line = Line{} + return nil + } + + if src, ok := src.(string); ok { + return scanPlanTextAnyToLineScanner{}.Scan([]byte(src), line) + } + + return fmt.Errorf("cannot scan %T", src) +} + +// Value implements the [database/sql/driver.Valuer] interface. +func (line Line) Value() (driver.Value, error) { + if !line.Valid { + return nil, nil + } + + buf, err := LineCodec{}.PlanEncode(nil, 0, TextFormatCode, line).Encode(line, nil) + if err != nil { + return nil, err + } + return string(buf), err +} + +type LineCodec struct{} + +func (LineCodec) FormatSupported(format int16) bool { + return format == TextFormatCode || format == BinaryFormatCode +} + +func (LineCodec) PreferredFormat() int16 { + return BinaryFormatCode +} + +func (LineCodec) PlanEncode(m *Map, oid uint32, format int16, value any) EncodePlan { + if _, ok := value.(LineValuer); !ok { + return nil + } + + switch format { + case BinaryFormatCode: + return encodePlanLineCodecBinary{} + case TextFormatCode: + return encodePlanLineCodecText{} + } + + return nil +} + +type encodePlanLineCodecBinary struct{} + +func (encodePlanLineCodecBinary) Encode(value any, buf []byte) (newBuf []byte, err error) { + line, err := value.(LineValuer).LineValue() + if err != nil { + return nil, err + } + + if !line.Valid { + return nil, nil + } + + buf = pgio.AppendUint64(buf, math.Float64bits(line.A)) + buf = pgio.AppendUint64(buf, math.Float64bits(line.B)) + buf = pgio.AppendUint64(buf, math.Float64bits(line.C)) + return buf, nil +} + +type encodePlanLineCodecText struct{} + +func (encodePlanLineCodecText) Encode(value any, buf []byte) (newBuf []byte, err error) { + line, err := value.(LineValuer).LineValue() + if err != nil { + return nil, err + } + + if !line.Valid { + return nil, nil + } + + buf = append(buf, fmt.Sprintf(`{%s,%s,%s}`, + strconv.FormatFloat(line.A, 'f', -1, 64), + strconv.FormatFloat(line.B, 'f', -1, 64), + strconv.FormatFloat(line.C, 'f', -1, 64), + )...) + return buf, nil +} + +func (LineCodec) PlanScan(m *Map, oid uint32, format int16, target any) ScanPlan { + switch format { + case BinaryFormatCode: + if _, ok := target.(LineScanner); ok { + return scanPlanBinaryLineToLineScanner{} + } + case TextFormatCode: + if _, ok := target.(LineScanner); ok { + return scanPlanTextAnyToLineScanner{} + } + } + + return nil +} + +type scanPlanBinaryLineToLineScanner struct{} + +func (scanPlanBinaryLineToLineScanner) Scan(src []byte, dst any) error { + scanner := (dst).(LineScanner) + + if src == nil { + return scanner.ScanLine(Line{}) + } + + if len(src) != 24 { + return fmt.Errorf("invalid length for line: %v", len(src)) + } + + a := binary.BigEndian.Uint64(src) + b := binary.BigEndian.Uint64(src[8:]) + c := binary.BigEndian.Uint64(src[16:]) + + return scanner.ScanLine(Line{ + A: math.Float64frombits(a), + B: math.Float64frombits(b), + C: math.Float64frombits(c), + Valid: true, + }) +} + +type scanPlanTextAnyToLineScanner struct{} + +func (scanPlanTextAnyToLineScanner) Scan(src []byte, dst any) error { + scanner := (dst).(LineScanner) + + if src == nil { + return scanner.ScanLine(Line{}) + } + + if len(src) < 7 { + return fmt.Errorf("invalid length for line: %v", len(src)) + } + + parts := strings.SplitN(string(src[1:len(src)-1]), ",", 3) + if len(parts) < 3 { + return fmt.Errorf("invalid format for line") + } + + a, err := strconv.ParseFloat(parts[0], 64) + if err != nil { + return err + } + + b, err := strconv.ParseFloat(parts[1], 64) + if err != nil { + return err + } + + c, err := strconv.ParseFloat(parts[2], 64) + if err != nil { + return err + } + + return scanner.ScanLine(Line{A: a, B: b, C: c, Valid: true}) +} + +func (c LineCodec) DecodeDatabaseSQLValue(m *Map, oid uint32, format int16, src []byte) (driver.Value, error) { + return codecDecodeToTextFormat(c, m, oid, format, src) +} + +func (c LineCodec) DecodeValue(m *Map, oid uint32, format int16, src []byte) (any, error) { + if src == nil { + return nil, nil + } + + var line Line + err := codecScan(c, m, oid, format, src, &line) + if err != nil { + return nil, err + } + return line, nil +} diff --git a/vendor/github.com/jackc/pgx/v5/pgtype/lseg.go b/vendor/github.com/jackc/pgx/v5/pgtype/lseg.go new file mode 100644 index 0000000000..438b45b622 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgtype/lseg.go @@ -0,0 +1,235 @@ +package pgtype + +import ( + "database/sql/driver" + "encoding/binary" + "fmt" + "math" + "strconv" + "strings" + + "github.com/jackc/pgx/v5/internal/pgio" +) + +type LsegScanner interface { + ScanLseg(v Lseg) error +} + +type LsegValuer interface { + LsegValue() (Lseg, error) +} + +type Lseg struct { + P [2]Vec2 + Valid bool +} + +// ScanLseg implements the [LsegScanner] interface. +func (lseg *Lseg) ScanLseg(v Lseg) error { + *lseg = v + return nil +} + +// LsegValue implements the [LsegValuer] interface. +func (lseg Lseg) LsegValue() (Lseg, error) { + return lseg, nil +} + +// Scan implements the [database/sql.Scanner] interface. +func (lseg *Lseg) Scan(src any) error { + if src == nil { + *lseg = Lseg{} + return nil + } + + if src, ok := src.(string); ok { + return scanPlanTextAnyToLsegScanner{}.Scan([]byte(src), lseg) + } + + return fmt.Errorf("cannot scan %T", src) +} + +// Value implements the [database/sql/driver.Valuer] interface. +func (lseg Lseg) Value() (driver.Value, error) { + if !lseg.Valid { + return nil, nil + } + + buf, err := LsegCodec{}.PlanEncode(nil, 0, TextFormatCode, lseg).Encode(lseg, nil) + if err != nil { + return nil, err + } + return string(buf), err +} + +type LsegCodec struct{} + +func (LsegCodec) FormatSupported(format int16) bool { + return format == TextFormatCode || format == BinaryFormatCode +} + +func (LsegCodec) PreferredFormat() int16 { + return BinaryFormatCode +} + +func (LsegCodec) PlanEncode(m *Map, oid uint32, format int16, value any) EncodePlan { + if _, ok := value.(LsegValuer); !ok { + return nil + } + + switch format { + case BinaryFormatCode: + return encodePlanLsegCodecBinary{} + case TextFormatCode: + return encodePlanLsegCodecText{} + } + + return nil +} + +type encodePlanLsegCodecBinary struct{} + +func (encodePlanLsegCodecBinary) Encode(value any, buf []byte) (newBuf []byte, err error) { + lseg, err := value.(LsegValuer).LsegValue() + if err != nil { + return nil, err + } + + if !lseg.Valid { + return nil, nil + } + + buf = pgio.AppendUint64(buf, math.Float64bits(lseg.P[0].X)) + buf = pgio.AppendUint64(buf, math.Float64bits(lseg.P[0].Y)) + buf = pgio.AppendUint64(buf, math.Float64bits(lseg.P[1].X)) + buf = pgio.AppendUint64(buf, math.Float64bits(lseg.P[1].Y)) + return buf, nil +} + +type encodePlanLsegCodecText struct{} + +func (encodePlanLsegCodecText) Encode(value any, buf []byte) (newBuf []byte, err error) { + lseg, err := value.(LsegValuer).LsegValue() + if err != nil { + return nil, err + } + + if !lseg.Valid { + return nil, nil + } + + buf = append(buf, fmt.Sprintf(`[(%s,%s),(%s,%s)]`, + strconv.FormatFloat(lseg.P[0].X, 'f', -1, 64), + strconv.FormatFloat(lseg.P[0].Y, 'f', -1, 64), + strconv.FormatFloat(lseg.P[1].X, 'f', -1, 64), + strconv.FormatFloat(lseg.P[1].Y, 'f', -1, 64), + )...) + return buf, nil +} + +func (LsegCodec) PlanScan(m *Map, oid uint32, format int16, target any) ScanPlan { + switch format { + case BinaryFormatCode: + if _, ok := target.(LsegScanner); ok { + return scanPlanBinaryLsegToLsegScanner{} + } + case TextFormatCode: + if _, ok := target.(LsegScanner); ok { + return scanPlanTextAnyToLsegScanner{} + } + } + + return nil +} + +type scanPlanBinaryLsegToLsegScanner struct{} + +func (scanPlanBinaryLsegToLsegScanner) Scan(src []byte, dst any) error { + scanner := (dst).(LsegScanner) + + if src == nil { + return scanner.ScanLseg(Lseg{}) + } + + if len(src) != 32 { + return fmt.Errorf("invalid length for lseg: %v", len(src)) + } + + x1 := binary.BigEndian.Uint64(src) + y1 := binary.BigEndian.Uint64(src[8:]) + x2 := binary.BigEndian.Uint64(src[16:]) + y2 := binary.BigEndian.Uint64(src[24:]) + + return scanner.ScanLseg(Lseg{ + P: [2]Vec2{ + {math.Float64frombits(x1), math.Float64frombits(y1)}, + {math.Float64frombits(x2), math.Float64frombits(y2)}, + }, + Valid: true, + }) +} + +type scanPlanTextAnyToLsegScanner struct{} + +func (scanPlanTextAnyToLsegScanner) Scan(src []byte, dst any) error { + scanner := (dst).(LsegScanner) + + if src == nil { + return scanner.ScanLseg(Lseg{}) + } + + if len(src) < 11 { + return fmt.Errorf("invalid length for lseg: %v", len(src)) + } + + // Expected format: [(x1,y1),(x2,y2)] + sp1, sp2, found := strings.Cut(string(src[2:len(src)-2]), "),(") + if !found { + return fmt.Errorf("invalid format for lseg") + } + + sx1, sy1, found := strings.Cut(sp1, ",") + if !found { + return fmt.Errorf("invalid format for lseg") + } + sx2, sy2, found := strings.Cut(sp2, ",") + if !found { + return fmt.Errorf("invalid format for lseg") + } + + x1, err := strconv.ParseFloat(sx1, 64) + if err != nil { + return err + } + y1, err := strconv.ParseFloat(sy1, 64) + if err != nil { + return err + } + x2, err := strconv.ParseFloat(sx2, 64) + if err != nil { + return err + } + y2, err := strconv.ParseFloat(sy2, 64) + if err != nil { + return err + } + + return scanner.ScanLseg(Lseg{P: [2]Vec2{{x1, y1}, {x2, y2}}, Valid: true}) +} + +func (c LsegCodec) DecodeDatabaseSQLValue(m *Map, oid uint32, format int16, src []byte) (driver.Value, error) { + return codecDecodeToTextFormat(c, m, oid, format, src) +} + +func (c LsegCodec) DecodeValue(m *Map, oid uint32, format int16, src []byte) (any, error) { + if src == nil { + return nil, nil + } + + var lseg Lseg + err := codecScan(c, m, oid, format, src, &lseg) + if err != nil { + return nil, err + } + return lseg, nil +} diff --git a/vendor/github.com/jackc/pgx/v5/pgtype/ltree.go b/vendor/github.com/jackc/pgx/v5/pgtype/ltree.go new file mode 100644 index 0000000000..6af3177944 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgtype/ltree.go @@ -0,0 +1,122 @@ +package pgtype + +import ( + "database/sql/driver" + "fmt" +) + +type LtreeCodec struct{} + +func (l LtreeCodec) FormatSupported(format int16) bool { + return format == TextFormatCode || format == BinaryFormatCode +} + +// PreferredFormat returns the preferred format. +func (l LtreeCodec) PreferredFormat() int16 { + return TextFormatCode +} + +// PlanEncode returns an EncodePlan for encoding value into PostgreSQL format for oid and format. If no plan can be +// found then nil is returned. +func (l LtreeCodec) PlanEncode(m *Map, oid uint32, format int16, value any) EncodePlan { + switch format { + case TextFormatCode: + return (TextCodec)(l).PlanEncode(m, oid, format, value) + case BinaryFormatCode: + switch value.(type) { + case string: + return encodeLtreeCodecBinaryString{} + case []byte: + return encodeLtreeCodecBinaryByteSlice{} + case TextValuer: + return encodeLtreeCodecBinaryTextValuer{} + } + } + + return nil +} + +type encodeLtreeCodecBinaryString struct{} + +func (encodeLtreeCodecBinaryString) Encode(value any, buf []byte) (newBuf []byte, err error) { + ltree := value.(string) + buf = append(buf, 1) + return append(buf, ltree...), nil +} + +type encodeLtreeCodecBinaryByteSlice struct{} + +func (encodeLtreeCodecBinaryByteSlice) Encode(value any, buf []byte) (newBuf []byte, err error) { + ltree := value.([]byte) + buf = append(buf, 1) + return append(buf, ltree...), nil +} + +type encodeLtreeCodecBinaryTextValuer struct{} + +func (encodeLtreeCodecBinaryTextValuer) Encode(value any, buf []byte) (newBuf []byte, err error) { + t, err := value.(TextValuer).TextValue() + if err != nil { + return nil, err + } + if !t.Valid { + return nil, nil + } + + buf = append(buf, 1) + return append(buf, t.String...), nil +} + +// PlanScan returns a ScanPlan for scanning a PostgreSQL value into a destination with the same type as target. If +// no plan can be found then nil is returned. +func (l LtreeCodec) PlanScan(m *Map, oid uint32, format int16, target any) ScanPlan { + switch format { + case TextFormatCode: + return (TextCodec)(l).PlanScan(m, oid, format, target) + case BinaryFormatCode: + switch target.(type) { + case *string: + return scanPlanBinaryLtreeToString{} + case TextScanner: + return scanPlanBinaryLtreeToTextScanner{} + } + } + + return nil +} + +type scanPlanBinaryLtreeToString struct{} + +func (scanPlanBinaryLtreeToString) Scan(src []byte, target any) error { + version := src[0] + if version != 1 { + return fmt.Errorf("unsupported ltree version %d", version) + } + + p := (target).(*string) + *p = string(src[1:]) + + return nil +} + +type scanPlanBinaryLtreeToTextScanner struct{} + +func (scanPlanBinaryLtreeToTextScanner) Scan(src []byte, target any) error { + version := src[0] + if version != 1 { + return fmt.Errorf("unsupported ltree version %d", version) + } + + scanner := (target).(TextScanner) + return scanner.ScanText(Text{String: string(src[1:]), Valid: true}) +} + +// DecodeDatabaseSQLValue returns src decoded into a value compatible with the sql.Scanner interface. +func (l LtreeCodec) DecodeDatabaseSQLValue(m *Map, oid uint32, format int16, src []byte) (driver.Value, error) { + return (TextCodec)(l).DecodeDatabaseSQLValue(m, oid, format, src) +} + +// DecodeValue returns src decoded into its default format. +func (l LtreeCodec) DecodeValue(m *Map, oid uint32, format int16, src []byte) (any, error) { + return (TextCodec)(l).DecodeValue(m, oid, format, src) +} diff --git a/vendor/github.com/jackc/pgx/v5/pgtype/macaddr.go b/vendor/github.com/jackc/pgx/v5/pgtype/macaddr.go new file mode 100644 index 0000000000..e913ec9034 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgtype/macaddr.go @@ -0,0 +1,162 @@ +package pgtype + +import ( + "database/sql/driver" + "net" +) + +type MacaddrCodec struct{} + +func (MacaddrCodec) FormatSupported(format int16) bool { + return format == TextFormatCode || format == BinaryFormatCode +} + +func (MacaddrCodec) PreferredFormat() int16 { + return BinaryFormatCode +} + +func (MacaddrCodec) PlanEncode(m *Map, oid uint32, format int16, value any) EncodePlan { + switch format { + case BinaryFormatCode: + switch value.(type) { + case net.HardwareAddr: + return encodePlanMacaddrCodecBinaryHardwareAddr{} + case TextValuer: + return encodePlanMacAddrCodecTextValuer{} + + } + case TextFormatCode: + switch value.(type) { + case net.HardwareAddr: + return encodePlanMacaddrCodecTextHardwareAddr{} + case TextValuer: + return encodePlanTextCodecTextValuer{} + } + } + + return nil +} + +type encodePlanMacaddrCodecBinaryHardwareAddr struct{} + +func (encodePlanMacaddrCodecBinaryHardwareAddr) Encode(value any, buf []byte) (newBuf []byte, err error) { + addr := value.(net.HardwareAddr) + if addr == nil { + return nil, nil + } + + return append(buf, addr...), nil +} + +type encodePlanMacAddrCodecTextValuer struct{} + +func (encodePlanMacAddrCodecTextValuer) Encode(value any, buf []byte) (newBuf []byte, err error) { + t, err := value.(TextValuer).TextValue() + if err != nil { + return nil, err + } + if !t.Valid { + return nil, nil + } + + addr, err := net.ParseMAC(t.String) + if err != nil { + return nil, err + } + + return append(buf, addr...), nil +} + +type encodePlanMacaddrCodecTextHardwareAddr struct{} + +func (encodePlanMacaddrCodecTextHardwareAddr) Encode(value any, buf []byte) (newBuf []byte, err error) { + addr := value.(net.HardwareAddr) + if addr == nil { + return nil, nil + } + + return append(buf, addr.String()...), nil +} + +func (MacaddrCodec) PlanScan(m *Map, oid uint32, format int16, target any) ScanPlan { + switch format { + case BinaryFormatCode: + switch target.(type) { + case *net.HardwareAddr: + return scanPlanBinaryMacaddrToHardwareAddr{} + case TextScanner: + return scanPlanBinaryMacaddrToTextScanner{} + } + case TextFormatCode: + switch target.(type) { + case *net.HardwareAddr: + return scanPlanTextMacaddrToHardwareAddr{} + case TextScanner: + return scanPlanTextAnyToTextScanner{} + } + } + + return nil +} + +type scanPlanBinaryMacaddrToHardwareAddr struct{} + +func (scanPlanBinaryMacaddrToHardwareAddr) Scan(src []byte, dst any) error { + dstBuf := dst.(*net.HardwareAddr) + if src == nil { + *dstBuf = nil + return nil + } + + *dstBuf = make([]byte, len(src)) + copy(*dstBuf, src) + return nil +} + +type scanPlanBinaryMacaddrToTextScanner struct{} + +func (scanPlanBinaryMacaddrToTextScanner) Scan(src []byte, dst any) error { + scanner := (dst).(TextScanner) + if src == nil { + return scanner.ScanText(Text{}) + } + + return scanner.ScanText(Text{String: net.HardwareAddr(src).String(), Valid: true}) +} + +type scanPlanTextMacaddrToHardwareAddr struct{} + +func (scanPlanTextMacaddrToHardwareAddr) Scan(src []byte, dst any) error { + p := dst.(*net.HardwareAddr) + + if src == nil { + *p = nil + return nil + } + + addr, err := net.ParseMAC(string(src)) + if err != nil { + return err + } + + *p = addr + + return nil +} + +func (c MacaddrCodec) DecodeDatabaseSQLValue(m *Map, oid uint32, format int16, src []byte) (driver.Value, error) { + return codecDecodeToTextFormat(c, m, oid, format, src) +} + +func (c MacaddrCodec) DecodeValue(m *Map, oid uint32, format int16, src []byte) (any, error) { + if src == nil { + return nil, nil + } + + var addr net.HardwareAddr + err := codecScan(c, m, oid, format, src, &addr) + if err != nil { + return nil, err + } + return addr, nil +} diff --git a/vendor/github.com/jackc/pgx/v5/pgtype/multirange.go b/vendor/github.com/jackc/pgx/v5/pgtype/multirange.go new file mode 100644 index 0000000000..11f30b448f --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgtype/multirange.go @@ -0,0 +1,453 @@ +package pgtype + +import ( + "bytes" + "database/sql/driver" + "encoding/binary" + "fmt" + "reflect" + + "github.com/jackc/pgx/v5/internal/pgio" +) + +// MultirangeGetter is a type that can be converted into a PostgreSQL multirange. +type MultirangeGetter interface { + // IsNull returns true if the value is SQL NULL. + IsNull() bool + + // Len returns the number of elements in the multirange. + Len() int + + // Index returns the element at i. + Index(i int) any + + // IndexType returns a non-nil scan target of the type Index will return. This is used by MultirangeCodec.PlanEncode. + IndexType() any +} + +// MultirangeSetter is a type can be set from a PostgreSQL multirange. +type MultirangeSetter interface { + // ScanNull sets the value to SQL NULL. + ScanNull() error + + // SetLen prepares the value such that ScanIndex can be called for each element. This will remove any existing + // elements. + SetLen(n int) error + + // ScanIndex returns a value usable as a scan target for i. SetLen must be called before ScanIndex. + ScanIndex(i int) any + + // ScanIndexType returns a non-nil scan target of the type ScanIndex will return. This is used by + // MultirangeCodec.PlanScan. + ScanIndexType() any +} + +// MultirangeCodec is a codec for any multirange type. +type MultirangeCodec struct { + ElementType *Type +} + +func (c *MultirangeCodec) FormatSupported(format int16) bool { + return c.ElementType.Codec.FormatSupported(format) +} + +func (c *MultirangeCodec) PreferredFormat() int16 { + return c.ElementType.Codec.PreferredFormat() +} + +func (c *MultirangeCodec) PlanEncode(m *Map, oid uint32, format int16, value any) EncodePlan { + multirangeValuer, ok := value.(MultirangeGetter) + if !ok { + return nil + } + + elementType := multirangeValuer.IndexType() + + elementEncodePlan := m.PlanEncode(c.ElementType.OID, format, elementType) + if elementEncodePlan == nil { + return nil + } + + switch format { + case BinaryFormatCode: + return &encodePlanMultirangeCodecBinary{ac: c, m: m, oid: oid} + case TextFormatCode: + return &encodePlanMultirangeCodecText{ac: c, m: m, oid: oid} + } + + return nil +} + +type encodePlanMultirangeCodecText struct { + ac *MultirangeCodec + m *Map + oid uint32 +} + +func (p *encodePlanMultirangeCodecText) Encode(value any, buf []byte) (newBuf []byte, err error) { + multirange := value.(MultirangeGetter) + + if multirange.IsNull() { + return nil, nil + } + + elementCount := multirange.Len() + + buf = append(buf, '{') + + var encodePlan EncodePlan + var lastElemType reflect.Type + inElemBuf := make([]byte, 0, 32) + for i := range elementCount { + if i > 0 { + buf = append(buf, ',') + } + + elem := multirange.Index(i) + var elemBuf []byte + if elem != nil { + elemType := reflect.TypeOf(elem) + if lastElemType != elemType { + lastElemType = elemType + encodePlan = p.m.PlanEncode(p.ac.ElementType.OID, TextFormatCode, elem) + if encodePlan == nil { + return nil, fmt.Errorf("unable to encode %v", multirange.Index(i)) + } + } + elemBuf, err = encodePlan.Encode(elem, inElemBuf) + if err != nil { + return nil, err + } + } + + if elemBuf == nil { + return nil, fmt.Errorf("multirange cannot contain NULL element") + } else { + buf = append(buf, elemBuf...) + } + } + + buf = append(buf, '}') + + return buf, nil +} + +type encodePlanMultirangeCodecBinary struct { + ac *MultirangeCodec + m *Map + oid uint32 +} + +func (p *encodePlanMultirangeCodecBinary) Encode(value any, buf []byte) (newBuf []byte, err error) { + multirange := value.(MultirangeGetter) + + if multirange.IsNull() { + return nil, nil + } + + elementCount := multirange.Len() + + buf = pgio.AppendInt32(buf, int32(elementCount)) + + var encodePlan EncodePlan + var lastElemType reflect.Type + for i := range elementCount { + sp := len(buf) + buf = pgio.AppendInt32(buf, -1) + + elem := multirange.Index(i) + var elemBuf []byte + if elem != nil { + elemType := reflect.TypeOf(elem) + if lastElemType != elemType { + lastElemType = elemType + encodePlan = p.m.PlanEncode(p.ac.ElementType.OID, BinaryFormatCode, elem) + if encodePlan == nil { + return nil, fmt.Errorf("unable to encode %v", multirange.Index(i)) + } + } + elemBuf, err = encodePlan.Encode(elem, buf) + if err != nil { + return nil, err + } + } + + if elemBuf == nil { + return nil, fmt.Errorf("multirange cannot contain NULL element") + } else { + buf = elemBuf + pgio.SetInt32(buf[sp:], int32(len(buf[sp:])-4)) + } + } + + return buf, nil +} + +func (c *MultirangeCodec) PlanScan(m *Map, oid uint32, format int16, target any) ScanPlan { + multirangeScanner, ok := target.(MultirangeSetter) + if !ok { + return nil + } + + elementType := multirangeScanner.ScanIndexType() + + elementScanPlan := m.PlanScan(c.ElementType.OID, format, elementType) + if _, ok := elementScanPlan.(*scanPlanFail); ok { + return nil + } + + return &scanPlanMultirangeCodec{ + multirangeCodec: c, + m: m, + oid: oid, + formatCode: format, + } +} + +func (c *MultirangeCodec) decodeBinary(m *Map, multirangeOID uint32, src []byte, multirange MultirangeSetter) error { + rp := 0 + + elementCount := int(binary.BigEndian.Uint32(src[rp:])) + rp += 4 + + // Each element requires at least 4 bytes for its length prefix. + if elementCount > len(src)/4 { + return fmt.Errorf("multirange element count %d exceeds available data", elementCount) + } + + err := multirange.SetLen(elementCount) + if err != nil { + return err + } + + if elementCount == 0 { + return nil + } + + elementScanPlan := c.ElementType.Codec.PlanScan(m, c.ElementType.OID, BinaryFormatCode, multirange.ScanIndex(0)) + if elementScanPlan == nil { + elementScanPlan = m.PlanScan(c.ElementType.OID, BinaryFormatCode, multirange.ScanIndex(0)) + } + + for i := range elementCount { + elem := multirange.ScanIndex(i) + if len(src[rp:]) < 4 { + return fmt.Errorf("multirange body truncated at element %d", i) + } + elemLen := int(int32(binary.BigEndian.Uint32(src[rp:]))) + rp += 4 + var elemSrc []byte + if elemLen >= 0 { + if len(src[rp:]) < elemLen { + return fmt.Errorf("multirange element %d length %d exceeds remaining %d bytes", i, elemLen, len(src[rp:])) + } + elemSrc = src[rp : rp+elemLen] + rp += elemLen + } + err = elementScanPlan.Scan(elemSrc, elem) + if err != nil { + return fmt.Errorf("failed to scan multirange element %d: %w", i, err) + } + } + + return nil +} + +func (c *MultirangeCodec) decodeText(m *Map, multirangeOID uint32, src []byte, multirange MultirangeSetter) error { + elements, err := parseUntypedTextMultirange(src) + if err != nil { + return err + } + + err = multirange.SetLen(len(elements)) + if err != nil { + return err + } + + if len(elements) == 0 { + return nil + } + + elementScanPlan := c.ElementType.Codec.PlanScan(m, c.ElementType.OID, TextFormatCode, multirange.ScanIndex(0)) + if elementScanPlan == nil { + elementScanPlan = m.PlanScan(c.ElementType.OID, TextFormatCode, multirange.ScanIndex(0)) + } + + for i, s := range elements { + elem := multirange.ScanIndex(i) + err = elementScanPlan.Scan([]byte(s), elem) + if err != nil { + return err + } + } + + return nil +} + +type scanPlanMultirangeCodec struct { + multirangeCodec *MultirangeCodec + m *Map + oid uint32 + formatCode int16 + elementScanPlan ScanPlan +} + +func (spac *scanPlanMultirangeCodec) Scan(src []byte, dst any) error { + c := spac.multirangeCodec + m := spac.m + oid := spac.oid + formatCode := spac.formatCode + + multirange := dst.(MultirangeSetter) + + if src == nil { + return multirange.ScanNull() + } + + switch formatCode { + case BinaryFormatCode: + return c.decodeBinary(m, oid, src, multirange) + case TextFormatCode: + return c.decodeText(m, oid, src, multirange) + default: + return fmt.Errorf("unknown format code %d", formatCode) + } +} + +func (c *MultirangeCodec) DecodeDatabaseSQLValue(m *Map, oid uint32, format int16, src []byte) (driver.Value, error) { + if src == nil { + return nil, nil + } + + switch format { + case TextFormatCode: + return string(src), nil + case BinaryFormatCode: + buf := make([]byte, len(src)) + copy(buf, src) + return buf, nil + default: + return nil, fmt.Errorf("unknown format code %d", format) + } +} + +func (c *MultirangeCodec) DecodeValue(m *Map, oid uint32, format int16, src []byte) (any, error) { + if src == nil { + return nil, nil + } + + var multirange Multirange[Range[any]] + err := m.PlanScan(oid, format, &multirange).Scan(src, &multirange) + return multirange, err +} + +func parseUntypedTextMultirange(src []byte) ([]string, error) { + elements := make([]string, 0) + + buf := bytes.NewBuffer(src) + + skipWhitespace(buf) + + r, _, err := buf.ReadRune() + if err != nil { + return nil, fmt.Errorf("invalid array: %w", err) + } + + if r != '{' { + return nil, fmt.Errorf("invalid multirange, expected '{' got %v", r) + } + +parseValueLoop: + for { + r, _, err = buf.ReadRune() + if err != nil { + return nil, fmt.Errorf("invalid multirange: %w", err) + } + + switch r { + case ',': // skip range separator + case '}': + break parseValueLoop + default: + buf.UnreadRune() + value, err := parseRange(buf) + if err != nil { + return nil, fmt.Errorf("invalid multirange value: %w", err) + } + elements = append(elements, value) + } + } + + skipWhitespace(buf) + + if buf.Len() > 0 { + return nil, fmt.Errorf("unexpected trailing data: %v", buf.String()) + } + + return elements, nil +} + +func parseRange(buf *bytes.Buffer) (string, error) { + s := &bytes.Buffer{} + + boundSepRead := false + for { + r, _, err := buf.ReadRune() + if err != nil { + return "", err + } + + switch r { + case ',', '}': + if r == ',' && !boundSepRead { + boundSepRead = true + break + } + buf.UnreadRune() + return s.String(), nil + } + + s.WriteRune(r) + } +} + +// Multirange is a generic multirange type. +// +// T should implement [RangeValuer] and *T should implement [RangeScanner]. However, there does not appear to be a way to +// enforce the [RangeScanner] constraint. +type Multirange[T RangeValuer] []T + +func (r Multirange[T]) IsNull() bool { + return r == nil +} + +func (r Multirange[T]) Len() int { + return len(r) +} + +func (r Multirange[T]) Index(i int) any { + return r[i] +} + +func (r Multirange[T]) IndexType() any { + var zero T + return zero +} + +func (r *Multirange[T]) ScanNull() error { + *r = nil + return nil +} + +func (r *Multirange[T]) SetLen(n int) error { + *r = make([]T, n) + return nil +} + +func (r Multirange[T]) ScanIndex(i int) any { + return &r[i] +} + +func (r Multirange[T]) ScanIndexType() any { + return new(T) +} diff --git a/vendor/github.com/jackc/pgx/v5/pgtype/numeric.go b/vendor/github.com/jackc/pgx/v5/pgtype/numeric.go new file mode 100644 index 0000000000..caf5ff17b6 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgtype/numeric.go @@ -0,0 +1,843 @@ +package pgtype + +import ( + "bytes" + "database/sql/driver" + "encoding/binary" + "fmt" + "math" + "math/big" + "strconv" + "strings" + + "github.com/jackc/pgx/v5/internal/pgio" +) + +// PostgreSQL internal numeric storage uses 16-bit "digits" with base of 10,000 +const nbase = 10_000 + +const ( + pgNumericNaN = 0x00000000c0000000 + pgNumericNaNSign = 0xc000 + + pgNumericPosInf = 0x00000000d0000000 + pgNumericPosInfSign = 0xd000 + + pgNumericNegInf = 0x00000000f0000000 + pgNumericNegInfSign = 0xf000 +) + +var ( + big1 *big.Int = big.NewInt(1) + big10 *big.Int = big.NewInt(10) + big100 *big.Int = big.NewInt(100) + big1000 *big.Int = big.NewInt(1000) +) + +var ( + bigNBase *big.Int = big.NewInt(nbase) + bigNBaseX2 *big.Int = big.NewInt(nbase * nbase) + bigNBaseX3 *big.Int = big.NewInt(nbase * nbase * nbase) + bigNBaseX4 *big.Int = big.NewInt(nbase * nbase * nbase * nbase) +) + +type NumericScanner interface { + ScanNumeric(v Numeric) error +} + +type NumericValuer interface { + NumericValue() (Numeric, error) +} + +type Numeric struct { + Int *big.Int + Exp int32 + NaN bool + InfinityModifier InfinityModifier + Valid bool +} + +// ScanNumeric implements the [NumericScanner] interface. +func (n *Numeric) ScanNumeric(v Numeric) error { + *n = v + return nil +} + +// NumericValue implements the [NumericValuer] interface. +func (n Numeric) NumericValue() (Numeric, error) { + return n, nil +} + +// Float64Value implements the [Float64Valuer] interface. +func (n Numeric) Float64Value() (Float8, error) { + switch { + case !n.Valid: + return Float8{}, nil + case n.NaN: + return Float8{Float64: math.NaN(), Valid: true}, nil + case n.InfinityModifier == Infinity: + return Float8{Float64: math.Inf(1), Valid: true}, nil + case n.InfinityModifier == NegativeInfinity: + return Float8{Float64: math.Inf(-1), Valid: true}, nil + } + + buf := make([]byte, 0, 32) + + if n.Int == nil { + buf = append(buf, '0') + } else { + buf = append(buf, n.Int.String()...) + } + buf = append(buf, 'e') + buf = append(buf, strconv.FormatInt(int64(n.Exp), 10)...) + + f, err := strconv.ParseFloat(string(buf), 64) + if err != nil { + return Float8{}, err + } + + return Float8{Float64: f, Valid: true}, nil +} + +// ScanInt64 implements the [Int64Scanner] interface. +func (n *Numeric) ScanInt64(v Int8) error { + if !v.Valid { + *n = Numeric{} + return nil + } + + *n = Numeric{Int: big.NewInt(v.Int64), Valid: true} + return nil +} + +// Int64Value implements the [Int64Valuer] interface. +func (n Numeric) Int64Value() (Int8, error) { + if !n.Valid { + return Int8{}, nil + } + + bi, err := n.toBigInt() + if err != nil { + return Int8{}, err + } + + if !bi.IsInt64() { + return Int8{}, fmt.Errorf("cannot convert %v to int64", n) + } + + return Int8{Int64: bi.Int64(), Valid: true}, nil +} + +func (n *Numeric) ScanScientific(src string) error { + if !strings.ContainsAny(src, "eE") { + return scanPlanTextAnyToNumericScanner{}.Scan([]byte(src), n) + } + + if bigF, ok := new(big.Float).SetString(src); ok { + smallF, _ := bigF.Float64() + src = strconv.FormatFloat(smallF, 'f', -1, 64) + } + + num, exp, err := parseNumericString(src) + if err != nil { + return err + } + + *n = Numeric{Int: num, Exp: exp, Valid: true} + + return nil +} + +func (n *Numeric) toBigInt() (*big.Int, error) { + if n.Exp == 0 { + return n.Int, nil + } + + num := &big.Int{} + num.Set(n.Int) + if n.Exp > 0 { + mul := &big.Int{} + mul.Exp(big10, big.NewInt(int64(n.Exp)), nil) + num.Mul(num, mul) + return num, nil + } + + div := &big.Int{} + div.Exp(big10, big.NewInt(int64(-n.Exp)), nil) + remainder := &big.Int{} + num.DivMod(num, div, remainder) + if remainder.Sign() != 0 { + return nil, fmt.Errorf("cannot convert %v to integer", n) + } + return num, nil +} + +func parseNumericString(str string) (n *big.Int, exp int32, err error) { + idx := strings.IndexByte(str, '.') + + if idx == -1 { + for len(str) > 1 && str[len(str)-1] == '0' && str[len(str)-2] != '-' { + str = str[:len(str)-1] + exp++ + } + } else { + exp = int32(-(len(str) - idx - 1)) + str = str[:idx] + str[idx+1:] + } + + accum := &big.Int{} + if _, ok := accum.SetString(str, 10); !ok { + return nil, 0, fmt.Errorf("%s is not a number", str) + } + + return accum, exp, nil +} + +func nbaseDigitsToInt64(src []byte) (accum int64, bytesRead, digitsRead int) { + digits := min(len(src)/2, 4) + + rp := 0 + + for i := range digits { + if i > 0 { + accum *= nbase + } + accum += int64(binary.BigEndian.Uint16(src[rp:])) + rp += 2 + } + + return accum, rp, digits +} + +// Scan implements the [database/sql.Scanner] interface. +func (n *Numeric) Scan(src any) error { + if src == nil { + *n = Numeric{} + return nil + } + + if src, ok := src.(string); ok { + return scanPlanTextAnyToNumericScanner{}.Scan([]byte(src), n) + } + + return fmt.Errorf("cannot scan %T", src) +} + +// Value implements the [database/sql/driver.Valuer] interface. +func (n Numeric) Value() (driver.Value, error) { + if !n.Valid { + return nil, nil + } + + buf, err := NumericCodec{}.PlanEncode(nil, 0, TextFormatCode, n).Encode(n, nil) + if err != nil { + return nil, err + } + return string(buf), err +} + +// MarshalJSON implements the [encoding/json.Marshaler] interface. +func (n Numeric) MarshalJSON() ([]byte, error) { + if !n.Valid { + return []byte("null"), nil + } + + if n.NaN { + return []byte(`"NaN"`), nil + } + + return n.numberTextBytes(), nil +} + +// UnmarshalJSON implements the [encoding/json.Unmarshaler] interface. +func (n *Numeric) UnmarshalJSON(src []byte) error { + if bytes.Equal(src, []byte(`null`)) { + *n = Numeric{} + return nil + } + if bytes.Equal(src, []byte(`"NaN"`)) { + *n = Numeric{NaN: true, Valid: true} + return nil + } + return scanPlanTextAnyToNumericScanner{}.Scan(src, n) +} + +// numberString returns a string of the number. undefined if NaN, infinite, or NULL +func (n Numeric) numberTextBytes() []byte { + if n.Int == nil { + return []byte("0") + } + + intStr := n.Int.String() + + buf := &bytes.Buffer{} + + if len(intStr) > 0 && intStr[:1] == "-" { + intStr = intStr[1:] + buf.WriteByte('-') + } + + exp := int(n.Exp) + switch { + case exp > 0: + buf.WriteString(intStr) + for range exp { + buf.WriteByte('0') + } + case exp < 0: + if len(intStr) <= -exp { + buf.WriteString("0.") + leadingZeros := -exp - len(intStr) + for range leadingZeros { + buf.WriteByte('0') + } + buf.WriteString(intStr) + } else if len(intStr) > -exp { + dpPos := len(intStr) + exp + buf.WriteString(intStr[:dpPos]) + buf.WriteByte('.') + buf.WriteString(intStr[dpPos:]) + } + default: + buf.WriteString(intStr) + } + + return buf.Bytes() +} + +type NumericCodec struct{} + +func (NumericCodec) FormatSupported(format int16) bool { + return format == TextFormatCode || format == BinaryFormatCode +} + +func (NumericCodec) PreferredFormat() int16 { + return BinaryFormatCode +} + +func (NumericCodec) PlanEncode(m *Map, oid uint32, format int16, value any) EncodePlan { + switch format { + case BinaryFormatCode: + switch value.(type) { + case NumericValuer: + return encodePlanNumericCodecBinaryNumericValuer{} + case Float64Valuer: + return encodePlanNumericCodecBinaryFloat64Valuer{} + case Int64Valuer: + return encodePlanNumericCodecBinaryInt64Valuer{} + } + case TextFormatCode: + switch value.(type) { + case NumericValuer: + return encodePlanNumericCodecTextNumericValuer{} + case Float64Valuer: + return encodePlanNumericCodecTextFloat64Valuer{} + case Int64Valuer: + return encodePlanNumericCodecTextInt64Valuer{} + } + } + + return nil +} + +type encodePlanNumericCodecBinaryNumericValuer struct{} + +func (encodePlanNumericCodecBinaryNumericValuer) Encode(value any, buf []byte) (newBuf []byte, err error) { + n, err := value.(NumericValuer).NumericValue() + if err != nil { + return nil, err + } + + return encodeNumericBinary(n, buf) +} + +type encodePlanNumericCodecBinaryFloat64Valuer struct{} + +func (encodePlanNumericCodecBinaryFloat64Valuer) Encode(value any, buf []byte) (newBuf []byte, err error) { + n, err := value.(Float64Valuer).Float64Value() + if err != nil { + return nil, err + } + + if !n.Valid { + return nil, nil + } + + switch { + case math.IsNaN(n.Float64): + return encodeNumericBinary(Numeric{NaN: true, Valid: true}, buf) + case math.IsInf(n.Float64, 1): + return encodeNumericBinary(Numeric{InfinityModifier: Infinity, Valid: true}, buf) + case math.IsInf(n.Float64, -1): + return encodeNumericBinary(Numeric{InfinityModifier: NegativeInfinity, Valid: true}, buf) + } + num, exp, err := parseNumericString(strconv.FormatFloat(n.Float64, 'f', -1, 64)) + if err != nil { + return nil, err + } + + return encodeNumericBinary(Numeric{Int: num, Exp: exp, Valid: true}, buf) +} + +type encodePlanNumericCodecBinaryInt64Valuer struct{} + +func (encodePlanNumericCodecBinaryInt64Valuer) Encode(value any, buf []byte) (newBuf []byte, err error) { + n, err := value.(Int64Valuer).Int64Value() + if err != nil { + return nil, err + } + + if !n.Valid { + return nil, nil + } + + return encodeNumericBinary(Numeric{Int: big.NewInt(n.Int64), Valid: true}, buf) +} + +func encodeNumericBinary(n Numeric, buf []byte) (newBuf []byte, err error) { + if !n.Valid { + return nil, nil + } + + switch { + case n.NaN: + buf = pgio.AppendUint64(buf, pgNumericNaN) + return buf, nil + case n.InfinityModifier == Infinity: + buf = pgio.AppendUint64(buf, pgNumericPosInf) + return buf, nil + case n.InfinityModifier == NegativeInfinity: + buf = pgio.AppendUint64(buf, pgNumericNegInf) + return buf, nil + } + + var sign int16 + if n.Int != nil && n.Int.Sign() < 0 { + sign = 16384 + } + + absInt := &big.Int{} + wholePart := &big.Int{} + fracPart := &big.Int{} + remainder := &big.Int{} + if n.Int != nil { + absInt.Abs(n.Int) + } + + // Normalize absInt and exp to where exp is always a multiple of 4. This makes + // converting to 16-bit base 10,000 digits easier. + var exp int32 + switch n.Exp % 4 { + case 1, -3: + exp = n.Exp - 1 + absInt.Mul(absInt, big10) + case 2, -2: + exp = n.Exp - 2 + absInt.Mul(absInt, big100) + case 3, -1: + exp = n.Exp - 3 + absInt.Mul(absInt, big1000) + default: + exp = n.Exp + } + + if exp < 0 { + divisor := &big.Int{} + divisor.Exp(big10, big.NewInt(int64(-exp)), nil) + wholePart.DivMod(absInt, divisor, fracPart) + fracPart.Add(fracPart, divisor) + } else { + wholePart = absInt + } + + var wholeDigits, fracDigits []int16 + + for wholePart.Sign() != 0 { + wholePart.DivMod(wholePart, bigNBase, remainder) + wholeDigits = append(wholeDigits, int16(remainder.Int64())) + } + + if fracPart.Sign() != 0 { + for fracPart.Cmp(big1) != 0 { + fracPart.DivMod(fracPart, bigNBase, remainder) + fracDigits = append(fracDigits, int16(remainder.Int64())) + } + } + + buf = pgio.AppendInt16(buf, int16(len(wholeDigits)+len(fracDigits))) + + var weight int16 + if len(wholeDigits) > 0 { + weight = int16(len(wholeDigits) - 1) + if exp > 0 { + weight += int16(exp / 4) + } + } else { + weight = int16(exp/4) - 1 + int16(len(fracDigits)) + } + buf = pgio.AppendInt16(buf, weight) + + buf = pgio.AppendInt16(buf, sign) + + var dscale int16 + if n.Exp < 0 { + dscale = int16(-n.Exp) + } + buf = pgio.AppendInt16(buf, dscale) + + for i := len(wholeDigits) - 1; i >= 0; i-- { + buf = pgio.AppendInt16(buf, wholeDigits[i]) + } + + for i := len(fracDigits) - 1; i >= 0; i-- { + buf = pgio.AppendInt16(buf, fracDigits[i]) + } + + return buf, nil +} + +type encodePlanNumericCodecTextNumericValuer struct{} + +func (encodePlanNumericCodecTextNumericValuer) Encode(value any, buf []byte) (newBuf []byte, err error) { + n, err := value.(NumericValuer).NumericValue() + if err != nil { + return nil, err + } + + return encodeNumericText(n, buf) +} + +type encodePlanNumericCodecTextFloat64Valuer struct{} + +func (encodePlanNumericCodecTextFloat64Valuer) Encode(value any, buf []byte) (newBuf []byte, err error) { + n, err := value.(Float64Valuer).Float64Value() + if err != nil { + return nil, err + } + + if !n.Valid { + return nil, nil + } + + switch { + case math.IsNaN(n.Float64): + buf = append(buf, "NaN"...) + case math.IsInf(n.Float64, 1): + buf = append(buf, "Infinity"...) + case math.IsInf(n.Float64, -1): + buf = append(buf, "-Infinity"...) + default: + buf = append(buf, strconv.FormatFloat(n.Float64, 'f', -1, 64)...) + } + return buf, nil +} + +type encodePlanNumericCodecTextInt64Valuer struct{} + +func (encodePlanNumericCodecTextInt64Valuer) Encode(value any, buf []byte) (newBuf []byte, err error) { + n, err := value.(Int64Valuer).Int64Value() + if err != nil { + return nil, err + } + + if !n.Valid { + return nil, nil + } + + buf = append(buf, strconv.FormatInt(n.Int64, 10)...) + return buf, nil +} + +func encodeNumericText(n Numeric, buf []byte) (newBuf []byte, err error) { + if !n.Valid { + return nil, nil + } + + switch { + case n.NaN: + buf = append(buf, "NaN"...) + return buf, nil + case n.InfinityModifier == Infinity: + buf = append(buf, "Infinity"...) + return buf, nil + case n.InfinityModifier == NegativeInfinity: + buf = append(buf, "-Infinity"...) + return buf, nil + } + + buf = append(buf, n.numberTextBytes()...) + + return buf, nil +} + +func (NumericCodec) PlanScan(m *Map, oid uint32, format int16, target any) ScanPlan { + switch format { + case BinaryFormatCode: + switch target.(type) { + case NumericScanner: + return scanPlanBinaryNumericToNumericScanner{} + case Float64Scanner: + return scanPlanBinaryNumericToFloat64Scanner{} + case Int64Scanner: + return scanPlanBinaryNumericToInt64Scanner{} + case TextScanner: + return scanPlanBinaryNumericToTextScanner{} + } + case TextFormatCode: + switch target.(type) { + case NumericScanner: + return scanPlanTextAnyToNumericScanner{} + case Float64Scanner: + return scanPlanTextAnyToFloat64Scanner{} + case Int64Scanner: + return scanPlanTextAnyToInt64Scanner{} + } + } + + return nil +} + +type scanPlanBinaryNumericToNumericScanner struct{} + +func (scanPlanBinaryNumericToNumericScanner) Scan(src []byte, dst any) error { + scanner := (dst).(NumericScanner) + + if src == nil { + return scanner.ScanNumeric(Numeric{}) + } + + if len(src) < 8 { + return fmt.Errorf("numeric incomplete %v", src) + } + + rp := 0 + ndigits := binary.BigEndian.Uint16(src[rp:]) + rp += 2 + weight := int16(binary.BigEndian.Uint16(src[rp:])) + rp += 2 + sign := binary.BigEndian.Uint16(src[rp:]) + rp += 2 + dscale := int16(binary.BigEndian.Uint16(src[rp:])) + rp += 2 + + switch sign { + case pgNumericNaNSign: + return scanner.ScanNumeric(Numeric{NaN: true, Valid: true}) + case pgNumericPosInfSign: + return scanner.ScanNumeric(Numeric{InfinityModifier: Infinity, Valid: true}) + case pgNumericNegInfSign: + return scanner.ScanNumeric(Numeric{InfinityModifier: NegativeInfinity, Valid: true}) + } + + if ndigits == 0 { + return scanner.ScanNumeric(Numeric{Int: big.NewInt(0), Valid: true}) + } + + if len(src[rp:]) < int(ndigits)*2 { + return fmt.Errorf("numeric incomplete %v", src) + } + + accum := &big.Int{} + + for i := 0; i < int(ndigits+3)/4; i++ { + int64accum, bytesRead, digitsRead := nbaseDigitsToInt64(src[rp:]) + rp += bytesRead + + if i > 0 { + var mul *big.Int + switch digitsRead { + case 1: + mul = bigNBase + case 2: + mul = bigNBaseX2 + case 3: + mul = bigNBaseX3 + case 4: + mul = bigNBaseX4 + default: + return fmt.Errorf("invalid digitsRead: %d (this can't happen)", digitsRead) + } + accum.Mul(accum, mul) + } + + accum.Add(accum, big.NewInt(int64accum)) + } + + exp := (int32(weight) - int32(ndigits) + 1) * 4 + + if dscale > 0 { + fracNBaseDigits := int(ndigits) - int(weight) - 1 + fracDecimalDigits := fracNBaseDigits * 4 + dscaleInt := int(dscale) + + if dscaleInt > fracDecimalDigits { + multCount := dscaleInt - fracDecimalDigits + for range multCount { + accum.Mul(accum, big10) + exp-- + } + } else if dscaleInt < fracDecimalDigits { + divCount := fracDecimalDigits - dscaleInt + for range divCount { + accum.Div(accum, big10) + exp++ + } + } + } + + reduced := &big.Int{} + remainder := &big.Int{} + if exp >= 0 { + for { + reduced.DivMod(accum, big10, remainder) + if remainder.Sign() != 0 { + break + } + accum.Set(reduced) + exp++ + } + } + + if sign != 0 { + accum.Neg(accum) + } + + return scanner.ScanNumeric(Numeric{Int: accum, Exp: exp, Valid: true}) +} + +type scanPlanBinaryNumericToFloat64Scanner struct{} + +func (scanPlanBinaryNumericToFloat64Scanner) Scan(src []byte, dst any) error { + scanner := (dst).(Float64Scanner) + + if src == nil { + return scanner.ScanFloat64(Float8{}) + } + + var n Numeric + + err := scanPlanBinaryNumericToNumericScanner{}.Scan(src, &n) + if err != nil { + return err + } + + f8, err := n.Float64Value() + if err != nil { + return err + } + + return scanner.ScanFloat64(f8) +} + +type scanPlanBinaryNumericToInt64Scanner struct{} + +func (scanPlanBinaryNumericToInt64Scanner) Scan(src []byte, dst any) error { + scanner := (dst).(Int64Scanner) + + if src == nil { + return scanner.ScanInt64(Int8{}) + } + + var n Numeric + + err := scanPlanBinaryNumericToNumericScanner{}.Scan(src, &n) + if err != nil { + return err + } + + bigInt, err := n.toBigInt() + if err != nil { + return err + } + + if !bigInt.IsInt64() { + return fmt.Errorf("%v is out of range for int64", bigInt) + } + + return scanner.ScanInt64(Int8{Int64: bigInt.Int64(), Valid: true}) +} + +type scanPlanBinaryNumericToTextScanner struct{} + +func (scanPlanBinaryNumericToTextScanner) Scan(src []byte, dst any) error { + scanner := (dst).(TextScanner) + + if src == nil { + return scanner.ScanText(Text{}) + } + + var n Numeric + + err := scanPlanBinaryNumericToNumericScanner{}.Scan(src, &n) + if err != nil { + return err + } + + sbuf, err := encodeNumericText(n, nil) + if err != nil { + return err + } + + return scanner.ScanText(Text{String: string(sbuf), Valid: true}) +} + +type scanPlanTextAnyToNumericScanner struct{} + +func (scanPlanTextAnyToNumericScanner) Scan(src []byte, dst any) error { + scanner := (dst).(NumericScanner) + + if src == nil { + return scanner.ScanNumeric(Numeric{}) + } + + switch string(src) { + case "NaN": + return scanner.ScanNumeric(Numeric{NaN: true, Valid: true}) + case "Infinity": + return scanner.ScanNumeric(Numeric{InfinityModifier: Infinity, Valid: true}) + case "-Infinity": + return scanner.ScanNumeric(Numeric{InfinityModifier: NegativeInfinity, Valid: true}) + } + + num, exp, err := parseNumericString(string(src)) + if err != nil { + return err + } + + return scanner.ScanNumeric(Numeric{Int: num, Exp: exp, Valid: true}) +} + +func (c NumericCodec) DecodeDatabaseSQLValue(m *Map, oid uint32, format int16, src []byte) (driver.Value, error) { + if src == nil { + return nil, nil + } + + if format == TextFormatCode { + return string(src), nil + } + + var n Numeric + err := codecScan(c, m, oid, format, src, &n) + if err != nil { + return nil, err + } + + buf, err := m.Encode(oid, TextFormatCode, n, nil) + if err != nil { + return nil, err + } + return string(buf), nil +} + +func (c NumericCodec) DecodeValue(m *Map, oid uint32, format int16, src []byte) (any, error) { + if src == nil { + return nil, nil + } + + var n Numeric + err := codecScan(c, m, oid, format, src, &n) + if err != nil { + return nil, err + } + return n, nil +} diff --git a/vendor/github.com/jackc/pgx/v5/pgtype/path.go b/vendor/github.com/jackc/pgx/v5/pgtype/path.go new file mode 100644 index 0000000000..6398b58152 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgtype/path.go @@ -0,0 +1,280 @@ +package pgtype + +import ( + "database/sql/driver" + "encoding/binary" + "fmt" + "math" + "strconv" + "strings" + + "github.com/jackc/pgx/v5/internal/pgio" +) + +type PathScanner interface { + ScanPath(v Path) error +} + +type PathValuer interface { + PathValue() (Path, error) +} + +type Path struct { + P []Vec2 + Closed bool + Valid bool +} + +// ScanPath implements the [PathScanner] interface. +func (path *Path) ScanPath(v Path) error { + *path = v + return nil +} + +// PathValue implements the [PathValuer] interface. +func (path Path) PathValue() (Path, error) { + return path, nil +} + +// Scan implements the [database/sql.Scanner] interface. +func (path *Path) Scan(src any) error { + if src == nil { + *path = Path{} + return nil + } + + if src, ok := src.(string); ok { + return scanPlanTextAnyToPathScanner{}.Scan([]byte(src), path) + } + + return fmt.Errorf("cannot scan %T", src) +} + +// Value implements the [database/sql/driver.Valuer] interface. +func (path Path) Value() (driver.Value, error) { + if !path.Valid { + return nil, nil + } + + buf, err := PathCodec{}.PlanEncode(nil, 0, TextFormatCode, path).Encode(path, nil) + if err != nil { + return nil, err + } + + return string(buf), err +} + +type PathCodec struct{} + +func (PathCodec) FormatSupported(format int16) bool { + return format == TextFormatCode || format == BinaryFormatCode +} + +func (PathCodec) PreferredFormat() int16 { + return BinaryFormatCode +} + +func (PathCodec) PlanEncode(m *Map, oid uint32, format int16, value any) EncodePlan { + if _, ok := value.(PathValuer); !ok { + return nil + } + + switch format { + case BinaryFormatCode: + return encodePlanPathCodecBinary{} + case TextFormatCode: + return encodePlanPathCodecText{} + } + + return nil +} + +type encodePlanPathCodecBinary struct{} + +func (encodePlanPathCodecBinary) Encode(value any, buf []byte) (newBuf []byte, err error) { + path, err := value.(PathValuer).PathValue() + if err != nil { + return nil, err + } + + if !path.Valid { + return nil, nil + } + + var closeByte byte + if path.Closed { + closeByte = 1 + } + buf = append(buf, closeByte) + + buf = pgio.AppendInt32(buf, int32(len(path.P))) + + for _, p := range path.P { + buf = pgio.AppendUint64(buf, math.Float64bits(p.X)) + buf = pgio.AppendUint64(buf, math.Float64bits(p.Y)) + } + + return buf, nil +} + +type encodePlanPathCodecText struct{} + +func (encodePlanPathCodecText) Encode(value any, buf []byte) (newBuf []byte, err error) { + path, err := value.(PathValuer).PathValue() + if err != nil { + return nil, err + } + + if !path.Valid { + return nil, nil + } + + var startByte, endByte byte + if path.Closed { + startByte = '(' + endByte = ')' + } else { + startByte = '[' + endByte = ']' + } + buf = append(buf, startByte) + + for i, p := range path.P { + if i > 0 { + buf = append(buf, ',') + } + buf = append(buf, fmt.Sprintf(`(%s,%s)`, + strconv.FormatFloat(p.X, 'f', -1, 64), + strconv.FormatFloat(p.Y, 'f', -1, 64), + )...) + } + + buf = append(buf, endByte) + + return buf, nil +} + +func (PathCodec) PlanScan(m *Map, oid uint32, format int16, target any) ScanPlan { + switch format { + case BinaryFormatCode: + if _, ok := target.(PathScanner); ok { + return scanPlanBinaryPathToPathScanner{} + } + case TextFormatCode: + if _, ok := target.(PathScanner); ok { + return scanPlanTextAnyToPathScanner{} + } + } + + return nil +} + +type scanPlanBinaryPathToPathScanner struct{} + +func (scanPlanBinaryPathToPathScanner) Scan(src []byte, dst any) error { + scanner := (dst).(PathScanner) + + if src == nil { + return scanner.ScanPath(Path{}) + } + + if len(src) < 5 { + return fmt.Errorf("invalid length for Path: %v", len(src)) + } + + closed := src[0] == 1 + pointCount := int(binary.BigEndian.Uint32(src[1:])) + + rp := 5 + + if 5+pointCount*16 != len(src) { + return fmt.Errorf("invalid length for Path with %d points: %v", pointCount, len(src)) + } + + points := make([]Vec2, pointCount) + for i := range points { + x := binary.BigEndian.Uint64(src[rp:]) + rp += 8 + y := binary.BigEndian.Uint64(src[rp:]) + rp += 8 + points[i] = Vec2{math.Float64frombits(x), math.Float64frombits(y)} + } + + return scanner.ScanPath(Path{ + P: points, + Closed: closed, + Valid: true, + }) +} + +type scanPlanTextAnyToPathScanner struct{} + +func (scanPlanTextAnyToPathScanner) Scan(src []byte, dst any) error { + scanner := (dst).(PathScanner) + + if src == nil { + return scanner.ScanPath(Path{}) + } + + if len(src) < 7 { + return fmt.Errorf("invalid length for Path: %v", len(src)) + } + + closed := src[0] == '(' + points := make([]Vec2, 0) + + // Expected format: ((x1,y1),...,(xn,yn)) or [(x1,y1),...,(xn,yn)] + str := string(src[1 : len(src)-1]) + + for { + if len(str) == 0 || str[0] != '(' { + return fmt.Errorf("invalid format for Path") + } + body, rest, found := strings.Cut(str[1:], ")") + if !found { + return fmt.Errorf("invalid format for Path") + } + + sx, sy, found := strings.Cut(body, ",") + if !found { + return fmt.Errorf("invalid format for Path") + } + x, err := strconv.ParseFloat(sx, 64) + if err != nil { + return err + } + y, err := strconv.ParseFloat(sy, 64) + if err != nil { + return err + } + + points = append(points, Vec2{x, y}) + + if rest == "" { + break + } + str, found = strings.CutPrefix(rest, ",") + if !found { + return fmt.Errorf("invalid format for Path") + } + } + + return scanner.ScanPath(Path{P: points, Closed: closed, Valid: true}) +} + +func (c PathCodec) DecodeDatabaseSQLValue(m *Map, oid uint32, format int16, src []byte) (driver.Value, error) { + return codecDecodeToTextFormat(c, m, oid, format, src) +} + +func (c PathCodec) DecodeValue(m *Map, oid uint32, format int16, src []byte) (any, error) { + if src == nil { + return nil, nil + } + + var path Path + err := codecScan(c, m, oid, format, src, &path) + if err != nil { + return nil, err + } + return path, nil +} diff --git a/vendor/github.com/jackc/pgx/v5/pgtype/pgtype.go b/vendor/github.com/jackc/pgx/v5/pgtype/pgtype.go new file mode 100644 index 0000000000..46b892bfb4 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgtype/pgtype.go @@ -0,0 +1,2067 @@ +package pgtype + +import ( + "database/sql" + "database/sql/driver" + "errors" + "fmt" + "net" + "net/netip" + "reflect" + "time" +) + +// PostgreSQL oids for common types +const ( + BoolOID = 16 + ByteaOID = 17 + QCharOID = 18 + NameOID = 19 + Int8OID = 20 + Int2OID = 21 + Int4OID = 23 + TextOID = 25 + OIDOID = 26 + TIDOID = 27 + XIDOID = 28 + CIDOID = 29 + JSONOID = 114 + XMLOID = 142 + XMLArrayOID = 143 + JSONArrayOID = 199 + XID8ArrayOID = 271 + PointOID = 600 + LsegOID = 601 + PathOID = 602 + BoxOID = 603 + PolygonOID = 604 + LineOID = 628 + LineArrayOID = 629 + CIDROID = 650 + CIDRArrayOID = 651 + Float4OID = 700 + Float8OID = 701 + CircleOID = 718 + CircleArrayOID = 719 + UnknownOID = 705 + Macaddr8OID = 774 + MacaddrOID = 829 + InetOID = 869 + BoolArrayOID = 1000 + QCharArrayOID = 1002 + NameArrayOID = 1003 + Int2ArrayOID = 1005 + Int4ArrayOID = 1007 + TextArrayOID = 1009 + TIDArrayOID = 1010 + ByteaArrayOID = 1001 + XIDArrayOID = 1011 + CIDArrayOID = 1012 + BPCharArrayOID = 1014 + VarcharArrayOID = 1015 + Int8ArrayOID = 1016 + PointArrayOID = 1017 + LsegArrayOID = 1018 + PathArrayOID = 1019 + BoxArrayOID = 1020 + Float4ArrayOID = 1021 + Float8ArrayOID = 1022 + PolygonArrayOID = 1027 + OIDArrayOID = 1028 + ACLItemOID = 1033 + ACLItemArrayOID = 1034 + MacaddrArrayOID = 1040 + InetArrayOID = 1041 + BPCharOID = 1042 + VarcharOID = 1043 + DateOID = 1082 + TimeOID = 1083 + TimestampOID = 1114 + TimestampArrayOID = 1115 + DateArrayOID = 1182 + TimeArrayOID = 1183 + TimestamptzOID = 1184 + TimestamptzArrayOID = 1185 + IntervalOID = 1186 + IntervalArrayOID = 1187 + NumericArrayOID = 1231 + TimetzOID = 1266 + TimetzArrayOID = 1270 + BitOID = 1560 + BitArrayOID = 1561 + VarbitOID = 1562 + VarbitArrayOID = 1563 + NumericOID = 1700 + RecordOID = 2249 + RecordArrayOID = 2287 + UUIDOID = 2950 + UUIDArrayOID = 2951 + TSVectorOID = 3614 + TSVectorArrayOID = 3643 + JSONBOID = 3802 + JSONBArrayOID = 3807 + DaterangeOID = 3912 + DaterangeArrayOID = 3913 + Int4rangeOID = 3904 + Int4rangeArrayOID = 3905 + NumrangeOID = 3906 + NumrangeArrayOID = 3907 + TsrangeOID = 3908 + TsrangeArrayOID = 3909 + TstzrangeOID = 3910 + TstzrangeArrayOID = 3911 + Int8rangeOID = 3926 + Int8rangeArrayOID = 3927 + JSONPathOID = 4072 + JSONPathArrayOID = 4073 + Int4multirangeOID = 4451 + NummultirangeOID = 4532 + TsmultirangeOID = 4533 + TstzmultirangeOID = 4534 + DatemultirangeOID = 4535 + Int8multirangeOID = 4536 + XID8OID = 5069 + Int4multirangeArrayOID = 6150 + NummultirangeArrayOID = 6151 + TsmultirangeArrayOID = 6152 + TstzmultirangeArrayOID = 6153 + DatemultirangeArrayOID = 6155 + Int8multirangeArrayOID = 6157 +) + +type InfinityModifier int8 + +const ( + Infinity InfinityModifier = 1 + Finite InfinityModifier = 0 + NegativeInfinity InfinityModifier = -Infinity +) + +func (im InfinityModifier) String() string { + switch im { + case Finite: + return "finite" + case Infinity: + return "infinity" + case NegativeInfinity: + return "-infinity" + default: + return "invalid" + } +} + +// PostgreSQL format codes +const ( + TextFormatCode = 0 + BinaryFormatCode = 1 +) + +// A Codec converts between Go and PostgreSQL values. A Codec must not be mutated after it is registered with a [Map]. +type Codec interface { + // FormatSupported returns true if the format is supported. + FormatSupported(int16) bool + + // PreferredFormat returns the preferred format. + PreferredFormat() int16 + + // PlanEncode returns an EncodePlan for encoding value into PostgreSQL format for oid and format. If no plan can be + // found then nil is returned. + PlanEncode(m *Map, oid uint32, format int16, value any) EncodePlan + + // PlanScan returns a ScanPlan for scanning a PostgreSQL value into a destination with the same type as target. If + // no plan can be found then nil is returned. + PlanScan(m *Map, oid uint32, format int16, target any) ScanPlan + + // DecodeDatabaseSQLValue returns src decoded into a value compatible with the sql.Scanner interface. + DecodeDatabaseSQLValue(m *Map, oid uint32, format int16, src []byte) (driver.Value, error) + + // DecodeValue returns src decoded into its default format. + DecodeValue(m *Map, oid uint32, format int16, src []byte) (any, error) +} + +type nullAssignmentError struct { + dst any +} + +func (e *nullAssignmentError) Error() string { + return fmt.Sprintf("cannot assign NULL to %T", e.dst) +} + +// Type represents a PostgreSQL data type. It must not be mutated after it is registered with a [Map]. +type Type struct { + Codec Codec + Name string + OID uint32 +} + +// Map is the mapping between PostgreSQL server types and Go type handling logic. It can encode values for +// transmission to a PostgreSQL server and scan received values. +type Map struct { + oidToType map[uint32]*Type + nameToType map[string]*Type + reflectTypeToName map[reflect.Type]string + oidToFormatCode map[uint32]int16 + + reflectTypeToType map[reflect.Type]*Type + + memoizedEncodePlans map[uint32]map[reflect.Type][2]EncodePlan + + // TryWrapEncodePlanFuncs is a slice of functions that will wrap a value that cannot be encoded by the Codec. Every + // time a wrapper is found the PlanEncode method will be recursively called with the new value. This allows several layers of wrappers + // to be built up. There are default functions placed in this slice by NewMap(). In most cases these functions + // should run last. i.e. Additional functions should typically be prepended not appended. + TryWrapEncodePlanFuncs []TryWrapEncodePlanFunc + + // TryWrapScanPlanFuncs is a slice of functions that will wrap a target that cannot be scanned into by the Codec. Every + // time a wrapper is found the PlanScan method will be recursively called with the new target. This allows several layers of wrappers + // to be built up. There are default functions placed in this slice by NewMap(). In most cases these functions + // should run last. i.e. Additional functions should typically be prepended not appended. + TryWrapScanPlanFuncs []TryWrapScanPlanFunc +} + +// Copy returns a new Map containing the same registered types. +func (m *Map) Copy() *Map { + newMap := NewMap() + for _, type_ := range m.oidToType { + newMap.RegisterType(type_) + } + return newMap +} + +func NewMap() *Map { + defaultMapInitOnce.Do(initDefaultMap) + + return &Map{ + oidToType: make(map[uint32]*Type), + nameToType: make(map[string]*Type), + reflectTypeToName: make(map[reflect.Type]string), + oidToFormatCode: make(map[uint32]int16), + + memoizedEncodePlans: make(map[uint32]map[reflect.Type][2]EncodePlan), + + TryWrapEncodePlanFuncs: []TryWrapEncodePlanFunc{ + TryWrapDerefPointerEncodePlan, + TryWrapBuiltinTypeEncodePlan, + TryWrapFindUnderlyingTypeEncodePlan, + TryWrapStringerEncodePlan, + TryWrapStructEncodePlan, + TryWrapSliceEncodePlan, + TryWrapMultiDimSliceEncodePlan, + TryWrapArrayEncodePlan, + }, + + TryWrapScanPlanFuncs: []TryWrapScanPlanFunc{ + TryPointerPointerScanPlan, + TryWrapBuiltinTypeScanPlan, + TryFindUnderlyingTypeScanPlan, + TryWrapStructScanPlan, + TryWrapPtrSliceScanPlan, + TryWrapPtrMultiDimSliceScanPlan, + TryWrapPtrArrayScanPlan, + }, + } +} + +// RegisterTypes registers multiple data types in the sequence they are provided. +func (m *Map) RegisterTypes(types []*Type) { + for _, t := range types { + m.RegisterType(t) + } +} + +// RegisterType registers a data type with the [Map]. t must not be mutated after it is registered. +func (m *Map) RegisterType(t *Type) { + m.oidToType[t.OID] = t + m.nameToType[t.Name] = t + m.oidToFormatCode[t.OID] = t.Codec.PreferredFormat() + + // Invalidated by type registration + m.reflectTypeToType = nil + for k := range m.memoizedEncodePlans { + delete(m.memoizedEncodePlans, k) + } +} + +// RegisterDefaultPgType registers a mapping of a Go type to a PostgreSQL type name. Typically the data type to be +// encoded or decoded is determined by the PostgreSQL OID. But if the OID of a value to be encoded or decoded is +// unknown, this additional mapping will be used by TypeForValue to determine a suitable data type. +func (m *Map) RegisterDefaultPgType(value any, name string) { + m.reflectTypeToName[reflect.TypeOf(value)] = name + + // Invalidated by type registration + m.reflectTypeToType = nil + for k := range m.memoizedEncodePlans { + delete(m.memoizedEncodePlans, k) + } +} + +// TypeForOID returns the [Type] registered for the given OID. The returned [Type] must not be mutated. +func (m *Map) TypeForOID(oid uint32) (*Type, bool) { + if dt, ok := m.oidToType[oid]; ok { + return dt, true + } + + dt, ok := defaultMap.oidToType[oid] + return dt, ok +} + +// TypeForName returns the [Type] registered for the given name. The returned [Type] must not be mutated. +func (m *Map) TypeForName(name string) (*Type, bool) { + if dt, ok := m.nameToType[name]; ok { + return dt, true + } + dt, ok := defaultMap.nameToType[name] + return dt, ok +} + +func (m *Map) buildReflectTypeToType() { + m.reflectTypeToType = make(map[reflect.Type]*Type) + + for reflectType, name := range m.reflectTypeToName { + if dt, ok := m.TypeForName(name); ok { + m.reflectTypeToType[reflectType] = dt + } + } +} + +// TypeForValue finds a data type suitable for v. Use [Map.RegisterType] to register types that can encode and decode +// themselves. Use [Map.RegisterDefaultPgType] to register that can be handled by a registered data type. The returned [Type] +// must not be mutated. +func (m *Map) TypeForValue(v any) (*Type, bool) { + if m.reflectTypeToType == nil { + m.buildReflectTypeToType() + } + + if dt, ok := m.reflectTypeToType[reflect.TypeOf(v)]; ok { + return dt, true + } + + dt, ok := defaultMap.reflectTypeToType[reflect.TypeOf(v)] + return dt, ok +} + +// FormatCodeForOID returns the preferred format code for type oid. If the type is not registered it returns the text +// format code. +func (m *Map) FormatCodeForOID(oid uint32) int16 { + if fc, ok := m.oidToFormatCode[oid]; ok { + return fc + } + + if fc, ok := defaultMap.oidToFormatCode[oid]; ok { + return fc + } + + return TextFormatCode +} + +// EncodePlan is a precompiled plan to encode a particular type into a particular OID and format. +type EncodePlan interface { + // Encode appends the encoded bytes of value to buf. If value is the SQL value NULL then append nothing and return + // (nil, nil). The caller of Encode is responsible for writing the correct NULL value or the length of the data + // written. + Encode(value any, buf []byte) (newBuf []byte, err error) +} + +// ScanPlan is a precompiled plan to scan into a type of destination. +type ScanPlan interface { + // Scan scans src into target. src is only valid during the call to Scan. The ScanPlan must not retain a reference to + // src. + Scan(src []byte, target any) error +} + +type scanPlanCodecSQLScanner struct { + c Codec + m *Map + oid uint32 + formatCode int16 +} + +func (plan *scanPlanCodecSQLScanner) Scan(src []byte, dst any) error { + value, err := plan.c.DecodeDatabaseSQLValue(plan.m, plan.oid, plan.formatCode, src) + if err != nil { + return err + } + + scanner := dst.(sql.Scanner) + return scanner.Scan(value) +} + +type scanPlanSQLScanner struct { + formatCode int16 +} + +func (plan *scanPlanSQLScanner) Scan(src []byte, dst any) error { + scanner := dst.(sql.Scanner) + + switch { + case src == nil: + // This is necessary because interface value []byte:nil does not equal nil:nil for the binary format path and the + // text format path would be converted to empty string. + return scanner.Scan(nil) + case plan.formatCode == BinaryFormatCode: + return scanner.Scan(src) + default: + return scanner.Scan(string(src)) + } +} + +type scanPlanString struct{} + +func (scanPlanString) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + p := (dst).(*string) + *p = string(src) + return nil +} + +type scanPlanAnyTextToBytes struct{} + +func (scanPlanAnyTextToBytes) Scan(src []byte, dst any) error { + dstBuf := dst.(*[]byte) + if src == nil { + *dstBuf = nil + return nil + } + + *dstBuf = make([]byte, len(src)) + copy(*dstBuf, src) + return nil +} + +type scanPlanFail struct { + m *Map + oid uint32 + formatCode int16 +} + +func (plan *scanPlanFail) Scan(src []byte, dst any) error { + // If src is NULL it might be possible to scan into dst even though it is the types are not compatible. While this + // may seem to be a contrived case it can occur when selecting NULL directly. PostgreSQL assigns it the type of text. + // It would be surprising to the caller to have to cast the NULL (e.g. `select null::int`). So try to figure out a + // compatible data type for dst and scan with that. + // + // See https://github.com/jackc/pgx/issues/1326 + if src == nil { + // As a horrible hack try all types to find anything that can scan into dst. + for oid := range plan.m.oidToType { + // using planScan instead of Scan or PlanScan to avoid polluting the planned scan cache. + plan := plan.m.planScan(oid, plan.formatCode, dst, 0) + if _, ok := plan.(*scanPlanFail); !ok { + return plan.Scan(src, dst) + } + } + for oid := range defaultMap.oidToType { + if _, ok := plan.m.oidToType[oid]; !ok { + plan := plan.m.planScan(oid, plan.formatCode, dst, 0) + if _, ok := plan.(*scanPlanFail); !ok { + return plan.Scan(src, dst) + } + } + } + } + + var format string + switch plan.formatCode { + case TextFormatCode: + format = "text" + case BinaryFormatCode: + format = "binary" + default: + format = fmt.Sprintf("unknown %d", plan.formatCode) + } + + var dataTypeName string + if t, ok := plan.m.TypeForOID(plan.oid); ok { + dataTypeName = t.Name + } else { + dataTypeName = "unknown type" + } + + return fmt.Errorf("cannot scan %s (OID %d) in %v format into %T", dataTypeName, plan.oid, format, dst) +} + +// TryWrapScanPlanFunc is a function that tries to create a wrapper plan for target. If successful it returns a plan +// that will convert the target passed to Scan and then call the next plan. nextTarget is target as it will be converted +// by plan. It must be used to find another suitable ScanPlan. When it is found SetNext must be called on plan for it +// to be usabled. ok indicates if a suitable wrapper was found. +type TryWrapScanPlanFunc func(target any) (plan WrappedScanPlanNextSetter, nextTarget any, ok bool) + +type pointerPointerScanPlan struct { + dstType reflect.Type + next ScanPlan +} + +func (plan *pointerPointerScanPlan) SetNext(next ScanPlan) { plan.next = next } + +func (plan *pointerPointerScanPlan) Scan(src []byte, dst any) error { + el := reflect.ValueOf(dst).Elem() + if src == nil { + el.Set(reflect.Zero(el.Type())) + return nil + } + + el.Set(reflect.New(el.Type().Elem())) + return plan.next.Scan(src, el.Interface()) +} + +// TryPointerPointerScanPlan handles a pointer to a pointer by setting the target to nil for SQL NULL and allocating and +// scanning for non-NULL. +func TryPointerPointerScanPlan(target any) (plan WrappedScanPlanNextSetter, nextTarget any, ok bool) { + if dstValue := reflect.ValueOf(target); dstValue.Kind() == reflect.Pointer { + elemValue := dstValue.Elem() + if elemValue.Kind() == reflect.Pointer { + plan = &pointerPointerScanPlan{dstType: dstValue.Type()} + return plan, reflect.Zero(elemValue.Type()).Interface(), true + } + } + + return nil, nil, false +} + +// SkipUnderlyingTypePlanner prevents PlanScan and PlanDecode from trying to use the underlying type. +type SkipUnderlyingTypePlanner interface { + SkipUnderlyingTypePlan() +} + +var elemKindToPointerTypes map[reflect.Kind]reflect.Type = map[reflect.Kind]reflect.Type{ + reflect.Int: reflect.TypeFor[*int](), + reflect.Int8: reflect.TypeFor[*int8](), + reflect.Int16: reflect.TypeFor[*int16](), + reflect.Int32: reflect.TypeFor[*int32](), + reflect.Int64: reflect.TypeFor[*int64](), + reflect.Uint: reflect.TypeFor[*uint](), + reflect.Uint8: reflect.TypeFor[*uint8](), + reflect.Uint16: reflect.TypeFor[*uint16](), + reflect.Uint32: reflect.TypeFor[*uint32](), + reflect.Uint64: reflect.TypeFor[*uint64](), + reflect.Float32: reflect.TypeFor[*float32](), + reflect.Float64: reflect.TypeFor[*float64](), + reflect.String: reflect.TypeFor[*string](), + reflect.Bool: reflect.TypeFor[*bool](), +} + +type underlyingTypeScanPlan struct { + dstType reflect.Type + nextDstType reflect.Type + next ScanPlan +} + +func (plan *underlyingTypeScanPlan) SetNext(next ScanPlan) { plan.next = next } + +func (plan *underlyingTypeScanPlan) Scan(src []byte, dst any) error { + return plan.next.Scan(src, reflect.ValueOf(dst).Convert(plan.nextDstType).Interface()) +} + +// TryFindUnderlyingTypeScanPlan tries to convert to a Go builtin type. e.g. If value was of type MyString and +// MyString was defined as a string then a wrapper plan would be returned that converts MyString to string. +func TryFindUnderlyingTypeScanPlan(dst any) (plan WrappedScanPlanNextSetter, nextDst any, ok bool) { + if _, ok := dst.(SkipUnderlyingTypePlanner); ok { + return nil, nil, false + } + + dstValue := reflect.ValueOf(dst) + + if dstValue.Kind() == reflect.Pointer { + var elemValue reflect.Value + if dstValue.IsNil() { + elemValue = reflect.New(dstValue.Type().Elem()).Elem() + } else { + elemValue = dstValue.Elem() + } + nextDstType := elemKindToPointerTypes[elemValue.Kind()] + if nextDstType == nil { + if elemValue.Kind() == reflect.Slice { + if elemValue.Type().Elem().Kind() == reflect.Uint8 { + var v *[]byte + nextDstType = reflect.TypeOf(v) + } + } + + // Get underlying type of any array. + // https://github.com/jackc/pgx/issues/2107 + if elemValue.Kind() == reflect.Array { + nextDstType = reflect.PointerTo(reflect.ArrayOf(elemValue.Len(), elemValue.Type().Elem())) + } + } + + if nextDstType != nil && dstValue.Type() != nextDstType && dstValue.CanConvert(nextDstType) { + return &underlyingTypeScanPlan{dstType: dstValue.Type(), nextDstType: nextDstType}, dstValue.Convert(nextDstType).Interface(), true + } + } + + return nil, nil, false +} + +type WrappedScanPlanNextSetter interface { + SetNext(ScanPlan) + ScanPlan +} + +// TryWrapBuiltinTypeScanPlan tries to wrap a builtin type with a wrapper that provides additional methods. e.g. If +// value was of type int32 then a wrapper plan would be returned that converts target to a value that implements +// Int64Scanner. +func TryWrapBuiltinTypeScanPlan(target any) (plan WrappedScanPlanNextSetter, nextDst any, ok bool) { + switch target := target.(type) { + case *int8: + return &wrapInt8ScanPlan{}, (*int8Wrapper)(target), true + case *int16: + return &wrapInt16ScanPlan{}, (*int16Wrapper)(target), true + case *int32: + return &wrapInt32ScanPlan{}, (*int32Wrapper)(target), true + case *int64: + return &wrapInt64ScanPlan{}, (*int64Wrapper)(target), true + case *int: + return &wrapIntScanPlan{}, (*intWrapper)(target), true + case *uint8: + return &wrapUint8ScanPlan{}, (*uint8Wrapper)(target), true + case *uint16: + return &wrapUint16ScanPlan{}, (*uint16Wrapper)(target), true + case *uint32: + return &wrapUint32ScanPlan{}, (*uint32Wrapper)(target), true + case *uint64: + return &wrapUint64ScanPlan{}, (*uint64Wrapper)(target), true + case *uint: + return &wrapUintScanPlan{}, (*uintWrapper)(target), true + case *float32: + return &wrapFloat32ScanPlan{}, (*float32Wrapper)(target), true + case *float64: + return &wrapFloat64ScanPlan{}, (*float64Wrapper)(target), true + case *string: + return &wrapStringScanPlan{}, (*stringWrapper)(target), true + case *time.Time: + return &wrapTimeScanPlan{}, (*timeWrapper)(target), true + case *time.Duration: + return &wrapDurationScanPlan{}, (*durationWrapper)(target), true + case *net.IPNet: + return &wrapNetIPNetScanPlan{}, (*netIPNetWrapper)(target), true + case *net.IP: + return &wrapNetIPScanPlan{}, (*netIPWrapper)(target), true + case *netip.Prefix: + return &wrapNetipPrefixScanPlan{}, (*netipPrefixWrapper)(target), true + case *netip.Addr: + return &wrapNetipAddrScanPlan{}, (*netipAddrWrapper)(target), true + case *map[string]*string: + return &wrapMapStringToPointerStringScanPlan{}, (*mapStringToPointerStringWrapper)(target), true + case *map[string]string: + return &wrapMapStringToStringScanPlan{}, (*mapStringToStringWrapper)(target), true + case *[16]byte: + return &wrapByte16ScanPlan{}, (*byte16Wrapper)(target), true + case *[]byte: + return &wrapByteSliceScanPlan{}, (*byteSliceWrapper)(target), true + } + + return nil, nil, false +} + +type wrapInt8ScanPlan struct { + next ScanPlan +} + +func (plan *wrapInt8ScanPlan) SetNext(next ScanPlan) { plan.next = next } + +func (plan *wrapInt8ScanPlan) Scan(src []byte, dst any) error { + return plan.next.Scan(src, (*int8Wrapper)(dst.(*int8))) +} + +type wrapInt16ScanPlan struct { + next ScanPlan +} + +func (plan *wrapInt16ScanPlan) SetNext(next ScanPlan) { plan.next = next } + +func (plan *wrapInt16ScanPlan) Scan(src []byte, dst any) error { + return plan.next.Scan(src, (*int16Wrapper)(dst.(*int16))) +} + +type wrapInt32ScanPlan struct { + next ScanPlan +} + +func (plan *wrapInt32ScanPlan) SetNext(next ScanPlan) { plan.next = next } + +func (plan *wrapInt32ScanPlan) Scan(src []byte, dst any) error { + return plan.next.Scan(src, (*int32Wrapper)(dst.(*int32))) +} + +type wrapInt64ScanPlan struct { + next ScanPlan +} + +func (plan *wrapInt64ScanPlan) SetNext(next ScanPlan) { plan.next = next } + +func (plan *wrapInt64ScanPlan) Scan(src []byte, dst any) error { + return plan.next.Scan(src, (*int64Wrapper)(dst.(*int64))) +} + +type wrapIntScanPlan struct { + next ScanPlan +} + +func (plan *wrapIntScanPlan) SetNext(next ScanPlan) { plan.next = next } + +func (plan *wrapIntScanPlan) Scan(src []byte, dst any) error { + return plan.next.Scan(src, (*intWrapper)(dst.(*int))) +} + +type wrapUint8ScanPlan struct { + next ScanPlan +} + +func (plan *wrapUint8ScanPlan) SetNext(next ScanPlan) { plan.next = next } + +func (plan *wrapUint8ScanPlan) Scan(src []byte, dst any) error { + return plan.next.Scan(src, (*uint8Wrapper)(dst.(*uint8))) +} + +type wrapUint16ScanPlan struct { + next ScanPlan +} + +func (plan *wrapUint16ScanPlan) SetNext(next ScanPlan) { plan.next = next } + +func (plan *wrapUint16ScanPlan) Scan(src []byte, dst any) error { + return plan.next.Scan(src, (*uint16Wrapper)(dst.(*uint16))) +} + +type wrapUint32ScanPlan struct { + next ScanPlan +} + +func (plan *wrapUint32ScanPlan) SetNext(next ScanPlan) { plan.next = next } + +func (plan *wrapUint32ScanPlan) Scan(src []byte, dst any) error { + return plan.next.Scan(src, (*uint32Wrapper)(dst.(*uint32))) +} + +type wrapUint64ScanPlan struct { + next ScanPlan +} + +func (plan *wrapUint64ScanPlan) SetNext(next ScanPlan) { plan.next = next } + +func (plan *wrapUint64ScanPlan) Scan(src []byte, dst any) error { + return plan.next.Scan(src, (*uint64Wrapper)(dst.(*uint64))) +} + +type wrapUintScanPlan struct { + next ScanPlan +} + +func (plan *wrapUintScanPlan) SetNext(next ScanPlan) { plan.next = next } + +func (plan *wrapUintScanPlan) Scan(src []byte, dst any) error { + return plan.next.Scan(src, (*uintWrapper)(dst.(*uint))) +} + +type wrapFloat32ScanPlan struct { + next ScanPlan +} + +func (plan *wrapFloat32ScanPlan) SetNext(next ScanPlan) { plan.next = next } + +func (plan *wrapFloat32ScanPlan) Scan(src []byte, dst any) error { + return plan.next.Scan(src, (*float32Wrapper)(dst.(*float32))) +} + +type wrapFloat64ScanPlan struct { + next ScanPlan +} + +func (plan *wrapFloat64ScanPlan) SetNext(next ScanPlan) { plan.next = next } + +func (plan *wrapFloat64ScanPlan) Scan(src []byte, dst any) error { + return plan.next.Scan(src, (*float64Wrapper)(dst.(*float64))) +} + +type wrapStringScanPlan struct { + next ScanPlan +} + +func (plan *wrapStringScanPlan) SetNext(next ScanPlan) { plan.next = next } + +func (plan *wrapStringScanPlan) Scan(src []byte, dst any) error { + return plan.next.Scan(src, (*stringWrapper)(dst.(*string))) +} + +type wrapTimeScanPlan struct { + next ScanPlan +} + +func (plan *wrapTimeScanPlan) SetNext(next ScanPlan) { plan.next = next } + +func (plan *wrapTimeScanPlan) Scan(src []byte, dst any) error { + return plan.next.Scan(src, (*timeWrapper)(dst.(*time.Time))) +} + +type wrapDurationScanPlan struct { + next ScanPlan +} + +func (plan *wrapDurationScanPlan) SetNext(next ScanPlan) { plan.next = next } + +func (plan *wrapDurationScanPlan) Scan(src []byte, dst any) error { + return plan.next.Scan(src, (*durationWrapper)(dst.(*time.Duration))) +} + +type wrapNetIPNetScanPlan struct { + next ScanPlan +} + +func (plan *wrapNetIPNetScanPlan) SetNext(next ScanPlan) { plan.next = next } + +func (plan *wrapNetIPNetScanPlan) Scan(src []byte, dst any) error { + return plan.next.Scan(src, (*netIPNetWrapper)(dst.(*net.IPNet))) +} + +type wrapNetIPScanPlan struct { + next ScanPlan +} + +func (plan *wrapNetIPScanPlan) SetNext(next ScanPlan) { plan.next = next } + +func (plan *wrapNetIPScanPlan) Scan(src []byte, dst any) error { + return plan.next.Scan(src, (*netIPWrapper)(dst.(*net.IP))) +} + +type wrapNetipPrefixScanPlan struct { + next ScanPlan +} + +func (plan *wrapNetipPrefixScanPlan) SetNext(next ScanPlan) { plan.next = next } + +func (plan *wrapNetipPrefixScanPlan) Scan(src []byte, dst any) error { + return plan.next.Scan(src, (*netipPrefixWrapper)(dst.(*netip.Prefix))) +} + +type wrapNetipAddrScanPlan struct { + next ScanPlan +} + +func (plan *wrapNetipAddrScanPlan) SetNext(next ScanPlan) { plan.next = next } + +func (plan *wrapNetipAddrScanPlan) Scan(src []byte, dst any) error { + return plan.next.Scan(src, (*netipAddrWrapper)(dst.(*netip.Addr))) +} + +type wrapMapStringToPointerStringScanPlan struct { + next ScanPlan +} + +func (plan *wrapMapStringToPointerStringScanPlan) SetNext(next ScanPlan) { plan.next = next } + +func (plan *wrapMapStringToPointerStringScanPlan) Scan(src []byte, dst any) error { + return plan.next.Scan(src, (*mapStringToPointerStringWrapper)(dst.(*map[string]*string))) +} + +type wrapMapStringToStringScanPlan struct { + next ScanPlan +} + +func (plan *wrapMapStringToStringScanPlan) SetNext(next ScanPlan) { plan.next = next } + +func (plan *wrapMapStringToStringScanPlan) Scan(src []byte, dst any) error { + return plan.next.Scan(src, (*mapStringToStringWrapper)(dst.(*map[string]string))) +} + +type wrapByte16ScanPlan struct { + next ScanPlan +} + +func (plan *wrapByte16ScanPlan) SetNext(next ScanPlan) { plan.next = next } + +func (plan *wrapByte16ScanPlan) Scan(src []byte, dst any) error { + return plan.next.Scan(src, (*byte16Wrapper)(dst.(*[16]byte))) +} + +type wrapByteSliceScanPlan struct { + next ScanPlan +} + +func (plan *wrapByteSliceScanPlan) SetNext(next ScanPlan) { plan.next = next } + +func (plan *wrapByteSliceScanPlan) Scan(src []byte, dst any) error { + return plan.next.Scan(src, (*byteSliceWrapper)(dst.(*[]byte))) +} + +type pointerEmptyInterfaceScanPlan struct { + codec Codec + m *Map + oid uint32 + formatCode int16 +} + +func (plan *pointerEmptyInterfaceScanPlan) Scan(src []byte, dst any) error { + value, err := plan.codec.DecodeValue(plan.m, plan.oid, plan.formatCode, src) + if err != nil { + return err + } + + ptrAny := dst.(*any) + *ptrAny = value + + return nil +} + +// TryWrapStructScanPlan tries to wrap a struct with a wrapper that implements CompositeIndexGetter. +func TryWrapStructScanPlan(target any) (plan WrappedScanPlanNextSetter, nextValue any, ok bool) { + targetValue := reflect.ValueOf(target) + if targetValue.Kind() != reflect.Pointer { + return nil, nil, false + } + + var targetElemValue reflect.Value + if targetValue.IsNil() { + targetElemValue = reflect.Zero(targetValue.Type().Elem()) + } else { + targetElemValue = targetValue.Elem() + } + targetElemType := targetElemValue.Type() + + if targetElemType.Kind() == reflect.Struct { + exportedFields := getExportedFieldValues(targetElemValue) + if len(exportedFields) == 0 { + return nil, nil, false + } + + w := ptrStructWrapper{ + s: target, + exportedFields: exportedFields, + } + return &wrapAnyPtrStructScanPlan{}, &w, true + } + + return nil, nil, false +} + +type wrapAnyPtrStructScanPlan struct { + next ScanPlan +} + +func (plan *wrapAnyPtrStructScanPlan) SetNext(next ScanPlan) { plan.next = next } + +func (plan *wrapAnyPtrStructScanPlan) Scan(src []byte, target any) error { + w := ptrStructWrapper{ + s: target, + exportedFields: getExportedFieldValues(reflect.ValueOf(target).Elem()), + } + + return plan.next.Scan(src, &w) +} + +// TryWrapPtrSliceScanPlan tries to wrap a pointer to a single dimension slice. +func TryWrapPtrSliceScanPlan(target any) (plan WrappedScanPlanNextSetter, nextValue any, ok bool) { + // Avoid using reflect path for common types. + switch target := target.(type) { + case *[]int16: + return &wrapPtrSliceScanPlan[int16]{}, (*FlatArray[int16])(target), true + case *[]int32: + return &wrapPtrSliceScanPlan[int32]{}, (*FlatArray[int32])(target), true + case *[]int64: + return &wrapPtrSliceScanPlan[int64]{}, (*FlatArray[int64])(target), true + case *[]float32: + return &wrapPtrSliceScanPlan[float32]{}, (*FlatArray[float32])(target), true + case *[]float64: + return &wrapPtrSliceScanPlan[float64]{}, (*FlatArray[float64])(target), true + case *[]string: + return &wrapPtrSliceScanPlan[string]{}, (*FlatArray[string])(target), true + case *[]time.Time: + return &wrapPtrSliceScanPlan[time.Time]{}, (*FlatArray[time.Time])(target), true + } + + targetType := reflect.TypeOf(target) + if targetType.Kind() != reflect.Pointer { + return nil, nil, false + } + + targetElemType := targetType.Elem() + + if targetElemType.Kind() == reflect.Slice { + slice := reflect.New(targetElemType).Elem() + return &wrapPtrSliceReflectScanPlan{}, &anySliceArrayReflect{slice: slice}, true + } + return nil, nil, false +} + +type wrapPtrSliceScanPlan[T any] struct { + next ScanPlan +} + +func (plan *wrapPtrSliceScanPlan[T]) SetNext(next ScanPlan) { plan.next = next } + +func (plan *wrapPtrSliceScanPlan[T]) Scan(src []byte, target any) error { + return plan.next.Scan(src, (*FlatArray[T])(target.(*[]T))) +} + +type wrapPtrSliceReflectScanPlan struct { + next ScanPlan +} + +func (plan *wrapPtrSliceReflectScanPlan) SetNext(next ScanPlan) { plan.next = next } + +func (plan *wrapPtrSliceReflectScanPlan) Scan(src []byte, target any) error { + return plan.next.Scan(src, &anySliceArrayReflect{slice: reflect.ValueOf(target).Elem()}) +} + +// TryWrapPtrMultiDimSliceScanPlan tries to wrap a pointer to a multi-dimension slice. +func TryWrapPtrMultiDimSliceScanPlan(target any) (plan WrappedScanPlanNextSetter, nextValue any, ok bool) { + targetValue := reflect.ValueOf(target) + if targetValue.Kind() != reflect.Pointer { + return nil, nil, false + } + + targetElemValue := targetValue.Elem() + + if targetElemValue.Kind() == reflect.Slice { + elemElemKind := targetElemValue.Type().Elem().Kind() + if elemElemKind == reflect.Slice { + if !isRagged(targetElemValue) { + return &wrapPtrMultiDimSliceScanPlan{}, &anyMultiDimSliceArray{slice: targetValue.Elem()}, true + } + } + } + + return nil, nil, false +} + +type wrapPtrMultiDimSliceScanPlan struct { + next ScanPlan +} + +func (plan *wrapPtrMultiDimSliceScanPlan) SetNext(next ScanPlan) { plan.next = next } + +func (plan *wrapPtrMultiDimSliceScanPlan) Scan(src []byte, target any) error { + return plan.next.Scan(src, &anyMultiDimSliceArray{slice: reflect.ValueOf(target).Elem()}) +} + +// TryWrapPtrArrayScanPlan tries to wrap a pointer to a single dimension array. +func TryWrapPtrArrayScanPlan(target any) (plan WrappedScanPlanNextSetter, nextValue any, ok bool) { + targetValue := reflect.ValueOf(target) + if targetValue.Kind() != reflect.Pointer { + return nil, nil, false + } + + targetElemValue := targetValue.Elem() + + if targetElemValue.Kind() == reflect.Array { + return &wrapPtrArrayReflectScanPlan{}, &anyArrayArrayReflect{array: targetElemValue}, true + } + return nil, nil, false +} + +type wrapPtrArrayReflectScanPlan struct { + next ScanPlan +} + +func (plan *wrapPtrArrayReflectScanPlan) SetNext(next ScanPlan) { plan.next = next } + +func (plan *wrapPtrArrayReflectScanPlan) Scan(src []byte, target any) error { + return plan.next.Scan(src, &anyArrayArrayReflect{array: reflect.ValueOf(target).Elem()}) +} + +// PlanScan prepares a plan to scan a value into target. +func (m *Map) PlanScan(oid uint32, formatCode int16, target any) ScanPlan { + return m.planScan(oid, formatCode, target, 0) +} + +func (m *Map) planScan(oid uint32, formatCode int16, target any, depth int) ScanPlan { + if depth > 8 { + return &scanPlanFail{m: m, oid: oid, formatCode: formatCode} + } + + if target == nil { + return &scanPlanFail{m: m, oid: oid, formatCode: formatCode} + } + + if _, ok := target.(*UndecodedBytes); ok { + return scanPlanAnyToUndecodedBytes{} + } + + switch formatCode { + case BinaryFormatCode: + if _, ok := target.(*string); ok { + switch oid { + case TextOID, VarcharOID: + return scanPlanString{} + } + } + case TextFormatCode: + switch target.(type) { + case *string: + return scanPlanString{} + case *[]byte: + if oid != ByteaOID { + return scanPlanAnyTextToBytes{} + } + case TextScanner: + return scanPlanTextAnyToTextScanner{} + } + } + + var dt *Type + + if dataType, ok := m.TypeForOID(oid); ok { + dt = dataType + } else if dataType, ok := m.TypeForValue(target); ok { + dt = dataType + oid = dt.OID // Preserve assumed OID in case we are recursively called below. + } + + if dt != nil { + if plan := dt.Codec.PlanScan(m, oid, formatCode, target); plan != nil { + return plan + } + } + + // This needs to happen before trying m.TryWrapScanPlanFuncs. Otherwise, a sql.Scanner would not get called if it was + // defined on a type that could be unwrapped such as `type myString string`. + // + // https://github.com/jackc/pgtype/issues/197 + if _, ok := target.(sql.Scanner); ok { + if dt == nil { + return &scanPlanSQLScanner{formatCode: formatCode} + } else { + return &scanPlanCodecSQLScanner{c: dt.Codec, m: m, oid: oid, formatCode: formatCode} + } + } + + for _, f := range m.TryWrapScanPlanFuncs { + if wrapperPlan, nextDst, ok := f(target); ok { + if nextPlan := m.planScan(oid, formatCode, nextDst, depth+1); nextPlan != nil { + if _, failed := nextPlan.(*scanPlanFail); !failed { + wrapperPlan.SetNext(nextPlan) + return wrapperPlan + } + } + } + } + + if _, ok := target.(*any); ok { + var codec Codec + if dt != nil { + codec = dt.Codec + } else { + if formatCode == TextFormatCode { + codec = TextCodec{} + } else { + codec = ByteaCodec{} + } + } + return &pointerEmptyInterfaceScanPlan{codec: codec, m: m, oid: oid, formatCode: formatCode} + } + + return &scanPlanFail{m: m, oid: oid, formatCode: formatCode} +} + +func (m *Map) Scan(oid uint32, formatCode int16, src []byte, dst any) error { + if dst == nil { + return nil + } + + plan := m.PlanScan(oid, formatCode, dst) + return plan.Scan(src, dst) +} + +var ErrScanTargetTypeChanged = errors.New("scan target type changed") + +func codecScan(codec Codec, m *Map, oid uint32, format int16, src []byte, dst any) error { + scanPlan := codec.PlanScan(m, oid, format, dst) + if scanPlan == nil { + return fmt.Errorf("PlanScan did not find a plan") + } + return scanPlan.Scan(src, dst) +} + +func codecDecodeToTextFormat(codec Codec, m *Map, oid uint32, format int16, src []byte) (driver.Value, error) { + if src == nil { + return nil, nil + } + + if format == TextFormatCode { + return string(src), nil + } else { + value, err := codec.DecodeValue(m, oid, format, src) + if err != nil { + return nil, err + } + buf, err := m.Encode(oid, TextFormatCode, value, nil) + if err != nil { + return nil, err + } + return string(buf), nil + } +} + +// PlanEncode returns an EncodePlan for encoding value into PostgreSQL format for oid and format. If no plan can be +// found then nil is returned. +func (m *Map) PlanEncode(oid uint32, format int16, value any) EncodePlan { + return m.planEncodeDepth(oid, format, value, 0) +} + +func (m *Map) planEncodeDepth(oid uint32, format int16, value any, depth int) EncodePlan { + // Guard against infinite recursion. + if depth > 8 { + return nil + } + + oidMemo := m.memoizedEncodePlans[oid] + if oidMemo == nil { + oidMemo = make(map[reflect.Type][2]EncodePlan) + m.memoizedEncodePlans[oid] = oidMemo + } + targetReflectType := reflect.TypeOf(value) + typeMemo := oidMemo[targetReflectType] + plan := typeMemo[format] + if plan == nil { + plan = m.planEncode(oid, format, value, depth) + typeMemo[format] = plan + oidMemo[targetReflectType] = typeMemo + } + + return plan +} + +func (m *Map) planEncode(oid uint32, format int16, value any, depth int) EncodePlan { + if format == TextFormatCode { + switch value.(type) { + case string: + return encodePlanStringToAnyTextFormat{} + case TextValuer: + return encodePlanTextValuerToAnyTextFormat{} + } + } + + var dt *Type + if dataType, ok := m.TypeForOID(oid); ok { + dt = dataType + } else { + // If no type for the OID was found, then either it is unknowable (e.g. the simple protocol) or it is an + // unregistered type. In either case try to find the type and OID that matches the value (e.g. a []byte would be + // registered to PostgreSQL bytea). + if dataType, ok := m.TypeForValue(value); ok { + dt = dataType + oid = dt.OID // Preserve assumed OID in case we are recursively called below. + } + } + + if dt != nil { + if plan := dt.Codec.PlanEncode(m, oid, format, value); plan != nil { + return plan + } + } + + for _, f := range m.TryWrapEncodePlanFuncs { + if wrapperPlan, nextValue, ok := f(value); ok { + if nextPlan := m.planEncodeDepth(oid, format, nextValue, depth+1); nextPlan != nil { + wrapperPlan.SetNext(nextPlan) + return wrapperPlan + } + } + } + + if _, ok := value.(driver.Valuer); ok { + return &encodePlanDriverValuer{m: m, oid: oid, formatCode: format} + } + + return nil +} + +type encodePlanStringToAnyTextFormat struct{} + +func (encodePlanStringToAnyTextFormat) Encode(value any, buf []byte) (newBuf []byte, err error) { + s := value.(string) + return append(buf, s...), nil +} + +type encodePlanTextValuerToAnyTextFormat struct{} + +func (encodePlanTextValuerToAnyTextFormat) Encode(value any, buf []byte) (newBuf []byte, err error) { + t, err := value.(TextValuer).TextValue() + if err != nil { + return nil, err + } + if !t.Valid { + return nil, nil + } + + return append(buf, t.String...), nil +} + +type encodePlanDriverValuer struct { + m *Map + oid uint32 + formatCode int16 +} + +func (plan *encodePlanDriverValuer) Encode(value any, buf []byte) (newBuf []byte, err error) { + dv := value.(driver.Valuer) + if dv == nil { + return nil, nil + } + v, err := dv.Value() + if err != nil { + return nil, err + } + if v == nil { + return nil, nil + } + + newBuf, err = plan.m.Encode(plan.oid, plan.formatCode, v, buf) + if err == nil { + return newBuf, nil + } + + s, ok := v.(string) + if !ok { + return nil, err + } + + var scannedValue any + scanErr := plan.m.Scan(plan.oid, TextFormatCode, []byte(s), &scannedValue) + if scanErr != nil { + return nil, err + } + + // Prevent infinite loop. We can't encode this. See https://github.com/jackc/pgx/issues/1331. + if reflect.TypeOf(value) == reflect.TypeOf(scannedValue) { + return nil, fmt.Errorf("tried to encode %v via encoding to text and scanning but failed due to receiving same type back", value) + } + + var err2 error + newBuf, err2 = plan.m.Encode(plan.oid, BinaryFormatCode, scannedValue, buf) + if err2 != nil { + return nil, err + } + + return newBuf, nil +} + +// TryWrapEncodePlanFunc is a function that tries to create a wrapper plan for value. If successful it returns a plan +// that will convert the value passed to Encode and then call the next plan. nextValue is value as it will be converted +// by plan. It must be used to find another suitable EncodePlan. When it is found SetNext must be called on plan for it +// to be usabled. ok indicates if a suitable wrapper was found. +type TryWrapEncodePlanFunc func(value any) (plan WrappedEncodePlanNextSetter, nextValue any, ok bool) + +type derefPointerEncodePlan struct { + next EncodePlan +} + +func (plan *derefPointerEncodePlan) SetNext(next EncodePlan) { plan.next = next } + +func (plan *derefPointerEncodePlan) Encode(value any, buf []byte) (newBuf []byte, err error) { + ptr := reflect.ValueOf(value) + + if ptr.IsNil() { + return nil, nil + } + + return plan.next.Encode(ptr.Elem().Interface(), buf) +} + +// TryWrapDerefPointerEncodePlan tries to dereference a pointer. e.g. If value was of type *string then a wrapper plan +// would be returned that dereferences the value. +func TryWrapDerefPointerEncodePlan(value any) (plan WrappedEncodePlanNextSetter, nextValue any, ok bool) { + if _, ok := value.(driver.Valuer); ok { + return nil, nil, false + } + + if valueType := reflect.TypeOf(value); valueType != nil && valueType.Kind() == reflect.Pointer { + return &derefPointerEncodePlan{}, reflect.New(valueType.Elem()).Elem().Interface(), true + } + + return nil, nil, false +} + +var kindToTypes map[reflect.Kind]reflect.Type = map[reflect.Kind]reflect.Type{ + reflect.Int: reflect.TypeFor[int](), + reflect.Int8: reflect.TypeFor[int8](), + reflect.Int16: reflect.TypeFor[int16](), + reflect.Int32: reflect.TypeFor[int32](), + reflect.Int64: reflect.TypeFor[int64](), + reflect.Uint: reflect.TypeFor[uint](), + reflect.Uint8: reflect.TypeFor[uint8](), + reflect.Uint16: reflect.TypeFor[uint16](), + reflect.Uint32: reflect.TypeFor[uint32](), + reflect.Uint64: reflect.TypeFor[uint64](), + reflect.Float32: reflect.TypeFor[float32](), + reflect.Float64: reflect.TypeFor[float64](), + reflect.String: reflect.TypeFor[string](), + reflect.Bool: reflect.TypeFor[bool](), +} + +var byteSliceType = reflect.TypeFor[[]byte]() + +type underlyingTypeEncodePlan struct { + nextValueType reflect.Type + next EncodePlan +} + +func (plan *underlyingTypeEncodePlan) SetNext(next EncodePlan) { plan.next = next } + +func (plan *underlyingTypeEncodePlan) Encode(value any, buf []byte) (newBuf []byte, err error) { + return plan.next.Encode(reflect.ValueOf(value).Convert(plan.nextValueType).Interface(), buf) +} + +// TryWrapFindUnderlyingTypeEncodePlan tries to convert to a Go builtin type. e.g. If value was of type MyString and +// MyString was defined as a string then a wrapper plan would be returned that converts MyString to string. +func TryWrapFindUnderlyingTypeEncodePlan(value any) (plan WrappedEncodePlanNextSetter, nextValue any, ok bool) { + if value == nil { + return nil, nil, false + } + + if _, ok := value.(driver.Valuer); ok { + return nil, nil, false + } + + if _, ok := value.(SkipUnderlyingTypePlanner); ok { + return nil, nil, false + } + + refValue := reflect.ValueOf(value) + + nextValueType := kindToTypes[refValue.Kind()] + if nextValueType != nil && refValue.Type() != nextValueType { + return &underlyingTypeEncodePlan{nextValueType: nextValueType}, refValue.Convert(nextValueType).Interface(), true + } + + // []byte is a special case. It is a slice but we treat it as a scalar type. In the case of a named type like + // json.RawMessage which is defined as []byte the underlying type should be considered as []byte. But any other slice + // does not have a special underlying type. + // + // https://github.com/jackc/pgx/issues/1763 + if refValue.Type() != byteSliceType && refValue.Type().AssignableTo(byteSliceType) { + return &underlyingTypeEncodePlan{nextValueType: byteSliceType}, refValue.Convert(byteSliceType).Interface(), true + } + + // Get underlying type of any array. + // https://github.com/jackc/pgx/issues/2107 + if refValue.Kind() == reflect.Array { + underlyingArrayType := reflect.ArrayOf(refValue.Len(), refValue.Type().Elem()) + if refValue.Type() != underlyingArrayType { + return &underlyingTypeEncodePlan{nextValueType: underlyingArrayType}, refValue.Convert(underlyingArrayType).Interface(), true + } + } + + return nil, nil, false +} + +// TryWrapStringerEncodePlan tries to wrap a fmt.Stringer type with a wrapper that provides TextValuer. This is +// intentionally a separate function from TryWrapBuiltinTypeEncodePlan so it can be ordered after +// TryWrapFindUnderlyingTypeEncodePlan. This ensures that named types with an underlying builtin type (e.g. type MyEnum +// int32 with a String() method) prefer encoding via the underlying type's codec (e.g. as an integer) rather than via +// Stringer. Stringer is only used as a fallback when no type-specific encoding plan succeeds. +// (https://github.com/jackc/pgx/discussions/2527) +func TryWrapStringerEncodePlan(value any) (plan WrappedEncodePlanNextSetter, nextValue any, ok bool) { + if _, ok := value.(driver.Valuer); ok { + return nil, nil, false + } + + if s, ok := value.(fmt.Stringer); ok { + return &wrapFmtStringerEncodePlan{}, fmtStringerWrapper{s}, true + } + + return nil, nil, false +} + +type WrappedEncodePlanNextSetter interface { + SetNext(EncodePlan) + EncodePlan +} + +// TryWrapBuiltinTypeEncodePlan tries to wrap a builtin type with a wrapper that provides additional methods. e.g. If +// value was of type int32 then a wrapper plan would be returned that converts value to a type that implements +// Int64Valuer. +func TryWrapBuiltinTypeEncodePlan(value any) (plan WrappedEncodePlanNextSetter, nextValue any, ok bool) { + if _, ok := value.(driver.Valuer); ok { + return nil, nil, false + } + + switch value := value.(type) { + case int8: + return &wrapInt8EncodePlan{}, int8Wrapper(value), true + case int16: + return &wrapInt16EncodePlan{}, int16Wrapper(value), true + case int32: + return &wrapInt32EncodePlan{}, int32Wrapper(value), true + case int64: + return &wrapInt64EncodePlan{}, int64Wrapper(value), true + case int: + return &wrapIntEncodePlan{}, intWrapper(value), true + case uint8: + return &wrapUint8EncodePlan{}, uint8Wrapper(value), true + case uint16: + return &wrapUint16EncodePlan{}, uint16Wrapper(value), true + case uint32: + return &wrapUint32EncodePlan{}, uint32Wrapper(value), true + case uint64: + return &wrapUint64EncodePlan{}, uint64Wrapper(value), true + case uint: + return &wrapUintEncodePlan{}, uintWrapper(value), true + case float32: + return &wrapFloat32EncodePlan{}, float32Wrapper(value), true + case float64: + return &wrapFloat64EncodePlan{}, float64Wrapper(value), true + case string: + return &wrapStringEncodePlan{}, stringWrapper(value), true + case time.Time: + return &wrapTimeEncodePlan{}, timeWrapper(value), true + case time.Duration: + return &wrapDurationEncodePlan{}, durationWrapper(value), true + case net.IPNet: + return &wrapNetIPNetEncodePlan{}, netIPNetWrapper(value), true + case net.IP: + return &wrapNetIPEncodePlan{}, netIPWrapper(value), true + case netip.Prefix: + return &wrapNetipPrefixEncodePlan{}, netipPrefixWrapper(value), true + case netip.Addr: + return &wrapNetipAddrEncodePlan{}, netipAddrWrapper(value), true + case map[string]*string: + return &wrapMapStringToPointerStringEncodePlan{}, mapStringToPointerStringWrapper(value), true + case map[string]string: + return &wrapMapStringToStringEncodePlan{}, mapStringToStringWrapper(value), true + case [16]byte: + return &wrapByte16EncodePlan{}, byte16Wrapper(value), true + case []byte: + return &wrapByteSliceEncodePlan{}, byteSliceWrapper(value), true + } + + return nil, nil, false +} + +type wrapInt8EncodePlan struct { + next EncodePlan +} + +func (plan *wrapInt8EncodePlan) SetNext(next EncodePlan) { plan.next = next } + +func (plan *wrapInt8EncodePlan) Encode(value any, buf []byte) (newBuf []byte, err error) { + return plan.next.Encode(int8Wrapper(value.(int8)), buf) +} + +type wrapInt16EncodePlan struct { + next EncodePlan +} + +func (plan *wrapInt16EncodePlan) SetNext(next EncodePlan) { plan.next = next } + +func (plan *wrapInt16EncodePlan) Encode(value any, buf []byte) (newBuf []byte, err error) { + return plan.next.Encode(int16Wrapper(value.(int16)), buf) +} + +type wrapInt32EncodePlan struct { + next EncodePlan +} + +func (plan *wrapInt32EncodePlan) SetNext(next EncodePlan) { plan.next = next } + +func (plan *wrapInt32EncodePlan) Encode(value any, buf []byte) (newBuf []byte, err error) { + return plan.next.Encode(int32Wrapper(value.(int32)), buf) +} + +type wrapInt64EncodePlan struct { + next EncodePlan +} + +func (plan *wrapInt64EncodePlan) SetNext(next EncodePlan) { plan.next = next } + +func (plan *wrapInt64EncodePlan) Encode(value any, buf []byte) (newBuf []byte, err error) { + return plan.next.Encode(int64Wrapper(value.(int64)), buf) +} + +type wrapIntEncodePlan struct { + next EncodePlan +} + +func (plan *wrapIntEncodePlan) SetNext(next EncodePlan) { plan.next = next } + +func (plan *wrapIntEncodePlan) Encode(value any, buf []byte) (newBuf []byte, err error) { + return plan.next.Encode(intWrapper(value.(int)), buf) +} + +type wrapUint8EncodePlan struct { + next EncodePlan +} + +func (plan *wrapUint8EncodePlan) SetNext(next EncodePlan) { plan.next = next } + +func (plan *wrapUint8EncodePlan) Encode(value any, buf []byte) (newBuf []byte, err error) { + return plan.next.Encode(uint8Wrapper(value.(uint8)), buf) +} + +type wrapUint16EncodePlan struct { + next EncodePlan +} + +func (plan *wrapUint16EncodePlan) SetNext(next EncodePlan) { plan.next = next } + +func (plan *wrapUint16EncodePlan) Encode(value any, buf []byte) (newBuf []byte, err error) { + return plan.next.Encode(uint16Wrapper(value.(uint16)), buf) +} + +type wrapUint32EncodePlan struct { + next EncodePlan +} + +func (plan *wrapUint32EncodePlan) SetNext(next EncodePlan) { plan.next = next } + +func (plan *wrapUint32EncodePlan) Encode(value any, buf []byte) (newBuf []byte, err error) { + return plan.next.Encode(uint32Wrapper(value.(uint32)), buf) +} + +type wrapUint64EncodePlan struct { + next EncodePlan +} + +func (plan *wrapUint64EncodePlan) SetNext(next EncodePlan) { plan.next = next } + +func (plan *wrapUint64EncodePlan) Encode(value any, buf []byte) (newBuf []byte, err error) { + return plan.next.Encode(uint64Wrapper(value.(uint64)), buf) +} + +type wrapUintEncodePlan struct { + next EncodePlan +} + +func (plan *wrapUintEncodePlan) SetNext(next EncodePlan) { plan.next = next } + +func (plan *wrapUintEncodePlan) Encode(value any, buf []byte) (newBuf []byte, err error) { + return plan.next.Encode(uintWrapper(value.(uint)), buf) +} + +type wrapFloat32EncodePlan struct { + next EncodePlan +} + +func (plan *wrapFloat32EncodePlan) SetNext(next EncodePlan) { plan.next = next } + +func (plan *wrapFloat32EncodePlan) Encode(value any, buf []byte) (newBuf []byte, err error) { + return plan.next.Encode(float32Wrapper(value.(float32)), buf) +} + +type wrapFloat64EncodePlan struct { + next EncodePlan +} + +func (plan *wrapFloat64EncodePlan) SetNext(next EncodePlan) { plan.next = next } + +func (plan *wrapFloat64EncodePlan) Encode(value any, buf []byte) (newBuf []byte, err error) { + return plan.next.Encode(float64Wrapper(value.(float64)), buf) +} + +type wrapStringEncodePlan struct { + next EncodePlan +} + +func (plan *wrapStringEncodePlan) SetNext(next EncodePlan) { plan.next = next } + +func (plan *wrapStringEncodePlan) Encode(value any, buf []byte) (newBuf []byte, err error) { + return plan.next.Encode(stringWrapper(value.(string)), buf) +} + +type wrapTimeEncodePlan struct { + next EncodePlan +} + +func (plan *wrapTimeEncodePlan) SetNext(next EncodePlan) { plan.next = next } + +func (plan *wrapTimeEncodePlan) Encode(value any, buf []byte) (newBuf []byte, err error) { + return plan.next.Encode(timeWrapper(value.(time.Time)), buf) +} + +type wrapDurationEncodePlan struct { + next EncodePlan +} + +func (plan *wrapDurationEncodePlan) SetNext(next EncodePlan) { plan.next = next } + +func (plan *wrapDurationEncodePlan) Encode(value any, buf []byte) (newBuf []byte, err error) { + return plan.next.Encode(durationWrapper(value.(time.Duration)), buf) +} + +type wrapNetIPNetEncodePlan struct { + next EncodePlan +} + +func (plan *wrapNetIPNetEncodePlan) SetNext(next EncodePlan) { plan.next = next } + +func (plan *wrapNetIPNetEncodePlan) Encode(value any, buf []byte) (newBuf []byte, err error) { + return plan.next.Encode(netIPNetWrapper(value.(net.IPNet)), buf) +} + +type wrapNetIPEncodePlan struct { + next EncodePlan +} + +func (plan *wrapNetIPEncodePlan) SetNext(next EncodePlan) { plan.next = next } + +func (plan *wrapNetIPEncodePlan) Encode(value any, buf []byte) (newBuf []byte, err error) { + return plan.next.Encode(netIPWrapper(value.(net.IP)), buf) +} + +type wrapNetipPrefixEncodePlan struct { + next EncodePlan +} + +func (plan *wrapNetipPrefixEncodePlan) SetNext(next EncodePlan) { plan.next = next } + +func (plan *wrapNetipPrefixEncodePlan) Encode(value any, buf []byte) (newBuf []byte, err error) { + return plan.next.Encode(netipPrefixWrapper(value.(netip.Prefix)), buf) +} + +type wrapNetipAddrEncodePlan struct { + next EncodePlan +} + +func (plan *wrapNetipAddrEncodePlan) SetNext(next EncodePlan) { plan.next = next } + +func (plan *wrapNetipAddrEncodePlan) Encode(value any, buf []byte) (newBuf []byte, err error) { + return plan.next.Encode(netipAddrWrapper(value.(netip.Addr)), buf) +} + +type wrapMapStringToPointerStringEncodePlan struct { + next EncodePlan +} + +func (plan *wrapMapStringToPointerStringEncodePlan) SetNext(next EncodePlan) { plan.next = next } + +func (plan *wrapMapStringToPointerStringEncodePlan) Encode(value any, buf []byte) (newBuf []byte, err error) { + return plan.next.Encode(mapStringToPointerStringWrapper(value.(map[string]*string)), buf) +} + +type wrapMapStringToStringEncodePlan struct { + next EncodePlan +} + +func (plan *wrapMapStringToStringEncodePlan) SetNext(next EncodePlan) { plan.next = next } + +func (plan *wrapMapStringToStringEncodePlan) Encode(value any, buf []byte) (newBuf []byte, err error) { + return plan.next.Encode(mapStringToStringWrapper(value.(map[string]string)), buf) +} + +type wrapByte16EncodePlan struct { + next EncodePlan +} + +func (plan *wrapByte16EncodePlan) SetNext(next EncodePlan) { plan.next = next } + +func (plan *wrapByte16EncodePlan) Encode(value any, buf []byte) (newBuf []byte, err error) { + return plan.next.Encode(byte16Wrapper(value.([16]byte)), buf) +} + +type wrapByteSliceEncodePlan struct { + next EncodePlan +} + +func (plan *wrapByteSliceEncodePlan) SetNext(next EncodePlan) { plan.next = next } + +func (plan *wrapByteSliceEncodePlan) Encode(value any, buf []byte) (newBuf []byte, err error) { + return plan.next.Encode(byteSliceWrapper(value.([]byte)), buf) +} + +type wrapFmtStringerEncodePlan struct { + next EncodePlan +} + +func (plan *wrapFmtStringerEncodePlan) SetNext(next EncodePlan) { plan.next = next } + +func (plan *wrapFmtStringerEncodePlan) Encode(value any, buf []byte) (newBuf []byte, err error) { + return plan.next.Encode(fmtStringerWrapper{value.(fmt.Stringer)}, buf) +} + +// TryWrapStructEncodePlan tries to wrap a struct with a wrapper that implements CompositeIndexGetter. +func TryWrapStructEncodePlan(value any) (plan WrappedEncodePlanNextSetter, nextValue any, ok bool) { + if _, ok := value.(driver.Valuer); ok { + return nil, nil, false + } + + if valueType := reflect.TypeOf(value); valueType != nil && valueType.Kind() == reflect.Struct { + exportedFields := getExportedFieldValues(reflect.ValueOf(value)) + if len(exportedFields) == 0 { + return nil, nil, false + } + + w := structWrapper{ + s: value, + exportedFields: exportedFields, + } + return &wrapAnyStructEncodePlan{}, w, true + } + + return nil, nil, false +} + +type wrapAnyStructEncodePlan struct { + next EncodePlan +} + +func (plan *wrapAnyStructEncodePlan) SetNext(next EncodePlan) { plan.next = next } + +func (plan *wrapAnyStructEncodePlan) Encode(value any, buf []byte) (newBuf []byte, err error) { + w := structWrapper{ + s: value, + exportedFields: getExportedFieldValues(reflect.ValueOf(value)), + } + + return plan.next.Encode(w, buf) +} + +func getExportedFieldValues(structValue reflect.Value) []reflect.Value { + structType := structValue.Type() + exportedFields := make([]reflect.Value, 0, structValue.NumField()) + for i := 0; i < structType.NumField(); i++ { + sf := structType.Field(i) + if sf.IsExported() { + exportedFields = append(exportedFields, structValue.Field(i)) + } + } + + return exportedFields +} + +func TryWrapSliceEncodePlan(value any) (plan WrappedEncodePlanNextSetter, nextValue any, ok bool) { + if _, ok := value.(driver.Valuer); ok { + return nil, nil, false + } + + // Avoid using reflect path for common types. + switch value := value.(type) { + case []int16: + return &wrapSliceEncodePlan[int16]{}, (FlatArray[int16])(value), true + case []int32: + return &wrapSliceEncodePlan[int32]{}, (FlatArray[int32])(value), true + case []int64: + return &wrapSliceEncodePlan[int64]{}, (FlatArray[int64])(value), true + case []float32: + return &wrapSliceEncodePlan[float32]{}, (FlatArray[float32])(value), true + case []float64: + return &wrapSliceEncodePlan[float64]{}, (FlatArray[float64])(value), true + case []string: + return &wrapSliceEncodePlan[string]{}, (FlatArray[string])(value), true + case []time.Time: + return &wrapSliceEncodePlan[time.Time]{}, (FlatArray[time.Time])(value), true + } + + if valueType := reflect.TypeOf(value); valueType != nil && valueType.Kind() == reflect.Slice { + w := anySliceArrayReflect{ + slice: reflect.ValueOf(value), + } + return &wrapSliceEncodeReflectPlan{}, w, true + } + + return nil, nil, false +} + +type wrapSliceEncodePlan[T any] struct { + next EncodePlan +} + +func (plan *wrapSliceEncodePlan[T]) SetNext(next EncodePlan) { plan.next = next } + +func (plan *wrapSliceEncodePlan[T]) Encode(value any, buf []byte) (newBuf []byte, err error) { + return plan.next.Encode((FlatArray[T])(value.([]T)), buf) +} + +type wrapSliceEncodeReflectPlan struct { + next EncodePlan +} + +func (plan *wrapSliceEncodeReflectPlan) SetNext(next EncodePlan) { plan.next = next } + +func (plan *wrapSliceEncodeReflectPlan) Encode(value any, buf []byte) (newBuf []byte, err error) { + w := anySliceArrayReflect{ + slice: reflect.ValueOf(value), + } + + return plan.next.Encode(w, buf) +} + +func TryWrapMultiDimSliceEncodePlan(value any) (plan WrappedEncodePlanNextSetter, nextValue any, ok bool) { + if _, ok := value.(driver.Valuer); ok { + return nil, nil, false + } + + sliceValue := reflect.ValueOf(value) + if sliceValue.Kind() == reflect.Slice { + valueElemType := sliceValue.Type().Elem() + + if valueElemType.Kind() == reflect.Slice { + if !isRagged(sliceValue) { + w := anyMultiDimSliceArray{ + slice: reflect.ValueOf(value), + } + return &wrapMultiDimSliceEncodePlan{}, &w, true + } + } + } + + return nil, nil, false +} + +type wrapMultiDimSliceEncodePlan struct { + next EncodePlan +} + +func (plan *wrapMultiDimSliceEncodePlan) SetNext(next EncodePlan) { plan.next = next } + +func (plan *wrapMultiDimSliceEncodePlan) Encode(value any, buf []byte) (newBuf []byte, err error) { + w := anyMultiDimSliceArray{ + slice: reflect.ValueOf(value), + } + + return plan.next.Encode(&w, buf) +} + +func TryWrapArrayEncodePlan(value any) (plan WrappedEncodePlanNextSetter, nextValue any, ok bool) { + if _, ok := value.(driver.Valuer); ok { + return nil, nil, false + } + + if valueType := reflect.TypeOf(value); valueType != nil && valueType.Kind() == reflect.Array { + w := anyArrayArrayReflect{ + array: reflect.ValueOf(value), + } + return &wrapArrayEncodeReflectPlan{}, w, true + } + + return nil, nil, false +} + +type wrapArrayEncodeReflectPlan struct { + next EncodePlan +} + +func (plan *wrapArrayEncodeReflectPlan) SetNext(next EncodePlan) { plan.next = next } + +func (plan *wrapArrayEncodeReflectPlan) Encode(value any, buf []byte) (newBuf []byte, err error) { + w := anyArrayArrayReflect{ + array: reflect.ValueOf(value), + } + + return plan.next.Encode(w, buf) +} + +func newEncodeError(value any, m *Map, oid uint32, formatCode int16, err error) error { + var format string + switch formatCode { + case TextFormatCode: + format = "text" + case BinaryFormatCode: + format = "binary" + default: + format = fmt.Sprintf("unknown (%d)", formatCode) + } + + var dataTypeName string + if t, ok := m.TypeForOID(oid); ok { + dataTypeName = t.Name + } else { + dataTypeName = "unknown type" + } + + return fmt.Errorf("unable to encode %#v into %s format for %s (OID %d): %w", value, format, dataTypeName, oid, err) +} + +// Encode appends the encoded bytes of value to buf. If value is the SQL value NULL then append nothing and return +// (nil, nil). The caller of Encode is responsible for writing the correct NULL value or the length of the data +// written. +func (m *Map) Encode(oid uint32, formatCode int16, value any, buf []byte) (newBuf []byte, err error) { + if isNil, callNilDriverValuer := isNilDriverValuer(value); isNil { + if callNilDriverValuer { + newBuf, err = (&encodePlanDriverValuer{m: m, oid: oid, formatCode: formatCode}).Encode(value, buf) + if err != nil { + return nil, newEncodeError(value, m, oid, formatCode, err) + } + + return newBuf, nil + } else { + return nil, nil + } + } + + plan := m.PlanEncode(oid, formatCode, value) + if plan == nil { + return nil, newEncodeError(value, m, oid, formatCode, errors.New("cannot find encode plan")) + } + + newBuf, err = plan.Encode(value, buf) + if err != nil { + return nil, newEncodeError(value, m, oid, formatCode, err) + } + + return newBuf, nil +} + +// SQLScanner returns a database/sql.Scanner for v. This is necessary for types like Array[T] and Range[T] where the +// type needs assistance from Map to implement the sql.Scanner interface. It is not necessary for types like Box that +// implement sql.Scanner directly. +// +// This uses the type of v to look up the PostgreSQL OID that v presumably came from. This means v must be registered +// with m by calling RegisterDefaultPgType. +func (m *Map) SQLScanner(v any) sql.Scanner { + if s, ok := v.(sql.Scanner); ok { + return s + } + + return &sqlScannerWrapper{m: m, v: v} +} + +type sqlScannerWrapper struct { + m *Map + v any +} + +func (w *sqlScannerWrapper) Scan(src any) error { + t, ok := w.m.TypeForValue(w.v) + if !ok { + return fmt.Errorf("cannot convert to sql.Scanner: cannot find registered type for %T", w.v) + } + + var bufSrc []byte + if src != nil { + switch src := src.(type) { + case string: + bufSrc = []byte(src) + case []byte: + bufSrc = src + default: + bufSrc = fmt.Append(nil, bufSrc) + } + } + + return w.m.Scan(t.OID, TextFormatCode, bufSrc, w.v) +} + +var valuerReflectType = reflect.TypeFor[driver.Valuer]() + +// isNilDriverValuer returns true if value is any type of nil unless it implements driver.Valuer. *T is not considered to implement +// driver.Valuer if it is only implemented by T. +func isNilDriverValuer(value any) (isNil, callNilDriverValuer bool) { + if value == nil { + return true, false + } + + refVal := reflect.ValueOf(value) + kind := refVal.Kind() + switch kind { + case reflect.Chan, reflect.Func, reflect.Map, reflect.Pointer, reflect.UnsafePointer, reflect.Interface, reflect.Slice: + if !refVal.IsNil() { + return false, false + } + + if _, ok := value.(driver.Valuer); ok { + if kind == reflect.Pointer { + // The type assertion will succeed if driver.Valuer is implemented on T or *T. Check if it is implemented on *T + // by checking if it is not implemented on *T. + return true, !refVal.Type().Elem().Implements(valuerReflectType) + } else { + return true, true + } + } + + return true, false + default: + return false, false + } +} diff --git a/vendor/github.com/jackc/pgx/v5/pgtype/pgtype_default.go b/vendor/github.com/jackc/pgx/v5/pgtype/pgtype_default.go new file mode 100644 index 0000000000..42b39d827d --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgtype/pgtype_default.go @@ -0,0 +1,251 @@ +package pgtype + +import ( + "encoding/json" + "encoding/xml" + "net" + "net/netip" + "reflect" + "sync" + "time" +) + +var ( + // defaultMap contains default mappings between PostgreSQL server types and Go type handling logic. + defaultMap *Map + defaultMapInitOnce = sync.Once{} +) + +func initDefaultMap() { + defaultMap = &Map{ + oidToType: make(map[uint32]*Type), + nameToType: make(map[string]*Type), + reflectTypeToName: make(map[reflect.Type]string), + oidToFormatCode: make(map[uint32]int16), + + memoizedEncodePlans: make(map[uint32]map[reflect.Type][2]EncodePlan), + + TryWrapEncodePlanFuncs: []TryWrapEncodePlanFunc{ + TryWrapDerefPointerEncodePlan, + TryWrapBuiltinTypeEncodePlan, + TryWrapFindUnderlyingTypeEncodePlan, + TryWrapStructEncodePlan, + TryWrapSliceEncodePlan, + TryWrapMultiDimSliceEncodePlan, + TryWrapArrayEncodePlan, + }, + + TryWrapScanPlanFuncs: []TryWrapScanPlanFunc{ + TryPointerPointerScanPlan, + TryWrapBuiltinTypeScanPlan, + TryFindUnderlyingTypeScanPlan, + TryWrapStructScanPlan, + TryWrapPtrSliceScanPlan, + TryWrapPtrMultiDimSliceScanPlan, + TryWrapPtrArrayScanPlan, + }, + } + + // Base types + defaultMap.RegisterType(&Type{Name: "aclitem", OID: ACLItemOID, Codec: &TextFormatOnlyCodec{TextCodec{}}}) + defaultMap.RegisterType(&Type{Name: "bit", OID: BitOID, Codec: BitsCodec{}}) + defaultMap.RegisterType(&Type{Name: "bool", OID: BoolOID, Codec: BoolCodec{}}) + defaultMap.RegisterType(&Type{Name: "box", OID: BoxOID, Codec: BoxCodec{}}) + defaultMap.RegisterType(&Type{Name: "bpchar", OID: BPCharOID, Codec: TextCodec{}}) + defaultMap.RegisterType(&Type{Name: "bytea", OID: ByteaOID, Codec: ByteaCodec{}}) + defaultMap.RegisterType(&Type{Name: "char", OID: QCharOID, Codec: QCharCodec{}}) + defaultMap.RegisterType(&Type{Name: "cid", OID: CIDOID, Codec: Uint32Codec{}}) + defaultMap.RegisterType(&Type{Name: "cidr", OID: CIDROID, Codec: InetCodec{}}) + defaultMap.RegisterType(&Type{Name: "circle", OID: CircleOID, Codec: CircleCodec{}}) + defaultMap.RegisterType(&Type{Name: "date", OID: DateOID, Codec: DateCodec{}}) + defaultMap.RegisterType(&Type{Name: "float4", OID: Float4OID, Codec: Float4Codec{}}) + defaultMap.RegisterType(&Type{Name: "float8", OID: Float8OID, Codec: Float8Codec{}}) + defaultMap.RegisterType(&Type{Name: "inet", OID: InetOID, Codec: InetCodec{}}) + defaultMap.RegisterType(&Type{Name: "int2", OID: Int2OID, Codec: Int2Codec{}}) + defaultMap.RegisterType(&Type{Name: "int4", OID: Int4OID, Codec: Int4Codec{}}) + defaultMap.RegisterType(&Type{Name: "int8", OID: Int8OID, Codec: Int8Codec{}}) + defaultMap.RegisterType(&Type{Name: "interval", OID: IntervalOID, Codec: IntervalCodec{}}) + defaultMap.RegisterType(&Type{Name: "json", OID: JSONOID, Codec: &JSONCodec{Marshal: json.Marshal, Unmarshal: json.Unmarshal}}) + defaultMap.RegisterType(&Type{Name: "jsonb", OID: JSONBOID, Codec: &JSONBCodec{Marshal: json.Marshal, Unmarshal: json.Unmarshal}}) + defaultMap.RegisterType(&Type{Name: "jsonpath", OID: JSONPathOID, Codec: &TextFormatOnlyCodec{TextCodec{}}}) + defaultMap.RegisterType(&Type{Name: "line", OID: LineOID, Codec: LineCodec{}}) + defaultMap.RegisterType(&Type{Name: "lseg", OID: LsegOID, Codec: LsegCodec{}}) + defaultMap.RegisterType(&Type{Name: "macaddr8", OID: Macaddr8OID, Codec: MacaddrCodec{}}) + defaultMap.RegisterType(&Type{Name: "macaddr", OID: MacaddrOID, Codec: MacaddrCodec{}}) + defaultMap.RegisterType(&Type{Name: "name", OID: NameOID, Codec: TextCodec{}}) + defaultMap.RegisterType(&Type{Name: "numeric", OID: NumericOID, Codec: NumericCodec{}}) + defaultMap.RegisterType(&Type{Name: "oid", OID: OIDOID, Codec: Uint32Codec{}}) + defaultMap.RegisterType(&Type{Name: "path", OID: PathOID, Codec: PathCodec{}}) + defaultMap.RegisterType(&Type{Name: "point", OID: PointOID, Codec: PointCodec{}}) + defaultMap.RegisterType(&Type{Name: "polygon", OID: PolygonOID, Codec: PolygonCodec{}}) + defaultMap.RegisterType(&Type{Name: "record", OID: RecordOID, Codec: RecordCodec{}}) + defaultMap.RegisterType(&Type{Name: "text", OID: TextOID, Codec: TextCodec{}}) + defaultMap.RegisterType(&Type{Name: "tid", OID: TIDOID, Codec: TIDCodec{}}) + defaultMap.RegisterType(&Type{Name: "tsvector", OID: TSVectorOID, Codec: TSVectorCodec{}}) + defaultMap.RegisterType(&Type{Name: "time", OID: TimeOID, Codec: TimeCodec{}}) + defaultMap.RegisterType(&Type{Name: "timestamp", OID: TimestampOID, Codec: &TimestampCodec{}}) + defaultMap.RegisterType(&Type{Name: "timestamptz", OID: TimestamptzOID, Codec: &TimestamptzCodec{}}) + defaultMap.RegisterType(&Type{Name: "unknown", OID: UnknownOID, Codec: TextCodec{}}) + defaultMap.RegisterType(&Type{Name: "uuid", OID: UUIDOID, Codec: UUIDCodec{}}) + defaultMap.RegisterType(&Type{Name: "varbit", OID: VarbitOID, Codec: BitsCodec{}}) + defaultMap.RegisterType(&Type{Name: "varchar", OID: VarcharOID, Codec: TextCodec{}}) + defaultMap.RegisterType(&Type{Name: "xid", OID: XIDOID, Codec: Uint32Codec{}}) + defaultMap.RegisterType(&Type{Name: "xid8", OID: XID8OID, Codec: Uint64Codec{}}) + defaultMap.RegisterType(&Type{Name: "xml", OID: XMLOID, Codec: &XMLCodec{ + Marshal: xml.Marshal, + // xml.Unmarshal does not support unmarshalling into *any. However, XMLCodec.DecodeValue calls Unmarshal with a + // *any. Wrap xml.Marshal with a function that copies the data into a new byte slice in this case. Not implementing + // directly in XMLCodec.DecodeValue to allow for the unlikely possibility that someone uses an alternative XML + // unmarshaler that does support unmarshalling into *any. + // + // https://github.com/jackc/pgx/issues/2227 + // https://github.com/jackc/pgx/pull/2228 + Unmarshal: func(data []byte, v any) error { + if v, ok := v.(*any); ok { + dstBuf := make([]byte, len(data)) + copy(dstBuf, data) + *v = dstBuf + return nil + } + return xml.Unmarshal(data, v) + }, + }}) + + // Range types + defaultMap.RegisterType(&Type{Name: "daterange", OID: DaterangeOID, Codec: &RangeCodec{ElementType: defaultMap.oidToType[DateOID]}}) + defaultMap.RegisterType(&Type{Name: "int4range", OID: Int4rangeOID, Codec: &RangeCodec{ElementType: defaultMap.oidToType[Int4OID]}}) + defaultMap.RegisterType(&Type{Name: "int8range", OID: Int8rangeOID, Codec: &RangeCodec{ElementType: defaultMap.oidToType[Int8OID]}}) + defaultMap.RegisterType(&Type{Name: "numrange", OID: NumrangeOID, Codec: &RangeCodec{ElementType: defaultMap.oidToType[NumericOID]}}) + defaultMap.RegisterType(&Type{Name: "tsrange", OID: TsrangeOID, Codec: &RangeCodec{ElementType: defaultMap.oidToType[TimestampOID]}}) + defaultMap.RegisterType(&Type{Name: "tstzrange", OID: TstzrangeOID, Codec: &RangeCodec{ElementType: defaultMap.oidToType[TimestamptzOID]}}) + + // Multirange types + defaultMap.RegisterType(&Type{Name: "datemultirange", OID: DatemultirangeOID, Codec: &MultirangeCodec{ElementType: defaultMap.oidToType[DaterangeOID]}}) + defaultMap.RegisterType(&Type{Name: "int4multirange", OID: Int4multirangeOID, Codec: &MultirangeCodec{ElementType: defaultMap.oidToType[Int4rangeOID]}}) + defaultMap.RegisterType(&Type{Name: "int8multirange", OID: Int8multirangeOID, Codec: &MultirangeCodec{ElementType: defaultMap.oidToType[Int8rangeOID]}}) + defaultMap.RegisterType(&Type{Name: "nummultirange", OID: NummultirangeOID, Codec: &MultirangeCodec{ElementType: defaultMap.oidToType[NumrangeOID]}}) + defaultMap.RegisterType(&Type{Name: "tsmultirange", OID: TsmultirangeOID, Codec: &MultirangeCodec{ElementType: defaultMap.oidToType[TsrangeOID]}}) + defaultMap.RegisterType(&Type{Name: "tstzmultirange", OID: TstzmultirangeOID, Codec: &MultirangeCodec{ElementType: defaultMap.oidToType[TstzrangeOID]}}) + + // Array types + defaultMap.RegisterType(&Type{Name: "_aclitem", OID: ACLItemArrayOID, Codec: &ArrayCodec{ElementType: defaultMap.oidToType[ACLItemOID]}}) + defaultMap.RegisterType(&Type{Name: "_bit", OID: BitArrayOID, Codec: &ArrayCodec{ElementType: defaultMap.oidToType[BitOID]}}) + defaultMap.RegisterType(&Type{Name: "_bool", OID: BoolArrayOID, Codec: &ArrayCodec{ElementType: defaultMap.oidToType[BoolOID]}}) + defaultMap.RegisterType(&Type{Name: "_box", OID: BoxArrayOID, Codec: &ArrayCodec{ElementType: defaultMap.oidToType[BoxOID]}}) + defaultMap.RegisterType(&Type{Name: "_bpchar", OID: BPCharArrayOID, Codec: &ArrayCodec{ElementType: defaultMap.oidToType[BPCharOID]}}) + defaultMap.RegisterType(&Type{Name: "_bytea", OID: ByteaArrayOID, Codec: &ArrayCodec{ElementType: defaultMap.oidToType[ByteaOID]}}) + defaultMap.RegisterType(&Type{Name: "_char", OID: QCharArrayOID, Codec: &ArrayCodec{ElementType: defaultMap.oidToType[QCharOID]}}) + defaultMap.RegisterType(&Type{Name: "_cid", OID: CIDArrayOID, Codec: &ArrayCodec{ElementType: defaultMap.oidToType[CIDOID]}}) + defaultMap.RegisterType(&Type{Name: "_cidr", OID: CIDRArrayOID, Codec: &ArrayCodec{ElementType: defaultMap.oidToType[CIDROID]}}) + defaultMap.RegisterType(&Type{Name: "_circle", OID: CircleArrayOID, Codec: &ArrayCodec{ElementType: defaultMap.oidToType[CircleOID]}}) + defaultMap.RegisterType(&Type{Name: "_date", OID: DateArrayOID, Codec: &ArrayCodec{ElementType: defaultMap.oidToType[DateOID]}}) + defaultMap.RegisterType(&Type{Name: "_daterange", OID: DaterangeArrayOID, Codec: &ArrayCodec{ElementType: defaultMap.oidToType[DaterangeOID]}}) + defaultMap.RegisterType(&Type{Name: "_float4", OID: Float4ArrayOID, Codec: &ArrayCodec{ElementType: defaultMap.oidToType[Float4OID]}}) + defaultMap.RegisterType(&Type{Name: "_float8", OID: Float8ArrayOID, Codec: &ArrayCodec{ElementType: defaultMap.oidToType[Float8OID]}}) + defaultMap.RegisterType(&Type{Name: "_inet", OID: InetArrayOID, Codec: &ArrayCodec{ElementType: defaultMap.oidToType[InetOID]}}) + defaultMap.RegisterType(&Type{Name: "_int2", OID: Int2ArrayOID, Codec: &ArrayCodec{ElementType: defaultMap.oidToType[Int2OID]}}) + defaultMap.RegisterType(&Type{Name: "_int4", OID: Int4ArrayOID, Codec: &ArrayCodec{ElementType: defaultMap.oidToType[Int4OID]}}) + defaultMap.RegisterType(&Type{Name: "_int4range", OID: Int4rangeArrayOID, Codec: &ArrayCodec{ElementType: defaultMap.oidToType[Int4rangeOID]}}) + defaultMap.RegisterType(&Type{Name: "_int8", OID: Int8ArrayOID, Codec: &ArrayCodec{ElementType: defaultMap.oidToType[Int8OID]}}) + defaultMap.RegisterType(&Type{Name: "_int8range", OID: Int8rangeArrayOID, Codec: &ArrayCodec{ElementType: defaultMap.oidToType[Int8rangeOID]}}) + defaultMap.RegisterType(&Type{Name: "_interval", OID: IntervalArrayOID, Codec: &ArrayCodec{ElementType: defaultMap.oidToType[IntervalOID]}}) + defaultMap.RegisterType(&Type{Name: "_json", OID: JSONArrayOID, Codec: &ArrayCodec{ElementType: defaultMap.oidToType[JSONOID]}}) + defaultMap.RegisterType(&Type{Name: "_jsonb", OID: JSONBArrayOID, Codec: &ArrayCodec{ElementType: defaultMap.oidToType[JSONBOID]}}) + defaultMap.RegisterType(&Type{Name: "_jsonpath", OID: JSONPathArrayOID, Codec: &ArrayCodec{ElementType: defaultMap.oidToType[JSONPathOID]}}) + defaultMap.RegisterType(&Type{Name: "_line", OID: LineArrayOID, Codec: &ArrayCodec{ElementType: defaultMap.oidToType[LineOID]}}) + defaultMap.RegisterType(&Type{Name: "_lseg", OID: LsegArrayOID, Codec: &ArrayCodec{ElementType: defaultMap.oidToType[LsegOID]}}) + defaultMap.RegisterType(&Type{Name: "_macaddr", OID: MacaddrArrayOID, Codec: &ArrayCodec{ElementType: defaultMap.oidToType[MacaddrOID]}}) + defaultMap.RegisterType(&Type{Name: "_name", OID: NameArrayOID, Codec: &ArrayCodec{ElementType: defaultMap.oidToType[NameOID]}}) + defaultMap.RegisterType(&Type{Name: "_numeric", OID: NumericArrayOID, Codec: &ArrayCodec{ElementType: defaultMap.oidToType[NumericOID]}}) + defaultMap.RegisterType(&Type{Name: "_numrange", OID: NumrangeArrayOID, Codec: &ArrayCodec{ElementType: defaultMap.oidToType[NumrangeOID]}}) + defaultMap.RegisterType(&Type{Name: "_oid", OID: OIDArrayOID, Codec: &ArrayCodec{ElementType: defaultMap.oidToType[OIDOID]}}) + defaultMap.RegisterType(&Type{Name: "_path", OID: PathArrayOID, Codec: &ArrayCodec{ElementType: defaultMap.oidToType[PathOID]}}) + defaultMap.RegisterType(&Type{Name: "_point", OID: PointArrayOID, Codec: &ArrayCodec{ElementType: defaultMap.oidToType[PointOID]}}) + defaultMap.RegisterType(&Type{Name: "_polygon", OID: PolygonArrayOID, Codec: &ArrayCodec{ElementType: defaultMap.oidToType[PolygonOID]}}) + defaultMap.RegisterType(&Type{Name: "_record", OID: RecordArrayOID, Codec: &ArrayCodec{ElementType: defaultMap.oidToType[RecordOID]}}) + defaultMap.RegisterType(&Type{Name: "_text", OID: TextArrayOID, Codec: &ArrayCodec{ElementType: defaultMap.oidToType[TextOID]}}) + defaultMap.RegisterType(&Type{Name: "_tid", OID: TIDArrayOID, Codec: &ArrayCodec{ElementType: defaultMap.oidToType[TIDOID]}}) + defaultMap.RegisterType(&Type{Name: "_tsvector", OID: TSVectorArrayOID, Codec: &ArrayCodec{ElementType: defaultMap.oidToType[TSVectorOID]}}) + defaultMap.RegisterType(&Type{Name: "_time", OID: TimeArrayOID, Codec: &ArrayCodec{ElementType: defaultMap.oidToType[TimeOID]}}) + defaultMap.RegisterType(&Type{Name: "_timestamp", OID: TimestampArrayOID, Codec: &ArrayCodec{ElementType: defaultMap.oidToType[TimestampOID]}}) + defaultMap.RegisterType(&Type{Name: "_timestamptz", OID: TimestamptzArrayOID, Codec: &ArrayCodec{ElementType: defaultMap.oidToType[TimestamptzOID]}}) + defaultMap.RegisterType(&Type{Name: "_tsrange", OID: TsrangeArrayOID, Codec: &ArrayCodec{ElementType: defaultMap.oidToType[TsrangeOID]}}) + defaultMap.RegisterType(&Type{Name: "_tstzrange", OID: TstzrangeArrayOID, Codec: &ArrayCodec{ElementType: defaultMap.oidToType[TstzrangeOID]}}) + defaultMap.RegisterType(&Type{Name: "_uuid", OID: UUIDArrayOID, Codec: &ArrayCodec{ElementType: defaultMap.oidToType[UUIDOID]}}) + defaultMap.RegisterType(&Type{Name: "_varbit", OID: VarbitArrayOID, Codec: &ArrayCodec{ElementType: defaultMap.oidToType[VarbitOID]}}) + defaultMap.RegisterType(&Type{Name: "_varchar", OID: VarcharArrayOID, Codec: &ArrayCodec{ElementType: defaultMap.oidToType[VarcharOID]}}) + defaultMap.RegisterType(&Type{Name: "_xid", OID: XIDArrayOID, Codec: &ArrayCodec{ElementType: defaultMap.oidToType[XIDOID]}}) + defaultMap.RegisterType(&Type{Name: "_xid8", OID: XID8ArrayOID, Codec: &ArrayCodec{ElementType: defaultMap.oidToType[XID8OID]}}) + defaultMap.RegisterType(&Type{Name: "_xml", OID: XMLArrayOID, Codec: &ArrayCodec{ElementType: defaultMap.oidToType[XMLOID]}}) + + // Integer types that directly map to a PostgreSQL type + registerDefaultPgTypeVariants[int16](defaultMap, "int2") + registerDefaultPgTypeVariants[int32](defaultMap, "int4") + registerDefaultPgTypeVariants[int64](defaultMap, "int8") + + // Integer types that do not have a direct match to a PostgreSQL type + registerDefaultPgTypeVariants[int8](defaultMap, "int8") + registerDefaultPgTypeVariants[int](defaultMap, "int8") + registerDefaultPgTypeVariants[uint8](defaultMap, "int8") + registerDefaultPgTypeVariants[uint16](defaultMap, "int8") + registerDefaultPgTypeVariants[uint32](defaultMap, "int8") + registerDefaultPgTypeVariants[uint64](defaultMap, "numeric") + registerDefaultPgTypeVariants[uint](defaultMap, "numeric") + + registerDefaultPgTypeVariants[float32](defaultMap, "float4") + registerDefaultPgTypeVariants[float64](defaultMap, "float8") + + registerDefaultPgTypeVariants[bool](defaultMap, "bool") + registerDefaultPgTypeVariants[time.Time](defaultMap, "timestamptz") + registerDefaultPgTypeVariants[time.Duration](defaultMap, "interval") + registerDefaultPgTypeVariants[string](defaultMap, "text") + registerDefaultPgTypeVariants[json.RawMessage](defaultMap, "json") + registerDefaultPgTypeVariants[[]byte](defaultMap, "bytea") + + registerDefaultPgTypeVariants[net.IP](defaultMap, "inet") + registerDefaultPgTypeVariants[net.IPNet](defaultMap, "cidr") + registerDefaultPgTypeVariants[netip.Addr](defaultMap, "inet") + registerDefaultPgTypeVariants[netip.Prefix](defaultMap, "cidr") + + // pgtype provided structs + registerDefaultPgTypeVariants[Bits](defaultMap, "varbit") + registerDefaultPgTypeVariants[Bool](defaultMap, "bool") + registerDefaultPgTypeVariants[Box](defaultMap, "box") + registerDefaultPgTypeVariants[Circle](defaultMap, "circle") + registerDefaultPgTypeVariants[Date](defaultMap, "date") + registerDefaultPgTypeVariants[Range[Date]](defaultMap, "daterange") + registerDefaultPgTypeVariants[Multirange[Range[Date]]](defaultMap, "datemultirange") + registerDefaultPgTypeVariants[Float4](defaultMap, "float4") + registerDefaultPgTypeVariants[Float8](defaultMap, "float8") + registerDefaultPgTypeVariants[Range[Float8]](defaultMap, "numrange") // There is no PostgreSQL builtin float8range so map it to numrange. + registerDefaultPgTypeVariants[Multirange[Range[Float8]]](defaultMap, "nummultirange") // There is no PostgreSQL builtin float8multirange so map it to nummultirange. + registerDefaultPgTypeVariants[Int2](defaultMap, "int2") + registerDefaultPgTypeVariants[Int4](defaultMap, "int4") + registerDefaultPgTypeVariants[Range[Int4]](defaultMap, "int4range") + registerDefaultPgTypeVariants[Multirange[Range[Int4]]](defaultMap, "int4multirange") + registerDefaultPgTypeVariants[Int8](defaultMap, "int8") + registerDefaultPgTypeVariants[Range[Int8]](defaultMap, "int8range") + registerDefaultPgTypeVariants[Multirange[Range[Int8]]](defaultMap, "int8multirange") + registerDefaultPgTypeVariants[Interval](defaultMap, "interval") + registerDefaultPgTypeVariants[Line](defaultMap, "line") + registerDefaultPgTypeVariants[Lseg](defaultMap, "lseg") + registerDefaultPgTypeVariants[Numeric](defaultMap, "numeric") + registerDefaultPgTypeVariants[Range[Numeric]](defaultMap, "numrange") + registerDefaultPgTypeVariants[Multirange[Range[Numeric]]](defaultMap, "nummultirange") + registerDefaultPgTypeVariants[Path](defaultMap, "path") + registerDefaultPgTypeVariants[Point](defaultMap, "point") + registerDefaultPgTypeVariants[Polygon](defaultMap, "polygon") + registerDefaultPgTypeVariants[TID](defaultMap, "tid") + registerDefaultPgTypeVariants[Text](defaultMap, "text") + registerDefaultPgTypeVariants[Time](defaultMap, "time") + registerDefaultPgTypeVariants[Timestamp](defaultMap, "timestamp") + registerDefaultPgTypeVariants[Timestamptz](defaultMap, "timestamptz") + registerDefaultPgTypeVariants[Range[Timestamp]](defaultMap, "tsrange") + registerDefaultPgTypeVariants[Multirange[Range[Timestamp]]](defaultMap, "tsmultirange") + registerDefaultPgTypeVariants[Range[Timestamptz]](defaultMap, "tstzrange") + registerDefaultPgTypeVariants[Multirange[Range[Timestamptz]]](defaultMap, "tstzmultirange") + registerDefaultPgTypeVariants[TSVector](defaultMap, "tsvector") + registerDefaultPgTypeVariants[UUID](defaultMap, "uuid") + + defaultMap.buildReflectTypeToType() +} diff --git a/vendor/github.com/jackc/pgx/v5/pgtype/point.go b/vendor/github.com/jackc/pgx/v5/pgtype/point.go new file mode 100644 index 0000000000..d90cb7033e --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgtype/point.go @@ -0,0 +1,266 @@ +package pgtype + +import ( + "bytes" + "database/sql/driver" + "encoding/binary" + "fmt" + "math" + "strconv" + "strings" + + "github.com/jackc/pgx/v5/internal/pgio" +) + +type Vec2 struct { + X float64 + Y float64 +} + +type PointScanner interface { + ScanPoint(v Point) error +} + +type PointValuer interface { + PointValue() (Point, error) +} + +type Point struct { + P Vec2 + Valid bool +} + +// ScanPoint implements the [PointScanner] interface. +func (p *Point) ScanPoint(v Point) error { + *p = v + return nil +} + +// PointValue implements the [PointValuer] interface. +func (p Point) PointValue() (Point, error) { + return p, nil +} + +func parsePoint(src []byte) (*Point, error) { + if src == nil || bytes.Equal(src, []byte("null")) { + return &Point{}, nil + } + + if len(src) < 5 { + return nil, fmt.Errorf("invalid length for point: %v", len(src)) + } + if src[0] == '"' && src[len(src)-1] == '"' { + src = src[1 : len(src)-1] + } + sx, sy, found := strings.Cut(string(src[1:len(src)-1]), ",") + if !found { + return nil, fmt.Errorf("invalid format for point") + } + + x, err := strconv.ParseFloat(sx, 64) + if err != nil { + return nil, err + } + + y, err := strconv.ParseFloat(sy, 64) + if err != nil { + return nil, err + } + + return &Point{P: Vec2{x, y}, Valid: true}, nil +} + +// Scan implements the [database/sql.Scanner] interface. +func (dst *Point) Scan(src any) error { + if src == nil { + *dst = Point{} + return nil + } + + if src, ok := src.(string); ok { + return scanPlanTextAnyToPointScanner{}.Scan([]byte(src), dst) + } + + return fmt.Errorf("cannot scan %T", src) +} + +// Value implements the [database/sql/driver.Valuer] interface. +func (src Point) Value() (driver.Value, error) { + if !src.Valid { + return nil, nil + } + + buf, err := PointCodec{}.PlanEncode(nil, 0, TextFormatCode, src).Encode(src, nil) + if err != nil { + return nil, err + } + return string(buf), err +} + +// MarshalJSON implements the [encoding/json.Marshaler] interface. +func (src Point) MarshalJSON() ([]byte, error) { + if !src.Valid { + return []byte("null"), nil + } + + var buff bytes.Buffer + buff.WriteByte('"') + buff.WriteString(fmt.Sprintf("(%g,%g)", src.P.X, src.P.Y)) + buff.WriteByte('"') + return buff.Bytes(), nil +} + +// UnmarshalJSON implements the [encoding/json.Unmarshaler] interface. +func (dst *Point) UnmarshalJSON(point []byte) error { + p, err := parsePoint(point) + if err != nil { + return err + } + *dst = *p + return nil +} + +type PointCodec struct{} + +func (PointCodec) FormatSupported(format int16) bool { + return format == TextFormatCode || format == BinaryFormatCode +} + +func (PointCodec) PreferredFormat() int16 { + return BinaryFormatCode +} + +func (PointCodec) PlanEncode(m *Map, oid uint32, format int16, value any) EncodePlan { + if _, ok := value.(PointValuer); !ok { + return nil + } + + switch format { + case BinaryFormatCode: + return encodePlanPointCodecBinary{} + case TextFormatCode: + return encodePlanPointCodecText{} + } + + return nil +} + +type encodePlanPointCodecBinary struct{} + +func (encodePlanPointCodecBinary) Encode(value any, buf []byte) (newBuf []byte, err error) { + point, err := value.(PointValuer).PointValue() + if err != nil { + return nil, err + } + + if !point.Valid { + return nil, nil + } + + buf = pgio.AppendUint64(buf, math.Float64bits(point.P.X)) + buf = pgio.AppendUint64(buf, math.Float64bits(point.P.Y)) + return buf, nil +} + +type encodePlanPointCodecText struct{} + +func (encodePlanPointCodecText) Encode(value any, buf []byte) (newBuf []byte, err error) { + point, err := value.(PointValuer).PointValue() + if err != nil { + return nil, err + } + + if !point.Valid { + return nil, nil + } + + return append(buf, fmt.Sprintf(`(%s,%s)`, + strconv.FormatFloat(point.P.X, 'f', -1, 64), + strconv.FormatFloat(point.P.Y, 'f', -1, 64), + )...), nil +} + +func (PointCodec) PlanScan(m *Map, oid uint32, format int16, target any) ScanPlan { + switch format { + case BinaryFormatCode: + if _, ok := target.(PointScanner); ok { + return scanPlanBinaryPointToPointScanner{} + } + case TextFormatCode: + if _, ok := target.(PointScanner); ok { + return scanPlanTextAnyToPointScanner{} + } + } + + return nil +} + +func (c PointCodec) DecodeDatabaseSQLValue(m *Map, oid uint32, format int16, src []byte) (driver.Value, error) { + return codecDecodeToTextFormat(c, m, oid, format, src) +} + +func (c PointCodec) DecodeValue(m *Map, oid uint32, format int16, src []byte) (any, error) { + if src == nil { + return nil, nil + } + + var point Point + err := codecScan(c, m, oid, format, src, &point) + if err != nil { + return nil, err + } + return point, nil +} + +type scanPlanBinaryPointToPointScanner struct{} + +func (scanPlanBinaryPointToPointScanner) Scan(src []byte, dst any) error { + scanner := (dst).(PointScanner) + + if src == nil { + return scanner.ScanPoint(Point{}) + } + + if len(src) != 16 { + return fmt.Errorf("invalid length for point: %v", len(src)) + } + + x := binary.BigEndian.Uint64(src) + y := binary.BigEndian.Uint64(src[8:]) + + return scanner.ScanPoint(Point{ + P: Vec2{math.Float64frombits(x), math.Float64frombits(y)}, + Valid: true, + }) +} + +type scanPlanTextAnyToPointScanner struct{} + +func (scanPlanTextAnyToPointScanner) Scan(src []byte, dst any) error { + scanner := (dst).(PointScanner) + + if src == nil { + return scanner.ScanPoint(Point{}) + } + + if len(src) < 5 { + return fmt.Errorf("invalid length for point: %v", len(src)) + } + + sx, sy, found := strings.Cut(string(src[1:len(src)-1]), ",") + if !found { + return fmt.Errorf("invalid format for point") + } + + x, err := strconv.ParseFloat(sx, 64) + if err != nil { + return err + } + + y, err := strconv.ParseFloat(sy, 64) + if err != nil { + return err + } + + return scanner.ScanPoint(Point{P: Vec2{x, y}, Valid: true}) +} diff --git a/vendor/github.com/jackc/pgx/v5/pgtype/polygon.go b/vendor/github.com/jackc/pgx/v5/pgtype/polygon.go new file mode 100644 index 0000000000..34aa0a6a5f --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgtype/polygon.go @@ -0,0 +1,261 @@ +package pgtype + +import ( + "database/sql/driver" + "encoding/binary" + "fmt" + "math" + "strconv" + "strings" + + "github.com/jackc/pgx/v5/internal/pgio" +) + +type PolygonScanner interface { + ScanPolygon(v Polygon) error +} + +type PolygonValuer interface { + PolygonValue() (Polygon, error) +} + +type Polygon struct { + P []Vec2 + Valid bool +} + +// ScanPolygon implements the [PolygonScanner] interface. +func (p *Polygon) ScanPolygon(v Polygon) error { + *p = v + return nil +} + +// PolygonValue implements the [PolygonValuer] interface. +func (p Polygon) PolygonValue() (Polygon, error) { + return p, nil +} + +// Scan implements the [database/sql.Scanner] interface. +func (p *Polygon) Scan(src any) error { + if src == nil { + *p = Polygon{} + return nil + } + + if src, ok := src.(string); ok { + return scanPlanTextAnyToPolygonScanner{}.Scan([]byte(src), p) + } + + return fmt.Errorf("cannot scan %T", src) +} + +// Value implements the [database/sql/driver.Valuer] interface. +func (p Polygon) Value() (driver.Value, error) { + if !p.Valid { + return nil, nil + } + + buf, err := PolygonCodec{}.PlanEncode(nil, 0, TextFormatCode, p).Encode(p, nil) + if err != nil { + return nil, err + } + + return string(buf), err +} + +type PolygonCodec struct{} + +func (PolygonCodec) FormatSupported(format int16) bool { + return format == TextFormatCode || format == BinaryFormatCode +} + +func (PolygonCodec) PreferredFormat() int16 { + return BinaryFormatCode +} + +func (PolygonCodec) PlanEncode(m *Map, oid uint32, format int16, value any) EncodePlan { + if _, ok := value.(PolygonValuer); !ok { + return nil + } + + switch format { + case BinaryFormatCode: + return encodePlanPolygonCodecBinary{} + case TextFormatCode: + return encodePlanPolygonCodecText{} + } + + return nil +} + +type encodePlanPolygonCodecBinary struct{} + +func (encodePlanPolygonCodecBinary) Encode(value any, buf []byte) (newBuf []byte, err error) { + polygon, err := value.(PolygonValuer).PolygonValue() + if err != nil { + return nil, err + } + + if !polygon.Valid { + return nil, nil + } + + buf = pgio.AppendInt32(buf, int32(len(polygon.P))) + + for _, p := range polygon.P { + buf = pgio.AppendUint64(buf, math.Float64bits(p.X)) + buf = pgio.AppendUint64(buf, math.Float64bits(p.Y)) + } + + return buf, nil +} + +type encodePlanPolygonCodecText struct{} + +func (encodePlanPolygonCodecText) Encode(value any, buf []byte) (newBuf []byte, err error) { + polygon, err := value.(PolygonValuer).PolygonValue() + if err != nil { + return nil, err + } + + if !polygon.Valid { + return nil, nil + } + + buf = append(buf, '(') + + for i, p := range polygon.P { + if i > 0 { + buf = append(buf, ',') + } + buf = append(buf, fmt.Sprintf(`(%s,%s)`, + strconv.FormatFloat(p.X, 'f', -1, 64), + strconv.FormatFloat(p.Y, 'f', -1, 64), + )...) + } + + buf = append(buf, ')') + + return buf, nil +} + +func (PolygonCodec) PlanScan(m *Map, oid uint32, format int16, target any) ScanPlan { + switch format { + case BinaryFormatCode: + if _, ok := target.(PolygonScanner); ok { + return scanPlanBinaryPolygonToPolygonScanner{} + } + case TextFormatCode: + if _, ok := target.(PolygonScanner); ok { + return scanPlanTextAnyToPolygonScanner{} + } + } + + return nil +} + +type scanPlanBinaryPolygonToPolygonScanner struct{} + +func (scanPlanBinaryPolygonToPolygonScanner) Scan(src []byte, dst any) error { + scanner := (dst).(PolygonScanner) + + if src == nil { + return scanner.ScanPolygon(Polygon{}) + } + + if len(src) < 5 { + return fmt.Errorf("invalid length for polygon: %v", len(src)) + } + + pointCount := int(binary.BigEndian.Uint32(src)) + rp := 4 + + if 4+pointCount*16 != len(src) { + return fmt.Errorf("invalid length for Polygon with %d points: %v", pointCount, len(src)) + } + + points := make([]Vec2, pointCount) + for i := range points { + x := binary.BigEndian.Uint64(src[rp:]) + rp += 8 + y := binary.BigEndian.Uint64(src[rp:]) + rp += 8 + points[i] = Vec2{math.Float64frombits(x), math.Float64frombits(y)} + } + + return scanner.ScanPolygon(Polygon{ + P: points, + Valid: true, + }) +} + +type scanPlanTextAnyToPolygonScanner struct{} + +func (scanPlanTextAnyToPolygonScanner) Scan(src []byte, dst any) error { + scanner := (dst).(PolygonScanner) + + if src == nil { + return scanner.ScanPolygon(Polygon{}) + } + + if len(src) < 7 { + return fmt.Errorf("invalid length for Polygon: %v", len(src)) + } + + points := make([]Vec2, 0) + + // Expected format: ((x1,y1),...,(xn,yn)) + str := string(src[1 : len(src)-1]) + + for { + if len(str) == 0 || str[0] != '(' { + return fmt.Errorf("invalid format for Polygon") + } + body, rest, found := strings.Cut(str[1:], ")") + if !found { + return fmt.Errorf("invalid format for Polygon") + } + + sx, sy, found := strings.Cut(body, ",") + if !found { + return fmt.Errorf("invalid format for Polygon") + } + x, err := strconv.ParseFloat(sx, 64) + if err != nil { + return err + } + y, err := strconv.ParseFloat(sy, 64) + if err != nil { + return err + } + + points = append(points, Vec2{x, y}) + + if rest == "" { + break + } + str, found = strings.CutPrefix(rest, ",") + if !found { + return fmt.Errorf("invalid format for Polygon") + } + } + + return scanner.ScanPolygon(Polygon{P: points, Valid: true}) +} + +func (c PolygonCodec) DecodeDatabaseSQLValue(m *Map, oid uint32, format int16, src []byte) (driver.Value, error) { + return codecDecodeToTextFormat(c, m, oid, format, src) +} + +func (c PolygonCodec) DecodeValue(m *Map, oid uint32, format int16, src []byte) (any, error) { + if src == nil { + return nil, nil + } + + var polygon Polygon + err := codecScan(c, m, oid, format, src, &polygon) + if err != nil { + return nil, err + } + return polygon, nil +} diff --git a/vendor/github.com/jackc/pgx/v5/pgtype/qchar.go b/vendor/github.com/jackc/pgx/v5/pgtype/qchar.go new file mode 100644 index 0000000000..3de0b01fcd --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgtype/qchar.go @@ -0,0 +1,163 @@ +package pgtype + +import ( + "database/sql/driver" + "fmt" + "math" +) + +// QCharCodec is for PostgreSQL's special 8-bit-only "char" type more akin to the C +// language's char type, or Go's byte type. (Note that the name in PostgreSQL +// itself is "char", in double-quotes, and not char.) It gets used a lot in +// PostgreSQL's system tables to hold a single ASCII character value (eg +// pg_class.relkind). It is named Qchar for quoted char to disambiguate from SQL +// standard type char. +type QCharCodec struct{} + +func (QCharCodec) FormatSupported(format int16) bool { + return format == TextFormatCode || format == BinaryFormatCode +} + +func (QCharCodec) PreferredFormat() int16 { + return BinaryFormatCode +} + +func (QCharCodec) PlanEncode(m *Map, oid uint32, format int16, value any) EncodePlan { + switch format { + case TextFormatCode, BinaryFormatCode: + switch value.(type) { + case byte: + return encodePlanQcharCodecByte{} + case rune: + return encodePlanQcharCodecRune{} + } + } + + return nil +} + +type encodePlanQcharCodecByte struct{} + +func (encodePlanQcharCodecByte) Encode(value any, buf []byte) (newBuf []byte, err error) { + b := value.(byte) + buf = append(buf, b) + return buf, nil +} + +type encodePlanQcharCodecRune struct{} + +func (encodePlanQcharCodecRune) Encode(value any, buf []byte) (newBuf []byte, err error) { + r := value.(rune) + if r > math.MaxUint8 { + return nil, fmt.Errorf(`%v cannot be encoded to "char"`, r) + } + b := byte(r) + buf = append(buf, b) + return buf, nil +} + +func (QCharCodec) PlanScan(m *Map, oid uint32, format int16, target any) ScanPlan { + switch format { + case TextFormatCode, BinaryFormatCode: + switch target.(type) { + case *byte: + return scanPlanQcharCodecByte{} + case *rune: + return scanPlanQcharCodecRune{} + case *string: + return scanPlanQcharCodecString{} + } + } + + return nil +} + +type scanPlanQcharCodecByte struct{} + +func (scanPlanQcharCodecByte) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + if len(src) > 1 { + return fmt.Errorf(`invalid length for "char": %v`, len(src)) + } + + b := dst.(*byte) + // In the text format the zero value is returned as a zero byte value instead of 0 + if len(src) == 0 { + *b = 0 + } else { + *b = src[0] + } + + return nil +} + +type scanPlanQcharCodecRune struct{} + +func (scanPlanQcharCodecRune) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + if len(src) > 1 { + return fmt.Errorf(`invalid length for "char": %v`, len(src)) + } + + r := dst.(*rune) + // In the text format the zero value is returned as a zero byte value instead of 0 + if len(src) == 0 { + *r = 0 + } else { + *r = rune(src[0]) + } + + return nil +} + +type scanPlanQcharCodecString struct{} + +func (scanPlanQcharCodecString) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + if len(src) > 1 { + return fmt.Errorf(`invalid length for "char": %v`, len(src)) + } + + p := dst.(*string) + // Copy the raw byte so the result matches the text-format *string scan path + // (string(src)) byte-for-byte. string(src[0]) would instead UTF-8-encode the + // byte as a code point, diverging for byte values >= 128. + *p = string(src) + + return nil +} + +func (c QCharCodec) DecodeDatabaseSQLValue(m *Map, oid uint32, format int16, src []byte) (driver.Value, error) { + if src == nil { + return nil, nil + } + + var r rune + err := codecScan(c, m, oid, format, src, &r) + if err != nil { + return nil, err + } + return string(r), nil +} + +func (c QCharCodec) DecodeValue(m *Map, oid uint32, format int16, src []byte) (any, error) { + if src == nil { + return nil, nil + } + + var r rune + err := codecScan(c, m, oid, format, src, &r) + if err != nil { + return nil, err + } + return r, nil +} diff --git a/vendor/github.com/jackc/pgx/v5/pgtype/range.go b/vendor/github.com/jackc/pgx/v5/pgtype/range.go new file mode 100644 index 0000000000..dec153e509 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgtype/range.go @@ -0,0 +1,331 @@ +package pgtype + +import ( + "bytes" + "encoding/binary" + "fmt" +) + +type BoundType byte + +const ( + Inclusive = BoundType('i') + Exclusive = BoundType('e') + Unbounded = BoundType('U') + Empty = BoundType('E') +) + +func (bt BoundType) String() string { + return string(bt) +} + +type untypedTextRange struct { + Lower string + Upper string + LowerType BoundType + UpperType BoundType +} + +func parseUntypedTextRange(src string) (*untypedTextRange, error) { + utr := &untypedTextRange{} + if src == "empty" { + utr.LowerType = Empty + utr.UpperType = Empty + return utr, nil + } + + buf := bytes.NewBufferString(src) + + skipWhitespace(buf) + + r, _, err := buf.ReadRune() + if err != nil { + return nil, fmt.Errorf("invalid lower bound: %w", err) + } + switch r { + case '(': + utr.LowerType = Exclusive + case '[': + utr.LowerType = Inclusive + default: + return nil, fmt.Errorf("missing lower bound, instead got: %v", string(r)) + } + + r, _, err = buf.ReadRune() + if err != nil { + return nil, fmt.Errorf("invalid lower value: %w", err) + } + buf.UnreadRune() + + if r == ',' { + utr.LowerType = Unbounded + } else { + utr.Lower, err = rangeParseValue(buf) + if err != nil { + return nil, fmt.Errorf("invalid lower value: %w", err) + } + } + + r, _, err = buf.ReadRune() + if err != nil { + return nil, fmt.Errorf("missing range separator: %w", err) + } + if r != ',' { + return nil, fmt.Errorf("missing range separator: %v", r) + } + + r, _, err = buf.ReadRune() + if err != nil { + return nil, fmt.Errorf("invalid upper value: %w", err) + } + + if r == ')' || r == ']' { + utr.UpperType = Unbounded + } else { + buf.UnreadRune() + utr.Upper, err = rangeParseValue(buf) + if err != nil { + return nil, fmt.Errorf("invalid upper value: %w", err) + } + + r, _, err = buf.ReadRune() + if err != nil { + return nil, fmt.Errorf("missing upper bound: %w", err) + } + switch r { + case ')': + utr.UpperType = Exclusive + case ']': + utr.UpperType = Inclusive + default: + return nil, fmt.Errorf("missing upper bound, instead got: %v", string(r)) + } + } + + skipWhitespace(buf) + + if buf.Len() > 0 { + return nil, fmt.Errorf("unexpected trailing data: %v", buf.String()) + } + + return utr, nil +} + +func rangeParseValue(buf *bytes.Buffer) (string, error) { + r, _, err := buf.ReadRune() + if err != nil { + return "", err + } + if r == '"' { + return rangeParseQuotedValue(buf) + } + buf.UnreadRune() + + s := &bytes.Buffer{} + + for { + r, _, err := buf.ReadRune() + if err != nil { + return "", err + } + + switch r { + case '\\': + r, _, err = buf.ReadRune() + if err != nil { + return "", err + } + case ',', '[', ']', '(', ')': + buf.UnreadRune() + return s.String(), nil + } + + s.WriteRune(r) + } +} + +func rangeParseQuotedValue(buf *bytes.Buffer) (string, error) { + s := &bytes.Buffer{} + + for { + r, _, err := buf.ReadRune() + if err != nil { + return "", err + } + + switch r { + case '\\': + r, _, err = buf.ReadRune() + if err != nil { + return "", err + } + case '"': + r, _, err = buf.ReadRune() + if err != nil { + return "", err + } + if r != '"' { + buf.UnreadRune() + return s.String(), nil + } + } + s.WriteRune(r) + } +} + +type untypedBinaryRange struct { + Lower []byte + Upper []byte + LowerType BoundType + UpperType BoundType +} + +// 0 = () = 00000 +// 1 = empty = 00001 +// 2 = [) = 00010 +// 4 = (] = 00100 +// 6 = [] = 00110 +// 8 = ) = 01000 +// 12 = ] = 01100 +// 16 = ( = 10000 +// 18 = [ = 10010 +// 24 = = 11000 + +const ( + emptyMask = 1 + lowerInclusiveMask = 2 + upperInclusiveMask = 4 + lowerUnboundedMask = 8 + upperUnboundedMask = 16 +) + +func parseUntypedBinaryRange(src []byte) (*untypedBinaryRange, error) { + ubr := &untypedBinaryRange{} + + if len(src) == 0 { + return nil, fmt.Errorf("range too short: %v", len(src)) + } + + rangeType := src[0] + rp := 1 + + if rangeType&emptyMask > 0 { + if len(src[rp:]) > 0 { + return nil, fmt.Errorf("unexpected trailing bytes parsing empty range: %v", len(src[rp:])) + } + ubr.LowerType = Empty + ubr.UpperType = Empty + return ubr, nil + } + + switch { + case rangeType&lowerInclusiveMask > 0: + ubr.LowerType = Inclusive + case rangeType&lowerUnboundedMask > 0: + ubr.LowerType = Unbounded + default: + ubr.LowerType = Exclusive + } + + switch { + case rangeType&upperInclusiveMask > 0: + ubr.UpperType = Inclusive + case rangeType&upperUnboundedMask > 0: + ubr.UpperType = Unbounded + default: + ubr.UpperType = Exclusive + } + + if ubr.LowerType == Unbounded && ubr.UpperType == Unbounded { + if len(src[rp:]) > 0 { + return nil, fmt.Errorf("unexpected trailing bytes parsing unbounded range: %v", len(src[rp:])) + } + return ubr, nil + } + + if len(src[rp:]) < 4 { + return nil, fmt.Errorf("too few bytes for size: %v", src[rp:]) + } + valueLen := int(binary.BigEndian.Uint32(src[rp:])) + rp += 4 + + if valueLen < 0 || len(src[rp:]) < valueLen { + return nil, fmt.Errorf("range lower bound length %d exceeds remaining %d bytes", valueLen, len(src[rp:])) + } + val := src[rp : rp+valueLen] + rp += valueLen + + if ubr.LowerType != Unbounded { + ubr.Lower = val + } else { + ubr.Upper = val + if len(src[rp:]) > 0 { + return nil, fmt.Errorf("unexpected trailing bytes parsing range: %v", len(src[rp:])) + } + return ubr, nil + } + + if ubr.UpperType != Unbounded { + if len(src[rp:]) < 4 { + return nil, fmt.Errorf("too few bytes for size: %v", src[rp:]) + } + valueLen := int(binary.BigEndian.Uint32(src[rp:])) + rp += 4 + if valueLen < 0 || len(src[rp:]) < valueLen { + return nil, fmt.Errorf("range upper bound length %d exceeds remaining %d bytes", valueLen, len(src[rp:])) + } + ubr.Upper = src[rp : rp+valueLen] + rp += valueLen + } + + if len(src[rp:]) > 0 { + return nil, fmt.Errorf("unexpected trailing bytes parsing range: %v", len(src[rp:])) + } + + return ubr, nil +} + +// Range is a generic range type. +type Range[T any] struct { + Lower T + Upper T + LowerType BoundType + UpperType BoundType + Valid bool +} + +func (r Range[T]) IsNull() bool { + return !r.Valid +} + +func (r Range[T]) BoundTypes() (lower, upper BoundType) { + return r.LowerType, r.UpperType +} + +func (r Range[T]) Bounds() (lower, upper any) { + return &r.Lower, &r.Upper +} + +func (r *Range[T]) ScanNull() error { + *r = Range[T]{} + return nil +} + +func (r *Range[T]) ScanBounds() (lowerTarget, upperTarget any) { + return &r.Lower, &r.Upper +} + +func (r *Range[T]) SetBoundTypes(lower, upper BoundType) error { + if lower == Unbounded || lower == Empty { + var zero T + r.Lower = zero + } + if upper == Unbounded || upper == Empty { + var zero T + r.Upper = zero + } + r.LowerType = lower + r.UpperType = upper + r.Valid = true + return nil +} diff --git a/vendor/github.com/jackc/pgx/v5/pgtype/range_codec.go b/vendor/github.com/jackc/pgx/v5/pgtype/range_codec.go new file mode 100644 index 0000000000..dc1ac8b06b --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgtype/range_codec.go @@ -0,0 +1,377 @@ +package pgtype + +import ( + "database/sql/driver" + "fmt" + + "github.com/jackc/pgx/v5/internal/pgio" +) + +// RangeValuer is a type that can be converted into a PostgreSQL range. +type RangeValuer interface { + // IsNull returns true if the value is SQL NULL. + IsNull() bool + + // BoundTypes returns the lower and upper bound types. + BoundTypes() (lower, upper BoundType) + + // Bounds returns the lower and upper range values. + Bounds() (lower, upper any) +} + +// RangeScanner is a type can be scanned from a PostgreSQL range. +type RangeScanner interface { + // ScanNull sets the value to SQL NULL. + ScanNull() error + + // ScanBounds returns values usable as a scan target. The returned values may not be scanned if the range is empty or + // the bound type is unbounded. + ScanBounds() (lowerTarget, upperTarget any) + + // SetBoundTypes sets the lower and upper bound types. ScanBounds will be called and the returned values scanned + // (if appropriate) before SetBoundTypes is called. If the bound types are unbounded or empty this method must + // also set the bound values. + SetBoundTypes(lower, upper BoundType) error +} + +// RangeCodec is a codec for any range type. +type RangeCodec struct { + ElementType *Type +} + +func (c *RangeCodec) FormatSupported(format int16) bool { + return c.ElementType.Codec.FormatSupported(format) +} + +func (c *RangeCodec) PreferredFormat() int16 { + if c.FormatSupported(BinaryFormatCode) { + return BinaryFormatCode + } + return TextFormatCode +} + +func (c *RangeCodec) PlanEncode(m *Map, oid uint32, format int16, value any) EncodePlan { + if _, ok := value.(RangeValuer); !ok { + return nil + } + + switch format { + case BinaryFormatCode: + return &encodePlanRangeCodecRangeValuerToBinary{rc: c, m: m} + case TextFormatCode: + return &encodePlanRangeCodecRangeValuerToText{rc: c, m: m} + } + + return nil +} + +type encodePlanRangeCodecRangeValuerToBinary struct { + rc *RangeCodec + m *Map +} + +func (plan *encodePlanRangeCodecRangeValuerToBinary) Encode(value any, buf []byte) (newBuf []byte, err error) { + getter := value.(RangeValuer) + + if getter.IsNull() { + return nil, nil + } + + lowerType, upperType := getter.BoundTypes() + lower, upper := getter.Bounds() + + var rangeType byte + switch lowerType { + case Inclusive: + rangeType |= lowerInclusiveMask + case Unbounded: + rangeType |= lowerUnboundedMask + case Exclusive: + case Empty: + return append(buf, emptyMask), nil + default: + return nil, fmt.Errorf("unknown LowerType: %v", lowerType) + } + + switch upperType { + case Inclusive: + rangeType |= upperInclusiveMask + case Unbounded: + rangeType |= upperUnboundedMask + case Exclusive: + default: + return nil, fmt.Errorf("unknown UpperType: %v", upperType) + } + + buf = append(buf, rangeType) + + if lowerType != Unbounded { + if lower == nil { + return nil, fmt.Errorf("Lower cannot be NULL unless LowerType is Unbounded") + } + + sp := len(buf) + buf = pgio.AppendInt32(buf, -1) + + lowerPlan := plan.m.PlanEncode(plan.rc.ElementType.OID, BinaryFormatCode, lower) + if lowerPlan == nil { + return nil, fmt.Errorf("cannot encode %v as element of range", lower) + } + + buf, err = lowerPlan.Encode(lower, buf) + if err != nil { + return nil, fmt.Errorf("failed to encode %v as element of range: %w", lower, err) + } + if buf == nil { + return nil, fmt.Errorf("Lower cannot be NULL unless LowerType is Unbounded") + } + + pgio.SetInt32(buf[sp:], int32(len(buf[sp:])-4)) + } + + if upperType != Unbounded { + if upper == nil { + return nil, fmt.Errorf("Upper cannot be NULL unless UpperType is Unbounded") + } + + sp := len(buf) + buf = pgio.AppendInt32(buf, -1) + + upperPlan := plan.m.PlanEncode(plan.rc.ElementType.OID, BinaryFormatCode, upper) + if upperPlan == nil { + return nil, fmt.Errorf("cannot encode %v as element of range", upper) + } + + buf, err = upperPlan.Encode(upper, buf) + if err != nil { + return nil, fmt.Errorf("failed to encode %v as element of range: %w", upper, err) + } + if buf == nil { + return nil, fmt.Errorf("Upper cannot be NULL unless UpperType is Unbounded") + } + + pgio.SetInt32(buf[sp:], int32(len(buf[sp:])-4)) + } + + return buf, nil +} + +type encodePlanRangeCodecRangeValuerToText struct { + rc *RangeCodec + m *Map +} + +func (plan *encodePlanRangeCodecRangeValuerToText) Encode(value any, buf []byte) (newBuf []byte, err error) { + getter := value.(RangeValuer) + + if getter.IsNull() { + return nil, nil + } + + lowerType, upperType := getter.BoundTypes() + lower, upper := getter.Bounds() + + switch lowerType { + case Exclusive, Unbounded: + buf = append(buf, '(') + case Inclusive: + buf = append(buf, '[') + case Empty: + return append(buf, "empty"...), nil + default: + return nil, fmt.Errorf("unknown lower bound type %v", lowerType) + } + + if lowerType != Unbounded { + if lower == nil { + return nil, fmt.Errorf("Lower cannot be NULL unless LowerType is Unbounded") + } + + lowerPlan := plan.m.PlanEncode(plan.rc.ElementType.OID, TextFormatCode, lower) + if lowerPlan == nil { + return nil, fmt.Errorf("cannot encode %v as element of range", lower) + } + + buf, err = lowerPlan.Encode(lower, buf) + if err != nil { + return nil, fmt.Errorf("failed to encode %v as element of range: %w", lower, err) + } + if buf == nil { + return nil, fmt.Errorf("Lower cannot be NULL unless LowerType is Unbounded") + } + } + + buf = append(buf, ',') + + if upperType != Unbounded { + if upper == nil { + return nil, fmt.Errorf("Upper cannot be NULL unless UpperType is Unbounded") + } + + upperPlan := plan.m.PlanEncode(plan.rc.ElementType.OID, TextFormatCode, upper) + if upperPlan == nil { + return nil, fmt.Errorf("cannot encode %v as element of range", upper) + } + + buf, err = upperPlan.Encode(upper, buf) + if err != nil { + return nil, fmt.Errorf("failed to encode %v as element of range: %w", upper, err) + } + if buf == nil { + return nil, fmt.Errorf("Upper cannot be NULL unless UpperType is Unbounded") + } + } + + switch upperType { + case Exclusive, Unbounded: + buf = append(buf, ')') + case Inclusive: + buf = append(buf, ']') + default: + return nil, fmt.Errorf("unknown upper bound type %v", upperType) + } + + return buf, nil +} + +func (c *RangeCodec) PlanScan(m *Map, oid uint32, format int16, target any) ScanPlan { + switch format { + case BinaryFormatCode: + if _, ok := target.(RangeScanner); ok { + return &scanPlanBinaryRangeToRangeScanner{rc: c, m: m} + } + case TextFormatCode: + if _, ok := target.(RangeScanner); ok { + return &scanPlanTextRangeToRangeScanner{rc: c, m: m} + } + } + + return nil +} + +type scanPlanBinaryRangeToRangeScanner struct { + rc *RangeCodec + m *Map +} + +func (plan *scanPlanBinaryRangeToRangeScanner) Scan(src []byte, target any) error { + rangeScanner := (target).(RangeScanner) + + if src == nil { + return rangeScanner.ScanNull() + } + + ubr, err := parseUntypedBinaryRange(src) + if err != nil { + return err + } + + if ubr.LowerType == Empty { + return rangeScanner.SetBoundTypes(ubr.LowerType, ubr.UpperType) + } + + lowerTarget, upperTarget := rangeScanner.ScanBounds() + + if ubr.LowerType == Inclusive || ubr.LowerType == Exclusive { + lowerPlan := plan.m.PlanScan(plan.rc.ElementType.OID, BinaryFormatCode, lowerTarget) + if lowerPlan == nil { + return fmt.Errorf("cannot scan into %v from range element", lowerTarget) + } + + err = lowerPlan.Scan(ubr.Lower, lowerTarget) + if err != nil { + return fmt.Errorf("cannot scan into %v from range element: %w", lowerTarget, err) + } + } + + if ubr.UpperType == Inclusive || ubr.UpperType == Exclusive { + upperPlan := plan.m.PlanScan(plan.rc.ElementType.OID, BinaryFormatCode, upperTarget) + if upperPlan == nil { + return fmt.Errorf("cannot scan into %v from range element", upperTarget) + } + + err = upperPlan.Scan(ubr.Upper, upperTarget) + if err != nil { + return fmt.Errorf("cannot scan into %v from range element: %w", upperTarget, err) + } + } + + return rangeScanner.SetBoundTypes(ubr.LowerType, ubr.UpperType) +} + +type scanPlanTextRangeToRangeScanner struct { + rc *RangeCodec + m *Map +} + +func (plan *scanPlanTextRangeToRangeScanner) Scan(src []byte, target any) error { + rangeScanner := (target).(RangeScanner) + + if src == nil { + return rangeScanner.ScanNull() + } + + utr, err := parseUntypedTextRange(string(src)) + if err != nil { + return err + } + + if utr.LowerType == Empty { + return rangeScanner.SetBoundTypes(utr.LowerType, utr.UpperType) + } + + lowerTarget, upperTarget := rangeScanner.ScanBounds() + + if utr.LowerType == Inclusive || utr.LowerType == Exclusive { + lowerPlan := plan.m.PlanScan(plan.rc.ElementType.OID, TextFormatCode, lowerTarget) + if lowerPlan == nil { + return fmt.Errorf("cannot scan into %v from range element", lowerTarget) + } + + err = lowerPlan.Scan([]byte(utr.Lower), lowerTarget) + if err != nil { + return fmt.Errorf("cannot scan into %v from range element: %w", lowerTarget, err) + } + } + + if utr.UpperType == Inclusive || utr.UpperType == Exclusive { + upperPlan := plan.m.PlanScan(plan.rc.ElementType.OID, TextFormatCode, upperTarget) + if upperPlan == nil { + return fmt.Errorf("cannot scan into %v from range element", upperTarget) + } + + err = upperPlan.Scan([]byte(utr.Upper), upperTarget) + if err != nil { + return fmt.Errorf("cannot scan into %v from range element: %w", upperTarget, err) + } + } + + return rangeScanner.SetBoundTypes(utr.LowerType, utr.UpperType) +} + +func (c *RangeCodec) DecodeDatabaseSQLValue(m *Map, oid uint32, format int16, src []byte) (driver.Value, error) { + if src == nil { + return nil, nil + } + + switch format { + case TextFormatCode: + return string(src), nil + case BinaryFormatCode: + buf := make([]byte, len(src)) + copy(buf, src) + return buf, nil + default: + return nil, fmt.Errorf("unknown format code %d", format) + } +} + +func (c *RangeCodec) DecodeValue(m *Map, oid uint32, format int16, src []byte) (any, error) { + if src == nil { + return nil, nil + } + + var r Range[any] + err := c.PlanScan(m, oid, format, &r).Scan(src, &r) + return r, err +} diff --git a/vendor/github.com/jackc/pgx/v5/pgtype/record_codec.go b/vendor/github.com/jackc/pgx/v5/pgtype/record_codec.go new file mode 100644 index 0000000000..a663e4dc22 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgtype/record_codec.go @@ -0,0 +1,123 @@ +package pgtype + +import ( + "database/sql/driver" + "fmt" +) + +// ArrayGetter is a type that can be converted into a PostgreSQL array. + +// RecordCodec is a codec for the generic PostgreSQL record type such as is created with the "row" function. Record can +// only decode the binary format. The text format output format from PostgreSQL does not include type information and +// is therefore impossible to decode. Encoding is impossible because PostgreSQL does not support input of generic +// records. +type RecordCodec struct{} + +func (RecordCodec) FormatSupported(format int16) bool { + return format == BinaryFormatCode +} + +func (RecordCodec) PreferredFormat() int16 { + return BinaryFormatCode +} + +func (RecordCodec) PlanEncode(m *Map, oid uint32, format int16, value any) EncodePlan { + return nil +} + +func (RecordCodec) PlanScan(m *Map, oid uint32, format int16, target any) ScanPlan { + if format == BinaryFormatCode { + if _, ok := target.(CompositeIndexScanner); ok { + return &scanPlanBinaryRecordToCompositeIndexScanner{m: m} + } + } + + return nil +} + +type scanPlanBinaryRecordToCompositeIndexScanner struct { + m *Map +} + +func (plan *scanPlanBinaryRecordToCompositeIndexScanner) Scan(src []byte, target any) error { + targetScanner := (target).(CompositeIndexScanner) + + if src == nil { + return targetScanner.ScanNull() + } + + scanner := NewCompositeBinaryScanner(plan.m, src) + for i := 0; scanner.Next(); i++ { + fieldTarget := targetScanner.ScanIndex(i) + if fieldTarget != nil { + fieldPlan := plan.m.PlanScan(scanner.OID(), BinaryFormatCode, fieldTarget) + if fieldPlan == nil { + return fmt.Errorf("unable to scan OID %d in binary format into %v", scanner.OID(), fieldTarget) + } + + err := fieldPlan.Scan(scanner.Bytes(), fieldTarget) + if err != nil { + return err + } + } + } + + if err := scanner.Err(); err != nil { + return err + } + + return nil +} + +func (RecordCodec) DecodeDatabaseSQLValue(m *Map, oid uint32, format int16, src []byte) (driver.Value, error) { + if src == nil { + return nil, nil + } + + switch format { + case TextFormatCode: + return string(src), nil + case BinaryFormatCode: + buf := make([]byte, len(src)) + copy(buf, src) + return buf, nil + default: + return nil, fmt.Errorf("unknown format code %d", format) + } +} + +func (RecordCodec) DecodeValue(m *Map, oid uint32, format int16, src []byte) (any, error) { + if src == nil { + return nil, nil + } + + switch format { + case TextFormatCode: + return string(src), nil + case BinaryFormatCode: + scanner := NewCompositeBinaryScanner(m, src) + values := make([]any, scanner.FieldCount()) + for i := 0; scanner.Next(); i++ { + var v any + fieldPlan := m.PlanScan(scanner.OID(), BinaryFormatCode, &v) + if fieldPlan == nil { + return nil, fmt.Errorf("unable to scan OID %d in binary format into %v", scanner.OID(), v) + } + + err := fieldPlan.Scan(scanner.Bytes(), &v) + if err != nil { + return nil, err + } + + values[i] = v + } + + if err := scanner.Err(); err != nil { + return nil, err + } + + return values, nil + default: + return nil, fmt.Errorf("unknown format code %d", format) + } +} diff --git a/vendor/github.com/jackc/pgx/v5/pgtype/register_default_pg_types.go b/vendor/github.com/jackc/pgx/v5/pgtype/register_default_pg_types.go new file mode 100644 index 0000000000..be1ca4a189 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgtype/register_default_pg_types.go @@ -0,0 +1,35 @@ +//go:build !nopgxregisterdefaulttypes + +package pgtype + +func registerDefaultPgTypeVariants[T any](m *Map, name string) { + arrayName := "_" + name + + var value T + m.RegisterDefaultPgType(value, name) // T + m.RegisterDefaultPgType(&value, name) // *T + + var sliceT []T + m.RegisterDefaultPgType(sliceT, arrayName) // []T + m.RegisterDefaultPgType(&sliceT, arrayName) // *[]T + + var slicePtrT []*T + m.RegisterDefaultPgType(slicePtrT, arrayName) // []*T + m.RegisterDefaultPgType(&slicePtrT, arrayName) // *[]*T + + var arrayOfT Array[T] + m.RegisterDefaultPgType(arrayOfT, arrayName) // Array[T] + m.RegisterDefaultPgType(&arrayOfT, arrayName) // *Array[T] + + var arrayOfPtrT Array[*T] + m.RegisterDefaultPgType(arrayOfPtrT, arrayName) // Array[*T] + m.RegisterDefaultPgType(&arrayOfPtrT, arrayName) // *Array[*T] + + var flatArrayOfT FlatArray[T] + m.RegisterDefaultPgType(flatArrayOfT, arrayName) // FlatArray[T] + m.RegisterDefaultPgType(&flatArrayOfT, arrayName) // *FlatArray[T] + + var flatArrayOfPtrT FlatArray[*T] + m.RegisterDefaultPgType(flatArrayOfPtrT, arrayName) // FlatArray[*T] + m.RegisterDefaultPgType(&flatArrayOfPtrT, arrayName) // *FlatArray[*T] +} diff --git a/vendor/github.com/jackc/pgx/v5/pgtype/register_default_pg_types_disabled.go b/vendor/github.com/jackc/pgx/v5/pgtype/register_default_pg_types_disabled.go new file mode 100644 index 0000000000..56fe7c226a --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgtype/register_default_pg_types_disabled.go @@ -0,0 +1,6 @@ +//go:build nopgxregisterdefaulttypes + +package pgtype + +func registerDefaultPgTypeVariants[T any](m *Map, name string) { +} diff --git a/vendor/github.com/jackc/pgx/v5/pgtype/text.go b/vendor/github.com/jackc/pgx/v5/pgtype/text.go new file mode 100644 index 0000000000..e08b12549e --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgtype/text.go @@ -0,0 +1,226 @@ +package pgtype + +import ( + "database/sql/driver" + "encoding/json" + "fmt" +) + +type TextScanner interface { + ScanText(v Text) error +} + +type TextValuer interface { + TextValue() (Text, error) +} + +type Text struct { + String string + Valid bool +} + +// ScanText implements the [TextScanner] interface. +func (t *Text) ScanText(v Text) error { + *t = v + return nil +} + +// TextValue implements the [TextValuer] interface. +func (t Text) TextValue() (Text, error) { + return t, nil +} + +// Scan implements the [database/sql.Scanner] interface. +func (dst *Text) Scan(src any) error { + if src == nil { + *dst = Text{} + return nil + } + + switch src := src.(type) { + case string: + *dst = Text{String: src, Valid: true} + return nil + case []byte: + *dst = Text{String: string(src), Valid: true} + return nil + } + + return fmt.Errorf("cannot scan %T", src) +} + +// Value implements the [database/sql/driver.Valuer] interface. +func (src Text) Value() (driver.Value, error) { + if !src.Valid { + return nil, nil + } + return src.String, nil +} + +// MarshalJSON implements the [encoding/json.Marshaler] interface. +func (src Text) MarshalJSON() ([]byte, error) { + if !src.Valid { + return []byte("null"), nil + } + + return json.Marshal(src.String) +} + +// UnmarshalJSON implements the [encoding/json.Unmarshaler] interface. +func (dst *Text) UnmarshalJSON(b []byte) error { + var s *string + err := json.Unmarshal(b, &s) + if err != nil { + return err + } + + if s == nil { + *dst = Text{} + } else { + *dst = Text{String: *s, Valid: true} + } + + return nil +} + +type TextCodec struct{} + +func (TextCodec) FormatSupported(format int16) bool { + return format == TextFormatCode || format == BinaryFormatCode +} + +func (TextCodec) PreferredFormat() int16 { + return TextFormatCode +} + +func (TextCodec) PlanEncode(m *Map, oid uint32, format int16, value any) EncodePlan { + switch format { + case TextFormatCode, BinaryFormatCode: + switch value.(type) { + case string: + return encodePlanTextCodecString{} + case []byte: + return encodePlanTextCodecByteSlice{} + case TextValuer: + return encodePlanTextCodecTextValuer{} + } + } + + return nil +} + +type encodePlanTextCodecString struct{} + +func (encodePlanTextCodecString) Encode(value any, buf []byte) (newBuf []byte, err error) { + s := value.(string) + buf = append(buf, s...) + return buf, nil +} + +type encodePlanTextCodecByteSlice struct{} + +func (encodePlanTextCodecByteSlice) Encode(value any, buf []byte) (newBuf []byte, err error) { + s := value.([]byte) + buf = append(buf, s...) + return buf, nil +} + +type encodePlanTextCodecStringer struct{} + +func (encodePlanTextCodecStringer) Encode(value any, buf []byte) (newBuf []byte, err error) { + s := value.(fmt.Stringer) + buf = append(buf, s.String()...) + return buf, nil +} + +type encodePlanTextCodecTextValuer struct{} + +func (encodePlanTextCodecTextValuer) Encode(value any, buf []byte) (newBuf []byte, err error) { + text, err := value.(TextValuer).TextValue() + if err != nil { + return nil, err + } + + if !text.Valid { + return nil, nil + } + + buf = append(buf, text.String...) + return buf, nil +} + +func (TextCodec) PlanScan(m *Map, oid uint32, format int16, target any) ScanPlan { + switch format { + case TextFormatCode, BinaryFormatCode: + switch target.(type) { + case *string: + return scanPlanTextAnyToString{} + case *[]byte: + return scanPlanAnyToNewByteSlice{} + case BytesScanner: + return scanPlanAnyToByteScanner{} + case TextScanner: + return scanPlanTextAnyToTextScanner{} + } + } + + return nil +} + +func (c TextCodec) DecodeDatabaseSQLValue(m *Map, oid uint32, format int16, src []byte) (driver.Value, error) { + return c.DecodeValue(m, oid, format, src) +} + +func (c TextCodec) DecodeValue(m *Map, oid uint32, format int16, src []byte) (any, error) { + if src == nil { + return nil, nil + } + + return string(src), nil +} + +type scanPlanTextAnyToString struct{} + +func (scanPlanTextAnyToString) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + p := (dst).(*string) + *p = string(src) + + return nil +} + +type scanPlanAnyToNewByteSlice struct{} + +func (scanPlanAnyToNewByteSlice) Scan(src []byte, dst any) error { + p := (dst).(*[]byte) + if src == nil { + *p = nil + } else { + *p = make([]byte, len(src)) + copy(*p, src) + } + + return nil +} + +type scanPlanAnyToByteScanner struct{} + +func (scanPlanAnyToByteScanner) Scan(src []byte, dst any) error { + p := (dst).(BytesScanner) + return p.ScanBytes(src) +} + +type scanPlanTextAnyToTextScanner struct{} + +func (scanPlanTextAnyToTextScanner) Scan(src []byte, dst any) error { + scanner := (dst).(TextScanner) + + if src == nil { + return scanner.ScanText(Text{}) + } + + return scanner.ScanText(Text{String: string(src), Valid: true}) +} diff --git a/vendor/github.com/jackc/pgx/v5/pgtype/text_format_only_codec.go b/vendor/github.com/jackc/pgx/v5/pgtype/text_format_only_codec.go new file mode 100644 index 0000000000..d5e4cdb381 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgtype/text_format_only_codec.go @@ -0,0 +1,13 @@ +package pgtype + +type TextFormatOnlyCodec struct { + Codec +} + +func (c *TextFormatOnlyCodec) FormatSupported(format int16) bool { + return format == TextFormatCode && c.Codec.FormatSupported(format) +} + +func (TextFormatOnlyCodec) PreferredFormat() int16 { + return TextFormatCode +} diff --git a/vendor/github.com/jackc/pgx/v5/pgtype/tid.go b/vendor/github.com/jackc/pgx/v5/pgtype/tid.go new file mode 100644 index 0000000000..98d067a652 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgtype/tid.go @@ -0,0 +1,240 @@ +package pgtype + +import ( + "database/sql/driver" + "encoding/binary" + "fmt" + "strconv" + "strings" + + "github.com/jackc/pgx/v5/internal/pgio" +) + +type TIDScanner interface { + ScanTID(v TID) error +} + +type TIDValuer interface { + TIDValue() (TID, error) +} + +// TID is PostgreSQL's Tuple Identifier type. +// +// When one does +// +// select ctid, * from some_table; +// +// it is the data type of the ctid hidden system column. +// +// It is currently implemented as a pair unsigned two byte integers. +// Its conversion functions can be found in src/backend/utils/adt/tid.c +// in the PostgreSQL sources. +type TID struct { + BlockNumber uint32 + OffsetNumber uint16 + Valid bool +} + +// ScanTID implements the [TIDScanner] interface. +func (b *TID) ScanTID(v TID) error { + *b = v + return nil +} + +// TIDValue implements the [TIDValuer] interface. +func (b TID) TIDValue() (TID, error) { + return b, nil +} + +// Scan implements the [database/sql.Scanner] interface. +func (dst *TID) Scan(src any) error { + if src == nil { + *dst = TID{} + return nil + } + + if src, ok := src.(string); ok { + return scanPlanTextAnyToTIDScanner{}.Scan([]byte(src), dst) + } + + return fmt.Errorf("cannot scan %T", src) +} + +// Value implements the [database/sql/driver.Valuer] interface. +func (src TID) Value() (driver.Value, error) { + if !src.Valid { + return nil, nil + } + + buf, err := TIDCodec{}.PlanEncode(nil, 0, TextFormatCode, src).Encode(src, nil) + if err != nil { + return nil, err + } + return string(buf), err +} + +type TIDCodec struct{} + +func (TIDCodec) FormatSupported(format int16) bool { + return format == TextFormatCode || format == BinaryFormatCode +} + +func (TIDCodec) PreferredFormat() int16 { + return BinaryFormatCode +} + +func (TIDCodec) PlanEncode(m *Map, oid uint32, format int16, value any) EncodePlan { + if _, ok := value.(TIDValuer); !ok { + return nil + } + + switch format { + case BinaryFormatCode: + return encodePlanTIDCodecBinary{} + case TextFormatCode: + return encodePlanTIDCodecText{} + } + + return nil +} + +type encodePlanTIDCodecBinary struct{} + +func (encodePlanTIDCodecBinary) Encode(value any, buf []byte) (newBuf []byte, err error) { + tid, err := value.(TIDValuer).TIDValue() + if err != nil { + return nil, err + } + + if !tid.Valid { + return nil, nil + } + + buf = pgio.AppendUint32(buf, tid.BlockNumber) + buf = pgio.AppendUint16(buf, tid.OffsetNumber) + return buf, nil +} + +type encodePlanTIDCodecText struct{} + +func (encodePlanTIDCodecText) Encode(value any, buf []byte) (newBuf []byte, err error) { + tid, err := value.(TIDValuer).TIDValue() + if err != nil { + return nil, err + } + + if !tid.Valid { + return nil, nil + } + + buf = append(buf, fmt.Sprintf(`(%d,%d)`, tid.BlockNumber, tid.OffsetNumber)...) + return buf, nil +} + +func (TIDCodec) PlanScan(m *Map, oid uint32, format int16, target any) ScanPlan { + switch format { + case BinaryFormatCode: + switch target.(type) { + case TIDScanner: + return scanPlanBinaryTIDToTIDScanner{} + case TextScanner: + return scanPlanBinaryTIDToTextScanner{} + } + case TextFormatCode: + if _, ok := target.(TIDScanner); ok { + return scanPlanTextAnyToTIDScanner{} + } + } + + return nil +} + +type scanPlanBinaryTIDToTIDScanner struct{} + +func (scanPlanBinaryTIDToTIDScanner) Scan(src []byte, dst any) error { + scanner := (dst).(TIDScanner) + + if src == nil { + return scanner.ScanTID(TID{}) + } + + if len(src) != 6 { + return fmt.Errorf("invalid length for tid: %v", len(src)) + } + + return scanner.ScanTID(TID{ + BlockNumber: binary.BigEndian.Uint32(src), + OffsetNumber: binary.BigEndian.Uint16(src[4:]), + Valid: true, + }) +} + +type scanPlanBinaryTIDToTextScanner struct{} + +func (scanPlanBinaryTIDToTextScanner) Scan(src []byte, dst any) error { + scanner := (dst).(TextScanner) + + if src == nil { + return scanner.ScanText(Text{}) + } + + if len(src) != 6 { + return fmt.Errorf("invalid length for tid: %v", len(src)) + } + + blockNumber := binary.BigEndian.Uint32(src) + offsetNumber := binary.BigEndian.Uint16(src[4:]) + + return scanner.ScanText(Text{ + String: fmt.Sprintf(`(%d,%d)`, blockNumber, offsetNumber), + Valid: true, + }) +} + +type scanPlanTextAnyToTIDScanner struct{} + +func (scanPlanTextAnyToTIDScanner) Scan(src []byte, dst any) error { + scanner := (dst).(TIDScanner) + + if src == nil { + return scanner.ScanTID(TID{}) + } + + if len(src) < 5 { + return fmt.Errorf("invalid length for tid: %v", len(src)) + } + + block, offset, found := strings.Cut(string(src[1:len(src)-1]), ",") + if !found { + return fmt.Errorf("invalid format for tid") + } + + blockNumber, err := strconv.ParseUint(block, 10, 32) + if err != nil { + return err + } + + offsetNumber, err := strconv.ParseUint(offset, 10, 16) + if err != nil { + return err + } + + return scanner.ScanTID(TID{BlockNumber: uint32(blockNumber), OffsetNumber: uint16(offsetNumber), Valid: true}) +} + +func (c TIDCodec) DecodeDatabaseSQLValue(m *Map, oid uint32, format int16, src []byte) (driver.Value, error) { + return codecDecodeToTextFormat(c, m, oid, format, src) +} + +func (c TIDCodec) DecodeValue(m *Map, oid uint32, format int16, src []byte) (any, error) { + if src == nil { + return nil, nil + } + + var tid TID + err := codecScan(c, m, oid, format, src, &tid) + if err != nil { + return nil, err + } + return tid, nil +} diff --git a/vendor/github.com/jackc/pgx/v5/pgtype/time.go b/vendor/github.com/jackc/pgx/v5/pgtype/time.go new file mode 100644 index 0000000000..72cdb50035 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgtype/time.go @@ -0,0 +1,273 @@ +package pgtype + +import ( + "database/sql/driver" + "encoding/binary" + "fmt" + "strconv" + + "github.com/jackc/pgx/v5/internal/pgio" +) + +type TimeScanner interface { + ScanTime(v Time) error +} + +type TimeValuer interface { + TimeValue() (Time, error) +} + +// Time represents the PostgreSQL time type. The PostgreSQL time is a time of day without time zone. +// +// Time is represented as the number of microseconds since midnight in the same way that PostgreSQL does. Other time and +// date types in pgtype can use time.Time as the underlying representation. However, pgtype.Time type cannot due to +// needing to handle 24:00:00. time.Time converts that to 00:00:00 on the following day. +// +// The time with time zone type is not supported. Use of time with time zone is discouraged by the PostgreSQL documentation. +type Time struct { + Microseconds int64 // Number of microseconds since midnight + Valid bool +} + +// ScanTime implements the [TimeScanner] interface. +func (t *Time) ScanTime(v Time) error { + *t = v + return nil +} + +// TimeValue implements the [TimeValuer] interface. +func (t Time) TimeValue() (Time, error) { + return t, nil +} + +// Scan implements the [database/sql.Scanner] interface. +func (t *Time) Scan(src any) error { + if src == nil { + *t = Time{} + return nil + } + + if src, ok := src.(string); ok { + err := scanPlanTextAnyToTimeScanner{}.Scan([]byte(src), t) + if err != nil { + t.Microseconds = 0 + t.Valid = false + } + return err + } + + return fmt.Errorf("cannot scan %T", src) +} + +// Value implements the [database/sql/driver.Valuer] interface. +func (t Time) Value() (driver.Value, error) { + if !t.Valid { + return nil, nil + } + + buf, err := TimeCodec{}.PlanEncode(nil, 0, TextFormatCode, t).Encode(t, nil) + if err != nil { + return nil, err + } + return string(buf), err +} + +type TimeCodec struct{} + +func (TimeCodec) FormatSupported(format int16) bool { + return format == TextFormatCode || format == BinaryFormatCode +} + +func (TimeCodec) PreferredFormat() int16 { + return BinaryFormatCode +} + +func (TimeCodec) PlanEncode(m *Map, oid uint32, format int16, value any) EncodePlan { + if _, ok := value.(TimeValuer); !ok { + return nil + } + + switch format { + case BinaryFormatCode: + return encodePlanTimeCodecBinary{} + case TextFormatCode: + return encodePlanTimeCodecText{} + } + + return nil +} + +type encodePlanTimeCodecBinary struct{} + +func (encodePlanTimeCodecBinary) Encode(value any, buf []byte) (newBuf []byte, err error) { + t, err := value.(TimeValuer).TimeValue() + if err != nil { + return nil, err + } + + if !t.Valid { + return nil, nil + } + + return pgio.AppendInt64(buf, t.Microseconds), nil +} + +type encodePlanTimeCodecText struct{} + +func (encodePlanTimeCodecText) Encode(value any, buf []byte) (newBuf []byte, err error) { + t, err := value.(TimeValuer).TimeValue() + if err != nil { + return nil, err + } + + if !t.Valid { + return nil, nil + } + + usec := t.Microseconds + hours := usec / microsecondsPerHour + usec -= hours * microsecondsPerHour + minutes := usec / microsecondsPerMinute + usec -= minutes * microsecondsPerMinute + seconds := usec / microsecondsPerSecond + usec -= seconds * microsecondsPerSecond + + s := fmt.Sprintf("%02d:%02d:%02d.%06d", hours, minutes, seconds, usec) + + return append(buf, s...), nil +} + +func (TimeCodec) PlanScan(m *Map, oid uint32, format int16, target any) ScanPlan { + switch format { + case BinaryFormatCode: + switch target.(type) { + case TimeScanner: + return scanPlanBinaryTimeToTimeScanner{} + case TextScanner: + return scanPlanBinaryTimeToTextScanner{} + } + case TextFormatCode: + if _, ok := target.(TimeScanner); ok { + return scanPlanTextAnyToTimeScanner{} + } + } + + return nil +} + +type scanPlanBinaryTimeToTimeScanner struct{} + +func (scanPlanBinaryTimeToTimeScanner) Scan(src []byte, dst any) error { + scanner := (dst).(TimeScanner) + + if src == nil { + return scanner.ScanTime(Time{}) + } + + if len(src) != 8 { + return fmt.Errorf("invalid length for time: %v", len(src)) + } + + usec := int64(binary.BigEndian.Uint64(src)) + + return scanner.ScanTime(Time{Microseconds: usec, Valid: true}) +} + +type scanPlanBinaryTimeToTextScanner struct{} + +func (scanPlanBinaryTimeToTextScanner) Scan(src []byte, dst any) error { + ts, ok := (dst).(TextScanner) + if !ok { + return ErrScanTargetTypeChanged + } + + if src == nil { + return ts.ScanText(Text{}) + } + + if len(src) != 8 { + return fmt.Errorf("invalid length for time: %v", len(src)) + } + + usec := int64(binary.BigEndian.Uint64(src)) + + tim := Time{Microseconds: usec, Valid: true} + + buf, err := TimeCodec{}.PlanEncode(nil, 0, TextFormatCode, tim).Encode(tim, nil) + if err != nil { + return err + } + + return ts.ScanText(Text{String: string(buf), Valid: true}) +} + +type scanPlanTextAnyToTimeScanner struct{} + +func (scanPlanTextAnyToTimeScanner) Scan(src []byte, dst any) error { + scanner := (dst).(TimeScanner) + + if src == nil { + return scanner.ScanTime(Time{}) + } + + s := string(src) + + if len(s) < 8 || s[2] != ':' || s[5] != ':' { + return fmt.Errorf("cannot decode %v into Time", s) + } + + hours, err := strconv.ParseInt(s[0:2], 10, 64) + if err != nil { + return fmt.Errorf("cannot decode %v into Time", s) + } + usec := hours * microsecondsPerHour + + minutes, err := strconv.ParseInt(s[3:5], 10, 64) + if err != nil { + return fmt.Errorf("cannot decode %v into Time", s) + } + usec += minutes * microsecondsPerMinute + + seconds, err := strconv.ParseInt(s[6:8], 10, 64) + if err != nil { + return fmt.Errorf("cannot decode %v into Time", s) + } + usec += seconds * microsecondsPerSecond + + if len(s) > 9 { + if s[8] != '.' || len(s) > 15 { + return fmt.Errorf("cannot decode %v into Time", s) + } + + fraction := s[9:] + n, err := strconv.ParseInt(fraction, 10, 64) + if err != nil { + return fmt.Errorf("cannot decode %v into Time", s) + } + + for i := len(fraction); i < 6; i++ { + n *= 10 + } + + usec += n + } + + return scanner.ScanTime(Time{Microseconds: usec, Valid: true}) +} + +func (c TimeCodec) DecodeDatabaseSQLValue(m *Map, oid uint32, format int16, src []byte) (driver.Value, error) { + return codecDecodeToTextFormat(c, m, oid, format, src) +} + +func (c TimeCodec) DecodeValue(m *Map, oid uint32, format int16, src []byte) (any, error) { + if src == nil { + return nil, nil + } + + var t Time + err := codecScan(c, m, oid, format, src, &t) + if err != nil { + return nil, err + } + return t, nil +} diff --git a/vendor/github.com/jackc/pgx/v5/pgtype/timestamp.go b/vendor/github.com/jackc/pgx/v5/pgtype/timestamp.go new file mode 100644 index 0000000000..405c77e96b --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgtype/timestamp.go @@ -0,0 +1,368 @@ +package pgtype + +import ( + "database/sql/driver" + "encoding/binary" + "encoding/json" + "fmt" + "strings" + "time" + + "github.com/jackc/pgx/v5/internal/pgio" +) + +const ( + pgTimestampFormat = "2006-01-02 15:04:05.999999999" + jsonISO8601 = "2006-01-02T15:04:05.999999999" +) + +type TimestampScanner interface { + ScanTimestamp(v Timestamp) error +} + +type TimestampValuer interface { + TimestampValue() (Timestamp, error) +} + +// Timestamp represents the PostgreSQL timestamp type. +type Timestamp struct { + Time time.Time // Time zone will be ignored when encoding to PostgreSQL. + InfinityModifier InfinityModifier + Valid bool +} + +// ScanTimestamp implements the [TimestampScanner] interface. +func (ts *Timestamp) ScanTimestamp(v Timestamp) error { + *ts = v + return nil +} + +// TimestampValue implements the [TimestampValuer] interface. +func (ts Timestamp) TimestampValue() (Timestamp, error) { + return ts, nil +} + +// Scan implements the [database/sql.Scanner] interface. +func (ts *Timestamp) Scan(src any) error { + if src == nil { + *ts = Timestamp{} + return nil + } + + switch src := src.(type) { + case string: + return (&scanPlanTextTimestampToTimestampScanner{}).Scan([]byte(src), ts) + case time.Time: + *ts = Timestamp{Time: src, Valid: true} + return nil + } + + return fmt.Errorf("cannot scan %T", src) +} + +// Value implements the [database/sql/driver.Valuer] interface. +func (ts Timestamp) Value() (driver.Value, error) { + if !ts.Valid { + return nil, nil + } + + if ts.InfinityModifier != Finite { + return ts.InfinityModifier.String(), nil + } + return ts.Time, nil +} + +// MarshalJSON implements the [encoding/json.Marshaler] interface. +func (ts Timestamp) MarshalJSON() ([]byte, error) { + if !ts.Valid { + return []byte("null"), nil + } + + var s string + + switch ts.InfinityModifier { + case Finite: + s = ts.Time.Format(jsonISO8601) + case Infinity: + s = "infinity" + case NegativeInfinity: + s = "-infinity" + } + + return json.Marshal(s) +} + +// UnmarshalJSON implements the [encoding/json.Unmarshaler] interface. +func (ts *Timestamp) UnmarshalJSON(b []byte) error { + var s *string + err := json.Unmarshal(b, &s) + if err != nil { + return err + } + + if s == nil { + *ts = Timestamp{} + return nil + } + + switch *s { + case "infinity": + *ts = Timestamp{Valid: true, InfinityModifier: Infinity} + case "-infinity": + *ts = Timestamp{Valid: true, InfinityModifier: -Infinity} + default: + // Parse time with or without timezone + tss := *s + // PostgreSQL uses ISO 8601 without timezone for to_json function and casting from a string to timestamp + tim, err := time.Parse(time.RFC3339Nano, tss) + if err == nil { + *ts = Timestamp{Time: tim, Valid: true} + return nil + } + tim, err = time.ParseInLocation(jsonISO8601, tss, time.UTC) + if err == nil { + *ts = Timestamp{Time: tim, Valid: true} + return nil + } + ts.Valid = false + return fmt.Errorf("cannot unmarshal %s to timestamp with layout %s or %s (%w)", + *s, time.RFC3339Nano, jsonISO8601, err) + } + return nil +} + +type TimestampCodec struct { + // ScanLocation is the location that the time is assumed to be in for scanning. This is different from + // TimestamptzCodec.ScanLocation in that this setting does change the instant in time that the timestamp represents. + ScanLocation *time.Location +} + +func (*TimestampCodec) FormatSupported(format int16) bool { + return format == TextFormatCode || format == BinaryFormatCode +} + +func (*TimestampCodec) PreferredFormat() int16 { + return BinaryFormatCode +} + +func (*TimestampCodec) PlanEncode(m *Map, oid uint32, format int16, value any) EncodePlan { + if _, ok := value.(TimestampValuer); !ok { + return nil + } + + switch format { + case BinaryFormatCode: + return encodePlanTimestampCodecBinary{} + case TextFormatCode: + return encodePlanTimestampCodecText{} + } + + return nil +} + +type encodePlanTimestampCodecBinary struct{} + +func (encodePlanTimestampCodecBinary) Encode(value any, buf []byte) (newBuf []byte, err error) { + ts, err := value.(TimestampValuer).TimestampValue() + if err != nil { + return nil, err + } + + if !ts.Valid { + return nil, nil + } + + var microsecSinceY2K int64 + switch ts.InfinityModifier { + case Finite: + t := discardTimeZone(ts.Time) + microsecSinceUnixEpoch := t.Unix()*1_000_000 + int64(t.Nanosecond())/1000 + microsecSinceY2K = microsecSinceUnixEpoch - microsecFromUnixEpochToY2K + case Infinity: + microsecSinceY2K = infinityMicrosecondOffset + case NegativeInfinity: + microsecSinceY2K = negativeInfinityMicrosecondOffset + } + + buf = pgio.AppendInt64(buf, microsecSinceY2K) + + return buf, nil +} + +type encodePlanTimestampCodecText struct{} + +func (encodePlanTimestampCodecText) Encode(value any, buf []byte) (newBuf []byte, err error) { + ts, err := value.(TimestampValuer).TimestampValue() + if err != nil { + return nil, err + } + + if !ts.Valid { + return nil, nil + } + + var s string + + switch ts.InfinityModifier { + case Finite: + t := discardTimeZone(ts.Time) + + // Year 0000 is 1 BC + bc := false + if year := t.Year(); year <= 0 { + year = -year + 1 + t = time.Date(year, t.Month(), t.Day(), t.Hour(), t.Minute(), t.Second(), t.Nanosecond(), time.UTC) + bc = true + } + + s = t.Truncate(time.Microsecond).Format(pgTimestampFormat) + + if bc { + s += " BC" + } + case Infinity: + s = "infinity" + case NegativeInfinity: + s = "-infinity" + } + + buf = append(buf, s...) + + return buf, nil +} + +func discardTimeZone(t time.Time) time.Time { + if t.Location() != time.UTC { + return time.Date(t.Year(), t.Month(), t.Day(), t.Hour(), t.Minute(), t.Second(), t.Nanosecond(), time.UTC) + } + + return t +} + +func (c *TimestampCodec) PlanScan(m *Map, oid uint32, format int16, target any) ScanPlan { + switch format { + case BinaryFormatCode: + if _, ok := target.(TimestampScanner); ok { + return &scanPlanBinaryTimestampToTimestampScanner{location: c.ScanLocation} + } + case TextFormatCode: + if _, ok := target.(TimestampScanner); ok { + return &scanPlanTextTimestampToTimestampScanner{location: c.ScanLocation} + } + } + + return nil +} + +type scanPlanBinaryTimestampToTimestampScanner struct{ location *time.Location } + +func (plan *scanPlanBinaryTimestampToTimestampScanner) Scan(src []byte, dst any) error { + scanner := (dst).(TimestampScanner) + + if src == nil { + return scanner.ScanTimestamp(Timestamp{}) + } + + if len(src) != 8 { + return fmt.Errorf("invalid length for timestamp: %v", len(src)) + } + + var ts Timestamp + microsecSinceY2K := int64(binary.BigEndian.Uint64(src)) + + switch microsecSinceY2K { + case infinityMicrosecondOffset: + ts = Timestamp{Valid: true, InfinityModifier: Infinity} + case negativeInfinityMicrosecondOffset: + ts = Timestamp{Valid: true, InfinityModifier: -Infinity} + default: + tim := time.Unix( + microsecFromUnixEpochToY2K/1_000_000+microsecSinceY2K/1_000_000, + (microsecFromUnixEpochToY2K%1_000_000*1_000)+(microsecSinceY2K%1_000_000*1000), + ).UTC() + if plan.location != nil { + tim = time.Date(tim.Year(), tim.Month(), tim.Day(), tim.Hour(), tim.Minute(), tim.Second(), tim.Nanosecond(), plan.location) + } + ts = Timestamp{Time: tim, Valid: true} + } + + return scanner.ScanTimestamp(ts) +} + +type scanPlanTextTimestampToTimestampScanner struct{ location *time.Location } + +func (plan *scanPlanTextTimestampToTimestampScanner) Scan(src []byte, dst any) error { + scanner := (dst).(TimestampScanner) + + if src == nil { + return scanner.ScanTimestamp(Timestamp{}) + } + + var ts Timestamp + sbuf := string(src) + switch sbuf { + case "infinity": + ts = Timestamp{Valid: true, InfinityModifier: Infinity} + case "-infinity": + ts = Timestamp{Valid: true, InfinityModifier: -Infinity} + default: + bc := false + if strings.HasSuffix(sbuf, " BC") { + sbuf = sbuf[:len(sbuf)-3] + bc = true + } + tim, err := time.Parse(pgTimestampFormat, sbuf) + if err != nil { + return err + } + + if bc { + year := -tim.Year() + 1 + tim = time.Date(year, tim.Month(), tim.Day(), tim.Hour(), tim.Minute(), tim.Second(), tim.Nanosecond(), tim.Location()) + } + + if plan.location != nil { + tim = time.Date(tim.Year(), tim.Month(), tim.Day(), tim.Hour(), tim.Minute(), tim.Second(), tim.Nanosecond(), plan.location) + } + + ts = Timestamp{Time: tim, Valid: true} + } + + return scanner.ScanTimestamp(ts) +} + +func (c *TimestampCodec) DecodeDatabaseSQLValue(m *Map, oid uint32, format int16, src []byte) (driver.Value, error) { + if src == nil { + return nil, nil + } + + var ts Timestamp + err := codecScan(c, m, oid, format, src, &ts) + if err != nil { + return nil, err + } + + if ts.InfinityModifier != Finite { + return ts.InfinityModifier.String(), nil + } + + return ts.Time, nil +} + +func (c *TimestampCodec) DecodeValue(m *Map, oid uint32, format int16, src []byte) (any, error) { + if src == nil { + return nil, nil + } + + var ts Timestamp + err := codecScan(c, m, oid, format, src, &ts) + if err != nil { + return nil, err + } + + if ts.InfinityModifier != Finite { + return ts.InfinityModifier, nil + } + + return ts.Time, nil +} diff --git a/vendor/github.com/jackc/pgx/v5/pgtype/timestamptz.go b/vendor/github.com/jackc/pgx/v5/pgtype/timestamptz.go new file mode 100644 index 0000000000..139312a591 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgtype/timestamptz.go @@ -0,0 +1,370 @@ +package pgtype + +import ( + "database/sql/driver" + "encoding/binary" + "encoding/json" + "fmt" + "strings" + "time" + + "github.com/jackc/pgx/v5/internal/pgio" +) + +const ( + pgTimestamptzHourFormat = "2006-01-02 15:04:05.999999999Z07" + pgTimestamptzMinuteFormat = "2006-01-02 15:04:05.999999999Z07:00" + pgTimestamptzSecondFormat = "2006-01-02 15:04:05.999999999Z07:00:00" + microsecFromUnixEpochToY2K = 946_684_800 * 1_000_000 +) + +const ( + negativeInfinityMicrosecondOffset = -9223372036854775808 + infinityMicrosecondOffset = 9223372036854775807 +) + +type TimestamptzScanner interface { + ScanTimestamptz(v Timestamptz) error +} + +type TimestamptzValuer interface { + TimestamptzValue() (Timestamptz, error) +} + +// Timestamptz represents the PostgreSQL timestamptz type. +type Timestamptz struct { + Time time.Time + InfinityModifier InfinityModifier + Valid bool +} + +// ScanTimestamptz implements the [TimestamptzScanner] interface. +func (tstz *Timestamptz) ScanTimestamptz(v Timestamptz) error { + *tstz = v + return nil +} + +// TimestamptzValue implements the [TimestamptzValuer] interface. +func (tstz Timestamptz) TimestamptzValue() (Timestamptz, error) { + return tstz, nil +} + +// Scan implements the [database/sql.Scanner] interface. +func (tstz *Timestamptz) Scan(src any) error { + if src == nil { + *tstz = Timestamptz{} + return nil + } + + switch src := src.(type) { + case string: + return (&scanPlanTextTimestamptzToTimestamptzScanner{}).Scan([]byte(src), tstz) + case time.Time: + *tstz = Timestamptz{Time: src, Valid: true} + return nil + } + + return fmt.Errorf("cannot scan %T", src) +} + +// Value implements the [database/sql/driver.Valuer] interface. +func (tstz Timestamptz) Value() (driver.Value, error) { + if !tstz.Valid { + return nil, nil + } + + if tstz.InfinityModifier != Finite { + return tstz.InfinityModifier.String(), nil + } + return tstz.Time, nil +} + +// MarshalJSON implements the [encoding/json.Marshaler] interface. +func (tstz Timestamptz) MarshalJSON() ([]byte, error) { + if !tstz.Valid { + return []byte("null"), nil + } + + var s string + + switch tstz.InfinityModifier { + case Finite: + s = tstz.Time.Format(time.RFC3339Nano) + case Infinity: + s = "infinity" + case NegativeInfinity: + s = "-infinity" + } + + return json.Marshal(s) +} + +// UnmarshalJSON implements the [encoding/json.Unmarshaler] interface. +func (tstz *Timestamptz) UnmarshalJSON(b []byte) error { + var s *string + err := json.Unmarshal(b, &s) + if err != nil { + return err + } + + if s == nil { + *tstz = Timestamptz{} + return nil + } + + switch *s { + case "infinity": + *tstz = Timestamptz{Valid: true, InfinityModifier: Infinity} + case "-infinity": + *tstz = Timestamptz{Valid: true, InfinityModifier: -Infinity} + default: + // PostgreSQL uses ISO 8601 for to_json function and casting from a string to timestamptz + tim, err := time.Parse(time.RFC3339Nano, *s) + if err != nil { + return err + } + + *tstz = Timestamptz{Time: tim, Valid: true} + } + + return nil +} + +type TimestamptzCodec struct { + // ScanLocation is the location to return scanned timestamptz values in. This does not change the instant in time that + // the timestamptz represents. + ScanLocation *time.Location +} + +func (*TimestamptzCodec) FormatSupported(format int16) bool { + return format == TextFormatCode || format == BinaryFormatCode +} + +func (*TimestamptzCodec) PreferredFormat() int16 { + return BinaryFormatCode +} + +func (*TimestamptzCodec) PlanEncode(m *Map, oid uint32, format int16, value any) EncodePlan { + if _, ok := value.(TimestamptzValuer); !ok { + return nil + } + + switch format { + case BinaryFormatCode: + return encodePlanTimestamptzCodecBinary{} + case TextFormatCode: + return encodePlanTimestamptzCodecText{} + } + + return nil +} + +type encodePlanTimestamptzCodecBinary struct{} + +func (encodePlanTimestamptzCodecBinary) Encode(value any, buf []byte) (newBuf []byte, err error) { + ts, err := value.(TimestamptzValuer).TimestamptzValue() + if err != nil { + return nil, err + } + + if !ts.Valid { + return nil, nil + } + + var microsecSinceY2K int64 + switch ts.InfinityModifier { + case Finite: + microsecSinceUnixEpoch := ts.Time.Unix()*1000000 + int64(ts.Time.Nanosecond())/1000 + microsecSinceY2K = microsecSinceUnixEpoch - microsecFromUnixEpochToY2K + case Infinity: + microsecSinceY2K = infinityMicrosecondOffset + case NegativeInfinity: + microsecSinceY2K = negativeInfinityMicrosecondOffset + } + + buf = pgio.AppendInt64(buf, microsecSinceY2K) + + return buf, nil +} + +type encodePlanTimestamptzCodecText struct{} + +func (encodePlanTimestamptzCodecText) Encode(value any, buf []byte) (newBuf []byte, err error) { + ts, err := value.(TimestamptzValuer).TimestamptzValue() + if err != nil { + return nil, err + } + + if !ts.Valid { + return nil, nil + } + + var s string + + switch ts.InfinityModifier { + case Finite: + + t := ts.Time.UTC().Truncate(time.Microsecond) + + // Year 0000 is 1 BC + bc := false + if year := t.Year(); year <= 0 { + year = -year + 1 + t = time.Date(year, t.Month(), t.Day(), t.Hour(), t.Minute(), t.Second(), t.Nanosecond(), time.UTC) + bc = true + } + + s = t.Format(pgTimestamptzSecondFormat) + + if bc { + s += " BC" + } + case Infinity: + s = "infinity" + case NegativeInfinity: + s = "-infinity" + } + + buf = append(buf, s...) + + return buf, nil +} + +func (c *TimestamptzCodec) PlanScan(m *Map, oid uint32, format int16, target any) ScanPlan { + switch format { + case BinaryFormatCode: + if _, ok := target.(TimestamptzScanner); ok { + return &scanPlanBinaryTimestamptzToTimestamptzScanner{location: c.ScanLocation} + } + case TextFormatCode: + if _, ok := target.(TimestamptzScanner); ok { + return &scanPlanTextTimestamptzToTimestamptzScanner{location: c.ScanLocation} + } + } + + return nil +} + +type scanPlanBinaryTimestamptzToTimestamptzScanner struct{ location *time.Location } + +func (plan *scanPlanBinaryTimestamptzToTimestamptzScanner) Scan(src []byte, dst any) error { + scanner := (dst).(TimestamptzScanner) + + if src == nil { + return scanner.ScanTimestamptz(Timestamptz{}) + } + + if len(src) != 8 { + return fmt.Errorf("invalid length for timestamptz: %v", len(src)) + } + + var tstz Timestamptz + microsecSinceY2K := int64(binary.BigEndian.Uint64(src)) + + switch microsecSinceY2K { + case infinityMicrosecondOffset: + tstz = Timestamptz{Valid: true, InfinityModifier: Infinity} + case negativeInfinityMicrosecondOffset: + tstz = Timestamptz{Valid: true, InfinityModifier: -Infinity} + default: + tim := time.Unix( + microsecFromUnixEpochToY2K/1_000_000+microsecSinceY2K/1_000_000, + (microsecFromUnixEpochToY2K%1_000_000*1_000)+(microsecSinceY2K%1_000_000*1_000), + ) + if plan.location != nil { + tim = tim.In(plan.location) + } + tstz = Timestamptz{Time: tim, Valid: true} + } + + return scanner.ScanTimestamptz(tstz) +} + +type scanPlanTextTimestamptzToTimestamptzScanner struct{ location *time.Location } + +func (plan *scanPlanTextTimestamptzToTimestamptzScanner) Scan(src []byte, dst any) error { + scanner := (dst).(TimestamptzScanner) + + if src == nil { + return scanner.ScanTimestamptz(Timestamptz{}) + } + + var tstz Timestamptz + sbuf := string(src) + switch sbuf { + case "infinity": + tstz = Timestamptz{Valid: true, InfinityModifier: Infinity} + case "-infinity": + tstz = Timestamptz{Valid: true, InfinityModifier: -Infinity} + default: + bc := false + if strings.HasSuffix(sbuf, " BC") { + sbuf = sbuf[:len(sbuf)-3] + bc = true + } + + var format string + switch { + case len(sbuf) >= 9 && (sbuf[len(sbuf)-9] == '-' || sbuf[len(sbuf)-9] == '+'): + format = pgTimestamptzSecondFormat + case len(sbuf) >= 6 && (sbuf[len(sbuf)-6] == '-' || sbuf[len(sbuf)-6] == '+'): + format = pgTimestamptzMinuteFormat + default: + format = pgTimestamptzHourFormat + } + + tim, err := time.Parse(format, sbuf) + if err != nil { + return err + } + + if bc { + year := -tim.Year() + 1 + tim = time.Date(year, tim.Month(), tim.Day(), tim.Hour(), tim.Minute(), tim.Second(), tim.Nanosecond(), tim.Location()) + } + + if plan.location != nil { + tim = tim.In(plan.location) + } + + tstz = Timestamptz{Time: tim, Valid: true} + } + + return scanner.ScanTimestamptz(tstz) +} + +func (c *TimestamptzCodec) DecodeDatabaseSQLValue(m *Map, oid uint32, format int16, src []byte) (driver.Value, error) { + if src == nil { + return nil, nil + } + + var tstz Timestamptz + err := codecScan(c, m, oid, format, src, &tstz) + if err != nil { + return nil, err + } + + if tstz.InfinityModifier != Finite { + return tstz.InfinityModifier.String(), nil + } + + return tstz.Time, nil +} + +func (c *TimestamptzCodec) DecodeValue(m *Map, oid uint32, format int16, src []byte) (any, error) { + if src == nil { + return nil, nil + } + + var tstz Timestamptz + err := codecScan(c, m, oid, format, src, &tstz) + if err != nil { + return nil, err + } + + if tstz.InfinityModifier != Finite { + return tstz.InfinityModifier, nil + } + + return tstz.Time, nil +} diff --git a/vendor/github.com/jackc/pgx/v5/pgtype/tsvector.go b/vendor/github.com/jackc/pgx/v5/pgtype/tsvector.go new file mode 100644 index 0000000000..cc7b831671 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgtype/tsvector.go @@ -0,0 +1,514 @@ +package pgtype + +import ( + "bytes" + "database/sql/driver" + "encoding/binary" + "fmt" + "strconv" + "strings" + + "github.com/jackc/pgx/v5/internal/pgio" +) + +type TSVectorScanner interface { + ScanTSVector(TSVector) error +} + +type TSVectorValuer interface { + TSVectorValue() (TSVector, error) +} + +// TSVector represents a PostgreSQL tsvector value. +type TSVector struct { + Lexemes []TSVectorLexeme + Valid bool +} + +// TSVectorLexeme represents a lexeme within a tsvector, consisting of a word and its positions. +type TSVectorLexeme struct { + Word string + Positions []TSVectorPosition +} + +// ScanTSVector implements the [TSVectorScanner] interface. +func (t *TSVector) ScanTSVector(v TSVector) error { + *t = v + return nil +} + +// TSVectorValue implements the [TSVectorValuer] interface. +func (t TSVector) TSVectorValue() (TSVector, error) { + return t, nil +} + +func (t TSVector) String() string { + buf, _ := encodePlanTSVectorCodecText{}.Encode(t, nil) + return string(buf) +} + +// Scan implements the [database/sql.Scanner] interface. +func (t *TSVector) Scan(src any) error { + if src == nil { + *t = TSVector{} + return nil + } + + if src, ok := src.(string); ok { + return scanPlanTextAnyToTSVectorScanner{}.scanString(src, t) + } + + return fmt.Errorf("cannot scan %T", src) +} + +// Value implements the [database/sql/driver.Valuer] interface. +func (t TSVector) Value() (driver.Value, error) { + if !t.Valid { + return nil, nil + } + + buf, err := TSVectorCodec{}.PlanEncode(nil, 0, TextFormatCode, t).Encode(t, nil) + if err != nil { + return nil, err + } + + return string(buf), nil +} + +// TSVectorWeight represents the weight label of a lexeme position in a tsvector. +type TSVectorWeight byte + +const ( + TSVectorWeightA = TSVectorWeight('A') + TSVectorWeightB = TSVectorWeight('B') + TSVectorWeightC = TSVectorWeight('C') + TSVectorWeightD = TSVectorWeight('D') +) + +// tsvectorWeightToBinary converts a TSVectorWeight to the 2-bit binary encoding used by PostgreSQL. +func tsvectorWeightToBinary(w TSVectorWeight) uint16 { + switch w { + case TSVectorWeightA: + return 3 + case TSVectorWeightB: + return 2 + case TSVectorWeightC: + return 1 + default: + return 0 // D or unset + } +} + +// tsvectorWeightFromBinary converts a 2-bit binary weight value to a TSVectorWeight. +func tsvectorWeightFromBinary(b uint16) TSVectorWeight { + switch b { + case 3: + return TSVectorWeightA + case 2: + return TSVectorWeightB + case 1: + return TSVectorWeightC + default: + return TSVectorWeightD + } +} + +// TSVectorPosition represents a lexeme position and its optional weight within a tsvector. +type TSVectorPosition struct { + Position uint16 + Weight TSVectorWeight +} + +func (p TSVectorPosition) String() string { + s := strconv.FormatUint(uint64(p.Position), 10) + if p.Weight != 0 && p.Weight != TSVectorWeightD { + s += string(p.Weight) + } + return s +} + +type TSVectorCodec struct{} + +func (TSVectorCodec) FormatSupported(format int16) bool { + return format == TextFormatCode || format == BinaryFormatCode +} + +func (TSVectorCodec) PreferredFormat() int16 { + return BinaryFormatCode +} + +func (TSVectorCodec) PlanEncode(m *Map, oid uint32, format int16, value any) EncodePlan { + if _, ok := value.(TSVectorValuer); !ok { + return nil + } + + switch format { + case BinaryFormatCode: + return encodePlanTSVectorCodecBinary{} + case TextFormatCode: + return encodePlanTSVectorCodecText{} + } + + return nil +} + +type encodePlanTSVectorCodecBinary struct{} + +func (encodePlanTSVectorCodecBinary) Encode(value any, buf []byte) ([]byte, error) { + tsv, err := value.(TSVectorValuer).TSVectorValue() + if err != nil { + return nil, err + } + + if !tsv.Valid { + return nil, nil + } + + buf = pgio.AppendInt32(buf, int32(len(tsv.Lexemes))) + + for _, entry := range tsv.Lexemes { + buf = append(buf, entry.Word...) + buf = append(buf, 0x00) + buf = pgio.AppendUint16(buf, uint16(len(entry.Positions))) + + // Each position is a uint16: weight (2 bits) | position (14 bits) + for _, pos := range entry.Positions { + packed := tsvectorWeightToBinary(pos.Weight)<<14 | pos.Position&0x3FFF + buf = pgio.AppendUint16(buf, packed) + } + } + + return buf, nil +} + +type scanPlanBinaryTSVectorToTSVectorScanner struct{} + +func (scanPlanBinaryTSVectorToTSVectorScanner) Scan(src []byte, dst any) error { + scanner := (dst).(TSVectorScanner) + + if src == nil { + return scanner.ScanTSVector(TSVector{}) + } + + rp := 0 + + const ( + uint16Len = 2 + uint32Len = 4 + ) + + if len(src[rp:]) < uint32Len { + return fmt.Errorf("tsvector incomplete %v", src) + } + entryCount := int(int32(binary.BigEndian.Uint32(src[rp:]))) + rp += uint32Len + + if entryCount < 0 { + return fmt.Errorf("tsvector invalid lexeme count: %d", entryCount) + } + // Each lexeme carries at minimum a 1-byte NUL terminator and a 2-byte position count, so + // entryCount cannot exceed remaining/3. This bounds the up-front make() against a malicious + // server claiming a huge lexeme count in a small message. + if maxEntries := len(src[rp:]) / 3; entryCount > maxEntries { + return fmt.Errorf("tsvector invalid lexeme count %d for %d remaining bytes", entryCount, len(src[rp:])) + } + + var tsv TSVector + if entryCount > 0 { + tsv.Lexemes = make([]TSVectorLexeme, entryCount) + } + + for i := range entryCount { + nullIndex := bytes.IndexByte(src[rp:], 0x00) + if nullIndex == -1 { + return fmt.Errorf("invalid tsvector binary format: missing null terminator") + } + + lexeme := TSVectorLexeme{Word: string(src[rp : rp+nullIndex])} + rp += nullIndex + 1 // skip past null terminator + + // Read position count. + if len(src[rp:]) < uint16Len { + return fmt.Errorf("invalid tsvector binary format: incomplete position count") + } + + numPositions := int(binary.BigEndian.Uint16(src[rp:])) + rp += uint16Len + + // Read each packed position: weight (2 bits) | position (14 bits) + if len(src[rp:]) < numPositions*uint16Len { + return fmt.Errorf("invalid tsvector binary format: incomplete positions") + } + + if numPositions > 0 { + lexeme.Positions = make([]TSVectorPosition, numPositions) + for pos := range numPositions { + packed := binary.BigEndian.Uint16(src[rp:]) + rp += uint16Len + lexeme.Positions[pos] = TSVectorPosition{ + Position: packed & 0x3FFF, + Weight: tsvectorWeightFromBinary(packed >> 14), + } + } + } + + tsv.Lexemes[i] = lexeme + } + tsv.Valid = true + + return scanner.ScanTSVector(tsv) +} + +var tsvectorLexemeReplacer = strings.NewReplacer( + `\`, `\\`, + `'`, `\'`, +) + +type encodePlanTSVectorCodecText struct{} + +func (encodePlanTSVectorCodecText) Encode(value any, buf []byte) ([]byte, error) { + tsv, err := value.(TSVectorValuer).TSVectorValue() + if err != nil { + return nil, err + } + + if !tsv.Valid { + return nil, nil + } + + if buf == nil { + buf = []byte{} + } + + for i, lex := range tsv.Lexemes { + if i > 0 { + buf = append(buf, ' ') + } + + buf = append(buf, '\'') + buf = append(buf, tsvectorLexemeReplacer.Replace(lex.Word)...) + buf = append(buf, '\'') + + sep := byte(':') + for _, p := range lex.Positions { + buf = append(buf, sep) + buf = append(buf, p.String()...) + sep = ',' + } + } + + return buf, nil +} + +func (TSVectorCodec) PlanScan(m *Map, oid uint32, format int16, target any) ScanPlan { + switch format { + case BinaryFormatCode: + if _, ok := target.(TSVectorScanner); ok { + return scanPlanBinaryTSVectorToTSVectorScanner{} + } + case TextFormatCode: + if _, ok := target.(TSVectorScanner); ok { + return scanPlanTextAnyToTSVectorScanner{} + } + } + + return nil +} + +type scanPlanTextAnyToTSVectorScanner struct{} + +func (s scanPlanTextAnyToTSVectorScanner) Scan(src []byte, dst any) error { + scanner := (dst).(TSVectorScanner) + + if src == nil { + return scanner.ScanTSVector(TSVector{}) + } + + return s.scanString(string(src), scanner) +} + +func (scanPlanTextAnyToTSVectorScanner) scanString(src string, scanner TSVectorScanner) error { + tsv, err := parseTSVector(src) + if err != nil { + return err + } + return scanner.ScanTSVector(tsv) +} + +func (c TSVectorCodec) DecodeDatabaseSQLValue(m *Map, oid uint32, format int16, src []byte) (driver.Value, error) { + return codecDecodeToTextFormat(c, m, oid, format, src) +} + +func (c TSVectorCodec) DecodeValue(m *Map, oid uint32, format int16, src []byte) (any, error) { + if src == nil { + return nil, nil + } + + var tsv TSVector + err := codecScan(c, m, oid, format, src, &tsv) + if err != nil { + return nil, err + } + return tsv, nil +} + +type tsvectorParser struct { + str string + pos int +} + +func (p *tsvectorParser) atEnd() bool { + return p.pos >= len(p.str) +} + +func (p *tsvectorParser) peek() byte { + return p.str[p.pos] +} + +func (p *tsvectorParser) consume() (byte, bool) { + if p.pos >= len(p.str) { + return 0, true + } + b := p.str[p.pos] + p.pos++ + return b, false +} + +func (p *tsvectorParser) consumeSpaces() { + for !p.atEnd() && p.peek() == ' ' { + p.consume() + } +} + +// consumeLexeme consumes a single-quoted lexeme, handling single quotes and backslash escapes. +func (p *tsvectorParser) consumeLexeme() (string, error) { + ch, end := p.consume() + if end || ch != '\'' { + return "", fmt.Errorf("invalid tsvector format: lexeme must start with a single quote") + } + + var buf strings.Builder + for { + ch, end := p.consume() + if end { + return "", fmt.Errorf("invalid tsvector format: unterminated quoted lexeme") + } + + switch ch { + case '\'': + // Escaped quote ('') — write a literal single quote + if !p.atEnd() && p.peek() == '\'' { + p.consume() + buf.WriteByte('\'') + } else { + // Closing quote — lexeme is complete + return buf.String(), nil + } + case '\\': + next, end := p.consume() + if end { + return "", fmt.Errorf("invalid tsvector format: unexpected end after backslash") + } + buf.WriteByte(next) + default: + buf.WriteByte(ch) + } + } +} + +// consumePositions consumes a comma-separated list of position[weight] values. +func (p *tsvectorParser) consumePositions() ([]TSVectorPosition, error) { + var positions []TSVectorPosition + + for { + pos, err := p.consumePosition() + if err != nil { + return nil, err + } + positions = append(positions, pos) + + if p.atEnd() || p.peek() != ',' { + break + } + + p.consume() // skip ',' + } + + return positions, nil +} + +// consumePosition consumes a single position number with optional weight letter. +func (p *tsvectorParser) consumePosition() (TSVectorPosition, error) { + start := p.pos + + for !p.atEnd() && p.peek() >= '0' && p.peek() <= '9' { + p.consume() + } + + if p.pos == start { + return TSVectorPosition{}, fmt.Errorf("invalid tsvector format: expected position number") + } + + num, err := strconv.ParseUint(p.str[start:p.pos], 10, 16) + if err != nil { + return TSVectorPosition{}, fmt.Errorf("invalid tsvector format: invalid position number %q", p.str[start:p.pos]) + } + + pos := TSVectorPosition{Position: uint16(num), Weight: TSVectorWeightD} + + // Check for optional weight letter + if !p.atEnd() { + switch p.peek() { + case 'A', 'a': + pos.Weight = TSVectorWeightA + case 'B', 'b': + pos.Weight = TSVectorWeightB + case 'C', 'c': + pos.Weight = TSVectorWeightC + case 'D', 'd': + pos.Weight = TSVectorWeightD + default: + return pos, nil + } + p.consume() + } + + return pos, nil +} + +// parseTSVector parses a PostgreSQL tsvector text representation. +func parseTSVector(s string) (TSVector, error) { + result := TSVector{} + p := &tsvectorParser{str: strings.TrimSpace(s), pos: 0} + + for !p.atEnd() { + p.consumeSpaces() + if p.atEnd() { + break + } + + word, err := p.consumeLexeme() + if err != nil { + return TSVector{}, err + } + + entry := TSVectorLexeme{Word: word} + + // Check for optional positions after ':' + if !p.atEnd() && p.peek() == ':' { + p.consume() // skip ':' + + positions, err := p.consumePositions() + if err != nil { + return TSVector{}, err + } + entry.Positions = positions + } + + result.Lexemes = append(result.Lexemes, entry) + } + + result.Valid = true + + return result, nil +} diff --git a/vendor/github.com/jackc/pgx/v5/pgtype/uint32.go b/vendor/github.com/jackc/pgx/v5/pgtype/uint32.go new file mode 100644 index 0000000000..e6d4b1cf6e --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgtype/uint32.go @@ -0,0 +1,352 @@ +package pgtype + +import ( + "database/sql/driver" + "encoding/binary" + "encoding/json" + "fmt" + "math" + "strconv" + + "github.com/jackc/pgx/v5/internal/pgio" +) + +type Uint32Scanner interface { + ScanUint32(v Uint32) error +} + +type Uint32Valuer interface { + Uint32Value() (Uint32, error) +} + +// Uint32 is the core type that is used to represent PostgreSQL types such as OID, CID, and XID. +type Uint32 struct { + Uint32 uint32 + Valid bool +} + +// ScanUint32 implements the [Uint32Scanner] interface. +func (n *Uint32) ScanUint32(v Uint32) error { + *n = v + return nil +} + +// Uint32Value implements the [Uint32Valuer] interface. +func (n Uint32) Uint32Value() (Uint32, error) { + return n, nil +} + +// Scan implements the [database/sql.Scanner] interface. +func (dst *Uint32) Scan(src any) error { + if src == nil { + *dst = Uint32{} + return nil + } + + var n int64 + + switch src := src.(type) { + case int64: + n = src + case string: + un, err := strconv.ParseUint(src, 10, 32) + if err != nil { + return err + } + n = int64(un) + default: + return fmt.Errorf("cannot scan %T", src) + } + + if n < 0 { + return fmt.Errorf("%d is less than the minimum value for Uint32", n) + } + if n > math.MaxUint32 { + return fmt.Errorf("%d is greater than maximum value for Uint32", n) + } + + *dst = Uint32{Uint32: uint32(n), Valid: true} + + return nil +} + +// Value implements the [database/sql/driver.Valuer] interface. +func (src Uint32) Value() (driver.Value, error) { + if !src.Valid { + return nil, nil + } + return int64(src.Uint32), nil +} + +// MarshalJSON implements the [encoding/json.Marshaler] interface. +func (src Uint32) MarshalJSON() ([]byte, error) { + if !src.Valid { + return []byte("null"), nil + } + return json.Marshal(src.Uint32) +} + +// UnmarshalJSON implements the [encoding/json.Unmarshaler] interface. +func (dst *Uint32) UnmarshalJSON(b []byte) error { + var n *uint32 + err := json.Unmarshal(b, &n) + if err != nil { + return err + } + + if n == nil { + *dst = Uint32{} + } else { + *dst = Uint32{Uint32: *n, Valid: true} + } + + return nil +} + +type Uint32Codec struct{} + +func (Uint32Codec) FormatSupported(format int16) bool { + return format == TextFormatCode || format == BinaryFormatCode +} + +func (Uint32Codec) PreferredFormat() int16 { + return BinaryFormatCode +} + +func (Uint32Codec) PlanEncode(m *Map, oid uint32, format int16, value any) EncodePlan { + switch format { + case BinaryFormatCode: + switch value.(type) { + case uint32: + return encodePlanUint32CodecBinaryUint32{} + case Uint32Valuer: + return encodePlanUint32CodecBinaryUint32Valuer{} + case Int64Valuer: + return encodePlanUint32CodecBinaryInt64Valuer{} + } + case TextFormatCode: + switch value.(type) { + case uint32: + return encodePlanUint32CodecTextUint32{} + case Int64Valuer: + return encodePlanUint32CodecTextInt64Valuer{} + } + } + + return nil +} + +type encodePlanUint32CodecBinaryUint32 struct{} + +func (encodePlanUint32CodecBinaryUint32) Encode(value any, buf []byte) (newBuf []byte, err error) { + v := value.(uint32) + return pgio.AppendUint32(buf, v), nil +} + +type encodePlanUint32CodecBinaryUint32Valuer struct{} + +func (encodePlanUint32CodecBinaryUint32Valuer) Encode(value any, buf []byte) (newBuf []byte, err error) { + v, err := value.(Uint32Valuer).Uint32Value() + if err != nil { + return nil, err + } + + if !v.Valid { + return nil, nil + } + + return pgio.AppendUint32(buf, v.Uint32), nil +} + +type encodePlanUint32CodecBinaryInt64Valuer struct{} + +func (encodePlanUint32CodecBinaryInt64Valuer) Encode(value any, buf []byte) (newBuf []byte, err error) { + v, err := value.(Int64Valuer).Int64Value() + if err != nil { + return nil, err + } + + if !v.Valid { + return nil, nil + } + + if v.Int64 < 0 { + return nil, fmt.Errorf("%d is less than minimum value for uint32", v.Int64) + } + if v.Int64 > math.MaxUint32 { + return nil, fmt.Errorf("%d is greater than maximum value for uint32", v.Int64) + } + + return pgio.AppendUint32(buf, uint32(v.Int64)), nil +} + +type encodePlanUint32CodecTextUint32 struct{} + +func (encodePlanUint32CodecTextUint32) Encode(value any, buf []byte) (newBuf []byte, err error) { + v := value.(uint32) + return append(buf, strconv.FormatUint(uint64(v), 10)...), nil +} + +type encodePlanUint32CodecTextUint32Valuer struct{} + +func (encodePlanUint32CodecTextUint32Valuer) Encode(value any, buf []byte) (newBuf []byte, err error) { + v, err := value.(Uint32Valuer).Uint32Value() + if err != nil { + return nil, err + } + + if !v.Valid { + return nil, nil + } + + return append(buf, strconv.FormatUint(uint64(v.Uint32), 10)...), nil +} + +type encodePlanUint32CodecTextInt64Valuer struct{} + +func (encodePlanUint32CodecTextInt64Valuer) Encode(value any, buf []byte) (newBuf []byte, err error) { + v, err := value.(Int64Valuer).Int64Value() + if err != nil { + return nil, err + } + + if !v.Valid { + return nil, nil + } + + if v.Int64 < 0 { + return nil, fmt.Errorf("%d is less than minimum value for uint32", v.Int64) + } + if v.Int64 > math.MaxUint32 { + return nil, fmt.Errorf("%d is greater than maximum value for uint32", v.Int64) + } + + return append(buf, strconv.FormatInt(v.Int64, 10)...), nil +} + +func (Uint32Codec) PlanScan(m *Map, oid uint32, format int16, target any) ScanPlan { + switch format { + case BinaryFormatCode: + switch target.(type) { + case *uint32: + return scanPlanBinaryUint32ToUint32{} + case Uint32Scanner: + return scanPlanBinaryUint32ToUint32Scanner{} + case TextScanner: + return scanPlanBinaryUint32ToTextScanner{} + } + case TextFormatCode: + switch target.(type) { + case *uint32: + return scanPlanTextAnyToUint32{} + case Uint32Scanner: + return scanPlanTextAnyToUint32Scanner{} + } + } + + return nil +} + +func (c Uint32Codec) DecodeDatabaseSQLValue(m *Map, oid uint32, format int16, src []byte) (driver.Value, error) { + if src == nil { + return nil, nil + } + + var n uint32 + err := codecScan(c, m, oid, format, src, &n) + if err != nil { + return nil, err + } + return int64(n), nil +} + +func (c Uint32Codec) DecodeValue(m *Map, oid uint32, format int16, src []byte) (any, error) { + if src == nil { + return nil, nil + } + + var n uint32 + err := codecScan(c, m, oid, format, src, &n) + if err != nil { + return nil, err + } + return n, nil +} + +type scanPlanBinaryUint32ToUint32 struct{} + +func (scanPlanBinaryUint32ToUint32) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + if len(src) != 4 { + return fmt.Errorf("invalid length for uint32: %v", len(src)) + } + + p := (dst).(*uint32) + *p = binary.BigEndian.Uint32(src) + + return nil +} + +type scanPlanBinaryUint32ToUint32Scanner struct{} + +func (scanPlanBinaryUint32ToUint32Scanner) Scan(src []byte, dst any) error { + s, ok := (dst).(Uint32Scanner) + if !ok { + return ErrScanTargetTypeChanged + } + + if src == nil { + return s.ScanUint32(Uint32{}) + } + + if len(src) != 4 { + return fmt.Errorf("invalid length for uint32: %v", len(src)) + } + + n := binary.BigEndian.Uint32(src) + + return s.ScanUint32(Uint32{Uint32: n, Valid: true}) +} + +type scanPlanBinaryUint32ToTextScanner struct{} + +func (scanPlanBinaryUint32ToTextScanner) Scan(src []byte, dst any) error { + s, ok := (dst).(TextScanner) + if !ok { + return ErrScanTargetTypeChanged + } + + if src == nil { + return s.ScanText(Text{}) + } + + if len(src) != 4 { + return fmt.Errorf("invalid length for uint32: %v", len(src)) + } + + n := uint64(binary.BigEndian.Uint32(src)) + return s.ScanText(Text{String: strconv.FormatUint(n, 10), Valid: true}) +} + +type scanPlanTextAnyToUint32Scanner struct{} + +func (scanPlanTextAnyToUint32Scanner) Scan(src []byte, dst any) error { + s, ok := (dst).(Uint32Scanner) + if !ok { + return ErrScanTargetTypeChanged + } + + if src == nil { + return s.ScanUint32(Uint32{}) + } + + n, err := strconv.ParseUint(string(src), 10, 32) + if err != nil { + return err + } + + return s.ScanUint32(Uint32{Uint32: uint32(n), Valid: true}) +} diff --git a/vendor/github.com/jackc/pgx/v5/pgtype/uint64.go b/vendor/github.com/jackc/pgx/v5/pgtype/uint64.go new file mode 100644 index 0000000000..fc407bdb62 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgtype/uint64.go @@ -0,0 +1,323 @@ +package pgtype + +import ( + "database/sql/driver" + "encoding/binary" + "fmt" + "math" + "strconv" + + "github.com/jackc/pgx/v5/internal/pgio" +) + +type Uint64Scanner interface { + ScanUint64(v Uint64) error +} + +type Uint64Valuer interface { + Uint64Value() (Uint64, error) +} + +// Uint64 is the core type that is used to represent PostgreSQL types such as XID8. +type Uint64 struct { + Uint64 uint64 + Valid bool +} + +// ScanUint64 implements the [Uint64Scanner] interface. +func (n *Uint64) ScanUint64(v Uint64) error { + *n = v + return nil +} + +// Uint64Value implements the [Uint64Valuer] interface. +func (n Uint64) Uint64Value() (Uint64, error) { + return n, nil +} + +// Scan implements the [database/sql.Scanner] interface. +func (dst *Uint64) Scan(src any) error { + if src == nil { + *dst = Uint64{} + return nil + } + + var n uint64 + + switch src := src.(type) { + case int64: + if src < 0 { + return fmt.Errorf("%d is less than the minimum value for Uint64", src) + } + n = uint64(src) + case string: + un, err := strconv.ParseUint(src, 10, 64) + if err != nil { + return err + } + n = un + default: + return fmt.Errorf("cannot scan %T", src) + } + + *dst = Uint64{Uint64: n, Valid: true} + + return nil +} + +// Value implements the [database/sql/driver.Valuer] interface. +func (src Uint64) Value() (driver.Value, error) { + if !src.Valid { + return nil, nil + } + + // If the value is greater than the maximum value for int64, return it as a string instead of losing data or returning + // an error. + if src.Uint64 > math.MaxInt64 { + return strconv.FormatUint(src.Uint64, 10), nil + } + + return int64(src.Uint64), nil +} + +type Uint64Codec struct{} + +func (Uint64Codec) FormatSupported(format int16) bool { + return format == TextFormatCode || format == BinaryFormatCode +} + +func (Uint64Codec) PreferredFormat() int16 { + return BinaryFormatCode +} + +func (Uint64Codec) PlanEncode(m *Map, oid uint32, format int16, value any) EncodePlan { + switch format { + case BinaryFormatCode: + switch value.(type) { + case uint64: + return encodePlanUint64CodecBinaryUint64{} + case Uint64Valuer: + return encodePlanUint64CodecBinaryUint64Valuer{} + case Int64Valuer: + return encodePlanUint64CodecBinaryInt64Valuer{} + } + case TextFormatCode: + switch value.(type) { + case uint64: + return encodePlanUint64CodecTextUint64{} + case Int64Valuer: + return encodePlanUint64CodecTextInt64Valuer{} + } + } + + return nil +} + +type encodePlanUint64CodecBinaryUint64 struct{} + +func (encodePlanUint64CodecBinaryUint64) Encode(value any, buf []byte) (newBuf []byte, err error) { + v := value.(uint64) + return pgio.AppendUint64(buf, v), nil +} + +type encodePlanUint64CodecBinaryUint64Valuer struct{} + +func (encodePlanUint64CodecBinaryUint64Valuer) Encode(value any, buf []byte) (newBuf []byte, err error) { + v, err := value.(Uint64Valuer).Uint64Value() + if err != nil { + return nil, err + } + + if !v.Valid { + return nil, nil + } + + return pgio.AppendUint64(buf, v.Uint64), nil +} + +type encodePlanUint64CodecBinaryInt64Valuer struct{} + +func (encodePlanUint64CodecBinaryInt64Valuer) Encode(value any, buf []byte) (newBuf []byte, err error) { + v, err := value.(Int64Valuer).Int64Value() + if err != nil { + return nil, err + } + + if !v.Valid { + return nil, nil + } + + if v.Int64 < 0 { + return nil, fmt.Errorf("%d is less than minimum value for uint64", v.Int64) + } + + return pgio.AppendUint64(buf, uint64(v.Int64)), nil +} + +type encodePlanUint64CodecTextUint64 struct{} + +func (encodePlanUint64CodecTextUint64) Encode(value any, buf []byte) (newBuf []byte, err error) { + v := value.(uint64) + return append(buf, strconv.FormatUint(v, 10)...), nil +} + +type encodePlanUint64CodecTextUint64Valuer struct{} + +func (encodePlanUint64CodecTextUint64Valuer) Encode(value any, buf []byte) (newBuf []byte, err error) { + v, err := value.(Uint64Valuer).Uint64Value() + if err != nil { + return nil, err + } + + if !v.Valid { + return nil, nil + } + + return append(buf, strconv.FormatUint(v.Uint64, 10)...), nil +} + +type encodePlanUint64CodecTextInt64Valuer struct{} + +func (encodePlanUint64CodecTextInt64Valuer) Encode(value any, buf []byte) (newBuf []byte, err error) { + v, err := value.(Int64Valuer).Int64Value() + if err != nil { + return nil, err + } + + if !v.Valid { + return nil, nil + } + + if v.Int64 < 0 { + return nil, fmt.Errorf("%d is less than minimum value for uint64", v.Int64) + } + + return append(buf, strconv.FormatInt(v.Int64, 10)...), nil +} + +func (Uint64Codec) PlanScan(m *Map, oid uint32, format int16, target any) ScanPlan { + switch format { + case BinaryFormatCode: + switch target.(type) { + case *uint64: + return scanPlanBinaryUint64ToUint64{} + case Uint64Scanner: + return scanPlanBinaryUint64ToUint64Scanner{} + case TextScanner: + return scanPlanBinaryUint64ToTextScanner{} + } + case TextFormatCode: + switch target.(type) { + case *uint64: + return scanPlanTextAnyToUint64{} + case Uint64Scanner: + return scanPlanTextAnyToUint64Scanner{} + } + } + + return nil +} + +func (c Uint64Codec) DecodeDatabaseSQLValue(m *Map, oid uint32, format int16, src []byte) (driver.Value, error) { + if src == nil { + return nil, nil + } + + var n uint64 + err := codecScan(c, m, oid, format, src, &n) + if err != nil { + return nil, err + } + return int64(n), nil +} + +func (c Uint64Codec) DecodeValue(m *Map, oid uint32, format int16, src []byte) (any, error) { + if src == nil { + return nil, nil + } + + var n uint64 + err := codecScan(c, m, oid, format, src, &n) + if err != nil { + return nil, err + } + return n, nil +} + +type scanPlanBinaryUint64ToUint64 struct{} + +func (scanPlanBinaryUint64ToUint64) Scan(src []byte, dst any) error { + if src == nil { + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + if len(src) != 8 { + return fmt.Errorf("invalid length for uint64: %v", len(src)) + } + + p := (dst).(*uint64) + *p = binary.BigEndian.Uint64(src) + + return nil +} + +type scanPlanBinaryUint64ToUint64Scanner struct{} + +func (scanPlanBinaryUint64ToUint64Scanner) Scan(src []byte, dst any) error { + s, ok := (dst).(Uint64Scanner) + if !ok { + return ErrScanTargetTypeChanged + } + + if src == nil { + return s.ScanUint64(Uint64{}) + } + + if len(src) != 8 { + return fmt.Errorf("invalid length for uint64: %v", len(src)) + } + + n := binary.BigEndian.Uint64(src) + + return s.ScanUint64(Uint64{Uint64: n, Valid: true}) +} + +type scanPlanBinaryUint64ToTextScanner struct{} + +func (scanPlanBinaryUint64ToTextScanner) Scan(src []byte, dst any) error { + s, ok := (dst).(TextScanner) + if !ok { + return ErrScanTargetTypeChanged + } + + if src == nil { + return s.ScanText(Text{}) + } + + if len(src) != 8 { + return fmt.Errorf("invalid length for uint64: %v", len(src)) + } + + n := binary.BigEndian.Uint64(src) + return s.ScanText(Text{String: strconv.FormatUint(n, 10), Valid: true}) +} + +type scanPlanTextAnyToUint64Scanner struct{} + +func (scanPlanTextAnyToUint64Scanner) Scan(src []byte, dst any) error { + s, ok := (dst).(Uint64Scanner) + if !ok { + return ErrScanTargetTypeChanged + } + + if src == nil { + return s.ScanUint64(Uint64{}) + } + + n, err := strconv.ParseUint(string(src), 10, 64) + if err != nil { + return err + } + + return s.ScanUint64(Uint64{Uint64: n, Valid: true}) +} diff --git a/vendor/github.com/jackc/pgx/v5/pgtype/uuid.go b/vendor/github.com/jackc/pgx/v5/pgtype/uuid.go new file mode 100644 index 0000000000..476889a8db --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgtype/uuid.go @@ -0,0 +1,291 @@ +package pgtype + +import ( + "bytes" + "database/sql/driver" + "encoding/hex" + "fmt" +) + +type UUIDScanner interface { + ScanUUID(v UUID) error +} + +type UUIDValuer interface { + UUIDValue() (UUID, error) +} + +type UUID struct { + Bytes [16]byte + Valid bool +} + +// ScanUUID implements the [UUIDScanner] interface. +func (b *UUID) ScanUUID(v UUID) error { + *b = v + return nil +} + +// UUIDValue implements the [UUIDValuer] interface. +func (b UUID) UUIDValue() (UUID, error) { + return b, nil +} + +// parseUUID converts a string UUID in standard form to a byte array. +func parseUUID(src string) (dst [16]byte, err error) { + switch len(src) { + case 36: + src = src[0:8] + src[9:13] + src[14:18] + src[19:23] + src[24:] + case 32: + // dashes already stripped, assume valid + default: + // assume invalid. + return dst, fmt.Errorf("cannot parse UUID %v", src) + } + + buf, err := hex.DecodeString(src) + if err != nil { + return dst, err + } + + copy(dst[:], buf) + return dst, err +} + +// encodeUUID converts a uuid byte array to UUID standard string form. +func encodeUUID(src [16]byte) string { + var buf [36]byte + + hex.Encode(buf[0:8], src[:4]) + buf[8] = '-' + hex.Encode(buf[9:13], src[4:6]) + buf[13] = '-' + hex.Encode(buf[14:18], src[6:8]) + buf[18] = '-' + hex.Encode(buf[19:23], src[8:10]) + buf[23] = '-' + hex.Encode(buf[24:], src[10:]) + + return string(buf[:]) +} + +// Scan implements the [database/sql.Scanner] interface. +func (dst *UUID) Scan(src any) error { + if src == nil { + *dst = UUID{} + return nil + } + + if src, ok := src.(string); ok { + buf, err := parseUUID(src) + if err != nil { + return err + } + *dst = UUID{Bytes: buf, Valid: true} + return nil + } + + return fmt.Errorf("cannot scan %T", src) +} + +// Value implements the [database/sql/driver.Valuer] interface. +func (src UUID) Value() (driver.Value, error) { + if !src.Valid { + return nil, nil + } + + return encodeUUID(src.Bytes), nil +} + +func (src UUID) String() string { + if !src.Valid { + return "" + } + + return encodeUUID(src.Bytes) +} + +// MarshalJSON implements the [encoding/json.Marshaler] interface. +func (src UUID) MarshalJSON() ([]byte, error) { + if !src.Valid { + return []byte("null"), nil + } + + var buff bytes.Buffer + buff.WriteByte('"') + buff.WriteString(encodeUUID(src.Bytes)) + buff.WriteByte('"') + return buff.Bytes(), nil +} + +// UnmarshalJSON implements the [encoding/json.Unmarshaler] interface. +func (dst *UUID) UnmarshalJSON(src []byte) error { + if bytes.Equal(src, []byte("null")) { + *dst = UUID{} + return nil + } + if len(src) != 38 { + return fmt.Errorf("invalid length for UUID: %v", len(src)) + } + buf, err := parseUUID(string(src[1 : len(src)-1])) + if err != nil { + return err + } + *dst = UUID{Bytes: buf, Valid: true} + return nil +} + +type UUIDCodec struct{} + +func (UUIDCodec) FormatSupported(format int16) bool { + return format == TextFormatCode || format == BinaryFormatCode +} + +func (UUIDCodec) PreferredFormat() int16 { + return BinaryFormatCode +} + +func (UUIDCodec) PlanEncode(m *Map, oid uint32, format int16, value any) EncodePlan { + if _, ok := value.(UUIDValuer); !ok { + return nil + } + + switch format { + case BinaryFormatCode: + return encodePlanUUIDCodecBinaryUUIDValuer{} + case TextFormatCode: + return encodePlanUUIDCodecTextUUIDValuer{} + } + + return nil +} + +type encodePlanUUIDCodecBinaryUUIDValuer struct{} + +func (encodePlanUUIDCodecBinaryUUIDValuer) Encode(value any, buf []byte) (newBuf []byte, err error) { + uuid, err := value.(UUIDValuer).UUIDValue() + if err != nil { + return nil, err + } + + if !uuid.Valid { + return nil, nil + } + + return append(buf, uuid.Bytes[:]...), nil +} + +type encodePlanUUIDCodecTextUUIDValuer struct{} + +func (encodePlanUUIDCodecTextUUIDValuer) Encode(value any, buf []byte) (newBuf []byte, err error) { + uuid, err := value.(UUIDValuer).UUIDValue() + if err != nil { + return nil, err + } + + if !uuid.Valid { + return nil, nil + } + + return append(buf, encodeUUID(uuid.Bytes)...), nil +} + +func (UUIDCodec) PlanScan(m *Map, oid uint32, format int16, target any) ScanPlan { + switch format { + case BinaryFormatCode: + switch target.(type) { + case UUIDScanner: + return scanPlanBinaryUUIDToUUIDScanner{} + case TextScanner: + return scanPlanBinaryUUIDToTextScanner{} + } + case TextFormatCode: + if _, ok := target.(UUIDScanner); ok { + return scanPlanTextAnyToUUIDScanner{} + } + } + + return nil +} + +type scanPlanBinaryUUIDToUUIDScanner struct{} + +func (scanPlanBinaryUUIDToUUIDScanner) Scan(src []byte, dst any) error { + scanner := (dst).(UUIDScanner) + + if src == nil { + return scanner.ScanUUID(UUID{}) + } + + if len(src) != 16 { + return fmt.Errorf("invalid length for UUID: %v", len(src)) + } + + uuid := UUID{Valid: true} + copy(uuid.Bytes[:], src) + + return scanner.ScanUUID(uuid) +} + +type scanPlanBinaryUUIDToTextScanner struct{} + +func (scanPlanBinaryUUIDToTextScanner) Scan(src []byte, dst any) error { + scanner := (dst).(TextScanner) + + if src == nil { + return scanner.ScanText(Text{}) + } + + if len(src) != 16 { + return fmt.Errorf("invalid length for UUID: %v", len(src)) + } + + var buf [16]byte + copy(buf[:], src) + + return scanner.ScanText(Text{String: encodeUUID(buf), Valid: true}) +} + +type scanPlanTextAnyToUUIDScanner struct{} + +func (scanPlanTextAnyToUUIDScanner) Scan(src []byte, dst any) error { + scanner := (dst).(UUIDScanner) + + if src == nil { + return scanner.ScanUUID(UUID{}) + } + + buf, err := parseUUID(string(src)) + if err != nil { + return err + } + + return scanner.ScanUUID(UUID{Bytes: buf, Valid: true}) +} + +func (c UUIDCodec) DecodeDatabaseSQLValue(m *Map, oid uint32, format int16, src []byte) (driver.Value, error) { + if src == nil { + return nil, nil + } + + var uuid UUID + err := codecScan(c, m, oid, format, src, &uuid) + if err != nil { + return nil, err + } + + return encodeUUID(uuid.Bytes), nil +} + +func (c UUIDCodec) DecodeValue(m *Map, oid uint32, format int16, src []byte) (any, error) { + if src == nil { + return nil, nil + } + + var uuid UUID + err := codecScan(c, m, oid, format, src, &uuid) + if err != nil { + return nil, err + } + return uuid.Bytes, nil +} diff --git a/vendor/github.com/jackc/pgx/v5/pgtype/xml.go b/vendor/github.com/jackc/pgx/v5/pgtype/xml.go new file mode 100644 index 0000000000..66e6dffdae --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgtype/xml.go @@ -0,0 +1,198 @@ +package pgtype + +import ( + "database/sql" + "database/sql/driver" + "encoding/xml" + "fmt" + "reflect" +) + +type XMLCodec struct { + Marshal func(v any) ([]byte, error) + Unmarshal func(data []byte, v any) error +} + +func (*XMLCodec) FormatSupported(format int16) bool { + return format == TextFormatCode || format == BinaryFormatCode +} + +func (*XMLCodec) PreferredFormat() int16 { + return TextFormatCode +} + +func (c *XMLCodec) PlanEncode(m *Map, oid uint32, format int16, value any) EncodePlan { + switch value.(type) { + case string: + return encodePlanXMLCodecEitherFormatString{} + case []byte: + return encodePlanXMLCodecEitherFormatByteSlice{} + + // Cannot rely on driver.Valuer being handled later because anything can be marshalled. + // + // https://github.com/jackc/pgx/issues/1430 + // + // Check for driver.Valuer must come before xml.Marshaler so that it is guaranteed to be used + // when both are implemented https://github.com/jackc/pgx/issues/1805 + case driver.Valuer: + return &encodePlanDriverValuer{m: m, oid: oid, formatCode: format} + + // Must come before trying wrap encode plans because a pointer to a struct may be unwrapped to a struct that can be + // marshalled. + // + // https://github.com/jackc/pgx/issues/1681 + case xml.Marshaler: + return &encodePlanXMLCodecEitherFormatMarshal{ + marshal: c.Marshal, + } + } + + // Because anything can be marshalled the normal wrapping in Map.PlanScan doesn't get a chance to run. So try the + // appropriate wrappers here. + for _, f := range []TryWrapEncodePlanFunc{ + TryWrapDerefPointerEncodePlan, + TryWrapFindUnderlyingTypeEncodePlan, + } { + if wrapperPlan, nextValue, ok := f(value); ok { + if nextPlan := c.PlanEncode(m, oid, format, nextValue); nextPlan != nil { + wrapperPlan.SetNext(nextPlan) + return wrapperPlan + } + } + } + + return &encodePlanXMLCodecEitherFormatMarshal{ + marshal: c.Marshal, + } +} + +type encodePlanXMLCodecEitherFormatString struct{} + +func (encodePlanXMLCodecEitherFormatString) Encode(value any, buf []byte) (newBuf []byte, err error) { + xmlString := value.(string) + buf = append(buf, xmlString...) + return buf, nil +} + +type encodePlanXMLCodecEitherFormatByteSlice struct{} + +func (encodePlanXMLCodecEitherFormatByteSlice) Encode(value any, buf []byte) (newBuf []byte, err error) { + xmlBytes := value.([]byte) + if xmlBytes == nil { + return nil, nil + } + + buf = append(buf, xmlBytes...) + return buf, nil +} + +type encodePlanXMLCodecEitherFormatMarshal struct { + marshal func(v any) ([]byte, error) +} + +func (e *encodePlanXMLCodecEitherFormatMarshal) Encode(value any, buf []byte) (newBuf []byte, err error) { + xmlBytes, err := e.marshal(value) + if err != nil { + return nil, err + } + + buf = append(buf, xmlBytes...) + return buf, nil +} + +func (c *XMLCodec) PlanScan(m *Map, oid uint32, format int16, target any) ScanPlan { + switch target.(type) { + case *string: + return scanPlanAnyToString{} + + case **string: + // This is to fix **string scanning. It seems wrong to special case **string, but it's not clear what a better + // solution would be. + // + // https://github.com/jackc/pgx/issues/1470 -- **string + // https://github.com/jackc/pgx/issues/1691 -- ** anything else + + if wrapperPlan, nextDst, ok := TryPointerPointerScanPlan(target); ok { + if nextPlan := m.planScan(oid, format, nextDst, 0); nextPlan != nil { + if _, failed := nextPlan.(*scanPlanFail); !failed { + wrapperPlan.SetNext(nextPlan) + return wrapperPlan + } + } + } + + case *[]byte: + return scanPlanXMLToByteSlice{} + case BytesScanner: + return scanPlanBinaryBytesToBytesScanner{} + + // Cannot rely on sql.Scanner being handled later because scanPlanXMLToXMLUnmarshal will take precedence. + // + // https://github.com/jackc/pgx/issues/1418 + case sql.Scanner: + return &scanPlanSQLScanner{formatCode: format} + } + + return &scanPlanXMLToXMLUnmarshal{ + unmarshal: c.Unmarshal, + } +} + +type scanPlanXMLToByteSlice struct{} + +func (scanPlanXMLToByteSlice) Scan(src []byte, dst any) error { + dstBuf := dst.(*[]byte) + if src == nil { + *dstBuf = nil + return nil + } + + *dstBuf = make([]byte, len(src)) + copy(*dstBuf, src) + return nil +} + +type scanPlanXMLToXMLUnmarshal struct { + unmarshal func(data []byte, v any) error +} + +func (s *scanPlanXMLToXMLUnmarshal) Scan(src []byte, dst any) error { + if src == nil { + dstValue := reflect.ValueOf(dst) + if dstValue.Kind() == reflect.Pointer { + el := dstValue.Elem() + switch el.Kind() { + case reflect.Pointer, reflect.Slice, reflect.Map, reflect.Interface, reflect.Struct: + el.Set(reflect.Zero(el.Type())) + return nil + } + } + + return fmt.Errorf("cannot scan NULL into %T", dst) + } + + elem := reflect.ValueOf(dst).Elem() + elem.Set(reflect.Zero(elem.Type())) + + return s.unmarshal(src, dst) +} + +func (c *XMLCodec) DecodeDatabaseSQLValue(m *Map, oid uint32, format int16, src []byte) (driver.Value, error) { + if src == nil { + return nil, nil + } + + dstBuf := make([]byte, len(src)) + copy(dstBuf, src) + return dstBuf, nil +} + +func (c *XMLCodec) DecodeValue(m *Map, oid uint32, format int16, src []byte) (any, error) { + if src == nil { + return nil, nil + } + + var dst any + err := c.Unmarshal(src, &dst) + return dst, err +} diff --git a/vendor/github.com/jackc/pgx/v5/pgxpool/batch_results.go b/vendor/github.com/jackc/pgx/v5/pgxpool/batch_results.go new file mode 100644 index 0000000000..5d5c681d5a --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgxpool/batch_results.go @@ -0,0 +1,52 @@ +package pgxpool + +import ( + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgconn" +) + +type errBatchResults struct { + err error +} + +func (br errBatchResults) Exec() (pgconn.CommandTag, error) { + return pgconn.CommandTag{}, br.err +} + +func (br errBatchResults) Query() (pgx.Rows, error) { + return errRows{err: br.err}, br.err +} + +func (br errBatchResults) QueryRow() pgx.Row { + return errRow{err: br.err} +} + +func (br errBatchResults) Close() error { + return br.err +} + +type poolBatchResults struct { + br pgx.BatchResults + c *Conn +} + +func (br *poolBatchResults) Exec() (pgconn.CommandTag, error) { + return br.br.Exec() +} + +func (br *poolBatchResults) Query() (pgx.Rows, error) { + return br.br.Query() +} + +func (br *poolBatchResults) QueryRow() pgx.Row { + return br.br.QueryRow() +} + +func (br *poolBatchResults) Close() error { + err := br.br.Close() + if br.c != nil { + br.c.Release() + br.c = nil + } + return err +} diff --git a/vendor/github.com/jackc/pgx/v5/pgxpool/conn.go b/vendor/github.com/jackc/pgx/v5/pgxpool/conn.go new file mode 100644 index 0000000000..b4f90605be --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgxpool/conn.go @@ -0,0 +1,133 @@ +package pgxpool + +import ( + "context" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgconn" + "github.com/jackc/puddle/v2" +) + +// Conn is an acquired *pgx.Conn from a Pool. +type Conn struct { + res *puddle.Resource[*connResource] + p *Pool +} + +// Release returns c to the pool it was acquired from. Once Release has been called, other methods must not be called. +// However, it is safe to call Release multiple times. Subsequent calls after the first will be ignored. +func (c *Conn) Release() { + if c.res == nil { + return + } + + conn := c.Conn() + res := c.res + c.res = nil + + if c.p.releaseTracer != nil { + c.p.releaseTracer.TraceRelease(c.p, TraceReleaseData{Conn: conn}) + } + + if conn.IsClosed() || conn.PgConn().IsBusy() || conn.PgConn().TxStatus() != 'I' { + res.Destroy() + // Signal to the health check to run since we just destroyed a connections + // and we might be below minConns now + c.p.triggerHealthCheck() + return + } + + // If the pool is consistently being used, we might never get to check the + // lifetime of a connection since we only check idle connections in checkConnsHealth + // so we also check the lifetime here and force a health check + if c.p.isExpired(res) { + c.p.lifetimeDestroyCount.Add(1) + res.Destroy() + // Signal to the health check to run since we just destroyed a connections + // and we might be below minConns now + c.p.triggerHealthCheck() + return + } + + if c.p.afterRelease == nil { + res.Release() + return + } + + go func() { + if c.p.afterRelease(conn) { + res.Release() + } else { + res.Destroy() + // Signal to the health check to run since we just destroyed a connections + // and we might be below minConns now + c.p.triggerHealthCheck() + } + }() +} + +// Hijack assumes ownership of the connection from the pool. Caller is responsible for closing the connection. Hijack +// will panic if called on an already released or hijacked connection. +func (c *Conn) Hijack() *pgx.Conn { + if c.res == nil { + panic("cannot hijack already released or hijacked connection") + } + + conn := c.Conn() + res := c.res + c.res = nil + + res.Hijack() + + return conn +} + +func (c *Conn) Exec(ctx context.Context, sql string, arguments ...any) (pgconn.CommandTag, error) { + return c.Conn().Exec(ctx, sql, arguments...) +} + +func (c *Conn) Query(ctx context.Context, sql string, args ...any) (pgx.Rows, error) { + return c.Conn().Query(ctx, sql, args...) +} + +func (c *Conn) QueryRow(ctx context.Context, sql string, args ...any) pgx.Row { + return c.Conn().QueryRow(ctx, sql, args...) +} + +func (c *Conn) SendBatch(ctx context.Context, b *pgx.Batch) pgx.BatchResults { + return c.Conn().SendBatch(ctx, b) +} + +func (c *Conn) CopyFrom(ctx context.Context, tableName pgx.Identifier, columnNames []string, rowSrc pgx.CopyFromSource) (int64, error) { + return c.Conn().CopyFrom(ctx, tableName, columnNames, rowSrc) +} + +// Begin starts a transaction block from the *Conn without explicitly setting a transaction mode (see BeginTx with TxOptions if transaction mode is required). +func (c *Conn) Begin(ctx context.Context) (pgx.Tx, error) { + return c.Conn().Begin(ctx) +} + +// BeginTx starts a transaction block from the *Conn with txOptions determining the transaction mode. +func (c *Conn) BeginTx(ctx context.Context, txOptions pgx.TxOptions) (pgx.Tx, error) { + return c.Conn().BeginTx(ctx, txOptions) +} + +func (c *Conn) Ping(ctx context.Context) error { + return c.Conn().Ping(ctx) +} + +func (c *Conn) Conn() *pgx.Conn { + return c.connResource().conn +} + +func (c *Conn) connResource() *connResource { + return c.res.Value() +} + +func (c *Conn) getPoolRow(r pgx.Row) *poolRow { + return c.connResource().getPoolRow(c, r) +} + +func (c *Conn) getPoolRows(r pgx.Rows) *poolRows { + return c.connResource().getPoolRows(c, r) +} diff --git a/vendor/github.com/jackc/pgx/v5/pgxpool/doc.go b/vendor/github.com/jackc/pgx/v5/pgxpool/doc.go new file mode 100644 index 0000000000..099443bca8 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgxpool/doc.go @@ -0,0 +1,27 @@ +// Package pgxpool is a concurrency-safe connection pool for pgx. +/* +pgxpool implements a nearly identical interface to pgx connections. + +Creating a Pool + +The primary way of creating a pool is with [pgxpool.New]: + + pool, err := pgxpool.New(context.Background(), os.Getenv("DATABASE_URL")) + +The database connection string can be in URL or keyword/value format. PostgreSQL settings, pgx settings, and pool settings can be +specified here. In addition, a config struct can be created by [ParseConfig]. + + config, err := pgxpool.ParseConfig(os.Getenv("DATABASE_URL")) + if err != nil { + // ... + } + config.AfterConnect = func(ctx context.Context, conn *pgx.Conn) error { + // do something with every new connection + } + + pool, err := pgxpool.NewWithConfig(context.Background(), config) + +A pool returns without waiting for any connections to be established. Acquire a connection immediately after creating +the pool to check if a connection can successfully be established. +*/ +package pgxpool diff --git a/vendor/github.com/jackc/pgx/v5/pgxpool/pool.go b/vendor/github.com/jackc/pgx/v5/pgxpool/pool.go new file mode 100644 index 0000000000..5c740b1f2f --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgxpool/pool.go @@ -0,0 +1,840 @@ +package pgxpool + +import ( + "context" + "errors" + "math/rand/v2" + "runtime" + "strconv" + "sync" + "sync/atomic" + "time" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgconn" + "github.com/jackc/puddle/v2" +) + +var ( + defaultMaxConns = int32(4) + defaultMinConns = int32(0) + defaultMinIdleConns = int32(0) + defaultMaxConnLifetime = time.Hour + defaultMaxConnIdleTime = time.Minute * 30 + defaultHealthCheckPeriod = time.Minute +) + +type connResource struct { + conn *pgx.Conn + conns []Conn + poolRows []poolRow + poolRowss []poolRows + maxAgeTime time.Time +} + +func (cr *connResource) getConn(p *Pool, res *puddle.Resource[*connResource]) *Conn { + if len(cr.conns) == 0 { + cr.conns = make([]Conn, 128) + } + + c := &cr.conns[len(cr.conns)-1] + cr.conns = cr.conns[0 : len(cr.conns)-1] + + c.res = res + c.p = p + + return c +} + +func (cr *connResource) getPoolRow(c *Conn, r pgx.Row) *poolRow { + if len(cr.poolRows) == 0 { + cr.poolRows = make([]poolRow, 128) + } + + pr := &cr.poolRows[len(cr.poolRows)-1] + cr.poolRows = cr.poolRows[0 : len(cr.poolRows)-1] + + pr.c = c + pr.r = r + + return pr +} + +func (cr *connResource) getPoolRows(c *Conn, r pgx.Rows) *poolRows { + if len(cr.poolRowss) == 0 { + cr.poolRowss = make([]poolRows, 128) + } + + pr := &cr.poolRowss[len(cr.poolRowss)-1] + cr.poolRowss = cr.poolRowss[0 : len(cr.poolRowss)-1] + + pr.c = c + pr.r = r + + return pr +} + +// Pool allows for connection reuse. +type Pool struct { + newConnsCount atomic.Int64 + lifetimeDestroyCount atomic.Int64 + idleDestroyCount atomic.Int64 + + p *puddle.Pool[*connResource] + config *Config + beforeConnect func(context.Context, *pgx.ConnConfig) error + afterConnect func(context.Context, *pgx.Conn) error + prepareConn func(context.Context, *pgx.Conn) (bool, error) + afterRelease func(*pgx.Conn) bool + beforeClose func(*pgx.Conn) + shouldPing func(context.Context, ShouldPingParams) bool + minConns int32 + minIdleConns int32 + maxConns int32 + maxConnLifetime time.Duration + maxConnLifetimeJitter time.Duration + maxConnIdleTime time.Duration + healthCheckPeriod time.Duration + pingTimeout time.Duration + + healthCheckMu sync.Mutex + healthCheckTimer *time.Timer + + healthCheckChan chan struct{} + + acquireTracer AcquireTracer + releaseTracer ReleaseTracer + + closeOnce sync.Once + closeChan chan struct{} +} + +// ShouldPingParams are the parameters passed to ShouldPing. +type ShouldPingParams struct { + Conn *pgx.Conn + IdleDuration time.Duration +} + +// Config is the configuration struct for creating a pool. It must be created by [ParseConfig] and then it can be +// modified. +type Config struct { + ConnConfig *pgx.ConnConfig + + // BeforeConnect is called before a new connection is made. It is passed a copy of the underlying [pgx.ConnConfig] and + // will not impact any existing open connections. + BeforeConnect func(context.Context, *pgx.ConnConfig) error + + // AfterConnect is called after a connection is established, but before it is added to the pool. + AfterConnect func(context.Context, *pgx.Conn) error + + // BeforeAcquire is called before a connection is acquired from the pool. It must return true to allow the + // acquisition or false to indicate that the connection should be destroyed and a different connection should be + // acquired. + // + // Deprecated: Use PrepareConn instead. If both PrepareConn and BeforeAcquire are set, PrepareConn will take + // precedence, ignoring BeforeAcquire. + BeforeAcquire func(context.Context, *pgx.Conn) bool + + // PrepareConn is called before a connection is acquired from the pool. If this function returns true, the connection + // is considered valid, otherwise the connection is destroyed. If the function returns a non-nil error, the instigating + // query will fail with the returned error. + // + // Specifically, this means that: + // + // - If it returns true and a nil error, the query proceeds as normal. + // - If it returns true and an error, the connection will be returned to the pool, and the instigating query will fail with the returned error. + // - If it returns false, and an error, the connection will be destroyed, and the query will fail with the returned error. + // - If it returns false and a nil error, the connection will be destroyed, and the instigating query will be retried on a new connection. + PrepareConn func(context.Context, *pgx.Conn) (bool, error) + + // AfterRelease is called after a connection is released, but before it is returned to the pool. It must return true to + // return the connection to the pool or false to destroy the connection. + AfterRelease func(*pgx.Conn) bool + + // BeforeClose is called right before a connection is closed and removed from the pool. + BeforeClose func(*pgx.Conn) + + // ShouldPing is called after a connection is acquired from the pool. If it returns true, the connection is pinged to check for liveness. + // If this func is not set, the default behavior is to ping connections that have been idle for at least 1 second. + ShouldPing func(context.Context, ShouldPingParams) bool + + // MaxConnLifetime is the duration since creation after which a connection will be automatically closed. + MaxConnLifetime time.Duration + + // MaxConnLifetimeJitter is the duration after MaxConnLifetime to randomly decide to close a connection. + // This helps prevent all connections from being closed at the exact same time, starving the pool. + MaxConnLifetimeJitter time.Duration + + // MaxConnIdleTime is the duration after which an idle connection will be automatically closed by the health check. + MaxConnIdleTime time.Duration + + // PingTimeout is the maximum amount of time to wait for a connection to pong before considering it as unhealthy and + // destroying it. If zero, the default is no timeout. + PingTimeout time.Duration + + // MaxConns is the maximum size of the pool. The default is the greater of 4 or runtime.NumCPU(). + MaxConns int32 + + // MinConns is the minimum size of the pool. After connection closes, the pool might dip below MinConns. A low + // number of MinConns might mean the pool is empty after MaxConnLifetime until the health check has a chance + // to create new connections. + MinConns int32 + + // MinIdleConns is the minimum number of idle connections in the pool. You can increase this to ensure that + // there are always idle connections available. This can help reduce tail latencies during request processing, + // as you can avoid the latency of establishing a new connection while handling requests. It is superior + // to MinConns for this purpose. + // Similar to MinConns, the pool might temporarily dip below MinIdleConns after connection closes. + MinIdleConns int32 + + // HealthCheckPeriod is the duration between checks of the health of idle connections. + HealthCheckPeriod time.Duration + + createdByParseConfig bool // Used to enforce created by ParseConfig rule. +} + +// Copy returns a deep copy of the config that is safe to use and modify. +// The only exception is the tls.Config: +// according to the tls.Config docs it must not be modified after creation. +func (c *Config) Copy() *Config { + newConfig := new(Config) + *newConfig = *c + newConfig.ConnConfig = c.ConnConfig.Copy() + return newConfig +} + +// ConnString returns the connection string as parsed by pgxpool.ParseConfig into pgxpool.Config. +func (c *Config) ConnString() string { return c.ConnConfig.ConnString() } + +// New creates a new Pool. See [ParseConfig] for information on connString format. +func New(ctx context.Context, connString string) (*Pool, error) { + config, err := ParseConfig(connString) + if err != nil { + return nil, err + } + + return NewWithConfig(ctx, config) +} + +// NewWithConfig creates a new [Pool]. config must have been created by [ParseConfig]. +func NewWithConfig(ctx context.Context, config *Config) (*Pool, error) { + // Default values are set in ParseConfig. Enforce initial creation by ParseConfig rather than setting defaults from + // zero values. + if !config.createdByParseConfig { + panic("config must be created by ParseConfig") + } + + prepareConn := config.PrepareConn + if prepareConn == nil && config.BeforeAcquire != nil { + prepareConn = func(ctx context.Context, conn *pgx.Conn) (bool, error) { + return config.BeforeAcquire(ctx, conn), nil + } + } + + p := &Pool{ + config: config, + beforeConnect: config.BeforeConnect, + afterConnect: config.AfterConnect, + prepareConn: prepareConn, + afterRelease: config.AfterRelease, + beforeClose: config.BeforeClose, + minConns: config.MinConns, + minIdleConns: config.MinIdleConns, + maxConns: config.MaxConns, + maxConnLifetime: config.MaxConnLifetime, + maxConnLifetimeJitter: config.MaxConnLifetimeJitter, + maxConnIdleTime: config.MaxConnIdleTime, + pingTimeout: config.PingTimeout, + healthCheckPeriod: config.HealthCheckPeriod, + healthCheckChan: make(chan struct{}, 1), + closeChan: make(chan struct{}), + } + + if t, ok := config.ConnConfig.Tracer.(AcquireTracer); ok { + p.acquireTracer = t + } + + if t, ok := config.ConnConfig.Tracer.(ReleaseTracer); ok { + p.releaseTracer = t + } + + if config.ShouldPing != nil { + p.shouldPing = config.ShouldPing + } else { + p.shouldPing = func(ctx context.Context, params ShouldPingParams) bool { + return params.IdleDuration > time.Second + } + } + + var err error + p.p, err = puddle.NewPool( + &puddle.Config[*connResource]{ + Constructor: func(ctx context.Context) (*connResource, error) { + p.newConnsCount.Add(1) + connConfig := p.config.ConnConfig.Copy() + + // Connection will continue in background even if Acquire is canceled. Ensure that a connect won't hang forever. + if connConfig.ConnectTimeout <= 0 { + connConfig.ConnectTimeout = 2 * time.Minute + } + + if p.beforeConnect != nil { + if err := p.beforeConnect(ctx, connConfig); err != nil { + return nil, err + } + } + + conn, err := pgx.ConnectConfig(ctx, connConfig) + if err != nil { + return nil, err + } + + if p.afterConnect != nil { + err = p.afterConnect(ctx, conn) + if err != nil { + conn.Close(ctx) + return nil, err + } + } + + jitterSecs := rand.Float64() * config.MaxConnLifetimeJitter.Seconds() + maxAgeTime := time.Now().Add(config.MaxConnLifetime).Add(time.Duration(jitterSecs) * time.Second) + + cr := &connResource{ + conn: conn, + conns: make([]Conn, 64), + poolRows: make([]poolRow, 64), + poolRowss: make([]poolRows, 64), + maxAgeTime: maxAgeTime, + } + + return cr, nil + }, + Destructor: func(value *connResource) { + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + conn := value.conn + if p.beforeClose != nil { + p.beforeClose(conn) + } + conn.Close(ctx) + select { + case <-conn.PgConn().CleanupDone(): + case <-ctx.Done(): + } + cancel() + }, + MaxSize: config.MaxConns, + }, + ) + if err != nil { + return nil, err + } + + go func() { + targetIdleResources := max(int(p.minConns), int(p.minIdleConns)) + p.createIdleResources(ctx, targetIdleResources) + p.backgroundHealthCheck() + }() + + return p, nil +} + +// ParseConfig builds a Config from connString. It parses connString with the same behavior as [pgx.ParseConfig] with the +// addition of the following variables: +// +// - pool_max_conns: integer greater than 0 (default 4) +// - pool_min_conns: integer 0 or greater (default 0) +// - pool_max_conn_lifetime: duration string (default 1 hour) +// - pool_max_conn_idle_time: duration string (default 30 minutes) +// - pool_health_check_period: duration string (default 1 minute) +// - pool_max_conn_lifetime_jitter: duration string (default 0) +// +// See Config for definitions of these arguments. +// +// # Example Keyword/Value +// user=jack password=secret host=pg.example.com port=5432 dbname=mydb sslmode=verify-ca pool_max_conns=10 pool_max_conn_lifetime=1h30m +// +// # Example URL +// postgres://jack:secret@pg.example.com:5432/mydb?sslmode=verify-ca&pool_max_conns=10&pool_max_conn_lifetime=1h30m +func ParseConfig(connString string) (*Config, error) { + connConfig, err := pgx.ParseConfig(connString) + if err != nil { + return nil, err + } + + config := &Config{ + ConnConfig: connConfig, + createdByParseConfig: true, + } + + if s, ok := config.ConnConfig.Config.RuntimeParams["pool_max_conns"]; ok { + delete(connConfig.Config.RuntimeParams, "pool_max_conns") + n, err := strconv.ParseInt(s, 10, 32) + if err != nil { + return nil, pgconn.NewParseConfigError(connString, "cannot parse pool_max_conns", err) + } + if n < 1 { + return nil, pgconn.NewParseConfigError(connString, "pool_max_conns too small", err) + } + config.MaxConns = int32(n) + } else { + config.MaxConns = defaultMaxConns + if numCPU := int32(runtime.NumCPU()); numCPU > config.MaxConns { + config.MaxConns = numCPU + } + } + + if s, ok := config.ConnConfig.Config.RuntimeParams["pool_min_conns"]; ok { + delete(connConfig.Config.RuntimeParams, "pool_min_conns") + n, err := strconv.ParseInt(s, 10, 32) + if err != nil { + return nil, pgconn.NewParseConfigError(connString, "cannot parse pool_min_conns", err) + } + config.MinConns = int32(n) + } else { + config.MinConns = defaultMinConns + } + + if s, ok := config.ConnConfig.Config.RuntimeParams["pool_min_idle_conns"]; ok { + delete(connConfig.Config.RuntimeParams, "pool_min_idle_conns") + n, err := strconv.ParseInt(s, 10, 32) + if err != nil { + return nil, pgconn.NewParseConfigError(connString, "cannot parse pool_min_idle_conns", err) + } + config.MinIdleConns = int32(n) + } else { + config.MinIdleConns = defaultMinIdleConns + } + + if s, ok := config.ConnConfig.Config.RuntimeParams["pool_max_conn_lifetime"]; ok { + delete(connConfig.Config.RuntimeParams, "pool_max_conn_lifetime") + d, err := time.ParseDuration(s) + if err != nil { + return nil, pgconn.NewParseConfigError(connString, "cannot parse pool_max_conn_lifetime", err) + } + config.MaxConnLifetime = d + } else { + config.MaxConnLifetime = defaultMaxConnLifetime + } + + if s, ok := config.ConnConfig.Config.RuntimeParams["pool_max_conn_idle_time"]; ok { + delete(connConfig.Config.RuntimeParams, "pool_max_conn_idle_time") + d, err := time.ParseDuration(s) + if err != nil { + return nil, pgconn.NewParseConfigError(connString, "cannot parse pool_max_conn_idle_time", err) + } + config.MaxConnIdleTime = d + } else { + config.MaxConnIdleTime = defaultMaxConnIdleTime + } + + if s, ok := config.ConnConfig.Config.RuntimeParams["pool_health_check_period"]; ok { + delete(connConfig.Config.RuntimeParams, "pool_health_check_period") + d, err := time.ParseDuration(s) + if err != nil { + return nil, pgconn.NewParseConfigError(connString, "cannot parse pool_health_check_period", err) + } + config.HealthCheckPeriod = d + } else { + config.HealthCheckPeriod = defaultHealthCheckPeriod + } + + if s, ok := config.ConnConfig.Config.RuntimeParams["pool_max_conn_lifetime_jitter"]; ok { + delete(connConfig.Config.RuntimeParams, "pool_max_conn_lifetime_jitter") + d, err := time.ParseDuration(s) + if err != nil { + return nil, pgconn.NewParseConfigError(connString, "cannot parse pool_max_conn_lifetime_jitter", err) + } + config.MaxConnLifetimeJitter = d + } + + return config, nil +} + +// Close closes all connections in the pool and rejects future [Pool.Acquire] calls. Blocks until all connections are returned +// to pool and closed. +func (p *Pool) Close() { + p.closeOnce.Do(func() { + close(p.closeChan) + p.p.Close() + }) +} + +func (p *Pool) isExpired(res *puddle.Resource[*connResource]) bool { + return time.Now().After(res.Value().maxAgeTime) +} + +func (p *Pool) triggerHealthCheck() { + const healthCheckDelay = 500 * time.Millisecond + + p.healthCheckMu.Lock() + defer p.healthCheckMu.Unlock() + + if p.healthCheckTimer == nil { + // Destroy is asynchronous so we give it time to actually remove itself from + // the pool otherwise we might try to check the pool size too soon + p.healthCheckTimer = time.AfterFunc(healthCheckDelay, func() { + select { + case <-p.closeChan: + case p.healthCheckChan <- struct{}{}: + default: + } + }) + return + } + + p.healthCheckTimer.Reset(healthCheckDelay) +} + +func (p *Pool) backgroundHealthCheck() { + ticker := time.NewTicker(p.healthCheckPeriod) + defer ticker.Stop() + for { + select { + case <-p.closeChan: + return + case <-p.healthCheckChan: + p.checkHealth() + case <-ticker.C: + p.checkHealth() + } + } +} + +func (p *Pool) checkHealth() { + for { + // If checkMinConns failed we don't destroy any connections since we couldn't + // even get to minConns + if err := p.checkMinConns(); err != nil { + // Should we log this error somewhere? + break + } + if !p.checkConnsHealth() { + // Since we didn't destroy any connections we can stop looping + break + } + // Technically Destroy is asynchronous but 500ms should be enough for it to + // remove it from the underlying pool + select { + case <-p.closeChan: + return + case <-time.After(500 * time.Millisecond): + } + } +} + +// checkConnsHealth will check all idle connections, destroy a connection if +// it's idle or too old, and returns true if any were destroyed +func (p *Pool) checkConnsHealth() bool { + var destroyed bool + totalConns := p.Stat().TotalConns() + resources := p.p.AcquireAllIdle() + for _, res := range resources { + switch { + // We're okay going under minConns if the lifetime is up + case p.isExpired(res) && totalConns >= p.minConns: + p.lifetimeDestroyCount.Add(1) + res.Destroy() + destroyed = true + // Since Destroy is async we manually decrement totalConns. + totalConns-- + case res.IdleDuration() > p.maxConnIdleTime && totalConns > p.minConns: + p.idleDestroyCount.Add(1) + res.Destroy() + destroyed = true + // Since Destroy is async we manually decrement totalConns. + totalConns-- + default: + res.ReleaseUnused() + } + } + return destroyed +} + +func (p *Pool) checkMinConns() error { + // TotalConns can include ones that are being destroyed but we should have + // sleep(500ms) around all of the destroys to help prevent that from throwing + // off this check + + // Create the number of connections needed to get to both minConns and minIdleConns + stat := p.Stat() + toCreate := max(p.minConns-stat.TotalConns(), p.minIdleConns-stat.IdleConns()) + if toCreate > 0 { + return p.createIdleResources(context.Background(), int(toCreate)) + } + return nil +} + +func (p *Pool) createIdleResources(parentCtx context.Context, targetResources int) error { + ctx, cancel := context.WithCancel(parentCtx) + defer cancel() + + errs := make(chan error, targetResources) + + for range targetResources { + go func() { + err := p.p.CreateResource(ctx) + // Ignore ErrNotAvailable since it means that the pool has become full since we started creating resource. + if err == puddle.ErrNotAvailable { + err = nil + } + errs <- err + }() + } + + var firstError error + for range targetResources { + err := <-errs + if err != nil && firstError == nil { + cancel() + firstError = err + } + } + + return firstError +} + +// Acquire returns a connection ([Conn]) from the [Pool]. +func (p *Pool) Acquire(ctx context.Context) (c *Conn, err error) { + if p.acquireTracer != nil { + ctx = p.acquireTracer.TraceAcquireStart(ctx, p, TraceAcquireStartData{}) + defer func() { + var conn *pgx.Conn + if c != nil { + conn = c.Conn() + } + p.acquireTracer.TraceAcquireEnd(ctx, p, TraceAcquireEndData{Conn: conn, Err: err}) + }() + } + + // Try to acquire from the connection pool up to maxConns + 1 times, so that + // any that fatal errors would empty the pool and still at least try 1 fresh + // connection. + for range int(p.maxConns) + 1 { + res, err := p.p.Acquire(ctx) + if err != nil { + return nil, err + } + + cr := res.Value() + + // Destroy expired connections before doing any further work (such as + // pinging) on them. This enforces MaxConnLifetime at acquire time so that + // a connection that expired while idle on a busy pool is not handed out. + if p.isExpired(res) { + p.lifetimeDestroyCount.Add(1) + res.Destroy() + continue + } + + shouldPingParams := ShouldPingParams{Conn: cr.conn, IdleDuration: res.IdleDuration()} + if p.shouldPing(ctx, shouldPingParams) { + err := func() error { + pingCtx := ctx + if p.pingTimeout > 0 { + var cancel context.CancelFunc + pingCtx, cancel = context.WithTimeout(ctx, p.pingTimeout) + defer cancel() + } + return cr.conn.Ping(pingCtx) + }() + if err != nil { + res.Destroy() + continue + } + } + + if p.prepareConn != nil { + ok, err := p.prepareConn(ctx, cr.conn) + if !ok { + res.Destroy() + } + if err != nil { + if ok { + res.Release() + } + return nil, err + } + if !ok { + continue + } + } + + return cr.getConn(p, res), nil + } + return nil, errors.New("pgxpool: too many failed attempts acquiring connection; likely bug in PrepareConn, BeforeAcquire, or ShouldPing hook") +} + +// AcquireFunc acquires a [Conn] and calls f with that [Conn]. ctx will only affect the [Pool.Acquire]. It has no effect on the +// call of f. The return value is either an error acquiring the [Conn] or the return value of f. The [Conn] is +// automatically released after the call of f. +func (p *Pool) AcquireFunc(ctx context.Context, f func(*Conn) error) error { + conn, err := p.Acquire(ctx) + if err != nil { + return err + } + defer conn.Release() + + return f(conn) +} + +// AcquireAllIdle atomically acquires all currently idle connections. Its intended use is for health check and +// keep-alive functionality. It does not update pool statistics. +func (p *Pool) AcquireAllIdle(ctx context.Context) []*Conn { + resources := p.p.AcquireAllIdle() + conns := make([]*Conn, 0, len(resources)) + for _, res := range resources { + cr := res.Value() + if p.prepareConn != nil { + ok, err := p.prepareConn(ctx, cr.conn) + if !ok || err != nil { + res.Destroy() + continue + } + } + conns = append(conns, cr.getConn(p, res)) + } + + return conns +} + +// Reset closes all connections, but leaves the pool open. It is intended for use when an error is detected that would +// disrupt all connections (such as a network interruption or a server state change). +// +// It is safe to reset a pool while connections are checked out. Those connections will be closed when they are returned +// to the pool. +func (p *Pool) Reset() { + p.p.Reset() +} + +// Config returns a copy of config that was used to initialize this [Pool]. +func (p *Pool) Config() *Config { return p.config.Copy() } + +// Stat returns a pgxpool.Stat struct with a snapshot of Pool statistics. +func (p *Pool) Stat() *Stat { + return &Stat{ + s: p.p.Stat(), + newConnsCount: p.newConnsCount.Load(), + lifetimeDestroyCount: p.lifetimeDestroyCount.Load(), + idleDestroyCount: p.idleDestroyCount.Load(), + } +} + +// Exec acquires a connection from the [Pool] and executes the given SQL. +// SQL can be either a prepared statement name or an SQL string. +// Arguments should be referenced positionally from the SQL string as $1, $2, etc. +// The acquired connection is returned to the pool when the [Pool.Exec] function returns. +func (p *Pool) Exec(ctx context.Context, sql string, arguments ...any) (pgconn.CommandTag, error) { + c, err := p.Acquire(ctx) + if err != nil { + return pgconn.CommandTag{}, err + } + defer c.Release() + + return c.Exec(ctx, sql, arguments...) +} + +// Query acquires a connection and executes a query that returns [pgx.Rows]. +// Arguments should be referenced positionally from the SQL string as $1, $2, etc. +// See [pgx.Rows] documentation to close the returned [pgx.Rows] and return the acquired connection to the [Pool]. +// +// If there is an error, the returned [pgx.Rows] will be returned in an error state. +// If preferred, ignore the error returned from [Pool.Query] and handle errors using the returned [pgx.Rows]. +// +// For extra control over how the query is executed, the types [pgx.QueryExecMode], [pgx.QueryResultFormats], and +// [pgx.QueryResultFormatsByOID] may be used as the first args to control exactly how the query is executed. This is rarely +// needed. See the documentation for those types for details. +func (p *Pool) Query(ctx context.Context, sql string, args ...any) (pgx.Rows, error) { + c, err := p.Acquire(ctx) + if err != nil { + return errRows{err: err}, err + } + + rows, err := c.Query(ctx, sql, args...) + if err != nil { + c.Release() + return errRows{err: err}, err + } + + return c.getPoolRows(rows), nil +} + +// QueryRow acquires a connection and executes a query that is expected +// to return at most one row ([pgx.Row]). Errors are deferred until [pgx.Row]'s +// Scan method is called. If the query selects no rows, [pgx.Row]'s Scan will +// return [pgx.ErrNoRows]. Otherwise, [pgx.Row]'s Scan scans the first selected row +// and discards the rest. The acquired connection is returned to the [Pool] when +// [pgx.Row]'s Scan method is called. +// +// Arguments should be referenced positionally from the SQL string as $1, $2, etc. +// +// For extra control over how the query is executed, the types [pgx.QueryExecMode], [pgx.QueryResultFormats], and +// [pgx.QueryResultFormatsByOID] may be used as the first args to control exactly how the query is executed. This is rarely +// needed. See the documentation for those types for details. +func (p *Pool) QueryRow(ctx context.Context, sql string, args ...any) pgx.Row { + c, err := p.Acquire(ctx) + if err != nil { + return errRow{err: err} + } + + row := c.QueryRow(ctx, sql, args...) + return c.getPoolRow(row) +} + +func (p *Pool) SendBatch(ctx context.Context, b *pgx.Batch) pgx.BatchResults { + c, err := p.Acquire(ctx) + if err != nil { + return errBatchResults{err: err} + } + + br := c.SendBatch(ctx, b) + return &poolBatchResults{br: br, c: c} +} + +// Begin acquires a connection from the [Pool] and starts a transaction. Unlike [database/sql], the context only affects the begin command. i.e. there is no +// auto-rollback on context cancellation. Begin initiates a transaction block without explicitly setting a transaction mode for the block (see [Pool.BeginTx] with [pgx.TxOptions] if transaction mode is required). +// [*Tx] is returned, which implements the [pgx.Tx] interface. +// [Tx.Commit] or [Tx.Rollback] must be called on the returned transaction to finalize the transaction block. +func (p *Pool) Begin(ctx context.Context) (pgx.Tx, error) { + return p.BeginTx(ctx, pgx.TxOptions{}) +} + +// BeginTx acquires a connection from the [Pool] and starts a transaction with [pgx.TxOptions] determining the transaction mode. +// Unlike [database/sql], the context only affects the begin command. i.e. there is no auto-rollback on context cancellation. +// [*Tx] is returned, which implements the [pgx.Tx] interface. +// [Tx.Commit] or [Tx.Rollback] must be called on the returned transaction to finalize the transaction block. +func (p *Pool) BeginTx(ctx context.Context, txOptions pgx.TxOptions) (pgx.Tx, error) { + c, err := p.Acquire(ctx) + if err != nil { + return nil, err + } + + t, err := c.BeginTx(ctx, txOptions) + if err != nil { + c.Release() + return nil, err + } + + return &Tx{t: t, c: c}, nil +} + +func (p *Pool) CopyFrom(ctx context.Context, tableName pgx.Identifier, columnNames []string, rowSrc pgx.CopyFromSource) (int64, error) { + c, err := p.Acquire(ctx) + if err != nil { + return 0, err + } + defer c.Release() + + return c.Conn().CopyFrom(ctx, tableName, columnNames, rowSrc) +} + +// Ping acquires a connection from the [Pool] and executes an empty sql statement against it. +// If the sql returns without error, the database [Pool.Ping] is considered successful, otherwise, the error is returned. +func (p *Pool) Ping(ctx context.Context) error { + c, err := p.Acquire(ctx) + if err != nil { + return err + } + defer c.Release() + return c.Ping(ctx) +} diff --git a/vendor/github.com/jackc/pgx/v5/pgxpool/rows.go b/vendor/github.com/jackc/pgx/v5/pgxpool/rows.go new file mode 100644 index 0000000000..f834b7ec30 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgxpool/rows.go @@ -0,0 +1,116 @@ +package pgxpool + +import ( + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgconn" +) + +type errRows struct { + err error +} + +func (errRows) Close() {} +func (e errRows) Err() error { return e.err } +func (errRows) CommandTag() pgconn.CommandTag { return pgconn.CommandTag{} } +func (errRows) FieldDescriptions() []pgconn.FieldDescription { return nil } +func (errRows) Next() bool { return false } +func (e errRows) Scan(dest ...any) error { return e.err } +func (e errRows) Values() ([]any, error) { return nil, e.err } +func (e errRows) RawValues() [][]byte { return nil } +func (e errRows) Conn() *pgx.Conn { return nil } + +type errRow struct { + err error +} + +func (e errRow) Scan(dest ...any) error { return e.err } + +type poolRows struct { + r pgx.Rows + c *Conn + err error +} + +func (rows *poolRows) Close() { + rows.r.Close() + if rows.c != nil { + rows.c.Release() + rows.c = nil + } +} + +func (rows *poolRows) Err() error { + if rows.err != nil { + return rows.err + } + return rows.r.Err() +} + +func (rows *poolRows) CommandTag() pgconn.CommandTag { + return rows.r.CommandTag() +} + +func (rows *poolRows) FieldDescriptions() []pgconn.FieldDescription { + return rows.r.FieldDescriptions() +} + +func (rows *poolRows) Next() bool { + if rows.err != nil { + return false + } + + n := rows.r.Next() + if !n { + rows.Close() + } + return n +} + +func (rows *poolRows) Scan(dest ...any) error { + err := rows.r.Scan(dest...) + if err != nil { + rows.Close() + } + return err +} + +func (rows *poolRows) Values() ([]any, error) { + values, err := rows.r.Values() + if err != nil { + rows.Close() + } + return values, err +} + +func (rows *poolRows) RawValues() [][]byte { + return rows.r.RawValues() +} + +func (rows *poolRows) Conn() *pgx.Conn { + return rows.r.Conn() +} + +type poolRow struct { + r pgx.Row + c *Conn + err error +} + +func (row *poolRow) Scan(dest ...any) error { + if row.err != nil { + return row.err + } + + panicked := true + defer func() { + if panicked && row.c != nil { + row.c.Release() + } + }() + err := row.r.Scan(dest...) + panicked = false + if row.c != nil { + row.c.Release() + } + return err +} diff --git a/vendor/github.com/jackc/pgx/v5/pgxpool/stat.go b/vendor/github.com/jackc/pgx/v5/pgxpool/stat.go new file mode 100644 index 0000000000..e02b6ac39e --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgxpool/stat.go @@ -0,0 +1,91 @@ +package pgxpool + +import ( + "time" + + "github.com/jackc/puddle/v2" +) + +// Stat is a snapshot of Pool statistics. +type Stat struct { + s *puddle.Stat + newConnsCount int64 + lifetimeDestroyCount int64 + idleDestroyCount int64 +} + +// AcquireCount returns the cumulative count of successful acquires from the pool. +func (s *Stat) AcquireCount() int64 { + return s.s.AcquireCount() +} + +// AcquireDuration returns the total duration of all successful acquires from +// the pool. +func (s *Stat) AcquireDuration() time.Duration { + return s.s.AcquireDuration() +} + +// AcquiredConns returns the number of currently acquired connections in the pool. +func (s *Stat) AcquiredConns() int32 { + return s.s.AcquiredResources() +} + +// CanceledAcquireCount returns the cumulative count of acquires from the pool +// that were canceled by a context. +func (s *Stat) CanceledAcquireCount() int64 { + return s.s.CanceledAcquireCount() +} + +// ConstructingConns returns the number of conns with construction in progress in +// the pool. +func (s *Stat) ConstructingConns() int32 { + return s.s.ConstructingResources() +} + +// EmptyAcquireCount returns the cumulative count of successful acquires from the pool +// that waited for a resource to be released or constructed because the pool was +// empty. +func (s *Stat) EmptyAcquireCount() int64 { + return s.s.EmptyAcquireCount() +} + +// IdleConns returns the number of currently idle conns in the pool. +func (s *Stat) IdleConns() int32 { + return s.s.IdleResources() +} + +// MaxConns returns the maximum size of the pool. +func (s *Stat) MaxConns() int32 { + return s.s.MaxResources() +} + +// TotalConns returns the total number of resources currently in the pool. +// The value is the sum of ConstructingConns, AcquiredConns, and +// IdleConns. +func (s *Stat) TotalConns() int32 { + return s.s.TotalResources() +} + +// NewConnsCount returns the cumulative count of new connections opened. +func (s *Stat) NewConnsCount() int64 { + return s.newConnsCount +} + +// MaxLifetimeDestroyCount returns the cumulative count of connections destroyed +// because they exceeded MaxConnLifetime. +func (s *Stat) MaxLifetimeDestroyCount() int64 { + return s.lifetimeDestroyCount +} + +// MaxIdleDestroyCount returns the cumulative count of connections destroyed because +// they exceeded MaxConnIdleTime. +func (s *Stat) MaxIdleDestroyCount() int64 { + return s.idleDestroyCount +} + +// EmptyAcquireWaitTime returns the cumulative time waited for successful acquires +// from the pool for a resource to be released or constructed because the pool was +// empty. +func (s *Stat) EmptyAcquireWaitTime() time.Duration { + return s.s.EmptyAcquireWaitTime() +} diff --git a/vendor/github.com/jackc/pgx/v5/pgxpool/tracer.go b/vendor/github.com/jackc/pgx/v5/pgxpool/tracer.go new file mode 100644 index 0000000000..78b9d15a2d --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgxpool/tracer.go @@ -0,0 +1,33 @@ +package pgxpool + +import ( + "context" + + "github.com/jackc/pgx/v5" +) + +// AcquireTracer traces Acquire. +type AcquireTracer interface { + // TraceAcquireStart is called at the beginning of Acquire. + // The returned context is used for the rest of the call and will be passed to the TraceAcquireEnd. + TraceAcquireStart(ctx context.Context, pool *Pool, data TraceAcquireStartData) context.Context + // TraceAcquireEnd is called when a connection has been acquired. + TraceAcquireEnd(ctx context.Context, pool *Pool, data TraceAcquireEndData) +} + +type TraceAcquireStartData struct{} + +type TraceAcquireEndData struct { + Conn *pgx.Conn + Err error +} + +// ReleaseTracer traces Release. +type ReleaseTracer interface { + // TraceRelease is called at the beginning of Release. + TraceRelease(pool *Pool, data TraceReleaseData) +} + +type TraceReleaseData struct { + Conn *pgx.Conn +} diff --git a/vendor/github.com/jackc/pgx/v5/pgxpool/tx.go b/vendor/github.com/jackc/pgx/v5/pgxpool/tx.go new file mode 100644 index 0000000000..b49e7f4d96 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/pgxpool/tx.go @@ -0,0 +1,83 @@ +package pgxpool + +import ( + "context" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgconn" +) + +// Tx represents a database transaction acquired from a Pool. +type Tx struct { + t pgx.Tx + c *Conn +} + +// Begin starts a pseudo nested transaction implemented with a savepoint. +func (tx *Tx) Begin(ctx context.Context) (pgx.Tx, error) { + return tx.t.Begin(ctx) +} + +// Commit commits the transaction and returns the associated connection back to the Pool. Commit will return an error +// where errors.Is(ErrTxClosed) is true if the Tx is already closed, but is otherwise safe to call multiple times. If +// the commit fails with a rollback status (e.g. the transaction was already in a broken state) then ErrTxCommitRollback +// will be returned. +func (tx *Tx) Commit(ctx context.Context) error { + err := tx.t.Commit(ctx) + if tx.c != nil { + tx.c.Release() + tx.c = nil + } + return err +} + +// Rollback rolls back the transaction and returns the associated connection back to the Pool. Rollback will return +// where an error where errors.Is(ErrTxClosed) is true if the Tx is already closed, but is otherwise safe to call +// multiple times. Hence, defer tx.Rollback() is safe even if tx.Commit() will be called first in a non-error condition. +func (tx *Tx) Rollback(ctx context.Context) error { + err := tx.t.Rollback(ctx) + if tx.c != nil { + tx.c.Release() + tx.c = nil + } + return err +} + +func (tx *Tx) CopyFrom(ctx context.Context, tableName pgx.Identifier, columnNames []string, rowSrc pgx.CopyFromSource) (int64, error) { + return tx.t.CopyFrom(ctx, tableName, columnNames, rowSrc) +} + +func (tx *Tx) SendBatch(ctx context.Context, b *pgx.Batch) pgx.BatchResults { + return tx.t.SendBatch(ctx, b) +} + +func (tx *Tx) LargeObjects() pgx.LargeObjects { + return tx.t.LargeObjects() +} + +// Prepare creates a prepared statement with name and sql. If the name is empty, +// an anonymous prepared statement will be used. sql can contain placeholders +// for bound parameters. These placeholders are referenced positionally as $1, $2, etc. +// +// Prepare is idempotent; i.e. it is safe to call Prepare multiple times with the same +// name and sql arguments. This allows a code path to Prepare and Query/Exec without +// needing to first check whether the statement has already been prepared. +func (tx *Tx) Prepare(ctx context.Context, name, sql string) (*pgconn.StatementDescription, error) { + return tx.t.Prepare(ctx, name, sql) +} + +func (tx *Tx) Exec(ctx context.Context, sql string, arguments ...any) (pgconn.CommandTag, error) { + return tx.t.Exec(ctx, sql, arguments...) +} + +func (tx *Tx) Query(ctx context.Context, sql string, args ...any) (pgx.Rows, error) { + return tx.t.Query(ctx, sql, args...) +} + +func (tx *Tx) QueryRow(ctx context.Context, sql string, args ...any) pgx.Row { + return tx.t.QueryRow(ctx, sql, args...) +} + +func (tx *Tx) Conn() *pgx.Conn { + return tx.t.Conn() +} diff --git a/vendor/github.com/jackc/pgx/v5/rows.go b/vendor/github.com/jackc/pgx/v5/rows.go new file mode 100644 index 0000000000..4e5cf95d05 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/rows.go @@ -0,0 +1,871 @@ +package pgx + +import ( + "context" + "errors" + "fmt" + "reflect" + "strings" + "sync" + "time" + + "github.com/jackc/pgx/v5/pgconn" + "github.com/jackc/pgx/v5/pgtype" +) + +// Rows is the result set returned from [Conn.Query]. Rows must be closed before +// the [Conn] can be used again. Rows are closed by explicitly calling [Rows.Close], +// calling [Rows.Next] until it returns false, or when a fatal error occurs. +// +// Once a Rows is closed the only methods that may be called are [Rows.Close], [Rows.Err], +// and [Rows.CommandTag]. +// +// Rows is an interface instead of a struct to allow tests to mock Query. However, +// adding a method to an interface is technically a breaking change. Because of this +// the Rows interface is partially excluded from semantic version requirements. +// Methods will not be removed or changed, but new methods may be added. +type Rows interface { + // Close closes the rows, making the connection ready for use again. It is safe + // to call Close after rows is already closed. + Close() + + // Err returns any error that occurred while executing a query or reading its results. Err must be called after the + // Rows is closed (either by calling Close or by Next returning false) to check if the query was successful. If it is + // called before the Rows is closed it may return nil even if the query failed on the server. + Err() error + + // CommandTag returns the command tag from this query. It is only available after Rows is closed. + CommandTag() pgconn.CommandTag + + // FieldDescriptions returns the field descriptions of the columns. It may return nil. In particular this can occur + // when there was an error executing the query. + FieldDescriptions() []pgconn.FieldDescription + + // Next prepares the next row for reading. It returns true if there is another row and false if no more rows are + // available or a fatal error has occurred. It automatically closes rows upon returning false (whether due to all rows + // having been read or due to an error). + // + // Callers should check rows.Err() after rows.Next() returns false to detect whether result-set reading ended + // prematurely due to an error. See [Conn.Query] for details. + // + // For simpler error handling, consider using the higher-level pgx v5 [CollectRows()] and [ForEachRow()] helpers instead. + Next() bool + + // Scan reads the values from the current row into dest values positionally. dest can include pointers to core types, + // values implementing the Scanner interface, and nil. nil will skip the value entirely. It is an error to call Scan + // without first calling Next() and checking that it returned true. Rows is automatically closed upon error. + Scan(dest ...any) error + + // Values returns the decoded row values. As with Scan(), it is an error to + // call Values without first calling Next() and checking that it returned + // true. + Values() ([]any, error) + + // RawValues returns the unparsed bytes of the row values. The returned data is only valid until the next Next + // call or the Rows is closed. + RawValues() [][]byte + + // Conn returns the underlying *Conn on which the query was executed. This may return nil if Rows did not come from a + // *Conn (e.g. if it was created by RowsFromResultReader) + Conn() *Conn +} + +// Row is a convenience wrapper over [Rows] that is returned by [Conn.QueryRow]. +// +// Row is an interface instead of a struct to allow tests to mock QueryRow. However, +// adding a method to an interface is technically a breaking change. Because of this +// the Row interface is partially excluded from semantic version requirements. +// Methods will not be removed or changed, but new methods may be added. +type Row interface { + // Scan works the same as Rows. with the following exceptions. If no + // rows were found it returns ErrNoRows. If multiple rows are returned it + // ignores all but the first. + Scan(dest ...any) error +} + +// RowScanner scans an entire row at a time into the RowScanner. +type RowScanner interface { + // ScanRows scans the row. + ScanRow(rows Rows) error +} + +// connRow implements the Row interface for Conn.QueryRow. +type connRow baseRows + +func (r *connRow) Scan(dest ...any) (err error) { + rows := (*baseRows)(r) + + if rows.Err() != nil { + return rows.Err() + } + + for _, d := range dest { + if _, ok := d.(*pgtype.DriverBytes); ok { + rows.Close() + return fmt.Errorf("cannot scan into *pgtype.DriverBytes from QueryRow") + } + } + + if !rows.Next() { + if rows.Err() == nil { + return ErrNoRows + } + return rows.Err() + } + + rows.Scan(dest...) + rows.Close() + return rows.Err() +} + +// baseRows implements the Rows interface for Conn.Query. +type baseRows struct { + typeMap *pgtype.Map + resultReader *pgconn.ResultReader + + values [][]byte + + commandTag pgconn.CommandTag + err error + closed bool + + scanPlans []pgtype.ScanPlan + scanTypes []reflect.Type + + conn *Conn + multiResultReader *pgconn.MultiResultReader + + queryTracer QueryTracer + batchTracer BatchTracer + ctx context.Context + startTime time.Time + sql string + args []any + rowCount int +} + +func (rows *baseRows) FieldDescriptions() []pgconn.FieldDescription { + return rows.resultReader.FieldDescriptions() +} + +func (rows *baseRows) Close() { + if rows.closed { + return + } + + rows.closed = true + + if rows.resultReader != nil { + var closeErr error + rows.commandTag, closeErr = rows.resultReader.Close() + if rows.err == nil { + rows.err = closeErr + } + } + + if rows.multiResultReader != nil { + closeErr := rows.multiResultReader.Close() + if rows.err == nil { + rows.err = closeErr + } + } + + if rows.err != nil && rows.conn != nil && rows.sql != "" { + if sc := rows.conn.statementCache; sc != nil { + sc.Invalidate(rows.sql) + } + + if sc := rows.conn.descriptionCache; sc != nil { + sc.Invalidate(rows.sql) + } + } + + if rows.batchTracer != nil { + rows.batchTracer.TraceBatchQuery(rows.ctx, rows.conn, TraceBatchQueryData{SQL: rows.sql, Args: rows.args, CommandTag: rows.commandTag, Err: rows.err}) + } else if rows.queryTracer != nil { + rows.queryTracer.TraceQueryEnd(rows.ctx, rows.conn, TraceQueryEndData{rows.commandTag, rows.err}) + } + + // Zero references to other memory allocations. This allows them to be GC'd even when the Rows still referenced. In + // particular, when using pgxpool GC could be delayed as pgxpool.poolRows are allocated in large slices. + // + // https://github.com/jackc/pgx/pull/2269 + rows.values = nil + rows.scanPlans = nil + rows.scanTypes = nil + rows.ctx = nil + rows.sql = "" + rows.args = nil +} + +func (rows *baseRows) CommandTag() pgconn.CommandTag { + return rows.commandTag +} + +func (rows *baseRows) Err() error { + return rows.err +} + +// fatal signals an error occurred after the query was sent to the server. It +// closes the rows automatically. +func (rows *baseRows) fatal(err error) { + if rows.err != nil { + return + } + + rows.err = err + rows.Close() +} + +func (rows *baseRows) Next() bool { + if rows.closed { + return false + } + + if rows.resultReader.NextRow() { + rows.rowCount++ + rows.values = rows.resultReader.Values() + return true + } else { + rows.Close() + return false + } +} + +func (rows *baseRows) Scan(dest ...any) error { + m := rows.typeMap + fieldDescriptions := rows.FieldDescriptions() + values := rows.values + + if len(fieldDescriptions) != len(values) { + err := fmt.Errorf("number of field descriptions must equal number of values, got %d and %d", len(fieldDescriptions), len(values)) + rows.fatal(err) + return err + } + + if len(dest) == 1 { + if rc, ok := dest[0].(RowScanner); ok { + err := rc.ScanRow(rows) + if err != nil { + rows.fatal(err) + } + return err + } + } + + if len(fieldDescriptions) != len(dest) { + err := fmt.Errorf("number of field descriptions must equal number of destinations, got %d and %d", len(fieldDescriptions), len(dest)) + rows.fatal(err) + return err + } + + if rows.scanPlans == nil { + rows.scanPlans = make([]pgtype.ScanPlan, len(values)) + rows.scanTypes = make([]reflect.Type, len(values)) + for i := range dest { + rows.scanPlans[i] = m.PlanScan(fieldDescriptions[i].DataTypeOID, fieldDescriptions[i].Format, dest[i]) + rows.scanTypes[i] = reflect.TypeOf(dest[i]) + } + } + + for i, dst := range dest { + if dst == nil { + continue + } + + if rows.scanTypes[i] != reflect.TypeOf(dst) { + rows.scanPlans[i] = m.PlanScan(fieldDescriptions[i].DataTypeOID, fieldDescriptions[i].Format, dest[i]) + rows.scanTypes[i] = reflect.TypeOf(dest[i]) + } + + err := rows.scanPlans[i].Scan(values[i], dst) + if err != nil { + err = ScanArgError{ColumnIndex: i, FieldName: fieldDescriptions[i].Name, Err: err} + rows.fatal(err) + return err + } + } + + return nil +} + +func (rows *baseRows) Values() ([]any, error) { + if rows.closed { + return nil, errors.New("rows is closed") + } + + values := make([]any, 0, len(rows.FieldDescriptions())) + + for i := range rows.FieldDescriptions() { + buf := rows.values[i] + fd := &rows.FieldDescriptions()[i] + + if buf == nil { + values = append(values, nil) + continue + } + + if dt, ok := rows.typeMap.TypeForOID(fd.DataTypeOID); ok { + value, err := dt.Codec.DecodeValue(rows.typeMap, fd.DataTypeOID, fd.Format, buf) + if err != nil { + rows.fatal(err) + } + values = append(values, value) + } else { + switch fd.Format { + case TextFormatCode: + values = append(values, string(buf)) + case BinaryFormatCode: + newBuf := make([]byte, len(buf)) + copy(newBuf, buf) + values = append(values, newBuf) + default: + rows.fatal(errors.New("unknown format code")) + } + } + + if rows.Err() != nil { + return nil, rows.Err() + } + } + + return values, rows.Err() +} + +func (rows *baseRows) RawValues() [][]byte { + return rows.values +} + +func (rows *baseRows) Conn() *Conn { + return rows.conn +} + +type ScanArgError struct { + ColumnIndex int + FieldName string + Err error +} + +func (e ScanArgError) Error() string { + if e.FieldName == "?column?" { // Don't include the fieldname if it's unknown + return fmt.Sprintf("can't scan into dest[%d]: %v", e.ColumnIndex, e.Err) + } + + return fmt.Sprintf("can't scan into dest[%d] (col: %s): %v", e.ColumnIndex, e.FieldName, e.Err) +} + +func (e ScanArgError) Unwrap() error { + return e.Err +} + +// ScanRow decodes raw row data into dest. It can be used to scan rows read from the lower level [pgconn] interface. +// +// typeMap - OID to Go type mapping. +// fieldDescriptions - OID and format of values +// values - the raw data as returned from the PostgreSQL server +// dest - the destination that values will be decoded into +func ScanRow(typeMap *pgtype.Map, fieldDescriptions []pgconn.FieldDescription, values [][]byte, dest ...any) error { + if len(fieldDescriptions) != len(values) { + return fmt.Errorf("number of field descriptions must equal number of values, got %d and %d", len(fieldDescriptions), len(values)) + } + if len(fieldDescriptions) != len(dest) { + return fmt.Errorf("number of field descriptions must equal number of destinations, got %d and %d", len(fieldDescriptions), len(dest)) + } + + for i, d := range dest { + if d == nil { + continue + } + + err := typeMap.Scan(fieldDescriptions[i].DataTypeOID, fieldDescriptions[i].Format, values[i], d) + if err != nil { + return ScanArgError{ColumnIndex: i, FieldName: fieldDescriptions[i].Name, Err: err} + } + } + + return nil +} + +// RowsFromResultReader returns a [Rows] that will read from values resultReader and decode with typeMap. It can be used +// to read from the lower level [pgconn] interface. +func RowsFromResultReader(typeMap *pgtype.Map, resultReader *pgconn.ResultReader) Rows { + return &baseRows{ + typeMap: typeMap, + resultReader: resultReader, + } +} + +// ForEachRow iterates through rows. For each row it scans into the elements of scans and calls fn. If any row +// fails to scan or fn returns an error the query will be aborted and the error will be returned. Rows will be closed +// when ForEachRow returns. +func ForEachRow(rows Rows, scans []any, fn func() error) (pgconn.CommandTag, error) { + defer rows.Close() + + for rows.Next() { + err := rows.Scan(scans...) + if err != nil { + return pgconn.CommandTag{}, err + } + + err = fn() + if err != nil { + return pgconn.CommandTag{}, err + } + } + + if err := rows.Err(); err != nil { + return pgconn.CommandTag{}, err + } + + return rows.CommandTag(), nil +} + +// CollectableRow is the subset of Rows methods that a RowToFunc is allowed to call. +type CollectableRow interface { + FieldDescriptions() []pgconn.FieldDescription + Scan(dest ...any) error + Values() ([]any, error) + RawValues() [][]byte +} + +// RowToFunc is a function that scans or otherwise converts row to a T. +type RowToFunc[T any] func(row CollectableRow) (T, error) + +// AppendRows iterates through rows, calling fn for each row, and appending the results into a slice of T. +// +// This function closes the rows automatically on return. +func AppendRows[T any, S ~[]T](slice S, rows Rows, fn RowToFunc[T]) (S, error) { + defer rows.Close() + + for rows.Next() { + value, err := fn(rows) + if err != nil { + return nil, err + } + slice = append(slice, value) + } + + if err := rows.Err(); err != nil { + return nil, err + } + + return slice, nil +} + +// CollectRows iterates through rows, calling fn for each row, and collecting the results into a slice of T. +// +// This function closes the rows automatically on return. +func CollectRows[T any](rows Rows, fn RowToFunc[T]) ([]T, error) { + return AppendRows([]T{}, rows, fn) +} + +// CollectOneRow calls fn for the first row in rows and returns the result. If no rows are found returns an error where errors.Is(ErrNoRows) is true. +// CollectOneRow is to [CollectRows] as [Conn.QueryRow] is to [Conn.Query]. +// +// This function closes the rows automatically on return. +func CollectOneRow[T any](rows Rows, fn RowToFunc[T]) (T, error) { + defer rows.Close() + + var value T + var err error + + if !rows.Next() { + if err = rows.Err(); err != nil { + return value, err + } + return value, ErrNoRows + } + + value, err = fn(rows) + if err != nil { + return value, err + } + + // The defer rows.Close() won't have executed yet. If the query returned more than one row, rows would still be open. + // rows.Close() must be called before rows.Err() so we explicitly call it here. + rows.Close() + return value, rows.Err() +} + +// CollectExactlyOneRow calls fn for the first row in rows and returns the result. +// - If no rows are found returns an error where errors.Is(ErrNoRows) is true. +// - If more than 1 row is found returns an error where errors.Is(ErrTooManyRows) is true. +// +// This function closes the rows automatically on return. +func CollectExactlyOneRow[T any](rows Rows, fn RowToFunc[T]) (T, error) { + defer rows.Close() + + var ( + err error + value T + ) + + if !rows.Next() { + if err = rows.Err(); err != nil { + return value, err + } + + return value, ErrNoRows + } + + value, err = fn(rows) + if err != nil { + return value, err + } + + if rows.Next() { + var zero T + + return zero, ErrTooManyRows + } + + return value, rows.Err() +} + +// RowTo returns a T scanned from row. +func RowTo[T any](row CollectableRow) (T, error) { + var value T + err := row.Scan(&value) + return value, err +} + +// RowToAddrOf returns the address of a T scanned from row. +func RowToAddrOf[T any](row CollectableRow) (*T, error) { + var value T + err := row.Scan(&value) + return &value, err +} + +// RowToMap returns a map scanned from row. +func RowToMap(row CollectableRow) (map[string]any, error) { + var value map[string]any + err := row.Scan((*mapRowScanner)(&value)) + return value, err +} + +type mapRowScanner map[string]any + +func (rs *mapRowScanner) ScanRow(rows Rows) error { + values, err := rows.Values() + if err != nil { + return err + } + + *rs = make(mapRowScanner, len(values)) + + for i := range values { + (*rs)[rows.FieldDescriptions()[i].Name] = values[i] + } + + return nil +} + +// RowToStructByPos returns a T scanned from row. T must be a struct. T must have the same number of public fields as row +// has fields. The row and T fields will be matched by position. If the "db" struct tag is "-" then the field will be +// ignored. +func RowToStructByPos[T any](row CollectableRow) (T, error) { + var value T + err := (&positionalStructRowScanner{ptrToStruct: &value}).ScanRow(row) + return value, err +} + +// RowToAddrOfStructByPos returns the address of a T scanned from row. T must be a struct. T must have the same number a +// public fields as row has fields. The row and T fields will be matched by position. If the "db" struct tag is "-" then +// the field will be ignored. +func RowToAddrOfStructByPos[T any](row CollectableRow) (*T, error) { + var value T + err := (&positionalStructRowScanner{ptrToStruct: &value}).ScanRow(row) + return &value, err +} + +type positionalStructRowScanner struct { + ptrToStruct any +} + +func (rs *positionalStructRowScanner) ScanRow(rows CollectableRow) error { + typ := reflect.TypeOf(rs.ptrToStruct).Elem() + fields := lookupStructFields(typ) + if len(rows.RawValues()) > len(fields) { + return fmt.Errorf( + "got %d values, but dst struct has only %d fields", + len(rows.RawValues()), + len(fields), + ) + } + scanTargets := setupStructScanTargets(rs.ptrToStruct, fields) + return rows.Scan(scanTargets...) +} + +// Map from reflect.Type -> []structRowField +var positionalStructFieldMap sync.Map + +func lookupStructFields(t reflect.Type) []structRowField { + if cached, ok := positionalStructFieldMap.Load(t); ok { + return cached.([]structRowField) + } + + fieldStack := make([]int, 0, 1) + fields := computeStructFields(t, make([]structRowField, 0, t.NumField()), &fieldStack) + fieldsIface, _ := positionalStructFieldMap.LoadOrStore(t, fields) + return fieldsIface.([]structRowField) +} + +func computeStructFields( + t reflect.Type, + fields []structRowField, + fieldStack *[]int, +) []structRowField { + tail := len(*fieldStack) + *fieldStack = append(*fieldStack, 0) + for i := 0; i < t.NumField(); i++ { + sf := t.Field(i) + (*fieldStack)[tail] = i + // Handle anonymous struct embedding, but do not try to handle embedded pointers. + if sf.Anonymous && sf.Type.Kind() == reflect.Struct { + fields = computeStructFields(sf.Type, fields, fieldStack) + } else if sf.PkgPath == "" { + dbTag, _ := sf.Tag.Lookup(structTagKey) + if dbTag == "-" { + // Field is ignored, skip it. + continue + } + fields = append(fields, structRowField{ + path: append([]int(nil), *fieldStack...), + }) + } + } + *fieldStack = (*fieldStack)[:tail] + return fields +} + +// RowToStructByName returns a T scanned from row. T must be a struct. T must have the same number of named public +// fields as row has fields. The row and T fields will be matched by name. The match is case-insensitive. The database +// column name can be overridden with a "db" struct tag. If the "db" struct tag is "-" then the field will be ignored. +func RowToStructByName[T any](row CollectableRow) (T, error) { + var value T + err := (&namedStructRowScanner{ptrToStruct: &value}).ScanRow(row) + return value, err +} + +// RowToAddrOfStructByName returns the address of a T scanned from row. T must be a struct. T must have the same number +// of named public fields as row has fields. The row and T fields will be matched by name. The match is +// case-insensitive. The database column name can be overridden with a "db" struct tag. If the "db" struct tag is "-" +// then the field will be ignored. +func RowToAddrOfStructByName[T any](row CollectableRow) (*T, error) { + var value T + err := (&namedStructRowScanner{ptrToStruct: &value}).ScanRow(row) + return &value, err +} + +// RowToStructByNameLax returns a T scanned from row. T must be a struct. T must have greater than or equal number of named public +// fields as row has fields. The row and T fields will be matched by name. The match is case-insensitive. The database +// column name can be overridden with a "db" struct tag. If the "db" struct tag is "-" then the field will be ignored. +func RowToStructByNameLax[T any](row CollectableRow) (T, error) { + var value T + err := (&namedStructRowScanner{ptrToStruct: &value, lax: true}).ScanRow(row) + return value, err +} + +// RowToAddrOfStructByNameLax returns the address of a T scanned from row. T must be a struct. T must have greater than or +// equal number of named public fields as row has fields. The row and T fields will be matched by name. The match is +// case-insensitive. The database column name can be overridden with a "db" struct tag. If the "db" struct tag is "-" +// then the field will be ignored. +func RowToAddrOfStructByNameLax[T any](row CollectableRow) (*T, error) { + var value T + err := (&namedStructRowScanner{ptrToStruct: &value, lax: true}).ScanRow(row) + return &value, err +} + +type namedStructRowScanner struct { + ptrToStruct any + lax bool +} + +func (rs *namedStructRowScanner) ScanRow(rows CollectableRow) error { + typ := reflect.TypeOf(rs.ptrToStruct).Elem() + fldDescs := rows.FieldDescriptions() + namedStructFields, err := lookupNamedStructFields(typ, fldDescs) + if err != nil { + return err + } + if !rs.lax && namedStructFields.missingField != "" { + return fmt.Errorf("cannot find field %s in returned row", namedStructFields.missingField) + } + fields := namedStructFields.fields + scanTargets := setupStructScanTargets(rs.ptrToStruct, fields) + return rows.Scan(scanTargets...) +} + +// Map from namedStructFieldMap -> *namedStructFields +var namedStructFieldMap sync.Map + +type namedStructFieldsKey struct { + t reflect.Type + colNames string +} + +type namedStructFields struct { + fields []structRowField + // missingField is the first field from the struct without a corresponding row field. + // This is used to construct the correct error message for non-lax queries. + missingField string +} + +func lookupNamedStructFields( + t reflect.Type, + fldDescs []pgconn.FieldDescription, +) (*namedStructFields, error) { + key := namedStructFieldsKey{ + t: t, + colNames: joinFieldNames(fldDescs), + } + if cached, ok := namedStructFieldMap.Load(key); ok { + return cached.(*namedStructFields), nil + } + + // We could probably do two-levels of caching, where we compute the key -> fields mapping + // for a type only once, cache it by type, then use that to compute the column -> fields + // mapping for a given set of columns. + fieldStack := make([]int, 0, 1) + fields, missingField := computeNamedStructFields( + fldDescs, + t, + make([]structRowField, len(fldDescs)), + &fieldStack, + ) + for i, f := range fields { + if f.path == nil { + return nil, fmt.Errorf( + "struct doesn't have corresponding row field %s", + fldDescs[i].Name, + ) + } + } + + fieldsIface, _ := namedStructFieldMap.LoadOrStore( + key, + &namedStructFields{fields: fields, missingField: missingField}, + ) + return fieldsIface.(*namedStructFields), nil +} + +func joinFieldNames(fldDescs []pgconn.FieldDescription) string { + switch len(fldDescs) { + case 0: + return "" + case 1: + return fldDescs[0].Name + } + + totalSize := len(fldDescs) - 1 // Space for separator bytes. + for _, d := range fldDescs { + totalSize += len(d.Name) + } + var b strings.Builder + b.Grow(totalSize) + b.WriteString(fldDescs[0].Name) + for _, d := range fldDescs[1:] { + b.WriteByte(0) // Join with NUL byte as it's (presumably) not a valid column character. + b.WriteString(d.Name) + } + return b.String() +} + +func computeNamedStructFields( + fldDescs []pgconn.FieldDescription, + t reflect.Type, + fields []structRowField, + fieldStack *[]int, +) ([]structRowField, string) { + var missingField string + tail := len(*fieldStack) + *fieldStack = append(*fieldStack, 0) + for i := 0; i < t.NumField(); i++ { + sf := t.Field(i) + (*fieldStack)[tail] = i + if sf.PkgPath != "" && !sf.Anonymous { + // Field is unexported, skip it. + continue + } + // Handle anonymous struct embedding, but do not try to handle embedded pointers. + if sf.Anonymous && sf.Type.Kind() == reflect.Struct { + var missingSubField string + fields, missingSubField = computeNamedStructFields( + fldDescs, + sf.Type, + fields, + fieldStack, + ) + if missingField == "" { + missingField = missingSubField + } + } else { + dbTag, dbTagPresent := sf.Tag.Lookup(structTagKey) + if dbTagPresent { + dbTag, _, _ = strings.Cut(dbTag, ",") + } + if dbTag == "-" { + // Field is ignored, skip it. + continue + } + colName := dbTag + if !dbTagPresent { + colName = sf.Name + } + fpos := fieldPosByName(fldDescs, colName, !dbTagPresent) + if fpos == -1 { + if missingField == "" { + missingField = colName + } + continue + } + fields[fpos] = structRowField{ + path: append([]int(nil), *fieldStack...), + } + } + } + *fieldStack = (*fieldStack)[:tail] + + return fields, missingField +} + +const structTagKey = "db" + +func fieldPosByName(fldDescs []pgconn.FieldDescription, field string, normalize bool) (i int) { + i = -1 + + if normalize { + field = strings.ReplaceAll(field, "_", "") + } + for i, desc := range fldDescs { + if normalize { + if strings.EqualFold(strings.ReplaceAll(desc.Name, "_", ""), field) { + return i + } + } else { + if desc.Name == field { + return i + } + } + } + return i +} + +// structRowField describes a field of a struct. +// +// TODO: It would be a bit more efficient to track the path using the pointer +// offset within the (outermost) struct and use unsafe.Pointer arithmetic to +// construct references when scanning rows. However, it's not clear it's worth +// using unsafe for this. +type structRowField struct { + path []int +} + +func setupStructScanTargets(receiver any, fields []structRowField) []any { + scanTargets := make([]any, len(fields)) + v := reflect.ValueOf(receiver).Elem() + for i, f := range fields { + scanTargets[i] = v.FieldByIndex(f.path).Addr().Interface() + } + return scanTargets +} diff --git a/vendor/github.com/jackc/pgx/v5/test.sh b/vendor/github.com/jackc/pgx/v5/test.sh new file mode 100644 index 0000000000..8bab2d280c --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/test.sh @@ -0,0 +1,170 @@ +#!/usr/bin/env bash +set -euo pipefail + +# test.sh - Run pgx tests against specific database targets +# +# Usage: +# ./test.sh [target] [go test flags...] +# +# Targets: +# pg14 - PostgreSQL 14 (port 5414) +# pg15 - PostgreSQL 15 (port 5415) +# pg16 - PostgreSQL 16 (port 5416) +# pg17 - PostgreSQL 17 (port 5417) +# pg18 - PostgreSQL 18 (port 5432) [default] +# crdb - CockroachDB (port 26257) +# all - Run against all targets sequentially +# +# Examples: +# ./test.sh # Test against PG18 +# ./test.sh pg14 # Test against PG14 +# ./test.sh crdb # Test against CockroachDB +# ./test.sh all # Test against all targets +# ./test.sh pg16 -run TestConnect # Test specific test against PG16 +# ./test.sh pg18 -count=1 -v # Verbose, no cache, PG18 + +# Color output (disabled if not a terminal) +if [ -t 1 ]; then + GREEN='\033[0;32m' + RED='\033[0;31m' + BLUE='\033[0;34m' + NC='\033[0m' +else + GREEN='' + RED='' + BLUE='' + NC='' +fi + +log_info() { echo -e "${BLUE}==> $*${NC}"; } +log_ok() { echo -e "${GREEN}==> $*${NC}"; } +log_err() { echo -e "${RED}==> $*${NC}" >&2; } + +# Wait for a database to accept connections +wait_for_ready() { + local connstr="$1" + local label="$2" + local max_attempts=30 + local attempt=0 + + log_info "Waiting for $label to be ready..." + while ! psql "$connstr" -c "SELECT 1" > /dev/null 2>&1; do + attempt=$((attempt + 1)) + if [ "$attempt" -ge "$max_attempts" ]; then + log_err "$label did not become ready after $max_attempts attempts" + return 1 + fi + sleep 1 + done + log_ok "$label is ready" +} + +# Directory containing this script (used to locate testsetup/) +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +CERTS_DIR="$SCRIPT_DIR/testsetup/certs" + +# Copy client certificates to /tmp for TLS tests +setup_client_certs() { + if [ -d "$CERTS_DIR" ]; then + base64 -d "$CERTS_DIR/ca.pem.b64" > /tmp/ca.pem + base64 -d "$CERTS_DIR/pgx_sslcert.crt.b64" > /tmp/pgx_sslcert.crt + base64 -d "$CERTS_DIR/pgx_sslcert.key.b64" > /tmp/pgx_sslcert.key + fi +} + +# Initialize CockroachDB (create database if not exists) +init_crdb() { + local connstr="postgresql://root@localhost:26257/?sslmode=disable" + wait_for_ready "$connstr" "CockroachDB" + log_info "Ensuring pgx_test database exists on CockroachDB..." + psql "$connstr" -c "CREATE DATABASE IF NOT EXISTS pgx_test" 2>/dev/null || true +} + +# Run tests against a single target +run_tests() { + local target="$1" + shift + local extra_args=("$@") + + local label="" + local port="" + + case "$target" in + pg14) label="PostgreSQL 14"; port=5414 ;; + pg15) label="PostgreSQL 15"; port=5415 ;; + pg16) label="PostgreSQL 16"; port=5416 ;; + pg17) label="PostgreSQL 17"; port=5417 ;; + pg18) label="PostgreSQL 18"; port=5432 ;; + crdb) + label="CockroachDB (port 26257)" + init_crdb + log_info "Testing against $label" + if ! PGX_TEST_DATABASE="postgresql://root@localhost:26257/pgx_test?sslmode=disable&experimental_enable_temp_tables=on" \ + go test -count=1 "${extra_args[@]}" ./...; then + log_err "Tests FAILED against $label" + return 1 + fi + log_ok "Tests passed against $label" + return 0 + ;; + *) + log_err "Unknown target: $target" + log_err "Valid targets: pg14, pg15, pg16, pg17, pg18, crdb, all" + return 1 + ;; + esac + + setup_client_certs + + log_info "Testing against $label (port $port)" + if ! PGX_TEST_DATABASE="host=localhost port=$port user=postgres password=postgres dbname=pgx_test" \ + PGX_TEST_UNIX_SOCKET_CONN_STRING="host=/var/run/postgresql port=$port user=postgres dbname=pgx_test" \ + PGX_TEST_TCP_CONN_STRING="host=127.0.0.1 port=$port user=pgx_md5 password=secret dbname=pgx_test" \ + PGX_TEST_MD5_PASSWORD_CONN_STRING="host=127.0.0.1 port=$port user=pgx_md5 password=secret dbname=pgx_test" \ + PGX_TEST_SCRAM_PASSWORD_CONN_STRING="host=127.0.0.1 port=$port user=pgx_scram password=secret dbname=pgx_test channel_binding=disable" \ + PGX_TEST_SCRAM_PLUS_CONN_STRING="host=localhost port=$port user=pgx_ssl password=secret sslmode=verify-full sslrootcert=/tmp/ca.pem dbname=pgx_test channel_binding=require" \ + PGX_TEST_PLAIN_PASSWORD_CONN_STRING="host=127.0.0.1 port=$port user=pgx_pw password=secret dbname=pgx_test" \ + PGX_TEST_TLS_CONN_STRING="host=localhost port=$port user=pgx_ssl password=secret sslmode=verify-full sslrootcert=/tmp/ca.pem dbname=pgx_test channel_binding=disable" \ + PGX_TEST_TLS_CLIENT_CONN_STRING="host=localhost port=$port user=pgx_sslcert sslmode=verify-full sslrootcert=/tmp/ca.pem sslcert=/tmp/pgx_sslcert.crt sslkey=/tmp/pgx_sslcert.key dbname=pgx_test" \ + PGX_SSL_PASSWORD=certpw \ + go test -count=1 "${extra_args[@]}" ./...; then + log_err "Tests FAILED against $label" + return 1 + fi + log_ok "Tests passed against $label" +} + +# Main +main() { + local target="${1:-pg18}" + + if [ "$target" = "all" ]; then + shift || true + local targets=(pg14 pg15 pg16 pg17 pg18 crdb) + local failed=() + + for t in "${targets[@]}"; do + echo "" + log_info "==========================================" + log_info "Target: $t" + log_info "==========================================" + if ! run_tests "$t" "$@"; then + failed+=("$t") + log_err "FAILED: $t" + fi + done + + echo "" + if [ ${#failed[@]} -gt 0 ]; then + log_err "Failed targets: ${failed[*]}" + return 1 + else + log_ok "All targets passed" + fi + else + shift || true + run_tests "$target" "$@" + fi +} + +main "$@" diff --git a/vendor/github.com/jackc/pgx/v5/tracer.go b/vendor/github.com/jackc/pgx/v5/tracer.go new file mode 100644 index 0000000000..58ca99f7e0 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/tracer.go @@ -0,0 +1,107 @@ +package pgx + +import ( + "context" + + "github.com/jackc/pgx/v5/pgconn" +) + +// 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 *Conn, data TraceQueryStartData) context.Context + + TraceQueryEnd(ctx context.Context, conn *Conn, data TraceQueryEndData) +} + +type TraceQueryStartData struct { + SQL string + Args []any +} + +type TraceQueryEndData struct { + CommandTag pgconn.CommandTag + Err error +} + +// BatchTracer traces SendBatch. +type BatchTracer interface { + // TraceBatchStart is called at the beginning of SendBatch calls. The returned context is used for the + // rest of the call and will be passed to TraceBatchQuery and TraceBatchEnd. + TraceBatchStart(ctx context.Context, conn *Conn, data TraceBatchStartData) context.Context + + TraceBatchQuery(ctx context.Context, conn *Conn, data TraceBatchQueryData) + TraceBatchEnd(ctx context.Context, conn *Conn, data TraceBatchEndData) +} + +type TraceBatchStartData struct { + Batch *Batch +} + +type TraceBatchQueryData struct { + SQL string + Args []any + CommandTag pgconn.CommandTag + Err error +} + +type TraceBatchEndData struct { + Err error +} + +// CopyFromTracer traces CopyFrom. +type CopyFromTracer interface { + // TraceCopyFromStart is called at the beginning of CopyFrom calls. The returned context is used for the + // rest of the call and will be passed to TraceCopyFromEnd. + TraceCopyFromStart(ctx context.Context, conn *Conn, data TraceCopyFromStartData) context.Context + + TraceCopyFromEnd(ctx context.Context, conn *Conn, data TraceCopyFromEndData) +} + +type TraceCopyFromStartData struct { + TableName Identifier + ColumnNames []string +} + +type TraceCopyFromEndData struct { + CommandTag pgconn.CommandTag + Err error +} + +// PrepareTracer traces Prepare. +type PrepareTracer interface { + // TracePrepareStart is called at the beginning of Prepare calls. The returned context is used for the + // rest of the call and will be passed to TracePrepareEnd. + TracePrepareStart(ctx context.Context, conn *Conn, data TracePrepareStartData) context.Context + + TracePrepareEnd(ctx context.Context, conn *Conn, data TracePrepareEndData) +} + +type TracePrepareStartData struct { + Name string + SQL string +} + +type TracePrepareEndData struct { + AlreadyPrepared bool + Err error +} + +// ConnectTracer traces Connect and ConnectConfig. +type ConnectTracer interface { + // TraceConnectStart is called at the beginning of Connect and ConnectConfig calls. The returned context is used for + // the rest of the call and will be passed to TraceConnectEnd. + TraceConnectStart(ctx context.Context, data TraceConnectStartData) context.Context + + TraceConnectEnd(ctx context.Context, data TraceConnectEndData) +} + +type TraceConnectStartData struct { + ConnConfig *ConnConfig +} + +type TraceConnectEndData struct { + Conn *Conn + Err error +} diff --git a/vendor/github.com/jackc/pgx/v5/tx.go b/vendor/github.com/jackc/pgx/v5/tx.go new file mode 100644 index 0000000000..3f93a6f247 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/tx.go @@ -0,0 +1,442 @@ +package pgx + +import ( + "context" + "errors" + "strconv" + "strings" + + "github.com/jackc/pgx/v5/pgconn" +) + +// TxIsoLevel is the transaction isolation level (serializable, repeatable read, read committed or read uncommitted) +type TxIsoLevel string + +// Transaction isolation levels +const ( + Serializable TxIsoLevel = "serializable" + RepeatableRead TxIsoLevel = "repeatable read" + ReadCommitted TxIsoLevel = "read committed" + ReadUncommitted TxIsoLevel = "read uncommitted" +) + +// TxAccessMode is the transaction access mode (read write or read only) +type TxAccessMode string + +// Transaction access modes +const ( + ReadWrite TxAccessMode = "read write" + ReadOnly TxAccessMode = "read only" +) + +// TxDeferrableMode is the transaction deferrable mode (deferrable or not deferrable) +type TxDeferrableMode string + +// Transaction deferrable modes +const ( + Deferrable TxDeferrableMode = "deferrable" + NotDeferrable TxDeferrableMode = "not deferrable" +) + +// TxOptions are transaction modes within a transaction block +type TxOptions struct { + IsoLevel TxIsoLevel + AccessMode TxAccessMode + DeferrableMode TxDeferrableMode + + // BeginQuery is the SQL query that will be executed to begin the transaction. This allows using non-standard syntax + // such as BEGIN PRIORITY HIGH with CockroachDB. If set this will override the other settings. + BeginQuery string + // CommitQuery is the SQL query that will be executed to commit the transaction. + CommitQuery string +} + +var emptyTxOptions TxOptions + +func (txOptions TxOptions) beginSQL() string { + if txOptions == emptyTxOptions { + return "begin" + } + + if txOptions.BeginQuery != "" { + return txOptions.BeginQuery + } + + var buf strings.Builder + buf.Grow(64) // 64 - maximum length of string with available options + buf.WriteString("begin") + + if txOptions.IsoLevel != "" { + buf.WriteString(" isolation level ") + buf.WriteString(string(txOptions.IsoLevel)) + } + if txOptions.AccessMode != "" { + buf.WriteByte(' ') + buf.WriteString(string(txOptions.AccessMode)) + } + if txOptions.DeferrableMode != "" { + buf.WriteByte(' ') + buf.WriteString(string(txOptions.DeferrableMode)) + } + + return buf.String() +} + +var ErrTxClosed = errors.New("tx is closed") + +// ErrTxCommitRollback occurs when an error has occurred in a transaction and +// Commit() is called. PostgreSQL accepts COMMIT on aborted transactions, but +// it is treated as ROLLBACK. +var ErrTxCommitRollback = errors.New("commit unexpectedly resulted in rollback") + +// Begin starts a transaction. Unlike [database/sql], the context only affects the begin command. i.e. there is no +// auto-rollback on context cancellation. +func (c *Conn) Begin(ctx context.Context) (Tx, error) { + return c.BeginTx(ctx, TxOptions{}) +} + +// BeginTx starts a transaction with txOptions determining the transaction mode. Unlike [database/sql], the context only +// affects the begin command. i.e. there is no auto-rollback on context cancellation. +func (c *Conn) BeginTx(ctx context.Context, txOptions TxOptions) (Tx, error) { + _, err := c.Exec(ctx, txOptions.beginSQL()) + if err != nil { + // begin should never fail unless there is an underlying connection issue or + // a context timeout. In either case, the connection is possibly broken. + c.die() + return nil, err + } + + return &dbTx{ + conn: c, + commitQuery: txOptions.CommitQuery, + }, nil +} + +// Tx represents a database transaction. +// +// Tx is an interface instead of a struct to enable connection pools to be implemented without relying on internal pgx +// state, to support pseudo-nested transactions with savepoints, and to allow tests to mock transactions. However, +// adding a method to an interface is technically a breaking change. If new methods are added to Conn it may be +// desirable to add them to Tx as well. Because of this the Tx interface is partially excluded from semantic version +// requirements. Methods will not be removed or changed, but new methods may be added. +type Tx interface { + // Begin starts a pseudo nested transaction. + Begin(ctx context.Context) (Tx, error) + + // Commit commits the transaction if this is a real transaction or releases the savepoint if this is a pseudo nested + // transaction. Commit will return an error where errors.Is(ErrTxClosed) is true if the Tx is already closed, but is + // otherwise safe to call multiple times. If the commit fails with a rollback status (e.g. the transaction was already + // in a broken state) then an error where errors.Is(ErrTxCommitRollback) is true will be returned. + Commit(ctx context.Context) error + + // Rollback rolls back the transaction if this is a real transaction or rolls back to the savepoint if this is a + // pseudo nested transaction. Rollback will return an error where errors.Is(ErrTxClosed) is true if the Tx is already + // closed, but is otherwise safe to call multiple times. Hence, a defer tx.Rollback() is safe even if tx.Commit() will + // be called first in a non-error condition. Any other failure of a real transaction will result in the connection + // being closed. + Rollback(ctx context.Context) error + + CopyFrom(ctx context.Context, tableName Identifier, columnNames []string, rowSrc CopyFromSource) (int64, error) + SendBatch(ctx context.Context, b *Batch) BatchResults + LargeObjects() LargeObjects + + Prepare(ctx context.Context, name, sql string) (*pgconn.StatementDescription, error) + + Exec(ctx context.Context, sql string, arguments ...any) (commandTag pgconn.CommandTag, err error) + Query(ctx context.Context, sql string, args ...any) (Rows, error) + QueryRow(ctx context.Context, sql string, args ...any) Row + + // Conn returns the underlying *Conn that on which this transaction is executing. + Conn() *Conn +} + +// dbTx represents a database transaction. +// +// All dbTx methods return ErrTxClosed if Commit or Rollback has already been +// called on the dbTx. +type dbTx struct { + conn *Conn + savepointNum int64 + closed bool + commitQuery string +} + +// Begin starts a pseudo nested transaction implemented with a savepoint. +func (tx *dbTx) Begin(ctx context.Context) (Tx, error) { + if tx.closed { + return nil, ErrTxClosed + } + + tx.savepointNum++ + _, err := tx.conn.Exec(ctx, "savepoint sp_"+strconv.FormatInt(tx.savepointNum, 10)) + if err != nil { + return nil, err + } + + return &dbSimulatedNestedTx{tx: tx, savepointNum: tx.savepointNum}, nil +} + +// Commit commits the transaction. +func (tx *dbTx) Commit(ctx context.Context) error { + if tx.closed { + return ErrTxClosed + } + + commandSQL := "commit" + if tx.commitQuery != "" { + commandSQL = tx.commitQuery + } + + commandTag, err := tx.conn.Exec(ctx, commandSQL) + tx.closed = true + if err != nil { + if tx.conn.PgConn().TxStatus() != 'I' { + _ = tx.conn.Close(ctx) // already have error to return + } + return err + } + if commandTag.String() == "ROLLBACK" { + return ErrTxCommitRollback + } + + return nil +} + +// Rollback rolls back the transaction. Rollback will return ErrTxClosed if the +// Tx is already closed, but is otherwise safe to call multiple times. Hence, a +// defer tx.Rollback() is safe even if tx.Commit() will be called first in a +// non-error condition. +func (tx *dbTx) Rollback(ctx context.Context) error { + if tx.closed { + return ErrTxClosed + } + + _, err := tx.conn.Exec(ctx, "rollback") + tx.closed = true + if err != nil { + // A rollback failure leaves the connection in an undefined state + tx.conn.die() + return err + } + + return nil +} + +// Exec delegates to the underlying *Conn +func (tx *dbTx) Exec(ctx context.Context, sql string, arguments ...any) (commandTag pgconn.CommandTag, err error) { + if tx.closed { + return pgconn.CommandTag{}, ErrTxClosed + } + + return tx.conn.Exec(ctx, sql, arguments...) +} + +// Prepare delegates to the underlying *Conn +func (tx *dbTx) Prepare(ctx context.Context, name, sql string) (*pgconn.StatementDescription, error) { + if tx.closed { + return nil, ErrTxClosed + } + + return tx.conn.Prepare(ctx, name, sql) +} + +// Query delegates to the underlying *Conn +func (tx *dbTx) Query(ctx context.Context, sql string, args ...any) (Rows, error) { + if tx.closed { + // Because checking for errors can be deferred to the *Rows, build one with the error + err := ErrTxClosed + return &baseRows{closed: true, err: err}, err + } + + return tx.conn.Query(ctx, sql, args...) +} + +// QueryRow delegates to the underlying *Conn +func (tx *dbTx) QueryRow(ctx context.Context, sql string, args ...any) Row { + rows, _ := tx.Query(ctx, sql, args...) + return (*connRow)(rows.(*baseRows)) +} + +// CopyFrom delegates to the underlying *Conn +func (tx *dbTx) CopyFrom(ctx context.Context, tableName Identifier, columnNames []string, rowSrc CopyFromSource) (int64, error) { + if tx.closed { + return 0, ErrTxClosed + } + + return tx.conn.CopyFrom(ctx, tableName, columnNames, rowSrc) +} + +// SendBatch delegates to the underlying *Conn +func (tx *dbTx) SendBatch(ctx context.Context, b *Batch) BatchResults { + if tx.closed { + return &batchResults{err: ErrTxClosed} + } + + return tx.conn.SendBatch(ctx, b) +} + +// LargeObjects returns a LargeObjects instance for the transaction. +func (tx *dbTx) LargeObjects() LargeObjects { + return LargeObjects{tx: tx} +} + +func (tx *dbTx) Conn() *Conn { + return tx.conn +} + +// dbSimulatedNestedTx represents a simulated nested transaction implemented by a savepoint. +type dbSimulatedNestedTx struct { + tx Tx + savepointNum int64 + closed bool +} + +// Begin starts a pseudo nested transaction implemented with a savepoint. +func (sp *dbSimulatedNestedTx) Begin(ctx context.Context) (Tx, error) { + if sp.closed { + return nil, ErrTxClosed + } + + return sp.tx.Begin(ctx) +} + +// Commit releases the savepoint essentially committing the pseudo nested transaction. +func (sp *dbSimulatedNestedTx) Commit(ctx context.Context) error { + if sp.closed { + return ErrTxClosed + } + + _, err := sp.Exec(ctx, "release savepoint sp_"+strconv.FormatInt(sp.savepointNum, 10)) + sp.closed = true + return err +} + +// Rollback rolls back to the savepoint essentially rolling back the pseudo nested transaction. Rollback will return +// ErrTxClosed if the dbSavepoint is already closed, but is otherwise safe to call multiple times. Hence, a defer sp.Rollback() +// is safe even if sp.Commit() will be called first in a non-error condition. +func (sp *dbSimulatedNestedTx) Rollback(ctx context.Context) error { + if sp.closed { + return ErrTxClosed + } + + _, err := sp.Exec(ctx, "rollback to savepoint sp_"+strconv.FormatInt(sp.savepointNum, 10)) + sp.closed = true + return err +} + +// Exec delegates to the underlying Tx +func (sp *dbSimulatedNestedTx) Exec(ctx context.Context, sql string, arguments ...any) (commandTag pgconn.CommandTag, err error) { + if sp.closed { + return pgconn.CommandTag{}, ErrTxClosed + } + + return sp.tx.Exec(ctx, sql, arguments...) +} + +// Prepare delegates to the underlying Tx +func (sp *dbSimulatedNestedTx) Prepare(ctx context.Context, name, sql string) (*pgconn.StatementDescription, error) { + if sp.closed { + return nil, ErrTxClosed + } + + return sp.tx.Prepare(ctx, name, sql) +} + +// Query delegates to the underlying Tx +func (sp *dbSimulatedNestedTx) Query(ctx context.Context, sql string, args ...any) (Rows, error) { + if sp.closed { + // Because checking for errors can be deferred to the *Rows, build one with the error + err := ErrTxClosed + return &baseRows{closed: true, err: err}, err + } + + return sp.tx.Query(ctx, sql, args...) +} + +// QueryRow delegates to the underlying Tx +func (sp *dbSimulatedNestedTx) QueryRow(ctx context.Context, sql string, args ...any) Row { + rows, _ := sp.Query(ctx, sql, args...) + return (*connRow)(rows.(*baseRows)) +} + +// CopyFrom delegates to the underlying *Conn +func (sp *dbSimulatedNestedTx) CopyFrom(ctx context.Context, tableName Identifier, columnNames []string, rowSrc CopyFromSource) (int64, error) { + if sp.closed { + return 0, ErrTxClosed + } + + return sp.tx.CopyFrom(ctx, tableName, columnNames, rowSrc) +} + +// SendBatch delegates to the underlying *Conn +func (sp *dbSimulatedNestedTx) SendBatch(ctx context.Context, b *Batch) BatchResults { + if sp.closed { + return &batchResults{err: ErrTxClosed} + } + + return sp.tx.SendBatch(ctx, b) +} + +func (sp *dbSimulatedNestedTx) LargeObjects() LargeObjects { + return LargeObjects{tx: sp} +} + +func (sp *dbSimulatedNestedTx) Conn() *Conn { + return sp.tx.Conn() +} + +// BeginFunc calls Begin on db and then calls fn. If fn does not return an error then it calls [Tx.Commit] on db. If fn +// returns an error it calls [Tx.Rollback] on db. The context will be used when executing the transaction control statements +// (BEGIN, ROLLBACK, and COMMIT) but does not otherwise affect the execution of fn. +func BeginFunc( + ctx context.Context, + db interface { + Begin(ctx context.Context) (Tx, error) + }, + fn func(Tx) error, +) (err error) { + var tx Tx + tx, err = db.Begin(ctx) + if err != nil { + return err + } + + return beginFuncExec(ctx, tx, fn) +} + +// BeginTxFunc calls BeginTx on db and then calls fn. If fn does not return an error then it calls [Tx.Commit] on db. If fn +// returns an error it calls [Tx.Rollback] on db. The context will be used when executing the transaction control statements +// (BEGIN, ROLLBACK, and COMMIT) but does not otherwise affect the execution of fn. +func BeginTxFunc( + ctx context.Context, + db interface { + BeginTx(ctx context.Context, txOptions TxOptions) (Tx, error) + }, + txOptions TxOptions, + fn func(Tx) error, +) (err error) { + var tx Tx + tx, err = db.BeginTx(ctx, txOptions) + if err != nil { + return err + } + + return beginFuncExec(ctx, tx, fn) +} + +func beginFuncExec(ctx context.Context, tx Tx, fn func(Tx) error) (err error) { + defer func() { + rollbackErr := tx.Rollback(ctx) + if rollbackErr != nil && !errors.Is(rollbackErr, ErrTxClosed) { + err = rollbackErr + } + }() + + fErr := fn(tx) + if fErr != nil { + _ = tx.Rollback(ctx) // ignore rollback error as there is already an error to return + return fErr + } + + return tx.Commit(ctx) +} diff --git a/vendor/github.com/jackc/pgx/v5/values.go b/vendor/github.com/jackc/pgx/v5/values.go new file mode 100644 index 0000000000..6e2ff30030 --- /dev/null +++ b/vendor/github.com/jackc/pgx/v5/values.go @@ -0,0 +1,63 @@ +package pgx + +import ( + "errors" + + "github.com/jackc/pgx/v5/internal/pgio" + "github.com/jackc/pgx/v5/pgtype" +) + +// PostgreSQL format codes +const ( + TextFormatCode = 0 + BinaryFormatCode = 1 +) + +func convertSimpleArgument(m *pgtype.Map, arg any) (any, error) { + buf, err := m.Encode(0, TextFormatCode, arg, []byte{}) + if err != nil { + return nil, err + } + if buf == nil { + return nil, nil + } + return string(buf), nil +} + +func encodeCopyValue(m *pgtype.Map, buf []byte, oid uint32, arg any) ([]byte, error) { + sp := len(buf) + buf = pgio.AppendInt32(buf, -1) + argBuf, err := m.Encode(oid, BinaryFormatCode, arg, buf) + if err != nil { + if argBuf2, err2 := tryScanStringCopyValueThenEncode(m, buf, oid, arg); err2 == nil { + argBuf = argBuf2 + } else { + return nil, err + } + } + + if argBuf != nil { + buf = argBuf + pgio.SetInt32(buf[sp:], int32(len(buf[sp:])-4)) + } + return buf, nil +} + +func tryScanStringCopyValueThenEncode(m *pgtype.Map, buf []byte, oid uint32, arg any) ([]byte, error) { + s, ok := arg.(string) + if !ok { + textBuf, err := m.Encode(oid, TextFormatCode, arg, nil) + if err != nil { + return nil, errors.New("not a string and cannot be encoded as text") + } + s = string(textBuf) + } + + var v any + err := m.Scan(oid, TextFormatCode, []byte(s), &v) + if err != nil { + return nil, err + } + + return m.Encode(oid, BinaryFormatCode, v, buf) +} diff --git a/vendor/github.com/jackc/puddle/v2/CHANGELOG.md b/vendor/github.com/jackc/puddle/v2/CHANGELOG.md new file mode 100644 index 0000000000..d0d202c74a --- /dev/null +++ b/vendor/github.com/jackc/puddle/v2/CHANGELOG.md @@ -0,0 +1,79 @@ +# 2.2.2 (September 10, 2024) + +* Add empty acquire time to stats (Maxim Ivanov) +* Stop importing nanotime from runtime via linkname (maypok86) + +# 2.2.1 (July 15, 2023) + +* Fix: CreateResource cannot overflow pool. This changes documented behavior of CreateResource. Previously, + CreateResource could create a resource even if the pool was full. This could cause the pool to overflow. While this + was documented, it was documenting incorrect behavior. CreateResource now returns an error if the pool is full. + +# 2.2.0 (February 11, 2023) + +* Use Go 1.19 atomics and drop go.uber.org/atomic dependency + +# 2.1.2 (November 12, 2022) + +* Restore support to Go 1.18 via go.uber.org/atomic + +# 2.1.1 (November 11, 2022) + +* Fix create resource concurrently with Stat call race + +# 2.1.0 (October 28, 2022) + +* Concurrency control is now implemented with a semaphore. This simplifies some internal logic, resolves a few error conditions (including a deadlock), and improves performance. (Jan Dubsky) +* Go 1.19 is now required for the improved atomic support. + +# 2.0.1 (October 28, 2022) + +* Fix race condition when Close is called concurrently with multiple constructors + +# 2.0.0 (September 17, 2022) + +* Use generics instead of interface{} (Столяров Владимир Алексеевич) +* Add Reset +* Do not cancel resource construction when Acquire is canceled +* NewPool takes Config + +# 1.3.0 (August 27, 2022) + +* Acquire creates resources in background to allow creation to continue after Acquire is canceled (James Hartig) + +# 1.2.1 (December 2, 2021) + +* TryAcquire now does not block when background constructing resource + +# 1.2.0 (November 20, 2021) + +* Add TryAcquire (A. Jensen) +* Fix: remove memory leak / unintentionally pinned memory when shrinking slices (Alexander Staubo) +* Fix: Do not leave pool locked after panic from nil context + +# 1.1.4 (September 11, 2021) + +* Fix: Deadlock in CreateResource if pool was closed during resource acquisition (Dmitriy Matrenichev) + +# 1.1.3 (December 3, 2020) + +* Fix: Failed resource creation could cause concurrent Acquire to hang. (Evgeny Vanslov) + +# 1.1.2 (September 26, 2020) + +* Fix: Resource.Destroy no longer removes itself from the pool before its destructor has completed. +* Fix: Prevent crash when pool is closed while resource is being created. + +# 1.1.1 (April 2, 2020) + +* Pool.Close can be safely called multiple times +* AcquireAllIDle immediately returns nil if pool is closed +* CreateResource checks if pool is closed before taking any action +* Fix potential race condition when CreateResource and Close are called concurrently. CreateResource now checks if pool is closed before adding newly created resource to pool. + +# 1.1.0 (February 5, 2020) + +* Use runtime.nanotime for faster tracking of acquire time and last usage time. +* Track resource idle time to enable client health check logic. (Patrick Ellul) +* Add CreateResource to construct a new resource without acquiring it. (Patrick Ellul) +* Fix deadlock race when acquire is cancelled. (Michael Tharp) diff --git a/vendor/github.com/jackc/puddle/v2/LICENSE b/vendor/github.com/jackc/puddle/v2/LICENSE new file mode 100644 index 0000000000..bcc286c54d --- /dev/null +++ b/vendor/github.com/jackc/puddle/v2/LICENSE @@ -0,0 +1,22 @@ +Copyright (c) 2018 Jack Christensen + +MIT License + +Permission is hereby granted, free of charge, to any person obtaining +a copy of this software and associated documentation files (the +"Software"), to deal in the Software without restriction, including +without limitation the rights to use, copy, modify, merge, publish, +distribute, sublicense, and/or sell copies of the Software, and to +permit persons to whom the Software is furnished to do so, subject to +the following conditions: + +The above copyright notice and this permission notice shall be +included in all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, +EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF +MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND +NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE +LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION +OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION +WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. diff --git a/vendor/github.com/jackc/puddle/v2/README.md b/vendor/github.com/jackc/puddle/v2/README.md new file mode 100644 index 0000000000..fa82a9d46f --- /dev/null +++ b/vendor/github.com/jackc/puddle/v2/README.md @@ -0,0 +1,80 @@ +[![Go Reference](https://pkg.go.dev/badge/github.com/jackc/puddle/v2.svg)](https://pkg.go.dev/github.com/jackc/puddle/v2) +![Build Status](https://github.com/jackc/puddle/actions/workflows/ci.yml/badge.svg) + +# Puddle + +Puddle is a tiny generic resource pool library for Go that uses the standard +context library to signal cancellation of acquires. It is designed to contain +the minimum functionality required for a resource pool. It can be used directly +or it can be used as the base for a domain specific resource pool. For example, +a database connection pool may use puddle internally and implement health checks +and keep-alive behavior without needing to implement any concurrent code of its +own. + +## Features + +* Acquire cancellation via context standard library +* Statistics API for monitoring pool pressure +* No dependencies outside of standard library and golang.org/x/sync +* High performance +* 100% test coverage of reachable code + +## Example Usage + +```go +package main + +import ( + "context" + "log" + "net" + + "github.com/jackc/puddle/v2" +) + +func main() { + constructor := func(context.Context) (net.Conn, error) { + return net.Dial("tcp", "127.0.0.1:8080") + } + destructor := func(value net.Conn) { + value.Close() + } + maxPoolSize := int32(10) + + pool, err := puddle.NewPool(&puddle.Config[net.Conn]{Constructor: constructor, Destructor: destructor, MaxSize: maxPoolSize}) + if err != nil { + log.Fatal(err) + } + + // Acquire resource from the pool. + res, err := pool.Acquire(context.Background()) + if err != nil { + log.Fatal(err) + } + + // Use resource. + _, err = res.Value().Write([]byte{1}) + if err != nil { + log.Fatal(err) + } + + // Release when done. + res.Release() +} +``` + +## Status + +Puddle is stable and feature complete. + +* Bug reports and fixes are welcome. +* New features will usually not be accepted if they can be feasibly implemented in a wrapper. +* Performance optimizations will usually not be accepted unless the performance issue rises to the level of a bug. + +## Supported Go Versions + +puddle supports the same versions of Go that are supported by the Go project. For [Go](https://golang.org/doc/devel/release.html#policy) that is the two most recent major releases. This means puddle supports Go 1.19 and higher. + +## License + +MIT diff --git a/vendor/github.com/jackc/puddle/v2/context.go b/vendor/github.com/jackc/puddle/v2/context.go new file mode 100644 index 0000000000..e19d2a609b --- /dev/null +++ b/vendor/github.com/jackc/puddle/v2/context.go @@ -0,0 +1,24 @@ +package puddle + +import ( + "context" + "time" +) + +// valueCancelCtx combines two contexts into one. One context is used for values and the other is used for cancellation. +type valueCancelCtx struct { + valueCtx context.Context + cancelCtx context.Context +} + +func (ctx *valueCancelCtx) Deadline() (time.Time, bool) { return ctx.cancelCtx.Deadline() } +func (ctx *valueCancelCtx) Done() <-chan struct{} { return ctx.cancelCtx.Done() } +func (ctx *valueCancelCtx) Err() error { return ctx.cancelCtx.Err() } +func (ctx *valueCancelCtx) Value(key any) any { return ctx.valueCtx.Value(key) } + +func newValueCancelCtx(valueCtx, cancelContext context.Context) context.Context { + return &valueCancelCtx{ + valueCtx: valueCtx, + cancelCtx: cancelContext, + } +} diff --git a/vendor/github.com/jackc/puddle/v2/doc.go b/vendor/github.com/jackc/puddle/v2/doc.go new file mode 100644 index 0000000000..818e4a6982 --- /dev/null +++ b/vendor/github.com/jackc/puddle/v2/doc.go @@ -0,0 +1,11 @@ +// Package puddle is a generic resource pool with type-parametrized api. +/* + +Puddle is a tiny generic resource pool library for Go that uses the standard +context library to signal cancellation of acquires. It is designed to contain +the minimum functionality a resource pool needs that cannot be implemented +without concurrency concerns. For example, a database connection pool may use +puddle internally and implement health checks and keep-alive behavior without +needing to implement any concurrent code of its own. +*/ +package puddle diff --git a/vendor/github.com/jackc/puddle/v2/internal/genstack/gen_stack.go b/vendor/github.com/jackc/puddle/v2/internal/genstack/gen_stack.go new file mode 100644 index 0000000000..7e4660c8c0 --- /dev/null +++ b/vendor/github.com/jackc/puddle/v2/internal/genstack/gen_stack.go @@ -0,0 +1,85 @@ +package genstack + +// GenStack implements a generational stack. +// +// GenStack works as common stack except for the fact that all elements in the +// older generation are guaranteed to be popped before any element in the newer +// generation. New elements are always pushed to the current (newest) +// generation. +// +// We could also say that GenStack behaves as a stack in case of a single +// generation, but it behaves as a queue of individual generation stacks. +type GenStack[T any] struct { + // We can represent arbitrary number of generations using 2 stacks. The + // new stack stores all new pushes and the old stack serves all reads. + // Old stack can represent multiple generations. If old == new, then all + // elements pushed in previous (not current) generations have already + // been popped. + + old *stack[T] + new *stack[T] +} + +// NewGenStack creates a new empty GenStack. +func NewGenStack[T any]() *GenStack[T] { + s := &stack[T]{} + return &GenStack[T]{ + old: s, + new: s, + } +} + +func (s *GenStack[T]) Pop() (T, bool) { + // Pushes always append to the new stack, so if the old once becomes + // empty, it will remail empty forever. + if s.old.len() == 0 && s.old != s.new { + s.old = s.new + } + + if s.old.len() == 0 { + var zero T + return zero, false + } + + return s.old.pop(), true +} + +// Push pushes a new element at the top of the stack. +func (s *GenStack[T]) Push(v T) { s.new.push(v) } + +// NextGen starts a new stack generation. +func (s *GenStack[T]) NextGen() { + if s.old == s.new { + s.new = &stack[T]{} + return + } + + // We need to pop from the old stack to the top of the new stack. Let's + // have an example: + // + // Old: 4 3 2 1 + // New: 8 7 6 5 + // PopOrder: 1 2 3 4 5 6 7 8 + // + // + // To preserve pop order, we have to take all elements from the old + // stack and push them to the top of new stack: + // + // New: 8 7 6 5 4 3 2 1 + // + s.new.push(s.old.takeAll()...) + + // We have the old stack allocated and empty, so why not to reuse it as + // new new stack. + s.old, s.new = s.new, s.old +} + +// Len returns number of elements in the stack. +func (s *GenStack[T]) Len() int { + l := s.old.len() + if s.old != s.new { + l += s.new.len() + } + + return l +} diff --git a/vendor/github.com/jackc/puddle/v2/internal/genstack/stack.go b/vendor/github.com/jackc/puddle/v2/internal/genstack/stack.go new file mode 100644 index 0000000000..dbced0c724 --- /dev/null +++ b/vendor/github.com/jackc/puddle/v2/internal/genstack/stack.go @@ -0,0 +1,39 @@ +package genstack + +// stack is a wrapper around an array implementing a stack. +// +// We cannot use slice to represent the stack because append might change the +// pointer value of the slice. That would be an issue in GenStack +// implementation. +type stack[T any] struct { + arr []T +} + +// push pushes a new element at the top of a stack. +func (s *stack[T]) push(vs ...T) { s.arr = append(s.arr, vs...) } + +// pop pops the stack top-most element. +// +// If stack length is zero, this method panics. +func (s *stack[T]) pop() T { + idx := s.len() - 1 + val := s.arr[idx] + + // Avoid memory leak + var zero T + s.arr[idx] = zero + + s.arr = s.arr[:idx] + return val +} + +// takeAll returns all elements in the stack in order as they are stored - i.e. +// the top-most stack element is the last one. +func (s *stack[T]) takeAll() []T { + arr := s.arr + s.arr = nil + return arr +} + +// len returns number of elements in the stack. +func (s *stack[T]) len() int { return len(s.arr) } diff --git a/vendor/github.com/jackc/puddle/v2/log.go b/vendor/github.com/jackc/puddle/v2/log.go new file mode 100644 index 0000000000..b21b946305 --- /dev/null +++ b/vendor/github.com/jackc/puddle/v2/log.go @@ -0,0 +1,32 @@ +package puddle + +import "unsafe" + +type ints interface { + int | int8 | int16 | int32 | int64 | uint | uint8 | uint16 | uint32 | uint64 +} + +// log2Int returns log2 of an integer. This function panics if val < 0. For val +// == 0, returns 0. +func log2Int[T ints](val T) uint8 { + if val <= 0 { + panic("log2 of non-positive number does not exist") + } + + return log2IntRange(val, 0, uint8(8*unsafe.Sizeof(val))) +} + +func log2IntRange[T ints](val T, begin, end uint8) uint8 { + length := end - begin + if length == 1 { + return begin + } + + delim := begin + length/2 + mask := T(1) << delim + if mask > val { + return log2IntRange(val, begin, delim) + } else { + return log2IntRange(val, delim, end) + } +} diff --git a/vendor/github.com/jackc/puddle/v2/nanotime.go b/vendor/github.com/jackc/puddle/v2/nanotime.go new file mode 100644 index 0000000000..8a5351a0df --- /dev/null +++ b/vendor/github.com/jackc/puddle/v2/nanotime.go @@ -0,0 +1,16 @@ +package puddle + +import "time" + +// nanotime returns the time in nanoseconds since process start. +// +// This approach, described at +// https://github.com/golang/go/issues/61765#issuecomment-1672090302, +// is fast, monotonic, and portable, and avoids the previous +// dependence on runtime.nanotime using the (unsafe) linkname hack. +// In particular, time.Since does less work than time.Now. +func nanotime() int64 { + return time.Since(globalStart).Nanoseconds() +} + +var globalStart = time.Now() diff --git a/vendor/github.com/jackc/puddle/v2/pool.go b/vendor/github.com/jackc/puddle/v2/pool.go new file mode 100644 index 0000000000..c411d2f6ef --- /dev/null +++ b/vendor/github.com/jackc/puddle/v2/pool.go @@ -0,0 +1,710 @@ +package puddle + +import ( + "context" + "errors" + "sync" + "sync/atomic" + "time" + + "github.com/jackc/puddle/v2/internal/genstack" + "golang.org/x/sync/semaphore" +) + +const ( + resourceStatusConstructing = 0 + resourceStatusIdle = iota + resourceStatusAcquired = iota + resourceStatusHijacked = iota +) + +// ErrClosedPool occurs on an attempt to acquire a connection from a closed pool +// or a pool that is closed while the acquire is waiting. +var ErrClosedPool = errors.New("closed pool") + +// ErrNotAvailable occurs on an attempt to acquire a resource from a pool +// that is at maximum capacity and has no available resources. +var ErrNotAvailable = errors.New("resource not available") + +// Constructor is a function called by the pool to construct a resource. +type Constructor[T any] func(ctx context.Context) (res T, err error) + +// Destructor is a function called by the pool to destroy a resource. +type Destructor[T any] func(res T) + +// Resource is the resource handle returned by acquiring from the pool. +type Resource[T any] struct { + value T + pool *Pool[T] + creationTime time.Time + lastUsedNano int64 + poolResetCount int + status byte +} + +// Value returns the resource value. +func (res *Resource[T]) Value() T { + if !(res.status == resourceStatusAcquired || res.status == resourceStatusHijacked) { + panic("tried to access resource that is not acquired or hijacked") + } + return res.value +} + +// Release returns the resource to the pool. res must not be subsequently used. +func (res *Resource[T]) Release() { + if res.status != resourceStatusAcquired { + panic("tried to release resource that is not acquired") + } + res.pool.releaseAcquiredResource(res, nanotime()) +} + +// ReleaseUnused returns the resource to the pool without updating when it was last used used. i.e. LastUsedNanotime +// will not change. res must not be subsequently used. +func (res *Resource[T]) ReleaseUnused() { + if res.status != resourceStatusAcquired { + panic("tried to release resource that is not acquired") + } + res.pool.releaseAcquiredResource(res, res.lastUsedNano) +} + +// Destroy returns the resource to the pool for destruction. res must not be +// subsequently used. +func (res *Resource[T]) Destroy() { + if res.status != resourceStatusAcquired { + panic("tried to destroy resource that is not acquired") + } + go res.pool.destroyAcquiredResource(res) +} + +// Hijack assumes ownership of the resource from the pool. Caller is responsible +// for cleanup of resource value. +func (res *Resource[T]) Hijack() { + if res.status != resourceStatusAcquired { + panic("tried to hijack resource that is not acquired") + } + res.pool.hijackAcquiredResource(res) +} + +// CreationTime returns when the resource was created by the pool. +func (res *Resource[T]) CreationTime() time.Time { + if !(res.status == resourceStatusAcquired || res.status == resourceStatusHijacked) { + panic("tried to access resource that is not acquired or hijacked") + } + return res.creationTime +} + +// LastUsedNanotime returns when Release was last called on the resource measured in nanoseconds from an arbitrary time +// (a monotonic time). Returns creation time if Release has never been called. This is only useful to compare with +// other calls to LastUsedNanotime. In almost all cases, IdleDuration should be used instead. +func (res *Resource[T]) LastUsedNanotime() int64 { + if !(res.status == resourceStatusAcquired || res.status == resourceStatusHijacked) { + panic("tried to access resource that is not acquired or hijacked") + } + + return res.lastUsedNano +} + +// IdleDuration returns the duration since Release was last called on the resource. This is equivalent to subtracting +// LastUsedNanotime to the current nanotime. +func (res *Resource[T]) IdleDuration() time.Duration { + if !(res.status == resourceStatusAcquired || res.status == resourceStatusHijacked) { + panic("tried to access resource that is not acquired or hijacked") + } + + return time.Duration(nanotime() - res.lastUsedNano) +} + +// Pool is a concurrency-safe resource pool. +type Pool[T any] struct { + // mux is the pool internal lock. Any modification of shared state of + // the pool (but Acquires of acquireSem) must be performed only by + // holder of the lock. Long running operations are not allowed when mux + // is held. + mux sync.Mutex + // acquireSem provides an allowance to acquire a resource. + // + // Releases are allowed only when caller holds mux. Acquires have to + // happen before mux is locked (doesn't apply to semaphore.TryAcquire in + // AcquireAllIdle). + acquireSem *semaphore.Weighted + destructWG sync.WaitGroup + + allResources resList[T] + idleResources *genstack.GenStack[*Resource[T]] + + constructor Constructor[T] + destructor Destructor[T] + maxSize int32 + + acquireCount int64 + acquireDuration time.Duration + emptyAcquireCount int64 + emptyAcquireWaitTime time.Duration + canceledAcquireCount atomic.Int64 + + resetCount int + + baseAcquireCtx context.Context + cancelBaseAcquireCtx context.CancelFunc + closed bool +} + +type Config[T any] struct { + Constructor Constructor[T] + Destructor Destructor[T] + MaxSize int32 +} + +// NewPool creates a new pool. Returns an error iff MaxSize is less than 1. +func NewPool[T any](config *Config[T]) (*Pool[T], error) { + if config.MaxSize < 1 { + return nil, errors.New("MaxSize must be >= 1") + } + + baseAcquireCtx, cancelBaseAcquireCtx := context.WithCancel(context.Background()) + + return &Pool[T]{ + acquireSem: semaphore.NewWeighted(int64(config.MaxSize)), + idleResources: genstack.NewGenStack[*Resource[T]](), + maxSize: config.MaxSize, + constructor: config.Constructor, + destructor: config.Destructor, + baseAcquireCtx: baseAcquireCtx, + cancelBaseAcquireCtx: cancelBaseAcquireCtx, + }, nil +} + +// Close destroys all resources in the pool and rejects future Acquire calls. +// Blocks until all resources are returned to pool and destroyed. +func (p *Pool[T]) Close() { + defer p.destructWG.Wait() + + p.mux.Lock() + defer p.mux.Unlock() + + if p.closed { + return + } + p.closed = true + p.cancelBaseAcquireCtx() + + for res, ok := p.idleResources.Pop(); ok; res, ok = p.idleResources.Pop() { + p.allResources.remove(res) + go p.destructResourceValue(res.value) + } +} + +// Stat is a snapshot of Pool statistics. +type Stat struct { + constructingResources int32 + acquiredResources int32 + idleResources int32 + maxResources int32 + acquireCount int64 + acquireDuration time.Duration + emptyAcquireCount int64 + emptyAcquireWaitTime time.Duration + canceledAcquireCount int64 +} + +// TotalResources returns the total number of resources currently in the pool. +// The value is the sum of ConstructingResources, AcquiredResources, and +// IdleResources. +func (s *Stat) TotalResources() int32 { + return s.constructingResources + s.acquiredResources + s.idleResources +} + +// ConstructingResources returns the number of resources with construction in progress in +// the pool. +func (s *Stat) ConstructingResources() int32 { + return s.constructingResources +} + +// AcquiredResources returns the number of currently acquired resources in the pool. +func (s *Stat) AcquiredResources() int32 { + return s.acquiredResources +} + +// IdleResources returns the number of currently idle resources in the pool. +func (s *Stat) IdleResources() int32 { + return s.idleResources +} + +// MaxResources returns the maximum size of the pool. +func (s *Stat) MaxResources() int32 { + return s.maxResources +} + +// AcquireCount returns the cumulative count of successful acquires from the pool. +func (s *Stat) AcquireCount() int64 { + return s.acquireCount +} + +// AcquireDuration returns the total duration of all successful acquires from +// the pool. +func (s *Stat) AcquireDuration() time.Duration { + return s.acquireDuration +} + +// EmptyAcquireCount returns the cumulative count of successful acquires from the pool +// that waited for a resource to be released or constructed because the pool was +// empty. +func (s *Stat) EmptyAcquireCount() int64 { + return s.emptyAcquireCount +} + +// EmptyAcquireWaitTime returns the cumulative time waited for successful acquires +// from the pool for a resource to be released or constructed because the pool was +// empty. +func (s *Stat) EmptyAcquireWaitTime() time.Duration { + return s.emptyAcquireWaitTime +} + +// CanceledAcquireCount returns the cumulative count of acquires from the pool +// that were canceled by a context. +func (s *Stat) CanceledAcquireCount() int64 { + return s.canceledAcquireCount +} + +// Stat returns the current pool statistics. +func (p *Pool[T]) Stat() *Stat { + p.mux.Lock() + defer p.mux.Unlock() + + s := &Stat{ + maxResources: p.maxSize, + acquireCount: p.acquireCount, + emptyAcquireCount: p.emptyAcquireCount, + emptyAcquireWaitTime: p.emptyAcquireWaitTime, + canceledAcquireCount: p.canceledAcquireCount.Load(), + acquireDuration: p.acquireDuration, + } + + for _, res := range p.allResources { + switch res.status { + case resourceStatusConstructing: + s.constructingResources += 1 + case resourceStatusIdle: + s.idleResources += 1 + case resourceStatusAcquired: + s.acquiredResources += 1 + } + } + + return s +} + +// tryAcquireIdleResource checks if there is any idle resource. If there is +// some, this method removes it from idle list and returns it. If the idle pool +// is empty, this method returns nil and doesn't modify the idleResources slice. +// +// WARNING: Caller of this method must hold the pool mutex! +func (p *Pool[T]) tryAcquireIdleResource() *Resource[T] { + res, ok := p.idleResources.Pop() + if !ok { + return nil + } + + res.status = resourceStatusAcquired + return res +} + +// createNewResource creates a new resource and inserts it into list of pool +// resources. +// +// WARNING: Caller of this method must hold the pool mutex! +func (p *Pool[T]) createNewResource() *Resource[T] { + res := &Resource[T]{ + pool: p, + creationTime: time.Now(), + lastUsedNano: nanotime(), + poolResetCount: p.resetCount, + status: resourceStatusConstructing, + } + + p.allResources.append(res) + p.destructWG.Add(1) + + return res +} + +// Acquire gets a resource from the pool. If no resources are available and the pool is not at maximum capacity it will +// create a new resource. If the pool is at maximum capacity it will block until a resource is available. ctx can be +// used to cancel the Acquire. +// +// If Acquire creates a new resource the resource constructor function will receive a context that delegates Value() to +// ctx. Canceling ctx will cause Acquire to return immediately but it will not cancel the resource creation. This avoids +// the problem of it being impossible to create resources when the time to create a resource is greater than any one +// caller of Acquire is willing to wait. +func (p *Pool[T]) Acquire(ctx context.Context) (_ *Resource[T], err error) { + select { + case <-ctx.Done(): + p.canceledAcquireCount.Add(1) + return nil, ctx.Err() + default: + } + + return p.acquire(ctx) +} + +// acquire is a continuation of Acquire function that doesn't check context +// validity. +// +// This function exists solely only for benchmarking purposes. +func (p *Pool[T]) acquire(ctx context.Context) (*Resource[T], error) { + startNano := nanotime() + + var waitedForLock bool + if !p.acquireSem.TryAcquire(1) { + waitedForLock = true + err := p.acquireSem.Acquire(ctx, 1) + if err != nil { + p.canceledAcquireCount.Add(1) + return nil, err + } + } + + p.mux.Lock() + if p.closed { + p.acquireSem.Release(1) + p.mux.Unlock() + return nil, ErrClosedPool + } + + // If a resource is available in the pool. + if res := p.tryAcquireIdleResource(); res != nil { + waitTime := time.Duration(nanotime() - startNano) + if waitedForLock { + p.emptyAcquireCount += 1 + p.emptyAcquireWaitTime += waitTime + } + p.acquireCount += 1 + p.acquireDuration += waitTime + p.mux.Unlock() + return res, nil + } + + if len(p.allResources) >= int(p.maxSize) { + // Unreachable code. + panic("bug: semaphore allowed more acquires than pool allows") + } + + // The resource is not idle, but there is enough space to create one. + res := p.createNewResource() + p.mux.Unlock() + + res, err := p.initResourceValue(ctx, res) + if err != nil { + return nil, err + } + + p.mux.Lock() + defer p.mux.Unlock() + + p.emptyAcquireCount += 1 + p.acquireCount += 1 + waitTime := time.Duration(nanotime() - startNano) + p.acquireDuration += waitTime + p.emptyAcquireWaitTime += waitTime + + return res, nil +} + +func (p *Pool[T]) initResourceValue(ctx context.Context, res *Resource[T]) (*Resource[T], error) { + // Create the resource in a goroutine to immediately return from Acquire + // if ctx is canceled without also canceling the constructor. + // + // See: + // - https://github.com/jackc/pgx/issues/1287 + // - https://github.com/jackc/pgx/issues/1259 + constructErrChan := make(chan error) + go func() { + constructorCtx := newValueCancelCtx(ctx, p.baseAcquireCtx) + value, err := p.constructor(constructorCtx) + if err != nil { + p.mux.Lock() + p.allResources.remove(res) + p.destructWG.Done() + + // The resource won't be acquired because its + // construction failed. We have to allow someone else to + // take that resouce. + p.acquireSem.Release(1) + p.mux.Unlock() + + select { + case constructErrChan <- err: + case <-ctx.Done(): + // The caller is cancelled, so no-one awaits the + // error. This branch avoid goroutine leak. + } + return + } + + // The resource is already in p.allResources where it might be read. So we need to acquire the lock to update its + // status. + p.mux.Lock() + res.value = value + res.status = resourceStatusAcquired + p.mux.Unlock() + + // This select works because the channel is unbuffered. + select { + case constructErrChan <- nil: + case <-ctx.Done(): + p.releaseAcquiredResource(res, res.lastUsedNano) + } + }() + + select { + case <-ctx.Done(): + p.canceledAcquireCount.Add(1) + return nil, ctx.Err() + case err := <-constructErrChan: + if err != nil { + return nil, err + } + return res, nil + } +} + +// TryAcquire gets a resource from the pool if one is immediately available. If not, it returns ErrNotAvailable. If no +// resources are available but the pool has room to grow, a resource will be created in the background. ctx is only +// used to cancel the background creation. +func (p *Pool[T]) TryAcquire(ctx context.Context) (*Resource[T], error) { + if !p.acquireSem.TryAcquire(1) { + return nil, ErrNotAvailable + } + + p.mux.Lock() + defer p.mux.Unlock() + + if p.closed { + p.acquireSem.Release(1) + return nil, ErrClosedPool + } + + // If a resource is available now + if res := p.tryAcquireIdleResource(); res != nil { + p.acquireCount += 1 + return res, nil + } + + if len(p.allResources) >= int(p.maxSize) { + // Unreachable code. + panic("bug: semaphore allowed more acquires than pool allows") + } + + res := p.createNewResource() + go func() { + value, err := p.constructor(ctx) + + p.mux.Lock() + defer p.mux.Unlock() + // We have to create the resource and only then release the + // semaphore - For the time being there is no resource that + // someone could acquire. + defer p.acquireSem.Release(1) + + if err != nil { + p.allResources.remove(res) + p.destructWG.Done() + return + } + + res.value = value + res.status = resourceStatusIdle + p.idleResources.Push(res) + }() + + return nil, ErrNotAvailable +} + +// acquireSemAll tries to acquire num free tokens from sem. This function is +// guaranteed to acquire at least the lowest number of tokens that has been +// available in the semaphore during runtime of this function. +// +// For the time being, semaphore doesn't allow to acquire all tokens atomically +// (see https://github.com/golang/sync/pull/19). We simulate this by trying all +// powers of 2 that are less or equal to num. +// +// For example, let's immagine we have 19 free tokens in the semaphore which in +// total has 24 tokens (i.e. the maxSize of the pool is 24 resources). Then if +// num is 24, the log2Uint(24) is 4 and we try to acquire 16, 8, 4, 2 and 1 +// tokens. Out of those, the acquire of 16, 2 and 1 tokens will succeed. +// +// Naturally, Acquires and Releases of the semaphore might take place +// concurrently. For this reason, it's not guaranteed that absolutely all free +// tokens in the semaphore will be acquired. But it's guaranteed that at least +// the minimal number of tokens that has been present over the whole process +// will be acquired. This is sufficient for the use-case we have in this +// package. +// +// TODO: Replace this with acquireSem.TryAcquireAll() if it gets to +// upstream. https://github.com/golang/sync/pull/19 +func acquireSemAll(sem *semaphore.Weighted, num int) int { + if sem.TryAcquire(int64(num)) { + return num + } + + var acquired int + for i := int(log2Int(num)); i >= 0; i-- { + val := 1 << i + if sem.TryAcquire(int64(val)) { + acquired += val + } + } + + return acquired +} + +// AcquireAllIdle acquires all currently idle resources. Its intended use is for +// health check and keep-alive functionality. It does not update pool +// statistics. +func (p *Pool[T]) AcquireAllIdle() []*Resource[T] { + p.mux.Lock() + defer p.mux.Unlock() + + if p.closed { + return nil + } + + numIdle := p.idleResources.Len() + if numIdle == 0 { + return nil + } + + // In acquireSemAll we use only TryAcquire and not Acquire. Because + // TryAcquire cannot block, the fact that we hold mutex locked and try + // to acquire semaphore cannot result in dead-lock. + // + // Because the mutex is locked, no parallel Release can run. This + // implies that the number of tokens can only decrease because some + // Acquire/TryAcquire call can consume the semaphore token. Consequently + // acquired is always less or equal to numIdle. Moreover if acquired < + // numIdle, then there are some parallel Acquire/TryAcquire calls that + // will take the remaining idle connections. + acquired := acquireSemAll(p.acquireSem, numIdle) + + idle := make([]*Resource[T], acquired) + for i := range idle { + res, _ := p.idleResources.Pop() + res.status = resourceStatusAcquired + idle[i] = res + } + + // We have to bump the generation to ensure that Acquire/TryAcquire + // calls running in parallel (those which caused acquired < numIdle) + // will consume old connections and not freshly released connections + // instead. + p.idleResources.NextGen() + + return idle +} + +// CreateResource constructs a new resource without acquiring it. It goes straight in the IdlePool. If the pool is full +// it returns an error. It can be useful to maintain warm resources under little load. +func (p *Pool[T]) CreateResource(ctx context.Context) error { + if !p.acquireSem.TryAcquire(1) { + return ErrNotAvailable + } + + p.mux.Lock() + if p.closed { + p.acquireSem.Release(1) + p.mux.Unlock() + return ErrClosedPool + } + + if len(p.allResources) >= int(p.maxSize) { + p.acquireSem.Release(1) + p.mux.Unlock() + return ErrNotAvailable + } + + res := p.createNewResource() + p.mux.Unlock() + + value, err := p.constructor(ctx) + p.mux.Lock() + defer p.mux.Unlock() + defer p.acquireSem.Release(1) + if err != nil { + p.allResources.remove(res) + p.destructWG.Done() + return err + } + + res.value = value + res.status = resourceStatusIdle + + // If closed while constructing resource then destroy it and return an error + if p.closed { + go p.destructResourceValue(res.value) + return ErrClosedPool + } + + p.idleResources.Push(res) + + return nil +} + +// Reset destroys all resources, but leaves the pool open. It is intended for use when an error is detected that would +// disrupt all resources (such as a network interruption or a server state change). +// +// It is safe to reset a pool while resources are checked out. Those resources will be destroyed when they are returned +// to the pool. +func (p *Pool[T]) Reset() { + p.mux.Lock() + defer p.mux.Unlock() + + p.resetCount++ + + for res, ok := p.idleResources.Pop(); ok; res, ok = p.idleResources.Pop() { + p.allResources.remove(res) + go p.destructResourceValue(res.value) + } +} + +// releaseAcquiredResource returns res to the the pool. +func (p *Pool[T]) releaseAcquiredResource(res *Resource[T], lastUsedNano int64) { + p.mux.Lock() + defer p.mux.Unlock() + defer p.acquireSem.Release(1) + + if p.closed || res.poolResetCount != p.resetCount { + p.allResources.remove(res) + go p.destructResourceValue(res.value) + } else { + res.lastUsedNano = lastUsedNano + res.status = resourceStatusIdle + p.idleResources.Push(res) + } +} + +// Remove removes res from the pool and closes it. If res is not part of the +// pool Remove will panic. +func (p *Pool[T]) destroyAcquiredResource(res *Resource[T]) { + p.destructResourceValue(res.value) + + p.mux.Lock() + defer p.mux.Unlock() + defer p.acquireSem.Release(1) + + p.allResources.remove(res) +} + +func (p *Pool[T]) hijackAcquiredResource(res *Resource[T]) { + p.mux.Lock() + defer p.mux.Unlock() + defer p.acquireSem.Release(1) + + p.allResources.remove(res) + res.status = resourceStatusHijacked + p.destructWG.Done() // not responsible for destructing hijacked resources +} + +func (p *Pool[T]) destructResourceValue(value T) { + p.destructor(value) + p.destructWG.Done() +} diff --git a/vendor/github.com/jackc/puddle/v2/resource_list.go b/vendor/github.com/jackc/puddle/v2/resource_list.go new file mode 100644 index 0000000000..b2430959bf --- /dev/null +++ b/vendor/github.com/jackc/puddle/v2/resource_list.go @@ -0,0 +1,28 @@ +package puddle + +type resList[T any] []*Resource[T] + +func (l *resList[T]) append(val *Resource[T]) { *l = append(*l, val) } + +func (l *resList[T]) popBack() *Resource[T] { + idx := len(*l) - 1 + val := (*l)[idx] + (*l)[idx] = nil // Avoid memory leak + *l = (*l)[:idx] + + return val +} + +func (l *resList[T]) remove(val *Resource[T]) { + for i, elem := range *l { + if elem == val { + lastIdx := len(*l) - 1 + (*l)[i] = (*l)[lastIdx] + (*l)[lastIdx] = nil // Avoid memory leak + (*l) = (*l)[:lastIdx] + return + } + } + + panic("BUG: removeResource could not find res in slice") +} diff --git a/vendor/github.com/riverqueue/river/.gitignore b/vendor/github.com/riverqueue/river/.gitignore new file mode 100644 index 0000000000..94b8808680 --- /dev/null +++ b/vendor/github.com/riverqueue/river/.gitignore @@ -0,0 +1,6 @@ +/.envrc +/go.work.sum +/internal/cmd/riverbench/riverbench +/river +/riverdriver/riverdrivertest/example_libsql_test.libsql +/sqlite/ diff --git a/vendor/github.com/riverqueue/river/.golangci.yaml b/vendor/github.com/riverqueue/river/.golangci.yaml new file mode 100644 index 0000000000..80bea4c39e --- /dev/null +++ b/vendor/github.com/riverqueue/river/.golangci.yaml @@ -0,0 +1,130 @@ +version: "2" + +linters: + default: all + + disable: + # disabled, but which we should enable with discussion + - wrapcheck # checks that errors are wrapped; currently not done anywhere + + # disabled because we're not compliant, but which we should think about + - exhaustruct # checks that properties in structs are exhaustively defined; may be a good idea + - testpackage # requires tests in test packages like `river_test` + + # disabled because it's deprecated, and `default: all` already enables its + # replacement (`gomodguard_v2`) + - gomodguard + + # disabled because they're annoying/bad + - cyclop # screams into the void at "cyclomatic complexity" + - funcorder # very particular about where unexported functions can go, lots of churn. + - funlen # screams when functions are more than 60 lines long; what are we even doing here guys + - goconst # wants repeated test strings and other obvious literals to be constants; lots of churn. + - interfacebloat # we do in fact want >10 methods on the Adapter interface or wherever we see fit. + - gocognit # yells that "cognitive complexity" is too high; why + - gocyclo # ANOTHER "cyclomatic complexity" checker (see also "cyclop" and "gocyclo") + - godox # bans TODO statements; total non-starter at the moment + - err113 # wants all errors to be defined as variables at the package level; quite obnoxious + - maintidx # ANOTHER ANOTHER "cyclomatic complexity" lint (see also "cyclop" and "gocyclo") + - mnd # detects "magic numbers", which it defines as any number; annoying + - nestif # yells when if blocks are nested; what planet do these people come from? + - noinlineerr # disallows `if err := ...`; because why miss an opportunity to leak variables out of scope? + - ireturn # bans returning interfaces; questionable as is, but also buggy as hell; very, very annoying + - lll # restricts maximum line length; annoying + - nlreturn # requires a blank line before returns; annoying + - unqueryvet # bans all use of `SELECT *`; just ... sigh + - wsl # a bunch of style/whitespace stuff; annoying + - wsl_v5 # a second version of the first annoying wsl; how nice + + settings: + depguard: + rules: + all: + files: ["$all"] + deny: + - desc: Use `github.com/google/uuid` package for UUIDs instead. + pkg: github.com/xtgo/uuid + not-test: + files: ["!$test"] + deny: + - desc: Don't use `dbadaptertest` package outside of test environments. + pkg: github.com/riverqueue/river/internal/dbadaptertest + - desc: Don't use `riverinternaltest` package outside of test environments. + pkg: github.com/riverqueue/river/internal/riverinternaltest + + forbidigo: + forbid: + - msg: Use `require` variants instead. + pattern: ^assert\. + - msg: Use `Func` suffix for function variables instead. + pattern: Fn\b + - msg: Use built-in `max` function instead. + pattern: \bmath\.Max\b + - msg: Use built-in `min` function instead. + pattern: \bmath\.Min\b + + gomoddirectives: + replace-local: true + + gosec: + excludes: + - G404 # use of non-crypto random; overly broad for our use case + + revive: + rules: + - name: unused-parameter + disabled: true + + tagliatelle: + case: + rules: + json: snake + + testifylint: + enable-all: true + disable: + - go-require + + varnamelen: + ignore-names: + - db + - eg + - f + - i + - id + - j + - mu + - r + - sb # common convention for string builder + - t + - tb + - tt # common convention for table tests + - tx + - w + - wg + + exclusions: + generated: lax + presets: + - comments + - common-false-positives + - legacy + - std-error-handling + rules: + - path: (.+)\.go$ + text: Error return value of .(\w+\.Rollback(.*)). is not checked + +formatters: + enable: + - gci + - gofmt + - gofumpt + - goimports + + settings: + gci: + sections: + - Standard + - Default + - Prefix(github.com/riverqueue) + - Prefix(riverqueue.com/riverpro) diff --git a/vendor/github.com/riverqueue/river/AGENTS.md b/vendor/github.com/riverqueue/river/AGENTS.md new file mode 100644 index 0000000000..4712c57579 --- /dev/null +++ b/vendor/github.com/riverqueue/river/AGENTS.md @@ -0,0 +1,128 @@ +# River Coding Guidelines + +## Running Tests and Lint + +- **Tests**: use `make test` as the default. Only use `go test ...` when you must pass specific flags (e.g. `-run`, `-count`, `-race`, build tags, etc.) or need to debug a specific test in isolation. +- **Lint**: use `make lint` as the default. Only call `golangci-lint ...` directly when you must pass specific flags. +- Always run test and lint prior to considering a task complete, unless told otherwise (or if you did not touch any Go code/tests). +- If your execution environment has a sandbox/permission model, run `make test` and `make lint` unsandboxed (full permissions) so results match local dev and CI. + +## Other Build/Test Commands + +- **Run all tests with race detector**: `make test/race` +- **Run single test**: `cd $MODULE && go test ./path/to/package -run TestName` +- **Run benchmark**: `make bench` +- **Generate sqlc**: `make generate`. Use this any time a `.sql` file has been modified and we need to then regenerate `.sql.go` files from it. +- **Run tidy when deps change**: `make tidy` + +## Code Style Guidelines + +- **Imports**: use gci sections - Standard, Default, github.com/riverqueue. +- **Formatting**: use gofmt, gofumpt, goimports. +- **JSON tags**: use snake_case for JSON field tags. +- **Dependencies**: minimize external dependencies beyond standard library and pgx. +- **SQL access (non-test code)**: avoid ad-hoc SQL strings in library/runtime code. Add or extend a sqlc query, regenerate with `make generate`, and expose it through the driver interface. Keep direct SQL for tests and benchmark/admin utilities only (for example, `pg_stat_statements` and `VACUUM`). +- **Driver interface stability**: treat `riverdriver` as an internal adapter seam, not as an official external API. Its package comments explicitly say it should not be implemented or invoked by user code, and changes there are not considered semver-breaking. Do not preserve driver-interface methods or semantics for outside consumers; add, remove, or reshape them as needed to preserve or improve user-facing functionality. +- **Cross-driver driver tests**: treat `riverdriver/riverdrivertest` as the shared conformance suite for driver behavior. When changing `riverdriver` or a concrete driver, update or extend `riverdrivertest` so the intended semantics are exercised across drivers, not only in a single driver-specific test. +- **Error handling**: prefer context-rich errors; review linting rules before disabling them. +- **Testing**: use require variants instead of assert. +- **Helpers**: use `Func` suffix for function variables, not `Fn`. +- **Documentation**: include comments for exported functions and types. +- **Naming**: use idiomatic Go names; see `.golangci.yaml` for allowed short variable names. + +## Package Naming and Organization + +- Package names are lowercase, short, and representative; avoid `common`, `util`, or overly broad names. +- Use singular package names; avoid plurals like `httputils`. +- Keep import paths clean; avoid `src/`, `pkg/`, or other repo-structure leakage. +- Organize by responsibility instead of `models`/`types` buckets; keep types close to usage. +- Do not export identifiers from `main` packages that only build binaries. +- Add package docs, and use `doc.go` when documentation is long. + +## Code Organization + +Alphabetization is important when adding new code (do not reorganize existing code unless asked). + +- Types should be sorted alphabetically by name. +- Struct field definitions on a type should be sorted alphabetically by name, unless there is a good reason to deviate (examples: ID fields first, grouping mutexed fields after a mutex, etc.). +- When declaring an instance of a struct, fields should be sorted alphabetically by name unless a similar deviation is justified. +- When defining methods on a type, they should be sorted alphabetically by name. +- Constructors should come immediately after the type definition. +- Keep all methods for a type grouped together, immediately after the type definition and any constructor(s), organized alphabetically by name. Do not intersperse methods with other types or functions, except in special cases where a small utility type is needed to support a method and not used elsewhere. +- In unit tests, the outer test blocks should be sorted alphabetically by name. Inner test blocks should also be sorted alphabetically by name within the outer block. + +## Go Testing Conventions: Parallel Test Bundle + Setup Helpers + +This repo uses a parallel test bundle pattern (inspired by Brandur's write-up: https://brandur.org/fragments/parallel-test-bundle) to keep parallel subtests isolated and setup/fixtures DRY. + +- **Always opt into parallel**: + - **Top-level tests**: the first statement in every `TestXxx` should be `t.Parallel()`. + - **Subtests**: the first statement in every `t.Run(..., func(t *testing.T) { ... })` should be `t.Parallel()`, unless the subtest is intentionally non-parallel and includes a short comment explaining why. +- **Statement ordering and spacing**: + - **Top-level tests**: use `t.Parallel()`, then a blank line, then test preamble (`ctx`, `setup`, helpers), then a blank line before subtests/assertions. + - **Subtests**: use `t.Parallel()`, then a blank line, then subtest preamble. + - If a subtest calls `setup(...)`, prefer one blank line after the setup assignment before assertions/actions. + - For tiny/obvious subtests (one short statement after `t.Parallel()` or `setup(...)`), omitting one of these blank lines is acceptable. +- **Context and setup ordering**: + - If `setup` needs `ctx` (`setup(ctx, t)`), assign/derive `ctx` before calling `setup`. + - If `setup` does not need `ctx` (`setup(t)`), call `setup` first and derive specialized contexts (`WithCancel`, `WithTimeout`) close to where they are used. + - Avoid creating/deriving `ctx` far from usage unless shared setup requires it. +- **Prefer local bundles**: + - Define a `type testBundle struct { ... }` inside the `TestXxx` function containing the system under test and any fixtures frequently used across subtests. + - Each parallel subtest should call `setup(t)` to get a fresh bundle. Avoid sharing mutable state across parallel subtests. +- **`setup` helper rules**: + - Define `setup` as a local closure in the test: + - `setup := func(t *testing.T) *testBundle { ... }` + - Always call `t.Helper()` at the top of `setup`. + - **Only accept a context parameter if it is needed**: + - **Default**: `setup(t)` should not take `ctx`. + - **If setup must derive/seed a context**: prefer returning it: `setup := func(t *testing.T) (*testBundle, context.Context)`. + - **If setup must be passed an existing context**: accept `ctx` as the first parameter: `setup := func(ctx context.Context, t *testing.T) *testBundle`. + - Keep `setup` deterministic and self-contained; it should only use the `*testing.T` (and `ctx` if explicitly required) passed in. +- **Test signal instrumentation**: + - Prefer `rivershared/testsignal.TestSignal` in a `...TestSignals` struct with an `Init(tb)` helper over ad hoc `chan struct{}` fields in spies/fakes. + - Keep test-signal structs zero-value by default and call `Init(t)` only in tests that need to observe those signals. + - Wait for async events with `WaitOrTimeout()`. For negative assertions, use `RequireEmpty()` or `WaitC()` with a timeout select. + - Avoid custom channel signaling helpers and hand-managed channel capacities unless there is a specific, documented reason. + +Template: + +```go +func TestThing(t *testing.T) { + t.Parallel() + + type testBundle struct { + // Put SUT + common fixtures here. + } + + setup := func(t *testing.T) *testBundle { + t.Helper() + + return &testBundle{} + } + + t.Run("CaseName", func(t *testing.T) { + t.Parallel() + + bundle := setup(t) + + // ... use `bundle` in assertions/actions ... + }) + + t.Run("CaseNameWithCtxRequiredBySetup", func(t *testing.T) { + t.Parallel() + + setupWithCtx := func(ctx context.Context, t *testing.T) *testBundle { + t.Helper() + + _ = ctx + return &testBundle{} + } + + ctx := context.Background() + bundle := setupWithCtx(ctx, t) + + // ... use `bundle` in assertions/actions ... + }) +} +``` diff --git a/vendor/github.com/riverqueue/river/CHANGELOG.md b/vendor/github.com/riverqueue/river/CHANGELOG.md new file mode 100644 index 0000000000..865e427105 --- /dev/null +++ b/vendor/github.com/riverqueue/river/CHANGELOG.md @@ -0,0 +1,1165 @@ +# Changelog + +All notable changes to this project will be documented in this file. + +The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), +and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). + +## [Unreleased] + +## [0.43.0] - 2026-08-05 + +### Added + +- Added `JobListParams.TagsAll` and `JobListParams.TagsAny` for filtering jobs that match every or any exact tag, respectively. [PR #1339](https://github.com/riverqueue/river/pull/1339). + +### Fixed + +- `testsignal` no longer imports `riversharedtest`, so testify, go-spew, goleak, and yaml are no longer linked into production binaries that use River. `WaitTimeout` moved to `rivershared/util/testutil`, with `riversharedtest.WaitTimeout` kept as a wrapper around it. [PR #1342](https://github.com/riverqueue/river/pull/1342). + +## [0.42.0] - 2026-07-31 + +### Added + +- Added `JobArgsWithPlugins` for installing plugins that apply only to a specific job type. [PR #1337](https://github.com/riverqueue/river/pull/1337). + +## [0.41.1] - 2026-07-29 + +### Fixed + +- Fixed `JobRescuer` incorrectly rescuing jobs whose worker timeout of zero inherits a client-level `Config.JobTimeout` of `-1` (no timeout). [PR #1331](https://github.com/riverqueue/river/pull/1331). + +## [0.41.0] - 2026-07-23 + +### Added + +- Added `rivertype.HookMetricEmit` for receiving metrics emitted by River. Initial metrics report the duration of successful job fetches with `JobGetAvailableDurationMetric` and the number of jobs fetched with `JobGetAvailableCountMetric`. [PR #1285](https://github.com/riverqueue/river/pull/1285). +- The `riversqlite` driver is now tested against Turso, an in-process SQLite-compatible database written in Rust. [PR #1311](https://github.com/riverqueue/river/pull/1311). + +### Changed + +- Added `Config.Plugins` for extensions that should be installed as both hooks and middleware. `Config.Hooks` and `Config.Middleware` remain available for hook-only and middleware-only registration. [PR #1284](https://github.com/riverqueue/river/pull/1284). +- Reduce producer keep alive interval from 1 minute to 30 seconds. [PR #1319](https://github.com/riverqueue/river/pull/1319). +- Add context helpers that name timeouts for easy attribution on where they happened. [PR #1329](https://github.com/riverqueue/river/pull/1329). + +### Fixed + +- Guard against empty job slice returned by `JobSetStateIfRunningMany` when a job has been deleted mid-run. [PR #1308](https://github.com/riverqueue/river/pull/1308). +- Fixed `JobRescuer` pagination so a full batch of running jobs with disabled or longer worker-specific timeouts can't prevent later stuck jobs from being rescued. [PR #1318](https://github.com/riverqueue/river/pull/1318). +- If a job fails to unmarshal from JSON during job rescue or job execution, back off using the retry schedule and eventually discard it, similar to any other error that might occur. [PR #1324](https://github.com/riverqueue/river/pull/1324). +- Fixed `JobListOrderByFinalizedAt` validation so finalized states are accepted while non-finalized states are rejected. [PR #1327](https://github.com/riverqueue/river/issues/1327). + +## [0.40.0] - 2026-07-02 + +⚠️ Version 0.40.0 contains a new database migration, version 7, that rolls up some database cleanups and a few SQLite features: + +- Drop tables `river_client` and `river_client_queue`. These were added prospectively, but in the end were never used for anything. [PR #1115](https://github.com/riverqueue/river/pull/1115). +- Add a default value of 25 to `river_job.max_attempts`. Go code was previously injecting a max value, so this has no functional effect on existing behavior. [PR #1115](https://github.com/riverqueue/river/pull/1115). +- Add a default value of `CURRENT_TIMESTAMP` to `river_queue.updated_at`. Go code was previously injecting the current time, so this has no functional effect on existing behavior. [PR #1115](https://github.com/riverqueue/river/pull/1115). +- SQLite only: Convert `json` columns to `jsonb`. [PR #1224](https://github.com/riverqueue/river/pull/1224). +- SQLite only: Add pseudo listen/notify mechanism in a new `river_notification` table. [PR #1275](https://github.com/riverqueue/river/pull/1275). + +For SQLite, running River apps must be stopped briefly while the migration is run and their code upgrade to 0.40.0 so they start reading and inserting new values in `jsonb` instead of `json`. + +See [documentation on running River migrations](https://riverqueue.com/docs/migrations). If migrating with the CLI, make sure to update it to its latest version: + +```shell +go install github.com/riverqueue/river/cmd/river@latest +river migrate-up --database-url "$DATABASE_URL" +``` + +If not using River's internal migration system, the raw SQL can alternatively be dumped with (or change `--database-url` to a Postgres URI for Postgres versions): + +```shell +go install github.com/riverqueue/river/cmd/river@latest +river migrate-get --database-url sqlite:// --version 7 --up > river7.up.sql +river migrate-get --database-url sqlite:// --version 7 --down > river7.down.sql +``` + +### Added + +- SQLite picks up a new `river_notification` table that allows River to provide listen/notify-like functionality despite these functions not being supported outside of Postgres. [PR #1275](https://github.com/riverqueue/river/pull/1275). +- Added `JobStuckHandler`, giving clients a hook to handle "stuck" jobs (i.e. ones which are passed timeout and haven't responded to context cancellation) and potentially open a new worker slot if so desired. [PR #1291](https://github.com/riverqueue/river/pull/1291). + +### Changed + +- Convert SQLite JSON columns to JSONB (including migration). [PR #1224](https://github.com/riverqueue/river/pull/1224). +- Change SQLite driver operations over to use bulk inserts where possible now that sqlc has better support for `json_each`. [PR #1276](https://github.com/riverqueue/river/pull/1276) +- Detect duplicate step names across `river.ResumableStep` and return a validation error. [PR #1281](https://github.com/riverqueue/river/pull/1281) +- Earlier backpressure from `BatchCompleter` when it's throughput is saturated with fewer warnings to console. [PR #1292](https://github.com/riverqueue/river/pull/1292) +- Series of minor optimizations in `BatchCompleter` raising throughput ~20% when it's the bottleneck in job processing (e.g. in benchmarks). [PR #1293](https://github.com/riverqueue/river/pull/1293) + +### Fixed + +- Fix `JobCancel` having no effect on running jobs when using a poll-only driver (e.g. `riverdatabasesql`). The `controlActionCancel` event was silently dropped in `fetchAndRunLoop`'s `queueControlCh` handler instead of being forwarded to `maybeCancelJob`. Note: this fix only works within a single process; cross-process cancels in poll-only setups must wait for the next poll cycle. [PR #1245](https://github.com/riverqueue/river/pull/1245). +- Ensure jobs that return a custom timeout of -1 (no timeout) are never rescued. [PR #1288](https://github.com/riverqueue/river/pull/1288). +- Detect numbered PostgreSQL `REINDEX INDEX CONCURRENTLY` artifacts like `_ccnew1` and `_ccold2` so the reindexer does not keep accumulating failed artifact indexes. Fixes [#1296](https://github.com/riverqueue/river/issues/1296). [PR #1297](https://github.com/riverqueue/river/pull/1297). + +## [0.39.0] - 2026-06-03 + +⚠️ **Breaking API change:** `rivermigrate.Migrator.Validate` and `rivermigrate.Migrator.ValidateTx` now take a `*rivermigrate.ValidateOpts` parameter. Pass `nil` to preserve previous behavior. We normally endeavor not to make any breaking API changes, but this one will keep the API in a much nicer state, and is on an ancillary function that most installations won't be using. [PR #1259](https://github.com/riverqueue/river/pull/1259) + +### Added + +- Added `MetadataSet` to stage job metadata updates from worker middleware, `HookWorkBegin`, workers, or `HookWorkEnd`, with changes persisted when the job is completed. [PR #1269](https://github.com/riverqueue/river/pull/1269) + +### Changed + +- Add `rivermigrate.ValidateOpts.TargetVersion` so validation can check migrations up to a specific target version, matching the target-version behavior available on `Migrate` and `MigrateTx`. Notably, this is a breaking API change as the validate functions previously didn't take any options. [PR #1259](https://github.com/riverqueue/river/pull/1259) +- When using `(*Migrator[TTx]).Migrate` with a `TargetVersion` that's already applied, River now no-ops idempotently instead of returning an error as a user convenience. [PR #1260](https://github.com/riverqueue/river/pull/1260) +- Add logging statement for dropped job and queue subscription events at warn level when a subscriber buffer is full. [PR #1271](https://github.com/riverqueue/river/pull/1271) + +### Fixed + +- Add a 10-second timeout around `StandardPilot.JobGetAvailable` so a stalled standard-pilot fetch no longer hangs a producer indefinitely. [PR #1255](https://github.com/riverqueue/river/pull/1255) +- Fixed `rivertest.Worker.Work` and `WorkJob` to honor a configured custom `Config.Schema` when transitioning a job to its running state. Previously, the running-state update ran unqualified and could fail on a connection whose `search_path` didn't include the configured schema. [PR #1262](https://github.com/riverqueue/river/pull/1262) + +## [0.38.0] - 2026-05-22 + +### Added + +- Added new configuration `Config.SoftStopTimeout` to provide a cleaner way to gracefully stop a client. [PR #1239](https://github.com/riverqueue/river/pull/1239) + +## [0.37.1] - 2026-05-15 + +### Fixed + +- Wrap `PeriodicJobEnqueuer.insertBatch` database calls in a 30-second timeout. Previously, a stalled pgx `Begin`/`Insert`/`Commit` could hang the periodic enqueuer indefinitely, halting all periodic job insertion until the process was restarted or leader re-elected. [PR #1251](https://github.com/riverqueue/river/pull/1251) + +## [0.37.0] - 2026-05-11 + +### Added + +- Added "resumable jobs" that can be broken down into multiple steps and with a step persisted after it finishes that lets them skip work that's already been done. This is particularly useful for long running jobs that may experience a cancellation (like in the event of a deploy) during the span of their run. [PR #1226](https://github.com/riverqueue/river/pull/1226). + +## [0.36.0] - 2026-05-09 + +### Added + +- Add `QeueueBundle.Remove` to remove an already added queue/producer. [PR #1235](https://github.com/riverqueue/river/pull/1235) and [PR #1240](https://github.com/riverqueue/river/pull/1240). + +### Fixed + +- Fix unsafe concurrent producer map access in client. [PR #1236](https://github.com/riverqueue/river/pull/1236). +- Mark schema replacements as `Stable` in sqlc templates, preventing query SQL from having to be reallocated over and over again.. [PR #1242](https://github.com/riverqueue/river/pull/1242). +- Fix bug in `sqltemplate` cached path in order in which named args are passed to a query (previously, the order was unstable). [PR #1243](https://github.com/riverqueue/river/pull/1243). + +## [0.35.1] - 2026-04-26 + +### Fixed + +- Fix accidentally inverted conditional on notifier error log check. [PR #1231](https://github.com/riverqueue/river/pull/1231). + +## [0.35.0] - 2026-04-18 + +### Changed + +- Ignore errors like `tls: failed to send closeNotify alert (but connection was closed anyway)` when closing listeners. [PR #1216](https://github.com/riverqueue/river/pull/1216). + +### Fixed + +- Fixed leader election to track explicit database-issued leadership terms, reducing handoff flakiness and same-client reacquisition edge cases while making reelection and resign target the current leadership lease instead of a stale one. [PR #1213](https://github.com/riverqueue/river/pull/1213). + +## [0.34.0] - 2026-04-08 + +### Added + +- Added `Config.ReindexerIndexNames` and `ReindexerIndexNamesDefault()` so the reindexer's target indexes can be customized from the public API. [PR #1194](https://github.com/riverqueue/river/pull/1194). + +### Fixed + +- Upon a client gaining leadership, its queue maintainer is given more than one opportunity to start. [PR #1184](https://github.com/riverqueue/river/pull/1184). + +## [0.33.0] - 2026-04-03 + +### Changed + +- Jobs erroring or panicking no longer logs at the error/warn level because this is not indicative of a problem inside of River itself. These log statements have been demoted to info. [PR #1190](https://github.com/riverqueue/river/pull/1190). + +### Fixed + +- Fix in `Client.Start` where previously it was possible for a River client that only partially started before erroring to not try to start on subsequent `Start` invocations. [PR #1187](https://github.com/riverqueue/river/pull/1187). + +## [0.32.0] - 2026-03-23 + +### Added + +- `riverlog.Middleware` now supports `MiddlewareConfig.MaxTotalBytes` (default 8 MB) to cap total persisted `river:log` history per job. When the cap is exceeded, oldest log entries are dropped first while retaining the newest entry. Values over 64 MB are clamped to 64 MB. [PR #1157](https://github.com/riverqueue/river/pull/1157). + +### Changed + +- Improved `riverlog` performance and reduced memory amplification when appending to large persisted `river:log` histories. [PR #1157](https://github.com/riverqueue/river/pull/1157). +- Reduced snooze-path memory amplification by setting `snoozes` in metadata updates before marshaling, avoiding an extra full-payload JSON rewrite. [PR #1159](https://github.com/riverqueue/river/pull/1159). +- Schema names are now quoted in SQL operations, enabling the use of spaces and other odd characters. [PR #1175](https://github.com/riverqueue/river/pull/1175). + +### Fixed + +- `riverpgxv5` now adapts JSON parameters for `simple protocol` / `exec` query modes so `[]byte` JSON payloads are not encoded as `bytea` in pgx text-mode execution paths. This fixes invalid JSON syntax errors when running through protocol-constrained setups like PgBouncer transaction pooling while preserving normal behavior for explicit `bytea` parameters. Fixes [#1153](https://github.com/riverqueue/river/issues/1153). [PR #1155](https://github.com/riverqueue/river/pull/1155). + +## [0.31.0] - 2026-02-21 + +### Added + +- Added root River CLI flag `--statement-timeout` so Postgres session statement timeout can be set explicitly for commands like migrations. Explicit flag values take priority over database URL query params, and query params still take priority over built-in defaults. [PR #1142](https://github.com/riverqueue/river/pull/1142). + +### Fixed + +- Fix connection leak in `Listener.Connect` in case where `afterConnectExec` failed. Thanks Johan Kjölhede ([@GiGurra](https://github.com/GiGurra))! [PR #1147](https://github.com/riverqueue/river/pull/1147). +- Fix missing `ticker.Stop` in producer's `pollForSettingChanges` ([@GiGurra](https://github.com/GiGurra)). [PR #1148](https://github.com/riverqueue/river/pull/1148). +- Fix accidental use of cancelled context for `Notifier.Ping` ([@GiGurra](https://github.com/GiGurra)). [PR #1149](https://github.com/riverqueue/river/pull/1149). +- Add jitter to fetch poll loop to prevent producer stampeding ([@GiGurra](https://github.com/GiGurra)). [PR #1150](https://github.com/riverqueue/river/pull/1150). + +### Changed + +- Upgrade supported Go versions to 1.25 and 1.26, and update CI accordingly. [PR #1144](https://github.com/riverqueue/river/pull/1144). + +### Fixed + +- `JobCountByQueueAndState` now returns consistent results across drivers, including requested queues with zero jobs, and deduplicates repeated queue names in input. This resolves an issue with the sqlite driver in River UI reported in [riverqueue/riverui#496](https://github.com/riverqueue/riverui#496). [PR #1140](https://github.com/riverqueue/river/pull/1140). + +## [0.30.2] - 2026-01-26 + +### Fixed + +- Fix bug in worker-level stuck job detection. [PR #1133](https://github.com/riverqueue/river/pull/1133). + +## [0.30.1] - 2026-01-19 + +### Fixed + +- Stuck job detection now accounts for worker-level timeouts as well as client-level timeouts. [PR #1125](https://github.com/riverqueue/river/pull/1125). + +## [0.30.0] - 2026-01-11 + +### Fixed + +- Fix possible nil pointer panic when using nil `opts` in `Migrator.MigrateTx`. [PR #1117](https://github.com/riverqueue/river/pull/1117). + +## [0.29.0] - 2025-12-22 + +### Added + +- Added `HookPeriodicJobsStart` that can be used to run custom logic when a periodic job enqueuer starts up on a new leader. [PR #1084](https://github.com/riverqueue/river/pull/1084). +- Added `Client.Notify().RequestResign` and `Client.Notify().RequestResignTx` functions allowing any client to request that the current leader resign. [PR #1085](https://github.com/riverqueue/river/pull/1085). +- Basic stuck detection after a job's exceeded its timeout and still not returned after the executor's initiated context cancellation and waited a short margin for the cancellation to take effect. [PR #1097](https://github.com/riverqueue/river/pull/1097). +- Added `Client.JobUpdate` which can be used to persist job output partway through a running work function instead of having to wait until the job is completed. [PR #1098](https://github.com/riverqueue/river/pull/1098). + +### Changed + +- Add a little more error flavor for when encountering a deadline exceeded error on leadership election suggesting that the user may want to try increasing their database pool size. [PR #1101](https://github.com/riverqueue/river/pull/1101). +- When migrating without an outer transaction, insert/delete version rows immediately after executing migration SQL so that in case a later migration fails, the migrator knows where to restart from. [PR #1106](https://github.com/riverqueue/river/pull/1106). + +## [0.28.0] - 2025-11-23 + +### Added + +- Added `riverlog.LoggerSafely` which provides a non-panic variant of `riverlog.Logger` for use when code may or may not have a context logger available. [PR #1093](https://github.com/riverqueue/river/pull/1093). + +## [0.27.0] - 2025-11-14 + +### Added + +- Periodic jobs with IDs may now be removed by ID using the new `PeriodicJobBundle.RemoveByID` and `PeriodicJobBundle.RemoveManyByID`. [PR #1071](https://github.com/riverqueue/river/pull/1071). + +### Changed + +- Decrease `serviceutil.MaxAttemptsBeforeResetDefault` from 10 to 7, lowering the effective limit on most internal exponential backoffs from ~512 seconds to 64 seconds. Further lowered the leader elector's keep leadership backoff interval to cap out at 4 seconds since leadership without a successful heartbeat will be lost soon after that anyway. [PR #1079](https://github.com/riverqueue/river/pull/1079). + +### Fixed + +- Fix snoozed events emitted from `rivertest.Worker` when snooze duration is zero seconds. [PR #1057](https://github.com/riverqueue/river/pull/1057). +- Rollbacks now use an uncancelled context so as to not leave transactions in an ambiguous state if a transaction in them fails due to context cancellation. [PR #1062](https://github.com/riverqueue/river/pull/1062). +- Removing periodic jobs with IDs assigned also remove them from ID map. [PR #1070](https://github.com/riverqueue/river/pull/1070). +- Clear periodic jobs also fully clears all those assigned with an ID. [PR #1083](https://github.com/riverqueue/river/pull/1083). +- `river:"unique"` annotations on substructs within `JobArgs` structs are now factored into uniqueness `ByArgs` calculations. [PR #1076](https://github.com/riverqueue/river/pull/1076). +- Stop subservices and embedded `baseservice.Service` on error in the event of a periodic job enqueuer start error. [PR #1081](https://github.com/riverqueue/river/pull/1081). + +## [0.26.0] - 2025-10-07 + +⚠️ Internal APIs used for communication between River and River Pro have changed. If using River Pro, make sure to update River and River Pro to latest at the same time to get compatible versions. River v0.26.0 is compatible with River Pro v0.19.0. + +### Added + +- The job rescuer now sets `river:rescue_count` with an integer count of how many times the job has been rescued by the `JobRescuer` maintenance process when it's considered stuck. [PR #1047](https://github.com/riverqueue/river/pull/1047). + +### Changed + +- Errors returned from job workers are now logged in full using a `slog.Any` attribute. Previously, only their error text was logged. [PR #1051](https://github.com/riverqueue/river/pull/1051). + +### Fixed + +- Set `updated_at` when invoking pilot `PeriodicJobUpsert`. [PR #1045](https://github.com/riverqueue/river/pull/1045). + +## [0.25.0] - 2025-09-14 + +⚠️ Internal APIs used for communication between River and River Pro have changed. If using River Pro, make sure to update River and River Pro to latest at the same time to get compatible versions. River v0.25.0 is compatible with River Pro v0.18.0. + +### Changed + +- Set minimum Go version to Go 1.24. [PR #1032](https://github.com/riverqueue/river/pull/1032). +- **Breaking change:** `Client.JobDeleteMany` now requires the use of `JobDeleteManyParams.UnsafeAll` to delete all jobs without a filter applied. This is a safety feature to make it more difficult to accidentally delete all non-running jobs. This is a minor breaking change, but on a fairly new feature that's not likely to be used on purpose by very many people yet. [PR #1033](https://github.com/riverqueue/river/pull/1033). + +### Fixed + +- Don't double log fetch errors. [PR #1025](https://github.com/riverqueue/river/pull/1025). +- When snoozing a job with zero duration so that it's retried immediately, subscription events no longer appear incorrectly with a kind of `rivertype.EventKindJobFailed`. Instead they're assigned `rivertype.EventKindJobSnoozed` just like they would have with a non-zero snooze duration. [PR #1037](https://github.com/riverqueue/river/pull/1037). + +## [0.24.0] - 2025-08-16 + +⚠️ Version 0.24.0 has a breaking change in `HookWorkEnd.WorkEnd` in that a new `JobRow` parameter has been added to the function's signature. Any intergration defining a custom `HookWorkEnd` hook should update its implementation so the hook continues to be called correctly. + +⚠️ Internal APIs used for communication between River and River Pro have changed. If using River Pro, make sure to update River and River Pro to latest at the same time to get compatible versions. River v0.24.0 is compatible with River Pro v0.16.0. + +### Added + +- The project now tests against [libSQL](https://github.com/tursodatabase/libsql), a popular SQLite fork. It's used through the same `riversqlite` driver that SQLite uses. [PR #957](https://github.com/riverqueue/river/pull/957) +- Added `JobDeleteMany` operations that remove many jobs in a single operation according to input criteria. [PR #962](https://github.com/riverqueue/river/pull/962) +- Added `Client.Schema()` method to return a client's configured schema. [PR #983](https://github.com/riverqueue/river/pull/983). +- Integrated riverui queries into the driver system to pave the way for multi-driver UI support. [PR #983](https://github.com/riverqueue/river/pull/983). +- Added `QueueConfig` level `FetchCooldown` and `FetchPollInterval` settings to enable queue-specific job fetch intervals. For example, a queue of high-priority jobs could be checked more often to improve responsiveness, while one with slow or time-insensitive tasks could be checked infrequently to reduce database load. [PR #994](https://github.com/riverqueue/river/pull/994). + +### Changed + +- Remove unecessary transactions where a single database operation will do. This reduces the number of subtransactions created which can be an operational benefit it many cases. [PR #950](https://github.com/riverqueue/river/pull/950) +- Bring all driver tests into separate package so they don't leak dependencies. This removes dependencies from the top level `river` package that most River installations won't need, thereby reducing the transitive dependency load of most River installations. [PR #955](https://github.com/riverqueue/river/pull/955). +- The reindexer maintenance service now reindexes all `river_job` indexes, including its primary key. This is expected to help in situations where the jobs table has in the past expanded to a very large size (which makes most indexes larger), is now a much more modest size, but has left the indexes in their expanded state. [PR #963](https://github.com/riverqueue/river/pull/963). +- The River CLI now accepts a `--target-version` of 0 with `river migrate-down` to run all down migrations and remove all River tables (previously, -1 was used for this; -1 still works, but now 0 also works). [PR #966](https://github.com/riverqueue/river/pull/966). +- **Breaking change:** The `HookWorkEnd` interface's `WorkEnd` function now receives a `JobRow` parameter in addition to the `error` it received before. Having a `JobRow` to work with is fairly crucial to most functionality that a hook would implement, and its previous omission was entirely an error. [PR #970](https://github.com/riverqueue/river/pull/970). +- Add maximum bound to each job's `attempted_by` array so that in degenerate cases where a job is run many, many times (say it's snoozed hundreds of times), it doesn't grow to unlimited bounds. [PR #974](https://github.com/riverqueue/river/pull/974). +- A logger passed in via `river.Config` now overrides the default test-based logger when using `rivertest.NewWorker`. [PR #980](https://github.com/riverqueue/river/pull/980). +- Cleaner retention periods (`CancelledJobRetentionPeriod`, `CompletedJobRetentionPeriod`, `DiscardedJobRetentionPeriod`) can be configured to -1 to disable them so that the corresponding type of job is retained indefinitely. [PR #990](https://github.com/riverqueue/river/pull/990). +- Jobs inserted from periodic jobs with IDs now have metadata `river:periodic_job_id` set so they can be traced back to the periodic job that inserted them. [PR #992](https://github.com/riverqueue/river/pull/992). +- The unused function `WorkerDefaults.Hooks` has been removed. This is technically a breaking change, but this function was a vestigal refactoring artifact that was never used by anything, so in practice it shouldn't be breaking. [PR #997](https://github.com/riverqueue/river/pull/997). +- Periodic job records are upserted immediately through a pilot when a client is started rather than the first time their associated job would run. This doesn't mean they're run immediately (they'll only run if `RunOnStart` is enabled), but rather just tracked immediately. [PR #998](https://github.com/riverqueue/river/pull/998). +- The job scheduler still schedules jobs in batches of up to 10,000, but when it encounters a series of consecutive timeouts it assumes that the database is in a degraded state and switches to doing work in a smaller batch size of 1,000 jobs. [PR #1013](https://github.com/riverqueue/river/pull/1013). +- Other maintenance services including the job cleaner, job rescuer, and queue cleaner also prefer a batch size of 10,000, but will fall back to smaller batches of 1,000 on consecutive database timeouts. [PR #1016](https://github.com/riverqueue/river/pull/1016). + +### Fixed + +- Cleanly error on invalid schema names in `Config.Schema`. [PR #952](https://github.com/riverqueue/river/pull/952). +- Jobs rescued by `JobRescuer` no longer have their trace set to "TODO". This becomes an empty string instead. [PR #1010](https://github.com/riverqueue/river/pull/1010). + +## [0.23.1] - 2025-06-04 + +This includes a minor CLI bugfix for riverpro and no other changes, see the v0.23.0 notes for major changes. + +### Fixed + +- Fixed a riverpro CLI integration point broken in v0.23.0. [PR #945](https://github.com/riverqueue/river/pull/945) + +## [0.23.0] - 2025-06-04 + +⚠️ Internal APIs used for communication between River and River Pro have changed. If using River Pro, make sure to update River and River Pro to latest at the same time to get compatible versions. River v0.23.0 is compatible with River Pro v0.15.0. + +**Terminal UI:** @almottier wrote a very cool [terminal UI for River](https://github.com/almottier/rivertui) featuring real-time job monitoring with automatic refresh, job filtering, a job details view providing detailed information (plus look up by ID in the UI or by command line argument), and job actions like retry and cancellation. And as good as all that might sound, go take a look because it's even better in person. + +### Added + +- Preliminary River driver for SQLite (`riverdriver/riversqlite`). This driver seems to produce good results as judged by the test suite, but so far has minimal real world vetting. Try it and let us know how it works out. [PR #870](https://github.com/riverqueue/river/pull/870). +- CLI `river migrate-get` now takes a `--schema` option to inject a custom schema into dumped migrations and schema comments are hidden if `--schema` option isn't provided. [PR #903](https://github.com/riverqueue/river/pull/903). +- Added `riverlog.NewMiddlewareCustomContext` that makes the use of `riverlog` job-persisted logging possible with non-slog loggers. [PR #919](https://github.com/riverqueue/river/pull/919). +- Added `RequireInsertedOpts.Schema`, allowing an explicit schema to be set when asserting on job inserts with `rivertest`. [PR #926](https://github.com/riverqueue/river/pull/926). +- When using a driver that doesn't support listen/notify, producers within same process are notified immediately of new job inserts and queue changes (e.g. pause/resume) without having to poll when non-transactional variants are used (i.e. `Insert` instead of `InsertTx`). [PR #928](https://github.com/riverqueue/river/pull/928). +- Added `JobListParams.Where`, which provides an escape hatch for job listing that runs arbitrary SQL with named parameters. [PR #933](https://github.com/riverqueue/river/pull/933). + +### Changed + +- Optimized the job completer's query `JobSetStateIfRunningMany`, resulting in an approximately 15% reduction in its duration when completing 2000 jobs, and around a 15-20% increase in `riverbench` throughput. [PR #904](https://github.com/riverqueue/river/pull/904). +- `TimeStub` has been removed from the `rivertest` package. Its original inclusion was entirely accidentally and it should be considered entirely an internal API. [PR #912](https://github.com/riverqueue/river/pull/912). +- When storing job-persisted logging with `riverlog`, if a work run's logging was completely empty, no metadata value is stored at all (previously, an empty value was stored). [PR #919](https://github.com/riverqueue/river/pull/919). +- Changed the internal integration APIs for River Pro. River Pro users must upgrade both libraries as part of this update. [PR #929](https://github.com/riverqueue/river/pull/929). + +### Fixed + +- Resuming an already unpaused queue is now fully an no-op, and won't touch the row's `updated_at` like it (unintentionally) did before. [PR #870](https://github.com/riverqueue/river/pull/870). +- Suppress an error log line from the producer that may occur on normal shutdown when operating in poll-only mode. [PR #896](https://github.com/riverqueue/river/pull/896). +- Added missing help documentation for CLI command `river migrate-list`. [PR #903](https://github.com/riverqueue/river/pull/903). +- Correct handling an explicit schema in the reindexer maintenance service. [PR #916](https://github.com/riverqueue/river/pull/916). +- Return specific explanatory error when attempting to use `JobListParams.Metadata` with `JobListTx` on SQLite. [PR #924](https://github.com/riverqueue/river/pull/924). +- The reindexer now skips work if artifacts from a failed reindex are present under the assumption that if they are, a new reindex build is likely to fail again. Context cancel timeout is increased from 15 seconds to 1 minute, allowing more time for reindexes to finish. Timeout becomes configurable with `Config.ReindexerTimeout`. [PR #935](https://github.com/riverqueue/river/pull/935). +- Accessing `Client.PeriodicJobs()` on an insert-only client now panics with a more helpful explanatory error message rather than an unhelpful nil pointer panic. [PR #938](https://github.com/riverqueue/river/pull/938). +- Return an error when adding a new queue at runtime via the `QueueBundle` if that queue was already added. [PR #929](https://github.com/riverqueue/river/pull/929). + +## [0.22.0] - 2025-05-10 + +### Added + +- A new `JobArgsWithKindAliases` interface lets job args implement `KindAliases` to register a second kind that their worker will respond to. This provides a way to safely rename job kinds even with jobs using the original kind already in the database. [PR #880](https://github.com/riverqueue/river/pull/880). + +### Changed + +- Job kinds must comply to a format of `\A[\w][\w\-\[\]<>\/.·:+]+\z`, mainly in an attempt to eliminate commas and spaces to make format more predictable for an upcoming search UI. This check can be disabled for now using `Config.SkipJobKindValidation`, but this option will likely be removed in a future version of River. The new `JobArgsWithKindAliases` interface (see above) can be used to rename non-compliant kinds. [PR #879](https://github.com/riverqueue/river/pull/879). + +### Fixed + +- The `riverdatabasesql` now fully supports raw connections through [`lib/pq`](https://github.com/lib/pq) rather than just `database/sql` through Pgx. We don't recommend the use of `lib/pq` as it's an unmaintained project, but this change should help with compatibility for older projects. [PR #883](https://github.com/riverqueue/river/pull/883). + +## [0.21.0] - 2025-05-02 + +⚠️ Internal APIs used for communication between River and River Pro have changed. If using River Pro, make sure to update River and River Pro to latest at the same time to get compatible versions. River v0.21.0 is compatible with River Pro v0.13.0. + +### Added + +- Added `river/riverlog` containing middleware that injects a context logger to workers that collates log output and persists it with job metadata. [PR #844](https://github.com/riverqueue/river/pull/844). +- Added `JobInsertMiddlewareFunc` and `WorkerMiddlewareFunc` to easily implement middleware with a function instead of a struct. [PR #844](https://github.com/riverqueue/river/pull/844). +- Added `Config.Schema` which lets a non-default schema be injected explicitly into a River client that'll be used for all database operations. This may be particularly useful for proxies like PgBouncer that may not respect a schema configured in `search_path`. [PR #848](https://github.com/riverqueue/river/pull/848). +- Added `rivertype.HookWorkEnd` hook interface that runs after a job has been worked. [PR #863](https://github.com/riverqueue/river/pull/863). +- Added support for filtering jobs by a list of job IDs and by priorities in `JobList` and `JobListParams`. For more flexible job listing. [PR #871](https://github.com/riverqueue/river/pull/871). + +### Changed + +- Client no longer returns an error if stopped before startup could complete (previously, it returned the unexported `ErrShutdown`). [PR #841](https://github.com/riverqueue/river/pull/841). + +### Fixed + +- A queue unpausing triggers an immediate fetch so that available jobs in the paused queue may be started faster than before. [PR #854](https://github.com/riverqueue/river/pull/854). + +## [0.20.2] - 2025-04-08 + +### Added + +- Added `QueueUpdateTx` API so there's a transactional variant of the `QueueUpdate` API from [PR #834](https://github.com/riverqueue/river/pull/834). [PR #838](https://github.com/riverqueue/river/pull/838). + +## [0.20.1] - 2025-04-05 + +### Fixed + +- Corrected the serialization of queue control event payloads emitted by `QueueUpdate`. [PR #834](https://github.com/riverqueue/river/pull/834). + +## [0.20.0] - 2025-04-04 + +### Added + +- Added a `QueueUpdate` API to the `Client` which will be used for upcoming functionality. [PR #822](https://github.com/riverqueue/river/pull/822). + +### Changed + +- Set minimum Go version to Go 1.23. [PR #811](https://github.com/riverqueue/river/pull/811). +- Deprecate `river.JobInsertMiddlewareDefaults` and `river.WorkerMiddlewareDefaults` in favor of the more general `river.MiddlewareDefaults` embeddable struct. The two former structs will be removed in a future version. [PR #815](https://github.com/riverqueue/river/pull/815). + +### Fixed + +- Cleanly error when attempting to add a queue at runtime to a `Client` which was not configured to run jobs (no `Workers`). [PR #826](https://github.com/riverqueue/river/pull/826). + +## [0.19.0] - 2025-03-16 + +⚠️ Version 0.19.0 has minor breaking changes for the `Worker.Middleware`, introduced fairly recently in 0.17.0 that has a worker's `Middleware` function now taking a non-generic `JobRow` parameter instead of a generic `Job[T]`. We tried not to make this change, but found the existing middleware interface insufficient to provide the necessary range of functionality we wanted, and this is a secondary middleware facility that won't be in use for many users, so it seemed worthwhile. + +### Added + +- Added a new "hooks" API for tying into River functionality at various points like job inserts or working. Differs from middleware in that it doesn't go on the stack and can't modify context, but in some cases is able to run at a more granular level (e.g. for each job insert rather than each _batch_ of inserts). [PR #789](https://github.com/riverqueue/river/pull/789). +- `river.Config` has a generic `Middleware` setting that can be used as a convenient way to configure middlewares that implement multiple middleware interfaces (e.g. `JobInsertMiddleware` _and_ `WorkerMiddleware`). Use of this setting is preferred over `Config.JobInsertMiddleware` and `Config.WorkerMiddleware`, which have been deprecated. [PR #804](https://github.com/riverqueue/river/pull/804). + +### Changed + +- The `river.RecordOutput` function now returns an error if the output is too large. The output is limited to 32MB in size. [PR #782](https://github.com/riverqueue/river/pull/782). +- **Breaking change:** The `Worker` interface's `Middleware` function now takes a `JobRow` parameter instead of a generic `Job[T]`. This was necessary to expand the potential of what middleware can do: by letting the executor extract a middleware stack from a worker before a job is fully unmarshaled, the middleware can also participate in the unmarshaling process. [PR #783](https://github.com/riverqueue/river/pull/783). +- `JobList` has been reimplemented to use sqlc. [PR #795](https://github.com/riverqueue/river/pull/795). + +## [0.18.0] - 2025-02-20 + +⚠️ Version 0.18.0 has breaking changes for the `rivertest.Worker` type that was just introduced. While attempting to round out some edge cases with its design, we realized some of them simply couldn't be solved adequately without changing the overall design such that all tested jobs are inserted into the database. Given the short duration since it was released (over a weekend) it's unlikely many users have adopted it and it seemed best to rip off the bandaid to fix it before it gets widely used. + +### Added + +- Jobs can now store a recorded "output" value, a JSON-encoded payload set by the job during execution and stored in the job's metadata. The `river.RecordOutput` function makes it easy to use the job row to store transient/temporary values that are needed for introspection or for other downstream jobs. The output can be accessed using the `JobRow.Output()` helper method. + + This output is stored at the same time as the job is completed following execution, so it does not require additional database calls or overhead. Output can be anything that can be stored in a Postgres JSONB field, though for performance reasons it should be limited in size. [PR #758](https://github.com/riverqueue/river/pull/758). + +### Changed + +- **Breaking change:** The `rivertest.Worker` type now requires all jobs to be inserted into the database. The original design allowed workers to be tested without hitting the database at all. Ultimately this design made it hard to correctly simulate features like `JobCompleteTx` and the other potential solutions seemed undesirable. + + As part of this change, the `Work` and `WorkJob` methods now take a transaction argument. The expectation is that a transaction will be opened by the caller and rolled back after test completion. Additionally, the return signature was changed to return a `WorkResult` struct alongside the error. The struct includes the post-execution job row as well as the event kind that occurred, making it easy to inspect the job's state after execution. + + Finally, the implementation was refactored so that it uses the _real_ `river.Client` insert path, and also uses the same job execution path as real execution. This minimizes the potential for differences in behavior between testing and real execution. + [PR #766](https://github.com/riverqueue/river/pull/766). + +- Adjusted panic stack traces to filter out irrelevant frames like the ones generated by the runtime package that constructed the trace, or River's internal rescuing code. This makes the first panic frame reflect the actual panic origin for easier debugging. [PR #774](https://github.com/riverqueue/river/pull/774). + +### Fixed + +- Fix error message on unsuccessful client subscribe that erroneously referred to "Workers" not configured. [PR #771](https://github.com/riverqueue/river/pull/771). +- Fix an issue with encoding unique keys in riverdatabasesql driver. [PR #777](https://github.com/riverqueue/river/pull/777). + +## [0.17.0] - 2025-02-16 + +### Added + +- Exposed `TestConfig` struct on `Config` under the `Test` field for configuration that is specific to test environments. For now, the only field on this type is `Time`, which can be used to set a synthetic `TimeGenerator` for tests. A stubbable time generator was added as `rivertest.TimeStub` to allow time to be easily stubbed in tests. [PR #754](https://github.com/riverqueue/river/pull/754). +- New `rivertest.Worker` type to make it significantly easier to test River workers. Either real or synthetic jobs can be worked using this interface, generally without requiring any database interactions. The `Worker` type provides a realistic execution environment with access to the full range of River features, including `river.ClientFromContext`, middleware (both global and per-worker), and timeouts. [PR #753](https://github.com/riverqueue/river/pull/753). + +### Changed + +- Errors returned from retryable jobs are now logged with warning logs instead of error logs. Error logs are still used for jobs that error after reaching `max_attempts`. [PR #743](https://github.com/riverqueue/river/pull/743). +- Remove range variable capture in `for` loops and use simplified `range` syntax. Each of these requires Go 1.22 or later, which was already our minimum required version since Go 1.23 was released. [PR #755](https://github.com/riverqueue/river/pull/755). + +### Fixed + +- `riverdatabasesql` driver: properly handle `nil` values in `bytea[]` inputs. This fixes the driver's handling of empty unique keys on insert for non-unique jobs with the newer unique jobs implementation. [PR #739](https://github.com/riverqueue/river/pull/739). +- `JobCompleteTx` now returns `rivertype.ErrNotFound` if the job doesn't exist instead of panicking. [PR #753](https://github.com/riverqueue/river/pull/753). +- - `NeverSchedule.Next` now returns the correct maximum time value, ensuring that the periodic job truly never runs. This fixes an issue where an incorrect maximum timestamp was previously used. Thanks Hubert Krauze ([@krhubert](https://github.com/krhubert))! [PR #760](https://github.com/riverqueue/river/pull/760) + +## [0.16.0] - 2024-01-27 + +### Added + +- `NeverSchedule` returns a `PeriodicSchedule` that never runs. This can be used to effectively disable the reindexer or any other maintenance service. [PR #718](https://github.com/riverqueue/river/pull/718). +- Add `SkipUnknownJobCheck` client config option to skip job arg worker validation. [PR #731](https://github.com/riverqueue/river/pull/731). + +### Changed + +- The reindexer maintenance process has been enabled. As of now, it will reindex only the `river_job_args_index` and `river_jobs_metadata_index` `GIN` indexes, which are more prone to bloat than b-tree indexes. By default it runs daily at midnight UTC, but can be customized on the `river.Config` type via `ReindexerSchedule`. Most installations will benefit from this process, but it can be disabled altogether using `NeverSchedule`. [PR #718](https://github.com/riverqueue/river/pull/718). +- Periodic jobs now have a `"periodic": true` attribute set in their metadata to make them more easily distinguishable from other types of jobs. [PR #728](https://github.com/riverqueue/river/pull/728). +- Snoozing a job now causes its `attempt` to be _decremented_, whereas previously the `max_attempts` would be incremented. In either case, this avoids allowing a snooze to exhaust a job's retries; however the new behavior also avoids potential issues with wrapping the `max_attempts` value, and makes it simpler to implement a `RetryPolicy` based on either `attempt` or `max_attempts`. The number of snoozes is also tracked in the job's metadata as `snoozes` for debugging purposes. + + The implementation of the builtin `RetryPolicy` implementations is not changed, so this change should not cause any user-facing breakage unless you're relying on `attempt - len(errors)` for some reason. [PR #730](https://github.com/riverqueue/river/pull/730). + +- `ByPeriod` uniqueness is now based off a job's `ScheduledAt` instead of the current time if it has a value. [PR #734](https://github.com/riverqueue/river/pull/734). + +## [0.15.0] - 2024-12-26 + +### Added + +- The River CLI will now respect the standard set of `PG*` environment variables like `PGHOST`, `PGPORT`, `PGDATABASE`, `PGUSER`, `PGPASSWORD`, and `PGSSLMODE` to configure a target database when the `--database-url` parameter is omitted. [PR #702](https://github.com/riverqueue/river/pull/702). +- Add missing doc for `JobRow.UniqueStates` + reveal `rivertype.UniqueOptsByStateDefault()` to provide access to the default set of unique job states. [PR #707](https://github.com/riverqueue/river/pull/707). + +### Changed + +- Sleep durations are now logged as Go-like duration strings (e.g. "10s") in either text or JSON instead of duration strings in text and nanoseconds in JSON. [PR #699](https://github.com/riverqueue/river/pull/699). +- Altered the migration comments from `river migrate-get` to include the "line" of the migration being run (`main`, or for River Pro `workflow` and `sequence`) to make them more distinguishable. [PR #703](https://github.com/riverqueue/river/pull/703). +- Fewer slice allocations during unique insertions. [PR #705](https://github.com/riverqueue/river/pull/705). + +### Fixed + +- Exponential backoffs at degenerately high job attempts (>= 310) no longer risk overflowing `time.Duration`. [PR #698](https://github.com/riverqueue/river/pull/698). + +## [0.14.3] - 2024-12-14 + +### Changed + +- Dropped internal random generators in favor of `math/rand/v2`, which will have the effect of making code fully incompatible with Go 1.21 (`go.mod` has specified a minimum of 1.22 for some time already though). [PR #691](https://github.com/riverqueue/river/pull/691). + +### Fixed + +- 006 migration now tolerates previous existence of a `unique_states` column in case it was added separately so that the new index could be raised with `CONCURRENTLY`. [PR #690](https://github.com/riverqueue/river/pull/690). + +## [0.14.2] - 2024-11-16 + +### Fixed + +- Cancellation of running jobs relied on a channel that was only being received when in the job fetch routine, meaning that jobs which were cancelled would not be cancelled until the next scheduled fetch. This was fixed by also receiving from the job cancellation channel when in the main producer loop, even if no fetches are happening. [PR #678](https://github.com/riverqueue/river/pull/678). +- Job insert middleware were not being utilized for periodic jobs. This insertion path has been refactored to rely on the unified insertion path from the client. Fixes #675. [PR #679](https://github.com/riverqueue/river/pull/679). + +## [0.14.1] - 2024-11-04 + +### Fixed + +- In [PR #663](https://github.com/riverqueue/river/pull/663) the client was changed to be more aggressive about re-fetching when it had previously fetched a full batch. Unfortunately a clause was missed, which resulted in the client being more aggressive any time even a single job was fetched on the previous attempt. This was corrected with a conditional to ensure it only happens when the last fetch was full. [PR #668](https://github.com/riverqueue/river/pull/668). + +## [0.14.0] - 2024-11-03 + +### Added + +- Expose `JobCancelError` and `JobSnoozeError` types to more easily facilitate testing. [PR #665](https://github.com/riverqueue/river/pull/665). + +### Changed + +- Tune the client to be more aggressive about fetching when it just fetched a full batch of jobs, or when it skipped its previous triggered fetch because it was already full. This should bring more consistent throughput to poll-only mode and in cases where there is a backlog of existing jobs but new ones aren't being actively inserted. This will result in increased fetch load on many installations, with the benefit of increased throughput. As before, `FetchCooldown` still limits how frequently these fetches can occur on each client and can be increased to reduce the amount of fetch querying. Thanks Chris Gaffney ([@gaffneyc](https://github.com/gaffneyc)) for the idea, initial implementation, and benchmarks. [PR #663](https://github.com/riverqueue/river/pull/663). + +### Fixed + +- `riverpgxv5` driver: `Hijack()` the underlying listener connection as soon as it is acquired from the `pgxpool.Pool` in order to prevent the pool from automatically closing it after it reaches its max age. A max lifetime makes sense in the context of a pool with many conns, but a long-lived listener does not need a max lifetime as long as it can ensure the conn remains healthy. [PR #661](https://github.com/riverqueue/river/pull/661). + +## [0.13.0] - 2024-10-07 + +⚠️ Version 0.13.0 removes the original advisory lock based unique jobs implementation that was deprecated in v0.12.0. See details in the note below or the v0.12.0 release notes. + +### Added + +- A middleware system was added for job insertion and execution, providing the ability to extract shared functionality across workers. Both `JobInsertMiddleware` and `WorkerMiddleware` can be configured globally on the `Client`, and `WorkerMiddleware` can also be added on a per-worker basis using the new `Middleware` method on `Worker[T]`. Middleware can be useful for logging, telemetry, or for building higher level abstractions on top of base River functionality. + + Despite the interface expansion, users should not encounter any breakage if they're embedding the `WorkerDefaults` type in their workers as recommended. [PR #632](https://github.com/riverqueue/river/pull/632). + +### Changed + +- **Breaking change:** The advisory lock unique jobs implementation which was deprecated in v0.12.0 has been removed. Users of that feature should first upgrade to v0.12.1 to ensure they don't see any warning logs about using the deprecated advisory lock uniqueness. The new, faster unique implementation will be used automatically as long as the `UniqueOpts.ByState` list hasn't been customized to remove [required states](https://riverqueue.com/docs/unique-jobs#unique-by-state) (`pending`, `scheduled`, `available`, and `running`). As of this release, customizing `ByState` without these required states returns an error. [PR #614](https://github.com/riverqueue/river/pull/614). +- Single job inserts are now unified under the hood to use the `InsertMany` bulk insert query. This should not be noticeable to users, and the unified code path will make it easier to build new features going forward. [PR #614](https://github.com/riverqueue/river/pull/614). + +### Fixed + +- Allow `river.JobCancel` to accept a `nil` error as input without panicking. [PR #634](https://github.com/riverqueue/river/pull/634). + +## [0.12.1] - 2024-09-26 + +### Changed + +- The `BatchCompleter` that marks jobs as completed can now batch database updates for _all_ states of jobs that have finished execution. Prior to this change, only `completed` jobs were batched into a single `UPDATE` call, while jobs moving to any other state used a single `UPDATE` per job. This change should significantly reduce database and pool contention on high volume system when jobs get retried, snoozed, cancelled, or discarded following execution. [PR #617](https://github.com/riverqueue/river/pull/617). + +### Fixed + +- Unique job changes from v0.12.0 / [PR #590](https://github.com/riverqueue/river/pull/590) introduced a bug with scheduled or retryable unique jobs where they could be considered in conflict with themselves and moved to `discarded` by mistake. There was also a possibility of a broken job scheduler if duplicate `retryable` unique jobs were attempted to be scheduled at the same time. The job scheduling query was corrected to address these issues along with missing test coverage. [PR #619](https://github.com/riverqueue/river/pull/619). + +## [0.12.0] - 2024-09-23 + +⚠️ Version 0.12.0 contains a new database migration, version 6. See [documentation on running River migrations](https://riverqueue.com/docs/migrations). If migrating with the CLI, make sure to update it to its latest version: + +```shell +go install github.com/riverqueue/river/cmd/river@latest +river migrate-up --database-url "$DATABASE_URL" +``` + +If not using River's internal migration system, the raw SQL can alternatively be dumped with: + +```shell +go install github.com/riverqueue/river/cmd/river@latest +river migrate-get --version 6 --up > river6.up.sql +river migrate-get --version 6 --down > river6.down.sql +``` + +The migration **includes a new index**. Users with a very large job table may want to consider raising the index separately using `CONCURRENTLY` (which must be run outside of a transaction), then run `river migrate-up` to finalize the process (it will tolerate an index that already exists): + +```sql +ALTER TABLE river_job ADD COLUMN unique_states BIT(8); + +CREATE UNIQUE INDEX CONCURRENTLY river_job_unique_idx ON river_job (unique_key) + WHERE unique_key IS NOT NULL + AND unique_states IS NOT NULL + AND river_job_state_in_bitmask(unique_states, state); +``` + +```shell +go install github.com/riverqueue/river/cmd/river@latest +river migrate-up --database-url "$DATABASE_URL" +``` + +## Added + +- `rivertest.WorkContext`, a test function that can be used to initialize a context to test a `JobArgs.Work` implementation that will have a client set to context for use with `river.ClientFromContext`. [PR #526](https://github.com/riverqueue/river/pull/526). +- A new `river migrate-list` command is available which lists available migrations and which version a target database is migrated to. [PR #534](https://github.com/riverqueue/river/pull/534). +- `river version` or `river --version` now prints River version information. [PR #537](https://github.com/riverqueue/river/pull/537). +- `Config.JobCleanerTimeout` was added to allow configuration of the job cleaner query timeout. In some deployments with millions of stale jobs, the cleaner may not be able to complete its query within the default 30 seconds. [PR #576](https://github.com/riverqueue/river/pull/576). + +### Changed + +⚠️ Version 0.12.0 has two small breaking changes, one for `InsertMany` and one in `rivermigrate`. As before, we try never to make breaking changes, but these ones were deemed worth it because of minimal impact and to help avoid panics. + +- **Breaking change:** `Client.InsertMany` / `InsertManyTx` now return the inserted rows rather than merely returning a count of the inserted rows. The new implementations no longer use Postgres' `COPY FROM` protocol in order to facilitate return values. + + Users who relied on the return count can merely wrap the returned rows in a `len()` to return to that behavior, or you can continue using the old APIs using their new names `InsertManyFast` and `InsertManyFastTx`. [PR #589](https://github.com/riverqueue/river/pull/589). + +- **Breaking change:** `rivermigrate.New` now returns a possible error along with a migrator. An error may be returned, for example, when a migration line is configured that doesn't exist. [PR #558](https://github.com/riverqueue/river/pull/558). + + ```go + # before + migrator := rivermigrate.New(riverpgxv5.New(dbPool), nil) + + # after + migrator, err := rivermigrate.New(riverpgxv5.New(dbPool), nil) + if err != nil { + // handle error + } + ``` + +- Unique jobs have been improved to allow bulk insertion of unique jobs via `InsertMany` / `InsertManyTx`, and to allow customizing the `ByState` list to add or remove certain states. This enables users to expand the set of unique states to also include `cancelled` and `discarded` jobs, or to remove `retryable` from uniqueness consideration. This updated implementation maintains the speed advantage of the newer index-backed uniqueness system, while allowing some flexibility in which job states. + + Unique jobs utilizing `ByArgs` can now also opt to have a subset of the job's arguments considered for uniqueness. For example, you could choose to consider only the `customer_id` field while ignoring the `trace_id` field: + + ```go + type MyJobArgs { + CustomerID string `json:"customer_id" river:"unique` + TraceID string `json:"trace_id"` + } + ``` + + Any fields considered in uniqueness are also sorted alphabetically in order to guarantee a consistent result, even if the encoded JSON isn't sorted consistently. For example `encoding/json` encodes struct fields in their defined order, so merely reordering struct fields would previously have been enough to cause a new job to not be considered identical to a pre-existing one with different JSON order. + + The `UniqueOpts` type also gains an `ExcludeKind` option for cases where uniqueness needs to be guaranteed across multiple job types. + + In-flight unique jobs using the previous designs will continue to be executed successfully with these changes, so there should be no need for downtime as part of the migration. However the v6 migration adds a new unique job index while also removing the old one, so users with in-flight unique jobs may also wish to avoid removing the old index until the new River release has been deployed in order to guarantee that jobs aren't duplicated by old River code once that index is removed. + + **Deprecated**: The original unique jobs implementation which relied on advisory locks has been deprecated, but not yet removed. The only way to trigger this old code path is with a single insert (`Insert`/`InsertTx`) and using `UniqueOpts.ByState` with a custom list of states that omits some of the now-required states for unique jobs. Specifically, `pending`, `scheduled`, `available`, and `running` can not be removed from the `ByState` list with the new implementation. These are included in the default list so only the places which customize this attribute need to be updated to opt into the new (much faster) unique jobs. The advisory lock unique implementation will be removed in an upcoming release, and until then emits warning level logs when it's used. + + [PR #590](https://github.com/riverqueue/river/pull/590). + +- **Deprecated**: The `MigrateTx` method of `rivermigrate` has been deprecated. It turns out there are certain combinations of schema changes which cannot be run within a single transaction, and the migrator now prefers to run each migration in its own transaction, one-at-a-time. `MigrateTx` will be removed in future version. + +- The migrator now produces a better error in case of a non-existent migration line including suggestions for known migration lines that are similar in name to the invalid one. [PR #558](https://github.com/riverqueue/river/pull/558). + +## Fixed + +- Fixed a panic that'd occur if `StopAndCancel` was invoked before a client was started. [PR #557](https://github.com/riverqueue/river/pull/557). +- A `PeriodicJobConstructor` should be able to return `nil` `JobArgs` if it wishes to not have any job inserted. However, this was either never working or was broken at some point. It's now fixed. Thanks [@semanser](https://github.com/semanser)! [PR #572](https://github.com/riverqueue/river/pull/572). +- Fixed a nil pointer exception if `Client.Subscribe` was called when the client had no configured workers (it still, panics with a more instructive error message now). [PR #599](https://github.com/riverqueue/river/pull/599). + +## [0.11.4] - 2024-08-20 + +### Fixed + +- Fixed release script that caused CLI to become uninstallable because its reference to `rivershared` wasn't updated. [PR #541](https://github.com/riverqueue/river/pull/541). + +## [0.11.3] - 2024-08-19 + +### Changed + +- Producer's logs are quieter unless jobs are actively being worked. [PR #529](https://github.com/riverqueue/river/pull/529). + +### Fixed + +- River CLI now accepts `postgresql://` URL schemes in addition to `postgres://`. [PR #532](https://github.com/riverqueue/river/pull/532). + +## [0.11.2] - 2024-08-08 + +### Fixed + +- Derive all internal contexts from user-provided `Client` context. This includes the job fetch context, notifier unlisten, and completer. [PR #514](https://github.com/riverqueue/river/pull/514). +- Lowered the `go` directives in `go.mod` to Go 1.21, which River aims to support. A more modern version of Go is specified with the `toolchain` directive. This should provide more flexibility on the minimum required Go version for programs importing River. [PR #522](https://github.com/riverqueue/river/pull/522). + +## [0.11.1] - 2024-08-05 + +### Fixed + +- `database/sql` driver: fix default value of `scheduled_at` for `InsertManyTx` when it is not specified in `InsertOpts`. [PR #504](https://github.com/riverqueue/river/pull/504). +- Change `ColumnExists` query to respect `search_path`, thereby allowing migrations to be runnable outside of default schema. [PR #505](https://github.com/riverqueue/river/pull/505). + +## [0.11.0] - 2024-08-02 + +### Added + +- Expose `Driver` on `Client` for additional River Pro integrations. This is not a stable API and should generally not be used by others. [PR #497](https://github.com/riverqueue/river/pull/497). + +## [0.10.2] - 2024-07-31 + +### Fixed + +- Include `pending` state in `JobListParams` by default so pending jobs are included in `JobList` / `JobListTx` results. [PR #477](https://github.com/riverqueue/river/pull/477). +- Quote strings when using `Client.JobList` functions with the `database/sql` driver. [PR #481](https://github.com/riverqueue/river/pull/481). +- Remove use of `filepath` for interacting with embedded migration files, fixing the migration CLI for Windows. [PR #485](https://github.com/riverqueue/river/pull/485). +- Respect `ScheduledAt` if set to a non-zero value by `JobArgsWithInsertOpts`. This allows for job arg definitions to utilize custom logic at the args level for determining when the job should be scheduled. [PR #487](https://github.com/riverqueue/river/pull/487). + +## [0.10.1] - 2024-07-23 + +### Fixed + +- Migration version 005 has been altered so that it can run even if the `river_migration` table isn't present, making it more friendly for projects that aren't using River's internal migration system. [PR #465](https://github.com/riverqueue/river/pull/465). + +## [0.10.0] - 2024-07-19 + +⚠️ Version 0.10.0 contains a new database migration, version 5. See [documentation on running River migrations](https://riverqueue.com/docs/migrations). If migrating with the CLI, make sure to update it to its latest version: + +```shell +go install github.com/riverqueue/river/cmd/river@latest +river migrate-up --database-url "$DATABASE_URL" +``` + +If not using River's internal migration system, the raw SQL can alternatively be dumped with: + +```shell +go install github.com/riverqueue/river/cmd/river@latest +river migrate-get --version 5 --up > river5.up.sql +river migrate-get --version 5 --down > river5.down.sql +``` + +The migration **includes a new index**. Users with a very large job table may want to consider raising the index separately using `CONCURRENTLY` (which must be run outside of a transaction), then run `river migrate-up` to finalize the process (it will tolerate an index that already exists): + +```sql +ALTER TABLE river_job + ADD COLUMN unique_key bytea; + +CREATE UNIQUE INDEX CONCURRENTLY river_job_kind_unique_key_idx ON river_job (kind, unique_key) WHERE unique_key IS NOT NULL; +``` + +```shell +go install github.com/riverqueue/river/cmd/river@latest +river migrate-up --database-url "$DATABASE_URL" +``` + +### Added + +- Fully functional driver for `database/sql` for use with packages like Bun and GORM. [PR #351](https://github.com/riverqueue/river/pull/351). +- Queues can be added after a client is initialized using `client.Queues().Add(queueName string, queueConfig QueueConfig)`. [PR #410](https://github.com/riverqueue/river/pull/410). +- Migration that adds a `line` column to the `river_migration` table so that it can support multiple migration lines. [PR #435](https://github.com/riverqueue/river/pull/435). +- `--line` flag added to the River CLI. [PR #454](https://github.com/riverqueue/river/pull/454). + +### Changed + +- Tags are now limited to 255 characters in length, and should match the regex `\A[\w][\w\-]+[\w]\z` (importantly, they can't contain commas). [PR #351](https://github.com/riverqueue/river/pull/351). +- Many info logging statements have been demoted to debug level. [PR #452](https://github.com/riverqueue/river/pull/452). +- `pending` is now part of the default set of unique job states. [PR #461](https://github.com/riverqueue/river/pull/461). + +## [0.9.0] - 2024-07-04 + +### Added + +- `Config.TestOnly` has been added. It disables various features in the River client like staggered maintenance service start that are useful in production, but may be somewhat harmful in tests because they make start/stop slower. [PR #414](https://github.com/riverqueue/river/pull/414). + +### Changed + +⚠️ Version 0.9.0 has a small breaking change in `ErrorHandler`. As before, we try never to make breaking changes, but this one was deemed quite important because `ErrorHandler` was fundamentally lacking important functionality. + +- **Breaking change:** Add stack trace to `ErrorHandler.HandlePanicFunc`. Fixing code only requires adding a new `trace string` argument to `HandlePanicFunc`. [PR #423](https://github.com/riverqueue/river/pull/423). + + ```go + # before + HandlePanic(ctx context.Context, job *rivertype.JobRow, panicVal any) *ErrorHandlerResult + + # after + HandlePanic(ctx context.Context, job *rivertype.JobRow, panicVal any, trace string) *ErrorHandlerResult + ``` + +### Fixed + +- Pausing or resuming a queue that was already paused or not paused respectively no longer returns `rivertype.ErrNotFound`. The same goes for pausing or resuming using the all queues string (`*`) when no queues are in the database (previously that also returned `rivertype.ErrNotFound`). [PR #408](https://github.com/riverqueue/river/pull/408). +- Fix a bug where periodic job constructors were only called once when adding the periodic job rather than being invoked every time the periodic job is scheduled. [PR #420](https://github.com/riverqueue/river/pull/420). + +## [0.8.0] - 2024-06-25 + +### Added + +- Add transaction variants for queue-related client functions: `QueueGetTx`, `QueueListTx`, `QueuePauseTx`, and `QueueResumeTx`. [PR #402](https://github.com/riverqueue/river/pull/402). + +### Fixed + +- Fix possible Client shutdown panics if the user-provided context is cancelled while jobs are still running. [PR #401](https://github.com/riverqueue/river/pull/401). + +## [0.7.0] - 2024-06-13 + +### Added + +- The default max attempts of 25 can now be customized on a per-client basis using `Config.MaxAttempts`. This is in addition to the ability to customize at the job type level with `JobArgs`, or on a per-job basis using `InsertOpts`. [PR #383](https://github.com/riverqueue/river/pull/383). +- Add `JobDelete` / `JobDeleteTx` APIs on `Client` to allow permanently deleting any job that's not currently running. [PR #390](https://github.com/riverqueue/river/pull/390). + +### Fixed + +- Fix `StopAndCancel` to not hang if called in parallel to an ongoing `Stop` call. [PR #376](https://github.com/riverqueue/river/pull/376). + +## [0.6.1] - 2024-05-21 + +### Fixed + +- River now considers per-worker timeout overrides when rescuing jobs so that jobs with a long custom timeout won't be rescued prematurely. [PR #350](https://github.com/riverqueue/river/pull/350). +- River CLI now exits with status 1 in the case of a problem with commands or flags, like an unknown command or missing required flag. [PR #363](https://github.com/riverqueue/river/pull/363). +- Fix migration version 4 (from 0.5.0) so that the up migration can be re-run after it was originally rolled back. [PR #364](https://github.com/riverqueue/river/pull/364). + +## [0.6.0] - 2024-05-08 + +### Added + +- `RequireNotInserted` test helper (in addition to the existing `RequireInserted`) that verifies that a job with matching conditions was _not_ inserted. [PR #237](https://github.com/riverqueue/river/pull/237). + +### Changed + +- The periodic job enqueuer now sets `scheduled_at` of inserted jobs to the more precise time of when they were scheduled to run, as opposed to when they were inserted. [PR #341](https://github.com/riverqueue/river/pull/341). + +### Fixed + +- Remove use of `github.com/lib/pq`, making it once again a test-only dependency. [PR #337](https://github.com/riverqueue/river/pull/337). + +## [0.5.0] - 2024-05-03 + +⚠️ Version 0.5.0 contains a new database migration, version 4. This migration is backward compatible with any River installation running the v3 migration. Be sure to run the v4 migration prior to deploying the code from this release. + +### Added + +- Add `pending` job state. This is currently unused, but will be used to build higher level functionality for staging jobs that are not yet ready to run (for some reason other than their scheduled time being in the future). Pending jobs will never be run or deleted and must first be moved to another state by external code. [PR #301](https://github.com/riverqueue/river/pull/301). +- Queue status tracking, pause and resume. [PR #301](https://github.com/riverqueue/river/pull/301). + + A useful operational lever is the ability to pause and resume a queue without shutting down clients. In addition to pause/resume being a feature request from [#54](https://github.com/riverqueue/river/pull/54), as part of the work on River's UI it's been useful to list out the active queues so that they can be displayed and manipulated. + + A new `river_queue` table is introduced in the v4 migration for this purpose. Upon startup, every producer in each River `Client` will make an `UPSERT` query to the database to either register the queue as being active, or if it already exists it will instead bump the timestamp to keep it active. This query will be run periodically in each producer as long as the `Client` is alive, even if the queue is paused. A separate query will delete/purge any queues which have not been active in awhile (currently fixed to 24 hours). + + `QueuePause` and `QueueResume` APIs have been introduced to `Client` pause and resume a single queue by name, or _all_ queues using the special `*` value. Each producer will watch for notifications on the relevant `LISTEN/NOTIFY` topic unless operating in poll-only mode, in which case they will periodically poll for changes to their queue record in the database. + +### Changed + +- Job insert notifications are now handled within application code rather than within the database using triggers. [PR #301](https://github.com/riverqueue/river/pull/301). + + The initial design for River utilized a trigger on job insert that issued notifications (`NOTIFY`) so that listening clients could quickly pick up the work if they were idle. While this is good for lowering latency, it does have the side effect of emitting a large amount of notifications any time there are lots of jobs being inserted. This adds overhead, particularly to high-throughput installations. + + To improve this situation and reduce overhead in high-throughput installations, the notifications have been refactored to be emitted at the application level. A client-level debouncer ensures that these notifications are not emitted more often than they could be useful. If a queue is due for an insert notification (on a particular Postgres schema), the notification is piggy-backed onto the insert query within the transaction. While this has the impact of increasing insert latency for a certain percentage of cases, the effect should be small. + + Additionally, initial releases of River did not properly scope notification topics within the global `LISTEN/NOTIFY` namespace. If two River installations were operating on the same Postgres database but within different schemas (search paths), their notifications would be emitted on a shared topic name. This is no longer the case and all notifications are prefixed with a `{schema_name}.` string. + +- Add `NOT NULL` constraints to the database for `river_job.args` and `river_job.metadata`. Normal code paths should never have allowed for null values any way, but this constraint further strengthens the guarantee. [PR #301](https://github.com/riverqueue/river/pull/301). +- Stricter constraint on `river_job.finalized_at` to ensure it is only set when paired with a finalized state (completed, discarded, cancelled). Normal code paths should never have allowed for invalid values any way, but this constraint further strengthens the guarantee. [PR #301](https://github.com/riverqueue/river/pull/301). + +## [0.4.1] - 2024-04-22 + +### Fixed + +- Update job state references in `./cmd/river` and some documentation to `rivertype`. Thanks Danny Hermes (@dhermes)! 🙏🏻 [PR #315](https://github.com/riverqueue/river/pull/315). + +## [0.4.0] - 2024-04-20 + +### Changed + +⚠️ Version 0.4.0 has a number of small breaking changes which we've decided to release all as part of a single version. More breaking changes in one release is inconvenient, but we've tried to coordinate them in hopes that any future breaking changes will be non-existent or very rare. All changes will get picked up by the Go compiler, and each one should be quite easy to fix. The changes don't apply to any of the most common core APIs, and likely many projects won't have to change any code. + +- **Breaking change:** There are a number of small breaking changes in the job list API using `JobList`/`JobListTx`: + - Now support querying jobs by a list of Job Kinds and States. Also allows for filtering by specific timestamp values. Thank you Jos Kraaijeveld (@thatjos)! 🙏🏻 [PR #236](https://github.com/riverqueue/river/pull/236). + - Job listing now defaults to ordering by job ID (`JobListOrderByID`) instead of a job timestamp dependent on requested job state. The previous ordering behavior is still available with `NewJobListParams().OrderBy(JobListOrderByTime, SortOrderAsc)`. [PR #307](https://github.com/riverqueue/river/pull/307). + - The function `JobListCursorFromJob` no longer needs a sort order parameter. Instead, sort order is determined based on the job list parameters that the cursor is subsequently used with. [PR #307](https://github.com/riverqueue/river/pull/307). +- **Breaking change:** Client `Insert` and `InsertTx` functions now return a `JobInsertResult` struct instead of a `JobRow`. This allows the result to include metadata like the new `UniqueSkippedAsDuplicate` property, so callers can tell whether an inserted job was skipped due to unique constraint. [PR #292](https://github.com/riverqueue/river/pull/292). +- **Breaking change:** Client `InsertMany` and `InsertManyTx` now return number of jobs inserted as `int` instead of `int64`. This change was made to make the type in use a little more idiomatic. [PR #293](https://github.com/riverqueue/river/pull/293). +- **Breaking change:** `river.JobState*` type aliases have been removed. All job state constants should be accessed through `rivertype.JobState*` instead. [PR #300](https://github.com/riverqueue/river/pull/300). + +See also the [0.4.0 release blog post](https://riverqueue.com/blog/a-few-breaking-changes) with code samples and rationale behind various changes. + +## [0.3.0] - 2024-04-15 + +### Added + +- The River client now supports "poll only" mode with `Config.PollOnly` which makes it avoid issuing `LISTEN` statements to wait for new events like a leadership resignation or new job available. The program instead polls periodically to look for changes. A leader resigning or a new job being available will be noticed less quickly, but `PollOnly` potentially makes River operable on systems without listen/notify support, like PgBouncer operating in transaction pooling mode. [PR #281](https://github.com/riverqueue/river/pull/281). +- Added `rivertype.JobStates()` that returns the full list of possible job states. [PR #297](https://github.com/riverqueue/river/pull/297). + +## [0.2.0] - 2024-03-28 + +### Added + +- New periodic jobs can now be added after a client's already started using `Client.PeriodicJobs().Add()` and removed with `Remove()`. [PR #288](https://github.com/riverqueue/river/pull/288). + +### Changed + +- The level of some of River's common log statements has changed, most often demoting `info` statements to `debug` so that `info`-level logging is overall less verbose. [PR #275](https://github.com/riverqueue/river/pull/275). + +### Fixed + +- Fixed a bug in the (log-only for now) reindexer service in which it might repeat its work loop multiple times unexpectedly while stopping. [PR #280](https://github.com/riverqueue/river/pull/280). +- Periodic job enqueuer now bases next run times on each periodic job's last target run time, instead of the time at which the enqueuer is currently running. This is a small difference that will be unnoticeable for most purposes, but makes scheduling of jobs with short cron frequencies a little more accurate. [PR #284](https://github.com/riverqueue/river/pull/284). +- Fixed a bug in the elector in which it was possible for a resigning, but not completely stopped, elector to reelect despite having just resigned. [PR #286](https://github.com/riverqueue/river/pull/286). + +## [0.1.0] - 2024-03-17 + +Although it comes with a number of improvements, there's nothing particularly notable about version 0.1.0. Until now we've only been incrementing the patch version given the project's nascent nature, but from here on we'll try to adhere more closely to semantic versioning, using the patch version for bug fixes, and incrementing the minor version when new functionality is added. + +### Added + +- The River CLI now supports `river bench` to benchmark River's job throughput against a database. [PR #254](https://github.com/riverqueue/river/pull/254). +- The River CLI now has a `river migrate-get` command to dump SQL for River migrations for use in alternative migration frameworks. Use it like `river migrate-get --up --version 3 > version3.up.sql`. [PR #273](https://github.com/riverqueue/river/pull/273). +- The River CLI's `migrate-down` and `migrate-up` options get two new options for `--dry-run` and `--show-sql`. They can be combined to easily run a preflight check on a River upgrade to see which migration commands would be run on a database, but without actually running them. [PR #273](https://github.com/riverqueue/river/pull/273). +- The River client gets a new `Client.SubscribeConfig` function that lets a subscriber specify the maximum size of their subscription channel. [PR #258](https://github.com/riverqueue/river/pull/258). + +### Changed + +- River uses a new job completer that batches up completion work so that large numbers of them can be performed more efficiently. In a purely synthetic (i.e. mostly unrealistic) benchmark, River's job throughput increases ~4.5x. [PR #258](https://github.com/riverqueue/river/pull/258). +- Changed default client IDs to be a combination of hostname and the time which the client started. This can still be changed by specifying `Config.ID`. [PR #255](https://github.com/riverqueue/river/pull/255). +- Notifier refactored for better robustness and testability. [PR #253](https://github.com/riverqueue/river/pull/253). + +## [0.0.25] - 2024-03-01 + +### Fixed + +- Fixed a problem in `riverpgxv5`'s `Listener` where it wouldn't unset an internal connection if `Close` returned an error, making the listener not reusable. Thanks @mfrister for pointing this one out! [PR #246](https://github.com/riverqueue/river/pull/246). + +## [0.0.24] - 2024-02-29 + +### Fixed + +- Fixed a memory leak caused by not always cancelling the context used to enable jobs to be cancelled remotely. [PR #243](https://github.com/riverqueue/river/pull/243). + +## [0.0.23] - 2024-02-29 + +### Added + +- `JobListParams.Kinds()` has been added so that jobs can now be listed by kind. [PR #212](https://github.com/riverqueue/river/pull/212). + +### Changed + +- The underlying driver system's been entirely revamped so that River's non-test code is now decoupled from `pgx/v5`. This will allow additional drivers to be implemented, although there are no additional ones for now. [PR #212](https://github.com/riverqueue/river/pull/212). + +### Fixed + +- Fixed a memory leak caused by allocating a new random source on every job execution. Thank you @shawnstephens for reporting ❤️ [PR #240](https://github.com/riverqueue/river/pull/240). +- Fix a problem where `JobListParams.Queues()` didn't filter correctly based on its arguments. [PR #212](https://github.com/riverqueue/river/pull/212). +- Fix a problem in `DebouncedChan` where it would fire on its "out" channel too often when it was being signaled continuously on its "in" channel. This would have caused work to be fetched more often than intended in busy systems. [PR #222](https://github.com/riverqueue/river/pull/222). + +## [0.0.22] - 2024-02-19 + +### Fixed + +- Brings in another leadership election fix similar to #217 in which a TTL equal to the elector's run interval plus a configured TTL padding is also used for the initial attempt to gain leadership (#217 brought it in for reelection only). [PR #219](https://github.com/riverqueue/river/pull/219). + +## [0.0.21] - 2024-02-19 + +### Changed + +- Tweaked behavior of `JobRetry` so that it does actually update the `ScheduledAt` time of the job in all cases where the job is actually being rescheduled. As before, jobs which are already available with a past `ScheduledAt` will not be touched by this query so that they retain their place in line. [PR #211](https://github.com/riverqueue/river/pull/211). + +### Fixed + +- Fixed a leadership re-election issue that was exposed by the fix in #199. Because we were internally using the same TTL for both an internal timer/ticker and the database update to set the new leader expiration time, a leader wasn't guaranteed to successfully re-elect itself even under normal operation. [PR #217](https://github.com/riverqueue/river/pull/217). + +## [0.0.20] - 2024-02-14 + +### Added + +- Added an `ID` setting to the `Client` `Config` type to allow users to override client IDs with their own naming convention. Expose the client ID programmatically (in case it's generated) in a new `Client.ID()` method. [PR #206](https://github.com/riverqueue/river/pull/206). + +### Fixed + +- Fix a leadership re-election query bug that would cause past leaders to think they were continuing to win elections. [PR #199](https://github.com/riverqueue/river/pull/199). + +## [0.0.19] - 2024-02-10 + +### Added + +- Added `JobGet` and `JobGetTx` to the `Client` to enable easily fetching a single job row from code for introspection. [PR #186]. +- Added `JobRetry` and `JobRetryTx` to the `Client` to enable a job to be retried immediately, even if it has already completed, been cancelled, or been discarded. [PR #190]. + +### Changed + +- Validate queue name on job insertion. Allow queue names with hyphen separators in addition to underscore. [PR #184](https://github.com/riverqueue/river/pull/184). + +## [0.0.18] - 2024-01-25 + +### Fixed + +- Remove a debug statement from periodic job enqueuer that was accidentally left in. [PR #176](https://github.com/riverqueue/river/pull/176). + +## [0.0.17] - 2024-01-22 + +### Added + +- Added `JobCancel` and `JobCancelTx` to the `Client` to enable cancellation of jobs. [PR #141](https://github.com/riverqueue/river/pull/141) and [PR #152](https://github.com/riverqueue/river/pull/152). +- Added `ClientFromContext` and `ClientFromContextSafely` helpers to extract the `Client` from the worker's context where it is now available to workers. This simplifies making the River client available within your workers for i.e. enqueueing additional jobs. [PR #145](https://github.com/riverqueue/river/pull/145). +- Add `JobList` API for listing jobs. [PR #117](https://github.com/riverqueue/river/pull/117). +- Added `river validate` command which fails with a non-zero exit code unless all migrations are applied. [PR #170](https://github.com/riverqueue/river/pull/170). + +### Changed + +- For short `JobSnooze` times (smaller than the scheduler's run interval) put the job straight into an `available` state with the specified `scheduled_at` time. This avoids an artificially long delay waiting for the next scheduler run. [PR #162](https://github.com/riverqueue/river/pull/162). + +### Fixed + +- Fixed incorrect default value handling for `ScheduledAt` option with `InsertMany` / `InsertManyTx`. [PR #149](https://github.com/riverqueue/river/pull/149). +- Add missing `t.Helper()` calls in `rivertest` internal functions that caused it to report itself as the site of a test failure. [PR #151](https://github.com/riverqueue/river/pull/151). +- Fixed problem where job uniqueness wasn't being respected when used in conjunction with periodic jobs. [PR #168](https://github.com/riverqueue/river/pull/168). + +## [0.0.16] - 2024-01-06 + +### Changed + +- Calls to `Stop` error if the client hasn't been started yet. [PR #138](https://github.com/riverqueue/river/pull/138). + +### Fixed + +- Fix typo in leadership resignation query to ensure faster new leader takeover. [PR #134](https://github.com/riverqueue/river/pull/134). +- Elector now uses the same `log/slog` instance configured by its parent client. [PR #137](https://github.com/riverqueue/river/pull/137). +- Notifier now uses the same `log/slog` instance configured by its parent client. [PR #140](https://github.com/riverqueue/river/pull/140). + +## [0.0.15] - 2023-12-21 + +### Fixed + +- Ensure `ScheduledAt` is respected on `InsertManyTx`. [PR #121](https://github.com/riverqueue/river/pull/121). + +## [0.0.14] - 2023-12-13 + +### Fixed + +- River CLI `go.sum` entries fixed for 0.0.13 release. + +## [0.0.13] - 2023-12-12 + +### Added + +- Added `riverdriver/riverdatabasesql` driver to enable River Go migrations through Go's built in `database/sql` package. [PR #98](https://github.com/riverqueue/river/pull/98). + +### Changed + +- Errored jobs that have a very short duration before their next retry (<5 seconds) are set to `available` immediately instead of being made `scheduled` and having to wait for the scheduler to make a pass to make them workable. [PR #105](https://github.com/riverqueue/river/pull/105). +- `riverdriver` becomes its own submodule. It contains types that `riverdriver/riverdatabasesql` and `riverdriver/riverpgxv5` need to reference. [PR #98](https://github.com/riverqueue/river/pull/98). +- The `river/cmd/river` CLI has been made its own Go module. This is possible now that it uses the exported `river/rivermigrate` API, and will help with project maintainability. [PR #107](https://github.com/riverqueue/river/pull/107). + +## [0.0.12] - 2023-12-02 + +### Added + +- Added `river/rivermigrate` package to enable migrations from Go code as an alternative to using the CLI. PR #67. + +## [0.0.11] - 2023-12-02 + +### Added + +- `Stop` and `StopAndCancel` have been changed to respect the provided context argument. When that context is cancelled or times out, those methods will now immediately return with the context's error, even if the Client's shutdown has not yet completed. Apps may need to adjust their graceful shutdown logic to account for this. PR #79. + +### Changed + +- `NewClient` no longer errors if it was provided a workers bundle with zero workers. Instead, that check's been moved to `Client.Start` instead. This allows adding workers to a bundle that'd like to reference a River client by letting `AddWorker` be invoked after a client reference is available from `NewClient`. PR #87. + +## [0.0.10] - 2023-11-26 + +### Added + +- Added `Example_scheduledJob`, demonstrating how to schedule a job to be run in the future. +- Added `Stopped` method to `Client` to make it easier to wait for graceful shutdown to complete. + +### Fixed + +- Fixed a panic in the periodic job enqueuer caused by sometimes trying to reset a `time.Ticker` with a negative or zero duration. Fixed in PR #73. + +### Changed + +- `DefaultClientRetryPolicy`: calculate the next attempt based on the current time instead of the time the prior attempt began. + +## [0.0.9] - 2023-11-23 + +### Fixed + +- **DATABASE MIGRATION**: Database schema v3 was introduced in v0.0.8 and contained an obvious flaw preventing it from running against existing tables. This migration was altered to execute the migration in multiple steps. + +## [0.0.8] - 2023-11-21 + +### Changed + +- License changed from LGPLv3 to MPL-2.0. +- **DATABASE MIGRATION**: Database schema v3, alter river_job tags column to set a default of `[]` and add not null constraint. + +## [0.0.7] - 2023-11-20 + +### Changed + +- Constants renamed so that adjectives like `Default` and `Min` become suffixes instead of prefixes. So for example, `DefaultFetchCooldown` becomes `FetchCooldownDefault`. +- Rename `AttemptError.Num` to `AttemptError.Attempt` to better fit with the name of `JobRow.Attempt`. +- Document `JobState`, `AttemptError`, and all fields its fields. +- A `NULL` tags value read from a database job is left as `[]string(nil)` on `JobRow.Tags` rather than a zero-element slice of `[]string{}`. `append` and `len` both work on a `nil` slice, so this should be functionally identical. + +## [0.0.6] - 2023-11-19 + +### Changed + +- `JobRow`, `JobState`, and other related types move into `river/rivertype` so they can more easily be shared amongst packages. Most of the River API doesn't change because `JobRow` is embedded on `river.Job`, which doesn't move. + +## [0.0.5] - 2023-11-19 + +### Changed + +- Remove `replace` directive from the project's `go.mod` so that it's possible to install River CLI with `@latest`. + +## [0.0.4] - 2023-11-17 + +### Changed + +- Allow River clients to be created with a driver with `nil` database pool for use in testing. +- Update River test helpers API to use River drivers like `riverdriver/riverpgxv5` to make them agnostic to the third party database package in use. +- Document `Config.JobTimeout`'s default value. +- Functionally disable the `Reindexer` queue maintenance service. It'd previously only operated on currently unused indexes anyway, indexes probably do _not_ need to be rebuilt except under fairly rare circumstances, and it needs more work to make sure it's shored up against edge cases like indexes that fail to rebuild before a client restart. + +## [0.0.3] - 2023-11-13 + +### Changed + +- Fix license detection issues with `riverdriver/riverpgxv5` submodule. +- Ensure that river requires the `riverpgxv5` module with the same version. + +## [0.0.2] - 2023-11-13 + +### Changed + +- Pin own `riverpgxv5` dependency to v0.0.1 and make it a direct locally-replaced dependency. This should allow projects to import versioned deps of both river and `riverpgxv5`. + +## [0.0.1] - 2023-11-12 + +### Added + +- This is the initial prerelease of River. diff --git a/vendor/github.com/riverqueue/river/LICENSE b/vendor/github.com/riverqueue/river/LICENSE new file mode 100644 index 0000000000..2f8ed188e8 --- /dev/null +++ b/vendor/github.com/riverqueue/river/LICENSE @@ -0,0 +1,374 @@ +Mozilla Public License Version 2.0 +================================== + +1. Definitions +-------------- + +1.1. "Contributor" + means each individual or legal entity that creates, contributes to + the creation of, or owns Covered Software. + +1.2. "Contributor Version" + means the combination of the Contributions of others (if any) used + by a Contributor and that particular Contributor's Contribution. + +1.3. "Contribution" + means Covered Software of a particular Contributor. + +1.4. "Covered Software" + means Source Code Form to which the initial Contributor has attached + the notice in Exhibit A, the Executable Form of such Source Code + Form, and Modifications of such Source Code Form, in each case + including portions thereof. + +1.5. "Incompatible With Secondary Licenses" + means + + (a) that the initial Contributor has attached the notice described + in Exhibit B to the Covered Software; or + + (b) that the Covered Software was made available under the terms of + version 1.1 or earlier of the License, but not also under the + terms of a Secondary License. + +1.6. "Executable Form" + means any form of the work other than Source Code Form. + +1.7. "Larger Work" + means a work that combines Covered Software with other material, in + a separate file or files, that is not Covered Software. + +1.8. "License" + means this document. + +1.9. "Licensable" + means having the right to grant, to the maximum extent possible, + whether at the time of the initial grant or subsequently, any and + all of the rights conveyed by this License. + +1.10. "Modifications" + means any of the following: + + (a) any file in Source Code Form that results from an addition to, + deletion from, or modification of the contents of Covered + Software; or + + (b) any new file in Source Code Form that contains any Covered + Software. + +1.11. "Patent Claims" of a Contributor + means any patent claim(s), including without limitation, method, + process, and apparatus claims, in any patent Licensable by such + Contributor that would be infringed, but for the grant of the + License, by the making, using, selling, offering for sale, having + made, import, or transfer of either its Contributions or its + Contributor Version. + +1.12. "Secondary License" + means either the GNU General Public License, Version 2.0, the GNU + Lesser General Public License, Version 2.1, the GNU Affero General + Public License, Version 3.0, or any later versions of those + licenses. + +1.13. "Source Code Form" + means the form of the work preferred for making modifications. + +1.14. "You" (or "Your") + means an individual or a legal entity exercising rights under this + License. For legal entities, "You" includes any entity that + controls, is controlled by, or is under common control with You. For + purposes of this definition, "control" means (a) the power, direct + or indirect, to cause the direction or management of such entity, + whether by contract or otherwise, or (b) ownership of more than + fifty percent (50%) of the outstanding shares or beneficial + ownership of such entity. + +2. License Grants and Conditions +-------------------------------- + +2.1. Grants + +Each Contributor hereby grants You a world-wide, royalty-free, +non-exclusive license: + +(a) under intellectual property rights (other than patent or trademark) + Licensable by such Contributor to use, reproduce, make available, + modify, display, perform, distribute, and otherwise exploit its + Contributions, either on an unmodified basis, with Modifications, or + as part of a Larger Work; and + +(b) under Patent Claims of such Contributor to make, use, sell, offer + for sale, have made, import, and otherwise transfer either its + Contributions or its Contributor Version. + +2.2. Effective Date + +The licenses granted in Section 2.1 with respect to any Contribution +become effective for each Contribution on the date the Contributor first +distributes such Contribution. + +2.3. Limitations on Grant Scope + +The licenses granted in this Section 2 are the only rights granted under +this License. No additional rights or licenses will be implied from the +distribution or licensing of Covered Software under this License. +Notwithstanding Section 2.1(b) above, no patent license is granted by a +Contributor: + +(a) for any code that a Contributor has removed from Covered Software; + or + +(b) for infringements caused by: (i) Your and any other third party's + modifications of Covered Software, or (ii) the combination of its + Contributions with other software (except as part of its Contributor + Version); or + +(c) under Patent Claims infringed by Covered Software in the absence of + its Contributions. + +This License does not grant any rights in the trademarks, service marks, +or logos of any Contributor (except as may be necessary to comply with +the notice requirements in Section 3.4). + +2.4. Subsequent Licenses + +No Contributor makes additional grants as a result of Your choice to +distribute the Covered Software under a subsequent version of this +License (see Section 10.2) or under the terms of a Secondary License (if +permitted under the terms of Section 3.3). + +2.5. Representation + +Each Contributor represents that the Contributor believes its +Contributions are its original creation(s) or it has sufficient rights +to grant the rights to its Contributions conveyed by this License. + +2.6. Fair Use + +This License is not intended to limit any rights You have under +applicable copyright doctrines of fair use, fair dealing, or other +equivalents. + +2.7. Conditions + +Sections 3.1, 3.2, 3.3, and 3.4 are conditions of the licenses granted +in Section 2.1. + +3. Responsibilities +------------------- + +3.1. Distribution of Source Form + +All distribution of Covered Software in Source Code Form, including any +Modifications that You create or to which You contribute, must be under +the terms of this License. You must inform recipients that the Source +Code Form of the Covered Software is governed by the terms of this +License, and how they can obtain a copy of this License. You may not +attempt to alter or restrict the recipients' rights in the Source Code +Form. + +3.2. Distribution of Executable Form + +If You distribute Covered Software in Executable Form then: + +(a) such Covered Software must also be made available in Source Code + Form, as described in Section 3.1, and You must inform recipients of + the Executable Form how they can obtain a copy of such Source Code + Form by reasonable means in a timely manner, at a charge no more + than the cost of distribution to the recipient; and + +(b) You may distribute such Executable Form under the terms of this + License, or sublicense it under different terms, provided that the + license for the Executable Form does not attempt to limit or alter + the recipients' rights in the Source Code Form under this License. + +3.3. Distribution of a Larger Work + +You may create and distribute a Larger Work under terms of Your choice, +provided that You also comply with the requirements of this License for +the Covered Software. If the Larger Work is a combination of Covered +Software with a work governed by one or more Secondary Licenses, and the +Covered Software is not Incompatible With Secondary Licenses, this +License permits You to additionally distribute such Covered Software +under the terms of such Secondary License(s), so that the recipient of +the Larger Work may, at their option, further distribute the Covered +Software under the terms of either this License or such Secondary +License(s). + +3.4. Notices + +You may not remove or alter the substance of any license notices +(including copyright notices, patent notices, disclaimers of warranty, +or limitations of liability) contained within the Source Code Form of +the Covered Software, except that You may alter any license notices to +the extent required to remedy known factual inaccuracies. + +3.5. Application of Additional Terms + +You may choose to offer, and to charge a fee for, warranty, support, +indemnity or liability obligations to one or more recipients of Covered +Software. However, You may do so only on Your own behalf, and not on +behalf of any Contributor. You must make it absolutely clear that any +such warranty, support, indemnity, or liability obligation is offered by +You alone, and You hereby agree to indemnify every Contributor for any +liability incurred by such Contributor as a result of warranty, support, +indemnity or liability terms You offer. You may include additional +disclaimers of warranty and limitations of liability specific to any +jurisdiction. + +4. Inability to Comply Due to Statute or Regulation +--------------------------------------------------- + +If it is impossible for You to comply with any of the terms of this +License with respect to some or all of the Covered Software due to +statute, judicial order, or regulation then You must: (a) comply with +the terms of this License to the maximum extent possible; and (b) +describe the limitations and the code they affect. Such description must +be placed in a text file included with all distributions of the Covered +Software under this License. Except to the extent prohibited by statute +or regulation, such description must be sufficiently detailed for a +recipient of ordinary skill to be able to understand it. + +5. Termination +-------------- + +5.1. The rights granted under this License will terminate automatically +if You fail to comply with any of its terms. However, if You become +compliant, then the rights granted under this License from a particular +Contributor are reinstated (a) provisionally, unless and until such +Contributor explicitly and finally terminates Your grants, and (b) on an +ongoing basis, if such Contributor fails to notify You of the +non-compliance by some reasonable means prior to 60 days after You have +come back into compliance. Moreover, Your grants from a particular +Contributor are reinstated on an ongoing basis if such Contributor +notifies You of the non-compliance by some reasonable means, this is the +first time You have received notice of non-compliance with this License +from such Contributor, and You become compliant prior to 30 days after +Your receipt of the notice. + +5.2. If You initiate litigation against any entity by asserting a patent +infringement claim (excluding declaratory judgment actions, +counter-claims, and cross-claims) alleging that a Contributor Version +directly or indirectly infringes any patent, then the rights granted to +You by any and all Contributors for the Covered Software under Section +2.1 of this License shall terminate. + +5.3. In the event of termination under Sections 5.1 or 5.2 above, all +end user license agreements (excluding distributors and resellers) which +have been validly granted by You or Your distributors under this License +prior to termination shall survive termination. + +************************************************************************ +* * +* 6. Disclaimer of Warranty * +* ------------------------- * +* * +* Covered Software is provided under this License on an "as is" * +* basis, without warranty of any kind, either expressed, implied, or * +* statutory, including, without limitation, warranties that the * +* Covered Software is free of defects, merchantable, fit for a * +* particular purpose or non-infringing. The entire risk as to the * +* quality and performance of the Covered Software is with You. * +* Should any Covered Software prove defective in any respect, You * +* (not any Contributor) assume the cost of any necessary servicing, * +* repair, or correction. This disclaimer of warranty constitutes an * +* essential part of this License. No use of any Covered Software is * +* authorized under this License except under this disclaimer. * +* * +************************************************************************ + +************************************************************************ +* * +* 7. Limitation of Liability * +* -------------------------- * +* * +* Under no circumstances and under no legal theory, whether tort * +* (including negligence), contract, or otherwise, shall any * +* Contributor, or anyone who distributes Covered Software as * +* permitted above, be liable to You for any direct, indirect, * +* special, incidental, or consequential damages of any character * +* including, without limitation, damages for lost profits, loss of * +* goodwill, work stoppage, computer failure or malfunction, or any * +* and all other commercial damages or losses, even if such party * +* shall have been informed of the possibility of such damages. This * +* limitation of liability shall not apply to liability for death or * +* personal injury resulting from such party's negligence to the * +* extent applicable law prohibits such limitation. Some * +* jurisdictions do not allow the exclusion or limitation of * +* incidental or consequential damages, so this exclusion and * +* limitation may not apply to You. * +* * +************************************************************************ + +8. Litigation +------------- + +Any litigation relating to this License may be brought only in the +courts of a jurisdiction where the defendant maintains its principal +place of business and such litigation shall be governed by laws of that +jurisdiction, without reference to its conflict-of-law provisions. +Nothing in this Section shall prevent a party's ability to bring +cross-claims or counter-claims. + +9. Miscellaneous +---------------- + +This License represents the complete agreement concerning the subject +matter hereof. If any provision of this License is held to be +unenforceable, such provision shall be reformed only to the extent +necessary to make it enforceable. Any law or regulation which provides +that the language of a contract shall be construed against the drafter +shall not be used to construe this License against a Contributor. + +10. Versions of the License +--------------------------- + +10.1. New Versions + +Mozilla Foundation is the license steward. Except as provided in Section +10.3, no one other than the license steward has the right to modify or +publish new versions of this License. Each version will be given a +distinguishing version number. + +10.2. Effect of New Versions + +You may distribute the Covered Software under the terms of the version +of the License under which You originally received the Covered Software, +or under the terms of any subsequent version published by the license +steward. + +10.3. Modified Versions + +If you create software not governed by this License, and you want to +create a new license for such software, you may create and use a +modified version of this License if you rename the license and remove +any references to the name of the license steward (except to note that +such modified license differs from this License). + +10.4. Distributing Source Code Form that is Incompatible With Secondary +Licenses + +If You choose to distribute Source Code Form that is Incompatible With +Secondary Licenses under the terms of this version of the License, the +notice described in Exhibit B of this License must be attached. + +Exhibit A - Source Code Form License Notice +------------------------------------------- + + This Source Code Form is subject to the terms of the Mozilla Public + License, v. 2.0. If a copy of the MPL was not distributed with this + file, You can obtain one at http://mozilla.org/MPL/2.0/. + +If it is not possible or desirable to put the notice in a particular +file, then You may include the notice in a location (such as a LICENSE +file in a relevant directory) where a recipient would be likely to look +for such a notice. + +You may add additional accurate notices of copyright ownership. + +Exhibit B - "Incompatible With Secondary Licenses" Notice +--------------------------------------------------------- + + This Source Code Form is "Incompatible With Secondary Licenses", as + defined by the Mozilla Public License, v. 2.0. + diff --git a/vendor/github.com/riverqueue/river/Makefile b/vendor/github.com/riverqueue/river/Makefile new file mode 100644 index 0000000000..ce82cd4fdf --- /dev/null +++ b/vendor/github.com/riverqueue/river/Makefile @@ -0,0 +1,105 @@ +.DEFAULT_GOAL := help + +.PHONY: db/reset +db/reset: ## Drop, create, and migrate dev and test databases +db/reset: db/reset/dev +db/reset: db/reset/test + +.PHONY: db/reset/dev +db/reset/dev: ## Drop, create, and migrate dev database + dropdb river_dev --force --if-exists + createdb river_dev + cd cmd/river && go run . migrate-up --database-url "postgres://localhost/river_dev" + +.PHONY: db/reset/test +db/reset/test: ## Drop, create, and migrate test databases + go run ./internal/cmd/testdbman reset + +.PHONY: generate +generate: ## Generate generated artifacts +generate: generate/migrations +generate: generate/sqlc + +.PHONY: generate/migrations +generate/migrations: ## Sync changes of pgxv5 migrations to database/sql + rsync -au --delete "riverdriver/riverpgxv5/migration/" "riverdriver/riverdatabasesql/migration/" + +.PHONY: generate/sqlc +generate/sqlc: ## Generate sqlc + cd riverdriver/riverdatabasesql/internal/dbsqlc && sqlc generate + cd riverdriver/riverpgxv5/internal/dbsqlc && sqlc generate + cd riverdriver/riversqlite/internal/dbsqlc && sqlc generate + +# Looks at comments using ## on targets and uses them to produce a help output. +.PHONY: help +help: ALIGN=22 +help: ## Print this message + @awk -F '::? .*## ' -- "/^[^':]+::? .*## /"' { printf "'$$(tput bold)'%-$(ALIGN)s'$$(tput sgr0)' %s\n", $$1, $$2 }' $(MAKEFILE_LIST) + +# Each directory of a submodule in the Go workspace. Go commands provide no +# built-in way to run for all workspace submodules. Add a new submodule to the +# workspace with `go work use ./driver/new`. +submodules := $(shell go list -f '{{.Dir}}' -m) + +ITERATIONS ?= 100 + +# Definitions of following tasks look ugly, but they're done this way because to +# produce the best/most comprehensible output by far (e.g. compared to a shell +# loop). +.PHONY: lint +lint:: ## Run linter (golangci-lint) for all submodules +define lint-target + lint:: ; cd $1 && golangci-lint run --fix +endef +$(foreach mod,$(submodules),$(eval $(call lint-target,$(mod)))) + +.PHONY: test +test:: ## Run test suite for all submodules +define test-target + test:: ; cd $1 && go test ./... -timeout 2m +endef +$(foreach mod,$(submodules),$(eval $(call test-target,$(mod)))) + +.PHONY: test/race +test/race:: ## Run test suite for all submodules with race detector +define test-race-target + test/race:: ; cd $1 && go test ./... -race -timeout 2m +endef +$(foreach mod,$(submodules),$(eval $(call test-race-target,$(mod)))) + +.PHONY: bench +bench:: ## Run benchmarks in each submodule (ITERATIONS=100) +define bench-target + bench:: ; cd $1 && go test -bench=. -benchtime=$(ITERATIONS)x -run=a^ ./... +endef +$(foreach mod,$(submodules),$(eval $(call bench-target,$(mod)))) + +.PHONY: tidy +tidy:: ## Run `go mod tidy` for all submodules +define tidy-target + tidy:: ; cd $1 && go mod tidy +endef +$(foreach mod,$(submodules),$(eval $(call tidy-target,$(mod)))) + +.PHONY: update-mod-go +update-mod-go: ## Update `go`/`toolchain` directives in all submodules to match `go.work` + go run ./rivershared/cmd/update-mod-go ./go.work + +.PHONY: update-mod-version +update-mod-version: ## Update River packages in all submodules to $VERSION + PACKAGE_PREFIX="github.com/riverqueue/river" go run ./rivershared/cmd/update-mod-version ./go.work + +.PHONY: verify +verify: ## Verify generated artifacts +verify: verify/migrations +verify: verify/sqlc + +.PHONY: verify/migrations +verify/migrations: ## Verify synced migrations + diff -qr riverdriver/riverpgxv5/migration riverdriver/riverdatabasesql/migration + +.PHONY: verify/sqlc +verify/sqlc: ## Verify generated sqlc + cd riverdriver/riverdatabasesql/internal/dbsqlc && sqlc diff + cd riverdriver/riverpgxv5/internal/dbsqlc && sqlc diff + cd riverdriver/riversqlite/internal/dbsqlc && sqlc diff diff --git a/vendor/github.com/riverqueue/river/client.go b/vendor/github.com/riverqueue/river/client.go new file mode 100644 index 0000000000..bf07bfc0f3 --- /dev/null +++ b/vendor/github.com/riverqueue/river/client.go @@ -0,0 +1,3025 @@ +package river + +import ( + "cmp" + "context" + "encoding/json" + "errors" + "fmt" + "log/slog" + "os" + "regexp" + "slices" + "strings" + "sync" + "sync/atomic" + "time" + + "github.com/riverqueue/river/internal/dblist" + "github.com/riverqueue/river/internal/dbunique" + "github.com/riverqueue/river/internal/jobcompleter" + "github.com/riverqueue/river/internal/jobexecutor" + "github.com/riverqueue/river/internal/leadership" + "github.com/riverqueue/river/internal/maintenance" + "github.com/riverqueue/river/internal/notifier" + "github.com/riverqueue/river/internal/notifylimiter" + "github.com/riverqueue/river/internal/pluginconfig" + "github.com/riverqueue/river/internal/pluginlookup" + "github.com/riverqueue/river/internal/retrypolicy" + "github.com/riverqueue/river/internal/rivercommon" + "github.com/riverqueue/river/internal/riverplugin" + "github.com/riverqueue/river/internal/workunit" + "github.com/riverqueue/river/riverdriver" + "github.com/riverqueue/river/rivershared/baseservice" + "github.com/riverqueue/river/rivershared/riverpilot" + "github.com/riverqueue/river/rivershared/riversharedmaintenance" + "github.com/riverqueue/river/rivershared/startstop" + "github.com/riverqueue/river/rivershared/util/dbutil" + "github.com/riverqueue/river/rivershared/util/maputil" + "github.com/riverqueue/river/rivershared/util/sliceutil" + "github.com/riverqueue/river/rivershared/util/testutil" + "github.com/riverqueue/river/rivershared/util/valutil" + "github.com/riverqueue/river/rivertype" +) + +const ( + FetchCooldownDefault = 100 * time.Millisecond + FetchCooldownMin = 1 * time.Millisecond + + FetchPollIntervalDefault = 1 * time.Second + FetchPollIntervalMin = 1 * time.Millisecond + + JobStuckThresholdDefault = 10 * time.Second + JobTimeoutDefault = 1 * time.Minute + MaxAttemptsDefault = rivercommon.MaxAttemptsDefault + PriorityDefault = rivercommon.PriorityDefault + QueueDefault = rivercommon.QueueDefault + QueueNumWorkersMax = 10_000 +) + +var ( + postgresSchemaNameRE = regexp.MustCompile(`^[a-zA-Z_][a-zA-Z0-9_]*$`) + + reindexerIndexNamesDefault = []string{ //nolint:gochecknoglobals + "river_job_args_index", + "river_job_kind", + "river_job_metadata_index", + "river_job_pkey", + "river_job_prioritized_fetching_index", + "river_job_state_and_finalized_at_index", + "river_job_unique_idx", + } +) + +// TestConfig contains configuration specific to test environments. +type TestConfig struct { + // DisableUniqueEnforcement disables the application of unique job + // constraints. This is useful for testing scenarios when testing a worker + // that typically uses uniqueness, but where enforcing uniqueness would cause + // conflicts with parallel test execution. + // + // The [rivertest.Worker] type automatically disables uniqueness enforcement + // when creating jobs. + DisableUniqueEnforcement bool + + // Time is a time generator to make time stubbable in tests. + Time rivertype.TimeGenerator +} + +// Config is the configuration for a Client. +// +// Both Queues and Workers are required for a client to work jobs, but an +// insert-only client can be initialized by omitting Queues, and not calling +// Start for the client. Workers can also be omitted, but it's better to include +// it so River can check that inserted job kinds have a worker that can run +// them. +type Config struct { + // AdvisoryLockPrefix is a configurable 32-bit prefix that River will use + // when generating any key to acquire a Postgres advisory lock. All advisory + // locks share the same 64-bit number space, so this allows a calling + // application to guarantee that a River advisory lock will never conflict + // with one of its own by cordoning each type to its own prefix. + // + // If this value isn't set, River defaults to generating key hashes across + // the entire 64-bit advisory lock number space, which is large enough that + // conflicts are exceedingly unlikely. If callers don't strictly need this + // option then it's recommended to leave it unset because the prefix leaves + // only 32 bits of number space for advisory lock hashes, so it makes + // internally conflicting River-generated keys more likely. + // + // Advisory locks are currently only used for the deprecated fallback/slow + // path of unique job insertion when pending, scheduled, available, or running + // are omitted from a customized ByState configuration. + AdvisoryLockPrefix int32 + + // CancelledJobRetentionPeriod is the amount of time to keep cancelled jobs + // around before they're removed permanently. + // + // The special value -1 disables deletion of cancelled jobs. + // + // Defaults to 24 hours. + CancelledJobRetentionPeriod time.Duration + + // CompletedJobRetentionPeriod is the amount of time to keep completed jobs + // around before they're removed permanently. + // + // The special value -1 disables deletion of completed jobs. + // + // Defaults to 24 hours. + CompletedJobRetentionPeriod time.Duration + + // DiscardedJobRetentionPeriod is the amount of time to keep discarded jobs + // around before they're removed permanently. + // + // The special value -1 disables deletion of discarded jobs. + // + // Defaults to 7 days. + DiscardedJobRetentionPeriod time.Duration + + // ErrorHandler can be configured to be invoked in case of an error or panic + // occurring in a job. This is often useful for logging and exception + // tracking, but can also be used to customize retry behavior. + ErrorHandler ErrorHandler + + // FetchCooldown is the minimum amount of time to wait between fetches of new + // jobs. Jobs will only be fetched *at most* this often, but if no new jobs + // are coming in via LISTEN/NOTIFY then fetches may be delayed as long as + // FetchPollInterval. + // + // Throughput is limited by this value. + // + // Individual QueueConfig structs may override this for a specific queue. + // + // Defaults to 100 ms. + FetchCooldown time.Duration + + // FetchPollInterval is the amount of time between periodic fetches for new + // jobs. Typically new jobs will be picked up ~immediately after insert via + // LISTEN/NOTIFY, but this provides a fallback. + // + // Individual QueueConfig structs may override this for a specific queue. + // + // Defaults to 1 second. + FetchPollInterval time.Duration + + // ID is the unique identifier for this client. If not set, a random + // identifier will be generated. + // + // This is used to identify the client in job attempts and for leader election. + // This value must be unique across all clients in the same database and + // schema and there must not be more than one process running with the same + // ID at the same time. + // + // A client ID should differ between different programs and must be unique + // across all clients in the same database and schema. There must not be + // more than one process running with the same ID at the same time. + // Duplicate IDs between processes will lead to facilities like leader + // election or client statistics to fail in novel ways. However, the client + // ID is shared by all executors within any given client. (i.e. different + // Go processes have different IDs, but IDs are shared within any given + // process.) + // + // If in doubt, leave this property empty. + ID string + + // JobCleanerTimeout is the timeout of the individual queries within the job + // cleaner. + // + // Defaults to 30 seconds, which should be more than enough time for most + // deployments. + JobCleanerTimeout time.Duration + + // JobInsertMiddleware are optional functions that can be called around job + // insertion. + // + // Deprecated: Prefer the use of Plugins instead (which may contain + // instances of rivertype.JobInsertMiddleware). + JobInsertMiddleware []rivertype.JobInsertMiddleware + + // JobStuckHandler is invoked when a producer detects that a job exceeded + // its timeout and did not return from context cancellation within the + // allotted JobStuckThreshold (and if it didn't, we usually assume it won't + // return at all). The handler receives minimal information about the stuck + // job and the total number of jobs currently considered stuck across the + // client. + // + // JobStuckHandler lets an implementation indicate that a new worker slot + // should be opened to replace the one now occupied by a stuck job. It can + // also be used (for example) to stop and exit the program if too many jobs + // have been reported stuck. + JobStuckHandler JobStuckHandler + + // JobStuckThreshold is the amount of time after JobTimeout elapses to + // wait before a still-running job is considered stuck. + // + // Defaults to 10 seconds. + JobStuckThreshold time.Duration + + // JobTimeout is the maximum amount of time a job is allowed to run before its + // context is cancelled. A timeout of zero means JobTimeoutDefault will be + // used, whereas a value of -1 means the job's context will not be cancelled + // unless the Client is shutting down. + // + // Defaults to 1 minute. + JobTimeout time.Duration + + // Hooks are functions that may activate at certain points during a job's + // lifecycle (see rivertype.Hook), installed globally. + // + // The effect of hooks in this list will depend on the specific hook + // interfaces they implement. rivertype.HookInsertBegin will cause the hook + // to be invoked before a job is inserted. rivertype.HookMetricEmit will + // cause the hook to be invoked when River emits a metric. Implementing + // rivertype.HookWorkBegin will cause it to be invoked before a job is + // worked. Hook structs may implement multiple hook interfaces. + // + // Order in this list is significant. A hook that appears first will be + // entered before a hook that appears later. For any particular phase, order + // is relevant only for hooks that will run for that phase. For example, if + // two rivertype.HookInsertBegin are separated by a rivertype.HookWorkBegin, + // during job insertion those two outer hooks will run one after another, + // and the work hook between them will not run. When a job is worked, the + // work hook runs and the insertion hooks on either side of it are skipped. + // + // Jobs may have their own specific hooks by implementing JobArgsWithHooks or + // JobArgsWithPlugins. + // + // Entries in Hooks are installed only as hooks, even if they also implement + // rivertype.Middleware. Use Plugins for an extension that should act as + // both a hook and middleware. + Hooks []rivertype.Hook + + // Logger is the structured logger to use for logging purposes. If none is + // specified, logs will be emitted to STDOUT with messages at warn level + // or higher. + Logger *slog.Logger + + // MaxAttempts is the default number of times a job will be retried before + // being discarded. This value is applied to all jobs by default, and can be + // overridden on individual job types on the JobArgs or on a per-job basis at + // insertion time. + // + // If not specified, defaults to 25 (MaxAttemptsDefault). + MaxAttempts int + + // Middleware contains middleware that may activate at certain points during + // a job's lifecycle (see rivertype.Middleware), installed globally. + // + // The effect of middleware in this list will depend on the specific + // middleware interfaces they implement, so for example implementing + // rivertype.JobInsertMiddleware will cause the middleware to be invoked + // when jobs are inserted, and implementing rivertype.WorkerMiddleware will + // cause it to be invoked when a job is worked. Middleware structs may + // implement multiple middleware interfaces. + // + // Order in this list is significant. Middleware that appears first will be + // entered before middleware that appears later. For any particular phase, + // order is relevant only for middlewares that will run for that phase. For + // example, if two rivertype.JobInsertMiddleware are separated by a + // rivertype.WorkerMiddleware, during job insertion those two outer + // middlewares will run one after another, and the work middleware between + // them will not run. When a job is worked, the work middleware runs and the + // insertion middlewares on either side of it are skipped. + // + // Entries in Middleware are installed only as middleware, even if they also + // implement rivertype.Hook. Use Plugins for an extension that should act as + // both middleware and a hook. + Middleware []rivertype.Middleware + + // Plugins contains extensions installed globally as hooks, middleware, or + // both. + // + // A type qualifies as a plugin by implementing [rivertype.Plugin]. Most + // existing hook and middleware implementations already do this by embedding + // HookDefaults or MiddlewareDefaults. If a type participates on both sides, + // it may embed PluginDefaults, or embed both HookDefaults and + // MiddlewareDefaults directly and define its own IsPlugin method, then + // implement any operation-specific hook or middleware interfaces it needs. + // + // Use Hooks or Middleware when an extension should be installed only as the + // corresponding kind. Use Plugins when it should be eligible as both. + // Jobs may have their own specific plugins by implementing + // JobArgsWithPlugins. + Plugins []rivertype.Plugin + + // PeriodicJobs are a set of periodic jobs to run at the specified intervals + // in the client. + PeriodicJobs []*PeriodicJob + + // PollOnly starts the client in "poll only" mode, which avoids issuing + // `LISTEN` statements to wait for events like a leadership resignation or + // new job available. The program instead polls periodically to look for + // changes (checking for new jobs on the period in FetchPollInterval). + // + // The downside of this mode of operation is that events will usually be + // noticed less quickly. A new job in the queue may have to wait up to + // FetchPollInterval to be locked for work. When a leader resigns, it will + // be up to five seconds before a new one elects itself. + // + // The upside is that it makes River compatible with systems where + // listen/notify isn't available. For example, PgBouncer in transaction + // pooling mode. + PollOnly bool + + // Queues is a list of queue names for this client to operate on along with + // configuration for the queue like the maximum number of workers to run for + // each queue. + // + // This field may be omitted for a program that's only queueing jobs rather + // than working them. If it's specified, then Workers must also be given. + Queues map[string]QueueConfig + + // ReindexerSchedule is the schedule for running the reindexer. If nil, the + // reindexer will run at midnight UTC every day. + ReindexerSchedule PeriodicSchedule + + // ReindexerIndexNames customizes which indexes River periodically reindexes. + // If nil, River uses [ReindexerIndexNamesDefault]. If non-nil, the provided + // slice is used as the exact list. + ReindexerIndexNames []string + + // ReindexerTimeout is the amount of time to wait for the reindexer to run a + // single reindex operation before cancelling it via context. Set to -1 to + // disable the timeout. + // + // Defaults to 1 minute. + ReindexerTimeout time.Duration + + // RescueStuckJobsAfter is the amount of time a job can be running before it + // is considered stuck. A stuck job which has not yet reached its max attempts + // will be scheduled for a retry, while one which has exhausted its attempts + // will be discarded. This prevents jobs from being stuck forever if a worker + // crashes or is killed. + // + // Note that this can result in repeat or duplicate execution of a job that is + // not actually stuck but is still working. The value should be set higher + // than the maximum duration you expect your jobs to run. Setting a value too + // low will result in more duplicate executions, whereas too high of a value + // will result in jobs being stuck for longer than necessary before they are + // retried. + // + // RescueStuckJobsAfter must be greater than JobTimeout. Otherwise, jobs + // would become eligible for rescue while they're still running. + // + // Defaults to 1 hour, or in cases where JobTimeout has been configured and + // is greater than 1 hour, JobTimeout + 1 hour. + RescueStuckJobsAfter time.Duration + + // RetryPolicy is a configurable retry policy for the client. + // + // Defaults to DefaultRetryPolicy. + RetryPolicy ClientRetryPolicy + + // Schema is a non-standard Schema where River tables are located. All table + // references in database queries will use this value as a prefix. + // + // Defaults to empty, which causes the client to look for tables using the + // setting of Postgres `search_path`. + Schema string + + // SoftStopTimeout is the maximum amount of time that the client will wait + // for running jobs to finish during a stop before their contexts are + // cancelled. After the timeout elapses, the client escalates to a hard stop + // by cancelling the context of all running jobs. This applies regardless of + // how stop is initiated — whether by calling Stop, StopAndCancel, or by + // cancelling the context passed to Start. + // + // In combination with signal.NotifyContext on the context passed to Start, + // this can simplify graceful stop to: + // + // ctx, stop := signal.NotifyContext(ctx, syscall.SIGINT, syscall.SIGTERM) + // defer stop() + // + // if err := client.Start(ctx); err != nil { ... } + // <-client.Stopped() + // + // The signal cancels the Start context, which initiates a soft stop. If + // running jobs haven't finished after SoftStopTimeout, their contexts are + // automatically cancelled to trigger a hard stop. + // + // StopAndCancel bypasses the timeout entirely and cancels job contexts + // immediately. + // + // Defaults to no timeout (wait indefinitely for jobs to finish). + SoftStopTimeout time.Duration + + // SkipJobKindValidation causes the job kind format validation check to be + // skipped. This is available as an interim stopgap for users that have + // invalid job kind names, but would rather disable the check rather than + // fix them immediately. + // + // Deprecated: This option will be removed in a future versions so that job + // kinds will always have to have a valid format. + SkipJobKindValidation bool + + // SkipUnknownJobCheck is a flag to control whether the client should skip + // checking to see if a registered worker exists in the client's worker bundle + // for a job arg prior to insertion. + // + // This can be set to true to allow a client to insert jobs which are + // intended to be worked by a different client which effectively makes + // the client's insertion behavior mimic that of an insert-only client. + // + // Defaults to false. + SkipUnknownJobCheck bool + + // Test holds configuration specific to test environments. + Test TestConfig + + // TestOnly can be set to true to disable certain features that are useful + // in production, but which may be harmful to tests, in ways like having the + // effect of making them slower. It should not be used outside of test + // suites. + // + // For example, queue maintenance services normally stagger their startup + // with a random jittered sleep so they don't all try to work at the same + // time. This is nice in production, but makes starting and stopping the + // client in a test case slower. + TestOnly bool + + // Workers is a bundle of registered job workers. + // + // This field may be omitted for a program that's only enqueueing jobs + // rather than working them, but if it is configured the client can validate + // ahead of time that a worker is properly registered for an inserted job. + // (i.e. That it wasn't forgotten by accident.) + Workers *Workers + + // WorkerMiddleware are optional functions that can be called around + // all job executions. + // + // Deprecated: Prefer the use of Plugins instead (which may contain + // instances of rivertype.WorkerMiddleware). + WorkerMiddleware []rivertype.WorkerMiddleware + + // queuePollInterval is the amount of time between periodic checks for queue + // setting changes. This is only used in poll-only mode (when no notifier is + // provided). + // + // This is internal for the time being as it hasn't had any major demand to + // be exposed, but it's needed to make sure that our poll-only tests can + // finish in a timely manner. + queuePollInterval time.Duration + + // Scheduler run interval. Shared between the scheduler and producer/job + // executors, but not currently exposed for configuration. + schedulerInterval time.Duration +} + +// ReindexerIndexNamesDefault returns the default set of indexes reindexed by River. +func ReindexerIndexNamesDefault() []string { + indexNames := make([]string, len(reindexerIndexNamesDefault)) + copy(indexNames, reindexerIndexNamesDefault) + + return indexNames +} + +// WithDefaults returns a copy of the Config with all default values applied. +func (c *Config) WithDefaults() *Config { + if c == nil { + c = &Config{} + } + + reindexerIndexNames := ReindexerIndexNamesDefault() + if c.ReindexerIndexNames != nil { + reindexerIndexNames = make([]string, len(c.ReindexerIndexNames)) + copy(reindexerIndexNames, c.ReindexerIndexNames) + } + + // Use the existing logger if set, otherwise create a default one. + logger := c.Logger + if logger == nil { + logger = slog.New(slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{ + Level: slog.LevelWarn, + })) + } + + // Compute the default rescue value. For convenience, if JobTimeout is specified + // but RescueStuckJobsAfter is not set (or less than 1) and JobTimeout is large, use + // JobTimeout + maintenance.JobRescuerRescueAfterDefault as the default. + rescueAfter := maintenance.JobRescuerRescueAfterDefault + if c.JobTimeout > 0 && c.RescueStuckJobsAfter < 1 && c.JobTimeout > c.RescueStuckJobsAfter { + rescueAfter = c.JobTimeout + maintenance.JobRescuerRescueAfterDefault + } + + // Set default retry policy if none is provided. + retryPolicy := c.RetryPolicy + if retryPolicy == nil { + retryPolicy = &DefaultClientRetryPolicy{} + } + + return &Config{ + AdvisoryLockPrefix: c.AdvisoryLockPrefix, + CancelledJobRetentionPeriod: cmp.Or(c.CancelledJobRetentionPeriod, riversharedmaintenance.CancelledJobRetentionPeriodDefault), + CompletedJobRetentionPeriod: cmp.Or(c.CompletedJobRetentionPeriod, riversharedmaintenance.CompletedJobRetentionPeriodDefault), + DiscardedJobRetentionPeriod: cmp.Or(c.DiscardedJobRetentionPeriod, riversharedmaintenance.DiscardedJobRetentionPeriodDefault), + ErrorHandler: c.ErrorHandler, + FetchCooldown: cmp.Or(c.FetchCooldown, FetchCooldownDefault), + FetchPollInterval: cmp.Or(c.FetchPollInterval, FetchPollIntervalDefault), + ID: valutil.ValOrDefaultFunc(c.ID, func() string { return defaultClientID(time.Now().UTC()) }), + Hooks: c.Hooks, + JobInsertMiddleware: c.JobInsertMiddleware, + JobStuckHandler: c.JobStuckHandler, + JobStuckThreshold: cmp.Or(c.JobStuckThreshold, JobStuckThresholdDefault), + JobTimeout: cmp.Or(c.JobTimeout, JobTimeoutDefault), + Logger: logger, + MaxAttempts: cmp.Or(c.MaxAttempts, MaxAttemptsDefault), + Middleware: c.Middleware, + PeriodicJobs: c.PeriodicJobs, + Plugins: c.Plugins, + PollOnly: c.PollOnly, + Queues: c.Queues, + ReindexerIndexNames: reindexerIndexNames, + ReindexerSchedule: c.ReindexerSchedule, + ReindexerTimeout: cmp.Or(c.ReindexerTimeout, maintenance.ReindexerTimeoutDefault), + RescueStuckJobsAfter: cmp.Or(c.RescueStuckJobsAfter, rescueAfter), + RetryPolicy: retryPolicy, + Schema: c.Schema, + SoftStopTimeout: c.SoftStopTimeout, + SkipJobKindValidation: c.SkipJobKindValidation, + SkipUnknownJobCheck: c.SkipUnknownJobCheck, + Test: c.Test, + TestOnly: c.TestOnly, + WorkerMiddleware: c.WorkerMiddleware, + Workers: c.Workers, + queuePollInterval: c.queuePollInterval, + schedulerInterval: cmp.Or(c.schedulerInterval, maintenance.JobSchedulerIntervalDefault), + } +} + +func (c *Config) validate() error { + if c.CancelledJobRetentionPeriod < -1 { + return errors.New("CancelledJobRetentionPeriod time cannot be less than zero, except for -1 (infinite)") + } + if c.CompletedJobRetentionPeriod < -1 { + return errors.New("CompletedJobRetentionPeriod cannot be less than zero, except for -1 (infinite)") + } + if c.DiscardedJobRetentionPeriod < -1 { + return errors.New("DiscardedJobRetentionPeriod cannot be less than zero, except for -1 (infinite)") + } + if c.FetchCooldown < FetchCooldownMin { + return fmt.Errorf("FetchCooldown must be at least %s", FetchCooldownMin) + } + if c.FetchPollInterval < FetchPollIntervalMin { + return fmt.Errorf("FetchPollInterval must be at least %s", FetchPollIntervalMin) + } + if c.FetchPollInterval < c.FetchCooldown { + return fmt.Errorf("FetchPollInterval cannot be shorter than FetchCooldown (%s)", c.FetchCooldown) + } + if len(c.ID) > 100 { + return errors.New("ID cannot be longer than 100 characters") + } + if c.JobTimeout < -1 { + return errors.New("JobTimeout cannot be negative, except for -1 (infinite)") + } + if c.JobStuckThreshold < 0 { + return errors.New("JobStuckThreshold cannot be less than zero") + } + if c.MaxAttempts < 0 { + return errors.New("MaxAttempts cannot be less than zero") + } + if len(c.Middleware) > 0 && (len(c.JobInsertMiddleware) > 0 || len(c.WorkerMiddleware) > 0) { + return errors.New("only one of the pair JobInsertMiddleware/WorkerMiddleware or Middleware may be provided (Middleware is recommended, and may contain both job insert and worker middleware)") + } + if c.ReindexerTimeout < -1 { + return errors.New("ReindexerTimeout cannot be negative, except for -1 (infinite)") + } + if c.RescueStuckJobsAfter < 0 { + return errors.New("RescueStuckJobsAfter cannot be less than zero") + } + if c.RescueStuckJobsAfter < c.JobTimeout { + return errors.New("RescueStuckJobsAfter cannot be less than JobTimeout") + } + + // Max Postgres notification topic length is 63 and we prefix schema to + // notification topic, so whatever schema the user specifies must fit inside + // this convention. + maxSchemaLength := 63 - 1 - len(string(notifier.NotificationTopicLongest)) // -1 for the dot in `.` + if len(c.Schema) > maxSchemaLength { + return fmt.Errorf("Schema length must be less than or equal to %d characters", maxSchemaLength) + } + if c.Schema != "" && !postgresSchemaNameRE.MatchString(c.Schema) { + return errors.New("Schema name can only contain letters, numbers, and underscores, and must start with a letter or underscore") + } + + for queue, queueConfig := range c.Queues { + if err := queueConfig.validate(queue, c.FetchCooldown, c.FetchPollInterval); err != nil { + return err + } + } + + if c.Workers == nil && c.Queues != nil { + return errors.New("Workers must be set if Queues is set") + } + + if c.Workers != nil { + for _, workerInfo := range c.Workers.workersMap { + kind := workerInfo.jobArgs.Kind() + if !rivercommon.UserSpecifiedIDOrKindRE.MatchString(kind) { + if c.SkipJobKindValidation { + c.Logger.Warn("job kind should match regex; this will be an error in future versions", + slog.String("kind", kind), + slog.String("regex", rivercommon.UserSpecifiedIDOrKindRE.String()), + ) + } else { + return fmt.Errorf("job kind %q should match regex %s", kind, rivercommon.UserSpecifiedIDOrKindRE.String()) + } + } + } + } + + return nil +} + +// Indicates whether with the given configuration, this client will be expected +// to execute jobs (rather than just being used to enqueue them). Executing jobs +// requires a set of configured queues. +func (c *Config) willExecuteJobs() bool { + return len(c.Queues) > 0 +} + +// QueueConfig contains queue-specific configuration. +type QueueConfig struct { + // FetchCooldown is the minimum amount of time to wait between fetches of new + // jobs. Jobs will only be fetched *at most* this often, but if no new jobs + // are coming in via LISTEN/NOTIFY then fetches may be delayed as long as + // FetchPollInterval. + // + // Throughput is limited by this value. + // + // If non-zero, this overrides the FetchCooldown setting in the Client's + // Config. + FetchCooldown time.Duration + + // FetchPollInterval is the amount of time between periodic fetches for new + // jobs. Typically new jobs will be picked up ~immediately after insert via + // LISTEN/NOTIFY, but this provides a fallback. + // + // If non-zero, this overrides the FetchCooldown setting in the Client's + // Config. + FetchPollInterval time.Duration + + // MaxWorkers is the maximum number of workers to run for the queue, or put + // otherwise, the maximum parallelism to run. + // + // This is the maximum number of workers within this particular client + // instance, but note that it doesn't control the total number of workers + // across parallel processes. Installations will want to calculate their + // total number by multiplying this number by the number of parallel nodes + // running River clients configured to the same database and queue. + // + // Requires a minimum of 1, and a maximum of 10,000. + MaxWorkers int +} + +func (c QueueConfig) validate(queueName string, clientFetchCooldown time.Duration, clientFetchPollInterval time.Duration) error { + if c.FetchCooldown < 0 { + return errors.New("FetchCooldown cannot be less than zero") + } + if c.FetchPollInterval < 0 { + return errors.New("FetchPollInterval cannot be less than zero") + } + + resolvedFetchCooldown := cmp.Or(c.FetchCooldown, clientFetchCooldown) + resolvedFetchPollInterval := cmp.Or(c.FetchPollInterval, clientFetchPollInterval) + if resolvedFetchPollInterval < resolvedFetchCooldown { + return errors.New("FetchPollInterval cannot be less than FetchCooldown") + } + + if c.MaxWorkers < 1 || c.MaxWorkers > QueueNumWorkersMax { + return fmt.Errorf("invalid number of workers for queue %q: %d", queueName, c.MaxWorkers) + } + if err := validateQueueName(queueName); err != nil { + return err + } + + return nil +} + +// Client is a single isolated instance of River. Your application may use +// multiple instances operating on different databases or Postgres schemas +// within a single database. +type Client[TTx any] struct { + // BaseService and BaseStartStop can't be embedded like on other services + // because their properties would leak to the external API. + baseService baseservice.BaseService + baseStartStop startstop.BaseStartStop + + clientNotifyBundle *ClientNotifyBundle[TTx] + completer jobcompleter.JobCompleter + config *Config + driver riverdriver.Driver[TTx] + elector *leadership.Elector + pluginLookupByJob *pluginlookup.JobPluginLookup + pluginLookupGlobal *pluginlookup.PluginLookup + insertNotifyLimiter *notifylimiter.Limiter + notifier *notifier.Notifier // may be nil in poll-only mode + periodicJobs *PeriodicJobBundle + pilot riverpilot.Pilot + producersByQueueName map[string]*producer + producersMu sync.RWMutex + queueMaintainer *maintenance.QueueMaintainer + queueMaintainerLeader *maintenance.QueueMaintainerLeader + queues *QueueBundle + services []startstop.Service + stopped <-chan struct{} + stuckJobCount atomic.Int32 + subscriptionManager *subscriptionManager + testSignals clientTestSignals + + // workCancel cancels the context used for all work goroutines. Normal Stop + // does not cancel that context. + workCancel context.CancelCauseFunc +} + +// Test-only signals. +type clientTestSignals struct { + jobCleaner *maintenance.JobCleanerTestSignals + jobRescuer *maintenance.JobRescuerTestSignals + jobScheduler *maintenance.JobSchedulerTestSignals + periodicJobEnqueuer *maintenance.PeriodicJobEnqueuerTestSignals + queueCleaner *maintenance.QueueCleanerTestSignals + queueMaintainerLeader *maintenance.QueueMaintainerLeaderTestSignals + reindexer *maintenance.ReindexerTestSignals +} + +func (ts *clientTestSignals) Init(tb testutil.TestingTB) { + if ts.jobCleaner != nil { + ts.jobCleaner.Init(tb) + } + if ts.jobRescuer != nil { + ts.jobRescuer.Init(tb) + } + if ts.jobScheduler != nil { + ts.jobScheduler.Init(tb) + } + if ts.periodicJobEnqueuer != nil { + ts.periodicJobEnqueuer.Init(tb) + } + if ts.queueCleaner != nil { + ts.queueCleaner.Init(tb) + } + if ts.queueMaintainerLeader != nil { + ts.queueMaintainerLeader.Init(tb) + } + if ts.reindexer != nil { + ts.reindexer.Init(tb) + } +} + +var ( + // ErrNotFound is returned when a query by ID does not match any existing + // rows. For example, attempting to cancel a job that doesn't exist will + // return this error. + ErrNotFound = rivertype.ErrNotFound + + errMissingConfig = errors.New("missing config") + errMissingDatabasePoolWithQueues = errors.New("must have a non-nil database pool to execute jobs (either use a driver with database pool or don't configure Queues)") + errMissingDriver = errors.New("missing database driver (try wrapping a Pgx pool with river/riverdriver/riverpgxv5.New)") +) + +// NewClient creates a new Client with the given database driver and +// configuration. +// +// Currently only one driver is supported, which is Pgx v5. See package +// riverpgxv5. +// +// The function takes a generic parameter TTx representing a transaction type, +// but it can be omitted because it'll generally always be inferred from the +// driver. For example: +// +// import "github.com/riverqueue/river" +// import "github.com/riverqueue/river/riverdriver/riverpgxv5" +// +// ... +// +// dbPool, err := pgxpool.New(ctx, os.Getenv("DATABASE_URL")) +// if err != nil { +// // handle error +// } +// defer dbPool.Close() +// +// riverClient, err := river.NewClient(riverpgxv5.New(dbPool), &river.Config{ +// ... +// }) +// if err != nil { +// // handle error +// } +func NewClient[TTx any](driver riverdriver.Driver[TTx], config *Config) (*Client[TTx], error) { + if driver == nil { + return nil, errMissingDriver + } + if config == nil { + return nil, errMissingConfig + } + + config = config.WithDefaults() + + if err := config.validate(); err != nil { + return nil, err + } + + archetype := baseservice.NewArchetype(config.Logger) + if config.Test.Time != nil { + if withStub, ok := config.Test.Time.(baseservice.TimeGeneratorWithStub); ok { + archetype.Time = withStub + } else { + archetype.Time = &baseservice.TimeGeneratorWithStubWrapper{TimeGenerator: config.Test.Time} + } + } + if _, ok := config.RetryPolicy.(*DefaultClientRetryPolicy); ok { + config.RetryPolicy = retrypolicy.NewDefault(archetype.Time) + } + + var ( + middleware = pluginconfig.CombinedMiddleware(config.Middleware, config.JobInsertMiddleware, config.WorkerMiddleware) + plugins = append(riverplugin.DefaultPlugins(), config.Plugins...) + ) + pluginLookupByJob := pluginlookup.NewJobPluginLookup(archetype) + pluginLookupGlobal := pluginlookup.NewPluginLookupFromConfig(archetype, config.Hooks, middleware, plugins) + + client := &Client[TTx]{ + clientNotifyBundle: &ClientNotifyBundle[TTx]{ + config: config, + driver: driver, + }, + config: config, + driver: driver, + pluginLookupByJob: pluginLookupByJob, + pluginLookupGlobal: pluginLookupGlobal, + producersByQueueName: make(map[string]*producer), + testSignals: clientTestSignals{}, + workCancel: func(cause error) {}, // replaced on start, but here in case StopAndCancel is called before start up + } + + client.queues = &QueueBundle{ + clientFetchCooldown: config.FetchCooldown, + clientFetchPollInterval: config.FetchPollInterval, + clientWillExecuteJobs: config.willExecuteJobs(), + producerAdd: client.producerAdd, + producerRemove: client.producerRemove, + } + + baseservice.Init(archetype, &client.baseService) + client.baseService.Name = "Client" // Have to correct the name because base service isn't embedded like it usually is + client.insertNotifyLimiter = notifylimiter.NewLimiter(archetype, config.FetchCooldown) + + pluginDriver, _ := driver.(driverPlugin[TTx]) + if pluginDriver != nil { + pluginDriver.PluginInit(archetype) + client.pilot = pluginDriver.PluginPilot() + } + + var workerMetadata []*rivertype.WorkerMetadata + if config.Workers != nil { + workerMetadata = make([]*rivertype.WorkerMetadata, 0, len(config.Workers.workersMap)) + for kind, workerInfo := range config.Workers.workersMap { + workerMetadata = append(workerMetadata, &rivertype.WorkerMetadata{ + JobArgHooks: pluginLookupByJob.ByJobArgs(workerInfo.jobArgs).Hooks(), + Kind: kind, + }) + } + } + + if client.pilot == nil { + client.pilot = &riverpilot.StandardPilot{} + } + client.pilot.PilotInit(archetype, (&riverpilot.PilotInitParams{ + Insert: client.insertMany, + NotifyNonTxJobInsert: client.notifyProducerWithoutListenerJobFetch, + ProducerReportInterval: producerReportIntervalDefault, + WorkerMetadata: workerMetadata, + }).Validate()) + pluginPilot, _ := client.pilot.(pilotPlugin) + + if withBaseService, ok := config.RetryPolicy.(baseservice.WithBaseService); ok { + baseservice.Init(archetype, withBaseService) + } + + // There are a number of internal components that are only needed/desired if + // we're actually going to be working jobs (as opposed to just enqueueing + // them): + if config.willExecuteJobs() { + if !driver.PoolIsSet() { + return nil, errMissingDatabasePoolWithQueues + } + + client.completer = jobcompleter.NewBatchCompleter(archetype, config.Schema, driver.GetExecutor(), client.pilot, nil) + client.subscriptionManager = newSubscriptionManager(archetype, nil) + client.services = append(client.services, client.completer, client.subscriptionManager) + + if driver.SupportsListener() { + // In poll only mode, we don't try to initialize a notifier that + // uses listen/notify. Instead, each service polls for changes it's + // interested in. e.g. Elector polls to see if leader has expired. + if !config.PollOnly { + client.notifier = notifier.New(archetype, driver.GetListener(&riverdriver.GetListenenerParams{Schema: config.Schema})) + client.services = append(client.services, client.notifier) + } + } else { + config.Logger.Info("Driver does not support listener; entering poll only mode") + } + + client.elector = leadership.NewElector(archetype, driver.GetExecutor(), client.notifier, &leadership.Config{ + ClientID: config.ID, + Schema: config.Schema, + }) + client.services = append(client.services, client.elector) + + for queue, queueConfig := range config.Queues { + if _, err := client.producerAdd(queue, queueConfig); err != nil { + return nil, err + } + } + + client.services = append(client.services, + startstop.StartStopFunc(client.logStatsLoop)) + + if pluginPilot != nil { + client.services = append(client.services, pluginPilot.PluginServices()...) + } + + // + // Maintenance services + // + + maintenanceServices := []startstop.Service{} + + { + jobCleaner := maintenance.NewJobCleaner(archetype, &maintenance.JobCleanerConfig{ + CancelledJobRetentionPeriod: config.CancelledJobRetentionPeriod, + CompletedJobRetentionPeriod: config.CompletedJobRetentionPeriod, + DiscardedJobRetentionPeriod: config.DiscardedJobRetentionPeriod, + QueuesExcluded: client.pilot.JobCleanerQueuesExcluded(), + Schema: config.Schema, + Timeout: config.JobCleanerTimeout, + }, driver.GetExecutor()) + maintenanceServices = append(maintenanceServices, jobCleaner) + client.testSignals.jobCleaner = &jobCleaner.TestSignals + } + + { + jobRescuer := maintenance.NewRescuer(archetype, &maintenance.JobRescuerConfig{ + ClientJobTimeout: config.JobTimeout, + ClientRetryPolicy: config.RetryPolicy, + Pilot: client.pilot, + RescueAfter: config.RescueStuckJobsAfter, + Schema: config.Schema, + WorkUnitFactoryFunc: func(kind string) workunit.WorkUnitFactory { + if workerInfo, ok := config.Workers.workersMap[kind]; ok { + return workerInfo.workUnitFactory + } + return nil + }, + }, driver.GetExecutor()) + maintenanceServices = append(maintenanceServices, jobRescuer) + client.testSignals.jobRescuer = &jobRescuer.TestSignals + } + + { + jobScheduler := maintenance.NewJobScheduler(archetype, &maintenance.JobSchedulerConfig{ + Interval: config.schedulerInterval, + NotifyInsert: client.maybeNotifyInsertForQueues, + Schema: config.Schema, + }, driver.GetExecutor()) + maintenanceServices = append(maintenanceServices, jobScheduler) + client.testSignals.jobScheduler = &jobScheduler.TestSignals + } + + { + periodicJobEnqueuer, err := maintenance.NewPeriodicJobEnqueuer(archetype, &maintenance.PeriodicJobEnqueuerConfig{ + AdvisoryLockPrefix: config.AdvisoryLockPrefix, + PluginLookupGlobal: client.pluginLookupGlobal, + Insert: client.insertMany, + Pilot: client.pilot, + Schema: config.Schema, + }, driver.GetExecutor()) + if err != nil { + return nil, err + } + maintenanceServices = append(maintenanceServices, periodicJobEnqueuer) + client.testSignals.periodicJobEnqueuer = &periodicJobEnqueuer.TestSignals + + client.periodicJobs = newPeriodicJobBundle(client.config, periodicJobEnqueuer) + client.periodicJobs.AddMany(config.PeriodicJobs) + } + + { + queueCleaner := maintenance.NewQueueCleaner(archetype, &maintenance.QueueCleanerConfig{ + RetentionPeriod: maintenance.QueueRetentionPeriodDefault, + Schema: config.Schema, + }, driver.GetExecutor()) + maintenanceServices = append(maintenanceServices, queueCleaner) + client.testSignals.queueCleaner = &queueCleaner.TestSignals + } + + if driver.DatabaseName() == riverdriver.DatabaseNameSQLite { + sqliteNotificationCleaner := maintenance.NewSQLiteNotificationCleaner(archetype, &maintenance.SQLiteNotificationCleanerConfig{ + Schema: config.Schema, + }, driver.GetExecutor()) + maintenanceServices = append(maintenanceServices, sqliteNotificationCleaner) + } + + { + var scheduleFunc func(time.Time) time.Time + if config.ReindexerSchedule != nil { + scheduleFunc = config.ReindexerSchedule.Next + } + + reindexer := maintenance.NewReindexer(archetype, &maintenance.ReindexerConfig{ + IndexNames: config.ReindexerIndexNames, + ScheduleFunc: scheduleFunc, + Schema: config.Schema, + Timeout: config.ReindexerTimeout, + }, driver.GetExecutor()) + maintenanceServices = append(maintenanceServices, reindexer) + client.testSignals.reindexer = &reindexer.TestSignals + } + + if pluginPilot != nil { + maintenanceServices = append(maintenanceServices, pluginPilot.PluginMaintenanceServices()...) + } + + // Not added to the main services list because the queue maintainer is + // started conditionally based on whether the client is the leader. + client.queueMaintainer = maintenance.NewQueueMaintainer(archetype, maintenanceServices) + + if config.TestOnly { + client.queueMaintainer.StaggerStartupDisable(true) + } + + client.queueMaintainerLeader = maintenance.NewQueueMaintainerLeader(archetype, &maintenance.QueueMaintainerLeaderConfig{ + ClientID: config.ID, + Elector: client.elector, + QueueMaintainer: client.queueMaintainer, + RequestResignFunc: client.clientNotifyBundle.RequestResign, + }) + client.services = append(client.services, client.queueMaintainerLeader) + client.testSignals.queueMaintainerLeader = &client.queueMaintainerLeader.TestSignals + } + + return client, nil +} + +// Start starts the client's job fetching and working loops. Once this is called, +// the client will run in a background goroutine until stopped. All jobs are +// run with a context inheriting from the provided context, but with a timeout +// deadline applied based on the job's settings. +// +// A graceful shutdown stops fetching new jobs but allows any previously fetched +// jobs to complete. This can be initiated with the Stop method. +// +// A more abrupt shutdown can be achieved by either cancelling the provided +// context or by calling StopAndCancel. This will not only stop fetching new +// jobs, but will also cancel the context for any currently-running jobs. If +// using StopAndCancel, there's no need to also call Stop. +func (c *Client[TTx]) Start(ctx context.Context) error { + fetchCtx, shouldStart, started, stopped := c.baseStartStop.StartInit(ctx) + if !shouldStart { + return nil + } + + c.queues.startStopMu.Lock() + defer c.queues.startStopMu.Unlock() + + // BaseStartStop will set its stopped channel to nil after it stops, so make + // sure to take a channel reference before finishing stopped. + c.stopped = c.baseStartStop.StoppedUnsafe() + + producersAsServices := func() []startstop.Service { + return sliceutil.Map( + maputil.Values(c.producersByQueueName), + func(p *producer) startstop.Service { return p }, + ) + } + + // Startup code. Wrapped in a closure so it doesn't have to remember to + // close the stopped channel if returning with an error. + if err := func() error { + if !c.config.willExecuteJobs() { + return errors.New("client Queues and Workers must be configured for a client to start working") + } + if c.config.Workers != nil && len(c.config.Workers.workersMap) < 1 { + return errors.New("at least one Worker must be added to the Workers bundle") + } + + // Before doing anything else, make an initial connection to the database to + // verify that it appears healthy. Many of the subcomponents below start up + // in a goroutine and in case of initial failure, only produce a log line, + // so even in the case of a fundamental failure like the database not being + // available, the client appears to have started even though it's completely + // non-functional. Here we try to make an initial assessment of health and + // return quickly in case of an apparent problem. + if err := c.driver.GetExecutor().Exec(fetchCtx, "SELECT 1"); err != nil { + return fmt.Errorf("error making initial connection to database: %w", err) + } + + // Each time we start, we need a fresh completer subscribe channel to + // send job completion events on, because the completer will close it + // each time it shuts down. + completerSubscribeCh := make(chan []jobcompleter.CompleterJobUpdated, 10) + c.completer.ResetSubscribeChan(completerSubscribeCh) + c.subscriptionManager.ResetSubscribeChan(completerSubscribeCh) + + // In case of error, stop any services that might have started. This + // is safe because even services that were never started will still + // tolerate being stopped. + stopServicesOnError := func() { + startstop.StopAllParallel(c.services...) + } + + // The completer is part of the services list below, but although it can + // stop gracefully along with all the other services, it needs to be + // started with a context that's _not_ cancelled if the user-provided + // context is cancelled. This ensures that even when fetch is cancelled on + // shutdown, the completer is still given a separate opportunity to start + // stopping only after the producers have finished up and returned. + if err := c.completer.Start(context.WithoutCancel(ctx)); err != nil { + stopServicesOnError() + return err + } + + // We use separate contexts for fetching and working to allow for a + // graceful stop. When SoftStopTimeout is configured, the work context + // is detached from the start context so that cancelling the start + // context initiates a soft stop (with timeout escalation) rather than + // an immediate hard stop. When SoftStopTimeout is not configured, the + // work context inherits from the start context to preserve the + // existing behavior where cancelling the start context is equivalent + // to StopAndCancel. + workParentCtx := ctx + if c.config.SoftStopTimeout > 0 { + workParentCtx = context.WithoutCancel(ctx) + } + workCtx, workCancel := context.WithCancelCause(workParentCtx) + + // Client available to executors and to various service hooks. + fetchCtx := withClient(fetchCtx, c) + workCtx = withClient(workCtx, c) + + if err := startstop.StartAll(fetchCtx, c.services...); err != nil { + workCancel(err) + stopServicesOnError() + return err + } + + for _, producer := range c.producersByQueueName { + if err := producer.StartWorkContext(fetchCtx, workCtx); err != nil { + workCancel(err) + startstop.StopAllParallel(producersAsServices()...) + stopServicesOnError() + return err + } + } + + c.queues.fetchCtx = fetchCtx + c.queues.workCtx = workCtx + c.workCancel = workCancel + + return nil + }(); err != nil { + defer stopped() + if errors.Is(context.Cause(fetchCtx), startstop.ErrStop) { + return nil + } + return err + } + + // Generate producer services while c.queues.startStopMu.Lock() is still + // held. This is used for WaitAllStarted below, but don't use it elsewhere + // because new producers may have been added while the client is running. + producerServices := producersAsServices() + + go func() { + // Wait for all subservices to start up before signaling our own start. + // This isn't strictly needed, but gives tests a way to fully confirm + // that all goroutines for subservices are spun up before continuing. + // + // Stop also cancels the "started" channel, so in case of a context + // cancellation, this statement will fall through. The client will + // briefly start, but then immediately stop again. + startstop.WaitAllStarted(append( + c.services, + producerServices..., // see comment on this variable + )...) + + started() + defer stopped() + + c.baseService.Logger.InfoContext(ctx, "River client started", slog.String("client_id", c.ID())) + defer c.baseService.Logger.InfoContext(ctx, "River client stopped", slog.String("client_id", c.ID())) + + // The call to Stop cancels this context. Block here until shutdown. + <-fetchCtx.Done() + + c.queues.startStopMu.Lock() + defer c.queues.startStopMu.Unlock() + + // If SoftStopTimeout is configured, start a timer that will cancel + // the work context (escalating to a hard stop) if producers don't + // finish in time. StopAndCancel also calls workCancel, in which case + // this timer is a harmless no-op because the context is already done. + if c.config.SoftStopTimeout > 0 { + softStopTimer := time.AfterFunc(c.config.SoftStopTimeout, func() { + c.baseService.Logger.WarnContext(ctx, c.baseService.Name+": Soft stop timeout; cancelling remaining job contexts", slog.Duration("soft_stop_timeout", c.config.SoftStopTimeout)) + c.workCancel(rivercommon.ErrStop) + }) + defer softStopTimer.Stop() + } + + // On stop, have the producers stop fetching first of all. + c.baseService.Logger.DebugContext(ctx, c.baseService.Name+": Stopping producers") + startstop.StopAllParallel(producersAsServices()...) + c.baseService.Logger.DebugContext(ctx, c.baseService.Name+": All producers stopped") + + c.workCancel(rivercommon.ErrStop) + + // Stop all mainline services where stop order isn't important. + startstop.StopAllParallel(append( + // This list of services contains the completer, which should always + // stop after the producers so that any remaining work that was enqueued + // will have a chance to have its state completed as it finishes. + // + // TODO: there's a risk here that the completer is stuck on a job that + // won't complete. We probably need a timeout or way to move on in those + // cases. + c.services, + + // Will only be started if this client was leader, but can tolerate a + // stop without having been started. + c.queueMaintainer, + )...) + }() + + return nil +} + +// Stop performs a graceful shutdown of the Client. It signals all producers +// to stop fetching new jobs and waits for any fetched or in-progress jobs to +// complete before exiting. If the provided context is done before shutdown has +// completed, Stop will return immediately with the context's error. +// +// If SoftStopTimeout is configured, running job contexts will be automatically +// cancelled after the timeout elapses, escalating to a hard stop. This also +// applies when stop is initiated by cancelling the context passed to Start. +// +// There's no need to call this method if a hard stop has already been initiated +// by cancelling the context passed to Start or by calling StopAndCancel. +func (c *Client[TTx]) Stop(ctx context.Context) error { + shouldStop, stopped, finalizeStop := c.baseStartStop.StopInit() + if !shouldStop { + return nil + } + + select { + case <-ctx.Done(): // stop context cancelled + finalizeStop(false) // not stopped; allow Stop to be called again + return ctx.Err() + case <-stopped: + finalizeStop(true) + return nil + } +} + +// StopAndCancel shuts down the client and cancels all work in progress. It is a +// more aggressive stop than Stop because the contexts for any in-progress jobs +// are cancelled. However, it still waits for jobs to complete before returning, +// even though their contexts are cancelled. If the provided context is done +// before shutdown has completed, StopAndCancel will return immediately with the +// context's error. +// +// This can also be initiated by cancelling the context passed to Start. There is +// no need to call this method if the context passed to Start is cancelled +// instead. +// +// In most cases, using Stop with SoftStopTimeout configured is preferable to +// calling StopAndCancel directly. SoftStopTimeout gives running jobs a chance +// to finish before automatically escalating to context cancellation, providing +// graceful stop semantics without requiring manual orchestration of Stop and +// StopAndCancel. +func (c *Client[TTx]) StopAndCancel(ctx context.Context) error { + c.baseService.Logger.InfoContext(ctx, c.baseService.Name+": Hard stop started; cancelling all work") + c.workCancel(rivercommon.ErrStop) + + shouldStop, stopped, finalizeStop := c.baseStartStop.StopInit() + if !shouldStop { + return nil + } + + select { + case <-ctx.Done(): // stop context cancelled + finalizeStop(false) // not stopped; allow Stop to be called again + return ctx.Err() + case <-stopped: + finalizeStop(true) + return nil + } +} + +// Stopped returns a channel that will be closed when the Client has stopped. +// It can be used to wait for a graceful shutdown to complete. +// +// It is not affected by any contexts passed to Stop or StopAndCancel. +func (c *Client[TTx]) Stopped() <-chan struct{} { + return c.stopped +} + +// Subscribe subscribes to the provided kinds of events that occur within the +// client, like EventKindJobCompleted for when a job completes. +// +// Returns a channel over which to receive events along with a cancel function +// that can be used to cancel and tear down resources associated with the +// subscription. It's recommended but not necessary to invoke the cancel +// function. Resources will be freed when the client stops in case it's not. +// +// The event channel is buffered and sends on it are non-blocking. Consumers +// must process events in a timely manner or it's possible for events to be +// dropped. Any slow operations performed in a response to a receipt (e.g. +// persisting to a database) should be made asynchronous to avoid event loss. +// +// Callers must specify the kinds of events they're interested in. This allows +// for forward compatibility in case new kinds of events are added in future +// versions. If new event kinds are added, callers will have to explicitly add +// them to their requested list and ensure they can be handled correctly. +func (c *Client[TTx]) Subscribe(kinds ...EventKind) (<-chan *Event, func()) { + return c.SubscribeConfig(&SubscribeConfig{Kinds: kinds}) +} + +// The default maximum size of the subscribe channel. Events that would overflow +// it will be dropped. +const subscribeChanSizeDefault = 1_000 + +// SubscribeConfig is more thorough subscription configuration used for +// Client.SubscribeConfig. +type SubscribeConfig struct { + // ChanSize is the size of the buffered channel that will be created for the + // subscription. Incoming events that would overflow this buffer because a + // listener isn't reading from the channel in a timely manner will be dropped. + // + // Defaults to 1000. + ChanSize int + + // Kinds are the kinds of events that the subscription will receive. + // Requiring that kinds are specified explicitly allows for forward + // compatibility in case new kinds of events are added in future versions. + // If new event kinds are added, callers will have to explicitly add them to + // their requested list and ensure they can be handled correctly. + Kinds []EventKind +} + +// SubscribeConfig is a special internal variant of Subscribe that lets us +// inject an overridden channel size. +func (c *Client[TTx]) SubscribeConfig(config *SubscribeConfig) (<-chan *Event, func()) { + if c.subscriptionManager == nil { + panic("created a subscription on a client that will never work jobs (Queues not configured)") + } + + return c.subscriptionManager.SubscribeConfig(config) +} + +// Dump aggregate stats from job completions to logs periodically. These +// numbers don't mean much in themselves, but can give a rough idea of the +// proportions of each compared to each other, and may help flag outlying values +// indicative of a problem. +func (c *Client[TTx]) logStatsLoop(ctx context.Context, shouldStart bool, started, stopped func()) error { + if !shouldStart { + return nil + } + + go func() { + started() + defer stopped() // this defer should come first so it's last out + + ticker := time.NewTicker(5 * time.Second) + defer ticker.Stop() + + for { + select { + case <-ctx.Done(): + return + + case <-ticker.C: + c.subscriptionManager.logStats(ctx, c.baseService.Name) + } + } + }() + + return nil +} + +// Driver exposes the underlying driver used by the client. +// +// API is not stable. DO NOT USE. +func (c *Client[TTx]) Driver() riverdriver.Driver[TTx] { + return c.driver +} + +// JobCancel cancels the job with the given ID. If possible, the job is +// cancelled immediately and will not be retried. The provided context is used +// for the underlying Postgres update and can be used to cancel the operation or +// apply a timeout. +// +// If the job is still in the queue (available, scheduled, or retryable), it is +// immediately marked as cancelled and will not be retried. +// +// If the job is already finalized (cancelled, completed, or discarded), no +// changes are made. +// +// If the job is currently running, it is not immediately cancelled, but is +// instead marked for cancellation. The client running the job will also be +// notified (via LISTEN/NOTIFY) to cancel the running job's context. Although +// the job's context will be cancelled, since Go does not provide a mechanism to +// interrupt a running goroutine the job will continue running until it returns. +// As always, it is important for workers to respect context cancellation and +// return promptly when the job context is done. +// +// Once the cancellation signal is received by the client running the job, any +// error returned by that job will result in it being cancelled permanently and +// not retried. However if the job returns no error, it will be completed as +// usual. +// +// In the event the running job finishes executing _before_ the cancellation +// signal is received but _after_ this update was made, the behavior depends on +// which state the job is being transitioned into (based on its return error): +// +// - If the job completed successfully, was cancelled from within, or was +// discarded due to exceeding its max attempts, the job will be updated as +// usual. +// - If the job was snoozed to run again later or encountered a retryable error, +// the job will be marked as cancelled and will not be attempted again. +// +// Returns the up-to-date JobRow for the specified jobID if it exists. Returns +// ErrNotFound if the job doesn't exist. +func (c *Client[TTx]) JobCancel(ctx context.Context, jobID int64) (*rivertype.JobRow, error) { + job, err := c.jobCancel(ctx, c.driver.GetExecutor(), jobID) + if err != nil { + return nil, err + } + + c.notifyProducerWithoutListenerQueueControlEvent(job.Queue, &controlEventPayload{ + Action: controlActionCancel, + JobID: job.ID, + Queue: job.Queue, + }) + + return job, nil +} + +// JobCancelTx cancels the job with the given ID within the specified +// transaction. This variant lets a caller cancel a job atomically alongside +// other database changes. A cancelled job doesn't take effect until the +// transaction commits, and if the transaction rolls back, so too is the +// cancelled job. +// +// If possible, the job is cancelled immediately and will not be retried. The +// provided context is used for the underlying Postgres update and can be used +// to cancel the operation or apply a timeout. +// +// If the job is still in the queue (available, scheduled, or retryable), it is +// immediately marked as cancelled and will not be retried. +// +// If the job is already finalized (cancelled, completed, or discarded), no +// changes are made. +// +// If the job is currently running, it is not immediately cancelled, but is +// instead marked for cancellation. The client running the job will also be +// notified (via LISTEN/NOTIFY) to cancel the running job's context. Although +// the job's context will be cancelled, since Go does not provide a mechanism to +// interrupt a running goroutine the job will continue running until it returns. +// As always, it is important for workers to respect context cancellation and +// return promptly when the job context is done. +// +// Once the cancellation signal is received by the client running the job, any +// error returned by that job will result in it being cancelled permanently and +// not retried. However if the job returns no error, it will be completed as +// usual. +// +// In the event the running job finishes executing _before_ the cancellation +// signal is received but _after_ this update was made, the behavior depends on +// which state the job is being transitioned into (based on its return error): +// +// - If the job completed successfully, was cancelled from within, or was +// discarded due to exceeding its max attempts, the job will be updated as +// usual. +// - If the job was snoozed to run again later or encountered a retryable error, +// the job will be marked as cancelled and will not be attempted again. +// +// Returns the up-to-date JobRow for the specified jobID if it exists. Returns +// ErrNotFound if the job doesn't exist. +func (c *Client[TTx]) JobCancelTx(ctx context.Context, tx TTx, jobID int64) (*rivertype.JobRow, error) { + return c.jobCancel(ctx, c.driver.UnwrapExecutor(tx), jobID) +} + +func (c *Client[TTx]) jobCancel(ctx context.Context, exec riverdriver.Executor, jobID int64) (*rivertype.JobRow, error) { + return c.pilot.JobCancel(ctx, exec, &riverdriver.JobCancelParams{ + ID: jobID, + CancelAttemptedAt: c.baseService.Time.Now(), + ControlTopic: string(notifier.NotificationTopicControl), + Now: c.baseService.Time.NowOrNil(), + Schema: c.config.Schema, + }) +} + +// JobDelete deletes the job with the given ID from the database, returning the +// deleted row if it was deleted. Jobs in the running state are not deleted, +// instead returning rivertype.ErrJobRunning. +func (c *Client[TTx]) JobDelete(ctx context.Context, id int64) (*rivertype.JobRow, error) { + return c.driver.GetExecutor().JobDelete(ctx, &riverdriver.JobDeleteParams{ + ID: id, + Schema: c.config.Schema, + }) +} + +// JobDeleteTx deletes the job with the given ID from the database, returning the +// deleted row if it was deleted. Jobs in the running state are not deleted, +// instead returning rivertype.ErrJobRunning. This variant lets a caller retry a +// job atomically alongside other database changes. A deleted job isn't deleted +// until the transaction commits, and if the transaction rolls back, so too is +// the deleted job. +func (c *Client[TTx]) JobDeleteTx(ctx context.Context, tx TTx, id int64) (*rivertype.JobRow, error) { + return c.driver.UnwrapExecutor(tx).JobDelete(ctx, &riverdriver.JobDeleteParams{ + ID: id, + Schema: c.config.Schema, + }) +} + +// JobGet fetches a single job by its ID. Returns the up-to-date JobRow for the +// specified jobID if it exists. Returns ErrNotFound if the job doesn't exist. +func (c *Client[TTx]) JobGet(ctx context.Context, id int64) (*rivertype.JobRow, error) { + return c.driver.GetExecutor().JobGetByID(ctx, &riverdriver.JobGetByIDParams{ + ID: id, + Schema: c.config.Schema, + }) +} + +// JobGetTx fetches a single job by its ID, within a transaction. Returns the +// up-to-date JobRow for the specified jobID if it exists. Returns ErrNotFound +// if the job doesn't exist. +func (c *Client[TTx]) JobGetTx(ctx context.Context, tx TTx, id int64) (*rivertype.JobRow, error) { + return c.driver.UnwrapExecutor(tx).JobGetByID(ctx, &riverdriver.JobGetByIDParams{ + ID: id, + Schema: c.config.Schema, + }) +} + +// JobRetry updates the job with the given ID to make it immediately available +// to be retried. Jobs in the running state are not touched, while jobs in any +// other state are made available. To prevent jobs already waiting in the queue +// from being set back in line, the job's scheduled_at field is set to the +// current time only if it's not already in the past. +// +// MaxAttempts is also incremented by one if the job has already exhausted its +// max attempts. +func (c *Client[TTx]) JobRetry(ctx context.Context, id int64) (*rivertype.JobRow, error) { + return c.jobRetry(ctx, c.driver.GetExecutor(), id) +} + +// JobRetryTx updates the job with the given ID to make it immediately available +// to be retried, within the specified transaction. This variant lets a caller +// retry a job atomically alongside other database changes. A retried job isn't +// visible to be worked until the transaction commits, and if the transaction +// rolls back, so too is the retried job. +// +// Jobs in the running state are not touched, while jobs in any other state are +// made available. To prevent jobs already waiting in the queue from being set +// back in line, the job's scheduled_at field is set to the current time only if +// it's not already in the past. +// +// MaxAttempts is also incremented by one if the job has already exhausted its +// max attempts. +func (c *Client[TTx]) JobRetryTx(ctx context.Context, tx TTx, id int64) (*rivertype.JobRow, error) { + return c.jobRetry(ctx, c.driver.UnwrapExecutor(tx), id) +} + +func (c *Client[TTx]) jobRetry(ctx context.Context, exec riverdriver.Executor, id int64) (*rivertype.JobRow, error) { + return c.pilot.JobRetry(ctx, exec, &riverdriver.JobRetryParams{ + ID: id, + Now: c.baseService.Time.NowOrNil(), + Schema: c.config.Schema, + }) +} + +// JobUpdateParams contains parameters for Client.JobUpdate and Client.JobUpdateTx. +type JobUpdateParams struct { + // Output is a new output value for a job. + // + // If not set, and a job is updated from inside a work function, the job's + // output is set based on output recorded so far using RecordOutput. + Output any +} + +// JobUpdate updates the job with the given ID. +// +// If JobUpdateParams.Output is not set, this function may be used inside a job +// work function to set a job's output based on output recorded so far using +// RecordOutput. +func (c *Client[TTx]) JobUpdate(ctx context.Context, id int64, params *JobUpdateParams) (*rivertype.JobRow, error) { + return c.jobUpdate(ctx, c.driver.GetExecutor(), id, params) +} + +// JobUpdateTx updates the job with the given ID. +// +// If JobUpdateParams.Output is not set, this function may be used inside a job +// work function to set a job's output based on output recorded so far using +// RecordOutput. +// +// This variant updates the job inside of a transaction. +func (c *Client[TTx]) JobUpdateTx(ctx context.Context, tx TTx, id int64, params *JobUpdateParams) (*rivertype.JobRow, error) { + return c.jobUpdate(ctx, c.driver.UnwrapExecutor(tx), id, params) +} + +func (c *Client[TTx]) jobUpdate(ctx context.Context, exec riverdriver.Executor, id int64, params *JobUpdateParams) (*rivertype.JobRow, error) { + if params == nil { + params = &JobUpdateParams{} + } + + outputFromWorkContext := func() json.RawMessage { + metadataUpdates, hasMetadataUpdates := jobexecutor.MetadataUpdatesFromWorkContext(ctx) + if !hasMetadataUpdates { + return nil + } + + if val, ok := metadataUpdates[rivertype.MetadataKeyOutput]; ok { + return val.(json.RawMessage) //nolint:forcetypeassert + } + + return nil + }() + + var ( + metadataDoMerge bool + metadataUpdatesBytes = []byte("{}") // even in the event of no update, still valid jsonb + ) + if outputFromWorkContext != nil || params.Output != nil { + metadataDoMerge = true + + var outputBytes json.RawMessage + + switch { + // comes first because params takes precedence over context output + case params.Output != nil: + var err error + outputBytes, err = json.Marshal(params.Output) + if err != nil { + return nil, err + } + + if err := checkOutputSize(outputBytes); err != nil { + return nil, err + } + + case outputFromWorkContext != nil: + // no size check necessary here because it's already been checked in RecordOutput + outputBytes = outputFromWorkContext + } + + var err error + metadataUpdatesBytes, err = json.Marshal(map[string]json.RawMessage{ + rivertype.MetadataKeyOutput: outputBytes, + }) + if err != nil { + return nil, fmt.Errorf("error marshaling metadata updates to JSON: %w", err) + } + } + + return exec.JobUpdate(ctx, &riverdriver.JobUpdateParams{ + ID: id, + MetadataDoMerge: metadataDoMerge, + Metadata: metadataUpdatesBytes, + Schema: c.config.Schema, + }) +} + +// ID returns the unique ID of this client as set in its config or +// auto-generated if not specified. +func (c *Client[TTx]) ID() string { + return c.config.ID +} + +func insertParamsFromConfigArgsAndOptions(archetype *baseservice.Archetype, config *Config, args JobArgs, insertOpts *InsertOpts) (*rivertype.JobInsertParams, error) { + encodedArgs, err := json.Marshal(args) + if err != nil { + return nil, fmt.Errorf("error marshaling args to JSON: %w", err) + } + + if insertOpts == nil { + insertOpts = &InsertOpts{} + } + + var jobInsertOpts InsertOpts + if argsWithOpts, ok := args.(JobArgsWithInsertOpts); ok { + jobInsertOpts = argsWithOpts.InsertOpts() + } + + // If the time is stubbed (in a test), use that for `created_at`. Otherwise, + // leave an empty value which will either use the database's `now()` or be defaulted + // by drivers as necessary. + createdAt := archetype.Time.NowOrNil() + + maxAttempts := cmp.Or(insertOpts.MaxAttempts, jobInsertOpts.MaxAttempts, config.MaxAttempts) + priority := cmp.Or(insertOpts.Priority, jobInsertOpts.Priority, rivercommon.PriorityDefault) + queue := cmp.Or(insertOpts.Queue, jobInsertOpts.Queue, rivercommon.QueueDefault) + + if err := validateQueueName(queue); err != nil { + return nil, err + } + + tags := insertOpts.Tags + if insertOpts.Tags == nil { + tags = jobInsertOpts.Tags + } + if tags == nil { + tags = []string{} + } else { + for _, tag := range tags { + if len(tag) > 255 { + return nil, errors.New("tags should be a maximum of 255 characters long") + } + if !tagRE.MatchString(tag) { + return nil, errors.New("tags should match regex " + tagRE.String()) + } + } + } + + if priority < 1 || priority > 4 { + return nil, errors.New("priority must be between 1 and 4") + } + + var uniqueOpts UniqueOpts + if !config.Test.DisableUniqueEnforcement { + uniqueOpts = insertOpts.UniqueOpts + if uniqueOpts.isEmpty() { + uniqueOpts = jobInsertOpts.UniqueOpts + } + } + if err := uniqueOpts.validate(); err != nil { + return nil, err + } + + metadata := insertOpts.Metadata + if len(metadata) == 0 { + metadata = []byte("{}") + } + + insertParams := &rivertype.JobInsertParams{ + Args: args, + CreatedAt: createdAt, + EncodedArgs: encodedArgs, + Kind: args.Kind(), + MaxAttempts: maxAttempts, + Metadata: metadata, + Priority: priority, + Queue: queue, + State: rivertype.JobStateAvailable, + Tags: tags, + } + if !uniqueOpts.isEmpty() { + internalUniqueOpts := (*dbunique.UniqueOpts)(&uniqueOpts) + insertParams.UniqueKey, err = dbunique.UniqueKey(archetype.Time, internalUniqueOpts, insertParams) + if err != nil { + return nil, err + } + insertParams.UniqueStates = internalUniqueOpts.StateBitmask() + } + + switch { + case !insertOpts.ScheduledAt.IsZero(): + insertParams.ScheduledAt = &insertOpts.ScheduledAt + insertParams.State = rivertype.JobStateScheduled + case !jobInsertOpts.ScheduledAt.IsZero(): + insertParams.ScheduledAt = &jobInsertOpts.ScheduledAt + insertParams.State = rivertype.JobStateScheduled + default: + // Use a stubbed time if there was one, but otherwise prefer the value + // generated by the database. createdAt is nil unless time is stubbed. + insertParams.ScheduledAt = createdAt + } + + if insertOpts.Pending { + insertParams.State = rivertype.JobStatePending + } + + return insertParams, nil +} + +var errNoDriverDBPool = errors.New("driver must have non-nil database pool to use non-transactional methods like Insert and InsertMany (try InsertTx or InsertManyTx instead") + +// Insert inserts a new job with the provided args. Job opts can be used to +// override any defaults that may have been provided by an implementation of +// JobArgsWithInsertOpts.InsertOpts, as well as any global defaults. The +// provided context is used for the underlying Postgres insert and can be used +// to cancel the operation or apply a timeout. +// +// jobRow, err := client.Insert(insertCtx, MyArgs{}, nil) +// if err != nil { +// // handle error +// } +func (c *Client[TTx]) Insert(ctx context.Context, args JobArgs, opts *InsertOpts) (*rivertype.JobInsertResult, error) { + if !c.driver.PoolIsSet() { + return nil, errNoDriverDBPool + } + + res, err := dbutil.WithTxV(ctx, c.driver.GetExecutor(), func(ctx context.Context, execTx riverdriver.ExecutorTx) ([]*rivertype.JobInsertResult, error) { + return c.validateParamsAndInsertMany(ctx, execTx, []InsertManyParams{{Args: args, InsertOpts: opts}}) + }) + if err != nil { + return nil, err + } + + c.notifyProducerWithoutListenerJobFetch(ctx, res) + + return res[0], nil +} + +// InsertTx inserts a new job with the provided args on the given transaction. +// Job opts can be used to override any defaults that may have been provided by +// an implementation of JobArgsWithInsertOpts.InsertOpts, as well as any global +// defaults. The provided context is used for the underlying Postgres insert and +// can be used to cancel the operation or apply a timeout. +// +// jobRow, err := client.InsertTx(insertCtx, tx, MyArgs{}, nil) +// if err != nil { +// // handle error +// } +// +// This variant lets a caller insert jobs atomically alongside other database +// changes. It's also possible to insert a job outside a transaction, but this +// usage is recommended to ensure that all data a job needs to run is available +// by the time it starts. Because of snapshot visibility guarantees across +// transactions, the job will not be worked until the transaction has committed, +// and if the transaction rolls back, so too is the inserted job. +func (c *Client[TTx]) InsertTx(ctx context.Context, tx TTx, args JobArgs, opts *InsertOpts) (*rivertype.JobInsertResult, error) { + res, err := c.validateParamsAndInsertMany(ctx, c.driver.UnwrapExecutor(tx), []InsertManyParams{{Args: args, InsertOpts: opts}}) + if err != nil { + return nil, err + } + return res[0], nil +} + +// InsertManyParams encapsulates a single job combined with insert options for +// use with batch insertion. +type InsertManyParams struct { + // Args are the arguments of the job to insert. + Args JobArgs + + // InsertOpts are insertion options for this job. + InsertOpts *InsertOpts +} + +// InsertMany inserts many jobs at once. Each job is inserted as an +// InsertManyParams tuple, which takes job args along with an optional set of +// insert options, which override insert options provided by an +// JobArgsWithInsertOpts.InsertOpts implementation or any client-level defaults. +// The provided context is used for the underlying Postgres inserts and can be +// used to cancel the operation or apply a timeout. +// +// count, err := client.InsertMany(ctx, []river.InsertManyParams{ +// {Args: BatchInsertArgs{}}, +// {Args: BatchInsertArgs{}, InsertOpts: &river.InsertOpts{Priority: 3}}, +// }) +// if err != nil { +// // handle error +// } +func (c *Client[TTx]) InsertMany(ctx context.Context, params []InsertManyParams) ([]*rivertype.JobInsertResult, error) { + if !c.driver.PoolIsSet() { + return nil, errNoDriverDBPool + } + + res, err := dbutil.WithTxV(ctx, c.driver.GetExecutor(), func(ctx context.Context, execTx riverdriver.ExecutorTx) ([]*rivertype.JobInsertResult, error) { + return c.validateParamsAndInsertMany(ctx, execTx, params) + }) + if err != nil { + return nil, err + } + + c.notifyProducerWithoutListenerJobFetch(ctx, res) + + return res, nil +} + +// InsertManyTx inserts many jobs at once. Each job is inserted as an +// InsertManyParams tuple, which takes job args along with an optional set of +// insert options, which override insert options provided by an +// JobArgsWithInsertOpts.InsertOpts implementation or any client-level defaults. +// The provided context is used for the underlying Postgres inserts and can be +// used to cancel the operation or apply a timeout. +// +// count, err := client.InsertManyTx(ctx, tx, []river.InsertManyParams{ +// {Args: BatchInsertArgs{}}, +// {Args: BatchInsertArgs{}, InsertOpts: &river.InsertOpts{Priority: 3}}, +// }) +// if err != nil { +// // handle error +// } +// +// This variant lets a caller insert jobs atomically alongside other database +// changes. An inserted job isn't visible to be worked until the transaction +// commits, and if the transaction rolls back, so too is the inserted job. +func (c *Client[TTx]) InsertManyTx(ctx context.Context, tx TTx, params []InsertManyParams) ([]*rivertype.JobInsertResult, error) { + res, err := c.validateParamsAndInsertMany(ctx, c.driver.UnwrapExecutor(tx), params) + if err != nil { + return nil, err + } + return res, nil +} + +// validateParamsAndInsertMany is a helper method that wraps the insertMany +// method to provide param validation and conversion prior to calling the actual +// insertMany method. This allows insertMany to be reused by the +// PeriodicJobEnqueuer which cannot reference top-level river package types. +func (c *Client[TTx]) validateParamsAndInsertMany(ctx context.Context, execTx riverdriver.ExecutorTx, params []InsertManyParams) ([]*rivertype.JobInsertResult, error) { + insertParams, err := c.insertManyParams(params) + if err != nil { + return nil, err + } + + return c.insertMany(ctx, execTx, insertParams) +} + +// insertMany is a shared code path for InsertMany and InsertManyTx, also used +// by the PeriodicJobEnqueuer. +func (c *Client[TTx]) insertMany(ctx context.Context, execTx riverdriver.ExecutorTx, insertParams []*rivertype.JobInsertParams) ([]*rivertype.JobInsertResult, error) { + return c.insertManyShared(ctx, execTx, insertParams, func(ctx context.Context, insertParams []*riverdriver.JobInsertFastParams) ([]*rivertype.JobInsertResult, error) { + results, err := c.pilot.JobInsertMany(ctx, execTx, &riverdriver.JobInsertFastManyParams{ + Jobs: insertParams, + Schema: c.config.Schema, + }) + if err != nil { + return nil, err + } + + return sliceutil.Map(results, + func(result *riverdriver.JobInsertFastResult) *rivertype.JobInsertResult { + return (*rivertype.JobInsertResult)(result) + }, + ), nil + }) +} + +// The shared code path for all Insert and InsertMany methods. It takes a +// function that executes the actual insert operation and allows for different +// implementations of the insert query to be passed in, each mapping their +// results back to a common result type. +func (c *Client[TTx]) insertManyShared( + ctx context.Context, + tx riverdriver.ExecutorTx, + insertParams []*rivertype.JobInsertParams, + execute func(context.Context, []*riverdriver.JobInsertFastParams) ([]*rivertype.JobInsertResult, error), +) ([]*rivertype.JobInsertResult, error) { + doInner := func(ctx context.Context) ([]*rivertype.JobInsertResult, error) { + for _, params := range insertParams { + for _, hook := range append( + c.pluginLookupGlobal.ByKind(pluginlookup.PluginKindHookInsertBegin), + c.pluginLookupByJob.ByJobArgs(params.Args).ByKind(pluginlookup.PluginKindHookInsertBegin)..., + ) { + if err := hook.(rivertype.HookInsertBegin).InsertBegin(ctx, params); err != nil { //nolint:forcetypeassert + return nil, err + } + } + } + + finalInsertParams := sliceutil.Map(insertParams, func(params *rivertype.JobInsertParams) *riverdriver.JobInsertFastParams { + return (*riverdriver.JobInsertFastParams)(params) + }) + + insertResults, err := execute(ctx, finalInsertParams) + if err != nil { + return insertResults, err + } + + queues := make([]string, 0, 10) + for _, params := range insertParams { + if params.State == rivertype.JobStateAvailable { + queues = append(queues, params.Queue) + } + } + + if err = c.maybeNotifyInsertForQueues(ctx, tx, queues); err != nil { + return nil, err + } + + return insertResults, nil + } + + jobInsertMiddleware := append([]any(nil), c.pluginLookupGlobal.ByKind(pluginlookup.PluginKindMiddlewareJobInsert)...) + jobKindsSeen := make(map[string]struct{}, len(insertParams)) + for _, params := range insertParams { + kind := params.Args.Kind() + if _, ok := jobKindsSeen[kind]; ok { + continue + } + jobKindsSeen[kind] = struct{}{} + + jobInsertMiddleware = append( + jobInsertMiddleware, + c.pluginLookupByJob.ByJobArgs(params.Args).ByKind(pluginlookup.PluginKindMiddlewareJobInsert)..., + ) + } + if len(jobInsertMiddleware) > 0 { + // Wrap middlewares in reverse order so the one defined first is wrapped + // as the outermost function and is first to receive the operation. + for _, v := range slices.Backward(jobInsertMiddleware) { + middlewareItem := v.(rivertype.JobInsertMiddleware) //nolint:forcetypeassert // capture the current middleware item + previousDoInner := doInner // Capture the current doInner function + doInner = func(ctx context.Context) ([]*rivertype.JobInsertResult, error) { + return middlewareItem.InsertMany(ctx, insertParams, previousDoInner) + } + } + } + + return doInner(ctx) +} + +// Validates input parameters for a batch insert operation and generates a set +// of batch insert parameters. +func (c *Client[TTx]) insertManyParams(params []InsertManyParams) ([]*rivertype.JobInsertParams, error) { + if len(params) < 1 { + return nil, errors.New("no jobs to insert") + } + + insertParams := make([]*rivertype.JobInsertParams, len(params)) + for i, param := range params { + if err := c.validateJobArgs(param.Args); err != nil { + return nil, err + } + + insertParamsItem, err := insertParamsFromConfigArgsAndOptions(&c.baseService.Archetype, c.config, param.Args, param.InsertOpts) + if err != nil { + return nil, err + } + + insertParams[i] = insertParamsItem + } + + return insertParams, nil +} + +// Notifies an internal producer of new jobs being queued for work. Only +// invoked if the client's driver doesn't support a listener. If a listener is +// supported, job notifications go out via listen/notify instead. +// +// Should only ever be invoked *outside* a transaction. If invoked within a +// transaction, the producer wouldn't yet be able to access the new jobs that +// triggered the notification because they're not committed yet. +func (c *Client[TTx]) notifyProducerWithoutListenerJobFetch(_ context.Context, res []*rivertype.JobInsertResult) { + if c.driver.SupportsListener() { + return + } + + c.producersMu.RLock() + defer c.producersMu.RUnlock() + + if len(c.producersByQueueName) < 1 { + return + } + + // Special case for when we were handling exactly one job, which is a very + // common case. Acts as a minor optimization by avoiding the map allocation. + if len(res) == 1 { + if producer, ok := c.producersByQueueName[res[0].Job.Queue]; ok { + producer.TriggerJobFetch() + } + + return + } + + queuesTriggered := make(map[string]struct{}) + + for _, insertRes := range res { + if _, ok := queuesTriggered[insertRes.Job.Queue]; ok { + continue + } + queuesTriggered[insertRes.Job.Queue] = struct{}{} + + if producer, ok := c.producersByQueueName[insertRes.Job.Queue]; ok { + producer.TriggerJobFetch() + } + } +} + +// InsertManyFast inserts many jobs at once using Postgres' `COPY FROM` mechanism, +// making the operation quite fast and memory efficient. Each job is inserted as +// an InsertManyParams tuple, which takes job args along with an optional set of +// insert options, which override insert options provided by an +// JobArgsWithInsertOpts.InsertOpts implementation or any client-level defaults. +// The provided context is used for the underlying Postgres inserts and can be +// used to cancel the operation or apply a timeout. +// +// count, err := client.InsertMany(ctx, []river.InsertManyParams{ +// {Args: BatchInsertArgs{}}, +// {Args: BatchInsertArgs{}, InsertOpts: &river.InsertOpts{Priority: 3}}, +// }) +// if err != nil { +// // handle error +// } +// +// Unlike with `InsertMany`, unique conflicts cannot be handled gracefully. If a +// unique constraint is violated, the operation will fail and no jobs will be inserted. +func (c *Client[TTx]) InsertManyFast(ctx context.Context, params []InsertManyParams) (int, error) { + if !c.driver.PoolIsSet() { + return 0, errNoDriverDBPool + } + + // Wrap in a transaction in case we need to notify about inserts. + res, err := dbutil.WithTxV(ctx, c.driver.GetExecutor(), func(ctx context.Context, execTx riverdriver.ExecutorTx) ([]*rivertype.JobInsertResult, error) { + return c.insertManyFast(ctx, execTx, params) + }) + if err != nil { + return 0, err + } + + c.notifyProducerWithoutListenerJobFetch(ctx, res) + + return len(res), nil +} + +// InsertManyFastTx inserts many jobs at once using Postgres' `COPY FROM` +// mechanism, making the operation quite fast and memory efficient. Each job is +// inserted as an InsertManyParams tuple, which takes job args along with an +// optional set of insert options, which override insert options provided by an +// JobArgsWithInsertOpts.InsertOpts implementation or any client-level defaults. +// The provided context is used for the underlying Postgres inserts and can be +// used to cancel the operation or apply a timeout. +// +// count, err := client.InsertManyTx(ctx, tx, []river.InsertManyParams{ +// {Args: BatchInsertArgs{}}, +// {Args: BatchInsertArgs{}, InsertOpts: &river.InsertOpts{Priority: 3}}, +// }) +// if err != nil { +// // handle error +// } +// +// This variant lets a caller insert jobs atomically alongside other database +// changes. An inserted job isn't visible to be worked until the transaction +// commits, and if the transaction rolls back, so too is the inserted job. +// +// Unlike with `InsertManyTx`, unique conflicts cannot be handled gracefully. If +// a unique constraint is violated, the operation will fail and no jobs will be +// inserted. +func (c *Client[TTx]) InsertManyFastTx(ctx context.Context, tx TTx, params []InsertManyParams) (int, error) { + res, err := c.insertManyFast(ctx, c.driver.UnwrapExecutor(tx), params) + if err != nil { + return 0, err + } + return len(res), nil +} + +func (c *Client[TTx]) insertManyFast(ctx context.Context, execTx riverdriver.ExecutorTx, params []InsertManyParams) ([]*rivertype.JobInsertResult, error) { + insertParams, err := c.insertManyParams(params) + if err != nil { + return nil, err + } + + return c.insertManyShared(ctx, execTx, insertParams, func(ctx context.Context, insertParams []*riverdriver.JobInsertFastParams) ([]*rivertype.JobInsertResult, error) { + count, err := execTx.JobInsertFastManyNoReturning(ctx, &riverdriver.JobInsertFastManyParams{ + Jobs: insertParams, + Schema: c.config.Schema, + }) + if err != nil { + return nil, err + } + return make([]*rivertype.JobInsertResult, count), nil + }) +} + +// Notify the given queues that new jobs are available. The queues list will be +// deduplicated and each will be checked to see if it is due for an insert +// notification from this client. +func (c *Client[TTx]) maybeNotifyInsertForQueues(ctx context.Context, tx riverdriver.ExecutorTx, queues []string) error { + if len(queues) < 1 { + return nil + } + + var ( + queuesDeduped = sliceutil.Uniq(queues) + payloads = make([]string, 0, len(queuesDeduped)) + queuesTriggered = make([]string, 0, len(queuesDeduped)) + ) + + for _, queue := range queuesDeduped { + if c.insertNotifyLimiter.ShouldTrigger(queue) { + payloads = append(payloads, fmt.Sprintf("{\"queue\": %q}", queue)) + queuesTriggered = append(queuesTriggered, queue) + } + } + + if len(payloads) < 1 { + return nil + } + + if c.driver.SupportsListenNotify() { + err := tx.NotifyMany(ctx, &riverdriver.NotifyManyParams{ + Payload: payloads, + Schema: c.config.Schema, + Topic: string(notifier.NotificationTopicInsert), + }) + if err != nil { + c.baseService.Logger.ErrorContext( + ctx, + c.baseService.Name+": Failed to send job insert notification", + slog.String("queues", strings.Join(queuesTriggered, ",")), + slog.String("err", err.Error()), + ) + return err + } + } + + return nil +} + +// emit a notification about a queue being paused or resumed. +func (c *Client[TTx]) notifyQueuePauseOrResume(ctx context.Context, tx riverdriver.ExecutorTx, action controlAction, queue string, opts *QueuePauseOpts) (*controlEventPayload, error) { + c.baseService.Logger.DebugContext(ctx, + c.baseService.Name+": Notifying about queue state change", + slog.String("action", string(action)), + slog.String("queue", queue), + slog.String("opts", fmt.Sprintf("%+v", opts)), + ) + + controlEvent := &controlEventPayload{Action: action, Queue: queue} + + payload, err := json.Marshal(controlEvent) + if err != nil { + return nil, err + } + + if c.driver.SupportsListenNotify() { + err = tx.NotifyMany(ctx, &riverdriver.NotifyManyParams{ + Payload: []string{string(payload)}, + Schema: c.config.Schema, + Topic: string(notifier.NotificationTopicControl), + }) + if err != nil { + c.baseService.Logger.ErrorContext( + ctx, + c.baseService.Name+": Failed to send queue state change notification", + slog.String("err", err.Error()), + ) + return nil, err + } + } + + return controlEvent, nil +} + +// Validates job args prior to insertion. Currently, verifies that a worker to +// handle the kind is registered in the configured workers bundle. +// This validation is skipped if the client is configured as an insert-only (with no workers) +// or if the client is configured to skip unknown job kinds. +func (c *Client[TTx]) validateJobArgs(args JobArgs) error { + if c.config.Workers == nil || c.config.SkipUnknownJobCheck { + return nil + } + + if _, ok := c.config.Workers.workersMap[args.Kind()]; !ok { + return &UnknownJobKindError{Kind: args.Kind()} + } + + return nil +} + +func (c *Client[TTx]) producerAdd(queueName string, queueConfig QueueConfig) (*producer, error) { + c.producersMu.Lock() + defer c.producersMu.Unlock() + + if _, alreadyExists := c.producersByQueueName[queueName]; alreadyExists { + return nil, &QueueAlreadyAddedError{Name: queueName} + } + + producer := newProducer(&c.baseService.Archetype, c.driver.GetExecutor(), c.pilot, &producerConfig{ + ClientID: c.config.ID, + Completer: c.completer, + ErrorHandler: c.config.ErrorHandler, + FetchCooldown: cmp.Or(queueConfig.FetchCooldown, c.config.FetchCooldown), + FetchPollInterval: cmp.Or(queueConfig.FetchPollInterval, c.config.FetchPollInterval), + PluginLookupByJob: c.pluginLookupByJob, + PluginLookupGlobal: c.pluginLookupGlobal, + JobStuckHandler: c.config.JobStuckHandler, + JobStuckCount: &c.stuckJobCount, + JobStuckThreshold: c.config.JobStuckThreshold, + JobTimeout: c.config.JobTimeout, + MaxWorkers: queueConfig.MaxWorkers, + Notifier: c.notifier, + Queue: queueName, + QueueEventCallback: c.subscriptionManager.distributeQueueEvent, + QueuePollInterval: c.config.queuePollInterval, + RetryPolicy: c.config.RetryPolicy, + SchedulerInterval: c.config.schedulerInterval, + Schema: c.config.Schema, + StaleProducerRetentionPeriod: 5 * time.Minute, + Workers: c.config.Workers, + }) + c.producersByQueueName[queueName] = producer + return producer, nil +} + +func (c *Client[TTx]) producerRemove(ctx context.Context, queueName string) error { + c.producersMu.Lock() + defer c.producersMu.Unlock() + + producer, ok := c.producersByQueueName[queueName] + if !ok { + return &QueueNotFoundError{Name: queueName} + } + + shouldStop, stopped, finalizeStop := producer.StopInit() + if shouldStop { + select { + case <-ctx.Done(): + finalizeStop(false) + return ctx.Err() + case <-stopped: + finalizeStop(true) + } + } + + delete(c.producersByQueueName, queueName) + + return nil +} + +var nameRegex = regexp.MustCompile(`^(?:[a-z0-9])+(?:[_|\-]?[a-z0-9]+)*$`) + +func validateQueueName(queueName string) error { + if queueName == "" { + return errors.New("queue name cannot be empty") + } + if len(queueName) > 64 { + return errors.New("queue name cannot be longer than 64 characters") + } + if !nameRegex.MatchString(queueName) { + return fmt.Errorf("queue name is invalid, expected letters and numbers separated by underscores or hyphens: %q", queueName) + } + return nil +} + +// JobDeleteManyResult is the result of a job list operation. It contains a list of +// jobs and a cursor for fetching the next page of results. +type JobDeleteManyResult struct { + // Jobs is a slice of job returned as part of the list operation. + Jobs []*rivertype.JobRow +} + +// JobDeleteMany deletes many jobs at once based on the conditions defined by +// JobDeleteManyParams. Running jobs are always ignored. +// +// params := river.NewJobDeleteManyParams().First(10).State(rivertype.JobStateCompleted) +// jobRows, err := client.JobDeleteMany(ctx, params) +// if err != nil { +// // handle error +// } +func (c *Client[TTx]) JobDeleteMany(ctx context.Context, params *JobDeleteManyParams) (*JobDeleteManyResult, error) { + if !c.driver.PoolIsSet() { + return nil, errNoDriverDBPool + } + + return c.jobDeleteMany(ctx, c.driver.GetExecutor(), params) +} + +// JobDeleteManyTx deletes many jobs at once based on the conditions defined by +// JobDeleteManyParams. Running jobs are always ignored. +// +// params := river.NewJobDeleteManyParams().First(10).States(river.JobStateCompleted) +// jobRows, err := client.JobDeleteManyTx(ctx, tx, params) +// if err != nil { +// // handle error +// } +func (c *Client[TTx]) JobDeleteManyTx(ctx context.Context, tx TTx, params *JobDeleteManyParams) (*JobDeleteManyResult, error) { + return c.jobDeleteMany(ctx, c.driver.UnwrapExecutor(tx), params) +} + +func (c *Client[TTx]) jobDeleteMany(ctx context.Context, exec riverdriver.Executor, params *JobDeleteManyParams) (*JobDeleteManyResult, error) { + if params == nil { + params = NewJobDeleteManyParams() + } + params.schema = c.config.Schema + + if params.filtersEmpty() && !params.unsafeAll { + return nil, errors.New("delete with no filters not allowed to prevent accidental deletion of all jobs; either specify a predicate (e.g. JobDeleteManyParams.IDs, JobDeleteManyParams.Kinds, ...) or call JobDeleteManyParams.All") + } + + listParams, err := dblist.JobMakeDriverParams(ctx, params.toDBParams(), c.driver) + if err != nil { + return nil, err + } + + jobs, err := exec.JobDeleteMany(ctx, (*riverdriver.JobDeleteManyParams)(listParams)) + if err != nil { + return nil, err + } + + return &JobDeleteManyResult{Jobs: jobs}, nil +} + +// JobListResult is the result of a job list operation. It contains a list of +// jobs and a cursor for fetching the next page of results. +type JobListResult struct { + // Jobs is a slice of job returned as part of the list operation. + Jobs []*rivertype.JobRow + + // LastCursor is a cursor that can be used to list the next page of jobs. + LastCursor *JobListCursor +} + +var errJobListParamsMetadataNotSupportedSQLite = errors.New("JobListParams.Metadata is not supported on SQLite") + +// JobList returns a paginated list of jobs matching the provided filters. The +// provided context is used for the underlying Postgres query and can be used to +// cancel the operation or apply a timeout. +// +// params := river.NewJobListParams().First(10).State(rivertype.JobStateCompleted) +// jobRows, err := client.JobList(ctx, params) +// if err != nil { +// // handle error +// } +func (c *Client[TTx]) JobList(ctx context.Context, params *JobListParams) (*JobListResult, error) { + if !c.driver.PoolIsSet() { + return nil, errNoDriverDBPool + } + + if params == nil { + params = NewJobListParams() + } + params.schema = c.config.Schema + + if c.driver.DatabaseName() == riverdriver.DatabaseNameSQLite && params.metadataCalled { + return nil, errJobListParamsMetadataNotSupportedSQLite + } + + dbParams, err := params.toDBParams() + if err != nil { + return nil, err + } + + listParams, err := dblist.JobMakeDriverParams(ctx, dbParams, c.driver) + if err != nil { + return nil, err + } + + jobs, err := c.driver.GetExecutor().JobList(ctx, listParams) + if err != nil { + return nil, err + } + + res := &JobListResult{Jobs: jobs} + if len(jobs) > 0 { + res.LastCursor = jobListCursorFromJobAndParams(jobs[len(jobs)-1], params) + } + return res, nil +} + +// JobListTx returns a paginated list of jobs matching the provided filters. The +// provided context is used for the underlying Postgres query and can be used to +// cancel the operation or apply a timeout. +// +// params := river.NewJobListParams().First(10).States(river.JobStateCompleted) +// jobRows, err := client.JobListTx(ctx, tx, params) +// if err != nil { +// // handle error +// } +func (c *Client[TTx]) JobListTx(ctx context.Context, tx TTx, params *JobListParams) (*JobListResult, error) { + if params == nil { + params = NewJobListParams() + } + params.schema = c.config.Schema + + if c.driver.DatabaseName() == riverdriver.DatabaseNameSQLite && params.metadataCalled { + return nil, errJobListParamsMetadataNotSupportedSQLite + } + + dbParams, err := params.toDBParams() + if err != nil { + return nil, err + } + + listParams, err := dblist.JobMakeDriverParams(ctx, dbParams, c.driver) + if err != nil { + return nil, err + } + + jobs, err := c.driver.UnwrapExecutor(tx).JobList(ctx, listParams) + if err != nil { + return nil, err + } + + res := &JobListResult{Jobs: jobs} + if len(jobs) > 0 { + res.LastCursor = jobListCursorFromJobAndParams(jobs[len(jobs)-1], params) + } + return res, nil +} + +// Notify retrieves a notification bundle for the client (in the sense of +// Postgres listen/notify) used to send notifications of various kinds. +func (c *Client[TTx]) Notify() *ClientNotifyBundle[TTx] { + return c.clientNotifyBundle +} + +// ClientNotifyBundle sends various notifications for a client (in the sense of +// Postgres listen/notify). Functions are on this bundle struct instead of the +// top-level client to keep them grouped together and better organized. +type ClientNotifyBundle[TTx any] struct { + config *Config + driver riverdriver.Driver[TTx] +} + +// RequestResign sends a notification requesting that the current leader resign. +// This usually causes the resignation of the current leader, but may have no +// effect if no leader is currently elected. +func (c *ClientNotifyBundle[TTx]) RequestResign(ctx context.Context) error { + return dbutil.WithTx(ctx, c.driver.GetExecutor(), func(ctx context.Context, execTx riverdriver.ExecutorTx) error { + return c.requestResignTx(ctx, execTx) + }) +} + +// RequestResignTx sends a notification requesting that the current leader +// resign. This usually causes the resignation of the current leader, but may +// have no effect if no leader is currently elected. +// +// This variant sends a notification in a transaction, which means that no +// notification is sent until the transaction commits. +func (c *ClientNotifyBundle[TTx]) RequestResignTx(ctx context.Context, tx TTx) error { + return c.requestResignTx(ctx, c.driver.UnwrapExecutor(tx)) +} + +// notifyExecTx is a shared helper between Notify and NotifyTx that sends a +// notification. +func (c *ClientNotifyBundle[TTx]) requestResignTx(ctx context.Context, execTx riverdriver.ExecutorTx) error { + payloadStr, err := json.Marshal(&leadership.DBNotification{ + Action: leadership.DBNotificationKindRequestResign, + }) + if err != nil { + return err + } + + return execTx.NotifyMany(ctx, &riverdriver.NotifyManyParams{ + Payload: []string{string(payloadStr)}, + Schema: c.config.Schema, + Topic: string(notifier.NotificationTopicLeadership), + }) +} + +// PeriodicJobs returns the currently configured set of periodic jobs for the +// client, and can be used to add new or remove existing ones. +// +// This function should only be invoked on clients capable of running perioidc +// jobs. Running periodic jobs requires that the client be electable as leader +// to run maintenance services, and being electable as leader requires that a +// client be started. To be startable, a client must have Queues and Workers +// configured. Invoking this function will panic if these conditions aren't met. +func (c *Client[TTx]) PeriodicJobs() *PeriodicJobBundle { + if !c.config.willExecuteJobs() { + panic("client Queues and Workers must be configured to modify periodic jobs (otherwise, they'll have no effect because a client not configured to work jobs can't be started)") + } + + return c.periodicJobs +} + +// Pilot returns the pilot in use by the pilot. If not configured, this is often +// simply StandardPilot. +// +// API is not stable. DO NOT USE. +func (c *Client[TTx]) Pilot() riverpilot.Pilot { + return c.pilot +} + +// Queues returns the currently configured set of queues for the client, and can +// be used to add new ones. +func (c *Client[TTx]) Queues() *QueueBundle { return c.queues } + +// QueueGet returns the queue with the given name. If the queue has not recently +// been active or does not exist, returns ErrNotFound. +// +// The provided context is used for the underlying Postgres query and can be +// used to cancel the operation or apply a timeout. +func (c *Client[TTx]) QueueGet(ctx context.Context, name string) (*rivertype.Queue, error) { + return c.driver.GetExecutor().QueueGet(ctx, &riverdriver.QueueGetParams{ + Name: name, + Schema: c.config.Schema, + }) +} + +// QueueGetTx returns the queue with the given name. If the queue has not recently +// been active or does not exist, returns ErrNotFound. +// +// The provided context is used for the underlying Postgres query and can be +// used to cancel the operation or apply a timeout. +func (c *Client[TTx]) QueueGetTx(ctx context.Context, tx TTx, name string) (*rivertype.Queue, error) { + return c.driver.UnwrapExecutor(tx).QueueGet(ctx, &riverdriver.QueueGetParams{ + Name: name, + Schema: c.config.Schema, + }) +} + +// QueueListResult is the result of a job list operation. It contains a list of +// jobs and leaves room for future cursor functionality. +type QueueListResult struct { + // Queues is a slice of queues returned as part of the list operation. + Queues []*rivertype.Queue +} + +// QueueList returns a list of all queues that are currently active or were +// recently active. Limit and offset can be used to paginate the results. +// +// The provided context is used for the underlying Postgres query and can be +// used to cancel the operation or apply a timeout. +// +// params := river.NewQueueListParams().First(10) +// queueRows, err := client.QueueListTx(ctx, tx, params) +// if err != nil { +// // handle error +// } +func (c *Client[TTx]) QueueList(ctx context.Context, params *QueueListParams) (*QueueListResult, error) { + if params == nil { + params = NewQueueListParams() + } + + queues, err := c.driver.GetExecutor().QueueList(ctx, &riverdriver.QueueListParams{ + Max: int(params.paginationCount), + Schema: c.config.Schema, + }) + if err != nil { + return nil, err + } + + return &QueueListResult{Queues: queues}, nil +} + +// QueueListTx returns a list of all queues that are currently active or were +// recently active. Limit and offset can be used to paginate the results. +// +// The provided context is used for the underlying Postgres query and can be +// used to cancel the operation or apply a timeout. +// +// params := river.NewQueueListParams().First(10) +// queueRows, err := client.QueueListTx(ctx, tx, params) +// if err != nil { +// // handle error +// } +func (c *Client[TTx]) QueueListTx(ctx context.Context, tx TTx, params *QueueListParams) (*QueueListResult, error) { + if params == nil { + params = NewQueueListParams() + } + + queues, err := c.driver.UnwrapExecutor(tx).QueueList(ctx, &riverdriver.QueueListParams{ + Max: int(params.paginationCount), + Schema: c.config.Schema, + }) + if err != nil { + return nil, err + } + + return &QueueListResult{Queues: queues}, nil +} + +// QueuePause pauses the queue with the given name. When a queue is paused, +// clients will not fetch any more jobs for that particular queue. To pause all +// queues at once, use the special queue name "*". +// +// Clients with a configured notifier should receive a notification about the +// paused queue(s) within a few milliseconds of the transaction commit. Clients +// in poll-only mode will pause after their next poll for queue configuration. +// +// The provided context is used for the underlying Postgres update and can be +// used to cancel the operation or apply a timeout. The opts are reserved for +// future functionality. +func (c *Client[TTx]) QueuePause(ctx context.Context, name string, opts *QueuePauseOpts) error { + tx, err := c.driver.GetExecutor().Begin(ctx) + if err != nil { + return err + } + defer dbutil.RollbackWithoutCancel(ctx, tx) + + if err := tx.QueuePause(ctx, &riverdriver.QueuePauseParams{ + Name: name, + Now: c.baseService.Time.NowOrNil(), + Schema: c.config.Schema, + }); err != nil { + return err + } + + controlEvent, err := c.notifyQueuePauseOrResume(ctx, tx, controlActionPause, name, opts) + if err != nil { + return err + } + + if err = tx.Commit(ctx); err != nil { + return err + } + + c.notifyProducerWithoutListenerQueueControlEvent(name, controlEvent) + + return nil +} + +// QueuePauseTx pauses the queue with the given name. When a queue is paused, +// clients will not fetch any more jobs for that particular queue. To pause all +// queues at once, use the special queue name "*". +// +// Clients with a configured notifier should receive a notification about the +// paused queue(s) within a few milliseconds of the transaction commit. Clients +// in poll-only mode will pause after their next poll for queue configuration. +// +// The provided context is used for the underlying Postgres update and can be +// used to cancel the operation or apply a timeout. The opts are reserved for +// future functionality. +func (c *Client[TTx]) QueuePauseTx(ctx context.Context, tx TTx, name string, opts *QueuePauseOpts) error { + executorTx := c.driver.UnwrapExecutor(tx) + + if err := executorTx.QueuePause(ctx, &riverdriver.QueuePauseParams{ + Name: name, + Now: c.baseService.Time.NowOrNil(), + Schema: c.config.Schema, + }); err != nil { + return err + } + + if _, err := c.notifyQueuePauseOrResume(ctx, executorTx, controlActionPause, name, opts); err != nil { + return err + } + + return nil +} + +// QueueResume resumes the queue with the given name. If the queue was +// previously paused, any clients configured to work that queue will resume +// fetching additional jobs. To resume all queues at once, use the special queue +// name "*". +// +// Clients with a configured notifier should receive a notification about the +// resumed queue(s) within a few milliseconds of the transaction commit. Clients +// in poll-only mode will resume after their next poll for queue configuration. +// +// The provided context is used for the underlying Postgres update and can be +// used to cancel the operation or apply a timeout. The opts are reserved for +// future functionality. +func (c *Client[TTx]) QueueResume(ctx context.Context, name string, opts *QueuePauseOpts) error { + tx, err := c.driver.GetExecutor().Begin(ctx) + if err != nil { + return err + } + defer dbutil.RollbackWithoutCancel(ctx, tx) + + if err := tx.QueueResume(ctx, &riverdriver.QueueResumeParams{ + Name: name, + Now: c.baseService.Time.NowOrNil(), + Schema: c.config.Schema, + }); err != nil { + return err + } + + controlEvent, err := c.notifyQueuePauseOrResume(ctx, tx, controlActionResume, name, opts) + if err != nil { + return err + } + + if err = tx.Commit(ctx); err != nil { + return err + } + + c.notifyProducerWithoutListenerQueueControlEvent(name, controlEvent) + + return nil +} + +// QueueResumeTx resumes the queue with the given name. If the queue was +// previously paused, any clients configured to work that queue will resume +// fetching additional jobs. To resume all queues at once, use the special queue +// name "*". +// +// Clients with a configured notifier should receive a notification about the +// resumed queue(s) within a few milliseconds of the transaction commit. Clients +// in poll-only mode will resume after their next poll for queue configuration. +// +// The provided context is used for the underlying Postgres update and can be +// used to cancel the operation or apply a timeout. The opts are reserved for +// future functionality. +func (c *Client[TTx]) QueueResumeTx(ctx context.Context, tx TTx, name string, opts *QueuePauseOpts) error { + executorTx := c.driver.UnwrapExecutor(tx) + + if err := executorTx.QueueResume(ctx, &riverdriver.QueueResumeParams{ + Name: name, + Now: c.baseService.Time.NowOrNil(), + Schema: c.config.Schema, + }); err != nil { + return err + } + + if _, err := c.notifyQueuePauseOrResume(ctx, executorTx, controlActionResume, name, opts); err != nil { + return err + } + + return nil +} + +// QueueUpdateParams are the parameters for a QueueUpdate operation. +type QueueUpdateParams struct { + // Metadata is the new metadata for the queue. If nil or empty, the metadata + // will not be changed. + Metadata []byte +} + +// QueueUpdate updates a queue's settings in the database. These settings +// override the settings in the client (if applied). +func (c *Client[TTx]) QueueUpdate(ctx context.Context, name string, params *QueueUpdateParams) (*rivertype.Queue, error) { + tx, err := c.driver.GetExecutor().Begin(ctx) + if err != nil { + return nil, err + } + defer dbutil.RollbackWithoutCancel(ctx, tx) + + queue, controlEvent, err := c.queueUpdate(ctx, tx, name, params) + if err != nil { + return nil, err + } + + if err := tx.Commit(ctx); err != nil { + return nil, err + } + + c.notifyProducerWithoutListenerQueueControlEvent(name, controlEvent) + + return queue, nil +} + +// QueueUpdateTx updates a queue's settings in the database. These settings +// override the settings in the client (if applied). +func (c *Client[TTx]) QueueUpdateTx(ctx context.Context, tx TTx, name string, params *QueueUpdateParams) (*rivertype.Queue, error) { + queue, _, err := c.queueUpdate(ctx, c.driver.UnwrapExecutor(tx), name, params) + if err != nil { + return nil, err + } + return queue, nil +} + +// Notifies an internal producer of a queue control event like pause/resume. +// Only invoked if the client's driver doesn't support a listener. If a listener +// is supported, control events go out via listen/notify instead. +// +// Should only ever be invoked *outside* a transaction. If invoked within a +// transaction, the producer wouldn't yet be able to access the state that +// triggered the notification because it's not committed yet. +func (c *Client[TTx]) notifyProducerWithoutListenerQueueControlEvent(queue string, controlEvent *controlEventPayload) { + if c.driver.SupportsListener() { + return + } + + c.producersMu.RLock() + defer c.producersMu.RUnlock() + + if len(c.producersByQueueName) < 1 { + return + } + + if producer, ok := c.producersByQueueName[queue]; ok { + producer.TriggerQueueControlEvent(controlEvent) + } +} + +func (c *Client[TTx]) queueUpdate(ctx context.Context, executorTx riverdriver.ExecutorTx, name string, params *QueueUpdateParams) (*rivertype.Queue, *controlEventPayload, error) { + updateMetadata := len(params.Metadata) > 0 + + queue, err := executorTx.QueueUpdate(ctx, &riverdriver.QueueUpdateParams{ + Metadata: params.Metadata, + MetadataDoUpdate: updateMetadata, + Name: name, + Schema: c.config.Schema, + }) + if err != nil { + return nil, nil, err + } + + if !updateMetadata { + return queue, nil, err + } + + controlEvent := &controlEventPayload{ + Action: controlActionMetadataChanged, + Metadata: params.Metadata, + Queue: queue.Name, + } + + payload, err := json.Marshal(controlEvent) + if err != nil { + return nil, nil, err + } + + if c.driver.SupportsListenNotify() { + if err := executorTx.NotifyMany(ctx, &riverdriver.NotifyManyParams{ + Payload: []string{string(payload)}, + Schema: c.config.Schema, + Topic: string(notifier.NotificationTopicControl), + }); err != nil { + return nil, nil, err + } + } + + return queue, controlEvent, nil +} + +// Schema returns the configured schema for the client. +func (c *Client[TTx]) Schema() string { return c.config.Schema } + +// QueueBundle is a bundle for adding additional queues. It's made accessible +// through Client.Queues. +type QueueBundle struct { + clientFetchCooldown time.Duration + clientFetchPollInterval time.Duration + + clientWillExecuteJobs bool + + fetchCtx context.Context //nolint:containedctx + producerAdd func(queueName string, queueConfig QueueConfig) (*producer, error) // add producer to associated client + producerRemove func(ctx context.Context, queueName string) error // remove producer from associated client + + // Mutex that's acquired when client is starting and stopping and when a + // queue is being added so that we can be sure that a client is fully + // stopped or fully started when adding a new queue. + startStopMu sync.Mutex + + workCtx context.Context //nolint:containedctx +} + +// Add adds a new queue to the client. If the client is already started, a +// producer for the queue is started. Context is inherited from the one given to +// Client.Start. +func (b *QueueBundle) Add(queueName string, queueConfig QueueConfig) error { + if !b.clientWillExecuteJobs { + return errors.New("client is not configured to execute jobs, cannot add queue") + } + + if err := queueConfig.validate(queueName, b.clientFetchCooldown, b.clientFetchPollInterval); err != nil { + return err + } + + b.startStopMu.Lock() + defer b.startStopMu.Unlock() + + producer, err := b.producerAdd(queueName, queueConfig) + if err != nil { + return err + } + + // Start the queue if the client is already started. + if b.fetchCtx != nil && b.fetchCtx.Err() == nil { + if err := producer.StartWorkContext(b.fetchCtx, b.workCtx); err != nil { + return err + } + } + + return nil +} + +// Remove removes a queue from the client, stopping the producer if the client +// is running. It waits for any jobs currently being worked in the queue to +// complete before returning. +// +// If the provided context is done before the producer has fully stopped, Remove +// returns the context's error and does not fully remove the queue, though the +// queue's producer may have started stopping, leaving it in a stopped state +// where it doesn't work new jobs. Call Remove again with a new context to +// remove it completely. +// +// Returns an error if the client is not configured to execute jobs or if the +// specified queue does not exist. +func (b *QueueBundle) Remove(ctx context.Context, queueName string) error { + if !b.clientWillExecuteJobs { + return errors.New("client is not configured to execute jobs, cannot remove queue") + } + + b.startStopMu.Lock() + defer b.startStopMu.Unlock() + + return b.producerRemove(ctx, queueName) +} + +// Generates a default client ID using the current hostname and time. +func defaultClientID(startedAt time.Time) string { + host, _ := os.Hostname() + if host == "" { + host = "unknown_host" + } + + return defaultClientIDWithHost(startedAt, host) +} + +// Same as the above, but allows host injection for testability. +func defaultClientIDWithHost(startedAt time.Time, host string) string { + const maxHostLength = 60 + + // Truncate degenerately long host names. + host = strings.ReplaceAll(host, ".", "_") + if len(host) > maxHostLength { + host = host[0:maxHostLength] + } + + // Dots, hyphens, and colons aren't particularly friendly for double click + // to select (depends on application and configuration), so avoid them all + // in favor of underscores. + // + // Go's time package is really dumb and can't format subseconds without + // using a dot. So use the dot, then replace it with an underscore below. + const rfc3339Compact = "2006_01_02T15_04_05.000000" + + return host + "_" + strings.Replace(startedAt.Format(rfc3339Compact), ".", "_", 1) +} diff --git a/vendor/github.com/riverqueue/river/client_context.go b/vendor/github.com/riverqueue/river/client_context.go new file mode 100644 index 0000000000..4bbded2ee9 --- /dev/null +++ b/vendor/github.com/riverqueue/river/client_context.go @@ -0,0 +1,54 @@ +package river + +import ( + "context" + "errors" + + "github.com/riverqueue/river/internal/rivercommon" +) + +var errClientNotInContext = errors.New("river: client not found in context, can only be used in a Worker") + +func withClient[TTx any](ctx context.Context, client *Client[TTx]) context.Context { + return context.WithValue(ctx, rivercommon.ContextKeyClient{}, client) +} + +// ClientFromContext returns the Client from the context. This function can +// only be used within a Worker's Work() method because that is the only place +// River sets the Client on the context. +// +// It panics if the context does not contain a Client, which will never happen +// from the context provided to a Worker's Work() method. +// +// When testing JobArgs.Work implementations, it might be useful to use +// rivertest.WorkContext to initialize a context that has an available client. +// +// The type parameter TTx is the transaction type used by the [Client], +// pgx.Tx for the pgx driver, and *sql.Tx for the [database/sql] driver. +func ClientFromContext[TTx any](ctx context.Context) *Client[TTx] { + client, err := ClientFromContextSafely[TTx](ctx) + if err != nil { + panic(err) + } + return client +} + +// ClientFromContextSafely returns the Client from the context. This function +// can only be used within a Worker's Work() method because that is the only +// place River sets the Client on the context. +// +// It returns an error if the context does not contain a Client, which will +// never happen from the context provided to a Worker's Work() method. +// +// When testing JobArgs.Work implementations, it might be useful to use +// rivertest.WorkContext to initialize a context that has an available client. +// +// See the examples for [ClientFromContext] to understand how to use this +// function. +func ClientFromContextSafely[TTx any](ctx context.Context) (*Client[TTx], error) { + client, exists := ctx.Value(rivercommon.ContextKeyClient{}).(*Client[TTx]) + if !exists || client == nil { + return nil, errClientNotInContext + } + return client, nil +} diff --git a/vendor/github.com/riverqueue/river/delete_many_params.go b/vendor/github.com/riverqueue/river/delete_many_params.go new file mode 100644 index 0000000000..9cca79416a --- /dev/null +++ b/vendor/github.com/riverqueue/river/delete_many_params.go @@ -0,0 +1,150 @@ +package river + +import ( + "github.com/riverqueue/river/internal/dblist" + "github.com/riverqueue/river/rivertype" +) + +// JobDeleteManyParams specifies the parameters for a JobDeleteMany query. It +// must be initialized with NewJobDeleteManyParams. Params can be built by +// chaining methods on the JobDeleteManyParams object: +// +// params := NewJobDeleteManyParams().First(100).States(river.JobStateCompleted) +type JobDeleteManyParams struct { + ids []int64 + kinds []string + limit int32 + priorities []int16 + queues []string + schema string + states []rivertype.JobState + unsafeAll bool +} + +// NewJobDeleteManyParams creates a new JobDeleteManyParams to delete jobs +// sorted by ID in ascending order, deleting 100 jobs at most. +func NewJobDeleteManyParams() *JobDeleteManyParams { + return &JobDeleteManyParams{ + limit: 100, + } +} + +func (p *JobDeleteManyParams) copy() *JobDeleteManyParams { + return &JobDeleteManyParams{ + ids: append([]int64(nil), p.ids...), + kinds: append([]string(nil), p.kinds...), + limit: p.limit, + priorities: append([]int16(nil), p.priorities...), + queues: append([]string(nil), p.queues...), + schema: p.schema, + states: append([]rivertype.JobState(nil), p.states...), + unsafeAll: p.unsafeAll, + } +} + +func (p *JobDeleteManyParams) filtersEmpty() bool { + return len(p.ids) < 1 && + len(p.kinds) < 1 && + len(p.priorities) < 1 && + len(p.queues) < 1 && + len(p.states) < 1 +} + +func (p *JobDeleteManyParams) toDBParams() *dblist.JobListParams { + return &dblist.JobListParams{ + IDs: p.ids, + Kinds: p.kinds, + LimitCount: p.limit, + OrderBy: []dblist.JobListOrderBy{{Expr: "id", Order: dblist.SortOrderAsc}}, + Priorities: p.priorities, + Queues: p.queues, + Schema: p.schema, + States: p.states, + } +} + +// First returns an updated filter set that will only delete the first +// count jobs. +// +// Count must be between 1 and 10_000, inclusive, or this will panic. +func (p *JobDeleteManyParams) First(count int) *JobDeleteManyParams { + if count <= 0 { + panic("count must be > 0") + } + if count > 10000 { + panic("count must be <= 10000") + } + paramsCopy := p.copy() + paramsCopy.limit = int32(count) + return paramsCopy +} + +// IDs returns an updated filter set that will only delete jobs with the given +// IDs. +func (p *JobDeleteManyParams) IDs(ids ...int64) *JobDeleteManyParams { + paramsCopy := p.copy() + paramsCopy.ids = make([]int64, len(ids)) + copy(paramsCopy.ids, ids) + return paramsCopy +} + +// Kinds returns an updated filter set that will only delete jobs of the given +// kinds. +func (p *JobDeleteManyParams) Kinds(kinds ...string) *JobDeleteManyParams { + paramsCopy := p.copy() + paramsCopy.kinds = make([]string, len(kinds)) + copy(paramsCopy.kinds, kinds) + return paramsCopy +} + +// Priorities returns an updated filter set that will only delete jobs with the +// given priorities. +func (p *JobDeleteManyParams) Priorities(priorities ...int16) *JobDeleteManyParams { + paramsCopy := p.copy() + paramsCopy.priorities = make([]int16, len(priorities)) + copy(paramsCopy.priorities, priorities) + return paramsCopy +} + +// Queues returns an updated filter set that will only delete jobs from the +// given queues. +func (p *JobDeleteManyParams) Queues(queues ...string) *JobDeleteManyParams { + paramsCopy := p.copy() + paramsCopy.queues = make([]string, len(queues)) + copy(paramsCopy.queues, queues) + return paramsCopy +} + +// States returns an updated filter set that will only delete jobs in the given +// states. +func (p *JobDeleteManyParams) States(states ...rivertype.JobState) *JobDeleteManyParams { + paramsCopy := p.copy() + paramsCopy.states = make([]rivertype.JobState, len(states)) + copy(paramsCopy.states, states) + return paramsCopy +} + +// UnsafeAll is a special directive that allows unbounded job deletion without +// any filters. Normally, filters like IDs or Kinds is required to scope down +// the deletion so that the caller doesn't accidentally delete all non-running +// jobs. Invoking UnsafeAll removes this safety guard so that all jobs can be +// removed arbitrarily. +// +// Example of use: +// +// deleteRes, err = client.JobDeleteMany(ctx, NewJobDeleteManyParams().UnsafeAll()) +// if err != nil { +// // handle error +// } +// +// It only makes sense to call this function if no filters have yet been applied +// on the parameters object. If some have already, calling it will panic. +func (p *JobDeleteManyParams) UnsafeAll() *JobDeleteManyParams { + if !p.filtersEmpty() { + panic("UnsafeAll no longer meaningful with non-default filters applied") + } + + paramsCopy := p.copy() + paramsCopy.unsafeAll = true + return paramsCopy +} diff --git a/vendor/github.com/riverqueue/river/doc.go b/vendor/github.com/riverqueue/river/doc.go new file mode 100644 index 0000000000..ac64df1f4e --- /dev/null +++ b/vendor/github.com/riverqueue/river/doc.go @@ -0,0 +1,195 @@ +/* +Package river is a robust high-performance job processing system for Go and +Postgres. + +See [homepage], [docs], and [godoc], as well as the [River UI]. + +Being built for Postgres, River encourages the use of the same database for +application data and job queue. By enqueueing jobs transactionally along with +other database changes, whole classes of distributed systems problems are +avoided. Jobs are guaranteed to be enqueued if their transaction commits, are +removed if their transaction rolls back, and aren't visible for work _until_ +commit. See [transactional enqueueing] for more background on this philosophy. + +# Job args and workers + +Jobs are defined in struct pairs, with an implementation of [`JobArgs`] and one +of [`Worker`]. + +Job args contain `json` annotations and define how jobs are serialized to and +from the database, along with a "kind", a stable string that uniquely identifies +the job. + + type SortArgs struct { + // Strings is a slice of strings to sort. + Strings []string `json:"strings"` + } + + func (SortArgs) Kind() string { return "sort" } + +Workers expose a `Work` function that dictates how jobs run. + + type SortWorker struct { + // An embedded WorkerDefaults sets up default methods to fulfill the rest of + // the Worker interface: + river.WorkerDefaults[SortArgs] + } + + func (w *SortWorker) Work(ctx context.Context, job *river.Job[SortArgs]) error { + sort.Strings(job.Args.Strings) + fmt.Printf("Sorted strings: %+v\n", job.Args.Strings) + return nil + } + +# Registering workers + +Jobs are uniquely identified by their "kind" string. Workers are registered on +start up so that River knows how to assign jobs to workers: + + workers := river.NewWorkers() + // AddWorker panics if the worker is already registered or invalid: + river.AddWorker(workers, &SortWorker{}) + +# Starting a client + +A River [`Client`] provides an interface for job insertion and manages job +processing and [maintenance services]. A client's created with a database pool, +[driver], and config struct containing a `Workers` bundle and other settings. +Here's a client `Client` working one queue (`"default"`) with up to 100 worker +goroutines at a time: + + riverClient, err := river.NewClient(riverpgxv5.New(dbPool), &river.Config{ + Queues: map[string]river.QueueConfig{ + river.QueueDefault: {MaxWorkers: 100}, + }, + Workers: workers, + }) + if err != nil { + panic(err) + } + + // Run the client inline. All executed jobs will inherit from ctx: + if err := riverClient.Start(ctx); err != nil { + panic(err) + } + +## Insert-only clients + +It's often desirable to have a client that'll be used for inserting jobs, but +not working them. This is possible by omitting the `Queues` configuration, and +skipping the call to `Start`: + + riverClient, err := river.NewClient(riverpgxv5.New(dbPool), &river.Config{ + Workers: workers, + }) + if err != nil { + panic(err) + } + +`Workers` can also be omitted, but it's better to include it so River can check +that inserted job kinds have a worker that can run them. + +## Stopping + +The client should also be stopped on program shutdown: + + // Stop fetching new work and wait for active jobs to finish. + if err := riverClient.Stop(ctx); err != nil { + panic(err) + } + +There are some complexities around ensuring clients stop cleanly, but also in a +timely manner. See [graceful shutdown] for more details on River's stop modes. + +# Inserting jobs + +[`Client.InsertTx`] is used in conjunction with an instance of job args to +insert a job to work on a transaction: + + _, err = riverClient.InsertTx(ctx, tx, SortArgs{ + Strings: []string{ + "whale", "tiger", "bear", + }, + }, nil) + + if err != nil { + panic(err) + } + +See the [`InsertAndWork` example] for complete code. + +# Other features + + - [Batch job insertion] for efficiently inserting many jobs at once using + Postgres `COPY FROM`. + + - [Cancelling jobs] from inside a work function. + + - [Error and panic handling]. + + - [Multiple queues] to better guarantee job throughput, worker availability, + and isolation between components. + + - [Periodic and cron jobs]. + + - [Scheduled jobs] that run automatically at their scheduled time in the + future. + + - [Snoozing jobs] from inside a work function. + + - [Subscriptions] to queue activity and statistics, providing easy hooks for + telemetry like logging and metrics. + + - [Test helpers] to verify that jobs are inserted as expected. + + - [Transactional job completion] to guarantee job completion commits with + other changes in a transaction. + + - [Unique jobs] by args, period, queue, and state. + + - [Web UI] for inspecting and interacting with jobs and queues. + + - [Work functions] for simplified worker implementation. + +## Cross language enqueueing + +River supports inserting jobs in some non-Go languages which are then worked by Go implementations. This may be desirable in performance sensitive cases so that jobs can take advantage of Go's fast runtime. + + - [Inserting jobs from Python]. + - [Inserting jobs from Ruby]. + +# Development + +See [developing River]. + +[`Client`]: https://pkg.go.dev/github.com/riverqueue/river#Client +[`Client.InsertTx`]: https://pkg.go.dev/github.com/riverqueue/river#Client.InsertTx +[`InsertAndWork` example]: https://pkg.go.dev/github.com/riverqueue/river#example-package-InsertAndWork +[`JobArgs`]: https://pkg.go.dev/github.com/riverqueue/river#JobArgs +[`Worker`]: https://pkg.go.dev/github.com/riverqueue/river#Worker +[Batch job insertion]: https://riverqueue.com/docs/batch-job-insertion +[Cancelling jobs]: https://riverqueue.com/docs/cancelling-jobs +[Error and panic handling]: https://riverqueue.com/docs/error-handling +[Inserting jobs from Python]: https://riverqueue.com/docs/python +[Inserting jobs from Ruby]: https://riverqueue.com/docs/ruby +[Multiple queues]: https://riverqueue.com/docs/multiple-queues +[Periodic and cron jobs]: https://riverqueue.com/docs/periodic-jobs +[River UI]: https://github.com/riverqueue/riverui +[Scheduled jobs]: https://riverqueue.com/docs/scheduled-jobs +[Snoozing jobs]: https://riverqueue.com/docs/snoozing-jobs +[Subscriptions]: https://riverqueue.com/docs/subscriptions +[Test helpers]: https://riverqueue.com/docs/testing +[Transactional job completion]: https://riverqueue.com/docs/transactional-job-completion +[Unique jobs]: https://riverqueue.com/docs/unique-jobs +[Web UI]: https://github.com/riverqueue/riverui +[Work functions]: https://riverqueue.com/docs/work-functions +[docs]: https://riverqueue.com/docs +[developing River]: https://github.com/riverqueue/river/blob/master/docs/development.md +[driver]: https://riverqueue.com/docs/database-drivers +[godoc]: https://pkg.go.dev/github.com/riverqueue/river +[graceful shutdown]: https://riverqueue.com/docs/graceful-shutdown +[homepage]: https://riverqueue.com +[maintenance services]: https://riverqueue.com/docs/maintenance-services +[transactional enqueueing]: https://riverqueue.com/docs/transactional-enqueueing +*/ +package river diff --git a/vendor/github.com/riverqueue/river/error.go b/vendor/github.com/riverqueue/river/error.go new file mode 100644 index 0000000000..7e634d65ab --- /dev/null +++ b/vendor/github.com/riverqueue/river/error.go @@ -0,0 +1,79 @@ +package river + +import ( + "fmt" + "time" + + "github.com/riverqueue/river/rivertype" +) + +// ErrJobCancelledRemotely is a sentinel error indicating that the job was cancelled remotely. +var ErrJobCancelledRemotely = rivertype.ErrJobCancelledRemotely + +// JobCancelError is the error type returned by JobCancel. It should not be +// initialized directly, but is returned from the [JobCancel] function and can +// be used for test assertions. +type JobCancelError = rivertype.JobCancelError + +// JobCancel wraps err and can be returned from a Worker's Work method to cancel +// the job at the end of execution. Regardless of whether or not the job has any +// remaining attempts, this will ensure the job does not execute again. +func JobCancel(err error) error { + return rivertype.JobCancel(err) +} + +// JobSnoozeError is the error type returned by JobSnooze. It should not be +// initialized directly, but is returned from the [JobSnooze] function and can +// be used for test assertions. +type JobSnoozeError = rivertype.JobSnoozeError + +// JobSnooze can be returned from a Worker's Work method to cause the job to be +// tried again after the specified duration. This will not increment the job's +// Attempt count, meaning that jobs can be repeatedly snoozed without ever being +// discarded. +// +// A special duration of zero can be used to make the job immediately available +// to be reworked. This may be useful in cases like where a long-running job is +// being interrupted on shutdown. Instead of returning a context cancelled error +// that'd schedule a retry for the future and count towards maximum attempts, +// the work function can return JobSnooze(0) and the job will be retried +// immediately the next time a client starts up. +// +// Panics if duration is < 0. +func JobSnooze(duration time.Duration) error { + return &rivertype.JobSnoozeError{Duration: duration} +} + +// QueueAlreadyAddedError is returned when attempting to add a queue that has +// already been added to the Client. +type QueueAlreadyAddedError struct { + Name string +} + +func (e *QueueAlreadyAddedError) Error() string { + return fmt.Sprintf("queue %q already added", e.Name) +} + +func (e *QueueAlreadyAddedError) Is(target error) bool { + _, ok := target.(*QueueAlreadyAddedError) + return ok +} + +// QueueNotFoundError is returned when attempting to remove a queue that does +// not exist on the Client. +type QueueNotFoundError struct { + Name string +} + +func (e *QueueNotFoundError) Error() string { + return fmt.Sprintf("queue %q not found", e.Name) +} + +func (e *QueueNotFoundError) Is(target error) bool { + _, ok := target.(*QueueNotFoundError) + return ok +} + +// UnknownJobKindError is returned when a Client fetches and attempts to +// work a job that has not been registered on the Client's Workers bundle (using AddWorker). +type UnknownJobKindError = rivertype.UnknownJobKindError diff --git a/vendor/github.com/riverqueue/river/error_handler.go b/vendor/github.com/riverqueue/river/error_handler.go new file mode 100644 index 0000000000..694bdae574 --- /dev/null +++ b/vendor/github.com/riverqueue/river/error_handler.go @@ -0,0 +1,35 @@ +package river + +import ( + "context" + + "github.com/riverqueue/river/rivertype" +) + +// ErrorHandler provides an interface that will be invoked in case of an error +// or panic occurring in the job. This is often useful for logging and exception +// tracking, but can also be used to customize retry behavior. +type ErrorHandler interface { + // HandleError is invoked in case of an error occurring in a job. + // + // Context is descended from the one used to start the River client that + // worked the job. Errors are handled above all middleware, so changes made + // to context by a middleware are not available in the context. + HandleError(ctx context.Context, job *rivertype.JobRow, err error) *ErrorHandlerResult + + // HandlePanic is invoked in case of a panic occurring in a job. + // + // Context is descended from the one used to start the River client that + // worked the job. Panics are handled above all middleware, so changes made + // to context by a middleware are not available in the context (however, + // panics can be recovered from in any middleware where middleware context + // is available). + HandlePanic(ctx context.Context, job *rivertype.JobRow, panicVal any, trace string) *ErrorHandlerResult +} + +type ErrorHandlerResult struct { + // SetCancelled can be set to true to fail the job immediately and + // permanently. By default it'll continue to follow the configured retry + // schedule. + SetCancelled bool +} diff --git a/vendor/github.com/riverqueue/river/event.go b/vendor/github.com/riverqueue/river/event.go new file mode 100644 index 0000000000..b1694b8891 --- /dev/null +++ b/vendor/github.com/riverqueue/river/event.go @@ -0,0 +1,91 @@ +package river + +import ( + "time" + + "github.com/riverqueue/river/internal/jobstats" + "github.com/riverqueue/river/rivertype" +) + +// EventKind is a kind of event to subscribe to from a client. +type EventKind string + +const ( + // EventKindJobCancelled occurs when a job is cancelled. + EventKindJobCancelled EventKind = "job_cancelled" + + // EventKindJobCompleted occurs when a job is completed. + EventKindJobCompleted EventKind = "job_completed" + + // EventKindJobFailed occurs when a job fails. Occurs both when a job fails + // and will be retried and when a job fails for the last time and will be + // discarded. Callers can use job fields like `Attempt` and `State` to + // differentiate each type of occurrence. + EventKindJobFailed EventKind = "job_failed" + + // EventKindJobSnoozed occurs when a job is snoozed. + EventKindJobSnoozed EventKind = "job_snoozed" + + // EventKindQueuePaused occurs when a queue is paused. + EventKindQueuePaused EventKind = "queue_paused" + + // EventKindQueueResumed occurs when a queue is resumed. + EventKindQueueResumed EventKind = "queue_resumed" +) + +// All known event kinds, used to validate incoming kinds. This is purposely not +// exported because end users should have no way of subscribing to all known +// kinds for forward compatibility reasons. +var allKinds = map[EventKind]struct{}{ //nolint:gochecknoglobals + EventKindJobCancelled: {}, + EventKindJobCompleted: {}, + EventKindJobFailed: {}, + EventKindJobSnoozed: {}, + EventKindQueuePaused: {}, + EventKindQueueResumed: {}, +} + +// Event wraps an event that occurred within a River client, like a job being +// completed. +type Event struct { + // Kind is the kind of event. Receivers should read this field and respond + // accordingly. Subscriptions will only receive event kinds that they + // requested when creating a subscription with Subscribe. + Kind EventKind + + // Job contains job-related information. + Job *rivertype.JobRow + + // JobStats are statistics about the run of a job. + JobStats *JobStatistics + + // Queue contains queue-related information. + Queue *rivertype.Queue +} + +// JobStatistics contains information about a single execution of a job. +type JobStatistics struct { + CompleteDuration time.Duration // Time it took to set the job completed, discarded, or errored. + QueueWaitDuration time.Duration // Time the job spent waiting in available state before starting execution. + RunDuration time.Duration // Time job spent running (measured around job worker.) +} + +func jobStatisticsFromInternal(stats *jobstats.JobStatistics) *JobStatistics { + return &JobStatistics{ + CompleteDuration: stats.CompleteDuration, + QueueWaitDuration: stats.QueueWaitDuration, + RunDuration: stats.RunDuration, + } +} + +// eventSubscription is an active subscription for events being produced by a +// client, created with Client.Subscribe. +type eventSubscription struct { + Chan chan *Event + Kinds map[EventKind]struct{} +} + +func (s *eventSubscription) ListensFor(kind EventKind) bool { + _, ok := s.Kinds[kind] + return ok +} diff --git a/vendor/github.com/riverqueue/river/go.work b/vendor/github.com/riverqueue/river/go.work new file mode 100644 index 0000000000..9a9927febb --- /dev/null +++ b/vendor/github.com/riverqueue/river/go.work @@ -0,0 +1,15 @@ +go 1.25.0 + +toolchain go1.25.7 + +use ( + . + ./cmd/river + ./riverdriver + ./riverdriver/riverdatabasesql + ./riverdriver/riverdrivertest + ./riverdriver/riverpgxv5 + ./riverdriver/riversqlite + ./rivershared + ./rivertype +) diff --git a/vendor/github.com/riverqueue/river/hook_defaults_funcs.go b/vendor/github.com/riverqueue/river/hook_defaults_funcs.go new file mode 100644 index 0000000000..0cec92209c --- /dev/null +++ b/vendor/github.com/riverqueue/river/hook_defaults_funcs.go @@ -0,0 +1,81 @@ +package river + +import ( + "context" + + "github.com/riverqueue/river/rivertype" +) + +// HookDefaults should be embedded on any hooks implementation. It helps +// identify a struct as hooks and a plugin, and guarantee forward compatibility +// in case additions are necessary to the rivertype.Hook interface. +type HookDefaults struct{} + +func (d *HookDefaults) IsHook() bool { return true } + +func (d *HookDefaults) IsPlugin() bool { return true } + +// HookInsertBeginFunc is a convenience helper for implementing +// rivertype.HookInsertBegin using a simple function instead of a struct. +type HookInsertBeginFunc func(ctx context.Context, params *rivertype.JobInsertParams) error + +func (f HookInsertBeginFunc) InsertBegin(ctx context.Context, params *rivertype.JobInsertParams) error { + return f(ctx, params) +} + +func (f HookInsertBeginFunc) IsHook() bool { return true } + +func (f HookInsertBeginFunc) IsPlugin() bool { return true } + +// HookMetricEmitFunc is a convenience helper for implementing +// rivertype.HookMetricEmit using a simple function instead of a struct. +// +// Notably, this function is invoked each time River emits a metric. Metrics are +// emitted in very hot paths like job fetching, and should therefore not block +// on network I/O or anything else, and should usually pass metrics through to +// an asynchronous instrumentation package like OpenTelemetry. +type HookMetricEmitFunc func(ctx context.Context, params *rivertype.HookMetricEmitParams) + +func (f HookMetricEmitFunc) IsHook() bool { return true } + +func (f HookMetricEmitFunc) IsPlugin() bool { return true } + +func (f HookMetricEmitFunc) MetricEmit(ctx context.Context, params *rivertype.HookMetricEmitParams) { + f(ctx, params) +} + +// HookPeriodicJobsStartFunc is a convenience helper for implementing +// rivertype.HookPeriodicJobsStart using a simple function instead of a struct. +type HookPeriodicJobsStartFunc func(ctx context.Context, params *rivertype.HookPeriodicJobsStartParams) error + +func (f HookPeriodicJobsStartFunc) IsHook() bool { return true } + +func (f HookPeriodicJobsStartFunc) IsPlugin() bool { return true } + +func (f HookPeriodicJobsStartFunc) Start(ctx context.Context, params *rivertype.HookPeriodicJobsStartParams) error { + return f(ctx, params) +} + +// HookWorkBeginFunc is a convenience helper for implementing +// rivertype.HookWorkBegin using a simple function instead of a struct. +type HookWorkBeginFunc func(ctx context.Context, job *rivertype.JobRow) error + +func (f HookWorkBeginFunc) IsHook() bool { return true } + +func (f HookWorkBeginFunc) IsPlugin() bool { return true } + +func (f HookWorkBeginFunc) WorkBegin(ctx context.Context, job *rivertype.JobRow) error { + return f(ctx, job) +} + +// HookWorkEndFunc is a convenience helper for implementing +// rivertype.HookWorkEnd using a simple function instead of a struct. +type HookWorkEndFunc func(ctx context.Context, job *rivertype.JobRow, err error) error + +func (f HookWorkEndFunc) IsHook() bool { return true } + +func (f HookWorkEndFunc) IsPlugin() bool { return true } + +func (f HookWorkEndFunc) WorkEnd(ctx context.Context, job *rivertype.JobRow, err error) error { + return f(ctx, job, err) +} diff --git a/vendor/github.com/riverqueue/river/insert_opts.go b/vendor/github.com/riverqueue/river/insert_opts.go new file mode 100644 index 0000000000..64fb770d50 --- /dev/null +++ b/vendor/github.com/riverqueue/river/insert_opts.go @@ -0,0 +1,287 @@ +package river + +import ( + "errors" + "fmt" + "regexp" + "slices" + "strings" + "time" + + "github.com/riverqueue/river/rivertype" +) + +// Regular expression to which the format of tags must comply. Mainly, no +// special characters, and with hyphens in the middle. +// +// A key property here (in case this is relaxed in the future) is that commas +// must never be allowed because they're used as a delimiter during batch job +// insertion for the `riverdatabasesql` driver. +var tagRE = regexp.MustCompile(`\A[\w][\w\-]+[\w]\z`) + +// InsertOpts are optional settings for a new job which can be provided at job +// insertion time. These will override any default InsertOpts settings provided +// by JobArgsWithInsertOpts, as well as any global defaults. +type InsertOpts struct { + // MaxAttempts is the maximum number of total attempts (including both the + // original run and all retries) before a job is abandoned and set as + // discarded. + MaxAttempts int + + // Metadata is a JSON object blob of arbitrary data that will be stored with + // the job. Users should not overwrite or remove anything stored in this + // field by River. + Metadata []byte + + // Pending indicates that the job should be inserted in the `pending` state. + // Pending jobs are not immediately available to be worked and are never + // deleted, but they can be used to indicate work which should be performed in + // the future once they are made available (or scheduled) by some external + // update. + Pending bool + + // Priority is the priority of the job, with 1 being the highest priority and + // 4 being the lowest. When fetching available jobs to work, the highest + // priority jobs will always be fetched before any lower priority jobs are + // fetched. Note that if your workers are swamped with more high-priority jobs + // then they can handle, lower priority jobs may not be fetched. + // + // Defaults to PriorityDefault. + Priority int + + // Queue is the name of the job queue in which to insert the job. + // + // Defaults to the job kind's default queue if set via + // `JobArgsWithInsertOpts`, or QueueDefault if not. + Queue string + + // ScheduledAt is a time in future at which to schedule the job (i.e. in + // cases where it shouldn't be run immediately). The job is guaranteed not + // to run before this time, but may run slightly after depending on the + // number of other scheduled jobs and how busy the queue is. + // + // Use of this option generally only makes sense when passing options into + // Insert rather than when a job args struct is implementing + // JobArgsWithInsertOpts, however, it will work in both cases. + ScheduledAt time.Time + + // Tags are an arbitrary list of keywords to add to the job. They don't + // affect job execution, but can be used with JobListParams.TagsAll and + // JobListParams.TagsAny to group and filter jobs. + // + // Tags should conform to the regex `\A[\w][\w\-]+[\w]\z` and be a maximum + // of 255 characters long. No special characters are allowed. + // + // If tags are specified from both a job args override and from options on + // Insert, the latter takes precedence. Tags are not merged. + Tags []string + + // UniqueOpts returns options relating to job uniqueness. An empty struct + // avoids setting any worker-level unique options. + UniqueOpts UniqueOpts +} + +// UniqueOpts contains parameters for uniqueness for a job. +// +// When the options struct is uninitialized (its zero value) no uniqueness at is +// enforced. As each property is initialized, it's added as a dimension on the +// uniqueness matrix. When any property has a non-zero value specified, the +// job's kind automatically counts toward uniqueness, but can be excluded by +// setting ExcludeKind to true. +// +// So for example, if only ByQueue is on, then for the given job kind, only a +// single instance is allowed in any given queue, regardless of other properties +// on the job. If both ByArgs and ByQueue are on, then for the given job kind, a +// single instance is allowed for each combination of args and queues. If either +// args or queue is changed on a new job, it's allowed to be inserted as a new +// job. +// +// Uniqueness relies on a hash of the job kind and any unique properties along +// with a database unique constraint. See the note on ByState for more details +// including about the fallback to a deprecated advisory lock method. +type UniqueOpts struct { + // ByArgs indicates that uniqueness should be enforced for any specific + // instance of encoded args for a job. + // + // Default is false, meaning that as long as any other unique property is + // enabled, uniqueness will be enforced for a kind regardless of input args. + // + // When set to true, the entire encoded args field will be included in the + // uniqueness hash, which requires care to ensure that no irrelevant args are + // factored into the uniqueness check. It is also possible to use a subset of + // the args by indicating on the `JobArgs` struct which fields should be + // included in the uniqueness check using struct tags: + // + // type MyJobArgs struct { + // CustomerID string `json:"customer_id" river:"unique"` + // TraceID string `json:"trace_id" + // } + // + // In this example, only the encoded `customer_id` key will be included in the + // uniqueness check and the `trace_id` key will be ignored. + // + // All keys are sorted alphabetically before hashing to ensure consistent + // results. + // + // River recurses into embedded structs and fields with struct values and + // looks for `river:"unique"` annotations on them as well: + // + // type MyJobArgs struct { + // Customer *Customer `json:"customer"` + // TraceID string `json:"trace_id" + // } + // + // type Customer struct { + // ID string `json:"id" river:"unique"` + // } + // + // In this example, the `id` value inside a `customer` subboject is used in + // the uniqueness check. It'd be the same story if Customer was embedded on + // MyJobArgs instead: + // + // type MyJobArgs struct { + // Customer + // TraceID string `json:"trace_id" + // } + // + // If the struct field itself has a `river:"unique"` annotation, but none on + // any fields in the substruct, then the entire JSON encoded value of the + // struct is used as a unique value: + // + // type MyJobArgs struct { + // Customer *Customer `json:"customer" river:"unique"` + // TraceID string `json:"trace_id" + // } + // + // type Customer struct { + // ID string `json:"id"` + // } + ByArgs bool + + // ByPeriod defines uniqueness within a given period. On an insert time is + // rounded down to the nearest multiple of the given period, and a job is + // only inserted if there isn't an existing job that will run between then + // and the next multiple of the period. + // + // Default is no unique period, meaning that as long as any other unique + // property is enabled, uniqueness will be enforced across all jobs of the + // kind in the database, regardless of when they were scheduled. + ByPeriod time.Duration + + // ByQueue indicates that uniqueness should be enforced within each queue. + // + // Default is false, meaning that as long as any other unique property is + // enabled, uniqueness will be enforced for a kind across all queues. + ByQueue bool + + // ByState indicates that uniqueness should be enforced across any of the + // states in the given set. Unlike other unique options, ByState gets a + // default when it's not set for user convenience. The default is equivalent + // to: + // + // ByState: []rivertype.JobState{rivertype.JobStateAvailable, rivertype.JobStateCompleted, rivertype.JobStatePending, rivertype.JobStateRunning, rivertype.JobStateRetryable, rivertype.JobStateScheduled} + // + // Or more succinctly: + // + // ByState: rivertype.UniqueOptsByStateDefault() + // + // With this setting, any jobs of the same kind that have been completed or + // discarded, but not yet cleaned out by the system, will still prevent a + // duplicate unique job from being inserted. For example, with the default + // states, if a unique job is actively `running`, a duplicate cannot be + // inserted. Likewise, if a unique job has `completed`, you still can't + // insert a duplicate, at least not until the job cleaner maintenance process + // eventually removes the completed job from the `river_job` table. + // + // The list may be safely customized to _add_ additional states (`cancelled` + // or `discarded`), though only `retryable` may be safely _removed_ from the + // list. + // + // The following states must be included in ByState if set: + // + // - rivertype.JobStateAvailable + // - rivertype.JobStatePending + // - rivertype.JobStateRunning + // - rivertype.JobStateScheduled + // + // These states being required is an implementation detail, but not + // requiring them would put River in a difficult position when moving jobs + // between common parts of the state machine and finding a unique conflict + // already there. Resolving this isn't completely intractable, but'll + // require some in-depth thinking and designing around every possible edge. + ByState []rivertype.JobState + + // ExcludeKind indicates that the job kind should not be included in the + // uniqueness check. This is useful when you want to enforce uniqueness + // across all jobs regardless of kind. + ExcludeKind bool +} + +// isEmpty returns true for an empty, uninitialized options struct. +// +// This is required because we can't check against `UniqueOpts{}` because slices +// aren't comparable. Unfortunately it makes things a little more brittle +// comparatively because any new options must also be considered here for things +// to work. +func (o *UniqueOpts) isEmpty() bool { + return !o.ByArgs && + o.ByPeriod == time.Duration(0) && + !o.ByQueue && + o.ByState == nil +} + +var jobStateAll = rivertype.JobStates() //nolint:gochecknoglobals + +// Required unique states. Requiring states like this isn't necessary +// fundamental for correctness, but doing so avoids some gnarly problems that +// don't have a clear error action otherwise. For example, if `available` was +// omittable and a producer tried to transition an `available` job to `running` +// but found another unique contender already there, we'd have to figure out +// what to do about the job that can't be scheduled. We can't send feedback to +// the caller at this point, so probably the best we could do is leave it in +// this untransitionable state until the `running` job finished, which isn't +// particularly satisfactory. +var requiredV3states = []rivertype.JobState{ //nolint:gochecknoglobals + rivertype.JobStateAvailable, + rivertype.JobStatePending, + rivertype.JobStateRunning, + rivertype.JobStateScheduled, +} + +func (o *UniqueOpts) validate() error { + if o.isEmpty() { + return nil + } + + if o.ByPeriod != time.Duration(0) && o.ByPeriod < 1*time.Second { + return errors.New("UniqueOpts.ByPeriod should not be less than 1 second") + } + + // Job states are typed, but since the underlying type is a string, users + // can put anything they want in there. + for _, state := range o.ByState { + // This could be turned to a map lookup, but last I checked the speed + // difference for tiny slice sizes is negligible, and map lookup might + // even be slower. + if !slices.Contains(jobStateAll, state) { + return fmt.Errorf("UniqueOpts.ByState contains invalid state %q", state) + } + } + + // Skip required states validation if no custom states were provided. + if len(o.ByState) == 0 { + return nil + } + + var missingStates []string + for _, state := range requiredV3states { + if !slices.Contains(o.ByState, state) { + missingStates = append(missingStates, string(state)) + } + } + if len(missingStates) > 0 { + return fmt.Errorf("UniqueOpts.ByState must contain all required states, missing: %s", strings.Join(missingStates, ", ")) + } + + return nil +} diff --git a/vendor/github.com/riverqueue/river/internal/dblist/db_list.go b/vendor/github.com/riverqueue/river/internal/dblist/db_list.go new file mode 100644 index 0000000000..dba6780dc2 --- /dev/null +++ b/vendor/github.com/riverqueue/river/internal/dblist/db_list.go @@ -0,0 +1,229 @@ +package dblist + +import ( + "context" + "errors" + "fmt" + "strings" + + "github.com/riverqueue/river/riverdriver" + "github.com/riverqueue/river/rivershared/util/sliceutil" + "github.com/riverqueue/river/rivertype" +) + +type SortOrder int + +const ( + SortOrderUnspecified SortOrder = iota + SortOrderAsc + SortOrderDesc +) + +type JobListOrderBy struct { + Expr string + Order SortOrder +} + +type JobListParams struct { + IDs []int64 + Kinds []string + LimitCount int32 + OrderBy []JobListOrderBy + Priorities []int16 + Queues []string + Schema string + States []rivertype.JobState + TagsAll []string + TagsAny []string + Where []WherePredicate +} + +type WherePredicate struct { + NamedArgs map[string]any + SQL string +} + +type sqlFragmentBuilder interface { + SQLFragmentColumnContainsAll(column, namedArg string, values []string) (string, any, error) + SQLFragmentColumnContainsAny(column, namedArg string, values []string) (string, any, error) + SQLFragmentColumnIn(column string, values any) (string, any, error) +} + +// JobMakeDriverParams converts client-level parameters for job and delete to +// driver-level parameters for use with an executor, which generally goes by +// converting typed fields for IDs, kinds, queues, etc. to lower-level SQL. +// +// This was originally implemented for listing jobs, but since the logic is so +// similar, it also performs the same function for JobDeleteMany. This works +// because `riverdriver.JobDeleteManyParams` has `JobListParams` as its +// underlying type and therefore pointer-level converts to it. +func JobMakeDriverParams(ctx context.Context, params *JobListParams, sqlFragmentBuilder sqlFragmentBuilder) (*riverdriver.JobListParams, error) { + var ( + namedArgs = make(map[string]any) + whereBuilder strings.Builder + ) + + orderBy := make([]JobListOrderBy, len(params.OrderBy)) + for i, o := range params.OrderBy { + orderBy[i] = JobListOrderBy{ + Expr: o.Expr, + Order: o.Order, + } + } + + // Writes an `AND` to connect SQL predicates as long as this isn't the first + // predicate. + writeAndAfterFirst := func() { + if whereBuilder.Len() != 0 { + whereBuilder.WriteString("\n AND ") + } + } + + if len(params.IDs) > 0 { + writeAndAfterFirst() + + const column = "id" + sqlFragment, arg, err := sqlFragmentBuilder.SQLFragmentColumnIn(column, params.IDs) + if err != nil { + return nil, fmt.Errorf("error building SQL fragment for %q: %w", column, err) + } + whereBuilder.WriteString(sqlFragment) + namedArgs[column] = arg + } + + if len(params.Kinds) > 0 { + writeAndAfterFirst() + + const column = "kind" + sqlFragment, arg, err := sqlFragmentBuilder.SQLFragmentColumnIn(column, params.Kinds) + if err != nil { + return nil, fmt.Errorf("error building SQL fragment for %q: %w", column, err) + } + whereBuilder.WriteString(sqlFragment) + namedArgs[column] = arg + } + + if len(params.Priorities) > 0 { + writeAndAfterFirst() + + const column = "priority" + sqlFragment, arg, err := sqlFragmentBuilder.SQLFragmentColumnIn(column, params.Priorities) + if err != nil { + return nil, fmt.Errorf("error building SQL fragment for %q: %w", column, err) + } + whereBuilder.WriteString(sqlFragment) + namedArgs[column] = arg + } + + if len(params.Queues) > 0 { + writeAndAfterFirst() + + const column = "queue" + sqlFragment, arg, err := sqlFragmentBuilder.SQLFragmentColumnIn(column, params.Queues) + if err != nil { + return nil, fmt.Errorf("error building SQL fragment for %q: %w", column, err) + } + whereBuilder.WriteString(sqlFragment) + namedArgs[column] = arg + } + + if len(params.States) > 0 { + writeAndAfterFirst() + + const column = "state" + sqlFragment, arg, err := sqlFragmentBuilder.SQLFragmentColumnIn(column, + sliceutil.Map(params.States, func(v rivertype.JobState) string { return string(v) })) + if err != nil { + return nil, fmt.Errorf("error building SQL fragment for %q: %w", column, err) + } + whereBuilder.WriteString(sqlFragment) + namedArgs[column] = arg + } + + if len(params.TagsAll) > 0 { + writeAndAfterFirst() + + const ( + column = "tags" + namedArg = "tags_all" + ) + sqlFragment, arg, err := sqlFragmentBuilder.SQLFragmentColumnContainsAll(column, namedArg, params.TagsAll) + if err != nil { + return nil, fmt.Errorf("error building SQL fragment for %q: %w", namedArg, err) + } + whereBuilder.WriteString(sqlFragment) + namedArgs[namedArg] = arg + } + + if len(params.TagsAny) > 0 { + writeAndAfterFirst() + + const ( + column = "tags" + namedArg = "tags_any" + ) + sqlFragment, arg, err := sqlFragmentBuilder.SQLFragmentColumnContainsAny(column, namedArg, params.TagsAny) + if err != nil { + return nil, fmt.Errorf("error building SQL fragment for %q: %w", namedArg, err) + } + whereBuilder.WriteString(sqlFragment) + namedArgs[namedArg] = arg + } + + for _, where := range params.Where { + writeAndAfterFirst() + + whereBuilder.WriteString(where.SQL) + for name, val := range where.NamedArgs { + expectedSymbol := "@" + name + if !strings.Contains(where.SQL, expectedSymbol) { + return nil, fmt.Errorf("expected %q to contain named arg symbol %s", where.SQL, expectedSymbol) + } + + if _, ok := namedArgs[name]; ok { + return nil, fmt.Errorf("named argument %s already registered", expectedSymbol) + } + + namedArgs[name] = val + } + } + + // A condition of some kind is needed, so given no others write one that'll + // always return true. + if whereBuilder.Len() < 1 { + whereBuilder.WriteString("true") + } + + if params.LimitCount < 1 { + return nil, errors.New("required parameter 'Count' in JobList must be greater than zero") + } + + if len(params.OrderBy) == 0 { + return nil, errors.New("sort order is required") + } + + var orderByBuilder strings.Builder + + for i, orderBy := range params.OrderBy { + orderByBuilder.WriteString(orderBy.Expr) + switch orderBy.Order { + case SortOrderAsc: + orderByBuilder.WriteString(" ASC") + case SortOrderDesc: + orderByBuilder.WriteString(" DESC") + case SortOrderUnspecified: + return nil, errors.New("should not have gotten SortOrderUnspecified by this point before executing list (bug?)") + } + if i < len(params.OrderBy)-1 { + orderByBuilder.WriteString(", ") + } + } + + return &riverdriver.JobListParams{ + Max: params.LimitCount, + NamedArgs: namedArgs, + OrderByClause: orderByBuilder.String(), + Schema: params.Schema, + WhereClause: whereBuilder.String(), + }, nil +} diff --git a/vendor/github.com/riverqueue/river/internal/dbunique/db_unique.go b/vendor/github.com/riverqueue/river/internal/dbunique/db_unique.go new file mode 100644 index 0000000000..bbf4866165 --- /dev/null +++ b/vendor/github.com/riverqueue/river/internal/dbunique/db_unique.go @@ -0,0 +1,126 @@ +package dbunique + +import ( + "crypto/sha256" + "slices" + "strings" + "time" + + "github.com/tidwall/gjson" + "github.com/tidwall/sjson" + + "github.com/riverqueue/river/rivershared/structtag" + "github.com/riverqueue/river/rivershared/uniquestates" + "github.com/riverqueue/river/rivershared/util/ptrutil" + "github.com/riverqueue/river/rivershared/util/sliceutil" + "github.com/riverqueue/river/rivertype" +) + +// Default job states for UniqueOpts.ByState. Stored here to a variable so we +// don't have to reallocate a slice over and over again. +var uniqueOptsByStateDefault = rivertype.UniqueOptsByStateDefault() //nolint:gochecknoglobals + +type UniqueOpts struct { + ByArgs bool + ByPeriod time.Duration + ByQueue bool + ByState []rivertype.JobState + ExcludeKind bool +} + +func (o *UniqueOpts) IsEmpty() bool { + return !o.ByArgs && + o.ByPeriod == time.Duration(0) && + !o.ByQueue && + o.ByState == nil && + !o.ExcludeKind +} + +func (o *UniqueOpts) StateBitmask() byte { + states := uniqueOptsByStateDefault + if len(o.ByState) > 0 { + states = o.ByState + } + return uniquestates.UniqueStatesToBitmask(states) +} + +func UniqueKey(timeGen rivertype.TimeGenerator, uniqueOpts *UniqueOpts, params *rivertype.JobInsertParams) ([]byte, error) { + uniqueKeyString, err := buildUniqueKeyString(timeGen, uniqueOpts, params) + if err != nil { + return nil, err + } + uniqueKeyHash := sha256.Sum256([]byte(uniqueKeyString)) + return uniqueKeyHash[:], nil +} + +// Builds a unique key made up of the unique options in place. The key is hashed +// to become a value for `unique_key`. +func buildUniqueKeyString(timeGen rivertype.TimeGenerator, uniqueOpts *UniqueOpts, params *rivertype.JobInsertParams) (string, error) { + var sb strings.Builder + + if !uniqueOpts.ExcludeKind { + sb.WriteString("&kind=" + params.Kind) + } + + if uniqueOpts.ByArgs { + var encodedArgsForUnique []byte + // Get unique JSON keys from the JobArgs struct: + uniqueFields, err := structtag.SortedFieldsWithTag(params.Args, "unique") + if err != nil { + return "", err + } + + if len(uniqueFields) > 0 { + // Extract unique values from the EncodedArgs JSON + uniqueValues := structtag.ExtractValues(params.EncodedArgs, uniqueFields) + + // Assemble the JSON object using bytes.Buffer + // Better to overallocate a bit than to allocate multiple times, so just + // assume we'll cap out at the length of the full encoded args. + sortedJSONWithOnlyUniqueValues := make([]byte, 0, len(params.EncodedArgs)) + + sjsonOpts := &sjson.Options{ReplaceInPlace: true} + for i, key := range uniqueFields { + if uniqueValues[i] == "undefined" { + continue + } + sortedJSONWithOnlyUniqueValues, err = sjson.SetRawBytesOptions(sortedJSONWithOnlyUniqueValues, key, []byte(uniqueValues[i]), sjsonOpts) + if err != nil { + // Should not happen unless key was invalid + return "", err + } + } + encodedArgsForUnique = sortedJSONWithOnlyUniqueValues + } else { + // Use all keys from EncodedArgs sorted alphabetically + keys := sliceutil.Map(gjson.GetBytes(params.EncodedArgs, "@keys").Array(), func(v gjson.Result) string { return v.String() }) + slices.Sort(keys) + + sortedJSON := make([]byte, 0, len(params.EncodedArgs)) + sortedJSON = append(sortedJSON, "{}"...) + sjsonOpts := &sjson.Options{ReplaceInPlace: true} + for _, key := range keys { + sortedJSON, err = sjson.SetRawBytesOptions(sortedJSON, key, []byte(gjson.GetBytes(params.EncodedArgs, key).Raw), sjsonOpts) + if err != nil { + // Should not happen unless key was invalid + return "", err + } + } + encodedArgsForUnique = sortedJSON + } + + sb.WriteString("&args=") + sb.Write(encodedArgsForUnique) + } + + if uniqueOpts.ByPeriod != time.Duration(0) { + lowerPeriodBound := ptrutil.ValOrDefaultFunc(params.ScheduledAt, timeGen.Now).Truncate(uniqueOpts.ByPeriod) + sb.WriteString("&period=" + lowerPeriodBound.Format(time.RFC3339)) + } + + if uniqueOpts.ByQueue { + sb.WriteString("&queue=" + params.Queue) + } + + return sb.String(), nil +} diff --git a/vendor/github.com/riverqueue/river/internal/execution/execution.go b/vendor/github.com/riverqueue/river/internal/execution/execution.go new file mode 100644 index 0000000000..9345d6b019 --- /dev/null +++ b/vendor/github.com/riverqueue/river/internal/execution/execution.go @@ -0,0 +1,43 @@ +package execution + +import ( + "context" + "slices" + + "github.com/riverqueue/river/rivertype" +) + +// ContextKeyInsideTestWorker is an internal context key that indicates whether +// the worker is running inside a [rivertest.Worker]. +type ContextKeyInsideTestWorker struct{} + +type Func func(ctx context.Context) error + +// MiddlewareChain chains together the given middleware functions, returning a +// single function that applies them all in reverse order. +func MiddlewareChain(globalMiddleware []rivertype.Middleware, workerMiddleware []rivertype.WorkerMiddleware, doInner Func, jobRow *rivertype.JobRow) Func { + // Quick return for no middleware, which will often be the case. + if len(globalMiddleware) < 1 && len(workerMiddleware) < 1 { + return doInner + } + + // Wrap middlewares in reverse order so the one defined first is wrapped + // as the outermost function and is first to receive the operation. + for _, v := range slices.Backward(globalMiddleware) { + middlewareItem := v.(rivertype.WorkerMiddleware) //nolint:forcetypeassert // capture the current middleware item + previousDoInner := doInner // capture the current doInner function + doInner = func(ctx context.Context) error { + return middlewareItem.Work(ctx, jobRow, previousDoInner) + } + } + + for _, v := range slices.Backward(workerMiddleware) { + middlewareItem := v // capture the current middleware item + previousDoInner := doInner // capture the current doInner function + doInner = func(ctx context.Context) error { + return middlewareItem.Work(ctx, jobRow, previousDoInner) + } + } + + return doInner +} diff --git a/vendor/github.com/riverqueue/river/internal/jobcompleter/job_completer.go b/vendor/github.com/riverqueue/river/internal/jobcompleter/job_completer.go new file mode 100644 index 0000000000..5e6c7c74c4 --- /dev/null +++ b/vendor/github.com/riverqueue/river/internal/jobcompleter/job_completer.go @@ -0,0 +1,714 @@ +package jobcompleter + +import ( + "context" + "errors" + "log/slog" + "sync" + "time" + + "golang.org/x/sync/errgroup" + + "github.com/riverqueue/river/internal/jobstats" + "github.com/riverqueue/river/internal/rivercommon" + "github.com/riverqueue/river/riverdriver" + "github.com/riverqueue/river/rivershared/baseservice" + "github.com/riverqueue/river/rivershared/riverpilot" + "github.com/riverqueue/river/rivershared/startstop" + "github.com/riverqueue/river/rivershared/util/serviceutil" + "github.com/riverqueue/river/rivershared/util/timeoututil" + "github.com/riverqueue/river/rivertype" +) + +// JobCompleter is an interface to a service that "completes" jobs by marking +// them with an appropriate state and any other necessary metadata in the +// database. It's a generic interface to let us experiment with the speed of a +// number of implementations, although River will likely always prefer our most +// optimized one. +type JobCompleter interface { + startstop.Service + + // JobSetState sets a new state for the given job, as long as it's + // still running (i.e. its state has not changed to something else already). + JobSetStateIfRunning(ctx context.Context, stats *jobstats.JobStatistics, params *riverdriver.JobSetStateIfRunningParams) error + + // ResetSubscribeChan resets the subscription channel for the completer. It + // must only be called when the completer is stopped. + ResetSubscribeChan(subscribeCh SubscribeChan) +} + +type SubscribeChan chan<- []CompleterJobUpdated + +// SubscribeFunc will be invoked whenever a job is updated. +type SubscribeFunc func(update CompleterJobUpdated) + +type CompleterJobUpdated struct { + Job *rivertype.JobRow + JobStats *jobstats.JobStatistics + Snoozed bool +} + +type InlineCompleter struct { + baseservice.BaseService + startstop.BaseStartStop + + disableSleep bool // disable sleep in testing + exec riverdriver.Executor + pilot riverpilot.Pilot + schema string + subscribeCh SubscribeChan + + // A waitgroup is not actually needed for the inline completer because as + // long as the caller is waiting on each function call, completion is + // guaranteed to be done by the time Wait is called. However, we use a + // generic test helper for all completers that starts goroutines, so this + // left in for now for the benefit of the test suite. + wg sync.WaitGroup +} + +func NewInlineCompleter(archetype *baseservice.Archetype, schema string, exec riverdriver.Executor, pilot riverpilot.Pilot, subscribeCh SubscribeChan) *InlineCompleter { + return baseservice.Init(archetype, &InlineCompleter{ + exec: exec, + pilot: pilot, + schema: schema, + subscribeCh: subscribeCh, + }) +} + +func (c *InlineCompleter) JobSetStateIfRunning(ctx context.Context, stats *jobstats.JobStatistics, params *riverdriver.JobSetStateIfRunningParams) error { + c.wg.Add(1) + defer c.wg.Done() + + start := c.Time.Now() + + jobs, err := withRetries(ctx, &c.BaseService, c.disableSleep, func(ctx context.Context) ([]*rivertype.JobRow, error) { + jobs, err := c.pilot.JobSetStateIfRunningMany(ctx, c.exec, setStateParamsToMany(c.Time.NowOrNil(), c.schema, params)) + if err != nil { + return nil, err + } + + return jobs, nil + }) + if err != nil { + return err + } + + // The driver intentionally returns 0 rows when a job is deleted while the + // completer is finalizing it (see UnknownJobIgnored shared driver test). + // Guard against an index-out-of-range panic in that case. + if len(jobs) < 1 { + return nil + } + + stats.CompleteDuration = c.Time.Now().Sub(start) + c.subscribeCh <- []CompleterJobUpdated{{ + Job: jobs[0], + JobStats: stats, + Snoozed: params.Snoozed, + }} + + return nil +} + +func (c *InlineCompleter) ResetSubscribeChan(subscribeCh SubscribeChan) { + c.subscribeCh = subscribeCh +} + +func (c *InlineCompleter) Start(ctx context.Context) error { + ctx, shouldStart, started, stopped := c.StartInit(ctx) + if !shouldStart { + return nil + } + + if c.subscribeCh == nil { + panic("subscribeCh must be non-nil") + } + + go func() { + started() + defer stopped() + defer close(c.subscribeCh) + + <-ctx.Done() + + c.wg.Wait() + }() + + return nil +} + +func setStateParamsToMany(now *time.Time, schema string, params *riverdriver.JobSetStateIfRunningParams) *riverdriver.JobSetStateIfRunningManyParams { + return &riverdriver.JobSetStateIfRunningManyParams{ + Attempt: []*int{params.Attempt}, + ErrData: [][]byte{params.ErrData}, + FinalizedAt: []*time.Time{params.FinalizedAt}, + ID: []int64{params.ID}, + MetadataDoMerge: []bool{params.MetadataDoMerge}, + MetadataUpdates: [][]byte{params.MetadataUpdates}, + Now: now, + ScheduledAt: []*time.Time{params.ScheduledAt}, + Schema: schema, + State: []rivertype.JobState{params.State}, + } +} + +// A default concurrency of 100 seems to perform better a much smaller number +// like 10, but it's quite dependent on environment (10 and 100 bench almost +// identically on MBA when it's on battery power). This number should represent +// our best known default for most use cases, but don't consider its choice to +// be particularly well informed at this point. +const asyncCompleterDefaultConcurrency = 100 + +type AsyncCompleter struct { + baseservice.BaseService + startstop.BaseStartStop + + concurrency int + disableSleep bool // disable sleep in testing + errGroup *errgroup.Group + exec riverdriver.Executor + pilot riverpilot.Pilot + schema string + subscribeCh SubscribeChan +} + +func NewAsyncCompleter(archetype *baseservice.Archetype, schema string, exec riverdriver.Executor, pilot riverpilot.Pilot, subscribeCh SubscribeChan) *AsyncCompleter { + return newAsyncCompleterWithConcurrency(archetype, schema, exec, pilot, asyncCompleterDefaultConcurrency, subscribeCh) +} + +func newAsyncCompleterWithConcurrency(archetype *baseservice.Archetype, schema string, exec riverdriver.Executor, pilot riverpilot.Pilot, concurrency int, subscribeCh SubscribeChan) *AsyncCompleter { + errGroup := &errgroup.Group{} + errGroup.SetLimit(concurrency) + + return baseservice.Init(archetype, &AsyncCompleter{ + concurrency: concurrency, + errGroup: errGroup, + exec: exec, + pilot: pilot, + schema: schema, + subscribeCh: subscribeCh, + }) +} + +func (c *AsyncCompleter) JobSetStateIfRunning(ctx context.Context, stats *jobstats.JobStatistics, params *riverdriver.JobSetStateIfRunningParams) error { + // Start clock outside of goroutine so that the time spent blocking waiting + // for an errgroup slot is accurately measured. + start := c.Time.Now() + + c.errGroup.Go(func() error { + jobs, err := withRetries(ctx, &c.BaseService, c.disableSleep, func(ctx context.Context) ([]*rivertype.JobRow, error) { + rows, err := c.pilot.JobSetStateIfRunningMany(ctx, c.exec, setStateParamsToMany(c.Time.NowOrNil(), c.schema, params)) + if err != nil { + return nil, err + } + + return rows, nil + }) + if err != nil { + return err + } + + // The driver intentionally returns 0 rows when a job is deleted while the + // completer is finalizing it (see UnknownJobIgnored shared driver test). + // Guard against an index-out-of-range panic in that case. + if len(jobs) < 1 { + return nil + } + + stats.CompleteDuration = c.Time.Now().Sub(start) + c.subscribeCh <- []CompleterJobUpdated{{ + Job: jobs[0], + JobStats: stats, + Snoozed: params.Snoozed, + }} + + return nil + }) + return nil +} + +func (c *AsyncCompleter) ResetSubscribeChan(subscribeCh SubscribeChan) { + c.subscribeCh = subscribeCh +} + +func (c *AsyncCompleter) Start(ctx context.Context) error { + ctx, shouldStart, started, stopped := c.StartInit(ctx) + if !shouldStart { + return nil + } + + if c.subscribeCh == nil { + panic("subscribeCh must be non-nil") + } + + go func() { + started() + defer stopped() // this defer should come first so it's first out + defer close(c.subscribeCh) + + <-ctx.Done() + + if err := c.errGroup.Wait(); err != nil { + c.Logger.ErrorContext(ctx, "Error waiting on async completer", "err", err) + } + }() + + return nil +} + +type batchCompleterSetState struct { + Params *riverdriver.JobSetStateIfRunningParams + StartTime time.Time + Stats *jobstats.JobStatistics +} + +// BatchCompleter accumulates incoming completions, and instead of completing +// them immediately, every so often complete many of them as a single efficient +// batch. To minimize the amount of driver surface area we need, the batching is +// only performed for jobs being changed to a `completed` state, which we expect +// to the vast common case under normal operation. The completer embeds an +// AsyncCompleter to perform other non-`completed` state completions. +type BatchCompleter struct { + baseservice.BaseService + startstop.BaseStartStop + + backlogWaitThreshold int // configurable for testing purposes; backlog at which completions start waiting for the completer to catch up + batchReadyChan chan struct{} + completionMaxSize int // configurable for testing purposes; max jobs to complete in single database operation + disableSleep bool // disable sleep in testing + maxBacklog int // configurable for testing purposes; emergency backlog threshold where a warning is logged + exec riverdriver.Executor + pilot riverpilot.Pilot + schema string + setStateParams map[int64]batchCompleterSetState + setStateParamsMu sync.RWMutex + subscribeCh SubscribeChan + waitOnBacklogChan chan struct{} + waitOnBacklogWaiting bool +} + +func NewBatchCompleter(archetype *baseservice.Archetype, schema string, exec riverdriver.Executor, pilot riverpilot.Pilot, subscribeCh SubscribeChan) *BatchCompleter { + const ( + completionMaxSize = 5_000 + backlogWaitThreshold = completionMaxSize * 2 + maxBacklog = 20_000 + ) + + return baseservice.Init(archetype, &BatchCompleter{ + backlogWaitThreshold: backlogWaitThreshold, + batchReadyChan: make(chan struct{}, 1), + completionMaxSize: completionMaxSize, + exec: exec, + maxBacklog: maxBacklog, + pilot: pilot, + schema: schema, + setStateParams: make(map[int64]batchCompleterSetState), + subscribeCh: subscribeCh, + }) +} + +func (c *BatchCompleter) ResetSubscribeChan(subscribeCh SubscribeChan) { + c.subscribeCh = subscribeCh +} + +func (c *BatchCompleter) Start(ctx context.Context) error { + stopCtx, shouldStart, started, stopped := c.StartInit(ctx) + if !shouldStart { + return nil + } + + if c.subscribeCh == nil { + panic("subscribeCh must be non-nil") + } + + go func() { + started() + defer stopped() // this defer should come first so it's first out + defer close(c.subscribeCh) + + c.Logger.DebugContext(ctx, c.Name+": Run loop started") + defer c.Logger.DebugContext(ctx, c.Name+": Run loop stopped") + + ticker := time.NewTicker(50 * time.Millisecond) + defer ticker.Stop() + + backlogSize := func() int { + c.setStateParamsMu.RLock() + defer c.setStateParamsMu.RUnlock() + return len(c.setStateParams) + } + + for numTicks := 0; ; numTicks++ { + select { + case <-stopCtx.Done(): + // Try to insert last batch before leaving. Note we use the + // original context so operations aren't immediately cancelled. + if err := c.handleBatch(ctx); err != nil { + c.Logger.ErrorContext(ctx, c.Name+": Error completing batch", "err", err) + } + return + + case <-c.batchReadyChan: + case <-ticker.C: + } + + // The ticker fires quite often to make sure that given a huge glut + // of jobs, we don't accidentally build up too much of a backlog by + // waiting too long. However, don't start a complete operation until + // we reach a minimum threshold unless we're on a tick that's a + // multiple of 5. So, jobs will be completed every 250ms even if the + // threshold hasn't been met. + const batchCompleterStartThreshold = 100 + if backlogSize() < min(c.backlogWaitThresholdEffective(), batchCompleterStartThreshold) && numTicks != 0 && numTicks%5 != 0 { + continue + } + + for { + if err := c.handleBatch(ctx); err != nil { + c.Logger.ErrorContext(ctx, c.Name+": Error completing batch", "err", err) + } + + // New jobs to complete may have come in while working the batch + // above. If enough have to bring us above the minimum complete + // threshold, loop again and do another batch. Otherwise, break + // and listen for a new tick. + if backlogSize() < batchCompleterStartThreshold { + break + } + } + } + }() + + return nil +} + +func (c *BatchCompleter) handleBatch(ctx context.Context) error { + var setStateBatch map[int64]batchCompleterSetState + func() { + c.setStateParamsMu.Lock() + defer c.setStateParamsMu.Unlock() + + setStateBatch = c.setStateParams + + // Don't bother resetting the map if there's nothing to process, + // allowing the completer to idle efficiently. + if len(setStateBatch) > 0 { + c.setStateParams = make(map[int64]batchCompleterSetState) + } else { + // Set nil to avoid a data race below in case the map is set as a + // new job comes in. + setStateBatch = nil + } + }() + + if len(setStateBatch) < 1 { + return nil + } + + handleBatchError := func(err error) error { + if isNonRetryableCompleterError(err) { + c.releaseBacklogWaitIfReady(ctx) + return err + } + + c.requeueBatch(ctx, setStateBatch) + return err + } + + // Complete a sub-batch with retries. Also helps reduce visual noise and + // increase readability of loop below. + completeSubBatch := func(batchParams *riverdriver.JobSetStateIfRunningManyParams) ([]*rivertype.JobRow, error) { + start := time.Now() + defer func() { + c.Logger.DebugContext(ctx, c.Name+": Completed sub-batch of job(s)", "duration", time.Since(start), "num_jobs", len(batchParams.ID)) + }() + + return withRetries(ctx, &c.BaseService, c.disableSleep, func(ctx context.Context) ([]*rivertype.JobRow, error) { + rows, err := c.pilot.JobSetStateIfRunningMany(ctx, c.exec, batchParams) + if err != nil { + return nil, err + } + + return rows, nil + }) + } + + // This could be written more simply using multiple map helpers, but it's + // done this way to allocate as few new slices as necessary. + mapBatch := func(setStateBatch map[int64]batchCompleterSetState) *riverdriver.JobSetStateIfRunningManyParams { + params := &riverdriver.JobSetStateIfRunningManyParams{ + ID: make([]int64, len(setStateBatch)), + Attempt: make([]*int, len(setStateBatch)), + ErrData: make([][]byte, len(setStateBatch)), + FinalizedAt: make([]*time.Time, len(setStateBatch)), + MetadataDoMerge: make([]bool, len(setStateBatch)), + MetadataUpdates: make([][]byte, len(setStateBatch)), + ScheduledAt: make([]*time.Time, len(setStateBatch)), + State: make([]rivertype.JobState, len(setStateBatch)), + } + var i int + for _, setState := range setStateBatch { + params.ID[i] = setState.Params.ID + params.Attempt[i] = setState.Params.Attempt + params.ErrData[i] = setState.Params.ErrData + params.FinalizedAt[i] = setState.Params.FinalizedAt + params.MetadataDoMerge[i] = setState.Params.MetadataDoMerge + params.MetadataUpdates[i] = setState.Params.MetadataUpdates + params.ScheduledAt[i] = setState.Params.ScheduledAt + params.State[i] = setState.Params.State + i++ + } + params.Schema = c.schema + return params + } + + // Tease apart enormous batches into sub-batches. + // + // All the code below is concerned with doing that, with a fast loop that + // doesn't allocate any additional memory in case the entire batch is + // smaller than the sub-batch maximum size (which will be the common case). + var ( + params = mapBatch(setStateBatch) + jobRows []*rivertype.JobRow + ) + c.Logger.DebugContext(ctx, c.Name+": Completing batch of job(s)", "num_jobs", len(setStateBatch)) + if len(setStateBatch) > c.completionMaxSize { + jobRows = make([]*rivertype.JobRow, 0, len(setStateBatch)) + for i := 0; i < len(setStateBatch); i += c.completionMaxSize { + endIndex := min(i+c.completionMaxSize, len(params.ID)) // beginning of next sub-batch or end of slice + subBatch := &riverdriver.JobSetStateIfRunningManyParams{ + ID: params.ID[i:endIndex], + Attempt: params.Attempt[i:endIndex], + ErrData: params.ErrData[i:endIndex], + FinalizedAt: params.FinalizedAt[i:endIndex], + MetadataDoMerge: params.MetadataDoMerge[i:endIndex], + MetadataUpdates: params.MetadataUpdates[i:endIndex], + ScheduledAt: params.ScheduledAt[i:endIndex], + Schema: params.Schema, + State: params.State[i:endIndex], + } + jobRowsSubBatch, err := completeSubBatch(subBatch) + if err != nil { + return handleBatchError(err) + } + jobRows = append(jobRows, jobRowsSubBatch...) + } + } else { + var err error + jobRows, err = completeSubBatch(params) + if err != nil { + return handleBatchError(err) + } + } + + var ( + completeTime = c.Time.Now() + events = make([]CompleterJobUpdated, len(jobRows)) + ) + for i, jobRow := range jobRows { + setState := setStateBatch[jobRow.ID] + setState.Stats.CompleteDuration = completeTime.Sub(setState.StartTime) + events[i] = CompleterJobUpdated{ + Job: jobRow, + JobStats: setState.Stats, + Snoozed: setState.Params.Snoozed, + } + } + + c.subscribeCh <- events + + func() { + c.setStateParamsMu.Lock() + defer c.setStateParamsMu.Unlock() + + if c.waitOnBacklogWaiting && len(c.setStateParams) < c.backlogResumeThreshold() { + c.Logger.DebugContext(ctx, c.Name+": Disabling waitOnBacklog; ready to complete more jobs") + close(c.waitOnBacklogChan) + c.waitOnBacklogWaiting = false + } + }() + + return nil +} + +func (c *BatchCompleter) releaseBacklogWaitIfReady(ctx context.Context) { + c.setStateParamsMu.Lock() + defer c.setStateParamsMu.Unlock() + + if c.waitOnBacklogWaiting && len(c.setStateParams) < c.backlogResumeThreshold() { + c.Logger.DebugContext(ctx, c.Name+": Disabling waitOnBacklog; ready to complete more jobs") + close(c.waitOnBacklogChan) + c.waitOnBacklogWaiting = false + } +} + +func (c *BatchCompleter) requeueBatch(ctx context.Context, setStateBatch map[int64]batchCompleterSetState) { + c.setStateParamsMu.Lock() + for id, setState := range setStateBatch { + if _, exists := c.setStateParams[id]; exists { + continue + } + c.setStateParams[id] = setState + } + backlogSize := len(c.setStateParams) + if c.waitOnBacklogWaiting && backlogSize < c.backlogResumeThreshold() { + c.Logger.DebugContext(ctx, c.Name+": Disabling waitOnBacklog; ready to complete more jobs") + close(c.waitOnBacklogChan) + c.waitOnBacklogWaiting = false + } + c.setStateParamsMu.Unlock() + + if backlogSize >= c.batchReadyThreshold() { + c.signalBatchReady() + } + + c.Logger.DebugContext(ctx, c.Name+": Requeued failed batch of job(s)", "num_jobs", len(setStateBatch)) +} + +func (c *BatchCompleter) JobSetStateIfRunning(ctx context.Context, stats *jobstats.JobStatistics, params *riverdriver.JobSetStateIfRunningParams) error { + now := c.Time.Now() + + var backlogSize int + for { + // Keep the common enqueue path to one lock acquisition. If the + // completer is behind, wait for the current backlog gate to open and + // retry so the threshold is checked against fresh state. + var waitChan <-chan struct{} + backlogSize, waitChan = c.tryEnqueueSetState(ctx, now, stats, params) + if waitChan != nil { + <-waitChan + continue + } + break + } + + if backlogSize >= c.batchReadyThreshold() { + c.signalBatchReady() + } + + return nil +} + +func (c *BatchCompleter) tryEnqueueSetState(ctx context.Context, now time.Time, stats *jobstats.JobStatistics, params *riverdriver.JobSetStateIfRunningParams) (int, <-chan struct{}) { + c.setStateParamsMu.Lock() + defer c.setStateParamsMu.Unlock() + + if c.waitOnBacklogWaiting { + return 0, c.waitOnBacklogChan + } + + var ( + backlogSize = len(c.setStateParams) + waitAt = c.backlogWaitThresholdEffective() + ) + if backlogSize >= waitAt { + c.initBacklogWaitLocked(ctx, backlogSize, waitAt) + } + + statsSnapshot := *stats + c.setStateParams[params.ID] = batchCompleterSetState{Params: params, StartTime: now, Stats: &statsSnapshot} + + return len(c.setStateParams), nil +} + +// backlogResumeThreshold returns the low-water mark below which waiting +// completers are released. Keeping this below the wait threshold avoids rapidly +// cycling between waiting and not waiting when the completer is near capacity. +func (c *BatchCompleter) backlogResumeThreshold() int { + return max(c.backlogWaitThresholdEffective()/2, 1) +} + +// backlogWaitThresholdEffective returns the backlog size at which new +// completions should wait for the batch completer to catch up. It's capped at +// maxBacklog so tests and future configuration can't set a normal wait +// threshold beyond the emergency warning threshold. +func (c *BatchCompleter) backlogWaitThresholdEffective() int { + if c.backlogWaitThreshold <= 0 { + return c.maxBacklog + } + return min(c.backlogWaitThreshold, c.maxBacklog) +} + +// batchReadyThreshold returns the backlog size at which the run loop should be +// nudged to process a batch immediately instead of waiting for its next ticker. +// It aims for a full database batch while still respecting low test thresholds. +func (c *BatchCompleter) batchReadyThreshold() int { + return min(c.completionMaxSize, c.backlogWaitThresholdEffective()) +} + +func (c *BatchCompleter) signalBatchReady() { + select { + case c.batchReadyChan <- struct{}{}: + default: + } +} + +// initBacklogWaitLocked starts a backlog wait gate and must be called with +// setStateParamsMu held. +func (c *BatchCompleter) initBacklogWaitLocked(ctx context.Context, backlogSize, waitAt int) chan struct{} { + c.waitOnBacklogChan = make(chan struct{}) + c.waitOnBacklogWaiting = true + if backlogSize >= c.maxBacklog { + c.Logger.WarnContext(ctx, c.Name+": Hit maximum backlog; completions will wait until below threshold", + "backlog_size", backlogSize, + "backlog_wait_threshold", waitAt, + "max_backlog", c.maxBacklog, + ) + } else { + c.Logger.DebugContext(ctx, c.Name+": Applying completion backlog pressure", + "backlog_resume_threshold", c.backlogResumeThreshold(), + "backlog_size", backlogSize, + "backlog_wait_threshold", waitAt, + ) + } + return c.waitOnBacklogChan +} + +func isNonRetryableCompleterError(err error) bool { + return errors.Is(err, context.Canceled) || errors.Is(err, riverdriver.ErrClosedPool) +} + +// As configured, total time asleep from initial attempt is ~7 seconds (1 + 2 + +// 4) (not including jitter). However, if each attempt times out, that's up to +// ~37 seconds (7 seconds + 3 * 10 seconds). +const numRetries = 3 + +func withRetries[T any](logCtx context.Context, baseService *baseservice.BaseService, disableSleep bool, retryFunc func(ctx context.Context) (T, error)) (T, error) { + uncancelledCtx := context.WithoutCancel(logCtx) + + var ( + defaultVal T + lastErr error + ) + + for attempt := 1; attempt <= numRetries; attempt++ { + // I've found that we want at least ten seconds for a large batch, + // although it usually doesn't need that long. + retVal, err := timeoututil.WithTimeoutV(uncancelledCtx, rivercommon.HotOperationTimeout, baseService.Name+".withRetries", retryFunc) + if err != nil { + // A cancelled context or a closed pool will never succeed. + if isNonRetryableCompleterError(err) { + return defaultVal, err + } + + lastErr = err + sleepDuration := serviceutil.ExponentialBackoff(attempt, serviceutil.MaxAttemptsBeforeResetDefault) + baseService.Logger.ErrorContext(logCtx, baseService.Name+": Completer error (will retry after sleep)", + slog.Int("attempt", attempt), + slog.String("err", err.Error()), + slog.String("sleep_duration", sleepDuration.String()), + slog.String("timeout", rivercommon.HotOperationTimeout.String()), + ) + if !disableSleep { + serviceutil.CancellableSleep(logCtx, sleepDuration) + } + continue + } + + return retVal, nil + } + + baseService.Logger.ErrorContext(logCtx, baseService.Name+": Too many errors; giving up") + + return defaultVal, lastErr +} diff --git a/vendor/github.com/riverqueue/river/internal/jobexecutor/job_executor.go b/vendor/github.com/riverqueue/river/internal/jobexecutor/job_executor.go new file mode 100644 index 0000000000..c668001445 --- /dev/null +++ b/vendor/github.com/riverqueue/river/internal/jobexecutor/job_executor.go @@ -0,0 +1,567 @@ +package jobexecutor + +import ( + "cmp" + "context" + "encoding/json" + "errors" + "fmt" + "log/slog" + "runtime" + "strings" + "sync/atomic" + "time" + + "github.com/tidwall/gjson" + + "github.com/riverqueue/river/internal/execution" + "github.com/riverqueue/river/internal/jobcompleter" + "github.com/riverqueue/river/internal/jobstats" + "github.com/riverqueue/river/internal/pluginlookup" + "github.com/riverqueue/river/internal/workunit" + "github.com/riverqueue/river/riverdriver" + "github.com/riverqueue/river/rivershared/baseservice" + "github.com/riverqueue/river/rivertype" +) + +type ClientRetryPolicy interface { + NextRetry(job *rivertype.JobRow) time.Time +} + +// ErrorHandler provides an interface that will be invoked in case of an error +// or panic occurring in the job. This is often useful for logging and exception +// tracking, but can also be used to customize retry behavior. +type ErrorHandler interface { + // HandleError is invoked in case of an error occurring in a job. + // + // Context is descended from the one used to start the River client that + // worked the job. + HandleError(ctx context.Context, job *rivertype.JobRow, err error) *ErrorHandlerResult + + // HandlePanic is invoked in case of a panic occurring in a job. + // + // Context is descended from the one used to start the River client that + // worked the job. + HandlePanic(ctx context.Context, job *rivertype.JobRow, panicVal any, trace string) *ErrorHandlerResult +} + +type ErrorHandlerResult struct { + // SetCancelled can be set to true to fail the job immediately and + // permanently. By default it'll continue to follow the configured retry + // schedule. + SetCancelled bool +} + +// Error used in CancelFunc in cases where the job was not cancelled for +// purposes of resource cleanup. Should never be user visible. +var errExecutorDefaultCancel = errors.New("context cancelled as executor finished") + +type contextKey string + +// ContextKeyMetadataUpdates is the context key for the metadata updates map +// stored in the context. It's exposed from this internal package solely so +// that it can be used in tests for JobCompleteTx. +const ContextKeyMetadataUpdates contextKey = "river_metadata_updates" + +// MetadataUpdatesFromWorkContext returns the metadata updates stored in the +// work context, if any. +// +// When run on a non-work context, it returns nil, false. +func MetadataUpdatesFromWorkContext(ctx context.Context) (map[string]any, bool) { + metadataUpdates := ctx.Value(ContextKeyMetadataUpdates) + if metadataUpdates == nil { + return nil, false + } + typedMetadataUpdates, ok := metadataUpdates.(map[string]any) + if !ok { + return nil, false + } + return typedMetadataUpdates, true +} + +type jobExecutorResult struct { + Err error + JobArgsUnmarshaled bool + MetadataUpdates map[string]any + NextRetry time.Time + PanicTrace string + PanicVal any +} + +// ErrorStr returns an appropriate string to persist to the database based on +// the type of internal failure (i.e. error or panic). Panics if called on a +// non-errored result. +func (r *jobExecutorResult) ErrorStr() string { + switch { + case r.Err != nil: + return r.Err.Error() + case r.PanicVal != nil: + return fmt.Sprintf("%v", r.PanicVal) + } + + panic("ErrorStr should not be called on non-errored result") +} + +type JobExecutor struct { + baseservice.BaseService + + CancelFunc context.CancelCauseFunc + ClientJobTimeout time.Duration + Completer jobcompleter.JobCompleter + ClientRetryPolicy ClientRetryPolicy + DefaultClientRetryPolicy ClientRetryPolicy + ErrorHandler ErrorHandler + PluginLookupByJob *pluginlookup.JobPluginLookup + PluginLookupGlobal *pluginlookup.PluginLookup + JobRow *rivertype.JobRow + ProducerCallbacks struct { + JobDone func(jobRow *rivertype.JobRow) + Stuck func(ctx context.Context, jobRow *rivertype.JobRow) + Unstuck func() + } + SchedulerInterval time.Duration + StuckThresholdOverride time.Duration + WorkerMiddleware []rivertype.WorkerMiddleware + WorkUnit workunit.WorkUnit + + // Meant to be used from within the job executor only. + slotClosed atomic.Bool + start time.Time + stats *jobstats.JobStatistics // initialized by the executor, and handed off to completer +} + +// TryCloseSlot marks this executor's producer slot as closed. A closed slot +// means the producer has already stopped counting this executor against its +// active worker capacity, although the executor goroutine may still be running. +// It returns true only the first time the slot is closed. +func (e *JobExecutor) TryCloseSlot() bool { + return e.slotClosed.CompareAndSwap(false, true) +} + +func (e *JobExecutor) Cancel(ctx context.Context) { + e.Logger.WarnContext(ctx, e.Name+": job cancelled remotely", slog.Int64("job_id", e.JobRow.ID)) + e.CancelFunc(rivertype.ErrJobCancelledRemotely) +} + +func (e *JobExecutor) Execute(ctx context.Context) { + // Ensure that the context is cancelled no matter what, or it will leak: + defer e.CancelFunc(errExecutorDefaultCancel) + + e.start = e.Time.Now() + e.stats = &jobstats.JobStatistics{ + QueueWaitDuration: e.start.Sub(e.JobRow.ScheduledAt), + } + + res := e.execute(ctx) + if res.Err != nil && errors.Is(context.Cause(ctx), rivertype.ErrJobCancelledRemotely) { + res.Err = context.Cause(ctx) + } + + var multiJobErrors withJobsAndErrorsByID + if res.Err != nil { + multiJobErrors, _ = res.Err.(withJobsAndErrorsByID) + } + + if multiJobErrors == nil { + e.reportResult(ctx, e.JobRow, res) + } else { + errorsByID := multiJobErrors.ErrorsByID() + for _, jobRow := range multiJobErrors.Jobs() { + jobSpecificRes := *res + jobSpecificRes.Err = errorsByID[jobRow.ID] + e.reportResult(ctx, jobRow, &jobSpecificRes) + } + } + + e.ProducerCallbacks.JobDone(e.JobRow) +} + +// Executes the job, handling a panic if necessary (and various other error +// conditions). The named return value is so that we can still return a value in +// case of a panic. +// +//nolint:nonamedreturns +func (e *JobExecutor) execute(ctx context.Context) (res *jobExecutorResult) { + metadataUpdates := make(map[string]any) + ctx = context.WithValue(ctx, ContextKeyMetadataUpdates, metadataUpdates) + jobArgsUnmarshaled := false + + defer func() { + if recovery := recover(); recovery != nil { + e.Logger.ErrorContext(ctx, e.Name+": panic recovery; possible bug with Worker", + slog.Int64("job_id", e.JobRow.ID), + slog.String("kind", e.JobRow.Kind), + slog.String("panic_val", fmt.Sprintf("%v", recovery)), + ) + + res = &jobExecutorResult{ + JobArgsUnmarshaled: jobArgsUnmarshaled, + MetadataUpdates: metadataUpdates, + // Skip the first 4 frames which are: + // + // 1. The `runtime.Callers` function. + // 2. The `captureStackTraceSkipFrames` function. + // 3. The current recovery defer function. + // 4. The `JobExecutor.execute` method working the job. + PanicTrace: captureStackTraceSkipFrames(4), + PanicVal: recovery, + } + } + e.stats.RunDuration = e.Time.Now().Sub(e.start) + }() + + if e.WorkUnit == nil { + e.Logger.ErrorContext(ctx, e.Name+": Unhandled job kind", + slog.String("kind", e.JobRow.Kind), + slog.Int64("job_id", e.JobRow.ID), + ) + return &jobExecutorResult{Err: &rivertype.UnknownJobKindError{Kind: e.JobRow.Kind}, MetadataUpdates: metadataUpdates} + } + pluginLookupByJob := e.WorkUnit.PluginLookup(e.PluginLookupByJob) + + doInner := execution.Func(func(ctx context.Context) error { + { + for _, hook := range append( + e.PluginLookupGlobal.ByKind(pluginlookup.PluginKindHookWorkBegin), + pluginLookupByJob.ByKind(pluginlookup.PluginKindHookWorkBegin)..., + ) { + if err := hook.(rivertype.HookWorkBegin).WorkBegin(ctx, e.JobRow); err != nil { //nolint:forcetypeassert + return err + } + } + } + + if err := e.WorkUnit.UnmarshalJob(); err != nil { + return err + } + jobArgsUnmarshaled = true + + jobTimeout := cmp.Or(e.WorkUnit.Timeout(), e.ClientJobTimeout) + + if jobTimeout > 0 { + var timeoutCancel context.CancelFunc + ctx, timeoutCancel = context.WithTimeout(ctx, jobTimeout) + defer timeoutCancel() + + watchStuckCancel := e.watchStuck(ctx, jobTimeout) + defer watchStuckCancel() + } + + err := e.WorkUnit.Work(ctx) + + { + for _, hook := range append( + e.PluginLookupGlobal.ByKind(pluginlookup.PluginKindHookWorkEnd), + pluginLookupByJob.ByKind(pluginlookup.PluginKindHookWorkEnd)..., + ) { + err = hook.(rivertype.HookWorkEnd).WorkEnd(ctx, e.JobRow, err) //nolint:forcetypeassert + } + } + + return err + }) + + pluginMiddleware := make([]rivertype.Middleware, 0, + len(e.PluginLookupGlobal.ByKind(pluginlookup.PluginKindMiddlewareWorker))+ + len(pluginLookupByJob.ByKind(pluginlookup.PluginKindMiddlewareWorker)), + ) + for _, plugin := range e.PluginLookupGlobal.ByKind(pluginlookup.PluginKindMiddlewareWorker) { + pluginMiddleware = append(pluginMiddleware, plugin.(rivertype.Middleware)) //nolint:forcetypeassert + } + for _, plugin := range pluginLookupByJob.ByKind(pluginlookup.PluginKindMiddlewareWorker) { + pluginMiddleware = append(pluginMiddleware, plugin.(rivertype.Middleware)) //nolint:forcetypeassert + } + + executeFunc := execution.MiddlewareChain( + pluginMiddleware, + e.WorkUnit.Middleware(), + doInner, + e.JobRow, + ) + + return &jobExecutorResult{ + Err: executeFunc(ctx), + JobArgsUnmarshaled: jobArgsUnmarshaled, + MetadataUpdates: metadataUpdates, + } +} + +// Watches for jobs that may have become stuck. i.e. They've run longer than +// their job timeout (plus a small margin) and don't appear to be responding to +// context cancellation (unfortunately, quite an easy error to make in Go). +// +// Producers use stuck-job notifications for periodic stats and optional user +// handlers. +func (e *JobExecutor) watchStuck(ctx context.Context, jobTimeout time.Duration) context.CancelFunc { + // We add a WithoutCancel here so that this inner goroutine becomes + // immune to all context cancellations _except_ the one where it's + // cancelled because we leave JobExecutor.execute. + // + // This shadows the context outside the e.ClientJobTimeout > 0 check. + ctx, cancel := context.WithCancel(context.WithoutCancel(ctx)) + + go func() { + const stuckThresholdDefault = 5 * time.Second + + select { + case <-ctx.Done(): + // context cancelled as we leave JobExecutor.execute + + case <-time.After(jobTimeout + cmp.Or(e.StuckThresholdOverride, stuckThresholdDefault)): + e.ProducerCallbacks.Stuck(ctx, e.JobRow) + + e.Logger.WarnContext(ctx, e.Name+": Job appears to be stuck", + slog.Int64("job_id", e.JobRow.ID), + slog.String("kind", e.JobRow.Kind), + slog.Duration("timeout", e.ClientJobTimeout), + ) + + // context cancelled as we leave JobExecutor.execute + <-ctx.Done() + + // In case the executor ever becomes unstuck, inform the + // producer. However, if we got all the way here there's a good + // chance this will never happen (the worker is really stuck and + // will never return). + defer e.ProducerCallbacks.Unstuck() + + defer func() { + e.Logger.InfoContext(ctx, e.Name+": Job became unstuck", + slog.Duration("duration", time.Since(e.start)), + slog.Int64("job_id", e.JobRow.ID), + slog.String("kind", e.JobRow.Kind), + ) + }() + } + }() + + return cancel +} + +func (e *JobExecutor) invokeErrorHandler(ctx context.Context, res *jobExecutorResult) bool { + invokeAndHandlePanic := func(funcName string, errorHandler func() *ErrorHandlerResult) *ErrorHandlerResult { + defer func() { + if panicVal := recover(); panicVal != nil { + e.Logger.ErrorContext(ctx, e.Name+": ErrorHandler invocation panicked", + slog.String("function_name", funcName), + slog.String("panic_val", fmt.Sprintf("%v", panicVal)), + ) + } + }() + + return errorHandler() + } + + var errorHandlerRes *ErrorHandlerResult + switch { + case res.Err != nil: + errorHandlerRes = invokeAndHandlePanic("HandleError", func() *ErrorHandlerResult { + return e.ErrorHandler.HandleError(ctx, e.JobRow, res.Err) + }) + + case res.PanicVal != nil: + errorHandlerRes = invokeAndHandlePanic("HandlePanic", func() *ErrorHandlerResult { + return e.ErrorHandler.HandlePanic(ctx, e.JobRow, res.PanicVal, res.PanicTrace) + }) + } + + return errorHandlerRes != nil && errorHandlerRes.SetCancelled +} + +func (e *JobExecutor) reportResult(ctx context.Context, jobRow *rivertype.JobRow, res *jobExecutorResult) { + var snoozeErr *rivertype.JobSnoozeError + + marshalMetadataUpdates := func(metadataUpdates map[string]any) ([]byte, error) { + if len(metadataUpdates) == 0 { + return nil, nil + } + + metadataUpdatesBytes, err := json.Marshal(metadataUpdates) + if err != nil { + return nil, err + } + + return metadataUpdatesBytes, nil + } + + if res.Err != nil && errors.As(res.Err, &snoozeErr) { + e.Logger.DebugContext(ctx, e.Name+": Job snoozed", + slog.Int64("job_id", jobRow.ID), + slog.String("job_kind", jobRow.Kind), + slog.Duration("duration", snoozeErr.Duration), + ) + nextAttemptScheduledAt := e.Time.Now().Add(snoozeErr.Duration) + + snoozesValue := gjson.GetBytes(jobRow.Metadata, "snoozes").Int() + if res.MetadataUpdates == nil { + res.MetadataUpdates = make(map[string]any) + } + // Set snooze count in the metadata map before marshaling so we avoid + // rewriting a potentially large encoded metadata payload. + res.MetadataUpdates["snoozes"] = snoozesValue + 1 + + metadataUpdatesBytes, err := marshalMetadataUpdates(res.MetadataUpdates) + if err != nil { + e.Logger.ErrorContext(ctx, e.Name+": Failed to marshal metadata updates", slog.String("error", err.Error())) + return + } + + // Normally, snoozed jobs are set `scheduled` for the future and it's the + // scheduler's job to set them back to `available` so they can be reworked. + // Just as with retryable jobs, this isn't friendly for short snooze times + // so we instead make the job immediately `available` if the snooze time is + // smaller than the scheduler's run interval. + var params *riverdriver.JobSetStateIfRunningParams + if nextAttemptScheduledAt.Sub(e.Time.Now()) <= e.SchedulerInterval { + params = riverdriver.JobSetStateSnoozedAvailable(jobRow.ID, nextAttemptScheduledAt, jobRow.Attempt-1, metadataUpdatesBytes) + } else { + params = riverdriver.JobSetStateSnoozed(jobRow.ID, nextAttemptScheduledAt, jobRow.Attempt-1, metadataUpdatesBytes) + } + if err := e.Completer.JobSetStateIfRunning(ctx, e.stats, params); err != nil { + e.Logger.ErrorContext(ctx, e.Name+": Error snoozing job", + slog.Int64("job_id", jobRow.ID), + ) + } + return + } + + metadataUpdatesBytes, err := marshalMetadataUpdates(res.MetadataUpdates) + if err != nil { + e.Logger.ErrorContext(ctx, e.Name+": Failed to marshal metadata updates", slog.String("error", err.Error())) + return + } + + if res.Err != nil || res.PanicVal != nil { + e.reportError(ctx, jobRow, res, metadataUpdatesBytes) + return + } + + if err := e.Completer.JobSetStateIfRunning(ctx, e.stats, riverdriver.JobSetStateCompleted(jobRow.ID, e.Time.Now(), metadataUpdatesBytes)); err != nil { + e.Logger.ErrorContext(ctx, e.Name+": Error completing job", + slog.String("err", err.Error()), + slog.Int64("job_id", jobRow.ID), + ) + return + } +} + +func (e *JobExecutor) reportError(ctx context.Context, jobRow *rivertype.JobRow, res *jobExecutorResult, metadataUpdates []byte) { + var ( + cancelJob bool + cancelErr *rivertype.JobCancelError + ) + + logAttrs := []any{ + slog.String("error", res.ErrorStr()), + slog.Int64("job_id", jobRow.ID), + slog.String("job_kind", jobRow.Kind), + } + + switch { + case errors.As(res.Err, &cancelErr): + cancelJob = true + e.Logger.DebugContext(ctx, e.Name+": Job cancelled explicitly", logAttrs...) + case res.Err != nil: + if jobRow.Attempt >= jobRow.MaxAttempts { + e.Logger.InfoContext(ctx, e.Name+": Job errored", logAttrs...) + } else { + e.Logger.InfoContext(ctx, e.Name+": Job errored; retrying", logAttrs...) + } + case res.PanicVal != nil: + e.Logger.InfoContext(ctx, e.Name+": Job panicked", logAttrs...) + } + + if e.ErrorHandler != nil && !cancelJob { + // Error handlers also have an opportunity to cancel the job. + cancelJob = e.invokeErrorHandler(ctx, res) + } + + attemptErr := rivertype.AttemptError{ + At: e.start, + Attempt: jobRow.Attempt, + Error: res.ErrorStr(), + Trace: res.PanicTrace, + } + + errData, err := json.Marshal(attemptErr) + if err != nil { + e.Logger.ErrorContext(ctx, e.Name+": Failed to marshal attempt error", logAttrs...) + return + } + + now := e.Time.Now() + + if cancelJob { + if err := e.Completer.JobSetStateIfRunning(ctx, e.stats, riverdriver.JobSetStateCancelled(jobRow.ID, now, errData, metadataUpdates)); err != nil { + e.Logger.ErrorContext(ctx, e.Name+": Failed to cancel job and report error", logAttrs...) + } + return + } + + if jobRow.Attempt >= jobRow.MaxAttempts { + if err := e.Completer.JobSetStateIfRunning(ctx, e.stats, riverdriver.JobSetStateDiscarded(jobRow.ID, now, errData, metadataUpdates)); err != nil { + e.Logger.ErrorContext(ctx, e.Name+": Failed to discard job and report error", logAttrs...) + } + return + } + + var nextRetryScheduledAt time.Time + if e.WorkUnit != nil && res.JobArgsUnmarshaled { + nextRetryScheduledAt = e.WorkUnit.NextRetry() + } + if nextRetryScheduledAt.IsZero() { + nextRetryScheduledAt = e.ClientRetryPolicy.NextRetry(jobRow) + } + if nextRetryScheduledAt.Before(now) { + e.Logger.WarnContext(ctx, + e.Name+": Retry policy returned invalid next retry before current time; using default retry policy instead", + slog.Int("error_count", len(jobRow.Errors)+1), + slog.Time("next_retry_scheduled_at", nextRetryScheduledAt), + slog.Time("now", now), + ) + nextRetryScheduledAt = e.DefaultClientRetryPolicy.NextRetry(jobRow) + } + + // Normally, errored jobs are set `retryable` for the future and it's the + // scheduler's job to set them back to `available` so they can be reworked. + // This isn't friendly for smaller retry times though because it means that + // effectively no retry time smaller than the scheduler's run interval is + // respected. Here, we offset that with a branch that makes jobs immediately + // `available` if their retry was smaller than the scheduler's run interval. + var params *riverdriver.JobSetStateIfRunningParams + if nextRetryScheduledAt.Sub(e.Time.Now()) <= e.SchedulerInterval { + params = riverdriver.JobSetStateErrorAvailable(jobRow.ID, nextRetryScheduledAt, errData, metadataUpdates) + } else { + params = riverdriver.JobSetStateErrorRetryable(jobRow.ID, nextRetryScheduledAt, errData, metadataUpdates) + } + if err := e.Completer.JobSetStateIfRunning(ctx, e.stats, params); err != nil { + e.Logger.ErrorContext(ctx, e.Name+": Failed to report error for job", logAttrs...) + } +} + +type withJobsAndErrorsByID interface { + ErrorsByID() map[int64]error + Jobs() []*rivertype.JobRow +} + +// captureStackTrace returns a formatted stack trace string starting after +// skipping the specified number of frames. The skip parameter should be +// adjusted so that frames you want to hide (like the ones generated by the +// tracing functions themselves) are excluded. +func captureStackTraceSkipFrames(skip int) string { + // Allocate room for up to 100 callers; adjust as needed. + pcs := make([]uintptr, 100) + // Skip the specified number of frames. + n := runtime.Callers(skip, pcs) + frames := runtime.CallersFrames(pcs[:n]) + + var stackTraceSB strings.Builder + for { + frame, more := frames.Next() + fmt.Fprintf(&stackTraceSB, "%s\n\t%s:%d\n", frame.Function, frame.File, frame.Line) + if !more { + break + } + } + return stackTraceSB.String() +} diff --git a/vendor/github.com/riverqueue/river/internal/jobstats/job_statistics.go b/vendor/github.com/riverqueue/river/internal/jobstats/job_statistics.go new file mode 100644 index 0000000000..0023886e71 --- /dev/null +++ b/vendor/github.com/riverqueue/river/internal/jobstats/job_statistics.go @@ -0,0 +1,14 @@ +package jobstats + +import "time" + +// JobStatistics contains information about a single execution of a job. +// +// This type has an identical one in the top-level package. The reason for that +// is so that we can use statistics from subpackages, but can reveal all +// public-facing River types in a single public-facing package. +type JobStatistics struct { + CompleteDuration time.Duration // Time it took to set the job completed, discarded, or errored. + QueueWaitDuration time.Duration // Time the job spent waiting in available state before starting execution. + RunDuration time.Duration // Time job spent running (measured around job worker.) +} diff --git a/vendor/github.com/riverqueue/river/internal/leadership/doc.go b/vendor/github.com/riverqueue/river/internal/leadership/doc.go new file mode 100644 index 0000000000..63ff7226dc --- /dev/null +++ b/vendor/github.com/riverqueue/river/internal/leadership/doc.go @@ -0,0 +1,117 @@ +// Package leadership implements leader election for River clients sharing a +// database schema. +// +// The database records at most one current leadership term at a time. The +// elected client runs distributed maintenance work such as queue management, +// job scheduling, and reindexing that should not be duplicated across clients. +// +// # Overview +// +// Leadership is modeled as a database-backed lease with an explicit term +// identity. +// +// A term is identified by: +// - `leader_id`: the stable client identity +// - `elected_at`: the database-issued timestamp for that specific term +// +// The database is authoritative for: +// - which client currently holds the leadership row +// - whether a term can be renewed +// - whether a term has already been replaced +// +// The process uses local time only to bound how long it trusts its last +// successful elect or reelect result. If it cannot renew a term in time, it +// steps down conservatively instead of continuing to act as leader on stale +// information. +// +// # State Model +// +// At a high level, an elector alternates between follower and leader states: +// +// Start +// │ +// ▼ +// ┌─────────────────────────────────────┐ +// │ Follower │ +// │ │ +// │ Retries election on timer or wakeup │ +// └──────────────────┬──────────────────┘ +// │ won election +// ▼ +// ┌─────────────────────────────────────┐ +// │ Leader │ +// │ │ +// │ Renews before trust window expires │ +// └──────────────────┬──────────────────┘ +// │ replaced / expired / resign requested / +// │ renewal failed for too long / shutdown +// ▼ +// Follower +// +// Followers attempt election periodically and can wake early when they learn +// that the previous leader resigned. Leaders renew their current term +// periodically. If renewal fails, the term is replaced, or the local trust +// window expires, the process stops acting as leader and returns to follower +// behavior. +// +// # Trust Window +// +// After each successful election or renewal, the elector computes a local +// trust deadline: +// +// trustedUntil = attemptStarted + TTL - safetyMargin +// +// This trust window has two important properties: +// - it is anchored to when the elect or reelect attempt started, so a slow +// successful database round trip cannot stretch leadership longer than the +// attempt budget allows +// - it ends before the database lease should expire, giving the process time +// to step down before it risks acting on a stale term +// +// The local trust window is a conservative stop condition, not an alternative +// source of truth. A client may step down while the database row is still +// present, but it should not continue acting as leader after it no longer +// trusts its last successful renewal. +// +// # Term-Scoped Operations +// +// Renewing and resigning are scoped to the exact term identified by +// `(leader_id, elected_at)`. +// +// That means: +// - an old term cannot accidentally renew a newer term for the same client +// - a delayed resign from an old term cannot delete a newer term for the +// same client +// - when the database says a term is gone, the elector can step down without +// ambiguity about which term it held +// +// # Notifications and Subscribers +// +// When a notifier is available, the elector listens for leadership-related +// events so followers can wake promptly and leaders can honor explicit +// resignation requests. +// +// Notification delivery is intentionally non-blocking: +// - wakeups may coalesce, because multiple rapid resignations only need to +// prompt another election attempt +// - polling remains the fallback when notifications are unavailable or missed +// +// Consumers inside the process can subscribe to leadership transitions. Those +// subscriptions preserve ordered `true`/`false` transitions so downstream +// maintenance components can reliably start and stop work, while still keeping +// slow subscribers from blocking the elector itself. +// +// # Failure Handling +// +// The system is intentionally conservative under failures: +// - if renewal errors persist until the trust window is exhausted, the leader +// steps down +// - if the database reports that the current term no longer exists, the +// leader steps down immediately +// - if resignation fails during shutdown or after a local timeout, the +// database lease expiry remains the safety net that eventually allows a new +// election +// +// This design keeps leadership decisions centered on the database while using +// local time only to stop trusting stale state sooner rather than later. +package leadership diff --git a/vendor/github.com/riverqueue/river/internal/leadership/elector.go b/vendor/github.com/riverqueue/river/internal/leadership/elector.go new file mode 100644 index 0000000000..c79391a7cf --- /dev/null +++ b/vendor/github.com/riverqueue/river/internal/leadership/elector.go @@ -0,0 +1,808 @@ +package leadership + +import ( + "cmp" + "context" + "encoding/json" + "errors" + "fmt" + "log/slog" + "strings" + "sync" + "time" + + "github.com/riverqueue/river/internal/notifier" + "github.com/riverqueue/river/riverdriver" + "github.com/riverqueue/river/rivershared/baseservice" + "github.com/riverqueue/river/rivershared/startstop" + "github.com/riverqueue/river/rivershared/testsignal" + "github.com/riverqueue/river/rivershared/util/dbutil" + "github.com/riverqueue/river/rivershared/util/randutil" + "github.com/riverqueue/river/rivershared/util/serviceutil" + "github.com/riverqueue/river/rivershared/util/testutil" + "github.com/riverqueue/river/rivershared/util/timeoututil" + "github.com/riverqueue/river/rivertype" +) + +const ( + electIntervalDefault = 5 * time.Second + electIntervalJitterDefault = 1 * time.Second + electIntervalTTLPaddingDefault = 10 * time.Second + leaderLocalDeadlineSafetyMargin = 1 * time.Second +) + +type DBNotification struct { + Action DBNotificationKind `json:"action"` + LeaderID string `json:"leader_id"` +} + +type DBNotificationKind string + +const ( + DBNotificationKindRequestResign DBNotificationKind = "request_resign" + DBNotificationKindResigned DBNotificationKind = "resigned" +) + +type Notification struct { + IsLeader bool + Timestamp time.Time +} + +// Subscription is a client-facing stream of leadership transitions. +// +// The elector publishes every transition (`false -> true -> false -> ...`) to +// each subscription. Delivery is delegated to a subscriptionRelay so the +// elector never blocks on a slow listener. +type Subscription struct { + creationTime time.Time + relay *subscriptionRelay + + unlistenOnce *sync.Once + e *Elector +} + +func (s *Subscription) C() <-chan *Notification { + return s.relay.C() +} + +func (s *Subscription) enqueue(notification *Notification) { + s.relay.enqueue(notification) +} + +func (s *Subscription) stop() { + s.relay.stop() +} + +func (s *Subscription) Unlisten() { + s.unlistenOnce.Do(func() { + s.e.unlisten(s) + }) +} + +// subscriptionRelay decouples elector publication from subscriber consumption. +// +// The elector may need to publish `true` and `false` transitions promptly while +// maintenance components are still busy reacting to the previous one. A plain +// buffered channel would either block the elector or force us to drop +// transitions when the buffer filled. The relay solves that by: +// - appending every Notification to an in-memory FIFO queue +// - waking a dedicated goroutine via `pendingChan` +// - letting that goroutine drain queued notifications into the subscriber's +// public channel `ch` +// +// `pendingChan` is only a wakeup signal. It does not carry the notifications +// themselves, so multiple sends may coalesce while the goroutine is already +// awake. That is safe because the authoritative queue is pendingNotifications. +type subscriptionRelay struct { + ch chan *Notification // public per-subscription delivery channel + done chan struct{} // closes the relay goroutine during Unlisten/stop + pendingChan chan struct{} // coalesced wakeup signal that queued work exists + + // pendingNotifications preserves every leadership transition in order. + // A dedicated goroutine drains it into `ch` so slow subscribers cannot + // block the elector, but consumers like QueueMaintainerLeader still see + // each `false` transition instead of only the latest state. + pendingMu sync.Mutex + pendingNotifications []*Notification +} + +func newSubscriptionRelay() *subscriptionRelay { + relay := &subscriptionRelay{ + ch: make(chan *Notification, 1), + done: make(chan struct{}), + pendingChan: make(chan struct{}, 1), + } + + go relay.run() + + return relay +} + +func (r *subscriptionRelay) C() <-chan *Notification { + return r.ch +} + +// enqueue appends a transition to the pending FIFO, then nudges the relay +// goroutine. The notification argument is the exact transition to preserve in +// order; unlike the notifier wakeup path elsewhere in the elector, these items +// must not be coalesced or replaced. +func (r *subscriptionRelay) enqueue(notification *Notification) { + r.pendingMu.Lock() + r.pendingNotifications = append(r.pendingNotifications, notification) + r.pendingMu.Unlock() + + select { + case r.pendingChan <- struct{}{}: + default: + } +} + +// nextPending pops the next queued notification for the relay goroutine. +func (r *subscriptionRelay) nextPending() (*Notification, bool) { + r.pendingMu.Lock() + defer r.pendingMu.Unlock() + + if len(r.pendingNotifications) == 0 { + return nil, false + } + + notification := r.pendingNotifications[0] + r.pendingNotifications = r.pendingNotifications[1:] + return notification, true +} + +// run waits until queued work exists, then drains as many pending +// notifications as possible into the subscriber channel before sleeping again. +// It exits promptly when stop closes done. +func (r *subscriptionRelay) run() { + for { + select { + case <-r.done: + return + + case <-r.pendingChan: + } + + for { + notification, ok := r.nextPending() + if !ok { + break + } + + select { + case <-r.done: + return + case r.ch <- notification: + } + } + } +} + +// stop terminates the relay goroutine. Callers must ensure they stop enqueueing +// through the owning Subscription afterwards. +func (r *subscriptionRelay) stop() { + close(r.done) +} + +// Test-only properties. +type electorTestSignals struct { + DeniedLeadership testsignal.TestSignal[struct{}] // notifies when elector fails to gain leadership + GainedLeadership testsignal.TestSignal[struct{}] // notifies when elector gains leadership + LostLeadership testsignal.TestSignal[struct{}] // notifies when an elected leader loses leadership + MaintainedLeadership testsignal.TestSignal[struct{}] // notifies when elector maintains leadership + ResignedLeadership testsignal.TestSignal[struct{}] // notifies when elector resigns leadership +} + +func (ts *electorTestSignals) Init(tb testutil.TestingTB) { + ts.DeniedLeadership.Init(tb) + ts.GainedLeadership.Init(tb) + ts.LostLeadership.Init(tb) + ts.MaintainedLeadership.Init(tb) + ts.ResignedLeadership.Init(tb) +} + +type Config struct { + ClientID string + ElectInterval time.Duration // period on which each elector attempts elect even without having received a resignation notification + ElectIntervalJitter time.Duration + Schema string +} + +func (c *Config) mustValidate() *Config { + if c.ClientID == "" { + panic("Config.ClientID must be non-empty") + } + if c.ElectInterval <= 0 { + panic("Config.ElectInterval must be above zero") + } + + return c +} + +type Elector struct { + baseservice.BaseService + startstop.BaseStartStop + + config *Config + exec riverdriver.Executor + notifier *notifier.Notifier + testSignals electorTestSignals + wakeupChan chan struct{} + + mu sync.Mutex + isLeader bool + pendingRequestResign bool + subscriptions []*Subscription +} + +type leadershipTerm struct { + clientID string + electedAt time.Time + trustedUntil time.Time +} + +func newLeadershipTerm(clientID string, electedAt, attemptStarted time.Time, ttl time.Duration) leadershipTerm { + term := leadershipTerm{ + clientID: clientID, + electedAt: electedAt, + } + + trustDuration := ttl - leaderLocalDeadlineSafetyMargin + if trustDuration <= 0 { + term.trustedUntil = attemptStarted + return term + } + + term.trustedUntil = attemptStarted.Add(trustDuration) + return term +} + +func (t leadershipTerm) remaining(now time.Time) time.Duration { + if !t.trustedUntil.After(now) { + return 0 + } + + return t.trustedUntil.Sub(now) +} + +func (t leadershipTerm) reelectAttemptTimeout(now time.Time) time.Duration { + remainingDuration := t.remaining(now) + if remainingDuration <= 0 { + return 0 + } + if remainingDuration < deadlineTimeout { + return remainingDuration + } + + return deadlineTimeout +} + +// NewElector returns an Elector using the given adapter. The name should correspond +// to the name of the database + schema combo and should be shared across all Clients +// running with that combination. The id should be unique to the Client. +func NewElector(archetype *baseservice.Archetype, exec riverdriver.Executor, notifier *notifier.Notifier, config *Config) *Elector { + return baseservice.Init(archetype, &Elector{ + config: (&Config{ + ClientID: config.ClientID, + ElectInterval: cmp.Or(config.ElectInterval, electIntervalDefault), + ElectIntervalJitter: cmp.Or(config.ElectIntervalJitter, electIntervalJitterDefault), + Schema: config.Schema, + }).mustValidate(), + exec: exec, + notifier: notifier, + }) +} + +func trySendWakeup(ctx context.Context, wakeupChan chan struct{}) { + if ctx.Err() != nil { + return + } + + select { + case <-ctx.Done(): + case wakeupChan <- struct{}{}: + default: + } +} + +func (e *Elector) Start(ctx context.Context) error { + ctx, shouldStart, started, stopped := e.StartInit(ctx) + if !shouldStart { + return nil + } + + // Buffered to 1 so notifications coalesce instead of blocking the elector. + e.wakeupChan = make(chan struct{}, 1) + + var sub *notifier.Subscription + if e.notifier == nil { + e.Logger.DebugContext(ctx, e.Name+": No notifier configured; starting in poll mode", "client_id", e.config.ClientID) + } else { + e.Logger.DebugContext(ctx, e.Name+": Listening for leadership changes", "client_id", e.config.ClientID, "topic", notifier.NotificationTopicLeadership) + var err error + sub, err = e.notifier.Listen(ctx, notifier.NotificationTopicLeadership, func(topic notifier.NotificationTopic, payload string) { + e.handleLeadershipNotification(ctx, topic, payload) + }) + if err != nil { + stopped() + if strings.HasSuffix(err.Error(), "conn closed") || errors.Is(err, context.Canceled) { + return nil + } + return err + } + } + + go func() { + started() + defer stopped() // this defer should come first so it's last out + + e.Logger.DebugContext(ctx, e.Name+": Run loop started") + defer e.Logger.DebugContext(ctx, e.Name+": Run loop stopped") + + if sub != nil { + defer sub.Unlisten(ctx) + } + + for { + term, err := e.runFollowerState(ctx) + if err != nil { + // Function above only returns an error if context was cancelled + // or overall context is done. + if !errors.Is(err, context.Canceled) && ctx.Err() == nil { + panic(err) + } + return + } + + e.publishLeadershipState(true) + e.Logger.DebugContext(ctx, e.Name+": Gained leadership", "client_id", e.config.ClientID) + e.testSignals.GainedLeadership.Signal(struct{}{}) + + err = e.runLeaderState(ctx, term) + if err != nil { + if errors.Is(err, context.Canceled) { + return + } + + e.Logger.ErrorContext(ctx, e.Name+": Error keeping leadership", "client_id", e.config.ClientID, "err", err) + } + } + }() + + return nil +} + +// runFollowerState is the follower side of the elector state machine. It keeps +// attempting election until this client becomes leader or the elector stops. +func (e *Elector) runFollowerState(ctx context.Context) (leadershipTerm, error) { + var attempt int + for { + attempt++ + e.Logger.DebugContext(ctx, e.Name+": Attempting to gain leadership", "client_id", e.config.ClientID) + // Use the local monotonic-bearing clock for the trust window. The + // DB-facing timestamp path stays on NowUTCOrNil below. + attemptStarted := e.Time.Now() + + leader, err := attemptElect(ctx, e.exec, &riverdriver.LeaderElectParams{ + LeaderID: e.config.ClientID, + Now: e.Time.NowOrNil(), + Schema: e.config.Schema, + TTL: e.leaderTTL(), + }) + if err != nil { + if errors.Is(err, context.Canceled) || ctx.Err() != nil { + return leadershipTerm{}, err + } + if !errors.Is(err, rivertype.ErrNotFound) { + sleepDuration := serviceutil.ExponentialBackoff(attempt, serviceutil.MaxAttemptsBeforeResetDefault) + e.Logger.ErrorContext(ctx, e.Name+": Error attempting to elect", e.errorSlogArgs(err, attempt, sleepDuration)...) + serviceutil.CancellableSleep(ctx, sleepDuration) + continue + } + } + + if leader != nil { + return newLeadershipTerm(leader.LeaderID, leader.ElectedAt, attemptStarted, e.leaderTTL()), nil + } + + attempt = 0 + + e.Logger.DebugContext(ctx, e.Name+": Leadership bid was unsuccessful (not an error)", "client_id", e.config.ClientID) + e.testSignals.DeniedLeadership.Signal(struct{}{}) + + select { + case <-serviceutil.CancellableSleepC(ctx, randutil.DurationBetween(e.config.ElectInterval, e.config.ElectInterval+e.config.ElectIntervalJitter)): + if ctx.Err() != nil { // context done + return leadershipTerm{}, ctx.Err() + } + + case <-e.wakeupChan: + // Somebody just resigned, try to win the next election after a very + // short random interval (to prevent all clients from bidding at once). + serviceutil.CancellableSleep(ctx, randutil.DurationBetween(0, 50*time.Millisecond)) + } + } +} + +// Handles a leadership notification from the notifier. +func (e *Elector) handleLeadershipNotification(ctx context.Context, topic notifier.NotificationTopic, payload string) { + if topic != notifier.NotificationTopicLeadership { + // This should not happen unless the notifier is broken. + e.Logger.ErrorContext(ctx, e.Name+": Received unexpected notification", "client_id", e.config.ClientID, "topic", topic, "payload", payload) + return + } + + notification := DBNotification{} + if err := json.Unmarshal([]byte(payload), ¬ification); err != nil { + e.Logger.ErrorContext(ctx, e.Name+": Unable to unmarshal leadership notification", "client_id", e.config.ClientID, "err", err) + return + } + + e.Logger.DebugContext(ctx, e.Name+": Received notification from notifier", "action", notification.Action, "client_id", e.config.ClientID) + + // Do an initial context check so in case context is done, it always takes + // precedence over sending a leadership notification. + if ctx.Err() != nil { + return + } + + switch notification.Action { + case DBNotificationKindRequestResign: + if !e.markPendingRequestResign() { + return + } + + trySendWakeup(ctx, e.wakeupChan) + case DBNotificationKindResigned: + // If this a resignation from _this_ client, ignore the change. + if notification.LeaderID == e.config.ClientID { + return + } + + trySendWakeup(ctx, e.wakeupChan) + } +} + +// runLeaderState is the leader side of the elector state machine. It waits for +// either a reelection interval, a forced resignation, or shutdown. +func (e *Elector) runLeaderState(ctx context.Context, term leadershipTerm) error { + defer e.clearPendingRequestResign() + defer e.publishLeadershipState(false) + + shouldResign := true + + // Before the elector returns, run a delete with NOTIFY to give up any + // leadership that we have. If we do that here, we guarantee that any locks + // we have will be released (even if they were acquired in + // attemptGainLeadership but we didn't wait for the response) + // + // This doesn't use ctx because it runs *after* the ctx is done. + defer func() { + if shouldResign { + e.attemptResignLoop(ctx, term) // will resign using WithoutCancel context, but ctx sent for logging + } + }() + + timer := time.NewTimer(0) + defer timer.Stop() + + numErrors := 0 + waitDuration := e.config.ElectInterval + + for { + resetTimer(timer, waitDuration) + + select { + case <-ctx.Done(): + return ctx.Err() + + case <-e.wakeupChan: + if !e.takePendingRequestResign() { + continue + } + + e.Logger.InfoContext(ctx, e.Name+": Current leader received forced resignation", "client_id", e.config.ClientID) + + // This client may win leadership again, but drop out of this + // function and make it start all over. + return nil + + case <-timer.C: + // Reelect timer expired; attempt reelection below. + } + + e.Logger.DebugContext(ctx, e.Name+": Current leader attempting reelect", "client_id", e.config.ClientID) + + // Use the local monotonic-bearing clock for the trust window. The + // DB-facing timestamp path stays on NowOrNil below. + attemptStarted := e.Time.Now() + attemptTimeout := term.reelectAttemptTimeout(attemptStarted) + if attemptTimeout <= 0 { + e.Logger.WarnContext(ctx, e.Name+": Current leader stepping down because the reelection deadline elapsed", "client_id", e.config.ClientID) + e.testSignals.LostLeadership.Signal(struct{}{}) + return nil + } + + leader, err := attemptReelectWithTimeout(ctx, e.exec, &riverdriver.LeaderReelectParams{ + ElectedAt: term.electedAt, + LeaderID: term.clientID, + Now: e.Time.NowOrNil(), + Schema: e.config.Schema, + TTL: e.leaderTTL(), + }, attemptTimeout) + if err != nil { + if errors.Is(err, context.Canceled) { + return err + } + if errors.Is(err, rivertype.ErrNotFound) { + shouldResign = false + e.testSignals.LostLeadership.Signal(struct{}{}) + return nil + } + + numErrors++ + sleepDuration := serviceutil.ExponentialBackoff(numErrors, 3) + remainingDuration := term.remaining(e.Time.Now()) + if remainingDuration <= 0 { + e.Logger.WarnContext(ctx, e.Name+": Current leader stepping down because the reelection deadline elapsed after an error", "client_id", e.config.ClientID) + e.testSignals.LostLeadership.Signal(struct{}{}) + return nil + } + + e.Logger.ErrorContext(ctx, e.Name+": Error attempting reelection", e.errorSlogArgs(err, numErrors, sleepDuration)...) + if remainingDuration < sleepDuration { + sleepDuration = remainingDuration + } + serviceutil.CancellableSleep(ctx, sleepDuration) + if ctx.Err() != nil { + return ctx.Err() + } + + // Retry immediately after the backoff because the time budget for this + // lease has already been reduced by the failed attempt above. + waitDuration = 0 + continue + } + + numErrors = 0 + term = newLeadershipTerm(leader.LeaderID, leader.ElectedAt, attemptStarted, e.leaderTTL()) + e.testSignals.MaintainedLeadership.Signal(struct{}{}) + waitDuration = e.config.ElectInterval + } +} + +// Try up to 3 times to give up any currently held leadership. +// +// The context received is used for logging purposes, but the function actually +// makes use of a background context to try and guarantee that leadership is +// always surrendered in a timely manner so it can be picked up quickly by +// another client, even in the event of a cancellation. +func (e *Elector) attemptResignLoop(ctx context.Context, term leadershipTerm) { + e.Logger.DebugContext(ctx, e.Name+": Attempting to resign leadership", "client_id", e.config.ClientID) + + // Make a good faith attempt to resign, even in the presence of errors, but + // don't keep hammering if it doesn't work. In case a resignation failure, + // leader TTLs will act as an additional hedge to ensure a new leader can + // still be elected. + const maxNumErrors = 3 + + // This does not inherit the parent context's cancellation because we want to + // give up leadership even during a shutdown. There is no way to short-circuit + // this, though there are timeouts per call within attemptResign. + ctx = context.WithoutCancel(ctx) + + for attempt := 1; attempt <= maxNumErrors; attempt++ { + if err := e.attemptResign(ctx, attempt, term); err != nil { + sleepDuration := serviceutil.ExponentialBackoff(attempt, maxNumErrors) + e.Logger.ErrorContext(ctx, e.Name+": Error attempting to resign", e.errorSlogArgs(err, attempt, sleepDuration)...) + serviceutil.CancellableSleep(ctx, sleepDuration) + + continue + } + + return + } +} + +// attemptResign attempts to resign any currently held leaderships for the +// elector's name and leader ID. +func (e *Elector) attemptResign(ctx context.Context, attempt int, term leadershipTerm) error { + // Wait one second longer each time we try to resign: + timeout := time.Duration(attempt) * time.Second + + return timeoututil.WithTimeout(ctx, timeout, e.Name+".attemptResign", func(ctx context.Context) error { + resigned, err := e.exec.LeaderResign(ctx, &riverdriver.LeaderResignParams{ + ElectedAt: term.electedAt, + LeaderID: term.clientID, + LeadershipTopic: string(notifier.NotificationTopicLeadership), + Schema: e.config.Schema, + }) + if err != nil { + return err + } + + if resigned { + e.Logger.DebugContext(ctx, e.Name+": Resigned leadership successfully", "client_id", e.config.ClientID) + e.testSignals.ResignedLeadership.Signal(struct{}{}) + } + + return nil + }) +} + +// Produces a common set of key/value pairs for logging when an error occurs. +// +// Refactored out because we had three repeats of identical information in this +// file, but if it causes things to get messy, may want to refactor again. +func (e *Elector) errorSlogArgs(err error, attempt int, sleepDuration time.Duration) []any { + return []any{ + slog.Int("attempt", attempt), + slog.String("client_id", e.config.ClientID), + slog.String("err", err.Error()), + slog.String("sleep_duration", sleepDuration.String()), + } +} + +func (e *Elector) Listen() *Subscription { + sub := &Subscription{ + creationTime: time.Now().UTC(), + e: e, + relay: newSubscriptionRelay(), + unlistenOnce: &sync.Once{}, + } + + e.mu.Lock() + defer e.mu.Unlock() + + initialNotification := &Notification{ + IsLeader: e.isLeader, + Timestamp: sub.creationTime, + } + sub.enqueue(initialNotification) + + e.subscriptions = append(e.subscriptions, sub) + return sub +} + +func (e *Elector) unlisten(sub *Subscription) { + success := e.tryUnlisten(sub) + if !success { + panic("BUG: tried to unlisten for subscription not in list") + } + + sub.stop() +} + +// needs to be in a separate method so the defer will cleanly unlock the mutex, +// even if we panic. +func (e *Elector) tryUnlisten(sub *Subscription) bool { + e.mu.Lock() + defer e.mu.Unlock() + + for i, s := range e.subscriptions { + if s == sub { + e.subscriptions = append(e.subscriptions[:i], e.subscriptions[i+1:]...) + return true + } + } + return false +} + +// leaderTTL is at least the reelect run interval used by clients to try and gain +// leadership or reelect themselves as leader, plus a little padding to account +// to give the leader a little breathing room in its reelection loop. +func (e *Elector) leaderTTL() time.Duration { + return e.config.ElectInterval + electIntervalTTLPaddingDefault +} + +func (e *Elector) markPendingRequestResign() bool { + e.mu.Lock() + defer e.mu.Unlock() + + if !e.isLeader { + return false + } + + e.pendingRequestResign = true + return true +} + +func (e *Elector) publishLeadershipState(isLeader bool) { + notifyTime := time.Now().UTC() + e.mu.Lock() + defer e.mu.Unlock() + + e.isLeader = isLeader + if !isLeader { + e.pendingRequestResign = false + } + + notification := &Notification{ + IsLeader: isLeader, + Timestamp: notifyTime, + } + + for _, s := range e.subscriptions { + s.enqueue(notification) + } +} + +func (e *Elector) clearPendingRequestResign() { + e.mu.Lock() + defer e.mu.Unlock() + + e.pendingRequestResign = false +} + +func (e *Elector) takePendingRequestResign() bool { + e.mu.Lock() + defer e.mu.Unlock() + + if !e.pendingRequestResign { + return false + } + + e.pendingRequestResign = false + return true +} + +const deadlineTimeout = 5 * time.Second + +func resetTimer(timer *time.Timer, duration time.Duration) { + if !timer.Stop() { + select { + case <-timer.C: + default: + } + } + + timer.Reset(duration) +} + +// attemptElect attempts to elect a leader for the given name. If there is no +// current leader or the previous leader expired, the provided leader ID is set +// as the new leader with a TTL of `params.TTL`. +func attemptElect(ctx context.Context, exec riverdriver.Executor, params *riverdriver.LeaderElectParams) (*riverdriver.Leader, error) { + return attemptElectWithTimeout(ctx, exec, params, deadlineTimeout) +} + +func attemptElectWithTimeout(ctx context.Context, exec riverdriver.Executor, params *riverdriver.LeaderElectParams, timeout time.Duration) (*riverdriver.Leader, error) { + return timeoututil.WithTimeoutV(ctx, timeout, "leadership.attemptElect", func(ctx context.Context) (*riverdriver.Leader, error) { + execTx, err := exec.Begin(ctx) + if err != nil { + var additionalDetail string + if errors.Is(err, context.DeadlineExceeded) { + additionalDetail = " (a common cause of this is a database pool that's at its connection limit; you may need to increase maximum connections)" + } + + return nil, fmt.Errorf("error beginning transaction: %w%s", err, additionalDetail) + } + defer dbutil.RollbackWithoutCancel(ctx, execTx) + + if _, err := execTx.LeaderDeleteExpired(ctx, &riverdriver.LeaderDeleteExpiredParams{ + Now: params.Now, + Schema: params.Schema, + }); err != nil { + return nil, err + } + + leader, err := execTx.LeaderAttemptElect(ctx, params) + if err != nil && !errors.Is(err, rivertype.ErrNotFound) { + return nil, err + } + if err := execTx.Commit(ctx); err != nil { + return nil, fmt.Errorf("error committing transaction: %w", err) + } + if err != nil { + return nil, err + } + + return leader, nil + }) +} + +func attemptReelectWithTimeout(ctx context.Context, exec riverdriver.Executor, params *riverdriver.LeaderReelectParams, timeout time.Duration) (*riverdriver.Leader, error) { + return timeoututil.WithTimeoutV(ctx, timeout, "leadership.attemptReelect", func(ctx context.Context) (*riverdriver.Leader, error) { + return exec.LeaderAttemptReelect(ctx, params) + }) +} diff --git a/vendor/github.com/riverqueue/river/internal/maintenance/job_cleaner.go b/vendor/github.com/riverqueue/river/internal/maintenance/job_cleaner.go new file mode 100644 index 0000000000..e70fe30041 --- /dev/null +++ b/vendor/github.com/riverqueue/river/internal/maintenance/job_cleaner.go @@ -0,0 +1,239 @@ +package maintenance + +import ( + "cmp" + "context" + "errors" + "fmt" + "log/slog" + "time" + + "github.com/riverqueue/river/riverdriver" + "github.com/riverqueue/river/rivershared/baseservice" + "github.com/riverqueue/river/rivershared/circuitbreaker" + "github.com/riverqueue/river/rivershared/riversharedmaintenance" + "github.com/riverqueue/river/rivershared/startstop" + "github.com/riverqueue/river/rivershared/testsignal" + "github.com/riverqueue/river/rivershared/util/randutil" + "github.com/riverqueue/river/rivershared/util/serviceutil" + "github.com/riverqueue/river/rivershared/util/testutil" + "github.com/riverqueue/river/rivershared/util/timeoututil" + "github.com/riverqueue/river/rivershared/util/timeutil" +) + +// JobCleanerTestSignals are internal signals used exclusively in tests. +type JobCleanerTestSignals struct { + DeletedBatch testsignal.TestSignal[struct{}] // notifies when runOnce finishes a pass +} + +func (ts *JobCleanerTestSignals) Init(tb testutil.TestingTB) { + ts.DeletedBatch.Init(tb) +} + +type JobCleanerConfig struct { + riversharedmaintenance.BatchSizes + + // CancelledJobRetentionPeriod is the amount of time to keep cancelled jobs + // around before they're removed permanently. + // + // The special value -1 disables deletion of cancelled jobs. + CancelledJobRetentionPeriod time.Duration + + // CompletedJobRetentionPeriod is the amount of time to keep completed jobs + // around before they're removed permanently. + // + // The special value -1 disables deletion of completed jobs. + CompletedJobRetentionPeriod time.Duration + + // DiscardedJobRetentionPeriod is the amount of time to keep cancelled jobs + // around before they're removed permanently. + // + // The special value -1 disables deletion of discarded jobs. + DiscardedJobRetentionPeriod time.Duration + + // Interval is the amount of time to wait between runs of the cleaner. + Interval time.Duration + + // QueuesExcluded are queues that'll be excluded from cleaning. + QueuesExcluded []string + + // Schema where River tables are located. Empty string omits schema, causing + // Postgres to default to `search_path`. + Schema string + + // Timeout of the individual queries in the job cleaner. + Timeout time.Duration +} + +func (c *JobCleanerConfig) mustValidate() *JobCleanerConfig { + c.MustValidate() + + if c.CancelledJobRetentionPeriod < -1 { + panic("JobCleanerConfig.CancelledJobRetentionPeriod must be above zero") + } + if c.CompletedJobRetentionPeriod < -1 { + panic("JobCleanerConfig.CompletedJobRetentionPeriod must be above zero") + } + if c.DiscardedJobRetentionPeriod < -1 { + panic("JobCleanerConfig.DiscardedJobRetentionPeriod must be above zero") + } + if c.Interval <= 0 { + panic("JobCleanerConfig.Interval must be above zero") + } + if c.Timeout <= 0 { + panic("JobCleanerConfig.Timeout must be above zero") + } + + return c +} + +// JobCleaner periodically removes finalized jobs that are cancelled, completed, +// or discarded. Each state's retention time can be configured individually. +type JobCleaner struct { + riversharedmaintenance.QueueMaintainerServiceBase + startstop.BaseStartStop + + // exported for test purposes + Config *JobCleanerConfig + TestSignals JobCleanerTestSignals + + exec riverdriver.Executor + + // Circuit breaker that tracks consecutive timeout failures from the central + // query. The query starts by using the full/default batch size, but after + // this breaker trips (after N consecutive timeouts occur in a row), it + // switches to a smaller batch. We assume that a database that's degraded is + // likely to stay degraded over a longer term, so after the circuit breaks, + // it stays broken until the program is restarted. + reducedBatchSizeBreaker *circuitbreaker.CircuitBreaker +} + +func NewJobCleaner(archetype *baseservice.Archetype, config *JobCleanerConfig, exec riverdriver.Executor) *JobCleaner { + batchSizes := config.WithDefaults() + + return baseservice.Init(archetype, &JobCleaner{ + Config: (&JobCleanerConfig{ + BatchSizes: batchSizes, + CancelledJobRetentionPeriod: cmp.Or(config.CancelledJobRetentionPeriod, riversharedmaintenance.CancelledJobRetentionPeriodDefault), + CompletedJobRetentionPeriod: cmp.Or(config.CompletedJobRetentionPeriod, riversharedmaintenance.CompletedJobRetentionPeriodDefault), + DiscardedJobRetentionPeriod: cmp.Or(config.DiscardedJobRetentionPeriod, riversharedmaintenance.DiscardedJobRetentionPeriodDefault), + QueuesExcluded: config.QueuesExcluded, + Interval: cmp.Or(config.Interval, riversharedmaintenance.JobCleanerIntervalDefault), + Schema: config.Schema, + Timeout: cmp.Or(config.Timeout, riversharedmaintenance.JobCleanerTimeoutDefault), + }).mustValidate(), + exec: exec, + reducedBatchSizeBreaker: riversharedmaintenance.ReducedBatchSizeBreaker(batchSizes), + }) +} + +func (s *JobCleaner) Start(ctx context.Context) error { //nolint:dupl + ctx, shouldStart, started, stopped := s.StartInit(ctx) + if !shouldStart { + return nil + } + + s.StaggerStart(ctx) + + go func() { + started() + defer stopped() // this defer should come first so it's last out + + s.Logger.DebugContext(ctx, s.Name+riversharedmaintenance.LogPrefixRunLoopStarted) + defer s.Logger.DebugContext(ctx, s.Name+riversharedmaintenance.LogPrefixRunLoopStopped) + + ticker := timeutil.NewTickerWithInitialTick(ctx, s.Config.Interval) + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + } + + res, err := s.runOnce(ctx) + if err != nil { + if !errors.Is(err, context.Canceled) { + s.Logger.ErrorContext(ctx, s.Name+": Error cleaning jobs", slog.String("error", err.Error())) + } + continue + } + + if res.NumJobsDeleted > 0 { + s.Logger.InfoContext(ctx, s.Name+riversharedmaintenance.LogPrefixRanSuccessfully, + slog.Int("num_jobs_deleted", res.NumJobsDeleted), + ) + } + } + }() + + return nil +} + +func (s *JobCleaner) batchSize() int { + if s.reducedBatchSizeBreaker.Open() { + return s.Config.Reduced + } + return s.Config.Default +} + +type jobCleanerRunOnceResult struct { + NumJobsDeleted int +} + +func (s *JobCleaner) runOnce(ctx context.Context) (*jobCleanerRunOnceResult, error) { + res := &jobCleanerRunOnceResult{} + + for { + numDeleted, err := timeoututil.WithTimeoutV(ctx, s.Config.Timeout, s.Name+".runOnce", func(ctx context.Context) (int, error) { + // In the special case that all retentions are indefinite, don't + // bother issuing the query at all as an optimization. + if s.Config.CompletedJobRetentionPeriod == -1 && + s.Config.CancelledJobRetentionPeriod == -1 && + s.Config.DiscardedJobRetentionPeriod == -1 { + return 0, nil + } + + numDeleted, err := s.exec.JobDeleteBefore(ctx, &riverdriver.JobDeleteBeforeParams{ + CancelledDoDelete: s.Config.CancelledJobRetentionPeriod != -1, + CancelledFinalizedAtHorizon: time.Now().Add(-s.Config.CancelledJobRetentionPeriod), + CompletedDoDelete: s.Config.CompletedJobRetentionPeriod != -1, + CompletedFinalizedAtHorizon: time.Now().Add(-s.Config.CompletedJobRetentionPeriod), + DiscardedDoDelete: s.Config.DiscardedJobRetentionPeriod != -1, + DiscardedFinalizedAtHorizon: time.Now().Add(-s.Config.DiscardedJobRetentionPeriod), + Max: s.batchSize(), + QueuesExcluded: s.Config.QueuesExcluded, + Schema: s.Config.Schema, + }) + if err != nil { + return 0, fmt.Errorf("error cleaning jobs: %w", err) + } + + s.reducedBatchSizeBreaker.ResetIfNotOpen() + + return numDeleted, nil + }) + if err != nil { + if errors.Is(err, context.DeadlineExceeded) { + s.reducedBatchSizeBreaker.Trip() + } + + return nil, err + } + + s.TestSignals.DeletedBatch.Signal(struct{}{}) + + res.NumJobsDeleted += numDeleted + // Deleted was less than query `LIMIT` which means work is done. + if numDeleted < s.batchSize() { + break + } + + s.Logger.DebugContext(ctx, s.Name+": Deleted batch of jobs", + slog.Int("num_jobs_deleted", numDeleted), + ) + + serviceutil.CancellableSleep(ctx, randutil.DurationBetween(riversharedmaintenance.BatchBackoffMin, riversharedmaintenance.BatchBackoffMax)) + } + + return res, nil +} diff --git a/vendor/github.com/riverqueue/river/internal/maintenance/job_rescuer.go b/vendor/github.com/riverqueue/river/internal/maintenance/job_rescuer.go new file mode 100644 index 0000000000..f35ae51acf --- /dev/null +++ b/vendor/github.com/riverqueue/river/internal/maintenance/job_rescuer.go @@ -0,0 +1,383 @@ +package maintenance + +import ( + "cmp" + "context" + "encoding/json" + "errors" + "fmt" + "log/slog" + "time" + + "github.com/riverqueue/river/internal/jobexecutor" + "github.com/riverqueue/river/internal/workunit" + "github.com/riverqueue/river/riverdriver" + "github.com/riverqueue/river/rivershared/baseservice" + "github.com/riverqueue/river/rivershared/circuitbreaker" + "github.com/riverqueue/river/rivershared/riverpilot" + "github.com/riverqueue/river/rivershared/riversharedmaintenance" + "github.com/riverqueue/river/rivershared/startstop" + "github.com/riverqueue/river/rivershared/testsignal" + "github.com/riverqueue/river/rivershared/util/randutil" + "github.com/riverqueue/river/rivershared/util/serviceutil" + "github.com/riverqueue/river/rivershared/util/testutil" + "github.com/riverqueue/river/rivershared/util/timeoututil" + "github.com/riverqueue/river/rivershared/util/timeutil" + "github.com/riverqueue/river/rivertype" +) + +const ( + JobRescuerRescueAfterDefault = time.Hour + JobRescuerIntervalDefault = 30 * time.Second +) + +// JobRescuerTestSignals are internal signals used exclusively in tests. +type JobRescuerTestSignals struct { + FetchedBatch testsignal.TestSignal[struct{}] // notifies when runOnce has fetched a batch of jobs + UpdatedBatch testsignal.TestSignal[struct{}] // notifies when runOnce has updated rescued jobs from a batch +} + +func (ts *JobRescuerTestSignals) Init(tb testutil.TestingTB) { + ts.FetchedBatch.Init(tb) + ts.UpdatedBatch.Init(tb) +} + +type JobRescuerConfig struct { + riversharedmaintenance.BatchSizes + + // ClientJobTimeout is the default job timeout used when a worker returns a + // timeout of zero. + ClientJobTimeout time.Duration + + // ClientRetryPolicy is the default retry policy to use for workers that don't + // override NextRetry. + ClientRetryPolicy jobexecutor.ClientRetryPolicy + + // Interval is the amount of time to wait between runs of the rescuer. + Interval time.Duration + + // Pilot controls driver-level behavior that can be customized by plugins. + Pilot riverpilot.Pilot + + // RescueAfter is the amount of time for a job to be active before it is + // considered stuck and should be rescued. + RescueAfter time.Duration + + // Schema where River tables are located. Empty string omits schema, causing + // Postgres to default to `search_path`. + Schema string + + WorkUnitFactoryFunc func(kind string) workunit.WorkUnitFactory +} + +func (c *JobRescuerConfig) mustValidate() *JobRescuerConfig { + c.MustValidate() + + if c.ClientRetryPolicy == nil { + panic("RescuerConfig.ClientRetryPolicy must be set") + } + if c.Interval <= 0 { + panic("RescuerConfig.Interval must be above zero") + } + if c.Pilot == nil { + panic("RescuerConfig.Pilot must be set") + } + if c.RescueAfter <= 0 { + panic("RescuerConfig.JobDuration must be above zero") + } + if c.WorkUnitFactoryFunc == nil { + panic("RescuerConfig.WorkUnitFactoryFunc must be set") + } + + return c +} + +// JobRescuer periodically rescues jobs that have been executing for too long +// and are considered to be "stuck". +type JobRescuer struct { + riversharedmaintenance.QueueMaintainerServiceBase + startstop.BaseStartStop + + // exported for test purposes + Config *JobRescuerConfig + TestSignals JobRescuerTestSignals + + exec riverdriver.Executor + + // Circuit breaker that tracks consecutive timeout failures from the central + // query. The query starts by using the full/default batch size, but after + // this breaker trips (after N consecutive timeouts occur in a row), it + // switches to a smaller batch. We assume that a database that's degraded is + // likely to stay degraded over a longer term, so after the circuit breaks, + // it stays broken until the program is restarted. + reducedBatchSizeBreaker *circuitbreaker.CircuitBreaker +} + +func NewRescuer(archetype *baseservice.Archetype, config *JobRescuerConfig, exec riverdriver.Executor) *JobRescuer { + batchSizes := config.WithDefaults() + pilot := config.Pilot + if pilot == nil { + pilot = &riverpilot.StandardPilot{} + } + + return baseservice.Init(archetype, &JobRescuer{ + Config: (&JobRescuerConfig{ + BatchSizes: batchSizes, + ClientJobTimeout: config.ClientJobTimeout, + ClientRetryPolicy: config.ClientRetryPolicy, + Interval: cmp.Or(config.Interval, JobRescuerIntervalDefault), + Pilot: pilot, + RescueAfter: cmp.Or(config.RescueAfter, JobRescuerRescueAfterDefault), + Schema: config.Schema, + WorkUnitFactoryFunc: config.WorkUnitFactoryFunc, + }).mustValidate(), + exec: exec, + reducedBatchSizeBreaker: riversharedmaintenance.ReducedBatchSizeBreaker(batchSizes), + }) +} + +func (s *JobRescuer) Start(ctx context.Context) error { + ctx, shouldStart, started, stopped := s.StartInit(ctx) + if !shouldStart { + return nil + } + + s.StaggerStart(ctx) + + go func() { + started() + defer stopped() // this defer should come first so it's last out + + s.Logger.DebugContext(ctx, s.Name+riversharedmaintenance.LogPrefixRunLoopStarted) + defer s.Logger.DebugContext(ctx, s.Name+riversharedmaintenance.LogPrefixRunLoopStopped) + + ticker := timeutil.NewTickerWithInitialTick(ctx, s.Config.Interval) + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + } + + res, err := s.runOnce(ctx) + if err != nil { + if !errors.Is(err, context.Canceled) { + s.Logger.ErrorContext(ctx, s.Name+": Error rescuing jobs", slog.String("error", err.Error())) + } + continue + } + + if res.NumJobsDiscarded > 0 || res.NumJobsRetried > 0 { + s.Logger.InfoContext(ctx, s.Name+riversharedmaintenance.LogPrefixRanSuccessfully, + slog.Int64("num_jobs_discarded", res.NumJobsDiscarded), + slog.Int64("num_jobs_retry_scheduled", res.NumJobsRetried), + ) + } + } + }() + + return nil +} + +func (s *JobRescuer) batchSize() int { + if s.reducedBatchSizeBreaker.Open() { + return s.Config.Reduced + } + return s.Config.Default +} + +type rescuerRunOnceResult struct { + NumJobsCancelled int64 + NumJobsDiscarded int64 + NumJobsRetried int64 +} + +type metadataWithCancelAttemptedAt struct { + CancelAttemptedAt time.Time `json:"cancel_attempted_at"` +} + +func (s *JobRescuer) runOnce(ctx context.Context) (*rescuerRunOnceResult, error) { + var afterID int64 + + res := &rescuerRunOnceResult{} + stuckHorizon := time.Now().Add(-s.Config.RescueAfter) + + for { + batchSize := s.batchSize() + stuckJobs, err := s.getStuckJobs(ctx, afterID, batchSize, stuckHorizon) + if err != nil { + if errors.Is(err, context.DeadlineExceeded) { + s.reducedBatchSizeBreaker.Trip() + } + + return nil, fmt.Errorf("error fetching stuck jobs: %w", err) + } + + s.reducedBatchSizeBreaker.ResetIfNotOpen() + + s.TestSignals.FetchedBatch.Signal(struct{}{}) + if len(stuckJobs) > 0 { + afterID = stuckJobs[len(stuckJobs)-1].ID + } + + now := time.Now().UTC() + + rescueManyParams := riverdriver.JobRescueManyParams{ + ID: make([]int64, 0, len(stuckJobs)), + Error: make([][]byte, 0, len(stuckJobs)), + FinalizedAt: make([]*time.Time, 0, len(stuckJobs)), + ScheduledAt: make([]time.Time, 0, len(stuckJobs)), + Schema: s.Config.Schema, + State: make([]string, 0, len(stuckJobs)), + StuckHorizon: stuckHorizon, + } + + for _, job := range stuckJobs { + var metadata metadataWithCancelAttemptedAt + if err := json.Unmarshal(job.Metadata, &metadata); err != nil { + return nil, fmt.Errorf("error unmarshaling job metadata: %w", err) + } + + errorData, err := json.Marshal(rivertype.AttemptError{ + At: now, + Attempt: max(job.Attempt, 0), + Error: "Stuck job rescued by JobRescuer", + Trace: "", + }) + if err != nil { + return nil, fmt.Errorf("error marshaling error JSON: %w", err) + } + + addRescueParam := func(state rivertype.JobState, finalizedAt *time.Time, scheduledAt time.Time) { + rescueManyParams.ID = append(rescueManyParams.ID, job.ID) + rescueManyParams.Error = append(rescueManyParams.Error, errorData) + rescueManyParams.FinalizedAt = append(rescueManyParams.FinalizedAt, finalizedAt) + rescueManyParams.ScheduledAt = append(rescueManyParams.ScheduledAt, scheduledAt) + rescueManyParams.State = append(rescueManyParams.State, string(state)) + } + + if !metadata.CancelAttemptedAt.IsZero() { + res.NumJobsCancelled++ + addRescueParam(rivertype.JobStateCancelled, &now, job.ScheduledAt) // reused previous scheduled value + continue + } + + retryDecision, retryAt := s.makeRetryDecision(ctx, job, now) + + switch retryDecision { + case jobRetryDecisionDiscard: + res.NumJobsDiscarded++ + addRescueParam(rivertype.JobStateDiscarded, &now, job.ScheduledAt) // reused previous scheduled value + + case jobRetryDecisionIgnore: + // job not timed out yet due to kind-specific timeout value; ignore + + case jobRetryDecisionRetry: + res.NumJobsRetried++ + addRescueParam(rivertype.JobStateRetryable, nil, retryAt) + } + } + + if len(rescueManyParams.ID) > 0 { + _, err = s.rescueMany(ctx, &rescueManyParams) + if err != nil { + return nil, fmt.Errorf("error rescuing stuck jobs: %w", err) + } + } + + s.TestSignals.UpdatedBatch.Signal(struct{}{}) + + // Number of rows fetched was less than query `LIMIT` which means work is + // done for this round: + if len(stuckJobs) < batchSize { + break + } + + serviceutil.CancellableSleep(ctx, randutil.DurationBetween(riversharedmaintenance.BatchBackoffMin, riversharedmaintenance.BatchBackoffMax)) + } + + return res, nil +} + +func (s *JobRescuer) getStuckJobs(ctx context.Context, afterID int64, batchSize int, stuckHorizon time.Time) ([]*rivertype.JobRow, error) { + params := &riverdriver.JobGetStuckParams{ + AfterID: afterID, + Max: batchSize, + Schema: s.Config.Schema, + StuckHorizon: stuckHorizon, + } + + return timeoututil.WithTimeoutV(ctx, riversharedmaintenance.TimeoutDefault, s.Name+".getStuckJobs", func(ctx context.Context) ([]*rivertype.JobRow, error) { + if pilot, ok := s.Config.Pilot.(riverpilot.PilotJobRescuer); ok { + return pilot.JobGetStuck(ctx, s.exec, params) + } + + // Compatibility fallback for Pilot implementations from before + // PilotJobRescuer. Once Pilot embeds PilotJobRescuer, replace the assertion + // above and this fallback with a direct call to s.Config.Pilot.JobGetStuck. + return s.exec.JobGetStuck(ctx, params) + }) +} + +// jobRetryDecision is a signal from makeRetryDecision as to what to do with a +// particular job that appears to be eligible for rescue. +type jobRetryDecision int + +const ( + jobRetryDecisionDiscard jobRetryDecision = iota // discard the job + jobRetryDecisionIgnore // don't retry or discard the job + jobRetryDecisionRetry // retry the job +) + +// makeRetryDecision decides whether or not a rescued job should be retried, and if so, +// when. +func (s *JobRescuer) makeRetryDecision(ctx context.Context, job *rivertype.JobRow, now time.Time) (jobRetryDecision, time.Time) { + workUnitFactory := s.Config.WorkUnitFactoryFunc(job.Kind) + if workUnitFactory == nil { + s.Logger.ErrorContext(ctx, s.Name+": Attempted to rescue unhandled job kind, discarding", + slog.String("job_kind", job.Kind), slog.Int64("job_id", job.ID)) + return jobRetryDecisionDiscard, time.Time{} + } + + workUnit := workUnitFactory.MakeUnit(job) + if err := workUnit.UnmarshalJob(); err != nil { + s.Logger.ErrorContext(ctx, s.Name+": Error unmarshaling job args", + slog.String("error", err.Error()), + slog.String("job_kind", job.Kind), + slog.Int64("job_id", job.ID), + ) + + if job.Attempt < max(job.MaxAttempts, 0) { + return jobRetryDecisionRetry, s.Config.ClientRetryPolicy.NextRetry(job) + } + + return jobRetryDecisionDiscard, time.Time{} + } + + timeout := cmp.Or(workUnit.Timeout(), s.Config.ClientJobTimeout) + if timeout < 0 || timeout > 0 && now.Sub(*job.AttemptedAt) < timeout { + return jobRetryDecisionIgnore, time.Time{} + } + + nextRetry := workUnit.NextRetry() + if nextRetry.IsZero() { + nextRetry = s.Config.ClientRetryPolicy.NextRetry(job) + } + + if job.Attempt < max(job.MaxAttempts, 0) { + return jobRetryDecisionRetry, nextRetry + } + + return jobRetryDecisionDiscard, time.Time{} +} + +func (s *JobRescuer) rescueMany(ctx context.Context, params *riverdriver.JobRescueManyParams) (*struct{}, error) { + if pilot, ok := s.Config.Pilot.(riverpilot.PilotJobRescuer); ok { + return pilot.JobRescueMany(ctx, s.exec, params) + } + + // Compatibility fallback for Pilot implementations from before + // PilotJobRescuer. Once Pilot embeds PilotJobRescuer, replace the assertion + // above and this fallback with a direct call to s.Config.Pilot.JobRescueMany. + return s.exec.JobRescueMany(ctx, params) +} diff --git a/vendor/github.com/riverqueue/river/internal/maintenance/job_scheduler.go b/vendor/github.com/riverqueue/river/internal/maintenance/job_scheduler.go new file mode 100644 index 0000000000..a2d289a244 --- /dev/null +++ b/vendor/github.com/riverqueue/river/internal/maintenance/job_scheduler.go @@ -0,0 +1,233 @@ +package maintenance + +import ( + "cmp" + "context" + "errors" + "fmt" + "log/slog" + "time" + + "github.com/riverqueue/river/riverdriver" + "github.com/riverqueue/river/rivershared/baseservice" + "github.com/riverqueue/river/rivershared/circuitbreaker" + "github.com/riverqueue/river/rivershared/riversharedmaintenance" + "github.com/riverqueue/river/rivershared/startstop" + "github.com/riverqueue/river/rivershared/testsignal" + "github.com/riverqueue/river/rivershared/util/dbutil" + "github.com/riverqueue/river/rivershared/util/randutil" + "github.com/riverqueue/river/rivershared/util/serviceutil" + "github.com/riverqueue/river/rivershared/util/testutil" + "github.com/riverqueue/river/rivershared/util/timeoututil" + "github.com/riverqueue/river/rivershared/util/timeutil" +) + +const ( + JobSchedulerIntervalDefault = 5 * time.Second +) + +// JobSchedulerTestSignals are internal signals used exclusively in tests. +type JobSchedulerTestSignals struct { + NotifiedQueues testsignal.TestSignal[[]string] // notifies when queues are sent an insert notification + ScheduledBatch testsignal.TestSignal[struct{}] // notifies when runOnce finishes a pass +} + +func (ts *JobSchedulerTestSignals) Init(tb testutil.TestingTB) { + ts.NotifiedQueues.Init(tb) + ts.ScheduledBatch.Init(tb) +} + +// NotifyInsertFunc is a function to call to emit notifications for queues where +// jobs were scheduled. +type NotifyInsertFunc func(ctx context.Context, execTx riverdriver.ExecutorTx, queues []string) error + +type JobSchedulerConfig struct { + riversharedmaintenance.BatchSizes + + // Interval is the amount of time between periodic checks for jobs to + // be moved from "scheduled" to "available". + Interval time.Duration + + // NotifyInsert is a function to call to emit notifications for queues + // where jobs were scheduled. + NotifyInsert NotifyInsertFunc + + // Schema where River tables are located. Empty string omits schema, causing + // Postgres to default to `search_path`. + Schema string +} + +func (c *JobSchedulerConfig) mustValidate() *JobSchedulerConfig { + c.MustValidate() + + if c.Interval <= 0 { + panic("SchedulerConfig.Interval must be above zero") + } + if c.Default <= 0 { + panic("SchedulerConfig.Limit must be above zero") + } + + return c +} + +// JobScheduler periodically moves jobs in `scheduled` or `retryable` state and +// which are ready to run over to `available` so that they're eligible to be +// worked. +type JobScheduler struct { + riversharedmaintenance.QueueMaintainerServiceBase + startstop.BaseStartStop + + // exported for test purposes + TestSignals JobSchedulerTestSignals + + config *JobSchedulerConfig + exec riverdriver.Executor + + // Circuit breaker that tracks consecutive timeout failures from the central + // query. The query starts by using the full/default batch size, but after + // this breaker trips (after N consecutive timeouts occur in a row), it + // switches to a smaller batch. We assume that a database that's degraded is + // likely to stay degraded over a longer term, so after the circuit breaks, + // it stays broken until the program is restarted. + reducedBatchSizeBreaker *circuitbreaker.CircuitBreaker +} + +func NewJobScheduler(archetype *baseservice.Archetype, config *JobSchedulerConfig, exec riverdriver.Executor) *JobScheduler { + batchSizes := config.WithDefaults() + + return baseservice.Init(archetype, &JobScheduler{ + config: (&JobSchedulerConfig{ + BatchSizes: batchSizes, + Interval: cmp.Or(config.Interval, JobSchedulerIntervalDefault), + NotifyInsert: config.NotifyInsert, + Schema: config.Schema, + }).mustValidate(), + exec: exec, + reducedBatchSizeBreaker: riversharedmaintenance.ReducedBatchSizeBreaker(batchSizes), + }) +} + +func (s *JobScheduler) Start(ctx context.Context) error { //nolint:dupl + ctx, shouldStart, started, stopped := s.StartInit(ctx) + if !shouldStart { + return nil + } + + s.StaggerStart(ctx) + + go func() { + started() + defer stopped() // this defer should come first so it's last out + + s.Logger.DebugContext(ctx, s.Name+riversharedmaintenance.LogPrefixRunLoopStarted) + defer s.Logger.DebugContext(ctx, s.Name+riversharedmaintenance.LogPrefixRunLoopStopped) + + ticker := timeutil.NewTickerWithInitialTick(ctx, s.config.Interval) + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + } + + res, err := s.runOnce(ctx) + if err != nil { + if !errors.Is(err, context.Canceled) { + s.Logger.ErrorContext(ctx, s.Name+": Error scheduling jobs", slog.String("error", err.Error())) + } + continue + } + + if res.NumCompletedJobsScheduled > 0 { + s.Logger.InfoContext(ctx, s.Name+riversharedmaintenance.LogPrefixRanSuccessfully, + slog.Int("num_jobs_scheduled", res.NumCompletedJobsScheduled), + ) + } + } + }() + + return nil +} + +func (s *JobScheduler) batchSize() int { + if s.reducedBatchSizeBreaker.Open() { + return s.config.Reduced + } + return s.config.Default +} + +type schedulerRunOnceResult struct { + NumCompletedJobsScheduled int +} + +func (s *JobScheduler) runOnce(ctx context.Context) (*schedulerRunOnceResult, error) { + res := &schedulerRunOnceResult{} + + for { + numScheduled, err := timeoututil.WithTimeoutV(ctx, riversharedmaintenance.TimeoutDefault, s.Name+".runOnce", func(ctx context.Context) (int, error) { + execTx, err := s.exec.Begin(ctx) + if err != nil { + return 0, fmt.Errorf("error starting transaction: %w", err) + } + defer dbutil.RollbackWithoutCancel(ctx, execTx) + + now := s.Time.Now() + nowWithLookAhead := now.Add(s.config.Interval) + + scheduledJobResults, err := execTx.JobSchedule(ctx, &riverdriver.JobScheduleParams{ + Max: s.batchSize(), + Now: &nowWithLookAhead, + Schema: s.config.Schema, + }) + if err != nil { + return 0, fmt.Errorf("error scheduling jobs: %w", err) + } + + s.reducedBatchSizeBreaker.ResetIfNotOpen() + + queues := make([]string, 0, len(scheduledJobResults)) + + // Notify about scheduled jobs with a scheduled_at in the past, or just + // slightly in the future (this loop, the notify, and tx commit will take + // a small amount of time). This isn't going to be perfect, but the goal + // is to roughly try to guess when the clients will attempt to fetch jobs. + notificationHorizon := s.Time.Now().Add(5 * time.Millisecond) + + for _, result := range scheduledJobResults { + if result.Job.ScheduledAt.After(notificationHorizon) { + continue + } + + queues = append(queues, result.Job.Queue) + } + + if len(queues) > 0 { + if err := s.config.NotifyInsert(ctx, execTx, queues); err != nil { + return 0, fmt.Errorf("error notifying insert: %w", err) + } + s.TestSignals.NotifiedQueues.Signal(queues) + } + + return len(scheduledJobResults), execTx.Commit(ctx) + }) + if err != nil { + if errors.Is(err, context.DeadlineExceeded) { + s.reducedBatchSizeBreaker.Trip() + } + + return nil, err + } + + s.TestSignals.ScheduledBatch.Signal(struct{}{}) + + res.NumCompletedJobsScheduled += numScheduled + // Scheduled was less than query `LIMIT` which means work is done. + if numScheduled < s.batchSize() { + break + } + + serviceutil.CancellableSleep(ctx, randutil.DurationBetween(riversharedmaintenance.BatchBackoffMin, riversharedmaintenance.BatchBackoffMax)) + } + + return res, nil +} diff --git a/vendor/github.com/riverqueue/river/internal/maintenance/periodic_job_enqueuer.go b/vendor/github.com/riverqueue/river/internal/maintenance/periodic_job_enqueuer.go new file mode 100644 index 0000000000..6f5a5c9058 --- /dev/null +++ b/vendor/github.com/riverqueue/river/internal/maintenance/periodic_job_enqueuer.go @@ -0,0 +1,693 @@ +package maintenance + +import ( + "context" + "errors" + "fmt" + "slices" + "sync" + "time" + + "github.com/tidwall/sjson" + + "github.com/riverqueue/river/internal/pluginlookup" + "github.com/riverqueue/river/internal/rivercommon" + "github.com/riverqueue/river/riverdriver" + "github.com/riverqueue/river/rivershared/baseservice" + "github.com/riverqueue/river/rivershared/riverpilot" + "github.com/riverqueue/river/rivershared/riversharedmaintenance" + "github.com/riverqueue/river/rivershared/startstop" + "github.com/riverqueue/river/rivershared/testsignal" + "github.com/riverqueue/river/rivershared/util/maputil" + "github.com/riverqueue/river/rivershared/util/sliceutil" + "github.com/riverqueue/river/rivershared/util/testutil" + "github.com/riverqueue/river/rivershared/util/timeutil" + "github.com/riverqueue/river/rivertype" +) + +// ErrNoJobToInsert can be returned by a PeriodicJob's JobToInsertFunc to +// signal that there's no job to insert at this time. +var ErrNoJobToInsert = errors.New("a nil job was returned, nothing to insert") + +// PeriodicJobEnqueuerTestSignals are internal signals used exclusively in tests. +type PeriodicJobEnqueuerTestSignals struct { + EnteredLoop testsignal.TestSignal[struct{}] // notifies when the enqueuer finishes start up and enters its initial run loop + InsertedJobs testsignal.TestSignal[struct{}] // notifies when a batch of jobs is inserted + PeriodicJobKeepAliveAndReap testsignal.TestSignal[struct{}] // notifies when the background services that runs keep alive and reap on periodic jobs ticks + PeriodicJobUpserted testsignal.TestSignal[struct{}] // notifies when a batch of periodic job records are upserted to pilot + SkippedJob testsignal.TestSignal[struct{}] // notifies when a job is skipped because of nil JobInsertParams +} + +func (ts *PeriodicJobEnqueuerTestSignals) Init(tb testutil.TestingTB) { + ts.EnteredLoop.Init(tb) + ts.InsertedJobs.Init(tb) + ts.PeriodicJobKeepAliveAndReap.Init(tb) + ts.PeriodicJobUpserted.Init(tb) + ts.SkippedJob.Init(tb) +} + +// PeriodicJob is a periodic job to be run. It's similar to the top-level +// river.PeriodicJobArgs, but needs a separate type because the enqueuer is in a +// subpackage. +type PeriodicJob struct { + ID string + ConstructorFunc func() (*rivertype.JobInsertParams, error) + RunOnStart bool + ScheduleFunc func(time.Time) time.Time + + nextRunAt time.Time // set on service start +} + +func (j *PeriodicJob) mustValidate() *PeriodicJob { + if err := j.validate(); err != nil { + panic(err) + } + return j +} + +func (j *PeriodicJob) validate() error { + if j.ID != "" { + if len(j.ID) >= 128 { + return errors.New("PeriodicJob.ID must be less than 128 characters") + } + if !rivercommon.UserSpecifiedIDOrKindRE.MatchString(j.ID) { + return fmt.Errorf("PeriodicJob.ID %q should match regex %s", j.ID, rivercommon.UserSpecifiedIDOrKindRE.String()) + } + } + if j.ConstructorFunc == nil { + return errors.New("PeriodicJob.ConstructorFunc must be set") + } + if j.ScheduleFunc == nil { + return errors.New("PeriodicJob.ScheduleFunc must be set") + } + + return nil +} + +type InsertFunc func(ctx context.Context, tx riverdriver.ExecutorTx, insertParams []*rivertype.JobInsertParams) ([]*rivertype.JobInsertResult, error) + +type PeriodicJobEnqueuerConfig struct { + AdvisoryLockPrefix int32 + + PluginLookupGlobal *pluginlookup.PluginLookup + + // Insert is the function to call to insert jobs into the database. + Insert InsertFunc + + // PeriodicJobs are the periodic jobs with which to configure the enqueuer. + PeriodicJobs []*PeriodicJob + + // Pilot is a plugin module providing additional non-standard functionality. + Pilot riverpilot.PilotPeriodicJob + + // Schema where River tables are located. Empty string omits schema, causing + // Postgres to default to `search_path`. + Schema string +} + +func (c *PeriodicJobEnqueuerConfig) mustValidate() *PeriodicJobEnqueuerConfig { + // no validations currently + return c +} + +// PeriodicJobEnqueuer inserts jobs configured to run periodically as unique +// jobs to make sure they'll run as frequently as their period dictates. +type PeriodicJobEnqueuer struct { + riversharedmaintenance.QueueMaintainerServiceBase + startstop.BaseStartStop + + // exported for test purposes + Config *PeriodicJobEnqueuerConfig + TestSignals PeriodicJobEnqueuerTestSignals + + exec riverdriver.Executor + mu sync.RWMutex + nextHandle rivertype.PeriodicJobHandle + periodicJobIDs map[string]rivertype.PeriodicJobHandle + periodicJobs map[rivertype.PeriodicJobHandle]*PeriodicJob + recalculateNextRun chan struct{} +} + +func NewPeriodicJobEnqueuer(archetype *baseservice.Archetype, config *PeriodicJobEnqueuerConfig, exec riverdriver.Executor) (*PeriodicJobEnqueuer, error) { + var ( + nextHandle rivertype.PeriodicJobHandle + periodicJobIDs = make(map[string]rivertype.PeriodicJobHandle) + periodicJobs = make(map[rivertype.PeriodicJobHandle]*PeriodicJob, len(config.PeriodicJobs)) + ) + + for _, periodicJob := range config.PeriodicJobs { + if err := periodicJob.validate(); err != nil { + return nil, err + } + + handle := nextHandle + + if err := addUniqueID(periodicJobIDs, periodicJob.ID, handle); err != nil { + return nil, err + } + + periodicJobs[handle] = periodicJob + nextHandle++ + } + + pluginLookupGlobal := config.PluginLookupGlobal + if pluginLookupGlobal == nil { + pluginLookupGlobal = pluginlookup.NewPluginLookup(nil) + } + + pilot := config.Pilot + if pilot == nil { + pilot = &riverpilot.StandardPilot{} + } + + svc := baseservice.Init(archetype, &PeriodicJobEnqueuer{ + Config: (&PeriodicJobEnqueuerConfig{ + AdvisoryLockPrefix: config.AdvisoryLockPrefix, + PluginLookupGlobal: pluginLookupGlobal, + Insert: config.Insert, + PeriodicJobs: config.PeriodicJobs, + Pilot: pilot, + Schema: config.Schema, + }).mustValidate(), + + exec: exec, + nextHandle: nextHandle, + periodicJobIDs: periodicJobIDs, + periodicJobs: periodicJobs, + recalculateNextRun: make(chan struct{}, 1), + }) + + return svc, nil +} + +// AddSafely adds a new periodic job to the enqueuer. The service's run loop is +// woken immediately so that the job is scheduled appropriately, and inserted if +// its RunOnStart flag is set to true. +func (s *PeriodicJobEnqueuer) AddSafely(periodicJob *PeriodicJob) (rivertype.PeriodicJobHandle, error) { + s.mu.Lock() + defer s.mu.Unlock() + + if err := periodicJob.validate(); err != nil { + return 0, err + } + + handle := s.nextHandle + + if err := addUniqueID(s.periodicJobIDs, periodicJob.ID, handle); err != nil { + return 0, err + } + + s.periodicJobs[handle] = periodicJob + s.nextHandle++ + + select { + case s.recalculateNextRun <- struct{}{}: + default: + } + + return handle, nil +} + +// AddManySafely adds many new periodic job to the enqueuer. The service's run loop is +// woken immediately so that the job is scheduled appropriately, and inserted if +// any RunOnStart flags are set to true. +func (s *PeriodicJobEnqueuer) AddManySafely(periodicJobs []*PeriodicJob) ([]rivertype.PeriodicJobHandle, error) { + s.mu.Lock() + defer s.mu.Unlock() + + handles := make([]rivertype.PeriodicJobHandle, len(periodicJobs)) + + for i, periodicJob := range periodicJobs { + if err := periodicJob.validate(); err != nil { + return nil, err + } + + handles[i] = s.nextHandle + + if err := addUniqueID(s.periodicJobIDs, periodicJob.ID, handles[i]); err != nil { + return nil, err + } + + s.periodicJobs[handles[i]] = periodicJob + s.nextHandle++ + } + + select { + case s.recalculateNextRun <- struct{}{}: + default: + } + + return handles, nil +} + +// Clear clears all periodic jobs from the enqueuer. +func (s *PeriodicJobEnqueuer) Clear() { + s.mu.Lock() + defer s.mu.Unlock() + + s.periodicJobIDs = make(map[string]rivertype.PeriodicJobHandle) + + // `nextHandle` is _not_ reset so that even across multiple generations of + // jobs, handles aren't reused. + s.periodicJobs = make(map[rivertype.PeriodicJobHandle]*PeriodicJob) +} + +// Remove removes a periodic job from the enqueuer. Its current target run time +// and all future runs are cancelled. +func (s *PeriodicJobEnqueuer) Remove(periodicJobHandle rivertype.PeriodicJobHandle) { + s.mu.Lock() + defer s.mu.Unlock() + + if periodicJob, ok := s.periodicJobs[periodicJobHandle]; ok { + s.removeJobLockFree(periodicJob, periodicJobHandle) + } +} + +// RemoveByID removes a periodic job from the enqueuer by ID. Its current target run +// time and all future runs are cancelled. +func (s *PeriodicJobEnqueuer) RemoveByID(id string) bool { + s.mu.Lock() + defer s.mu.Unlock() + + if handle, ok := s.periodicJobIDs[id]; ok { + delete(s.periodicJobIDs, id) + delete(s.periodicJobs, handle) + return true + } + + return false +} + +// RemoveMany removes many periodic jobs from the enqueuer. Their current target +// run time and all future runs are cancelled. +func (s *PeriodicJobEnqueuer) RemoveMany(periodicJobHandles []rivertype.PeriodicJobHandle) { + s.mu.Lock() + defer s.mu.Unlock() + + for _, periodicJobHandle := range periodicJobHandles { + if periodicJob, ok := s.periodicJobs[periodicJobHandle]; ok { + s.removeJobLockFree(periodicJob, periodicJobHandle) + } + } +} + +// RemoveManyByID removes many periodic jobs from the enqueuer by ID. Their +// current target run time and all future runs are cancelled. +func (s *PeriodicJobEnqueuer) RemoveManyByID(ids []string) { + s.mu.Lock() + defer s.mu.Unlock() + + for _, id := range ids { + if handle, ok := s.periodicJobIDs[id]; ok { + delete(s.periodicJobIDs, id) + delete(s.periodicJobs, handle) + } + } +} + +func (s *PeriodicJobEnqueuer) removeJobLockFree(periodicJob *PeriodicJob, periodicJobHandle rivertype.PeriodicJobHandle) { + delete(s.periodicJobs, periodicJobHandle) + delete(s.periodicJobIDs, periodicJob.ID) +} + +func (s *PeriodicJobEnqueuer) Start(ctx context.Context) error { + ctx, shouldStart, started, stopped := s.StartInit(ctx) + if !shouldStart { + return nil + } + + s.StaggerStart(ctx) + + var ( + initialPeriodicJobs []*riverpilot.PeriodicJob + subServices []startstop.Service + ) + if err := func() error { + var err error + initialPeriodicJobs, err = s.Config.Pilot.PeriodicJobGetAll(ctx, s.exec, &riverpilot.PeriodicJobGetAllParams{ + Schema: s.Config.Schema, + }) + if err != nil { + return err + } + + for _, hook := range s.Config.PluginLookupGlobal.ByKind(pluginlookup.PluginKindHookPeriodicJobsStart) { + if err := hook.(rivertype.HookPeriodicJobsStart).Start(ctx, &rivertype.HookPeriodicJobsStartParams{ //nolint:forcetypeassert + DurableJobs: sliceutil.Map(initialPeriodicJobs, func(job *riverpilot.PeriodicJob) *rivertype.DurablePeriodicJob { + return (*rivertype.DurablePeriodicJob)(job) + }), + }); err != nil { + return err + } + } + + subServices = []startstop.Service{ + startstop.StartStopFunc(s.periodicJobKeepAliveAndReapPeriodically), + } + stopServicesOnError := func() { + startstop.StopAllParallel(subServices...) + } + if err := startstop.StartAll(ctx, subServices...); err != nil { + stopServicesOnError() + return err + } + + return nil + }(); err != nil { + stopped() + return err + } + + go func() { + started() + defer stopped() // this defer should come first so it's last out + + s.Logger.DebugContext(ctx, s.Name+riversharedmaintenance.LogPrefixRunLoopStarted) + defer s.Logger.DebugContext(ctx, s.Name+riversharedmaintenance.LogPrefixRunLoopStopped) + + defer startstop.StopAllParallel(subServices...) + + // Drain the signal to recalculate next run if it's been sent (i.e. Add + // or AddMany called before Start). We're about to schedule jobs from + // scratch, and therefore don't need to immediately do so again. + select { + case <-s.recalculateNextRun: + default: + } + + // Initial set of periodic job IDs mapped to next run at times fetched + // from a configured pilot. Not used in most cases. + initialPeriodicJobsMap := sliceutil.KeyBy(initialPeriodicJobs, + func(j *riverpilot.PeriodicJob) (string, time.Time) { return j.ID, j.NextRunAt }) + + var lastHandleSeen rivertype.PeriodicJobHandle = -1 // so handle 0 is considered + + validateInsertRunOnStartAndScheduleNewlyAdded := func() { + s.mu.RLock() + defer s.mu.RUnlock() + + var ( + insertParamsMany []*rivertype.JobInsertParams + now = s.Time.Now() + periodicJobUpsertParams = &riverpilot.PeriodicJobUpsertManyParams{Schema: s.Config.Schema} + ) + + // Handle periodic jobs in sorted order so we can correctly account + // for the most recently added one that we've seen. + sortedPeriodicJobHandles := maputil.Keys(s.periodicJobs) + slices.Sort(sortedPeriodicJobHandles) + + for _, handle := range sortedPeriodicJobHandles { + if handle <= lastHandleSeen { + continue + } + + lastHandleSeen = handle + + periodicJob := s.periodicJobs[handle].mustValidate() + + if nextRunAt, ok := initialPeriodicJobsMap[periodicJob.ID]; periodicJob.ID != "" && ok { + periodicJob.nextRunAt = nextRunAt + delete(initialPeriodicJobsMap, periodicJob.ID) + } else { + periodicJob.nextRunAt = periodicJob.ScheduleFunc(now) + } + + if periodicJob.ID != "" { + periodicJobUpsertParams.Jobs = append(periodicJobUpsertParams.Jobs, &riverpilot.PeriodicJobUpsertParams{ + ID: periodicJob.ID, + NextRunAt: periodicJob.nextRunAt, + UpdatedAt: s.Time.Now(), + }) + } + + if !periodicJob.RunOnStart { + continue + } + + if insertParams, ok := s.insertParamsFromConstructor(ctx, periodicJob.ID, periodicJob.ConstructorFunc, now); ok { + insertParamsMany = append(insertParamsMany, insertParams) + } + } + + s.insertBatch(ctx, insertParamsMany, periodicJobUpsertParams) + + if len(insertParamsMany) > 0 { + s.Logger.DebugContext(ctx, s.Name+": Inserted RunOnStart jobs", "num_jobs", len(insertParamsMany)) + } + } + + // Run any jobs that need to run on start and calculate initial runs. + validateInsertRunOnStartAndScheduleNewlyAdded() + + s.TestSignals.EnteredLoop.Signal(struct{}{}) + + timerUntilNextRun := time.NewTimer(s.timeUntilNextRun()) + + for { + select { + case <-timerUntilNextRun.C: + var ( + insertParamsMany []*rivertype.JobInsertParams + periodicJobUpsertParams = &riverpilot.PeriodicJobUpsertManyParams{Schema: s.Config.Schema} + ) + + now := s.Time.Now() + + // Add a small margin to the current time so we're not only + // running jobs that are already ready, but also ones ready at + // this exact moment or ready in the very near future. + nowWithMargin := now.Add(100 * time.Millisecond) + + func() { + s.mu.RLock() + defer s.mu.RUnlock() + + for _, periodicJob := range s.periodicJobs { + if periodicJob.nextRunAt.IsZero() || !periodicJob.nextRunAt.Before(nowWithMargin) { + continue + } + + if insertParams, ok := s.insertParamsFromConstructor(ctx, periodicJob.ID, periodicJob.ConstructorFunc, periodicJob.nextRunAt); ok { + insertParamsMany = append(insertParamsMany, insertParams) + } + + // Although we may have inserted a new job a little + // preemptively due to the margin applied above, try to stay + // as true as possible to the original schedule by using the + // original run time when calculating the next one. + periodicJob.nextRunAt = periodicJob.ScheduleFunc(periodicJob.nextRunAt) + + if periodicJob.ID != "" { + periodicJobUpsertParams.Jobs = append(periodicJobUpsertParams.Jobs, &riverpilot.PeriodicJobUpsertParams{ + ID: periodicJob.ID, + NextRunAt: periodicJob.nextRunAt, + UpdatedAt: s.Time.Now(), + }) + } + } + }() + + s.insertBatch(ctx, insertParamsMany, periodicJobUpsertParams) + + case <-s.recalculateNextRun: + if !timerUntilNextRun.Stop() { + <-timerUntilNextRun.C + } + + case <-ctx.Done(): + // Clean up timer resources. We know it has _not_ received from the + // timer since its last reset because that would have led us to the case + // above instead of here. + if !timerUntilNextRun.Stop() { + <-timerUntilNextRun.C + } + return + } + + // Insert any RunOnStart initial runs for new jobs that've been + // added since the last run loop. + validateInsertRunOnStartAndScheduleNewlyAdded() + + // Reset the timer after the insert loop has finished so it's + // paused during work. Makes its firing more deterministic. + timerUntilNextRun.Reset(s.timeUntilNextRun()) + } + }() + + return nil +} + +func (s *PeriodicJobEnqueuer) insertBatch(ctx context.Context, insertParamsMany []*rivertype.JobInsertParams, periodicJobUpsertParams *riverpilot.PeriodicJobUpsertManyParams) { + if len(insertParamsMany) < 1 && len(periodicJobUpsertParams.Jobs) < 1 { + return + } + + ctx, cancel := context.WithTimeout(ctx, riversharedmaintenance.TimeoutDefault) + defer cancel() + + tx, err := s.exec.Begin(ctx) + if err != nil { + s.Logger.ErrorContext(ctx, s.Name+": Error starting transaction", "error", err.Error()) + return + } + defer tx.Rollback(ctx) + + if len(insertParamsMany) > 0 { + if _, err := s.Config.Insert(ctx, tx, insertParamsMany); err != nil { + s.Logger.ErrorContext(ctx, s.Name+": Error inserting periodic jobs", + "error", err.Error(), "num_jobs", len(insertParamsMany)) + } + } + + if len(periodicJobUpsertParams.Jobs) > 0 { + if _, err = s.Config.Pilot.PeriodicJobUpsertMany(ctx, tx, periodicJobUpsertParams); err != nil { + s.Logger.ErrorContext(ctx, s.Name+": Error upserting periodic job next run times", + "error", err.Error(), "num_jobs", len(insertParamsMany), "num_next_run_at_upserts", len(periodicJobUpsertParams.Jobs)) + return + } + } + + if err := tx.Commit(ctx); err != nil { + s.Logger.ErrorContext(ctx, s.Name+": Error committing transaction", "error", err.Error()) + return + } + + if len(insertParamsMany) > 0 { + s.TestSignals.InsertedJobs.Signal(struct{}{}) + } + if len(periodicJobUpsertParams.Jobs) > 0 { + s.TestSignals.PeriodicJobUpserted.Signal(struct{}{}) + } +} + +func (s *PeriodicJobEnqueuer) insertParamsFromConstructor(ctx context.Context, periodicJobID string, constructorFunc func() (*rivertype.JobInsertParams, error), scheduledAt time.Time) (*rivertype.JobInsertParams, bool) { + insertParams, err := constructorFunc() + if err != nil { + if errors.Is(err, ErrNoJobToInsert) { + s.Logger.InfoContext(ctx, s.Name+": nil returned from periodic job constructor, skipping") + s.TestSignals.SkippedJob.Signal(struct{}{}) + return nil, false + } + s.Logger.ErrorContext(ctx, s.Name+": Internal error generating periodic job", "error", err.Error()) + return nil, false + } + + if insertParams.ScheduledAt == nil { + insertParams.ScheduledAt = &scheduledAt + } + + if periodicJobID != "" { + var err error + if insertParams.Metadata, err = sjson.SetBytes(insertParams.Metadata, rivercommon.MetadataKeyPeriodicJobID, periodicJobID); err != nil { + s.Logger.ErrorContext(ctx, s.Name+": Error setting periodic metadata", "error", err.Error()) + } + } + + if insertParams.Metadata, err = sjson.SetBytes(insertParams.Metadata, "periodic", true); err != nil { + s.Logger.ErrorContext(ctx, s.Name+": Error setting periodic metadata", "error", err.Error()) + return nil, false + } + + return insertParams, true +} + +func (s *PeriodicJobEnqueuer) periodicJobKeepAliveAndReapPeriodically(ctx context.Context, shouldStart bool, started, stopped func()) error { + if !shouldStart { + return nil + } + + go func() { + started() + defer stopped() // this defer should come first so it's last out + + ticker := timeutil.NewTickerWithInitialTick(ctx, 10*time.Minute) + for { + select { + case <-ctx.Done(): + return + + case <-ticker.C: + func() { + s.mu.RLock() + defer s.mu.RUnlock() + + if len(s.periodicJobIDs) > 0 { + if _, err := s.Config.Pilot.PeriodicJobKeepAliveAndReap(ctx, s.exec, &riverpilot.PeriodicJobKeepAliveAndReapParams{ + ID: maputil.Keys(s.periodicJobIDs), + Schema: s.Config.Schema, + }); err != nil { + s.Logger.ErrorContext(ctx, s.Name+": Error executing periodic job keep alive and reap", "error", err.Error()) + return + } + } + + s.TestSignals.PeriodicJobKeepAliveAndReap.Signal(struct{}{}) + }() + } + } + }() + + return nil +} + +const periodicJobEnqueuerVeryLongDuration = 24 * time.Hour + +func (s *PeriodicJobEnqueuer) timeUntilNextRun() time.Duration { + s.mu.RLock() + defer s.mu.RUnlock() + + // With no configured jobs, just return a big duration for the loop to block + // on. + if len(s.periodicJobs) < 1 { + return periodicJobEnqueuerVeryLongDuration + } + + var ( + firstNextRunAt time.Time + now = s.Time.Now() + ) + + for _, periodicJob := range s.periodicJobs { + // Jobs may have been added after service start, but before this + // function runs for the first time. They're not scheduled properly yet, + // but they will be soon, at which point this function will run again. + // Skip them for now. + if periodicJob.nextRunAt.IsZero() { + continue + } + + // In case we detect a job that should've run before now, immediately short + // circuit with a 0 duration. This avoids needlessly iterating through the + // rest of the loop when we already know we're overdue for the next job. + if periodicJob.nextRunAt.Before(now) { + return 0 + } + + if firstNextRunAt.IsZero() || periodicJob.nextRunAt.Before(firstNextRunAt) { + firstNextRunAt = periodicJob.nextRunAt + } + } + + // Only encountered unscheduled jobs (see comment above). Don't schedule + // anything for now. + if firstNextRunAt.IsZero() { + return periodicJobEnqueuerVeryLongDuration + } + + return firstNextRunAt.Sub(now) +} + +// Adds a unique ID to known periodic job IDs, erroring in case of a duplicate. +func addUniqueID(periodicJobIDs map[string]rivertype.PeriodicJobHandle, id string, handle rivertype.PeriodicJobHandle) error { + if id == "" { + return nil + } + + if _, ok := periodicJobIDs[id]; ok { + return errors.New("periodic job with ID already registered: " + id) + } + + periodicJobIDs[id] = handle + return nil +} diff --git a/vendor/github.com/riverqueue/river/internal/maintenance/queue_cleaner.go b/vendor/github.com/riverqueue/river/internal/maintenance/queue_cleaner.go new file mode 100644 index 0000000000..e417558b9f --- /dev/null +++ b/vendor/github.com/riverqueue/river/internal/maintenance/queue_cleaner.go @@ -0,0 +1,194 @@ +package maintenance + +import ( + "cmp" + "context" + "errors" + "fmt" + "log/slog" + "strings" + "time" + + "github.com/riverqueue/river/riverdriver" + "github.com/riverqueue/river/rivershared/baseservice" + "github.com/riverqueue/river/rivershared/circuitbreaker" + "github.com/riverqueue/river/rivershared/riversharedmaintenance" + "github.com/riverqueue/river/rivershared/startstop" + "github.com/riverqueue/river/rivershared/testsignal" + "github.com/riverqueue/river/rivershared/util/randutil" + "github.com/riverqueue/river/rivershared/util/serviceutil" + "github.com/riverqueue/river/rivershared/util/testutil" + "github.com/riverqueue/river/rivershared/util/timeoututil" + "github.com/riverqueue/river/rivershared/util/timeutil" +) + +const ( + queueCleanerIntervalDefault = time.Hour + QueueRetentionPeriodDefault = 24 * time.Hour +) + +// QueueCleanerTestSignals are internal signals used exclusively in tests. +type QueueCleanerTestSignals struct { + DeletedBatch testsignal.TestSignal[struct{}] // notifies when runOnce finishes a pass +} + +func (ts *QueueCleanerTestSignals) Init(tb testutil.TestingTB) { + ts.DeletedBatch.Init(tb) +} + +type QueueCleanerConfig struct { + riversharedmaintenance.BatchSizes + + // Interval is the amount of time to wait between runs of the cleaner. + Interval time.Duration + + // RetentionPeriod is the amount of time to keep queues around before they're + // removed. + RetentionPeriod time.Duration + + // Schema where River tables are located. Empty string omits schema, causing + // Postgres to default to `search_path`. + Schema string +} + +func (c *QueueCleanerConfig) mustValidate() *QueueCleanerConfig { + c.MustValidate() + + if c.Interval <= 0 { + panic("QueueCleanerConfig.Interval must be above zero") + } + if c.RetentionPeriod <= 0 { + panic("QueueCleanerConfig.RetentionPeriod must be above zero") + } + + return c +} + +// QueueCleaner periodically removes queues from the river_queue table that have +// not been updated in a while, indicating that they are no longer active. +type QueueCleaner struct { + riversharedmaintenance.QueueMaintainerServiceBase + startstop.BaseStartStop + + // exported for test purposes + Config *QueueCleanerConfig + TestSignals QueueCleanerTestSignals + + exec riverdriver.Executor + + // Circuit breaker that tracks consecutive timeout failures from the central + // query. The query starts by using the full/default batch size, but after + // this breaker trips (after N consecutive timeouts occur in a row), it + // switches to a smaller batch. We assume that a database that's degraded is + // likely to stay degraded over a longer term, so after the circuit breaks, + // it stays broken until the program is restarted. + reducedBatchSizeBreaker *circuitbreaker.CircuitBreaker +} + +func NewQueueCleaner(archetype *baseservice.Archetype, config *QueueCleanerConfig, exec riverdriver.Executor) *QueueCleaner { + batchSizes := config.WithDefaults() + + return baseservice.Init(archetype, &QueueCleaner{ + Config: (&QueueCleanerConfig{ + BatchSizes: batchSizes, + Interval: cmp.Or(config.Interval, queueCleanerIntervalDefault), + RetentionPeriod: cmp.Or(config.RetentionPeriod, QueueRetentionPeriodDefault), + Schema: config.Schema, + }).mustValidate(), + exec: exec, + reducedBatchSizeBreaker: riversharedmaintenance.ReducedBatchSizeBreaker(batchSizes), + }) +} + +func (s *QueueCleaner) Start(ctx context.Context) error { + ctx, shouldStart, started, stopped := s.StartInit(ctx) + if !shouldStart { + return nil + } + + s.StaggerStart(ctx) + + go func() { + started() + defer stopped() // this defer should come first so it's last out + + s.Logger.DebugContext(ctx, s.Name+riversharedmaintenance.LogPrefixRunLoopStarted) + defer s.Logger.DebugContext(ctx, s.Name+riversharedmaintenance.LogPrefixRunLoopStopped) + + ticker := timeutil.NewTickerWithInitialTick(ctx, s.Config.Interval) + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + } + + res, err := s.runOnce(ctx) + if err != nil { + if !errors.Is(err, context.Canceled) { + s.Logger.ErrorContext(ctx, s.Name+": Error cleaning queues", slog.String("error", err.Error())) + } + continue + } + + if len(res.QueuesDeleted) > 0 { + s.Logger.InfoContext(ctx, s.Name+riversharedmaintenance.LogPrefixRanSuccessfully, + slog.String("queues_deleted", strings.Join(res.QueuesDeleted, ",")), + ) + } + } + }() + + return nil +} + +func (s *QueueCleaner) batchSize() int { + if s.reducedBatchSizeBreaker.Open() { + return s.Config.Reduced + } + return s.Config.Default +} + +type queueCleanerRunOnceResult struct { + QueuesDeleted []string +} + +func (s *QueueCleaner) runOnce(ctx context.Context) (*queueCleanerRunOnceResult, error) { + res := &queueCleanerRunOnceResult{QueuesDeleted: make([]string, 0, 10)} + + for { + queuesDeleted, err := timeoututil.WithTimeoutV(ctx, riversharedmaintenance.TimeoutDefault, s.Name+".runOnce", func(ctx context.Context) ([]string, error) { + queuesDeleted, err := s.exec.QueueDeleteExpired(ctx, &riverdriver.QueueDeleteExpiredParams{ + Max: s.batchSize(), + Schema: s.Config.Schema, + UpdatedAtHorizon: time.Now().Add(-s.Config.RetentionPeriod), + }) + if err != nil { + return nil, fmt.Errorf("error deleting expired queues: %w", err) + } + + s.reducedBatchSizeBreaker.ResetIfNotOpen() + + return queuesDeleted, nil + }) + if err != nil { + if errors.Is(err, context.DeadlineExceeded) { + s.reducedBatchSizeBreaker.Trip() + } + + return nil, err + } + + s.TestSignals.DeletedBatch.Signal(struct{}{}) + + res.QueuesDeleted = append(res.QueuesDeleted, queuesDeleted...) + // Deleted was less than query `LIMIT` which means work is done. + if len(queuesDeleted) < s.batchSize() { + break + } + + serviceutil.CancellableSleep(ctx, randutil.DurationBetween(riversharedmaintenance.BatchBackoffMin, riversharedmaintenance.BatchBackoffMax)) + } + + return res, nil +} diff --git a/vendor/github.com/riverqueue/river/internal/maintenance/queue_maintainer.go b/vendor/github.com/riverqueue/river/internal/maintenance/queue_maintainer.go new file mode 100644 index 0000000000..142bf283a0 --- /dev/null +++ b/vendor/github.com/riverqueue/river/internal/maintenance/queue_maintainer.go @@ -0,0 +1,94 @@ +package maintenance + +import ( + "context" + "reflect" + + "github.com/riverqueue/river/rivershared/baseservice" + "github.com/riverqueue/river/rivershared/startstop" + "github.com/riverqueue/river/rivershared/util/maputil" +) + +// QueueMaintainer runs regular maintenance operations against job queues, like +// pruning completed jobs. It runs only on the client which has been elected +// leader at any given time. +// +// Its methods are not safe for concurrent usage. +type QueueMaintainer struct { + baseservice.BaseService + startstop.BaseStartStop + + servicesByName map[string]startstop.Service +} + +func NewQueueMaintainer(archetype *baseservice.Archetype, services []startstop.Service) *QueueMaintainer { + servicesByName := make(map[string]startstop.Service, len(services)) + for _, service := range services { + servicesByName[serviceName(service)] = service + } + return baseservice.Init(archetype, &QueueMaintainer{ + servicesByName: servicesByName, + }) +} + +// StaggerStartupDisable sets whether the short staggered sleep on start up +// is disabled. This is useful in tests where the extra sleep involved in a +// staggered start up is not helpful for test run time. +func (m *QueueMaintainer) StaggerStartupDisable(disabled bool) { + for _, svc := range m.servicesByName { + if svcWithDisable, ok := svc.(withStaggerStartupDisable); ok { + svcWithDisable.StaggerStartupDisable(disabled) + } + } +} + +func (m *QueueMaintainer) Start(ctx context.Context) error { + ctx, shouldStart, started, stopped := m.StartInit(ctx) + if !shouldStart { + return nil + } + + for _, service := range m.servicesByName { + if err := service.Start(ctx); err != nil { + startstop.StopAllParallel(maputil.Values(m.servicesByName)...) + stopped() + return err + } + } + + go func() { + // Wait for all subservices to start up before signaling our own start. + startstop.WaitAllStarted(maputil.Values(m.servicesByName)...) + + started() + defer stopped() // this defer should come first so it's last out + + <-ctx.Done() + + startstop.StopAllParallel(maputil.Values(m.servicesByName)...) + }() + + return nil +} + +// GetService is a convenience method for getting a service by name and casting +// it to the desired type. It should only be used in tests due to its use of +// reflection and potential for panics. +func GetService[T startstop.Service](maintainer *QueueMaintainer) T { + var kindPtr T + return maintainer.servicesByName[serviceName(kindPtr)].(T) //nolint:forcetypeassert +} + +func serviceName(service startstop.Service) string { + elem := reflect.TypeOf(service).Elem() + return elem.PkgPath() + "." + elem.Name() +} + +// withStaggerStartupDisable is an interface to a service whose stagger startup +// sleep can be disabled. +type withStaggerStartupDisable interface { + // StaggerStartupDisable sets whether the short staggered sleep on start up + // is disabled. This is useful in tests where the extra sleep involved in a + // staggered start up is not helpful for test run time. + StaggerStartupDisable(disabled bool) +} diff --git a/vendor/github.com/riverqueue/river/internal/maintenance/queue_maintainer_leader.go b/vendor/github.com/riverqueue/river/internal/maintenance/queue_maintainer_leader.go new file mode 100644 index 0000000000..c196e90e9d --- /dev/null +++ b/vendor/github.com/riverqueue/river/internal/maintenance/queue_maintainer_leader.go @@ -0,0 +1,189 @@ +package maintenance + +import ( + "context" + "log/slog" + "sync" + + "github.com/riverqueue/river/internal/leadership" + "github.com/riverqueue/river/rivershared/baseservice" + "github.com/riverqueue/river/rivershared/startstop" + "github.com/riverqueue/river/rivershared/testsignal" + "github.com/riverqueue/river/rivershared/util/serviceutil" + "github.com/riverqueue/river/rivershared/util/testutil" +) + +const queueMaintainerMaxStartAttempts = 3 + +// QueueMaintainerLeaderTestSignals are internal signals used exclusively in tests. +type QueueMaintainerLeaderTestSignals struct { + ElectedLeader testsignal.TestSignal[struct{}] // notifies when elected leader + StartError testsignal.TestSignal[error] // notifies on each failed queue maintainer start attempt + StartRetriesExhausted testsignal.TestSignal[struct{}] // notifies when all start retries have been exhausted +} + +func (ts *QueueMaintainerLeaderTestSignals) Init(tb testutil.TestingTB) { + ts.ElectedLeader.Init(tb) + ts.StartError.Init(tb) + ts.StartRetriesExhausted.Init(tb) +} + +// QueueMaintainerLeaderConfig is the configuration for QueueMaintainerLeader. +type QueueMaintainerLeaderConfig struct { + // ClientID is used for logging on leadership changes. + ClientID string + + // Elector provides leadership change notifications. + Elector *leadership.Elector + + // QueueMaintainer is the underlying maintainer to start/stop on leadership + // changes. + QueueMaintainer *QueueMaintainer + + // RequestResignFunc sends a notification requesting leader resignation. + // It's injected from the client because the notification mechanism depends + // on the driver, which the maintenance package doesn't know about. + RequestResignFunc func(ctx context.Context) error +} + +// QueueMaintainerLeader listens for leadership changes and starts/stops the +// queue maintainer accordingly. It handles retries with exponential backoff on +// start failures, and requests leader resignation when all retries are +// exhausted. This is extracted to a separate struct because to get all the edge +// cases right, it ends up being a fair bit of code that would otherwise make +// Client fairly heavy. +type QueueMaintainerLeader struct { + baseservice.BaseService + startstop.BaseStartStop + + // exported for test purposes + TestSignals QueueMaintainerLeaderTestSignals + + config *QueueMaintainerLeaderConfig + + // epoch is incremented each time leadership is gained, giving each start + // goroutine a term number. mu serializes epoch checks with Stop calls so + // a stale goroutine cannot tear down a newer term's maintainer. + epoch int64 + mu sync.Mutex +} + +func NewQueueMaintainerLeader(archetype *baseservice.Archetype, config *QueueMaintainerLeaderConfig) *QueueMaintainerLeader { + return baseservice.Init(archetype, &QueueMaintainerLeader{ + config: config, + }) +} + +func (s *QueueMaintainerLeader) Start(ctx context.Context) error { + ctx, shouldStart, started, stopped := s.StartInit(ctx) + if !shouldStart { + return nil + } + + go func() { + started() + defer stopped() // this defer should come first so it's last out + + sub := s.config.Elector.Listen() + defer sub.Unlisten() + + // Cancel function for an in-progress start attempt. If leadership is + // lost while the start process is still retrying, used to abort it + // promptly instead of waiting for retries to finish. + var cancelStart context.CancelCauseFunc = func(_ error) {} + + // Tracks in-flight tryStart goroutines so we can wait for them to + // finish before returning, preventing logging on a dead test. + var startWg sync.WaitGroup + defer startWg.Wait() + + for { + select { + case <-ctx.Done(): + cancelStart(context.Cause(ctx)) + return + + case notification := <-sub.C(): + s.Logger.DebugContext(ctx, s.Name+": Election change received", + slog.String("client_id", s.config.ClientID), slog.Bool("is_leader", notification.IsLeader)) + + switch { + case notification.IsLeader: + s.TestSignals.ElectedLeader.Signal(struct{}{}) + + // Start with retries in a separate goroutine so the + // leadership change loop remains responsive. + var startCtx context.Context + startCtx, cancelStart = context.WithCancelCause(ctx) + + s.mu.Lock() + s.epoch++ + epoch := s.epoch + s.mu.Unlock() + + startWg.Go(func() { + s.tryStart(startCtx, epoch) + }) + + default: + // Cancel any in-progress start attempts before stopping. + // Send ErrStop so services like Reindexer run cleanup. + cancelStart(startstop.ErrStop) + cancelStart = func(_ error) {} + + s.config.QueueMaintainer.Stop() + } + } + } + }() + + return nil +} + +func (s *QueueMaintainerLeader) tryStart(ctx context.Context, epoch int64) { + var lastErr error + for attempt := 1; attempt <= queueMaintainerMaxStartAttempts; attempt++ { + if ctx.Err() != nil { + return + } + + if lastErr = s.config.QueueMaintainer.Start(ctx); lastErr == nil { + return + } + + s.Logger.ErrorContext(ctx, s.Name+": Error starting queue maintainer", + slog.String("err", lastErr.Error()), slog.Int("attempt", attempt)) + + s.TestSignals.StartError.Signal(lastErr) + + // Stop to fully reset state before retrying. The mutex serializes + // the epoch check with the increment in Start so a stale goroutine + // cannot tear down a newer term's maintainer. + s.mu.Lock() + stale := s.epoch != epoch + if !stale { + s.config.QueueMaintainer.Stop() + } + s.mu.Unlock() + if stale { + return + } + + if attempt < queueMaintainerMaxStartAttempts { + serviceutil.CancellableSleep(ctx, serviceutil.ExponentialBackoff(attempt, serviceutil.MaxAttemptsBeforeResetDefault)) + } + } + + if ctx.Err() != nil { + return + } + + s.Logger.ErrorContext(ctx, s.Name+": Queue maintainer failed to start after all attempts, requesting leader resignation", + slog.String("err", lastErr.Error())) + + s.TestSignals.StartRetriesExhausted.Signal(struct{}{}) + + if err := s.config.RequestResignFunc(ctx); err != nil { + s.Logger.ErrorContext(ctx, s.Name+": Error requesting leader resignation", slog.String("err", err.Error())) + } +} diff --git a/vendor/github.com/riverqueue/river/internal/maintenance/reindexer.go b/vendor/github.com/riverqueue/river/internal/maintenance/reindexer.go new file mode 100644 index 0000000000..cc167164e3 --- /dev/null +++ b/vendor/github.com/riverqueue/river/internal/maintenance/reindexer.go @@ -0,0 +1,297 @@ +package maintenance + +import ( + "cmp" + "context" + "errors" + "log/slog" + "time" + + "github.com/riverqueue/river/riverdriver" + "github.com/riverqueue/river/rivershared/baseservice" + "github.com/riverqueue/river/rivershared/riversharedmaintenance" + "github.com/riverqueue/river/rivershared/startstop" + "github.com/riverqueue/river/rivershared/testsignal" + "github.com/riverqueue/river/rivershared/util/testutil" +) + +const ( + // ReindexerTimeoutDefault is the default timeout of the reindexer. + // + // We've had user reports of builds taking 45 seconds on large tables, so + // set a timeout of that plus a little margin. Use of `CONCURRENTLY` should + // prevent index operations that run a little long from impacting work from + // an operational standpoint. + // + // https://github.com/riverqueue/river/issues/909#issuecomment-2909949466 + ReindexerTimeoutDefault = 1 * time.Minute +) + +// ReindexerTestSignals are internal signals used exclusively in tests. +type ReindexerTestSignals struct { + Reindexed testsignal.TestSignal[struct{}] // notifies when a run finishes executing reindexes for all indexes +} + +func (ts *ReindexerTestSignals) Init(tb testutil.TestingTB) { + ts.Reindexed.Init(tb) +} + +type ReindexerConfig struct { + // IndexNames is the exact list of indexes to reindex on each run. It must + // be non-nil. An empty slice disables reindex work. + IndexNames []string + + // ScheduleFunc returns the next scheduled run time for the reindexer given the + // current time. + ScheduleFunc func(time.Time) time.Time + + // Schema where River tables are located. Empty string omits schema, causing + // Postgres to default to `search_path`. + Schema string + + // Timeout is the amount of time to wait for a single reindex query to run + // before cancelling it via context. + Timeout time.Duration +} + +func (c *ReindexerConfig) mustValidate() *ReindexerConfig { + if c.IndexNames == nil { + panic("ReindexerConfig.IndexNames must be set") + } + if c.ScheduleFunc == nil { + panic("ReindexerConfig.ScheduleFunc must be set") + } + if c.Timeout < -1 { + panic("ReindexerConfig.Timeout must be above zero") + } + + return c +} + +// Reindexer periodically executes a REINDEX command on the important job +// indexes to rebuild them and fix bloat issues. +type Reindexer struct { + riversharedmaintenance.QueueMaintainerServiceBase + startstop.BaseStartStop + + // exported for test purposes + Config *ReindexerConfig + TestSignals ReindexerTestSignals + + exec riverdriver.Executor // driver executor + skipReindexArtifactCheck bool // lets the reindex artifact check be skipped for test purposes +} + +func NewReindexer(archetype *baseservice.Archetype, config *ReindexerConfig, exec riverdriver.Executor) *Reindexer { + if config.IndexNames == nil { + panic("ReindexerConfig.IndexNames must be set") + } + + indexNames := make([]string, len(config.IndexNames)) + copy(indexNames, config.IndexNames) + + scheduleFunc := config.ScheduleFunc + if scheduleFunc == nil { + scheduleFunc = (&DefaultReindexerSchedule{}).Next + } + + return baseservice.Init(archetype, &Reindexer{ + Config: (&ReindexerConfig{ + IndexNames: indexNames, + ScheduleFunc: scheduleFunc, + Schema: config.Schema, + Timeout: cmp.Or(config.Timeout, ReindexerTimeoutDefault), + }).mustValidate(), + + exec: exec, + }) +} + +func (s *Reindexer) Start(ctx context.Context) error { + ctx, shouldStart, started, stopped := s.StartInit(ctx) + if !shouldStart { + return nil + } + + s.StaggerStart(ctx) + + go func() { + started() + defer stopped() // this defer should come first so it's last out + + s.Logger.DebugContext(ctx, s.Name+riversharedmaintenance.LogPrefixRunLoopStarted) + defer s.Logger.DebugContext(ctx, s.Name+riversharedmaintenance.LogPrefixRunLoopStopped) + + nextRunAt := s.Config.ScheduleFunc(time.Now().UTC()) + + s.Logger.DebugContext(ctx, s.Name+": Scheduling first run", slog.Time("next_run_at", nextRunAt)) + + timerUntilNextRun := time.NewTimer(time.Until(nextRunAt)) + scheduleNextRun := func() { + // Advance from the previous scheduled time, not "now", so retries + // stay aligned with the configured cadence and don't immediately + // refire after a timer that has already elapsed. + nextRunAt = s.Config.ScheduleFunc(nextRunAt) + timerUntilNextRun.Reset(time.Until(nextRunAt)) + } + + for { + select { + case <-timerUntilNextRun.C: + reindexableIndexNames, err := s.reindexableIndexNames(ctx) + if err != nil { + if !errors.Is(err, context.Canceled) { + s.Logger.ErrorContext(ctx, s.Name+": Error listing reindexable indexes", slog.String("error", err.Error())) + } + scheduleNextRun() + continue + } + + for _, indexName := range reindexableIndexNames { + if _, err := s.reindexOne(ctx, indexName); err != nil { + if !errors.Is(err, context.Canceled) { + s.Logger.ErrorContext(ctx, s.Name+": Error reindexing", slog.String("error", err.Error()), slog.String("index_name", indexName)) + } + continue + } + } + + s.TestSignals.Reindexed.Signal(struct{}{}) + + // On each run, we calculate the new schedule based on the + // previous run's start time. This ensures that we don't + // accidentally skip a run as time elapses during the run. + scheduleNextRun() + + // TODO: maybe we should log differently if some of these fail? + s.Logger.DebugContext(ctx, s.Name+riversharedmaintenance.LogPrefixRanSuccessfully, + slog.Time("next_run_at", nextRunAt), slog.Int("num_reindexes_initiated", len(reindexableIndexNames))) + + case <-ctx.Done(): + // Clean up timer resources. We know it has _not_ received from + // the timer since its last reset because that would have led us + // to the case above instead of here. + if !timerUntilNextRun.Stop() { + <-timerUntilNextRun.C + } + return + } + } + }() + + return nil +} + +func (s *Reindexer) reindexableIndexNames(ctx context.Context) ([]string, error) { + indexesExist, err := s.exec.IndexesExist(ctx, &riverdriver.IndexesExistParams{ + IndexNames: s.Config.IndexNames, + Schema: s.Config.Schema, + }) + if err != nil { + return nil, err + } + + indexNames := make([]string, 0, len(s.Config.IndexNames)) + missingIndexNames := make([]string, 0) + for _, indexName := range s.Config.IndexNames { + if indexesExist[indexName] { + indexNames = append(indexNames, indexName) + continue + } + + missingIndexNames = append(missingIndexNames, indexName) + } + + if len(missingIndexNames) > 0 { + s.Logger.WarnContext(ctx, s.Name+": Configured reindex indexes do not exist; run migrations or update ReindexerIndexNames", + slog.Any("index_names", missingIndexNames)) + } + + return indexNames, nil +} + +func (s *Reindexer) reindexOne(ctx context.Context, indexName string) (bool, error) { + var cancel func() + if s.Config.Timeout > -1 { + ctx, cancel = context.WithTimeout(ctx, s.Config.Timeout) + defer cancel() + } + + // Make sure that no `CONCURRENTLY` artifacts from a previous reindexing run + // exist before trying to reindex. When using `CONCURRENTLY`, Postgres + // creates a new index suffixed with `_ccnew` before swapping it in as the + // new index. The existing index is renamed `_ccold` before being dropped + // concurrently. If multiple failed artifacts exist, Postgres may add a + // numeric suffix like `_ccnew1` or `_ccold2` to keep names unique. + // + // If one of these artifacts exists, it probably means that a previous + // reindex attempt timed out, and attempting to reindex again is likely + // slated for the same fate. We opt to log a warning and no op instead of + // trying to clean up the artifacts of a previously failed run for the same + // reason: even with the artifacts removed, if a previous reindex failed + // then a new one is likely to as well, so cleaning up would result in a + // forever loop of failed index builds that'd put unnecessary pressure on + // the underlying database. + // + // https://www.postgresql.org/docs/current/sql-reindex.html#SQL-REINDEX-CONCURRENTLY + if !s.skipReindexArtifactCheck { + reindexArtifactNames, err := s.exec.IndexReindexArtifacts(ctx, &riverdriver.IndexReindexArtifactsParams{Index: indexName, Schema: s.Config.Schema}) + if err != nil { + return false, err + } + + if len(reindexArtifactNames) > 0 { + s.Logger.WarnContext(ctx, s.Name+": Found reindex artifact likely resulting from previous partially completed reindex attempt; skipping reindex", + slog.Any("artifact_names", reindexArtifactNames), slog.String("index_name", indexName), slog.Duration("timeout", s.Config.Timeout)) + return false, nil + } + } + + if err := s.exec.IndexReindex(ctx, &riverdriver.IndexReindexParams{Index: indexName, Schema: s.Config.Schema}); err != nil { + // This should be quite rare because the reindexer has a slow run + // period, but it's possible for the reindexer to be stopped while it's + // trying to rebuild an index, and doing so would normally put in the + // reindexer into permanent purgatory because the cancellation would + // leave a concurrent index artifact which would cause the reindexer to + // skip work on future runs. + // + // So here, in the case of a cancellation due to stop, take a little + // extra time to drop any artifacts that may result from the cancelled + // build. This will slow shutdown somewhat, but should still be + // reasonably fast since we're only dropping indexes rather than + // building them. + if errors.Is(context.Cause(ctx), startstop.ErrStop) { + ctx := context.WithoutCancel(ctx) + + ctx, cancel = context.WithTimeout(ctx, 15*time.Second) + defer cancel() + + s.Logger.InfoContext(ctx, s.Name+": Signaled to stop during index build; attempting to clean up concurrent artifacts") + + reindexArtifactNames, err := s.exec.IndexReindexArtifacts(ctx, &riverdriver.IndexReindexArtifactsParams{Index: indexName, Schema: s.Config.Schema}) + if err != nil { + s.Logger.ErrorContext(ctx, s.Name+": Error listing reindex artifacts", slog.String("error", err.Error())) + } + + for _, reindexArtifactName := range reindexArtifactNames { + if err := s.exec.IndexDropIfExists(ctx, &riverdriver.IndexDropIfExistsParams{Index: reindexArtifactName, Schema: s.Config.Schema}); err != nil { + s.Logger.ErrorContext(ctx, s.Name+": Error dropping reindex artifact", slog.String("artifact_name", reindexArtifactName), slog.String("error", err.Error())) + } + } + } + + return false, err + } + + s.Logger.InfoContext(ctx, s.Name+": Initiated reindex", slog.String("index_name", indexName)) + return true, nil +} + +// DefaultReindexerSchedule is a default schedule for the reindexer job which +// runs at midnight UTC daily. +type DefaultReindexerSchedule struct{} + +// Next returns the next scheduled time for the reindexer job. +func (s *DefaultReindexerSchedule) Next(t time.Time) time.Time { + return t.Add(24 * time.Hour).Truncate(24 * time.Hour) +} diff --git a/vendor/github.com/riverqueue/river/internal/maintenance/sqlite_notification_cleaner.go b/vendor/github.com/riverqueue/river/internal/maintenance/sqlite_notification_cleaner.go new file mode 100644 index 0000000000..de9e9526db --- /dev/null +++ b/vendor/github.com/riverqueue/river/internal/maintenance/sqlite_notification_cleaner.go @@ -0,0 +1,152 @@ +package maintenance + +import ( + "cmp" + "context" + "errors" + "log/slog" + "time" + + "github.com/riverqueue/river/riverdriver" + "github.com/riverqueue/river/rivershared/baseservice" + "github.com/riverqueue/river/rivershared/riversharedmaintenance" + "github.com/riverqueue/river/rivershared/startstop" + "github.com/riverqueue/river/rivershared/testsignal" + "github.com/riverqueue/river/rivershared/util/testutil" + "github.com/riverqueue/river/rivershared/util/timeoututil" + "github.com/riverqueue/river/rivershared/util/timeutil" +) + +const ( + SQLiteNotificationCleanerIntervalDefault = time.Minute + SQLiteNotificationCleanerRetentionPeriodDefault = 5 * time.Minute +) + +// SQLiteNotificationCleanerTestSignals are internal signals used exclusively in tests. +type SQLiteNotificationCleanerTestSignals struct { + DeletedBatch testsignal.TestSignal[struct{}] // notifies when runOnce finishes a pass +} + +func (ts *SQLiteNotificationCleanerTestSignals) Init(tb testutil.TestingTB) { + ts.DeletedBatch.Init(tb) +} + +type SQLiteNotificationCleanerConfig struct { + // Interval is the amount of time to wait between cleaner runs. + Interval time.Duration + + // RetentionPeriod is the amount of time to keep notification rows around + // before they're removed. + RetentionPeriod time.Duration + + // Schema where River tables are located. Empty string omits schema. + Schema string + + // Timeout is the timeout for each delete query. + Timeout time.Duration +} + +func (c *SQLiteNotificationCleanerConfig) mustValidate() *SQLiteNotificationCleanerConfig { + if c.Interval <= 0 { + panic("SQLiteNotificationCleanerConfig.Interval must be above zero") + } + if c.RetentionPeriod <= 0 { + panic("SQLiteNotificationCleanerConfig.RetentionPeriod must be above zero") + } + if c.Timeout <= 0 { + panic("SQLiteNotificationCleanerConfig.Timeout must be above zero") + } + + return c +} + +// SQLiteNotificationCleaner periodically removes old rows from SQLite's +// notification outbox. It is only needed for the SQLite driver's emulated +// listen/notify support. +type SQLiteNotificationCleaner struct { + riversharedmaintenance.QueueMaintainerServiceBase + startstop.BaseStartStop + + // exported for test purposes + Config *SQLiteNotificationCleanerConfig + TestSignals SQLiteNotificationCleanerTestSignals + + exec riverdriver.Executor +} + +// NewSQLiteNotificationCleaner returns a SQLite notification cleaner. +func NewSQLiteNotificationCleaner(archetype *baseservice.Archetype, config *SQLiteNotificationCleanerConfig, exec riverdriver.Executor) *SQLiteNotificationCleaner { + return baseservice.Init(archetype, &SQLiteNotificationCleaner{ + Config: (&SQLiteNotificationCleanerConfig{ + Interval: cmp.Or(config.Interval, SQLiteNotificationCleanerIntervalDefault), + RetentionPeriod: cmp.Or(config.RetentionPeriod, SQLiteNotificationCleanerRetentionPeriodDefault), + Schema: config.Schema, + Timeout: cmp.Or(config.Timeout, riversharedmaintenance.TimeoutDefault), + }).mustValidate(), + exec: exec, + }) +} + +func (s *SQLiteNotificationCleaner) Start(ctx context.Context) error { //nolint:dupl + ctx, shouldStart, started, stopped := s.StartInit(ctx) + if !shouldStart { + return nil + } + + s.StaggerStart(ctx) + + go func() { + started() + defer stopped() // this defer should come first so it's last out + + s.Logger.DebugContext(ctx, s.Name+riversharedmaintenance.LogPrefixRunLoopStarted) + defer s.Logger.DebugContext(ctx, s.Name+riversharedmaintenance.LogPrefixRunLoopStopped) + + ticker := timeutil.NewTickerWithInitialTick(ctx, s.Config.Interval) + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + } + + res, err := s.runOnce(ctx) + if err != nil { + if !errors.Is(err, context.Canceled) { + s.Logger.ErrorContext(ctx, s.Name+": Error cleaning SQLite notifications", slog.String("error", err.Error())) + } + continue + } + + if res.NumNotificationsDeleted > 0 { + s.Logger.InfoContext(ctx, s.Name+riversharedmaintenance.LogPrefixRanSuccessfully, + slog.Int("num_notifications_deleted", res.NumNotificationsDeleted), + ) + } + } + }() + + return nil +} + +type sqliteNotificationCleanerRunOnceResult struct { + NumNotificationsDeleted int +} + +func (s *SQLiteNotificationCleaner) runOnce(ctx context.Context) (*sqliteNotificationCleanerRunOnceResult, error) { + return timeoututil.WithTimeoutV(ctx, s.Config.Timeout, s.Name+".runOnce", func(ctx context.Context) (*sqliteNotificationCleanerRunOnceResult, error) { + numDeleted, err := s.exec.NotificationDeleteBefore(ctx, &riverdriver.NotificationDeleteBeforeParams{ + CreatedAtHorizon: time.Now().Add(-s.Config.RetentionPeriod), + Schema: s.Config.Schema, + }) + if err != nil { + return nil, err + } + + s.TestSignals.DeletedBatch.Signal(struct{}{}) + + return &sqliteNotificationCleanerRunOnceResult{ + NumNotificationsDeleted: numDeleted, + }, nil + }) +} diff --git a/vendor/github.com/riverqueue/river/internal/notifier/notifier.go b/vendor/github.com/riverqueue/river/internal/notifier/notifier.go new file mode 100644 index 0000000000..8b48fbdd88 --- /dev/null +++ b/vendor/github.com/riverqueue/river/internal/notifier/notifier.go @@ -0,0 +1,604 @@ +package notifier + +import ( + "cmp" + "context" + "errors" + "fmt" + "log/slog" + "slices" + "strings" + "sync" + "time" + + "github.com/riverqueue/river/riverdriver" + "github.com/riverqueue/river/rivershared/baseservice" + "github.com/riverqueue/river/rivershared/startstop" + "github.com/riverqueue/river/rivershared/testsignal" + "github.com/riverqueue/river/rivershared/util/maputil" + "github.com/riverqueue/river/rivershared/util/serviceutil" + "github.com/riverqueue/river/rivershared/util/sliceutil" + "github.com/riverqueue/river/rivershared/util/testutil" + "github.com/riverqueue/river/rivershared/util/timeoututil" +) + +type NotificationTopic string + +const ( + NotificationTopicControl NotificationTopic = "river_control" + NotificationTopicInsert NotificationTopic = "river_insert" + NotificationTopicLeadership NotificationTopic = "river_leadership" +) + +var notificationTopicAll = []NotificationTopic{ //nolint:gochecknoglobals + NotificationTopicControl, + NotificationTopicInsert, + NotificationTopicLeadership, +} + +// NotificationTopicLongest is just the longest notification topic. This is used +// to determine the maximum length of allowed custom schema names because +// schemas are prefixed to notification topic names and Postgres enforces a +// maximum topic length of 63 characters. +var NotificationTopicLongest = func() NotificationTopic { //nolint:gochecknoglobals + return slices.MaxFunc(notificationTopicAll, func(t1, t2 NotificationTopic) int { + return len(string(t1)) - len(string(t2)) + }) +}() + +type NotifyFunc func(topic NotificationTopic, payload string) + +type Subscription struct { + notifyFunc NotifyFunc + notifier *Notifier + topic NotificationTopic + unlistenOnce sync.Once +} + +func (s *Subscription) Unlisten(ctx context.Context) { + s.unlistenOnce.Do(func() { + // Unlisten strips cancellation from the parent context to ensure it runs: + if err := s.notifier.unlisten(context.WithoutCancel(ctx), s); err != nil { + s.notifier.Logger.ErrorContext(ctx, s.notifier.Name+": Error unlistening on topic", "err", err, "topic", s.topic) + } + }) +} + +// Test-only properties. +type notifierTestSignals struct { + BackoffError testsignal.TestSignal[error] // non-cancellation error received by main run loop + ListeningBegin testsignal.TestSignal[struct{}] // notifier has entered a listen loop + ListeningEnd testsignal.TestSignal[struct{}] // notifier has left a listen loop +} + +func (ts *notifierTestSignals) Init(tb testutil.TestingTB) { + ts.BackoffError.Init(tb) + ts.ListeningBegin.Init(tb) + ts.ListeningEnd.Init(tb) +} + +type Notifier struct { + baseservice.BaseService + startstop.BaseStartStop + + listener riverdriver.Listener + notificationBuf chan *riverdriver.Notification + testDisableSleep bool // for tests only; disable sleep on exponential backoff + testPingInterval time.Duration // for tests only; override the 5s ping interval + testSignals notifierTestSignals + waitInterruptChan chan func() + + mu sync.RWMutex + isConnected bool + isStarted bool + isWaiting bool + subscriptions map[NotificationTopic][]*Subscription + waitCancel context.CancelFunc +} + +func New(archetype *baseservice.Archetype, listener riverdriver.Listener) *Notifier { + notifier := baseservice.Init(archetype, &Notifier{ + listener: listener, + notificationBuf: make(chan *riverdriver.Notification, 1000), + waitInterruptChan: make(chan func(), 10), + + subscriptions: make(map[NotificationTopic][]*Subscription), + }) + return notifier +} + +func (n *Notifier) Start(ctx context.Context) error { + ctx, shouldStart, started, stopped := n.StartInit(ctx) + if !shouldStart { + return nil + } + + // The loop below will connect/close on every iteration, but do one initial + // connect so the notifier fails fast in case of an obvious problem. + if err := n.listenerConnect(ctx, false); err != nil { + stopped() + if errors.Is(err, context.Canceled) { + return nil + } + return err + } + + go func() { + started() + defer stopped() + + n.Logger.DebugContext(ctx, n.Name+": Run loop started") + defer n.Logger.DebugContext(ctx, n.Name+": Run loop stopped") + + n.withLock(func() { n.isStarted = true }) + defer n.withLock(func() { n.isStarted = false }) + + defer n.listenerClose(ctx, false) + + var wg sync.WaitGroup + + wg.Go(func() { + n.deliverNotifications(ctx) + }) + + for attempt := 0; ; attempt++ { + if err := n.listenAndWait(ctx); err != nil { + if errors.Is(err, context.Canceled) { + break + } + + sleepDuration := serviceutil.ExponentialBackoff(attempt, serviceutil.MaxAttemptsBeforeResetDefault) + n.Logger.ErrorContext(ctx, n.Name+": Error running listener (will attempt reconnect after backoff)", + slog.Int("attempt", attempt), + slog.String("err", err.Error()), + slog.String("sleep_duration", sleepDuration.String()), + ) + n.testSignals.BackoffError.Signal(err) + if !n.testDisableSleep { + serviceutil.CancellableSleep(ctx, sleepDuration) + } + } + } + + wg.Wait() + }() + + return nil +} + +func (n *Notifier) deliverNotifications(ctx context.Context) { + for { + select { + case <-ctx.Done(): + return + + case notification := <-n.notificationBuf: + notifyFuncs := func() []NotifyFunc { + n.mu.RLock() + defer n.mu.RUnlock() + + return sliceutil.Map(n.subscriptions[NotificationTopic(notification.Topic)], func(s *Subscription) NotifyFunc { return s.notifyFunc }) + }() + + for _, notifyFunc := range notifyFuncs { + // TODO: panic recovery on delivery attempts + notifyFunc(NotificationTopic(notification.Topic), notification.Payload) + } + } + } +} + +func (n *Notifier) listenAndWait(ctx context.Context) error { + if err := n.listenerConnect(ctx, false); err != nil { + return err + } + defer n.listenerClose(ctx, false) + + topics := func() []NotificationTopic { + n.mu.RLock() + defer n.mu.RUnlock() + + return maputil.Keys(n.subscriptions) + }() + + for _, topic := range topics { + if err := n.listenerListen(ctx, topic); err != nil { + return err + } + } + + n.Logger.DebugContext(ctx, n.Name+": Notifier healthy") + + n.testSignals.ListeningBegin.Signal(struct{}{}) + defer n.testSignals.ListeningEnd.Signal(struct{}{}) + + drainInterrupts := func() { + for { + select { + case interruptOperation := <-n.waitInterruptChan: + interruptOperation() + default: + return + } + } + } + + // Drain interrupts one last time before leaving to make sure we're not + // leaving any goroutines hanging anywhere. + defer drainInterrupts() + + for { + // Top level context is done, meaning we're shutting down. + if ctx.Err() != nil { + return ctx.Err() + } + + // Drain any and all interrupt operations before continuing back into a + // new wait to give any new subscribers a chance to listen/unlisten. + drainInterrupts() + + err := n.waitOnce(ctx) + if err != nil { + // On cancellation, reenter loop, but the check at the top on + // `ctx.Err()` will end it if the service is shutting down. + if errors.Is(err, context.Canceled) { + continue + } + + n.Logger.InfoContext(ctx, n.Name+": Notifier unhealthy") + + return err + } + } +} + +func (n *Notifier) listenerClose(ctx context.Context, skipLock bool) { + if !skipLock { + n.mu.Lock() + defer n.mu.Unlock() + } + + if !n.isConnected { + return + } + + n.Logger.DebugContext(ctx, n.Name+": Listener closing") + if err := n.listener.Close(ctx); err != nil { + if !shouldIgnoreListenerError(err) { + n.Logger.ErrorContext(ctx, n.Name+": Error closing listener", "err", err) + } + } + + n.isConnected = false +} + +// shouldIgnoreListenerError returns true if the error is a certain type of +// common error that we see when trying to close a listener. +// +// It's probably not strictly necessary to log errors on listener close, but it +// seems not ideal to ignore them completely either. If this function were ever +// to become to unwieldy though, we might want to go back to the drawing board. +func shouldIgnoreListenerError(err error) bool { + if errors.Is(err, context.Canceled) { + return true + } + + // In practice, this occurs a fair bit in some systems. See: + // + // https://github.com/riverqueue/river/issues/256 + // + // There's been an issue opened to make this a well-known error type, but + // it's not right now: + // + // https://github.com/golang/go/issues/75600 + if strings.Contains(err.Error(), "tls: failed to send closeNotify alert") { + return true + } + + return false +} + +const listenerTimeout = 10 * time.Second + +func (n *Notifier) listenerConnect(ctx context.Context, skipLock bool) error { + if !skipLock { + n.mu.Lock() + defer n.mu.Unlock() + } + + if n.isConnected { + return nil + } + + return timeoututil.WithTimeout(ctx, listenerTimeout, n.Name+".listenerConnect", func(ctx context.Context) error { + n.Logger.DebugContext(ctx, n.Name+": Listener connecting") + if err := n.listener.Connect(ctx); err != nil { + if !errors.Is(err, context.Canceled) { + n.Logger.ErrorContext(ctx, n.Name+": Error connecting listener", "err", err) + } + + return err + } + + n.isConnected = true + return nil + }) +} + +// Listens on a topic with an appropriate logging statement. Should be preferred +// to `listener.Listen` for improved logging/telemetry. +// +// Not protected by mutex because it doesn't modify any notifier state and the +// underlying listener has a mutex around its operations. +func (n *Notifier) listenerListen(ctx context.Context, topic NotificationTopic) error { + return timeoututil.WithTimeout(ctx, listenerTimeout, n.Name+".listenerListen", func(ctx context.Context) error { + n.Logger.DebugContext(ctx, n.Name+": Listening on topic", "topic", topic) + if err := n.listener.Listen(ctx, string(topic)); err != nil { + return fmt.Errorf("error listening on topic %q: %w", topic, err) + } + + return nil + }) +} + +// Unlistens on a topic with an appropriate logging statement. Should be +// preferred to `listener.Unlisten` for improved logging/telemetry. +// +// Not protected by mutex because it doesn't modify any notifier state and the +// underlying listener has a mutex around its operations. +func (n *Notifier) listenerUnlisten(ctx context.Context, topic NotificationTopic) error { + return timeoututil.WithTimeout(ctx, listenerTimeout, n.Name+".listenerUnlisten", func(ctx context.Context) error { + n.Logger.DebugContext(ctx, n.Name+": Unlistening on topic", "topic", topic) + if err := n.listener.Unlisten(ctx, string(topic)); err != nil { + return fmt.Errorf("error unlistening on topic %q: %w", topic, err) + } + + return nil + }) +} + +// Enters a single blocking wait for notifications on the underlying listener. +// Waiting for a notification locks an underlying connection, so infrastructure +// elsewhere in the notifier must preempt it by sending to `n.waitInterruptChan` +// and invoking `n.waitCancel()`. Cancelling the input context (as occurs during +// shutdown) also unblocks the wait. +func (n *Notifier) waitOnce(ctx context.Context) error { + n.withLock(func() { + n.isWaiting = true + ctx, n.waitCancel = context.WithCancel(ctx) //nolint:fatcontext + }) + defer n.withLock(func() { + n.isWaiting = false + n.waitCancel() + }) + + // Save a reference to the parent context before creating the inner + // cancellable context. The inner context is cancelled by drainErrChan to + // interrupt WaitForNotification, but we still need a live context for the + // Ping health check afterward. + pingCtx := ctx + + ctx, cancel := context.WithCancel(ctx) + defer cancel() + + errChan := make(chan error) + + go func() { + for { + notification, err := n.listener.WaitForNotification(ctx) + if err != nil { + errChan <- err + return + } + + select { + case n.notificationBuf <- notification: + default: + n.Logger.WarnContext(ctx, n.Name+": Dropping notification due to full buffer", "payload", notification.Payload) + } + } + }() + + drainErrChan := func() error { + cancel() + + // There's a chance we encounter some other error before the context.Canceled comes in: + err := <-errChan + if err != nil && !errors.Is(err, context.Canceled) { + // A non-cancel error means something went wrong with the conn, so we should bail. + n.Logger.ErrorContext(ctx, n.Name+": Error on draining notification wait", "err", err) + return err + } + // If we got a context cancellation error, it means we successfully + // interrupted the WaitForNotification so that we could make the + // subscription change. + return nil + } + + pingInterval := cmp.Or(n.testPingInterval, 5*time.Second) + needPingCtx, needPingCancel := context.WithTimeout(ctx, pingInterval) + defer needPingCancel() + + // * Wait for notifications + // * Ping conn if 5 seconds have elapsed between notifications to keep it alive + // * Manage listens/unlistens on conn (waitInterruptChan) + // * If any errors are encountered, return them so we can kill the conn and start over + select { + case <-ctx.Done(): + return <-errChan + + case <-needPingCtx.Done(): + if err := drainErrChan(); err != nil { + return err + } + // Ping the conn to see if it's still alive. Use pingCtx (the parent + // context) because the inner ctx was cancelled by drainErrChan above + // to interrupt WaitForNotification. + // + // Note: Previously this used the (already cancelled) inner ctx, making + // the ping a no-op that always returned context.Canceled. With the fix, + // dead or flaky connections are now actively detected, which may trigger + // reconnections that were previously silently swallowed. + if err := n.listener.Ping(pingCtx); err != nil { + return err + } + + case err := <-errChan: + if errors.Is(err, context.Canceled) { + return nil + } + if err != nil { + n.Logger.ErrorContext(ctx, n.Name+": Error from notification wait", "err", err) + return err + } + } + + return nil +} + +// Sends an interrupt operation to the main loop, waits on the result, and +// returns an error if there was one. +// +// MUST be called with the `n.mu` mutex already locked. +func (n *Notifier) sendInterruptAndReceiveResult(operation func() error) error { + errChan := make(chan error) + n.waitInterruptChan <- func() { + errChan <- operation() + } + + n.waitCancel() + + // Notably, these unlock then lock again, the reverse of what you'd normally + // expect in a mutex pattern. This is because this function is only expected + // to be called with the mutex already locked, but we need to unlock it to + // give the main loop a chance to run interrupt operations. + n.mu.Unlock() + defer n.mu.Lock() + + select { + case err := <-errChan: + return err + case <-time.After(5 * time.Second): + return errors.New("timed out waiting for interrupt operation") + } +} + +func (n *Notifier) Listen(ctx context.Context, topic NotificationTopic, notifyFunc NotifyFunc) (*Subscription, error) { + n.mu.Lock() + defer n.mu.Unlock() + + sub := &Subscription{ + notifyFunc: notifyFunc, + topic: topic, + notifier: n, + } + + existingSubs, existingTopic := n.subscriptions[topic] + if !existingTopic { + existingSubs = make([]*Subscription, 0, 10) + } + n.subscriptions[topic] = append(existingSubs, sub) + + n.Logger.DebugContext(ctx, n.Name+": Added subscription", "new_num_subscriptions", len(n.subscriptions[topic]), "topic", topic) + + // We add the new subscription to the subscription list optimistically, and + // it needs to be done this way in case of a restart after an interrupt + // below has been run, but after a return to this function (say we were to + // add the new sub at the end of this function, it would not be picked + // during the restart). But in case of an error subscribing, remove the sub. + // + // By the time this function is run (i.e. after an interrupt), a lock on + // `n.mu` has been reacquired, and modifying subscription state is safe. + removeSub := func() { n.removeSubscription(ctx, sub) } + + if !existingTopic { + // If already waiting, send an interrupt to the wait function to run a + // listen operation. If not, connect and listen directly, returning any + // errors as feedback to the caller. + if n.isWaiting { + if err := n.sendInterruptAndReceiveResult(func() error { return n.listenerListen(ctx, topic) }); err != nil { + removeSub() + return nil, err + } + } else { + var justConnected bool + + if !n.isConnected { + if err := n.listenerConnect(ctx, true); err != nil { + removeSub() + return nil, err + } + justConnected = true + } + + if err := n.listenerListen(ctx, topic); err != nil { + removeSub() + + // If we just connected above and the notifier hasn't started in + // the interim, also close the connection so we don't leave any + // resources hanging. + if justConnected && !n.isStarted { + n.listenerClose(ctx, true) + } + + return nil, err + } + } + } + + return sub, nil +} + +func (n *Notifier) unlisten(ctx context.Context, sub *Subscription) error { + n.mu.Lock() + defer n.mu.Unlock() + + subs := n.subscriptions[sub.topic] + + // If this is the last subscription on the topic, unlisten if we're connected. + if len(subs) <= 1 { + // If already waiting, send an interrupt to the wait function to run an + // unlisten operation. If not, if connected, unlisten directly. + if n.isWaiting { + if err := n.sendInterruptAndReceiveResult(func() error { return n.listenerUnlisten(ctx, sub.topic) }); err != nil { + return err + } + } else { + if n.isConnected { + if err := n.listenerUnlisten(ctx, sub.topic); err != nil { + return err + } + + // If this was the last subscription, we weren't in a wait loop, + // and the notifier never started, also clean up by closing the + // listener. + if !n.isStarted && len(n.subscriptions) <= 1 { + n.listenerClose(ctx, true) + } + } + } + } + + n.removeSubscription(ctx, sub) + + return nil +} + +// This function requires that the caller already has a lock on `n.mu`. +func (n *Notifier) removeSubscription(ctx context.Context, sub *Subscription) { + n.subscriptions[sub.topic] = slices.DeleteFunc(n.subscriptions[sub.topic], func(s *Subscription) bool { + return s == sub + }) + + if len(n.subscriptions[sub.topic]) < 1 { + delete(n.subscriptions, sub.topic) + } + + n.Logger.DebugContext(ctx, n.Name+": Removed subscription", "new_num_subscriptions", len(n.subscriptions[sub.topic]), "topic", sub.topic) +} + +func (n *Notifier) withLock(lockedFunc func()) { + n.mu.Lock() + defer n.mu.Unlock() + lockedFunc() +} diff --git a/vendor/github.com/riverqueue/river/internal/notifylimiter/limiter.go b/vendor/github.com/riverqueue/river/internal/notifylimiter/limiter.go new file mode 100644 index 0000000000..04598912be --- /dev/null +++ b/vendor/github.com/riverqueue/river/internal/notifylimiter/limiter.go @@ -0,0 +1,44 @@ +package notifylimiter + +import ( + "sync" + "time" + + "github.com/riverqueue/river/rivershared/baseservice" +) + +type NotifyFunc func(name string) + +type Limiter struct { + baseservice.BaseService + + waitDuration time.Duration + + mu sync.Mutex // protects lastSentByTopic + lastSentByTopic map[string]time.Time +} + +// NewLimiter creates a new Limiter, calling the NotifyFunc no more than once per waitDuration. +// The function must be fast as it is called within a mutex. +func NewLimiter(archetype *baseservice.Archetype, waitDuration time.Duration) *Limiter { + return baseservice.Init(archetype, &Limiter{ + lastSentByTopic: make(map[string]time.Time), + waitDuration: waitDuration, + }) +} + +func (l *Limiter) ShouldTrigger(topic string) bool { + // Calculate this beforehand to reduce mutex duration. + now := l.Time.Now() + lastSentHorizon := now.Add(-l.waitDuration) + + l.mu.Lock() + defer l.mu.Unlock() + + if l.lastSentByTopic[topic].Before(lastSentHorizon) { + l.lastSentByTopic[topic] = now + return true + } + + return false +} diff --git a/vendor/github.com/riverqueue/river/internal/pluginconfig/plugin_config.go b/vendor/github.com/riverqueue/river/internal/pluginconfig/plugin_config.go new file mode 100644 index 0000000000..4431180950 --- /dev/null +++ b/vendor/github.com/riverqueue/river/internal/pluginconfig/plugin_config.go @@ -0,0 +1,39 @@ +package pluginconfig + +import "github.com/riverqueue/river/rivertype" + +// CombinedMiddleware combines middleware from the current and legacy +// configuration fields into one list. Explicit middleware is preserved first, +// followed by legacy job insert and worker middleware. +// +// The reason this exists is that River originally had a seprate configuration +// field for each type of middleware (JobInsertMiddleware, WorkerMiddleware) +// before it was unified to a combined Middleware field so that a new one +// wouldn't be needed every time a new type of middleware was added. Middleware +// then became Plugins. The JobInsertMiddleware and WorkerMiddleware fields have +// been deprecated for quite some time and we should consider just removing them +// to simplify this set up. +func CombinedMiddleware(middleware []rivertype.Middleware, jobInsertMiddleware []rivertype.JobInsertMiddleware, workerMiddleware []rivertype.WorkerMiddleware) []rivertype.Middleware { + allMiddleware := make([]rivertype.Middleware, 0, + len(middleware)+len(jobInsertMiddleware)+len(workerMiddleware)) + allMiddleware = append(allMiddleware, middleware...) + + for _, jobInsertMiddlewareItem := range jobInsertMiddleware { + allMiddleware = append(allMiddleware, jobInsertMiddlewareItem) + } + +outerLoop: + for _, workerMiddlewareItem := range workerMiddleware { + if workerMiddlewareAsJobInsertMiddleware, ok := workerMiddlewareItem.(rivertype.JobInsertMiddleware); ok { + for _, jobInsertMiddlewareItem := range jobInsertMiddleware { + if workerMiddlewareAsJobInsertMiddleware == jobInsertMiddlewareItem { + continue outerLoop + } + } + } + + allMiddleware = append(allMiddleware, workerMiddlewareItem) + } + + return allMiddleware +} diff --git a/vendor/github.com/riverqueue/river/internal/pluginlookup/plugin_lookup.go b/vendor/github.com/riverqueue/river/internal/pluginlookup/plugin_lookup.go new file mode 100644 index 0000000000..5ce2753477 --- /dev/null +++ b/vendor/github.com/riverqueue/river/internal/pluginlookup/plugin_lookup.go @@ -0,0 +1,208 @@ +package pluginlookup + +import ( + "sync" + + "github.com/riverqueue/river/rivershared/baseservice" + "github.com/riverqueue/river/rivertype" +) + +// +// PluginKind +// + +type PluginKind string + +const ( + PluginKindHookInsertBegin PluginKind = "hook_insert_begin" + PluginKindHookMetricEmit PluginKind = "hook_metric_emit" + PluginKindHookPeriodicJobsStart PluginKind = "hook_periodic_jobs_start" + PluginKindHookWorkBegin PluginKind = "hook_work_begin" + PluginKindHookWorkEnd PluginKind = "hook_work_end" + PluginKindMiddlewareJobInsert PluginKind = "middleware_job_insert" + PluginKindMiddlewareWorker PluginKind = "middleware_worker" +) + +// +// PluginLookup +// + +// PluginLookup looks up plugins by kind. Its zero value is an empty lookup. +type PluginLookup struct { + hooks []rivertype.Hook + pluginsByKind map[PluginKind][]any +} + +// NewPluginLookup returns a new plugin lookup based on the given plugins. Each +// input is considered for every hook and middleware kind it implements. +// +// The plugins parameter is []any rather than []rivertype.Plugin because the +// lookup may contain legacy hooks and middleware that don't implement +// rivertype.Plugin. Keeping their original concrete values avoids compatibility +// wrappers that would need to forward every operation-specific interface. +func NewPluginLookup(plugins []any) *PluginLookup { + return newPluginLookup(plugins, plugins) +} + +// NewPluginLookupFromConfig returns a plugin lookup from separately configured +// hooks, middleware, and plugins. Explicit plugins may participate as either +// hooks or middleware, while entries from the legacy Hooks and Middleware +// configuration fields participate only as the kind they were configured as. +// Base services embedded in any configured extension are initialized with +// archetype when it's non-nil. +func NewPluginLookupFromConfig(archetype *baseservice.Archetype, hooks []rivertype.Hook, middlewares []rivertype.Middleware, plugins []rivertype.Plugin) *PluginLookup { + if archetype != nil { + initBaseServices(archetype, hooks) + initBaseServices(archetype, middlewares) + initBaseServices(archetype, plugins) + } + + pluginValues := toAnySlice(plugins) + + hookValues := make([]any, 0, len(plugins)+len(hooks)) + hookValues = append(hookValues, pluginValues...) + hookValues = append(hookValues, toAnySlice(hooks)...) + + middlewareValues := make([]any, 0, len(plugins)+len(middlewares)) + middlewareValues = append(middlewareValues, pluginValues...) + middlewareValues = append(middlewareValues, toAnySlice(middlewares)...) + + return newPluginLookup(hookValues, middlewareValues) +} + +func (c *PluginLookup) ByKind(kind PluginKind) []any { + return c.pluginsByKind[kind] +} + +// Hooks returns all the hooks in the lookup in configuration order. +func (c *PluginLookup) Hooks() []rivertype.Hook { + return c.hooks +} + +func initBaseServices[T any](archetype *baseservice.Archetype, plugins []T) { + for _, plugin := range plugins { + if withBaseService, ok := any(plugin).(baseservice.WithBaseService); ok { + baseservice.Init(archetype, withBaseService) + } + } +} + +func newPluginLookup(hooks, middlewares []any) *PluginLookup { + lookup := &PluginLookup{} + if len(hooks) < 1 && len(middlewares) < 1 { + return lookup + } + + lookup.pluginsByKind = make(map[PluginKind][]any) + + for _, plugin := range hooks { + if plugin == nil { + continue + } + + if hook, ok := plugin.(rivertype.Hook); ok { + lookup.hooks = append(lookup.hooks, hook) + } + if _, ok := plugin.(rivertype.HookInsertBegin); ok { + lookup.pluginsByKind[PluginKindHookInsertBegin] = append(lookup.pluginsByKind[PluginKindHookInsertBegin], plugin) + } + if _, ok := plugin.(rivertype.HookMetricEmit); ok { + lookup.pluginsByKind[PluginKindHookMetricEmit] = append(lookup.pluginsByKind[PluginKindHookMetricEmit], plugin) + } + if _, ok := plugin.(rivertype.HookPeriodicJobsStart); ok { + lookup.pluginsByKind[PluginKindHookPeriodicJobsStart] = append(lookup.pluginsByKind[PluginKindHookPeriodicJobsStart], plugin) + } + if _, ok := plugin.(rivertype.HookWorkBegin); ok { + lookup.pluginsByKind[PluginKindHookWorkBegin] = append(lookup.pluginsByKind[PluginKindHookWorkBegin], plugin) + } + if _, ok := plugin.(rivertype.HookWorkEnd); ok { + lookup.pluginsByKind[PluginKindHookWorkEnd] = append(lookup.pluginsByKind[PluginKindHookWorkEnd], plugin) + } + } + + for _, plugin := range middlewares { + if plugin == nil { + continue + } + + if _, ok := plugin.(rivertype.JobInsertMiddleware); ok { + lookup.pluginsByKind[PluginKindMiddlewareJobInsert] = append(lookup.pluginsByKind[PluginKindMiddlewareJobInsert], plugin) + } + if _, ok := plugin.(rivertype.WorkerMiddleware); ok { + lookup.pluginsByKind[PluginKindMiddlewareWorker] = append(lookup.pluginsByKind[PluginKindMiddlewareWorker], plugin) + } + } + + return lookup +} + +func toAnySlice[T any](values []T) []any { + plugins := make([]any, 0, len(values)) + for _, value := range values { + plugins = append(plugins, value) + } + return plugins +} + +// +// JobPluginLookup +// + +type JobPluginLookup struct { + archetype *baseservice.Archetype + + mu sync.RWMutex + pluginLookupByKind map[string]*PluginLookup +} + +func NewJobPluginLookup(archetype *baseservice.Archetype) *JobPluginLookup { + return &JobPluginLookup{ + archetype: archetype, + pluginLookupByKind: make(map[string]*PluginLookup), + } +} + +// ByJobArgs returns a plugin lookup for the given job args. +func (c *JobPluginLookup) ByJobArgs(args rivertype.JobArgs) *PluginLookup { + kind := args.Kind() + + c.mu.RLock() + entry, ok := c.pluginLookupByKind[kind] + c.mu.RUnlock() + if ok { + return entry + } + + c.mu.Lock() + defer c.mu.Unlock() + if entry, ok := c.pluginLookupByKind[kind]; ok { + return entry + } + + var ( + hooks []rivertype.Hook + plugins []rivertype.Plugin + ) + if argsWithHooks, ok := args.(jobArgsWithHooks); ok { + hooks = argsWithHooks.Hooks() + } + if argsWithPlugins, ok := args.(jobArgsWithPlugins); ok { + plugins = argsWithPlugins.Plugins() + } + + entry = NewPluginLookupFromConfig(c.archetype, hooks, nil, plugins) + c.pluginLookupByKind[kind] = entry + return entry +} + +// Same as river.JobArgsWithHooks, but duplicated here so that can still live in +// the top level package. +type jobArgsWithHooks interface { + Hooks() []rivertype.Hook +} + +// Same as river.JobArgsWithPlugins, but duplicated here so that can still live +// in the top level package. +type jobArgsWithPlugins interface { + Plugins() []rivertype.Plugin +} diff --git a/vendor/github.com/riverqueue/river/internal/retrypolicy/default.go b/vendor/github.com/riverqueue/river/internal/retrypolicy/default.go new file mode 100644 index 0000000000..96b16b9a4e --- /dev/null +++ b/vendor/github.com/riverqueue/river/internal/retrypolicy/default.go @@ -0,0 +1,77 @@ +// Package retrypolicy contains River's internal retry policy implementations. +package retrypolicy + +import ( + "math" + "math/rand/v2" + "time" + + "github.com/riverqueue/river/rivershared/util/timeutil" + "github.com/riverqueue/river/rivertype" +) + +// Default is River's clock-aware default retry policy for internal use. +type Default struct { + timeGenerator rivertype.TimeGenerator +} + +// NewDefault returns a default retry policy that derives retries from the +// given time generator. +func NewDefault(timeGenerator rivertype.TimeGenerator) *Default { + return &Default{timeGenerator: timeGenerator} +} + +// NextRetry calculates when the next retry for a failed job should take place. +func (p *Default) NextRetry(job *rivertype.JobRow) time.Time { + return NextRetryAt(p.timeGenerator.Now().UTC(), job) +} + +// NextRetryAt calculates when the next retry for a failed job should take +// place relative to now. +func NextRetryAt(now time.Time, job *rivertype.JobRow) time.Time { + // In modern versions of River `len(job.Errors)` is the same number as + // `attempt`. However, in older versions snoozing a job wouldn't restore its + // attempt count to the pre-fetch value, and that would lead to incorrect + // retry durations when jobs are first snoozed, then retried. To avoid this + // and keep backward compatibility, the number of errors are used instead. + errorCount := len(job.Errors) + 1 + + return now.Add(timeutil.SecondsAsDuration(retrySeconds(errorCount))) +} + +// The maximum value of a duration before it overflows. About 292 years. +const maxDuration time.Duration = 1<<63 - 1 + +// Same as the above, but changed to a float represented in seconds. +var maxDurationSeconds = maxDuration.Seconds() //nolint:gochecknoglobals + +// Gets a number of retry seconds for the given attempt, random jitter included. +func retrySeconds(attempt int) float64 { + retrySeconds := retrySecondsWithoutJitter(attempt) + + // After hitting maximum retry durations jitter is no longer applied because + // it might overflow time.Duration. That's okay though because so much + // jitter will already have been applied up to this point (jitter measured + // in decades) that jobs will no longer run anywhere near contemporaneously + // unless there's been considerable manual intervention. + if retrySeconds == maxDurationSeconds { + return maxDurationSeconds + } + + // Jitter number of seconds +/- 10%. + retrySeconds += retrySeconds * (rand.Float64()*0.2 - 0.1) + + // Cap retrySeconds once more in case adding random jitter pushed it over + // maxDurationSeconds. (This should never realistically happen, but protect + // against it just in case.) + return min(retrySeconds, maxDurationSeconds) +} + +// Gets a base number of retry seconds for the given attempt, jitter excluded. +// If the number of seconds returned would overflow time.Duration if it were to +// be made one, returns the maximum number of seconds that can fit in a +// time.Duration instead, approximately 292 years. +func retrySecondsWithoutJitter(attempt int) float64 { + retrySeconds := math.Pow(float64(attempt), 4) + return min(retrySeconds, maxDurationSeconds) +} diff --git a/vendor/github.com/riverqueue/river/internal/rivercommon/river_common.go b/vendor/github.com/riverqueue/river/internal/rivercommon/river_common.go new file mode 100644 index 0000000000..4f769b391a --- /dev/null +++ b/vendor/github.com/riverqueue/river/internal/rivercommon/river_common.go @@ -0,0 +1,66 @@ +package rivercommon + +import ( + "errors" + "regexp" + "time" +) + +// These constants are made available in rivercommon so that they're accessible +// by internal packages, but the top-level river package re-exports them, and +// all user code must use that set instead. +const ( + // AllQueuesString is a special string that can be used to indicate all + // queues in some operations, particularly pause and resume. + AllQueuesString = "*" + MaxAttemptsDefault = 25 + PriorityDefault = 1 + QueueDefault = "default" +) + +// HotOperationTimeout attempts to standardize timeouts for some "hot" +// operations like locking available jobs or completing finished jobs. It's +// somewhat questionable whether it makes sense to share timing on these +// queries, but for the time being it makes more sense than each part of the +// code randomly choosing its own timing. +// +// We probably want to have another look at this in the not-too-distant future +// to make sure we can't do anything a bit smarter when it comes to timeouts. +const HotOperationTimeout = 10 * time.Second + +const ( + // MetadataKeyPeriodicJobID is a metadata key inserted with a periodic job + // when a configured periodic job has its ID property set. This lets + // inserted jobs easily be traced back to the periodic job that created + // them. + MetadataKeyPeriodicJobID = "river:periodic_job_id" + + // MetadataKeyResumableStep records the last successfully completed step for + // a resumable job so later attempts can skip ahead. + MetadataKeyResumableStep = "river:resumable_step" + + // MetadataKeyResumableCursor records a resumable step cursor so a later + // attempt can resume a partially completed step. + MetadataKeyResumableCursor = "river:resumable_cursor" + + // MetadataKeyRescueCount records how many times the job has been rescued. + MetadataKeyRescueCount = "river:rescue_count" + + // MetadataKeyUniqueNonce is a special metadata key used by the SQLite driver to + // determine whether an upsert is was skipped or not because the `(xmax != 0)` + // trick we use in Postgres doesn't work in SQLite. + MetadataKeyUniqueNonce = "river:unique_nonce" +) + +type ContextKeyClient struct{} + +// ErrStop is a special error injected by the client into its fetch and work +// CancelCauseFuncs when it's stopping. It may be used by components for such +// cases like avoiding logging an error during a normal shutdown procedure. +var ErrStop = errors.New("stop initiated") + +// UserSpecifiedIDOrKindRE is a regular expression to which the format of job +// kinds and some other user-specified IDs (e.g. periodic job names) must +// comply. Mainly, minimal special characters, and excluding spaces and commas +// which are problematic for the search UI. +var UserSpecifiedIDOrKindRE = regexp.MustCompile(`\A[\w][\w\-\[\]<>\/.·:+]+\z`) diff --git a/vendor/github.com/riverqueue/river/internal/riverplugin/plugin.go b/vendor/github.com/riverqueue/river/internal/riverplugin/plugin.go new file mode 100644 index 0000000000..9cce6d3910 --- /dev/null +++ b/vendor/github.com/riverqueue/river/internal/riverplugin/plugin.go @@ -0,0 +1,95 @@ +package riverplugin + +import ( + "context" + "encoding/json" + "errors" + "fmt" + + "github.com/tidwall/gjson" + + "github.com/riverqueue/river/internal/jobexecutor" + "github.com/riverqueue/river/internal/rivercommon" + "github.com/riverqueue/river/rivertype" +) + +// DefaultPlugins returns the default plugins that River applies to all jobs. +// This includes internal middleware like the resumable step middleware. +func DefaultPlugins() []rivertype.Plugin { + return []rivertype.Plugin{&ResumableMiddleware{}} +} + +// ResumableMiddleware is internal middleware that enables resumable step +// functionality. It reads the last completed step and cursor data from job +// metadata, injects them into the context, and persists updated step/cursor +// state back to metadata when a job errors after making progress. +type ResumableMiddleware struct{} + +func (*ResumableMiddleware) IsMiddleware() bool { return true } +func (*ResumableMiddleware) IsPlugin() bool { return true } + +func (*ResumableMiddleware) Work(ctx context.Context, job *rivertype.JobRow, doInner func(ctx context.Context) error) error { + metadataUpdates, hasMetadataUpdates := jobexecutor.MetadataUpdatesFromWorkContext(ctx) + if !hasMetadataUpdates { + return errors.New("expected to find metadata updates in context, but didn't") + } + + state := &ResumableState{ + AllStepNames: make(map[string]struct{}), + Cursors: make(map[string]json.RawMessage), + ResumeMatched: true, + ResumeStep: gjson.GetBytes(job.Metadata, rivercommon.MetadataKeyResumableStep).Str, + } + if state.ResumeStep != "" { + state.ResumeMatched = false + } + + hadCursors := false + if cursorJSON := gjson.GetBytes(job.Metadata, rivercommon.MetadataKeyResumableCursor); cursorJSON.Exists() && cursorJSON.Type == gjson.JSON { + if err := json.Unmarshal([]byte(cursorJSON.Raw), &state.Cursors); err != nil { + return fmt.Errorf("river: unmarshal resumable cursors: %w", err) + } + hadCursors = len(state.Cursors) > 0 + } + + ctx = context.WithValue(ctx, ResumableContextKey{}, state) + + err := doInner(ctx) + if err == nil { + switch { + case state.Err != nil: + err = state.Err + case state.ResumeStep != "" && !state.ResumeMatched: + err = fmt.Errorf("river: resumable step %q not found in Worker", state.ResumeStep) + } + } + + if err != nil && state.CompletedStep != "" { + if len(state.Cursors) > 0 { + metadataUpdates[rivercommon.MetadataKeyResumableCursor] = state.Cursors + } else if hadCursors { + // All cursors were consumed (their steps completed). Write null + // to clear the stale cursor metadata, since the metadata merge + // is additive and wouldn't remove the old key on its own. + metadataUpdates[rivercommon.MetadataKeyResumableCursor] = nil + } + metadataUpdates[rivercommon.MetadataKeyResumableStep] = state.CompletedStep + } + + return err +} + +// ResumableState holds the state for a resumable job execution. It is stored in +// the context and accessed by ResumableStep and ResumableStepCursor. +type ResumableState struct { + AllStepNames map[string]struct{} + CompletedStep string + Cursors map[string]json.RawMessage + Err error + ResumeMatched bool + ResumeStep string + StepName string +} + +// ResumableContextKey is the context key for ResumableState. +type ResumableContextKey struct{} diff --git a/vendor/github.com/riverqueue/river/internal/util/chanutil/debounced_chan.go b/vendor/github.com/riverqueue/river/internal/util/chanutil/debounced_chan.go new file mode 100644 index 0000000000..7e685805c5 --- /dev/null +++ b/vendor/github.com/riverqueue/river/internal/util/chanutil/debounced_chan.go @@ -0,0 +1,145 @@ +package chanutil + +import ( + "context" + "sync" + "time" +) + +// DebouncedChan is a channel that will only fire once per cooldown period, at +// the leading edge. If it is called again during the cooldown, the subsequent +// calls are delayed until the cooldown period has elapsed and are also +// coalesced into a single call. +type DebouncedChan struct { + c chan struct{} + cooldown time.Duration + ctxDone <-chan struct{} + sendLeading bool + + // mu protects variables in group below + mu sync.Mutex + sendOnTimerExpired bool + timer *time.Timer + timerDone bool +} + +// NewDebouncedChan returns a new DebouncedChan which sends on the channel no +// more often than the cooldown period. +// +// If sendLeading is true, the channel will signal once on C the first time it +// receives a signal, then again once per cooldown period. If sendLeading is +// false, the initial signal isn't sent. +func NewDebouncedChan(ctx context.Context, cooldown time.Duration, sendLeading bool) *DebouncedChan { + return &DebouncedChan{ + ctxDone: ctx.Done(), + c: make(chan struct{}, 1), + cooldown: cooldown, + sendLeading: sendLeading, + } +} + +// C is the debounced channel. Multiple invocations to Call during the cooldown +// period will deduplicate to a single emission on this channel on the period's +// leading edge (if sendLeading was enabled), and one more on the trailing edge +// for as many periods as invocations continue to come in. +func (d *DebouncedChan) C() <-chan struct{} { + return d.c +} + +// Call invokes the debounced channel, and is the call which will be debounced. +// If multiple invocations of this function are made during the cooldown period, +// they'll be debounced to a single emission on C on the period's leading edge +// (if sendLeading is enabled), and then one fire on the trailing edge of each +// period for as long as Call continues to be invoked. If a timer period elapses +// without an invocation on Call, the timer is stopped and behavior resets the +// next time Call is invoked again. +func (d *DebouncedChan) Call() { + d.mu.Lock() + defer d.mu.Unlock() + + // A timer has already been initialized and hasn't already expired. (If it + // has expired, we'll reset it below.) Set to signal when it does expire. + if d.timer != nil && !d.timerDone { + d.sendOnTimerExpired = true + return + } + + // No timer had been started yet, or the last one running was expired and + // will be reset. Send immediately (i.e. n the leading edge of the + // debounce period), if sendLeading is enabled. + if d.sendLeading { + d.nonBlockingSendOnC() + } else { + d.sendOnTimerExpired = true + } + + // Next, start the timer, during which we'll monitor for additional calls, + // and send at the end of the period if any came in. Create a new timer if + // this is the first run. Otherwise, reset an existing one. + if d.timer == nil { + d.timer = time.NewTimer(d.cooldown) + } else { + d.timer.Reset(d.cooldown) + } + d.timerDone = false + + go d.waitForTimerLoop() +} + +func (d *DebouncedChan) nonBlockingSendOnC() { + select { + case d.c <- struct{}{}: + default: + } +} + +// Waits for the timer to be fired, and loops as long as Call invocations come +// in. If a period elapses without a new Call coming in, the loop returns, and +// DebouncedChan returns to its initial state, waiting for a new Call. +// +// The loop also stops if context becomes done. +func (d *DebouncedChan) waitForTimerLoop() { + for { + if stopLoop := d.waitForTimerOnce(); stopLoop { + break + } + } +} + +// Waits for the timer to fire once or context becomes done. Returns true if the +// caller should stop looping (i.e. don't wait on the timer again), and false +// otherwise. +func (d *DebouncedChan) waitForTimerOnce() bool { + select { + case <-d.ctxDone: + d.mu.Lock() + defer d.mu.Unlock() + + if d.timer != nil { + if !d.timer.Stop() { + <-d.timer.C + } + } + + d.timerDone = true + + case <-d.timer.C: + d.mu.Lock() + defer d.mu.Unlock() + + if d.sendOnTimerExpired { + d.sendOnTimerExpired = false + d.nonBlockingSendOnC() + + // Wait for another timer expiry, which will fire again if another + // Call comes in during that time. If no Call comes in, the timer + // will stop on the next cycle and we return to initial state. + d.timer.Reset(d.cooldown) + return false // do _not_ stop looping + } + + d.timerDone = true + } + + return true // stop looping +} diff --git a/vendor/github.com/riverqueue/river/internal/workunit/work_unit.go b/vendor/github.com/riverqueue/river/internal/workunit/work_unit.go new file mode 100644 index 0000000000..8becb69799 --- /dev/null +++ b/vendor/github.com/riverqueue/river/internal/workunit/work_unit.go @@ -0,0 +1,40 @@ +package workunit + +import ( + "context" + "time" + + "github.com/riverqueue/river/internal/pluginlookup" + "github.com/riverqueue/river/rivertype" +) + +// WorkUnit provides an interface to a struct that wraps a job to be done +// combined with a work function that can execute it. Its main purpose is to +// wrap a struct that contains generic types (like a Worker[T] that needs to be +// invoked with a Job[T]) in such a way as to make it non-generic so that it can +// be used in other non-generic code like jobExecutor. +// +// Implemented by river.wrapperWorkUnit. +type WorkUnit interface { + // PluginLookup procures the a hook lookup bundle for the wrapped job using + // the given job hook lookup bundle. Hooks are looked up by job args and + // otherwise not available to jobexecutor. + PluginLookup(lookup *pluginlookup.JobPluginLookup) *pluginlookup.PluginLookup + + Middleware() []rivertype.WorkerMiddleware + NextRetry() time.Time + Timeout() time.Duration + UnmarshalJob() error + Work(ctx context.Context) error +} + +// WorkUnitFactory provides an interface to a struct that can generate a +// workUnit, a wrapper around a job to be done combined with a work function +// that can execute it. +// +// Implemented by river.workUnitFactoryWrapper. +type WorkUnitFactory interface { + // Make a workUnit, which wraps a job to be done and work function that can + // execute it. + MakeUnit(jobRow *rivertype.JobRow) WorkUnit +} diff --git a/vendor/github.com/riverqueue/river/job.go b/vendor/github.com/riverqueue/river/job.go new file mode 100644 index 0000000000..5a04bbda46 --- /dev/null +++ b/vendor/github.com/riverqueue/river/job.go @@ -0,0 +1,114 @@ +package river + +import ( + "github.com/riverqueue/river/rivertype" +) + +// Job represents a single unit of work, holding both the arguments and +// information for a job with args of type T. +type Job[T JobArgs] struct { + *rivertype.JobRow + + // Args are the arguments for the job. + Args T +} + +// JobArgs is an interface that represents the arguments for a job of type T. +// These arguments are serialized into JSON and stored in the database. +// +// The struct is serialized using `encoding/json`. All exported fields are +// serialized, unless skipped with a struct field tag. +type JobArgs interface { + // Kind is a string that uniquely identifies the type of job. This must be + // provided on your job arguments struct. Jobs are identified by a string + // instead of being based on type names so that previously inserted jobs + // can be worked across deploys even if job/worker types are renamed. + // + // Kinds should be formatted without spaces like `my_custom_job`, + // `mycustomjob`, or `my-custom-job`. Many special characters like colons, + // dots, hyphens, and underscores are allowed, but those like spaces and + // commas, which would interfere with UI functionality, are invalid. + // + // After initially deploying a job, it's generally not safe to rename its + // kind (unless the database is completely empty) because River won't know + // which worker should work the old kind. Job kinds can be renamed safely + // over multiple deploys using the JobArgsWithKindAliases interface. + Kind() string +} + +// JobArgsWithKindAliases is an interface that jobs args can implement to +// provide an alternate kind which a worker will be registered under in addition +// to the primary kind. This is useful for renaming a job kind in a safe manner +// so that any jobs already in the database aren't orphaned. +// +// Renaming a job is a three part process. To begin, a job args with its +// original name: +// +// type jobArgsBeingRenamed struct{} +// +// func (a jobArgsBeingRenamed) Kind() string { return "old_name" } +// +// Rename by putting the new name in Kind and moving the old name to +// KindAliases: +// +// type jobArgsBeingRenamed struct{} +// +// func (a jobArgsBeingRenamed) Kind() string { return "new_name" } +// func (a jobArgsBeingRenamed) KindAliases() []string { return []string{"old_name"} } +// +// After all jobs inserted under the original name have finished working +// (including all their possible retries, which notably might take up to three +// weeks on the default retry policy), remove KindAliases: +// +// type jobArgsBeingRenamed struct{} +// +// func (a jobArgsBeingRenamed) Kind() string { return "new_name" } +type JobArgsWithKindAliases interface { + // KindAliases returns alias kinds that an associated job args worker will + // respond to. + KindAliases() []string +} + +// JobArgsWithHooks is an interface that job args can implement to attach +// specific hooks (i.e. other than those globally installed to a client) to +// certain kinds of jobs. +type JobArgsWithHooks interface { + // Hooks returns specific hooks to run for this job type. These will run + // after the global hooks configured on the client. + // + // Warning: Hooks returned should be based on the job type only and be + // invariant of the specific contents of a job. Hooks are extracted by + // instantiating a generic instance of the job even when a specific instance + // is available, so any conditional logic within will be ignored. This is + // done because although specific job information may be available in some + // hook contexts like on InsertBegin, it won't be in others like WorkBegin. + Hooks() []rivertype.Hook +} + +// JobArgsWithInsertOpts is an extra interface that a job may implement on top +// of JobArgs to provide insertion-time options for all jobs of this type. +type JobArgsWithInsertOpts interface { + // InsertOpts returns options for all jobs of this job type, overriding any + // system defaults. These can also be overridden at insertion time. + InsertOpts() InsertOpts +} + +// JobArgsWithPlugins is an interface that job args can implement to attach +// specific plugins (i.e. other than those globally installed to a client) to +// certain kinds of jobs. +type JobArgsWithPlugins interface { + // Plugins returns specific plugins to run for this job type. Plugin hooks + // run after global hooks, and plugin middleware is combined with other + // middleware configured for the job. + // + // A plugin implementing rivertype.JobInsertMiddleware runs once around an + // insertion batch containing one or more jobs of this type. Its InsertMany + // method receives the complete batch, which may also contain other job + // types. + // + // Warning: Plugins returned should be based on the job type only and be + // invariant of the specific contents of a job. Plugins are extracted by + // instantiating a generic instance of the job even when a specific instance + // is available, so any conditional logic within will be ignored. + Plugins() []rivertype.Plugin +} diff --git a/vendor/github.com/riverqueue/river/job_complete_tx.go b/vendor/github.com/riverqueue/river/job_complete_tx.go new file mode 100644 index 0000000000..5c1a75eb09 --- /dev/null +++ b/vendor/github.com/riverqueue/river/job_complete_tx.go @@ -0,0 +1,85 @@ +package river + +import ( + "context" + "encoding/json" + "errors" + "time" + + "github.com/riverqueue/river/internal/execution" + "github.com/riverqueue/river/internal/jobexecutor" + "github.com/riverqueue/river/riverdriver" + "github.com/riverqueue/river/rivertype" +) + +// JobCompleteTx marks the job as completed as part of transaction tx. If tx is +// rolled back, the completion will be as well. +// +// The function needs to know the type of the River database driver, which is +// the same as the one in use by Client, but the other generic parameters can be +// inferred. An invocation should generally look like: +// +// _, err := river.JobCompleteTx[*riverpgxv5.Driver](ctx, tx, job) +// if err != nil { +// // handle error +// } +// +// Returns the updated, completed job. +func JobCompleteTx[TDriver riverdriver.Driver[TTx], TTx any, TArgs JobArgs](ctx context.Context, tx TTx, job *Job[TArgs]) (*Job[TArgs], error) { + if job.State != rivertype.JobStateRunning { + return nil, errors.New("job must be running") + } + + client := ClientFromContext[TTx](ctx) + if client == nil { + return nil, errors.New("client not found in context, can only work within a River worker") + } + + driver := client.Driver() + pilot := client.Pilot() + + // extract metadata updates from context + metadataUpdates, hasMetadataUpdates := jobexecutor.MetadataUpdatesFromWorkContext(ctx) + hasMetadataUpdates = hasMetadataUpdates && len(metadataUpdates) > 0 + var ( + metadataUpdatesBytes []byte + err error + ) + if hasMetadataUpdates { + metadataUpdatesBytes, err = json.Marshal(metadataUpdates) + if err != nil { + return nil, err + } + } + + execTx := driver.UnwrapExecutor(tx) + params := riverdriver.JobSetStateCompleted(job.ID, client.baseService.Time.Now(), nil) + rows, err := pilot.JobSetStateIfRunningMany(ctx, execTx, &riverdriver.JobSetStateIfRunningManyParams{ + ID: []int64{params.ID}, + Attempt: []*int{params.Attempt}, + ErrData: [][]byte{params.ErrData}, + FinalizedAt: []*time.Time{params.FinalizedAt}, + MetadataDoMerge: []bool{hasMetadataUpdates}, + MetadataUpdates: [][]byte{metadataUpdatesBytes}, + ScheduledAt: []*time.Time{params.ScheduledAt}, + Schema: client.config.Schema, + State: []rivertype.JobState{params.State}, + }) + if err != nil { + return nil, err + } + if len(rows) == 0 { + if _, isInsideTestWorker := ctx.Value(execution.ContextKeyInsideTestWorker{}).(bool); isInsideTestWorker { + panic("to use JobCompleteTx in a rivertest.Worker, the job must be inserted into the database first") + } + + return nil, rivertype.ErrNotFound + } + updatedJob := &Job[TArgs]{JobRow: rows[0]} + + if err := json.Unmarshal(updatedJob.EncodedArgs, &updatedJob.Args); err != nil { + return nil, err + } + + return updatedJob, nil +} diff --git a/vendor/github.com/riverqueue/river/job_list_params.go b/vendor/github.com/riverqueue/river/job_list_params.go new file mode 100644 index 0000000000..4c275857cc --- /dev/null +++ b/vendor/github.com/riverqueue/river/job_list_params.go @@ -0,0 +1,541 @@ +package river + +import ( + "encoding/base64" + "encoding/json" + "errors" + "fmt" + "maps" + "time" + + "github.com/riverqueue/river/internal/dblist" + "github.com/riverqueue/river/rivershared/util/ptrutil" + "github.com/riverqueue/river/rivertype" +) + +// JobListCursor is used to specify a starting point for a paginated +// job list query. +type JobListCursor struct { + id int64 + job *rivertype.JobRow // used for JobListCursorFromJob path; not serialized + kind string + queue string + sortField JobListOrderByField + time time.Time // may be empty +} + +// JobListCursorFromJob creates a JobListCursor from a JobRow. +func JobListCursorFromJob(job *rivertype.JobRow) *JobListCursor { + // Other fields are initialized when the cursor is used in After below. + return &JobListCursor{job: job} +} + +func jobListCursorFromJobAndParams(job *rivertype.JobRow, listParams *JobListParams) *JobListCursor { + // A pointer so that we can detect a condition where we accidentally left + // this value unset. + var cursorTime *time.Time + + // Don't include a `default` so `exhaustive` lint can detect omissions. + switch listParams.sortField { + case JobListOrderByID: + cursorTime = ptrutil.Ptr(time.Time{}) + case JobListOrderByTime: + cursorTime = ptrutil.Ptr(jobListTimeValue(job)) + case JobListOrderByFinalizedAt: + if job.FinalizedAt != nil { + cursorTime = job.FinalizedAt + } + case JobListOrderByScheduledAt: + cursorTime = &job.ScheduledAt + } + + if cursorTime == nil { + panic("invalid sort field") + } + + return &JobListCursor{ + id: job.ID, + kind: job.Kind, + queue: job.Queue, + sortField: listParams.sortField, + time: *cursorTime, + } +} + +// UnmarshalText implements encoding.TextUnmarshaler to decode the cursor from +// a previously marshaled string. +func (c *JobListCursor) UnmarshalText(text []byte) error { + dst := make([]byte, base64.StdEncoding.DecodedLen(len(text))) + n, err := base64.StdEncoding.Decode(dst, text) + if err != nil { + return err + } + dst = dst[:n] + + wrapperValue := jobListPaginationCursorJSON{} + if err := json.Unmarshal(dst, &wrapperValue); err != nil { + return err + } + *c = JobListCursor{ + id: wrapperValue.ID, + kind: wrapperValue.Kind, + queue: wrapperValue.Queue, + sortField: JobListOrderByField(wrapperValue.SortField), + time: wrapperValue.Time, + } + return nil +} + +// MarshalText implements encoding.TextMarshaler to encode the cursor as an +// opaque string. +func (c JobListCursor) MarshalText() ([]byte, error) { + if c.job != nil { + return nil, errors.New("cursor initialized with only a job can't be marshaled; try a cursor from JobListResult instead") + } + + wrapperValue := jobListPaginationCursorJSON{ + ID: c.id, + Kind: c.kind, + Queue: c.queue, + SortField: string(c.sortField), + Time: c.time, + } + data, err := json.Marshal(wrapperValue) + if err != nil { + return nil, err + } + dst := make([]byte, base64.URLEncoding.EncodedLen(len(data))) + base64.URLEncoding.Encode(dst, data) + return dst, nil +} + +type jobListPaginationCursorJSON struct { + ID int64 `json:"id"` + Kind string `json:"kind"` + Queue string `json:"queue"` + SortField string `json:"sort_field"` + Time time.Time `json:"time"` +} + +// SortOrder specifies the direction of a sort. +type SortOrder int + +const ( + // SortOrderAsc specifies that the sort should in ascending order. + SortOrderAsc SortOrder = iota + + // SortOrderDesc specifies that the sort should in descending order. + SortOrderDesc +) + +// JobListOrderByField specifies the field to sort by. +type JobListOrderByField string + +const ( + // JobListOrderByID specifies that the sort should be by job ID. + JobListOrderByID JobListOrderByField = "id" + + // JobListOrderByFinalizedAt specifies that the sort should be by + // `finalized_at`. + // + // This option must be used in conjunction with filtering by only finalized + // job states. + JobListOrderByFinalizedAt JobListOrderByField = "finalized_at" + + // JobListOrderByScheduledAt specifies that the sort should be by + // `scheduled_at`. + JobListOrderByScheduledAt JobListOrderByField = "scheduled_at" + + // JobListOrderByTime specifies that the sort should be by the "best fit" + // time field based on listed state. The best fit is determined by looking + // at the first value given to JobListParams.States. If multiple states are + // specified, the ones after the first will be ignored. + // + // The specific time field used for sorting depends on requested state: + // + // * States `available`, `retryable`, or `scheduled` use `scheduled_at`. + // * State `running` uses `attempted_at`. + // * States `cancelled`, `completed`, or `discarded` use `finalized_at`. + JobListOrderByTime JobListOrderByField = "time" +) + +// JobListParams specifies the parameters for a JobList query. It must be +// initialized with NewJobListParams. Params can be built by chaining methods on +// the JobListParams object: +// +// params := NewJobListParams().OrderBy(JobListOrderByTime, SortOrderAsc).First(100) +type JobListParams struct { + after *JobListCursor + ids []int64 + kinds []string + metadataCalled bool + overrodeState bool + limit int32 + priorities []int16 + queues []string + schema string + sortField JobListOrderByField + sortOrder SortOrder + states []rivertype.JobState + tagsAll []string + tagsAny []string + where []dblist.WherePredicate +} + +// NewJobListParams creates a new JobListParams to return available jobs sorted +// by time in ascending order, returning 100 jobs at most. +func NewJobListParams() *JobListParams { + return &JobListParams{ + limit: 100, + sortField: JobListOrderByID, + sortOrder: SortOrderAsc, + states: []rivertype.JobState{ + rivertype.JobStateAvailable, + rivertype.JobStateCancelled, + rivertype.JobStateCompleted, + rivertype.JobStateDiscarded, + rivertype.JobStatePending, + rivertype.JobStateRetryable, + rivertype.JobStateRunning, + rivertype.JobStateScheduled, + }, + } +} + +func (p *JobListParams) copy() *JobListParams { + return &JobListParams{ + after: p.after, + ids: append([]int64(nil), p.ids...), + kinds: append([]string(nil), p.kinds...), + metadataCalled: p.metadataCalled, + overrodeState: p.overrodeState, + limit: p.limit, + priorities: append([]int16(nil), p.priorities...), + queues: append([]string(nil), p.queues...), + sortField: p.sortField, + sortOrder: p.sortOrder, + schema: p.schema, + states: append([]rivertype.JobState(nil), p.states...), + tagsAll: append([]string(nil), p.tagsAll...), + tagsAny: append([]string(nil), p.tagsAny...), + where: append([]dblist.WherePredicate(nil), p.where...), + } +} + +func (p *JobListParams) toDBParams() (*dblist.JobListParams, error) { + orderBy := make([]dblist.JobListOrderBy, 0, 2) + + var sortOrder dblist.SortOrder + switch p.sortOrder { + case SortOrderAsc: + sortOrder = dblist.SortOrderAsc + case SortOrderDesc: + sortOrder = dblist.SortOrderDesc + default: + return nil, errors.New("invalid sort order") + } + + if p.sortField == JobListOrderByFinalizedAt { + if len(p.states) == 0 { + return nil, errors.New("cannot order by finalized_at without finalized state filters") + } + + currentNonFinalizedStates := make([]rivertype.JobState, 0, len(p.states)) + for _, state := range p.states { + switch state { + case rivertype.JobStateAvailable, rivertype.JobStatePending, rivertype.JobStateRetryable, rivertype.JobStateRunning, rivertype.JobStateScheduled: + currentNonFinalizedStates = append(currentNonFinalizedStates, state) + case rivertype.JobStateCancelled, rivertype.JobStateCompleted, rivertype.JobStateDiscarded: + } + } + // FinalizedAt ordering is only supported when filtering to finalized + // states because non-finalized jobs have no finalized_at value. + if len(currentNonFinalizedStates) > 0 { + return nil, fmt.Errorf("cannot order by finalized_at with non-finalized state filters %+v", currentNonFinalizedStates) + } + } + + var timeField string + switch { + case p.sortField == JobListOrderByID: + // no time field + + case len(p.states) > 0 && p.sortField == JobListOrderByTime: + timeField = jobListTimeFieldForState(p.states[0]) + orderBy = append(orderBy, dblist.JobListOrderBy{Expr: timeField, Order: sortOrder}) + + default: + timeField = string(p.sortField) + orderBy = append(orderBy, dblist.JobListOrderBy{Expr: timeField, Order: sortOrder}) + } + + orderBy = append(orderBy, dblist.JobListOrderBy{Expr: "id", Order: sortOrder}) + + if p.after != nil { + namedArgs := map[string]any{"after_id": p.after.id} + if p.after.time.IsZero() { // order by ID only + if sortOrder == dblist.SortOrderAsc { + p.where = append(p.where, dblist.WherePredicate{NamedArgs: namedArgs, SQL: "(id > @after_id)"}) + } else { + p.where = append(p.where, dblist.WherePredicate{NamedArgs: namedArgs, SQL: "(id < @after_id)"}) + } + } else { + namedArgs["cursor_time"] = p.after.time + if sortOrder == dblist.SortOrderAsc { + p.where = append(p.where, dblist.WherePredicate{NamedArgs: namedArgs, SQL: fmt.Sprintf(`("%s" > @cursor_time OR ("%s" = @cursor_time AND "id" > @after_id))`, timeField, timeField)}) + } else { + p.where = append(p.where, dblist.WherePredicate{NamedArgs: namedArgs, SQL: fmt.Sprintf(`("%s" < @cursor_time OR ("%s" = @cursor_time AND "id" < @after_id))`, timeField, timeField)}) + } + } + } + + return &dblist.JobListParams{ + IDs: p.ids, + Kinds: p.kinds, + LimitCount: p.limit, + OrderBy: orderBy, + Priorities: p.priorities, + Queues: p.queues, + Schema: p.schema, + States: p.states, + TagsAll: p.tagsAll, + TagsAny: p.tagsAny, + Where: p.where, + }, nil +} + +// After returns an updated filter set that will only return jobs +// after the given cursor. +func (p *JobListParams) After(cursor *JobListCursor) *JobListParams { + paramsCopy := p.copy() + if cursor.job == nil { + paramsCopy.after = cursor + } else { + paramsCopy.after = jobListCursorFromJobAndParams(cursor.job, paramsCopy) + } + return paramsCopy +} + +// First returns an updated filter set that will only return the first +// count jobs. +// +// Count must be between 1 and 10_000, inclusive, or this will panic. +func (p *JobListParams) First(count int) *JobListParams { + if count <= 0 { + panic("count must be > 0") + } + if count > 10_000 { + panic("count must be <= 10_000") + } + paramsCopy := p.copy() + paramsCopy.limit = int32(count) + return paramsCopy +} + +// IDs returns an updated filter set that will only return jobs with the given +// IDs. +func (p *JobListParams) IDs(ids ...int64) *JobListParams { + paramsCopy := p.copy() + paramsCopy.ids = make([]int64, len(ids)) + copy(paramsCopy.ids, ids) + return paramsCopy +} + +// Kinds returns an updated filter set that will only return jobs of the given +// kinds. +func (p *JobListParams) Kinds(kinds ...string) *JobListParams { + paramsCopy := p.copy() + paramsCopy.kinds = make([]string, len(kinds)) + copy(paramsCopy.kinds, kinds) + return paramsCopy +} + +// Metadata returns an updated filter set that will return only jobs that has +// metadata which contains the given JSON fragment at its top level. This is +// equivalent to the `@>` operator in Postgres: +// +// https://www.postgresql.org/docs/current/functions-json.html +// +// This function isn't supported in SQLite due to SQLite not having an +// equivalent operator to use, so there's no efficient way to implement it. We +// recommend the use of Where using a condition with a comparison on the `->>` +// operator instead. +func (p *JobListParams) Metadata(json string) *JobListParams { + paramsCopy := p.copy() + paramsCopy.metadataCalled = true + paramsCopy.where = append(paramsCopy.where, dblist.WherePredicate{ + NamedArgs: map[string]any{"metadata_fragment": json}, + SQL: `metadata @> @metadata_fragment::jsonb`, + }) + return paramsCopy +} + +// OrderBy returns an updated filter set that will sort the results using the +// specified field and direction. +// +// If ordering by FinalizedAt, the States filter will be set to only include +// finalized job states unless it has already been overridden. +func (p *JobListParams) OrderBy(field JobListOrderByField, direction SortOrder) *JobListParams { + paramsCopy := p.copy() + switch field { + case JobListOrderByID, JobListOrderByTime, JobListOrderByScheduledAt: + paramsCopy.sortField = field + case JobListOrderByFinalizedAt: + paramsCopy.sortField = field + if !p.overrodeState { + paramsCopy.states = []rivertype.JobState{ + rivertype.JobStateCancelled, + rivertype.JobStateCompleted, + rivertype.JobStateDiscarded, + } + } + default: + panic("invalid order by field") + } + paramsCopy.sortField = field + paramsCopy.sortOrder = direction + return paramsCopy +} + +// Priorities returns an updated filter set that will only return jobs with the +// given priorities. +func (p *JobListParams) Priorities(priorities ...int16) *JobListParams { + paramsCopy := p.copy() + paramsCopy.priorities = make([]int16, len(priorities)) + copy(paramsCopy.priorities, priorities) + return paramsCopy +} + +// Queues returns an updated filter set that will only return jobs from the +// given queues. +func (p *JobListParams) Queues(queues ...string) *JobListParams { + paramsCopy := p.copy() + paramsCopy.queues = make([]string, len(queues)) + copy(paramsCopy.queues, queues) + return paramsCopy +} + +// States returns an updated filter set that will only return jobs in the given +// states. +func (p *JobListParams) States(states ...rivertype.JobState) *JobListParams { + paramsCopy := p.copy() + paramsCopy.states = make([]rivertype.JobState, len(states)) + paramsCopy.overrodeState = true + copy(paramsCopy.states, states) + return paramsCopy +} + +// TagsAll returns an updated filter set that will only return jobs containing +// all of the given tags. Matching is exact and case-sensitive. TagsAll is +// combined with TagsAny and all other filters using AND. +// +// Calling TagsAll replaces any tags supplied to a previous TagsAll call. +// Calling it with no tags removes the filter. +func (p *JobListParams) TagsAll(tags ...string) *JobListParams { + paramsCopy := p.copy() + paramsCopy.tagsAll = make([]string, len(tags)) + copy(paramsCopy.tagsAll, tags) + return paramsCopy +} + +// TagsAny returns an updated filter set that will only return jobs containing +// at least one of the given tags. Matching is exact and case-sensitive. +// TagsAny is combined with TagsAll and all other filters using AND. +// +// Calling TagsAny replaces any tags supplied to a previous TagsAny call. +// Calling it with no tags removes the filter. +func (p *JobListParams) TagsAny(tags ...string) *JobListParams { + paramsCopy := p.copy() + paramsCopy.tagsAny = make([]string, len(tags)) + copy(paramsCopy.tagsAny, tags) + return paramsCopy +} + +// NamedArgs are named arguments for use with JobListParams.Where. Keys should +// look like "my_param", and map to parameters like "@my_param" in SQL queries. +// "@" are present in the SQL, but not in the keys of this map. +type NamedArgs map[string]any + +// Where is an all-encompassing query escape hatch that adds an arbitrary +// predicate after a list query's `WHERE ...` clause. Use of other JobListParams +// filters should be preferred where possible because they're safer and their +// compatibility between drivers is better guaranteed, but in case none is +// suitable, Where can be used as a last resort. +// +// For example, using Where to query with `jsonb_path_query_first(...)` using a +// JSON path, a function that's specific to Postgres: +// +// listParams = listParams.Where("jsonb_path_query_first(metadata, @json_path) = @json_val", NamedArgs{"json_path": "$.foo", "json_val": `"bar"`}) +// +// A JSON path can be used in a query in SQLite as well, but there the `->` or +// `->>` operators must be used instead: +// +// listParams = listParams.Where("metadata ->> @json_path = @json_val", NamedArgs{"json_path": "$.foo", "json_val": "bar"}) +// +// Arguments beyond the first are interpreted as named parameters. Each one +// should be present in the query SQL prefixed with a `@` symbol. Multiple sets +// of named parameters will be merged together, with values in later sets +// overwriting those in earlier ones. +// +// Calling Where multiple times will add multiple conditions separate by `AND`. +// Use `OR` instead by stuffing all conditions into a single Where invocation. +// +// Consider use of this function possibly hazardous! Any time raw SQL is in +// play, an application is opening itself up to SQL injection attacks. Never mix +// unsanitized user input into a SQL string, and use named parameters to curb +// the likelihood of injection. +func (p *JobListParams) Where(sql string, namedArgsMany ...NamedArgs) *JobListParams { + paramsCopy := p.copy() + + var allNamedArgs NamedArgs + if len(namedArgsMany) > 0 { + for i, namedArgs := range namedArgsMany { + if i == 0 { + allNamedArgs = namedArgs + } else { + maps.Copy(allNamedArgs, namedArgs) + } + } + } + + paramsCopy.where = append(paramsCopy.where, dblist.WherePredicate{NamedArgs: allNamedArgs, SQL: sql}) + return paramsCopy +} + +func jobListTimeFieldForState(state rivertype.JobState) string { + // Don't include a `default` so `exhaustive` lint can detect omissions. + switch state { + case rivertype.JobStateAvailable, rivertype.JobStatePending, rivertype.JobStateRetryable, rivertype.JobStateScheduled: + return "scheduled_at" + case rivertype.JobStateRunning: + return "attempted_at" + case rivertype.JobStateCancelled, rivertype.JobStateCompleted, rivertype.JobStateDiscarded: + return "finalized_at" + } + + return "created_at" // should never happen +} + +func jobListTimeValue(job *rivertype.JobRow) time.Time { + // Don't include a `default` so `exhaustive` lint can detect omissions. + switch job.State { + case rivertype.JobStateAvailable, rivertype.JobStatePending, rivertype.JobStateRetryable, rivertype.JobStateScheduled: + return job.ScheduledAt + + case rivertype.JobStateRunning: + if job.AttemptedAt == nil { + // This should never happen unless a job has been manually manipulated. + return job.CreatedAt + } + return *job.AttemptedAt + + case rivertype.JobStateCancelled, rivertype.JobStateCompleted, rivertype.JobStateDiscarded: + if job.FinalizedAt == nil { + // This should never happen unless a job has been manually manipulated. + return job.CreatedAt + } + return *job.FinalizedAt + } + + return job.CreatedAt // should never happen +} diff --git a/vendor/github.com/riverqueue/river/metadata.go b/vendor/github.com/riverqueue/river/metadata.go new file mode 100644 index 0000000000..eafba14d00 --- /dev/null +++ b/vendor/github.com/riverqueue/river/metadata.go @@ -0,0 +1,38 @@ +package river + +import ( + "context" + "errors" + "strings" + + "github.com/riverqueue/river/internal/jobexecutor" +) + +var errMetadataNotSettable = errors.New("MetadataSet must be called within a worker, worker middleware, or work hook") + +// MetadataSet records a metadata value to be merged into the job's metadata +// when the current work attempt finishes. +// +// This function is only valid from a worker, worker middleware, or work hook +// like rivertype.HookWorkBegin or rivertype.HookWorkEnd. +// +// Metadata updates are stored on the work context and merged into the job row +// when the current work attempt finishes, whether the attempt succeeds or +// errors. Values must be JSON marshalable because metadata is stored in a +// jsonb column, and setting a key replaces any existing value at that key. +// +// Keys prefixed with `river:` are reserved for internal use and may not be set +// by user code. +func MetadataSet(ctx context.Context, key string, value any) error { + if strings.HasPrefix(key, "river:") { + return errors.New("MetadataSet cannot be used with keys prefixed with `river:`") + } + + metadataUpdates, ok := jobexecutor.MetadataUpdatesFromWorkContext(ctx) + if !ok { + return errMetadataNotSettable + } + + metadataUpdates[key] = value + return nil +} diff --git a/vendor/github.com/riverqueue/river/middleware_defaults.go b/vendor/github.com/riverqueue/river/middleware_defaults.go new file mode 100644 index 0000000000..3633449b19 --- /dev/null +++ b/vendor/github.com/riverqueue/river/middleware_defaults.go @@ -0,0 +1,65 @@ +package river + +import ( + "context" + + "github.com/riverqueue/river/rivertype" +) + +// MiddlewareDefaults should be embedded on any middleware implementation. It +// helps identify a struct as middleware and a plugin, and guarantees forward +// compatibility in case additions are necessary to the rivertype.Middleware +// interface. +type MiddlewareDefaults struct{} + +func (d *MiddlewareDefaults) IsMiddleware() bool { return true } + +func (d *MiddlewareDefaults) IsPlugin() bool { return true } + +// JobInsertMiddlewareDefaults is an embeddable struct that provides default +// implementations for the rivertype.JobInsertMiddleware. Use of this struct is +// recommended in case rivertype.JobInsertMiddleware is expanded in the future +// so that existing code isn't unexpectedly broken during an upgrade. +// +// Deprecated: Prefer embedding the more general MiddlewareDefaults instead. +type JobInsertMiddlewareDefaults struct{ MiddlewareDefaults } + +func (d *JobInsertMiddlewareDefaults) InsertMany(ctx context.Context, manyParams []*rivertype.JobInsertParams, doInner func(ctx context.Context) ([]*rivertype.JobInsertResult, error)) ([]*rivertype.JobInsertResult, error) { + return doInner(ctx) +} + +// JobInsertMiddlewareFunc is a convenience helper for implementing +// rivertype.JobInsertMiddleware using a simple function instead of a struct. +type JobInsertMiddlewareFunc func(ctx context.Context, manyParams []*rivertype.JobInsertParams, doInner func(ctx context.Context) ([]*rivertype.JobInsertResult, error)) ([]*rivertype.JobInsertResult, error) + +func (f JobInsertMiddlewareFunc) InsertMany(ctx context.Context, manyParams []*rivertype.JobInsertParams, doInner func(ctx context.Context) ([]*rivertype.JobInsertResult, error)) ([]*rivertype.JobInsertResult, error) { + return f(ctx, manyParams, doInner) +} + +func (f JobInsertMiddlewareFunc) IsMiddleware() bool { return true } + +func (f JobInsertMiddlewareFunc) IsPlugin() bool { return true } + +// WorkerInsertMiddlewareDefaults is an embeddable struct that provides default +// implementations for the rivertype.WorkerMiddleware. Use of this struct is +// recommended in case rivertype.WorkerMiddleware is expanded in the future so +// that existing code isn't unexpectedly broken during an upgrade. +// +// Deprecated: Prefer embedding the more general MiddlewareDefaults instead. +type WorkerMiddlewareDefaults struct{ MiddlewareDefaults } + +func (d *WorkerMiddlewareDefaults) Work(ctx context.Context, job *rivertype.JobRow, doInner func(ctx context.Context) error) error { + return doInner(ctx) +} + +// WorkerMiddlewareFunc is a convenience helper for implementing +// rivertype.WorkerMiddleware using a simple function instead of a struct. +type WorkerMiddlewareFunc func(ctx context.Context, job *rivertype.JobRow, doInner func(ctx context.Context) error) error + +func (f WorkerMiddlewareFunc) IsMiddleware() bool { return true } + +func (f WorkerMiddlewareFunc) IsPlugin() bool { return true } + +func (f WorkerMiddlewareFunc) Work(ctx context.Context, job *rivertype.JobRow, doInner func(ctx context.Context) error) error { + return f(ctx, job, doInner) +} diff --git a/vendor/github.com/riverqueue/river/periodic_job.go b/vendor/github.com/riverqueue/river/periodic_job.go new file mode 100644 index 0000000000..a6fd87956a --- /dev/null +++ b/vendor/github.com/riverqueue/river/periodic_job.go @@ -0,0 +1,264 @@ +package river + +import ( + "time" + + "github.com/riverqueue/river/internal/maintenance" + "github.com/riverqueue/river/rivershared/baseservice" + "github.com/riverqueue/river/rivershared/util/sliceutil" + "github.com/riverqueue/river/rivertype" +) + +// PeriodicSchedule is a schedule for a periodic job. Periodic jobs should +// generally have an interval of at least 1 minute, and never less than one +// second. +type PeriodicSchedule interface { + // Next returns the next time at which the job should be run given the + // current time. + Next(current time.Time) time.Time +} + +// PeriodicJobConstructor is a function that gets called each time the paired +// PeriodicSchedule is triggered. +// +// A constructor must never block. It may return nil to indicate that no job +// should be inserted. +type PeriodicJobConstructor func() (JobArgs, *InsertOpts) + +// PeriodicJob is a configuration for a periodic job. +type PeriodicJob struct { + constructorFunc PeriodicJobConstructor + opts *PeriodicJobOpts + scheduleFunc PeriodicSchedule +} + +// PeriodicJobOpts are options for a periodic job. +type PeriodicJobOpts struct { + // ID is an optional identifier for the job. Identifiers must be unique + // between all periodic jobs and adding a periodic job will error if they're + // not. + ID string + + // RunOnStart can be used to indicate that a periodic job should insert an + // initial job as a new scheduler is started. This can be used as a hedge + // for jobs with longer scheduled durations that may not get to expiry + // before a new scheduler is elected. + // + // RunOnStart also applies when a new periodic job is added dynamically with + // `PeriodicJobs().Add` or `PeriodicJobs().AddMany`. Jobs added this way + // with RunOnStart set to true are inserted once, then continue with their + // normal run schedule. + RunOnStart bool +} + +// NewPeriodicJob returns a new PeriodicJob given a schedule and a constructor +// function. +// +// The schedule returns a time until the next time the periodic job should run. +// The helper PeriodicInterval is available for jobs that should run on simple, +// fixed intervals (e.g. every 15 minutes), and a custom schedule or third party +// cron package can be used for more complex scheduling (see the cron example). +// The constructor function is invoked each time a periodic job's schedule +// elapses, returning job arguments to insert along with optional insertion +// options. +// +// The periodic job scheduler is approximate and doesn't guarantee strong +// durability. It's started by the elected leader in a River cluster, and each +// periodic job is assigned an initial run time when that occurs. New run times +// are scheduled each time a job's target run time is reached and a new job +// inserted. However, each scheduler only retains in-memory state, so anytime a +// process quits or a new leader is elected, the whole process starts over +// without regard for the state of the last scheduler. The RunOnStart option +// can be used as a hedge to make sure that jobs with long run durations are +// guaranteed to occasionally run. +func NewPeriodicJob(scheduleFunc PeriodicSchedule, constructorFunc PeriodicJobConstructor, opts *PeriodicJobOpts) *PeriodicJob { + return &PeriodicJob{ + constructorFunc: constructorFunc, + opts: opts, + scheduleFunc: scheduleFunc, + } +} + +type neverSchedule struct{} + +func (s *neverSchedule) Next(t time.Time) time.Time { + // Return the maximum future time so that the schedule never runs. + return time.Unix(1<<63-62135596801, 999999999) +} + +// NeverSchedule returns a PeriodicSchedule that never runs. +func NeverSchedule() PeriodicSchedule { + return &neverSchedule{} +} + +type periodicIntervalSchedule struct { + interval time.Duration +} + +// PeriodicInterval returns a simple PeriodicSchedule that runs at the given +// interval. +func PeriodicInterval(interval time.Duration) PeriodicSchedule { + return &periodicIntervalSchedule{interval} +} + +func (s *periodicIntervalSchedule) Next(t time.Time) time.Time { + return t.Add(s.interval) +} + +// PeriodicJobBundle is a bundle of currently configured periodic jobs. It's +// made accessible through Client, where new periodic jobs can be configured, +// and old ones removed. +type PeriodicJobBundle struct { + mapper *periodicJobInternalMapper + periodicJobEnqueuer *maintenance.PeriodicJobEnqueuer +} + +func newPeriodicJobBundle(config *Config, periodicJobEnqueuer *maintenance.PeriodicJobEnqueuer) *PeriodicJobBundle { + return &PeriodicJobBundle{ + mapper: &periodicJobInternalMapper{archetype: &periodicJobEnqueuer.Archetype, config: config}, + periodicJobEnqueuer: periodicJobEnqueuer, + } +} + +// Add adds a new periodic job to the client. The job is queued immediately if +// RunOnStart is enabled, and then scheduled normally. +// +// Returns a periodic job handle which can be used to subsequently remove the +// job if desired. +// +// Adding or removing periodic jobs has no effect unless this client is elected +// leader because only the leader enqueues periodic jobs. To make sure that a +// new periodic job is fully enabled or disabled, it should be added or removed +// from _every_ active River client across all processes. +func (b *PeriodicJobBundle) Add(periodicJob *PeriodicJob) rivertype.PeriodicJobHandle { + handle, err := b.periodicJobEnqueuer.AddSafely(b.mapper.toInternal(periodicJob)) + if err != nil { + panic(err) + } + return handle +} + +// AddSafely is the same as Add, but it returns an error in the case of a +// validation problem or duplicate ID instead of panicking. +func (b *PeriodicJobBundle) AddSafely(periodicJob *PeriodicJob) (rivertype.PeriodicJobHandle, error) { + return b.periodicJobEnqueuer.AddSafely(b.mapper.toInternal(periodicJob)) +} + +// AddMany adds many new periodic jobs to the client. The jobs are queued +// immediately if their RunOnStart is enabled, and then scheduled normally. +// +// Returns a periodic job handle which can be used to subsequently remove the +// job if desired. +// +// Adding or removing periodic jobs has no effect unless this client is elected +// leader because only the leader enqueues periodic jobs. To make sure that a +// new periodic job is fully enabled or disabled, it should be added or removed +// from _every_ active River client across all processes. +func (b *PeriodicJobBundle) AddMany(periodicJobs []*PeriodicJob) []rivertype.PeriodicJobHandle { + handles, err := b.periodicJobEnqueuer.AddManySafely(sliceutil.Map(periodicJobs, b.mapper.toInternal)) + if err != nil { + panic(err) + } + return handles +} + +// AddManySafely is the same as AddMany, but it returns an error in the case of +// a validation problem or duplicate ID instead of panicking. +func (b *PeriodicJobBundle) AddManySafely(periodicJobs []*PeriodicJob) ([]rivertype.PeriodicJobHandle, error) { + return b.periodicJobEnqueuer.AddManySafely(sliceutil.Map(periodicJobs, b.mapper.toInternal)) +} + +// Clear clears all periodic jobs, cancelling all scheduled runs. +// +// Adding or removing periodic jobs has no effect unless this client is elected +// leader because only the leader enqueues periodic jobs. To make sure that a +// new periodic job is fully enabled or disabled, it should be added or removed +// from _every_ active River client across all processes. +func (b *PeriodicJobBundle) Clear() { + b.periodicJobEnqueuer.Clear() +} + +// Remove removes a periodic job, cancelling all scheduled runs. +// +// Requires the use of the periodic job handle that was returned when the job +// was added. +// +// Adding or removing periodic jobs has no effect unless this client is elected +// leader because only the leader enqueues periodic jobs. To make sure that a +// new periodic job is fully enabled or disabled, it should be added or removed +// from _every_ active River client across all processes. +func (b *PeriodicJobBundle) Remove(periodicJobHandle rivertype.PeriodicJobHandle) { + b.periodicJobEnqueuer.Remove(periodicJobHandle) +} + +// RemoveByID removes a periodic job by ID, cancelling all scheduled runs. +// +// Adding or removing periodic jobs has no effect unless this client is elected +// leader because only the leader enqueues periodic jobs. To make sure that a +// new periodic job is fully enabled or disabled, it should be added or removed +// from _every_ active River client across all processes. +// +// Has no effect if no jobs with the given ID is configured. +// +// Returns true if a job with the given ID existed (and was removed), and false +// otherwise. +func (b *PeriodicJobBundle) RemoveByID(id string) bool { + return b.periodicJobEnqueuer.RemoveByID(id) +} + +// RemoveMany removes many periodic jobs, cancelling all scheduled runs. +// +// Requires the use of the periodic job handles that were returned when the jobs +// were added. +// +// Adding or removing periodic jobs has no effect unless this client is elected +// leader because only the leader enqueues periodic jobs. To make sure that a +// new periodic job is fully enabled or disabled, it should be added or removed +// from _every_ active River client across all processes. +func (b *PeriodicJobBundle) RemoveMany(periodicJobHandles []rivertype.PeriodicJobHandle) { + b.periodicJobEnqueuer.RemoveMany(periodicJobHandles) +} + +// RemoveManyByID removes many periodic jobs by ID, cancelling all scheduled +// runs. +// +// Adding or removing periodic jobs has no effect unless this client is elected +// leader because only the leader enqueues periodic jobs. To make sure that a +// new periodic job is fully enabled or disabled, it should be added or removed +// from _every_ active River client across all processes. +// +// Has no effect if no jobs with the given IDs are configured. +func (b *PeriodicJobBundle) RemoveManyByID(ids []string) { + b.periodicJobEnqueuer.RemoveManyByID(ids) +} + +// An empty set of periodic job opts used as a default when none are specified. +var periodicJobEmptyOpts PeriodicJobOpts //nolint:gochecknoglobals + +type periodicJobInternalMapper struct { + archetype *baseservice.Archetype + config *Config +} + +// There are two separate periodic job structs so that the top-level River +// package can expose one while still containing most periodic job logic in a +// subpackage. This function converts a top-level periodic job struct (used for +// configuration) to an internal one. +func (m *periodicJobInternalMapper) toInternal(periodicJob *PeriodicJob) *maintenance.PeriodicJob { + opts := &periodicJobEmptyOpts + if periodicJob.opts != nil { + opts = periodicJob.opts + } + return &maintenance.PeriodicJob{ + ID: opts.ID, + ConstructorFunc: func() (*rivertype.JobInsertParams, error) { + args, options := periodicJob.constructorFunc() + if args == nil { + return nil, maintenance.ErrNoJobToInsert + } + return insertParamsFromConfigArgsAndOptions(m.archetype, m.config, args, options) + }, + RunOnStart: opts.RunOnStart, + ScheduleFunc: periodicJob.scheduleFunc.Next, + } +} diff --git a/vendor/github.com/riverqueue/river/plugin.go b/vendor/github.com/riverqueue/river/plugin.go new file mode 100644 index 0000000000..95becdc7ca --- /dev/null +++ b/vendor/github.com/riverqueue/river/plugin.go @@ -0,0 +1,36 @@ +package river + +import ( + "github.com/riverqueue/river/rivershared/baseservice" + "github.com/riverqueue/river/rivershared/riverpilot" + "github.com/riverqueue/river/rivershared/startstop" +) + +// A plugin API that drivers may implement to extend a River client. Driver +// plugins may, for example, add additional maintenance services. +// +// This should be considered a River internal API and its stability is not +// guaranteed. DO NOT USE. +type driverPlugin[TTx any] interface { + // PluginInit initializes a plugin with an archetype. It's invoked on + // Client.NewClient. + PluginInit(archetype *baseservice.Archetype) + + // PluginPilot returns a custom Pilot implementation. + PluginPilot() riverpilot.Pilot +} + +// A plugin API that pilots may implement to extend a River client. Pilot +// plugins may, for example, add additional maintenance services. +// +// This should be considered a River internal API and its stability is not +// guaranteed. DO NOT USE. +type pilotPlugin interface { + // PluginMaintenanceServices returns additional maintenance services (will + // only run on an elected leader) for a River client. + PluginMaintenanceServices() []startstop.Service + + // PluginServices returns additional non-maintenance services (will run on + // all clients) for a River client. + PluginServices() []startstop.Service +} diff --git a/vendor/github.com/riverqueue/river/plugin_defaults.go b/vendor/github.com/riverqueue/river/plugin_defaults.go new file mode 100644 index 0000000000..2930e479dd --- /dev/null +++ b/vendor/github.com/riverqueue/river/plugin_defaults.go @@ -0,0 +1,12 @@ +package river + +// PluginDefaults should be embedded on plugin implementations. It helps +// identify a struct as both hook and middleware, and guarantees forward +// compatibility in case additions are necessary to the rivertype.Hook or +// rivertype.Middleware interfaces. +type PluginDefaults struct { + HookDefaults + MiddlewareDefaults +} + +func (d *PluginDefaults) IsPlugin() bool { return true } diff --git a/vendor/github.com/riverqueue/river/producer.go b/vendor/github.com/riverqueue/river/producer.go new file mode 100644 index 0000000000..da11246e72 --- /dev/null +++ b/vendor/github.com/riverqueue/river/producer.go @@ -0,0 +1,1159 @@ +package river + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "log/slog" + "math" + "strings" + "sync" + "sync/atomic" + "time" + + "github.com/riverqueue/river/internal/jobcompleter" + "github.com/riverqueue/river/internal/jobexecutor" + "github.com/riverqueue/river/internal/notifier" + "github.com/riverqueue/river/internal/pluginlookup" + "github.com/riverqueue/river/internal/retrypolicy" + "github.com/riverqueue/river/internal/rivercommon" + "github.com/riverqueue/river/internal/util/chanutil" + "github.com/riverqueue/river/internal/workunit" + "github.com/riverqueue/river/riverdriver" + "github.com/riverqueue/river/rivershared/baseservice" + "github.com/riverqueue/river/rivershared/riverpilot" + "github.com/riverqueue/river/rivershared/startstop" + "github.com/riverqueue/river/rivershared/testsignal" + "github.com/riverqueue/river/rivershared/util/randutil" + "github.com/riverqueue/river/rivershared/util/serviceutil" + "github.com/riverqueue/river/rivershared/util/testutil" + "github.com/riverqueue/river/rivershared/util/timeoututil" + "github.com/riverqueue/river/rivershared/util/timeutil" + "github.com/riverqueue/river/rivertype" +) + +const ( + producerReportIntervalDefault = 30 * time.Second + queuePollIntervalDefault = 2 * time.Second + queueReportIntervalDefault = 10 * time.Minute +) + +// Test-only properties. +type producerTestSignals struct { + DeletedExpiredQueueRecords testsignal.TestSignal[struct{}] // notifies when the producer deletes expired queue records + JobFetchTriggered testsignal.TestSignal[struct{}] // notifies when the producer's fetch limiter is triggered via triggerJobFetch + MetadataChanged testsignal.TestSignal[struct{}] // notifies when the producer detects a metadata change + Paused testsignal.TestSignal[struct{}] // notifies when the producer is paused + PolledQueueConfig testsignal.TestSignal[struct{}] // notifies when the producer polls for queue settings + QueueControlEventTriggered testsignal.TestSignal[*controlEventPayload] // notifies when a queue control event is triggered via triggerQueueControlEvent + ReportedProducerStatus testsignal.TestSignal[struct{}] // notifies when the producer reports its own status + ReportedQueueStatus testsignal.TestSignal[struct{}] // notifies when the producer reports queue status + Resumed testsignal.TestSignal[struct{}] // notifies when the producer is resumed + StartedExecutors testsignal.TestSignal[struct{}] // notifies when runOnce finishes a pass +} + +func (ts *producerTestSignals) Init(tb testutil.TestingTB) { + ts.DeletedExpiredQueueRecords.Init(tb) + ts.JobFetchTriggered.Init(tb) + ts.MetadataChanged.Init(tb) + ts.Paused.Init(tb) + ts.PolledQueueConfig.Init(tb) + ts.QueueControlEventTriggered.Init(tb) + ts.ReportedQueueStatus.Init(tb) + ts.ReportedProducerStatus.Init(tb) + ts.Resumed.Init(tb) + ts.StartedExecutors.Init(tb) +} + +type producerConfig struct { + ClientID string + Completer jobcompleter.JobCompleter + ErrorHandler ErrorHandler + + // FetchCooldown is the minimum amount of time to wait between fetches of new + // jobs. Jobs will only be fetched *at most* this often, but if no new jobs + // are coming in via LISTEN/NOTIFY then fetches may be delayed as long as + // FetchPollInterval. + FetchCooldown time.Duration + + // FetchPollInterval is the amount of time between periodic fetches for new + // jobs. Typically new jobs will be picked up ~immediately after insert via + // LISTEN/NOTIFY, but this provides a fallback. + FetchPollInterval time.Duration + + PluginLookupByJob *pluginlookup.JobPluginLookup + PluginLookupGlobal *pluginlookup.PluginLookup + JobStuckHandler JobStuckHandler + JobStuckCount *atomic.Int32 + JobStuckThreshold time.Duration + JobTimeout time.Duration + MaxWorkers int + + // Notifier is a notifier for subscribing to new job inserts and job + // control. If nil, the producer will operate in poll-only mode. + Notifier *notifier.Notifier + // ProducerReportInterval is the amount of time between periodic reports + // of the producer status. + ProducerReportInterval time.Duration + + Queue string + // QueueEventCallback gets called when a queue's config changes (such as + // pausing or resuming) events can be emitted to subscriptions. + QueueEventCallback func(event *Event) + + // QueuePollInterval is the amount of time between periodic checks for + // queue setting changes. This is only used in poll-only mode (when no + // notifier is provided). + QueuePollInterval time.Duration + // QueueReportInterval is the amount of time between periodic reports + // of the queue status. + QueueReportInterval time.Duration + RetryPolicy ClientRetryPolicy + SchedulerInterval time.Duration + Schema string + StaleProducerRetentionPeriod time.Duration + Workers *Workers +} + +func (c *producerConfig) mustValidate() *producerConfig { + if c.Completer == nil { + panic("producerConfig.Completer is required") + } + if c.ClientID == "" { + panic("producerConfig.ClientID is required") + } + if c.FetchCooldown <= 0 { + panic("producerConfig.FetchCooldown must be great than zero") + } + if c.FetchPollInterval <= 0 { + panic("producerConfig.FetchPollInterval must be greater than zero") + } + if c.JobStuckCount == nil { + c.JobStuckCount = &atomic.Int32{} + } + if c.JobStuckThreshold == 0 { + c.JobStuckThreshold = JobStuckThresholdDefault + } + if c.JobStuckThreshold < 0 { + panic("producerConfig.JobStuckThreshold must be greater or equal to zero") + } + if c.JobTimeout < -1 { + panic("producerConfig.JobTimeout must be greater or equal to zero") + } + if c.MaxWorkers == 0 { + panic("producerConfig.MaxWorkers is required") + } + if c.ProducerReportInterval == 0 { + c.ProducerReportInterval = producerReportIntervalDefault + } + if c.Queue == "" { + panic("producerConfig.Queue is required") + } + if c.QueuePollInterval == 0 { + c.QueuePollInterval = queuePollIntervalDefault + } + if c.QueuePollInterval <= 0 { + panic("producerConfig.QueueSettingsPollInterval must be greater than zero") + } + if c.QueueReportInterval == 0 { + c.QueueReportInterval = queueReportIntervalDefault + } + if c.QueueReportInterval <= 0 { + panic("producerConfig.QueueSettingsReportInterval must be greater than zero") + } + if c.RetryPolicy == nil { + panic("producerConfig.RetryPolicy is required") + } + if c.SchedulerInterval == 0 { + panic("producerConfig.SchedulerInterval is required") + } + if c.StaleProducerRetentionPeriod <= 0 { + panic("producerConfig.StaleProducerRetentionPeriod must be greater than zero") + } + if c.Workers == nil { + panic("producerConfig.Workers is required") + } + + return c +} + +// producer manages a fleet of Workers up to a maximum size. It periodically fetches jobs +// from the adapter and dispatches them to Workers. It receives completed job results from Workers. +// +// The producer never fetches more jobs than the number of free Worker slots it +// has available. This is not optimal for throughput compared to pre-fetching +// extra jobs, but it is better for smaller job counts or slower jobs where even +// distribution and minimizing execution latency is more important. +type producer struct { + baseservice.BaseService + startstop.BaseStartStop + + // Jobs which are currently being worked. Only used by main goroutine. + activeJobs map[int64]*jobexecutor.JobExecutor + + completer jobcompleter.JobCompleter + config *producerConfig + id atomic.Int64 // atomic because it's written at startup and read during shutdown + exec riverdriver.Executor + errorHandler jobexecutor.ErrorHandler + fetchLimiter *chanutil.DebouncedChan + metricEmitHooks []rivertype.HookMetricEmit // memoized hooks of type HookMetricEmit for reuse in dispatchWork + state riverpilot.ProducerState + pilot riverpilot.Pilot + workers *Workers + + // Receives job IDs to cancel. Written by notifier goroutine, only read from + // main goroutine. + cancelCh chan int64 + + // Set to true when the producer thinks it should trigger another fetch as + // soon as slots are available. This is written and read by the main + // goroutine. + fetchWhenSlotsAreAvailable bool + + // Receives completed jobs from workers. Written by completed workers, only + // read from main goroutine. + jobResultCh chan *rivertype.JobRow + + jobTimeout time.Duration + + // An atomic count of the number of jobs actively being worked on. This is + // written to by the main goroutine, but read by the dispatcher. + numJobsActive atomic.Int32 + numJobsStuck atomic.Int32 + + numJobsRan atomic.Uint64 + paused bool + // Receives control messages from the notifier goroutine. Written by notifier + // goroutine, only read from main goroutine. + queueControlCh chan *controlEventPayload + retryPolicy ClientRetryPolicy + testSignals producerTestSignals +} + +func newProducer(archetype *baseservice.Archetype, exec riverdriver.Executor, pilot riverpilot.Pilot, config *producerConfig) *producer { + if archetype == nil { + panic("archetype is required") + } + if exec == nil { + panic("exec is required") + } + + var errorHandler jobexecutor.ErrorHandler + if config.ErrorHandler != nil { + errorHandler = &errorHandlerAdapter{config.ErrorHandler} + } + + producer := baseservice.Init(archetype, &producer{ + activeJobs: make(map[int64]*jobexecutor.JobExecutor), + cancelCh: make(chan int64, 1000), + completer: config.Completer, + config: config.mustValidate(), + exec: exec, + errorHandler: errorHandler, + jobResultCh: make(chan *rivertype.JobRow, config.MaxWorkers), + jobTimeout: config.JobTimeout, + pilot: pilot, + queueControlCh: make(chan *controlEventPayload, 100), + retryPolicy: config.RetryPolicy, + workers: config.Workers, + }) + + producer.metricEmitHooks = producer.metricEmitHooksFromLookup() + + return producer +} + +// Start starts the producer. It backgrounds a goroutine which is stopped when +// context is cancelled or Stop is invoked. +// +// This variant uses a single context as fetchCtx and workCtx, and is here to +// implement startstop.Service so that the producer can be stored as a service +// variable and used with various service utilities. StartWorkContext below +// should be preferred for production use. +func (p *producer) Start(ctx context.Context) error { + return p.StartWorkContext(ctx, ctx) +} + +func (p *producer) Stop() { + p.Logger.Debug(p.Name+": Stopping", slog.String("queue", p.config.Queue), slog.Int64("id", p.id.Load())) + p.BaseStartStop.Stop() + p.Logger.Debug(p.Name+": Stop returned", slog.String("queue", p.config.Queue), slog.Int64("id", p.id.Load())) +} + +// StartWorkContext starts the producer. It backgrounds a goroutine which is +// stopped when context is cancelled or Stop is invoked. +// +// When fetchCtx is cancelled, no more jobs will be fetched; however, if a fetch +// is already in progress, It will be allowed to complete and run any fetched +// jobs. When workCtx is cancelled, any in-progress jobs will have their +// contexts cancelled too. +func (p *producer) StartWorkContext(fetchCtx, workCtx context.Context) error { + fetchCtx, shouldStart, started, stopped := p.StartInit(fetchCtx) + if !shouldStart { + return nil + } + + isExpectedShutdownError := func(err error) bool { + return errors.Is(err, startstop.ErrStop) || strings.HasSuffix(err.Error(), "conn closed") || fetchCtx.Err() != nil + } + + fetchedQueue, err := timeoututil.WithTimeoutV(fetchCtx, 10*time.Second, p.Name+".StartWorkContext", func(ctx context.Context) (*rivertype.Queue, error) { + p.Logger.DebugContext(ctx, p.Name+": Fetching initial queue settings", slog.String("queue", p.config.Queue)) + return p.exec.QueueCreateOrSetUpdatedAt(ctx, &riverdriver.QueueCreateOrSetUpdatedAtParams{ + Metadata: []byte("{}"), + Name: p.config.Queue, + Now: p.Time.NowOrNil(), + Schema: p.config.Schema, + }) + }) + if err != nil { + stopped() + if isExpectedShutdownError(err) { + return nil + } + p.Logger.ErrorContext(fetchCtx, p.Name+": Error fetching initial queue settings", slog.String("err", err.Error())) + return err + } + + initiallyPaused := fetchedQueue != nil && (fetchedQueue.PausedAt != nil) + initialMetadata := []byte("{}") + if fetchedQueue != nil { + initialMetadata = fetchedQueue.Metadata + if err := p.pilot.QueueMetadataChanged(fetchCtx, p.exec, &riverpilot.QueueMetadataChangedParams{ + Queue: p.config.Queue, + Metadata: initialMetadata, + }); err != nil { + p.Logger.ErrorContext(fetchCtx, p.Name+": Error setting fetched queue metadata with pilot", slog.String("queue", p.config.Queue), slog.String("err", err.Error())) + } + } + p.paused = initiallyPaused + + id := p.id.Load() + id, p.state, err = p.pilot.ProducerInit(fetchCtx, p.exec, &riverpilot.ProducerInitParams{ + ClientID: p.config.ClientID, + ProducerID: id, + Queue: p.config.Queue, + Schema: p.config.Schema, + }) + if err != nil { + stopped() + if isExpectedShutdownError(err) { + return nil + } + p.Logger.ErrorContext(fetchCtx, p.Name+": Error initializing producer state", slog.String("err", err.Error())) + return err + } + p.id.Store(id) + + p.fetchLimiter = chanutil.NewDebouncedChan(fetchCtx, p.config.FetchCooldown, true) + + var ( + controlSub *notifier.Subscription + insertSub *notifier.Subscription + ) + if p.config.Notifier != nil { + var err error + + handleInsertNotification := func(topic notifier.NotificationTopic, payload string) { + var decoded insertPayload + if err := json.Unmarshal([]byte(payload), &decoded); err != nil { + p.Logger.ErrorContext(workCtx, p.Name+": Failed to unmarshal insert notification payload", slog.String("err", err.Error())) + return + } + if decoded.Queue != p.config.Queue { + return + } + p.Logger.DebugContext(workCtx, p.Name+": Received insert notification", slog.String("queue", decoded.Queue)) + p.fetchLimiter.Call() + } + insertSub, err = p.config.Notifier.Listen(fetchCtx, notifier.NotificationTopicInsert, handleInsertNotification) + if err != nil { + stopped() + if strings.HasSuffix(err.Error(), "conn closed") || errors.Is(err, context.Canceled) { + return nil + } + return err + } + + controlSub, err = p.config.Notifier.Listen(fetchCtx, notifier.NotificationTopicControl, p.handleControlNotification(workCtx)) + if err != nil { + stopped() + if strings.HasSuffix(err.Error(), "conn closed") || errors.Is(err, context.Canceled) { + return nil + } + return err + } + } + + go func() { + started() + defer stopped() // this defer should come first so it's last out + + p.Logger.DebugContext(fetchCtx, p.Name+": Run loop started", slog.String("queue", p.config.Queue), slog.Bool("paused", p.paused)) + defer func() { + p.Logger.DebugContext(fetchCtx, p.Name+": Run loop stopped", slog.String("queue", p.config.Queue), slog.Uint64("num_completed_jobs", p.numJobsRan.Load())) + }() + + if insertSub != nil { + defer insertSub.Unlisten(fetchCtx) + } + + if controlSub != nil { + defer controlSub.Unlisten(fetchCtx) + } + + var subroutineWG sync.WaitGroup + subroutineCtx, cancelSubroutines := context.WithCancelCause(context.WithoutCancel(fetchCtx)) + + subroutineWG.Add(1) + go p.fetchPollLoop(subroutineCtx, &subroutineWG) + + subroutineWG.Add(1) + go p.heartbeatLogLoop(subroutineCtx, &subroutineWG) + + subroutineWG.Add(1) + go p.reportQueueStatusLoop(subroutineCtx, &subroutineWG) + + subroutineWG.Add(1) + go p.reportProducerStatusLoop(subroutineCtx, &subroutineWG) + + if p.config.Notifier == nil { + p.Logger.DebugContext(subroutineCtx, p.Name+": No notifier configured; starting in poll mode", "client_id", p.config.ClientID) + + subroutineWG.Add(1) + go p.pollForSettingChanges(subroutineCtx, &subroutineWG, initiallyPaused, initialMetadata) + } + + p.fetchAndRunLoop(fetchCtx, workCtx) + p.Logger.DebugContext(workCtx, p.Name+": Entering shutdown loop", slog.String("queue", p.config.Queue), slog.Int64("id", p.id.Load())) + p.executorShutdownLoop() + + p.Logger.DebugContext(workCtx, p.Name+": Shutdown loop exited, awaiting subroutines", slog.String("queue", p.config.Queue), slog.Int64("id", p.id.Load())) + cancelSubroutines(fmt.Errorf("producer stopped: %w", startstop.ErrStop)) + subroutineWG.Wait() + p.Logger.DebugContext(workCtx, p.Name+": Shutdown subroutines completed, finalizing", slog.String("queue", p.config.Queue), slog.Int64("id", p.id.Load())) + + p.finalizeShutdown(context.WithoutCancel(fetchCtx)) + }() + + return nil +} + +// TriggerJobFetch manually triggers the producer to perform a job fetch +// (although it's debounced, so it may not happen immediately if a fetch was +// performed very recently). This is used by clients using drivers that don't +// support listeners to wake a producer immediately after a job insert was known +// to be performed so the producer doesn't have to wait on polling. +func (p *producer) TriggerJobFetch() { + if p.fetchLimiter != nil { + p.fetchLimiter.Call() + } + p.testSignals.JobFetchTriggered.Signal(struct{}{}) +} + +// TriggerQueueControlEvent manually injects a queue control event into the +// producer's queue control channel as if it'd been received through +// listen/notify. This is used by clients using drivers that don't support +// listeners to wake a producer immediately after a queue control event was +// known to be performed so the producer doesn't have to wait on polling. +func (p *producer) TriggerQueueControlEvent(controlEvent *controlEventPayload) { + p.queueControlCh <- controlEvent + p.testSignals.QueueControlEventTriggered.Signal(controlEvent) +} + +type controlAction string + +const ( + controlActionCancel controlAction = "cancel" + controlActionMetadataChanged controlAction = "metadata_changed" + controlActionPause controlAction = "pause" + controlActionResume controlAction = "resume" +) + +type controlEventPayload struct { + Action controlAction `json:"action"` + JobID int64 `json:"job_id,omitempty"` + Metadata json.RawMessage `json:"metadata,omitempty"` + Queue string `json:"queue"` +} + +type insertPayload struct { + Queue string `json:"queue"` +} + +func (p *producer) handleControlNotification(workCtx context.Context) func(notifier.NotificationTopic, string) { + return func(topic notifier.NotificationTopic, payload string) { + var decoded controlEventPayload + if err := json.Unmarshal([]byte(payload), &decoded); err != nil { + p.Logger.ErrorContext(workCtx, p.Name+": Failed to unmarshal job control notification payload", slog.String("err", err.Error())) + return + } + + switch decoded.Action { + case controlActionMetadataChanged, controlActionPause, controlActionResume: + if decoded.Queue != rivercommon.AllQueuesString && decoded.Queue != p.config.Queue { + p.Logger.DebugContext(workCtx, p.Name+": Queue control notification for other queue", slog.String("action", string(decoded.Action))) + return + } + select { + case <-workCtx.Done(): + case p.queueControlCh <- &decoded: + default: + p.Logger.WarnContext(workCtx, p.Name+": Queue control notification dropped due to full buffer", slog.String("action", string(decoded.Action))) + } + case controlActionCancel: + if decoded.Queue != p.config.Queue { + p.Logger.DebugContext(workCtx, p.Name+": Received job cancel notification for other queue", + slog.String("action", string(decoded.Action)), + slog.Int64("job_id", decoded.JobID), + slog.String("queue", decoded.Queue), + ) + return + } + select { + case <-workCtx.Done(): + case p.cancelCh <- decoded.JobID: + default: + p.Logger.WarnContext(workCtx, p.Name+": Job cancel notification dropped due to full buffer", slog.Int64("job_id", decoded.JobID)) + } + default: + p.Logger.DebugContext(workCtx, p.Name+": Received job control notification with unknown action", + slog.String("action", string(decoded.Action)), + slog.Int64("job_id", decoded.JobID), + slog.String("queue", decoded.Queue), + ) + } + } +} + +func (p *producer) fetchAndRunLoop(fetchCtx, workCtx context.Context) { + // Prime the fetchLimiter so we can make an initial fetch without waiting for + // an insert notification or a fetch poll. + p.fetchLimiter.Call() + + fetchResultCh := make(chan producerFetchResult) + for { + select { + case <-fetchCtx.Done(): + return + case msg := <-p.queueControlCh: + switch msg.Action { + case controlActionCancel: + // This path is only expected to take effect in poll-only mode, and + // only works for the case of a single process. Multi-process setups + // will have to wait for the next poll event for a cancel to take effect. + p.maybeCancelJob(workCtx, msg.JobID) + case controlActionMetadataChanged: + p.Logger.DebugContext(workCtx, p.Name+": Queue metadata changed", slog.String("queue", p.config.Queue), slog.String("queue_in_message", msg.Queue)) + p.testSignals.MetadataChanged.Signal(struct{}{}) + if err := p.pilot.QueueMetadataChanged(workCtx, p.exec, &riverpilot.QueueMetadataChangedParams{ + Queue: p.config.Queue, + Metadata: msg.Metadata, + }); err != nil { + p.Logger.ErrorContext(workCtx, p.Name+": Error updating queue metadata with pilot", slog.String("queue", p.config.Queue), slog.String("err", err.Error())) + } + case controlActionPause: + if p.paused { + continue + } + p.paused = true + p.Logger.DebugContext(workCtx, p.Name+": Paused", slog.String("queue", p.config.Queue), slog.String("queue_in_message", msg.Queue)) + p.testSignals.Paused.Signal(struct{}{}) + if p.config.QueueEventCallback != nil { + p.config.QueueEventCallback(&Event{Kind: EventKindQueuePaused, Queue: &rivertype.Queue{Name: p.config.Queue}}) + } + case controlActionResume: + if !p.paused { + continue + } + p.paused = false + p.Logger.DebugContext(workCtx, p.Name+": Resumed", slog.String("queue", p.config.Queue), slog.String("queue_in_message", msg.Queue)) + p.fetchLimiter.Call() // try another fetch because more jobs may be available to run which were gated behind the paused queue + p.testSignals.Resumed.Signal(struct{}{}) + if p.config.QueueEventCallback != nil { + p.config.QueueEventCallback(&Event{Kind: EventKindQueueResumed, Queue: &rivertype.Queue{Name: p.config.Queue}}) + } + default: + p.Logger.DebugContext(workCtx, p.Name+": Unknown queue control action", "action", msg.Action) + } + case jobID := <-p.cancelCh: + p.maybeCancelJob(workCtx, jobID) + case <-p.fetchLimiter.C(): + p.innerFetchLoop(workCtx, fetchResultCh) + // Ensure we can't start another fetch when fetchCtx is done, even if + // the fetchLimiter is also ready to fire: + select { + case <-fetchCtx.Done(): + return + default: + } + case result := <-p.jobResultCh: + p.removeActiveJob(result) + if p.fetchWhenSlotsAreAvailable { + // If we missed a fetch because all worker slots were full, or if we + // fetched the maximum number of jobs on the last attempt, get a little + // more aggressive triggering the fetch limiter now that we have a slot + // available. + p.fetchWhenSlotsAreAvailable = false + p.fetchLimiter.Call() + } + } + } +} + +// Loops every FetchPollInterval to check for jobs. This is meant as a back up +// in case something with listen/notify didn't work, or the fetch limiter was +// limited so there's still jobs to pick up, and it's also important in +// poll-only mode. +func (p *producer) fetchPollLoop(ctx context.Context, wg *sync.WaitGroup) { + defer wg.Done() + + fetchPollTimer := time.NewTimer(p.jitteredFetchPollInterval()) + for { + select { + case <-ctx.Done(): + // Stop fetch timer so no more fetches are triggered. + if !fetchPollTimer.Stop() { + <-fetchPollTimer.C + } + return + case <-fetchPollTimer.C: + p.fetchLimiter.Call() + fetchPollTimer.Reset(p.jitteredFetchPollInterval()) + } + } +} + +// jitteredFetchPollInterval returns FetchPollInterval with random jitter in +// [0, 10% of FetchPollInterval) added (minimum 10ms). This prevents multiple +// producers from synchronizing their fetches after a transient event (e.g. GC +// pause, network blip), which would cause periodic DB load spikes. +func (p *producer) jitteredFetchPollInterval() time.Duration { + jitterRange := max(p.config.FetchPollInterval/10, 10*time.Millisecond) + return randutil.DurationBetween(p.config.FetchPollInterval, p.config.FetchPollInterval+jitterRange) +} + +func (p *producer) innerFetchLoop(workCtx context.Context, fetchResultCh chan producerFetchResult) { + var limit int + if p.paused { + limit = 0 + } else { + limit = p.maxJobsToFetch() + if limit <= 0 { + // We have no slots for new jobs, so don't bother fetching. However, since + // we knew it was time to fetch, we keep track of what happened so we can + // trigger another fetch as soon as we have open slots. + p.fetchWhenSlotsAreAvailable = true + return + } + } + + go p.dispatchWork(workCtx, limit, fetchResultCh) + + for { + select { + case result := <-fetchResultCh: + if result.err != nil { + p.Logger.ErrorContext(workCtx, p.Name+": Error fetching jobs", slog.String("err", result.err.Error()), slog.String("queue", p.config.Queue)) + } else if len(result.jobs) > 0 { + p.startNewExecutors(workCtx, result.jobs) + + if len(result.jobs) == limit { + // Fetch returned the maximum number of jobs that were requested, + // implying there may be more in the queue. Trigger another fetch when + // slots are available. + p.fetchWhenSlotsAreAvailable = true + } + } + return + case result := <-p.jobResultCh: + p.removeActiveJob(result) + case jobID := <-p.cancelCh: + p.maybeCancelJob(workCtx, jobID) + } + } +} + +func (p *producer) executorShutdownLoop() { + // No more jobs will be fetched or executed. However, we must wait for all + // in-progress jobs to complete. + for len(p.activeJobs) != 0 { + result := <-p.jobResultCh + p.removeActiveJob(result) + } +} + +func (p *producer) finalizeShutdown(ctx context.Context) { + p.Logger.DebugContext(ctx, p.Name+": Finalizing shutdown") + + const ( + maxAttempts = 4 // Maximum number of shutdown attempts + baseTimeout = 100 * time.Millisecond // Base timeout for the first attempt + ) + + attemptShutdown := func(timeout time.Duration) error { + return timeoututil.WithTimeout(ctx, timeout, p.Name+".finalizeShutdown", func(ctx context.Context) error { + if err := p.pilot.ProducerShutdown(ctx, p.exec, &riverpilot.ProducerShutdownParams{ + ProducerID: p.id.Load(), + Queue: p.config.Queue, + Schema: p.config.Schema, + }); err != nil { + // Don't retry on these errors: + // - context.Canceled: parent context is canceled, so retrying with a new timeout won't help + // - ErrClosedPool: the database connection pool is closed, so retrying won't succeed + if errors.Is(err, context.Canceled) || errors.Is(err, riverdriver.ErrClosedPool) { + return nil + } + return err + } + return nil + }) + } + + // Progressive retry with increasing timeouts: + for attempt := 1; attempt <= maxAttempts; attempt++ { + // Exponential backoff with base 5 + // Attempt 1: 100ms, Attempt 2: 500ms, Attempt 3: 2.5s, Attempt 4: 12.5s + timeout := baseTimeout * time.Duration(math.Pow(5, float64(attempt-1))) + + if ctx.Err() != nil { + return // Don't retry if parent context is already done + } + + if err := attemptShutdown(timeout); err != nil { + p.Logger.ErrorContext(ctx, p.Name+": Error shutting down producer with pilot", + slog.String("err", err.Error()), + slog.Int("attempt", attempt), + slog.Duration("timeout", timeout)) + continue + } + return + } + + p.Logger.WarnContext(ctx, p.Name+": Failed to cleanly shutdown producer after all attempts") +} + +func (p *producer) addActiveJob(id int64, executor *jobexecutor.JobExecutor) { + p.numJobsActive.Add(1) + p.activeJobs[id] = executor +} + +func (p *producer) removeActiveJob(job *rivertype.JobRow) { + executor := p.activeJobs[job.ID] + delete(p.activeJobs, job.ID) + if executor == nil || executor.TryCloseSlot() { + p.numJobsActive.Add(-1) + } + p.numJobsRan.Add(1) + p.state.JobFinish(job) +} + +func (p *producer) handleWorkerStuck(ctx context.Context, executor *jobexecutor.JobExecutor, job *rivertype.JobRow) { + p.numJobsStuck.Add(1) + totalStuckJobs := int(p.config.JobStuckCount.Add(1)) + + if p.config.JobStuckHandler == nil { + return + } + + result := p.config.JobStuckHandler(ctx, JobStuckHandlerParams{ + ID: job.ID, + Kind: job.Kind, + Queue: job.Queue, + TotalStuckJobs: totalStuckJobs, + }) + if !result.AddWorkerSlot || !executor.TryCloseSlot() { + return + } + + p.numJobsActive.Add(-1) + if p.fetchLimiter != nil { + p.fetchLimiter.Call() + } +} + +func (p *producer) handleWorkerUnstuck() { + p.numJobsStuck.Add(-1) + p.config.JobStuckCount.Add(-1) +} + +func (p *producer) maybeCancelJob(ctx context.Context, id int64) { + executor, ok := p.activeJobs[id] + if !ok { + return + } + executor.Cancel(ctx) +} + +func (p *producer) metricEmitHooksFromLookup() []rivertype.HookMetricEmit { + pluginLookup := p.config.PluginLookupGlobal + if pluginLookup == nil { + return nil + } + + plugins := pluginLookup.ByKind(pluginlookup.PluginKindHookMetricEmit) + if len(plugins) < 1 { + return nil + } + + metricEmitHooks := make([]rivertype.HookMetricEmit, len(plugins)) + for i, plugin := range plugins { + metricEmitHooks[i] = plugin.(rivertype.HookMetricEmit) //nolint:forcetypeassert + } + + return metricEmitHooks +} + +func (p *producer) dispatchWork(workCtx context.Context, count int, fetchResultCh chan<- producerFetchResult) { + // When a queue is paused, innerFetchLoop dispatches with count zero so it can + // continue servicing state changes without attempting to lock jobs or emit metrics. + if count <= 0 { + fetchResultCh <- producerFetchResult{} + return + } + + // This intentionally removes any deadlines or cancellation from the parent + // context because we don't want it to get cancelled if the producer is asked + // to shut down. In that situation, we want to finish fetching any jobs we are + // in the midst of fetching, work them, and then stop. Otherwise we'd have a + // risk of shutting down when we had already fetched jobs in the database, + // leaving those jobs stranded. We'd then potentially have to release them + // back to the queue. + ctx := context.WithoutCancel(workCtx) + + // Maximum size of the `attempted_by` array on each job row. This maximum is + // rarely hit, but exists to protect against degenerate cases. + const maxAttemptedBy = 100 + + var startedAt time.Time + if len(p.metricEmitHooks) > 0 { + startedAt = time.Now() + } + + jobs, err := p.pilot.JobGetAvailable(ctx, p.exec, p.state, &riverdriver.JobGetAvailableParams{ + ClientID: p.config.ClientID, + MaxAttemptedBy: maxAttemptedBy, + MaxToLock: count, + Now: p.Time.NowOrNil(), + Queue: p.config.Queue, + ProducerID: p.id.Load(), + Schema: p.config.Schema, + }) + if err != nil { + fetchResultCh <- producerFetchResult{err: err} + return + } + + if len(p.metricEmitHooks) > 0 { + p.emitMetric(ctx, &rivertype.HookMetricEmitParams{ + Metric: &rivertype.JobGetAvailableDurationMetric{ + Duration: time.Since(startedAt), + Queue: p.config.Queue, + }, + }) + p.emitMetric(ctx, &rivertype.HookMetricEmitParams{ + Metric: &rivertype.JobGetAvailableCountMetric{ + Count: len(jobs), + Queue: p.config.Queue, + }, + }) + } + + fetchResultCh <- producerFetchResult{jobs: jobs} +} + +func (p *producer) emitMetric(ctx context.Context, params *rivertype.HookMetricEmitParams) { + for _, hook := range p.metricEmitHooks { + hook.MetricEmit(ctx, params) + } +} + +// Periodically logs an informational log line giving some insight into the +// current state of the producer. +func (p *producer) heartbeatLogLoop(ctx context.Context, wg *sync.WaitGroup) { + defer wg.Done() + + ticker := time.NewTicker(5 * time.Second) + defer ticker.Stop() + type jobCount struct { + active int + ran uint64 + stuck int + } + var prevCount jobCount + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + curCount := jobCount{ + active: int(p.numJobsActive.Load()), + ran: p.numJobsRan.Load(), + stuck: int(p.numJobsStuck.Load()), + } + if curCount != prevCount { + p.Logger.InfoContext(ctx, p.Name+": Producer job counts", + slog.Uint64("num_completed_jobs", curCount.ran), + slog.Int("num_jobs_running", curCount.active), + slog.Int("num_jobs_stuck", curCount.stuck), + slog.String("queue", p.config.Queue), + ) + } + prevCount = curCount + } + } +} + +func (p *producer) startNewExecutors(workCtx context.Context, jobs []*rivertype.JobRow) { + defaultClientRetryPolicy := retrypolicy.NewDefault(p.Time) + + for _, job := range jobs { + workInfo, ok := p.workers.workersMap[job.Kind] + + var workUnit workunit.WorkUnit + if ok { + workUnit = workInfo.workUnitFactory.MakeUnit(job) + } + + // jobCancel will always be called by the executor to prevent leaks. + jobCtx, jobCancel := context.WithCancelCause(workCtx) + + var executor *jobexecutor.JobExecutor + executor = baseservice.Init(&p.Archetype, &jobexecutor.JobExecutor{ + CancelFunc: jobCancel, + ClientJobTimeout: p.jobTimeout, + ClientRetryPolicy: p.retryPolicy, + Completer: p.completer, + DefaultClientRetryPolicy: defaultClientRetryPolicy, + ErrorHandler: p.errorHandler, + PluginLookupByJob: p.config.PluginLookupByJob, + PluginLookupGlobal: p.config.PluginLookupGlobal, + JobRow: job, + ProducerCallbacks: struct { + JobDone func(jobRow *rivertype.JobRow) + Stuck func(ctx context.Context, jobRow *rivertype.JobRow) + Unstuck func() + }{ + JobDone: p.handleWorkerDone, + Stuck: func(ctx context.Context, jobRow *rivertype.JobRow) { p.handleWorkerStuck(ctx, executor, jobRow) }, + Unstuck: p.handleWorkerUnstuck, + }, + SchedulerInterval: p.config.SchedulerInterval, + StuckThresholdOverride: p.config.JobStuckThreshold, + WorkUnit: workUnit, + }) + p.addActiveJob(job.ID, executor) + + go executor.Execute(jobCtx) + } + + p.Logger.DebugContext(workCtx, p.Name+": Distributed batch of jobs to executors", "num_jobs", len(jobs)) + + p.testSignals.StartedExecutors.Signal(struct{}{}) +} + +func (p *producer) maxJobsToFetch() int { + return p.config.MaxWorkers - int(p.numJobsActive.Load()) +} + +func (p *producer) handleWorkerDone(job *rivertype.JobRow) { + p.jobResultCh <- job +} + +func (p *producer) pollForSettingChanges(ctx context.Context, wg *sync.WaitGroup, lastPaused bool, lastMetadata []byte) { + defer wg.Done() + + ticker := time.NewTicker(p.config.QueuePollInterval) + defer ticker.Stop() + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + updatedQueue, err := timeoututil.WithTimeoutV(ctx, 10*time.Second, p.Name+".pollForSettingChanges", func(ctx context.Context) (*rivertype.Queue, error) { + return p.exec.QueueGet(ctx, &riverdriver.QueueGetParams{ + Name: p.config.Queue, + Schema: p.config.Schema, + }) + }) + if err != nil { + // Don't log if this is part of a standard shutdown. + if !errors.Is(context.Cause(ctx), startstop.ErrStop) { + p.Logger.ErrorContext(ctx, p.Name+": Error fetching queue settings", slog.String("err", err.Error())) + } + continue + } + + if updatedQueue == nil { + p.Logger.ErrorContext(ctx, p.Name+": Queue row not found when polling for setting changes", slog.String("queue", p.config.Queue)) + continue + } + + // Look for a change in the paused state: + shouldBePaused := (updatedQueue.PausedAt != nil) + if lastPaused != shouldBePaused { + action := controlActionPause + if !shouldBePaused { + action = controlActionResume + } + payload := &controlEventPayload{ + Action: action, + Queue: p.config.Queue, + } + p.Logger.DebugContext(ctx, p.Name+": Queue control state changed from polling", + slog.String("queue", p.config.Queue), + slog.String("action", string(action)), + slog.Bool("paused", shouldBePaused), + ) + + select { + case p.queueControlCh <- payload: + lastPaused = shouldBePaused + default: + p.Logger.WarnContext(ctx, p.Name+": Queue control notification dropped due to full buffer", slog.String("action", string(action))) + } + } + + // Look for a change in the queue's metadata: + if !metadataEqual(lastMetadata, updatedQueue.Metadata) { + payload := &controlEventPayload{ + Action: controlActionMetadataChanged, + Queue: p.config.Queue, + Metadata: updatedQueue.Metadata, + } + p.Logger.DebugContext(ctx, p.Name+": Queue metadata changed from polling", + slog.String("queue", p.config.Queue), + ) + + select { + case p.queueControlCh <- payload: + lastMetadata = updatedQueue.Metadata + default: + p.Logger.WarnContext(ctx, p.Name+": Queue control notification dropped due to full buffer", slog.String("action", string(controlActionMetadataChanged))) + } + } + + p.testSignals.PolledQueueConfig.Signal(struct{}{}) + } + } +} + +func (p *producer) reportProducerStatusLoop(ctx context.Context, wg *sync.WaitGroup) { + defer wg.Done() + + serviceutil.CancellableSleep(ctx, randutil.DurationBetween(0, time.Second)) + reportTicker := timeutil.NewTickerWithInitialTick(ctx, p.config.ProducerReportInterval) + for { + select { + case <-ctx.Done(): + return + case <-reportTicker.C: + p.reportProducerStatusOnce(ctx) + } + } +} + +func (p *producer) reportProducerStatusOnce(ctx context.Context) { + err := timeoututil.WithTimeout(ctx, 10*time.Second, p.Name+".reportProducerStatusOnce", func(ctx context.Context) error { + p.Logger.DebugContext(ctx, p.Name+": Reporting producer status", slog.Int64("id", p.id.Load()), slog.String("queue", p.config.Queue)) + return p.pilot.ProducerKeepAlive(ctx, p.exec, &riverdriver.ProducerKeepAliveParams{ + ID: p.id.Load(), + QueueName: p.config.Queue, + Schema: p.config.Schema, + StaleUpdatedAtHorizon: p.Time.Now().Add(-p.config.StaleProducerRetentionPeriod), + }) + }) + if err != nil && errors.Is(context.Cause(ctx), startstop.ErrStop) { + return + } + if err != nil { + p.Logger.ErrorContext(ctx, p.Name+": Producer status update, error updating in database", + slog.Int64("id", p.id.Load()), + slog.String("queue", p.config.Queue), + slog.String("err", err.Error()), + ) + return + } + p.testSignals.ReportedProducerStatus.Signal(struct{}{}) +} + +func (p *producer) reportQueueStatusLoop(ctx context.Context, wg *sync.WaitGroup) { + defer wg.Done() + + serviceutil.CancellableSleep(ctx, randutil.DurationBetween(0, time.Second)) + reportTicker := time.NewTicker(p.config.QueueReportInterval) + for { + select { + case <-ctx.Done(): + reportTicker.Stop() + return + case <-reportTicker.C: + p.reportQueueStatusOnce(ctx) + } + } +} + +func (p *producer) reportQueueStatusOnce(ctx context.Context) { + err := timeoututil.WithTimeout(ctx, 10*time.Second, p.Name+".reportQueueStatusOnce", func(ctx context.Context) error { + p.Logger.DebugContext(ctx, p.Name+": Reporting queue status", slog.String("queue", p.config.Queue)) + _, err := p.exec.QueueCreateOrSetUpdatedAt(ctx, &riverdriver.QueueCreateOrSetUpdatedAtParams{ + Metadata: []byte("{}"), + Name: p.config.Queue, + Now: p.Time.NowOrNil(), + Schema: p.config.Schema, + }) + return err + }) + if err != nil && errors.Is(context.Cause(ctx), startstop.ErrStop) { + return + } + if err != nil { + p.Logger.ErrorContext(ctx, p.Name+": Queue status update, error updating in database", slog.String("err", err.Error())) + return + } + p.testSignals.ReportedQueueStatus.Signal(struct{}{}) +} + +type producerFetchResult struct { + jobs []*rivertype.JobRow + err error +} + +type errorHandlerAdapter struct { + errorHandler ErrorHandler +} + +func (e *errorHandlerAdapter) HandleError(ctx context.Context, job *rivertype.JobRow, err error) *jobexecutor.ErrorHandlerResult { + result := e.errorHandler.HandleError(ctx, job, err) + return (*jobexecutor.ErrorHandlerResult)(result) +} + +func (e *errorHandlerAdapter) HandlePanic(ctx context.Context, job *rivertype.JobRow, panicVal any, trace string) *jobexecutor.ErrorHandlerResult { + result := e.errorHandler.HandlePanic(ctx, job, panicVal, trace) + return (*jobexecutor.ErrorHandlerResult)(result) +} + +// metadataEqual compares two JSON byte slices for semantic equality by parsing +// them into maps and re-marshaling them. This handles cases where the JSON is +// equivalent but formatted differently (whitespace, field order, etc). +func metadataEqual(a, b []byte) bool { + var unmarshaledA, unmarshaledB map[string]any + if err := json.Unmarshal(a, &unmarshaledA); err != nil { + return false + } + if err := json.Unmarshal(b, &unmarshaledB); err != nil { + return false + } + marshaledA, err := json.Marshal(unmarshaledA) + if err != nil { + return false + } + marshaledB, err := json.Marshal(unmarshaledB) + if err != nil { + return false + } + return bytes.Equal(marshaledA, marshaledB) +} diff --git a/vendor/github.com/riverqueue/river/queue_list_params.go b/vendor/github.com/riverqueue/river/queue_list_params.go new file mode 100644 index 0000000000..59d8bcc799 --- /dev/null +++ b/vendor/github.com/riverqueue/river/queue_list_params.go @@ -0,0 +1,40 @@ +package river + +// QueueListParams specifies the parameters for a QueueList query. It must be +// initialized with NewQueueListParams. Params can be built by chaining methods +// on the QueueListParams object: +// +// params := NewQueueListParams().First(100) +type QueueListParams struct { + paginationCount int32 +} + +// NewQueueListParams creates a new QueueListParams to return available queues +// sorted by time in ascending order, returning 100 jobs at most. +func NewQueueListParams() *QueueListParams { + return &QueueListParams{ + paginationCount: 100, + } +} + +func (p *QueueListParams) copy() *QueueListParams { + return &QueueListParams{ + paginationCount: p.paginationCount, + } +} + +// First returns an updated filter set that will only return the first count +// queues. +// +// Count must be between 1 and 10000, inclusive, or this will panic. +func (p *QueueListParams) First(count int) *QueueListParams { + if count <= 0 { + panic("count must be > 0") + } + if count > 10000 { + panic("count must be <= 10000") + } + result := p.copy() + result.paginationCount = int32(count) + return result +} diff --git a/vendor/github.com/riverqueue/river/queue_pause_opts.go b/vendor/github.com/riverqueue/river/queue_pause_opts.go new file mode 100644 index 0000000000..4e89689873 --- /dev/null +++ b/vendor/github.com/riverqueue/river/queue_pause_opts.go @@ -0,0 +1,4 @@ +package river + +// QueuePauseOpts are optional settings for pausing or resuming a queue. +type QueuePauseOpts struct{} diff --git a/vendor/github.com/riverqueue/river/recorded_output.go b/vendor/github.com/riverqueue/river/recorded_output.go new file mode 100644 index 0000000000..a7879e3cff --- /dev/null +++ b/vendor/github.com/riverqueue/river/recorded_output.go @@ -0,0 +1,79 @@ +package river + +import ( + "context" + "encoding/json" + "errors" + "fmt" + + "github.com/riverqueue/river/internal/jobexecutor" + "github.com/riverqueue/river/rivertype" +) + +const ( + maxOutputSizeMB = 32 + maxOutputSizeBytes = maxOutputSizeMB * 1024 * 1024 +) + +// RecordOutput records output JSON from a job. The "output" can be any +// JSON-encodable value and will be stored in the database on the job row after +// the current execution attempt completes. Output may be useful for debugging, +// or for storing the result of a job temporarily without needing to create a +// dedicated table to keep it in. +// +// For example, with workflows, it's common for subsequent task to depend on +// something done in an earlier dependency task. Consider the creation of an +// external resource in another API or in an database—it will typically have a +// unique ID that must be used to reference the resource later. A later step +// may require that info in order to complete its work, and the output can be +// a convenient way to store that info. +// +// Output is stored in the job's metadata under the `"output"` key +// ([github.com/riverqueue/river/rivertype.MetadataKeyOutput]). +// This function must be called within an Worker's Work function. It returns an +// error if called anywhere else. As with any stored value, care should be taken +// to ensure that the payload size is not too large. Output is limited to 32MB +// in size for safety, but should be kept much smaller than this. +// +// Only one output can be stored per job. If this function is called more than +// once, the output will be overwritten with the latest value. The output also +// must be recorded _before_ the job finishes executing so that it can be stored +// when the job's row is updated. +// +// Once recorded, the output is stored regardless of the outcome of the +// execution attempt (success, error, panic, etc.). +// +// RecordOutput always stores output lazily as a job is being completed (whether +// that's completion to success or failure). Client.JobUpdate and JobUpdateTx +// are available to store output eagerly at any time, including from inside a +// work function as the job is being executed. +// +// The output is marshalled to JSON as part of this function and it will return +// an error if the output is not JSON-encodable. +func RecordOutput(ctx context.Context, output any) error { + metadataUpdates, hasMetadataUpdates := jobexecutor.MetadataUpdatesFromWorkContext(ctx) + if !hasMetadataUpdates { + return errors.New("RecordOutput must be called within a Worker") + } + + outputBytes, err := json.Marshal(output) + if err != nil { + return err + } + + if err := checkOutputSize(outputBytes); err != nil { + return err + } + + metadataUpdates[rivertype.MetadataKeyOutput] = json.RawMessage(outputBytes) + return nil +} + +// Postgres JSONB is limited to 255MB, but it would be a bad idea to get +// anywhere close to that limit for output. +func checkOutputSize(outputBytes []byte) error { + if len(outputBytes) > maxOutputSizeBytes { + return fmt.Errorf("output is too large: %d bytes (max %d MB)", len(outputBytes), maxOutputSizeMB) + } + return nil +} diff --git a/vendor/github.com/riverqueue/river/resumable.go b/vendor/github.com/riverqueue/river/resumable.go new file mode 100644 index 0000000000..272b3d9c15 --- /dev/null +++ b/vendor/github.com/riverqueue/river/resumable.go @@ -0,0 +1,184 @@ +package river + +import ( + "context" + "encoding/json" + "errors" + "fmt" + + "github.com/riverqueue/river/internal/riverplugin" +) + +var ( + errResumableStepNotInWorker = errors.New("river: resumable step can only be used within a Worker") + errResumableCursorNotInStep = errors.New("river: resumable cursor can only be used within ResumableStepCursor") +) + +// ResumableSetCursor records a cursor for the current resumable cursor step. +// The cursor is stored only if the job attempt ends in an error, allowing a +// later retry to resume the same step from the recorded position. +// +// Alternatively, ResumableSetStepCursorTx is available to persist a step and +// cursor immediately as part of a transaction, guaranteeing that it's stored +// durably. +func ResumableSetCursor[TCursor any](ctx context.Context, cursor TCursor) error { + state := mustResumableState(ctx) + if state.StepName == "" { + return errResumableCursorNotInStep + } + + cursorBytes, err := json.Marshal(cursor) + if err != nil { + return err + } + + if state.Cursors == nil { + state.Cursors = make(map[string]json.RawMessage) + } + state.Cursors[state.StepName] = json.RawMessage(cursorBytes) + return nil +} + +// StepOpts are options for ResumableStep and ResumableStepCursor. There are +// currently no available options, but this space is reserved for future use. +type StepOpts struct{} + +// ResumableStep runs a resumable step, skipping the step on a later retry if +// an earlier attempt already completed it successfully. +// Step names must be unique across all ResumableStep and ResumableStepCursor +// calls in the same Worker execution. +// +// After a step returns an error, no subsequent steps will be run and the +// overall job will be marked as failed with that error. Be careful to put all +// executable code in steps, because any code outside of them will be run, even +// if a step returned an error. +// +// opts may be nil. +func ResumableStep(ctx context.Context, name string, opts *StepOpts, stepFunc func(ctx context.Context) error) { + state := mustResumableState(ctx) + if state.Err != nil { + return + } + if !registerResumableStepName(state, name) { + return + } + + if !state.ResumeMatched { + if name == state.ResumeStep { + state.CompletedStep = name + state.ResumeMatched = true + } + return + } + + previousStepName := state.StepName + state.StepName = name + defer func() { state.StepName = previousStepName }() + + if err := stepFunc(ctx); err != nil { + state.Err = err + return + } + + state.CompletedStep = name +} + +// ResumableStepCursor runs a resumable step that also receives a persisted +// cursor value from an earlier failed attempt, if one was recorded with +// ResumableSetCursor. +// Step names must be unique across all ResumableStep and ResumableStepCursor +// calls in the same Worker execution. +// +// The cursor type T is user-specified. It may be a primitive value like an +// integer ID, or a more complex type like a struct with multiple fields. It's +// stored in a job's metadata, so it needs to be marshable and unmarshable to +// and from JSON. +// +// Notably, it's the responsibility of the step function to call +// ResumableSetCursor with an updated cursor value as progress is made, and to +// check the cursor value before running to determine where to resume from. +// +// After a step returns an error, no subsequent steps will be run and the +// overall job will be marked as failed with that error. Be careful to put all +// executable code in steps, because any code outside of them will be run, even +// if a step returned an error. +// +// opts may be nil. +func ResumableStepCursor[TCursor any](ctx context.Context, name string, opts *StepOpts, stepFunc func(ctx context.Context, cursor TCursor) error) { + state := mustResumableState(ctx) + if state.Err != nil { + return + } + if !registerResumableStepName(state, name) { + return + } + + if !state.ResumeMatched { + if name == state.ResumeStep { + state.CompletedStep = name + state.ResumeMatched = true + + // If cursor data exists for this step, it was only partially + // completed on the previous attempt. Fall through to re-execute + // it with the cursor rather than skipping it. + if _, hasCursor := state.Cursors[name]; !hasCursor { + return + } + } else { + return + } + } + + var cursor TCursor + if cursorBytes, ok := state.Cursors[name]; ok && len(cursorBytes) > 0 { + if err := json.Unmarshal(cursorBytes, &cursor); err != nil { + state.Err = fmt.Errorf("river: unmarshal resumable cursor for step %q: %w", name, err) + return + } + } + + previousStepName := state.StepName + state.StepName = name + defer func() { state.StepName = previousStepName }() + + if err := stepFunc(ctx, cursor); err != nil { + state.Err = err + return + } + + state.CompletedStep = name + delete(state.Cursors, name) +} + +func mustResumableState(ctx context.Context) *riverplugin.ResumableState { + state, ok := resumableStateFromContext(ctx) + if !ok { + panic(errResumableStepNotInWorker) + } + + return state +} + +func registerResumableStepName(state *riverplugin.ResumableState, name string) bool { + if _, ok := state.AllStepNames[name]; ok { + state.Err = fmt.Errorf("river: duplicate resumable step name %q", name) + return false + } + + state.AllStepNames[name] = struct{}{} + return true +} + +func resumableStateFromContext(ctx context.Context) (*riverplugin.ResumableState, bool) { + state := ctx.Value(riverplugin.ResumableContextKey{}) + if state == nil { + return nil, false + } + + typedState, ok := state.(*riverplugin.ResumableState) + if !ok || typedState == nil { + return nil, false + } + + return typedState, true +} diff --git a/vendor/github.com/riverqueue/river/resumable_step_tx.go b/vendor/github.com/riverqueue/river/resumable_step_tx.go new file mode 100644 index 0000000000..fddb8c63af --- /dev/null +++ b/vendor/github.com/riverqueue/river/resumable_step_tx.go @@ -0,0 +1,124 @@ +package river + +import ( + "context" + "encoding/json" + "errors" + + "github.com/riverqueue/river/internal/execution" + "github.com/riverqueue/river/internal/jobexecutor" + "github.com/riverqueue/river/internal/rivercommon" + "github.com/riverqueue/river/riverdriver" + "github.com/riverqueue/river/rivertype" +) + +// ResumableSetStepTx immediately persists the current resumable step as +// part of transaction tx. If tx is rolled back, the step update will be as +// well. +// +// Normally, a resumable job's step progress is recorded after it runs along +// with its result status. This is normally sufficient, but because it happens +// out-of-transaction, there's a chance that it doesn't happen in case of panic +// or other abrupt termination. This function is useful in cases where a +// resumable worker needs a guarantee of a checkpoint being recorded durably, at +// the cost of an extra database operation. +// +// Must be called from within a ResumableStep or ResumableStepCursor callback. +// The current step name to persist is read from context. +func ResumableSetStepTx[TDriver riverdriver.Driver[TTx], TTx any, TArgs JobArgs](ctx context.Context, tx TTx, job *Job[TArgs]) (*Job[TArgs], error) { + return resumableSetStepTx(ctx, tx, job, nil) +} + +// ResumableSetStepCursorTx immediately persists the current resumable step and +// cursor as part of transaction tx. If tx is rolled back, the step and cursor +// update will be as well. +// +// Normally, a resumable job's step progress is recorded after it runs along +// with its result status. This is normally sufficient, but because it happens +// out-of-transaction, there's a chance that it doesn't happen in case of panic +// or other abrupt termination. This function is useful in cases where a +// resumable worker needs a guarantee of a checkpoint being recorded durably, at +// the cost of an extra database operation. +// +// Must be called from within a ResumableStepCursor callback. The current step +// name to persist is read from context. +func ResumableSetStepCursorTx[TDriver riverdriver.Driver[TTx], TTx any, TArgs JobArgs, TCursor any](ctx context.Context, tx TTx, job *Job[TArgs], cursor TCursor) (*Job[TArgs], error) { + cursorBytes, err := json.Marshal(cursor) + if err != nil { + return nil, err + } + + return resumableSetStepTx(ctx, tx, job, json.RawMessage(cursorBytes)) +} + +func resumableSetStepTx[TTx any, TArgs JobArgs](ctx context.Context, tx TTx, job *Job[TArgs], cursor json.RawMessage) (*Job[TArgs], error) { + if job.State != rivertype.JobStateRunning { + return nil, errors.New("job must be running") + } + + state, ok := resumableStateFromContext(ctx) + if !ok { + return nil, errors.New("not inside a resumable step; must be called from within ResumableStep or ResumableStepCursor") + } + if state.StepName == "" { + return nil, errors.New("not inside a resumable step; must be called from within ResumableStep or ResumableStepCursor") + } + + step := state.StepName + + client := ClientFromContext[TTx](ctx) + if client == nil { + return nil, errors.New("client not found in context, can only work within a River worker") + } + + metadataUpdates := map[string]any{ + rivercommon.MetadataKeyResumableStep: step, + } + + state.CompletedStep = step + if cursor != nil { + if state.Cursors == nil { + state.Cursors = make(map[string]json.RawMessage) + } + state.Cursors[step] = cursor + } + if len(state.Cursors) > 0 { + metadataUpdates[rivercommon.MetadataKeyResumableCursor] = state.Cursors + } + + workMetadataUpdates, hasWorkMetadataUpdates := jobexecutor.MetadataUpdatesFromWorkContext(ctx) + if hasWorkMetadataUpdates { + workMetadataUpdates[rivercommon.MetadataKeyResumableStep] = step + if resumableCursorMetadata, ok := metadataUpdates[rivercommon.MetadataKeyResumableCursor]; ok { + workMetadataUpdates[rivercommon.MetadataKeyResumableCursor] = resumableCursorMetadata + } + } + + metadataUpdatesBytes, err := json.Marshal(metadataUpdates) + if err != nil { + return nil, err + } + + updatedJob, err := client.Driver().UnwrapExecutor(tx).JobUpdate(ctx, &riverdriver.JobUpdateParams{ + ID: job.ID, + MetadataDoMerge: true, + Metadata: metadataUpdatesBytes, + Schema: client.config.Schema, + }) + if err != nil { + if errors.Is(err, rivertype.ErrNotFound) { + if _, isInsideTestWorker := ctx.Value(execution.ContextKeyInsideTestWorker{}).(bool); isInsideTestWorker { + panic("to use ResumableSetStepTx or ResumableSetStepCursorTx in a rivertest.Worker, the job must be inserted into the database first") + } + } + + return nil, err + } + + result := &Job[TArgs]{JobRow: updatedJob} + if err := json.Unmarshal(result.EncodedArgs, &result.Args); err != nil { + return nil, err + } + + return result, nil +} diff --git a/vendor/github.com/riverqueue/river/retry_policy.go b/vendor/github.com/riverqueue/river/retry_policy.go new file mode 100644 index 0000000000..021551f15e --- /dev/null +++ b/vendor/github.com/riverqueue/river/retry_policy.go @@ -0,0 +1,57 @@ +package river + +import ( + "time" + + "github.com/riverqueue/river/internal/retrypolicy" + "github.com/riverqueue/river/rivertype" +) + +// ClientRetryPolicy is an interface that can be implemented to provide a retry +// policy for how River deals with failed jobs at the client level (when a +// worker does not define an override for `NextRetry`). Jobs are scheduled to be +// retried in the future up until they've reached the job's max attempts, at +// which pointed they're set as discarded. +// +// The ClientRetryPolicy does not have access to generics and operates on the +// raw JobRow struct with encoded args. +type ClientRetryPolicy interface { + // NextRetry calculates when the next retry for a failed job should take place + // given when it was last attempted and its number of attempts, or any other + // of the job's properties a user-configured retry policy might want to + // consider. + NextRetry(job *rivertype.JobRow) time.Time +} + +// DefaultClientRetryPolicy is River's default retry policy. +type DefaultClientRetryPolicy struct { + timeNowFunc func() time.Time +} + +// NextRetry gets the next retry given for the given job, accounting for when it +// was last attempted and what attempt number that was. Reschedules using a +// basic exponential backoff of `ATTEMPT^4`, so after the first failure a new +// try will be scheduled in 1 seconds, 16 seconds after the second, 1 minute and +// 21 seconds after the third, etc. +// +// Snoozes do not count as attempts and do not influence retry behavior. +// Earlier versions of River would allow the attempt to increment each time a +// job was snoozed. Although this has been changed and snoozes now decrement the +// attempt count, we can maintain the same retry schedule even for pre-existing +// jobs by using the number of errors instead of the attempt count. This ensures +// consistent behavior across River versions. +// +// At degenerately high retry counts (>= 310) the policy starts adding the +// equivalent of the maximum of time.Duration to each retry, about 292 years. +// The schedule is no longer exponential past this point. +func (p *DefaultClientRetryPolicy) NextRetry(job *rivertype.JobRow) time.Time { + return retrypolicy.NextRetryAt(p.timeNowUTC(), job) +} + +func (p *DefaultClientRetryPolicy) timeNowUTC() time.Time { + if p.timeNowFunc != nil { + return p.timeNowFunc() + } + + return time.Now().UTC() +} diff --git a/vendor/github.com/riverqueue/river/riverdriver/LICENSE b/vendor/github.com/riverqueue/river/riverdriver/LICENSE new file mode 100644 index 0000000000..2f8ed188e8 --- /dev/null +++ b/vendor/github.com/riverqueue/river/riverdriver/LICENSE @@ -0,0 +1,374 @@ +Mozilla Public License Version 2.0 +================================== + +1. Definitions +-------------- + +1.1. "Contributor" + means each individual or legal entity that creates, contributes to + the creation of, or owns Covered Software. + +1.2. "Contributor Version" + means the combination of the Contributions of others (if any) used + by a Contributor and that particular Contributor's Contribution. + +1.3. "Contribution" + means Covered Software of a particular Contributor. + +1.4. "Covered Software" + means Source Code Form to which the initial Contributor has attached + the notice in Exhibit A, the Executable Form of such Source Code + Form, and Modifications of such Source Code Form, in each case + including portions thereof. + +1.5. "Incompatible With Secondary Licenses" + means + + (a) that the initial Contributor has attached the notice described + in Exhibit B to the Covered Software; or + + (b) that the Covered Software was made available under the terms of + version 1.1 or earlier of the License, but not also under the + terms of a Secondary License. + +1.6. "Executable Form" + means any form of the work other than Source Code Form. + +1.7. "Larger Work" + means a work that combines Covered Software with other material, in + a separate file or files, that is not Covered Software. + +1.8. "License" + means this document. + +1.9. "Licensable" + means having the right to grant, to the maximum extent possible, + whether at the time of the initial grant or subsequently, any and + all of the rights conveyed by this License. + +1.10. "Modifications" + means any of the following: + + (a) any file in Source Code Form that results from an addition to, + deletion from, or modification of the contents of Covered + Software; or + + (b) any new file in Source Code Form that contains any Covered + Software. + +1.11. "Patent Claims" of a Contributor + means any patent claim(s), including without limitation, method, + process, and apparatus claims, in any patent Licensable by such + Contributor that would be infringed, but for the grant of the + License, by the making, using, selling, offering for sale, having + made, import, or transfer of either its Contributions or its + Contributor Version. + +1.12. "Secondary License" + means either the GNU General Public License, Version 2.0, the GNU + Lesser General Public License, Version 2.1, the GNU Affero General + Public License, Version 3.0, or any later versions of those + licenses. + +1.13. "Source Code Form" + means the form of the work preferred for making modifications. + +1.14. "You" (or "Your") + means an individual or a legal entity exercising rights under this + License. For legal entities, "You" includes any entity that + controls, is controlled by, or is under common control with You. For + purposes of this definition, "control" means (a) the power, direct + or indirect, to cause the direction or management of such entity, + whether by contract or otherwise, or (b) ownership of more than + fifty percent (50%) of the outstanding shares or beneficial + ownership of such entity. + +2. License Grants and Conditions +-------------------------------- + +2.1. Grants + +Each Contributor hereby grants You a world-wide, royalty-free, +non-exclusive license: + +(a) under intellectual property rights (other than patent or trademark) + Licensable by such Contributor to use, reproduce, make available, + modify, display, perform, distribute, and otherwise exploit its + Contributions, either on an unmodified basis, with Modifications, or + as part of a Larger Work; and + +(b) under Patent Claims of such Contributor to make, use, sell, offer + for sale, have made, import, and otherwise transfer either its + Contributions or its Contributor Version. + +2.2. Effective Date + +The licenses granted in Section 2.1 with respect to any Contribution +become effective for each Contribution on the date the Contributor first +distributes such Contribution. + +2.3. Limitations on Grant Scope + +The licenses granted in this Section 2 are the only rights granted under +this License. No additional rights or licenses will be implied from the +distribution or licensing of Covered Software under this License. +Notwithstanding Section 2.1(b) above, no patent license is granted by a +Contributor: + +(a) for any code that a Contributor has removed from Covered Software; + or + +(b) for infringements caused by: (i) Your and any other third party's + modifications of Covered Software, or (ii) the combination of its + Contributions with other software (except as part of its Contributor + Version); or + +(c) under Patent Claims infringed by Covered Software in the absence of + its Contributions. + +This License does not grant any rights in the trademarks, service marks, +or logos of any Contributor (except as may be necessary to comply with +the notice requirements in Section 3.4). + +2.4. Subsequent Licenses + +No Contributor makes additional grants as a result of Your choice to +distribute the Covered Software under a subsequent version of this +License (see Section 10.2) or under the terms of a Secondary License (if +permitted under the terms of Section 3.3). + +2.5. Representation + +Each Contributor represents that the Contributor believes its +Contributions are its original creation(s) or it has sufficient rights +to grant the rights to its Contributions conveyed by this License. + +2.6. Fair Use + +This License is not intended to limit any rights You have under +applicable copyright doctrines of fair use, fair dealing, or other +equivalents. + +2.7. Conditions + +Sections 3.1, 3.2, 3.3, and 3.4 are conditions of the licenses granted +in Section 2.1. + +3. Responsibilities +------------------- + +3.1. Distribution of Source Form + +All distribution of Covered Software in Source Code Form, including any +Modifications that You create or to which You contribute, must be under +the terms of this License. You must inform recipients that the Source +Code Form of the Covered Software is governed by the terms of this +License, and how they can obtain a copy of this License. You may not +attempt to alter or restrict the recipients' rights in the Source Code +Form. + +3.2. Distribution of Executable Form + +If You distribute Covered Software in Executable Form then: + +(a) such Covered Software must also be made available in Source Code + Form, as described in Section 3.1, and You must inform recipients of + the Executable Form how they can obtain a copy of such Source Code + Form by reasonable means in a timely manner, at a charge no more + than the cost of distribution to the recipient; and + +(b) You may distribute such Executable Form under the terms of this + License, or sublicense it under different terms, provided that the + license for the Executable Form does not attempt to limit or alter + the recipients' rights in the Source Code Form under this License. + +3.3. Distribution of a Larger Work + +You may create and distribute a Larger Work under terms of Your choice, +provided that You also comply with the requirements of this License for +the Covered Software. If the Larger Work is a combination of Covered +Software with a work governed by one or more Secondary Licenses, and the +Covered Software is not Incompatible With Secondary Licenses, this +License permits You to additionally distribute such Covered Software +under the terms of such Secondary License(s), so that the recipient of +the Larger Work may, at their option, further distribute the Covered +Software under the terms of either this License or such Secondary +License(s). + +3.4. Notices + +You may not remove or alter the substance of any license notices +(including copyright notices, patent notices, disclaimers of warranty, +or limitations of liability) contained within the Source Code Form of +the Covered Software, except that You may alter any license notices to +the extent required to remedy known factual inaccuracies. + +3.5. Application of Additional Terms + +You may choose to offer, and to charge a fee for, warranty, support, +indemnity or liability obligations to one or more recipients of Covered +Software. However, You may do so only on Your own behalf, and not on +behalf of any Contributor. You must make it absolutely clear that any +such warranty, support, indemnity, or liability obligation is offered by +You alone, and You hereby agree to indemnify every Contributor for any +liability incurred by such Contributor as a result of warranty, support, +indemnity or liability terms You offer. You may include additional +disclaimers of warranty and limitations of liability specific to any +jurisdiction. + +4. Inability to Comply Due to Statute or Regulation +--------------------------------------------------- + +If it is impossible for You to comply with any of the terms of this +License with respect to some or all of the Covered Software due to +statute, judicial order, or regulation then You must: (a) comply with +the terms of this License to the maximum extent possible; and (b) +describe the limitations and the code they affect. Such description must +be placed in a text file included with all distributions of the Covered +Software under this License. Except to the extent prohibited by statute +or regulation, such description must be sufficiently detailed for a +recipient of ordinary skill to be able to understand it. + +5. Termination +-------------- + +5.1. The rights granted under this License will terminate automatically +if You fail to comply with any of its terms. However, if You become +compliant, then the rights granted under this License from a particular +Contributor are reinstated (a) provisionally, unless and until such +Contributor explicitly and finally terminates Your grants, and (b) on an +ongoing basis, if such Contributor fails to notify You of the +non-compliance by some reasonable means prior to 60 days after You have +come back into compliance. Moreover, Your grants from a particular +Contributor are reinstated on an ongoing basis if such Contributor +notifies You of the non-compliance by some reasonable means, this is the +first time You have received notice of non-compliance with this License +from such Contributor, and You become compliant prior to 30 days after +Your receipt of the notice. + +5.2. If You initiate litigation against any entity by asserting a patent +infringement claim (excluding declaratory judgment actions, +counter-claims, and cross-claims) alleging that a Contributor Version +directly or indirectly infringes any patent, then the rights granted to +You by any and all Contributors for the Covered Software under Section +2.1 of this License shall terminate. + +5.3. In the event of termination under Sections 5.1 or 5.2 above, all +end user license agreements (excluding distributors and resellers) which +have been validly granted by You or Your distributors under this License +prior to termination shall survive termination. + +************************************************************************ +* * +* 6. Disclaimer of Warranty * +* ------------------------- * +* * +* Covered Software is provided under this License on an "as is" * +* basis, without warranty of any kind, either expressed, implied, or * +* statutory, including, without limitation, warranties that the * +* Covered Software is free of defects, merchantable, fit for a * +* particular purpose or non-infringing. The entire risk as to the * +* quality and performance of the Covered Software is with You. * +* Should any Covered Software prove defective in any respect, You * +* (not any Contributor) assume the cost of any necessary servicing, * +* repair, or correction. This disclaimer of warranty constitutes an * +* essential part of this License. No use of any Covered Software is * +* authorized under this License except under this disclaimer. * +* * +************************************************************************ + +************************************************************************ +* * +* 7. Limitation of Liability * +* -------------------------- * +* * +* Under no circumstances and under no legal theory, whether tort * +* (including negligence), contract, or otherwise, shall any * +* Contributor, or anyone who distributes Covered Software as * +* permitted above, be liable to You for any direct, indirect, * +* special, incidental, or consequential damages of any character * +* including, without limitation, damages for lost profits, loss of * +* goodwill, work stoppage, computer failure or malfunction, or any * +* and all other commercial damages or losses, even if such party * +* shall have been informed of the possibility of such damages. This * +* limitation of liability shall not apply to liability for death or * +* personal injury resulting from such party's negligence to the * +* extent applicable law prohibits such limitation. Some * +* jurisdictions do not allow the exclusion or limitation of * +* incidental or consequential damages, so this exclusion and * +* limitation may not apply to You. * +* * +************************************************************************ + +8. Litigation +------------- + +Any litigation relating to this License may be brought only in the +courts of a jurisdiction where the defendant maintains its principal +place of business and such litigation shall be governed by laws of that +jurisdiction, without reference to its conflict-of-law provisions. +Nothing in this Section shall prevent a party's ability to bring +cross-claims or counter-claims. + +9. Miscellaneous +---------------- + +This License represents the complete agreement concerning the subject +matter hereof. If any provision of this License is held to be +unenforceable, such provision shall be reformed only to the extent +necessary to make it enforceable. Any law or regulation which provides +that the language of a contract shall be construed against the drafter +shall not be used to construe this License against a Contributor. + +10. Versions of the License +--------------------------- + +10.1. New Versions + +Mozilla Foundation is the license steward. Except as provided in Section +10.3, no one other than the license steward has the right to modify or +publish new versions of this License. Each version will be given a +distinguishing version number. + +10.2. Effect of New Versions + +You may distribute the Covered Software under the terms of the version +of the License under which You originally received the Covered Software, +or under the terms of any subsequent version published by the license +steward. + +10.3. Modified Versions + +If you create software not governed by this License, and you want to +create a new license for such software, you may create and use a +modified version of this License if you rename the license and remove +any references to the name of the license steward (except to note that +such modified license differs from this License). + +10.4. Distributing Source Code Form that is Incompatible With Secondary +Licenses + +If You choose to distribute Source Code Form that is Incompatible With +Secondary Licenses under the terms of this version of the License, the +notice described in Exhibit B of this License must be attached. + +Exhibit A - Source Code Form License Notice +------------------------------------------- + + This Source Code Form is subject to the terms of the Mozilla Public + License, v. 2.0. If a copy of the MPL was not distributed with this + file, You can obtain one at http://mozilla.org/MPL/2.0/. + +If it is not possible or desirable to put the notice in a particular +file, then You may include the notice in a location (such as a LICENSE +file in a relevant directory) where a recipient would be likely to look +for such a notice. + +You may add additional accurate notices of copyright ownership. + +Exhibit B - "Incompatible With Secondary Licenses" Notice +--------------------------------------------------------- + + This Source Code Form is "Incompatible With Secondary Licenses", as + defined by the Mozilla Public License, v. 2.0. + diff --git a/vendor/github.com/riverqueue/river/riverdriver/river_driver_interface.go b/vendor/github.com/riverqueue/river/riverdriver/river_driver_interface.go new file mode 100644 index 0000000000..6ab9e8127a --- /dev/null +++ b/vendor/github.com/riverqueue/river/riverdriver/river_driver_interface.go @@ -0,0 +1,931 @@ +// Package riverdriver exposes generic constructs to be implemented by specific +// drivers that wrap third party database packages, with the aim being to keep +// the main River interface decoupled from a specific database package so that +// other packages or other major versions of packages can be supported in future +// River versions. +// +// River currently only supports Pgx v5, and the interface here wrap it with +// only the thinnest possible layer. Adding support for alternate packages will +// require the interface to change substantially, and therefore it should not be +// implemented or invoked by user code. Changes to interfaces in this package +// WILL NOT be considered breaking changes for purposes of River's semantic +// versioning. +package riverdriver + +import ( + "context" + "errors" + "fmt" + "io/fs" + "time" + + "github.com/riverqueue/river/rivertype" +) + +const AllQueuesString = "*" + +const ( + DatabaseNamePostgres = "postgres" + DatabaseNameSQLite = "sqlite" +) + +const MigrationLineMain = "main" + +var ( + ErrClosedPool = errors.New("underlying driver pool is closed") + ErrNotImplemented = errors.New("driver does not implement this functionality") +) + +// Driver provides a database driver for use with river.Client. +// +// Its purpose is to wrap the interface of a third party database package, with +// the aim being to keep the main River interface decoupled from a specific +// database package so that other packages or major versions of packages can be +// supported in future River versions. +// +// River currently only supports Pgx v5, and this interface wraps it with only +// the thinnest possible layer. Adding support for alternate packages will +// require it to change substantially, and therefore it should not be +// implemented or invoked by user code. Changes to this interface WILL NOT be +// considered breaking changes for purposes of River's semantic versioning. +// +// API is not stable. DO NOT IMPLEMENT. +type Driver[TTx any] interface { + // ArgPlaceholder is the placeholder character used in query positional + // arguments, so "$" for "$1", "$2", "$3", etc. This is a "$" for Postgres + // and "?" for SQLite. + // + // API is not stable. DO NOT USE. + ArgPlaceholder() string + + // DatabaseName is the name of the database that the driver targets like + // "postgres" or "sqlite". This is used for purposes like a cache key prefix + // in riverdbtest so that multiple drivers may share schemas as long as they + // target the same database. + // + // API is not stable. DO NOT USE. + DatabaseName() string + + // GetExecutor gets an executor for the driver. + // + // API is not stable. DO NOT USE. + GetExecutor() Executor + + // GetListener gets a listener for purposes of receiving notifications. + // + // API is not stable. DO NOT USE. + GetListener(params *GetListenenerParams) Listener + + // GetMigrationDefaultLines gets default migration lines that should be + // applied when using this driver. This is mainly used by riverdbtest to + // figure out what migration lines should be available by default for new + // test schemas. + // + // API is not stable. DO NOT USE. + GetMigrationDefaultLines() []string + + // GetMigrationFS gets a filesystem containing migrations for the driver. + // + // Each set of migration files is expected to exist within the filesystem as + // `migration//`. For example: + // + // migration/main/001_create_river_migration.up.sql + // + // API is not stable. DO NOT USE. + GetMigrationFS(line string) fs.FS + + // GetMigrationLines gets supported migration lines from the driver. Most + // drivers will only support a single line: MigrationLineMain. + // + // API is not stable. DO NOT USE. + GetMigrationLines() []string + + // GetMigrationTruncateTables gets the tables that should be truncated + // before or after tests for a specific migration line returned by this + // driver. Tables to truncate doesn't need to consider intermediary states, + // and should return tables for the latest migration version. + // + // API is not stable. DO NOT USE. + GetMigrationTruncateTables(line string, version int) []string + + // PoolIsSet returns true if the driver is configured with a database pool. + // + // API is not stable. DO NOT USE. + PoolIsSet() bool + + // PoolSet sets a database pool into a driver will a nil pool. This is meant + // only for use in testing, and only in specific circumstances where it's + // needed. The pool in a driver should generally be treated as immutable + // because it's inherited by driver executors, and changing if when active + // executors exist will cause problems. + // + // Most drivers don't implement this function and return ErrNotImplemented. + // + // Drivers should only set a pool if the previous pool was nil (to help root + // out bugs where something unexpected is happening), and panic in case a + // pool is set to a driver twice. + // + // API is not stable. DO NOT USE. + PoolSet(dbPool any) error + + // SQLFragmentColumnContainsAll generates an SQL fragment to be included as + // a predicate in a `WHERE` query for a collection column containing all of + // the given values. PostgreSQL uses array containment while SQLite compares + // values from a JSON array. + // + // API is not stable. DO NOT USE. + SQLFragmentColumnContainsAll(column, namedArg string, values []string) (string, any, error) + + // SQLFragmentColumnContainsAny generates an SQL fragment to be included as + // a predicate in a `WHERE` query for a collection column containing at least + // one of the given values. PostgreSQL uses array overlap while SQLite + // compares values from a JSON array. + // + // API is not stable. DO NOT USE. + SQLFragmentColumnContainsAny(column, namedArg string, values []string) (string, any, error) + + // SQLFragmentColumnIn generates an SQL fragment to be included as a + // predicate in a `WHERE` query for the existence of a set of values in a + // column like `id IN (...)`. The actual implementation depends on support + // for specific data types. Postgres uses arrays while SQLite uses a JSON + // fragment with `json_each`. + // + // API is not stable. DO NOT USE. + SQLFragmentColumnIn(column string, values any) (string, any, error) + + // SupportsListener gets whether this driver supports a listener. Drivers + // that don't support a listener support poll only mode only. + // + // API is not stable. DO NOT USE. + SupportsListener() bool + + // SupportsListenNotify indicates whether the driver can broadcast + // notifications that a listener can receive, either through a native + // database mechanism like Postgres LISTEN/NOTIFY or a driver-specific + // emulation. This differs from SupportsListener in that even if a driver + // doesn't support a listener but the database supports the underlying + // notification mechanism, it will still broadcast in case there are other + // clients/drivers on the database that do support a listener. If + // notifications can't be supported at all, no broadcast attempt is made. + // + // API is not stable. DO NOT USE. + SupportsListenNotify() bool + + // TimePrecision returns the maximum time resolution supported by the + // database. This is used in test assertions when checking round trips on + // timestamps. + // + // API is not stable. DO NOT USE. + TimePrecision() time.Duration + + // UnwrapExecutor gets an executor from a driver transaction. + // + // API is not stable. DO NOT USE. + UnwrapExecutor(tx TTx) ExecutorTx + + // UnwrapTx gets a driver transaction from an executor. This is currently + // only needed for test transaction helpers. + // + // API is not stable. DO NOT USE. + UnwrapTx(execTx ExecutorTx) TTx +} + +// Executor provides River operations against a database. It may be a database +// pool or transaction. +// +// API is not stable. DO NOT IMPLEMENT. +type Executor interface { + // Begin begins a new subtransaction. ErrSubTxNotSupported may be returned + // if the executor is a transaction and the driver doesn't support + // subtransactions (like riverdriver/riverdatabasesql for database/sql). + Begin(ctx context.Context) (ExecutorTx, error) + + // ColumnExists checks whether a column for a particular table exists for + // the schema in the current search schema. + ColumnExists(ctx context.Context, params *ColumnExistsParams) (bool, error) + + // Exec executes raw SQL. Used for migrations. + Exec(ctx context.Context, sql string, args ...any) error + + // IndexDropIfExists drops a database index if exists. This abstraction is a + // little leaky right now because Postgres runs this `CONCURRENTLY` and + // that's not possible in SQLite. + // + // API is not stable. DO NOT USE. + IndexDropIfExists(ctx context.Context, params *IndexDropIfExistsParams) error + IndexExists(ctx context.Context, params *IndexExistsParams) (bool, error) + IndexesExist(ctx context.Context, params *IndexesExistParams) (map[string]bool, error) + + // IndexReindex reindexes a database index. This abstraction is a little + // leaky right now because Postgres runs this `CONCURRENTLY` and that's not + // possible in SQLite. + // + // API is not stable. DO NOT USE. + IndexReindex(ctx context.Context, params *IndexReindexParams) error + IndexReindexArtifacts(ctx context.Context, params *IndexReindexArtifactsParams) ([]string, error) + + JobCancel(ctx context.Context, params *JobCancelParams) (*rivertype.JobRow, error) + JobCountByAllStates(ctx context.Context, params *JobCountByAllStatesParams) (map[rivertype.JobState]int, error) + JobCountByQueueAndState(ctx context.Context, params *JobCountByQueueAndStateParams) ([]*JobCountByQueueAndStateResult, error) + JobCountByState(ctx context.Context, params *JobCountByStateParams) (int, error) + JobDelete(ctx context.Context, params *JobDeleteParams) (*rivertype.JobRow, error) + JobDeleteBefore(ctx context.Context, params *JobDeleteBeforeParams) (int, error) + JobDeleteMany(ctx context.Context, params *JobDeleteManyParams) ([]*rivertype.JobRow, error) + JobGetAvailable(ctx context.Context, params *JobGetAvailableParams) ([]*rivertype.JobRow, error) + JobGetByID(ctx context.Context, params *JobGetByIDParams) (*rivertype.JobRow, error) + JobGetByIDMany(ctx context.Context, params *JobGetByIDManyParams) ([]*rivertype.JobRow, error) + JobGetByKindMany(ctx context.Context, params *JobGetByKindManyParams) ([]*rivertype.JobRow, error) + JobGetStuck(ctx context.Context, params *JobGetStuckParams) ([]*rivertype.JobRow, error) + JobInsertFastMany(ctx context.Context, params *JobInsertFastManyParams) ([]*JobInsertFastResult, error) + JobInsertFastManyNoReturning(ctx context.Context, params *JobInsertFastManyParams) (int, error) + JobInsertFull(ctx context.Context, params *JobInsertFullParams) (*rivertype.JobRow, error) + JobInsertFullMany(ctx context.Context, jobs *JobInsertFullManyParams) ([]*rivertype.JobRow, error) + JobKindList(ctx context.Context, params *JobKindListParams) ([]string, error) + JobList(ctx context.Context, params *JobListParams) ([]*rivertype.JobRow, error) + JobRescueMany(ctx context.Context, params *JobRescueManyParams) (*struct{}, error) + JobRetry(ctx context.Context, params *JobRetryParams) (*rivertype.JobRow, error) + JobSchedule(ctx context.Context, params *JobScheduleParams) ([]*JobScheduleResult, error) + JobSetStateIfRunningMany(ctx context.Context, params *JobSetStateIfRunningManyParams) ([]*rivertype.JobRow, error) + JobUpdate(ctx context.Context, params *JobUpdateParams) (*rivertype.JobRow, error) + JobUpdateFull(ctx context.Context, params *JobUpdateFullParams) (*rivertype.JobRow, error) + LeaderAttemptElect(ctx context.Context, params *LeaderElectParams) (*Leader, error) + LeaderAttemptReelect(ctx context.Context, params *LeaderReelectParams) (*Leader, error) + LeaderDeleteExpired(ctx context.Context, params *LeaderDeleteExpiredParams) (int, error) + LeaderGetElectedLeader(ctx context.Context, params *LeaderGetElectedLeaderParams) (*Leader, error) + LeaderInsert(ctx context.Context, params *LeaderInsertParams) (*Leader, error) + LeaderResign(ctx context.Context, params *LeaderResignParams) (bool, error) + + // MigrationDeleteAssumingMainMany deletes many migrations assuming + // everything is on the main line. This is suitable for use in databases on + // a version before the `line` column exists. + MigrationDeleteAssumingMainMany(ctx context.Context, params *MigrationDeleteAssumingMainManyParams) ([]*Migration, error) + + // MigrationDeleteByLineAndVersionMany deletes many migration versions on a + // particular line. + MigrationDeleteByLineAndVersionMany(ctx context.Context, params *MigrationDeleteByLineAndVersionManyParams) ([]*Migration, error) + + // MigrationGetAllAssumingMain gets all migrations assuming everything is on + // the main line. This is suitable for use in databases on a version before + // the `line` column exists. + MigrationGetAllAssumingMain(ctx context.Context, params *MigrationGetAllAssumingMainParams) ([]*Migration, error) + + // MigrationGetByLine gets all currently applied migrations. + MigrationGetByLine(ctx context.Context, params *MigrationGetByLineParams) ([]*Migration, error) + + // MigrationInsertMany inserts many migration versions. + MigrationInsertMany(ctx context.Context, params *MigrationInsertManyParams) ([]*Migration, error) + + // MigrationInsertManyAssumingMain inserts many migrations, assuming they're + // on the main line. This operation is necessary for compatibility before + // the `line` column was added to the migrations table. + MigrationInsertManyAssumingMain(ctx context.Context, params *MigrationInsertManyAssumingMainParams) ([]*Migration, error) + + // NotificationDeleteBefore deletes notifications before a certain time + // horizon. + // + // A "notification" in this context refers to a row in `river_notification` + // which is a special table implemented in some databases (e.g. SQLite) that + // simulates Postgres' listen/notify when not available. + NotificationDeleteBefore(ctx context.Context, params *NotificationDeleteBeforeParams) (int, error) + + NotifyMany(ctx context.Context, params *NotifyManyParams) error + PGAdvisoryXactLock(ctx context.Context, key int64) (*struct{}, error) + + QueueCreateOrSetUpdatedAt(ctx context.Context, params *QueueCreateOrSetUpdatedAtParams) (*rivertype.Queue, error) + QueueDeleteExpired(ctx context.Context, params *QueueDeleteExpiredParams) ([]string, error) + QueueGet(ctx context.Context, params *QueueGetParams) (*rivertype.Queue, error) + QueueList(ctx context.Context, params *QueueListParams) ([]*rivertype.Queue, error) + QueueNameList(ctx context.Context, params *QueueNameListParams) ([]string, error) + QueuePause(ctx context.Context, params *QueuePauseParams) error + QueueResume(ctx context.Context, params *QueueResumeParams) error + QueueUpdate(ctx context.Context, params *QueueUpdateParams) (*rivertype.Queue, error) + QueryRow(ctx context.Context, sql string, args ...any) Row + + SchemaCreate(ctx context.Context, params *SchemaCreateParams) error + SchemaDrop(ctx context.Context, params *SchemaDropParams) error + SchemaGetExpired(ctx context.Context, params *SchemaGetExpiredParams) ([]string, error) + + // TableExists checks whether a table exists for the schema in the current + // search schema. + TableExists(ctx context.Context, params *TableExistsParams) (bool, error) + TableTruncate(ctx context.Context, params *TableTruncateParams) error +} + +// ExecutorTx is an executor which is a transaction. In addition to standard +// Executor operations, it may be committed or rolled back. +// +// API is not stable. DO NOT IMPLEMENT. +type ExecutorTx interface { + Executor + + // Commit commits the transaction. + // + // API is not stable. DO NOT USE. + Commit(ctx context.Context) error + + // Rollback rolls back the transaction. + // + // API is not stable. DO NOT USE. + Rollback(ctx context.Context) error +} + +type GetListenenerParams struct { + Schema string +} + +// Listener listens for notifications. In Postgres, this is a database +// connection where `LISTEN` has been run. +// +// API is not stable. DO NOT IMPLEMENT. +type Listener interface { + Close(ctx context.Context) error + Connect(ctx context.Context) error + Listen(ctx context.Context, topic string) error + Ping(ctx context.Context) error + Schema() string + SetAfterConnectExec(sql string) // should only ever be used in testing + Unlisten(ctx context.Context, topic string) error + WaitForNotification(ctx context.Context) (*Notification, error) +} + +type Notification struct { + Payload string + Topic string +} + +type ColumnExistsParams struct { + Column string + Schema string + Table string +} + +type IndexDropIfExistsParams struct { + Index string + Schema string +} + +type IndexExistsParams struct { + Index string + Schema string +} + +type IndexesExistParams struct { + IndexNames []string + Schema string +} + +type JobCancelParams struct { + ID int64 + CancelAttemptedAt time.Time + ControlTopic string + Now *time.Time + Schema string +} + +type JobCountByAllStatesParams struct { + Schema string +} + +type JobCountByQueueAndStateParams struct { + QueueNames []string + Schema string +} + +type JobCountByQueueAndStateResult struct { + CountAvailable int64 + CountRunning int64 + Queue string +} + +type JobCountByStateParams struct { + Schema string + State rivertype.JobState +} + +type JobDeleteParams struct { + ID int64 + Schema string +} + +type JobDeleteBeforeParams struct { + CancelledDoDelete bool + CancelledFinalizedAtHorizon time.Time + CompletedDoDelete bool + CompletedFinalizedAtHorizon time.Time + DiscardedDoDelete bool + DiscardedFinalizedAtHorizon time.Time + Max int + QueuesExcluded []string + QueuesIncluded []string + Schema string +} + +type JobDeleteManyParams JobListParams + +type JobGetAvailableParams struct { + ClientID string + MaxAttemptedBy int + MaxToLock int + Now *time.Time + ProducerID int64 + Queue string + Schema string +} + +type JobGetByIDParams struct { + ID int64 + Schema string +} + +type JobGetByIDManyParams struct { + ID []int64 + Schema string +} + +type JobGetByKindManyParams struct { + Kind []string + Schema string +} + +type JobGetStuckParams struct { + AfterID int64 + Max int + Schema string + StuckHorizon time.Time +} + +type JobInsertFastParams struct { + ID *int64 + // Args contains the raw underlying job arguments struct. It has already been + // encoded into EncodedArgs, but the original is kept here for to leverage its + // struct tags and interfaces, such as for use in unique key generation. + Args rivertype.JobArgs + CreatedAt *time.Time + EncodedArgs []byte + Kind string + MaxAttempts int + Metadata []byte + Priority int + Queue string + ScheduledAt *time.Time + State rivertype.JobState + Tags []string + UniqueKey []byte + UniqueStates byte +} + +type JobInsertFastManyParams struct { + Jobs []*JobInsertFastParams + Schema string +} + +type JobInsertFastResult struct { + Job *rivertype.JobRow + UniqueSkippedAsDuplicate bool +} + +type JobInsertFullParams struct { + Attempt int + AttemptedAt *time.Time + AttemptedBy []string + CreatedAt *time.Time + EncodedArgs []byte + Errors [][]byte + FinalizedAt *time.Time + Kind string + MaxAttempts int + Metadata []byte + Priority int + Queue string + ScheduledAt *time.Time + Schema string + State rivertype.JobState + Tags []string + UniqueKey []byte + UniqueStates byte +} + +type JobInsertFullManyParams struct { + Jobs []*JobInsertFullParams + Schema string +} + +type JobKindListParams struct { + After string + Exclude []string + Match string + Max int + Schema string +} + +type JobListParams struct { + Max int32 + NamedArgs map[string]any + OrderByClause string + Schema string + WhereClause string +} + +type JobRescueManyParams struct { + ID []int64 + Error [][]byte + FinalizedAt []*time.Time + ScheduledAt []time.Time + Schema string + State []string + StuckHorizon time.Time +} + +type JobRetryParams struct { + ID int64 + Now *time.Time + Schema string +} + +type JobScheduleParams struct { + Max int + Now *time.Time + Schema string +} + +type JobScheduleResult struct { + Job rivertype.JobRow + ConflictDiscarded bool +} + +// JobSetStateIfRunningParams are parameters to update the state of a currently +// running job. Use one of the constructors below to ensure a correct +// combination of parameters. +type JobSetStateIfRunningParams struct { + ID int64 + Attempt *int + ErrData []byte + FinalizedAt *time.Time + MetadataDoMerge bool + MetadataUpdates []byte + ScheduledAt *time.Time + Schema string // added by completer + Snoozed bool + State rivertype.JobState +} + +func JobSetStateCancelled(id int64, finalizedAt time.Time, errData []byte, metadataUpdates []byte) *JobSetStateIfRunningParams { + return &JobSetStateIfRunningParams{ + ID: id, + ErrData: errData, + MetadataDoMerge: len(metadataUpdates) > 0, + MetadataUpdates: metadataUpdates, + FinalizedAt: &finalizedAt, + State: rivertype.JobStateCancelled, + } +} + +func JobSetStateCompleted(id int64, finalizedAt time.Time, metadataUpdates []byte) *JobSetStateIfRunningParams { + return &JobSetStateIfRunningParams{ + FinalizedAt: &finalizedAt, + ID: id, + MetadataDoMerge: len(metadataUpdates) > 0, + MetadataUpdates: metadataUpdates, + State: rivertype.JobStateCompleted, + } +} + +func JobSetStateDiscarded(id int64, finalizedAt time.Time, errData []byte, metadataUpdates []byte) *JobSetStateIfRunningParams { + return &JobSetStateIfRunningParams{ + ID: id, + ErrData: errData, + MetadataDoMerge: len(metadataUpdates) > 0, + MetadataUpdates: metadataUpdates, + FinalizedAt: &finalizedAt, + State: rivertype.JobStateDiscarded, + } +} + +func JobSetStateErrorAvailable(id int64, scheduledAt time.Time, errData []byte, metadataUpdates []byte) *JobSetStateIfRunningParams { + return &JobSetStateIfRunningParams{ + ID: id, + ErrData: errData, + MetadataDoMerge: len(metadataUpdates) > 0, + MetadataUpdates: metadataUpdates, + ScheduledAt: &scheduledAt, + State: rivertype.JobStateAvailable, + } +} + +func JobSetStateErrorRetryable(id int64, scheduledAt time.Time, errData []byte, metadataUpdates []byte) *JobSetStateIfRunningParams { + return &JobSetStateIfRunningParams{ + ID: id, + ErrData: errData, + MetadataDoMerge: len(metadataUpdates) > 0, + MetadataUpdates: metadataUpdates, + ScheduledAt: &scheduledAt, + State: rivertype.JobStateRetryable, + } +} + +func JobSetStateSnoozed(id int64, scheduledAt time.Time, attempt int, metadataUpdates []byte) *JobSetStateIfRunningParams { + return &JobSetStateIfRunningParams{ + Attempt: &attempt, + ID: id, + MetadataDoMerge: len(metadataUpdates) > 0, + MetadataUpdates: metadataUpdates, + ScheduledAt: &scheduledAt, + Snoozed: true, + State: rivertype.JobStateScheduled, + } +} + +func JobSetStateSnoozedAvailable(id int64, scheduledAt time.Time, attempt int, metadataUpdates []byte) *JobSetStateIfRunningParams { + return &JobSetStateIfRunningParams{ + Attempt: &attempt, + ID: id, + MetadataDoMerge: len(metadataUpdates) > 0, + MetadataUpdates: metadataUpdates, + ScheduledAt: &scheduledAt, + Snoozed: true, + State: rivertype.JobStateAvailable, + } +} + +// JobSetStateIfRunningManyParams are parameters to update the state of +// currently running jobs. Use one of the constructors below to ensure a correct +// combination of parameters. +type JobSetStateIfRunningManyParams struct { + ID []int64 + Attempt []*int + ErrData [][]byte + FinalizedAt []*time.Time + MetadataDoMerge []bool + MetadataUpdates [][]byte + Now *time.Time + ScheduledAt []*time.Time + Schema string + State []rivertype.JobState +} + +type JobUpdateParams struct { + ID int64 + MetadataDoMerge bool + Metadata []byte + Schema string +} + +type JobUpdateFullParams struct { + ID int64 + AttemptDoUpdate bool + Attempt int + AttemptedAtDoUpdate bool + AttemptedAt *time.Time + AttemptedByDoUpdate bool + AttemptedBy []string + ErrorsDoUpdate bool + Errors [][]byte + FinalizedAtDoUpdate bool + FinalizedAt *time.Time + MaxAttemptsDoUpdate bool + MaxAttempts int + MetadataDoUpdate bool + Metadata []byte + Schema string + StateDoUpdate bool + State rivertype.JobState + // Deprecated and will be removed when advisory lock unique path is removed. + UniqueKeyDoUpdate bool + // Deprecated and will be removed when advisory lock unique path is removed. + UniqueKey []byte +} + +// Leader represents a River leader. +// +// API is not stable. DO NOT USE. +type Leader struct { + ElectedAt time.Time + ExpiresAt time.Time + LeaderID string +} + +type LeaderDeleteExpiredParams struct { + Now *time.Time + Schema string +} + +type LeaderGetElectedLeaderParams struct { + Schema string +} + +type LeaderInsertParams struct { + ElectedAt *time.Time + ExpiresAt *time.Time + LeaderID string + Now *time.Time + Schema string + TTL time.Duration +} + +type LeaderElectParams struct { + LeaderID string + Now *time.Time + Schema string + TTL time.Duration +} + +type LeaderReelectParams struct { + ElectedAt time.Time + LeaderID string + Now *time.Time + Schema string + TTL time.Duration +} + +type LeaderResignParams struct { + ElectedAt time.Time + LeaderID string + LeadershipTopic string + Schema string +} + +// Migration represents a River migration. +// +// API is not stable. DO NOT USE. +type Migration struct { + // CreatedAt is when the migration was initially created. + // + // API is not stable. DO NOT USE. + CreatedAt time.Time + + // Line is the migration line that the migration belongs to. + // + // API is not stable. DO NOT USE. + Line string + + // Version is the version of the migration. + // + // API is not stable. DO NOT USE. + Version int +} + +type MigrationDeleteAssumingMainManyParams struct { + Schema string + Versions []int +} + +type MigrationDeleteByLineAndVersionManyParams struct { + Line string + Schema string + Versions []int +} + +type MigrationGetAllAssumingMainParams struct { + Schema string +} + +type MigrationGetByLineParams struct { + Line string + Schema string +} + +type MigrationInsertManyParams struct { + Line string + Schema string + Versions []int +} + +type MigrationInsertManyAssumingMainParams struct { + Schema string + Versions []int +} + +// NotifyManyParams are parameters to issue many pubsub notifications all at +// once for a single topic. +type NotifyManyParams struct { + Payload []string + Topic string + Schema string +} + +type NotificationDeleteBeforeParams struct { + CreatedAtHorizon time.Time + Schema string +} + +type ProducerKeepAliveParams struct { + ID int64 + QueueName string + Schema string + StaleUpdatedAtHorizon time.Time +} + +type QueueCreateOrSetUpdatedAtParams struct { + Metadata []byte + Name string + Now *time.Time + PausedAt *time.Time + Schema string + UpdatedAt *time.Time +} + +type QueueDeleteExpiredParams struct { + Max int + Schema string + UpdatedAtHorizon time.Time +} + +type QueueGetParams struct { + Name string + Schema string +} + +type QueueListParams struct { + Max int + Schema string +} + +type QueueNameListParams struct { + After string + Exclude []string + Match string + Max int + Schema string +} + +type QueuePauseParams struct { + Name string + Now *time.Time + Schema string +} + +type QueueResumeParams struct { + Name string + Now *time.Time + Schema string +} + +type QueueUpdateParams struct { + Metadata []byte + MetadataDoUpdate bool + Name string + Schema string +} + +type Row interface { + Scan(dest ...any) error +} + +type IndexReindexParams struct { + Index string + Schema string +} + +type IndexReindexArtifactsParams struct { + Index string + Schema string +} + +type Schema struct { + Name string +} + +type SchemaCreateParams struct { + Schema string +} + +type SchemaDropParams struct { + Schema string +} + +type SchemaGetExpiredParams struct { + BeforeName string + Prefix string +} + +type TableExistsParams struct { + Schema string + Table string +} + +type TableTruncateParams struct { + Schema string + Table []string +} + +// MigrationLineMainTruncateTables is a shared helper that produces tables to +// truncate for the main migration line. It's reused across all drivers. +// +// API is not stable. DO NOT USE. +func MigrationLineMainTruncateTables(version int) []string { + // 0 value must be handled and should always point to latest migration version + switch version { + case 1: + return nil // don't truncate `river_migrate` + case 2, 3: + return []string{"river_job", "river_leader"} + case 4: + return []string{"river_job", "river_leader", "river_queue"} + case 5, 6: + return []string{"river_job", "river_leader", "river_queue", "river_client", "river_client_queue"} + case 0, 7: + return []string{"river_job", "river_leader", "river_queue", "river_notification"} + } + + panic(fmt.Sprintf("unrecognized migration version: %d", version)) +} diff --git a/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/LICENSE b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/LICENSE new file mode 100644 index 0000000000..2f8ed188e8 --- /dev/null +++ b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/LICENSE @@ -0,0 +1,374 @@ +Mozilla Public License Version 2.0 +================================== + +1. Definitions +-------------- + +1.1. "Contributor" + means each individual or legal entity that creates, contributes to + the creation of, or owns Covered Software. + +1.2. "Contributor Version" + means the combination of the Contributions of others (if any) used + by a Contributor and that particular Contributor's Contribution. + +1.3. "Contribution" + means Covered Software of a particular Contributor. + +1.4. "Covered Software" + means Source Code Form to which the initial Contributor has attached + the notice in Exhibit A, the Executable Form of such Source Code + Form, and Modifications of such Source Code Form, in each case + including portions thereof. + +1.5. "Incompatible With Secondary Licenses" + means + + (a) that the initial Contributor has attached the notice described + in Exhibit B to the Covered Software; or + + (b) that the Covered Software was made available under the terms of + version 1.1 or earlier of the License, but not also under the + terms of a Secondary License. + +1.6. "Executable Form" + means any form of the work other than Source Code Form. + +1.7. "Larger Work" + means a work that combines Covered Software with other material, in + a separate file or files, that is not Covered Software. + +1.8. "License" + means this document. + +1.9. "Licensable" + means having the right to grant, to the maximum extent possible, + whether at the time of the initial grant or subsequently, any and + all of the rights conveyed by this License. + +1.10. "Modifications" + means any of the following: + + (a) any file in Source Code Form that results from an addition to, + deletion from, or modification of the contents of Covered + Software; or + + (b) any new file in Source Code Form that contains any Covered + Software. + +1.11. "Patent Claims" of a Contributor + means any patent claim(s), including without limitation, method, + process, and apparatus claims, in any patent Licensable by such + Contributor that would be infringed, but for the grant of the + License, by the making, using, selling, offering for sale, having + made, import, or transfer of either its Contributions or its + Contributor Version. + +1.12. "Secondary License" + means either the GNU General Public License, Version 2.0, the GNU + Lesser General Public License, Version 2.1, the GNU Affero General + Public License, Version 3.0, or any later versions of those + licenses. + +1.13. "Source Code Form" + means the form of the work preferred for making modifications. + +1.14. "You" (or "Your") + means an individual or a legal entity exercising rights under this + License. For legal entities, "You" includes any entity that + controls, is controlled by, or is under common control with You. For + purposes of this definition, "control" means (a) the power, direct + or indirect, to cause the direction or management of such entity, + whether by contract or otherwise, or (b) ownership of more than + fifty percent (50%) of the outstanding shares or beneficial + ownership of such entity. + +2. License Grants and Conditions +-------------------------------- + +2.1. Grants + +Each Contributor hereby grants You a world-wide, royalty-free, +non-exclusive license: + +(a) under intellectual property rights (other than patent or trademark) + Licensable by such Contributor to use, reproduce, make available, + modify, display, perform, distribute, and otherwise exploit its + Contributions, either on an unmodified basis, with Modifications, or + as part of a Larger Work; and + +(b) under Patent Claims of such Contributor to make, use, sell, offer + for sale, have made, import, and otherwise transfer either its + Contributions or its Contributor Version. + +2.2. Effective Date + +The licenses granted in Section 2.1 with respect to any Contribution +become effective for each Contribution on the date the Contributor first +distributes such Contribution. + +2.3. Limitations on Grant Scope + +The licenses granted in this Section 2 are the only rights granted under +this License. No additional rights or licenses will be implied from the +distribution or licensing of Covered Software under this License. +Notwithstanding Section 2.1(b) above, no patent license is granted by a +Contributor: + +(a) for any code that a Contributor has removed from Covered Software; + or + +(b) for infringements caused by: (i) Your and any other third party's + modifications of Covered Software, or (ii) the combination of its + Contributions with other software (except as part of its Contributor + Version); or + +(c) under Patent Claims infringed by Covered Software in the absence of + its Contributions. + +This License does not grant any rights in the trademarks, service marks, +or logos of any Contributor (except as may be necessary to comply with +the notice requirements in Section 3.4). + +2.4. Subsequent Licenses + +No Contributor makes additional grants as a result of Your choice to +distribute the Covered Software under a subsequent version of this +License (see Section 10.2) or under the terms of a Secondary License (if +permitted under the terms of Section 3.3). + +2.5. Representation + +Each Contributor represents that the Contributor believes its +Contributions are its original creation(s) or it has sufficient rights +to grant the rights to its Contributions conveyed by this License. + +2.6. Fair Use + +This License is not intended to limit any rights You have under +applicable copyright doctrines of fair use, fair dealing, or other +equivalents. + +2.7. Conditions + +Sections 3.1, 3.2, 3.3, and 3.4 are conditions of the licenses granted +in Section 2.1. + +3. Responsibilities +------------------- + +3.1. Distribution of Source Form + +All distribution of Covered Software in Source Code Form, including any +Modifications that You create or to which You contribute, must be under +the terms of this License. You must inform recipients that the Source +Code Form of the Covered Software is governed by the terms of this +License, and how they can obtain a copy of this License. You may not +attempt to alter or restrict the recipients' rights in the Source Code +Form. + +3.2. Distribution of Executable Form + +If You distribute Covered Software in Executable Form then: + +(a) such Covered Software must also be made available in Source Code + Form, as described in Section 3.1, and You must inform recipients of + the Executable Form how they can obtain a copy of such Source Code + Form by reasonable means in a timely manner, at a charge no more + than the cost of distribution to the recipient; and + +(b) You may distribute such Executable Form under the terms of this + License, or sublicense it under different terms, provided that the + license for the Executable Form does not attempt to limit or alter + the recipients' rights in the Source Code Form under this License. + +3.3. Distribution of a Larger Work + +You may create and distribute a Larger Work under terms of Your choice, +provided that You also comply with the requirements of this License for +the Covered Software. If the Larger Work is a combination of Covered +Software with a work governed by one or more Secondary Licenses, and the +Covered Software is not Incompatible With Secondary Licenses, this +License permits You to additionally distribute such Covered Software +under the terms of such Secondary License(s), so that the recipient of +the Larger Work may, at their option, further distribute the Covered +Software under the terms of either this License or such Secondary +License(s). + +3.4. Notices + +You may not remove or alter the substance of any license notices +(including copyright notices, patent notices, disclaimers of warranty, +or limitations of liability) contained within the Source Code Form of +the Covered Software, except that You may alter any license notices to +the extent required to remedy known factual inaccuracies. + +3.5. Application of Additional Terms + +You may choose to offer, and to charge a fee for, warranty, support, +indemnity or liability obligations to one or more recipients of Covered +Software. However, You may do so only on Your own behalf, and not on +behalf of any Contributor. You must make it absolutely clear that any +such warranty, support, indemnity, or liability obligation is offered by +You alone, and You hereby agree to indemnify every Contributor for any +liability incurred by such Contributor as a result of warranty, support, +indemnity or liability terms You offer. You may include additional +disclaimers of warranty and limitations of liability specific to any +jurisdiction. + +4. Inability to Comply Due to Statute or Regulation +--------------------------------------------------- + +If it is impossible for You to comply with any of the terms of this +License with respect to some or all of the Covered Software due to +statute, judicial order, or regulation then You must: (a) comply with +the terms of this License to the maximum extent possible; and (b) +describe the limitations and the code they affect. Such description must +be placed in a text file included with all distributions of the Covered +Software under this License. Except to the extent prohibited by statute +or regulation, such description must be sufficiently detailed for a +recipient of ordinary skill to be able to understand it. + +5. Termination +-------------- + +5.1. The rights granted under this License will terminate automatically +if You fail to comply with any of its terms. However, if You become +compliant, then the rights granted under this License from a particular +Contributor are reinstated (a) provisionally, unless and until such +Contributor explicitly and finally terminates Your grants, and (b) on an +ongoing basis, if such Contributor fails to notify You of the +non-compliance by some reasonable means prior to 60 days after You have +come back into compliance. Moreover, Your grants from a particular +Contributor are reinstated on an ongoing basis if such Contributor +notifies You of the non-compliance by some reasonable means, this is the +first time You have received notice of non-compliance with this License +from such Contributor, and You become compliant prior to 30 days after +Your receipt of the notice. + +5.2. If You initiate litigation against any entity by asserting a patent +infringement claim (excluding declaratory judgment actions, +counter-claims, and cross-claims) alleging that a Contributor Version +directly or indirectly infringes any patent, then the rights granted to +You by any and all Contributors for the Covered Software under Section +2.1 of this License shall terminate. + +5.3. In the event of termination under Sections 5.1 or 5.2 above, all +end user license agreements (excluding distributors and resellers) which +have been validly granted by You or Your distributors under this License +prior to termination shall survive termination. + +************************************************************************ +* * +* 6. Disclaimer of Warranty * +* ------------------------- * +* * +* Covered Software is provided under this License on an "as is" * +* basis, without warranty of any kind, either expressed, implied, or * +* statutory, including, without limitation, warranties that the * +* Covered Software is free of defects, merchantable, fit for a * +* particular purpose or non-infringing. The entire risk as to the * +* quality and performance of the Covered Software is with You. * +* Should any Covered Software prove defective in any respect, You * +* (not any Contributor) assume the cost of any necessary servicing, * +* repair, or correction. This disclaimer of warranty constitutes an * +* essential part of this License. No use of any Covered Software is * +* authorized under this License except under this disclaimer. * +* * +************************************************************************ + +************************************************************************ +* * +* 7. Limitation of Liability * +* -------------------------- * +* * +* Under no circumstances and under no legal theory, whether tort * +* (including negligence), contract, or otherwise, shall any * +* Contributor, or anyone who distributes Covered Software as * +* permitted above, be liable to You for any direct, indirect, * +* special, incidental, or consequential damages of any character * +* including, without limitation, damages for lost profits, loss of * +* goodwill, work stoppage, computer failure or malfunction, or any * +* and all other commercial damages or losses, even if such party * +* shall have been informed of the possibility of such damages. This * +* limitation of liability shall not apply to liability for death or * +* personal injury resulting from such party's negligence to the * +* extent applicable law prohibits such limitation. Some * +* jurisdictions do not allow the exclusion or limitation of * +* incidental or consequential damages, so this exclusion and * +* limitation may not apply to You. * +* * +************************************************************************ + +8. Litigation +------------- + +Any litigation relating to this License may be brought only in the +courts of a jurisdiction where the defendant maintains its principal +place of business and such litigation shall be governed by laws of that +jurisdiction, without reference to its conflict-of-law provisions. +Nothing in this Section shall prevent a party's ability to bring +cross-claims or counter-claims. + +9. Miscellaneous +---------------- + +This License represents the complete agreement concerning the subject +matter hereof. If any provision of this License is held to be +unenforceable, such provision shall be reformed only to the extent +necessary to make it enforceable. Any law or regulation which provides +that the language of a contract shall be construed against the drafter +shall not be used to construe this License against a Contributor. + +10. Versions of the License +--------------------------- + +10.1. New Versions + +Mozilla Foundation is the license steward. Except as provided in Section +10.3, no one other than the license steward has the right to modify or +publish new versions of this License. Each version will be given a +distinguishing version number. + +10.2. Effect of New Versions + +You may distribute the Covered Software under the terms of the version +of the License under which You originally received the Covered Software, +or under the terms of any subsequent version published by the license +steward. + +10.3. Modified Versions + +If you create software not governed by this License, and you want to +create a new license for such software, you may create and use a +modified version of this License if you rename the license and remove +any references to the name of the license steward (except to note that +such modified license differs from this License). + +10.4. Distributing Source Code Form that is Incompatible With Secondary +Licenses + +If You choose to distribute Source Code Form that is Incompatible With +Secondary Licenses under the terms of this version of the License, the +notice described in Exhibit B of this License must be attached. + +Exhibit A - Source Code Form License Notice +------------------------------------------- + + This Source Code Form is subject to the terms of the Mozilla Public + License, v. 2.0. If a copy of the MPL was not distributed with this + file, You can obtain one at http://mozilla.org/MPL/2.0/. + +If it is not possible or desirable to put the notice in a particular +file, then You may include the notice in a location (such as a LICENSE +file in a relevant directory) where a recipient would be likely to look +for such a notice. + +You may add additional accurate notices of copyright ownership. + +Exhibit B - "Incompatible With Secondary Licenses" Notice +--------------------------------------------------------- + + This Source Code Form is "Incompatible With Secondary Licenses", as + defined by the Mozilla Public License, v. 2.0. + diff --git a/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/copyfrom.go b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/copyfrom.go new file mode 100644 index 0000000000..5523c792dd --- /dev/null +++ b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/copyfrom.go @@ -0,0 +1,53 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.31.0 +// source: copyfrom.go + +package dbsqlc + +import ( + "context" +) + +// iteratorForJobInsertFastManyCopyFrom implements pgx.CopyFromSource. +type iteratorForJobInsertFastManyCopyFrom struct { + rows []*JobInsertFastManyCopyFromParams + skippedFirstNextCall bool +} + +func (r *iteratorForJobInsertFastManyCopyFrom) Next() bool { + if len(r.rows) == 0 { + return false + } + if !r.skippedFirstNextCall { + r.skippedFirstNextCall = true + return true + } + r.rows = r.rows[1:] + return len(r.rows) > 0 +} + +func (r iteratorForJobInsertFastManyCopyFrom) Values() ([]interface{}, error) { + return []interface{}{ + r.rows[0].Args, + r.rows[0].CreatedAt, + r.rows[0].Kind, + r.rows[0].MaxAttempts, + r.rows[0].Metadata, + r.rows[0].Priority, + r.rows[0].Queue, + r.rows[0].ScheduledAt, + r.rows[0].State, + r.rows[0].Tags, + r.rows[0].UniqueKey, + r.rows[0].UniqueStates, + }, nil +} + +func (r iteratorForJobInsertFastManyCopyFrom) Err() error { + return nil +} + +func (q *Queries) JobInsertFastManyCopyFrom(ctx context.Context, db DBTX, arg []*JobInsertFastManyCopyFromParams) (int64, error) { + return db.CopyFrom(ctx, []string{"river_job"}, []string{"args", "created_at", "kind", "max_attempts", "metadata", "priority", "queue", "scheduled_at", "state", "tags", "unique_key", "unique_states"}, &iteratorForJobInsertFastManyCopyFrom{rows: arg}) +} diff --git a/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/db.go b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/db.go new file mode 100644 index 0000000000..eaa98cafb2 --- /dev/null +++ b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/db.go @@ -0,0 +1,26 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.31.0 + +package dbsqlc + +import ( + "context" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgconn" +) + +type DBTX interface { + Exec(context.Context, string, ...interface{}) (pgconn.CommandTag, error) + Query(context.Context, string, ...interface{}) (pgx.Rows, error) + QueryRow(context.Context, string, ...interface{}) pgx.Row + CopyFrom(ctx context.Context, tableName pgx.Identifier, columnNames []string, rowSrc pgx.CopyFromSource) (int64, error) +} + +func New() *Queries { + return &Queries{} +} + +type Queries struct { +} diff --git a/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/models.go b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/models.go new file mode 100644 index 0000000000..4ad18ccb97 --- /dev/null +++ b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/models.go @@ -0,0 +1,110 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.31.0 + +package dbsqlc + +import ( + "database/sql/driver" + "fmt" + "time" + + "github.com/jackc/pgx/v5/pgtype" +) + +type RiverJobState string + +const ( + RiverJobStateAvailable RiverJobState = "available" + RiverJobStateCancelled RiverJobState = "cancelled" + RiverJobStateCompleted RiverJobState = "completed" + RiverJobStateDiscarded RiverJobState = "discarded" + RiverJobStatePending RiverJobState = "pending" + RiverJobStateRetryable RiverJobState = "retryable" + RiverJobStateRunning RiverJobState = "running" + RiverJobStateScheduled RiverJobState = "scheduled" +) + +func (e *RiverJobState) Scan(src interface{}) error { + switch s := src.(type) { + case []byte: + *e = RiverJobState(s) + case string: + *e = RiverJobState(s) + default: + return fmt.Errorf("unsupported scan type for RiverJobState: %T", src) + } + return nil +} + +type NullRiverJobState struct { + RiverJobState RiverJobState + Valid bool // Valid is true if RiverJobState is not NULL +} + +// Scan implements the Scanner interface. +func (ns *NullRiverJobState) Scan(value interface{}) error { + if value == nil { + ns.RiverJobState, ns.Valid = "", false + return nil + } + ns.Valid = true + return ns.RiverJobState.Scan(value) +} + +// Value implements the driver Valuer interface. +func (ns NullRiverJobState) Value() (driver.Value, error) { + if !ns.Valid { + return nil, nil + } + return string(ns.RiverJobState), nil +} + +type RiverJob struct { + ID int64 + Args []byte + Attempt int16 + AttemptedAt *time.Time + AttemptedBy []string + CreatedAt time.Time + Errors [][]byte + FinalizedAt *time.Time + Kind string + MaxAttempts int16 + Metadata []byte + Priority int16 + Queue string + State RiverJobState + ScheduledAt time.Time + Tags []string + UniqueKey []byte + UniqueStates pgtype.Bits +} + +type RiverLeader struct { + ElectedAt time.Time + ExpiresAt time.Time + LeaderID string + Name string +} + +type RiverMigration struct { + Line string + Version int64 + CreatedAt time.Time +} + +type RiverNotification struct { + ID int64 + CreatedAt time.Time + Payload string + Topic string +} + +type RiverQueue struct { + Name string + CreatedAt time.Time + Metadata []byte + PausedAt *time.Time + UpdatedAt time.Time +} diff --git a/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/pg_misc.sql b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/pg_misc.sql new file mode 100644 index 0000000000..19a7b99f67 --- /dev/null +++ b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/pg_misc.sql @@ -0,0 +1,14 @@ +-- name: PGAdvisoryXactLock :exec +SELECT pg_advisory_xact_lock(@key); + +-- name: PGNotifyMany :exec +WITH topic_to_notify AS ( + SELECT + concat(coalesce(sqlc.narg('schema')::text, current_schema()), '.', @topic::text) AS topic, + unnest(@payload::text[]) AS payload +) +SELECT pg_notify( + topic_to_notify.topic, + topic_to_notify.payload + ) +FROM topic_to_notify; diff --git a/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/pg_misc.sql.go b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/pg_misc.sql.go new file mode 100644 index 0000000000..9215c089bd --- /dev/null +++ b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/pg_misc.sql.go @@ -0,0 +1,45 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.31.0 +// source: pg_misc.sql + +package dbsqlc + +import ( + "context" + + "github.com/jackc/pgx/v5/pgtype" +) + +const pGAdvisoryXactLock = `-- name: PGAdvisoryXactLock :exec +SELECT pg_advisory_xact_lock($1) +` + +func (q *Queries) PGAdvisoryXactLock(ctx context.Context, db DBTX, key int64) error { + _, err := db.Exec(ctx, pGAdvisoryXactLock, key) + return err +} + +const pGNotifyMany = `-- name: PGNotifyMany :exec +WITH topic_to_notify AS ( + SELECT + concat(coalesce($1::text, current_schema()), '.', $2::text) AS topic, + unnest($3::text[]) AS payload +) +SELECT pg_notify( + topic_to_notify.topic, + topic_to_notify.payload + ) +FROM topic_to_notify +` + +type PGNotifyManyParams struct { + Schema pgtype.Text + Topic string + Payload []string +} + +func (q *Queries) PGNotifyMany(ctx context.Context, db DBTX, arg *PGNotifyManyParams) error { + _, err := db.Exec(ctx, pGNotifyMany, arg.Schema, arg.Topic, arg.Payload) + return err +} diff --git a/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql new file mode 100644 index 0000000000..509e479be1 --- /dev/null +++ b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql @@ -0,0 +1,730 @@ +CREATE TYPE river_job_state AS ENUM( + 'available', + 'cancelled', + 'completed', + 'discarded', + 'pending', + 'retryable', + 'running', + 'scheduled' +); + +CREATE TABLE river_job ( + id bigserial PRIMARY KEY, + args jsonb NOT NULL DEFAULT '{}', + attempt smallint NOT NULL DEFAULT 0, + attempted_at timestamptz, + attempted_by text[], + created_at timestamptz NOT NULL DEFAULT now(), + errors jsonb[], + finalized_at timestamptz, + kind text NOT NULL, + max_attempts smallint NOT NULL DEFAULT 25, + metadata jsonb NOT NULL DEFAULT '{}', + priority smallint NOT NULL DEFAULT 1, + queue text NOT NULL DEFAULT 'default', + state river_job_state NOT NULL DEFAULT 'available', + scheduled_at timestamptz NOT NULL DEFAULT now(), + tags varchar(255)[] NOT NULL DEFAULT '{}', + unique_key bytea, + unique_states bit(8), + CONSTRAINT finalized_or_finalized_at_null CHECK ( + (finalized_at IS NULL AND state NOT IN ('cancelled', 'completed', 'discarded')) OR + (finalized_at IS NOT NULL AND state IN ('cancelled', 'completed', 'discarded')) + ), + CONSTRAINT priority_in_range CHECK (priority >= 1 AND priority <= 4), + CONSTRAINT queue_length CHECK (char_length(queue) > 0 AND char_length(queue) < 128), + CONSTRAINT kind_length CHECK (char_length(kind) > 0 AND char_length(kind) < 128) +); + +-- name: JobCancel :one +WITH locked_job AS ( + SELECT + id, queue, state, finalized_at + FROM /* TEMPLATE: schema */river_job + WHERE river_job.id = @id + FOR UPDATE +), +notification AS ( + SELECT + id, + pg_notify( + concat(coalesce(sqlc.narg('schema')::text, current_schema()), '.', @control_topic::text), + json_build_object('action', 'cancel', 'job_id', id, 'queue', queue)::text + ) + FROM + locked_job + WHERE + state NOT IN ('cancelled', 'completed', 'discarded') + AND finalized_at IS NULL +), +updated_job AS ( + UPDATE /* TEMPLATE: schema */river_job + SET + -- If the job is actively running, we want to let its current client and + -- producer handle the cancellation. Otherwise, immediately cancel it. + state = CASE WHEN state = 'running' THEN state ELSE 'cancelled' END, + finalized_at = CASE WHEN state = 'running' THEN finalized_at ELSE coalesce(sqlc.narg('now')::timestamptz, now()) END, + -- Mark the job as cancelled by query so that the rescuer knows not to + -- rescue it, even if it gets stuck in the running state: + metadata = jsonb_set(metadata, '{cancel_attempted_at}'::text[], @cancel_attempted_at::jsonb, true) + FROM notification + WHERE river_job.id = notification.id + RETURNING river_job.* +) +SELECT * +FROM /* TEMPLATE: schema */river_job +WHERE id = @id::bigint + AND id NOT IN (SELECT id FROM updated_job) +UNION +SELECT * +FROM updated_job; + +-- name: JobCountByAllStates :many +SELECT state, count(*) +FROM /* TEMPLATE: schema */ river_job +GROUP BY state; + +-- name: JobCountByQueueAndState :many +WITH all_queues AS ( + SELECT DISTINCT unnest(@queue_names::text[])::text AS queue +), + +running_job_counts AS ( + SELECT + queue, + COUNT(*) AS count + FROM /* TEMPLATE: schema */river_job + WHERE queue = ANY(@queue_names::text[]) + AND state = 'running' + GROUP BY queue +), + +available_job_counts AS ( + SELECT + queue, + COUNT(*) AS count + FROM + /* TEMPLATE: schema */river_job + WHERE queue = ANY(@queue_names::text[]) + AND state = 'available' + GROUP BY queue +) + +SELECT + all_queues.queue, + COALESCE(available_job_counts.count, 0) AS count_available, + COALESCE(running_job_counts.count, 0) AS count_running +FROM + all_queues +LEFT JOIN + running_job_counts ON all_queues.queue = running_job_counts.queue +LEFT JOIN + available_job_counts ON all_queues.queue = available_job_counts.queue +ORDER BY all_queues.queue ASC; + +-- name: JobCountByState :one +SELECT count(*) +FROM /* TEMPLATE: schema */river_job +WHERE state = @state; + +-- name: JobDelete :one +WITH job_to_delete AS ( + SELECT id + FROM /* TEMPLATE: schema */river_job + WHERE river_job.id = @id + FOR UPDATE +), +deleted_job AS ( + DELETE + FROM /* TEMPLATE: schema */river_job + USING job_to_delete + WHERE river_job.id = job_to_delete.id + -- Do not touch running jobs: + AND river_job.state != 'running' + RETURNING river_job.* +) +SELECT * +FROM /* TEMPLATE: schema */river_job +WHERE id = @id::bigint + AND id NOT IN (SELECT id FROM deleted_job) +UNION +SELECT * +FROM deleted_job; + +-- name: JobDeleteBefore :execresult +DELETE FROM /* TEMPLATE: schema */river_job +WHERE id IN ( + SELECT id + FROM /* TEMPLATE: schema */river_job + WHERE ( + (state = 'cancelled' AND @cancelled_do_delete AND finalized_at < @cancelled_finalized_at_horizon::timestamptz) OR + (state = 'completed' AND @completed_do_delete AND finalized_at < @completed_finalized_at_horizon::timestamptz) OR + (state = 'discarded' AND @discarded_do_delete AND finalized_at < @discarded_finalized_at_horizon::timestamptz) + ) + AND ( + @queues_excluded::text[] IS NULL + OR NOT (queue = any(@queues_excluded)) + ) + AND ( + @queues_included::text[] IS NULL + OR queue = any(@queues_included) + ) + ORDER BY id + LIMIT @max::bigint +); + +-- name: JobDeleteMany :many +WITH jobs_to_delete AS ( + SELECT * + FROM /* TEMPLATE: schema */river_job + WHERE /* TEMPLATE_BEGIN: where_clause */ true /* TEMPLATE_END */ + AND state != 'running' + ORDER BY /* TEMPLATE_BEGIN: order_by_clause */ id /* TEMPLATE_END */ + LIMIT @max::int + FOR UPDATE + SKIP LOCKED +), +deleted_jobs AS ( + DELETE FROM /* TEMPLATE: schema */river_job + WHERE id IN (SELECT id FROM jobs_to_delete) + RETURNING * +) +-- this last SELECT step is necessary because there's no other way to define +-- order records come back from a DELETE statement +SELECT * +FROM /* TEMPLATE: schema */river_job +WHERE id IN (SELECT id FROM deleted_jobs) +ORDER BY /* TEMPLATE_BEGIN: order_by_clause */ id /* TEMPLATE_END */; + +-- name: JobGetAvailable :many +WITH locked_jobs AS ( + SELECT + * + FROM + /* TEMPLATE: schema */river_job + WHERE + state = 'available' + AND queue = @queue::text + AND scheduled_at <= coalesce(sqlc.narg('now')::timestamptz, now()) + ORDER BY + priority ASC, + scheduled_at ASC, + id ASC + LIMIT @max_to_lock::integer + FOR UPDATE + SKIP LOCKED +) +UPDATE + /* TEMPLATE: schema */river_job +SET + state = 'running', + attempt = river_job.attempt + 1, + attempted_at = coalesce(sqlc.narg('now')::timestamptz, now()), + attempted_by = array_append( + CASE WHEN array_length(river_job.attempted_by, 1) >= @max_attempted_by::int + -- +2 instead of +1 because Postgres array indexing starts at 1, not 0. + THEN river_job.attempted_by[array_length(river_job.attempted_by, 1) + 2 - @max_attempted_by:] + ELSE river_job.attempted_by + END, + @attempted_by::text + ) +FROM + locked_jobs +WHERE + river_job.id = locked_jobs.id +RETURNING + river_job.*; + +-- name: JobGetByID :one +SELECT * +FROM /* TEMPLATE: schema */river_job +WHERE id = @id +LIMIT 1; + +-- name: JobGetByIDMany :many +SELECT * +FROM /* TEMPLATE: schema */river_job +WHERE id = any(@id::bigint[]) +ORDER BY id; + +-- name: JobGetByKindMany :many +SELECT * +FROM /* TEMPLATE: schema */river_job +WHERE kind = any(@kind::text[]) +ORDER BY id; + +-- name: JobGetStuck :many +SELECT * +FROM /* TEMPLATE: schema */river_job +WHERE state = 'running' + AND id > @after_id::bigint + AND attempted_at < @stuck_horizon::timestamptz +ORDER BY id +LIMIT @max; + +-- name: JobInsertFastMany :many +WITH raw_job_data AS ( + SELECT + unnest(@id::bigint[]) AS id, + unnest(@args::jsonb[]) AS args, + unnest(@created_at::timestamptz[]) AS created_at, + unnest(@kind::text[]) AS kind, + unnest(@max_attempts::smallint[]) AS max_attempts, + unnest(@metadata::jsonb[]) AS metadata, + unnest(@priority::smallint[]) AS priority, + unnest(@queue::text[]) AS queue, + unnest(@scheduled_at::timestamptz[]) AS scheduled_at, + unnest(@state::text[]) AS state, + unnest(@tags::text[]) AS tags, + unnest(@unique_key::bytea[]) AS unique_key, + unnest(@unique_states::integer[]) AS unique_states +) +INSERT INTO /* TEMPLATE: schema */river_job( + id, + args, + created_at, + kind, + max_attempts, + metadata, + priority, + queue, + scheduled_at, + state, + tags, + unique_key, + unique_states +) SELECT + coalesce(nullif(id, 0), nextval('/* TEMPLATE: schema */river_job_id_seq'::regclass)), + args, + coalesce(nullif(created_at, '0001-01-01 00:00:00 +0000'), now()) AS created_at, + kind, + max_attempts, + coalesce(metadata, '{}'::jsonb) AS metadata, + priority, + queue, + coalesce(nullif(scheduled_at, '0001-01-01 00:00:00 +0000'), now()) AS scheduled_at, + state::/* TEMPLATE: schema */river_job_state, + string_to_array(tags, ',')::varchar(255)[], + -- `nullif` is required for `lib/pq`, which doesn't do a good job of reading + -- `nil` into `bytea`. We use `text` because otherwise `lib/pq` will encode + -- to Postgres binary like `\xAAAA`. + nullif(unique_key, '')::bytea, + nullif(unique_states::integer, 0)::bit(8) +FROM raw_job_data +ON CONFLICT (unique_key) + WHERE unique_key IS NOT NULL + AND unique_states IS NOT NULL + AND /* TEMPLATE: schema */river_job_state_in_bitmask(unique_states, state) + -- Something needs to be updated for a row to be returned on a conflict. + DO UPDATE SET kind = EXCLUDED.kind +RETURNING sqlc.embed(river_job), (xmax != 0) AS unique_skipped_as_duplicate; + +-- name: JobInsertFastManyNoReturning :execrows +INSERT INTO /* TEMPLATE: schema */river_job( + args, + created_at, + kind, + max_attempts, + metadata, + priority, + queue, + scheduled_at, + state, + tags, + unique_key, + unique_states +) SELECT + unnest(@args::jsonb[]), + unnest(@created_at::timestamptz[]), + unnest(@kind::text[]), + unnest(@max_attempts::smallint[]), + unnest(@metadata::jsonb[]), + unnest(@priority::smallint[]), + unnest(@queue::text[]), + unnest(@scheduled_at::timestamptz[]), + unnest(@state::/* TEMPLATE: schema */river_job_state[]), + + -- lib/pq really, REALLY does not play nicely with multi-dimensional arrays, + -- so instead we pack each set of tags into a string, send them through, + -- then unpack them here into an array to put in each row. This isn't + -- necessary in the Pgx driver where copyfrom is used instead. + string_to_array(unnest(@tags::text[]), ','), + + nullif(unnest(@unique_key::bytea[]), ''), + nullif(unnest(@unique_states::integer[]), 0)::bit(8) +ON CONFLICT (unique_key) + WHERE unique_key IS NOT NULL + AND unique_states IS NOT NULL + AND /* TEMPLATE: schema */river_job_state_in_bitmask(unique_states, state) +DO NOTHING; + +-- name: JobInsertFull :one +INSERT INTO /* TEMPLATE: schema */river_job( + args, + attempt, + attempted_at, + attempted_by, + created_at, + errors, + finalized_at, + kind, + max_attempts, + metadata, + priority, + queue, + scheduled_at, + state, + tags, + unique_key, + unique_states +) VALUES ( + @args::jsonb, + coalesce(@attempt::smallint, 0), + @attempted_at, + @attempted_by, + coalesce(sqlc.narg('created_at')::timestamptz, now()), + @errors, + @finalized_at, + @kind, + @max_attempts::smallint, + coalesce(@metadata::jsonb, '{}'), + @priority, + @queue, + coalesce(sqlc.narg('scheduled_at')::timestamptz, now()), + @state::/* TEMPLATE: schema */river_job_state, + coalesce(@tags::varchar(255)[], '{}'), + -- `nullif` is required for `lib/pq`, which doesn't do a good job of reading + -- `nil` into `bytea`. We use `text` because otherwise `lib/pq` will encode + -- to Postgres binary like `\xAAAA`. + nullif(@unique_key::text, '')::bytea, + nullif(@unique_states::integer, 0)::bit(8) +) RETURNING *; + +-- name: JobInsertFullMany :many +WITH raw_job_data AS ( + SELECT + unnest(@args::jsonb[]) AS args, + unnest(@attempt::smallint[]) AS attempt, + unnest(@attempted_at::timestamptz[]) AS attempted_at, + unnest(@created_at::timestamptz[]) AS created_at, + unnest(@finalized_at::timestamptz[]) AS finalized_at, + unnest(@kind::text[]) AS kind, + unnest(@max_attempts::smallint[]) AS max_attempts, + unnest(@metadata::jsonb[]) AS metadata, + unnest(@priority::smallint[]) AS priority, + unnest(@queue::text[]) AS queue, + unnest(@scheduled_at::timestamptz[]) AS scheduled_at, + unnest(@state::text[]) AS state, + unnest(@tags::text[]) AS tags, + unnest(@unique_key::text[]) AS unique_key, + unnest(@unique_states::integer[]) AS unique_states +) +INSERT INTO /* TEMPLATE: schema */river_job( + args, + attempt, + attempted_at, + created_at, + finalized_at, + kind, + max_attempts, + metadata, + priority, + queue, + scheduled_at, + state, + tags, + unique_key, + unique_states +) +SELECT + args, + coalesce(attempt, 0) AS attempt, + coalesce(nullif(attempted_at, '0001-01-01 00:00:00 +0000'), now()) AS attempted_at, + coalesce(nullif(created_at, '0001-01-01 00:00:00 +0000'), now()) AS created_at, + nullif(finalized_at, '0001-01-01 00:00:00 +0000') AS finalized_at, + kind, + max_attempts, + coalesce(metadata, '{}'::jsonb) AS metadata, + priority, + queue, + coalesce(nullif(scheduled_at, '0001-01-01 00:00:00 +0000'), now()) AS scheduled_at, + state::/* TEMPLATE: schema */river_job_state, + string_to_array(tags, ',')::varchar(255)[], + -- `nullif` is required for `lib/pq`, which doesn't do a good job of reading + -- `nil` into `bytea`. We use `text` because otherwise `lib/pq` will encode + -- to Postgres binary like `\xAAAA`. + nullif(unique_key, '')::bytea, + nullif(unique_states::integer, 0)::bit(8) +FROM raw_job_data +RETURNING *; + +-- name: JobKindList :many +SELECT DISTINCT ON (kind) kind +FROM /* TEMPLATE: schema */river_job +WHERE (@match = '' OR kind ILIKE '%' || @match || '%') + AND (@after = '' OR kind > @after) + AND (@exclude::text[] IS NULL OR kind != ALL(@exclude)) +ORDER BY kind ASC +LIMIT @max; + +-- name: JobList :many +SELECT * +FROM /* TEMPLATE: schema */river_job +WHERE /* TEMPLATE_BEGIN: where_clause */ true /* TEMPLATE_END */ +ORDER BY /* TEMPLATE_BEGIN: order_by_clause */ id /* TEMPLATE_END */ +LIMIT @max::int; + +-- Run by the rescuer to queue for retry or discard depending on job state. +-- name: JobRescueMany :exec +UPDATE /* TEMPLATE: schema */river_job +SET + errors = array_append(errors, updated_job.error), + finalized_at = updated_job.finalized_at, + scheduled_at = updated_job.scheduled_at, + metadata = river_job.metadata || jsonb_build_object( + 'river:rescue_count', + coalesce( + CASE + WHEN jsonb_typeof(river_job.metadata -> 'river:rescue_count') = 'number' + THEN (river_job.metadata ->> 'river:rescue_count')::int + END, + 0 + ) + 1 + ), + state = updated_job.state +FROM ( + SELECT + unnest(@id::bigint[]) AS id, + unnest(@error::jsonb[]) AS error, + nullif(unnest(@finalized_at::timestamptz[]), '0001-01-01 00:00:00 +0000') AS finalized_at, + unnest(@scheduled_at::timestamptz[]) AS scheduled_at, + unnest(@state::text[])::/* TEMPLATE: schema */river_job_state AS state +) AS updated_job +WHERE river_job.id = updated_job.id; + +-- name: JobRetry :one +WITH job_to_update AS ( + SELECT id + FROM /* TEMPLATE: schema */river_job + WHERE river_job.id = @id + FOR UPDATE +), +updated_job AS ( + UPDATE /* TEMPLATE: schema */river_job + SET + state = 'available', + max_attempts = CASE WHEN attempt = max_attempts THEN max_attempts + 1 ELSE max_attempts END, + finalized_at = NULL, + scheduled_at = coalesce(sqlc.narg('now')::timestamptz, now()) + FROM job_to_update + WHERE river_job.id = job_to_update.id + -- Do not touch running jobs: + AND river_job.state != 'running' + -- If the job is already available with a prior scheduled_at, leave it alone. + AND NOT ( + river_job.state = 'available' + AND river_job.scheduled_at < coalesce(sqlc.narg('now')::timestamptz, now()) + ) + RETURNING river_job.* +) +SELECT * +FROM /* TEMPLATE: schema */river_job +WHERE id = @id::bigint + AND id NOT IN (SELECT id FROM updated_job) +UNION +SELECT * +FROM updated_job; + +-- name: JobSchedule :many +WITH jobs_to_schedule AS ( + SELECT + id, + unique_key, + unique_states, + priority, + scheduled_at + FROM /* TEMPLATE: schema */river_job + WHERE + state IN ('retryable', 'scheduled') + AND priority >= 0 + AND queue IS NOT NULL + AND scheduled_at <= coalesce(sqlc.narg('now')::timestamptz, now()) + ORDER BY + priority, + scheduled_at, + id + LIMIT @max::bigint + FOR UPDATE +), +jobs_with_rownum AS ( + SELECT + *, + CASE + WHEN unique_key IS NOT NULL AND unique_states IS NOT NULL THEN + ROW_NUMBER() OVER ( + PARTITION BY unique_key + ORDER BY priority, scheduled_at, id + ) + ELSE NULL + END AS row_num + FROM jobs_to_schedule +), +unique_conflicts AS ( + SELECT river_job.unique_key + FROM /* TEMPLATE: schema */river_job + JOIN jobs_with_rownum + ON river_job.unique_key = jobs_with_rownum.unique_key + AND river_job.id != jobs_with_rownum.id + WHERE + river_job.unique_key IS NOT NULL + AND river_job.unique_states IS NOT NULL + AND /* TEMPLATE: schema */river_job_state_in_bitmask(river_job.unique_states, river_job.state) +), +job_updates AS ( + SELECT + job.id, + job.unique_key, + job.unique_states, + CASE + WHEN job.row_num IS NULL THEN 'available'::/* TEMPLATE: schema */river_job_state + WHEN uc.unique_key IS NOT NULL THEN 'discarded'::/* TEMPLATE: schema */river_job_state + WHEN job.row_num = 1 THEN 'available'::/* TEMPLATE: schema */river_job_state + ELSE 'discarded'::/* TEMPLATE: schema */river_job_state + END AS new_state, + (job.row_num IS NOT NULL AND (uc.unique_key IS NOT NULL OR job.row_num > 1)) AS finalized_at_do_update, + (job.row_num IS NOT NULL AND (uc.unique_key IS NOT NULL OR job.row_num > 1)) AS metadata_do_update + FROM jobs_with_rownum job + LEFT JOIN unique_conflicts uc ON job.unique_key = uc.unique_key +), +updated_jobs AS ( + UPDATE /* TEMPLATE: schema */river_job + SET + state = job_updates.new_state, + finalized_at = CASE WHEN job_updates.finalized_at_do_update THEN coalesce(sqlc.narg('now')::timestamptz, now()) + ELSE river_job.finalized_at END, + metadata = CASE WHEN job_updates.metadata_do_update THEN river_job.metadata || '{"unique_key_conflict": "scheduler_discarded"}'::jsonb + ELSE river_job.metadata END + FROM job_updates + WHERE river_job.id = job_updates.id + RETURNING + river_job.id, + job_updates.new_state = 'discarded'::/* TEMPLATE: schema */river_job_state AS conflict_discarded +) +SELECT + sqlc.embed(river_job), + updated_jobs.conflict_discarded +FROM /* TEMPLATE: schema */river_job +JOIN updated_jobs ON river_job.id = updated_jobs.id; + +-- name: JobSetStateIfRunningMany :many +WITH job_input AS ( + SELECT + unnest(@ids::bigint[]) AS id, + unnest(@attempt_do_update::boolean[]) AS attempt_do_update, + unnest(@attempt::int[]) AS attempt, + unnest(@errors_do_update::boolean[]) AS errors_do_update, + unnest(@errors::jsonb[]) AS errors, + unnest(@finalized_at_do_update::boolean[]) AS finalized_at_do_update, + unnest(@finalized_at::timestamptz[]) AS finalized_at, + unnest(@metadata_do_merge::boolean[]) AS metadata_do_merge, + unnest(@metadata_updates::jsonb[]) AS metadata_updates, + unnest(@scheduled_at_do_update::boolean[]) AS scheduled_at_do_update, + unnest(@scheduled_at::timestamptz[]) AS scheduled_at, + -- To avoid requiring pgx users to register the OID of the river_job_state[] + -- type, we cast the array to text[] and then to river_job_state. + unnest(@state::text[])::/* TEMPLATE: schema */river_job_state AS state +), +updated AS ( + UPDATE /* TEMPLATE: schema */river_job + SET + attempt = CASE + WHEN river_job.state = 'running' + AND NOT (job_input.state IN ('retryable','scheduled') AND river_job.metadata ? 'cancel_attempted_at') + AND job_input.attempt_do_update + THEN job_input.attempt + ELSE river_job.attempt + END, + errors = CASE + WHEN river_job.state = 'running' + AND job_input.errors_do_update + THEN array_append(river_job.errors, job_input.errors) + ELSE river_job.errors + END, + finalized_at = CASE + WHEN river_job.state = 'running' + AND (job_input.state IN ('retryable','scheduled') AND river_job.metadata ? 'cancel_attempted_at') + THEN coalesce(sqlc.narg('now')::timestamptz, now()) + WHEN river_job.state = 'running' + AND job_input.finalized_at_do_update + THEN job_input.finalized_at + ELSE river_job.finalized_at + END, + metadata = CASE + WHEN job_input.metadata_do_merge + THEN river_job.metadata || job_input.metadata_updates + ELSE river_job.metadata + END, + scheduled_at = CASE + WHEN river_job.state = 'running' + AND NOT (job_input.state IN ('retryable','scheduled') AND river_job.metadata ? 'cancel_attempted_at') + AND job_input.scheduled_at_do_update + THEN job_input.scheduled_at + ELSE river_job.scheduled_at + END, + state = CASE + WHEN river_job.state = 'running' + AND (job_input.state IN ('retryable','scheduled') AND river_job.metadata ? 'cancel_attempted_at') + THEN 'cancelled'::/* TEMPLATE: schema */river_job_state + WHEN river_job.state = 'running' + THEN job_input.state + ELSE river_job.state + END + FROM job_input + WHERE river_job.id = job_input.id + AND (river_job.state = 'running' OR job_input.metadata_do_merge) + RETURNING river_job.* +) +SELECT river_job.* +FROM /* TEMPLATE: schema */river_job +JOIN job_input ON river_job.id = job_input.id +WHERE NOT EXISTS ( + SELECT 1 + FROM updated + WHERE updated.id = river_job.id +) +UNION ALL +SELECT * +FROM updated +ORDER BY id; + +-- name: JobUpdate :one +WITH locked_job AS ( + SELECT id + FROM /* TEMPLATE: schema */river_job + WHERE river_job.id = @id + FOR UPDATE +) +UPDATE /* TEMPLATE: schema */river_job +SET + metadata = CASE WHEN @metadata_do_merge::boolean THEN metadata || @metadata::jsonb ELSE metadata END +FROM + locked_job +WHERE river_job.id = locked_job.id +RETURNING river_job.*; + +-- A generalized update for any property on a job. This brings in a large number +-- of parameters and therefore may be more suitable for testing than production. +-- name: JobUpdateFull :one +UPDATE /* TEMPLATE: schema */river_job +SET + attempt = CASE WHEN @attempt_do_update::boolean THEN @attempt ELSE attempt END, + attempted_at = CASE WHEN @attempted_at_do_update::boolean THEN @attempted_at ELSE attempted_at END, + attempted_by = CASE WHEN @attempted_by_do_update::boolean THEN @attempted_by ELSE attempted_by END, + errors = CASE WHEN @errors_do_update::boolean THEN @errors::jsonb[] ELSE errors END, + finalized_at = CASE WHEN @finalized_at_do_update::boolean THEN @finalized_at ELSE finalized_at END, + max_attempts = CASE WHEN @max_attempts_do_update::boolean THEN @max_attempts ELSE max_attempts END, + metadata = CASE WHEN @metadata_do_update::boolean THEN @metadata::jsonb ELSE metadata END, + state = CASE WHEN @state_do_update::boolean THEN @state::/* TEMPLATE: schema */river_job_state ELSE state END +WHERE id = @id +RETURNING *; diff --git a/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql.go b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql.go new file mode 100644 index 0000000000..a361baacc2 --- /dev/null +++ b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql.go @@ -0,0 +1,1705 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.31.0 +// source: river_job.sql + +package dbsqlc + +import ( + "context" + "time" + + "github.com/jackc/pgx/v5/pgconn" + "github.com/jackc/pgx/v5/pgtype" +) + +const jobCancel = `-- name: JobCancel :one +WITH locked_job AS ( + SELECT + id, queue, state, finalized_at + FROM /* TEMPLATE: schema */river_job + WHERE river_job.id = $1 + FOR UPDATE +), +notification AS ( + SELECT + id, + pg_notify( + concat(coalesce($2::text, current_schema()), '.', $3::text), + json_build_object('action', 'cancel', 'job_id', id, 'queue', queue)::text + ) + FROM + locked_job + WHERE + state NOT IN ('cancelled', 'completed', 'discarded') + AND finalized_at IS NULL +), +updated_job AS ( + UPDATE /* TEMPLATE: schema */river_job + SET + -- If the job is actively running, we want to let its current client and + -- producer handle the cancellation. Otherwise, immediately cancel it. + state = CASE WHEN state = 'running' THEN state ELSE 'cancelled' END, + finalized_at = CASE WHEN state = 'running' THEN finalized_at ELSE coalesce($4::timestamptz, now()) END, + -- Mark the job as cancelled by query so that the rescuer knows not to + -- rescue it, even if it gets stuck in the running state: + metadata = jsonb_set(metadata, '{cancel_attempted_at}'::text[], $5::jsonb, true) + FROM notification + WHERE river_job.id = notification.id + RETURNING river_job.id, river_job.args, river_job.attempt, river_job.attempted_at, river_job.attempted_by, river_job.created_at, river_job.errors, river_job.finalized_at, river_job.kind, river_job.max_attempts, river_job.metadata, river_job.priority, river_job.queue, river_job.state, river_job.scheduled_at, river_job.tags, river_job.unique_key, river_job.unique_states +) +SELECT id, args, attempt, attempted_at, attempted_by, created_at, errors, finalized_at, kind, max_attempts, metadata, priority, queue, state, scheduled_at, tags, unique_key, unique_states +FROM /* TEMPLATE: schema */river_job +WHERE id = $1::bigint + AND id NOT IN (SELECT id FROM updated_job) +UNION +SELECT id, args, attempt, attempted_at, attempted_by, created_at, errors, finalized_at, kind, max_attempts, metadata, priority, queue, state, scheduled_at, tags, unique_key, unique_states +FROM updated_job +` + +type JobCancelParams struct { + ID int64 + Schema pgtype.Text + ControlTopic string + Now *time.Time + CancelAttemptedAt []byte +} + +func (q *Queries) JobCancel(ctx context.Context, db DBTX, arg *JobCancelParams) (*RiverJob, error) { + row := db.QueryRow(ctx, jobCancel, + arg.ID, + arg.Schema, + arg.ControlTopic, + arg.Now, + arg.CancelAttemptedAt, + ) + var i RiverJob + err := row.Scan( + &i.ID, + &i.Args, + &i.Attempt, + &i.AttemptedAt, + &i.AttemptedBy, + &i.CreatedAt, + &i.Errors, + &i.FinalizedAt, + &i.Kind, + &i.MaxAttempts, + &i.Metadata, + &i.Priority, + &i.Queue, + &i.State, + &i.ScheduledAt, + &i.Tags, + &i.UniqueKey, + &i.UniqueStates, + ) + return &i, err +} + +const jobCountByAllStates = `-- name: JobCountByAllStates :many +SELECT state, count(*) +FROM /* TEMPLATE: schema */ river_job +GROUP BY state +` + +type JobCountByAllStatesRow struct { + State RiverJobState + Count int64 +} + +func (q *Queries) JobCountByAllStates(ctx context.Context, db DBTX) ([]*JobCountByAllStatesRow, error) { + rows, err := db.Query(ctx, jobCountByAllStates) + if err != nil { + return nil, err + } + defer rows.Close() + var items []*JobCountByAllStatesRow + for rows.Next() { + var i JobCountByAllStatesRow + if err := rows.Scan(&i.State, &i.Count); err != nil { + return nil, err + } + items = append(items, &i) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const jobCountByQueueAndState = `-- name: JobCountByQueueAndState :many +WITH all_queues AS ( + SELECT DISTINCT unnest($1::text[])::text AS queue +), + +running_job_counts AS ( + SELECT + queue, + COUNT(*) AS count + FROM /* TEMPLATE: schema */river_job + WHERE queue = ANY($1::text[]) + AND state = 'running' + GROUP BY queue +), + +available_job_counts AS ( + SELECT + queue, + COUNT(*) AS count + FROM + /* TEMPLATE: schema */river_job + WHERE queue = ANY($1::text[]) + AND state = 'available' + GROUP BY queue +) + +SELECT + all_queues.queue, + COALESCE(available_job_counts.count, 0) AS count_available, + COALESCE(running_job_counts.count, 0) AS count_running +FROM + all_queues +LEFT JOIN + running_job_counts ON all_queues.queue = running_job_counts.queue +LEFT JOIN + available_job_counts ON all_queues.queue = available_job_counts.queue +ORDER BY all_queues.queue ASC +` + +type JobCountByQueueAndStateRow struct { + Queue string + CountAvailable int64 + CountRunning int64 +} + +func (q *Queries) JobCountByQueueAndState(ctx context.Context, db DBTX, queueNames []string) ([]*JobCountByQueueAndStateRow, error) { + rows, err := db.Query(ctx, jobCountByQueueAndState, queueNames) + if err != nil { + return nil, err + } + defer rows.Close() + var items []*JobCountByQueueAndStateRow + for rows.Next() { + var i JobCountByQueueAndStateRow + if err := rows.Scan(&i.Queue, &i.CountAvailable, &i.CountRunning); err != nil { + return nil, err + } + items = append(items, &i) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const jobCountByState = `-- name: JobCountByState :one +SELECT count(*) +FROM /* TEMPLATE: schema */river_job +WHERE state = $1 +` + +func (q *Queries) JobCountByState(ctx context.Context, db DBTX, state RiverJobState) (int64, error) { + row := db.QueryRow(ctx, jobCountByState, state) + var count int64 + err := row.Scan(&count) + return count, err +} + +const jobDelete = `-- name: JobDelete :one +WITH job_to_delete AS ( + SELECT id + FROM /* TEMPLATE: schema */river_job + WHERE river_job.id = $1 + FOR UPDATE +), +deleted_job AS ( + DELETE + FROM /* TEMPLATE: schema */river_job + USING job_to_delete + WHERE river_job.id = job_to_delete.id + -- Do not touch running jobs: + AND river_job.state != 'running' + RETURNING river_job.id, river_job.args, river_job.attempt, river_job.attempted_at, river_job.attempted_by, river_job.created_at, river_job.errors, river_job.finalized_at, river_job.kind, river_job.max_attempts, river_job.metadata, river_job.priority, river_job.queue, river_job.state, river_job.scheduled_at, river_job.tags, river_job.unique_key, river_job.unique_states +) +SELECT id, args, attempt, attempted_at, attempted_by, created_at, errors, finalized_at, kind, max_attempts, metadata, priority, queue, state, scheduled_at, tags, unique_key, unique_states +FROM /* TEMPLATE: schema */river_job +WHERE id = $1::bigint + AND id NOT IN (SELECT id FROM deleted_job) +UNION +SELECT id, args, attempt, attempted_at, attempted_by, created_at, errors, finalized_at, kind, max_attempts, metadata, priority, queue, state, scheduled_at, tags, unique_key, unique_states +FROM deleted_job +` + +func (q *Queries) JobDelete(ctx context.Context, db DBTX, id int64) (*RiverJob, error) { + row := db.QueryRow(ctx, jobDelete, id) + var i RiverJob + err := row.Scan( + &i.ID, + &i.Args, + &i.Attempt, + &i.AttemptedAt, + &i.AttemptedBy, + &i.CreatedAt, + &i.Errors, + &i.FinalizedAt, + &i.Kind, + &i.MaxAttempts, + &i.Metadata, + &i.Priority, + &i.Queue, + &i.State, + &i.ScheduledAt, + &i.Tags, + &i.UniqueKey, + &i.UniqueStates, + ) + return &i, err +} + +const jobDeleteBefore = `-- name: JobDeleteBefore :execresult +DELETE FROM /* TEMPLATE: schema */river_job +WHERE id IN ( + SELECT id + FROM /* TEMPLATE: schema */river_job + WHERE ( + (state = 'cancelled' AND $1 AND finalized_at < $2::timestamptz) OR + (state = 'completed' AND $3 AND finalized_at < $4::timestamptz) OR + (state = 'discarded' AND $5 AND finalized_at < $6::timestamptz) + ) + AND ( + $7::text[] IS NULL + OR NOT (queue = any($7)) + ) + AND ( + $8::text[] IS NULL + OR queue = any($8) + ) + ORDER BY id + LIMIT $9::bigint +) +` + +type JobDeleteBeforeParams struct { + CancelledDoDelete interface{} + CancelledFinalizedAtHorizon time.Time + CompletedDoDelete interface{} + CompletedFinalizedAtHorizon time.Time + DiscardedDoDelete interface{} + DiscardedFinalizedAtHorizon time.Time + QueuesExcluded []string + QueuesIncluded []string + Max int64 +} + +func (q *Queries) JobDeleteBefore(ctx context.Context, db DBTX, arg *JobDeleteBeforeParams) (pgconn.CommandTag, error) { + return db.Exec(ctx, jobDeleteBefore, + arg.CancelledDoDelete, + arg.CancelledFinalizedAtHorizon, + arg.CompletedDoDelete, + arg.CompletedFinalizedAtHorizon, + arg.DiscardedDoDelete, + arg.DiscardedFinalizedAtHorizon, + arg.QueuesExcluded, + arg.QueuesIncluded, + arg.Max, + ) +} + +const jobDeleteMany = `-- name: JobDeleteMany :many +WITH jobs_to_delete AS ( + SELECT id, args, attempt, attempted_at, attempted_by, created_at, errors, finalized_at, kind, max_attempts, metadata, priority, queue, state, scheduled_at, tags, unique_key, unique_states + FROM /* TEMPLATE: schema */river_job + WHERE /* TEMPLATE_BEGIN: where_clause */ true /* TEMPLATE_END */ + AND state != 'running' + ORDER BY /* TEMPLATE_BEGIN: order_by_clause */ id /* TEMPLATE_END */ + LIMIT $1::int + FOR UPDATE + SKIP LOCKED +), +deleted_jobs AS ( + DELETE FROM /* TEMPLATE: schema */river_job + WHERE id IN (SELECT id FROM jobs_to_delete) + RETURNING id, args, attempt, attempted_at, attempted_by, created_at, errors, finalized_at, kind, max_attempts, metadata, priority, queue, state, scheduled_at, tags, unique_key, unique_states +) +SELECT id, args, attempt, attempted_at, attempted_by, created_at, errors, finalized_at, kind, max_attempts, metadata, priority, queue, state, scheduled_at, tags, unique_key, unique_states +FROM /* TEMPLATE: schema */river_job +WHERE id IN (SELECT id FROM deleted_jobs) +ORDER BY /* TEMPLATE_BEGIN: order_by_clause */ id /* TEMPLATE_END */ +` + +// this last SELECT step is necessary because there's no other way to define +// order records come back from a DELETE statement +func (q *Queries) JobDeleteMany(ctx context.Context, db DBTX, max int32) ([]*RiverJob, error) { + rows, err := db.Query(ctx, jobDeleteMany, max) + if err != nil { + return nil, err + } + defer rows.Close() + var items []*RiverJob + for rows.Next() { + var i RiverJob + if err := rows.Scan( + &i.ID, + &i.Args, + &i.Attempt, + &i.AttemptedAt, + &i.AttemptedBy, + &i.CreatedAt, + &i.Errors, + &i.FinalizedAt, + &i.Kind, + &i.MaxAttempts, + &i.Metadata, + &i.Priority, + &i.Queue, + &i.State, + &i.ScheduledAt, + &i.Tags, + &i.UniqueKey, + &i.UniqueStates, + ); err != nil { + return nil, err + } + items = append(items, &i) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const jobGetAvailable = `-- name: JobGetAvailable :many +WITH locked_jobs AS ( + SELECT + id, args, attempt, attempted_at, attempted_by, created_at, errors, finalized_at, kind, max_attempts, metadata, priority, queue, state, scheduled_at, tags, unique_key, unique_states + FROM + /* TEMPLATE: schema */river_job + WHERE + state = 'available' + AND queue = $4::text + AND scheduled_at <= coalesce($1::timestamptz, now()) + ORDER BY + priority ASC, + scheduled_at ASC, + id ASC + LIMIT $5::integer + FOR UPDATE + SKIP LOCKED +) +UPDATE + /* TEMPLATE: schema */river_job +SET + state = 'running', + attempt = river_job.attempt + 1, + attempted_at = coalesce($1::timestamptz, now()), + attempted_by = array_append( + CASE WHEN array_length(river_job.attempted_by, 1) >= $2::int + -- +2 instead of +1 because Postgres array indexing starts at 1, not 0. + THEN river_job.attempted_by[array_length(river_job.attempted_by, 1) + 2 - $2:] + ELSE river_job.attempted_by + END, + $3::text + ) +FROM + locked_jobs +WHERE + river_job.id = locked_jobs.id +RETURNING + river_job.id, river_job.args, river_job.attempt, river_job.attempted_at, river_job.attempted_by, river_job.created_at, river_job.errors, river_job.finalized_at, river_job.kind, river_job.max_attempts, river_job.metadata, river_job.priority, river_job.queue, river_job.state, river_job.scheduled_at, river_job.tags, river_job.unique_key, river_job.unique_states +` + +type JobGetAvailableParams struct { + Now *time.Time + MaxAttemptedBy int32 + AttemptedBy string + Queue string + MaxToLock int32 +} + +func (q *Queries) JobGetAvailable(ctx context.Context, db DBTX, arg *JobGetAvailableParams) ([]*RiverJob, error) { + rows, err := db.Query(ctx, jobGetAvailable, + arg.Now, + arg.MaxAttemptedBy, + arg.AttemptedBy, + arg.Queue, + arg.MaxToLock, + ) + if err != nil { + return nil, err + } + defer rows.Close() + var items []*RiverJob + for rows.Next() { + var i RiverJob + if err := rows.Scan( + &i.ID, + &i.Args, + &i.Attempt, + &i.AttemptedAt, + &i.AttemptedBy, + &i.CreatedAt, + &i.Errors, + &i.FinalizedAt, + &i.Kind, + &i.MaxAttempts, + &i.Metadata, + &i.Priority, + &i.Queue, + &i.State, + &i.ScheduledAt, + &i.Tags, + &i.UniqueKey, + &i.UniqueStates, + ); err != nil { + return nil, err + } + items = append(items, &i) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const jobGetByID = `-- name: JobGetByID :one +SELECT id, args, attempt, attempted_at, attempted_by, created_at, errors, finalized_at, kind, max_attempts, metadata, priority, queue, state, scheduled_at, tags, unique_key, unique_states +FROM /* TEMPLATE: schema */river_job +WHERE id = $1 +LIMIT 1 +` + +func (q *Queries) JobGetByID(ctx context.Context, db DBTX, id int64) (*RiverJob, error) { + row := db.QueryRow(ctx, jobGetByID, id) + var i RiverJob + err := row.Scan( + &i.ID, + &i.Args, + &i.Attempt, + &i.AttemptedAt, + &i.AttemptedBy, + &i.CreatedAt, + &i.Errors, + &i.FinalizedAt, + &i.Kind, + &i.MaxAttempts, + &i.Metadata, + &i.Priority, + &i.Queue, + &i.State, + &i.ScheduledAt, + &i.Tags, + &i.UniqueKey, + &i.UniqueStates, + ) + return &i, err +} + +const jobGetByIDMany = `-- name: JobGetByIDMany :many +SELECT id, args, attempt, attempted_at, attempted_by, created_at, errors, finalized_at, kind, max_attempts, metadata, priority, queue, state, scheduled_at, tags, unique_key, unique_states +FROM /* TEMPLATE: schema */river_job +WHERE id = any($1::bigint[]) +ORDER BY id +` + +func (q *Queries) JobGetByIDMany(ctx context.Context, db DBTX, id []int64) ([]*RiverJob, error) { + rows, err := db.Query(ctx, jobGetByIDMany, id) + if err != nil { + return nil, err + } + defer rows.Close() + var items []*RiverJob + for rows.Next() { + var i RiverJob + if err := rows.Scan( + &i.ID, + &i.Args, + &i.Attempt, + &i.AttemptedAt, + &i.AttemptedBy, + &i.CreatedAt, + &i.Errors, + &i.FinalizedAt, + &i.Kind, + &i.MaxAttempts, + &i.Metadata, + &i.Priority, + &i.Queue, + &i.State, + &i.ScheduledAt, + &i.Tags, + &i.UniqueKey, + &i.UniqueStates, + ); err != nil { + return nil, err + } + items = append(items, &i) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const jobGetByKindMany = `-- name: JobGetByKindMany :many +SELECT id, args, attempt, attempted_at, attempted_by, created_at, errors, finalized_at, kind, max_attempts, metadata, priority, queue, state, scheduled_at, tags, unique_key, unique_states +FROM /* TEMPLATE: schema */river_job +WHERE kind = any($1::text[]) +ORDER BY id +` + +func (q *Queries) JobGetByKindMany(ctx context.Context, db DBTX, kind []string) ([]*RiverJob, error) { + rows, err := db.Query(ctx, jobGetByKindMany, kind) + if err != nil { + return nil, err + } + defer rows.Close() + var items []*RiverJob + for rows.Next() { + var i RiverJob + if err := rows.Scan( + &i.ID, + &i.Args, + &i.Attempt, + &i.AttemptedAt, + &i.AttemptedBy, + &i.CreatedAt, + &i.Errors, + &i.FinalizedAt, + &i.Kind, + &i.MaxAttempts, + &i.Metadata, + &i.Priority, + &i.Queue, + &i.State, + &i.ScheduledAt, + &i.Tags, + &i.UniqueKey, + &i.UniqueStates, + ); err != nil { + return nil, err + } + items = append(items, &i) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const jobGetStuck = `-- name: JobGetStuck :many +SELECT id, args, attempt, attempted_at, attempted_by, created_at, errors, finalized_at, kind, max_attempts, metadata, priority, queue, state, scheduled_at, tags, unique_key, unique_states +FROM /* TEMPLATE: schema */river_job +WHERE state = 'running' + AND id > $1::bigint + AND attempted_at < $2::timestamptz +ORDER BY id +LIMIT $3 +` + +type JobGetStuckParams struct { + AfterID int64 + StuckHorizon time.Time + Max int32 +} + +func (q *Queries) JobGetStuck(ctx context.Context, db DBTX, arg *JobGetStuckParams) ([]*RiverJob, error) { + rows, err := db.Query(ctx, jobGetStuck, arg.AfterID, arg.StuckHorizon, arg.Max) + if err != nil { + return nil, err + } + defer rows.Close() + var items []*RiverJob + for rows.Next() { + var i RiverJob + if err := rows.Scan( + &i.ID, + &i.Args, + &i.Attempt, + &i.AttemptedAt, + &i.AttemptedBy, + &i.CreatedAt, + &i.Errors, + &i.FinalizedAt, + &i.Kind, + &i.MaxAttempts, + &i.Metadata, + &i.Priority, + &i.Queue, + &i.State, + &i.ScheduledAt, + &i.Tags, + &i.UniqueKey, + &i.UniqueStates, + ); err != nil { + return nil, err + } + items = append(items, &i) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const jobInsertFastMany = `-- name: JobInsertFastMany :many +WITH raw_job_data AS ( + SELECT + unnest($1::bigint[]) AS id, + unnest($2::jsonb[]) AS args, + unnest($3::timestamptz[]) AS created_at, + unnest($4::text[]) AS kind, + unnest($5::smallint[]) AS max_attempts, + unnest($6::jsonb[]) AS metadata, + unnest($7::smallint[]) AS priority, + unnest($8::text[]) AS queue, + unnest($9::timestamptz[]) AS scheduled_at, + unnest($10::text[]) AS state, + unnest($11::text[]) AS tags, + unnest($12::bytea[]) AS unique_key, + unnest($13::integer[]) AS unique_states +) +INSERT INTO /* TEMPLATE: schema */river_job( + id, + args, + created_at, + kind, + max_attempts, + metadata, + priority, + queue, + scheduled_at, + state, + tags, + unique_key, + unique_states +) SELECT + coalesce(nullif(id, 0), nextval('/* TEMPLATE: schema */river_job_id_seq'::regclass)), + args, + coalesce(nullif(created_at, '0001-01-01 00:00:00 +0000'), now()) AS created_at, + kind, + max_attempts, + coalesce(metadata, '{}'::jsonb) AS metadata, + priority, + queue, + coalesce(nullif(scheduled_at, '0001-01-01 00:00:00 +0000'), now()) AS scheduled_at, + state::/* TEMPLATE: schema */river_job_state, + string_to_array(tags, ',')::varchar(255)[], + -- ` + "`" + `nullif` + "`" + ` is required for ` + "`" + `lib/pq` + "`" + `, which doesn't do a good job of reading + -- ` + "`" + `nil` + "`" + ` into ` + "`" + `bytea` + "`" + `. We use ` + "`" + `text` + "`" + ` because otherwise ` + "`" + `lib/pq` + "`" + ` will encode + -- to Postgres binary like ` + "`" + `\xAAAA` + "`" + `. + nullif(unique_key, '')::bytea, + nullif(unique_states::integer, 0)::bit(8) +FROM raw_job_data +ON CONFLICT (unique_key) + WHERE unique_key IS NOT NULL + AND unique_states IS NOT NULL + AND /* TEMPLATE: schema */river_job_state_in_bitmask(unique_states, state) + -- Something needs to be updated for a row to be returned on a conflict. + DO UPDATE SET kind = EXCLUDED.kind +RETURNING river_job.id, river_job.args, river_job.attempt, river_job.attempted_at, river_job.attempted_by, river_job.created_at, river_job.errors, river_job.finalized_at, river_job.kind, river_job.max_attempts, river_job.metadata, river_job.priority, river_job.queue, river_job.state, river_job.scheduled_at, river_job.tags, river_job.unique_key, river_job.unique_states, (xmax != 0) AS unique_skipped_as_duplicate +` + +type JobInsertFastManyParams struct { + ID []int64 + Args [][]byte + CreatedAt []time.Time + Kind []string + MaxAttempts []int16 + Metadata [][]byte + Priority []int16 + Queue []string + ScheduledAt []time.Time + State []string + Tags []string + UniqueKey [][]byte + UniqueStates []int32 +} + +type JobInsertFastManyRow struct { + RiverJob RiverJob + UniqueSkippedAsDuplicate bool +} + +func (q *Queries) JobInsertFastMany(ctx context.Context, db DBTX, arg *JobInsertFastManyParams) ([]*JobInsertFastManyRow, error) { + rows, err := db.Query(ctx, jobInsertFastMany, + arg.ID, + arg.Args, + arg.CreatedAt, + arg.Kind, + arg.MaxAttempts, + arg.Metadata, + arg.Priority, + arg.Queue, + arg.ScheduledAt, + arg.State, + arg.Tags, + arg.UniqueKey, + arg.UniqueStates, + ) + if err != nil { + return nil, err + } + defer rows.Close() + var items []*JobInsertFastManyRow + for rows.Next() { + var i JobInsertFastManyRow + if err := rows.Scan( + &i.RiverJob.ID, + &i.RiverJob.Args, + &i.RiverJob.Attempt, + &i.RiverJob.AttemptedAt, + &i.RiverJob.AttemptedBy, + &i.RiverJob.CreatedAt, + &i.RiverJob.Errors, + &i.RiverJob.FinalizedAt, + &i.RiverJob.Kind, + &i.RiverJob.MaxAttempts, + &i.RiverJob.Metadata, + &i.RiverJob.Priority, + &i.RiverJob.Queue, + &i.RiverJob.State, + &i.RiverJob.ScheduledAt, + &i.RiverJob.Tags, + &i.RiverJob.UniqueKey, + &i.RiverJob.UniqueStates, + &i.UniqueSkippedAsDuplicate, + ); err != nil { + return nil, err + } + items = append(items, &i) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const jobInsertFastManyNoReturning = `-- name: JobInsertFastManyNoReturning :execrows +INSERT INTO /* TEMPLATE: schema */river_job( + args, + created_at, + kind, + max_attempts, + metadata, + priority, + queue, + scheduled_at, + state, + tags, + unique_key, + unique_states +) SELECT + unnest($1::jsonb[]), + unnest($2::timestamptz[]), + unnest($3::text[]), + unnest($4::smallint[]), + unnest($5::jsonb[]), + unnest($6::smallint[]), + unnest($7::text[]), + unnest($8::timestamptz[]), + unnest($9::/* TEMPLATE: schema */river_job_state[]), + + -- lib/pq really, REALLY does not play nicely with multi-dimensional arrays, + -- so instead we pack each set of tags into a string, send them through, + -- then unpack them here into an array to put in each row. This isn't + -- necessary in the Pgx driver where copyfrom is used instead. + string_to_array(unnest($10::text[]), ','), + + nullif(unnest($11::bytea[]), ''), + nullif(unnest($12::integer[]), 0)::bit(8) +ON CONFLICT (unique_key) + WHERE unique_key IS NOT NULL + AND unique_states IS NOT NULL + AND /* TEMPLATE: schema */river_job_state_in_bitmask(unique_states, state) +DO NOTHING +` + +type JobInsertFastManyNoReturningParams struct { + Args [][]byte + CreatedAt []time.Time + Kind []string + MaxAttempts []int16 + Metadata [][]byte + Priority []int16 + Queue []string + ScheduledAt []time.Time + State []RiverJobState + Tags []string + UniqueKey [][]byte + UniqueStates []int32 +} + +func (q *Queries) JobInsertFastManyNoReturning(ctx context.Context, db DBTX, arg *JobInsertFastManyNoReturningParams) (int64, error) { + result, err := db.Exec(ctx, jobInsertFastManyNoReturning, + arg.Args, + arg.CreatedAt, + arg.Kind, + arg.MaxAttempts, + arg.Metadata, + arg.Priority, + arg.Queue, + arg.ScheduledAt, + arg.State, + arg.Tags, + arg.UniqueKey, + arg.UniqueStates, + ) + if err != nil { + return 0, err + } + return result.RowsAffected(), nil +} + +const jobInsertFull = `-- name: JobInsertFull :one +INSERT INTO /* TEMPLATE: schema */river_job( + args, + attempt, + attempted_at, + attempted_by, + created_at, + errors, + finalized_at, + kind, + max_attempts, + metadata, + priority, + queue, + scheduled_at, + state, + tags, + unique_key, + unique_states +) VALUES ( + $1::jsonb, + coalesce($2::smallint, 0), + $3, + $4, + coalesce($5::timestamptz, now()), + $6, + $7, + $8, + $9::smallint, + coalesce($10::jsonb, '{}'), + $11, + $12, + coalesce($13::timestamptz, now()), + $14::/* TEMPLATE: schema */river_job_state, + coalesce($15::varchar(255)[], '{}'), + -- ` + "`" + `nullif` + "`" + ` is required for ` + "`" + `lib/pq` + "`" + `, which doesn't do a good job of reading + -- ` + "`" + `nil` + "`" + ` into ` + "`" + `bytea` + "`" + `. We use ` + "`" + `text` + "`" + ` because otherwise ` + "`" + `lib/pq` + "`" + ` will encode + -- to Postgres binary like ` + "`" + `\xAAAA` + "`" + `. + nullif($16::text, '')::bytea, + nullif($17::integer, 0)::bit(8) +) RETURNING id, args, attempt, attempted_at, attempted_by, created_at, errors, finalized_at, kind, max_attempts, metadata, priority, queue, state, scheduled_at, tags, unique_key, unique_states +` + +type JobInsertFullParams struct { + Args []byte + Attempt int16 + AttemptedAt *time.Time + AttemptedBy []string + CreatedAt *time.Time + Errors [][]byte + FinalizedAt *time.Time + Kind string + MaxAttempts int16 + Metadata []byte + Priority int16 + Queue string + ScheduledAt *time.Time + State RiverJobState + Tags []string + UniqueKey string + UniqueStates int32 +} + +func (q *Queries) JobInsertFull(ctx context.Context, db DBTX, arg *JobInsertFullParams) (*RiverJob, error) { + row := db.QueryRow(ctx, jobInsertFull, + arg.Args, + arg.Attempt, + arg.AttemptedAt, + arg.AttemptedBy, + arg.CreatedAt, + arg.Errors, + arg.FinalizedAt, + arg.Kind, + arg.MaxAttempts, + arg.Metadata, + arg.Priority, + arg.Queue, + arg.ScheduledAt, + arg.State, + arg.Tags, + arg.UniqueKey, + arg.UniqueStates, + ) + var i RiverJob + err := row.Scan( + &i.ID, + &i.Args, + &i.Attempt, + &i.AttemptedAt, + &i.AttemptedBy, + &i.CreatedAt, + &i.Errors, + &i.FinalizedAt, + &i.Kind, + &i.MaxAttempts, + &i.Metadata, + &i.Priority, + &i.Queue, + &i.State, + &i.ScheduledAt, + &i.Tags, + &i.UniqueKey, + &i.UniqueStates, + ) + return &i, err +} + +const jobInsertFullMany = `-- name: JobInsertFullMany :many +WITH raw_job_data AS ( + SELECT + unnest($1::jsonb[]) AS args, + unnest($2::smallint[]) AS attempt, + unnest($3::timestamptz[]) AS attempted_at, + unnest($4::timestamptz[]) AS created_at, + unnest($5::timestamptz[]) AS finalized_at, + unnest($6::text[]) AS kind, + unnest($7::smallint[]) AS max_attempts, + unnest($8::jsonb[]) AS metadata, + unnest($9::smallint[]) AS priority, + unnest($10::text[]) AS queue, + unnest($11::timestamptz[]) AS scheduled_at, + unnest($12::text[]) AS state, + unnest($13::text[]) AS tags, + unnest($14::text[]) AS unique_key, + unnest($15::integer[]) AS unique_states +) +INSERT INTO /* TEMPLATE: schema */river_job( + args, + attempt, + attempted_at, + created_at, + finalized_at, + kind, + max_attempts, + metadata, + priority, + queue, + scheduled_at, + state, + tags, + unique_key, + unique_states +) +SELECT + args, + coalesce(attempt, 0) AS attempt, + coalesce(nullif(attempted_at, '0001-01-01 00:00:00 +0000'), now()) AS attempted_at, + coalesce(nullif(created_at, '0001-01-01 00:00:00 +0000'), now()) AS created_at, + nullif(finalized_at, '0001-01-01 00:00:00 +0000') AS finalized_at, + kind, + max_attempts, + coalesce(metadata, '{}'::jsonb) AS metadata, + priority, + queue, + coalesce(nullif(scheduled_at, '0001-01-01 00:00:00 +0000'), now()) AS scheduled_at, + state::/* TEMPLATE: schema */river_job_state, + string_to_array(tags, ',')::varchar(255)[], + -- ` + "`" + `nullif` + "`" + ` is required for ` + "`" + `lib/pq` + "`" + `, which doesn't do a good job of reading + -- ` + "`" + `nil` + "`" + ` into ` + "`" + `bytea` + "`" + `. We use ` + "`" + `text` + "`" + ` because otherwise ` + "`" + `lib/pq` + "`" + ` will encode + -- to Postgres binary like ` + "`" + `\xAAAA` + "`" + `. + nullif(unique_key, '')::bytea, + nullif(unique_states::integer, 0)::bit(8) +FROM raw_job_data +RETURNING id, args, attempt, attempted_at, attempted_by, created_at, errors, finalized_at, kind, max_attempts, metadata, priority, queue, state, scheduled_at, tags, unique_key, unique_states +` + +type JobInsertFullManyParams struct { + Args [][]byte + Attempt []int16 + AttemptedAt []time.Time + CreatedAt []time.Time + FinalizedAt []time.Time + Kind []string + MaxAttempts []int16 + Metadata [][]byte + Priority []int16 + Queue []string + ScheduledAt []time.Time + State []string + Tags []string + UniqueKey []string + UniqueStates []int32 +} + +func (q *Queries) JobInsertFullMany(ctx context.Context, db DBTX, arg *JobInsertFullManyParams) ([]*RiverJob, error) { + rows, err := db.Query(ctx, jobInsertFullMany, + arg.Args, + arg.Attempt, + arg.AttemptedAt, + arg.CreatedAt, + arg.FinalizedAt, + arg.Kind, + arg.MaxAttempts, + arg.Metadata, + arg.Priority, + arg.Queue, + arg.ScheduledAt, + arg.State, + arg.Tags, + arg.UniqueKey, + arg.UniqueStates, + ) + if err != nil { + return nil, err + } + defer rows.Close() + var items []*RiverJob + for rows.Next() { + var i RiverJob + if err := rows.Scan( + &i.ID, + &i.Args, + &i.Attempt, + &i.AttemptedAt, + &i.AttemptedBy, + &i.CreatedAt, + &i.Errors, + &i.FinalizedAt, + &i.Kind, + &i.MaxAttempts, + &i.Metadata, + &i.Priority, + &i.Queue, + &i.State, + &i.ScheduledAt, + &i.Tags, + &i.UniqueKey, + &i.UniqueStates, + ); err != nil { + return nil, err + } + items = append(items, &i) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const jobKindList = `-- name: JobKindList :many +SELECT DISTINCT ON (kind) kind +FROM /* TEMPLATE: schema */river_job +WHERE ($1 = '' OR kind ILIKE '%' || $1 || '%') + AND ($2 = '' OR kind > $2) + AND ($3::text[] IS NULL OR kind != ALL($3)) +ORDER BY kind ASC +LIMIT $4 +` + +type JobKindListParams struct { + Match interface{} + After interface{} + Exclude []string + Max int32 +} + +func (q *Queries) JobKindList(ctx context.Context, db DBTX, arg *JobKindListParams) ([]string, error) { + rows, err := db.Query(ctx, jobKindList, + arg.Match, + arg.After, + arg.Exclude, + arg.Max, + ) + if err != nil { + return nil, err + } + defer rows.Close() + var items []string + for rows.Next() { + var kind string + if err := rows.Scan(&kind); err != nil { + return nil, err + } + items = append(items, kind) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const jobList = `-- name: JobList :many +SELECT id, args, attempt, attempted_at, attempted_by, created_at, errors, finalized_at, kind, max_attempts, metadata, priority, queue, state, scheduled_at, tags, unique_key, unique_states +FROM /* TEMPLATE: schema */river_job +WHERE /* TEMPLATE_BEGIN: where_clause */ true /* TEMPLATE_END */ +ORDER BY /* TEMPLATE_BEGIN: order_by_clause */ id /* TEMPLATE_END */ +LIMIT $1::int +` + +func (q *Queries) JobList(ctx context.Context, db DBTX, max int32) ([]*RiverJob, error) { + rows, err := db.Query(ctx, jobList, max) + if err != nil { + return nil, err + } + defer rows.Close() + var items []*RiverJob + for rows.Next() { + var i RiverJob + if err := rows.Scan( + &i.ID, + &i.Args, + &i.Attempt, + &i.AttemptedAt, + &i.AttemptedBy, + &i.CreatedAt, + &i.Errors, + &i.FinalizedAt, + &i.Kind, + &i.MaxAttempts, + &i.Metadata, + &i.Priority, + &i.Queue, + &i.State, + &i.ScheduledAt, + &i.Tags, + &i.UniqueKey, + &i.UniqueStates, + ); err != nil { + return nil, err + } + items = append(items, &i) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const jobRescueMany = `-- name: JobRescueMany :exec +UPDATE /* TEMPLATE: schema */river_job +SET + errors = array_append(errors, updated_job.error), + finalized_at = updated_job.finalized_at, + scheduled_at = updated_job.scheduled_at, + metadata = river_job.metadata || jsonb_build_object( + 'river:rescue_count', + coalesce( + CASE + WHEN jsonb_typeof(river_job.metadata -> 'river:rescue_count') = 'number' + THEN (river_job.metadata ->> 'river:rescue_count')::int + END, + 0 + ) + 1 + ), + state = updated_job.state +FROM ( + SELECT + unnest($1::bigint[]) AS id, + unnest($2::jsonb[]) AS error, + nullif(unnest($3::timestamptz[]), '0001-01-01 00:00:00 +0000') AS finalized_at, + unnest($4::timestamptz[]) AS scheduled_at, + unnest($5::text[])::/* TEMPLATE: schema */river_job_state AS state +) AS updated_job +WHERE river_job.id = updated_job.id +` + +type JobRescueManyParams struct { + ID []int64 + Error [][]byte + FinalizedAt []time.Time + ScheduledAt []time.Time + State []string +} + +// Run by the rescuer to queue for retry or discard depending on job state. +func (q *Queries) JobRescueMany(ctx context.Context, db DBTX, arg *JobRescueManyParams) error { + _, err := db.Exec(ctx, jobRescueMany, + arg.ID, + arg.Error, + arg.FinalizedAt, + arg.ScheduledAt, + arg.State, + ) + return err +} + +const jobRetry = `-- name: JobRetry :one +WITH job_to_update AS ( + SELECT id + FROM /* TEMPLATE: schema */river_job + WHERE river_job.id = $1 + FOR UPDATE +), +updated_job AS ( + UPDATE /* TEMPLATE: schema */river_job + SET + state = 'available', + max_attempts = CASE WHEN attempt = max_attempts THEN max_attempts + 1 ELSE max_attempts END, + finalized_at = NULL, + scheduled_at = coalesce($2::timestamptz, now()) + FROM job_to_update + WHERE river_job.id = job_to_update.id + -- Do not touch running jobs: + AND river_job.state != 'running' + -- If the job is already available with a prior scheduled_at, leave it alone. + AND NOT ( + river_job.state = 'available' + AND river_job.scheduled_at < coalesce($2::timestamptz, now()) + ) + RETURNING river_job.id, river_job.args, river_job.attempt, river_job.attempted_at, river_job.attempted_by, river_job.created_at, river_job.errors, river_job.finalized_at, river_job.kind, river_job.max_attempts, river_job.metadata, river_job.priority, river_job.queue, river_job.state, river_job.scheduled_at, river_job.tags, river_job.unique_key, river_job.unique_states +) +SELECT id, args, attempt, attempted_at, attempted_by, created_at, errors, finalized_at, kind, max_attempts, metadata, priority, queue, state, scheduled_at, tags, unique_key, unique_states +FROM /* TEMPLATE: schema */river_job +WHERE id = $1::bigint + AND id NOT IN (SELECT id FROM updated_job) +UNION +SELECT id, args, attempt, attempted_at, attempted_by, created_at, errors, finalized_at, kind, max_attempts, metadata, priority, queue, state, scheduled_at, tags, unique_key, unique_states +FROM updated_job +` + +type JobRetryParams struct { + ID int64 + Now *time.Time +} + +func (q *Queries) JobRetry(ctx context.Context, db DBTX, arg *JobRetryParams) (*RiverJob, error) { + row := db.QueryRow(ctx, jobRetry, arg.ID, arg.Now) + var i RiverJob + err := row.Scan( + &i.ID, + &i.Args, + &i.Attempt, + &i.AttemptedAt, + &i.AttemptedBy, + &i.CreatedAt, + &i.Errors, + &i.FinalizedAt, + &i.Kind, + &i.MaxAttempts, + &i.Metadata, + &i.Priority, + &i.Queue, + &i.State, + &i.ScheduledAt, + &i.Tags, + &i.UniqueKey, + &i.UniqueStates, + ) + return &i, err +} + +const jobSchedule = `-- name: JobSchedule :many +WITH jobs_to_schedule AS ( + SELECT + id, + unique_key, + unique_states, + priority, + scheduled_at + FROM /* TEMPLATE: schema */river_job + WHERE + state IN ('retryable', 'scheduled') + AND priority >= 0 + AND queue IS NOT NULL + AND scheduled_at <= coalesce($1::timestamptz, now()) + ORDER BY + priority, + scheduled_at, + id + LIMIT $2::bigint + FOR UPDATE +), +jobs_with_rownum AS ( + SELECT + id, unique_key, unique_states, priority, scheduled_at, + CASE + WHEN unique_key IS NOT NULL AND unique_states IS NOT NULL THEN + ROW_NUMBER() OVER ( + PARTITION BY unique_key + ORDER BY priority, scheduled_at, id + ) + ELSE NULL + END AS row_num + FROM jobs_to_schedule +), +unique_conflicts AS ( + SELECT river_job.unique_key + FROM /* TEMPLATE: schema */river_job + JOIN jobs_with_rownum + ON river_job.unique_key = jobs_with_rownum.unique_key + AND river_job.id != jobs_with_rownum.id + WHERE + river_job.unique_key IS NOT NULL + AND river_job.unique_states IS NOT NULL + AND /* TEMPLATE: schema */river_job_state_in_bitmask(river_job.unique_states, river_job.state) +), +job_updates AS ( + SELECT + job.id, + job.unique_key, + job.unique_states, + CASE + WHEN job.row_num IS NULL THEN 'available'::/* TEMPLATE: schema */river_job_state + WHEN uc.unique_key IS NOT NULL THEN 'discarded'::/* TEMPLATE: schema */river_job_state + WHEN job.row_num = 1 THEN 'available'::/* TEMPLATE: schema */river_job_state + ELSE 'discarded'::/* TEMPLATE: schema */river_job_state + END AS new_state, + (job.row_num IS NOT NULL AND (uc.unique_key IS NOT NULL OR job.row_num > 1)) AS finalized_at_do_update, + (job.row_num IS NOT NULL AND (uc.unique_key IS NOT NULL OR job.row_num > 1)) AS metadata_do_update + FROM jobs_with_rownum job + LEFT JOIN unique_conflicts uc ON job.unique_key = uc.unique_key +), +updated_jobs AS ( + UPDATE /* TEMPLATE: schema */river_job + SET + state = job_updates.new_state, + finalized_at = CASE WHEN job_updates.finalized_at_do_update THEN coalesce($1::timestamptz, now()) + ELSE river_job.finalized_at END, + metadata = CASE WHEN job_updates.metadata_do_update THEN river_job.metadata || '{"unique_key_conflict": "scheduler_discarded"}'::jsonb + ELSE river_job.metadata END + FROM job_updates + WHERE river_job.id = job_updates.id + RETURNING + river_job.id, + job_updates.new_state = 'discarded'::/* TEMPLATE: schema */river_job_state AS conflict_discarded +) +SELECT + river_job.id, river_job.args, river_job.attempt, river_job.attempted_at, river_job.attempted_by, river_job.created_at, river_job.errors, river_job.finalized_at, river_job.kind, river_job.max_attempts, river_job.metadata, river_job.priority, river_job.queue, river_job.state, river_job.scheduled_at, river_job.tags, river_job.unique_key, river_job.unique_states, + updated_jobs.conflict_discarded +FROM /* TEMPLATE: schema */river_job +JOIN updated_jobs ON river_job.id = updated_jobs.id +` + +type JobScheduleParams struct { + Now *time.Time + Max int64 +} + +type JobScheduleRow struct { + RiverJob RiverJob + ConflictDiscarded bool +} + +func (q *Queries) JobSchedule(ctx context.Context, db DBTX, arg *JobScheduleParams) ([]*JobScheduleRow, error) { + rows, err := db.Query(ctx, jobSchedule, arg.Now, arg.Max) + if err != nil { + return nil, err + } + defer rows.Close() + var items []*JobScheduleRow + for rows.Next() { + var i JobScheduleRow + if err := rows.Scan( + &i.RiverJob.ID, + &i.RiverJob.Args, + &i.RiverJob.Attempt, + &i.RiverJob.AttemptedAt, + &i.RiverJob.AttemptedBy, + &i.RiverJob.CreatedAt, + &i.RiverJob.Errors, + &i.RiverJob.FinalizedAt, + &i.RiverJob.Kind, + &i.RiverJob.MaxAttempts, + &i.RiverJob.Metadata, + &i.RiverJob.Priority, + &i.RiverJob.Queue, + &i.RiverJob.State, + &i.RiverJob.ScheduledAt, + &i.RiverJob.Tags, + &i.RiverJob.UniqueKey, + &i.RiverJob.UniqueStates, + &i.ConflictDiscarded, + ); err != nil { + return nil, err + } + items = append(items, &i) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const jobSetStateIfRunningMany = `-- name: JobSetStateIfRunningMany :many +WITH job_input AS ( + SELECT + unnest($1::bigint[]) AS id, + unnest($2::boolean[]) AS attempt_do_update, + unnest($3::int[]) AS attempt, + unnest($4::boolean[]) AS errors_do_update, + unnest($5::jsonb[]) AS errors, + unnest($6::boolean[]) AS finalized_at_do_update, + unnest($7::timestamptz[]) AS finalized_at, + unnest($8::boolean[]) AS metadata_do_merge, + unnest($9::jsonb[]) AS metadata_updates, + unnest($10::boolean[]) AS scheduled_at_do_update, + unnest($11::timestamptz[]) AS scheduled_at, + -- To avoid requiring pgx users to register the OID of the river_job_state[] + -- type, we cast the array to text[] and then to river_job_state. + unnest($12::text[])::/* TEMPLATE: schema */river_job_state AS state +), +updated AS ( + UPDATE /* TEMPLATE: schema */river_job + SET + attempt = CASE + WHEN river_job.state = 'running' + AND NOT (job_input.state IN ('retryable','scheduled') AND river_job.metadata ? 'cancel_attempted_at') + AND job_input.attempt_do_update + THEN job_input.attempt + ELSE river_job.attempt + END, + errors = CASE + WHEN river_job.state = 'running' + AND job_input.errors_do_update + THEN array_append(river_job.errors, job_input.errors) + ELSE river_job.errors + END, + finalized_at = CASE + WHEN river_job.state = 'running' + AND (job_input.state IN ('retryable','scheduled') AND river_job.metadata ? 'cancel_attempted_at') + THEN coalesce($13::timestamptz, now()) + WHEN river_job.state = 'running' + AND job_input.finalized_at_do_update + THEN job_input.finalized_at + ELSE river_job.finalized_at + END, + metadata = CASE + WHEN job_input.metadata_do_merge + THEN river_job.metadata || job_input.metadata_updates + ELSE river_job.metadata + END, + scheduled_at = CASE + WHEN river_job.state = 'running' + AND NOT (job_input.state IN ('retryable','scheduled') AND river_job.metadata ? 'cancel_attempted_at') + AND job_input.scheduled_at_do_update + THEN job_input.scheduled_at + ELSE river_job.scheduled_at + END, + state = CASE + WHEN river_job.state = 'running' + AND (job_input.state IN ('retryable','scheduled') AND river_job.metadata ? 'cancel_attempted_at') + THEN 'cancelled'::/* TEMPLATE: schema */river_job_state + WHEN river_job.state = 'running' + THEN job_input.state + ELSE river_job.state + END + FROM job_input + WHERE river_job.id = job_input.id + AND (river_job.state = 'running' OR job_input.metadata_do_merge) + RETURNING river_job.id, river_job.args, river_job.attempt, river_job.attempted_at, river_job.attempted_by, river_job.created_at, river_job.errors, river_job.finalized_at, river_job.kind, river_job.max_attempts, river_job.metadata, river_job.priority, river_job.queue, river_job.state, river_job.scheduled_at, river_job.tags, river_job.unique_key, river_job.unique_states +) +SELECT river_job.id, river_job.args, river_job.attempt, river_job.attempted_at, river_job.attempted_by, river_job.created_at, river_job.errors, river_job.finalized_at, river_job.kind, river_job.max_attempts, river_job.metadata, river_job.priority, river_job.queue, river_job.state, river_job.scheduled_at, river_job.tags, river_job.unique_key, river_job.unique_states +FROM /* TEMPLATE: schema */river_job +JOIN job_input ON river_job.id = job_input.id +WHERE NOT EXISTS ( + SELECT 1 + FROM updated + WHERE updated.id = river_job.id +) +UNION ALL +SELECT id, args, attempt, attempted_at, attempted_by, created_at, errors, finalized_at, kind, max_attempts, metadata, priority, queue, state, scheduled_at, tags, unique_key, unique_states +FROM updated +ORDER BY id +` + +type JobSetStateIfRunningManyParams struct { + IDs []int64 + AttemptDoUpdate []bool + Attempt []int32 + ErrorsDoUpdate []bool + Errors [][]byte + FinalizedAtDoUpdate []bool + FinalizedAt []time.Time + MetadataDoMerge []bool + MetadataUpdates [][]byte + ScheduledAtDoUpdate []bool + ScheduledAt []time.Time + State []string + Now *time.Time +} + +func (q *Queries) JobSetStateIfRunningMany(ctx context.Context, db DBTX, arg *JobSetStateIfRunningManyParams) ([]*RiverJob, error) { + rows, err := db.Query(ctx, jobSetStateIfRunningMany, + arg.IDs, + arg.AttemptDoUpdate, + arg.Attempt, + arg.ErrorsDoUpdate, + arg.Errors, + arg.FinalizedAtDoUpdate, + arg.FinalizedAt, + arg.MetadataDoMerge, + arg.MetadataUpdates, + arg.ScheduledAtDoUpdate, + arg.ScheduledAt, + arg.State, + arg.Now, + ) + if err != nil { + return nil, err + } + defer rows.Close() + var items []*RiverJob + for rows.Next() { + var i RiverJob + if err := rows.Scan( + &i.ID, + &i.Args, + &i.Attempt, + &i.AttemptedAt, + &i.AttemptedBy, + &i.CreatedAt, + &i.Errors, + &i.FinalizedAt, + &i.Kind, + &i.MaxAttempts, + &i.Metadata, + &i.Priority, + &i.Queue, + &i.State, + &i.ScheduledAt, + &i.Tags, + &i.UniqueKey, + &i.UniqueStates, + ); err != nil { + return nil, err + } + items = append(items, &i) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const jobUpdate = `-- name: JobUpdate :one +WITH locked_job AS ( + SELECT id + FROM /* TEMPLATE: schema */river_job + WHERE river_job.id = $3 + FOR UPDATE +) +UPDATE /* TEMPLATE: schema */river_job +SET + metadata = CASE WHEN $1::boolean THEN metadata || $2::jsonb ELSE metadata END +FROM + locked_job +WHERE river_job.id = locked_job.id +RETURNING river_job.id, river_job.args, river_job.attempt, river_job.attempted_at, river_job.attempted_by, river_job.created_at, river_job.errors, river_job.finalized_at, river_job.kind, river_job.max_attempts, river_job.metadata, river_job.priority, river_job.queue, river_job.state, river_job.scheduled_at, river_job.tags, river_job.unique_key, river_job.unique_states +` + +type JobUpdateParams struct { + MetadataDoMerge bool + Metadata []byte + ID int64 +} + +func (q *Queries) JobUpdate(ctx context.Context, db DBTX, arg *JobUpdateParams) (*RiverJob, error) { + row := db.QueryRow(ctx, jobUpdate, arg.MetadataDoMerge, arg.Metadata, arg.ID) + var i RiverJob + err := row.Scan( + &i.ID, + &i.Args, + &i.Attempt, + &i.AttemptedAt, + &i.AttemptedBy, + &i.CreatedAt, + &i.Errors, + &i.FinalizedAt, + &i.Kind, + &i.MaxAttempts, + &i.Metadata, + &i.Priority, + &i.Queue, + &i.State, + &i.ScheduledAt, + &i.Tags, + &i.UniqueKey, + &i.UniqueStates, + ) + return &i, err +} + +const jobUpdateFull = `-- name: JobUpdateFull :one +UPDATE /* TEMPLATE: schema */river_job +SET + attempt = CASE WHEN $1::boolean THEN $2 ELSE attempt END, + attempted_at = CASE WHEN $3::boolean THEN $4 ELSE attempted_at END, + attempted_by = CASE WHEN $5::boolean THEN $6 ELSE attempted_by END, + errors = CASE WHEN $7::boolean THEN $8::jsonb[] ELSE errors END, + finalized_at = CASE WHEN $9::boolean THEN $10 ELSE finalized_at END, + max_attempts = CASE WHEN $11::boolean THEN $12 ELSE max_attempts END, + metadata = CASE WHEN $13::boolean THEN $14::jsonb ELSE metadata END, + state = CASE WHEN $15::boolean THEN $16::/* TEMPLATE: schema */river_job_state ELSE state END +WHERE id = $17 +RETURNING id, args, attempt, attempted_at, attempted_by, created_at, errors, finalized_at, kind, max_attempts, metadata, priority, queue, state, scheduled_at, tags, unique_key, unique_states +` + +type JobUpdateFullParams struct { + AttemptDoUpdate bool + Attempt int16 + AttemptedAtDoUpdate bool + AttemptedAt *time.Time + AttemptedByDoUpdate bool + AttemptedBy []string + ErrorsDoUpdate bool + Errors [][]byte + FinalizedAtDoUpdate bool + FinalizedAt *time.Time + MaxAttemptsDoUpdate bool + MaxAttempts int16 + MetadataDoUpdate bool + Metadata []byte + StateDoUpdate bool + State RiverJobState + ID int64 +} + +// A generalized update for any property on a job. This brings in a large number +// of parameters and therefore may be more suitable for testing than production. +func (q *Queries) JobUpdateFull(ctx context.Context, db DBTX, arg *JobUpdateFullParams) (*RiverJob, error) { + row := db.QueryRow(ctx, jobUpdateFull, + arg.AttemptDoUpdate, + arg.Attempt, + arg.AttemptedAtDoUpdate, + arg.AttemptedAt, + arg.AttemptedByDoUpdate, + arg.AttemptedBy, + arg.ErrorsDoUpdate, + arg.Errors, + arg.FinalizedAtDoUpdate, + arg.FinalizedAt, + arg.MaxAttemptsDoUpdate, + arg.MaxAttempts, + arg.MetadataDoUpdate, + arg.Metadata, + arg.StateDoUpdate, + arg.State, + arg.ID, + ) + var i RiverJob + err := row.Scan( + &i.ID, + &i.Args, + &i.Attempt, + &i.AttemptedAt, + &i.AttemptedBy, + &i.CreatedAt, + &i.Errors, + &i.FinalizedAt, + &i.Kind, + &i.MaxAttempts, + &i.Metadata, + &i.Priority, + &i.Queue, + &i.State, + &i.ScheduledAt, + &i.Tags, + &i.UniqueKey, + &i.UniqueStates, + ) + return &i, err +} diff --git a/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/river_job_copyfrom.sql b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/river_job_copyfrom.sql new file mode 100644 index 0000000000..54fdbeaa40 --- /dev/null +++ b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/river_job_copyfrom.sql @@ -0,0 +1,28 @@ +-- name: JobInsertFastManyCopyFrom :copyfrom +INSERT INTO river_job( + args, + created_at, + kind, + max_attempts, + metadata, + priority, + queue, + scheduled_at, + state, + tags, + unique_key, + unique_states +) VALUES ( + @args, + @created_at, + @kind, + @max_attempts, + @metadata, + @priority, + @queue, + @scheduled_at, + @state, + @tags, + @unique_key, + @unique_states +); diff --git a/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/river_job_copyfrom.sql.go b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/river_job_copyfrom.sql.go new file mode 100644 index 0000000000..3daff367dc --- /dev/null +++ b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/river_job_copyfrom.sql.go @@ -0,0 +1,27 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.31.0 +// source: river_job_copyfrom.sql + +package dbsqlc + +import ( + "time" + + "github.com/jackc/pgx/v5/pgtype" +) + +type JobInsertFastManyCopyFromParams struct { + Args []byte + CreatedAt time.Time + Kind string + MaxAttempts int16 + Metadata []byte + Priority int16 + Queue string + ScheduledAt time.Time + State RiverJobState + Tags []string + UniqueKey []byte + UniqueStates pgtype.Bits +} diff --git a/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/river_leader.sql b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/river_leader.sql new file mode 100644 index 0000000000..cea2195f14 --- /dev/null +++ b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/river_leader.sql @@ -0,0 +1,69 @@ +CREATE UNLOGGED TABLE river_leader( + elected_at timestamptz NOT NULL, + expires_at timestamptz NOT NULL, + leader_id text NOT NULL, + name text PRIMARY KEY DEFAULT 'default' CHECK (name = 'default'), + CONSTRAINT name_length CHECK (name = 'default'), + CONSTRAINT leader_id_length CHECK (char_length(leader_id) > 0 AND char_length(leader_id) < 128) +); + +-- name: LeaderAttemptElect :one +INSERT INTO /* TEMPLATE: schema */river_leader ( + leader_id, + elected_at, + expires_at +) VALUES ( + @leader_id, + coalesce(sqlc.narg('now')::timestamptz, now()), + -- @ttl is inserted as as seconds rather than a duration because `lib/pq` doesn't support the latter + coalesce(sqlc.narg('now')::timestamptz, now()) + make_interval(secs => @ttl) +) +ON CONFLICT (name) + DO NOTHING +RETURNING *; + +-- name: LeaderAttemptReelect :one +UPDATE /* TEMPLATE: schema */river_leader +SET expires_at = coalesce(sqlc.narg('now')::timestamptz, now()) + make_interval(secs => @ttl) +WHERE + elected_at = @elected_at::timestamptz + AND expires_at >= coalesce(sqlc.narg('now')::timestamptz, now()) + AND leader_id = @leader_id +RETURNING *; + +-- name: LeaderDeleteExpired :execrows +DELETE FROM /* TEMPLATE: schema */river_leader +WHERE expires_at < coalesce(sqlc.narg('now')::timestamptz, now()); + +-- name: LeaderGetElectedLeader :one +SELECT * +FROM /* TEMPLATE: schema */river_leader; + +-- name: LeaderInsert :one +INSERT INTO /* TEMPLATE: schema */river_leader( + elected_at, + expires_at, + leader_id +) VALUES ( + coalesce(sqlc.narg('elected_at')::timestamptz, coalesce(sqlc.narg('now')::timestamptz, now())), + coalesce(sqlc.narg('expires_at')::timestamptz, coalesce(sqlc.narg('now')::timestamptz, now()) + make_interval(secs => @ttl)), + @leader_id +) RETURNING *; + +-- name: LeaderResign :execrows +WITH currently_held_leaders AS ( + SELECT * + FROM /* TEMPLATE: schema */river_leader + WHERE + elected_at = @elected_at::timestamptz + AND leader_id = @leader_id::text + FOR UPDATE +), +notified_resignations AS ( + SELECT pg_notify( + concat(coalesce(sqlc.narg('schema')::text, current_schema()), '.', @leadership_topic::text), + json_build_object('leader_id', leader_id, 'action', 'resigned')::text + ) + FROM currently_held_leaders +) +DELETE FROM /* TEMPLATE: schema */river_leader USING notified_resignations; diff --git a/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/river_leader.sql.go b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/river_leader.sql.go new file mode 100644 index 0000000000..1976cf4d7b --- /dev/null +++ b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/river_leader.sql.go @@ -0,0 +1,188 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.31.0 +// source: river_leader.sql + +package dbsqlc + +import ( + "context" + "time" + + "github.com/jackc/pgx/v5/pgtype" +) + +const leaderAttemptElect = `-- name: LeaderAttemptElect :one +INSERT INTO /* TEMPLATE: schema */river_leader ( + leader_id, + elected_at, + expires_at +) VALUES ( + $1, + coalesce($2::timestamptz, now()), + -- @ttl is inserted as as seconds rather than a duration because ` + "`" + `lib/pq` + "`" + ` doesn't support the latter + coalesce($2::timestamptz, now()) + make_interval(secs => $3) +) +ON CONFLICT (name) + DO NOTHING +RETURNING elected_at, expires_at, leader_id, name +` + +type LeaderAttemptElectParams struct { + LeaderID string + Now *time.Time + TTL float64 +} + +func (q *Queries) LeaderAttemptElect(ctx context.Context, db DBTX, arg *LeaderAttemptElectParams) (*RiverLeader, error) { + row := db.QueryRow(ctx, leaderAttemptElect, arg.LeaderID, arg.Now, arg.TTL) + var i RiverLeader + err := row.Scan( + &i.ElectedAt, + &i.ExpiresAt, + &i.LeaderID, + &i.Name, + ) + return &i, err +} + +const leaderAttemptReelect = `-- name: LeaderAttemptReelect :one +UPDATE /* TEMPLATE: schema */river_leader +SET expires_at = coalesce($1::timestamptz, now()) + make_interval(secs => $2) +WHERE + elected_at = $3::timestamptz + AND expires_at >= coalesce($1::timestamptz, now()) + AND leader_id = $4 +RETURNING elected_at, expires_at, leader_id, name +` + +type LeaderAttemptReelectParams struct { + Now *time.Time + TTL float64 + ElectedAt time.Time + LeaderID string +} + +func (q *Queries) LeaderAttemptReelect(ctx context.Context, db DBTX, arg *LeaderAttemptReelectParams) (*RiverLeader, error) { + row := db.QueryRow(ctx, leaderAttemptReelect, + arg.Now, + arg.TTL, + arg.ElectedAt, + arg.LeaderID, + ) + var i RiverLeader + err := row.Scan( + &i.ElectedAt, + &i.ExpiresAt, + &i.LeaderID, + &i.Name, + ) + return &i, err +} + +const leaderDeleteExpired = `-- name: LeaderDeleteExpired :execrows +DELETE FROM /* TEMPLATE: schema */river_leader +WHERE expires_at < coalesce($1::timestamptz, now()) +` + +func (q *Queries) LeaderDeleteExpired(ctx context.Context, db DBTX, now *time.Time) (int64, error) { + result, err := db.Exec(ctx, leaderDeleteExpired, now) + if err != nil { + return 0, err + } + return result.RowsAffected(), nil +} + +const leaderGetElectedLeader = `-- name: LeaderGetElectedLeader :one +SELECT elected_at, expires_at, leader_id, name +FROM /* TEMPLATE: schema */river_leader +` + +func (q *Queries) LeaderGetElectedLeader(ctx context.Context, db DBTX) (*RiverLeader, error) { + row := db.QueryRow(ctx, leaderGetElectedLeader) + var i RiverLeader + err := row.Scan( + &i.ElectedAt, + &i.ExpiresAt, + &i.LeaderID, + &i.Name, + ) + return &i, err +} + +const leaderInsert = `-- name: LeaderInsert :one +INSERT INTO /* TEMPLATE: schema */river_leader( + elected_at, + expires_at, + leader_id +) VALUES ( + coalesce($1::timestamptz, coalesce($2::timestamptz, now())), + coalesce($3::timestamptz, coalesce($2::timestamptz, now()) + make_interval(secs => $4)), + $5 +) RETURNING elected_at, expires_at, leader_id, name +` + +type LeaderInsertParams struct { + ElectedAt *time.Time + Now *time.Time + ExpiresAt *time.Time + TTL float64 + LeaderID string +} + +func (q *Queries) LeaderInsert(ctx context.Context, db DBTX, arg *LeaderInsertParams) (*RiverLeader, error) { + row := db.QueryRow(ctx, leaderInsert, + arg.ElectedAt, + arg.Now, + arg.ExpiresAt, + arg.TTL, + arg.LeaderID, + ) + var i RiverLeader + err := row.Scan( + &i.ElectedAt, + &i.ExpiresAt, + &i.LeaderID, + &i.Name, + ) + return &i, err +} + +const leaderResign = `-- name: LeaderResign :execrows +WITH currently_held_leaders AS ( + SELECT elected_at, expires_at, leader_id, name + FROM /* TEMPLATE: schema */river_leader + WHERE + elected_at = $1::timestamptz + AND leader_id = $2::text + FOR UPDATE +), +notified_resignations AS ( + SELECT pg_notify( + concat(coalesce($3::text, current_schema()), '.', $4::text), + json_build_object('leader_id', leader_id, 'action', 'resigned')::text + ) + FROM currently_held_leaders +) +DELETE FROM /* TEMPLATE: schema */river_leader USING notified_resignations +` + +type LeaderResignParams struct { + ElectedAt time.Time + LeaderID string + Schema pgtype.Text + LeadershipTopic string +} + +func (q *Queries) LeaderResign(ctx context.Context, db DBTX, arg *LeaderResignParams) (int64, error) { + result, err := db.Exec(ctx, leaderResign, + arg.ElectedAt, + arg.LeaderID, + arg.Schema, + arg.LeadershipTopic, + ) + if err != nil { + return 0, err + } + return result.RowsAffected(), nil +} diff --git a/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/river_migration.sql b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/river_migration.sql new file mode 100644 index 0000000000..a9d9b03f42 --- /dev/null +++ b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/river_migration.sql @@ -0,0 +1,69 @@ +CREATE TABLE river_migration( + line text NOT NULL, + version bigint NOT NULL, + created_at timestamptz NOT NULL DEFAULT now(), + CONSTRAINT line_length CHECK (char_length(line) > 0 AND char_length(line) < 128), + CONSTRAINT version_gte_1 CHECK (version >= 1), + PRIMARY KEY (line, version) +); + +-- name: RiverMigrationDeleteAssumingMainMany :many +DELETE FROM /* TEMPLATE: schema */river_migration +WHERE version = any(@version::bigint[]) +RETURNING + created_at, + version; + +-- name: RiverMigrationDeleteByLineAndVersionMany :many +DELETE FROM /* TEMPLATE: schema */river_migration +WHERE line = @line + AND version = any(@version::bigint[]) +RETURNING *; + +-- This is a compatibility query for getting existing migrations before the +-- `line` column was added to the table in version 005. We need to make sure to +-- only select non-line properties so the query doesn't error on older schemas. +-- (Even if we use `SELECT *` below, sqlc materializes it to a list of column +-- names in the generated query.) +-- +-- name: RiverMigrationGetAllAssumingMain :many +SELECT + created_at, + version +FROM /* TEMPLATE: schema */river_migration +ORDER BY version; + +-- name: RiverMigrationGetByLine :many +SELECT * +FROM /* TEMPLATE: schema */river_migration +WHERE line = @line +ORDER BY version; + +-- name: RiverMigrationInsert :one +INSERT INTO /* TEMPLATE: schema */river_migration ( + line, + version +) VALUES ( + @line, + @version +) RETURNING *; + +-- name: RiverMigrationInsertMany :many +INSERT INTO /* TEMPLATE: schema */river_migration ( + line, + version +) +SELECT + @line, + unnest(@version::bigint[]) +RETURNING *; + +-- name: RiverMigrationInsertManyAssumingMain :many +INSERT INTO /* TEMPLATE: schema */river_migration ( + version +) +SELECT + unnest(@version::bigint[]) +RETURNING + created_at, + version; \ No newline at end of file diff --git a/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/river_migration.sql.go b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/river_migration.sql.go new file mode 100644 index 0000000000..bf20e03bda --- /dev/null +++ b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/river_migration.sql.go @@ -0,0 +1,235 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.31.0 +// source: river_migration.sql + +package dbsqlc + +import ( + "context" + "time" +) + +const riverMigrationDeleteAssumingMainMany = `-- name: RiverMigrationDeleteAssumingMainMany :many +DELETE FROM /* TEMPLATE: schema */river_migration +WHERE version = any($1::bigint[]) +RETURNING + created_at, + version +` + +type RiverMigrationDeleteAssumingMainManyRow struct { + CreatedAt time.Time + Version int64 +} + +func (q *Queries) RiverMigrationDeleteAssumingMainMany(ctx context.Context, db DBTX, version []int64) ([]*RiverMigrationDeleteAssumingMainManyRow, error) { + rows, err := db.Query(ctx, riverMigrationDeleteAssumingMainMany, version) + if err != nil { + return nil, err + } + defer rows.Close() + var items []*RiverMigrationDeleteAssumingMainManyRow + for rows.Next() { + var i RiverMigrationDeleteAssumingMainManyRow + if err := rows.Scan(&i.CreatedAt, &i.Version); err != nil { + return nil, err + } + items = append(items, &i) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const riverMigrationDeleteByLineAndVersionMany = `-- name: RiverMigrationDeleteByLineAndVersionMany :many +DELETE FROM /* TEMPLATE: schema */river_migration +WHERE line = $1 + AND version = any($2::bigint[]) +RETURNING line, version, created_at +` + +type RiverMigrationDeleteByLineAndVersionManyParams struct { + Line string + Version []int64 +} + +func (q *Queries) RiverMigrationDeleteByLineAndVersionMany(ctx context.Context, db DBTX, arg *RiverMigrationDeleteByLineAndVersionManyParams) ([]*RiverMigration, error) { + rows, err := db.Query(ctx, riverMigrationDeleteByLineAndVersionMany, arg.Line, arg.Version) + if err != nil { + return nil, err + } + defer rows.Close() + var items []*RiverMigration + for rows.Next() { + var i RiverMigration + if err := rows.Scan(&i.Line, &i.Version, &i.CreatedAt); err != nil { + return nil, err + } + items = append(items, &i) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const riverMigrationGetAllAssumingMain = `-- name: RiverMigrationGetAllAssumingMain :many +SELECT + created_at, + version +FROM /* TEMPLATE: schema */river_migration +ORDER BY version +` + +type RiverMigrationGetAllAssumingMainRow struct { + CreatedAt time.Time + Version int64 +} + +// This is a compatibility query for getting existing migrations before the +// `line` column was added to the table in version 005. We need to make sure to +// only select non-line properties so the query doesn't error on older schemas. +// (Even if we use `SELECT *` below, sqlc materializes it to a list of column +// names in the generated query.) +func (q *Queries) RiverMigrationGetAllAssumingMain(ctx context.Context, db DBTX) ([]*RiverMigrationGetAllAssumingMainRow, error) { + rows, err := db.Query(ctx, riverMigrationGetAllAssumingMain) + if err != nil { + return nil, err + } + defer rows.Close() + var items []*RiverMigrationGetAllAssumingMainRow + for rows.Next() { + var i RiverMigrationGetAllAssumingMainRow + if err := rows.Scan(&i.CreatedAt, &i.Version); err != nil { + return nil, err + } + items = append(items, &i) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const riverMigrationGetByLine = `-- name: RiverMigrationGetByLine :many +SELECT line, version, created_at +FROM /* TEMPLATE: schema */river_migration +WHERE line = $1 +ORDER BY version +` + +func (q *Queries) RiverMigrationGetByLine(ctx context.Context, db DBTX, line string) ([]*RiverMigration, error) { + rows, err := db.Query(ctx, riverMigrationGetByLine, line) + if err != nil { + return nil, err + } + defer rows.Close() + var items []*RiverMigration + for rows.Next() { + var i RiverMigration + if err := rows.Scan(&i.Line, &i.Version, &i.CreatedAt); err != nil { + return nil, err + } + items = append(items, &i) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const riverMigrationInsert = `-- name: RiverMigrationInsert :one +INSERT INTO /* TEMPLATE: schema */river_migration ( + line, + version +) VALUES ( + $1, + $2 +) RETURNING line, version, created_at +` + +type RiverMigrationInsertParams struct { + Line string + Version int64 +} + +func (q *Queries) RiverMigrationInsert(ctx context.Context, db DBTX, arg *RiverMigrationInsertParams) (*RiverMigration, error) { + row := db.QueryRow(ctx, riverMigrationInsert, arg.Line, arg.Version) + var i RiverMigration + err := row.Scan(&i.Line, &i.Version, &i.CreatedAt) + return &i, err +} + +const riverMigrationInsertMany = `-- name: RiverMigrationInsertMany :many +INSERT INTO /* TEMPLATE: schema */river_migration ( + line, + version +) +SELECT + $1, + unnest($2::bigint[]) +RETURNING line, version, created_at +` + +type RiverMigrationInsertManyParams struct { + Line string + Version []int64 +} + +func (q *Queries) RiverMigrationInsertMany(ctx context.Context, db DBTX, arg *RiverMigrationInsertManyParams) ([]*RiverMigration, error) { + rows, err := db.Query(ctx, riverMigrationInsertMany, arg.Line, arg.Version) + if err != nil { + return nil, err + } + defer rows.Close() + var items []*RiverMigration + for rows.Next() { + var i RiverMigration + if err := rows.Scan(&i.Line, &i.Version, &i.CreatedAt); err != nil { + return nil, err + } + items = append(items, &i) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const riverMigrationInsertManyAssumingMain = `-- name: RiverMigrationInsertManyAssumingMain :many +INSERT INTO /* TEMPLATE: schema */river_migration ( + version +) +SELECT + unnest($1::bigint[]) +RETURNING + created_at, + version +` + +type RiverMigrationInsertManyAssumingMainRow struct { + CreatedAt time.Time + Version int64 +} + +func (q *Queries) RiverMigrationInsertManyAssumingMain(ctx context.Context, db DBTX, version []int64) ([]*RiverMigrationInsertManyAssumingMainRow, error) { + rows, err := db.Query(ctx, riverMigrationInsertManyAssumingMain, version) + if err != nil { + return nil, err + } + defer rows.Close() + var items []*RiverMigrationInsertManyAssumingMainRow + for rows.Next() { + var i RiverMigrationInsertManyAssumingMainRow + if err := rows.Scan(&i.CreatedAt, &i.Version); err != nil { + return nil, err + } + items = append(items, &i) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} diff --git a/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/river_notification.sql b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/river_notification.sql new file mode 100644 index 0000000000..576d2444d9 --- /dev/null +++ b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/river_notification.sql @@ -0,0 +1,16 @@ +-- This table isn't used under Postgres currently, but we have it in place +-- because its useful for simulating under Postgres as if we were running +-- SQLite, and it may be useful as a good listen/notify alternative for Postgres +-- down the line instead of poll-only mode in cases like where a bouncer makes +-- listen/notify difficult to use. +CREATE TABLE river_notification ( + id bigserial PRIMARY KEY, + created_at timestamptz NOT NULL DEFAULT now(), + payload text NOT NULL, + topic text NOT NULL, + CONSTRAINT topic_length CHECK (length(topic) > 0 AND length(topic) < 128) +); + +-- name: NotificationDeleteBefore :execrows +DELETE FROM /* TEMPLATE: schema */river_notification +WHERE created_at < @created_at_horizon::timestamptz; diff --git a/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/river_notification.sql.go b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/river_notification.sql.go new file mode 100644 index 0000000000..cb460451cb --- /dev/null +++ b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/river_notification.sql.go @@ -0,0 +1,24 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.31.0 +// source: river_notification.sql + +package dbsqlc + +import ( + "context" + "time" +) + +const notificationDeleteBefore = `-- name: NotificationDeleteBefore :execrows +DELETE FROM /* TEMPLATE: schema */river_notification +WHERE created_at < $1::timestamptz +` + +func (q *Queries) NotificationDeleteBefore(ctx context.Context, db DBTX, createdAtHorizon time.Time) (int64, error) { + result, err := db.Exec(ctx, notificationDeleteBefore, createdAtHorizon) + if err != nil { + return 0, err + } + return result.RowsAffected(), nil +} diff --git a/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/river_queue.sql b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/river_queue.sql new file mode 100644 index 0000000000..abc45ed44c --- /dev/null +++ b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/river_queue.sql @@ -0,0 +1,78 @@ +CREATE TABLE river_queue ( + name text PRIMARY KEY NOT NULL, + created_at timestamptz NOT NULL DEFAULT now(), + metadata jsonb NOT NULL DEFAULT '{}' ::jsonb, + paused_at timestamptz, + updated_at timestamptz NOT NULL DEFAULT CURRENT_TIMESTAMP +); + +-- name: QueueCreateOrSetUpdatedAt :one +INSERT INTO /* TEMPLATE: schema */river_queue ( + created_at, + metadata, + name, + paused_at, + updated_at +) VALUES ( + coalesce(sqlc.narg('now')::timestamptz, now()), + coalesce(@metadata::jsonb, '{}'::jsonb), + @name::text, + coalesce(sqlc.narg('paused_at')::timestamptz, NULL), + coalesce(sqlc.narg('updated_at')::timestamptz, sqlc.narg('now')::timestamptz, now()) +) ON CONFLICT (name) DO UPDATE +SET + updated_at = EXCLUDED.updated_at +RETURNING *; + +-- name: QueueDeleteExpired :many +DELETE FROM /* TEMPLATE: schema */river_queue +WHERE name IN ( + SELECT name + FROM /* TEMPLATE: schema */river_queue + WHERE river_queue.updated_at < @updated_at_horizon + ORDER BY name ASC + LIMIT @max::bigint +) +RETURNING *; + +-- name: QueueGet :one +SELECT * +FROM /* TEMPLATE: schema */river_queue +WHERE name = @name::text; + +-- name: QueueList :many +SELECT * +FROM /* TEMPLATE: schema */river_queue +ORDER BY name ASC +LIMIT @max; + +-- name: QueueNameList :many +SELECT name +FROM /* TEMPLATE: schema */river_queue +WHERE name > @after::text + AND (@match::text = '' OR name ILIKE '%' || @match::text || '%') + AND (@exclude::text[] IS NULL OR name != ALL(@exclude)) +ORDER BY name +LIMIT @max::int; + +-- name: QueuePause :execrows +UPDATE /* TEMPLATE: schema */river_queue +SET + paused_at = CASE WHEN paused_at IS NULL THEN coalesce(sqlc.narg('now')::timestamptz, now()) ELSE paused_at END, + updated_at = CASE WHEN paused_at IS NULL THEN coalesce(sqlc.narg('now')::timestamptz, now()) ELSE updated_at END +WHERE CASE WHEN @name::text = '*' THEN true ELSE name = @name END; + +-- name: QueueResume :execrows +UPDATE /* TEMPLATE: schema */river_queue +SET + paused_at = NULL, + updated_at = CASE WHEN paused_at IS NOT NULL THEN coalesce(sqlc.narg('now')::timestamptz, now()) ELSE updated_at END +WHERE CASE WHEN @name::text = '*' THEN true ELSE name = @name END; + +-- name: QueueUpdate :one +UPDATE /* TEMPLATE: schema */river_queue +SET + metadata = CASE WHEN @metadata_do_update::boolean THEN @metadata::jsonb ELSE metadata END, + updated_at = now() +WHERE name = @name +RETURNING *; diff --git a/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/river_queue.sql.go b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/river_queue.sql.go new file mode 100644 index 0000000000..0988828b8a --- /dev/null +++ b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/river_queue.sql.go @@ -0,0 +1,264 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.31.0 +// source: river_queue.sql + +package dbsqlc + +import ( + "context" + "time" +) + +const queueCreateOrSetUpdatedAt = `-- name: QueueCreateOrSetUpdatedAt :one +INSERT INTO /* TEMPLATE: schema */river_queue ( + created_at, + metadata, + name, + paused_at, + updated_at +) VALUES ( + coalesce($1::timestamptz, now()), + coalesce($2::jsonb, '{}'::jsonb), + $3::text, + coalesce($4::timestamptz, NULL), + coalesce($5::timestamptz, $1::timestamptz, now()) +) ON CONFLICT (name) DO UPDATE +SET + updated_at = EXCLUDED.updated_at +RETURNING name, created_at, metadata, paused_at, updated_at +` + +type QueueCreateOrSetUpdatedAtParams struct { + Now *time.Time + Metadata []byte + Name string + PausedAt *time.Time + UpdatedAt *time.Time +} + +func (q *Queries) QueueCreateOrSetUpdatedAt(ctx context.Context, db DBTX, arg *QueueCreateOrSetUpdatedAtParams) (*RiverQueue, error) { + row := db.QueryRow(ctx, queueCreateOrSetUpdatedAt, + arg.Now, + arg.Metadata, + arg.Name, + arg.PausedAt, + arg.UpdatedAt, + ) + var i RiverQueue + err := row.Scan( + &i.Name, + &i.CreatedAt, + &i.Metadata, + &i.PausedAt, + &i.UpdatedAt, + ) + return &i, err +} + +const queueDeleteExpired = `-- name: QueueDeleteExpired :many +DELETE FROM /* TEMPLATE: schema */river_queue +WHERE name IN ( + SELECT name + FROM /* TEMPLATE: schema */river_queue + WHERE river_queue.updated_at < $1 + ORDER BY name ASC + LIMIT $2::bigint +) +RETURNING name, created_at, metadata, paused_at, updated_at +` + +type QueueDeleteExpiredParams struct { + UpdatedAtHorizon time.Time + Max int64 +} + +func (q *Queries) QueueDeleteExpired(ctx context.Context, db DBTX, arg *QueueDeleteExpiredParams) ([]*RiverQueue, error) { + rows, err := db.Query(ctx, queueDeleteExpired, arg.UpdatedAtHorizon, arg.Max) + if err != nil { + return nil, err + } + defer rows.Close() + var items []*RiverQueue + for rows.Next() { + var i RiverQueue + if err := rows.Scan( + &i.Name, + &i.CreatedAt, + &i.Metadata, + &i.PausedAt, + &i.UpdatedAt, + ); err != nil { + return nil, err + } + items = append(items, &i) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const queueGet = `-- name: QueueGet :one +SELECT name, created_at, metadata, paused_at, updated_at +FROM /* TEMPLATE: schema */river_queue +WHERE name = $1::text +` + +func (q *Queries) QueueGet(ctx context.Context, db DBTX, name string) (*RiverQueue, error) { + row := db.QueryRow(ctx, queueGet, name) + var i RiverQueue + err := row.Scan( + &i.Name, + &i.CreatedAt, + &i.Metadata, + &i.PausedAt, + &i.UpdatedAt, + ) + return &i, err +} + +const queueList = `-- name: QueueList :many +SELECT name, created_at, metadata, paused_at, updated_at +FROM /* TEMPLATE: schema */river_queue +ORDER BY name ASC +LIMIT $1 +` + +func (q *Queries) QueueList(ctx context.Context, db DBTX, max int32) ([]*RiverQueue, error) { + rows, err := db.Query(ctx, queueList, max) + if err != nil { + return nil, err + } + defer rows.Close() + var items []*RiverQueue + for rows.Next() { + var i RiverQueue + if err := rows.Scan( + &i.Name, + &i.CreatedAt, + &i.Metadata, + &i.PausedAt, + &i.UpdatedAt, + ); err != nil { + return nil, err + } + items = append(items, &i) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const queueNameList = `-- name: QueueNameList :many +SELECT name +FROM /* TEMPLATE: schema */river_queue +WHERE name > $1::text + AND ($2::text = '' OR name ILIKE '%' || $2::text || '%') + AND ($3::text[] IS NULL OR name != ALL($3)) +ORDER BY name +LIMIT $4::int +` + +type QueueNameListParams struct { + After string + Match string + Exclude []string + Max int32 +} + +func (q *Queries) QueueNameList(ctx context.Context, db DBTX, arg *QueueNameListParams) ([]string, error) { + rows, err := db.Query(ctx, queueNameList, + arg.After, + arg.Match, + arg.Exclude, + arg.Max, + ) + if err != nil { + return nil, err + } + defer rows.Close() + var items []string + for rows.Next() { + var name string + if err := rows.Scan(&name); err != nil { + return nil, err + } + items = append(items, name) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const queuePause = `-- name: QueuePause :execrows +UPDATE /* TEMPLATE: schema */river_queue +SET + paused_at = CASE WHEN paused_at IS NULL THEN coalesce($1::timestamptz, now()) ELSE paused_at END, + updated_at = CASE WHEN paused_at IS NULL THEN coalesce($1::timestamptz, now()) ELSE updated_at END +WHERE CASE WHEN $2::text = '*' THEN true ELSE name = $2 END +` + +type QueuePauseParams struct { + Now *time.Time + Name string +} + +func (q *Queries) QueuePause(ctx context.Context, db DBTX, arg *QueuePauseParams) (int64, error) { + result, err := db.Exec(ctx, queuePause, arg.Now, arg.Name) + if err != nil { + return 0, err + } + return result.RowsAffected(), nil +} + +const queueResume = `-- name: QueueResume :execrows +UPDATE /* TEMPLATE: schema */river_queue +SET + paused_at = NULL, + updated_at = CASE WHEN paused_at IS NOT NULL THEN coalesce($1::timestamptz, now()) ELSE updated_at END +WHERE CASE WHEN $2::text = '*' THEN true ELSE name = $2 END +` + +type QueueResumeParams struct { + Now *time.Time + Name string +} + +func (q *Queries) QueueResume(ctx context.Context, db DBTX, arg *QueueResumeParams) (int64, error) { + result, err := db.Exec(ctx, queueResume, arg.Now, arg.Name) + if err != nil { + return 0, err + } + return result.RowsAffected(), nil +} + +const queueUpdate = `-- name: QueueUpdate :one +UPDATE /* TEMPLATE: schema */river_queue +SET + metadata = CASE WHEN $1::boolean THEN $2::jsonb ELSE metadata END, + updated_at = now() +WHERE name = $3 +RETURNING name, created_at, metadata, paused_at, updated_at +` + +type QueueUpdateParams struct { + MetadataDoUpdate bool + Metadata []byte + Name string +} + +func (q *Queries) QueueUpdate(ctx context.Context, db DBTX, arg *QueueUpdateParams) (*RiverQueue, error) { + row := db.QueryRow(ctx, queueUpdate, arg.MetadataDoUpdate, arg.Metadata, arg.Name) + var i RiverQueue + err := row.Scan( + &i.Name, + &i.CreatedAt, + &i.Metadata, + &i.PausedAt, + &i.UpdatedAt, + ) + return &i, err +} diff --git a/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/schema.sql b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/schema.sql new file mode 100644 index 0000000000..a9953a447f --- /dev/null +++ b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/schema.sql @@ -0,0 +1,60 @@ +-- name: ColumnExists :one +SELECT EXISTS ( + SELECT column_name + FROM information_schema.columns + WHERE table_name = @table_name::text + AND table_schema = /* TEMPLATE_BEGIN: schema */ CURRENT_SCHEMA /* TEMPLATE_END */ + AND column_name = @column_name::text +); + +-- name: IndexExists :one +SELECT EXISTS ( + SELECT 1 + FROM pg_catalog.pg_class + JOIN pg_catalog.pg_namespace ON pg_namespace.oid = pg_class.relnamespace + WHERE pg_class.relname = @index::text + AND pg_namespace.nspname = coalesce(sqlc.narg('schema')::text, current_schema()) + AND pg_class.relkind = 'i' +); + +-- name: IndexReindexArtifacts :many +WITH index_artifacts AS ( + SELECT + c.relname::text AS index_name, + substring(c.relname FROM length(@index::text) + 1) AS suffix + FROM pg_catalog.pg_class c + JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace + WHERE n.nspname = coalesce(sqlc.narg('schema')::text, current_schema()) + AND c.relkind = 'i' + AND left(c.relname, length(@index::text)) = @index::text +) +SELECT index_name +FROM index_artifacts +WHERE suffix ~ '^_cc(new|old)[0-9]*$' +ORDER BY index_name; + +-- name: IndexesExist :many +WITH index_names AS ( + SELECT unnest(@index_names::text[]) as index_name +) +SELECT index_names.index_name::text AS index_name, + EXISTS ( + SELECT 1 + FROM pg_catalog.pg_class c + JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace + WHERE n.nspname = coalesce(sqlc.narg('schema')::text, current_schema()) + AND c.relname = index_names.index_name + AND c.relkind = 'i' + ) AS exists +FROM index_names; + +-- name: SchemaGetExpired :many +SELECT schema_name::text +FROM information_schema.schemata +WHERE schema_name LIKE @prefix + AND schema_name < @before_name +ORDER BY schema_name; + +-- name: TableExists :one +SELECT CASE WHEN to_regclass(@schema_and_table) IS NULL THEN false + ELSE true END; diff --git a/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/schema.sql.go b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/schema.sql.go new file mode 100644 index 0000000000..2500099743 --- /dev/null +++ b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/schema.sql.go @@ -0,0 +1,190 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.31.0 +// source: schema.sql + +package dbsqlc + +import ( + "context" + + "github.com/jackc/pgx/v5/pgtype" +) + +const columnExists = `-- name: ColumnExists :one +SELECT EXISTS ( + SELECT column_name + FROM information_schema.columns + WHERE table_name = $1::text + AND table_schema = /* TEMPLATE_BEGIN: schema */ CURRENT_SCHEMA /* TEMPLATE_END */ + AND column_name = $2::text +) +` + +type ColumnExistsParams struct { + TableName string + ColumnName string +} + +func (q *Queries) ColumnExists(ctx context.Context, db DBTX, arg *ColumnExistsParams) (bool, error) { + row := db.QueryRow(ctx, columnExists, arg.TableName, arg.ColumnName) + var exists bool + err := row.Scan(&exists) + return exists, err +} + +const indexExists = `-- name: IndexExists :one +SELECT EXISTS ( + SELECT 1 + FROM pg_catalog.pg_class + JOIN pg_catalog.pg_namespace ON pg_namespace.oid = pg_class.relnamespace + WHERE pg_class.relname = $1::text + AND pg_namespace.nspname = coalesce($2::text, current_schema()) + AND pg_class.relkind = 'i' +) +` + +type IndexExistsParams struct { + Index string + Schema pgtype.Text +} + +func (q *Queries) IndexExists(ctx context.Context, db DBTX, arg *IndexExistsParams) (bool, error) { + row := db.QueryRow(ctx, indexExists, arg.Index, arg.Schema) + var exists bool + err := row.Scan(&exists) + return exists, err +} + +const indexReindexArtifacts = `-- name: IndexReindexArtifacts :many +WITH index_artifacts AS ( + SELECT + c.relname::text AS index_name, + substring(c.relname FROM length($1::text) + 1) AS suffix + FROM pg_catalog.pg_class c + JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace + WHERE n.nspname = coalesce($2::text, current_schema()) + AND c.relkind = 'i' + AND left(c.relname, length($1::text)) = $1::text +) +SELECT index_name +FROM index_artifacts +WHERE suffix ~ '^_cc(new|old)[0-9]*$' +ORDER BY index_name +` + +type IndexReindexArtifactsParams struct { + Index string + Schema pgtype.Text +} + +func (q *Queries) IndexReindexArtifacts(ctx context.Context, db DBTX, arg *IndexReindexArtifactsParams) ([]string, error) { + rows, err := db.Query(ctx, indexReindexArtifacts, arg.Index, arg.Schema) + if err != nil { + return nil, err + } + defer rows.Close() + var items []string + for rows.Next() { + var index_name string + if err := rows.Scan(&index_name); err != nil { + return nil, err + } + items = append(items, index_name) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const indexesExist = `-- name: IndexesExist :many +WITH index_names AS ( + SELECT unnest($2::text[]) as index_name +) +SELECT index_names.index_name::text AS index_name, + EXISTS ( + SELECT 1 + FROM pg_catalog.pg_class c + JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace + WHERE n.nspname = coalesce($1::text, current_schema()) + AND c.relname = index_names.index_name + AND c.relkind = 'i' + ) AS exists +FROM index_names +` + +type IndexesExistParams struct { + Schema pgtype.Text + IndexNames []string +} + +type IndexesExistRow struct { + IndexName string + Exists bool +} + +func (q *Queries) IndexesExist(ctx context.Context, db DBTX, arg *IndexesExistParams) ([]*IndexesExistRow, error) { + rows, err := db.Query(ctx, indexesExist, arg.Schema, arg.IndexNames) + if err != nil { + return nil, err + } + defer rows.Close() + var items []*IndexesExistRow + for rows.Next() { + var i IndexesExistRow + if err := rows.Scan(&i.IndexName, &i.Exists); err != nil { + return nil, err + } + items = append(items, &i) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const schemaGetExpired = `-- name: SchemaGetExpired :many +SELECT schema_name::text +FROM information_schema.schemata +WHERE schema_name LIKE $1 + AND schema_name < $2 +ORDER BY schema_name +` + +type SchemaGetExpiredParams struct { + Prefix interface{} + BeforeName interface{} +} + +func (q *Queries) SchemaGetExpired(ctx context.Context, db DBTX, arg *SchemaGetExpiredParams) ([]string, error) { + rows, err := db.Query(ctx, schemaGetExpired, arg.Prefix, arg.BeforeName) + if err != nil { + return nil, err + } + defer rows.Close() + var items []string + for rows.Next() { + var schema_name string + if err := rows.Scan(&schema_name); err != nil { + return nil, err + } + items = append(items, schema_name) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const tableExists = `-- name: TableExists :one +SELECT CASE WHEN to_regclass($1) IS NULL THEN false + ELSE true END +` + +func (q *Queries) TableExists(ctx context.Context, db DBTX, schemaAndTable string) (bool, error) { + row := db.QueryRow(ctx, tableExists, schemaAndTable) + var column_1 bool + err := row.Scan(&column_1) + return column_1, err +} diff --git a/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/sqlc.yaml b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/sqlc.yaml new file mode 100644 index 0000000000..464423fe45 --- /dev/null +++ b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc/sqlc.yaml @@ -0,0 +1,46 @@ +version: "2" +sql: + - engine: "postgresql" + queries: + - pg_misc.sql + - river_job.sql + - river_job_copyfrom.sql + - river_leader.sql + - river_migration.sql + - river_notification.sql + - river_queue.sql + - schema.sql + schema: + - pg_misc.sql + - river_job.sql + - river_leader.sql + - river_migration.sql + - river_notification.sql + - river_queue.sql + - schema.sql + gen: + go: + package: "dbsqlc" + sql_package: "pgx/v5" + out: "." + emit_exact_table_names: true + emit_methods_with_db_argument: true + emit_params_struct_pointers: true + emit_result_struct_pointers: true + + rename: + ids: "IDs" + ttl: "TTL" + + overrides: + - db_type: "pg_catalog.interval" + go_type: "time.Duration" + + - db_type: "timestamptz" + go_type: "time.Time" + + - db_type: "timestamptz" + go_type: + type: "time.Time" + pointer: true + nullable: true diff --git a/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/json_text_mode_adaptation.go b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/json_text_mode_adaptation.go new file mode 100644 index 0000000000..13c0113dd1 --- /dev/null +++ b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/json_text_mode_adaptation.go @@ -0,0 +1,276 @@ +package riverpgxv5 + +import ( + "context" + "encoding/json" + "regexp" + "strconv" + "strings" + + "github.com/jackc/pgx/v5" +) + +// River commonly provides marshaled JSON to sqlc/pgx query inputs as +// `[]byte` for fast extended-protocol paths. In pgx text execution modes +// (simple protocol and exec), `[]byte` is encoded as `bytea`, which makes +// Postgres reject JSON/JSONB parameters with invalid JSON syntax errors. +// +// This adapter rewrites JSON-like `[]byte` and `[][]byte` args to JSON-aware +// types only in those text modes, while leaving normal extended-protocol +// behavior untouched. It uses explicit `::json`/`::jsonb` casts where +// available, plus a guarded fallback for uncast generated SQL. Args explicitly +// cast to `::bytea` are protected so intentional binary parameters are not +// changed. +// +// Query option parsing mirrors pgx's "options before first bind arg" behavior +// so per-query `QueryExecMode` overrides are respected. When a +// `QueryRewriter` is present, the driver wraps it so JSON adaptation runs after +// rewrite against the final SQL/args. + +var ( + jsonCastPlaceholderRegexp = regexp.MustCompile(`(?i)\$([0-9]+)\s*::\s*jsonb?\s*(\[\s*\])?`) + byteaTypecastPlaceholderRegexp = regexp.MustCompile(`(?i)\$([0-9]+)\s*::\s*bytea\s*(\[\s*\])?`) + byteaCastFunctionPlaceholderRegexp = regexp.MustCompile(`(?i)cast\s*\(\s*\$([0-9]+)\s+as\s+bytea\s*(\[\s*\])?\s*\)`) +) + +type jsonPlaceholderCast struct { + argIndex int + isArray bool +} + +func jsonPlaceholderCasts(sql string) []jsonPlaceholderCast { + matches := jsonCastPlaceholderRegexp.FindAllStringSubmatch(sql, -1) + casts := make([]jsonPlaceholderCast, 0, len(matches)) + seen := make(map[int]int, len(matches)) + + for _, match := range matches { + if len(match) < 3 { + continue + } + + placeholderNum, err := strconv.Atoi(match[1]) + if err != nil || placeholderNum < 1 { + continue + } + + cast := jsonPlaceholderCast{ + argIndex: placeholderNum - 1, + isArray: strings.TrimSpace(match[2]) != "", + } + + if priorIndex, found := seen[cast.argIndex]; found { + if cast.isArray { + casts[priorIndex].isArray = true + } + continue + } + + seen[cast.argIndex] = len(casts) + casts = append(casts, cast) + } + + return casts +} + +func adaptArgsForJSONTextModes(defaultMode pgx.QueryExecMode, sql string, args []any) []any { + queryOptions := parseQueryOptions(defaultMode, args) + if !isJSONTextMode(queryOptions.mode) { + return args + } + + // QueryRewriter can rewrite both SQL and args. Wrap it so JSON adaptation + // runs after rewrite against the final bind arguments. + if queryOptions.queryRewriterIndex >= 0 { + return wrapQueryRewriterForJSONTextMode(args, queryOptions.queryRewriterIndex, queryOptions.mode) + } + + return adaptBindArgsForJSONTextMode(sql, args, queryOptions.bindArgStart) +} + +func adaptBindArgsForJSONTextMode(sql string, args []any, bindArgStart int) []any { + casts := jsonPlaceholderCasts(sql) + if len(casts) == 0 { + casts = nil + } + + byteaArgIndices := byteaPlaceholderArgIndices(sql) + var updatedArgs []any + adaptedArgs := make(map[int]struct{}, len(casts)) + for _, cast := range casts { + argIndex := bindArgStart + cast.argIndex + if argIndex >= len(args) { + continue + } + + updatedArg, changed := adaptArgForJSONTextMode(cast, args[argIndex]) + if !changed { + continue + } + + updatedArgs = ensureMutableArgsCopy(args, updatedArgs) + updatedArgs[argIndex] = updatedArg + adaptedArgs[cast.argIndex] = struct{}{} + } + + // Caveat: some generated SQL leaves JSON columns uncast in VALUES/SET lists. + // In simple/exec modes, pgx assumes []byte is bytea, so these would fail. + // + // We adapt remaining []byte/[][]byte arguments unless the placeholder is + // explicitly cast to bytea. New SQL that intentionally expects binary data + // should always use an explicit bytea cast (`::bytea` or CAST(... AS bytea)). + for i := bindArgStart; i < len(args); i++ { + logicalIndex := i - bindArgStart + if _, isBytea := byteaArgIndices[logicalIndex]; isBytea { + continue + } + if _, alreadyAdapted := adaptedArgs[logicalIndex]; alreadyAdapted { + continue + } + + updatedArg, changed := adaptArgForJSONTextMode(jsonPlaceholderCast{isArray: false}, args[i]) + if !changed { + updatedArg, changed = adaptArgForJSONTextMode(jsonPlaceholderCast{isArray: true}, args[i]) + if !changed { + continue + } + } + + updatedArgs = ensureMutableArgsCopy(args, updatedArgs) + updatedArgs[i] = updatedArg + } + + if updatedArgs != nil { + return updatedArgs + } + return args +} + +func wrapQueryRewriterForJSONTextMode(args []any, queryRewriterIndex int, mode pgx.QueryExecMode) []any { + queryRewriter := args[queryRewriterIndex].(pgx.QueryRewriter) //nolint:forcetypeassert + if existingWrapper, ok := queryRewriter.(jsonTextModeAdaptingQueryRewriter); ok && existingWrapper.mode == mode { + return args + } + + updatedArgs := append([]any(nil), args...) + updatedArgs[queryRewriterIndex] = jsonTextModeAdaptingQueryRewriter{ + mode: mode, + inner: queryRewriter, + } + return updatedArgs +} + +type jsonTextModeAdaptingQueryRewriter struct { + mode pgx.QueryExecMode + inner pgx.QueryRewriter +} + +func (r jsonTextModeAdaptingQueryRewriter) RewriteQuery(ctx context.Context, conn *pgx.Conn, sql string, args []any) (string, []any, error) { + sql, args, err := r.inner.RewriteQuery(ctx, conn, sql, args) + if err != nil { + return "", nil, err + } + if !isJSONTextMode(r.mode) { + return sql, args, nil + } + return sql, adaptBindArgsForJSONTextMode(sql, args, 0), nil +} + +func isJSONTextMode(mode pgx.QueryExecMode) bool { + return mode == pgx.QueryExecModeSimpleProtocol || mode == pgx.QueryExecModeExec +} + +type queryOptions struct { + mode pgx.QueryExecMode + bindArgStart int + queryRewriterIndex int +} + +func parseQueryOptions(defaultMode pgx.QueryExecMode, args []any) queryOptions { + opts := queryOptions{ + mode: defaultMode, + queryRewriterIndex: -1, + } + + // pgx query options (including per-query QueryExecMode) are only recognized + // before the first bind argument. We mirror that parsing boundary here. + for i := range args { + switch arg := args[i].(type) { + case pgx.QueryResultFormats, pgx.QueryResultFormatsByOID: + continue + case pgx.QueryExecMode: + opts.mode = arg + case pgx.QueryRewriter: + opts.queryRewriterIndex = i + default: + opts.bindArgStart = i + return opts + } + } + + opts.bindArgStart = len(args) + return opts +} + +func ensureMutableArgsCopy(args, updatedArgs []any) []any { + if updatedArgs != nil { + return updatedArgs + } + return append([]any(nil), args...) +} + +func adaptArgForJSONTextMode(cast jsonPlaceholderCast, arg any) (any, bool) { + if cast.isArray { + switch arg := arg.(type) { + case [][]byte: + if arg == nil { + return []json.RawMessage(nil), true + } + out := make([]json.RawMessage, len(arg)) + for i := range arg { + out[i] = json.RawMessage(arg[i]) + } + return out, true + case []json.RawMessage: + return arg, false + default: + return arg, false + } + } + + switch arg := arg.(type) { + case []byte: + return json.RawMessage(arg), true + case json.RawMessage: + return arg, false + default: + return arg, false + } +} + +func byteaPlaceholderArgIndices(sql string) map[int]struct{} { + typecastMatches := byteaTypecastPlaceholderRegexp.FindAllStringSubmatch(sql, -1) + castFunctionMatches := byteaCastFunctionPlaceholderRegexp.FindAllStringSubmatch(sql, -1) + if len(typecastMatches) == 0 && len(castFunctionMatches) == 0 { + return nil + } + + argIndices := make(map[int]struct{}, len(typecastMatches)+len(castFunctionMatches)) + addPlaceholderArgIndices(typecastMatches, argIndices) + addPlaceholderArgIndices(castFunctionMatches, argIndices) + + return argIndices +} + +func addPlaceholderArgIndices(matches [][]string, argIndices map[int]struct{}) { + for _, match := range matches { + if len(match) < 2 { + continue + } + + placeholderNum, err := strconv.Atoi(match[1]) + if err != nil || placeholderNum < 1 { + continue + } + argIndices[placeholderNum-1] = struct{}{} + } +} diff --git a/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/migration/main/001_create_river_migration.down.sql b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/migration/main/001_create_river_migration.down.sql new file mode 100644 index 0000000000..8bfe820276 --- /dev/null +++ b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/migration/main/001_create_river_migration.down.sql @@ -0,0 +1 @@ +DROP TABLE /* TEMPLATE: schema */river_migration; \ No newline at end of file diff --git a/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/migration/main/001_create_river_migration.up.sql b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/migration/main/001_create_river_migration.up.sql new file mode 100644 index 0000000000..27006d5626 --- /dev/null +++ b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/migration/main/001_create_river_migration.up.sql @@ -0,0 +1,8 @@ +CREATE TABLE /* TEMPLATE: schema */river_migration( + id bigserial PRIMARY KEY, + created_at timestamptz NOT NULL DEFAULT NOW(), + version bigint NOT NULL, + CONSTRAINT version CHECK (version >= 1) +); + +CREATE UNIQUE INDEX ON /* TEMPLATE: schema */river_migration USING btree(version); \ No newline at end of file diff --git a/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/migration/main/002_initial_schema.down.sql b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/migration/main/002_initial_schema.down.sql new file mode 100644 index 0000000000..d334d8a65a --- /dev/null +++ b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/migration/main/002_initial_schema.down.sql @@ -0,0 +1,5 @@ +DROP TABLE /* TEMPLATE: schema */river_job; +DROP FUNCTION /* TEMPLATE: schema */river_job_notify; +DROP TYPE /* TEMPLATE: schema */river_job_state; + +DROP TABLE /* TEMPLATE: schema */river_leader; \ No newline at end of file diff --git a/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/migration/main/002_initial_schema.up.sql b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/migration/main/002_initial_schema.up.sql new file mode 100644 index 0000000000..7fbca71b41 --- /dev/null +++ b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/migration/main/002_initial_schema.up.sql @@ -0,0 +1,96 @@ +CREATE TYPE /* TEMPLATE: schema */river_job_state AS ENUM( + 'available', + 'cancelled', + 'completed', + 'discarded', + 'retryable', + 'running', + 'scheduled' +); + +CREATE TABLE /* TEMPLATE: schema */river_job( + -- 8 bytes + id bigserial PRIMARY KEY, + + -- 8 bytes (4 bytes + 2 bytes + 2 bytes) + -- + -- `state` is kept near the top of the table for operator convenience -- when + -- looking at jobs with `SELECT *` it'll appear first after ID. The other two + -- fields aren't as important but are kept adjacent to `state` for alignment + -- to get an 8-byte block. + state /* TEMPLATE: schema */river_job_state NOT NULL DEFAULT 'available', + attempt smallint NOT NULL DEFAULT 0, + max_attempts smallint NOT NULL, + + -- 8 bytes each (no alignment needed) + attempted_at timestamptz, + created_at timestamptz NOT NULL DEFAULT NOW(), + finalized_at timestamptz, + scheduled_at timestamptz NOT NULL DEFAULT NOW(), + + -- 2 bytes (some wasted padding probably) + priority smallint NOT NULL DEFAULT 1, + + -- types stored out-of-band + args jsonb, + attempted_by text[], + errors jsonb[], + kind text NOT NULL, + metadata jsonb NOT NULL DEFAULT '{}', + queue text NOT NULL DEFAULT 'default', + tags varchar(255)[], + + CONSTRAINT finalized_or_finalized_at_null CHECK ((state IN ('cancelled', 'completed', 'discarded') AND finalized_at IS NOT NULL) OR finalized_at IS NULL), + CONSTRAINT max_attempts_is_positive CHECK (max_attempts > 0), + CONSTRAINT priority_in_range CHECK (priority >= 1 AND priority <= 4), + CONSTRAINT queue_length CHECK (char_length(queue) > 0 AND char_length(queue) < 128), + CONSTRAINT kind_length CHECK (char_length(kind) > 0 AND char_length(kind) < 128) +); + +-- We may want to consider adding another property here after `kind` if it seems +-- like it'd be useful for something. +CREATE INDEX river_job_kind ON /* TEMPLATE: schema */river_job USING btree(kind); + +CREATE INDEX river_job_state_and_finalized_at_index ON /* TEMPLATE: schema */river_job USING btree(state, finalized_at) WHERE finalized_at IS NOT NULL; + +CREATE INDEX river_job_prioritized_fetching_index ON /* TEMPLATE: schema */river_job USING btree(state, queue, priority, scheduled_at, id); + +CREATE INDEX river_job_args_index ON /* TEMPLATE: schema */river_job USING GIN(args); + +CREATE INDEX river_job_metadata_index ON /* TEMPLATE: schema */river_job USING GIN(metadata); + +CREATE OR REPLACE FUNCTION /* TEMPLATE: schema */river_job_notify() + RETURNS TRIGGER + AS $$ +DECLARE + payload json; +BEGIN + IF NEW.state = 'available' THEN + -- Notify will coalesce duplicate notifications within a transaction, so + -- keep these payloads generalized: + payload = json_build_object('queue', NEW.queue); + PERFORM + pg_notify('river_insert', payload::text); + END IF; + RETURN NULL; +END; +$$ +LANGUAGE plpgsql; + +CREATE TRIGGER river_notify + AFTER INSERT ON /* TEMPLATE: schema */river_job + FOR EACH ROW + EXECUTE PROCEDURE /* TEMPLATE: schema */river_job_notify(); + +CREATE UNLOGGED TABLE /* TEMPLATE: schema */river_leader( + -- 8 bytes each (no alignment needed) + elected_at timestamptz NOT NULL, + expires_at timestamptz NOT NULL, + + -- types stored out-of-band + leader_id text NOT NULL, + name text PRIMARY KEY, + + CONSTRAINT name_length CHECK (char_length(name) > 0 AND char_length(name) < 128), + CONSTRAINT leader_id_length CHECK (char_length(leader_id) > 0 AND char_length(leader_id) < 128) +); diff --git a/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/migration/main/003_river_job_tags_non_null.down.sql b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/migration/main/003_river_job_tags_non_null.down.sql new file mode 100644 index 0000000000..acef65cb94 --- /dev/null +++ b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/migration/main/003_river_job_tags_non_null.down.sql @@ -0,0 +1,3 @@ +ALTER TABLE /* TEMPLATE: schema */river_job + ALTER COLUMN tags DROP NOT NULL, + ALTER COLUMN tags DROP DEFAULT; diff --git a/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/migration/main/003_river_job_tags_non_null.up.sql b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/migration/main/003_river_job_tags_non_null.up.sql new file mode 100644 index 0000000000..0a472dde45 --- /dev/null +++ b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/migration/main/003_river_job_tags_non_null.up.sql @@ -0,0 +1,3 @@ +ALTER TABLE /* TEMPLATE: schema */river_job ALTER COLUMN tags SET DEFAULT '{}'; +UPDATE /* TEMPLATE: schema */river_job SET tags = '{}' WHERE tags IS NULL; +ALTER TABLE /* TEMPLATE: schema */river_job ALTER COLUMN tags SET NOT NULL; diff --git a/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/migration/main/004_pending_and_more.down.sql b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/migration/main/004_pending_and_more.down.sql new file mode 100644 index 0000000000..1b7ec7e842 --- /dev/null +++ b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/migration/main/004_pending_and_more.down.sql @@ -0,0 +1,42 @@ +ALTER TABLE /* TEMPLATE: schema */river_job ALTER COLUMN args DROP NOT NULL; + +ALTER TABLE /* TEMPLATE: schema */river_job ALTER COLUMN metadata DROP NOT NULL; +ALTER TABLE /* TEMPLATE: schema */river_job ALTER COLUMN metadata DROP DEFAULT; + +-- It is not possible to safely remove 'pending' from the river_job_state enum, +-- so leave it in place. + +ALTER TABLE /* TEMPLATE: schema */river_job DROP CONSTRAINT finalized_or_finalized_at_null; +ALTER TABLE /* TEMPLATE: schema */river_job ADD CONSTRAINT finalized_or_finalized_at_null CHECK ( + (state IN ('cancelled', 'completed', 'discarded') AND finalized_at IS NOT NULL) OR finalized_at IS NULL +); + +CREATE OR REPLACE FUNCTION /* TEMPLATE: schema */river_job_notify() + RETURNS TRIGGER + AS $$ +DECLARE + payload json; +BEGIN + IF NEW.state = 'available' THEN + -- Notify will coalesce duplicate notifications within a transaction, so + -- keep these payloads generalized: + payload = json_build_object('queue', NEW.queue); + PERFORM + pg_notify('river_insert', payload::text); + END IF; + RETURN NULL; +END; +$$ +LANGUAGE plpgsql; + +CREATE TRIGGER river_notify + AFTER INSERT ON /* TEMPLATE: schema */river_job + FOR EACH ROW + EXECUTE PROCEDURE /* TEMPLATE: schema */river_job_notify(); + +DROP TABLE /* TEMPLATE: schema */river_queue; + +ALTER TABLE /* TEMPLATE: schema */river_leader + ALTER COLUMN name DROP DEFAULT, + DROP CONSTRAINT name_length, + ADD CONSTRAINT name_length CHECK (char_length(name) > 0 AND char_length(name) < 128); \ No newline at end of file diff --git a/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/migration/main/004_pending_and_more.up.sql b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/migration/main/004_pending_and_more.up.sql new file mode 100644 index 0000000000..9f5e47bb1f --- /dev/null +++ b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/migration/main/004_pending_and_more.up.sql @@ -0,0 +1,45 @@ +-- The args column never had a NOT NULL constraint or default value at the +-- database level, though we tried to ensure one at the application level. +ALTER TABLE /* TEMPLATE: schema */river_job ALTER COLUMN args SET DEFAULT '{}'; +UPDATE /* TEMPLATE: schema */river_job SET args = '{}' WHERE args IS NULL; +ALTER TABLE /* TEMPLATE: schema */river_job ALTER COLUMN args SET NOT NULL; +ALTER TABLE /* TEMPLATE: schema */river_job ALTER COLUMN args DROP DEFAULT; + +-- The metadata column never had a NOT NULL constraint or default value at the +-- database level, though we tried to ensure one at the application level. +ALTER TABLE /* TEMPLATE: schema */river_job ALTER COLUMN metadata SET DEFAULT '{}'; +UPDATE /* TEMPLATE: schema */river_job SET metadata = '{}' WHERE metadata IS NULL; +ALTER TABLE /* TEMPLATE: schema */river_job ALTER COLUMN metadata SET NOT NULL; + +-- The 'pending' job state will be used for upcoming functionality: +ALTER TYPE /* TEMPLATE: schema */river_job_state ADD VALUE IF NOT EXISTS 'pending' AFTER 'discarded'; + +ALTER TABLE /* TEMPLATE: schema */river_job DROP CONSTRAINT finalized_or_finalized_at_null; +ALTER TABLE /* TEMPLATE: schema */river_job ADD CONSTRAINT finalized_or_finalized_at_null CHECK ( + (finalized_at IS NULL AND state NOT IN ('cancelled', 'completed', 'discarded')) OR + (finalized_at IS NOT NULL AND state IN ('cancelled', 'completed', 'discarded')) +); + +DROP TRIGGER river_notify ON /* TEMPLATE: schema */river_job; +DROP FUNCTION /* TEMPLATE: schema */river_job_notify; + +-- +-- Create table `river_queue`. +-- + +CREATE TABLE /* TEMPLATE: schema */river_queue ( + name text PRIMARY KEY NOT NULL, + created_at timestamptz NOT NULL DEFAULT now(), + metadata jsonb NOT NULL DEFAULT '{}' ::jsonb, + paused_at timestamptz, + updated_at timestamptz NOT NULL +); + +-- +-- Alter `river_leader` to add a default value of 'default` to `name`. +-- + +ALTER TABLE /* TEMPLATE: schema */river_leader + ALTER COLUMN name SET DEFAULT 'default', + DROP CONSTRAINT name_length, + ADD CONSTRAINT name_length CHECK (name = 'default'); \ No newline at end of file diff --git a/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/migration/main/005_migration_unique_client.down.sql b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/migration/main/005_migration_unique_client.down.sql new file mode 100644 index 0000000000..b8e041d541 --- /dev/null +++ b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/migration/main/005_migration_unique_client.down.sql @@ -0,0 +1,57 @@ +-- +-- Revert to migration table based only on `(version)`. +-- +-- If any non-main migrations are present, 005 is considered irreversible. +-- + +DO +$body$ +BEGIN + -- Tolerate users who may be using their own migration system rather than + -- River's. If they are, they will have skipped version 001 containing + -- `CREATE TABLE river_migration`, so this table won't exist. + IF (SELECT to_regclass('/* TEMPLATE: schema */river_migration') IS NOT NULL) THEN + IF EXISTS ( + SELECT * + FROM /* TEMPLATE: schema */river_migration + WHERE line <> 'main' + ) THEN + RAISE EXCEPTION 'Found non-main migration lines in the database; version 005 migration is irreversible because it would result in loss of migration information.'; + END IF; + + ALTER TABLE /* TEMPLATE: schema */river_migration + RENAME TO river_migration_old; + + CREATE TABLE /* TEMPLATE: schema */river_migration( + id bigserial PRIMARY KEY, + created_at timestamptz NOT NULL DEFAULT NOW(), + version bigint NOT NULL, + CONSTRAINT version CHECK (version >= 1) + ); + + CREATE UNIQUE INDEX ON /* TEMPLATE: schema */river_migration USING btree(version); + + INSERT INTO /* TEMPLATE: schema */river_migration + (created_at, version) + SELECT created_at, version + FROM /* TEMPLATE: schema */river_migration_old; + + DROP TABLE /* TEMPLATE: schema */river_migration_old; + END IF; +END; +$body$ +LANGUAGE 'plpgsql'; + +-- +-- Drop `river_job.unique_key`. +-- + +ALTER TABLE /* TEMPLATE: schema */river_job + DROP COLUMN unique_key; + +-- +-- Drop `river_client` and derivative. +-- + +DROP TABLE /* TEMPLATE: schema */river_client_queue; +DROP TABLE /* TEMPLATE: schema */river_client; diff --git a/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/migration/main/005_migration_unique_client.up.sql b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/migration/main/005_migration_unique_client.up.sql new file mode 100644 index 0000000000..e0f1711ec2 --- /dev/null +++ b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/migration/main/005_migration_unique_client.up.sql @@ -0,0 +1,79 @@ +-- +-- Rebuild the migration table so it's based on `(line, version)`. +-- + +DO +$body$ +BEGIN + -- Tolerate users who may be using their own migration system rather than + -- River's. If they are, they will have skipped version 001 containing + -- `CREATE TABLE river_migration`, so this table won't exist. + IF (SELECT to_regclass('/* TEMPLATE: schema */river_migration') IS NOT NULL) THEN + ALTER TABLE /* TEMPLATE: schema */river_migration + RENAME TO river_migration_old; + + CREATE TABLE /* TEMPLATE: schema */river_migration( + line TEXT NOT NULL, + version bigint NOT NULL, + created_at timestamptz NOT NULL DEFAULT NOW(), + CONSTRAINT line_length CHECK (char_length(line) > 0 AND char_length(line) < 128), + CONSTRAINT version_gte_1 CHECK (version >= 1), + PRIMARY KEY (line, version) + ); + + INSERT INTO /* TEMPLATE: schema */river_migration + (created_at, line, version) + SELECT created_at, 'main', version + FROM /* TEMPLATE: schema */river_migration_old; + + DROP TABLE /* TEMPLATE: schema */river_migration_old; + END IF; +END; +$body$ +LANGUAGE 'plpgsql'; + +-- +-- Add `river_job.unique_key` and bring up an index on it. +-- + +-- These statements use `IF NOT EXISTS` to allow users with a `river_job` table +-- of non-trivial size to build the index `CONCURRENTLY` out of band of this +-- migration, then follow by completing the migration. +ALTER TABLE /* TEMPLATE: schema */river_job + ADD COLUMN IF NOT EXISTS unique_key bytea; + +CREATE UNIQUE INDEX IF NOT EXISTS river_job_kind_unique_key_idx ON /* TEMPLATE: schema */river_job (kind, unique_key) WHERE unique_key IS NOT NULL; + +-- +-- Create `river_client` and derivative. +-- +-- This feature hasn't quite yet been implemented, but we're taking advantage of +-- the migration to add the schema early so that we can add it later without an +-- additional migration. +-- + +CREATE UNLOGGED TABLE /* TEMPLATE: schema */river_client ( + id text PRIMARY KEY NOT NULL, + created_at timestamptz NOT NULL DEFAULT now(), + metadata jsonb NOT NULL DEFAULT '{}', + paused_at timestamptz, + updated_at timestamptz NOT NULL, + CONSTRAINT name_length CHECK (char_length(id) > 0 AND char_length(id) < 128) +); + +-- Differs from `river_queue` in that it tracks the queue state for a particular +-- active client. +CREATE UNLOGGED TABLE /* TEMPLATE: schema */river_client_queue ( + river_client_id text NOT NULL REFERENCES /* TEMPLATE: schema */river_client (id) ON DELETE CASCADE, + name text NOT NULL, + created_at timestamptz NOT NULL DEFAULT now(), + max_workers bigint NOT NULL DEFAULT 0, + metadata jsonb NOT NULL DEFAULT '{}', + num_jobs_completed bigint NOT NULL DEFAULT 0, + num_jobs_running bigint NOT NULL DEFAULT 0, + updated_at timestamptz NOT NULL, + PRIMARY KEY (river_client_id, name), + CONSTRAINT name_length CHECK (char_length(name) > 0 AND char_length(name) < 128), + CONSTRAINT num_jobs_completed_zero_or_positive CHECK (num_jobs_completed >= 0), + CONSTRAINT num_jobs_running_zero_or_positive CHECK (num_jobs_running >= 0) +); \ No newline at end of file diff --git a/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/migration/main/006_bulk_unique.down.sql b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/migration/main/006_bulk_unique.down.sql new file mode 100644 index 0000000000..26cd843451 --- /dev/null +++ b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/migration/main/006_bulk_unique.down.sql @@ -0,0 +1,16 @@ + +-- +-- Drop `river_job.unique_states` and its index. +-- + +DROP INDEX /* TEMPLATE: schema */river_job_unique_idx; + +ALTER TABLE /* TEMPLATE: schema */river_job + DROP COLUMN unique_states; + +CREATE UNIQUE INDEX IF NOT EXISTS river_job_kind_unique_key_idx ON /* TEMPLATE: schema */river_job (kind, unique_key) WHERE unique_key IS NOT NULL; + +-- +-- Drop `river_job_state_in_bitmask` function. +-- +DROP FUNCTION /* TEMPLATE: schema */river_job_state_in_bitmask; diff --git a/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/migration/main/006_bulk_unique.up.sql b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/migration/main/006_bulk_unique.up.sql new file mode 100644 index 0000000000..ef96a19f9e --- /dev/null +++ b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/migration/main/006_bulk_unique.up.sql @@ -0,0 +1,40 @@ +CREATE OR REPLACE FUNCTION /* TEMPLATE: schema */river_job_state_in_bitmask(bitmask BIT(8), state /* TEMPLATE: schema */river_job_state) +RETURNS boolean +LANGUAGE SQL +IMMUTABLE +AS $$ + SELECT CASE state + WHEN 'available' THEN get_bit(bitmask, 7) + WHEN 'cancelled' THEN get_bit(bitmask, 6) + WHEN 'completed' THEN get_bit(bitmask, 5) + WHEN 'discarded' THEN get_bit(bitmask, 4) + WHEN 'pending' THEN get_bit(bitmask, 3) + WHEN 'retryable' THEN get_bit(bitmask, 2) + WHEN 'running' THEN get_bit(bitmask, 1) + WHEN 'scheduled' THEN get_bit(bitmask, 0) + ELSE 0 + END = 1; +$$; + +-- +-- Add `river_job.unique_states` and bring up an index on it. +-- +-- This column may exist already if users manually created the column and index +-- as instructed in the changelog so the index could be created `CONCURRENTLY`. +-- +ALTER TABLE /* TEMPLATE: schema */river_job ADD COLUMN IF NOT EXISTS unique_states BIT(8); + +-- This statement uses `IF NOT EXISTS` to allow users with a `river_job` table +-- of non-trivial size to build the index `CONCURRENTLY` out of band of this +-- migration, then follow by completing the migration. +CREATE UNIQUE INDEX IF NOT EXISTS river_job_unique_idx ON /* TEMPLATE: schema */river_job (unique_key) + WHERE unique_key IS NOT NULL + AND unique_states IS NOT NULL + AND /* TEMPLATE: schema */river_job_state_in_bitmask(unique_states, state); + +-- Remove the old unique index. Users who are actively using the unique jobs +-- feature and who wish to avoid deploy downtime may want od drop this in a +-- subsequent migration once all jobs using the old unique system have been +-- completed (i.e. no more rows with non-null unique_key and null +-- unique_states). +DROP INDEX /* TEMPLATE: schema */river_job_kind_unique_key_idx; diff --git a/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/migration/main/007_notification_outbox_sqlite_jsonb_and_sql_cleanup.down.sql b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/migration/main/007_notification_outbox_sqlite_jsonb_and_sql_cleanup.down.sql new file mode 100644 index 0000000000..bed717f87d --- /dev/null +++ b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/migration/main/007_notification_outbox_sqlite_jsonb_and_sql_cleanup.down.sql @@ -0,0 +1,56 @@ +-- +-- SQL cleanup rollback. +-- + +-- +-- Add back unused tables `river_client` and `river_client_queue`. +-- + +CREATE UNLOGGED TABLE /* TEMPLATE: schema */river_client ( + id text PRIMARY KEY NOT NULL, + created_at timestamptz NOT NULL DEFAULT now(), + metadata jsonb NOT NULL DEFAULT '{}', + paused_at timestamptz, + updated_at timestamptz NOT NULL, + CONSTRAINT name_length CHECK (char_length(id) > 0 AND char_length(id) < 128) +); + +CREATE UNLOGGED TABLE /* TEMPLATE: schema */river_client_queue ( + river_client_id text NOT NULL REFERENCES /* TEMPLATE: schema */river_client (id) ON DELETE CASCADE, + name text NOT NULL, + created_at timestamptz NOT NULL DEFAULT now(), + max_workers bigint NOT NULL DEFAULT 0, + metadata jsonb NOT NULL DEFAULT '{}', + num_jobs_completed bigint NOT NULL DEFAULT 0, + num_jobs_running bigint NOT NULL DEFAULT 0, + updated_at timestamptz NOT NULL, + PRIMARY KEY (river_client_id, name), + CONSTRAINT name_length CHECK (char_length(name) > 0 AND char_length(name) < 128), + CONSTRAINT num_jobs_completed_zero_or_positive CHECK (num_jobs_completed >= 0), + CONSTRAINT num_jobs_running_zero_or_positive CHECK (num_jobs_running >= 0) +); + +-- +-- Revert addition of `DEFAULT 25` to `river_job.max_attempts`. +-- + +ALTER TABLE /* TEMPLATE: schema */river_job + ALTER COLUMN max_attempts DROP DEFAULT; + +-- +-- Changes `river_queue.updated_at` to revert the default of `CURRENT_TIMESTAMP`. +-- + +ALTER TABLE /* TEMPLATE: schema */river_queue + ALTER COLUMN updated_at DROP DEFAULT; + +-- +-- SQLite JSONB conversion rollback. +-- +-- No-op. PostgreSQL already stores River JSON columns as jsonb. + +-- +-- Notification outbox rollback. +-- + +DROP TABLE /* TEMPLATE: schema */river_notification; diff --git a/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/migration/main/007_notification_outbox_sqlite_jsonb_and_sql_cleanup.up.sql b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/migration/main/007_notification_outbox_sqlite_jsonb_and_sql_cleanup.up.sql new file mode 100644 index 0000000000..39e3249c9a --- /dev/null +++ b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/migration/main/007_notification_outbox_sqlite_jsonb_and_sql_cleanup.up.sql @@ -0,0 +1,44 @@ +-- +-- Notification outbox. +-- + +CREATE TABLE /* TEMPLATE: schema */river_notification ( + id bigserial PRIMARY KEY, + created_at timestamptz NOT NULL DEFAULT now(), + payload text NOT NULL, + topic text NOT NULL, + CONSTRAINT topic_length CHECK (length(topic) > 0 AND length(topic) < 128) +); + +CREATE INDEX river_notification_created_at_idx ON /* TEMPLATE: schema */river_notification (created_at); +CREATE INDEX river_notification_topic_id_idx ON /* TEMPLATE: schema */river_notification (topic, id); + +-- +-- SQLite JSONB conversion. +-- +-- No-op. PostgreSQL already stores River JSON columns as jsonb. + +-- +-- SQL cleanup. +-- + +-- +-- Drop unused tables `river_client` and `river_client_queue`. +-- + +DROP TABLE /* TEMPLATE: schema */river_client_queue; +DROP TABLE /* TEMPLATE: schema */river_client; + +-- +-- Adds `DEFAULT 25` to `river_job.max_attempts`. +-- + +ALTER TABLE /* TEMPLATE: schema */river_job + ALTER COLUMN max_attempts SET DEFAULT 25; + +-- +-- Changes `river_queue.updated_at` to have a default of `CURRENT_TIMESTAMP`. +-- + +ALTER TABLE /* TEMPLATE: schema */river_queue + ALTER COLUMN updated_at SET DEFAULT CURRENT_TIMESTAMP; diff --git a/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/river_pgx_v5_driver.go b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/river_pgx_v5_driver.go new file mode 100644 index 0000000000..71dcc6d55e --- /dev/null +++ b/vendor/github.com/riverqueue/river/riverdriver/riverpgxv5/river_pgx_v5_driver.go @@ -0,0 +1,1340 @@ +// Package riverpgxv5 provides a River driver implementation for Pgx v5. +// +// This is currently the only supported driver for River and will therefore be +// used by all projects using River, but the code is organized this way so that +// other database packages can be supported in future River versions. +package riverpgxv5 + +import ( + "cmp" + "context" + "embed" + "encoding/json" + "errors" + "fmt" + "io/fs" + "math" + "strings" + "sync" + "time" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgconn" + "github.com/jackc/pgx/v5/pgtype" + "github.com/jackc/pgx/v5/pgxpool" + "github.com/jackc/puddle/v2" + + "github.com/riverqueue/river/riverdriver" + "github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc" + "github.com/riverqueue/river/rivershared/sqlctemplate" + "github.com/riverqueue/river/rivershared/uniquestates" + "github.com/riverqueue/river/rivershared/util/dbutil" + "github.com/riverqueue/river/rivershared/util/ptrutil" + "github.com/riverqueue/river/rivershared/util/sliceutil" + "github.com/riverqueue/river/rivertype" +) + +//go:embed migration/*/*.sql +var migrationFS embed.FS + +// Driver is an implementation of riverdriver.Driver for Pgx v5. +type Driver struct { + dbPool *pgxpool.Pool + replacer sqlctemplate.Replacer +} + +// New returns a new Pgx v5 River driver for use with River. +// +// It takes a pgxpool.Pool to use for use with River. The pool should already be +// configured to use the schema specified in the client's Schema field. The pool +// must not be closed while associated River objects are running. +// +// The database pool may be nil. If it is, a client that it's sent into will not +// be able to start up (calls to Start will error) and the Insert and InsertMany +// functions will be disabled, but the transactional-variants InsertTx and +// InsertManyTx continue to function. This behavior may be particularly useful +// in testing so that inserts can be performed and verified on a test +// transaction that will be rolled back. +func New(dbPool *pgxpool.Pool) *Driver { + return &Driver{ + dbPool: dbPool, + } +} + +const argPlaceholder = "$" + +func (d *Driver) ArgPlaceholder() string { return argPlaceholder } +func (d *Driver) DatabaseName() string { return riverdriver.DatabaseNamePostgres } + +func (d *Driver) GetExecutor() riverdriver.Executor { + return &Executor{templateReplaceWrapper{d.dbPool, &d.replacer}, d} +} + +func (d *Driver) GetListener(params *riverdriver.GetListenenerParams) riverdriver.Listener { + return &Listener{dbPool: d.dbPool, schema: params.Schema} +} + +func (d *Driver) GetMigrationDefaultLines() []string { return []string{riverdriver.MigrationLineMain} } +func (d *Driver) GetMigrationFS(line string) fs.FS { + if line == riverdriver.MigrationLineMain { + return migrationFS + } + panic("migration line does not exist: " + line) +} +func (d *Driver) GetMigrationLines() []string { return []string{riverdriver.MigrationLineMain} } +func (d *Driver) GetMigrationTruncateTables(line string, version int) []string { + if line == riverdriver.MigrationLineMain { + return riverdriver.MigrationLineMainTruncateTables(version) + } + panic("migration line does not exist: " + line) +} + +func (d *Driver) PoolIsSet() bool { return d.dbPool != nil } +func (d *Driver) PoolSet(dbPool any) error { return riverdriver.ErrNotImplemented } + +func (d *Driver) SQLFragmentColumnContainsAll(column, namedArg string, values []string) (string, any, error) { + return fmt.Sprintf("%s @> @%s", column, namedArg), values, nil +} + +func (d *Driver) SQLFragmentColumnContainsAny(column, namedArg string, values []string) (string, any, error) { + return fmt.Sprintf("%s && @%s", column, namedArg), values, nil +} + +func (d *Driver) SQLFragmentColumnIn(column string, values any) (string, any, error) { + return fmt.Sprintf("%s = any(@%s)", column, column), values, nil +} + +func (d *Driver) SupportsListener() bool { return true } +func (d *Driver) SupportsListenNotify() bool { return true } +func (d *Driver) TimePrecision() time.Duration { return time.Microsecond } + +func (d *Driver) UnwrapExecutor(tx pgx.Tx) riverdriver.ExecutorTx { + // Allows UnwrapExecutor to be invoked even if driver is nil. + var replacer *sqlctemplate.Replacer + if d == nil { + replacer = &sqlctemplate.Replacer{} + } else { + replacer = &d.replacer + } + + return &ExecutorTx{Executor: Executor{templateReplaceWrapper{tx, replacer}, d}, tx: tx} +} + +func (d *Driver) UnwrapTx(execTx riverdriver.ExecutorTx) pgx.Tx { return execTx.(*ExecutorTx).tx } //nolint:forcetypeassert + +type Executor struct { + dbtx templateReplaceWrapper + driver *Driver +} + +func (e *Executor) Begin(ctx context.Context) (riverdriver.ExecutorTx, error) { + tx, err := e.dbtx.Begin(ctx) + if err != nil { + return nil, err + } + return &ExecutorTx{Executor: Executor{templateReplaceWrapper{tx, &e.driver.replacer}, e.driver}, tx: tx}, nil +} + +func (e *Executor) ColumnExists(ctx context.Context, params *riverdriver.ColumnExistsParams) (bool, error) { + // Schema injection is a bit different on this one because we're querying a table with a schema name. + schema := "CURRENT_SCHEMA" + if params.Schema != "" { + schema = "'" + params.Schema + "'" + } + ctx = sqlctemplate.WithReplacements(ctx, map[string]sqlctemplate.Replacement{ + "schema": {Value: schema}, + }, nil) + + exists, err := dbsqlc.New().ColumnExists(ctx, e.dbtx, &dbsqlc.ColumnExistsParams{ + ColumnName: params.Column, + TableName: params.Table, + }) + return exists, interpretError(err) +} + +func (e *Executor) Exec(ctx context.Context, sql string, args ...any) error { + _, err := e.dbtx.Exec(ctx, sql, args...) + return interpretError(err) +} + +func (e *Executor) IndexDropIfExists(ctx context.Context, params *riverdriver.IndexDropIfExistsParams) error { + var maybeSchema string + if params.Schema != "" { + maybeSchema = dbutil.SafeIdentifier(params.Schema) + "." + } + + _, err := e.dbtx.Exec(ctx, "DROP INDEX CONCURRENTLY IF EXISTS "+maybeSchema+dbutil.SafeIdentifier(params.Index)) + return interpretError(err) +} + +func (e *Executor) IndexExists(ctx context.Context, params *riverdriver.IndexExistsParams) (bool, error) { + exists, err := dbsqlc.New().IndexExists(ctx, e.dbtx, &dbsqlc.IndexExistsParams{ + Index: params.Index, + Schema: pgtype.Text{String: params.Schema, Valid: params.Schema != ""}, + }) + if err != nil { + return false, interpretError(err) + } + return exists, nil +} + +func (e *Executor) IndexReindex(ctx context.Context, params *riverdriver.IndexReindexParams) error { + var maybeSchema string + if params.Schema != "" { + maybeSchema = dbutil.SafeIdentifier(params.Schema) + "." + } + + _, err := e.dbtx.Exec(ctx, "REINDEX INDEX CONCURRENTLY "+maybeSchema+dbutil.SafeIdentifier(params.Index)) + return interpretError(err) +} + +func (e *Executor) IndexReindexArtifacts(ctx context.Context, params *riverdriver.IndexReindexArtifactsParams) ([]string, error) { + artifacts, err := dbsqlc.New().IndexReindexArtifacts(ctx, e.dbtx, &dbsqlc.IndexReindexArtifactsParams{ + Index: params.Index, + Schema: pgtype.Text{String: params.Schema, Valid: params.Schema != ""}, + }) + if err != nil { + return nil, interpretError(err) + } + return artifacts, nil +} + +func (e *Executor) IndexesExist(ctx context.Context, params *riverdriver.IndexesExistParams) (map[string]bool, error) { + rows, err := dbsqlc.New().IndexesExist(ctx, e.dbtx, &dbsqlc.IndexesExistParams{ + IndexNames: params.IndexNames, + Schema: pgtype.Text{String: params.Schema, Valid: params.Schema != ""}, + }) + if err != nil { + return nil, interpretError(err) + } + + exists := make(map[string]bool) + for _, row := range rows { + exists[row.IndexName] = row.Exists + } + return exists, nil +} + +func (e *Executor) JobCancel(ctx context.Context, params *riverdriver.JobCancelParams) (*rivertype.JobRow, error) { + cancelledAt, err := params.CancelAttemptedAt.MarshalJSON() + if err != nil { + return nil, err + } + + job, err := dbsqlc.New().JobCancel(schemaTemplateParam(ctx, params.Schema), e.dbtx, &dbsqlc.JobCancelParams{ + ID: params.ID, + CancelAttemptedAt: cancelledAt, + ControlTopic: params.ControlTopic, + Now: params.Now, + Schema: pgtype.Text{String: params.Schema, Valid: params.Schema != ""}, + }) + if err != nil { + return nil, interpretError(err) + } + return jobRowFromInternal(job) +} + +func (e *Executor) JobCountByAllStates(ctx context.Context, params *riverdriver.JobCountByAllStatesParams) (map[rivertype.JobState]int, error) { + counts, err := dbsqlc.New().JobCountByAllStates(schemaTemplateParam(ctx, params.Schema), e.dbtx) + if err != nil { + return nil, interpretError(err) + } + countsMap := make(map[rivertype.JobState]int) + for _, state := range rivertype.JobStates() { + countsMap[state] = 0 + } + for _, count := range counts { + countsMap[rivertype.JobState(count.State)] = int(count.Count) + } + return countsMap, nil +} + +func (e *Executor) JobCountByQueueAndState(ctx context.Context, params *riverdriver.JobCountByQueueAndStateParams) ([]*riverdriver.JobCountByQueueAndStateResult, error) { + rows, err := dbsqlc.New().JobCountByQueueAndState(schemaTemplateParam(ctx, params.Schema), e.dbtx, params.QueueNames) + if err != nil { + return nil, interpretError(err) + } + results := make([]*riverdriver.JobCountByQueueAndStateResult, len(rows)) + for i, row := range rows { + results[i] = &riverdriver.JobCountByQueueAndStateResult{ + CountAvailable: row.CountAvailable, + CountRunning: row.CountRunning, + Queue: row.Queue, + } + } + return results, nil +} + +func (e *Executor) JobCountByState(ctx context.Context, params *riverdriver.JobCountByStateParams) (int, error) { + numJobs, err := dbsqlc.New().JobCountByState(schemaTemplateParam(ctx, params.Schema), e.dbtx, dbsqlc.RiverJobState(params.State)) + if err != nil { + return 0, err + } + return int(numJobs), nil +} + +func (e *Executor) JobDelete(ctx context.Context, params *riverdriver.JobDeleteParams) (*rivertype.JobRow, error) { + job, err := dbsqlc.New().JobDelete(schemaTemplateParam(ctx, params.Schema), e.dbtx, params.ID) + if err != nil { + return nil, interpretError(err) + } + if job.State == dbsqlc.RiverJobStateRunning { + return nil, rivertype.ErrJobRunning + } + return jobRowFromInternal(job) +} + +func (e *Executor) JobDeleteBefore(ctx context.Context, params *riverdriver.JobDeleteBeforeParams) (int, error) { + res, err := dbsqlc.New().JobDeleteBefore(schemaTemplateParam(ctx, params.Schema), e.dbtx, &dbsqlc.JobDeleteBeforeParams{ + CancelledDoDelete: params.CancelledDoDelete, + CancelledFinalizedAtHorizon: params.CancelledFinalizedAtHorizon, + CompletedDoDelete: params.CompletedDoDelete, + CompletedFinalizedAtHorizon: params.CompletedFinalizedAtHorizon, + DiscardedDoDelete: params.DiscardedDoDelete, + DiscardedFinalizedAtHorizon: params.DiscardedFinalizedAtHorizon, + Max: int64(params.Max), + QueuesExcluded: params.QueuesExcluded, + QueuesIncluded: params.QueuesIncluded, + }) + if err != nil { + return 0, interpretError(err) + } + return int(res.RowsAffected()), nil +} + +func (e *Executor) JobDeleteMany(ctx context.Context, params *riverdriver.JobDeleteManyParams) ([]*rivertype.JobRow, error) { + ctx = sqlctemplate.WithReplacements(ctx, map[string]sqlctemplate.Replacement{ + "order_by_clause": {Value: params.OrderByClause}, + "where_clause": {Value: params.WhereClause}, + }, params.NamedArgs) + + jobs, err := dbsqlc.New().JobDeleteMany(schemaTemplateParam(ctx, params.Schema), e.dbtx, params.Max) + if err != nil { + return nil, interpretError(err) + } + return sliceutil.MapError(jobs, jobRowFromInternal) +} + +func (e *Executor) JobGetAvailable(ctx context.Context, params *riverdriver.JobGetAvailableParams) ([]*rivertype.JobRow, error) { + jobs, err := dbsqlc.New().JobGetAvailable(schemaTemplateParam(ctx, params.Schema), e.dbtx, &dbsqlc.JobGetAvailableParams{ + AttemptedBy: params.ClientID, + MaxAttemptedBy: int32(min(params.MaxAttemptedBy, math.MaxInt32)), //nolint:gosec + MaxToLock: int32(min(params.MaxToLock, math.MaxInt32)), //nolint:gosec + Now: params.Now, + Queue: params.Queue, + }) + if err != nil { + return nil, interpretError(err) + } + return sliceutil.MapError(jobs, jobRowFromInternal) +} + +func (e *Executor) JobGetByID(ctx context.Context, params *riverdriver.JobGetByIDParams) (*rivertype.JobRow, error) { + job, err := dbsqlc.New().JobGetByID(schemaTemplateParam(ctx, params.Schema), e.dbtx, params.ID) + if err != nil { + return nil, interpretError(err) + } + return jobRowFromInternal(job) +} + +func (e *Executor) JobGetByIDMany(ctx context.Context, params *riverdriver.JobGetByIDManyParams) ([]*rivertype.JobRow, error) { + jobs, err := dbsqlc.New().JobGetByIDMany(schemaTemplateParam(ctx, params.Schema), e.dbtx, params.ID) + if err != nil { + return nil, interpretError(err) + } + return sliceutil.MapError(jobs, jobRowFromInternal) +} + +func (e *Executor) JobGetByKindMany(ctx context.Context, params *riverdriver.JobGetByKindManyParams) ([]*rivertype.JobRow, error) { + jobs, err := dbsqlc.New().JobGetByKindMany(schemaTemplateParam(ctx, params.Schema), e.dbtx, params.Kind) + if err != nil { + return nil, interpretError(err) + } + return sliceutil.MapError(jobs, jobRowFromInternal) +} + +func (e *Executor) JobGetStuck(ctx context.Context, params *riverdriver.JobGetStuckParams) ([]*rivertype.JobRow, error) { + jobs, err := dbsqlc.New().JobGetStuck(schemaTemplateParam(ctx, params.Schema), e.dbtx, &dbsqlc.JobGetStuckParams{ + AfterID: params.AfterID, + Max: int32(min(params.Max, math.MaxInt32)), //nolint:gosec + StuckHorizon: params.StuckHorizon, + }) + if err != nil { + return nil, interpretError(err) + } + return sliceutil.MapError(jobs, jobRowFromInternal) +} + +func (e *Executor) JobInsertFastMany(ctx context.Context, params *riverdriver.JobInsertFastManyParams) ([]*riverdriver.JobInsertFastResult, error) { + insertJobsParams := &dbsqlc.JobInsertFastManyParams{ + ID: make([]int64, len(params.Jobs)), + Args: make([][]byte, len(params.Jobs)), + CreatedAt: make([]time.Time, len(params.Jobs)), + Kind: make([]string, len(params.Jobs)), + MaxAttempts: make([]int16, len(params.Jobs)), + Metadata: make([][]byte, len(params.Jobs)), + Priority: make([]int16, len(params.Jobs)), + Queue: make([]string, len(params.Jobs)), + ScheduledAt: make([]time.Time, len(params.Jobs)), + State: make([]string, len(params.Jobs)), + Tags: make([]string, len(params.Jobs)), + UniqueKey: make([][]byte, len(params.Jobs)), + UniqueStates: make([]int32, len(params.Jobs)), + } + now := time.Now().UTC() + for i := range len(params.Jobs) { + params := params.Jobs[i] + + createdAt := now + if params.CreatedAt != nil { + createdAt = *params.CreatedAt + } + + scheduledAt := now + if params.ScheduledAt != nil { + scheduledAt = *params.ScheduledAt + } + + tags := params.Tags + if tags == nil { + tags = []string{} + } + + defaultObject := []byte("{}") + + insertJobsParams.ID[i] = ptrutil.ValOrDefault(params.ID, 0) + insertJobsParams.Args[i] = sliceutil.FirstNonEmpty(params.EncodedArgs, defaultObject) + insertJobsParams.CreatedAt[i] = createdAt + insertJobsParams.Kind[i] = params.Kind + insertJobsParams.MaxAttempts[i] = int16(min(params.MaxAttempts, math.MaxInt16)) //nolint:gosec + insertJobsParams.Metadata[i] = sliceutil.FirstNonEmpty(params.Metadata, defaultObject) + insertJobsParams.Priority[i] = int16(min(params.Priority, math.MaxInt16)) //nolint:gosec + insertJobsParams.Queue[i] = params.Queue + insertJobsParams.ScheduledAt[i] = scheduledAt + insertJobsParams.State[i] = string(params.State) + insertJobsParams.Tags[i] = strings.Join(tags, ",") + insertJobsParams.UniqueKey[i] = sliceutil.FirstNonEmpty(params.UniqueKey) + insertJobsParams.UniqueStates[i] = int32(params.UniqueStates) + } + + items, err := dbsqlc.New().JobInsertFastMany(schemaTemplateParam(ctx, params.Schema), e.dbtx, insertJobsParams) + if err != nil { + return nil, interpretError(err) + } + + return sliceutil.MapError(items, func(row *dbsqlc.JobInsertFastManyRow) (*riverdriver.JobInsertFastResult, error) { + job, err := jobRowFromInternal(&row.RiverJob) + if err != nil { + return nil, err + } + return &riverdriver.JobInsertFastResult{Job: job, UniqueSkippedAsDuplicate: row.UniqueSkippedAsDuplicate}, nil + }) +} + +func (e *Executor) JobInsertFastManyNoReturning(ctx context.Context, params *riverdriver.JobInsertFastManyParams) (int, error) { + insertJobsParams := make([]*dbsqlc.JobInsertFastManyCopyFromParams, len(params.Jobs)) + now := time.Now().UTC() + + for i := range len(params.Jobs) { + params := params.Jobs[i] + + createdAt := now + if params.CreatedAt != nil { + createdAt = *params.CreatedAt + } + + metadata := params.Metadata + if metadata == nil { + metadata = []byte("{}") + } + + scheduledAt := now + if params.ScheduledAt != nil { + scheduledAt = *params.ScheduledAt + } + + tags := params.Tags + if tags == nil { + tags = []string{} + } + + insertJobsParams[i] = &dbsqlc.JobInsertFastManyCopyFromParams{ + Args: params.EncodedArgs, + CreatedAt: createdAt, + Kind: params.Kind, + MaxAttempts: int16(min(params.MaxAttempts, math.MaxInt16)), //nolint:gosec + Metadata: metadata, + Priority: int16(min(params.Priority, math.MaxInt16)), //nolint:gosec + Queue: params.Queue, + ScheduledAt: scheduledAt, + State: dbsqlc.RiverJobState(params.State), + Tags: tags, + UniqueKey: params.UniqueKey, + UniqueStates: pgtype.Bits{Bytes: []byte{params.UniqueStates}, Len: 8, Valid: params.UniqueStates != 0}, + } + } + + numInserted, err := dbsqlc.New().JobInsertFastManyCopyFrom(schemaCopyFrom(ctx, params.Schema), e.dbtx, insertJobsParams) + if err != nil { + return 0, interpretError(err) + } + + return int(numInserted), nil +} + +func (e *Executor) JobInsertFull(ctx context.Context, params *riverdriver.JobInsertFullParams) (*rivertype.JobRow, error) { + job, err := dbsqlc.New().JobInsertFull(schemaTemplateParam(ctx, params.Schema), e.dbtx, &dbsqlc.JobInsertFullParams{ + Attempt: int16(min(params.Attempt, math.MaxInt16)), //nolint:gosec + AttemptedAt: params.AttemptedAt, + AttemptedBy: params.AttemptedBy, + Args: params.EncodedArgs, + CreatedAt: params.CreatedAt, + Errors: params.Errors, + FinalizedAt: params.FinalizedAt, + Kind: params.Kind, + MaxAttempts: int16(min(params.MaxAttempts, math.MaxInt16)), //nolint:gosec + Metadata: params.Metadata, + Priority: int16(min(params.Priority, math.MaxInt16)), //nolint:gosec + Queue: params.Queue, + ScheduledAt: params.ScheduledAt, + State: dbsqlc.RiverJobState(params.State), + Tags: params.Tags, + UniqueKey: string(params.UniqueKey), + UniqueStates: int32(params.UniqueStates), + }) + if err != nil { + return nil, interpretError(err) + } + return jobRowFromInternal(job) +} + +func (e *Executor) JobInsertFullMany(ctx context.Context, params *riverdriver.JobInsertFullManyParams) ([]*rivertype.JobRow, error) { + insertJobsParams := &dbsqlc.JobInsertFullManyParams{ + Args: make([][]byte, len(params.Jobs)), + Attempt: make([]int16, len(params.Jobs)), + AttemptedAt: make([]time.Time, len(params.Jobs)), + CreatedAt: make([]time.Time, len(params.Jobs)), + FinalizedAt: make([]time.Time, len(params.Jobs)), + Kind: make([]string, len(params.Jobs)), + MaxAttempts: make([]int16, len(params.Jobs)), + Metadata: make([][]byte, len(params.Jobs)), + Priority: make([]int16, len(params.Jobs)), + Queue: make([]string, len(params.Jobs)), + ScheduledAt: make([]time.Time, len(params.Jobs)), + State: make([]string, len(params.Jobs)), + Tags: make([]string, len(params.Jobs)), + UniqueKey: make([]string, len(params.Jobs)), + UniqueStates: make([]int32, len(params.Jobs)), + } + now := time.Now().UTC() + + for i := range len(params.Jobs) { + jobParams := params.Jobs[i] + + insertJobsParams.Args[i] = sliceutil.FirstNonEmpty(jobParams.EncodedArgs, []byte("{}")) + insertJobsParams.Attempt[i] = int16(min(jobParams.Attempt, math.MaxInt16)) //nolint:gosec + insertJobsParams.AttemptedAt[i] = ptrutil.ValOrDefault(jobParams.AttemptedAt, time.Time{}) + insertJobsParams.CreatedAt[i] = ptrutil.ValOrDefault(jobParams.CreatedAt, now) + insertJobsParams.FinalizedAt[i] = ptrutil.ValOrDefault(jobParams.FinalizedAt, time.Time{}) + insertJobsParams.Kind[i] = jobParams.Kind + insertJobsParams.MaxAttempts[i] = int16(min(jobParams.MaxAttempts, math.MaxInt16)) //nolint:gosec + insertJobsParams.Metadata[i] = jobParams.Metadata + insertJobsParams.Priority[i] = int16(min(jobParams.Priority, math.MaxInt16)) //nolint:gosec + insertJobsParams.Queue[i] = jobParams.Queue + insertJobsParams.ScheduledAt[i] = ptrutil.ValOrDefault(jobParams.ScheduledAt, now) + insertJobsParams.State[i] = string(jobParams.State) + insertJobsParams.Tags[i] = strings.Join(sliceutil.FirstNonEmpty(jobParams.Tags, []string{}), ",") + insertJobsParams.UniqueKey[i] = string(jobParams.UniqueKey) + insertJobsParams.UniqueStates[i] = int32(jobParams.UniqueStates) + } + + items, err := dbsqlc.New().JobInsertFullMany(schemaTemplateParam(ctx, params.Schema), e.dbtx, insertJobsParams) + if err != nil { + return nil, interpretError(err) + } + + return sliceutil.MapError(items, jobRowFromInternal) +} + +func (e *Executor) JobKindList(ctx context.Context, params *riverdriver.JobKindListParams) ([]string, error) { + kinds, err := dbsqlc.New().JobKindList(schemaTemplateParam(ctx, params.Schema), e.dbtx, &dbsqlc.JobKindListParams{ + After: params.After, + Exclude: params.Exclude, + Match: params.Match, + Max: int32(params.Max), //nolint:gosec + }) + if err != nil { + return nil, interpretError(err) + } + return kinds, nil +} + +func (e *Executor) JobList(ctx context.Context, params *riverdriver.JobListParams) ([]*rivertype.JobRow, error) { + ctx = sqlctemplate.WithReplacements(ctx, map[string]sqlctemplate.Replacement{ + "order_by_clause": {Value: params.OrderByClause}, + "where_clause": {Value: params.WhereClause}, + }, params.NamedArgs) + + jobs, err := dbsqlc.New().JobList(schemaTemplateParam(ctx, params.Schema), e.dbtx, params.Max) + if err != nil { + return nil, interpretError(err) + } + return sliceutil.MapError(jobs, jobRowFromInternal) +} + +func (e *Executor) JobRescueMany(ctx context.Context, params *riverdriver.JobRescueManyParams) (*struct{}, error) { + err := dbsqlc.New().JobRescueMany(schemaTemplateParam(ctx, params.Schema), e.dbtx, &dbsqlc.JobRescueManyParams{ + ID: params.ID, + Error: params.Error, + FinalizedAt: sliceutil.Map(params.FinalizedAt, func(t *time.Time) time.Time { return ptrutil.ValOrDefault(t, time.Time{}) }), + ScheduledAt: params.ScheduledAt, + State: params.State, + }) + if err != nil { + return nil, interpretError(err) + } + return &struct{}{}, nil +} + +func (e *Executor) JobRetry(ctx context.Context, params *riverdriver.JobRetryParams) (*rivertype.JobRow, error) { + job, err := dbsqlc.New().JobRetry(schemaTemplateParam(ctx, params.Schema), e.dbtx, &dbsqlc.JobRetryParams{ + ID: params.ID, + Now: params.Now, + }) + if err != nil { + return nil, interpretError(err) + } + return jobRowFromInternal(job) +} + +func (e *Executor) JobSchedule(ctx context.Context, params *riverdriver.JobScheduleParams) ([]*riverdriver.JobScheduleResult, error) { + scheduleResults, err := dbsqlc.New().JobSchedule(schemaTemplateParam(ctx, params.Schema), e.dbtx, &dbsqlc.JobScheduleParams{ + Max: int64(params.Max), + Now: params.Now, + }) + if err != nil { + return nil, interpretError(err) + } + return sliceutil.MapError(scheduleResults, func(result *dbsqlc.JobScheduleRow) (*riverdriver.JobScheduleResult, error) { + job, err := jobRowFromInternal(&result.RiverJob) + if err != nil { + return nil, err + } + return &riverdriver.JobScheduleResult{ConflictDiscarded: result.ConflictDiscarded, Job: *job}, nil + }) +} + +func (e *Executor) JobSetStateIfRunningMany(ctx context.Context, params *riverdriver.JobSetStateIfRunningManyParams) ([]*rivertype.JobRow, error) { + setStateParams := &dbsqlc.JobSetStateIfRunningManyParams{ + IDs: params.ID, + Attempt: make([]int32, len(params.ID)), + AttemptDoUpdate: make([]bool, len(params.ID)), + Errors: params.ErrData, + ErrorsDoUpdate: make([]bool, len(params.ID)), + FinalizedAt: make([]time.Time, len(params.ID)), + FinalizedAtDoUpdate: make([]bool, len(params.ID)), + MetadataDoMerge: make([]bool, len(params.ID)), + MetadataUpdates: make([][]byte, len(params.ID)), + Now: params.Now, + ScheduledAt: make([]time.Time, len(params.ID)), + ScheduledAtDoUpdate: make([]bool, len(params.ID)), + State: make([]string, len(params.ID)), + } + + for i := range len(params.ID) { + if params.Attempt[i] != nil { + setStateParams.AttemptDoUpdate[i] = true + setStateParams.Attempt[i] = int32(*params.Attempt[i]) //nolint:gosec + } + if params.ErrData[i] != nil { + setStateParams.ErrorsDoUpdate[i] = true + } + if params.FinalizedAt[i] != nil { + setStateParams.FinalizedAtDoUpdate[i] = true + setStateParams.FinalizedAt[i] = *params.FinalizedAt[i] + } + if params.MetadataDoMerge[i] { + setStateParams.MetadataDoMerge[i] = true + setStateParams.MetadataUpdates[i] = params.MetadataUpdates[i] + } + if params.ScheduledAt[i] != nil { + setStateParams.ScheduledAtDoUpdate[i] = true + setStateParams.ScheduledAt[i] = *params.ScheduledAt[i] + } + setStateParams.State[i] = string(params.State[i]) + } + + jobs, err := dbsqlc.New().JobSetStateIfRunningMany(schemaTemplateParam(ctx, params.Schema), e.dbtx, setStateParams) + if err != nil { + return nil, interpretError(err) + } + return sliceutil.MapError(jobs, jobRowFromInternal) +} + +func (e *Executor) JobUpdate(ctx context.Context, params *riverdriver.JobUpdateParams) (*rivertype.JobRow, error) { + metadata := params.Metadata + if metadata == nil { + metadata = []byte("{}") + } + + job, err := dbsqlc.New().JobUpdate(schemaTemplateParam(ctx, params.Schema), e.dbtx, &dbsqlc.JobUpdateParams{ + ID: params.ID, + MetadataDoMerge: params.MetadataDoMerge, + Metadata: metadata, + }) + if err != nil { + return nil, interpretError(err) + } + + return jobRowFromInternal(job) +} + +func (e *Executor) JobUpdateFull(ctx context.Context, params *riverdriver.JobUpdateFullParams) (*rivertype.JobRow, error) { + metadata := params.Metadata + if metadata == nil { + metadata = []byte("{}") + } + + job, err := dbsqlc.New().JobUpdateFull(schemaTemplateParam(ctx, params.Schema), e.dbtx, &dbsqlc.JobUpdateFullParams{ + ID: params.ID, + AttemptedAtDoUpdate: params.AttemptedAtDoUpdate, + Attempt: int16(min(params.Attempt, math.MaxInt16)), //nolint:gosec + AttemptDoUpdate: params.AttemptDoUpdate, + AttemptedAt: params.AttemptedAt, + AttemptedBy: params.AttemptedBy, + AttemptedByDoUpdate: params.AttemptedByDoUpdate, + ErrorsDoUpdate: params.ErrorsDoUpdate, + Errors: params.Errors, + FinalizedAtDoUpdate: params.FinalizedAtDoUpdate, + FinalizedAt: params.FinalizedAt, + MaxAttemptsDoUpdate: params.MaxAttemptsDoUpdate, + MaxAttempts: int16(min(params.MaxAttempts, math.MaxInt16)), //nolint:gosec + MetadataDoUpdate: params.MetadataDoUpdate, + Metadata: metadata, + StateDoUpdate: params.StateDoUpdate, + State: dbsqlc.RiverJobState(cmp.Or(params.State, rivertype.JobStateAvailable)), // can't send empty job state, so provider default value that may not be set + }) + if err != nil { + return nil, interpretError(err) + } + + return jobRowFromInternal(job) +} + +func (e *Executor) LeaderAttemptElect(ctx context.Context, params *riverdriver.LeaderElectParams) (*riverdriver.Leader, error) { + leader, err := dbsqlc.New().LeaderAttemptElect(schemaTemplateParam(ctx, params.Schema), e.dbtx, &dbsqlc.LeaderAttemptElectParams{ + LeaderID: params.LeaderID, + Now: params.Now, + TTL: params.TTL.Seconds(), + }) + if err != nil { + return nil, interpretError(err) + } + return leaderFromInternal(leader), nil +} + +func (e *Executor) LeaderAttemptReelect(ctx context.Context, params *riverdriver.LeaderReelectParams) (*riverdriver.Leader, error) { + leader, err := dbsqlc.New().LeaderAttemptReelect(schemaTemplateParam(ctx, params.Schema), e.dbtx, &dbsqlc.LeaderAttemptReelectParams{ + ElectedAt: params.ElectedAt, + LeaderID: params.LeaderID, + Now: params.Now, + TTL: params.TTL.Seconds(), + }) + if err != nil { + return nil, interpretError(err) + } + return leaderFromInternal(leader), nil +} + +func (e *Executor) LeaderDeleteExpired(ctx context.Context, params *riverdriver.LeaderDeleteExpiredParams) (int, error) { + numDeleted, err := dbsqlc.New().LeaderDeleteExpired(schemaTemplateParam(ctx, params.Schema), e.dbtx, params.Now) + if err != nil { + return 0, interpretError(err) + } + return int(numDeleted), nil +} + +func (e *Executor) LeaderGetElectedLeader(ctx context.Context, params *riverdriver.LeaderGetElectedLeaderParams) (*riverdriver.Leader, error) { + leader, err := dbsqlc.New().LeaderGetElectedLeader(schemaTemplateParam(ctx, params.Schema), e.dbtx) + if err != nil { + return nil, interpretError(err) + } + return leaderFromInternal(leader), nil +} + +func (e *Executor) LeaderInsert(ctx context.Context, params *riverdriver.LeaderInsertParams) (*riverdriver.Leader, error) { + leader, err := dbsqlc.New().LeaderInsert(schemaTemplateParam(ctx, params.Schema), e.dbtx, &dbsqlc.LeaderInsertParams{ + ElectedAt: params.ElectedAt, + ExpiresAt: params.ExpiresAt, + LeaderID: params.LeaderID, + Now: params.Now, + TTL: params.TTL.Seconds(), + }) + if err != nil { + return nil, interpretError(err) + } + return leaderFromInternal(leader), nil +} + +func (e *Executor) LeaderResign(ctx context.Context, params *riverdriver.LeaderResignParams) (bool, error) { + numResigned, err := dbsqlc.New().LeaderResign(schemaTemplateParam(ctx, params.Schema), e.dbtx, &dbsqlc.LeaderResignParams{ + ElectedAt: params.ElectedAt, + LeaderID: params.LeaderID, + LeadershipTopic: params.LeadershipTopic, + Schema: pgtype.Text{String: params.Schema, Valid: params.Schema != ""}, + }) + if err != nil { + return false, interpretError(err) + } + return numResigned > 0, nil +} + +func (e *Executor) MigrationDeleteAssumingMainMany(ctx context.Context, params *riverdriver.MigrationDeleteAssumingMainManyParams) ([]*riverdriver.Migration, error) { + migrations, err := dbsqlc.New().RiverMigrationDeleteAssumingMainMany(schemaTemplateParam(ctx, params.Schema), e.dbtx, + sliceutil.Map(params.Versions, func(v int) int64 { return int64(v) })) + if err != nil { + return nil, interpretError(err) + } + return sliceutil.Map(migrations, func(internal *dbsqlc.RiverMigrationDeleteAssumingMainManyRow) *riverdriver.Migration { + return &riverdriver.Migration{ + CreatedAt: internal.CreatedAt.UTC(), + Line: riverdriver.MigrationLineMain, + Version: int(internal.Version), + } + }), nil +} + +func (e *Executor) MigrationDeleteByLineAndVersionMany(ctx context.Context, params *riverdriver.MigrationDeleteByLineAndVersionManyParams) ([]*riverdriver.Migration, error) { + migrations, err := dbsqlc.New().RiverMigrationDeleteByLineAndVersionMany(schemaTemplateParam(ctx, params.Schema), e.dbtx, &dbsqlc.RiverMigrationDeleteByLineAndVersionManyParams{ + Line: params.Line, + Version: sliceutil.Map(params.Versions, func(v int) int64 { return int64(v) }), + }) + if err != nil { + return nil, interpretError(err) + } + return sliceutil.Map(migrations, migrationFromInternal), nil +} + +func (e *Executor) MigrationGetAllAssumingMain(ctx context.Context, params *riverdriver.MigrationGetAllAssumingMainParams) ([]*riverdriver.Migration, error) { + migrations, err := dbsqlc.New().RiverMigrationGetAllAssumingMain(schemaTemplateParam(ctx, params.Schema), e.dbtx) + if err != nil { + return nil, interpretError(err) + } + return sliceutil.Map(migrations, func(internal *dbsqlc.RiverMigrationGetAllAssumingMainRow) *riverdriver.Migration { + return &riverdriver.Migration{ + CreatedAt: internal.CreatedAt.UTC(), + Line: riverdriver.MigrationLineMain, + Version: int(internal.Version), + } + }), nil +} + +func (e *Executor) MigrationGetByLine(ctx context.Context, params *riverdriver.MigrationGetByLineParams) ([]*riverdriver.Migration, error) { + migrations, err := dbsqlc.New().RiverMigrationGetByLine(schemaTemplateParam(ctx, params.Schema), e.dbtx, params.Line) + if err != nil { + return nil, interpretError(err) + } + return sliceutil.Map(migrations, migrationFromInternal), nil +} + +func (e *Executor) MigrationInsertMany(ctx context.Context, params *riverdriver.MigrationInsertManyParams) ([]*riverdriver.Migration, error) { + migrations, err := dbsqlc.New().RiverMigrationInsertMany(schemaTemplateParam(ctx, params.Schema), e.dbtx, &dbsqlc.RiverMigrationInsertManyParams{ + Line: params.Line, + Version: sliceutil.Map(params.Versions, func(v int) int64 { return int64(v) }), + }) + if err != nil { + return nil, interpretError(err) + } + return sliceutil.Map(migrations, migrationFromInternal), nil +} + +func (e *Executor) MigrationInsertManyAssumingMain(ctx context.Context, params *riverdriver.MigrationInsertManyAssumingMainParams) ([]*riverdriver.Migration, error) { + migrations, err := dbsqlc.New().RiverMigrationInsertManyAssumingMain(schemaTemplateParam(ctx, params.Schema), e.dbtx, + sliceutil.Map(params.Versions, func(v int) int64 { return int64(v) }), + ) + if err != nil { + return nil, interpretError(err) + } + return sliceutil.Map(migrations, func(internal *dbsqlc.RiverMigrationInsertManyAssumingMainRow) *riverdriver.Migration { + return &riverdriver.Migration{ + CreatedAt: internal.CreatedAt.UTC(), + Line: riverdriver.MigrationLineMain, + Version: int(internal.Version), + } + }), nil +} + +func (e *Executor) NotificationDeleteBefore(ctx context.Context, params *riverdriver.NotificationDeleteBeforeParams) (int, error) { + numDeleted, err := dbsqlc.New().NotificationDeleteBefore(schemaTemplateParam(ctx, params.Schema), e.dbtx, params.CreatedAtHorizon) + return int(numDeleted), interpretError(err) +} + +func (e *Executor) NotifyMany(ctx context.Context, params *riverdriver.NotifyManyParams) error { + return dbsqlc.New().PGNotifyMany(ctx, e.dbtx, &dbsqlc.PGNotifyManyParams{ + Payload: params.Payload, + Schema: pgtype.Text{String: params.Schema, Valid: params.Schema != ""}, + Topic: params.Topic, + }) +} + +func (e *Executor) PGAdvisoryXactLock(ctx context.Context, key int64) (*struct{}, error) { + err := dbsqlc.New().PGAdvisoryXactLock(ctx, e.dbtx, key) + return &struct{}{}, interpretError(err) +} + +func (e *Executor) QueueCreateOrSetUpdatedAt(ctx context.Context, params *riverdriver.QueueCreateOrSetUpdatedAtParams) (*rivertype.Queue, error) { + queue, err := dbsqlc.New().QueueCreateOrSetUpdatedAt(schemaTemplateParam(ctx, params.Schema), e.dbtx, &dbsqlc.QueueCreateOrSetUpdatedAtParams{ + Metadata: params.Metadata, + Name: params.Name, + Now: params.Now, + PausedAt: params.PausedAt, + UpdatedAt: params.UpdatedAt, + }) + if err != nil { + return nil, interpretError(err) + } + return queueFromInternal(queue), nil +} + +func (e *Executor) QueueDeleteExpired(ctx context.Context, params *riverdriver.QueueDeleteExpiredParams) ([]string, error) { + queues, err := dbsqlc.New().QueueDeleteExpired(schemaTemplateParam(ctx, params.Schema), e.dbtx, &dbsqlc.QueueDeleteExpiredParams{ + Max: int64(params.Max), + UpdatedAtHorizon: params.UpdatedAtHorizon, + }) + if err != nil { + return nil, interpretError(err) + } + queueNames := make([]string, len(queues)) + for i, q := range queues { + queueNames[i] = q.Name + } + return queueNames, nil +} + +func (e *Executor) QueueGet(ctx context.Context, params *riverdriver.QueueGetParams) (*rivertype.Queue, error) { + queue, err := dbsqlc.New().QueueGet(schemaTemplateParam(ctx, params.Schema), e.dbtx, params.Name) + if err != nil { + return nil, interpretError(err) + } + return queueFromInternal(queue), nil +} + +func (e *Executor) QueueList(ctx context.Context, params *riverdriver.QueueListParams) ([]*rivertype.Queue, error) { + queues, err := dbsqlc.New().QueueList(schemaTemplateParam(ctx, params.Schema), e.dbtx, int32(min(params.Max, math.MaxInt32))) //nolint:gosec + if err != nil { + return nil, interpretError(err) + } + return sliceutil.Map(queues, queueFromInternal), nil +} + +func (e *Executor) QueueNameList(ctx context.Context, params *riverdriver.QueueNameListParams) ([]string, error) { + queueNames, err := dbsqlc.New().QueueNameList(schemaTemplateParam(ctx, params.Schema), e.dbtx, &dbsqlc.QueueNameListParams{ + After: params.After, + Exclude: params.Exclude, + Match: params.Match, + Max: int32(min(params.Max, math.MaxInt32)), //nolint:gosec + }) + if err != nil { + return nil, interpretError(err) + } + return queueNames, nil +} + +func (e *Executor) QueuePause(ctx context.Context, params *riverdriver.QueuePauseParams) error { + rowsAffected, err := dbsqlc.New().QueuePause(schemaTemplateParam(ctx, params.Schema), e.dbtx, &dbsqlc.QueuePauseParams{ + Name: params.Name, + Now: params.Now, + }) + if err != nil { + return interpretError(err) + } + if rowsAffected < 1 && params.Name != riverdriver.AllQueuesString { + return rivertype.ErrNotFound + } + return nil +} + +func (e *Executor) QueueResume(ctx context.Context, params *riverdriver.QueueResumeParams) error { + rowsAffected, err := dbsqlc.New().QueueResume(schemaTemplateParam(ctx, params.Schema), e.dbtx, &dbsqlc.QueueResumeParams{ + Name: params.Name, + Now: params.Now, + }) + if err != nil { + return interpretError(err) + } + if rowsAffected < 1 && params.Name != riverdriver.AllQueuesString { + return rivertype.ErrNotFound + } + return nil +} + +func (e *Executor) QueueUpdate(ctx context.Context, params *riverdriver.QueueUpdateParams) (*rivertype.Queue, error) { + queue, err := dbsqlc.New().QueueUpdate(schemaTemplateParam(ctx, params.Schema), e.dbtx, &dbsqlc.QueueUpdateParams{ + Metadata: params.Metadata, + MetadataDoUpdate: params.MetadataDoUpdate, + Name: params.Name, + }) + if err != nil { + return nil, interpretError(err) + } + return queueFromInternal(queue), nil +} + +func (e *Executor) QueryRow(ctx context.Context, sql string, args ...any) riverdriver.Row { + return e.dbtx.QueryRow(ctx, sql, args...) +} + +func (e *Executor) SchemaCreate(ctx context.Context, params *riverdriver.SchemaCreateParams) error { + _, err := e.dbtx.Exec(ctx, "CREATE SCHEMA "+dbutil.SafeIdentifier(params.Schema)) + return interpretError(err) +} + +func (e *Executor) SchemaDrop(ctx context.Context, params *riverdriver.SchemaDropParams) error { + _, err := e.dbtx.Exec(ctx, "DROP SCHEMA "+dbutil.SafeIdentifier(params.Schema)+" CASCADE") + return interpretError(err) +} + +func (e *Executor) SchemaGetExpired(ctx context.Context, params *riverdriver.SchemaGetExpiredParams) ([]string, error) { + schemas, err := dbsqlc.New().SchemaGetExpired(ctx, e.dbtx, &dbsqlc.SchemaGetExpiredParams{ + BeforeName: params.BeforeName, + Prefix: params.Prefix + "%", + }) + if err != nil { + return nil, interpretError(err) + } + return schemas, nil +} + +func (e *Executor) TableExists(ctx context.Context, params *riverdriver.TableExistsParams) (bool, error) { + // Different from other operations because the schemaAndTable name is a parameter. + schemaAndTable := params.Table + if params.Schema != "" { + schemaAndTable = dbutil.SafeIdentifier(params.Schema) + "." + schemaAndTable + } + + exists, err := dbsqlc.New().TableExists(ctx, e.dbtx, schemaAndTable) + return exists, interpretError(err) +} + +func (e *Executor) TableTruncate(ctx context.Context, params *riverdriver.TableTruncateParams) error { + var maybeSchema string + if params.Schema != "" { + maybeSchema = dbutil.SafeIdentifier(params.Schema) + "." + } + + // Uses raw SQL so we can truncate multiple tables at once. + _, err := e.dbtx.Exec(ctx, "TRUNCATE TABLE "+ + strings.Join( + sliceutil.Map( + params.Table, + func(table string) string { return maybeSchema + table }, + ), + ", ", + ), + ) + return interpretError(err) +} + +type ExecutorTx struct { + Executor + + tx pgx.Tx +} + +func (t *ExecutorTx) Commit(ctx context.Context) error { + return t.tx.Commit(ctx) +} + +func (t *ExecutorTx) Rollback(ctx context.Context) error { + return t.tx.Rollback(ctx) +} + +type Listener struct { + afterConnectExec string // should only ever be used in testing + conn *pgx.Conn + dbPool *pgxpool.Pool + prefix string // schema with a dot on the end (very minor optimization) + mu sync.Mutex + schema string +} + +func (l *Listener) Close(ctx context.Context) error { + l.mu.Lock() + defer l.mu.Unlock() + + if l.conn == nil { + return nil + } + + // Release below would take care of cleanup and potentially put the + // connection back into rotation, but in case a Listen was invoked without a + // subsequent Unlisten on the same topic, close the connection explicitly to + // guarantee no other caller will receive a partially tainted connection. + err := l.conn.Close(ctx) + + // Even in the event of an error, make sure conn is set back to nil so that + // the listener can be reused. + l.conn = nil + + return err +} + +func (l *Listener) Connect(ctx context.Context) error { + l.mu.Lock() + defer l.mu.Unlock() + + if l.conn != nil { + return errors.New("connection already established") + } + + poolConn, err := l.dbPool.Acquire(ctx) + if err != nil { + return err + } + + if l.afterConnectExec != "" { + if _, err := poolConn.Exec(ctx, l.afterConnectExec); err != nil { + poolConn.Release() + return err + } + } + + // Use a configured schema if non-empty, otherwise try to select the current + // schema based on `search_path`. + schema := l.schema + if schema == "" { + // `current_schema` may be `NULL` if `search_path` is unset completely. + if err := poolConn.QueryRow(ctx, "SELECT coalesce(current_schema(), '');").Scan(&schema); err != nil { + poolConn.Release() + return err + } + l.schema = schema + } + + if schema != "" { + l.prefix = schema + "." + } + + // Assume full ownership of the conn so that it doesn't get released back to + // the pool or auto-closed by the pool. + l.conn = poolConn.Hijack() + + return nil +} + +func (l *Listener) Listen(ctx context.Context, topic string) error { + l.mu.Lock() + defer l.mu.Unlock() + + _, err := l.conn.Exec(ctx, "LISTEN \""+l.prefix+topic+"\"") + return err +} + +func (l *Listener) Ping(ctx context.Context) error { + l.mu.Lock() + defer l.mu.Unlock() + + return l.conn.Ping(ctx) +} + +func (l *Listener) Schema() string { + l.mu.Lock() + defer l.mu.Unlock() + + return l.schema +} + +func (l *Listener) SetAfterConnectExec(sql string) { + l.mu.Lock() + defer l.mu.Unlock() + + l.afterConnectExec = sql +} + +func (l *Listener) Unlisten(ctx context.Context, topic string) error { + l.mu.Lock() + defer l.mu.Unlock() + + _, err := l.conn.Exec(ctx, "UNLISTEN \""+l.prefix+topic+"\"") + return err +} + +func (l *Listener) WaitForNotification(ctx context.Context) (*riverdriver.Notification, error) { + l.mu.Lock() + defer l.mu.Unlock() + + notification, err := l.conn.WaitForNotification(ctx) + if err != nil { + return nil, err + } + + return &riverdriver.Notification{ + Topic: strings.TrimPrefix(notification.Channel, l.prefix), + Payload: notification.Payload, + }, nil +} + +type templateReplaceWrapper struct { + dbtx interface { + dbsqlc.DBTX + Begin(ctx context.Context) (pgx.Tx, error) + } + replacer *sqlctemplate.Replacer +} + +func (w templateReplaceWrapper) Begin(ctx context.Context) (pgx.Tx, error) { + return w.dbtx.Begin(ctx) +} + +func (w templateReplaceWrapper) defaultQueryExecMode() pgx.QueryExecMode { + if poolWithConfig, ok := any(w.dbtx).(interface{ Config() *pgxpool.Config }); ok { + if config := poolWithConfig.Config(); config != nil { + return config.ConnConfig.DefaultQueryExecMode + } + } + if txWithConn, ok := any(w.dbtx).(interface{ Conn() *pgx.Conn }); ok { + if conn := txWithConn.Conn(); conn != nil { + return conn.Config().DefaultQueryExecMode + } + } + return pgx.QueryExecModeCacheStatement +} + +func (w templateReplaceWrapper) Exec(ctx context.Context, sql string, args ...any) (pgconn.CommandTag, error) { + sql, args = w.replacer.Run(ctx, argPlaceholder, sql, args) + // Keep JSON/JSONB arguments valid in pgx text-only execution modes. + args = adaptArgsForJSONTextModes(w.defaultQueryExecMode(), sql, args) + return w.dbtx.Exec(ctx, sql, args...) +} + +func (w templateReplaceWrapper) Query(ctx context.Context, sql string, args ...any) (pgx.Rows, error) { + sql, args = w.replacer.Run(ctx, argPlaceholder, sql, args) + args = adaptArgsForJSONTextModes(w.defaultQueryExecMode(), sql, args) + return w.dbtx.Query(ctx, sql, args...) +} + +func (w templateReplaceWrapper) QueryRow(ctx context.Context, sql string, args ...any) pgx.Row { + sql, args = w.replacer.Run(ctx, argPlaceholder, sql, args) + args = adaptArgsForJSONTextModes(w.defaultQueryExecMode(), sql, args) + return w.dbtx.QueryRow(ctx, sql, args...) +} + +func (w templateReplaceWrapper) CopyFrom(ctx context.Context, tableName pgx.Identifier, columnNames []string, rowSrc pgx.CopyFromSource) (int64, error) { + if schema, ok := ctx.Value(schemaCopyFromContextKey{}).(string); ok { + tableName = append([]string{schema}, tableName...) + } + + return w.dbtx.CopyFrom(ctx, tableName, columnNames, rowSrc) +} + +func interpretError(err error) error { + if errors.Is(err, puddle.ErrClosedPool) { + return riverdriver.ErrClosedPool + } + if errors.Is(err, pgx.ErrNoRows) { + return rivertype.ErrNotFound + } + return err +} + +func jobRowFromInternal(internal *dbsqlc.RiverJob) (*rivertype.JobRow, error) { + var attemptedAt *time.Time + if internal.AttemptedAt != nil { + t := internal.AttemptedAt.UTC() + attemptedAt = &t + } + + errors := make([]rivertype.AttemptError, len(internal.Errors)) + for i, rawError := range internal.Errors { + if err := json.Unmarshal(rawError, &errors[i]); err != nil { + return nil, err + } + } + + var finalizedAt *time.Time + if internal.FinalizedAt != nil { + t := internal.FinalizedAt.UTC() + finalizedAt = &t + } + + var uniqueStatesByte byte + if internal.UniqueStates.Valid && len(internal.UniqueStates.Bytes) > 0 { + uniqueStatesByte = internal.UniqueStates.Bytes[0] + } + + return &rivertype.JobRow{ + ID: internal.ID, + Attempt: max(int(internal.Attempt), 0), + AttemptedAt: attemptedAt, + AttemptedBy: internal.AttemptedBy, + CreatedAt: internal.CreatedAt.UTC(), + EncodedArgs: internal.Args, + Errors: errors, + FinalizedAt: finalizedAt, + Kind: internal.Kind, + MaxAttempts: max(int(internal.MaxAttempts), 0), + Metadata: internal.Metadata, + Priority: max(int(internal.Priority), 0), + Queue: internal.Queue, + ScheduledAt: internal.ScheduledAt.UTC(), + State: rivertype.JobState(internal.State), + Tags: internal.Tags, + UniqueKey: internal.UniqueKey, + UniqueStates: uniquestates.UniqueBitmaskToStates(uniqueStatesByte), + }, nil +} + +func leaderFromInternal(internal *dbsqlc.RiverLeader) *riverdriver.Leader { + return &riverdriver.Leader{ + ElectedAt: internal.ElectedAt.UTC(), + ExpiresAt: internal.ExpiresAt.UTC(), + LeaderID: internal.LeaderID, + } +} + +func migrationFromInternal(internal *dbsqlc.RiverMigration) *riverdriver.Migration { + return &riverdriver.Migration{ + CreatedAt: internal.CreatedAt.UTC(), + Line: internal.Line, + Version: int(internal.Version), + } +} + +func queueFromInternal(internal *dbsqlc.RiverQueue) *rivertype.Queue { + var pausedAt *time.Time + if internal.PausedAt != nil { + t := internal.PausedAt.UTC() + pausedAt = &t + } + return &rivertype.Queue{ + CreatedAt: internal.CreatedAt.UTC(), + Metadata: internal.Metadata, + Name: internal.Name, + PausedAt: pausedAt, + UpdatedAt: internal.UpdatedAt.UTC(), + } +} + +// A special internal context key used only to set a schema for use in CopyFrom. +// If we end up eliminating the use of copyfrom functions (which can't use +// sqlctemplate because no SQL is executed at any time so there's nowhere to +// otherwise do a replacement), we can get rid of this completely. +type schemaCopyFromContextKey struct{} + +func schemaCopyFrom(ctx context.Context, schema string) context.Context { + if schema != "" { + ctx = context.WithValue(ctx, schemaCopyFromContextKey{}, schema) + } + + return ctx +} + +func schemaTemplateParam(ctx context.Context, schema string) context.Context { + if schema != "" { + schema = dbutil.SafeIdentifier(schema) + "." + } + + return sqlctemplate.WithReplacements(ctx, map[string]sqlctemplate.Replacement{ + "schema": {Value: schema, Stable: true}, + }, nil) +} diff --git a/vendor/github.com/riverqueue/river/rivermigrate/river_migrate.go b/vendor/github.com/riverqueue/river/rivermigrate/river_migrate.go new file mode 100644 index 0000000000..19d15fdf78 --- /dev/null +++ b/vendor/github.com/riverqueue/river/rivermigrate/river_migrate.go @@ -0,0 +1,898 @@ +// Package rivermigrate provides a Go API for running migrations as alternative +// to migrating via the bundled CLI. +package rivermigrate + +import ( + "cmp" + "context" + "errors" + "fmt" + "io" + "io/fs" + "log/slog" + "maps" + "os" + "slices" + "strconv" + "strings" + "time" + + "github.com/riverqueue/river/riverdriver" + "github.com/riverqueue/river/rivershared/baseservice" + "github.com/riverqueue/river/rivershared/levenshtein" + "github.com/riverqueue/river/rivershared/sqlctemplate" + "github.com/riverqueue/river/rivershared/util/dbutil" + "github.com/riverqueue/river/rivershared/util/maputil" + "github.com/riverqueue/river/rivershared/util/sliceutil" +) + +const ( + // The migrate version where the `line` column was added. Meaningful in that + // the migrator has to behave a little differently depending on whether it's + // working with versions before or after this boundary. + migrateVersionLineColumnAdded = 5 + + // The migration version where the `river_migration` table is added. This is + // used for one special case where we don't try to delete a version record + // after downmigrating version 1. + migrateVersionTableAdded = 1 +) + +// Migration is a bundled migration containing a version (e.g. 1, 2, 3), and SQL +// for up and down directions. +type Migration struct { + // Name is a human-friendly name for the migration derived from its + // filename. + Name string + + // SQLDown is the s SQL for the migration's down direction. + SQLDown string + + // SQLUp is the s SQL for the migration's up direction. + SQLUp string + + // Version is the integer version number of this migration. + Version int +} + +// Config contains configuration for Migrator. +type Config struct { + // Line is the migration line to use. Most drivers will only have a single + // line, which is `main`. + // + // Defaults to `main`. + Line string + + // Logger is the structured logger to use for logging purposes. If none is + // specified, logs will be emitted to STDOUT with messages at warn level + // or higher. + Logger *slog.Logger + + // Schema is the target schema to migrate. + // + // Defaults to empty, which means that no schema is explicitly targeted. In + // Postgres a schema will be selected based on what's set in `search_path`. + Schema string +} + +// Migrator is a database migration tool for River which can run up or down +// migrations in order to establish the schema that the queue needs to run. +type Migrator[TTx any] struct { + baseservice.BaseService + + driver riverdriver.Driver[TTx] + line string + migrations map[int]Migration // allows us to inject test migrations + replacer sqlctemplate.Replacer + schema string +} + +// New returns a new migrator with the given database driver and configuration. +// The config parameter may be omitted as nil. +// +// Two drivers are supported for migrations, one for Pgx v5 and one for the +// built-in database/sql package for use with migration frameworks like Goose. +// See packages riverpgxv5 and riverdatabasesql respectively. +// +// The function takes a generic parameter TTx representing a transaction type, +// but it can be omitted because it'll generally always be inferred from the +// driver. For example: +// +// import "github.com/riverqueue/river/riverdriver/riverpgxv5" +// import "github.com/riverqueue/rivermigrate" +// +// ... +// +// dbPool, err := pgxpool.New(ctx, os.Getenv("DATABASE_URL")) +// if err != nil { +// // handle error +// } +// defer dbPool.Close() +// +// migrator, err := rivermigrate.New(riverpgxv5.New(dbPool), nil) +// if err != nil { +// // handle error +// } +func New[TTx any](driver riverdriver.Driver[TTx], config *Config) (*Migrator[TTx], error) { + if config == nil { + config = &Config{} + } + + line := cmp.Or(config.Line, riverdriver.MigrationLineMain) + + logger := config.Logger + if logger == nil { + logger = slog.New(slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{ + Level: slog.LevelWarn, + })) + } + + archetype := &baseservice.Archetype{ + Logger: logger, + Time: &baseservice.UnStubbableTimeGenerator{}, + } + + if !slices.Contains(driver.GetMigrationLines(), line) { + const minLevenshteinDistance = 2 + + var suggestedLines []string + for _, existingLine := range driver.GetMigrationLines() { + if distance := levenshtein.ComputeDistance(existingLine, line); distance <= minLevenshteinDistance { + suggestedLines = append(suggestedLines, "`"+existingLine+"`") + } + } + + errorStr := "migration line does not exist: " + line + switch { + case len(suggestedLines) == 1: + errorStr += fmt.Sprintf(" (did you mean %s?)", suggestedLines[0]) + case len(suggestedLines) > 1: + errorStr += fmt.Sprintf(" (did you mean one of %v?)", strings.Join(suggestedLines, ", ")) + } + + return nil, errors.New(errorStr) + } + + riverMigrations, err := migrationsFromFS(driver.GetMigrationFS(line), line) + if err != nil { + // If there's ever a problem here, it's a very fundamental internal + // River one, so it's okay to panic. + panic(err) + } + + return baseservice.Init(archetype, &Migrator[TTx]{ + driver: driver, + line: line, + migrations: validateAndInit(riverMigrations), + schema: config.Schema, + }), nil +} + +// ExistingVersions gets the existing set of versions that have been migrated in +// the database, ordered by version. +func (m *Migrator[TTx]) ExistingVersions(ctx context.Context) ([]Migration, error) { + migrations, err := m.existingMigrations(ctx, m.driver.GetExecutor()) + if err != nil { + return nil, err + } + + versions, err := m.versionsFromDriver(migrations) + if err != nil { + return nil, err + } + + return versions, nil +} + +// ExistingVersionsTx gets the existing set of versions that have been migrated +// in the database, ordered by version. +// +// This variant checks for existing versions in a transaction. +func (m *Migrator[TTx]) ExistingVersionsTx(ctx context.Context, tx TTx) ([]Migration, error) { + migrations, err := m.existingMigrations(ctx, m.driver.UnwrapExecutor(tx)) + if err != nil { + return nil, err + } + + versions, err := m.versionsFromDriver(migrations) + if err != nil { + return nil, err + } + + return versions, nil +} + +func (m *Migrator[TTx]) versionsFromDriver(migrations []*riverdriver.Migration) ([]Migration, error) { + versions := make([]Migration, len(migrations)) + for i, existingMigration := range migrations { + migration, ok := m.migrations[existingMigration.Version] + if !ok { + return nil, fmt.Errorf("migration %d not found in migrator bundle", existingMigration.Version) + } + versions[i] = migration + } + return versions, nil +} + +// MigrateOpts are options for a migrate operation. +type MigrateOpts struct { + DryRun bool + + // MaxSteps is the maximum number of migrations to apply either up or down. + // When migrating in the up direction, migrates an unlimited number of steps + // by default. When migrating in the down direction, migrates only a single + // step by default (set TargetVersion to -1 to apply unlimited steps down). + // Set to -1 to apply no migrations (for testing/checking purposes). + MaxSteps int + + // TargetVersion is a specific migration version to apply migrations to. The + // version must exist and it must be in the possible list of migrations to + // apply. e.g. If requesting an up migration with version 3, version 3 must + // not already be applied. + // + // When applying migrations up, migrations are applied including the target + // version, so when starting at version 0 and requesting version 3, versions + // 1, 2, and 3 would be applied. When applying migrations down, down + // migrations are applied excluding the target version, so when starting at + // version 5 and requesting version 3, down migrations for versions 5 and 4 + // would be applied, leaving the final schema at version 3. + // + // When migrating down, TargetVersion can be set to the special value of -1 + // to apply all down migrations (i.e. River schema is removed completely). + TargetVersion int +} + +// MigrateResult is the result of a migrate operation. +type MigrateResult struct { + // Direction is the direction that migration occurred (up or down). + Direction Direction + + // Versions are migration versions that were added (for up migrations) or + // removed (for down migrations) for this run. + Versions []MigrateVersion +} + +// MigrateVersion is the result for a single applied migration. +type MigrateVersion struct { + // Duration is the amount of time it took to apply the migration. + Duration time.Duration + + // Name is a human-friendly name for the migration derived from its + // filename. + Name string + + // SQL is the SQL that was applied along with the migration. + SQL string + + // Version is the version of the migration applied. + Version int +} + +func migrateVersionToInt(version MigrateVersion) int { return version.Version } + +type Direction string + +const ( + DirectionDown Direction = "down" + DirectionUp Direction = "up" +) + +// AllVersions gets information on all known migration versions. +func (m *Migrator[TTx]) AllVersions() []Migration { + migrations := maputil.Values(m.migrations) + slices.SortFunc(migrations, func(v1, v2 Migration) int { return v1.Version - v2.Version }) + return migrations +} + +// GetVersion gets information about a specific migration version. An error is +// returned if a versions is requested that doesn't exist. +func (m *Migrator[TTx]) GetVersion(version int) (Migration, error) { + migration, ok := m.migrations[version] + if !ok { + availableVersions := maputil.Keys(m.migrations) + slices.Sort(availableVersions) + return Migration{}, fmt.Errorf("migration %d not found (available versions: %v)", version, availableVersions) + } + + return migration, nil +} + +// Migrate migrates the database in the given direction (up or down). The opts +// parameter may be omitted for convenience. +// +// By default, applies all outstanding migrations when moving in the up +// direction, but for safety, only one step when moving in the down direction. +// To migrate more than one step down, MigrateOpts.MaxSteps or +// MigrateOpts.TargetVersion are available. Setting MigrateOpts.TargetVersion to +// -1 will apply every available downstep so that River's schema is removed +// completely. +// +// res, err := migrator.Migrate(ctx, rivermigrate.DirectionUp, nil) +// if err != nil { +// // handle error +// } +func (m *Migrator[TTx]) Migrate(ctx context.Context, direction Direction, opts *MigrateOpts) (*MigrateResult, error) { + exec := m.driver.GetExecutor() + switch direction { + case DirectionDown: + return m.migrateDown(ctx, exec, direction, opts, false) + case DirectionUp: + return m.migrateUp(ctx, exec, direction, opts, false) + } + + panic("invalid direction: " + direction) +} + +// Migrate migrates the database in the given direction (up or down). The opts +// parameter may be omitted for convenience. +// +// By default, applies all outstanding migrations when moving in the up +// direction, but for safety, only one step when moving in the down direction. +// To migrate more than one step down, MigrateOpts.MaxSteps or +// MigrateOpts.TargetVersion are available. Setting MigrateOpts.TargetVersion to +// -1 will apply every available downstep so that River's schema is removed +// completely. +// +// res, err := migrator.MigrateTx(ctx, tx, rivermigrate.DirectionUp, nil) +// if err != nil { +// // handle error +// } +// +// This variant lets a caller run migrations within a transaction. Postgres DDL +// is transactional, so migration changes aren't visible until the transaction +// commits, and are rolled back if the transaction rolls back. +// +// Deprecated: Use Migrate instead. Certain migrations cannot be batched together +// in a single transaction, so this method is not recommended. +func (m *Migrator[TTx]) MigrateTx(ctx context.Context, tx TTx, direction Direction, opts *MigrateOpts) (*MigrateResult, error) { + switch direction { + case DirectionDown: + return m.migrateDown(ctx, m.driver.UnwrapExecutor(tx), direction, opts, true) + case DirectionUp: + return m.migrateUp(ctx, m.driver.UnwrapExecutor(tx), direction, opts, true) + } + + panic("invalid direction: " + direction) +} + +// ValidateResult is the result of a validation operation. +type ValidateResult struct { + // Messages contain informational messages of what wasn't valid in case of a + // failed validation. Always empty if OK is true. + Messages []string + + // OK is true if validation completed with no problems. + OK bool +} + +// ValidateOpts are options for a validate operation. +type ValidateOpts struct { + // TargetVersion is a specific migration version to validate up to. The + // version must exist. When set, validation only checks that migrations up + // to and including TargetVersion have been applied. + TargetVersion int +} + +// Validate validates the current state of migrations, returning an unsuccessful +// validation and usable message in case there are migrations that haven't yet +// been applied. +func (m *Migrator[TTx]) Validate(ctx context.Context, opts *ValidateOpts) (*ValidateResult, error) { + return dbutil.WithTxV(ctx, m.driver.GetExecutor(), func(ctx context.Context, tx riverdriver.ExecutorTx) (*ValidateResult, error) { + return m.validate(ctx, tx, opts) + }) +} + +// ValidateTx validates the current state of migrations, returning an unsuccessful +// validation and usable message in case there are migrations that haven't yet +// been applied. +// +// This variant lets a caller validate within a transaction. +func (m *Migrator[TTx]) ValidateTx(ctx context.Context, tx TTx, opts *ValidateOpts) (*ValidateResult, error) { + return m.validate(ctx, m.driver.UnwrapExecutor(tx), opts) +} + +// migrateDown runs down migrations. +func (m *Migrator[TTx]) migrateDown(ctx context.Context, exec riverdriver.Executor, direction Direction, opts *MigrateOpts, inOuterTx bool) (*MigrateResult, error) { + if opts == nil { + opts = &MigrateOpts{} + } + + existingMigrations, err := m.existingMigrations(ctx, exec) + if err != nil { + return nil, err + } + existingMigrationsMap := sliceutil.KeyBy(existingMigrations, + func(m *riverdriver.Migration) (int, struct{}) { return m.Version, struct{}{} }) + + targetMigrations := maps.Clone(m.migrations) + for version := range targetMigrations { + if _, ok := existingMigrationsMap[version]; !ok { + delete(targetMigrations, version) + } + } + + sortedTargetMigrations := maputil.Values(targetMigrations) + slices.SortFunc(sortedTargetMigrations, func(a, b Migration) int { return b.Version - a.Version }) // reverse order + + res, err := m.applyMigrations(ctx, exec, direction, opts, inOuterTx, sortedTargetMigrations) + if err != nil { + return nil, err + } + + // If we did no work, leave early. This allows a zero-migrated database + // that's being no-op downmigrated again to succeed because otherwise + // the delete below would cause it to error. + if len(res.Versions) < 1 { + return res, nil + } + + // Migration version 1 is special-cased because if it was downmigrated + // it means the `river_migration` table is no longer present so there's + // nothing to delete out of. + if slices.ContainsFunc(res.Versions, func(v MigrateVersion) bool { return v.Version == 1 }) { + return res, nil + } + + // When operating with an outer transaction, all versions are removed at + // once so we can save a few database operations. + if inOuterTx { + if err := m.versionsDelete(ctx, exec, opts, sliceutil.Map(res.Versions, migrateVersionToInt)...); err != nil { + return nil, err + } + } + + return res, nil +} + +// migrateUp runs up migrations. +func (m *Migrator[TTx]) migrateUp(ctx context.Context, exec riverdriver.Executor, direction Direction, opts *MigrateOpts, inOuterTx bool) (*MigrateResult, error) { + if opts == nil { + opts = &MigrateOpts{} + } + + existingMigrations, err := m.existingMigrations(ctx, exec) + if err != nil { + return nil, err + } + + targetMigrations := maps.Clone(m.migrations) + for _, migrateRow := range existingMigrations { + delete(targetMigrations, migrateRow.Version) + } + + sortedTargetMigrations := maputil.Values(targetMigrations) + slices.SortFunc(sortedTargetMigrations, func(a, b Migration) int { return a.Version - b.Version }) + + res, err := m.applyMigrations(ctx, exec, direction, opts, inOuterTx, sortedTargetMigrations) + if err != nil { + return nil, err + } + + // When operating with an outer transaction, all versions are added at once + // so we can save a few database operations. + if inOuterTx { + if err := m.versionsInsert(ctx, exec, opts, sliceutil.Map(res.Versions, migrateVersionToInt)...); err != nil { + return nil, err + } + } + + return res, nil +} + +// validate validates current migration state. +func (m *Migrator[TTx]) validate(ctx context.Context, exec riverdriver.Executor, opts *ValidateOpts) (*ValidateResult, error) { + if opts == nil { + opts = &ValidateOpts{} + } + + existingMigrations, err := m.existingMigrations(ctx, exec) + if err != nil { + return nil, err + } + + targetMigrations := maps.Clone(m.migrations) + for _, migrateRow := range existingMigrations { + delete(targetMigrations, migrateRow.Version) + } + + if opts.TargetVersion > 0 { + if _, ok := m.migrations[opts.TargetVersion]; !ok { + return nil, fmt.Errorf("version %d is not a valid River migration version", opts.TargetVersion) + } + + for version := range targetMigrations { + if version > opts.TargetVersion { + delete(targetMigrations, version) + } + } + } + + notOKWithMessage := func(message string) *ValidateResult { + m.Logger.InfoContext(ctx, m.Name+": "+message) + return &ValidateResult{Messages: []string{message}} + } + + if len(targetMigrations) > 0 { + sortedTargetMigrations := maputil.Keys(targetMigrations) + slices.Sort(sortedTargetMigrations) + + return notOKWithMessage(fmt.Sprintf("Unapplied migrations: %v", sortedTargetMigrations)), nil + } + + return &ValidateResult{OK: true}, nil +} + +// Common code shared between the up and down migration directions that walks +// through each target migration and applies it, logging appropriately. +func (m *Migrator[TTx]) applyMigrations(ctx context.Context, exec riverdriver.Executor, direction Direction, opts *MigrateOpts, inOuterTx bool, sortedTargetMigrations []Migration) (*MigrateResult, error) { + var maxSteps int + switch { + case opts.MaxSteps != 0: + maxSteps = opts.MaxSteps + case direction == DirectionDown && opts.TargetVersion == 0: + maxSteps = 1 + } + + switch { + case maxSteps < 0: + sortedTargetMigrations = []Migration{} + case maxSteps > 0: + sortedTargetMigrations = sortedTargetMigrations[0:min(maxSteps, len(sortedTargetMigrations))] + } + + if opts.TargetVersion > 0 { + if _, ok := m.migrations[opts.TargetVersion]; !ok { + return nil, fmt.Errorf("version %d is not a valid River migration version", opts.TargetVersion) + } + + targetIndex := slices.IndexFunc(sortedTargetMigrations, func(b Migration) bool { return b.Version == opts.TargetVersion }) + if targetIndex == -1 { + // Error, but only if the migration doesn't exist or was never + // applied on a down migration. Up migrations with TargetVersion + // that's already applied should fall through with a no-op. + if _, ok := m.migrations[opts.TargetVersion]; !ok || direction == DirectionDown { + return nil, fmt.Errorf("version %d is not in target list of valid migrations to apply", opts.TargetVersion) + } + } else { + // Replace target list with list up to target index. Migrations are + // sorted according to the direction we're migrating in, so when down + // migration, the list is already reversed, so this will truncate it so + // it's the most current migration down to the target. + sortedTargetMigrations = sortedTargetMigrations[0 : targetIndex+1] + + if direction == DirectionDown && len(sortedTargetMigrations) > 0 { + sortedTargetMigrations = sortedTargetMigrations[0 : len(sortedTargetMigrations)-1] + } + } + } + + res := &MigrateResult{Direction: direction, Versions: make([]MigrateVersion, 0, len(sortedTargetMigrations))} + + // Short circuit early if there's nothing to do. + if len(sortedTargetMigrations) < 1 { + m.Logger.InfoContext(ctx, m.Name+": No migrations to apply") + return res, nil + } + + var schema string + if m.schema != "" { + schema = dbutil.SafeIdentifier(m.schema) + "." + } + schemaReplacement := map[string]sqlctemplate.Replacement{ + "schema": {Value: schema}, + } + + for _, versionBundle := range sortedTargetMigrations { + var sql string + switch direction { + case DirectionDown: + sql = versionBundle.SQLDown + case DirectionUp: + sql = versionBundle.SQLUp + } + + // Most migrations contain schema in their SQL by necessity, but some of + // the test ones do not because they only run trivial operations. + if strings.Contains(sql, "/* TEMPLATE: schema */") { + ctx := sqlctemplate.WithReplacements(ctx, schemaReplacement, nil) + sql, _ = m.replacer.Run(ctx, m.driver.ArgPlaceholder(), sql, nil) + } + + var duration time.Duration + + if !opts.DryRun { + start := time.Now() + + // Similar to ActiveRecord migrations, we wrap each individual migration + // in its own transaction. Without this, certain migrations that require + // a commit on a preexisting operation (such as adding an enum value to be + // used in an immutable function) cannot succeed. + err := dbutil.WithTx(ctx, exec, func(ctx context.Context, exec riverdriver.ExecutorTx) error { + if err := exec.Exec(ctx, sql); err != nil { + return fmt.Errorf("error applying version %03d [%s]: %w", + versionBundle.Version, strings.ToUpper(string(direction)), err) + } + + // If operating without outer transaction, add/remove the + // migration version in the same transaction in which we + // executed the migration SQL. + if !inOuterTx { + switch direction { + case DirectionDown: + if err := m.versionsDelete(ctx, exec, opts, versionBundle.Version); err != nil { + return err + } + case DirectionUp: + if err := m.versionsInsert(ctx, exec, opts, versionBundle.Version); err != nil { + return err + } + } + } + + return nil + }) + if err != nil { + return nil, err + } + duration = time.Since(start) + } + + m.Logger.InfoContext(ctx, m.Name+": Applied migration", + slog.String("direction", string(direction)), + slog.Bool("dry_run", opts.DryRun), + slog.Duration("duration", duration), + slog.Int("version", versionBundle.Version), + ) + + res.Versions = append(res.Versions, MigrateVersion{Duration: duration, Name: versionBundle.Name, SQL: sql, Version: versionBundle.Version}) + } + + return res, nil +} + +// Get existing migrations that've already been run in the database. This is +// encapsulated to run a check in a subtransaction and the handle the case of +// the `river_migration` table not existing yet. (The subtransaction is needed +// because otherwise the existing transaction would become aborted on an +// unsuccessful `river_migration` check.) +func (m *Migrator[TTx]) existingMigrations(ctx context.Context, exec riverdriver.Executor) ([]*riverdriver.Migration, error) { + migrateTableExists, err := exec.TableExists(ctx, &riverdriver.TableExistsParams{ + Schema: m.schema, + Table: "river_migration", + }) + if err != nil { + return nil, fmt.Errorf("error checking if `%s` exists: %w", "river_migration", err) + } + if !migrateTableExists { + if m.line != riverdriver.MigrationLineMain { + return nil, errors.New("can't add a non-main migration line until `river_migration` is raised; fully migrate the main migration line and try again") + } + + return nil, nil + } + + lineColumnExists, err := exec.ColumnExists(ctx, &riverdriver.ColumnExistsParams{ + Column: "line", + Schema: m.schema, + Table: "river_migration", + }) + if err != nil { + return nil, fmt.Errorf("error checking if `%s.%s` exists: %w", "river_migration", "line", err) + } + + if !lineColumnExists { + if m.line != riverdriver.MigrationLineMain { + return nil, errors.New("can't add a non-main migration line until `river_migration.line` is raised; fully migrate the main migration line and try again") + } + + migrations, err := exec.MigrationGetAllAssumingMain(ctx, &riverdriver.MigrationGetAllAssumingMainParams{ + Schema: m.schema, + }) + if err != nil { + return nil, fmt.Errorf("error getting existing migrations: %w", err) + } + + return migrations, nil + } + + migrations, err := exec.MigrationGetByLine(ctx, &riverdriver.MigrationGetByLineParams{ + Line: m.line, + Schema: m.schema, + }) + if err != nil { + return nil, fmt.Errorf("error getting existing migrations for line %q: %w", m.line, err) + } + + return migrations, nil +} + +func (m *Migrator[TTx]) versionsDelete(ctx context.Context, exec riverdriver.Executor, opts *MigrateOpts, versions ...int) error { + if opts.DryRun || len(versions) < 1 { + return nil + } + + // Don't try to remove anything if we're migrating back below version 1, + // where `river_migration` was added. + if len(versions) == 1 && versions[0] <= migrateVersionTableAdded { + return nil + } + + // Version 005 is hard-coded here because that's the version in which + // the migration `line` comes in. If migration to a point equal or above + // 005, we can remove migrations with a line included, but otherwise we + // must omit the `line` column from queries because it doesn't exist. + if m.line == riverdriver.MigrationLineMain && slices.Min(versions) <= migrateVersionLineColumnAdded { + if _, err := exec.MigrationDeleteAssumingMainMany(ctx, &riverdriver.MigrationDeleteAssumingMainManyParams{ + Versions: versions, + Schema: m.schema, + }); err != nil { + return fmt.Errorf("error inserting migration rows for versions %+v assuming main: %w", versions, err) + } + } else { + if _, err := exec.MigrationDeleteByLineAndVersionMany(ctx, &riverdriver.MigrationDeleteByLineAndVersionManyParams{ + Line: m.line, + Schema: m.schema, + Versions: versions, + }); err != nil { + return fmt.Errorf("error deleting migration rows for versions %+v on line %q: %w", versions, m.line, err) + } + } + + return nil +} + +func (m *Migrator[TTx]) versionsInsert(ctx context.Context, exec riverdriver.Executor, opts *MigrateOpts, versions ...int) error { + if opts.DryRun || len(versions) < 1 { + return nil + } + + // Version 005 is hard-coded here because that's the version in which + // the migration `line` comes in. If migration to a point equal or above + // 005, we can insert migrations with a line included, but otherwise we + // must omit the `line` column from queries because it doesn't exist. + if m.line == riverdriver.MigrationLineMain && slices.Max(versions) < migrateVersionLineColumnAdded { + if _, err := exec.MigrationInsertManyAssumingMain(ctx, &riverdriver.MigrationInsertManyAssumingMainParams{ + Schema: m.schema, + Versions: versions, + }); err != nil { + return fmt.Errorf("error inserting migration rows for versions %+v assuming main: %w", versions, err) + } + } else { + if _, err := exec.MigrationInsertMany(ctx, &riverdriver.MigrationInsertManyParams{ + Line: m.line, + Schema: m.schema, + Versions: versions, + }); err != nil { + return fmt.Errorf("error inserting migration rows for versions %+v on line %q: %w", versions, m.line, err) + } + } + + return nil +} + +// Reads a series of migration bundles from a file system, which practically +// speaking will always be the embedded FS read from the contents of the +// `migration//` subdirectory. +func migrationsFromFS(migrationFS fs.FS, line string) ([]Migration, error) { + const subdir = "migration" + + var ( + lastBundle *Migration + migrations []Migration + ) + + err := fs.WalkDir(migrationFS, subdir, func(path string, entry fs.DirEntry, err error) error { + if err != nil { + return fmt.Errorf("error walking FS: %w", err) + } + + // The WalkDir callback is invoked for each embedded subdirectory and + // file. For our purposes here, we're only interested in files. + if entry.IsDir() { + return nil + } + + filename := path + + // Invoked with the full path name. Strip `migration/` from the front so + // we have a name that we can parse with. + if !strings.HasPrefix(filename, subdir) { + return fmt.Errorf("expected path %q to start with subdir %q", path, subdir) + } + filename = filename[len(subdir)+1:] + + // Ignore any migrations that don't belong to the line we're reading. + if !strings.HasPrefix(filename, line) { + return nil + } + filename = filename[len(line)+1:] + + versionStr, name, ok := strings.Cut(filename, "_") + if !ok { + return fmt.Errorf("expected name to start with version string like '001_': %q", filename) + } + + version, err := strconv.Atoi(versionStr) + if err != nil { + return fmt.Errorf("error parsing version %q: %w", versionStr, err) + } + + // Non-version name for the migration. So for `002_initial_schema` it + // would be `initial schema`. + name, _, _ = strings.Cut(name, ".") + name = strings.ReplaceAll(name, "_", " ") + + // This works because `fs.WalkDir` guarantees lexical order, so all 001* + // files always appear before all 002* files, etc. + if lastBundle == nil || lastBundle.Version != version { + migrations = append(migrations, Migration{Name: name, Version: version}) + lastBundle = &migrations[len(migrations)-1] + } + + file, err := migrationFS.Open(path) + if err != nil { + return fmt.Errorf("error opening file %q: %w", path, err) + } + + contents, err := io.ReadAll(file) + if err != nil { + return fmt.Errorf("error reading file %q: %w", path, err) + } + + switch { + case strings.HasSuffix(filename, ".down.sql"): + lastBundle.SQLDown = string(contents) + case strings.HasSuffix(filename, ".up.sql"): + lastBundle.SQLUp = string(contents) + default: + return fmt.Errorf("file %q should end with either '.down.sql' or '.up.sql'", filename) + } + + return nil + }) + if err != nil { + return nil, err + } + + if len(migrations) < 1 { + return nil, fmt.Errorf("no migrations found for line: %q", line) + } + + return migrations, nil +} + +// Validates and fully initializes a set of migrations to reduce the probability +// of configuration problems as new migrations are introduced. e.g. Checks for +// missing fields or accidentally duplicated version numbers from copy/pasta +// problems. +func validateAndInit(versions []Migration) map[int]Migration { + lastVersion := 0 + migrations := make(map[int]Migration, len(versions)) + + for _, versionBundle := range versions { + if versionBundle.SQLDown == "" { + panic(fmt.Sprintf("version bundle should specify Down: %+v", versionBundle)) + } + if versionBundle.SQLUp == "" { + panic(fmt.Sprintf("version bundle should specify Up: %+v", versionBundle)) + } + if versionBundle.Version == 0 { + panic(fmt.Sprintf("version bundle should specify Version: %+v", versionBundle)) + } + + if _, ok := migrations[versionBundle.Version]; ok { + panic(fmt.Sprintf("duplicate version: %03d", versionBundle.Version)) + } + if versionBundle.Version <= lastVersion { + panic(fmt.Sprintf("versions should be ascending; current: %03d, last: %03d", versionBundle.Version, lastVersion)) + } + if versionBundle.Version > lastVersion+1 { + panic(fmt.Sprintf("versions shouldn't skip a sequence number; current: %03d, last: %03d", versionBundle.Version, lastVersion)) + } + + lastVersion = versionBundle.Version + migrations[versionBundle.Version] = versionBundle + } + + return migrations +} diff --git a/vendor/github.com/riverqueue/river/rivershared/LICENSE b/vendor/github.com/riverqueue/river/rivershared/LICENSE new file mode 100644 index 0000000000..2f8ed188e8 --- /dev/null +++ b/vendor/github.com/riverqueue/river/rivershared/LICENSE @@ -0,0 +1,374 @@ +Mozilla Public License Version 2.0 +================================== + +1. Definitions +-------------- + +1.1. "Contributor" + means each individual or legal entity that creates, contributes to + the creation of, or owns Covered Software. + +1.2. "Contributor Version" + means the combination of the Contributions of others (if any) used + by a Contributor and that particular Contributor's Contribution. + +1.3. "Contribution" + means Covered Software of a particular Contributor. + +1.4. "Covered Software" + means Source Code Form to which the initial Contributor has attached + the notice in Exhibit A, the Executable Form of such Source Code + Form, and Modifications of such Source Code Form, in each case + including portions thereof. + +1.5. "Incompatible With Secondary Licenses" + means + + (a) that the initial Contributor has attached the notice described + in Exhibit B to the Covered Software; or + + (b) that the Covered Software was made available under the terms of + version 1.1 or earlier of the License, but not also under the + terms of a Secondary License. + +1.6. "Executable Form" + means any form of the work other than Source Code Form. + +1.7. "Larger Work" + means a work that combines Covered Software with other material, in + a separate file or files, that is not Covered Software. + +1.8. "License" + means this document. + +1.9. "Licensable" + means having the right to grant, to the maximum extent possible, + whether at the time of the initial grant or subsequently, any and + all of the rights conveyed by this License. + +1.10. "Modifications" + means any of the following: + + (a) any file in Source Code Form that results from an addition to, + deletion from, or modification of the contents of Covered + Software; or + + (b) any new file in Source Code Form that contains any Covered + Software. + +1.11. "Patent Claims" of a Contributor + means any patent claim(s), including without limitation, method, + process, and apparatus claims, in any patent Licensable by such + Contributor that would be infringed, but for the grant of the + License, by the making, using, selling, offering for sale, having + made, import, or transfer of either its Contributions or its + Contributor Version. + +1.12. "Secondary License" + means either the GNU General Public License, Version 2.0, the GNU + Lesser General Public License, Version 2.1, the GNU Affero General + Public License, Version 3.0, or any later versions of those + licenses. + +1.13. "Source Code Form" + means the form of the work preferred for making modifications. + +1.14. "You" (or "Your") + means an individual or a legal entity exercising rights under this + License. For legal entities, "You" includes any entity that + controls, is controlled by, or is under common control with You. For + purposes of this definition, "control" means (a) the power, direct + or indirect, to cause the direction or management of such entity, + whether by contract or otherwise, or (b) ownership of more than + fifty percent (50%) of the outstanding shares or beneficial + ownership of such entity. + +2. License Grants and Conditions +-------------------------------- + +2.1. Grants + +Each Contributor hereby grants You a world-wide, royalty-free, +non-exclusive license: + +(a) under intellectual property rights (other than patent or trademark) + Licensable by such Contributor to use, reproduce, make available, + modify, display, perform, distribute, and otherwise exploit its + Contributions, either on an unmodified basis, with Modifications, or + as part of a Larger Work; and + +(b) under Patent Claims of such Contributor to make, use, sell, offer + for sale, have made, import, and otherwise transfer either its + Contributions or its Contributor Version. + +2.2. Effective Date + +The licenses granted in Section 2.1 with respect to any Contribution +become effective for each Contribution on the date the Contributor first +distributes such Contribution. + +2.3. Limitations on Grant Scope + +The licenses granted in this Section 2 are the only rights granted under +this License. No additional rights or licenses will be implied from the +distribution or licensing of Covered Software under this License. +Notwithstanding Section 2.1(b) above, no patent license is granted by a +Contributor: + +(a) for any code that a Contributor has removed from Covered Software; + or + +(b) for infringements caused by: (i) Your and any other third party's + modifications of Covered Software, or (ii) the combination of its + Contributions with other software (except as part of its Contributor + Version); or + +(c) under Patent Claims infringed by Covered Software in the absence of + its Contributions. + +This License does not grant any rights in the trademarks, service marks, +or logos of any Contributor (except as may be necessary to comply with +the notice requirements in Section 3.4). + +2.4. Subsequent Licenses + +No Contributor makes additional grants as a result of Your choice to +distribute the Covered Software under a subsequent version of this +License (see Section 10.2) or under the terms of a Secondary License (if +permitted under the terms of Section 3.3). + +2.5. Representation + +Each Contributor represents that the Contributor believes its +Contributions are its original creation(s) or it has sufficient rights +to grant the rights to its Contributions conveyed by this License. + +2.6. Fair Use + +This License is not intended to limit any rights You have under +applicable copyright doctrines of fair use, fair dealing, or other +equivalents. + +2.7. Conditions + +Sections 3.1, 3.2, 3.3, and 3.4 are conditions of the licenses granted +in Section 2.1. + +3. Responsibilities +------------------- + +3.1. Distribution of Source Form + +All distribution of Covered Software in Source Code Form, including any +Modifications that You create or to which You contribute, must be under +the terms of this License. You must inform recipients that the Source +Code Form of the Covered Software is governed by the terms of this +License, and how they can obtain a copy of this License. You may not +attempt to alter or restrict the recipients' rights in the Source Code +Form. + +3.2. Distribution of Executable Form + +If You distribute Covered Software in Executable Form then: + +(a) such Covered Software must also be made available in Source Code + Form, as described in Section 3.1, and You must inform recipients of + the Executable Form how they can obtain a copy of such Source Code + Form by reasonable means in a timely manner, at a charge no more + than the cost of distribution to the recipient; and + +(b) You may distribute such Executable Form under the terms of this + License, or sublicense it under different terms, provided that the + license for the Executable Form does not attempt to limit or alter + the recipients' rights in the Source Code Form under this License. + +3.3. Distribution of a Larger Work + +You may create and distribute a Larger Work under terms of Your choice, +provided that You also comply with the requirements of this License for +the Covered Software. If the Larger Work is a combination of Covered +Software with a work governed by one or more Secondary Licenses, and the +Covered Software is not Incompatible With Secondary Licenses, this +License permits You to additionally distribute such Covered Software +under the terms of such Secondary License(s), so that the recipient of +the Larger Work may, at their option, further distribute the Covered +Software under the terms of either this License or such Secondary +License(s). + +3.4. Notices + +You may not remove or alter the substance of any license notices +(including copyright notices, patent notices, disclaimers of warranty, +or limitations of liability) contained within the Source Code Form of +the Covered Software, except that You may alter any license notices to +the extent required to remedy known factual inaccuracies. + +3.5. Application of Additional Terms + +You may choose to offer, and to charge a fee for, warranty, support, +indemnity or liability obligations to one or more recipients of Covered +Software. However, You may do so only on Your own behalf, and not on +behalf of any Contributor. You must make it absolutely clear that any +such warranty, support, indemnity, or liability obligation is offered by +You alone, and You hereby agree to indemnify every Contributor for any +liability incurred by such Contributor as a result of warranty, support, +indemnity or liability terms You offer. You may include additional +disclaimers of warranty and limitations of liability specific to any +jurisdiction. + +4. Inability to Comply Due to Statute or Regulation +--------------------------------------------------- + +If it is impossible for You to comply with any of the terms of this +License with respect to some or all of the Covered Software due to +statute, judicial order, or regulation then You must: (a) comply with +the terms of this License to the maximum extent possible; and (b) +describe the limitations and the code they affect. Such description must +be placed in a text file included with all distributions of the Covered +Software under this License. Except to the extent prohibited by statute +or regulation, such description must be sufficiently detailed for a +recipient of ordinary skill to be able to understand it. + +5. Termination +-------------- + +5.1. The rights granted under this License will terminate automatically +if You fail to comply with any of its terms. However, if You become +compliant, then the rights granted under this License from a particular +Contributor are reinstated (a) provisionally, unless and until such +Contributor explicitly and finally terminates Your grants, and (b) on an +ongoing basis, if such Contributor fails to notify You of the +non-compliance by some reasonable means prior to 60 days after You have +come back into compliance. Moreover, Your grants from a particular +Contributor are reinstated on an ongoing basis if such Contributor +notifies You of the non-compliance by some reasonable means, this is the +first time You have received notice of non-compliance with this License +from such Contributor, and You become compliant prior to 30 days after +Your receipt of the notice. + +5.2. If You initiate litigation against any entity by asserting a patent +infringement claim (excluding declaratory judgment actions, +counter-claims, and cross-claims) alleging that a Contributor Version +directly or indirectly infringes any patent, then the rights granted to +You by any and all Contributors for the Covered Software under Section +2.1 of this License shall terminate. + +5.3. In the event of termination under Sections 5.1 or 5.2 above, all +end user license agreements (excluding distributors and resellers) which +have been validly granted by You or Your distributors under this License +prior to termination shall survive termination. + +************************************************************************ +* * +* 6. Disclaimer of Warranty * +* ------------------------- * +* * +* Covered Software is provided under this License on an "as is" * +* basis, without warranty of any kind, either expressed, implied, or * +* statutory, including, without limitation, warranties that the * +* Covered Software is free of defects, merchantable, fit for a * +* particular purpose or non-infringing. The entire risk as to the * +* quality and performance of the Covered Software is with You. * +* Should any Covered Software prove defective in any respect, You * +* (not any Contributor) assume the cost of any necessary servicing, * +* repair, or correction. This disclaimer of warranty constitutes an * +* essential part of this License. No use of any Covered Software is * +* authorized under this License except under this disclaimer. * +* * +************************************************************************ + +************************************************************************ +* * +* 7. Limitation of Liability * +* -------------------------- * +* * +* Under no circumstances and under no legal theory, whether tort * +* (including negligence), contract, or otherwise, shall any * +* Contributor, or anyone who distributes Covered Software as * +* permitted above, be liable to You for any direct, indirect, * +* special, incidental, or consequential damages of any character * +* including, without limitation, damages for lost profits, loss of * +* goodwill, work stoppage, computer failure or malfunction, or any * +* and all other commercial damages or losses, even if such party * +* shall have been informed of the possibility of such damages. This * +* limitation of liability shall not apply to liability for death or * +* personal injury resulting from such party's negligence to the * +* extent applicable law prohibits such limitation. Some * +* jurisdictions do not allow the exclusion or limitation of * +* incidental or consequential damages, so this exclusion and * +* limitation may not apply to You. * +* * +************************************************************************ + +8. Litigation +------------- + +Any litigation relating to this License may be brought only in the +courts of a jurisdiction where the defendant maintains its principal +place of business and such litigation shall be governed by laws of that +jurisdiction, without reference to its conflict-of-law provisions. +Nothing in this Section shall prevent a party's ability to bring +cross-claims or counter-claims. + +9. Miscellaneous +---------------- + +This License represents the complete agreement concerning the subject +matter hereof. If any provision of this License is held to be +unenforceable, such provision shall be reformed only to the extent +necessary to make it enforceable. Any law or regulation which provides +that the language of a contract shall be construed against the drafter +shall not be used to construe this License against a Contributor. + +10. Versions of the License +--------------------------- + +10.1. New Versions + +Mozilla Foundation is the license steward. Except as provided in Section +10.3, no one other than the license steward has the right to modify or +publish new versions of this License. Each version will be given a +distinguishing version number. + +10.2. Effect of New Versions + +You may distribute the Covered Software under the terms of the version +of the License under which You originally received the Covered Software, +or under the terms of any subsequent version published by the license +steward. + +10.3. Modified Versions + +If you create software not governed by this License, and you want to +create a new license for such software, you may create and use a +modified version of this License if you rename the license and remove +any references to the name of the license steward (except to note that +such modified license differs from this License). + +10.4. Distributing Source Code Form that is Incompatible With Secondary +Licenses + +If You choose to distribute Source Code Form that is Incompatible With +Secondary Licenses under the terms of this version of the License, the +notice described in Exhibit B of this License must be attached. + +Exhibit A - Source Code Form License Notice +------------------------------------------- + + This Source Code Form is subject to the terms of the Mozilla Public + License, v. 2.0. If a copy of the MPL was not distributed with this + file, You can obtain one at http://mozilla.org/MPL/2.0/. + +If it is not possible or desirable to put the notice in a particular +file, then You may include the notice in a location (such as a LICENSE +file in a relevant directory) where a recipient would be likely to look +for such a notice. + +You may add additional accurate notices of copyright ownership. + +Exhibit B - "Incompatible With Secondary Licenses" Notice +--------------------------------------------------------- + + This Source Code Form is "Incompatible With Secondary Licenses", as + defined by the Mozilla Public License, v. 2.0. + diff --git a/vendor/github.com/riverqueue/river/rivershared/baseservice/base_service.go b/vendor/github.com/riverqueue/river/rivershared/baseservice/base_service.go new file mode 100644 index 0000000000..95d81f04d4 --- /dev/null +++ b/vendor/github.com/riverqueue/river/rivershared/baseservice/base_service.go @@ -0,0 +1,161 @@ +// Package baseservice contains structs and initialization functions for +// "service-like" objects that provide commonly needed facilities so that they +// don't have to be redefined on every struct. The word "service" is used quite +// loosely here in that it may be applied to many long-lived object that aren't +// strictly services (e.g. adapters). +package baseservice + +import ( + "log/slog" + "reflect" + "regexp" + "strings" + "time" + + "github.com/riverqueue/river/rivertype" +) + +// Archetype contains the set of base service properties that are immutable, or +// otherwise safe for services to copy from another service. The struct is also +// embedded in BaseService, so these properties are available on services +// directly. +type Archetype struct { + // Logger is a structured logger. + Logger *slog.Logger + + // Time returns a time generator for in-process deadline/duration math and + // optional stubbed wall-clock timestamps in tests. + // + // The production path intentionally uses `time.Now()` rather than + // `time.Now().UTC()`: per the Go time package's monotonic clock semantics, + // normalizing through `UTC()` strips the monotonic reading. Services should + // use this clock for local timing math and normalize to UTC only at + // database or serialization boundaries. + Time TimeGeneratorWithStub +} + +// NewArchetype returns a new archetype. This function is most suitable for +// non-test usage wherein nothing should be stubbed. +func NewArchetype(logger *slog.Logger) *Archetype { + return &Archetype{ + Logger: logger, + Time: &UnStubbableTimeGenerator{}, + } +} + +// BaseService is a struct that's meant to be embedded on "service-like" objects +// (e.g. client, producer, queue maintainer) and which provides a number of +// convenient properties that are widely needed so that they don't have to be +// defined on every individual service and can easily be copied from each other. +// +// An initial Archetype should be defined near the program's entrypoint +// (currently in Client), and then each service should invoke Init along with +// the archetype to initialize its own base service. This is often done in the +// service's constructor, but if it doesn't have one, it's the job of the caller +// which instantiates it to invoke Init. +type BaseService struct { + Archetype + + // Name is a name of the service. It should generally be used to prefix all + // log lines the service emits. + Name string +} + +func (s *BaseService) GetBaseService() *BaseService { return s } + +// WithBaseService is an interface to a struct that embeds BaseService. An +// implementation is provided automatically by BaseService, and it's largely +// meant for internal use. +type WithBaseService interface { + GetBaseService() *BaseService +} + +// Init initializes a base service from an archetype. It returns the same +// service that was passed into it for convenience. +func Init[TService WithBaseService](archetype *Archetype, service TService) TService { + var ( + baseService = service.GetBaseService() + serviceType = reflect.TypeOf(service).Elem() + ) + + baseService.Logger = archetype.Logger + baseService.Name = lastPkgPathSegmentIfNotRiver(serviceType.PkgPath()) + simplifyLogName(serviceType.Name()) + baseService.Time = archetype.Time + + return service +} + +type TimeGeneratorWithStub interface { + rivertype.TimeGenerator + + // StubNow stubs the current wall-clock time. It will panic if invoked + // outside of tests. Returns the same time passed as parameter for + // convenience. + StubNow(now time.Time) time.Time +} + +// TimeGeneratorWithStubWrapper provides a wrapper around TimeGenerator that +// implements missing TimeGeneratorWithStub functions. This is used so that we +// only need to expose the minimal TimeGenerator interface publicly, but can +// keep a stubbable version of widely available for internal use. +type TimeGeneratorWithStubWrapper struct { + rivertype.TimeGenerator +} + +func (g *TimeGeneratorWithStubWrapper) StubNow(now time.Time) time.Time { + panic("time not stubbable outside tests") +} + +// UnStubbableTimeGenerator is a TimeGenerator implementation that can't be +// stubbed. It's always the generator used outside of tests. +type UnStubbableTimeGenerator struct{} + +// Now intentionally returns `time.Now()` without calling `.UTC()`. River uses +// this clock for in-process duration and deadline math, and Go strips the +// monotonic clock reading when changing a Time's location with methods like +// `UTC()`. Normalize at database or serialization boundaries instead. +func (g *UnStubbableTimeGenerator) Now() time.Time { return time.Now() } +func (g *UnStubbableTimeGenerator) NowOrNil() *time.Time { return nil } + +func (g *UnStubbableTimeGenerator) StubNow(now time.Time) time.Time { + panic("time not stubbable outside tests") +} + +// Takes a package path and extracts the last part of it to use in a service +// name for logging purposes. If the package is the top-level `river` returns an +// empty string so that top-level structs aren't prefixed but sub-packages +// structs are. +// +// - github.com/riverqueue/river -> "" +// - github.com/riverqueue/river/riverlog -> "riverlog." +// - github.com/riverqueue/riverui -> "riverui." +// +// Helps produce log-friendly service names like `riverlog.Middleware`. +func lastPkgPathSegmentIfNotRiver(pkgPath string) string { + lastSlashIndex := strings.LastIndex(pkgPath, "/") + if lastSlashIndex == -1 { + return "" + } + + lastPart := pkgPath[lastSlashIndex+1:] + if lastPart == "" || lastPart == "river" { + return "" + } + + return lastPart + "." +} + +var stripGenericTypePathRE = regexp.MustCompile(`\[([\[\]\*]*).*/([^/]+)\]`) + +// Simplifies the name of a Go type that uses generics for cleaner logging output. +// +// So this: +// +// QueryCacher[[]*github.com/riverqueue/riverui/internal/dbsqlc.JobCountByStateRow] +// +// Becomes this: +// +// QueryCacher[[]*dbsqlc.JobCountByStateRow] +func simplifyLogName(name string) string { + return stripGenericTypePathRE.ReplaceAllString(name, `[$1$2]`) +} diff --git a/vendor/github.com/riverqueue/river/rivershared/circuitbreaker/circuit_breaker.go b/vendor/github.com/riverqueue/river/rivershared/circuitbreaker/circuit_breaker.go new file mode 100644 index 0000000000..51b9315019 --- /dev/null +++ b/vendor/github.com/riverqueue/river/rivershared/circuitbreaker/circuit_breaker.go @@ -0,0 +1,113 @@ +package circuitbreaker + +import ( + "slices" + "time" + + "github.com/riverqueue/river/rivershared/baseservice" + "github.com/riverqueue/river/rivertype" +) + +// CircuitBreakerOptions are options for CircuitBreaker. +type CircuitBreakerOptions struct { + // Limit is the maximum number of trips/actions allowed within Window + // before the circuit breaker opens. + Limit int + + // Window is the window of time during which Limit number of trips/actions + // can occur before the circuit breaker opens. The window is sliding, and + // actions are reaped as they fall outside of Window compared to the current + // time. + Window time.Duration +} + +func (o *CircuitBreakerOptions) mustValidate() *CircuitBreakerOptions { + if o.Limit < 1 { + panic("CircuitBreakerOptions.Limit must be greater than 0") + } + if o.Window < 1 { + panic("CircuitBreakerOptions.Window must be greater than 0") + } + return o +} + +// CircuitBreaker is a basic implementation of the circuit breaker pattern. Trip +// is called a number of times until the circuit breaker reaches its defined +// limit inside its allowed window at which point the circuit breaker +// opens. Once open, this version of the circuit breaker does not close again +// and stays open indefinitely. The window of time in question is sliding, and +// actions are repeaed as they fall outside of it compared to the current time +// (and don't count towards the breaker's limit). +type CircuitBreaker struct { + open bool + opts *CircuitBreakerOptions + timeGenerator rivertype.TimeGenerator + trips []time.Time +} + +func NewCircuitBreaker(opts *CircuitBreakerOptions) *CircuitBreaker { + return &CircuitBreaker{ + opts: opts.mustValidate(), + timeGenerator: &baseservice.UnStubbableTimeGenerator{}, + } +} + +// Limit returns the configured limit of the circuit breaker. +func (b *CircuitBreaker) Limit() int { + return b.opts.Limit +} + +// Open returns true if the circuit breaker is open (i.e. is broken). +func (b *CircuitBreaker) Open() bool { + return b.open +} + +// ResetIfNotOpen resets the circuit breaker to its empty state if it's not +// already open. i.e. As if no calls to Trip have been invoked at all. If the +// circuit breaker is open, it has no effect. +// +// This may be useful for example in cases where we're trying to track a number +// of consecutive failures before deciding to open. In the case of a success, we +// want to indicate that the series of consecutive failures has ended, so we +// call ResetIfNotOpen to trip it. +// +// Returns true if the breaker was reset (i.e. it had not been open), and false +// otherwise. +func (b *CircuitBreaker) ResetIfNotOpen() bool { + if !b.open { + b.trips = nil + } + return !b.open +} + +// Trip "trips" the circuit breaker by counting an action towards opening it. If +// the action causes the breaker to reach its limit within its window, the +// breaker opens and the function returns true. Subsequent calls to Trip after +// the breaker is open will also return true. +func (b *CircuitBreaker) Trip() bool { + if b.open { + return true + } + + var ( + horizonIndex = -1 + now = b.timeGenerator.Now() + ) + for i, v := range slices.Backward(b.trips) { + if v.Before(now.Add(-b.opts.Window)) { + horizonIndex = i + break + } + } + + if horizonIndex != -1 { + b.trips = b.trips[horizonIndex+1:] + } + + b.trips = append(b.trips, now) + if len(b.trips) >= b.opts.Limit { + b.open = true + } + + return b.open +} diff --git a/vendor/github.com/riverqueue/river/rivershared/levenshtein/License.txt b/vendor/github.com/riverqueue/river/rivershared/levenshtein/License.txt new file mode 100644 index 0000000000..a55defac43 --- /dev/null +++ b/vendor/github.com/riverqueue/river/rivershared/levenshtein/License.txt @@ -0,0 +1,21 @@ +The MIT License (MIT) + +Copyright (c) 2015 Agniva De Sarker + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. \ No newline at end of file diff --git a/vendor/github.com/riverqueue/river/rivershared/levenshtein/levenshtein.go b/vendor/github.com/riverqueue/river/rivershared/levenshtein/levenshtein.go new file mode 100644 index 0000000000..141d4e5a41 --- /dev/null +++ b/vendor/github.com/riverqueue/river/rivershared/levenshtein/levenshtein.go @@ -0,0 +1,86 @@ +// Package levenshtein is a Go implementation to calculate Levenshtein Distance. +// +// Vendored from this repository: +// https://github.com/agnivade/levenshtein +// +// Implementation taken from +// https://gist.github.com/andrei-m/982927#gistcomment-1931258 +package levenshtein + +import "unicode/utf8" + +// minLengthThreshold is the length of the string beyond which +// an allocation will be made. Strings smaller than this will be +// zero alloc. +const minLengthThreshold = 32 + +// ComputeDistance computes the levenshtein distance between the two +// strings passed as an argument. The return value is the levenshtein distance +// +// Works on runes (Unicode code points) but does not normalize +// the input strings. See https://blog.golang.org/normalization +// and the golang.org/x/text/unicode/norm package. +func ComputeDistance(str1, str2 string) int { + if len(str1) == 0 { + return utf8.RuneCountInString(str2) + } + + if len(str2) == 0 { + return utf8.RuneCountInString(str1) + } + + if str1 == str2 { + return 0 + } + + // We need to convert to []rune if the strings are non-ASCII. + // This could be avoided by using utf8.RuneCountInString + // and then doing some juggling with rune indices, + // but leads to far more bounds checks. It is a reasonable trade-off. + runeSlice1 := []rune(str1) + runeSlice2 := []rune(str2) + + // swap to save some memory O(min(a,b)) instead of O(a) + if len(runeSlice1) > len(runeSlice2) { + runeSlice1, runeSlice2 = runeSlice2, runeSlice1 + } + lenRuneSlice1 := len(runeSlice1) + lenRuneSlice2 := len(runeSlice2) + + // Init the row. + var distances []uint16 + if lenRuneSlice1+1 > minLengthThreshold { + distances = make([]uint16, lenRuneSlice1+1) + } else { + // We make a small optimization here for small strings. Because a slice + // of constant length is effectively an array, it does not allocate. So + // we can re-slice it to the right length as long as it is below a + // desired threshold. + distances = make([]uint16, minLengthThreshold) + distances = distances[:lenRuneSlice1+1] + } + + // we start from 1 because index 0 is already 0. + for i := 1; i < len(distances); i++ { + distances[i] = uint16(i) + } + + // Make a dummy bounds check to prevent the 2 bounds check down below. The + // one inside the loop is particularly costly. + _ = distances[lenRuneSlice1] + + // fill in the rest + for i := 1; i <= lenRuneSlice2; i++ { + prev := uint16(i) + for j := 1; j <= lenRuneSlice1; j++ { + current := distances[j-1] // match + if runeSlice2[i-1] != runeSlice1[j-1] { + current = min(min(distances[j-1]+1, prev+1), distances[j]+1) + } + distances[j-1] = prev + prev = current + } + distances[lenRuneSlice1] = prev + } + return int(distances[lenRuneSlice1]) +} diff --git a/vendor/github.com/riverqueue/river/rivershared/riverpilot/pilot.go b/vendor/github.com/riverqueue/river/rivershared/riverpilot/pilot.go new file mode 100644 index 0000000000..69d48d3b14 --- /dev/null +++ b/vendor/github.com/riverqueue/river/rivershared/riverpilot/pilot.go @@ -0,0 +1,188 @@ +package riverpilot + +import ( + "context" + "time" + + "github.com/riverqueue/river/riverdriver" + "github.com/riverqueue/river/rivershared/baseservice" + "github.com/riverqueue/river/rivertype" +) + +// A Pilot bridges the gap between the River client and the driver, implementing +// higher level functionality on top of the driver's underlying queries. It +// tracks closely to the underlying driver's API, but may add additional +// functionality or logic wrapping the queries. +// +// This should be considered a River internal API and its stability is not +// guaranteed. DO NOT USE. +type Pilot interface { + PilotPeriodicJob + + JobCancel(ctx context.Context, exec riverdriver.Executor, params *riverdriver.JobCancelParams) (*rivertype.JobRow, error) + + // JobCleanerQueuesExcluded returns queues that should be excluded from the + // main River client's JobCleaner. If no queues should be omitted, this + // function should return nil as opposed to an empty array. (Because the + // underlying database query uses an `IS NULL` check, though this could + // conceivably be changed.) + JobCleanerQueuesExcluded() []string + + JobGetAvailable( + ctx context.Context, + exec riverdriver.Executor, + state ProducerState, + params *riverdriver.JobGetAvailableParams, + ) ([]*rivertype.JobRow, error) + + JobInsertMany( + ctx context.Context, + exec riverdriver.Executor, + params *riverdriver.JobInsertFastManyParams, + ) ([]*riverdriver.JobInsertFastResult, error) + + JobRetry(ctx context.Context, exec riverdriver.Executor, params *riverdriver.JobRetryParams) (*rivertype.JobRow, error) + + JobSetStateIfRunningMany(ctx context.Context, exec riverdriver.Executor, params *riverdriver.JobSetStateIfRunningManyParams) ([]*rivertype.JobRow, error) + + PilotInit(archetype *baseservice.Archetype, params *PilotInitParams) + + // ProducerInit is called when a producer is started. It should return the ID + // of the new producer, a new state object that will be used to track the + // producer's state, and an error if the producer could not be initialized. + ProducerInit(ctx context.Context, exec riverdriver.Executor, params *ProducerInitParams) (int64, ProducerState, error) + + ProducerKeepAlive(ctx context.Context, exec riverdriver.Executor, params *riverdriver.ProducerKeepAliveParams) error + + ProducerShutdown(ctx context.Context, exec riverdriver.Executor, params *ProducerShutdownParams) error + + QueueMetadataChanged(ctx context.Context, exec riverdriver.Executor, params *QueueMetadataChangedParams) error +} + +// PilotInitParams are parameters for initializing a pilot. +// +// API is not stable. DO NOT USE. +type PilotInitParams struct { + // Insert is the insert implementation from the main client. This is + // used as a low-level insert that shouldn't be accessible via public API, + // but should be accessible to deep integrations. + Insert func(ctx context.Context, tx riverdriver.ExecutorTx, insertParams []*rivertype.JobInsertParams) ([]*rivertype.JobInsertResult, error) + + // NotifyNonTxJobInsert is a special function that should be invoked when a + // client knows that a job has become available and the transaction that + // committed it has finished so that it's possible for a producer to fetch + // it. This is used in special cases like poll-only clients to improve + // latency between job insert and when a job is worked. + NotifyNonTxJobInsert func(ctx context.Context, res []*rivertype.JobInsertResult) + + // ProducerReportInterval is the amount of time between periodic reports of + // producer status. + ProducerReportInterval time.Duration + + // WorkerMetadata is metadata about registered workers as received from the + // client's worker bundle. Only available when a client will work jobs (i.e. + // has Workers configured), so while it's safe to assume the presence of + // this value in places like maintenance services, it's not in all contexts. + WorkerMetadata []*rivertype.WorkerMetadata +} + +func (p *PilotInitParams) Validate() *PilotInitParams { + if p.Insert == nil { + panic("need PilotInitParams.Insert") + } + if p.NotifyNonTxJobInsert == nil { + panic("need PilotInitParams.NotifyNonTxJobInsert ") + } + return p +} + +// PilotJobRescuer contains optional Pilot functionality related to rescuing +// stuck jobs. Pilots that don't implement it fall back to the standard +// executor-backed behavior. +// +// This is temporarily separate from Pilot so implementations built against +// older River versions remain compatible. It can be embedded into Pilot after +// downstream implementations have had a release cycle to adopt it. +// +// Once all supported River Pro releases implement PilotJobRescuer, embed this +// interface into Pilot, replace JobRescuer's capability assertions with direct +// Pilot calls, and remove its executor fallbacks. +type PilotJobRescuer interface { + JobGetStuck(ctx context.Context, exec riverdriver.Executor, params *riverdriver.JobGetStuckParams) ([]*rivertype.JobRow, error) + JobRescueMany(ctx context.Context, exec riverdriver.Executor, params *riverdriver.JobRescueManyParams) (*struct{}, error) +} + +// PilotPeriodicJob contains pilot functions related to periodic jobs. This is +// extracted as its own interface so there's less surface area to mock in places +// like the periodic job enqueuer where that's needed. +type PilotPeriodicJob interface { + // PeriodicJobGetAll gets all currently known periodic jobs. + // + // API is not stable. DO NOT USE. + PeriodicJobGetAll(ctx context.Context, exec riverdriver.Executor, params *PeriodicJobGetAllParams) ([]*PeriodicJob, error) + + // PeriodicJobTouchMany updates the `updated_at` timestamp on many jobs at + // once to keep them alive and reaps any jobs that haven't been seen in some + // time. + // + // API is not stable. DO NOT USE. + PeriodicJobKeepAliveAndReap(ctx context.Context, exec riverdriver.Executor, params *PeriodicJobKeepAliveAndReapParams) ([]*PeriodicJob, error) + + // PeriodicJobUpsertMany upserts many periodic jobs. + // + // API is not stable. DO NOT USE. + PeriodicJobUpsertMany(ctx context.Context, exec riverdriver.Executor, params *PeriodicJobUpsertManyParams) ([]*PeriodicJob, error) +} + +// PeriodicJob represents a durable periodic job. +// +// TODO: Get rid of this in favor of rivertype.PeriodicJob the next time we're +// making River <-> River Pro API contract changes. +type PeriodicJob struct { + ID string + CreatedAt time.Time + NextRunAt time.Time + UpdatedAt time.Time +} + +type PeriodicJobGetAllParams struct { + Schema string +} + +type PeriodicJobKeepAliveAndReapParams struct { + ID []string + Schema string +} + +type PeriodicJobUpsertManyParams struct { + Jobs []*PeriodicJobUpsertParams + Schema string +} + +type PeriodicJobUpsertParams struct { + ID string + NextRunAt time.Time + UpdatedAt time.Time +} + +type ProducerState interface { + JobFinish(job *rivertype.JobRow) +} + +type ProducerInitParams struct { + ClientID string + ProducerID int64 + Queue string + Schema string +} + +type ProducerShutdownParams struct { + ProducerID int64 + Queue string + Schema string +} + +type QueueMetadataChangedParams struct { + Queue string + Metadata []byte +} diff --git a/vendor/github.com/riverqueue/river/rivershared/riverpilot/standard_pilot.go b/vendor/github.com/riverqueue/river/rivershared/riverpilot/standard_pilot.go new file mode 100644 index 0000000000..af2a92fe82 --- /dev/null +++ b/vendor/github.com/riverqueue/river/rivershared/riverpilot/standard_pilot.go @@ -0,0 +1,95 @@ +package riverpilot + +import ( + "context" + "sync/atomic" + + "github.com/riverqueue/river/internal/rivercommon" + "github.com/riverqueue/river/riverdriver" + "github.com/riverqueue/river/rivershared/baseservice" + "github.com/riverqueue/river/rivershared/util/timeoututil" + "github.com/riverqueue/river/rivertype" +) + +type StandardPilot struct { + seq atomic.Int64 +} + +func (p *StandardPilot) JobCleanerQueuesExcluded() []string { return nil } + +func (p *StandardPilot) JobGetAvailable(ctx context.Context, exec riverdriver.Executor, state ProducerState, params *riverdriver.JobGetAvailableParams) ([]*rivertype.JobRow, error) { + if params.MaxToLock <= 0 { + return nil, nil + } + + return timeoututil.WithTimeoutV(ctx, rivercommon.HotOperationTimeout, "StandardPilot.JobGetAvailable", func(ctx context.Context) ([]*rivertype.JobRow, error) { + return exec.JobGetAvailable(ctx, params) + }) +} + +func (p *StandardPilot) JobGetStuck(ctx context.Context, exec riverdriver.Executor, params *riverdriver.JobGetStuckParams) ([]*rivertype.JobRow, error) { + return exec.JobGetStuck(ctx, params) +} + +func (p *StandardPilot) JobCancel(ctx context.Context, exec riverdriver.Executor, params *riverdriver.JobCancelParams) (*rivertype.JobRow, error) { + return exec.JobCancel(ctx, params) +} + +func (p *StandardPilot) JobInsertMany( + ctx context.Context, + exec riverdriver.Executor, + params *riverdriver.JobInsertFastManyParams, +) ([]*riverdriver.JobInsertFastResult, error) { + return exec.JobInsertFastMany(ctx, params) +} + +func (p *StandardPilot) JobRescueMany(ctx context.Context, exec riverdriver.Executor, params *riverdriver.JobRescueManyParams) (*struct{}, error) { + return exec.JobRescueMany(ctx, params) +} + +func (p *StandardPilot) JobRetry(ctx context.Context, exec riverdriver.Executor, params *riverdriver.JobRetryParams) (*rivertype.JobRow, error) { + return exec.JobRetry(ctx, params) +} + +func (p *StandardPilot) JobSetStateIfRunningMany(ctx context.Context, exec riverdriver.Executor, params *riverdriver.JobSetStateIfRunningManyParams) ([]*rivertype.JobRow, error) { + return exec.JobSetStateIfRunningMany(ctx, params) +} + +func (p *StandardPilot) PeriodicJobKeepAliveAndReap(ctx context.Context, exec riverdriver.Executor, params *PeriodicJobKeepAliveAndReapParams) ([]*PeriodicJob, error) { + return nil, nil +} + +func (p *StandardPilot) PeriodicJobGetAll(ctx context.Context, exec riverdriver.Executor, params *PeriodicJobGetAllParams) ([]*PeriodicJob, error) { + return nil, nil +} + +func (p *StandardPilot) PeriodicJobUpsertMany(ctx context.Context, exec riverdriver.Executor, params *PeriodicJobUpsertManyParams) ([]*PeriodicJob, error) { + return nil, nil +} + +func (p *StandardPilot) PilotInit(archetype *baseservice.Archetype, params *PilotInitParams) { + // No-op +} + +func (p *StandardPilot) ProducerInit(ctx context.Context, exec riverdriver.Executor, params *ProducerInitParams) (int64, ProducerState, error) { + id := p.seq.Add(1) + return id, &standardProducerState{}, nil +} + +func (p *StandardPilot) ProducerKeepAlive(ctx context.Context, exec riverdriver.Executor, params *riverdriver.ProducerKeepAliveParams) error { + return nil +} + +func (p *StandardPilot) ProducerShutdown(ctx context.Context, exec riverdriver.Executor, params *ProducerShutdownParams) error { + return nil +} + +func (p *StandardPilot) QueueMetadataChanged(ctx context.Context, exec riverdriver.Executor, params *QueueMetadataChangedParams) error { + return nil +} + +type standardProducerState struct{} + +func (s *standardProducerState) JobFinish(job *rivertype.JobRow) { + // No-op +} diff --git a/vendor/github.com/riverqueue/river/rivershared/riversharedmaintenance/river_shared_maintenance.go b/vendor/github.com/riverqueue/river/rivershared/riversharedmaintenance/river_shared_maintenance.go new file mode 100644 index 0000000000..824873001f --- /dev/null +++ b/vendor/github.com/riverqueue/river/rivershared/riversharedmaintenance/river_shared_maintenance.go @@ -0,0 +1,145 @@ +package riversharedmaintenance + +import ( + "cmp" + "context" + "time" + + "github.com/riverqueue/river/rivershared/baseservice" + "github.com/riverqueue/river/rivershared/circuitbreaker" + "github.com/riverqueue/river/rivershared/util/randutil" + "github.com/riverqueue/river/rivershared/util/serviceutil" +) + +// Maintainers will sleep a brief period of time between batches to give the +// database some breathing room. +const ( + BatchBackoffMax = 1 * time.Second + BatchBackoffMin = 50 * time.Millisecond +) + +const ( + LogPrefixRanSuccessfully = ": Ran successfully" + LogPrefixRunLoopStarted = ": Run loop started" + LogPrefixRunLoopStopped = ": Run loop stopped" + + // TimeoutDefault is a reasonable timeout for any large maintenance-related + // queries. Some maintainers may opt to switch to their own timeout, but + // this one should generally be used unless there's a good reason to have a + // specific version. + TimeoutDefault = 30 * time.Second +) + +// Constants related to JobCleaner. +const ( + CancelledJobRetentionPeriodDefault = 24 * time.Hour + CompletedJobRetentionPeriodDefault = 24 * time.Hour + DiscardedJobRetentionPeriodDefault = 7 * 24 * time.Hour + + JobCleanerIntervalDefault = 30 * time.Second + JobCleanerTimeoutDefault = 30 * time.Second +) + +const ( + // BatchSizeDefault is the default batch size of most maintenance services. + // + // Bulk maintenance tasks like job removal operate in batches so that even + // in the event of an enormous backlog of work to do, transactions stay + // relatively short and aren't at risk of cancellation. This number is the + // batch size, or the number of rows that are handled at a time. + // + // The specific value is somewhat arbitrary as large enough to make good + // progress, but not so large as to make the operation overstay its welcome. + // For now it's not configurable because we can likely pick a number that's + // suitable for almost everyone. + // + // In case database degradation is detected, most maintenance services will + // back off to use the smaller batch size BatchSizeReduced. + BatchSizeDefault = 10_000 + + // BatchSizeReduced is the reduced batch size of most maintenance services. + // + // Services start out with a batch size of BatchSizeDefault, but as they + // detect a degraded database may switch to BatchSizeReduced instead so that + // they're trying to do less work per operation. + BatchSizeReduced = 1_000 +) + +// BatchSizes containing batch size information for maintenance services. It's +// mean to be embedded on each service's configuration struct so as to provide a +// common way of organizing and initializing batch sizes for improved +// succinctness. +type BatchSizes struct { + // Default is the maximum number of jobs to transition at once from + // "scheduled" to "available" during periodic scheduling checks. + Default int + + // Reduced is a considerably smaller batch size that the service + // uses after encountering 3 consecutive timeouts in a row. The idea behind + // this is that if it appears the database is degraded, then we start doing + // less work in the hope that it can better succeed. + Reduced int +} + +// MustValidate validates the struct and panics in case a value is invalid. +func (b BatchSizes) MustValidate() BatchSizes { + if b.Default <= 0 { + panic("BatchSizes.Default must be above zero") + } + if b.Reduced <= 0 { + panic("BatchSizes.Reduced must be above zero") + } + return b +} + +// WithDefaults returns the struct with any configuration overrides that were +// already set, but sets defaults for any that were zero values. +func (b BatchSizes) WithDefaults() BatchSizes { + return BatchSizes{ + Default: cmp.Or(b.Default, BatchSizeDefault), + Reduced: cmp.Or(b.Reduced, BatchSizeReduced), + } +} + +// ReducedBatchSizeBreaker returns a reduced batch circuit breaker suitable for +// use in most maintenance services. After being tripped three consecutive times +// inside a ten minute window it switches to a reduced batch size. A success at +// any point between failures will reset it, but after the circuit breaker has +// tripped, it stays tripped for the life time of the program. +func ReducedBatchSizeBreaker(batchSizes BatchSizes) *circuitbreaker.CircuitBreaker { + return circuitbreaker.NewCircuitBreaker(&circuitbreaker.CircuitBreakerOptions{ + Limit: 3, + Window: 10 * time.Minute, + }) +} + +// QueueMaintainerServiceBase is a struct that should be embedded on all queue +// maintainer services. Its main use is to provide a StaggerStart function that +// should be called on service start to avoid thundering herd problems. +type QueueMaintainerServiceBase struct { + baseservice.BaseService + + staggerStartupDisabled bool +} + +// StaggerStart is called when queue maintainer services start. It jitters by +// sleeping for a short random period so services don't all perform their first +// run at exactly the same time. +func (s *QueueMaintainerServiceBase) StaggerStart(ctx context.Context) { + if s.staggerStartupDisabled { + return + } + + serviceutil.CancellableSleep(ctx, randutil.DurationBetween(0*time.Second, 1*time.Second)) +} + +// StaggerStartupDisable sets whether the short staggered sleep on start up +// is disabled. This is useful in tests where the extra sleep involved in a +// staggered start up is not helpful for test run time. +func (s *QueueMaintainerServiceBase) StaggerStartupDisable(disabled bool) { + s.staggerStartupDisabled = disabled +} + +func (s *QueueMaintainerServiceBase) StaggerStartupIsDisabled() bool { + return s.staggerStartupDisabled +} diff --git a/vendor/github.com/riverqueue/river/rivershared/sqlctemplate/sqlc_template.go b/vendor/github.com/riverqueue/river/rivershared/sqlctemplate/sqlc_template.go new file mode 100644 index 0000000000..1151b7d13a --- /dev/null +++ b/vendor/github.com/riverqueue/river/rivershared/sqlctemplate/sqlc_template.go @@ -0,0 +1,324 @@ +// Package sqlctemplate provides a way of making arbitrary text replacement in +// sqlc queries which normally only allow parameters which are in places valid +// in a prepared statement. For example, it can be used to insert a schema name +// as a prefix to tables referenced in sqlc, which is otherwise impossible. +// +// Replacement is carried out from within invocations of sqlc's generated DBTX +// interface, after sqlc generated code runs, but before queries are executed. +// This is accomplished by implementing DBTX, calling Replacer.Run from within +// them, and injecting parameters in with WithReplacements (which is unfortunately +// the only way of injecting them). +// +// Templates are modeled as SQL comments so that they're still parseable as +// valid SQL. An example use of the basic /* TEMPLATE ... */ syntax: +// +// -- name: JobCountByState :one +// SELECT count(*) +// FROM /* TEMPLATE: schema */river_job +// WHERE state = @state; +// +// An open/close syntax is also available for when SQL is required before +// processing for the query to be valid. For example, a WHERE or ORDER BY clause +// can't be empty, so the SQL includes a sentinel value that's parseable which +// is then replaced later with template values: +// +// -- name: JobList :many +// SELECT * +// FROM river_job +// WHERE /* TEMPLATE_BEGIN: where_clause */ true /* TEMPLATE_END */ +// ORDER BY /* TEMPLATE_BEGIN: order_by_clause */ id /* TEMPLATE_END */ +// LIMIT @max::int; +// +// Be careful not to place a template on a line by itself because sqlc will +// strip any lines that start with a comment. For example, this does NOT work: +// +// -- name: JobList :many +// SELECT * +// FROM river_job +// /* TEMPLATE_BEGIN: where_clause */ +// LIMIT @max::int; +package sqlctemplate + +import ( + "context" + "errors" + "fmt" + "maps" + "regexp" + "slices" + "strconv" + "strings" + "sync" + + "github.com/riverqueue/river/rivershared/util/maputil" +) + +// Context container added by WithReplacements. +type contextContainer struct { + // NamedArgs and their values to be replaced after templates in Replacements + // are rendered. + NamedArgs map[string]any + + // Replacements maps template names to replacement values. + Replacements map[string]Replacement +} + +type contextKey struct{} + +// Replacement defines a replacement for a template value in some input SQL. +type Replacement struct { + // Stable is whether the replacement value is expected to be stable for any + // number of times Replacer.Run is called with the same given input SQL. If + // all replacements are stable, then the output of Replacer.Run is cached so + // that it doesn't have to be processed again. Replacements should be not be + // stable if they depend on input parameters. + Stable bool + + // Value is the value which the template should be replaced with. For a /* + // TEMPLATE ... */ tag, replaces template and the comment containing it. For + // a /* TEMPLATE_BEGIN ... */ ... /* TEMPLATE_END */ tag pair, replaces both + // templates, comments, and the value between them. + Value string +} + +// Replacer replaces templates with template values. As an optimization, it +// contains an internal cache to short circuit SQL that has entirely stable +// template replacements and whose output is invariant of input parameters. +// +// The struct is written so that it's safe to use as a value and doesn't need to +// be initialized with a constructor. This lets it default to a usable instance +// on drivers that may themselves not be initialized. +type Replacer struct { + cache map[replacerCacheKey]string + cacheMu sync.RWMutex +} + +var ( + templateBeginEndRE = regexp.MustCompile(`/\* TEMPLATE_BEGIN: (.*?) \*/ .*? /\* TEMPLATE_END \*/`) + templateRE = regexp.MustCompile(`/\* TEMPLATE: (.*?) \*/`) +) + +// Regex to search for in SQL after replacement has occurred and which probably +// represents a syntax error. sqlctemplate isn't a true compiler so if template +// REs don't match, we can be left with some subtle bugs where there's some +// minor problem like a missing semicolon that are hard to debug. +var postReplaceMistakeRE = regexp.MustCompile(`\/\*\s*TEMPLATE([A-Z0-9_]+)?`) // also finds "/* TEMPLATE_BEGIN" + +// Run replaces any tempates in input SQL with values from context added via +// WithReplacements. +// +// args aren't used for replacements in the input SQL, but are needed to +// determine which placeholder number (e.g. $1, $2, $3, ...) we should start +// with to replace any template named args. The returned args value should then +// be used as query input as named args from context may have been added to it. +// +// argPlaceholder is the character to use as a placeholder like "$" in "$1" or +// "$2". This should be a "$" for Postgres, but a "?" for SQLite. +func (r *Replacer) Run(ctx context.Context, argPlaceholder, sql string, args []any) (string, []any) { + sql, namedArgs, err := r.RunSafely(ctx, argPlaceholder, sql, args) + if err != nil { + panic(err) + } + return sql, namedArgs +} + +// RunSafely is the same as Run, but returns an error in case of missing or +// extra templates. +func (r *Replacer) RunSafely(ctx context.Context, argPlaceholder, sql string, args []any) (string, []any, error) { + var ( + container, containerOK = ctx.Value(contextKey{}).(*contextContainer) + sqlContainsTemplate = strings.Contains(sql, "/* TEMPLATE") + ) + switch { + case !containerOK && !sqlContainsTemplate: + // Neither context container or template in SQL; short circuit fast because there's no work to do. + return sql, args, nil + + case containerOK && !sqlContainsTemplate: + return "", nil, errors.New("sqlctemplate found context container but SQL contains no templates; bug?") + + case !containerOK && sqlContainsTemplate: + return "", nil, errors.New("sqlctemplate found template(s) in SQL, but no context container; bug?") + } + + cacheKey, cacheEligible := replacerCacheKeyFrom(sql, container) + if cacheEligible { + r.cacheMu.RLock() + var ( + cachedSQL string + cachedSQLOK bool + ) + if r.cache != nil { // protect against map not initialized yet + cachedSQL, cachedSQLOK = r.cache[cacheKey] + } + r.cacheMu.RUnlock() + + // If all input templates were stable, the finished SQL will have been cached. + if cachedSQLOK { + if len(container.NamedArgs) > 0 { + // Named args must be appended in sorted order to match the + // placeholder positions baked into the cached SQL during + // RunSafely's cache miss path. + sortedNamedArgs := maputil.Keys(container.NamedArgs) + slices.Sort(sortedNamedArgs) + for _, name := range sortedNamedArgs { + args = append(args, container.NamedArgs[name]) + } + } + return cachedSQL, args, nil + } + } + + var ( + templatesExpected = maputil.Keys(container.Replacements) + templatesMissing []string // not preallocated because we don't expect any missing parameters in the common case + ) + + replaceTemplate := func(sql string, templateRE *regexp.Regexp) string { + return templateRE.ReplaceAllStringFunc(sql, func(templateStr string) string { + // Really dumb, but Go doesn't provide any way to get submatches in a + // function, so we have to match twice. + // https://github.com/golang/go/issues/5690 + matches := templateRE.FindStringSubmatch(templateStr) + + template := matches[1] + + if replacement, ok := container.Replacements[template]; ok { + templatesExpected = slices.DeleteFunc(templatesExpected, func(p string) bool { return p == template }) + return replacement.Value + } else { + templatesMissing = append(templatesMissing, template) + } + + return templateStr + }) + } + + updatedSQL := sql + updatedSQL = replaceTemplate(updatedSQL, templateBeginEndRE) + updatedSQL = replaceTemplate(updatedSQL, templateRE) + + if len(templatesExpected) > 0 { + return "", nil, errors.New("sqlctemplate params present in context but missing in SQL: " + strings.Join(templatesExpected, ", ")) + } + + if len(templatesMissing) > 0 { + return "", nil, errors.New("sqlctemplate params present in SQL but missing in context: " + strings.Join(templatesMissing, ", ")) + } + + probableMistakes := postReplaceMistakeRE.FindAllString(updatedSQL, -1) + if len(probableMistakes) > 0 { + return "", nil, errors.New("sqlctemplate found template-like tag after replacements; probably syntax error or missing end tag: " + strings.Join(probableMistakes, ", ")) + } + + if len(container.NamedArgs) > 0 { + placeholderNum := len(args) + + // For the benefit of the test suite's output being predictable, sort + // named args before processing them. + sortedNamedArgs := maputil.Keys(container.NamedArgs) + slices.Sort(sortedNamedArgs) + for _, arg := range sortedNamedArgs { + placeholderNum++ + + var ( + symbol = "@" + arg + symbolIndex = strings.Index(updatedSQL, symbol) + val = container.NamedArgs[arg] + ) + + if symbolIndex == -1 { + return "", nil, fmt.Errorf("sqltemplate expected to find named arg %q, but it wasn't present", symbol) + } + + // ReplaceAll because an input parameter may appear multiple times. + updatedSQL = strings.ReplaceAll(updatedSQL, symbol, argPlaceholder+strconv.Itoa(placeholderNum)) + args = append(args, val) + } + } + + if cacheEligible { + r.cacheMu.Lock() + if r.cache == nil { + r.cache = make(map[replacerCacheKey]string) + } + r.cache[cacheKey] = updatedSQL + r.cacheMu.Unlock() + } + + return updatedSQL, args, nil +} + +// WithReplacements adds sqlctemplate templates to the given context (they go in +// context because it's the only way to get them down into the innards of sqlc). +// namedArgs can also be passed in to replace arguments found in +// +// If sqlctemplate params are already present in context, the two sets are +// merged, with the new params taking precedent. +func WithReplacements(ctx context.Context, replacements map[string]Replacement, namedArgs map[string]any) context.Context { + if container, ok := ctx.Value(contextKey{}).(*contextContainer); ok { + maps.Copy(container.NamedArgs, namedArgs) + maps.Copy(container.Replacements, replacements) + return ctx + } + + if namedArgs == nil { + namedArgs = make(map[string]any) + } + + return context.WithValue(ctx, contextKey{}, &contextContainer{ + NamedArgs: namedArgs, + Replacements: replacements, + }) +} + +// Comparable struct that's used as a key for template caching. +type replacerCacheKey struct { + namedArgs string // all arg names concatenated together + replacementValues string // all values concatenated together + sql string +} + +// Builds a cache key for the given SQL and context container. +// +// A key is only built if the given SQL/templates are cacheable, which means all +// template values must be stable. The second return value is a boolean +// indicating whether a cache key was built or not. If false, the input is not +// eligible for caching, and no check against the cache should be made. +func replacerCacheKeyFrom(sql string, container *contextContainer) (replacerCacheKey, bool) { + // Only eligible for caching if all replacements are stable. + for _, replacement := range container.Replacements { + if !replacement.Stable { + return replacerCacheKey{}, false + } + } + + var ( + namedArgsBuilder strings.Builder + + // Named args must be sorted for key stability because Go maps don't + // provide any ordering guarantees. + sortedNamedArgs = maputil.Keys(container.NamedArgs) + ) + slices.Sort(sortedNamedArgs) + for _, namedArg := range sortedNamedArgs { + namedArgsBuilder.WriteRune('@') // useful as separator because not valid in the name of a named arg + namedArgsBuilder.WriteString(namedArg) + } + + var ( + replacementValuesBuilder strings.Builder + sortedReplacements = maputil.Keys(container.Replacements) + ) + slices.Sort(sortedReplacements) + for _, template := range sortedReplacements { + replacementValuesBuilder.WriteRune('•') // use a separator that SQL would reject under most circumstances (this may be imperfect) + replacementValuesBuilder.WriteString(container.Replacements[template].Value) + } + + return replacerCacheKey{ + namedArgs: namedArgsBuilder.String(), + replacementValues: replacementValuesBuilder.String(), + sql: sql, + }, true +} diff --git a/vendor/github.com/riverqueue/river/rivershared/startstop/start_stop.go b/vendor/github.com/riverqueue/river/rivershared/startstop/start_stop.go new file mode 100644 index 0000000000..1ca9b51e97 --- /dev/null +++ b/vendor/github.com/riverqueue/river/rivershared/startstop/start_stop.go @@ -0,0 +1,344 @@ +package startstop + +import ( + "context" + "errors" + "sync" +) + +// ErrStop is an error injected into WithCancelCause when context is canceled +// because a service is stopping. Makes it possible to differentiate a +// controlled stop from a context cancellation. +var ErrStop = errors.New("service stopped") + +// Service is a generalized interface for a service that starts and stops, +// usually one backed by embedding BaseStartStop. +type Service interface { + // Start starts a service. Services are responsible for backgrounding + // themselves, so this function should be invoked synchronously. Services + // may return an error if they have trouble starting up, so the caller + // should wait and respond to the error if necessary. + Start(ctx context.Context) error + + // Started returns a channel that's closed when a service finishes starting, + // or if failed to start and is stopped instead. It can be used in + // conjunction with WaitAllStarted to verify startup of a constellation of + // services. + Started() <-chan struct{} + + // Stop stops a service. Services are responsible for making sure their stop + // is complete before returning so a caller can wait on this invocation + // synchronously and be guaranteed the service is fully stopped. Services + // are expected to be able to tolerate (1) being stopped without having been + // started, and (2) being double-stopped. + Stop() +} + +// ServiceWithStopped is a Service that can also return a Stopped channel. I've +// kept this as a separate interface for the time being because I'm not sure +// this is strictly necessary to be part of startstop. +type serviceWithStopped interface { + Service + + // Stopped returns a channel that can be waited on for the service to be + // stopped. This function is only safe to invoke after successfully waiting on a + // service's Start, and a reference to it must be taken _before_ invoking Stop. + Stopped() <-chan struct{} +} + +// BaseStartStop is a helper that can be embedded on a queue maintenance service +// and which will provide the basic necessities to safely implement the Service +// interface in a way that's not racy and can tolerate a number of edge cases. +// It's packaged separately so that it doesn't leak its internal variables into +// services that use it. +// +// Services should implement their own Start function which invokes StartInit +// first thing, return if told not to start, spawn a goroutine with their main +// run block otherwise, and make sure to defer a close on the stop channel +// returned by StartInit within that goroutine. +// +// A Stop implementation is provided automatically and it's not necessary to +// override it. +type BaseStartStop struct { + cancelFunc context.CancelCauseFunc + isRunning bool + mu sync.Mutex + started chan struct{} + stopped chan struct{} +} + +// StartInit should be invoked at the beginning of a service's Start function. +// It returns a context for the service to use, a boolean indicating whether it +// should start (which will be false if the service is already started), and a +// stopped channel. Services should defer a close on the stop channel in their +// main run loop. +// +// func (s *Service) Start(ctx context.Context) error { +// ctx, shouldStart, stopped := s.StartInit(ctx) +// if !shouldStart { +// return nil +// } +// +// go func() { +// defer close(stopped) +// +// <-ctx.Done() +// +// ... +// }() +// +// return nil +// } +// +// Be careful to also close it in the event of startup errors, otherwise a +// service that failed to start once will never be able to start up. +// +// ctx, shouldStart, stopped := s.StartInit(ctx) +// if !shouldStart { +// return nil +// } +// +// if err := possibleStartUpError(); err != nil { +// close(stopped) +// return err +// } +// +// ... +func (s *BaseStartStop) StartInit(ctx context.Context) (context.Context, bool, func(), func()) { + s.mu.Lock() + defer s.mu.Unlock() + + if s.isRunning { + // If stopped has already been closed (e.g. a previous Start failed and + // called stopped()), reset state so the service can start again. + // + // Notably, for this branch to be taken, Stop will not have been called. + // If it was, isRunning will have been set to false via finalizeStop. + if s.stopped != nil { + select { + case <-s.stopped: + s.isRunning = false + s.started = nil + s.stopped = nil + default: + } + } + + if s.isRunning { + return ctx, false, nil, nil + } + } + + s.isRunning = true + + // Only allocate a started or stopped channels when not preallocated by + // Started or Stopped. + if s.started == nil { + s.started = make(chan struct{}) + } + if s.stopped == nil { + s.stopped = make(chan struct{}) + } + + ctx, s.cancelFunc = context.WithCancelCause(ctx) + + closeStartedOnce := sync.OnceFunc(func() { close(s.started) }) + + return ctx, true, closeStartedOnce, func() { + // Also close the started channel (in case it wasn't already), just in + // case `started()` was never invoked and someone is waiting on it. + closeStartedOnce() + + close(s.stopped) + } +} + +// Started returns a channel that's closed when a service finishes starting, or +// if failed to start and is stopped instead. It can be used in conjunction with +// WaitAllStarted to verify startup of a constellation of services. +func (s *BaseStartStop) Started() <-chan struct{} { + s.mu.Lock() + defer s.mu.Unlock() + + // If the call to Started is before the service was actually started, + // preallocate the started channel so that regardless of whether the wait + // started before or after the service started, it will still do the right + // thing. + if s.started == nil { + s.started = make(chan struct{}) + } + + return s.started +} + +// Stop is an automatically provided implementation for the maintenance Service +// interface's Stop. +func (s *BaseStartStop) Stop() { + shouldStop, stopped, finalizeStop := s.StopInit() + if !shouldStop { + return + } + + <-stopped + finalizeStop(true) +} + +// StopInit provides a way to build a more customized Stop implementation. It +// should be avoided unless there's an exceptional reason not to because Stop +// should be fine in the vast majority of situations. +// +// It returns a boolean indicating whether the service should do any additional +// work to stop (false is returned if the service was never started), a stopped +// channel to wait on for full stop, and a finalizeStop function that should be +// deferred in the stop function to ensure that locks are cleaned up and the +// struct is reset after stopping. +// +// func (s *Service) Stop(ctx context.Context) error { +// shouldStop, stopped, finalizeStop := s.StopInit(ctx) +// if !shouldStop { +// return +// } +// +// defer finalizeStop(true) +// +// ... +// } +// +// finalizeStop takes a boolean which indicates where the service should indeed +// be considered stopped. This should usually be true, but callers can pass +// false to cancel the stop action, keeping the service from starting again, and +// potentially allowing the service to try another stop. +func (s *BaseStartStop) StopInit() (bool, <-chan struct{}, func(didStop bool)) { + s.mu.Lock() + + // Tolerate being told to stop without having been started. + if !s.isRunning { + s.mu.Unlock() + return false, nil, func(didStop bool) {} + } + + s.cancelFunc(ErrStop) + + return true, s.stopped, func(didStop bool) { + defer s.mu.Unlock() + if didStop { + s.isRunning = false + s.started = nil + s.stopped = nil + } + } +} + +// Stopped returns a channel that can be waited on for the service to be +// stopped. This function may be used to return a stopped channel before a +// service is started or while it's running, but a reference to it must be taken +// _before_ invoking Stop. +func (s *BaseStartStop) Stopped() <-chan struct{} { + s.mu.Lock() + defer s.mu.Unlock() + + // If the call to Stopped is before the service was actually started, + // preallocate the stopped channel so that regardless of whether the wait + // started before or after the service started, it will still do the right + // thing. + if s.stopped == nil { + s.stopped = make(chan struct{}) + } + + return s.stopped +} + +// StoppedUnsafe returns a channel that can be waited on for the service to be +// stopped. +// +// Unlike Stopped, this returns the struct's internal channel directly without +// preallocation and without taking a lock on the mutex (making it safe to call +// while StopInit is ongoing). Most users of BaseStartStop shouldn't use this +// variant and it basically exists for river.Client so it can provide slightly +// different stop channel semantics compared to BaseStartStop's. +func (s *BaseStartStop) StoppedUnsafe() <-chan struct{} { return s.stopped } + +type startStopFunc struct { + BaseStartStop + + startFunc func(ctx context.Context, shouldStart bool, started, stopped func()) error +} + +// StartStopFunc produces a `startstop.Service` from a function. It's useful for +// very small services that don't necessarily need a whole struct defined for +// them. +func StartStopFunc(startFunc func(ctx context.Context, shouldStart bool, started, stopped func()) error) *startStopFunc { + return &startStopFunc{ + startFunc: startFunc, + } +} + +func (s *startStopFunc) Start(ctx context.Context) error { + return s.startFunc(s.StartInit(ctx)) +} + +// StartAll starts all given services. If any service returns an error while +// being started, that error is returned, and any services that were started +// successfully up to that point are stopped. +func StartAll(ctx context.Context, services ...Service) error { + for i, service := range services { + if err := service.Start(ctx); err != nil { + StopAllParallel(services[0:i]...) + + return err + } + } + return nil +} + +// StopAllParallel stops all the given services in parallel and waits until +// they've all stopped successfully. +func StopAllParallel(services ...Service) { + var wg sync.WaitGroup + wg.Add(len(services)) + + for i := range services { + service := services[i] + go func() { + defer wg.Done() + service.Stop() + }() + } + + wg.Wait() +} + +// WaitAllStarted waits until all the given services are started (or stopped in +// a degenerate start scenario, like if context is cancelled while starting up). +// +// Unlike StopAllParallel, WaitAllStarted doesn't bother with parallelism +// because the services themselves have already backgrounded themselves, and we +// have to wait until the slowest service has started anyway. +func WaitAllStarted(services ...Service) { + <-WaitAllStartedC(services...) +} + +// WaitAllStartedC waits until all the given services are started (or stopped in +// a degenerate start scenario, like if context is cancelled while starting up). +// +// This variant returns a channel so that a caller can apply a timeout branch +// with `select` if they'd like. For the most part this shouldn't be needed +// though, as long as each service individually is confirmed to be able to start +// and stop itself in a healthy way. (i.e. Never dies for any reason before +// managing to call `started()` or `stopped()`). +// +// Unlike StopAllParallel, WaitAllStartedC doesn't bother with parallelism +// because the services themselves have already background themselves, and we +// have to wait until the slowest service has started anyway. +func WaitAllStartedC(services ...Service) <-chan struct{} { + allStarted := make(chan struct{}) + + go func() { + defer close(allStarted) + for _, service := range services { + <-service.Started() + } + }() + + return allStarted +} diff --git a/vendor/github.com/riverqueue/river/rivershared/structtag/struct_tag.go b/vendor/github.com/riverqueue/river/rivershared/structtag/struct_tag.go new file mode 100644 index 0000000000..433ee02fb6 --- /dev/null +++ b/vendor/github.com/riverqueue/river/rivershared/structtag/struct_tag.go @@ -0,0 +1,196 @@ +package structtag + +import ( + "fmt" + "reflect" + "sort" + "strings" + "sync" + + "github.com/tidwall/gjson" + + "github.com/riverqueue/river/rivertype" +) + +// ExtractValues extracts the raw JSON values of the specified keys from the JSON-encoded args. +func ExtractValues(encodedArgs []byte, uniqueKeys []string) []string { + // Use GetManyBytes to retrieve multiple values at once + results := gjson.GetManyBytes(encodedArgs, uniqueKeys...) + + uniqueValues := make([]string, len(results)) + for i, res := range results { + if res.Exists() { + uniqueValues[i] = res.Raw // Use Raw to get the JSON-encoded value + } else { + // Handle missing keys as "undefined" (they'll be skipped when + // building the key). We don't want to use "null" here because the + // JSON may actually contain "null" as a value. + uniqueValues[i] = "undefined" + } + } + + return uniqueValues +} + +type uniqueFieldCacheKey struct { + typ reflect.Type + tagValue string +} + +var ( + // uniqueFieldsCache caches the unique fields for each JobArgs type. These are + // global to ensure that each struct type's tags are only extracted once. + uniqueFieldsCache = make(map[uniqueFieldCacheKey][]string) //nolint:gochecknoglobals + cacheMutex sync.RWMutex //nolint:gochecknoglobals +) + +// SortedFieldsWithTag retrieves unique fields with caching to avoid +// extracting fields from the same struct type repeatedly. +func SortedFieldsWithTag(args rivertype.JobArgs, tagValue string) ([]string, error) { + var ( + typ = reflect.TypeOf(args) + cacheKey = uniqueFieldCacheKey{typ: typ, tagValue: tagValue} + ) + + // Check cache first + cacheMutex.RLock() + if fields, ok := uniqueFieldsCache[cacheKey]; ok { + cacheMutex.RUnlock() + return fields, nil + } + cacheMutex.RUnlock() + + // Not in cache; retrieve using reflection + fields, err := sortedFieldsWithTagUncached(reflect.TypeOf(args), tagValue, nil, make(map[reflect.Type]struct{})) + if err != nil { + return nil, err + } + + // Store in cache + cacheMutex.Lock() + uniqueFieldsCache[cacheKey] = fields + cacheMutex.Unlock() + + return fields, nil +} + +// sortedFieldsWithTagUncached uses reflection to retrieve the JSON keys of fields +// marked with `river:""` among potentially other comma-separated +// values. The return values are the JSON keys using the same logic as the +// `json` struct tag. +// +// typesSeen should be a map passed through to make sure that recursive types +// don't cause a stack overflow. +func sortedFieldsWithTagUncached(typ reflect.Type, tagValue string, path []string, typesSeen map[reflect.Type]struct{}) ([]string, error) { + // Handle pointer to struct + if typ.Kind() == reflect.Pointer { + typ = typ.Elem() + } + + // Ensure we're dealing with a struct + if typ.Kind() != reflect.Struct { + return nil, fmt.Errorf("expected struct, got %T", typ.Name()) + } + + // Stop when encountering a recursive type. This has the effect of the + // entire subfield's value being extracted by gjson, but this is about as + // right of a way to handle it as any other I can think of. + if _, ok := typesSeen[typ]; ok { + return nil, nil + } + typesSeen[typ] = struct{}{} + + var uniqueFields []string + + // Iterate over all fields + for i := range typ.NumField() { + field := typ.Field(i) + + if !field.IsExported() { + continue + } + + var uniqueName string + { + // Get the corresponding JSON key + jsonTag := field.Tag.Get("json") + + if jsonTag == "" { + // If no JSON tag, use the field name as-is + uniqueName = field.Name + } else { + // Handle cases like `json:"recipient,omitempty"` + uniqueName = parseJSONTag(jsonTag) + } + } + + // Check for `river:"unique"` tag, possibly among other comma-separated values + var hasUniqueTag bool + if riverTag, ok := field.Tag.Lookup("river"); ok { + tags := strings.SplitSeq(riverTag, ",") + for tag := range tags { + if strings.TrimSpace(tag) == tagValue { + hasUniqueTag = true + } + } + } + + if typeStructOrPointerToStruct(field.Type) { + // Append the JSON to the path (all path segments sent down + // recursively) unless we're looking at an anonymous struct, whose + // fields will be let at the top level. + fullPath := path + if !field.Anonymous { + fullPath = append(path, uniqueName) //nolint:gocritic + } + + uniqueSubFields, err := sortedFieldsWithTagUncached(field.Type, tagValue, fullPath, typesSeen) + if err != nil { + return nil, err + } + + if len(uniqueSubFields) > 0 { + uniqueFields = append(uniqueFields, uniqueSubFields...) + } else if hasUniqueTag { + // If a struct field is marked `river:""`, use its entire + // JSON serialization as a unique value. This may not be the + // greatest idea practically, but keeping it in place for + // backwards compatibility. + uniqueFields = append(uniqueFields, strings.Join(append(path, uniqueName), ".")) + } + + continue + } + + if hasUniqueTag { + uniqueFields = append(uniqueFields, strings.Join(append(path, uniqueName), ".")) + } + } + + // Sort the uniqueFields alphabetically for consistent ordering + sort.Strings(uniqueFields) + + return uniqueFields, nil +} + +// parseJSONTag extracts the JSON key from the struct tag. +// It handles tags with options, e.g., `json:"recipient,omitempty"`. +func parseJSONTag(tag string) string { + // Tags can be like "recipient,omitempty", so split by comma + if before, _, ok := strings.Cut(tag, ","); ok { + return before + } + return tag +} + +func typeStructOrPointerToStruct(typ reflect.Type) bool { + if typ.Kind() == reflect.Struct { + return true + } + + if typ.Kind() == reflect.Pointer && typ.Elem().Kind() == reflect.Struct { + return true + } + + return false +} diff --git a/vendor/github.com/riverqueue/river/rivershared/testsignal/test_signal.go b/vendor/github.com/riverqueue/river/rivershared/testsignal/test_signal.go new file mode 100644 index 0000000000..c46ca6102e --- /dev/null +++ b/vendor/github.com/riverqueue/river/rivershared/testsignal/test_signal.go @@ -0,0 +1,99 @@ +package testsignal + +import ( + "time" + + "github.com/riverqueue/river/rivershared/util/testutil" +) + +// TestSignalWaiter provides an interface for TestSignal which only exposes +// waiting on the signal. This is useful for minimizing functionality across +// package boundaries. +type TestSignalWaiter[T any] interface { + WaitOrTimeout() T +} + +// TestSignal is a channel wrapper designed to allow tests to wait on certain +// events (to test difficult concurrent conditions without intermittency) while +// also having minimal impact on the production code that calls into it. +// +// Its default value produces a state where its safe to call Signal to signal +// into it, but where doing so will have no effect. Entities that embed it +// should by convention provide a TestSignalsInit function that tests can invoke +// and which calls Init on all member test signals, after which it becomes +// possible for tests to WaitOrTimeout on them. +type TestSignal[T any] struct { + internalChan chan T + tb testutil.TestingTB +} + +const testSignalInternalChanSize = 50 + +// Init initializes the test signal for use. This should only ever be called +// from tests. +func (s *TestSignal[T]) Init(tb testutil.TestingTB) { + s.internalChan = make(chan T, testSignalInternalChanSize) + s.tb = tb +} + +// RequireEmpty requires that the test signal be empty (i.e. have not received +// any values). +func (s *TestSignal[T]) RequireEmpty() { + if s.internalChan == nil { + panic("test only signal is not initialized; called outside of tests?") + } + + select { + case val := <-s.internalChan: + s.tb.Errorf("test signal should be empty, but wasn't\ngot value: %v\n", val) + default: + } +} + +// Signal signals the test signal. In production where the signal hasn't been +// initialized, this no ops harmlessly. In tests, the value is written to an +// internal asynchronous channel which can be waited with WaitOrTimeout. +func (s *TestSignal[T]) Signal(val T) { + // Occurs in the case of a raw signal that hasn't been initialized (which is + // what should always be happening outside of tests). + if s.internalChan == nil { + return + } + + select { // never block on send + case s.internalChan <- val: + default: + s.tb.Errorf("test only signal channel is full") + } +} + +// WaitC returns a channel on which a value from the test signal can be waited +// upon. +func (s TestSignal[T]) WaitC() <-chan T { + if s.internalChan == nil { + panic("test only signal is not initialized; called outside of tests?") + } + + return s.internalChan +} + +// WaitOrTimeout waits on the next value injected by Signal. This should only be +// used in tests, and can only be used if Init has been invoked on the test +// signal. +func (s *TestSignal[T]) WaitOrTimeout() T { + if s.internalChan == nil { + panic("test only signal is not initialized; called outside of tests?") + } + + timeout := testutil.WaitTimeout() + + select { + case val := <-s.internalChan: + return val + case <-time.After(timeout): + s.tb.Errorf("timed out waiting on test signal after %s", timeout) + } + + var val T + return val +} diff --git a/vendor/github.com/riverqueue/river/rivershared/uniquestates/unique_states.go b/vendor/github.com/riverqueue/river/rivershared/uniquestates/unique_states.go new file mode 100644 index 0000000000..9e3db7dab8 --- /dev/null +++ b/vendor/github.com/riverqueue/river/rivershared/uniquestates/unique_states.go @@ -0,0 +1,47 @@ +package uniquestates + +import ( + "slices" + + "github.com/riverqueue/river/rivertype" +) + +func UniqueBitmaskToStates(mask byte) []rivertype.JobState { + var states []rivertype.JobState + + for state, bitIndex := range jobStateBitPositions { + bitPosition := 7 - (bitIndex % 8) + if mask&(1< 0 { + return input + } + } + return nil +} + +// GroupBy returns an object composed of keys generated from the results of +// running each element of collection through keyFunc. +func GroupBy[T any, U comparable](collection []T, keyFunc func(T) U) map[U][]T { + result := map[U][]T{} + + for _, item := range collection { + key := keyFunc(item) + + result[key] = append(result[key], item) + } + + return result +} + +// KeyBy converts a slice into a map using the key/value tuples returned by +// tupleFunc. If any two pairs would have the same key, the last one wins. Go +// maps are unordered and the order of the new map isn't guaranteed to the same +// as the original slice. +func KeyBy[T any, K comparable, V any](collection []T, tupleFunc func(item T) (K, V)) map[K]V { + result := make(map[K]V, len(collection)) + + for _, t := range collection { + k, v := tupleFunc(t) + result[k] = v + } + + return result +} + +// Map manipulates a slice and transforms it to a slice of another type. +func Map[T any, R any](collection []T, mapFunc func(T) R) []R { + result := make([]R, len(collection)) + + for i, item := range collection { + result[i] = mapFunc(item) + } + + return result +} + +// MapError manipulates a slice and transforms it to a slice of another type, +// returning the first error that occurred invoking the map function, if there +// was one. +func MapError[T any, R any](collection []T, mapFunc func(T) (R, error)) ([]R, error) { + result := make([]R, len(collection)) + + for i, item := range collection { + var err error + result[i], err = mapFunc(item) + if err != nil { + return nil, err + } + } + + return result, nil +} + +// Uniq returns a duplicate-free version of an array, in which only the first occurrence of each element is kept. +// The order of result values is determined by the order they occur in the array. +func Uniq[T comparable](collection []T) []T { + result := make([]T, 0, len(collection)) + seen := make(map[T]struct{}, len(collection)) + + for _, item := range collection { + if _, ok := seen[item]; ok { + continue + } + + seen[item] = struct{}{} + result = append(result, item) + } + + return result +} diff --git a/vendor/github.com/riverqueue/river/rivershared/util/testutil/job_args_reflect_kind.go b/vendor/github.com/riverqueue/river/rivershared/util/testutil/job_args_reflect_kind.go new file mode 100644 index 0000000000..bd8eb9e392 --- /dev/null +++ b/vendor/github.com/riverqueue/river/rivershared/util/testutil/job_args_reflect_kind.go @@ -0,0 +1,33 @@ +package testutil + +import "reflect" + +// JobArgsReflectKind can be embedded on a struct to implement JobArgs such that +// the job's kind will be the name of TKind. Typically, for convenience TKind +// will be the same type that does the embedding. Use of JobArgsReflectKind may +// not be typical, but in combination with WorkFunc, it allows the entirety of a +// job args and worker pair to be implemented inside the body of a function. +// +// type InFuncWorkFuncArgs struct { +// testutil.JobArgsReflectKind[InFuncWorkFuncArgs] +// Message `json:"message"` +// } +// +// AddWorker(client.config.Workers, WorkFunc(func(ctx context.Context, job *Job[WorkFuncArgs]) error { +// ... +// +// Its major downside compared to a normal JobArgs implementation is that it's +// possible to easily break things accidentally by renaming its type, deploying, +// and then finding that the worker will no longer work any jobs that were +// inserted before the deploy. It also depends on reflection, which likely makes +// it marginally slower. +// +// We're not sure yet whether it's appropriate to expose this publicly, so for +// now we've localized it to the test suite only. When a test case needs a job +// type that won't be reused, it's preferable to make use of JobArgsReflectKind +// so the type doesn't pollute the global namespace. +type JobArgsReflectKind[TKind any] struct{} + +func (a JobArgsReflectKind[TKind]) Kind() string { + return reflect.TypeFor[JobArgsReflectKind[TKind]]().Name() +} diff --git a/vendor/github.com/riverqueue/river/rivershared/util/testutil/test_util.go b/vendor/github.com/riverqueue/river/rivershared/util/testutil/test_util.go new file mode 100644 index 0000000000..b8f6d51f63 --- /dev/null +++ b/vendor/github.com/riverqueue/river/rivershared/util/testutil/test_util.go @@ -0,0 +1,141 @@ +package testutil + +import ( + "bytes" + "fmt" + "io" + "os" + "time" +) + +// See docs on PanicTB. +type panicTB struct{} + +// PanicTB is an implementation for testing.TB that panics when an error is +// logged or FailNow is called. This is useful to inject into test helpers in +// example tests where no *testing.T is available. +// +// If env is set with `RIVER_DEBUG=true`, output is logged to os.Stderr (Stderr +// instead of Stdout to not interfere with example test output). +// +// Doesn't fully implement testing.TB. Functions where it's used should take the +// more streamlined TestingTB instead. +func PanicTB() *panicTB { + return &panicTB{} +} + +func (tb *panicTB) Errorf(format string, args ...any) { + panic(fmt.Sprintf(format, args...)) +} + +func (tb *panicTB) FailNow() { + panic("FailNow invoked") +} + +func (tb *panicTB) Helper() {} + +func (tb *panicTB) Log(args ...any) { + logOut := tb.maybeDebugOut() + if logOut != nil { + fmt.Fprintln(logOut, args...) + } +} + +func (tb *panicTB) Logf(format string, args ...any) { + logOut := tb.maybeDebugOut() + if logOut != nil { + fmt.Fprintf(logOut, format+"\n", args...) + } +} + +func (tb *panicTB) Name() string { return "panicTB" } + +func (tb *panicTB) maybeDebugOut() io.Writer { + if os.Getenv("RIVER_DEBUG") == "1" || os.Getenv("RIVER_DEBUG") == "true" { + // Send output to stderr so it doesn't interfere with example tests. + return os.Stderr + } + + return nil +} + +// MockT mocks TestingTB. It's used to let us verify our test helpers. +type MockT struct { + Failed bool + logOutput bytes.Buffer + tb TestingTB +} + +// NewMockT initializes a new MockT. It takes another TestingTB which is usually +// something like a *testing.T and where logs are emitted to along with being +// internalized and retrievable on LogOutput. +func NewMockT(tb TestingTB) *MockT { + tb.Helper() + return &MockT{tb: tb} +} + +func (t *MockT) Errorf(format string, args ...any) { + // Errorf is equivalent to Log + Fail + fmt.Fprintf(&t.logOutput, format, args...) + t.logOutput.WriteString("\n") + t.Failed = true +} + +func (t *MockT) FailNow() { + t.Failed = true +} + +func (t *MockT) Helper() {} + +func (t *MockT) Log(args ...any) { + t.tb.Log(args...) + + fmt.Fprint(&t.logOutput, args...) + t.logOutput.WriteString("\n") +} + +func (t *MockT) Logf(format string, args ...any) { + t.tb.Logf(format, args...) + + fmt.Fprintf(&t.logOutput, format, args...) + t.logOutput.WriteString("\n") +} + +func (t *MockT) LogOutput() string { + return t.logOutput.String() +} + +func (t *MockT) Name() string { return "MockT" } + +// TestingTB is an interface wrapper around *testing.T that's implemented by all +// of *testing.T, *testing.F, and *testing.B. +// +// It's used internally to verify that River's test assertions are working as +// expected. +type TestingTB interface { + Errorf(format string, args ...any) + FailNow() + Helper() + Log(args ...any) + Logf(format string, args ...any) + Name() string +} + +// WaitTimeout returns a duration broadly appropriate for waiting on an expected +// event in a test, and which is used for `TestSignal.WaitOrTimeout` in +// testsignal and `WaitOrTimeout` in riversharedtest. Its main purpose is to +// allow a little extra leeway in GitHub Actions where we occasionally seem to +// observe subpar performance which leads to timeouts and test intermittency, +// while still keeping a tight a timeout for local test runs where this is never +// a problem. +// +// It lives here instead of riversharedtest so that testsignal, which is +// compiled into production binaries, can use it without importing +// riversharedtest and its heavier test-only dependencies. +func WaitTimeout() time.Duration { + if os.Getenv("GITHUB_ACTIONS") == "true" { + return 10 * time.Second + } + + return 3 * time.Second +} diff --git a/vendor/github.com/riverqueue/river/rivershared/util/timeoututil/timeout_util.go b/vendor/github.com/riverqueue/river/rivershared/util/timeoututil/timeout_util.go new file mode 100644 index 0000000000..1093da0113 --- /dev/null +++ b/vendor/github.com/riverqueue/river/rivershared/util/timeoututil/timeout_util.go @@ -0,0 +1,45 @@ +package timeoututil + +import ( + "context" + "errors" + "fmt" + "time" +) + +// WithTimeout runs innerFunc with a timeout. +// +// If innerFunc returns context.DeadlineExceeded because this helper's local +// timeout fired, WithTimeout returns an error that includes operation and wraps +// context.DeadlineExceeded. This makes timeout errors easier to trace back to +// the specific River operation that introduced the timeout instead of surfacing +// only the generic "context deadline exceeded" message. +func WithTimeout(ctx context.Context, timeout time.Duration, operation string, innerFunc func(ctx context.Context) error) error { + _, err := WithTimeoutV(ctx, timeout, operation, func(ctx context.Context) (struct{}, error) { + return struct{}{}, innerFunc(ctx) + }) + return err +} + +// WithTimeoutV runs innerFunc with a timeout and returns its value. +// +// If innerFunc returns context.DeadlineExceeded because this helper's local +// timeout fired, WithTimeoutV returns an error that includes operation and +// wraps context.DeadlineExceeded. This makes timeout errors easier to trace +// back to the specific River operation that introduced the timeout instead of +// surfacing only the generic "context deadline exceeded" message. +func WithTimeoutV[T any](ctx context.Context, timeout time.Duration, operation string, innerFunc func(ctx context.Context) (T, error)) (T, error) { + // need a specific, local error that we can recognize in case multiple + // levels of these helpers are wrapped within one another + timeoutErr := fmt.Errorf("timeoututil.WithTimeout: %w", context.DeadlineExceeded) + + ctx, cancel := context.WithTimeoutCause(ctx, timeout, timeoutErr) + defer cancel() + + ret, err := innerFunc(ctx) + if err != nil && errors.Is(err, context.DeadlineExceeded) && errors.Is(context.Cause(ctx), timeoutErr) { + var zero T + return zero, fmt.Errorf("%s timed out after %s: %w", operation, timeout, err) + } + return ret, err +} diff --git a/vendor/github.com/riverqueue/river/rivershared/util/timeutil/time_util.go b/vendor/github.com/riverqueue/river/rivershared/util/timeutil/time_util.go new file mode 100644 index 0000000000..d80fb1aff3 --- /dev/null +++ b/vendor/github.com/riverqueue/river/rivershared/util/timeutil/time_util.go @@ -0,0 +1,73 @@ +package timeutil + +import ( + "context" + "time" +) + +// SecondsAsDuration is a simple shortcut for converting seconds represented as +// a float to a `time.Duration`. +func SecondsAsDuration(seconds float64) time.Duration { + return time.Duration(seconds * float64(time.Second)) +} + +// TickerWithInitialTick is similar to `time.Ticker`, except that it fires once +// immediately upon initialization. It also respects context cancellation and +// prefers to be stopped that way rather than an explicit `Stop` function. +type TickerWithInitialTick struct { + // C fires once on initial startup, then after each interval has passed. + C <-chan time.Time + + interval time.Duration + tickChan chan time.Time +} + +// NewTickerWithInitialTick creates a new ticker similar to `time.Ticker`, +// except that it fires once immediately upon initialization. It also respects +// context cancellation and prefers to be stopped that way rather than an +// explicit `Stop` function. +func NewTickerWithInitialTick(ctx context.Context, interval time.Duration) *TickerWithInitialTick { + // Channel of size one combined with non-blocking send are modeled on how + // Go's internal ticker works. Ticks may be dropped if the caller falls behind. + tickChan := make(chan time.Time, 1) + + timer := &TickerWithInitialTick{ + C: tickChan, + interval: interval, + tickChan: tickChan, + } + go timer.runLoop(ctx) + return timer +} + +// Sends a non-blocking tick into the ticker's channel. Ticks may be dropped if +// the caller falls behind. +func (t *TickerWithInitialTick) nonBlockingTick(tm time.Time) { + select { + case t.tickChan <- tm: + default: + } +} + +func (t *TickerWithInitialTick) runLoop(ctx context.Context) { + // Return immediately if context is done. + select { + case <-ctx.Done(): + return + default: + } + + t.nonBlockingTick(time.Now()) + + ticker := time.NewTicker(t.interval) + for { + select { + case <-ctx.Done(): + ticker.Stop() + return + + case tm := <-ticker.C: + t.nonBlockingTick(tm) + } + } +} diff --git a/vendor/github.com/riverqueue/river/rivershared/util/valutil/val_util.go b/vendor/github.com/riverqueue/river/rivershared/util/valutil/val_util.go new file mode 100644 index 0000000000..aebeb479f9 --- /dev/null +++ b/vendor/github.com/riverqueue/river/rivershared/util/valutil/val_util.go @@ -0,0 +1,24 @@ +package valutil + +// ValOrDefault returns the given value if it's non-zero, and otherwise returns +// the default. +// +// Deprecated: Use `cmp.Or` instead. This function will be removed in a near +// future version. +func ValOrDefault[T comparable](val, defaultVal T) T { + var zero T + if val != zero { + return val + } + return defaultVal +} + +// ValOrDefaultFunc returns the given value if it's non-zero, and otherwise +// invokes defaultFunc to produce a default value. +func ValOrDefaultFunc[T comparable](val T, defaultFunc func() T) T { + var zero T + if val != zero { + return val + } + return defaultFunc() +} diff --git a/vendor/github.com/riverqueue/river/rivertype/LICENSE b/vendor/github.com/riverqueue/river/rivertype/LICENSE new file mode 100644 index 0000000000..2f8ed188e8 --- /dev/null +++ b/vendor/github.com/riverqueue/river/rivertype/LICENSE @@ -0,0 +1,374 @@ +Mozilla Public License Version 2.0 +================================== + +1. Definitions +-------------- + +1.1. "Contributor" + means each individual or legal entity that creates, contributes to + the creation of, or owns Covered Software. + +1.2. "Contributor Version" + means the combination of the Contributions of others (if any) used + by a Contributor and that particular Contributor's Contribution. + +1.3. "Contribution" + means Covered Software of a particular Contributor. + +1.4. "Covered Software" + means Source Code Form to which the initial Contributor has attached + the notice in Exhibit A, the Executable Form of such Source Code + Form, and Modifications of such Source Code Form, in each case + including portions thereof. + +1.5. "Incompatible With Secondary Licenses" + means + + (a) that the initial Contributor has attached the notice described + in Exhibit B to the Covered Software; or + + (b) that the Covered Software was made available under the terms of + version 1.1 or earlier of the License, but not also under the + terms of a Secondary License. + +1.6. "Executable Form" + means any form of the work other than Source Code Form. + +1.7. "Larger Work" + means a work that combines Covered Software with other material, in + a separate file or files, that is not Covered Software. + +1.8. "License" + means this document. + +1.9. "Licensable" + means having the right to grant, to the maximum extent possible, + whether at the time of the initial grant or subsequently, any and + all of the rights conveyed by this License. + +1.10. "Modifications" + means any of the following: + + (a) any file in Source Code Form that results from an addition to, + deletion from, or modification of the contents of Covered + Software; or + + (b) any new file in Source Code Form that contains any Covered + Software. + +1.11. "Patent Claims" of a Contributor + means any patent claim(s), including without limitation, method, + process, and apparatus claims, in any patent Licensable by such + Contributor that would be infringed, but for the grant of the + License, by the making, using, selling, offering for sale, having + made, import, or transfer of either its Contributions or its + Contributor Version. + +1.12. "Secondary License" + means either the GNU General Public License, Version 2.0, the GNU + Lesser General Public License, Version 2.1, the GNU Affero General + Public License, Version 3.0, or any later versions of those + licenses. + +1.13. "Source Code Form" + means the form of the work preferred for making modifications. + +1.14. "You" (or "Your") + means an individual or a legal entity exercising rights under this + License. For legal entities, "You" includes any entity that + controls, is controlled by, or is under common control with You. For + purposes of this definition, "control" means (a) the power, direct + or indirect, to cause the direction or management of such entity, + whether by contract or otherwise, or (b) ownership of more than + fifty percent (50%) of the outstanding shares or beneficial + ownership of such entity. + +2. License Grants and Conditions +-------------------------------- + +2.1. Grants + +Each Contributor hereby grants You a world-wide, royalty-free, +non-exclusive license: + +(a) under intellectual property rights (other than patent or trademark) + Licensable by such Contributor to use, reproduce, make available, + modify, display, perform, distribute, and otherwise exploit its + Contributions, either on an unmodified basis, with Modifications, or + as part of a Larger Work; and + +(b) under Patent Claims of such Contributor to make, use, sell, offer + for sale, have made, import, and otherwise transfer either its + Contributions or its Contributor Version. + +2.2. Effective Date + +The licenses granted in Section 2.1 with respect to any Contribution +become effective for each Contribution on the date the Contributor first +distributes such Contribution. + +2.3. Limitations on Grant Scope + +The licenses granted in this Section 2 are the only rights granted under +this License. No additional rights or licenses will be implied from the +distribution or licensing of Covered Software under this License. +Notwithstanding Section 2.1(b) above, no patent license is granted by a +Contributor: + +(a) for any code that a Contributor has removed from Covered Software; + or + +(b) for infringements caused by: (i) Your and any other third party's + modifications of Covered Software, or (ii) the combination of its + Contributions with other software (except as part of its Contributor + Version); or + +(c) under Patent Claims infringed by Covered Software in the absence of + its Contributions. + +This License does not grant any rights in the trademarks, service marks, +or logos of any Contributor (except as may be necessary to comply with +the notice requirements in Section 3.4). + +2.4. Subsequent Licenses + +No Contributor makes additional grants as a result of Your choice to +distribute the Covered Software under a subsequent version of this +License (see Section 10.2) or under the terms of a Secondary License (if +permitted under the terms of Section 3.3). + +2.5. Representation + +Each Contributor represents that the Contributor believes its +Contributions are its original creation(s) or it has sufficient rights +to grant the rights to its Contributions conveyed by this License. + +2.6. Fair Use + +This License is not intended to limit any rights You have under +applicable copyright doctrines of fair use, fair dealing, or other +equivalents. + +2.7. Conditions + +Sections 3.1, 3.2, 3.3, and 3.4 are conditions of the licenses granted +in Section 2.1. + +3. Responsibilities +------------------- + +3.1. Distribution of Source Form + +All distribution of Covered Software in Source Code Form, including any +Modifications that You create or to which You contribute, must be under +the terms of this License. You must inform recipients that the Source +Code Form of the Covered Software is governed by the terms of this +License, and how they can obtain a copy of this License. You may not +attempt to alter or restrict the recipients' rights in the Source Code +Form. + +3.2. Distribution of Executable Form + +If You distribute Covered Software in Executable Form then: + +(a) such Covered Software must also be made available in Source Code + Form, as described in Section 3.1, and You must inform recipients of + the Executable Form how they can obtain a copy of such Source Code + Form by reasonable means in a timely manner, at a charge no more + than the cost of distribution to the recipient; and + +(b) You may distribute such Executable Form under the terms of this + License, or sublicense it under different terms, provided that the + license for the Executable Form does not attempt to limit or alter + the recipients' rights in the Source Code Form under this License. + +3.3. Distribution of a Larger Work + +You may create and distribute a Larger Work under terms of Your choice, +provided that You also comply with the requirements of this License for +the Covered Software. If the Larger Work is a combination of Covered +Software with a work governed by one or more Secondary Licenses, and the +Covered Software is not Incompatible With Secondary Licenses, this +License permits You to additionally distribute such Covered Software +under the terms of such Secondary License(s), so that the recipient of +the Larger Work may, at their option, further distribute the Covered +Software under the terms of either this License or such Secondary +License(s). + +3.4. Notices + +You may not remove or alter the substance of any license notices +(including copyright notices, patent notices, disclaimers of warranty, +or limitations of liability) contained within the Source Code Form of +the Covered Software, except that You may alter any license notices to +the extent required to remedy known factual inaccuracies. + +3.5. Application of Additional Terms + +You may choose to offer, and to charge a fee for, warranty, support, +indemnity or liability obligations to one or more recipients of Covered +Software. However, You may do so only on Your own behalf, and not on +behalf of any Contributor. You must make it absolutely clear that any +such warranty, support, indemnity, or liability obligation is offered by +You alone, and You hereby agree to indemnify every Contributor for any +liability incurred by such Contributor as a result of warranty, support, +indemnity or liability terms You offer. You may include additional +disclaimers of warranty and limitations of liability specific to any +jurisdiction. + +4. Inability to Comply Due to Statute or Regulation +--------------------------------------------------- + +If it is impossible for You to comply with any of the terms of this +License with respect to some or all of the Covered Software due to +statute, judicial order, or regulation then You must: (a) comply with +the terms of this License to the maximum extent possible; and (b) +describe the limitations and the code they affect. Such description must +be placed in a text file included with all distributions of the Covered +Software under this License. Except to the extent prohibited by statute +or regulation, such description must be sufficiently detailed for a +recipient of ordinary skill to be able to understand it. + +5. Termination +-------------- + +5.1. The rights granted under this License will terminate automatically +if You fail to comply with any of its terms. However, if You become +compliant, then the rights granted under this License from a particular +Contributor are reinstated (a) provisionally, unless and until such +Contributor explicitly and finally terminates Your grants, and (b) on an +ongoing basis, if such Contributor fails to notify You of the +non-compliance by some reasonable means prior to 60 days after You have +come back into compliance. Moreover, Your grants from a particular +Contributor are reinstated on an ongoing basis if such Contributor +notifies You of the non-compliance by some reasonable means, this is the +first time You have received notice of non-compliance with this License +from such Contributor, and You become compliant prior to 30 days after +Your receipt of the notice. + +5.2. If You initiate litigation against any entity by asserting a patent +infringement claim (excluding declaratory judgment actions, +counter-claims, and cross-claims) alleging that a Contributor Version +directly or indirectly infringes any patent, then the rights granted to +You by any and all Contributors for the Covered Software under Section +2.1 of this License shall terminate. + +5.3. In the event of termination under Sections 5.1 or 5.2 above, all +end user license agreements (excluding distributors and resellers) which +have been validly granted by You or Your distributors under this License +prior to termination shall survive termination. + +************************************************************************ +* * +* 6. Disclaimer of Warranty * +* ------------------------- * +* * +* Covered Software is provided under this License on an "as is" * +* basis, without warranty of any kind, either expressed, implied, or * +* statutory, including, without limitation, warranties that the * +* Covered Software is free of defects, merchantable, fit for a * +* particular purpose or non-infringing. The entire risk as to the * +* quality and performance of the Covered Software is with You. * +* Should any Covered Software prove defective in any respect, You * +* (not any Contributor) assume the cost of any necessary servicing, * +* repair, or correction. This disclaimer of warranty constitutes an * +* essential part of this License. No use of any Covered Software is * +* authorized under this License except under this disclaimer. * +* * +************************************************************************ + +************************************************************************ +* * +* 7. Limitation of Liability * +* -------------------------- * +* * +* Under no circumstances and under no legal theory, whether tort * +* (including negligence), contract, or otherwise, shall any * +* Contributor, or anyone who distributes Covered Software as * +* permitted above, be liable to You for any direct, indirect, * +* special, incidental, or consequential damages of any character * +* including, without limitation, damages for lost profits, loss of * +* goodwill, work stoppage, computer failure or malfunction, or any * +* and all other commercial damages or losses, even if such party * +* shall have been informed of the possibility of such damages. This * +* limitation of liability shall not apply to liability for death or * +* personal injury resulting from such party's negligence to the * +* extent applicable law prohibits such limitation. Some * +* jurisdictions do not allow the exclusion or limitation of * +* incidental or consequential damages, so this exclusion and * +* limitation may not apply to You. * +* * +************************************************************************ + +8. Litigation +------------- + +Any litigation relating to this License may be brought only in the +courts of a jurisdiction where the defendant maintains its principal +place of business and such litigation shall be governed by laws of that +jurisdiction, without reference to its conflict-of-law provisions. +Nothing in this Section shall prevent a party's ability to bring +cross-claims or counter-claims. + +9. Miscellaneous +---------------- + +This License represents the complete agreement concerning the subject +matter hereof. If any provision of this License is held to be +unenforceable, such provision shall be reformed only to the extent +necessary to make it enforceable. Any law or regulation which provides +that the language of a contract shall be construed against the drafter +shall not be used to construe this License against a Contributor. + +10. Versions of the License +--------------------------- + +10.1. New Versions + +Mozilla Foundation is the license steward. Except as provided in Section +10.3, no one other than the license steward has the right to modify or +publish new versions of this License. Each version will be given a +distinguishing version number. + +10.2. Effect of New Versions + +You may distribute the Covered Software under the terms of the version +of the License under which You originally received the Covered Software, +or under the terms of any subsequent version published by the license +steward. + +10.3. Modified Versions + +If you create software not governed by this License, and you want to +create a new license for such software, you may create and use a +modified version of this License if you rename the license and remove +any references to the name of the license steward (except to note that +such modified license differs from this License). + +10.4. Distributing Source Code Form that is Incompatible With Secondary +Licenses + +If You choose to distribute Source Code Form that is Incompatible With +Secondary Licenses under the terms of this version of the License, the +notice described in Exhibit B of this License must be attached. + +Exhibit A - Source Code Form License Notice +------------------------------------------- + + This Source Code Form is subject to the terms of the Mozilla Public + License, v. 2.0. If a copy of the MPL was not distributed with this + file, You can obtain one at http://mozilla.org/MPL/2.0/. + +If it is not possible or desirable to put the notice in a particular +file, then You may include the notice in a location (such as a LICENSE +file in a relevant directory) where a recipient would be likely to look +for such a notice. + +You may add additional accurate notices of copyright ownership. + +Exhibit B - "Incompatible With Secondary Licenses" Notice +--------------------------------------------------------- + + This Source Code Form is "Incompatible With Secondary Licenses", as + defined by the Mozilla Public License, v. 2.0. + diff --git a/vendor/github.com/riverqueue/river/rivertype/execution_error.go b/vendor/github.com/riverqueue/river/rivertype/execution_error.go new file mode 100644 index 0000000000..fcbe724fef --- /dev/null +++ b/vendor/github.com/riverqueue/river/rivertype/execution_error.go @@ -0,0 +1,82 @@ +package rivertype + +import ( + "errors" + "fmt" + "time" +) + +var ErrJobCancelledRemotely = JobCancel(errors.New("job cancelled remotely")) + +// JobCancel wraps err and can be returned from a Worker's Work method to cancel +// the job at the end of execution. Regardless of whether or not the job has any +// remaining attempts, this will ensure the job does not execute again. +// +// This function primarily exists for cross module compatibility. Users should +// use river.JobCancel instead. +func JobCancel(err error) error { + return &JobCancelError{err: err} +} + +// JobCancelError is the error type returned by JobCancel. It should not be +// initialized directly, but is returned from the [JobCancel] function and can +// be used for test assertions. +type JobCancelError struct { + err error +} + +func (e *JobCancelError) Error() string { + if e.err == nil { + return "JobCancelError: " + } + // should not ever be called, but add a prefix just in case: + return "JobCancelError: " + e.err.Error() +} + +func (e *JobCancelError) Is(target error) bool { + _, ok := target.(*JobCancelError) + return ok +} + +func (e *JobCancelError) Unwrap() error { return e.err } + +// JobSnoozeError is the error type returned by JobSnooze. It should not be +// initialized directly, but is returned from the [JobSnooze] function and can +// be used for test assertions. +type JobSnoozeError struct { + Duration time.Duration +} + +func (e *JobSnoozeError) Error() string { + // should not ever be called, but add a prefix just in case: + return fmt.Sprintf("JobSnoozeError: %s", e.Duration) +} + +func (e *JobSnoozeError) Is(target error) bool { + _, ok := target.(*JobSnoozeError) + return ok +} + +// UnknownJobKindError is returned when a Client fetches and attempts to +// work a job that has not been registered on the Client's Workers bundle (using +// AddWorker). +type UnknownJobKindError struct { + // Kind is the string that was returned by the JobArgs Kind method. + Kind string +} + +// Error returns the error string. +func (e *UnknownJobKindError) Error() string { + return "job kind is not registered in the client's Workers bundle: " + e.Kind +} + +// Is implements the interface used by errors.Is to determine if errors are +// equivalent. It returns true for any other UnknownJobKindError without +// regard to the Kind string so it is possible to detect this type of error +// with: +// +// errors.Is(err, &UnknownJobKindError{}) +func (e *UnknownJobKindError) Is(target error) bool { + _, ok := target.(*UnknownJobKindError) + return ok +} diff --git a/vendor/github.com/riverqueue/river/rivertype/river_type.go b/vendor/github.com/riverqueue/river/rivertype/river_type.go new file mode 100644 index 0000000000..fe74d70763 --- /dev/null +++ b/vendor/github.com/riverqueue/river/rivertype/river_type.go @@ -0,0 +1,652 @@ +// Package rivertype stores some of the lowest level River primitives so they +// can be shared amongst a number of packages including the top-level river +// package, database drivers, and internal utilities. +package rivertype + +import ( + "context" + "encoding/json" + "errors" + "time" +) + +// MetadataKeyOutput is the metadata key used to store recorded job output. +const MetadataKeyOutput = "output" + +// ErrNotFound is returned when a query by ID does not match any existing +// rows. For example, attempting to cancel a job that doesn't exist will +// return this error. +var ErrNotFound = errors.New("not found") + +// ErrJobRunning is returned when a job is attempted to be deleted while it's +// running. +var ErrJobRunning = errors.New("running jobs cannot be deleted") + +// JobArgs is an interface that should be implemented by the arguments to a job. +// This definition duplicates the JobArgs interface in the river package so that +// it can be used in other packages without creating a circular dependency. +type JobArgs interface { + // Kind returns a unique string that identifies the type of job. It's used to + // determine which worker should work the job. + Kind() string +} + +// JobInsertResult is the result of a job insert, containing the inserted job +// along with some other useful metadata. +type JobInsertResult struct { + // Job is a struct containing the database persisted properties of the + // inserted job. + Job *JobRow + + // UniqueSkippedAsDuplicate is true if for a unique job, the insertion was + // skipped due to an equivalent job matching unique property already being + // present. + UniqueSkippedAsDuplicate bool +} + +// JobRow contains the properties of a job that are persisted to the database. +// Use of `Job[T]` will generally be preferred in user-facing code like worker +// interfaces. +type JobRow struct { + // ID of the job. Generated as part of a Postgres sequence and generally + // ascending in nature, but there may be gaps in it as transactions roll + // back. + ID int64 + + // Attempt is the attempt number of the job. Jobs are inserted at 0, the + // number is incremented to 1 the first time work its worked, and may + // increment further if it errors. Attempt will decrement on snooze so that + // repeated snoozes don't increment this value. + Attempt int + + // AttemptedAt is the time that the job was last worked. Starts out as `nil` + // on a new insert. + AttemptedAt *time.Time + + // AttemptedBy is the set of client IDs that have worked this job. + AttemptedBy []string + + // CreatedAt is when the job record was created. + CreatedAt time.Time + + // EncodedArgs is the job's JobArgs encoded as JSON. + EncodedArgs []byte + + // Errors is a set of errors that occurred when the job was worked, one for + // each attempt. Ordered from earliest error to the latest error. + Errors []AttemptError + + // FinalizedAt is the time at which the job was "finalized", meaning it was + // either completed successfully or errored for the last time such that + // it'll no longer be retried. + FinalizedAt *time.Time + + // Kind uniquely identifies the type of job and instructs which worker + // should work it. It is set at insertion time via `Kind()` on the + // `JobArgs`. + Kind string + + // MaxAttempts is the maximum number of attempts that the job will be tried + // before it errors for the last time and will no longer be worked. + // + // Extracted (in order of precedence) from job-specific InsertOpts + // on Insert, from the worker level InsertOpts from JobArgsWithInsertOpts, + // or from a client's default value. + MaxAttempts int + + // Metadata is a field for storing arbitrary metadata on a job. It should + // always be a valid JSON object payload, and users should not overwrite or + // remove anything stored in this field by River. + Metadata []byte + + // Priority is the priority of the job, with 1 being the highest priority and + // 4 being the lowest. When fetching available jobs to work, the highest + // priority jobs will always be fetched before any lower priority jobs are + // fetched. Note that if your workers are swamped with more high-priority jobs + // then they can handle, lower priority jobs may not be fetched. + Priority int + + // Queue is the name of the queue where the job will be worked. Queues can + // be configured independently and be used to isolate jobs. + // + // Extracted from either specific InsertOpts on Insert, or InsertOpts from + // JobArgsWithInsertOpts, or a client's default value. + Queue string + + // ScheduledAt is when the job is scheduled to become available to be + // worked. Jobs default to running immediately, but may be scheduled + // for the future when they're inserted. They may also be scheduled for + // later because they were snoozed or because they errored and have + // additional retry attempts remaining. + ScheduledAt time.Time + + // State is the state of job like `available` or `completed`. Jobs are + // `available` when they're first inserted. + State JobState + + // Tags are an arbitrary list of keywords attached to the job. They don't + // affect job execution, but clients can use them to group and filter jobs. + Tags []string + + // UniqueKey is a unique key for the job within its kind that's used for + // unique job insertions. It's generated by hashing an inserted job's unique + // opts configuration. + UniqueKey []byte + + // UniqueStates is the set of states where uniqueness is enforced for this + // job. Equivalent to the default set of unique states unless + // UniqueOpts.ByState was assigned a custom value. + UniqueStates []JobState +} + +// Output returns the previously recorded output for the job, if any. The return +// value is a raw JSON payload from the output that was recorded by the job, or +// nil if no output was recorded. +func (j *JobRow) Output() []byte { + type metadataWithOutput struct { + Output json.RawMessage `json:"output"` + } + + var metadata metadataWithOutput + if err := json.Unmarshal(j.Metadata, &metadata); err != nil { + return nil + } + + return metadata.Output +} + +// JobState is the state of a job. Jobs start their lifecycle as either +// JobStateAvailable or JobStateScheduled, and if all goes well, transition to +// JobStateCompleted after they're worked. +type JobState string + +const ( + // JobStateAvailable is the state for jobs that are immediately eligible to + // be worked. + JobStateAvailable JobState = "available" + + // JobStateCancelled is the state for jobs that have been manually cancelled + // by user request. + // + // Cancelled jobs are reaped by the job cleaner service after a configured + // amount of time (default 24 hours). + JobStateCancelled JobState = "cancelled" + + // JobStateCompleted is the state for jobs that have successfully run to + // completion. + // + // Completed jobs are reaped by the job cleaner service after a configured + // amount of time (default 24 hours). + JobStateCompleted JobState = "completed" + + // JobStateDiscarded is the state for jobs that have errored enough times + // that they're no longer eligible to be retried. Manual user intervention + // is required for them to be tried again. + // + // Discarded jobs are reaped by the job cleaner service after a configured + // amount of time (default 7 days). + JobStateDiscarded JobState = "discarded" + + // JobStatePending is a state for jobs to be parked while waiting for some + // external action before they can be worked. Jobs in pending will never be + // worked or deleted unless moved out of this state by the user. + JobStatePending JobState = "pending" + + // JobStateRetryable is the state for jobs that have errored, but will be + // retried. + // + // The job scheduler service changes them to JobStateAvailable when they're + // ready to be worked (their `scheduled_at` timestamp comes due). + // + // Jobs that will be retried very soon in the future may be changed to + // JobStateAvailable immediately instead of JobStateRetryable so that they + // don't have to wait for the job scheduler to run. + JobStateRetryable JobState = "retryable" + + // JobStateRunning are jobs which are actively running. + // + // If River can't update state of a running job (in the case of a program + // crash, underlying hardware failure, or job that doesn't return from its + // Work function), that job will be left as JobStateRunning, and will + // require a pass by the job rescuer service to be set back to + // JobStateAvailable and be eligible for another run attempt. + JobStateRunning JobState = "running" + + // JobStateScheduled is the state for jobs that are scheduled for the + // future. + // + // The job scheduler service changes them to JobStateAvailable when they're + // ready to be worked (their `scheduled_at` timestamp comes due). + JobStateScheduled JobState = "scheduled" +) + +// JobStates returns all possible job states. +func JobStates() []JobState { + return []JobState{ + JobStateAvailable, + JobStateCancelled, + JobStateCompleted, + JobStateDiscarded, + JobStatePending, + JobStateRetryable, + JobStateRunning, + JobStateScheduled, + } +} + +// MetricName identifies a metric emitted through HookMetricEmit. +type MetricName string + +const ( + // MetricNameJobGetAvailableDuration is the duration of a successful + // JobGetAvailable call. + MetricNameJobGetAvailableDuration MetricName = "job_get_available_duration" + + // MetricNameJobGetAvailableCount is the number of jobs locked by a + // successful JobGetAvailable call. + MetricNameJobGetAvailableCount MetricName = "job_get_available_count" +) + +// Metric is a strongly typed metric payload emitted through HookMetricEmit. +// +// River provides all Metric implementations. New metric types may be added in +// future versions without changing HookMetricEmit's method signature. +type Metric interface { + // Name identifies the emitted metric. + Name() MetricName + + isMetric() +} + +// AttemptError is an error from a single job attempt that failed due to an +// error or a panic. +type AttemptError struct { + // At is the time at which the error occurred. + At time.Time `json:"at"` + + // Attempt is the attempt number on which the error occurred (maps to + // Attempt on a job row). + Attempt int `json:"attempt"` + + // Error contains the stringified error of an error returned from a job or a + // panic value in case of a panic. + Error string `json:"error"` + + // Trace contains a stack trace from a job that panicked. The trace is + // produced by invoking `debug.Trace()`. + // + // In the case of a non-panic or an error produced as a stuck job was + // rescued, this value will be an empty string. + Trace string `json:"trace"` +} + +type JobInsertParams struct { + ID *int64 + Args JobArgs + CreatedAt *time.Time + EncodedArgs []byte + Kind string + MaxAttempts int + Metadata []byte + Priority int + Queue string + ScheduledAt *time.Time + State JobState + Tags []string + UniqueKey []byte + UniqueStates byte +} + +// Hook is an arbitrary interface for a plugin "hook" which will execute some +// arbitrary code at a predefined step in the job lifecycle. +// +// This interface is left purposely non-specific. Hook structs should embed +// river.HookDefaults to inherit an IsHook implementation, then implement one +// of the more specific hook interfaces like HookInsertBegin or HookWorkBegin. A +// hook struct may also implement multiple specific hook interfaces which are +// logically related and benefit from being grouped together. +// +// Hooks differ from middleware in that they're invoked at a specific lifecycle +// phase, but finish immediately instead of wrapping an inner call like a +// middleware does. One of the main ramifications of this different is that a +// hook cannot modify context in any useful way to pass down into the stack. +// Like a normal function, any changes it makes to its context are discarded on +// return. +// +// All else equal, hooks should generally be preferred over middleware because +// they don't add anything to the call stack. Call stacks that get overly deep +// can become a bit of an operational nightmare because they get hard to read. +// +// In a language with more specific type capabilities, this interface would be a +// union type. In Go we implement it somewhat awkwardly so that we can get +// future extensibility, but also some typing guarantees to prevent misuse (i.e. +// if Hook was an empty interface, then any object could be passed as a hook, +// but having a single function to implement forces the caller to make some +// token motions in the direction of implementing hooks). +// +// List of hook interfaces that may be implemented: +// - HookInsertBegin +// - HookMetricEmit +// - HookPeriodicJobsStart +// - HookWorkBegin +// - HookWorkEnd +// +// More operation-specific interfaces may be added in future versions. +type Hook interface { + // IsHook is a sentinel function to check that a type is implementing Hook + // on purpose and not by accident (Hook would otherwise be an empty + // interface). Hooks should embed river.HookDefaults to pick up an + // implementation for this function automatically. + IsHook() bool +} + +// HookInsertBegin is an interface to a hook that runs before job insertion. +type HookInsertBegin interface { + Hook + + // InsertBegin is invoked just before a job is inserted to the database. + InsertBegin(ctx context.Context, params *JobInsertParams) error +} + +// HookMetricEmit is an interface to a hook that receives metrics emitted by +// River. +type HookMetricEmit interface { + Hook + + // MetricEmit is invoked each time River emits a metric. Metrics are emitted + // in very hot paths like job fetching, and should therefore not block on + // network I/O or anything else, and should usually pass metrics through to + // an asynchronous instrumentation package like OpenTelemetry. + MetricEmit(ctx context.Context, params *HookMetricEmitParams) +} + +// HookMetricEmitParams are parameters for HookMetricEmit. +type HookMetricEmitParams struct { + // Metric is the emitted metric payload. Use a type switch to access + // metric-specific fields. + Metric Metric +} + +// JobGetAvailableDurationMetric is emitted after a successful JobGetAvailable +// call with the call's duration. +type JobGetAvailableDurationMetric struct { + // Duration is how long the JobGetAvailable call took. + Duration time.Duration + + // Queue is the queue that jobs were locked from. + Queue string +} + +func (m *JobGetAvailableDurationMetric) Name() MetricName { + return MetricNameJobGetAvailableDuration +} + +func (m *JobGetAvailableDurationMetric) isMetric() {} + +// JobGetAvailableCountMetric is emitted after a successful JobGetAvailable +// call with the number of jobs locked. +type JobGetAvailableCountMetric struct { + // Count is the number of jobs locked. + Count int + + // Queue is the queue that jobs were locked from. + Queue string +} + +func (m *JobGetAvailableCountMetric) Name() MetricName { + return MetricNameJobGetAvailableCount +} + +func (m *JobGetAvailableCountMetric) isMetric() {} + +// HookPeriodicJobsStart is an interface to a hook that runs when the periodic +// job enqueuer starts on a newly elected leader. +type HookPeriodicJobsStart interface { + Hook + + // Start is invoked when the periodic job enqueuer starts on a newly elected + // leader. + // + // Returning an error will cancel the periodic job enqueuer's start up + // routine. Be very careful with this because if the error is chronic, it + // will prevent any client from successfully starting as leader, thereby + // effectively disabling all maintenance services. + Start(ctx context.Context, params *HookPeriodicJobsStartParams) error +} + +// HookPeriodicJobsStartParams are parameters for HookPeriodicJobsStart. +type HookPeriodicJobsStartParams struct { + // DurableJobs contains a list of durable periodic job records that + // were found in the database. This includes durable jobs that have been + // recently active in an elected periodic job enqueuer, but may also contain + // jobs that've been previously removed, but for which their database record + // has not yet been reaped. + // + // This property will be empty unless durable jobs (a pro feature) are + // enabled. + DurableJobs []*DurablePeriodicJob +} + +// HookWorkBegin is an interface to a hook that runs after a job has been locked +// for work and before it's worked. +type HookWorkBegin interface { + Hook + + // WorkBegin is invoked after a job has been locked and assigned to a + // particular executor for work and just before the job is actually worked. + // + // Returning an error from any HookWorkBegin hook will abort the job early + // such that it has an error set and doesn't work, with a retry scheduled + // according to its retry policy. + // + // This function doesn't return a context so any context set in WorkBegin is + // discarded after the function returns. If persistent context needs to be + // set, middleware should be used instead. + WorkBegin(ctx context.Context, job *JobRow) error +} + +// HookWorkEnd is an interface to a hook that runs after a job has been worked. +type HookWorkEnd interface { + Hook + + // WorkEnd is invoked after a job has been worked with the error result of + // the worked job. It's invoked after any middleware has already run. + // + // WorkEnd may modify a returned work error or pass it through unchanged. + // Each returned error is passed through to the next hook and the final + // error result is returned from the job executor: + // + // err := e.WorkUnit.Work(ctx) + // for _, hook := range hooks { + // err = hook.(rivertype.HookWorkEnd).WorkEnd(ctx, e.JobRow, err) + // } + // return err + // + // If a hook does not want to modify an error result, it should make sure to + // return whatever error value it received as its argument whether that + // error is nil or not. + // + // The JobRow received by WorkEnd is the same one passed to HookWorkBegin's + // WorkBegin. Its state, errors, next scheduled at time, etc. have not yet + // been updated based on the latest work result. + // + // Will not receive a common context related to HookWorkBegin because + // WorkBegin doesn't return a context. Middleware should be used for this + // sort of shared context instead. + WorkEnd(ctx context.Context, job *JobRow, err error) error +} + +// Middleware is an arbitrary interface for a struct which will execute some +// arbitrary code at a predefined step in the job lifecycle. +// +// This interface is left purposely non-specific. Middleware structs should +// embed river.MiddlewareDefaults to inherit an IsMiddleware implementation, +// then implement a more specific hook interface like JobInsertMiddleware or +// WorkerMiddleware. A middleware struct may also implement multiple specific +// hook interfaces which are logically related and benefit from being grouped +// together. +// +// Hooks differ from middleware in that they're invoked at a specific lifecycle +// phase, but finish immediately instead of wrapping an inner call like a +// middleware does. One of the main ramifications of this different is that a +// hook cannot modify context in any useful way to pass down into the stack. +// Like a normal function, any changes it makes to its context are discarded on +// return. +// +// Middleware differs from hooks in that they wrap a specific lifecycle phase, +// staying on the callstack for the duration of the step while they call into a +// doInner function that executes the step and the rest of the middleware stack. +// The main ramification of this difference is that middleware can modify +// context for the step and any other middleware inner relative to it. +// +// All else equal, hooks should generally be preferred over middleware because +// they don't add anything to the call stack. Call stacks that get overly deep +// can become a bit of an operational nightmare because they get hard to read. +// +// In a language with more specific type capabilities, this interface would be a +// union type. In Go we implement it somewhat awkwardly so that we can get +// future extensibility, but also some typing guarantees to prevent misuse (i.e. +// if Hook was an empty interface, then any object could be passed as a hook, +// but having a single function to implement forces the caller to make some +// token motions in the direction of implementing hooks). +// +// List of middleware interfaces that may be implemented: +// - JobInsertMiddleware +// - WorkerMiddleware +// +// More operation-specific interfaces may be added in future versions. +type Middleware interface { + // IsMiddleware is a sentinel function to check that a type is implementing + // Middleware on purpose and not by accident (Middleware would otherwise be + // an empty interface). Middleware should embed river.MiddlewareDefaults to + // pick up an implementation for this function automatically. + IsMiddleware() bool +} + +// Plugin is a generic extension installed globally. +// +// Plugin structs should embed river.PluginDefaults, or embed either +// river.HookDefaults or river.MiddlewareDefaults, then implement any +// operation-specific hook or middleware interfaces they need. +// +// For example, a plugin that receives emitted metrics: +// +// type MetricsPlugin struct { +// river.PluginDefaults +// } +// +// func (p *MetricsPlugin) MetricEmit(ctx context.Context, params *rivertype.HookMetricEmitParams) { +// // Export params.Metric asynchronously. +// } +// +// config := &river.Config{ +// Plugins: []rivertype.Plugin{&MetricsPlugin{}}, +// } +type Plugin interface { + IsPlugin() bool +} + +// JobInsertMiddleware provides an interface for middleware that integrations +// can use to encapsulate common logic around job insertion. +// +// Implementations should embed river.JobMiddlewareDefaults to inherit default +// implementations for phases where no custom code is needed, and for forward +// compatibility in case new functions are added to this interface. +type JobInsertMiddleware interface { + Middleware + + // InsertMany is invoked around a batch insert operation. Implementations + // must always include a call to doInner to call down the middleware stack + // and perform the batch insertion, and may run custom code before and after. + // + // Returning an error from this function will fail the overarching insert + // operation, even if the inner insertion originally succeeded. + InsertMany(ctx context.Context, manyParams []*JobInsertParams, doInner func(context.Context) ([]*JobInsertResult, error)) ([]*JobInsertResult, error) +} + +// WorkerMiddleware provides an interface for middleware that integrations can +// use to encapsulate common logic when a job is worked. +type WorkerMiddleware interface { + Middleware + + // Work is invoked after a job's JSON args being unmarshaled and before the + // job is worked. Implementations must always include a call to doInner to + // call down the middleware stack and perform the batch insertion, and may + // run custom code before and after. + // + // Returning an error from this function will fail the overarching work + // operation, even if the inner work originally succeeded. + Work(ctx context.Context, job *JobRow, doInner func(context.Context) error) error +} + +// DurablePeriodicJob represents a durable periodic job. +type DurablePeriodicJob struct { + // ID is a unique identifier for the durable periodic job. + ID string + + // CreatedAt is when the database record was created. + CreatedAt time.Time + + // NextRunAt is when the periodic job is next scheduled to run. + NextRunAt time.Time + + // UpdatedAt is when the database record was last updated. + UpdatedAt time.Time +} + +// PeriodicJobHandle is a reference to a dynamically added periodic job +// (returned by the use of `Client.PeriodicJobs().Add()`) which can be used to +// subsequently remove the periodic job with `Remove()`. +type PeriodicJobHandle int + +// Queue is a configuration for a queue that is currently (or recently was) in +// use by a client. +type Queue struct { + // CreatedAt is the time at which the queue first began being worked by a + // client. Unused queues are deleted after a retention period, so this only + // reflects the most recent time the queue was created if there was a long + // gap. + CreatedAt time.Time + // Metadata is a field for storing arbitrary metadata on a queue. It is + // currently reserved for River's internal use and should not be modified by + // users. + Metadata []byte + // Name is the name of the queue. + Name string + // PausedAt is the time the queue was paused, if any. When a paused queue is + // resumed, this field is set to nil. + PausedAt *time.Time + // UpdatedAt is the last time the queue was updated. This field is updated + // periodically any time an active Client is configured to work the queue, + // even if the queue is paused. + // + // If UpdatedAt has not been updated for awhile, the queue record will be + // deleted from the table by a maintenance process. + UpdatedAt time.Time +} + +// UniqueOptsByStateDefault is the set of job states that are used to determine +// uniqueness unless unique job states have been overridden with +// UniqueOpts.ByState. So for example, with this default set a new unique job +// may be inserted even if another job already exists, as long as that other job +// is set `cancelled` or `discarded`. +func UniqueOptsByStateDefault() []JobState { + return []JobState{ + JobStateAvailable, + JobStateCompleted, + JobStatePending, + JobStateRetryable, + JobStateRunning, + JobStateScheduled, + } +} + +// WorkerMetadata is metadata about workers registered with a client. +type WorkerMetadata struct { + // JobArgHooks are job args specific hooks returned from JobArgsWithHooks or + // from plugins returned by JobArgsWithPlugins. + JobArgHooks []Hook + + // Kind is the kind returned from job args and recognized by worker to work. + Kind string +} diff --git a/vendor/github.com/riverqueue/river/rivertype/time_generator.go b/vendor/github.com/riverqueue/river/rivertype/time_generator.go new file mode 100644 index 0000000000..9144cf8e5c --- /dev/null +++ b/vendor/github.com/riverqueue/river/rivertype/time_generator.go @@ -0,0 +1,28 @@ +package rivertype + +import "time" + +// TimeGenerator generates current time values for in-process timing math and +// optional stubbed wall-clock timestamps in tests. In test environments it's +// implemented by riversharedtest.TimeStub which lets the current time be +// stubbed. Otherwise, it's implemented as UnStubbableTimeGenerator which +// doesn't allow stubbing. +type TimeGenerator interface { + // Now returns the current time. This may be a stubbed time if the time has + // been actively stubbed in a test. + // + // Production implementations should preserve Go's monotonic clock reading + // from time.Now for in-process duration and deadline math. Do not normalize + // through `t.UTC()` here: per the time package's monotonic clock semantics, + // changing location strips the monotonic reading. Normalize at database or + // serialization boundaries instead. + Now() time.Time + + // NowOrNil returns the currently stubbed time if the current time is + // stubbed, and nil otherwise. + // + // This is mainly for database-facing test paths that want to inject a + // deterministic wall-clock timestamp when time is stubbed, but to fall back + // to a database-side time default in production. + NowOrNil() *time.Time +} diff --git a/vendor/github.com/riverqueue/river/stuck_job.go b/vendor/github.com/riverqueue/river/stuck_job.go new file mode 100644 index 0000000000..9a2736a95b --- /dev/null +++ b/vendor/github.com/riverqueue/river/stuck_job.go @@ -0,0 +1,32 @@ +package river + +import "context" + +// JobStuckHandler is invoked when a producer detects that a job exceeded its +// timeout and did not return within the configured stuck-job timeout margin. +type JobStuckHandler func(ctx context.Context, params JobStuckHandlerParams) JobStuckHandlerResult + +// JobStuckHandlerParams are parameters passed to JobStuckHandler. +type JobStuckHandlerParams struct { + // ID is the ID of the stuck job. + ID int64 + + // Kind is the kind of the stuck job. + Kind string + + // Queue is the queue where the stuck job is running. + Queue string + + // TotalStuckJobs is the total number of jobs currently considered stuck + // across the client (includes all queues). + TotalStuckJobs int +} + +// JobStuckHandlerResult is the result returned by JobStuckHandler. +type JobStuckHandlerResult struct { + // AddWorkerSlot instructs River to treat the stuck job as no longer + // occupying a worker slot so another job can begin executing. This can be + // dangerous because the stuck job's goroutine is still running, so the + // queue may temporarily have more active job goroutines than MaxWorkers. + AddWorkerSlot bool +} diff --git a/vendor/github.com/riverqueue/river/subscription_manager.go b/vendor/github.com/riverqueue/river/subscription_manager.go new file mode 100644 index 0000000000..c74ef4d460 --- /dev/null +++ b/vendor/github.com/riverqueue/river/subscription_manager.go @@ -0,0 +1,267 @@ +package river + +import ( + "context" + "fmt" + "log/slog" + "sync" + "time" + + "github.com/riverqueue/river/internal/jobcompleter" + "github.com/riverqueue/river/internal/jobstats" + "github.com/riverqueue/river/rivershared/baseservice" + "github.com/riverqueue/river/rivershared/startstop" + "github.com/riverqueue/river/rivershared/util/sliceutil" + "github.com/riverqueue/river/rivertype" +) + +type subscriptionManager struct { + baseservice.BaseService + startstop.BaseStartStop + + subscribeCh <-chan []jobcompleter.CompleterJobUpdated + + statsMu sync.Mutex // protects stats fields + statsAggregate jobstats.JobStatistics + statsNumJobs int + + mu sync.Mutex // protects subscription fields + subscriptions map[int]*eventSubscription + subscriptionsSeq int // used for generating simple IDs +} + +func newSubscriptionManager(archetype *baseservice.Archetype, subscribeCh <-chan []jobcompleter.CompleterJobUpdated) *subscriptionManager { + return baseservice.Init(archetype, &subscriptionManager{ + subscribeCh: subscribeCh, + subscriptions: make(map[int]*eventSubscription), + }) +} + +// ResetSubscribeChan is used to change the channel that the subscription +// manager listens on. It must only be called when the subscription manager is +// stopped. +func (sm *subscriptionManager) ResetSubscribeChan(subscribeCh <-chan []jobcompleter.CompleterJobUpdated) { + sm.subscribeCh = subscribeCh +} + +func (sm *subscriptionManager) Start(ctx context.Context) error { + ctx, shouldStart, started, stopped := sm.StartInit(ctx) + if !shouldStart { + return nil + } + + go func() { + started() + defer stopped() // this defer should come first so it's last out + + sm.Logger.DebugContext(ctx, sm.Name+": Run loop started") + defer sm.Logger.DebugContext(ctx, sm.Name+": Run loop stopped") + + // On shutdown, close and remove all active subscriptions. + defer func() { + sm.mu.Lock() + defer sm.mu.Unlock() + + for subID, sub := range sm.subscriptions { + close(sub.Chan) + delete(sm.subscriptions, subID) + } + }() + + for { + select { + case <-ctx.Done(): + // Distribute remaining subscriptions until the channel is + // closed. This does make the subscription manager a little + // problematic in that it requires the subscription channel to + // be closed before it will fully stop. This always happens in + // the case of a real client by virtue of the completer always + // stopping at the same time as the subscription manager, but + // one has to be careful in tests. + sm.Logger.DebugContext(ctx, sm.Name+": Stopping; distributing subscriptions until channel is closed") + for updates := range sm.subscribeCh { + sm.distributeJobUpdates(ctx, updates) + } + + return + + case updates := <-sm.subscribeCh: + sm.distributeJobUpdates(ctx, updates) + } + } + }() + + return nil +} + +func (sm *subscriptionManager) logStats(ctx context.Context, svcName string) { + sm.statsMu.Lock() + defer sm.statsMu.Unlock() + + sm.Logger.DebugContext(ctx, svcName+": Job stats (since last stats line)", + "num_jobs_run", sm.statsNumJobs, + "average_complete_duration", sm.safeDurationAverage(sm.statsAggregate.CompleteDuration, sm.statsNumJobs), + "average_queue_wait_duration", sm.safeDurationAverage(sm.statsAggregate.QueueWaitDuration, sm.statsNumJobs), + "average_run_duration", sm.safeDurationAverage(sm.statsAggregate.RunDuration, sm.statsNumJobs)) + + sm.statsAggregate = jobstats.JobStatistics{} + sm.statsNumJobs = 0 +} + +// Handles a potential divide by zero. +func (sm *subscriptionManager) safeDurationAverage(d time.Duration, n int) time.Duration { + if n == 0 { + return 0 + } + return d / time.Duration(n) +} + +// Receives updates from the completer and prompts the client to update +// statistics and distribute jobs into any listening subscriber channels. +// (Subscriber channels are non-blocking so this should be quite fast.) +func (sm *subscriptionManager) distributeJobUpdates(ctx context.Context, updates []jobcompleter.CompleterJobUpdated) { + func() { + sm.statsMu.Lock() + defer sm.statsMu.Unlock() + + for _, update := range updates { + stats := update.JobStats + sm.statsAggregate.CompleteDuration += stats.CompleteDuration + sm.statsAggregate.QueueWaitDuration += stats.QueueWaitDuration + sm.statsAggregate.RunDuration += stats.RunDuration + sm.statsNumJobs++ + } + }() + + sm.mu.Lock() + defer sm.mu.Unlock() + + // Quick path so we don't need to allocate anything if no one is listening. + if len(sm.subscriptions) < 1 { + return + } + + for _, update := range updates { + sm.distributeJobEvent(ctx, update.Job, jobStatisticsFromInternal(update.JobStats), update.Snoozed) + } +} + +// Distribute a single event into any listening subscriber channels. +// +// Job events should specify the job and stats, while queue events should only specify +// the queue. +// +// MUST be called with sm.mu already held. +func (sm *subscriptionManager) distributeJobEvent(ctx context.Context, job *rivertype.JobRow, stats *JobStatistics, snoozed bool) { + var event *Event + if snoozed { + event = &Event{Kind: EventKindJobSnoozed, Job: job, JobStats: stats} + } else { + switch job.State { + case rivertype.JobStateCancelled: + event = &Event{Kind: EventKindJobCancelled, Job: job, JobStats: stats} + case rivertype.JobStateCompleted: + event = &Event{Kind: EventKindJobCompleted, Job: job, JobStats: stats} + case rivertype.JobStateAvailable, rivertype.JobStateDiscarded, rivertype.JobStateRetryable, rivertype.JobStateRunning: + event = &Event{Kind: EventKindJobFailed, Job: job, JobStats: stats} + case rivertype.JobStatePending, rivertype.JobStateScheduled: + // job state may be set to scheduled, but only for snoozed jobs, so + // the case at the top should always take precedence before this + panic(fmt.Sprintf("completion subscriber unexpectedly received job in %s state, river bug", job.State)) + default: + // linter exhaustive rule prevents this from being reached + panic("unreachable state to distribute, river bug") + } + } + + // All subscription channels are non-blocking so this is always fast and + // there's no risk of falling behind what producers are sending. + for _, sub := range sm.subscriptions { + if sub.ListensFor(event.Kind) { + // TODO: THIS IS UNSAFE AND WILL LEAD TO DROPPED EVENTS. + // + // We are allocating subscriber channels with a fixed size of 1000, but + // potentially processing job events in batches of 5000 (batch completer + // max batch size). It's probably not possible for the subscriber to keep + // up with these bursts. + select { + case sub.Chan <- event: + default: + sm.Logger.WarnContext(ctx, sm.Name+": Subscription event dropped due to full buffer", + slog.String("event_kind", string(event.Kind)), + ) + } + } + } +} + +func (sm *subscriptionManager) distributeQueueEvent(event *Event) { + sm.distributeQueueEventWithContext(context.Background(), event) +} + +func (sm *subscriptionManager) distributeQueueEventWithContext(ctx context.Context, event *Event) { + sm.mu.Lock() + defer sm.mu.Unlock() + + // All subscription channels are non-blocking so this is always fast and + // there's no risk of falling behind what producers are sending. + for _, sub := range sm.subscriptions { + if sub.ListensFor(event.Kind) { + select { + case sub.Chan <- event: + default: + sm.Logger.WarnContext(ctx, sm.Name+": Subscription event dropped due to full buffer", + slog.String("event_kind", string(event.Kind)), + ) + } + } + } +} + +// SubscribeConfig is a special internal Subscribe variant that lets us inject +// an overridden size. +func (sm *subscriptionManager) SubscribeConfig(config *SubscribeConfig) (<-chan *Event, func()) { + if config.ChanSize < 0 { + panic("SubscribeConfig.ChanSize must be greater or equal to 1") + } + if config.ChanSize == 0 { + config.ChanSize = subscribeChanSizeDefault + } + + for _, kind := range config.Kinds { + if _, ok := allKinds[kind]; !ok { + panic(fmt.Errorf("unknown event kind: %s", kind)) + } + } + + subChan := make(chan *Event, config.ChanSize) + + sm.mu.Lock() + defer sm.mu.Unlock() + + // Just gives us an easy way of removing the subscription again later. + subID := sm.subscriptionsSeq + sm.subscriptionsSeq++ + + sm.subscriptions[subID] = &eventSubscription{ + Chan: subChan, + Kinds: sliceutil.KeyBy(config.Kinds, func(k EventKind) (EventKind, struct{}) { return k, struct{}{} }), + } + + cancel := func() { + sm.mu.Lock() + defer sm.mu.Unlock() + + // May no longer be present in case this was called after a stop. + sub, ok := sm.subscriptions[subID] + if !ok { + return + } + + close(sub.Chan) + + delete(sm.subscriptions, subID) + } + + return subChan, cancel +} diff --git a/vendor/github.com/riverqueue/river/work_unit_wrapper.go b/vendor/github.com/riverqueue/river/work_unit_wrapper.go new file mode 100644 index 0000000000..7c1000fb6a --- /dev/null +++ b/vendor/github.com/riverqueue/river/work_unit_wrapper.go @@ -0,0 +1,47 @@ +package river + +import ( + "context" + "encoding/json" + "time" + + "github.com/riverqueue/river/internal/pluginlookup" + "github.com/riverqueue/river/internal/workunit" + "github.com/riverqueue/river/rivertype" +) + +// workUnitFactoryWrapper wraps a Worker to implement workUnitFactory. +type workUnitFactoryWrapper[T JobArgs] struct { + worker Worker[T] +} + +func (w *workUnitFactoryWrapper[T]) MakeUnit(jobRow *rivertype.JobRow) workunit.WorkUnit { + return &wrapperWorkUnit[T]{jobRow: jobRow, worker: w.worker} +} + +// wrapperWorkUnit implements workUnit for a job and Worker. +type wrapperWorkUnit[T JobArgs] struct { + job *Job[T] // not set until after UnmarshalJob is invoked + jobRow *rivertype.JobRow + worker Worker[T] +} + +func (w *wrapperWorkUnit[T]) PluginLookup(lookup *pluginlookup.JobPluginLookup) *pluginlookup.PluginLookup { + var job T + return lookup.ByJobArgs(job) +} + +func (w *wrapperWorkUnit[T]) Middleware() []rivertype.WorkerMiddleware { + return w.worker.Middleware(w.jobRow) +} +func (w *wrapperWorkUnit[T]) NextRetry() time.Time { return w.worker.NextRetry(w.job) } +func (w *wrapperWorkUnit[T]) Timeout() time.Duration { return w.worker.Timeout(w.job) } +func (w *wrapperWorkUnit[T]) Work(ctx context.Context) error { return w.worker.Work(ctx, w.job) } + +func (w *wrapperWorkUnit[T]) UnmarshalJob() error { + w.job = &Job[T]{ + JobRow: w.jobRow, + } + + return json.Unmarshal(w.jobRow.EncodedArgs, &w.job.Args) +} diff --git a/vendor/github.com/riverqueue/river/worker.go b/vendor/github.com/riverqueue/river/worker.go new file mode 100644 index 0000000000..f833b760de --- /dev/null +++ b/vendor/github.com/riverqueue/river/worker.go @@ -0,0 +1,223 @@ +package river + +import ( + "context" + "fmt" + "time" + + "github.com/riverqueue/river/internal/workunit" + "github.com/riverqueue/river/rivertype" +) + +// Worker is an interface that can perform a job with args of type T. A typical +// implementation will be a JSON-serializable `JobArgs` struct that implements +// `Kind()`, along with a Worker that embeds WorkerDefaults and implements `Work()`. +// Workers may optionally override other methods to provide job-specific +// configuration for all jobs of that type: +// +// type SleepArgs struct { +// Duration time.Duration `json:"duration"` +// } +// +// func (SleepArgs) Kind() string { return "sleep" } +// +// type SleepWorker struct { +// WorkerDefaults[SleepArgs] +// } +// +// func (w *SleepWorker) Work(ctx context.Context, job *Job[SleepArgs]) error { +// select { +// case <-ctx.Done(): +// return ctx.Err() +// case <-time.After(job.Args.Duration): +// return nil +// } +// } +// +// In addition to fulfilling the Worker interface, workers must be registered +// with the client using the AddWorker function. +type Worker[T JobArgs] interface { + // Middleware returns the type-specific middleware for this job. + Middleware(job *rivertype.JobRow) []rivertype.WorkerMiddleware + + // NextRetry calculates when the next retry for a failed job should take + // place given when it was last attempted and its number of attempts, or any + // other of the job's properties a user-configured retry policy might want + // to consider. + // + // Note that this method on a worker overrides any client-level retry policy. + // To use the client-level retry policy, return an empty `time.Time{}` or + // include WorkerDefaults to do this for you. + NextRetry(job *Job[T]) time.Time + + // Timeout is the maximum amount of time the job is allowed to run before + // its context is cancelled. A timeout of zero (the default) means the job + // will inherit the Client-level timeout. A timeout of -1 means the job's + // context will never time out. + Timeout(job *Job[T]) time.Duration + + // Work performs the job and returns an error if the job failed. The context + // will be configured with a timeout according to the worker settings and may + // be cancelled for other reasons. + // + // If no error is returned, the job is assumed to have succeeded and will be + // marked completed. + // + // It is important for any worker to respect context cancellation to enable + // the client to respond to shutdown requests. In particular, workers that + // wait on channels, timers, or network operations should prefer a `select` + // that also watches `ctx.Done()`. There is no way to cancel a running job + // that does not respect context cancellation, other than terminating the + // process. + // + // A worker that ignores cancellation may continue running even after the + // client has timed out the job or the job rescuer has moved it out of + // running. + Work(ctx context.Context, job *Job[T]) error +} + +// WorkerDefaults is an empty struct that can be embedded in your worker +// struct to make it fulfill the Worker interface with default values. +type WorkerDefaults[T JobArgs] struct{} + +func (w WorkerDefaults[T]) Middleware(*rivertype.JobRow) []rivertype.WorkerMiddleware { return nil } + +// NextRetry returns an empty time.Time{} to avoid setting any job or +// Worker-specific overrides on the next retry time. This means that the +// Client-level retry policy schedule will be used instead. +func (w WorkerDefaults[T]) NextRetry(*Job[T]) time.Time { return time.Time{} } + +// Timeout returns the job-specific timeout. Override this method to set a +// job-specific timeout, otherwise the Client-level timeout will be applied. +func (w WorkerDefaults[T]) Timeout(*Job[T]) time.Duration { return 0 } + +// AddWorker registers a Worker on the provided Workers bundle. Each Worker must +// be registered so that the Client knows it should handle a specific kind of +// job (as returned by its `Kind()` method). +// +// Use by explicitly specifying a JobArgs type and then passing an instance of a +// worker for the same type: +// +// river.AddWorker(workers, &SortWorker{}) +// +// Note that AddWorker can panic in some situations, such as if the worker is +// already registered or if its configuration is otherwise invalid. This default +// probably makes sense for most applications because you wouldn't want to start +// an application with invalid hardcoded runtime configuration. If you want to +// avoid panics, use AddWorkerSafely instead. +func AddWorker[T JobArgs](workers *Workers, worker Worker[T]) { + if err := AddWorkerSafely(workers, worker); err != nil { + panic(err) + } +} + +// AddWorkerArgs is the same as AddWorker except that it lets args be passed +// explicitly rather than being instantiated implicitly. We don't know of any +// use for this function beyond exercising some args-related edge cases in tests +// are difficult/impossible to exercise otherwise, and its use should be +// considered internal only. +func AddWorkerArgs[T JobArgs](workers *Workers, jobArgs T, worker Worker[T]) { + if err := workers.add(jobArgs, &workUnitFactoryWrapper[T]{worker: worker}); err != nil { + panic(err) + } +} + +// AddWorkerSafely registers a worker on the provided Workers bundle. Unlike AddWorker, +// AddWorkerSafely does not panic and instead returns an error if the worker +// is already registered or if its configuration is invalid. +// +// Use by explicitly specifying a JobArgs type and then passing an instance of a +// worker for the same type: +// +// river.AddWorkerSafely[SortArgs](workers, &SortWorker{}). +func AddWorkerSafely[T JobArgs](workers *Workers, worker Worker[T]) error { + var jobArgs T + return workers.add(jobArgs, &workUnitFactoryWrapper[T]{worker: worker}) +} + +// Workers is a list of available job workers. A Worker must be registered for +// each type of Job to be handled. +// +// Use the top-level AddWorker function combined with a Workers to register a +// worker. +type Workers struct { + workersMap map[string]workerInfo // job kind -> worker info +} + +// workerInfo bundles information about a registered worker for later lookup +// in a Workers bundle. +type workerInfo struct { + jobArgs JobArgs + workUnitFactory workunit.WorkUnitFactory +} + +// NewWorkers initializes a new registry of available job workers. +// +// Use the top-level AddWorker function combined with a Workers registry to +// register each available worker. +func NewWorkers() *Workers { + return &Workers{ + workersMap: make(map[string]workerInfo), + } +} + +func (w Workers) add(jobArgs JobArgs, workUnitFactory workunit.WorkUnitFactory) error { + checkRegistered := func(kind string) error { + if _, ok := w.workersMap[kind]; ok { + return fmt.Errorf("worker for kind %q is already registered", kind) + } + return nil + } + + workerInfo := workerInfo{ + jobArgs: jobArgs, + workUnitFactory: workUnitFactory, + } + + kind := jobArgs.Kind() + if err := checkRegistered(kind); err != nil { + return err + } + w.workersMap[kind] = workerInfo + + // Jobs can register an alternate kind to make renaming easier. + if jobArgsWithKindAliases, ok := jobArgs.(JobArgsWithKindAliases); ok { + for _, kind := range jobArgsWithKindAliases.KindAliases() { + if err := checkRegistered(kind); err != nil { + return err + } + w.workersMap[kind] = workerInfo + } + } + + return nil +} + +// workFunc implements JobArgs and is used to wrap a function given to WorkFunc. +type workFunc[T JobArgs] struct { + WorkerDefaults[T] + + kind string + f func(context.Context, *Job[T]) error +} + +func (wf *workFunc[T]) Kind() string { + return wf.kind +} + +func (wf *workFunc[T]) Work(ctx context.Context, job *Job[T]) error { + return wf.f(ctx, job) +} + +// WorkFunc wraps a function to implement the Worker interface. A job args +// struct implementing JobArgs will still be required to specify a Kind. +// +// For example: +// +// river.AddWorker(workers, river.WorkFunc(func(ctx context.Context, job *river.Job[WorkFuncArgs]) error { +// fmt.Printf("Message: %s", job.Args.Message) +// return nil +// })) +func WorkFunc[T JobArgs](f func(context.Context, *Job[T]) error) Worker[T] { + return &workFunc[T]{f: f, kind: (*new(T)).Kind()} +} diff --git a/vendor/github.com/tidwall/gjson/README.md b/vendor/github.com/tidwall/gjson/README.md index c8db11f147..f87b873776 100644 --- a/vendor/github.com/tidwall/gjson/README.md +++ b/vendor/github.com/tidwall/gjson/README.md @@ -1,7 +1,9 @@

-GJSON + + + + GJSON +
GoDoc GJSON Playground @@ -54,7 +56,7 @@ This will print: ``` Prichard ``` -*There's also the [GetMany](#get-multiple-values-at-once) function to get multiple values at once, and [GetBytes](#working-with-bytes) for working with JSON byte slices.* +*There's also [GetBytes](#working-with-bytes) for working with JSON byte slices.* ## Path Syntax @@ -211,6 +213,7 @@ There are currently the following built-in modifiers: - `@tostr`: Converts json to a string. Wraps a json string. - `@fromstr`: Converts a string from json. Unwraps a json string. - `@group`: Groups arrays of objects. See [e4fc67c](https://github.com/tidwall/gjson/commit/e4fc67c92aeebf2089fabc7872f010e340d105db). +- `@dig`: Search for a value without providing its entire path. See [e8e87f2](https://github.com/tidwall/gjson/commit/e8e87f2a00dc41f3aba5631094e21f59a8cf8cbf). ### Modifier arguments @@ -426,16 +429,6 @@ if result.Index > 0 { This is a best-effort no allocation sub slice of the original json. This method utilizes the `result.Index` field, which is the position of the raw data in the original json. It's possible that the value of `result.Index` equals zero, in which case the `result.Raw` is converted to a `[]byte`. -## Get multiple values at once - -The `GetMany` function can be used to get multiple values at the same time. - -```go -results := gjson.GetMany(json, "name.first", "name.last", "age") -``` - -The return value is a `[]Result`, which will always contain exactly the same number of items as the input paths. - ## Performance Benchmarks of GJSON alongside [encoding/json](https://golang.org/pkg/encoding/json/), @@ -445,15 +438,15 @@ Benchmarks of GJSON alongside [encoding/json](https://golang.org/pkg/encoding/js and [json-iterator](https://github.com/json-iterator/go) ``` -BenchmarkGJSONGet-16 11644512 311 ns/op 0 B/op 0 allocs/op -BenchmarkGJSONUnmarshalMap-16 1122678 3094 ns/op 1920 B/op 26 allocs/op -BenchmarkJSONUnmarshalMap-16 516681 6810 ns/op 2944 B/op 69 allocs/op -BenchmarkJSONUnmarshalStruct-16 697053 5400 ns/op 928 B/op 13 allocs/op -BenchmarkJSONDecoder-16 330450 10217 ns/op 3845 B/op 160 allocs/op -BenchmarkFFJSONLexer-16 1424979 2585 ns/op 880 B/op 8 allocs/op -BenchmarkEasyJSONLexer-16 3000000 729 ns/op 501 B/op 5 allocs/op -BenchmarkJSONParserGet-16 3000000 366 ns/op 21 B/op 0 allocs/op -BenchmarkJSONIterator-16 3000000 869 ns/op 693 B/op 14 allocs/op +BenchmarkGJSONGet-10 17893731 202.1 ns/op 0 B/op 0 allocs/op +BenchmarkGJSONUnmarshalMap-10 1663548 2157 ns/op 1920 B/op 26 allocs/op +BenchmarkJSONUnmarshalMap-10 832236 4279 ns/op 2920 B/op 68 allocs/op +BenchmarkJSONUnmarshalStruct-10 1076475 3219 ns/op 920 B/op 12 allocs/op +BenchmarkJSONDecoder-10 585729 6126 ns/op 3845 B/op 160 allocs/op +BenchmarkFFJSONLexer-10 2508573 1391 ns/op 880 B/op 8 allocs/op +BenchmarkEasyJSONLexer-10 3000000 537.9 ns/op 501 B/op 5 allocs/op +BenchmarkJSONParserGet-10 13707510 263.9 ns/op 21 B/op 0 allocs/op +BenchmarkJSONIterator-10 3000000 561.2 ns/op 693 B/op 14 allocs/op ``` JSON document used: @@ -494,4 +487,6 @@ widget.image.hOffset widget.text.onMouseUp ``` -*These benchmarks were run on a MacBook Pro 16" 2.4 GHz Intel Core i9 using Go 1.17 and can be found [here](https://github.com/tidwall/gjson-benchmarks).* +** + +*These benchmarks were run on a MacBook Pro M1 Max using Go 1.22 and can be found [here](https://github.com/tidwall/gjson-benchmarks).* diff --git a/vendor/github.com/tidwall/gjson/SYNTAX.md b/vendor/github.com/tidwall/gjson/SYNTAX.md index 7a9b6a2d7c..a3f0fac238 100644 --- a/vendor/github.com/tidwall/gjson/SYNTAX.md +++ b/vendor/github.com/tidwall/gjson/SYNTAX.md @@ -1,6 +1,6 @@ # GJSON Path Syntax -A GJSON Path is a text string syntax that describes a search pattern for quickly retreiving values from a JSON payload. +A GJSON Path is a text string syntax that describes a search pattern for quickly retrieving values from a JSON payload. This document is designed to explain the structure of a GJSON Path through examples. @@ -15,12 +15,12 @@ This document is designed to explain the structure of a GJSON Path through examp - [Multipaths](#multipaths) - [Literals](#literals) -The definitive implemenation is [github.com/tidwall/gjson](https://github.com/tidwall/gjson). +The definitive implementation is [github.com/tidwall/gjson](https://github.com/tidwall/gjson). Use the [GJSON Playground](https://gjson.dev) to experiment with the syntax online. ## Path structure -A GJSON Path is intended to be easily expressed as a series of components seperated by a `.` character. +A GJSON Path is intended to be easily expressed as a series of components separated by a `.` character. Along with `.` character, there are a few more that have special meaning, including `|`, `#`, `@`, `\`, `*`, `!`, and `?`. @@ -46,7 +46,7 @@ The following GJSON Paths evaluate to the accompanying values. ### Basic -In many cases you'll just want to retreive values by object name or array index. +In many cases you'll just want to retrieve values by object name or array index. ```go name.last "Anderson" @@ -137,12 +137,21 @@ next major release.* The `~` (tilde) operator will convert a value to a boolean before comparison. +Supported tilde comparison type are: + +``` +~true Converts true-ish values to true +~false Converts false-ish and non-existent values to true +~null Converts null and non-existent values to true +~* Converts any existing value to true +``` + For example, using the following JSON: ```json { "vals": [ - { "a": 1, "b": true }, + { "a": 1, "b": "data" }, { "a": 2, "b": true }, { "a": 3, "b": false }, { "a": 4, "b": "0" }, @@ -157,15 +166,23 @@ For example, using the following JSON: } ``` -You can now query for all true(ish) or false(ish) values: +To query for all true-ish or false-ish values: ``` -vals.#(b==~true)#.a >> [1,2,6,7,8] +vals.#(b==~true)#.a >> [2,6,7,8] vals.#(b==~false)#.a >> [3,4,5,9,10,11] ``` The last value which was non-existent is treated as `false` +To query for null and explicit value existence: + +``` +vals.#(b==~null)#.a >> [10,11] +vals.#(b==~*)#.a >> [1,2,3,4,5,6,7,8,9,10] +vals.#(b!=~*)#.a >> [11] +``` + ### Dot vs Pipe The `.` is standard separator, but it's also possible to use a `|`. @@ -241,6 +258,7 @@ There are currently the following built-in modifiers: - `@tostr`: Converts json to a string. Wraps a json string. - `@fromstr`: Converts a string from json. Unwraps a json string. - `@group`: Groups arrays of objects. See [e4fc67c](https://github.com/tidwall/gjson/commit/e4fc67c92aeebf2089fabc7872f010e340d105db). +- `@dig`: Search for a value without providing its entire path. See [e8e87f2](https://github.com/tidwall/gjson/commit/e8e87f2a00dc41f3aba5631094e21f59a8cf8cbf). #### Modifier arguments diff --git a/vendor/github.com/tidwall/gjson/gjson.go b/vendor/github.com/tidwall/gjson/gjson.go index 53cbd2363f..97d6516ae2 100644 --- a/vendor/github.com/tidwall/gjson/gjson.go +++ b/vendor/github.com/tidwall/gjson/gjson.go @@ -1,7 +1,14 @@ +// Copyright 2024 Joshua J Baker. All rights reserved. +// Use of this source code is governed by an MIT-style +// license that can be found in the LICENSE file. +// +// https://github.com/tidwall/gjson + // Package gjson provides searching for json strings. package gjson import ( + "iter" "strconv" "strings" "time" @@ -645,9 +652,9 @@ func tostr(json string) (raw string, str string) { // Exists returns true if value exists. // -// if gjson.Get(json, "name.last").Exists(){ -// println("value exists") -// } +// if gjson.Get(json, "name.last").Exists(){ +// println("value exists") +// } func (t Result) Exists() bool { return t.Type != Null || len(t.Raw) != 0 } @@ -661,7 +668,6 @@ func (t Result) Exists() bool { // nil, for JSON null // map[string]interface{}, for JSON objects // []interface{}, for JSON arrays -// func (t Result) Value() interface{} { if t.Type == String { return t.Str @@ -826,19 +832,28 @@ func parseArrayPath(path string) (r arrayPathResult) { } // splitQuery takes a query and splits it into three parts: -// path, op, middle, and right. +// +// path, op, middle, and right. +// // So for this query: -// #(first_name=="Murphy").last +// +// #(first_name=="Murphy").last +// // Becomes -// first_name # path -// =="Murphy" # middle -// .last # right +// +// first_name # path +// =="Murphy" # middle +// .last # right +// // Or, -// #(service_roles.#(=="one")).cap +// +// #(service_roles.#(=="one")).cap +// // Becomes -// service_roles.#(=="one") # path -// # middle -// .cap # right +// +// service_roles.#(=="one") # path +// # middle +// .cap # right func parseQuery(query string) ( path, op, value, remain string, i int, vesc, ok bool, ) { @@ -1032,6 +1047,10 @@ func parseObjectPath(path string) (r objectPathResult) { return } +var vchars = [256]byte{ + '"': 2, '{': 3, '(': 3, '[': 3, '}': 1, ')': 1, ']': 1, +} + func parseSquash(json string, i int) (int, string) { // expects that the lead character is a '[' or '{' or '(' // squash the value, ignoring all nested arrays and objects. @@ -1039,43 +1058,137 @@ func parseSquash(json string, i int) (int, string) { s := i i++ depth := 1 - for ; i < len(json); i++ { - if json[i] >= '"' && json[i] <= '}' { - switch json[i] { - case '"': + var c byte + for i < len(json) { + for i < len(json)-8 { + jslice := json[i : i+8] + c = vchars[jslice[0]] + if c != 0 { + i += 0 + goto token + } + c = vchars[jslice[1]] + if c != 0 { + i += 1 + goto token + } + c = vchars[jslice[2]] + if c != 0 { + i += 2 + goto token + } + c = vchars[jslice[3]] + if c != 0 { + i += 3 + goto token + } + c = vchars[jslice[4]] + if c != 0 { + i += 4 + goto token + } + c = vchars[jslice[5]] + if c != 0 { + i += 5 + goto token + } + c = vchars[jslice[6]] + if c != 0 { + i += 6 + goto token + } + c = vchars[jslice[7]] + if c != 0 { + i += 7 + goto token + } + i += 8 + } + c = vchars[json[i]] + if c == 0 { + i++ + continue + } + token: + if c == 2 { + // '"' string + i++ + s2 := i + nextquote: + for i < len(json)-8 { + jslice := json[i : i+8] + if jslice[0] == '"' { + i += 0 + goto strchkesc + } + if jslice[1] == '"' { + i += 1 + goto strchkesc + } + if jslice[2] == '"' { + i += 2 + goto strchkesc + } + if jslice[3] == '"' { + i += 3 + goto strchkesc + } + if jslice[4] == '"' { + i += 4 + goto strchkesc + } + if jslice[5] == '"' { + i += 5 + goto strchkesc + } + if jslice[6] == '"' { + i += 6 + goto strchkesc + } + if jslice[7] == '"' { + i += 7 + goto strchkesc + } + i += 8 + } + goto strchkstd + strchkesc: + if json[i-1] != '\\' { i++ - s2 := i - for ; i < len(json); i++ { - if json[i] > '\\' { - continue - } - if json[i] == '"' { - // look for an escaped slash - if json[i-1] == '\\' { - n := 0 - for j := i - 2; j > s2-1; j-- { - if json[j] != '\\' { - break - } - n++ - } - if n%2 == 0 { - continue - } + continue + } + strchkstd: + for i < len(json) { + if json[i] > '\\' || json[i] != '"' { + i++ + continue + } + // look for an escaped slash + if json[i-1] == '\\' { + n := 0 + for j := i - 2; j > s2-1; j-- { + if json[j] != '\\' { + break } - break + n++ + } + if n%2 == 0 { + i++ + goto nextquote } } - case '{', '[', '(': - depth++ - case '}', ']', ')': - depth-- - if depth == 0 { - i++ - return i, json[s:i] - } + break + } + } else { + // '{', '[', '(', '}', ']', ')' + // open close tokens + depth += int(c) - 2 + if depth == 0 { + i++ + return i, json[s:i] } } + i++ } return i, json[s:] } @@ -1244,22 +1357,81 @@ func parseObject(c *parseContext, i int, path string) (int, bool) { } // matchLimit will limit the complexity of the match operation to avoid ReDos -// attacks from arbritary inputs. +// attacks from arbitrary inputs. // See the github.com/tidwall/match.MatchLimit function for more information. func matchLimit(str, pattern string) bool { matched, _ := match.MatchLimit(str, pattern, 10000) return matched } +func falseish(t Result) bool { + switch t.Type { + case Null: + return true + case False: + return true + case String: + b, err := strconv.ParseBool(strings.ToLower(t.Str)) + if err != nil { + return false + } + return !b + case Number: + return t.Num == 0 + default: + return false + } +} + +func trueish(t Result) bool { + switch t.Type { + case True: + return true + case String: + b, err := strconv.ParseBool(strings.ToLower(t.Str)) + if err != nil { + return false + } + return b + case Number: + return t.Num != 0 + default: + return false + } +} + +func nullish(t Result) bool { + return t.Type == Null +} + func queryMatches(rp *arrayPathResult, value Result) bool { rpv := rp.query.value - if len(rpv) > 0 && rpv[0] == '~' { - // convert to bool - rpv = rpv[1:] - if value.Bool() { - value = Result{Type: True} - } else { - value = Result{Type: False} + if len(rpv) > 0 { + if rpv[0] == '~' { + // convert to bool + rpv = rpv[1:] + var ish, ok bool + switch rpv { + case "*": + ish, ok = value.Exists(), true + case "null": + ish, ok = nullish(value), true + case "true": + ish, ok = trueish(value), true + case "false": + ish, ok = falseish(value), true + } + if ok { + rpv = "true" + if ish { + value = Result{Type: True} + } else { + value = Result{Type: False} + } + } else { + rpv = "" + value = Result{} + } } } if !value.Exists() { @@ -1850,6 +2022,16 @@ func appendHex16(dst []byte, x uint16) []byte { ) } +// DisableEscapeHTML will disable the automatic escaping of certain +// "problamatic" HTML characters when encoding to JSON. +// These character include '>', '<' and '&', which get escaped to \u003e, +// \u0026, and \u003c respectively. +// +// This is a global flag and will affect all further gjson operations. +// Ideally, if used, it should be set one time before other gjson functions +// are called. +var DisableEscapeHTML = false + // AppendJSONString is a convenience function that converts the provided string // to a valid JSON string and appends it to dst. func AppendJSONString(dst []byte, s string) []byte { @@ -1859,6 +2041,10 @@ func AppendJSONString(dst []byte, s string) []byte { if s[i] < ' ' { dst = append(dst, '\\') switch s[i] { + case '\b': + dst = append(dst, 'b') + case '\f': + dst = append(dst, 'f') case '\n': dst = append(dst, 'n') case '\r': @@ -1869,7 +2055,8 @@ func AppendJSONString(dst []byte, s string) []byte { dst = append(dst, 'u') dst = appendHex16(dst, uint16(s[i])) } - } else if s[i] == '>' || s[i] == '<' || s[i] == '&' { + } else if !DisableEscapeHTML && + (s[i] == '>' || s[i] == '<' || s[i] == '&') { dst = append(dst, '\\', 'u') dst = appendHex16(dst, uint16(s[i])) } else if s[i] == '\\' { @@ -1918,23 +2105,23 @@ type parseContext struct { // the '#' character. // The dot and wildcard character can be escaped with '\'. // -// { -// "name": {"first": "Tom", "last": "Anderson"}, -// "age":37, -// "children": ["Sara","Alex","Jack"], -// "friends": [ -// {"first": "James", "last": "Murphy"}, -// {"first": "Roger", "last": "Craig"} -// ] -// } -// "name.last" >> "Anderson" -// "age" >> 37 -// "children" >> ["Sara","Alex","Jack"] -// "children.#" >> 3 -// "children.1" >> "Alex" -// "child*.2" >> "Jack" -// "c?ildren.0" >> "Sara" -// "friends.#.first" >> ["James","Roger"] +// { +// "name": {"first": "Tom", "last": "Anderson"}, +// "age":37, +// "children": ["Sara","Alex","Jack"], +// "friends": [ +// {"first": "James", "last": "Murphy"}, +// {"first": "Roger", "last": "Craig"} +// ] +// } +// "name.last" >> "Anderson" +// "age" >> 37 +// "children" >> ["Sara","Alex","Jack"] +// "children.#" >> 3 +// "children.1" >> "Alex" +// "child*.2" >> "Jack" +// "c?ildren.0" >> "Sara" +// "friends.#.first" >> ["James","Roger"] // // This function expects that the json is well-formed, and does not validate. // Invalid json will not panic, but it may return back unexpected results. @@ -2123,11 +2310,10 @@ func unescape(json string) string { } // Less return true if a token is less than another token. -// The caseSensitive paramater is used when the tokens are Strings. +// The caseSensitive parameter is used when the tokens are Strings. // The order when comparing two different type is: // -// Null < False < Number < String < True < JSON -// +// Null < False < Number < String < True < JSON func (t Result) Less(token Result, caseSensitive bool) bool { if t.Type < token.Type { return true @@ -2556,11 +2742,10 @@ func validnull(data []byte, i int) (outi int, ok bool) { // Valid returns true if the input is valid json. // -// if !gjson.Valid(json) { -// return errors.New("invalid json") -// } -// value := gjson.Get(json, "name.last") -// +// if !gjson.Valid(json) { +// return errors.New("invalid json") +// } +// value := gjson.Get(json, "name.last") func Valid(json string) bool { _, ok := validpayload(stringBytes(json), 0) return ok @@ -2568,13 +2753,12 @@ func Valid(json string) bool { // ValidBytes returns true if the input is valid json. // -// if !gjson.Valid(json) { -// return errors.New("invalid json") -// } -// value := gjson.Get(json, "name.last") +// if !gjson.Valid(json) { +// return errors.New("invalid json") +// } +// value := gjson.Get(json, "name.last") // // If working with bytes, this method preferred over ValidBytes(string(data)) -// func ValidBytes(json []byte) bool { _, ok := validpayload(json, 0) return ok @@ -2690,6 +2874,7 @@ func execModifier(json, path string) (pathOut, res string, ok bool) { var parsedArgs bool switch pathOut[0] { case '{', '[', '"': + // json arg res := Parse(pathOut) if res.Exists() { args = squash(pathOut) @@ -2698,14 +2883,20 @@ func execModifier(json, path string) (pathOut, res string, ok bool) { } } if !parsedArgs { - idx := strings.IndexByte(pathOut, '|') - if idx == -1 { - args = pathOut - pathOut = "" - } else { - args = pathOut[:idx] - pathOut = pathOut[idx:] + // simple arg + i := 0 + for ; i < len(pathOut); i++ { + if pathOut[i] == '|' { + break + } + switch pathOut[i] { + case '{', '[', '"', '(': + s := squash(pathOut[i:]) + i += len(s) - 1 + } } + args = pathOut[:i] + pathOut = pathOut[i:] } } return pathOut, fn(json, args), true @@ -2725,19 +2916,24 @@ func unwrap(json string) string { // DisableModifiers will disable the modifier syntax var DisableModifiers = false -var modifiers = map[string]func(json, arg string) string{ - "pretty": modPretty, - "ugly": modUgly, - "reverse": modReverse, - "this": modThis, - "flatten": modFlatten, - "join": modJoin, - "valid": modValid, - "keys": modKeys, - "values": modValues, - "tostr": modToStr, - "fromstr": modFromStr, - "group": modGroup, +var modifiers map[string]func(json, arg string) string + +func init() { + modifiers = map[string]func(json, arg string) string{ + "pretty": modPretty, + "ugly": modUgly, + "reverse": modReverse, + "this": modThis, + "flatten": modFlatten, + "join": modJoin, + "valid": modValid, + "keys": modKeys, + "values": modValues, + "tostr": modToStr, + "fromstr": modFromStr, + "group": modGroup, + "dig": modDig, + } } // AddModifier binds a custom modifier command to the GJSON syntax. @@ -2848,9 +3044,13 @@ func modReverse(json, arg string) string { } // @flatten an array with child arrays. -// [1,[2],[3,4],[5,[6,7]]] -> [1,2,3,4,5,[6,7]] +// +// [1,[2],[3,4],[5,[6,7]]] -> [1,2,3,4,5,[6,7]] +// // The {"deep":true} arg can be provide for deep flattening. -// [1,[2],[3,4],[5,[6,7]]] -> [1,2,3,4,5,6,7] +// +// [1,[2],[3,4],[5,[6,7]]] -> [1,2,3,4,5,6,7] +// // The original json is returned when the json is not an array. func modFlatten(json, arg string) string { res := Parse(json) @@ -2895,7 +3095,8 @@ func modFlatten(json, arg string) string { } // @keys extracts the keys from an object. -// {"first":"Tom","last":"Smith"} -> ["first","last"] +// +// {"first":"Tom","last":"Smith"} -> ["first","last"] func modKeys(json, arg string) string { v := Parse(json) if !v.Exists() { @@ -2922,7 +3123,8 @@ func modKeys(json, arg string) string { } // @values extracts the values from an object. -// {"first":"Tom","last":"Smith"} -> ["Tom","Smith"] +// +// {"first":"Tom","last":"Smith"} -> ["Tom","Smith"] func modValues(json, arg string) string { v := Parse(json) if !v.Exists() { @@ -2947,11 +3149,17 @@ func modValues(json, arg string) string { } // @join multiple objects into a single object. -// [{"first":"Tom"},{"last":"Smith"}] -> {"first","Tom","last":"Smith"} +// +// [{"first":"Tom"},{"last":"Smith"}] -> {"first","Tom","last":"Smith"} +// // The arg can be "true" to specify that duplicate keys should be preserved. -// [{"first":"Tom","age":37},{"age":41}] -> {"first","Tom","age":37,"age":41} +// +// [{"first":"Tom","age":37},{"age":41}] -> {"first","Tom","age":37,"age":41} +// // Without preserved keys: -// [{"first":"Tom","age":37},{"age":41}] -> {"first","Tom","age":41} +// +// [{"first":"Tom","age":37},{"age":41}] -> {"first","Tom","age":41} +// // The original json is returned when the json is not an object. func modJoin(json, arg string) string { res := Parse(json) @@ -3024,7 +3232,8 @@ func modValid(json, arg string) string { } // @fromstr converts a string to json -// "{\"id\":1023,\"name\":\"alert\"}" -> {"id":1023,"name":"alert"} +// +// "{\"id\":1023,\"name\":\"alert\"}" -> {"id":1023,"name":"alert"} func modFromStr(json, arg string) string { if !Valid(json) { return "" @@ -3033,7 +3242,8 @@ func modFromStr(json, arg string) string { } // @tostr converts a string to json -// {"id":1023,"name":"alert"} -> "{\"id\":1023,\"name\":\"alert\"}" +// +// {"id":1023,"name":"alert"} -> "{\"id\":1023,\"name\":\"alert\"}" func modToStr(str, arg string) string { return string(AppendJSONString(nil, str)) } @@ -3210,11 +3420,11 @@ func revSquash(json string) string { // Paths returns the original GJSON paths for a Result where the Result came // from a simple query path that returns an array, like: // -// gjson.Get(json, "friends.#.first") +// gjson.Get(json, "friends.#.first") // // The returned value will be in the form of a JSON array: // -// ["friends.0.first","friends.1.first","friends.2.first"] +// ["friends.0.first","friends.1.first","friends.2.first"] // // The param 'json' must be the original JSON used when calling Get. // @@ -3239,11 +3449,11 @@ func (t Result) Paths(json string) []string { // Path returns the original GJSON path for a Result where the Result came // from a simple path that returns a single value, like: // -// gjson.Get(json, "friends.#(last=Murphy)") +// gjson.Get(json, "friends.#(last=Murphy)") // // The returned value will be in the form of a JSON string: // -// "friends.0" +// "friends.0" // // The param 'json' must be the original JSON used when calling Get. // @@ -3259,7 +3469,7 @@ func (t Result) Path(json string) string { goto fail } if !strings.HasPrefix(json[t.Index:], t.Raw) { - // Result is not at the JSON index as exepcted. + // Result is not at the JSON index as expected. goto fail } for ; i >= 0; i-- { @@ -3320,7 +3530,7 @@ func (t Result) Path(json string) string { if !rcomp.Exists() { goto fail } - comp := escapeComp(rcomp.String()) + comp := Escape(rcomp.String()) path = append(path, '.') path = append(path, comp...) } @@ -3335,17 +3545,31 @@ fail: // isSafePathKeyChar returns true if the input character is safe for not // needing escaping. func isSafePathKeyChar(c byte) bool { - return c <= ' ' || c > '~' || c == '_' || c == '-' || c == ':' || - (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || - (c >= '0' && c <= '9') + return (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || + (c >= '0' && c <= '9') || c <= ' ' || c > '~' || c == '_' || + c == '-' || c == ':' } -// escapeComp escaped a path compontent, making it safe for generating a -// path for later use. -func escapeComp(comp string) string { +// Escape returns an escaped path component. +// +// json := `{ +// "user":{ +// "first.name": "Janet", +// "last.name": "Prichard" +// } +// }` +// user := gjson.Get(json, "user") +// println(user.Get(gjson.Escape("first.name")) +// println(user.Get(gjson.Escape("last.name")) +// // Output: +// // Janet +// // Prichard +func Escape(comp string) string { for i := 0; i < len(comp); i++ { if !isSafePathKeyChar(comp[i]) { - ncomp := []byte(comp[:i]) + ncomp := make([]byte, len(comp)+1) + copy(ncomp, comp[:i]) + ncomp = ncomp[:i] for ; i < len(comp); i++ { if !isSafePathKeyChar(comp[i]) { ncomp = append(ncomp, '\\') @@ -3357,3 +3581,70 @@ func escapeComp(comp string) string { } return comp } + +func parseRecursiveDescent(all []Result, parent Result, path string) []Result { + if res := parent.Get(path); res.Exists() { + all = append(all, res) + } + if parent.IsArray() || parent.IsObject() { + parent.ForEach(func(_, val Result) bool { + all = parseRecursiveDescent(all, val, path) + return true + }) + } + return all +} + +func modDig(json, arg string) string { + all := parseRecursiveDescent(nil, Parse(json), arg) + var out []byte + out = append(out, '[') + for i, res := range all { + if i > 0 { + out = append(out, ',') + } + out = append(out, res.Raw...) + } + out = append(out, ']') + return string(out) +} + +// All iterates over a json result. +// This works identically to ForEach, but allows modern Go loops: +// +// for key, value := range res.All() { +// fmt.Printf("%s %s\n", key, value) +// } +func (t Result) All() iter.Seq2[Result, Result] { + return func(yield func(Result, Result) bool) { + t.ForEach(yield) + } +} + +// Keys iterates over a json result. +// This works identically to ForEach, but allows modern Go loops: +// +// for key := range res.Keys() { +// fmt.Printf("%s\n", key) +// } +func (t Result) Keys() iter.Seq[Result] { + return func(yield func(Result) bool) { + t.ForEach(func(key, _ Result) bool { + return yield(key) + }) + } +} + +// Values iterates over a json result. +// This works identically to ForEach, but allows modern Go loops: +// +// for value := range res.Values() { +// fmt.Printf("%s\n", value) +// } +func (t Result) Values() iter.Seq[Result] { + return func(yield func(Result) bool) { + t.ForEach(func(_, value Result) bool { + return yield(value) + }) + } +} diff --git a/vendor/github.com/tidwall/gjson/logo.png b/vendor/github.com/tidwall/gjson/logo.png deleted file mode 100644 index 17a8bbe9d6..0000000000 Binary files a/vendor/github.com/tidwall/gjson/logo.png and /dev/null differ diff --git a/vendor/github.com/tidwall/match/README.md b/vendor/github.com/tidwall/match/README.md index 5fdd4cf63d..9134079ed8 100644 --- a/vendor/github.com/tidwall/match/README.md +++ b/vendor/github.com/tidwall/match/README.md @@ -26,4 +26,4 @@ Josh Baker [@tidwall](http://twitter.com/tidwall) ## License -Redcon source code is available under the MIT [License](/LICENSE). +Match source code is available under the MIT [License](/LICENSE). diff --git a/vendor/github.com/tidwall/match/match.go b/vendor/github.com/tidwall/match/match.go index 11da28f1b9..d855a03fc6 100644 --- a/vendor/github.com/tidwall/match/match.go +++ b/vendor/github.com/tidwall/match/match.go @@ -10,18 +10,30 @@ import ( // and '?' matches on any one character. // // pattern: -// { term } +// +// { term } +// // term: -// '*' matches any sequence of non-Separator characters -// '?' matches any single non-Separator character -// c matches character c (c != '*', '?', '\\') -// '\\' c matches character c // +// '*' matches any sequence of non-Separator characters +// '?' matches any single non-Separator character +// c matches character c (c != '*', '?', '\\') +// '\\' c matches character c func Match(str, pattern string) bool { + return match0(str, pattern, false) +} + +// MatchNoCase is the same as Match but performs a case-insensitive match. +// Such that string "Hello World" with match with lower case pattern "hello*" +func MatchNoCase(str, pattern string) bool { + return match0(str, pattern, true) +} + +func match0(str, pattern string, nocase bool) bool { if pattern == "*" { return true } - return match(str, pattern, 0, nil, -1) == rMatch + return match(str, pattern, 0, nil, -1, nocase) == rMatch } // MatchLimit is the same as Match but will limit the complexity of the match @@ -34,11 +46,21 @@ func Match(str, pattern string) bool { // Everytime it calls itself a counter is incremented. // The operation is stopped when counter > maxcomp*len(str). func MatchLimit(str, pattern string, maxcomp int) (matched, stopped bool) { + return matchLimit0(str, pattern, maxcomp, false) +} + +func MatchLimitNoCase(str, pattern string, maxcomp int, +) (matched, stopped bool) { + return matchLimit0(str, pattern, maxcomp, true) +} + +func matchLimit0(str, pattern string, maxcomp int, nocase bool, +) (matched, stopped bool) { if pattern == "*" { return true, false } counter := 0 - r := match(str, pattern, len(str), &counter, maxcomp) + r := match(str, pattern, len(str), &counter, maxcomp, nocase) if r == rStop { return false, true } @@ -53,7 +75,15 @@ const ( rStop ) -func match(str, pat string, slen int, counter *int, maxcomp int) result { +func tolower(r rune) rune { + if r >= 'A' && r <= 'Z' { + return r + 32 + } + return r +} + +func match(str, pat string, slen int, counter *int, maxcomp int, nocase bool, +) result { // check complexity limit if maxcomp > -1 { if *counter > slen*maxcomp { @@ -94,7 +124,7 @@ func match(str, pat string, slen int, counter *int, maxcomp int) result { // Match and trim any non-wildcard suffix characters. var ok bool - str, pat, ok = matchTrimSuffix(str, pat) + str, pat, ok = matchTrimSuffix(str, pat, nocase) if !ok { return rNoMatch } @@ -105,7 +135,7 @@ func match(str, pat string, slen int, counter *int, maxcomp int) result { } // Perform recursive wildcard search. - r := match(str, pat[1:], slen, counter, maxcomp) + r := match(str, pat[1:], slen, counter, maxcomp, nocase) if r != rNoMatch { return r } @@ -124,6 +154,9 @@ func match(str, pat string, slen int, counter *int, maxcomp int) result { return rNoMatch } } + if nocase { + sc, pc = tolower(sc), tolower(pc) + } if sc != pc { return rNoMatch } @@ -150,7 +183,7 @@ func match(str, pat string, slen int, counter *int, maxcomp int) result { // // Any matched characters will be trimmed from both the target // string and the pattern. -func matchTrimSuffix(str, pat string) (string, string, bool) { +func matchTrimSuffix(str, pat string, nocase bool) (string, string, bool) { // It's expected that the pattern has at least two bytes and the first byte // is a wildcard star '*' match := true @@ -171,6 +204,9 @@ func matchTrimSuffix(str, pat string) (string, string, bool) { break } sc, ss := utf8.DecodeLastRuneInString(str) + if nocase { + pc, sc = tolower(pc), tolower(sc) + } if !((pc == '?' && !esc) || pc == sc) { match = false break diff --git a/vendor/github.com/tidwall/sjson/LICENSE b/vendor/github.com/tidwall/sjson/LICENSE new file mode 100644 index 0000000000..89593c7c84 --- /dev/null +++ b/vendor/github.com/tidwall/sjson/LICENSE @@ -0,0 +1,21 @@ +The MIT License (MIT) + +Copyright (c) 2016 Josh Baker + +Permission is hereby granted, free of charge, to any person obtaining a copy of +this software and associated documentation files (the "Software"), to deal in +the Software without restriction, including without limitation the rights to +use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of +the Software, and to permit persons to whom the Software is furnished to do so, +subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS +FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR +COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER +IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN +CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + diff --git a/vendor/github.com/tidwall/sjson/README.md b/vendor/github.com/tidwall/sjson/README.md new file mode 100644 index 0000000000..4598424efa --- /dev/null +++ b/vendor/github.com/tidwall/sjson/README.md @@ -0,0 +1,278 @@ +

+SJSON +
+GoDoc +

+ +

set a json value quickly

+ +SJSON is a Go package that provides a [very fast](#performance) and simple way to set a value in a json document. +For quickly retrieving json values check out [GJSON](https://github.com/tidwall/gjson). + +For a command line interface check out [JJ](https://github.com/tidwall/jj). + +Getting Started +=============== + +Installing +---------- + +To start using SJSON, install Go and run `go get`: + +```sh +$ go get -u github.com/tidwall/sjson +``` + +This will retrieve the library. + +Set a value +----------- +Set sets the value for the specified path. +A path is in dot syntax, such as "name.last" or "age". +This function expects that the json is well-formed and validated. +Invalid json will not panic, but it may return back unexpected results. +Invalid paths may return an error. + +```go +package main + +import "github.com/tidwall/sjson" + +const json = `{"name":{"first":"Janet","last":"Prichard"},"age":47}` + +func main() { + value, _ := sjson.Set(json, "name.last", "Anderson") + println(value) +} +``` + +This will print: + +```json +{"name":{"first":"Janet","last":"Anderson"},"age":47} +``` + +Path syntax +----------- + +A path is a series of keys separated by a dot. +The dot and colon characters can be escaped with ``\``. + +```json +{ + "name": {"first": "Tom", "last": "Anderson"}, + "age":37, + "children": ["Sara","Alex","Jack"], + "fav.movie": "Deer Hunter", + "friends": [ + {"first": "James", "last": "Murphy"}, + {"first": "Roger", "last": "Craig"} + ] +} +``` +``` +"name.last" >> "Anderson" +"age" >> 37 +"children.1" >> "Alex" +"friends.1.last" >> "Craig" +``` + +The `-1` key can be used to append a value to an existing array: + +``` +"children.-1" >> appends a new value to the end of the children array +``` + +Normally number keys are used to modify arrays, but it's possible to force a numeric object key by using the colon character: + +```json +{ + "users":{ + "2313":{"name":"Sara"}, + "7839":{"name":"Andy"} + } +} +``` + +A colon path would look like: + +``` +"users.:2313.name" >> "Sara" +``` + +Supported types +--------------- + +Pretty much any type is supported: + +```go +sjson.Set(`{"key":true}`, "key", nil) +sjson.Set(`{"key":true}`, "key", false) +sjson.Set(`{"key":true}`, "key", 1) +sjson.Set(`{"key":true}`, "key", 10.5) +sjson.Set(`{"key":true}`, "key", "hello") +sjson.Set(`{"key":true}`, "key", []string{"hello", "world"}) +sjson.Set(`{"key":true}`, "key", map[string]interface{}{"hello":"world"}) +``` + +When a type is not recognized, SJSON will fallback to the `encoding/json` Marshaller. + + +Examples +-------- + +Set a value from empty document: +```go +value, _ := sjson.Set("", "name", "Tom") +println(value) + +// Output: +// {"name":"Tom"} +``` + +Set a nested value from empty document: +```go +value, _ := sjson.Set("", "name.last", "Anderson") +println(value) + +// Output: +// {"name":{"last":"Anderson"}} +``` + +Set a new value: +```go +value, _ := sjson.Set(`{"name":{"last":"Anderson"}}`, "name.first", "Sara") +println(value) + +// Output: +// {"name":{"first":"Sara","last":"Anderson"}} +``` + +Update an existing value: +```go +value, _ := sjson.Set(`{"name":{"last":"Anderson"}}`, "name.last", "Smith") +println(value) + +// Output: +// {"name":{"last":"Smith"}} +``` + +Set a new array value: +```go +value, _ := sjson.Set(`{"friends":["Andy","Carol"]}`, "friends.2", "Sara") +println(value) + +// Output: +// {"friends":["Andy","Carol","Sara"] +``` + +Append an array value by using the `-1` key in a path: +```go +value, _ := sjson.Set(`{"friends":["Andy","Carol"]}`, "friends.-1", "Sara") +println(value) + +// Output: +// {"friends":["Andy","Carol","Sara"] +``` + +Append an array value that is past the end: +```go +value, _ := sjson.Set(`{"friends":["Andy","Carol"]}`, "friends.4", "Sara") +println(value) + +// Output: +// {"friends":["Andy","Carol",null,null,"Sara"] +``` + +Delete a value: +```go +value, _ := sjson.Delete(`{"name":{"first":"Sara","last":"Anderson"}}`, "name.first") +println(value) + +// Output: +// {"name":{"last":"Anderson"}} +``` + +Delete an array value: +```go +value, _ := sjson.Delete(`{"friends":["Andy","Carol"]}`, "friends.1") +println(value) + +// Output: +// {"friends":["Andy"]} +``` + +Delete the last array value: +```go +value, _ := sjson.Delete(`{"friends":["Andy","Carol"]}`, "friends.-1") +println(value) + +// Output: +// {"friends":["Andy"]} +``` + +## Performance + +Benchmarks of SJSON alongside [encoding/json](https://golang.org/pkg/encoding/json/), +[ffjson](https://github.com/pquerna/ffjson), +[EasyJSON](https://github.com/mailru/easyjson), +and [Gabs](https://github.com/Jeffail/gabs) + +``` +Benchmark_SJSON-8 3000000 805 ns/op 1077 B/op 3 allocs/op +Benchmark_SJSON_ReplaceInPlace-8 3000000 449 ns/op 0 B/op 0 allocs/op +Benchmark_JSON_Map-8 300000 21236 ns/op 6392 B/op 150 allocs/op +Benchmark_JSON_Struct-8 300000 14691 ns/op 1789 B/op 24 allocs/op +Benchmark_Gabs-8 300000 21311 ns/op 6752 B/op 150 allocs/op +Benchmark_FFJSON-8 300000 17673 ns/op 3589 B/op 47 allocs/op +Benchmark_EasyJSON-8 1500000 3119 ns/op 1061 B/op 13 allocs/op +``` + +JSON document used: + +```json +{ + "widget": { + "debug": "on", + "window": { + "title": "Sample Konfabulator Widget", + "name": "main_window", + "width": 500, + "height": 500 + }, + "image": { + "src": "Images/Sun.png", + "hOffset": 250, + "vOffset": 250, + "alignment": "center" + }, + "text": { + "data": "Click Here", + "size": 36, + "style": "bold", + "vOffset": 100, + "alignment": "center", + "onMouseUp": "sun1.opacity = (sun1.opacity / 100) * 90;" + } + } +} +``` + +Each operation was rotated though one of the following search paths: + +``` +widget.window.name +widget.image.hOffset +widget.text.onMouseUp +``` + +*These benchmarks were run on a MacBook Pro 15" 2.8 GHz Intel Core i7 using Go 1.7 and can be be found [here](https://github.com/tidwall/sjson-benchmarks)*. + +## Contact +Josh Baker [@tidwall](http://twitter.com/tidwall) + +## License + +SJSON source code is available under the MIT [License](/LICENSE). diff --git a/vendor/github.com/tidwall/sjson/logo.png b/vendor/github.com/tidwall/sjson/logo.png new file mode 100644 index 0000000000..b5aa257b6b Binary files /dev/null and b/vendor/github.com/tidwall/sjson/logo.png differ diff --git a/vendor/github.com/tidwall/sjson/sjson.go b/vendor/github.com/tidwall/sjson/sjson.go new file mode 100644 index 0000000000..a55eef3fdb --- /dev/null +++ b/vendor/github.com/tidwall/sjson/sjson.go @@ -0,0 +1,737 @@ +// Package sjson provides setting json values. +package sjson + +import ( + jsongo "encoding/json" + "sort" + "strconv" + "unsafe" + + "github.com/tidwall/gjson" +) + +type errorType struct { + msg string +} + +func (err *errorType) Error() string { + return err.msg +} + +// Options represents additional options for the Set and Delete functions. +type Options struct { + // Optimistic is a hint that the value likely exists which + // allows for the sjson to perform a fast-track search and replace. + Optimistic bool + // ReplaceInPlace is a hint to replace the input json rather than + // allocate a new json byte slice. When this field is specified + // the input json will not longer be valid and it should not be used + // In the case when the destination slice doesn't have enough free + // bytes to replace the data in place, a new bytes slice will be + // created under the hood. + // The Optimistic flag must be set to true and the input must be a + // byte slice in order to use this field. + ReplaceInPlace bool +} + +type pathResult struct { + part string // current key part + gpart string // gjson get part + path string // remaining path + force bool // force a string key + more bool // there is more path to parse +} + +func isSimpleChar(ch byte) bool { + switch ch { + case '|', '#', '@', '*', '?': + return false + default: + return true + } +} + +func parsePath(path string) (res pathResult, simple bool) { + var r pathResult + if len(path) > 0 && path[0] == ':' { + r.force = true + path = path[1:] + } + for i := 0; i < len(path); i++ { + if path[i] == '.' { + r.part = path[:i] + r.gpart = path[:i] + r.path = path[i+1:] + r.more = true + return r, true + } + if !isSimpleChar(path[i]) { + return r, false + } + if path[i] == '\\' { + // go into escape mode. this is a slower path that + // strips off the escape character from the part. + epart := []byte(path[:i]) + gpart := []byte(path[:i+1]) + i++ + if i < len(path) { + epart = append(epart, path[i]) + gpart = append(gpart, path[i]) + i++ + for ; i < len(path); i++ { + if path[i] == '\\' { + gpart = append(gpart, '\\') + i++ + if i < len(path) { + epart = append(epart, path[i]) + gpart = append(gpart, path[i]) + } + continue + } else if path[i] == '.' { + r.part = string(epart) + r.gpart = string(gpart) + r.path = path[i+1:] + r.more = true + return r, true + } else if !isSimpleChar(path[i]) { + return r, false + } + epart = append(epart, path[i]) + gpart = append(gpart, path[i]) + } + } + // append the last part + r.part = string(epart) + r.gpart = string(gpart) + return r, true + } + } + r.part = path + r.gpart = path + return r, true +} + +func mustMarshalString(s string) bool { + for i := 0; i < len(s); i++ { + if s[i] < ' ' || s[i] > 0x7f || s[i] == '"' || s[i] == '\\' { + return true + } + } + return false +} + +// appendStringify makes a json string and appends to buf. +func appendStringify(buf []byte, s string) []byte { + if mustMarshalString(s) { + b, _ := jsongo.Marshal(s) + return append(buf, b...) + } + buf = append(buf, '"') + buf = append(buf, s...) + buf = append(buf, '"') + return buf +} + +// appendBuild builds a json block from a json path. +func appendBuild(buf []byte, array bool, paths []pathResult, raw string, + stringify bool) []byte { + if !array { + buf = appendStringify(buf, paths[0].part) + buf = append(buf, ':') + } + if len(paths) > 1 { + n, numeric := atoui(paths[1]) + if numeric || (!paths[1].force && paths[1].part == "-1") { + buf = append(buf, '[') + buf = appendRepeat(buf, "null,", n) + buf = appendBuild(buf, true, paths[1:], raw, stringify) + buf = append(buf, ']') + } else { + buf = append(buf, '{') + buf = appendBuild(buf, false, paths[1:], raw, stringify) + buf = append(buf, '}') + } + } else { + if stringify { + buf = appendStringify(buf, raw) + } else { + buf = append(buf, raw...) + } + } + return buf +} + +// atoui does a rip conversion of string -> unigned int. +func atoui(r pathResult) (n int, ok bool) { + if r.force { + return 0, false + } + for i := 0; i < len(r.part); i++ { + if r.part[i] < '0' || r.part[i] > '9' { + return 0, false + } + n = n*10 + int(r.part[i]-'0') + } + return n, true +} + +// appendRepeat repeats string "n" times and appends to buf. +func appendRepeat(buf []byte, s string, n int) []byte { + for i := 0; i < n; i++ { + buf = append(buf, s...) + } + return buf +} + +// trim does a rip trim +func trim(s string) string { + for len(s) > 0 { + if s[0] <= ' ' { + s = s[1:] + continue + } + break + } + for len(s) > 0 { + if s[len(s)-1] <= ' ' { + s = s[:len(s)-1] + continue + } + break + } + return s +} + +// deleteTailItem deletes the previous key or comma. +func deleteTailItem(buf []byte) ([]byte, bool) { +loop: + for i := len(buf) - 1; i >= 0; i-- { + // look for either a ',',':','[' + switch buf[i] { + case '[': + return buf, true + case ',': + return buf[:i], false + case ':': + // delete tail string + i-- + for ; i >= 0; i-- { + if buf[i] == '"' { + i-- + for ; i >= 0; i-- { + if buf[i] == '"' { + i-- + if i >= 0 && buf[i] == '\\' { + i-- + continue + } + for ; i >= 0; i-- { + // look for either a ',','{' + switch buf[i] { + case '{': + return buf[:i+1], true + case ',': + return buf[:i], false + } + } + } + } + break + } + } + break loop + } + } + return buf, false +} + +var errNoChange = &errorType{"no change"} + +func appendRawPaths(buf []byte, jstr string, paths []pathResult, raw string, + stringify, del bool) ([]byte, error) { + var err error + var res gjson.Result + var found bool + if del { + if paths[0].part == "-1" && !paths[0].force { + res = gjson.Get(jstr, "#") + if res.Int() > 0 { + res = gjson.Get(jstr, strconv.FormatInt(int64(res.Int()-1), 10)) + found = true + } + } + } + if !found { + res = gjson.Get(jstr, paths[0].gpart) + } + if res.Index > 0 { + if len(paths) > 1 { + buf = append(buf, jstr[:res.Index]...) + buf, err = appendRawPaths(buf, res.Raw, paths[1:], raw, + stringify, del) + if err != nil { + return nil, err + } + buf = append(buf, jstr[res.Index+len(res.Raw):]...) + return buf, nil + } + buf = append(buf, jstr[:res.Index]...) + var exidx int // additional forward stripping + if del { + var delNextComma bool + buf, delNextComma = deleteTailItem(buf) + if delNextComma { + i, j := res.Index+len(res.Raw), 0 + for ; i < len(jstr); i, j = i+1, j+1 { + if jstr[i] <= ' ' { + continue + } + if jstr[i] == ',' { + exidx = j + 1 + } + break + } + } + } else { + if stringify { + buf = appendStringify(buf, raw) + } else { + buf = append(buf, raw...) + } + } + buf = append(buf, jstr[res.Index+len(res.Raw)+exidx:]...) + return buf, nil + } + if del { + return nil, errNoChange + } + n, numeric := atoui(paths[0]) + isempty := true + for i := 0; i < len(jstr); i++ { + if jstr[i] > ' ' { + isempty = false + break + } + } + if isempty { + if numeric { + jstr = "[]" + } else { + jstr = "{}" + } + } + jsres := gjson.Parse(jstr) + if jsres.Type != gjson.JSON { + if numeric { + jstr = "[]" + } else { + jstr = "{}" + } + jsres = gjson.Parse(jstr) + } + var comma bool + for i := 1; i < len(jsres.Raw); i++ { + if jsres.Raw[i] <= ' ' { + continue + } + if jsres.Raw[i] == '}' || jsres.Raw[i] == ']' { + break + } + comma = true + break + } + switch jsres.Raw[0] { + default: + return nil, &errorType{"json must be an object or array"} + case '{': + end := len(jsres.Raw) - 1 + for ; end > 0; end-- { + if jsres.Raw[end] == '}' { + break + } + } + buf = append(buf, jsres.Raw[:end]...) + if comma { + buf = append(buf, ',') + } + buf = appendBuild(buf, false, paths, raw, stringify) + buf = append(buf, '}') + return buf, nil + case '[': + var appendit bool + if !numeric { + if paths[0].part == "-1" && !paths[0].force { + appendit = true + } else { + return nil, &errorType{ + "cannot set array element for non-numeric key '" + + paths[0].part + "'"} + } + } + if appendit { + njson := trim(jsres.Raw) + if njson[len(njson)-1] == ']' { + njson = njson[:len(njson)-1] + } + buf = append(buf, njson...) + if comma { + buf = append(buf, ',') + } + + buf = appendBuild(buf, true, paths, raw, stringify) + buf = append(buf, ']') + return buf, nil + } + buf = append(buf, '[') + ress := jsres.Array() + for i := 0; i < len(ress); i++ { + if i > 0 { + buf = append(buf, ',') + } + buf = append(buf, ress[i].Raw...) + } + if len(ress) == 0 { + buf = appendRepeat(buf, "null,", n-len(ress)) + } else { + buf = appendRepeat(buf, ",null", n-len(ress)) + if comma { + buf = append(buf, ',') + } + } + buf = appendBuild(buf, true, paths, raw, stringify) + buf = append(buf, ']') + return buf, nil + } +} + +func isOptimisticPath(path string) bool { + for i := 0; i < len(path); i++ { + if path[i] < '.' || path[i] > 'z' { + return false + } + if path[i] > '9' && path[i] < 'A' { + return false + } + if path[i] > 'z' { + return false + } + } + return true +} + +// Set sets a json value for the specified path. +// A path is in dot syntax, such as "name.last" or "age". +// This function expects that the json is well-formed, and does not validate. +// Invalid json will not panic, but it may return back unexpected results. +// An error is returned if the path is not valid. +// +// A path is a series of keys separated by a dot. +// +// { +// "name": {"first": "Tom", "last": "Anderson"}, +// "age":37, +// "children": ["Sara","Alex","Jack"], +// "friends": [ +// {"first": "James", "last": "Murphy"}, +// {"first": "Roger", "last": "Craig"} +// ] +// } +// "name.last" >> "Anderson" +// "age" >> 37 +// "children.1" >> "Alex" +// +func Set(json, path string, value interface{}) (string, error) { + return SetOptions(json, path, value, nil) +} + +// SetBytes sets a json value for the specified path. +// If working with bytes, this method preferred over +// Set(string(data), path, value) +func SetBytes(json []byte, path string, value interface{}) ([]byte, error) { + return SetBytesOptions(json, path, value, nil) +} + +// SetRaw sets a raw json value for the specified path. +// This function works the same as Set except that the value is set as a +// raw block of json. This allows for setting premarshalled json objects. +func SetRaw(json, path, value string) (string, error) { + return SetRawOptions(json, path, value, nil) +} + +// SetRawOptions sets a raw json value for the specified path with options. +// This furnction works the same as SetOptions except that the value is set +// as a raw block of json. This allows for setting premarshalled json objects. +func SetRawOptions(json, path, value string, opts *Options) (string, error) { + var optimistic bool + if opts != nil { + optimistic = opts.Optimistic + } + res, err := set(json, path, value, false, false, optimistic, false) + if err == errNoChange { + return json, nil + } + return string(res), err +} + +// SetRawBytes sets a raw json value for the specified path. +// If working with bytes, this method preferred over +// SetRaw(string(data), path, value) +func SetRawBytes(json []byte, path string, value []byte) ([]byte, error) { + return SetRawBytesOptions(json, path, value, nil) +} + +type dtype struct{} + +// Delete deletes a value from json for the specified path. +func Delete(json, path string) (string, error) { + return Set(json, path, dtype{}) +} + +// DeleteBytes deletes a value from json for the specified path. +func DeleteBytes(json []byte, path string) ([]byte, error) { + return SetBytes(json, path, dtype{}) +} + +type stringHeader struct { + data unsafe.Pointer + len int +} + +type sliceHeader struct { + data unsafe.Pointer + len int + cap int +} + +func set(jstr, path, raw string, + stringify, del, optimistic, inplace bool) ([]byte, error) { + if path == "" { + return []byte(jstr), &errorType{"path cannot be empty"} + } + if !del && optimistic && isOptimisticPath(path) { + res := gjson.Get(jstr, path) + if res.Exists() && res.Index > 0 { + sz := len(jstr) - len(res.Raw) + len(raw) + if stringify { + sz += 2 + } + if inplace && sz <= len(jstr) { + if !stringify || !mustMarshalString(raw) { + jsonh := *(*stringHeader)(unsafe.Pointer(&jstr)) + jsonbh := sliceHeader{ + data: jsonh.data, len: jsonh.len, cap: jsonh.len} + jbytes := *(*[]byte)(unsafe.Pointer(&jsonbh)) + if stringify { + jbytes[res.Index] = '"' + copy(jbytes[res.Index+1:], []byte(raw)) + jbytes[res.Index+1+len(raw)] = '"' + copy(jbytes[res.Index+1+len(raw)+1:], + jbytes[res.Index+len(res.Raw):]) + } else { + copy(jbytes[res.Index:], []byte(raw)) + copy(jbytes[res.Index+len(raw):], + jbytes[res.Index+len(res.Raw):]) + } + return jbytes[:sz], nil + } + return []byte(jstr), nil + } + buf := make([]byte, 0, sz) + buf = append(buf, jstr[:res.Index]...) + if stringify { + buf = appendStringify(buf, raw) + } else { + buf = append(buf, raw...) + } + buf = append(buf, jstr[res.Index+len(res.Raw):]...) + return buf, nil + } + } + var paths []pathResult + r, simple := parsePath(path) + if simple { + paths = append(paths, r) + for r.more { + r, simple = parsePath(r.path) + if !simple { + break + } + paths = append(paths, r) + } + } + if !simple { + if del { + return []byte(jstr), + &errorType{"cannot delete value from a complex path"} + } + return setComplexPath(jstr, path, raw, stringify) + } + njson, err := appendRawPaths(nil, jstr, paths, raw, stringify, del) + if err != nil { + return []byte(jstr), err + } + return njson, nil +} + +func setComplexPath(jstr, path, raw string, stringify bool) ([]byte, error) { + res := gjson.Get(jstr, path) + if !res.Exists() || !(res.Index != 0 || len(res.Indexes) != 0) { + return []byte(jstr), errNoChange + } + if res.Index != 0 { + njson := []byte(jstr[:res.Index]) + if stringify { + njson = appendStringify(njson, raw) + } else { + njson = append(njson, raw...) + } + njson = append(njson, jstr[res.Index+len(res.Raw):]...) + jstr = string(njson) + } + if len(res.Indexes) > 0 { + type val struct { + index int + res gjson.Result + } + vals := make([]val, 0, len(res.Indexes)) + res.ForEach(func(_, vres gjson.Result) bool { + vals = append(vals, val{res: vres}) + return true + }) + if len(res.Indexes) != len(vals) { + return []byte(jstr), errNoChange + } + for i := 0; i < len(res.Indexes); i++ { + vals[i].index = res.Indexes[i] + } + sort.SliceStable(vals, func(i, j int) bool { + return vals[i].index > vals[j].index + }) + for _, val := range vals { + vres := val.res + index := val.index + njson := []byte(jstr[:index]) + if stringify { + njson = appendStringify(njson, raw) + } else { + njson = append(njson, raw...) + } + njson = append(njson, jstr[index+len(vres.Raw):]...) + jstr = string(njson) + } + } + return []byte(jstr), nil +} + +// SetOptions sets a json value for the specified path with options. +// A path is in dot syntax, such as "name.last" or "age". +// This function expects that the json is well-formed, and does not validate. +// Invalid json will not panic, but it may return back unexpected results. +// An error is returned if the path is not valid. +func SetOptions(json, path string, value interface{}, + opts *Options) (string, error) { + if opts != nil { + if opts.ReplaceInPlace { + // it's not safe to replace bytes in-place for strings + // copy the Options and set options.ReplaceInPlace to false. + nopts := *opts + opts = &nopts + opts.ReplaceInPlace = false + } + } + jsonh := *(*stringHeader)(unsafe.Pointer(&json)) + jsonbh := sliceHeader{data: jsonh.data, len: jsonh.len, cap: jsonh.len} + jsonb := *(*[]byte)(unsafe.Pointer(&jsonbh)) + res, err := SetBytesOptions(jsonb, path, value, opts) + return string(res), err +} + +// SetBytesOptions sets a json value for the specified path with options. +// If working with bytes, this method preferred over +// SetOptions(string(data), path, value) +func SetBytesOptions(json []byte, path string, value interface{}, + opts *Options) ([]byte, error) { + var optimistic, inplace bool + if opts != nil { + optimistic = opts.Optimistic + inplace = opts.ReplaceInPlace + } + jstr := *(*string)(unsafe.Pointer(&json)) + var res []byte + var err error + switch v := value.(type) { + default: + b, merr := jsongo.Marshal(value) + if merr != nil { + return nil, merr + } + raw := *(*string)(unsafe.Pointer(&b)) + res, err = set(jstr, path, raw, false, false, optimistic, inplace) + case dtype: + res, err = set(jstr, path, "", false, true, optimistic, inplace) + case string: + res, err = set(jstr, path, v, true, false, optimistic, inplace) + case []byte: + raw := *(*string)(unsafe.Pointer(&v)) + res, err = set(jstr, path, raw, true, false, optimistic, inplace) + case bool: + if v { + res, err = set(jstr, path, "true", false, false, optimistic, inplace) + } else { + res, err = set(jstr, path, "false", false, false, optimistic, inplace) + } + case int8: + res, err = set(jstr, path, strconv.FormatInt(int64(v), 10), + false, false, optimistic, inplace) + case int16: + res, err = set(jstr, path, strconv.FormatInt(int64(v), 10), + false, false, optimistic, inplace) + case int32: + res, err = set(jstr, path, strconv.FormatInt(int64(v), 10), + false, false, optimistic, inplace) + case int64: + res, err = set(jstr, path, strconv.FormatInt(int64(v), 10), + false, false, optimistic, inplace) + case uint8: + res, err = set(jstr, path, strconv.FormatUint(uint64(v), 10), + false, false, optimistic, inplace) + case uint16: + res, err = set(jstr, path, strconv.FormatUint(uint64(v), 10), + false, false, optimistic, inplace) + case uint32: + res, err = set(jstr, path, strconv.FormatUint(uint64(v), 10), + false, false, optimistic, inplace) + case uint64: + res, err = set(jstr, path, strconv.FormatUint(uint64(v), 10), + false, false, optimistic, inplace) + case float32: + res, err = set(jstr, path, strconv.FormatFloat(float64(v), 'f', -1, 64), + false, false, optimistic, inplace) + case float64: + res, err = set(jstr, path, strconv.FormatFloat(float64(v), 'f', -1, 64), + false, false, optimistic, inplace) + } + if err == errNoChange { + return json, nil + } + return res, err +} + +// SetRawBytesOptions sets a raw json value for the specified path with options. +// If working with bytes, this method preferred over +// SetRawOptions(string(data), path, value, opts) +func SetRawBytesOptions(json []byte, path string, value []byte, + opts *Options) ([]byte, error) { + jstr := *(*string)(unsafe.Pointer(&json)) + vstr := *(*string)(unsafe.Pointer(&value)) + var optimistic, inplace bool + if opts != nil { + optimistic = opts.Optimistic + inplace = opts.ReplaceInPlace + } + res, err := set(jstr, path, vstr, false, false, optimistic, inplace) + if err == errNoChange { + return json, nil + } + return res, err +} diff --git a/vendor/modules.txt b/vendor/modules.txt index 1efcc98c76..95aec35697 100644 --- a/vendor/modules.txt +++ b/vendor/modules.txt @@ -384,6 +384,23 @@ github.com/jackc/pgtype github.com/jackc/pgx/v4 github.com/jackc/pgx/v4/internal/sanitize github.com/jackc/pgx/v4/stdlib +# github.com/jackc/pgx/v5 v5.10.0 +## explicit; go 1.25.0 +github.com/jackc/pgx/v5 +github.com/jackc/pgx/v5/internal/iobufpool +github.com/jackc/pgx/v5/internal/pgio +github.com/jackc/pgx/v5/internal/sanitize +github.com/jackc/pgx/v5/internal/stmtcache +github.com/jackc/pgx/v5/pgconn +github.com/jackc/pgx/v5/pgconn/ctxwatch +github.com/jackc/pgx/v5/pgconn/internal/bgreader +github.com/jackc/pgx/v5/pgproto3 +github.com/jackc/pgx/v5/pgtype +github.com/jackc/pgx/v5/pgxpool +# github.com/jackc/puddle/v2 v2.2.2 +## explicit; go 1.19 +github.com/jackc/puddle/v2 +github.com/jackc/puddle/v2/internal/genstack # github.com/jferrl/go-githubauth v1.1.0 ## explicit; go 1.22 github.com/jferrl/go-githubauth @@ -552,6 +569,59 @@ github.com/prometheus/common/model github.com/prometheus/procfs github.com/prometheus/procfs/internal/fs github.com/prometheus/procfs/internal/util +# github.com/riverqueue/river v0.43.0 +## explicit; go 1.25.0 +github.com/riverqueue/river +github.com/riverqueue/river/internal/dblist +github.com/riverqueue/river/internal/dbunique +github.com/riverqueue/river/internal/execution +github.com/riverqueue/river/internal/jobcompleter +github.com/riverqueue/river/internal/jobexecutor +github.com/riverqueue/river/internal/jobstats +github.com/riverqueue/river/internal/leadership +github.com/riverqueue/river/internal/maintenance +github.com/riverqueue/river/internal/notifier +github.com/riverqueue/river/internal/notifylimiter +github.com/riverqueue/river/internal/pluginconfig +github.com/riverqueue/river/internal/pluginlookup +github.com/riverqueue/river/internal/retrypolicy +github.com/riverqueue/river/internal/rivercommon +github.com/riverqueue/river/internal/riverplugin +github.com/riverqueue/river/internal/util/chanutil +github.com/riverqueue/river/internal/workunit +github.com/riverqueue/river/rivermigrate +# github.com/riverqueue/river/riverdriver v0.43.0 +## explicit; go 1.25.0 +github.com/riverqueue/river/riverdriver +# github.com/riverqueue/river/riverdriver/riverpgxv5 v0.43.0 +## explicit; go 1.25.0 +github.com/riverqueue/river/riverdriver/riverpgxv5 +github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc +# github.com/riverqueue/river/rivershared v0.43.0 +## explicit; go 1.25.0 +github.com/riverqueue/river/rivershared/baseservice +github.com/riverqueue/river/rivershared/circuitbreaker +github.com/riverqueue/river/rivershared/levenshtein +github.com/riverqueue/river/rivershared/riverpilot +github.com/riverqueue/river/rivershared/riversharedmaintenance +github.com/riverqueue/river/rivershared/sqlctemplate +github.com/riverqueue/river/rivershared/startstop +github.com/riverqueue/river/rivershared/structtag +github.com/riverqueue/river/rivershared/testsignal +github.com/riverqueue/river/rivershared/uniquestates +github.com/riverqueue/river/rivershared/util/dbutil +github.com/riverqueue/river/rivershared/util/maputil +github.com/riverqueue/river/rivershared/util/ptrutil +github.com/riverqueue/river/rivershared/util/randutil +github.com/riverqueue/river/rivershared/util/serviceutil +github.com/riverqueue/river/rivershared/util/sliceutil +github.com/riverqueue/river/rivershared/util/testutil +github.com/riverqueue/river/rivershared/util/timeoututil +github.com/riverqueue/river/rivershared/util/timeutil +github.com/riverqueue/river/rivershared/util/valutil +# github.com/riverqueue/river/rivertype v0.43.0 +## explicit; go 1.25.0 +github.com/riverqueue/river/rivertype # github.com/shirou/gopsutil/v4 v4.26.5 ## explicit; go 1.24.0 github.com/shirou/gopsutil/v4/common @@ -612,15 +682,18 @@ github.com/testcontainers/testcontainers-go/wait # github.com/testcontainers/testcontainers-go/modules/postgres v0.43.0 ## explicit; go 1.25.0 github.com/testcontainers/testcontainers-go/modules/postgres -# github.com/tidwall/gjson v1.14.4 -## explicit; go 1.12 +# github.com/tidwall/gjson v1.19.0 +## explicit; go 1.23 github.com/tidwall/gjson -# github.com/tidwall/match v1.1.1 +# github.com/tidwall/match v1.2.0 ## explicit; go 1.15 github.com/tidwall/match # github.com/tidwall/pretty v1.2.1 ## explicit; go 1.16 github.com/tidwall/pretty +# github.com/tidwall/sjson v1.2.5 +## explicit; go 1.14 +github.com/tidwall/sjson # github.com/tklauser/go-sysconf v0.3.16 ## explicit; go 1.24.0 github.com/tklauser/go-sysconf @@ -739,7 +812,7 @@ golang.org/x/crypto/ssh/internal/bcrypt_pbkdf ## explicit; go 1.25.0 golang.org/x/exp/constraints golang.org/x/exp/maps -# golang.org/x/mod v0.37.0 +# golang.org/x/mod v0.38.0 ## explicit; go 1.25.0 golang.org/x/mod/internal/lazyregexp golang.org/x/mod/module