wal 0.0.0 → 0.0.1

data/lib/wal/replicator.rb

@@ -0,0 +1,205 @@
+# typed: false
+
+require "ostruct"
+
+module Wal
+  # Responsible to hook into a Postgres logical replication slot and stream the changes to a specific `Watcher`.
+  # Also it supports inject "contexts" into the replication events.
+  class Replicator
+    extend T::Sig
+    include PG::Replication::Protocol
+
+    sig do
+      params(
+        replication_slot: String,
+        use_temporary_slot: T::Boolean,
+        db_config: T::Hash[Symbol, T.untyped],
+      ).void
+    end
+    def initialize(
+      replication_slot:,
+      use_temporary_slot: false,
+      db_config: ActiveRecord::Base.configurations.configs_for(name: "primary").configuration_hash
+    )
+      @db_config = db_config
+      @replication_slot = replication_slot
+      @use_temporary_slot = use_temporary_slot
+    end
+
+    sig { params(watcher: Watcher, publications: T::Array[String]).void }
+    def replicate_forever(watcher, publications:)
+      replication = replicate(watcher, publications:)
+      loop { replication.next }
+    rescue StopIteration
+      nil
+    end
+
+    sig { params(watcher: Watcher, publications: T::Array[String]).returns(T::Enumerator::Lazy[Event]) }
+    def replicate(watcher, publications:)
+      watch_conn = PG.connect(
+        dbname: @db_config[:database],
+        host: @db_config[:host],
+        user: @db_config[:username],
+        password: @db_config[:password].presence,
+        port: @db_config[:port].presence,
+        replication: "database",
+      )
+
+      begin
+        watch_conn.query(<<~SQL)
+          CREATE_REPLICATION_SLOT #{@replication_slot} #{@use_temporary_slot ? "TEMPORARY" : ""} LOGICAL "pgoutput"
+        SQL
+      rescue PG::DuplicateObject
+        # We are fine, we already have created this slot in a previous run
+      end
+
+      tables = {}
+      context = {}
+      transaction_id = nil
+
+      watch_conn.start_pgoutput_replication_slot(@replication_slot, publications, messages: true).filter_map do |msg|
+        case msg
+        in XLogData(data: PG::Replication::PGOutput::Relation(oid:, name:, columns:))
+          tables[oid] = Table.new(
+            # TODO: for now we are forcing an id column here, but that is not really correct
+            primary_key_colums: columns.any? { |col| col.name == "id" } ? ["id"] : [],
+            name:,
+            columns: columns.map do |col|
+              Column.new(
+                name: col.name,
+                decoder: ActiveRecord::Base.connection.lookup_cast_type_from_column(
+                  # We have to create this OpenStruct because weird AR API reasons...
+                  # And the `sql_type` param luckly doesn't really matter for our use case
+                  ::OpenStruct.new(oid: col.oid, fmod: col.modifier, sql_type: "")
+                ),
+              )
+            end
+          )
+          next
+
+        in XLogData(lsn:, data: PG::Replication::PGOutput::Begin(xid:, timestamp:, final_lsn:))
+          transaction_id = xid
+          context = {}
+          BeginTransactionEvent.new(
+            transaction_id:,
+            lsn:,
+            final_lsn:,
+            timestamp:,
+          ).tap { |event| watcher.on_event(event) }
+
+        in XLogData(lsn:, data: PG::Replication::PGOutput::Commit(timestamp:))
+          CommitTransactionEvent.new(
+            transaction_id:,
+            lsn:,
+            context:,
+            timestamp:,
+          ).tap do |event|
+            watcher.on_event(event)
+            watch_conn.standby_status_update(write_lsn: lsn)
+          end
+
+        in XLogData(data: PG::Replication::PGOutput::Message(prefix:, content:)) if watcher.valid_context_prefix? prefix
+          begin
+            context = JSON.parse(content).presence || {}
+            next
+          rescue JSON::ParserError
+            # Invalid context received, just ignore
+          end
+
+        in XLogData(lsn:, data: PG::Replication::PGOutput::Insert(oid:, new:))
+          table = tables[oid]
+          next unless watcher.should_watch_table? table.name
+          new_data = table.decode_row(new)
+          record_id = table.primary_key(new_data)
+          next unless record_id
+
+          InsertEvent.new(
+            transaction_id:,
+            lsn:,
+            context:,
+            table: table.name,
+            primary_key: record_id,
+            new: new_data,
+          ).tap { |event| watcher.on_event(event) }
+
+        in XLogData(lsn:, data: PG::Replication::PGOutput::Update(oid:, new:, old:))
+          table = tables[oid]
+          next unless watcher.should_watch_table? table.name
+          old_data = table.decode_row(old)
+          new_data = table.decode_row(new)
+          record_id = table.primary_key(new_data)
+          next unless record_id
+
+          UpdateEvent.new(
+            transaction_id:,
+            lsn:,
+            context:,
+            table: table.name,
+            primary_key: record_id,
+            old: old_data,
+            new: new_data,
+          ).tap { |event| watcher.on_event(event) }
+
+        in XLogData(lsn:, data: PG::Replication::PGOutput::Delete(oid:, old:, key:))
+          table = tables[oid]
+          next unless watcher.should_watch_table? table.name
+          old_data = table.decode_row(old.presence || key)
+          record_id = table.primary_key(old_data)
+          next unless record_id
+
+          DeleteEvent.new(
+            transaction_id:,
+            lsn:,
+            context:,
+            table: table.name,
+            primary_key: record_id,
+            old: old_data,
+          ).tap { |event| watcher.on_event(event) }
+
+        else
+          next
+        end
+      end
+    end
+
+    class Column < T::Struct
+      const :name, String
+      const :decoder, T.untyped
+
+      def decode(value)
+        decoder.deserialize(value)
+      end
+    end
+
+    class Table < T::Struct
+      const :name, String
+      const :primary_key_colums, T::Array[String]
+      const :columns, T::Array[Column]
+
+      def primary_key(decoded_row)
+        case primary_key_colums
+        in [key]
+          case decoded_row[key]
+          in Integer => id
+            id
+          in String => id
+            id
+          else
+            # Only supporting string and integer primary keys for now
+            nil
+          end
+        else
+          # Not supporting coumpound primary keys
+          nil
+        end
+      end
+
+      def decode_row(values)
+        values
+          .zip(columns)
+          .map { |tuple, col| [col.name, col.decode(tuple.data)] }
+          .to_h
+      end
+    end
+  end
+end

