riverqueue-activerecord 0.10.0 → 0.11.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 2c36ab74d951b35c50d30d615bcbf4676ef3babd6d8a7d620688430949a77b96
-  data.tar.gz: 2835b1da28dcd38c0dce0a0dc3665995ac556c24c2e215e9c565fdf2f4ffcac6
+  metadata.gz: 8316b332b5bc9bb1eae57d81751a31ef42342da5ef9a1914b3bceccb222041ac
+  data.tar.gz: 49784779614e66d985af52516698d1766980ec43af019c62f20a8cd8459d1a69
 SHA512:
-  metadata.gz: c622b956af517b97dde111c6314b71a6688120e69d3f7785d21d3a464d9fd503f461370dcc96ee8d91c990c1c9df110e86917dd92da532250f79d222b021a1f2
-  data.tar.gz: 16f7ba56b217cd8fff977718cabc4b9bf3c2d7e5eedcdff671b72cd7726029bf8f2f2cf35371f6c4a91fc6a812c2f4c7c198d467d417d0c469c0f6e732bf2839
+  metadata.gz: d3b77f7ad4752bad4401ccaad39bf66513ed84e9aa498bfe8f3f8552fd7f7c6a63a60ed4b2b8cdb3ac33ea8367698c61c35a7f790ed43d9f6e67e1729cd8ede0
+  data.tar.gz: 39992e6b530ed49a19deb622e2e689c0a221902fc23ef1ec8c7ca7bc54bbe9eb79e0bb2094aede12eafce8fe08497e248139260e2b9410e06eb43d85ae33bfbb

data/lib/driver.rb

@@ -1,13 +1,65 @@
+require "securerandom"
+
 module River::Driver
-  # Provides a ActiveRecord driver for River.
+  # Provides an ActiveRecord driver for River that supports both PostgreSQL
+  # and SQLite.
   #
   # Used in conjunction with a River client like:
   #
-  #   DB = ActiveRecord.connect("postgres://...")
-  #   client = River::Client.new(River::Driver::ActiveRecord.new(DB))
+  #   ActiveRecord::Base.establish_connection("postgres://...")
+  #   client = River::Client.new(River::Driver::ActiveRecord.new)
   #
   class ActiveRecord
+    SQLITE_CONFLICT_WHERE = <<~SQL.chomp
+      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
+    SQL
+    private_constant :SQLITE_CONFLICT_WHERE
+
+    # SQLite 3.45+ may store JSON as binary JSONB. Always project JSON columns
+    # through json() so this driver can read both the current JSONB format and
+    # the text JSON used by River migrations through version 006. Cast times to
+    # text so ActiveRecord doesn't interpret timezone-less SQLite timestamps in
+    # the process timezone.
+    SQLITE_JOB_COLUMNS = <<~SQL.chomp
+      id,
+      json(args) AS args,
+      attempt,
+      CAST(attempted_at AS text) AS attempted_at,
+      json(attempted_by) AS attempted_by,
+      CAST(created_at AS text) AS created_at,
+      json(errors) AS errors,
+      CAST(finalized_at AS text) AS finalized_at,
+      kind,
+      max_attempts,
+      json(metadata) AS metadata,
+      priority,
+      queue,
+      state,
+      CAST(scheduled_at AS text) AS scheduled_at,
+      json(tags) AS tags,
+      unique_key,
+      unique_states
+    SQL
+    private_constant :SQLITE_JOB_COLUMNS
+
+    SQLITE_UNIQUE_NONCE_KEY = "river:unique_nonce"
+    private_constant :SQLITE_UNIQUE_NONCE_KEY
+
     def initialize
+      @is_sqlite = ::ActiveRecord::Base.connection.adapter_name.downcase.include?("sqlite")
+
       # It's Ruby, so we can only define a model after ActiveRecord's established a
       # connection because it's all dynamic.
       if !River::Driver::ActiveRecord.const_defined?(:RiverJob)
