riverqueue-activerecord 0.10.1 → 0.12.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: b74ec767cc9eb7216aaad6ea31922d9c7ba2d19fd9f49700c230092001f38b6a
-  data.tar.gz: aee6b04ec4bcfa62ab555480fc2a6c476cccb165c18392f61d697f86a03d0f1e
+  metadata.gz: fe1889a91869f37c4e65eb7d68f94ca9413e1c27c35b6388bfc032cdc7042efd
+  data.tar.gz: 538826dffc04cf44bddb4a45deec552f99f3c71d3059d42453929a05383c6995
 SHA512:
-  metadata.gz: 9d06c2301b8d231bd75c97c64e960b7fccdc13a9969a4a3f2fee945de12c0f08c3c028f89e16f67600e422a9b1d64cb05e513bc80f25f34af45f6fff01866981
-  data.tar.gz: 7e5489860c578f9eabab7c7c8d81ce52446b35f6512e7bf2711121c6e1981ab79ffaf52b127998ad0135c74b796e64961acdbbe74e17f7cfc71185618651967f
+  metadata.gz: ea37012be253b41ee43e17f43837ec35809171c6668bd367faf9716c930e1321521206757636013e135dbfd3f13e677b8ec0076f245be2327ddded4bfa9c09c8
+  data.tar.gz: dcf1f5c99e8ebef5f2900f702b3714b26d1de04ad63a7b4e938c9e499ecd1ac8f51ad8a6d43b2dd942d58b34ca223dc4714ed3901aba27ada01c586b5e75e20c

data/lib/driver.rb

@@ -1,39 +1,61 @@
+# frozen_string_literal: true
+
+require "securerandom"
+
 module River::Driver
-  # Provides a ActiveRecord driver for River.
+  # Provides an ActiveRecord driver for River supporting PostgreSQL, YugabyteDB,
+  # 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
-    def initialize
-      # 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)
-        River::Driver::ActiveRecord.const_set(:RiverJob, Class.new(::ActiveRecord::Base) do
-          self.table_name = "river_job"
-
-          # Unfortunately, Rails errors if you have a column called `errors` and
-          # provides no way to remap names (beyond ignoring a column, which we
-          # really don't want). This patch is in place so we can hydrate this
-          # model at all without ActiveRecord self-immolating.
-          def self.dangerous_attribute_method?(method_name)
-            return false if method_name == "errors"
-            super
-          end
-
-          # See comment above, but since we force allowed `errors` as an
-          # attribute name, ActiveRecord would otherwise fail to save a row as
-          # it checked for its own `errors` hash and finding no values.
-          def errors = {}
-        end)
+    include River::Driver::Runtime
+
+    # Connection class whose pool and transaction context this driver uses.
+    attr_reader :connection_class
+
+    # Uses an established PostgreSQL or SQLite connection. The class must be
+    # ActiveRecord::Base or an abstract Active Record class. Routing follows its
+    # current role/shard; configure consumers explicitly for each database.
+    def initialize(connection_class: ::ActiveRecord::Base)
+      unless connection_class.is_a?(Class) && connection_class <= ::ActiveRecord::Base
+        raise ArgumentError, "connection_class must be an Active Record class"
+      end
+
+      unless connection_class == ::ActiveRecord::Base || connection_class.abstract_class?
+        raise ArgumentError, "connection_class must be abstract"
+      end
+
+      @connection_class = connection_class
+      @postgres_capabilities_cache = PostgresCapabilities::Cache.new
+      @is_sqlite = connection_class.connection.adapter_name.downcase.include?("sqlite")
+      # Do not inherit application scopes, callbacks, or schema caches. Each
+      # driver owns its model, but all operations use the selected class's pool.
+      @job_model = Class.new(::ActiveRecord::Base) do
+        self.table_name = "river_job"
+        define_singleton_method(:connection_pool) { connection_class.connection_pool }
+
+        def self.dangerous_attribute_method?(method_name)
+          return false if method_name == "errors"
+
+          super
+        end
+
+        def errors = {}
       end
     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
+        row = @job_model.find_by(id: id)
+        row ? to_job_row_from_model(row) : nil
+      end
     end
 
     def job_insert(insert_params)
@@ -41,138 +63,312 @@ module River::Driver
     end
 
     def job_insert_many(insert_params_many)