data/lib/wal/streaming_watcher.rb

@@ -0,0 +1,74 @@
+# typed: true
+
+module Wal
+  # A watcher that streams all the events of each WAL transaction on a separate thread.
+  #
+  # Useful to improve the throughput, as it will allow you to process events while fetching for more in parallel.
+  #
+  # Example:
+  #
+  # Watcher that persists all delete events as it arrives using a single database transaction, and without waiting
+  # for the full WAL log transaction to be finished.
+  #
+  # ```ruby
+  # class RegisterDeletesWalWatcher < Wal::StreamingWalWatcher
+  #   sig { override.params(events: T::Enumerator[Event]).void }
+  #   def on_transaction_events(events)
+  #     DeletedApplicationRecord.transaction do
+  #       events
+  #         .lazy
+  #         .filter { |event| event.is_a? DeleteEvent }
+  #         .each { |event| DeletedApplicationRecord.create_from_event(event) }
+  #     end
+  #   end
+  # end
+  # ```
+  class StreamingWatcher
+    extend T::Sig
+    extend T::Helpers
+    include Wal::Watcher
+    abstract!
+
+    sig { abstract.params(events: T::Enumerator[Event]).void }
+    def on_transaction_events(events); end
+
+    sig { params(event: BeginTransactionEvent).returns(Integer) }
+    def queue_size(event)
+      5_000
+    end
+
+    sig { override.params(event: Event).void }
+    def on_event(event)
+      case event
+      when BeginTransactionEvent
+        @queue = SizedQueue.new(queue_size(event))
+
+        event_stream = Enumerator.new do |y|
+          while (item = @queue.pop)
+            case item
+            when CommitTransactionEvent
+              y << item
+              break
+            else
+              y << item
+            end
+          end
+        end
+        @worker = Thread.new { on_transaction_events(event_stream) }
+
+        @queue << event
+
+      when CommitTransactionEvent
+        @queue << event
+        @worker.join
+
+        # We are cleaning this up to hint to Ruby GC that this can be freed before the next begin transaction arrives
+        @queue.clear
+        @queue = nil
+
+      else
+        @queue << event
+      end
+    end
+  end
+end