@@ -32,8 +84,13 @@ module River::Driver
     end
 
     def job_get_by_id(id)
-      data_set = RiverJob.where(id: id)
-      data_set.first ? to_job_row_from_model(data_set.first) : nil
+      if @is_sqlite
+        row = sqlite_job_rows("WHERE id = ? LIMIT 1", [id]).first
+        row ? sqlite_to_job_row_from_raw(row) : nil
+      else
+        data_set = RiverJob.where(id: id)
+        data_set.first ? to_job_row_from_model(data_set.first) : nil
+      end
     end
 
     def job_insert(insert_params)
@@ -41,8 +98,28 @@ module River::Driver
     end
 
     def job_insert_many(insert_params_many)
+      @is_sqlite ? sqlite_job_insert_many(insert_params_many) : postgres_job_insert_many(insert_params_many)
+    end
+
+    def job_list
+      if @is_sqlite
+        sqlite_job_rows("ORDER BY id").map { |row| sqlite_to_job_row_from_raw(row) }
+      else
+        RiverJob.order(:id).all.map { |job| to_job_row_from_model(job) }
+      end
+    end
+
+    def rollback_exception
+      ::ActiveRecord::Rollback
+    end
+
+    def transaction(&)
+      ::ActiveRecord::Base.transaction(requires_new: true, &)
+    end
+
+    private def postgres_job_insert_many(insert_params_many)
       res = RiverJob.upsert_all(
-        insert_params_many.map { |param| insert_params_to_hash(param) },
+        insert_params_many.map { |param| postgres_insert_params_to_hash(param) },
         on_duplicate: Arel.sql("kind = EXCLUDED.kind"),
         returning: Arel.sql("*, (xmax != 0) AS unique_skipped_as_duplicate"),
 
@@ -53,25 +130,70 @@ module River::Driver
         # clause. The workaround is to target the index name instead of columns.
         unique_by: "river_job_unique_idx"
       )
-      to_insert_results(res)
+      postgres_to_insert_results(res)
     end
 
-    def job_list
-      data_set = RiverJob.order(:id)
-      data_set.all.map { |job| to_job_row_from_model(job) }
-    end
+    # River's current SQLite driver uses json_each to make a batch a single,
+    # atomic statement. The JSON columns are converted to SQLite JSONB here,
+    # matching migration 007 and newer River databases.
+    private def sqlite_job_insert_many(insert_params_many)
+      return [] if insert_params_many.empty?
 
-    def rollback_exception
-      ::ActiveRecord::Rollback
-    end
+      ::ActiveRecord::Base.transaction(requires_new: true) do
+        nonce = SecureRandom.hex(8)
+        jobs = insert_params_many.map { |param| sqlite_insert_params_to_hash(param, nonce) }
 
-    def transaction(&)
-      ::ActiveRecord::Base.transaction(requires_new: true, &)
+        sql = <<~SQL
+          INSERT INTO river_job (
+            args,
+            created_at,
+            kind,
+            max_attempts,
+            metadata,
+            priority,
+            queue,
+            scheduled_at,
+            state,
+            tags,
+            unique_key,
+            unique_states
+          )
+          SELECT
+            jsonb(json_extract(value, '$.args')),
+            datetime('now', 'subsec'),
+            cast(json_extract(value, '$.kind') AS text),
+            cast(json_extract(value, '$.max_attempts') AS integer),
+            jsonb(json_extract(value, '$.metadata')),
+            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),
+            jsonb(json_extract(value, '$.tags')),
+            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(? AS blob))
+          WHERE true
+          ON CONFLICT (unique_key) WHERE #{SQLITE_CONFLICT_WHERE}
+          DO UPDATE SET kind = EXCLUDED.kind
+          RETURNING #{SQLITE_JOB_COLUMNS}
+        SQL
+
+        rows = ::ActiveRecord::Base.connection.raw_connection.execute(sql, [JSON.dump(jobs)])
+        sqlite_notify_insert(insert_params_many)
+
+        rows.map do |row|
+          metadata = JSON.parse(row["metadata"])
+          [sqlite_to_job_row_from_raw(row), metadata[SQLITE_UNIQUE_NONCE_KEY] != nonce]
+        end
+      end
     end
 
