feat: load tenant database dump (#1998)

* feat: load tenant database dump

Shortcut to load the entire Tenant DB schema at once instead of running
migrations sequentially that adds overhead to tenant creatiion workflow.
Applies only to new tenants.

Ref REAL-861

Author: leandrocp

.github/workflows/update-tenant-db-catalog.yml

@@ -6,6 +6,7 @@ on:
       - "lib/realtime/tenants/repo/migrations/**"
       - "lib/realtime/tenants/migrations.ex"
       - "lib/mix/tasks/realtime.export_tenant_db_catalog.ex"
+      - "lib/mix/tasks/realtime.export_tenant_db_dump.ex"
       - "mix.lock"
       - "mise.toml"
       - "compose.dbs.yml"
@@ -20,8 +21,8 @@ concurrency:
   cancel-in-progress: true
 
 jobs:
-  update-tenant-db-catalog:
-    name: Export catalog (pg${{ matrix.pg_major }})
+  update-tenant-dump:
+    name: Export catalog & dump (pg${{ matrix.pg_major }})
     runs-on: blacksmith-2vcpu-ubuntu-2404
     if: github.event.pull_request.head.repo.full_name == github.repository
     strategy:
@@ -54,6 +55,18 @@ jobs:
       - name: Install dependencies
         run: mix deps.get
 
+      - name: Install PostgreSQL client
+        run: |
+          sudo install -d /usr/share/postgresql-common/pgdg
+          sudo curl -o /usr/share/postgresql-common/pgdg/apt.postgresql.org.asc --fail \
+            https://www.postgresql.org/media/keys/ACCC4CF8.asc
+          . /etc/os-release
+          echo "deb [signed-by=/usr/share/postgresql-common/pgdg/apt.postgresql.org.asc] https://apt.postgresql.org/pub/repos/apt ${VERSION_CODENAME}-pgdg main" \
+            | sudo tee /etc/apt/sources.list.d/pgdg.list
+          sudo apt-get update -y
+          sudo apt-get install -y "postgresql-client-${{ matrix.pg_major }}"
+          echo "/usr/lib/postgresql/${{ matrix.pg_major }}/bin" >> "$GITHUB_PATH"
+
       - name: Cache Docker images
         uses: actions/cache@668228422ae6a00e4ad889ee87cd7109ec5666a7 # v5.0.4
         id: docker-cache
@@ -75,22 +88,31 @@ jobs:
       - name: Start Postgres
         run: docker compose -f compose.dbs.yml up -d --wait
 
+      # Force re-generating fresh dump and catalogs
+      - name: Remove committed db dump
+        run: rm -f priv/repo/tenant_db_dump_${{ matrix.pg_major }}.sql
+
       - name: Set up realtime DB and migrate tenant DB
         run: mix ecto.setup
 
-      - name: Export catalog snapshot
+      - name: Export Tenant DB catalog
         run: mix realtime.export_tenant_db_catalog
 
-      - name: Upload catalog artifact
+      - name: Export Tenant DB dump
+        run: mix realtime.export_tenant_db_dump --pg-major ${{ matrix.pg_major }}
+
+      - name: Upload artifact
         uses: actions/upload-artifact@330a01c490aca151604b8cf639adc76d48f6c5d4 # v5.0.0
         with:
-          name: tenant-catalog-${{ matrix.pg_major }}
-          path: priv/repo/tenant_db_catalog_${{ matrix.pg_major }}.json
-          if-no-files-found: error
+          name: tenant-dump-${{ matrix.pg_major }}
+          path: |
+            priv/repo/tenant_db_dump_${{ matrix.pg_major }}.sql
+            priv/repo/tenant_db_catalog_${{ matrix.pg_major }}.json
+          if-no-files-found: ignore
 
   commit-changes:
-    name: Commit catalog snapshots
-    needs: update-tenant-db-catalog
+    name: Commit catalog and dump snapshots
+    needs: update-tenant-dump
     runs-on: blacksmith-2vcpu-ubuntu-2404
     if: github.event.pull_request.head.repo.full_name == github.repository
 
@@ -108,31 +130,32 @@ jobs:
           repository: ${{ github.event.pull_request.head.repo.full_name }}
           token: ${{ steps.app-token.outputs.token }}
 
-      - name: Download catalog artifacts
+      - name: Download artifacts
         uses: actions/download-artifact@634f93cb2916e3fdff6788551b99b062d0335ce0 # v5.0.0
         with:
-          pattern: tenant-catalog-*
+          pattern: tenant-dump-*
           path: artifacts
 
       - name: Move artifacts into place
-        run: find artifacts -type f -name '*.json' -exec cp {} priv/repo/ \;
+        run: find artifacts -type f \( -name '*.sql' -o -name '*.json' \) -exec cp {} priv/repo/ \;
 
-      - name: Check if catalog changed
+      - name: Check if catalog or dumps changed
         id: check-changes
         run: |
-          git add priv/repo/tenant_db_catalog_15.json priv/repo/tenant_db_catalog_17.json
+          git add priv/repo/tenant_db_catalog_15.json priv/repo/tenant_db_catalog_17.json \
+            priv/repo/tenant_db_dump_15.sql priv/repo/tenant_db_dump_17.sql
           if git diff --cached --quiet; then
             echo "changed=false" >> "$GITHUB_OUTPUT"
           else
             echo "changed=true" >> "$GITHUB_OUTPUT"
           fi
 
-      - name: Commit and push catalog
+      - name: Commit and push catalog and dumps
         if: steps.check-changes.outputs.changed == 'true'
         env:
           PR_HEAD_REF: ${{ github.event.pull_request.head.ref }}
         run: |
           git config user.name "github-actions[bot]"
           git config user.email "41898282+github-actions[bot]@users.noreply.github.com"
-          git commit -m "chore: update tenant db catalog snapshots"
+          git commit -m "chore: update tenant db catalog and dumps"
           git push origin "HEAD:$PR_HEAD_REF"

lib/mix/tasks/realtime.export_tenant_db_dump.ex

@@ -0,0 +1,109 @@
+defmodule Mix.Tasks.Realtime.ExportTenantDbDump do
+  @shortdoc "Regenerate priv/repo/tenant_db_dump_<pg_major>.sql"
+
+  @moduledoc """
+  Dumps the tenant database's `realtime` schema to `priv/repo/tenant_db_dump_<pg_major>.sql`
+  and the `realtime.schema_migrations` rows.
+
+  Usage:
+
+      mix realtime.export_tenant_db_dump --pg-major 17
+
+  The target tenant DB is expected to already have all tenant migrations applied,
+  so make sure it is in a good state before generating it:
+
+      mise task run db-rm
+      mise task run db-start
+      mix setup
+
+  The target DB is read from `DB_HOST` / `DB_PORT` / `DB_NAME` / `DB_USER` / `DB_PASSWORD` env vars.
+
+  Requires `pg_dump` matching the target's major version on `$PATH`.
+  """
+  use Mix.Task
+
+  @impl Mix.Task
+  def run(args) do
+    {:ok, _} = Application.ensure_all_started(:postgrex)
+
+    {opts, _, _} = OptionParser.parse(args, strict: [pg_major: :integer])
+    pg_major = opts[:pg_major] || Mix.raise("--pg-major is required, e.g. --pg-major 17")
+
+    host = System.get_env("DB_HOST", "127.0.0.1")
+    port = Realtime.Env.get_integer("DB_PORT", 5433)
+    database = System.get_env("DB_NAME", "postgres")
+    user = System.get_env("DB_USER", "supabase_admin")
+    password = System.get_env("DB_PASSWORD", "postgres")
+    path = dump_path(pg_major)
+
+    Mix.shell().info("[export_tenant_db_dump] target: #{host}:#{port}/#{database} (pg#{pg_major})")
+
+    pg_dump!(host, port, database, user, password, path)
+    append_schema_migrations!(host, port, database, user, password, path)
+    postprocess!(path)
+
+    Mix.shell().info("[export_tenant_db_dump] wrote #{path}")
+  end
+
+  defp dump_path(pg_major), do: Application.app_dir(:realtime, "priv/repo/tenant_db_dump_#{pg_major}.sql")
+
+  defp pg_dump!(host, port, database, user, password, path) do
+    pg_dump = System.find_executable("pg_dump") || Mix.raise("pg_dump not found on $PATH")
+
+    args = [
+      "--host",
+      host,
+      "--port",
+      to_string(port),
+      "--username",
+      user,
+      "--dbname",
+      database,
+      "--schema-only",
+      "--schema",
+      "realtime",
+      "--file",
+      path
+    ]
+
+    case System.cmd(pg_dump, args, env: [{"PGPASSWORD", password}], stderr_to_stdout: true) do
+      {_output, 0} -> :ok
+      {output, code} -> Mix.raise("pg_dump exited #{code}:\n#{output}")
+    end
+  end
+
+  defp append_schema_migrations!(host, port, database, user, password, path) do
+    {:ok, conn} =
+      Postgrex.start_link(hostname: host, port: port, database: database, username: user, password: password)
+
+    {:ok, %{rows: rows}} =
+      Postgrex.query(conn, ~s(SELECT version FROM realtime."schema_migrations" ORDER BY version), [])
+
+    GenServer.stop(conn)
+
+    inserts =
+      Enum.map_join(rows, fn [version] ->
+        "INSERT INTO realtime.\"schema_migrations\" (version) VALUES (#{version});\n"
+      end)
+
+    sql = "ALTER TABLE realtime.schema_migrations ALTER COLUMN inserted_at SET DEFAULT now();\n" <> inserts
+
+    File.write!(path, sql, [:append])
+  end
+
+  defp postprocess!(path) do
+    tmp_path = path <> ".tmp"
+
+    path
+    |> File.stream!()
+    |> Stream.reject(&String.starts_with?(&1, ["\\restrict ", "\\unrestrict "]))
+    |> Stream.map(fn
+      "CREATE SCHEMA realtime;\n" -> "CREATE SCHEMA IF NOT EXISTS realtime;\n"
+      line -> line
+    end)
+    |> Stream.into(File.stream!(tmp_path))
+    |> Stream.run()
+
+    File.rename!(tmp_path, path)
+  end
+end

lib/realtime/api.ex

@@ -258,7 +258,6 @@ defmodule Realtime.Api do
     end
   end
 
-  @spec preload_counters(nil | Realtime.Api.Tenant.t(), any()) :: nil | Realtime.Api.Tenant.t()
   @doc """
   Updates the migrations_ran field for a tenant.
   """
@@ -285,6 +284,7 @@ defmodule Realtime.Api do
     end
   end
 
+  @spec preload_counters(nil | Realtime.Api.Tenant.t(), any()) :: nil | Realtime.Api.Tenant.t()
   def preload_counters(nil), do: nil
 
   def preload_counters(%Tenant{} = tenant) do

lib/realtime/telemetry/logger.ex

@@ -42,23 +42,31 @@ defmodule Realtime.Telemetry.Logger do
   def handle_event([:realtime, :tenants, :migrations, :start], _measurements, metadata, _config) do
     Logger.info(
       "Applying migrations to #{metadata.hostname}",
+      external_id: metadata.external_id,
       project: metadata.external_id
     )
   end
 
   def handle_event([:realtime, :tenants, :migrations, :stop], measurements, metadata, _config) do
     duration_ms = System.convert_time_unit(measurements.duration, :native, :millisecond)
 
-    Logger.info(
-      "Finished applying #{metadata.migrations_executed} migrations for tenant #{metadata.external_id} in #{duration_ms}ms",
-      project: metadata.external_id
-    )
+    message =
+      case metadata.source do
+        :dump ->
+          "Finished loading the dump of #{metadata.migrations_executed} migrations for tenant #{metadata.external_id} in #{duration_ms}ms"
+
+        :migrator ->
+          "Finished applying #{metadata.migrations_executed} migrations for tenant #{metadata.external_id} in #{duration_ms}ms"
+      end
+
+    Logger.info(message, external_id: metadata.external_id, project: metadata.external_id)
   end
 
   def handle_event([:realtime, :tenants, :migrations, :exception], _measurements, metadata, _config) do
     log_error(
       "MigrationsFailedToRun",
       metadata.reason,
+      external_id: metadata.external_id,
       project: metadata.external_id,
       error_code: metadata.error_code
     )
@@ -68,6 +76,7 @@ defmodule Realtime.Telemetry.Logger do
     log_warning(
       "MigrationCountMismatch",
       "Reconciling migrations_ran for tenant #{metadata.external_id} cached=#{metadata.cached_migrations_ran} database=#{metadata.database_migrations_ran}",
+      external_id: metadata.external_id,
       project: metadata.external_id
     )
   end
@@ -76,6 +85,7 @@ defmodule Realtime.Telemetry.Logger do
     log_error(
       "MigrationCountMismatchReconcileFailed",
       metadata.reason,
+      external_id: metadata.external_id,
       project: metadata.external_id
     )
   end