Use json_each to make SQLite single inserts batch inserts (#1276)

Back when I was first implementing SQLite, the sqlc support was so
catastrophically bad that many, many basic utilities just could not be
used in queries. I tried like a half dozen different ways, but couldn't
find anything that was both supported by SQLite and sqlc, so I ended up
making a number of batch inserts into single inserts that get looped
over in the driver. This isn't as bad in SQLite because usually your
inserting to a local database (i.e. no network connections), but it's
still not ideal in that it's probably somewhat slower, and deviates from
the Postgres implementation.

Sqlc's made some SQLite improvements since then, and I found that I'm
now able to get a batch insert working `json_each` and `json_extract`.

Here, go through the existing SQLite driver function and update things
to use batch inserts wherever possible. Fixed operations:

* `JobInsertFastMany`
* `JobInsertFastManyNoReturning`
* `JobInsertFullMany`
* `MigrationInsertMany`
* `MigrationInsertManyAssumingMain`

Author: brandur

CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Changed
+
+- 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)
+
 ### 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).

riverdriver/riversqlite/internal/dbsqlc/river_job.sql

@@ -189,15 +189,6 @@ WHERE state = 'running'
 ORDER BY id
 LIMIT @max;
 
--- Insert a job.
---
--- This is supposed to be a batch insert, but various limitations of the
--- combined SQLite + sqlc has left me unable to find a way of injecting many
--- arguments en masse (like how we slightly abuse arrays to pull it off for the
--- Postgres drivers), so we loop over many insert operations instead, with the
--- expectation that this may be fixable in the future. Because SQLite targets
--- will often be local and therefore with a very minimal round trip compared to
--- a network, looping over operations is probably okay performance-wise.
 -- name: JobInsertFast :one
 INSERT INTO /* TEMPLATE: schema */river_job(
     id,
@@ -246,6 +237,56 @@ ON CONFLICT (unique_key)
     DO UPDATE SET kind = EXCLUDED.kind
 RETURNING *;
 
+-- name: JobInsertFastMany :many
+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
+    cast(json_extract(value, '$.id') AS integer),
+    json(cast(json_extract(value, '$.args') AS blob)),
+    coalesce(cast(json_extract(value, '$.created_at') AS text), datetime('now', 'subsec')),
+    cast(json_extract(value, '$.kind') AS text),
+    cast(json_extract(value, '$.max_attempts') AS integer),
+    json(cast(json_extract(value, '$.metadata') AS blob)),
+    cast(json_extract(value, '$.priority') AS integer),
+    cast(json_extract(value, '$.queue') AS text),
+    coalesce(cast(json_extract(value, '$.scheduled_at') AS text), datetime('now', 'subsec')),
+    cast(json_extract(value, '$.state') AS text),
+    json(cast(json_extract(value, '$.tags') AS blob)),
+    CASE WHEN length(cast(json_extract(value, '$.unique_key') AS text)) = 0 THEN NULL ELSE unhex(cast(json_extract(value, '$.unique_key') AS text)) END,
+    nullif(cast(json_extract(value, '$.unique_states') AS integer), 0)
+FROM json_each(cast(@jobs AS blob))
+WHERE true
+ON CONFLICT (unique_key)
+    WHERE unique_key IS NOT NULL
+        AND unique_states IS NOT NULL
+        AND CASE state
+                WHEN 'available' THEN unique_states & (1 << 0)
+                WHEN 'cancelled' THEN unique_states & (1 << 1)
+                WHEN 'completed' THEN unique_states & (1 << 2)
+                WHEN 'discarded' THEN unique_states & (1 << 3)
+                WHEN 'pending'   THEN unique_states & (1 << 4)
+                WHEN 'retryable' THEN unique_states & (1 << 5)
+                WHEN 'running'   THEN unique_states & (1 << 6)
+                WHEN 'scheduled' THEN unique_states & (1 << 7)
+                ELSE 0
+            END >= 1
+    -- Something needs to be updated for a row to be returned on a conflict.
+    DO UPDATE SET kind = EXCLUDED.kind
+RETURNING *;
+
 -- name: JobInsertFastNoReturning :execrows
 INSERT INTO /* TEMPLATE: schema */river_job(
     args,
@@ -290,6 +331,52 @@ ON CONFLICT (unique_key)
             END >= 1
 DO NOTHING;
 
+-- 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
+    json(cast(json_extract(value, '$.args') AS blob)),
+    coalesce(cast(json_extract(value, '$.created_at') AS text), datetime('now', 'subsec')),
+    cast(json_extract(value, '$.kind') AS text),
+    cast(json_extract(value, '$.max_attempts') AS integer),
+    json(cast(json_extract(value, '$.metadata') AS blob)),
+    cast(json_extract(value, '$.priority') AS integer),
+    cast(json_extract(value, '$.queue') AS text),
+    coalesce(cast(json_extract(value, '$.scheduled_at') AS text), datetime('now', 'subsec')),
+    cast(json_extract(value, '$.state') AS text),
+    json(cast(json_extract(value, '$.tags') AS blob)),
+    CASE WHEN length(cast(json_extract(value, '$.unique_key') AS text)) = 0 THEN NULL ELSE unhex(cast(json_extract(value, '$.unique_key') AS text)) END,
+    nullif(cast(json_extract(value, '$.unique_states') AS integer), 0)
+FROM json_each(cast(@jobs AS blob))
+WHERE true
+ON CONFLICT (unique_key)
+    WHERE unique_key IS NOT NULL
+        AND unique_states IS NOT NULL
+        AND CASE state
+                WHEN 'available' THEN unique_states & (1 << 0)
+                WHEN 'cancelled' THEN unique_states & (1 << 1)
+                WHEN 'completed' THEN unique_states & (1 << 2)
+                WHEN 'discarded' THEN unique_states & (1 << 3)
+                WHEN 'pending'   THEN unique_states & (1 << 4)
+                WHEN 'retryable' THEN unique_states & (1 << 5)
+                WHEN 'running'   THEN unique_states & (1 << 6)
+                WHEN 'scheduled' THEN unique_states & (1 << 7)
+                ELSE 0
+            END >= 1
+DO NOTHING;
+
 -- name: JobInsertFull :one
 INSERT INTO /* TEMPLATE: schema */river_job(
     args,
@@ -329,6 +416,47 @@ INSERT INTO /* TEMPLATE: schema */river_job(
     @unique_states
 ) RETURNING *;
 
+-- name: JobInsertFullMany :many
+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
+)
+SELECT
+    json(cast(json_extract(value, '$.args') AS blob)),
+    cast(json_extract(value, '$.attempt') AS integer),
+    cast(json_extract(value, '$.attempted_at') AS text),
+    CASE WHEN json_type(value, '$.attempted_by') IS NULL THEN NULL ELSE json(cast(json_extract(value, '$.attempted_by') AS blob)) END,
+    coalesce(cast(json_extract(value, '$.created_at') AS text), datetime('now', 'subsec')),
+    CASE WHEN json_type(value, '$.errors') IS NULL THEN NULL ELSE json(cast(json_extract(value, '$.errors') AS blob)) END,
+    cast(json_extract(value, '$.finalized_at') AS text),
+    cast(json_extract(value, '$.kind') AS text),
+    cast(json_extract(value, '$.max_attempts') AS integer),
+    json(cast(json_extract(value, '$.metadata') AS blob)),
+    cast(json_extract(value, '$.priority') AS integer),
+    cast(json_extract(value, '$.queue') AS text),
+    coalesce(cast(json_extract(value, '$.scheduled_at') AS text), datetime('now', 'subsec')),
+    cast(json_extract(value, '$.state') AS text),
+    json(cast(json_extract(value, '$.tags') AS blob)),
+    CASE WHEN length(cast(json_extract(value, '$.unique_key') AS text)) = 0 THEN NULL ELSE unhex(cast(json_extract(value, '$.unique_key') AS text)) END,
+    nullif(cast(json_extract(value, '$.unique_states') AS integer), 0)
+FROM json_each(cast(@jobs AS blob))
+RETURNING *;
+
 -- name: JobKindList :many
 SELECT DISTINCT kind
 FROM /* TEMPLATE: schema */river_job
@@ -513,4 +641,4 @@ SET
     metadata = CASE WHEN cast(@metadata_do_update AS boolean) THEN json(cast(@metadata AS blob)) ELSE metadata END,
     state = CASE WHEN cast(@state_do_update AS boolean) THEN @state ELSE state END
 WHERE id = @id
-RETURNING *;
+RETURNING *;