-      res = RiverJob.upsert_all(
-        insert_params_many.map { |param| insert_params_to_hash(param) },
-        on_duplicate: Arel.sql("kind = EXCLUDED.kind"),
-        returning: Arel.sql("*, (xmax != 0) AS unique_skipped_as_duplicate"),
+      return [] if insert_params_many.empty?
 
-        # It'd be nice to specify this as `(kind, unique_key) WHERE unique_key
-        # IS NOT NULL` like we do elsewhere, but in its pure ingenuity, fucking
-        # ActiveRecord tries to look up a unique index instead of letting
-        # Postgres handle that, and of course it doesn't support a `WHERE`
-        # clause. The workaround is to target the index name instead of columns.
-        unique_by: "river_job_unique_idx"
-      )
-      to_insert_results(res)
+      return sqlite_job_insert_many(insert_params_many) if @is_sqlite
+
+      transaction do
+        results = postgres_job_insert_many(insert_params_many)
+        postgres_notify_insert(insert_params_many)
+        results
+      end
     end
 
-    def job_list
-      data_set = RiverJob.order(:id)
-      data_set.all.map { |job| to_job_row_from_model(job) }
+    def job_list(params = :all)
+      return super unless params == :all
+
+      runtime_job_rows("ORDER BY id")
     end
 
     def rollback_exception
       ::ActiveRecord::Rollback
     end
 
+    # Backend used by River::Migrator.
+    def migration_backend = @is_sqlite ? :sqlite : :postgresql
+
+    # Pins a raw connection for migration SQL and refreshes cached model columns.
+    def migration_connection
+      @connection_class.connection_pool.with_connection do |connection|
+        raise River::Error, "migrations cannot run inside an application transaction" if connection.transaction_open?
+
+        begin
+          yield connection.raw_connection
+        ensure
+          @job_model.reset_column_information
+        end
+      end
+    end
+
+    # Runs the block in a new Active Record transaction or savepoint and returns
+    # the block's result.
     def transaction(&)
-      ::ActiveRecord::Base.transaction(requires_new: true, &)
+      @connection_class.transaction(requires_new: true, &)
+    end
+
+    SQLITE_CONFLICT_WHERE = <<~SQL.chomp.freeze
+      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
+
+    # 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. Flag 10 validates JSONB contents, not just its header,
+    # so damaged blobs reach the row decoder without aborting the SQL query.
+    SQLITE_JOB_COLUMNS = <<~SQL.chomp.freeze
+      id,
+      CASE WHEN json_valid(args, 10) THEN json(args) ELSE args END AS args,
+      attempt,
+      CAST(attempted_at AS text) AS attempted_at,
+      CASE WHEN json_valid(attempted_by, 10) THEN json(attempted_by) ELSE attempted_by END AS attempted_by,
+      CAST(created_at AS text) AS created_at,
+      CASE WHEN json_valid(errors, 10) THEN json(errors) ELSE errors END AS errors,
+      CAST(finalized_at AS text) AS finalized_at,
+      kind,
+      max_attempts,
+      CASE WHEN json_valid(metadata, 10) THEN json(metadata) ELSE metadata END AS metadata,
+      priority,
+      queue,
+      state,
+      CAST(scheduled_at AS text) AS scheduled_at,
+      CASE WHEN json_valid(tags, 10) THEN json(tags) ELSE tags END AS tags,
+      unique_key,
+      unique_states
+    SQL
+
+    UNIQUE_INSERT_METADATA_KEY = "river:unique_nonce"
+
+    private_constant :SQLITE_CONFLICT_WHERE, :SQLITE_JOB_COLUMNS, :UNIQUE_INSERT_METADATA_KEY
+
+    private def format_time(time)
+      time.getutc.round(3).strftime("%Y-%m-%d %H:%M:%S.%3N")
     end
 
-    private def insert_params_to_hash(insert_params)
+    private def postgres_insert_params_to_hash(insert_params, nonce)
+      metadata = insert_params.metadata || {}
+      metadata = metadata.merge(UNIQUE_INSERT_METADATA_KEY => nonce) if nonce
       {
         args: JSON.parse(insert_params.encoded_args),
         kind: insert_params.kind,
         max_attempts: insert_params.max_attempts,
+        metadata: metadata,
         priority: insert_params.priority,
         queue: insert_params.queue,
-        state: insert_params.state,
         scheduled_at: insert_params.scheduled_at,
+        state: insert_params.state,
         tags: insert_params.tags || [],
         unique_key: insert_params.unique_key,
         unique_states: insert_params.unique_states
       }
     end
 
