🗂️ SinkConsumer SchemaFilter

Author: RTLS

lib/sequin/consumers/consumers.ex

@@ -12,6 +12,7 @@ defmodule Sequin.Consumers do
   alias Sequin.Consumers.ConsumerRecord
   alias Sequin.Consumers.Function
   alias Sequin.Consumers.HttpEndpoint
+  alias Sequin.Consumers.SchemaFilter
   alias Sequin.Consumers.SequenceFilter
   alias Sequin.Consumers.SequenceFilter.CiStringValue
   alias Sequin.Consumers.SequenceFilter.ColumnFilter
@@ -1098,12 +1099,24 @@ defmodule Sequin.Consumers do
     end
   end
 
-  # Source Table Matching
+  # Schema Matching
+  def matches_message?(
+        %SinkConsumer{schema_filter: %SchemaFilter{} = schema_filter} = consumer,
+        %SlotProcessor.Message{} = message
+      ) do
+    matches? = message_matches_schema?(schema_filter, message)
+
+    Health.put_event(consumer, %Event{slug: :messages_filtered, status: :success})
+
+    matches?
+  end
+
+  # Sequence Matching
   def matches_message?(
         %{sequence: %Sequence{} = sequence, sequence_filter: %SequenceFilter{} = sequence_filter} = consumer,
         %SlotProcessor.Message{} = message
       ) do
-    matches? = matches_message?(sequence, sequence_filter, message)
+    matches? = message_matches_sequence?(sequence, sequence_filter, message)
 
     Health.put_event(consumer, %Event{slug: :messages_filtered, status: :success})
 
@@ -1123,27 +1136,14 @@ defmodule Sequin.Consumers do
       reraise error, __STACKTRACE__
   end
 
+  # Source Table Matching
   def matches_message?(consumer_or_wal_pipeline, %SlotProcessor.Message{} = message) do
     matches? =
       Enum.any?(consumer_or_wal_pipeline.source_tables, fn %SourceTable{} = source_table ->
         table_matches = source_table.oid == message.table_oid
         action_matches = action_matches?(source_table.actions, message.action)
         column_filters_match = column_filters_match_message?(source_table.column_filters, message)
 
-        # Logger.debug("""
-        # [Consumers]
-        #   matches?: #{table_matches && action_matches && column_filters_match}
-        #     table_matches: #{table_matches}
-        #     action_matches: #{action_matches}
-        #     column_filters_match: #{column_filters_match}
-
-        #   consumer_or_wal_pipeline:
-        #     #{inspect(consumer_or_wal_pipeline, pretty: true)}
-
-        #   message:
-        #     #{inspect(message, pretty: true)}
-        # """)
-
         table_matches && action_matches && column_filters_match
       end)
 
@@ -1168,13 +1168,20 @@ defmodule Sequin.Consumers do
       reraise error, __STACKTRACE__
   end
 
-  def matches_message?(%Sequence{} = _sequence, %SequenceFilter{} = sequence_filter, %SlotProcessor.Message{} = message) do
-    # table_matches? = sequence.table_oid == message.table_oid
+  defp message_matches_schema?(%SchemaFilter{} = schema_filter, %SlotProcessor.Message{} = message) do
+    message.table_schema == schema_filter.schema
+  end
+
+  defp message_matches_sequence?(
+         %Sequence{} = sequence,
+         %SequenceFilter{} = sequence_filter,
+         %SlotProcessor.Message{} = message
+       ) do
+    table_matches? = sequence.table_oid == message.table_oid
     actions_match? = action_matches?(sequence_filter.actions, message.action)
     column_filters_match? = column_filters_match_message?(sequence_filter.column_filters, message)
 