data/lib/wal/version.rb

@@ -1,5 +1,6 @@
+# typed: strict
 # frozen_string_literal: true
 
 module Wal
-  VERSION = "0.0.0"
+  VERSION = "0.0.1"
 end

data/lib/wal/watcher.rb

@@ -0,0 +1,95 @@
+# typed: strict
+
+module Wal
+  # Watcher is the core API used to hook into Postgres WAL log.
+  # The only required method on the API is the `on_event`, which will receive a WAL entry of the following events:
+  # - Transaction started: `Wal::BeginTransactionEvent`
+  # - Row inserted: `Wal::InsertEvent`
+  # - Row updated: `Wal::UpdateEvent`
+  # - Row deleted: `Wal::DeleteEvent`
+  # - Transaction committed: `Wal::CommitTransactionEvent`
+  #
+  # The `on_event` method will be called without any buffering, so it is up to implementators to aggregate them if
+  # desired. In practice, it is rarelly useful to implement this module directly for application level business logic,
+  # and instead it is more recomended using more specific ones, such as the `RecordWatcher` and `StreamingWalWatcher`.
+  module Watcher
+    extend T::Sig
+    extend T::Helpers
+    include Wal
+    abstract!
+
+    sig { abstract.params(event: Event).void }
+    def on_event(event); end
+
+    # Allows dropping the processing of any table
+    sig { params(table: String).returns(T::Boolean) }
+    def should_watch_table?(table)
+      true
+    end
+
+    # Check if the given context prefix should be allowed for this watcher
+    sig { params(prefix: String).returns(T::Boolean) }
+    def valid_context_prefix?(prefix)
+      true
+    end
+
+    # Include this module if you prefer to work with each event having its own method.
+    # This might be useful when you always want to process each type of event in a different way.
+    #
+    # Example:
+    #
+    # Watcher that calculates how much time passed between the begin and commit of a WAL transaction.
+    #
+    # ```ruby
+    # class MeasureTransactionTimeWatcher
+    #   extend T::Sig
+    #   include Wal::Watcher
+    #   include Wal::Watcher::SeparatedEvents
+    #
+    #   sig { params(event: BeginTransactionEvent).void }
+    #   def on_begin(event)
+    #     @start_time = Time.current
+    #   end
+    #
+    #   sig { params(event: CommitTransactionEvent).void }
+    #   def on_commit(event)
+    #     puts "Transaction processing time: #{Time.current - @start_time}"
+    #   end
+    # end
+    # ```
+    module SeparatedEvents
+      extend T::Sig
+
+      sig { params(event: Event).void }
+      def on_event(event)
+        case event
+        when BeginTransactionEvent
+          on_begin(event)
+        when CommitTransactionEvent
+          on_commit(event)
+        when InsertEvent
+          on_insert(event)
+        when UpdateEvent
+          on_update(event)
+        when DeleteEvent
+          on_delete(event)
+        end
+      end
+
+      sig { params(event: BeginTransactionEvent).void }
+      def on_begin(event); end
+
+      sig { params(event: InsertEvent).void }
+      def on_insert(event); end
+
+      sig { params(event: UpdateEvent).void }
+      def on_update(event); end
+
+      sig { params(event: DeleteEvent).void }
+      def on_delete(event); end
+
+      sig { params(event: CommitTransactionEvent).void }
+      def on_commit(event); end
+    end
+  end
+end

