wal 0.0.0 → 0.0.2

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 8bf05ae2e837a2307b76eb41a17343f41c3f40340b935d681b955fad9ee48db2
-  data.tar.gz: 5111d2c221b9f4151da3e5fe73d66cff1c87ead67df9e537f00544fbba118e70
+  metadata.gz: cc5df8a8a67077b8c9a4da905a102e72bf4b1e6e41ad2569a4662038d81d123a
+  data.tar.gz: dc5dd65b5b87f43b0b1b4bc7d7d66b2b97fa92fd04963bbbdb32f2e781f4bcaf
 SHA512:
-  metadata.gz: bdb71b6b2dc3e1773dffe0dec0c1d4d61aabab6faf76cade6277b8b65354e69795ed8a6a32ca8968e1b348641f6b53c5ca38faea621249adfc6d7cf8760870b8
-  data.tar.gz: 983cd7e80bec90f766341bae9b8fdd1c1d3f641d56ad687eec541f8d52a27e2de023a17d447c2fe75c171bc6d8d12399c619d49a83a9cb508821ebf48af8ffec
+  metadata.gz: 0e1b415b21cef101a2493ed922b95e8e70e972575bb6bfc190200aee567a51a4582a26b29f07eca316621ae48ec4b63e995e718cec3a963ffd42fb6774c40ee1
+  data.tar.gz: 940d945bafe3f36dd303a427bbe25524f50a98ac88a76f2ef02e842cfe30b190ba87c047a15256c5b7f80408ee211bda08beb7b6bb37bb93325540fd5ad6e4e7

data/.rspec

@@ -0,0 +1 @@
+--require spec_helper

data/README.md

@@ -0,0 +1,49 @@
+# Wal
+
+Easily hook into Postgres WAL event log from your Rails app.
+
+Proper documentation TBD
+
+## Examples
+
+### Watch for model changes using the RecordWatcher DSL
+
+```ruby
+class ProductAvailabilityWatcher < Wal::RecordWatcher
+  on_save Product, changed: %w[price] do |event|
+    recalculate_inventory_price(event.primary_key, event.new["price"])
+  end
+
+  on_destroy Product do |event|
+    clear_product_inventory(event.primary_key)
+  end
+
+  on_save Sales, changed: %w[status] do |event|
+    recalculate_inventory_quantity(event.primary_key)
+  end
+
+  def recalculate_inventory_price(product_id, new_price)
+    # ...
+  end
+
+  def clear_product_inventory(product_id)
+    # ...
+  end
+
+  def recalculate_inventory_quantity(sales_id)
+    # ...
+  end
+end
+```
+
+### Basic watcher implementation
+
+```ruby
+class LogWatcher
+  include Wal::Watcher
+
+  def on_event(event)
+    puts "Wal event received #{event}"
+  end
+end
+```

data/Rakefile

@@ -1,4 +1,12 @@
 # frozen_string_literal: true
 
 require "bundler/gem_tasks"
-task default: %i[]
+
+task(:test) { sh "bundle exec rspec" }
+
+task default: %i[build]
+
+task("sig/wal.rbi") { sh "bundle exec parlour" }
+task("rbi/wal.rbs") { sh "rbs prototype rbi rbi/wal.rbi > sig/wal.rbs" }
+
+Rake::Task["build"].enhance(["sig/wal.rbi", "rbi/wal.rbs"])

data/exe/wal