-    private def insert_params_to_hash(insert_params)
+    private def postgres_insert_params_to_hash(insert_params)
       {
-        args: insert_params.encoded_args,
+        args: JSON.parse(insert_params.encoded_args),
         kind: insert_params.kind,
         max_attempts: insert_params.max_attempts,
         priority: insert_params.priority,
@@ -84,14 +206,34 @@ module River::Driver
       }
     end
 
+    private def sqlite_insert_params_to_hash(insert_params, nonce)
+      {
+        args: JSON.parse(insert_params.encoded_args),
+        kind: insert_params.kind,
+        max_attempts: insert_params.max_attempts,
+        metadata: {SQLITE_UNIQUE_NONCE_KEY => nonce},
+        priority: insert_params.priority,
+        queue: insert_params.queue,
+        scheduled_at: insert_params.scheduled_at ? format_time(insert_params.scheduled_at) : nil,
+        state: insert_params.state,
+        tags: insert_params.tags || [],
+        unique_key: insert_params.unique_key&.unpack1("H*"),
+        unique_states: insert_params.unique_states&.to_i(2)
+      }
+    end
+
     private def to_job_row_from_model(river_job)
+      @is_sqlite ? sqlite_to_job_row_from_model(river_job) : postgres_to_job_row_from_model(river_job)
+    end
+
+    private def postgres_to_job_row_from_model(river_job)
       # needs to be accessed through values because `errors` is shadowed by both
       # ActiveRecord and the patch above
       errors = river_job.attributes["errors"]
 
       River::JobRow.new(
         id: river_job.id,
-        args: JSON.parse(river_job.args),
+        args: river_job.args,
         attempt: river_job.attempt,
         attempted_at: river_job.attempted_at&.getutc,
         attempted_by: river_job.attempted_by,
@@ -120,9 +262,46 @@ module River::Driver
       )
     end
 
-    private def to_insert_results(res)
+    private def sqlite_to_job_row_from_model(river_job)
+      row = sqlite_job_rows("WHERE id = ? LIMIT 1", [river_job.id]).first
+      sqlite_to_job_row_from_raw(row)
+    end
+
+    private def sqlite_to_job_row_from_raw(row)
+      errors = row["errors"] ? JSON.parse(row["errors"]) : []
+
+      River::JobRow.new(
+        id: row["id"],
+        args: JSON.parse(row["args"]),
+        attempt: row["attempt"],
+        attempted_at: parse_sqlite_time(row["attempted_at"]),
+        attempted_by: row["attempted_by"] ? JSON.parse(row["attempted_by"]) : nil,
+        created_at: parse_sqlite_time(row["created_at"]),
+        errors: errors.map { |e|
+          River::AttemptError.new(
+            at: Time.parse(e["at"]),
+            attempt: e["attempt"],
+            error: e["error"],
+            trace: e["trace"]
+          )
+        },
+        finalized_at: parse_sqlite_time(row["finalized_at"]),
+        kind: row["kind"],
+        max_attempts: row["max_attempts"],
+        metadata: JSON.parse(row["metadata"]),
+        priority: row["priority"],
+        queue: row["queue"],
+        scheduled_at: parse_sqlite_time(row["scheduled_at"]),
+        state: row["state"],
+        tags: JSON.parse(row["tags"]),
+        unique_key: row["unique_key"]&.to_s,
+        unique_states: row["unique_states"] ? ::River::UniqueBitmask.to_states(row["unique_states"]) : nil
+      )
+    end
+
+    private def postgres_to_insert_results(res)
       res.rows.map do |row|
-        to_job_row_from_raw(row, res.columns, res.column_types)
+        postgres_to_job_row_from_raw(row, res.columns, res.column_types)
       end
     end
 
