Simplify restore

Author: robacourt

packages/sync-service/lib/electric/replication/shape_log_collector.ex

@@ -270,23 +270,22 @@ defmodule Electric.Replication.ShapeLogCollector do
 
         {partitions, event_router, layers, count} =
           state.stack_id
-          |> Electric.ShapeCache.ShapeStatus.list_shapes()
+          |> Electric.ShapeCache.ShapeStatus.list_shapes_with_ids()
           |> Enum.reduce(
             {state.partitions, state.event_router, state.dependency_layers, 0},
-            fn {shape_handle, shape}, {partitions, event_router, layers, count} = acc ->
-              with {:ok, shape_id} <-
-                     Electric.ShapeCache.ShapeStatus.id_for_handle(state.stack_id, shape_handle),
-                   {:ok, dependency_ids} <- dependency_ids(state.stack_id, shape) do
-                restore_shape(
-                  {shape_handle, shape_id, shape, dependency_ids},
-                  {partitions, event_router, layers, count}
-                )
-              else
-                # The shape (or one of its dependencies) has no id mapping in
-                # ShapeStatus, e.g. it's mid-removal. Skip it.
+            fn {shape_handle, shape_id, shape}, {partitions, event_router, layers, count} = acc ->
+              case dependency_ids(state.stack_id, shape) do
+                {:ok, dependency_ids} ->
+                  restore_shape(
+                    {shape_handle, shape_id, shape, dependency_ids},
+                    {partitions, event_router, layers, count}
+                  )
+
+                # A dependency has no id mapping in ShapeStatus, e.g. it's
+                # mid-removal. Skip this shape.
                 :error ->
                   Logger.warning(
-                    "Skipping shape during restore: no id mapping",
+                    "Skipping shape during restore: dependency has no id mapping",
                     shape_handle: shape_handle
                   )
 

packages/sync-service/lib/electric/shape_cache/shape_status.ex

@@ -154,6 +154,21 @@ defmodule Electric.ShapeCache.ShapeStatus do
     end)
   end
 
+  @doc """
+  Like `list_shapes/1` but also returns each shape's in-memory `shape_id`.
+  """
+  @spec list_shapes_with_ids(stack_id()) :: [{shape_handle(), shape_id(), Shape.t()}]
+  def list_shapes_with_ids(stack_id) when is_stack_id(stack_id) do
+    stack_id
+    |> list_shapes()
+    |> Enum.flat_map(fn {handle, shape} ->
+      case id_for_handle(stack_id, handle) do
+        {:ok, id} -> [{handle, id, shape}]
+        :error -> []
+      end
+    end)
+  end
+
   @spec topological_sort([{shape_handle(), Shape.t()}]) :: [{shape_handle(), Shape.t()}]
   defp topological_sort(handles_and_shapes, acc \\ [], visited \\ MapSet.new())
   defp topological_sort([], acc, _visited), do: Enum.reverse(acc) |> List.flatten()

packages/sync-service/test/electric/shape_cache/shape_status_test.exs

@@ -449,6 +449,15 @@ defmodule Electric.ShapeCache.ShapeStatusTest do
       assert is_integer(restored_id)
       assert {:ok, ^handle} = ShapeStatus.handle_for_id(stack_id, restored_id)
     end
+
+    test "list_shapes_with_ids returns each shape's handle, id and shape", %{stack_id: stack_id} do
+      {:ok, {handle1, id1}} = ShapeStatus.add_shape(stack_id, shape!())
+      {:ok, {handle2, id2}} = ShapeStatus.add_shape(stack_id, shape2!())
+
+      result = ShapeStatus.list_shapes_with_ids(stack_id) |> Enum.sort_by(&elem(&1, 1))
+
+      assert [{^handle1, ^id1, %Shape{}}, {^handle2, ^id2, %Shape{}}] = result
+    end
   end
 
   defp shape!, do: shape!("test")