fix: Run migrations only when necessary (#1203)

lib/realtime/tenants.ex

@@ -67,44 +67,13 @@ defmodule Realtime.Tenants do
 
       {:ok, health_conn} ->
         connected_cluster = UsersCounter.tenant_users(external_id)
-        %{extensions: [%{settings: settings} | _]} = Cache.get_tenant_by_external_id(external_id)
-
-        query =
-          "select * from pg_catalog.pg_tables where schemaname = 'realtime' and tablename = 'schema_migrations';"
-
-        Database.transaction(health_conn, fn transaction_conn ->
-          res = Postgrex.query!(transaction_conn, query, [])
-
-          if res.rows == [] do
-            Migrations.run_migrations(%Migrations{
-              tenant_external_id: external_id,
-              settings: settings
-            })
-          end
-        end)
-
+        Migrations.maybe_run_migrations(health_conn, external_id)
         {:ok, %{healthy: true, db_connected: true, connected_cluster: connected_cluster}}
 
       connected_cluster when is_integer(connected_cluster) ->
         tenant = Cache.get_tenant_by_external_id(external_id)
         {:ok, db_conn} = Database.connect(tenant, "realtime_health_check", 1)
-
-        Database.transaction(db_conn, fn transaction_conn ->
-          query =
-            "select * from pg_catalog.pg_tables where schemaname = 'realtime' and tablename = 'schema_migrations';"
-
-          res = Postgrex.query!(transaction_conn, query, [])
-
-          if res.rows == [] do
-            %{extensions: [%{settings: settings} | _]} = tenant
-
-            Migrations.run_migrations(%Migrations{
-              tenant_external_id: external_id,
-              settings: settings
-            })
-          end
-        end)
-
+        Migrations.maybe_run_migrations(db_conn, external_id)
         Process.alive?(db_conn) && GenServer.stop(db_conn)
         {:ok, %{healthy: true, db_connected: false, connected_cluster: connected_cluster}}
     end

lib/realtime/tenants/connect/migrations.ex

@@ -4,17 +4,10 @@ defmodule Realtime.Tenants.Connect.Migrations do
   """
   @behaviour Realtime.Tenants.Connect.Piper
   alias Realtime.Tenants.Migrations
-  alias Realtime.Tenants.Cache
-  @impl true
-  def run(acc) do
-    %{tenant_id: tenant_id} = acc
-    tenant = Cache.get_tenant_by_external_id(tenant_id)
-    [%{settings: settings} | _] = tenant.extensions
-    migrations = %Migrations{tenant_external_id: tenant.external_id, settings: settings}
 
-    case Migrations.run_migrations(migrations) do
-      :ok -> {:ok, acc}
-      {:error, error} -> {:error, error}
-    end
+  @impl true
+  def run(%{db_conn_pid: db_conn_pid, tenant_id: tenant_id} = acc) do
+    {:ok, _} = Migrations.maybe_run_migrations(db_conn_pid, tenant_id)
+    {:ok, acc}
   end
 end

lib/realtime/tenants/migrations.ex

@@ -10,8 +10,9 @@ defmodule Realtime.Tenants.Migrations do
 
   alias Realtime.Crypto
   alias Realtime.Database
-  alias Realtime.Repo
   alias Realtime.Registry.Unique
+  alias Realtime.Repo
+  alias Realtime.Tenants.Cache
 
   alias Realtime.Tenants.Migrations.{
     CreateRealtimeSubscriptionTable,
@@ -121,6 +122,9 @@ defmodule Realtime.Tenants.Migrations do
     {20_241_030_150_047, MessagesPartitioning},
     {20_241_108_114_728, MessagesUsingUuid}
   ]
+
+  @expected_migration_count length(@migrations)
+
   defstruct [:tenant_external_id, :settings]
   @spec run_migrations(map()) :: :ok | {:error, any()}
   def run_migrations(%__MODULE__{} = attrs) do
@@ -186,4 +190,26 @@ defmodule Realtime.Tenants.Migrations do
       end
     end)
   end
+
+  @doc """
+  Checks if the number of migrations ran in the database is equal to the expected number of migrations.
+
+  If not all migrations have been run, it will run the missing migrations.
+  """
+  @spec maybe_run_migrations(pid(), String.t()) :: {:ok, any()} | {:error, any()}
+  def maybe_run_migrations(db_conn, tenant_external_id) do
+    query =
+      "select * from pg_catalog.pg_tables where schemaname = 'realtime' and tablename = 'schema_migrations';"
+
+    %{extensions: [%{settings: settings} | _]} =
+      Cache.get_tenant_by_external_id(tenant_external_id)
+
+    Database.transaction(db_conn, fn transaction_conn ->
+      %{num_rows: num_rows} = Postgrex.query!(transaction_conn, query, [])
+
+      if num_rows < @expected_migration_count do
+        run_migrations(%__MODULE__{tenant_external_id: tenant_external_id, settings: settings})
+      end
+    end)
+  end
 end

mix.exs

@@ -4,7 +4,7 @@ defmodule Realtime.MixProject do
   def project do
     [
       app: :realtime,
-      version: "2.33.27",
+      version: "2.33.28",
       elixir: "~> 1.16.0",
       elixirc_paths: elixirc_paths(Mix.env()),
       start_permanent: Mix.env() == :prod,

test/realtime/tenants/connect_test.exs

@@ -181,11 +181,12 @@ defmodule Realtime.Tenants.ConnectTest do
     end
 
     test "on migrations failure, stop the process", %{tenant: tenant} do
-      with_mock Realtime.Tenants.Migrations, [], run_migrations: fn _ -> raise("error") end do
+      with_mock Realtime.Tenants.Migrations, [],
+        maybe_run_migrations: fn _, _ -> raise("error") end do
         assert {:error, :tenant_database_unavailable} =
                  Connect.lookup_or_start_connection(tenant.external_id)
 
-        assert_called(Realtime.Tenants.Migrations.run_migrations(:_))
+        assert_called(Realtime.Tenants.Migrations.maybe_run_migrations(:_, :_))
       end
     end