@@ -132,7 +311,7 @@ module River::Driver
     # searched long and hard for a way to have the former type of method return
     # raw or the latter type of method return a model, but was unable to find
     # anything.
-    private def to_job_row_from_raw(row, columns, column_types)
+    private def postgres_to_job_row_from_raw(row, columns, column_types)
       river_job = {}
 
       row.each_with_index do |val, i|
@@ -153,7 +332,7 @@ module River::Driver
       [
         River::JobRow.new(
           id: river_job["id"],
-          args: JSON.parse(river_job["args"]),
+          args: river_job["args"],
           attempt: river_job["attempt"],
           attempted_at: river_job["attempted_at"]&.getutc,
           attempted_by: river_job["attempted_by"],
@@ -174,5 +353,42 @@ module River::Driver
         river_job["unique_skipped_as_duplicate"]
       ]
     end
+
+    private def format_time(time)
+      time.getutc.round(3).strftime("%Y-%m-%d %H:%M:%S.%3N")
+    end
+
+    private def parse_sqlite_time(value)
+      return nil unless value
+
+      value = value.to_s
+      value += " UTC" unless value.match?(/(?:Z|[+-]\d{2}:?\d{2})\z/)
+      Time.parse(value).utc
+    end
+
+    private def sqlite_job_rows(suffix, binds = [])
+      sql = "SELECT #{SQLITE_JOB_COLUMNS} FROM river_job #{suffix}"
+      ::ActiveRecord::Base.connection.raw_connection.execute(sql, binds)
+    end
+
+    private def sqlite_notify_insert(insert_params_many)
+      queues = insert_params_many
+        .select { |param| param.state == ::River::JOB_STATE_AVAILABLE }
+        .map(&:queue)
+        .uniq
+      return if queues.empty?
+
+      notifications = queues.map do |queue|
+        {payload: JSON.dump({queue: queue}), topic: "insert"}
+      end
+
+      ::ActiveRecord::Base.connection.raw_connection.execute(<<~SQL, [JSON.dump(notifications)])
+        INSERT INTO river_notification (payload, topic)
+        SELECT
+          json_extract(value, '$.payload'),
+          json_extract(value, '$.topic')
+        FROM json_each(cast(? AS blob))
+      SQL
+    end
   end
 end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: riverqueue-activerecord
 version: !ruby/object:Gem::Version
-  version: 0.10.0
+  version: 0.11.0
 platform: ruby
 authors:
 - Blake Gentry
@@ -50,28 +50,8 @@ dependencies:
     - - "<"
       - !ruby/object:Gem::Version
         version: '1000'
-- !ruby/object:Gem::Dependency
-  name: pg
-  requirement: !ruby/object:Gem::Requirement
-    requirements:
-    - - ">"
-      - !ruby/object:Gem::Version
-        version: '0'
-    - - "<"
-      - !ruby/object:Gem::Version
-        version: '1000'
-  type: :runtime
-  prerelease: false
-  version_requirements: !ruby/object:Gem::Requirement
-    requirements:
-    - - ">"
-      - !ruby/object:Gem::Version
-        version: '0'
-    - - "<"
-      - !ruby/object:Gem::Version
-        version: '1000'
-description: ActiveRecord driver for the River Ruby gem. Use in conjunction with the
-  riverqueue gem to insert jobs that are worked in Go.
+description: ActiveRecord PostgreSQL and SQLite driver for the River Ruby gem. Use
+  in conjunction with the riverqueue gem to insert jobs that are worked in Go.
 email: brandur@brandur.org
 executables: []
 extensions: []
@@ -99,5 +79,5 @@ required_rubygems_version: !ruby/object:Gem::Requirement
 requirements: []
 rubygems_version: 4.0.9
 specification_version: 4
-summary: ActiveRecord driver for the River Ruby gem.
+summary: ActiveRecord PostgreSQL and SQLite driver for the River Ruby gem.
 test_files: []