@@ -0,0 +1,64 @@
+#!/usr/bin/env ruby
+
+require "docopt"
+
+begin
+  cli = Docopt.docopt(<<~DOCOPT)
+    Usage:
+      wal watch --watcher <watcher-class> (--slot <replication-slot> | --tmp-slot) [--publication=<publication>...] [--with-timings]
+      wal start <config-file>
+
+    Options:
+      -h --help                         Show this screen.
+      --watcher=<watcher-class>         The watcher class to be used to listen for WAL changes.
+      --slot=<replication-slot>         The replication slot that will be used.
+      --tmp-slot                        Use a temporary replication slot.
+      [--publication=<publication>...]  Force using the informed Postgres publications.
+      [--with-timings]                  Add timing logs to the output.
+  DOCOPT
+rescue Docopt::Exit => err
+  puts err.message
+  exit
+end
+
+require "./config/environment"
+
+if cli["watch"]
+  watcher = cli["--watcher"].constantize.new
+  watcher = Wal::MonitoringWatcher.new(watcher) if cli["--with-timings"]
+
+  use_temporary_slot = cli["--tmp-slot"] || false
+  replication_slot = cli["--slot"]
+  replication_slot = replication_slot.presence || "wal_watcher_#{SecureRandom.alphanumeric(4)}" if use_temporary_slot
+
+  Wal::Replicator
+    .new(replication_slot:, use_temporary_slot:)
+    .replicate_forever(watcher, publications: cli["--publication"])
+
+elsif cli["start"]
+  workers = YAML.load_file(cli["<config-file>"])["slots"].map do |slot, config|
+    watcher = config["watcher"].constantize.new
+    watcher = Wal::MonitoringWatcher.new(watcher) if config["log_execution_time"]
+    temporary = config["temporary"] || false
+    publications = config["publications"] || []
+
+    Thread.new(slot, watcher, temporary, publications) do |replication_slot, watcher, use_temporary_slot, publications|
+      replication_slot = "#{replication_slot}_#{SecureRandom.alphanumeric(4)}" if use_temporary_slot
+      puts "Watcher started for #{replication_slot} slot (#{publications.join(", ")})"
+
+      Wal::Replicator
+        .new(replication_slot:, use_temporary_slot:)
+        .replicate_forever(watcher, publications:)
+
+      puts "Watcher finished for #{replication_slot}"
+    end
+  end
+
+  Signal.trap("INT") do
+    puts "Stopping WAL workers..."
+    workers.each(&:kill)
+    puts "WAL workers stopped"
+  end
+
+  workers.each(&:join)
+end

data/lib/wal/active_record_context_extension.rb

@@ -0,0 +1,15 @@
+require "active_record/connection_adapters/postgresql_adapter"
+
+module Wal
+  module ActiveRecordContextExtension
+    def set_wal_watcher_context(context, prefix: "")
+      execute "SELECT pg_logical_emit_message(true, #{quote(prefix)}, #{quote(context.to_json)})" if transaction_open?
+    end
+  end
+end
+
+if defined? ActiveRecord::ConnectionAdapters::PostgreSQLAdapter
+  class ActiveRecord::ConnectionAdapters::PostgreSQLAdapter
+    include Wal::ActiveRecordContextExtension
+  end
+end

data/lib/wal/noop_watcher.rb

@@ -0,0 +1,12 @@
+# typed: strict
+
+module Wal
+  # A watcher that does nothing. Just for performance testing in general. Useful in testing aswell.
+  class NoopWatcher
+    extend T::Sig
+    include Wal::Watcher
+
+    sig { override.params(event: Event).void }
+    def on_event(event); end
+  end
+end

data/lib/wal/record_watcher.rb