-    # table_matches? and actions_match? and column_filters_match?
-    actions_match? and column_filters_match?
+    table_matches? and actions_match? and column_filters_match?
   end
 
   def matches_record?(

lib/sequin/consumers/schema_filter.ex

@@ -0,0 +1,24 @@
+defmodule Sequin.Consumers.SchemaFilter do
+  @moduledoc false
+  use Ecto.Schema
+
+  import Ecto.Changeset
+
+  @derive {Jason.Encoder, only: [:schema]}
+
+  @type t :: %__MODULE__{
+          schema: String.t()
+        }
+
+  @primary_key false
+  embedded_schema do
+    field :schema, :string
+  end
+
+  def create_changeset(schema_filter, attrs) do
+    schema_filter
+    |> cast(attrs, [:schema])
+    |> validate_required([:schema])
+    |> validate_format(:schema, ~r/^[a-zA-Z_][a-zA-Z0-9_]*$/, message: "must be a valid PostgreSQL identifier")
+  end
+end

lib/sequin/consumers/sink_consumer.ex

@@ -20,6 +20,7 @@ defmodule Sequin.Consumers.SinkConsumer do
   alias Sequin.Consumers.RabbitMqSink
   alias Sequin.Consumers.RedisStreamSink
   alias Sequin.Consumers.RedisStringSink
+  alias Sequin.Consumers.SchemaFilter
   alias Sequin.Consumers.SequenceFilter
   alias Sequin.Consumers.SequinStreamSink
   alias Sequin.Consumers.SnsSink
@@ -103,6 +104,7 @@ defmodule Sequin.Consumers.SinkConsumer do
     # FIXME: Refactor / remove both of these?
     belongs_to :sequence, Sequence
     embeds_one :sequence_filter, SequenceFilter, on_replace: :delete
+    embeds_one :schema_filter, SchemaFilter, on_replace: :delete
 
     belongs_to :account, Account
     belongs_to :replication_slot, PostgresReplicationSlot
@@ -150,6 +152,7 @@ defmodule Sequin.Consumers.SinkConsumer do
     ])
     |> changeset(attrs)
     |> cast_embed(:sequence_filter, with: &SequenceFilter.create_changeset/2)
+    |> cast_embed(:schema_filter, with: &SchemaFilter.create_changeset/2)
     |> foreign_key_constraint(:sequence_id)
     |> foreign_key_constraint(:transform_id)
     |> foreign_key_constraint(:routing_id)
@@ -160,17 +163,20 @@ defmodule Sequin.Consumers.SinkConsumer do
       name: "ensure_batch_size_one",
       message: "batch_size must be 1 when batch is false for webhook sinks"
     )
+    |> validate_filter_constraints()
     |> Sequin.Changeset.validate_name()
   end
 
   def update_changeset(consumer, attrs) do
     consumer
     |> changeset(attrs)
     |> cast_embed(:sequence_filter, with: &SequenceFilter.create_changeset/2)
+    |> cast_embed(:schema_filter, with: &SchemaFilter.create_changeset/2)
     |> check_constraint(:batch_size,
       name: "ensure_batch_size_one",
       message: "batch_size must be 1 when batch is false for webhook sinks"
     )
+    |> validate_filter_constraints()
   end
 
   def changeset(consumer, attrs) do
@@ -236,6 +242,22 @@ defmodule Sequin.Consumers.SinkConsumer do
     end
   end
 
+  defp validate_filter_constraints(changeset) do
+    sequence_filter = get_field(changeset, :sequence_filter)
+    schema_filter = get_field(changeset, :schema_filter)
+
+    cond do
+      is_nil(sequence_filter) and is_nil(schema_filter) ->
+        add_error(changeset, :base, "either sequence_filter or schema_filter must be set")
+
+      not is_nil(sequence_filter) and not is_nil(schema_filter) ->
+        add_error(changeset, :base, "cannot set both sequence_filter and schema_filter")
+
+      true ->
+        changeset
+    end
+  end
+
   defp put_defaults(changeset) do
     changeset
     |> put_change(:batch_size, get_field(changeset, :batch_size) || 1)

lib/sequin/runtime/message_handler.ex

@@ -527,14 +527,17 @@ defmodule Sequin.Runtime.MessageHandler do
   defp generate_group_id(consumer, message) do
     # This should be way more assertive - we should error if we don't find the source table
     # We have a lot of tests that do not line up consumer source_tables with the message table oid