data/lib/wal.rb

@@ -1,6 +1,133 @@
-# frozen_string_literal: true
+# typed: strict
 
+require "pg"
+require "pg/replication"
+require "active_support"
+require "active_record"
+require_relative "wal/watcher"
+require_relative "wal/noop_watcher"
+require_relative "wal/record_watcher"
+require_relative "wal/streaming_watcher"
+require_relative "wal/replicator"
+require_relative "wal/active_record_context_extension"
 require_relative "wal/version"
 
 module Wal
+  Event = T.type_alias do
+    T.any(
+      BeginTransactionEvent,
+      CommitTransactionEvent,
+      InsertEvent,
+      UpdateEvent,
+      DeleteEvent,
+    )
+  end
+
+  class BeginTransactionEvent < T::Struct
+    extend T::Sig
+
+    const :transaction_id, Integer
+    const :lsn, Integer
+    const :final_lsn, Integer
+    const :timestamp, Time
+
+    sig { returns(Integer) }
+    def estimated_size
+      final_lsn - lsn
+    end
+  end
+
+  class CommitTransactionEvent < T::Struct
+    const :transaction_id, Integer
+    const :lsn, Integer
+    const :context, T::Hash[String, T.untyped]
+    const :timestamp, Time
+  end
+
+  module DiffedEvent
+    extend T::Sig
+
+    sig { returns(T::Hash[String, [T.untyped, T.untyped]]) }
+    def diff
+      {}
+    end
+
+    sig { params(attribute: T.any(Symbol, String)).returns(T::Boolean) }
+    def changed_attribute?(attribute)
+      diff.key? attribute.to_s
+    end
+
+    sig { params(attribute: T.any(Symbol, String)).returns(T.untyped) }
+    def attribute(attribute)
+      if (changes = diff[attribute.to_s])
+        changes[1]
+      end
+    end
+
+    sig { params(attribute: T.any(Symbol, String)).returns(T.nilable([T.untyped, T.untyped])) }
+    def attribute_changes(attribute)
+      diff[attribute.to_s]
+    end
+
+    sig { params(attribute: T.any(Symbol, String)).returns(T.untyped) }
+    def attribute_was(attribute)
+      if (changes = diff[attribute.to_s])
+        changes[0]
+      end
+    end
+  end
+
+  class InsertEvent < T::Struct
+    extend T::Sig
+    include DiffedEvent
+
+    const :transaction_id, Integer
+    const :lsn, Integer
+    const :context, T::Hash[String, T.untyped]
+    const :table, String
+    const :primary_key, T.untyped
+    const :new, T::Hash[String, T.untyped]
+
+    sig { returns(T::Hash[String, [T.untyped, T.untyped]]) }
+    def diff
+      new.transform_values { |val| [nil, val] }
+    end
+  end
+
+  class UpdateEvent < T::Struct
+    extend T::Sig
+    include DiffedEvent
+
+    const :transaction_id, Integer
+    const :lsn, Integer
+    const :context, T::Hash[String, T.untyped]
+    const :table, String
+    const :primary_key, T.untyped
+    const :old, T::Hash[String, T.untyped]
+    const :new, T::Hash[String, T.untyped]
+
+    sig { returns(T::Hash[String, [T.untyped, T.untyped]]) }
+    def diff
+      (old.keys | new.keys).reduce({}) do |diff, key|
+        old[key] != new[key] ? diff.merge(key => [old[key], new[key]]) : diff
+      end
+    end
+  end
+
+  class DeleteEvent < T::Struct
+    extend T::Sig
+    include DiffedEvent
+
+    const :transaction_id, Integer
+    const :lsn, Integer
+    const :context, T::Hash[String, T.untyped]
+    const :table, String
+    const :primary_key, T.untyped
+    const :old, T::Hash[String, T.untyped]
+
+    sig { returns(T::Hash[String, [T.untyped, T.untyped]]) }
+    def diff
+      old.transform_values { |val| [val, nil] }
+    end
+  end
 end