@@ -0,0 +1,389 @@
+# typed: true
+
+module Wal
+  # Watcher that process records at the end of a transaction, keeping only its final state.
+  #
+  # Example:
+  #
+  # ```ruby
+  # class InventoryAvailabilityWatcher < Wal::RecordWatcher
+  #   on_save Item, changed: %w[weight_unid_id] do |event|
+  #     recalculate_inventory_availability(event.primary_key)
+  #   end
+  #
+  #   on_save SalesOrder, changed: %w[status] do |event|
+  #     next unless event.attributes_changes(:status).one? "filled"
+  #
+  #     OrderItem
+  #       .where(sales_order_id: event.primary_key)
+  #       .pluck(:item_id)
+  #       .each { |item_id| recalculate_inventory_availability(item_id) }
+  #   end
+  #
+  #   on_save OrderItem, changed: %w[item_id weight_unit weight_unid_id] do |event|
+  #     if (old_item_id, new_item_id = event.attributes_changes("item_id"))
+  #       recalculate_inventory_availability(old_item_id)
+  #       recalculate_inventory_availability(new_item_id)
+  #     else
+  #       recalculate_inventory_availability(event.attribute(:item_id))
+  #     end
+  #   end
+  #
+  #   on_destroy OrderItem do |event|
+  #     recalculate_inventory_availability(event.attribute(:item_id))
+  #   end
+  #
+  #   def recalculate_inventory_availability(item_id)
+  #     ...
+  #   end
+  # end
+  # ```
+  class RecordWatcher
+    extend T::Sig
+    extend T::Helpers
+    include Wal::Watcher
+    abstract!
+
+    RecordEvent = T.type_alias { T.any(InsertEvent, UpdateEvent, DeleteEvent) }
+
+    def self.inherited(subclass)
+      super
+      @@change_callbacks = Hash.new { |hash, key| hash[key] = [] }
+      @@delete_callbacks = Hash.new { |hash, key| hash[key] = [] }
+    end
+
+    sig do
+      params(
+        table: T.any(String, T.class_of(::ActiveRecord::Base)),
+        block: T.proc.bind(T.attached_class).params(event: InsertEvent).void,
+      ).void
+    end
+    def self.on_insert(table, &block)
+      table = table.is_a?(String) ? table : table.table_name
+      @@change_callbacks[table].push(only: [:create], block: block)
+    end
+
+    sig do
+      params(
+        table: T.any(String, T.class_of(::ActiveRecord::Base)),
+        changed: T.nilable(T::Array[T.any(String, Symbol)]),
+        block: T.proc.bind(T.attached_class).params(event: UpdateEvent).void,
+      ).void
+    end
+    def self.on_update(table, changed: nil, &block)
+      table = table.is_a?(String) ? table : table.table_name
+      @@change_callbacks[table].push(only: [:update], changed: changed&.map(&:to_s), block: block)
+    end
+
+    sig do
+      params(
+        table: T.any(String, T.class_of(::ActiveRecord::Base)),
+        changed: T.nilable(T::Array[T.any(String, Symbol)]),
+        block: T.proc.bind(T.attached_class).params(event: T.any(InsertEvent, UpdateEvent)).void,
+      ).void
+    end
+    def self.on_save(table, changed: nil, &block)
+      table = table.is_a?(String) ? table : table.table_name
+      @@change_callbacks[table].push(only: [:create, :update], changed: changed&.map(&:to_s), block: block)
+    end
+
+    sig do
+      params(
+        table: T.any(String, T.class_of(::ActiveRecord::Base)),
+        block: T.proc.bind(T.attached_class).params(event: DeleteEvent).void,
+      ).void
+    end
+    def self.on_destroy(table, &block)
+      table = table.is_a?(String) ? table : table.table_name
+      @@delete_callbacks[table].push(block: block)
+    end
+
+    sig { params(event: RecordEvent).void }
+    def on_record_changed(event)
+      case event
+      when InsertEvent
+        @@change_callbacks[event.table]
+          .filter { |callback| callback[:only].include? :create }
+          .each { |callback| instance_exec(event, &callback[:block]) }
+
+      when UpdateEvent
+        @@change_callbacks[event.table]
+          .filter { |callback| callback[:only].include? :update }
+          .each do |callback|
+            if (attributes = callback[:changed])
+              instance_exec(event, &callback[:block]) unless (event.diff.keys & attributes).empty?
+            else
+              instance_exec(event, &callback[:block])
+            end
+          end
+
+      when DeleteEvent
+        @@delete_callbacks[event.table].each do |callback|
+          instance_exec(event, &callback[:block])
+        end
+      end
+    end
+
+    sig { params(table: String).returns(T::Boolean) }
+    def should_watch_table?(table)
+      (@@change_callbacks.keys | @@delete_callbacks.keys).include? table
+    end
+
+    # `RecordWatcher` supports two processing strategies:
+    #
+    # `:memory`: Stores and aggregates records from a single transaction in memory. This has better performance but uses
+    #  more memory, as at least one event for each record must be stored in memory until the end of a transaction
+    #
+    # `:temporary_table`: Offloads the record aggregation to a temporary table on the database. This is useful when you
+    # are processing very large transactions that can't fit in memory. The tradeoff is obviously a worse performance.
+    #
+    # These strategies can be defined per transaction, and by default it will uses the memory one, and only fallback
+    # to the temporary table if the transaction size is roughly 2 gigabytes or more.
+    sig { params(event: BeginTransactionEvent).returns(Symbol) }
+    def aggregation_strategy(event)
+      if event.estimated_size > 1024.pow(3) * 2
+        :temporary_table
+      else
+        :memory
+      end
+    end
+
+    sig { override.params(event: Event).void }
+    def on_event(event)
+      if event.is_a? BeginTransactionEvent
+        @current_record_watcher = case (strategy = aggregation_strategy(event))
+        when :memory
+          MemoryRecordWatcher.new(self)
+        when :temporary_table
+          TemporaryTableRecordWatcher.new(self)
+        else
+          raise "Invalid aggregation strategy: #{strategy}"
+        end
+      end
+      @current_record_watcher.on_event(event)
+    end
+
+    class MemoryRecordWatcher
+      extend T::Sig
+      extend T::Helpers
+      include Wal::Watcher
+      include Wal::Watcher::SeparatedEvents
+
+      # Records indexed by table and primary key
+      RecordsStorage = T.type_alias { T::Hash[[String, Integer], T.nilable(RecordEvent)] }
+
+      def initialize(watcher)
+        @watcher = watcher
+      end
+
+      sig { params(event: BeginTransactionEvent).void }
+      def on_begin(event)
+        @records = T.let({}, T.nilable(RecordsStorage))
+      end
+
+      def on_commit(_event)
+        @records
+          &.values
+          &.lazy
+          &.each { |event| @watcher.on_record_changed(event) if event }
+      end
+
+      sig { params(event: InsertEvent).void }
+      def on_insert(event)
+        if (id = event.primary_key)
+          @records ||= {}
+          @records[[event.table, id]] = event
+        end
+      end
+
+      sig { params(event: UpdateEvent).void }
+      def on_update(event)
+        if (id = event.primary_key)
+          @records ||= {}
+          @records[[event.table, id]] = case (existing_event = @records[[event.table, id]])
+          when InsertEvent
+            # A record inserted on this transaction is being updated, which means it should still reflect as a insert
+            # event, we just change the information to reflect the most current data that was just updated.
+            existing_event.with(new: event.new)
+
+          when UpdateEvent
+            # We are updating again a event that was already updated on this transaction.
+            # Same as the insert, we keep the old data from the previous update and the new data from the new one.
+            existing_event.with(new: event.new)
+
+          else
+            event
+          end
+        end
+      end
+
+      sig { params(event: DeleteEvent).void }
+      def on_delete(event)
+        if (id = event.primary_key)
+          @records ||= {}
+          @records[[event.table, id]] = case (existing_event = @records[[event.table, id]])
+          when InsertEvent
+            # We are removing a record that was inserted on this transaction, we should not even report this change, as
+            # this record never existed outside this transaction anyways.
+            nil
+
+          when UpdateEvent
+            # Deleting a record that was previously updated by this transaction. Just store the previous data while
+            # keeping the record as deleted.
+            event.with(old: existing_event.old)
+
+          else
+            event
+          end
+        end
+      end
+    end
+
+    class TemporaryTableRecordWatcher
+      extend T::Sig
+      extend T::Helpers
+      include Wal::Watcher
+      include Wal::Watcher::SeparatedEvents
+
+      # ActiveRecord base class used to persist the temporary table. Defaults to `ActiveRecord::Base`, but can be
+      # changed if you want, for example, to use a different database for Wal processing.
+      # Note that the class specified here must be a `abstract_class`.
+      mattr_accessor :base_active_record_class
+
+      def initialize(watcher, batch_size: 5_000)
+        @watcher = watcher
+        @batch_size = 5_000
+      end
+
+      sig { params(event: BeginTransactionEvent).void }
+      def on_begin(event)
+        @table = begin
+          table_name = "temp_record_watcher_#{SecureRandom.alphanumeric(10).downcase}"
+
+          base_class.connection.create_table(table_name, temporary: true) do |t|
+            t.bigint :transaction_id, null: false
+            t.bigint :lsn, null: false
+            t.column :action, :string, null: false
+            t.string :table_name, null: false
+            t.bigint :primary_key
+            t.jsonb :old, null: false, default: {}
+            t.jsonb :new, null: false, default: {}
+            t.jsonb :context, null: false, default: {}
+          end
+
+          unique_index = %i[table_name primary_key]
+
+          base_class.connection.add_index table_name, unique_index, unique: true
+
+          Class.new(base_class) do
+            # Using cast here since Sorbet bugs when we don't pass a explicit class to `Class.new`
+            T.cast(self, T.class_of(::ActiveRecord::Base)).table_name = table_name
+
+            # All this sh$#1t was necessary because AR schema cache doesn't work with temporary tables...
+            insert_all_class = Class.new(::ActiveRecord::InsertAll) do
+              unique_index_definition = ::ActiveRecord::ConnectionAdapters::IndexDefinition.new(
+                table_name, "unique_index", true, unique_index
+              )
+              define_method(:find_unique_index_for) { |_| unique_index_definition }
+            end
+
+            define_singleton_method(:upsert) do |attributes, update_only: nil|
+              insert_all_class
+                .new(
+                  T.cast(self, T.class_of(::ActiveRecord::Base)).none,
+                  T.cast(self, T.class_of(::ActiveRecord::Base)).connection,
+                  [attributes],
+                  on_duplicate: :update,
+                  unique_by: unique_index,
+                  update_only:,
+                  returning: nil,
+                  record_timestamps: nil,
+                )
+                .execute
+            end
+          end
+        end
+      end
+
+      def on_commit(_event)
+        @table
+          .in_batches(of: @batch_size)
+          .each_record
+          .lazy
+          .filter_map { |persisted_event| deserialize(persisted_event) }
+          .each { |event| @watcher.on_record_changed(event) }
+
+        base_class.connection.drop_table @table.table_name
+      end
+
+      sig { params(event: InsertEvent).void }
+      def on_insert(event)
+        @table.upsert(serialize(event))
+      end
+
+      sig { params(event: UpdateEvent).void }
+      def on_update(event)
+        @table.upsert(serialize(event), update_only: %w[new])
+      end
+
+      sig { params(event: DeleteEvent).void }
+      def on_delete(event)
+        case @table.where(table_name: event.table, primary_key: event.primary_key).pluck(:action, :old).first
+        in ["insert", _]
+          @table.where(table_name: event.table, primary_key: event.primary_key).delete_all
+        in ["update", old]
+          @table.upsert(serialize(event).merge(old:))
+        in ["delete", _]
+          # We don't need to store another delete
+        else
+          @table.upsert(serialize(event))
+        end
+      end
+
+      private
+
+      sig { returns(T.class_of(::ActiveRecord::Base)) }
+      def base_class
+        self.class.base_active_record_class || ::ActiveRecord::Base
+      end
+
+      def serialize(event)
+        serialized = {
+          transaction_id: event.transaction_id,
+          lsn: event.lsn,
+          table_name: event.table,
+          primary_key: event.primary_key,
+          context: event.context,
+        }
+        case event
+        when InsertEvent
+          serialized.merge(action: "insert", new: event.new)
+        when UpdateEvent
+          serialized.merge(action: "update", old: event.old, new: event.new)
+        when DeleteEvent
+          serialized.merge(action: "delete", old: event.old)
+        else
+          serialized
+        end
+      end
+
+      def deserialize(persisted_event)
+        deserialized = {
+          transaction_id: persisted_event.transaction_id,
+          lsn: persisted_event.lsn,
+          table: persisted_event.table_name,
+          primary_key: persisted_event.primary_key,
+          context: persisted_event.context,
+        }
+        case persisted_event.action
+        when "insert"
+          InsertEvent.new(**deserialized, new: persisted_event.new)
+        when "update"
+          UpdateEvent.new(**deserialized, old: persisted_event.old, new: persisted_event.new)
+        when "delete"
+          DeleteEvent.new(**deserialized, old: persisted_event.old)
+        end
+      end
+    end
+  end
+end