-    private def 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: river_job.args,
-        attempt: river_job.attempt,
-        attempted_at: river_job.attempted_at&.getutc,
-        attempted_by: river_job.attempted_by,
-        created_at: river_job.created_at.getutc,
-        errors: errors&.map { |e|
-          deserialized_error = JSON.parse(e, symbolize_names: true)
-
-          River::AttemptError.new(
-            at: Time.parse(deserialized_error[:at]),
-            attempt: deserialized_error[:attempt],
-            error: deserialized_error[:error],
-            trace: deserialized_error[:trace]
-          )
-        },
-        finalized_at: river_job.finalized_at&.getutc,
-        kind: river_job.kind,
-        max_attempts: river_job.max_attempts,
-        metadata: river_job.metadata,
-        priority: river_job.priority,
-        queue: river_job.queue,
-        scheduled_at: river_job.scheduled_at.getutc,
-        state: river_job.state,
-        tags: river_job.tags,
-        unique_key: river_job.unique_key,
-        unique_states: river_job.unique_states
+    private def postgres_job_insert_many(insert_params_many)
+      capabilities = postgres_capabilities
+      nonce = SecureRandom.hex(8) if capabilities.unique_insert_mode == :metadata_nonce
+      res = @job_model.upsert_all(
+        insert_params_many.map { |param| postgres_insert_params_to_hash(param, nonce) },
+        on_duplicate: Arel.sql("kind = EXCLUDED.kind"),
+        returning: Arel.sql("*, #{capabilities.unique_insert_sql} AS unique_skipped_as_duplicate"),
+
+        # It'd be nice to specify this as `(kind, unique_key) WHERE unique_key
+        # IS NOT NULL` like we do elsewhere, but in its pure ingenuity, fucking
+        # ActiveRecord tries to look up a unique index instead of letting
+        # Postgres handle that, and of course it doesn't support a `WHERE`
+        # clause. The workaround is to target the index name instead of columns.
+        unique_by: "river_job_unique_idx"
       )
+      postgres_to_insert_results(res, nonce)
     end
 
-    private def to_insert_results(res)
+    private def postgres_to_insert_results(res, nonce)
       res.rows.map do |row|
-        to_job_row_from_raw(row, res.columns, res.column_types)
+        job, duplicate = postgres_to_job_row_from_raw(row, res.columns, res.column_types)
+        [job, nonce ? job.metadata[UNIQUE_INSERT_METADATA_KEY] != nonce : duplicate]
       end
     end
 
-    # This is really awful, but some of ActiveRecord's methods (e.g. `.create`)
-    # return a model, and others (e.g. `.upsert`) return raw values, and
-    # therefore this second version from unmarshaling a job row exists. I
-    # 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_model(river_job)
+      # Read attributes directly because Active Record shadows the errors column.
+      runtime_decode_postgres_job(river_job.attributes)
+    end
+
+    # Upserts bypass model casting; normalize their values before shared decoding.
+    private def postgres_to_job_row_from_raw(row, columns, column_types)
       river_job = {}
 
       row.each_with_index do |val, i|
         river_job[columns[i]] = column_types[i].deserialize(val)
       end
 