-    if consumer.sequence_filter.group_column_attnums do
-      Enum.map_join(consumer.sequence_filter.group_column_attnums, ":", fn attnum ->
-        fields = if message.action == :delete, do: message.old_fields, else: message.fields
-        field = Sequin.Enum.find!(fields, &(&1.column_attnum == attnum))
-        to_string(field.value)
-      end)
-    else
-      Enum.map_join(message.ids, ":", &to_string/1)
+    case consumer do
+      %SinkConsumer{sequence_filter: %SequenceFilter{group_column_attnums: group_column_attnums}}
+      when not is_nil(group_column_attnums) ->
+        Enum.map_join(consumer.sequence_filter.group_column_attnums, ":", fn attnum ->
+          fields = if message.action == :delete, do: message.old_fields, else: message.fields
+          field = Sequin.Enum.find!(fields, &(&1.column_attnum == attnum))
+          to_string(field.value)
+        end)
+
+      _ ->
+        Enum.map_join(message.ids, ":", &to_string/1)
     end
   end
 

priv/repo/migrations/20250523211710_add_schema_filter_to_sink_consumers.exs

@@ -0,0 +1,11 @@
+defmodule Sequin.Repo.Migrations.AddSchemaFilterToSinkConsumers do
+  use Ecto.Migration
+
+  @config_schema Application.compile_env(:sequin, [Sequin.Repo, :config_schema_prefix])
+
+  def change do
+    alter table(:sink_consumers, prefix: @config_schema) do
+      add :schema_filter, :jsonb, null: true
+    end
+  end
+end

test/sequin/message_handler_test.exs

@@ -165,14 +165,22 @@ defmodule Sequin.MessageHandlerTest do
              }
     end
 
-    @tag skip: "FIXME: This test is broken while we work on table fan-in"
     test "fans out messages correctly for mixed message_kind consumers and wal_pipelines" do
       account = AccountsFactory.insert_account!()
       database = DatabasesFactory.insert_postgres_database!(account_id: account.id)
 
       field = ReplicationFactory.field()
-      message1 = ReplicationFactory.postgres_message(table_oid: 123, action: :insert, fields: [field])
-      message2 = ReplicationFactory.postgres_message(table_oid: 456, action: :update, fields: [field])
+      table_schema1 = Factory.postgres_object()
+      table_schema2 = Factory.postgres_object()
+
+      message1 =
+        ReplicationFactory.postgres_message(table_oid: 123, action: :insert, fields: [field], table_schema: table_schema1)
+
+      message2 =
+        ReplicationFactory.postgres_message(table_oid: 456, action: :update, fields: [field], table_schema: table_schema2)
+
+      message3 =
+        ReplicationFactory.postgres_message(table_oid: 789, action: :insert, fields: [field], table_schema: table_schema1)
 
       sequence1 =
         DatabasesFactory.insert_sequence!(
@@ -218,11 +226,24 @@ defmodule Sequin.MessageHandlerTest do
           source_tables: []
         )
 
+      schema_filter = ConsumersFactory.schema_filter_attrs(schema: table_schema1)
+
+      consumer3 =
+        ConsumersFactory.insert_sink_consumer!(
+          account_id: account.id,
+          sequence_id: nil,
+          sequence_filter: nil,
+          source_tables: [],
+          schema_filter: schema_filter
+        )
+
       start_supervised!({SlotMessageStoreSupervisor, [consumer: consumer1, test_pid: self(), persisted_mode?: false]})
       start_supervised!({SlotMessageStoreSupervisor, [consumer: consumer2, test_pid: self(), persisted_mode?: false]})
+      start_supervised!({SlotMessageStoreSupervisor, [consumer: consumer3, test_pid: self(), persisted_mode?: false]})
 
       consumer1 = Repo.preload(consumer1, [:postgres_database, :sequence, :filter])
       consumer2 = Repo.preload(consumer2, [:postgres_database, :sequence, :filter])
+      consumer3 = Repo.preload(consumer3, [:postgres_database, :sequence, :filter])
 
       database = Repo.preload(database, :sequences)
 
@@ -236,16 +257,17 @@ defmodule Sequin.MessageHandlerTest do
         )
 
       context = %MessageHandler.Context{
-        consumers: [consumer1, consumer2],
+        consumers: [consumer1, consumer2, consumer3],
         wal_pipelines: [wal_pipeline],
         replication_slot_id: UUID.uuid4(),
         postgres_database: database
       }
 