-      errors = river_job["errors"]&.map do |e|
-        deserialized_error = JSON.parse(e)
+      [runtime_decode_postgres_job(river_job), river_job["unique_skipped_as_duplicate"]]
+    end
+
+    private def runtime_connection_pool
+      @connection_class.connection_pool
+    end
+
+    private def runtime_execute(sql)
+      @connection_class.connection.execute(sql)
+    end
+
+    private def runtime_job_list_without_params
+      job_list(:all)
+    end
+
+    private def runtime_job_rows(suffix)
+      if @is_sqlite
+        sqlite_job_rows(suffix).map { |row| sqlite_to_job_row_from_raw(row) }
+      else
+        @job_model.find_by_sql("SELECT * FROM river_job #{suffix}").map { |row| to_job_row_from_model(row) }
+      end
+    end
+
+    private def runtime_postgres?
+      !@is_sqlite
+    end
+
+    private def runtime_query_rows(sql)
+      if @is_sqlite
+        @connection_class.connection.raw_connection.execute(sql).to_a
+      else
+        @connection_class.connection.select_all(sql).to_a
+      end
+    end
+
+    private def runtime_quote(value)
+      @connection_class.connection.quote(value)
+    end
+
+    private def runtime_unique_violation_class
+      ::ActiveRecord::RecordNotUnique
+    end
+
+    private def runtime_value(row, key)
+      row[key.to_s] || row[key]
+    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: {UNIQUE_INSERT_METADATA_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)
+      }.tap { |values| values[:metadata] = (insert_params.metadata || {}).merge(values[:metadata]) }
+    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)
+      nonces = sqlite_insert_nonces(insert_params_many)
+      @connection_class.transaction(requires_new: true) do
+        jobs = insert_params_many.zip(nonces).map { |param, nonce| sqlite_insert_params_to_hash(param, nonce) }
+        inserted_nonces = nonces.to_h { |nonce| [nonce, 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 = @connection_class.connection.raw_connection.execute(sql, [JSON.generate(jobs)])
+        sqlite_notify_insert(insert_params_many)
+
+        rows.map do |row|
+          metadata = JSON.parse(row["metadata"])
+          [sqlite_to_job_row_from_raw(row), !inserted_nonces.key?(metadata[UNIQUE_INSERT_METADATA_KEY])]
+        end
+      end
+    end
+
+    private def sqlite_job_rows(suffix, binds = [])
+      sql = "SELECT #{SQLITE_JOB_COLUMNS} FROM river_job #{suffix}"
+      @connection_class.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?
 
-        River::AttemptError.new(
-          at: Time.parse(deserialized_error["at"]),
-          attempt: deserialized_error["attempt"],
-          error: deserialized_error["error"],
-          trace: deserialized_error["trace"]
-        )
+      notifications = queues.map do |queue|
+        {payload: JSON.generate({queue: queue}), topic: "river_insert"}
       end
 
-      [
-        River::JobRow.new(
-          id: river_job["id"],
-          args: river_job["args"],
-          attempt: river_job["attempt"],
-          attempted_at: river_job["attempted_at"]&.getutc,
-          attempted_by: river_job["attempted_by"],
-          created_at: river_job["created_at"].getutc,
-          errors: errors,
-          finalized_at: river_job["finalized_at"]&.getutc,
-          kind: river_job["kind"],
-          max_attempts: river_job["max_attempts"],
-          metadata: river_job["metadata"],
-          priority: river_job["priority"],
-          queue: river_job["queue"],
-          scheduled_at: river_job["scheduled_at"].getutc,
-          state: river_job["state"],
-          tags: river_job["tags"],
-          unique_key: river_job["unique_key"],
-          unique_states: ::River::UniqueBitmask.to_states(river_job["unique_states"]&.to_i(2))
-        ),
-        river_job["unique_skipped_as_duplicate"]
-      ]
+      @connection_class.connection.raw_connection.execute(<<~SQL, [JSON.generate(notifications)])
+        INSERT INTO river_notification (payload, topic)
+        SELECT
+          json_extract(value, '$.payload'),
+          json_extract(value, '$.topic')
+        FROM json_each(cast(? AS blob))
+      SQL
+    end
+
+    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)
+      runtime_decode_sqlite_job(row)
+    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
   end
 end

data/lib/riverqueue-activerecord.rb

@@ -1,4 +1,7 @@
+# frozen_string_literal: true
+
 require "active_record"
+require "riverqueue"
 
 require_relative "driver"
 

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: riverqueue-activerecord
 version: !ruby/object:Gem::Version
-  version: 0.10.1
+  version: 0.12.0
 platform: ruby
 authors:
 - Blake Gentry
@@ -51,27 +51,21 @@ dependencies:
       - !ruby/object:Gem::Version
         version: '1000'
 - !ruby/object:Gem::Dependency
-  name: pg
+  name: riverqueue
   requirement: !ruby/object:Gem::Requirement
     requirements:
-    - - ">"
-      - !ruby/object:Gem::Version
-        version: '0'
-    - - "<"
+    - - '='
       - !ruby/object:Gem::Version
-        version: '1000'
+        version: 0.12.0
   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.
+        version: 0.12.0
+description: ActiveRecord PostgreSQL and SQLite driver for inserting and working River
+  jobs in Ruby.
 email: brandur@brandur.org
 executables: []
 extensions: []
@@ -81,7 +75,7 @@ files:
 - lib/riverqueue-activerecord.rb
 homepage: https://riverqueue.com
 licenses:
-- LGPL-3.0-or-later
+- MPL-2.0
 metadata: {}
 rdoc_options: []
 require_paths:
@@ -90,7 +84,7 @@ required_ruby_version: !ruby/object:Gem::Requirement
   requirements:
   - - ">="
     - !ruby/object:Gem::Version
-      version: '0'
+      version: '3.2'
 required_rubygems_version: !ruby/object:Gem::Requirement
   requirements:
   - - ">="
@@ -99,5 +93,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: []