-      {:ok, 4} = MessageHandler.handle_messages(context, [message1, message2])
+      {:ok, 6} = MessageHandler.handle_messages(context, [message1, message2, message3])
 
       consumer1_messages = list_messages(consumer1)
       consumer2_messages = list_messages(consumer2)
+      consumer3_messages = list_messages(consumer3)
       wal_events = Replication.list_wal_events(wal_pipeline.id)
 
       assert length(consumer1_messages) == 1
@@ -254,6 +276,9 @@ defmodule Sequin.MessageHandlerTest do
       assert length(consumer2_messages) == 1
       assert hd(consumer2_messages).table_oid == 456
 
+      assert length(consumer3_messages) == 2
+      assert Enum.all?(consumer3_messages, &(&1.data.metadata.table_schema == table_schema1))
+
       assert length(wal_events) == 2
     end
 

test/sequin/postgres_replication_test.exs

@@ -15,6 +15,7 @@ defmodule Sequin.PostgresReplicationTest do
   import ExUnit.CaptureLog
 
   alias Sequin.Consumers
+  alias Sequin.Consumers.SchemaFilter
   alias Sequin.Consumers.SequenceFilter
   alias Sequin.Databases.ConnectionCache
   alias Sequin.Databases.DatabaseUpdateWorker
@@ -1492,6 +1493,16 @@ defmodule Sequin.PostgresReplicationTest do
       # Randomly select a consumer
       consumer = Enum.random([event_consumer, record_consumer])
 
+      # Attach a schema filter to the consumer
+      consumer
+      |> Ecto.Changeset.cast(%{schema_filter: ConsumersFactory.schema_filter_attrs(schema: "public")}, [])
+      |> Ecto.Changeset.cast_embed(:schema_filter, with: &SchemaFilter.create_changeset/2)
+      |> Repo.update!()
+
+      # Restart the consumer to apply the changes
+      Consumers.update_sink_consumer(consumer, %{})
+      Runtime.Supervisor.refresh_message_handler_ctx(consumer.replication_slot_id)
+
       # Verify no consumer messages yet
       assert list_messages(consumer) == []
 

test/support/factory/consumers_factory.ex

@@ -18,6 +18,7 @@ defmodule Sequin.Factory.ConsumersFactory do
   alias Sequin.Consumers.RabbitMqSink
   alias Sequin.Consumers.RedisStreamSink
   alias Sequin.Consumers.RedisStringSink
+  alias Sequin.Consumers.SchemaFilter
   alias Sequin.Consumers.SequenceFilter
   alias Sequin.Consumers.SequenceFilter.ColumnFilter
   alias Sequin.Consumers.SequinStreamSink
@@ -133,7 +134,6 @@ defmodule Sequin.Factory.ConsumersFactory do
   def sink_consumer_attrs(attrs \\ []) do
     attrs
     |> sink_consumer()
-    |> Sequin.Map.from_ecto()
     |> Map.update!(:sink, &Sequin.Map.from_ecto/1)
     |> Map.update!(:source_tables, fn source_tables ->
       Enum.map(source_tables, fn source_table ->
@@ -153,6 +153,7 @@ defmodule Sequin.Factory.ConsumersFactory do
         end)
       end
     end)
+    |> Sequin.Map.from_ecto()
   end
 
   def insert_sink_consumer!(attrs \\ []) do
@@ -815,6 +816,23 @@ defmodule Sequin.Factory.ConsumersFactory do
     |> Sequin.Map.from_ecto()
   end
 
+  def schema_filter(attrs \\ []) do
+    attrs = Map.new(attrs)
+
+    merge_attributes(
+      %SchemaFilter{
+        schema: Factory.postgres_object()
+      },
+      attrs
+    )
+  end
+
+  def schema_filter_attrs(attrs \\ []) do
+    attrs
+    |> schema_filter()
+    |> Sequin.Map.from_ecto(keep_nils: true)
+  end
+
   def backfill(attrs \\ []) do
     attrs = Map.new(attrs)